기술 요약: 세브라스 웨이퍼 스케일 엔진에서의 스텐실 연산

문제 제기

스텐실 연산은 유체 역학, 열 확산, 파동 전파와 같은 물리적 현상 시뮬레이션에 필수적인 과학적 컴퓨팅의 핵심 커널입니다. 이러한 알고리즘은 전통적인 고성능 컴퓨팅 (HPC) 아키텍처, 특히 GPU 에서 "메모리 벽 (Memory Wall)"으로 인해 주로 메모리 병목 현상에 직면합니다. 이는 높은 연산 처리량과 제한된 칩 외부 메모리 대역폭 사이의 불일치 때문입니다. 반면, 세브라스 웨이퍼 스케일 엔진 (WSE) 과 같은 AI 지향 하드웨어는 대규모 코어 병렬성과 높은 대역폭의 온칩 메모리를 제공하지만, 역사적으로는 저정밀도 머신러닝 워크로드에 최적화되어 왔습니다. 이 연구가 다루는 핵심 과제는 기존 시스템에 내재된 메모리 병목 현상을 극복하기 위해, 안정성에 민감한 전통적인 과학 알고리즘 (최소 단일 정밀도 FP32 필요) 을 이러한 특수 AI 하드웨어에 매핑할 수 있는 타당성입니다.

방법론

저자들은 세브라스 웨이퍼 스케일 엔진 3(WSE-3) 에서 2 차원 스텐실 연산을 구현하기 위해 설계된 새로운 프레임워크인 CStencil을 소개합니다. 방법론은 WSE-3 의 고유한 데이터 흐름 아키텍처에 일반적인 반복형 스텐실 솔버인 야코비 (Jacobi) 방법을 엄격하게 매핑하는 과정을 포함합니다.

아키텍처 및 매핑

WSE-3 은 2 차원 메시로 상호 연결된 90 만 개의 처리 요소 (PE) 를 특징으로 하며, 각 PE 는 48KB 의 로컬 SRAM 을 포함합니다. CStencil 은 이 분산 메모리 모델을 활용하여 전체 계산 영역을 연산 로직과 물리적으로 인접하게 유지함으로써 반복 과정 중 칩 외부 메모리 전송을 제거합니다.

• 데이터 준비: 입력 행렬은 전역적으로 패딩되고 타일로 분할되며, 각 PE 는 경계 조건을 위한 로컬 헤일로 (halo) 와 함께 하나의 타일을 수신합니다.
• 통신 (헤일로 교환): 이 프레임워크는 데이터 구조 디스크립터 (DSD) 를 사용하여 비동기 데이터 교환을 구현합니다.
  • 별 (Star) 패턴의 경우, PE 들은 8 개의 서로 다른 통신 색상을 사용하여 주축 방향 이웃과 경계 데이터를 교환합니다.
  • 상자 (Box) 패턴의 경우, 대각선 이웃이 필요하므로 2 단계 "저장 및 전달 (store-and-forward)" 전략이 사용됩니다. 주축 방향 이웃이 먼저 모서리를 포함한 가장자리를 교환한 후, 특정 모서리 데이터를 대각선 이웃으로 전달합니다. 이는 복잡한 라우터 전달 로직을 피하면서도 기존 마이크로스레드를 효율적으로 활용합니다.
• 연산: CStencil 은 스칼라 루프 대신 DSD 를 사용하여 요소별 행렬 연산을 수행합니다. DSD 기본 포인터를 이동시킴으로써 프레임워크는 단일 명령어로 이웃 하위 행렬에 접근하여 스텐실 업데이트를 벡터화된 곱셈 - 누적 연산의 시퀀스로 축소합니다.

베이스라인 및 평가

공정하고 엄격한 비교를 보장하기 위해 저자들은 이중 정밀도 (FP64) 텐서 코어를 위해 원래 설계된 최첨단 GPU 스텐실 솔버인 ConvStencil을 적용했습니다.

• 이식: ConvStencil 은 WSE-3 의 정밀도 기능과 일치하도록 단일 정밀도 (FP32/TF32) 로 이식되어 NVIDIA A100 GPU 에서 실행되도록 했습니다.
• 최적화: 이 이식 작업은 TF32 텐서 코어에 최적화된 "stencil2row" 변환 및 듀얼 테셀레이션 (Dual Tessellation) 알고리즘을 활용했습니다.
• 시뮬레이션 대 하드웨어: WSE-3 의 성능은 주로 사이클 정확도 시뮬레이터를 사용하여 평가되었으며, Star2d-1r 커널의 실제 하드웨어 실행과 비교하여 검증되었습니다. 그 결과 편차는 5.02% 미만으로 나타났습니다.

