TLX — Triton 의 hardware-near 확장

강의의 출발점. Triton 의 lowering 은 “compiler heuristic 이 충분히 똑똑하다” 는 가정 위에 산다. 그러나 H100 의 진짜 성능은 — warp specialization · async tensor core · TMA · multi-stage pipelining 같은 새 기능을 직접 조합해야 나온다. 이게 compiler 자동화의 한계.

강의에서 자주 인용되는 예 — flash attention 3. CUTLASS 로 짠 hand-tuned 버전이 H100 에서 700+ TFLOP/s. Triton 으로 같은 로직을 짜면 60–70% 수준. 그 차이가 warp specialization · async tensor core 의 직접 조작에서 온다. TLX 가 그걸 Triton 안에서 가능하게 한다.

TLX repo 의 README 가 정리하는 4영역. Triton 의 high-level 추상 위에 추가되는 hardware-near 도구들.

이 4영역의 공통점 — Triton 의 자동화를 부분적으로 끄고, 사용자가 명시적으로 통제. 그래서 TLX 코드는 vanilla Triton 보다 길고 복잡하지만, compiler 가 못 잡는 자리를 잡는다.

TLX 가 가능하게 하는 핵심 — H100 의 hardware feature 를 Triton 안에서 직접 호출. CUTLASS 의 정신을 Python 으로.

• TMA (Tensor Memory Accelerator) — Hopper 의 새 DMA 엔진. 큰 tile 을 한 명령으로 global → shared 로. CPU 가 아닌 hardware descriptor 가 multi-D 의 indexing 자동 풀이.
• WGMMA (Warp Group MMA) — H100 의 async tensor core. async 로 시작해서 token 으로 완료 대기. CUDA 의 일반적 mma 보다 훨씬 큰 tile 한 번에.
• warp specialization — block 의 warp 를 그룹으로 나눠 다른 일. 한 그룹은 load (memory unit), 다른 그룹은 mma (tensor core), 또 다른 그룹은 epilogue. memory 와 compute 의 overlap 을 hardware 레벨로.
• multi-stage pipeline — 같은 buffer 를 N 개로 두고, 각 stage 가 다른 위치를 처리. async transfer 의 latency 가 다음 stage 의 compute 로 hide.

TLX repo 의 첫 example — pipelined GEMM. vanilla Triton 으로 짜면 compiler 가 일부 pipeline 을 만들지만, TLX 로 직접 짜면 N-stage 의 timing 을 사용자가 통제.

