[triton] Consumer Blackwell(sm_120)에서 PTX Codegen Segfault 수정

PR 링크: triton-lang/triton#9734 상태: Merged | 변경: +63 / -2

들어가며

NVIDIA의 Blackwell 아키텍처는 데이터센터용(sm_100)과 consumer용(sm_120, RTX 5070 Ti/5080/5090 등)으로 나뉩니다. consumer Blackwell은 tensor memory와 같은 데이터센터 전용 기능을 지원하지 않습니다. 기존 Triton 코드는 capability >= 90이면 일괄적으로 "a" suffix를 붙여 arch-accelerated 모드로 컴파일했는데, 이는 consumer Blackwell에서 존재하지 않는 명령어를 생성하여 런타임 segfault를 유발했습니다.

핵심 코드 분석

Before:

def sm_arch_from_capability(capability: int):
    # TODO: Handle non-"a" sms
    suffix = "a" if capability >= 90 else ""
    return f"sm_{capability}{suffix}"

After:

def sm_arch_from_capability(capability: int):
    # sm_90a  - Hopper datacenter (H100, H200)
    # sm_100a - Blackwell datacenter (B100, B200)
    # Consumer Blackwell (sm_120) does NOT have an "a" variant
    suffix = "a" if capability >= 90 and capability != 120 else ""
    return f"sm_{capability}{suffix}"

테스트에서 핵심 검증:

def test_compile_only_sm120():
    k = triton.compile(..., target=GPUTarget("cuda", 120, 32))
    ptx = k.asm["ptx"]
    assert ".target sm_120" in ptx
    assert ".target sm_120a" not in ptx
    # tensor memory 명령어가 생성되지 않아야 함
    assert "tcgen05" not in ptx

왜 이게 좋은가

이 수정은 consumer Blackwell GPU 사용자들이 Triton을 사용할 때 겪는 크래시 수준의 버그를 해결합니다. sm_120a로 잘못 컴파일하면 LLVM/ptxas가 tensor memory(tcgen05) 명령어를 생성하는데, 이는 consumer Blackwell에 존재하지 않아 런타임에 segfault가 발생합니다. 테스트에서 tl.dot 커널을 사용하여 matmul 파이프라인까지 검증하는 것이 인상적이며, PTX와 TTGIR 모두에서 금지된 명령어가 없음을 확인합니다.

정리

sm_arch_from_capability 함수에서 sm_120(consumer Blackwell)에 "a" suffix를 붙이지 않도록 수정하고, sm_120 타겟의 컴파일 정합성을 검증하는 테스트를 추가했습니다.

참고 자료

• triton-lang/triton#9734

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