주요 기여

1. CStencil 프레임워크: WSE-3 을 위해 특별히 설계된 최초의 2 차원 스텐실 솔버 구현으로, 네이티브 2 차원 워크로드를 위해 장치의 연산 밀도와 온칩 메모리 용량을 완전히 활용하는 최적화된 매핑 전략을 특징으로 합니다.
2. 공정한 베이스라인 적용: ConvStencil 을 단일 정밀도로 포괄적으로 이식하여 WSE-3 과 텐서 코어에 의존하는 현재 최첨단 GPU 접근 방식 간의 직접적인 비교를 가능하게 했습니다.
3. 아키텍처 분석: WSE-3 의 분산 SRAM 과 메시 인터커넥트가 스텐실 연산에서 GPU 성능을 제약하는 메모리 병목 현상을 어떻게 제거하는지에 대한 상세한 조사입니다.

결과

실험 결과는 NVIDIA A100 에서 적용된 ConvStencil 대비 WSE-3 이 상당한 성능 우위를 보임을 입증합니다.

• 속도 향상: CStencil 은 최대 규모 (44,800 × 44,800 그리드의 Star2d-1r 패턴) 에서 적용된 ConvStencil 버전 대비 최대 342 배의 속도 향상을 달성했습니다.
• 확장성: WSE-3 은 다양한 문제 크기와 스텐실 복잡도 (별 및 상자 패턴) 에서 거의 완벽한 약한 확장성 (weak scaling) 을 보입니다. 이러한 선형적 성장은 문제 크기가 증가함에 따라 반복당 실행 시간을 일정하게 유지하는 대규모 온칩 메모리 대역폭과 효율적인 PE 간 통신에 기인합니다.
• 루프라인 분석: 분석 결과 GPU 구현은 HBM 대역폭에 의해 제한되는 엄격한 메모리 병목 상태임을 확인한 반면, WSE-3 의 CStencil 은 이론적 연산 피크에 훨씬 더 가깝게 작동함이 확인되었습니다. WSE-3 은 사용 가능한 연산 및 메모리 자원을 효과적으로 포화시켜 전통적인 메모리 벽을 극복합니다.
• 정밀도 트레이드오프: ConvStencil 의 단일 정밀도 이식 버전은 이중 정밀도 원본과 거의 동일한 성능을 보여주었으며, 이는 GPU 구현이 더 빠른 TF32 유닛을 활용하더라도 연산 처리량이 아닌 메모리 공급에 의해 제한됨을 나타냅니다.

중요성 및 주장

이 논문은 WSE-3 의 데이터 흐름 아키텍처가 전통적인 과학 알고리즘으로 성공적으로 재사용될 수 있으며, 현재 HPC 아키텍처의 메모리 제한을 극복할 수 있는 유망한 경로를 제시한다고 주장합니다.

• 재수렴: 이 연구는 특수 AI 하드웨어와 전통적인 HPC 워크로드 간의 재수렴 가능성을 강조하며, AI 전용 가속기가 메모리 병목 과학 애플리케이션에서 기존 시스템으로는 달성할 수 없는 하드웨어 활용 수준을 제공할 수 있음을 보여줍니다.
• 확장성: 이 연구는 웨이퍼 스케일 통합이 다중 GPU 구성에서 PCIe 또는 NVLink 병목 현상으로 인해 제한될 수 있는 유체 역학 및 기후 모델링과 같은 극대규모 시뮬레이션에 특히 효과적임을 시사합니다.
• 주의사항: 저자들은 이러한 성능 향상이 증가된 프로그래밍 복잡성을 수반함을 인정합니다. WSE-3 을 위한 개발은 CUDA 와 같은 확립된 생태계의 추상화 및 도구 성숙도가 부족한 Cerebras 소프트웨어 언어 (CSL) 를 통한 저수준 수동 데이터 라우팅 및 명시적 동기화가 필요합니다. 논문은 WSE-3 이 범용 작업을 위해 GPU 를 대체하지는 않지만, 대규모 메모리 병목 과학 시뮬레이션을 위한 구조적으로 건전한 접근 방식을 제공한다고 결론지었습니다.