# pseudo-code — TLX pipelined GEMM 의 골격
@triton.jit
def gemm_kernel(...):
    # 3-stage shared memory buffer
    a_buf = tlx.local_alloc(
        shape=(3, BM, BK),
        dtype=tl.float16,
    )
    b_buf = tlx.local_alloc(
        shape=(3, BK, BN),
        dtype=tl.float16,
    )

    # warmup — 첫 2 tile 미리 load
    tok0 = tlx.async_copy(a_ptr, a_buf[0])
    tok1 = tlx.async_copy(a_ptr+BK, a_buf[1])

    acc = tl.zeros(...)

    for k in range(K // BK):
        slot = k % 3
        next_slot = (k + 2) % 3

        # 다음 tile async load 시작
        tok = tlx.async_copy(
            a_ptr + (k+2)*BK, a_buf[next_slot]
        )

        # 현재 tile 의 load 완료 대기
        tlx.wait(toks[slot])

        # mma — async tensor core
        acc = tlx.async_mma(
            a_buf[slot], b_buf[slot], acc
        )

    tlx.wait_all()
    tl.store(c_ptr, acc)

이 패턴이 가르치는 것 — async + buffer rotation이 multi-stage pipelining 의 본질. compiler 가 못 보던 timing 을 사용자가 직접 표현. 실제 코드는 더 길지만 (boundary 처리, prologue/epilogue) 핵심은 위 골격.

TLX 의 가장 중요한 example. flash attention 의 새 변형이 — TLX 의 async_tasks 로 자연스럽게 표현. 한 block 의 warp 들을 두 그룹으로 분리.

• load warpgroup — Q, K, V tile 을 global → shared 로 async 가져옴. TMA 사용. tensor core 는 안 만짐.
• compute warpgroup — load warpgroup 이 가져온 tile 위에서 QK^T → softmax → AV 의 attention. WGMMA 호출. load 는 안 만짐.
• epilogue warpgroup (옵션) — compute 결과를 global 로 store. compute 와 동시에 도는 게 가능.

이 분리가 만들어내는 효과 — memory 와 compute 가 다른 hardware unit 에서 동시에. CUTLASS 의 producer-consumer pattern 을 Triton 안에서.

# TLX warp specialization 의 골격
@triton.jit
def attention(...):
    # shared memory + barrier 셋업
    q_buf = tlx.local_alloc(...)
    k_buf = tlx.local_alloc(...)
    v_buf = tlx.local_alloc(...)
    barrier = tlx.named_barrier()

    with tlx.async_tasks() as tasks:
        @tasks.task(num_warps=4)
        def load_loop():
            for i in range(N_BLOCKS):
                tlx.tma_load(k_ptr, k_buf[i])
                tlx.tma_load(v_ptr, v_buf[i])
                barrier.arrive(i)

        @tasks.task(num_warps=4)
        def compute_loop():
            for i in range(N_BLOCKS):
                barrier.wait(i)
                qk = tlx.async_mma(q_buf, k_buf[i])
                p  = softmax(qk)
                acc = tlx.async_mma(p, v_buf[i], acc)

    tl.store(out_ptr, acc)

TLX 가 facebookexperimental 이라는 이름에서 드러나듯 아직 실험. 그러나 frontier kernel 영역에서는 흥미로운 자리.

• flash attention 변형 — TLX 로 짠 attention 이 vanilla Triton 대비 큰 성능 향상. CUTLASS 의 hand-tuned 와 거의 동등.
• pipelined GEMM — kernel 라이브러리의 핵심 op. cuBLAS 와 경쟁 가능한 자리.
• MoE GEMM — sparse routing 이 결합된 MoE 의 GEMM. expert 별 batch 가 다양한 자리에서 TLX 의 명시적 통제가 도움.
• 학술 연구 — Triton + TLX 의 lowering 자체에 대한 논문이 등장 시작.

• arch-specific — Hopper/Blackwell 의 새 feature 에 의존. A100 위에서는 일부 기능이 emulation 또는 미동작. AMD 위에서는 거의 다른 코드.
• 학습 곡선 — async / sync / barrier / warp group 의 mental model 이 vanilla Triton 보다 한 단계 더. CUTLASS 사용 경험이 있으면 친숙.
• deadlock 가능성 — barrier 와 wait 가 순환 의존을 만들 수 있음. 디버깅이 어려움 — interpret 모드도 한계.
• portability 의 손실 — TLX 코드가 다른 GPU 에서 안 돌면 fallback 코드를 별도로 유지해야. “하나의 코드가 모든 GPU 에서” 의 Triton 약속이 깨짐.
• upstream 미통합 — Triton 본가에 들어가지 않은 실험. Triton 본가의 다른 변경과 충돌 가능.
• 학습 자료 부족 — vanilla Triton 의 풍부한 예제와 비교해 example/tutorial 이 적다.

현재 TLX 는 facebookexperimental fork. 그러나 Triton 본가가 Hopper/Blackwell 지원을 깊게 가져가면 — TLX 의 일부 추상 (특히 async tensor core, warp specialization) 이 흡수될 가능성.

• warp specialization 의 표준화 — Triton 본가가 tl.async_tasks 같은 형태로 native 지원할 가능성.
• TMA / WGMMA — 이미 Triton 3.x 에서 일부 들어옴. TLX 의 token-based async 가 표준에 합류할 가능성.
• compiler heuristic 의 학습 — TLX 코드의 패턴을 compiler 가 자동 인식해 vanilla Triton 코드에서도 같은 lowering 을 만들 가능성.
• CUTLASS 와의 수렴 — Triton + TLX 와 CUTLASS Python 이 같은 영역으로 수렴. 둘이 서로 영향.

실습 시퀀스

1. repo clone + build — facebookexperimental/triton 의 tlx branch. cmake build, H100 에서 hello world example 실행.
2. vanilla GEMM vs TLX pipelined GEMM — 같은 size 의 GEMM 을 두 버전으로. NCU 로 SM occupancy / tensor core utilization 비교.
3. flash attention 비교 — TLX 의 attention example 과 vanilla Triton 의 attention 의 throughput. CUTLASS hand-tuned 와도 비교.
4. warp specialization 직접 짜기 — 단순 vector add 를 load/compute warpgroup 으로 분리. 작은 예제로 mental model 정착.
5. NCU 로 hardware utilization — TLX 코드와 vanilla 의 NCU dump 비교. tensor core idle time, memory unit busy time 의 차이.

• 스피커 명단 확인 필요 — notes 에 “speaker missing”. 영상에서 발표자 확인.
• TLX 의 정확한 풀네임 — “Triton Low-level eXtensions” vs “Tile-Level eXtensions”. 강의에서 어떤 표현 썼는지.
• TLX 코드의 정확한 API — 본 노트의 예시는 README + 일반 패턴의 추정. 실제 함수 시그니처는 repo example 에서.
• Triton 본가와의 관계 — 강의 시점에 TLX 의 어떤 부분이 본가에 들어갈 계획이었는지.
• vanilla Triton 대비 정확한 throughput 비율 — flash attention 의 60–70% 같은 비율은 대략적 추정. 강의에서 정확한 숫자가 있었을 가능성.
• Blackwell 지원 — 강의 시점에 Blackwell (B100/B200) 의 새 기능 — fp8 tensor core, 더 큰 register file — 이 어떻게 다뤄졌는지.
• 다른 lower-level Triton 변형 — Triton 의 다른 fork (예: AITemplate 등) 과의 비교가 있었는지.