[sglang] AMD MI355X 환경에서 Triton 3.7 레지스터 스필링 최적화

PR 링크: sgl-project/sglang#34741 상태: Merged | 변경: +80 / -5

들어가며

최근 SGLang 프로젝트에서 AMD Instinct MI355X(gfx950) 환경에 Triton 3.7을 도입한 후, 특정 어텐션 커널(extend_attention)에서 심각한 성능 저하가 발견되었습니다. 분석 결과, Triton 3.7의 코드 생성 방식이 레지스터 압박을 가중시켜 과도한 레지스터 스필링(Register Spilling)을 유발하고 있었으며, 이로 인해 프리필(prefill) 성능이 크게 하락했습니다. 본 글에서는 이 문제를 해결하기 위해 적용된 타일 사이즈 최적화 전략을 살펴봅니다.

코드 분석

1. Triton 버전 감지 및 조건부 로직 추가 (python/sglang/kernels/ops/attention/extend_attention.py)

가장 먼저 Triton 버전을 동적으로 감지하여, 문제가 발생하는 환경에서만 최적화된 설정을 적용하도록 수정했습니다.

# Before
# (버전 감지 로직 없음)

# After
try:
    _triton_version_parts = tuple(
        int(part) for part in triton.__version__.split(".")[:2]
    )
except (AttributeError, ValueError):
    _triton_version_parts = (0, 0)
_is_triton_ge_37 = _triton_version_parts >= (3, 7)

2. BLOCK_N 타일 사이즈 최적화

기존 (64, 64) 타일 구성이 gfx950 + Triton 3.7 조합에서 레지스터 압박을 유발하므로, BLOCK_N을 32로 줄여 레지스터 사용량을 낮췄습니다.

# After
if _is_gfx95 and _is_triton_ge_37 and Lq == 576 and Lv == 512:
    BLOCK_M, BLOCK_N = (64, 32)
    num_warps = 4
elif _is_gfx95 and 128 < Lq <= 256:
    # ... 기존 로직

왜 이게 좋은가

이번 최적화의 핵심은 레지스터 압박(Register Pressure) 완화입니다.

• 성능 수치: N64 설정 대비 N32 설정으로 변경 시, VGPR 사용량이 512에서 392~433으로 감소했습니다. 그 결과, 스크래치 메모리 사용이 0으로 제거되었고, isolated p50 레이턴시가 12.57ms에서 5.24ms로 약 58.3% 개선되었습니다.
• 교훈: 컴파일러 버전이 업데이트될 때 생성되는 커널의 레지스터 할당 전략이 변경될 수 있습니다. 특히 고성능 연산 커널에서는 타일 사이즈(BLOCK_M, BLOCK_N)를 하드웨어 아키텍처와 컴파일러 버전에 맞춰 세밀하게 튜닝하는 것이 성능 유지의 핵심입니다.

또한, test_extend_attention_triton37_lq576_n32 테스트를 통해 N64와 N32 결과값의 오차 범위가 허용치 내에 있음을 검증하여 정확성을 보장했습니다.

참고 자료

• https://triton-lang.org/main/index.html
• https://docs.sglang.ai/developer_guide/contribution_guide.html

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