[cutlass] NVIDIA CUTLASS CuTeDSL: SM103 Grouped Block-Scaled GEMM 최적화 분석

PR 링크: NVIDIA/cutlass#3124 상태: Merged | 변경: +3600 / -0

들어가며

NVIDIA의 CUTLASS 라이브러리는 고성능 텐서 연산을 위한 핵심 도구입니다. 최근 cutlass 레포지토리에 추가된 Sm103GroupedBlockScaledGemmKernel은 Blackwell(SM103) 아키텍처의 강력한 연산 성능과 SM100 스타일의 Grouped Scheduling을 결합한 결과물입니다. 이 커널은 서로 다른 M/N/K 형태를 가진 여러 GEMM 연산을 단일 커널 실행으로 처리할 수 있게 하여, 특히 MoE(Mixture of Experts) 모델의 추론 효율을 극대화합니다. 본 글에서는 이 커널의 설계 핵심과, 대규모 연산 시 발생했던 동기화 버그를 어떻게 해결했는지 분석합니다.

코드 분석

1. SM103 아키텍처 최적화

SM103 커널은 7개의 워프(Warp)를 사용하여 연산을 분담합니다. 특히 주목할 점은 스케일 팩터(Scale-factor)를 처리하는 전용 워프(Warp 6)를 A/B 행렬 처리 워프(Warp 5)와 분리했다는 점입니다. 이를 통해 메모리 대역폭 활용을 최적화했습니다.

# SM103 전용 워프 할당
self.mma_warp_id = 4
self.tma_ab_warp_id = 5
self.tma_sf_warp_id = 6
self.threads_per_cta = 224  # 7 warps * 32

2. 파이프라인 동기화 버그 수정

이 PR의 가장 중요한 기술적 기여는 대규모 연산 시 발생하던 결과 오류를 수정한 것입니다. 기존에는 TMA(Tensor Memory Accelerator)의 생산자 토큰이 루프 시작 전에 단 한 번만 초기화되었습니다.

Before (오류 발생 코드):

# 루프 진입 전 단 한 번만 초기화되어 stale token이 루프 내부에 잔류함
ab_producer.try_acquire()
for tile in range(num_tiles):
    # ... MMA 연산 ...

After (수정된 코드):

# 각 타일 경계마다 생산자 토큰을 리셋하여 동기화 보장
for tile in range(num_tiles):
    ab_producer.reset()
    ab_producer.try_acquire()
    # ... MMA 연산 ...

각 타일의 마지막 단계에서 try_acquire()를 건너뛰는 로직 때문에 이전 타일의 토큰이 다음 타일로 넘어가는 문제가 있었습니다. 이를 각 타일 경계에서 리셋하도록 수정함으로써, 대규모 데이터 처리 시 발생하는 메모리 덮어쓰기 문제를 해결했습니다.

왜 이게 좋은가

이번 최적화는 단순히 기능을 추가하는 것을 넘어, '지속적 커널(Persistent Kernel)에서의 상태 관리'에 대한 중요한 교훈을 줍니다.

1. 성능: GB300 환경에서 8개 전문가(Experts) 모델 기준 약 4,000~6,000 TFLOPS의 높은 처리량을 달성했습니다.
2. 안정성: compute-sanitizer를 통해 8K×8K×8K 대규모 연산에서도 메모리 무결성을 검증했습니다.
3. 교훈: 파이프라인 병렬 처리를 구현할 때, 공유 자원(TMA 토큰 등)의 생명주기를 루프 전체가 아닌 '작업 단위(Work-tile)'로 관리해야 함을 보여줍니다.

결론

이번 Sm103GroupedBlockScaledGemmKernel 추가는 Blackwell 아키텍처의 하드웨어 가속 기능을 소프트웨어 수준에서 완벽하게 활용하는 사례입니다. 특히 동기화 로직의 정교한 수정은 고성능 커널 개발에서 '디테일'이 성능과 정확도에 얼마나 큰 영향을 미치는지 잘 보여줍니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#tma-tensor-memory-accelerator

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.