[Triton] Hopper에서 소규모 배치 크기 벤치마크 수정

PR 링크: triton-lang/triton#8877 상태: Merged | 변경: +14 / -4

들어가며

Triton의 MLP 벤치마크에서 Hopper GPU에 대한 num_warps 설정이 소규모 배치에서 잘못 적용되는 문제가 있었다. 이 PR은 Hopper에서 weight이 scaled인 경우에만 8 warps를 사용하도록 조건을 수정하고, small batch 테스트 케이스를 추가한다.

핵심 코드 분석

Before: batch 크기만으로 num_warps 결정

num_warps = 4 if batch <= 512 else 8

After: GPU 아키텍처도 고려

# on hopper we only use 8 warps when weight is scaled
num_warps = 4 if batch <= 512 and cuda_capability_geq(10, 0) else 8

테스트 케이스 추가

# Before: large batch만 테스트
[(1024, 1024, 1024, 1, 1, "bf16", "bf16", 1, 1), ...]

# After: small batch도 추가
[
    (128, 1024, 1024, 1, 1, "bf16", "bf16", 1, 1),
    (128, 1024, 1024, 1, 1, "fp8", "fp8", 1, 1),
    (1024, 1024, 1024, 1, 1, "bf16", "bf16", 1, 1),
    (1024, 1024, 1024, 1, 1, "fp8", "fp8", 1, 1),
]

왜 이게 좋은가

1. 아키텍처 인식: Hopper(sm_90+)와 이전 아키텍처에서 최적의 warp 수가 다른 점을 반영한다.
2. small batch 커버리지: batch=128인 테스트 케이스를 추가하여 edge case를 검증한다.
3. 간단하고 명확한 수정: 조건 한 줄과 테스트 추가로 문제를 해결한다.

정리

이 PR은 Hopper GPU에서 소규모 배치 MLP 벤치마크의 num_warps 설정을 수정하고, batch=128 테스트 케이스를 추가하여 소규모 배치에서도 올바른 성능 측정을 보장한다.

참고 자료

• triton-lang/triton#8877

---

이 글은 AI를 활용하여 PR의 핵심 변경사항을 분석하고 정리한 것입니다. 실제 코드의 맥락은 원본 PR을 참고해 주세요.