[triton] Triton 커널 최적화: High Occupancy Persistent Matmul 구현을 통한 성능 향상

PR 링크: triton-lang/triton#9248 상태: Merged | 변경: +41 / -20

들어가며

최신 GPU 아키텍처에서 고성능 행렬 곱셈(Matmul)을 구현할 때, 단순히 연산 효율을 높이는 것뿐만 아니라 GPU의 SM(Streaming Multiprocessor) 활용도를 극대화하는 것이 중요합니다. 이번 Triton PR은 Persistent Matmul 커널 내에서 독립적인 작업을 병렬로 배치하여 워프 스케줄러가 파이프라인의 유휴 시간(bubble)을 숨길 수 있도록 개선했습니다. 결과적으로 H200 환경에서 bf16 x mxfp4 MoE 연산 기준 2290 GBps에서 2640 GBps로 약 15%의 성능 향상을 기록했습니다.

코드 분석

1. matmul.py: 그리드 계산 로직 최적화

기존에는 단순히 max_grid를 기준으로 커널을 실행했으나, 이제 occupancy_target을 도입하여 SM 점유율을 제어합니다.

# Before
grid = min(target_info.num_sms() - opt_flags.idle_sms, max_grid) if opt_flags.is_persistent else max_grid

# After
grid = batch_size * grid_m * grid_n * opt_flags.split_k
if opt_flags.is_persistent:
    available_sms = target_info.num_sms() - opt_flags.idle_sms
    grid = min(opt_flags.occupancy_target * available_sms, grid)

이 변경을 통해 커널이 단순히 전체 SM을 점유하는 것이 아니라, 타겟 점유율에 맞춰 스케줄링되도록 유연성을 확보했습니다.

2. opt_flags.py: 레지스터 압력 제어 (maxnreg)

레지스터 압력으로 인한 점유율 저하를 방지하기 위해 maxnreg를 명시적으로 계산하여 커널 인자로 전달합니다.

# After (추가된 로직)
if is_persistent and not is_blackwell_or_newer:
    maxnreg = reg_per_sm // (num_warps * threads_per_warp * occupancy_target)
    maxnreg = min(max_reg_per_thread, maxnreg)
else:
    maxnreg = None

리뷰어 peterbell10은 이 방식이 레지스터 스필링(spilling) 없이 점유율 제한 문제를 해결하는 효과적인 방법임을 확인했습니다.

3. opt_flags_nvidia.py: 공유 메모리(Smem) 용량 조정

occupancy_target에 따라 공유 메모리 할당량을 동적으로 조정하여 더 많은 워프가 동시에 실행될 수 있는 환경을 조성했습니다.

# After
smem_capacity = device_props.shared_memory_per_block_optin
smem_capacity //= occupancy_target

왜 이게 좋은가

이번 최적화의 핵심은 **'파이프라인 버블 최소화'**와 **'리소스 기반 점유율 제어'**입니다.

1. 워프 스케줄링 효율화: 단일 SM 내에서 독립적인 작업을 수행하게 함으로써, 메모리 로드나 연산 대기 시간 동안 다른 워프가 작업을 수행할 수 있게 되어 하드웨어 유닛의 가동률이 높아집니다.
2. 레지스터 및 공유 메모리 관리: maxnreg와 occupancy_target을 통해 하드웨어 리소스 사용량을 정밀하게 제어함으로써, 레지스터 부족으로 인한 성능 저하를 방지하고 더 높은 동시 실행 워프 수를 확보했습니다.

일반적 교훈: GPU 커널 최적화 시 단순히 연산 루프를 최적화하는 것을 넘어, 하드웨어의 리소스 제한(레지스터, Smem)과 스케줄러의 동작 방식을 고려한 'Occupancy-aware' 설계가 대규모 행렬 연산에서 결정적인 성능 차이를 만든다는 점을 시사합니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://docs.nvidia.com/cuda/parallel-thread-execution/#packing-format-used-for-matrix-a-and-b-by-kind-mxf8f6f4-in-shared-memory

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