Domain Specific Languages for GPU Kernels

2024-2025 의 GPU DSL 풍경은 이미 복잡하다 — Triton, Mosaic, ThunderKittens, CuTe DSL, Pallas, Helion, JAX Pallas. 각자가 자기 자리에서 최선이라고 주장하지만, 사용자 입장에서는 "내 문제에 어떤 걸 골라야 하는가" 의 답이 명확하지 않다. Tri Dao 가 FlashAttention 저자의 입장에서 그 지도를 직접 그린다.

강의의 frame 이 명확하다.

1. 같은 산 (efficient GPU kernel) 을 오르는 다른 길들 — 모두 비슷한 산봉우리를 노리지만 출발점과 도구가 다르다.
2. 두 축으로 정렬한다 — 추상 레벨 (얼마나 적게 짜도 되는가) 와 통제 가능 자유도 (필요할 때 얼마나 깊이 내려갈 수 있는가). 두 축이 통상 trade-off.
3. "어느 게 best 인가" 가 아니라 "어느 자리에 어느 것이 자연스러운가" — 답은 작업의 hot path 깊이 + 팀의 학습 곡선.

Tri 가 슬라이드 한 장으로 정리한 정렬 — 추상 레벨이 높으면 코드가 짧다, 통제 가능 자유도가 높으면 hardware 의 모든 칸을 직접 짤 수 있다. 둘은 통상 trade-off.

강의에서 가장 진하게 다뤄진 비교. Tri Dao 가 같은 GEMM 작업을 5 개 DSL 로 어떻게 표현하는지 줄별로 짚는다.

같은 작업의 코드 길이도 큰 차이가 있다 (단순 GEMM, BF16, 4096³, 개념적).

• Helion — ~30 줄. PyTorch 호출처럼 보인다.
• Triton — ~80 줄. autotuner config 까지 포함.
• Mosaic-GPU — ~120 줄. JAX 의 Pallas 위에 spec.
• ThunderKittens — ~100 줄. WGMMA 명시.
• CuTe DSL — ~250+ 줄. Layout 정의 + Atom 조합.
• raw CUDA — ~500 줄. 모든 boilerplate 포함.

강의의 핵심 한 장 — 모든 DSL 이 추상 레벨 × 통제 자유도의 trade-off 위에 있고, 결정 변수는 "필요할 때 어디까지 내려갈 수 있는가".

강의에서 Tri 가 짚은 핵심 — FA-2 → FA-3 의 전환은 escape hatch 의 깊이 부족 때문이었다. Triton 으로는 FA-3 의 ping-pong, async warpgroup specialization 을 표현할 도구가 부족했다. 그래서 ThunderKittens / 직접 CUDA 로 내려갔다.

DSL 별 lowering pipeline 의 갈래. 어디서 어떤 결정이 내려지는지가 성능 의 결정 변수.

이 갈래를 보면 "어디서 결정이 내려지는가" 가 명확해진다.

• Triton — 대부분의 결정 (layout, register allocation, vectorization, async pipeline) 이 MLIR pass 안에서. 사용자는 BLOCK_SIZE, num_warps, num_stages 정도만 통제.
• CuTe — Layout 과 Atom 의 조합으로 사용자가 결정. compiler 는 그것을 기계어로 옮길 뿐.
• ThunderKittens — 중간. tile 형태는 사용자가, swizzling 의 hex detail 은 라이브러리가.

DSL 의 가치는 행복한 path 의 코드 길이가 아니라 버그 났을 때 빠져나오는 시간. Tri 가 강조한 ergonomics 의 두 면.

H100 / B200 같은 새 칩이 나오면 새 명령 (TMA, WGMMA, tcgen05.mma) 을 DSL 이 흡수해야 한다. 흡수 속도가 DSL 의 살아남는 변수.

큰 패턴 — 새 명령은 항상 raw CUDA / CUTLASS 가 먼저 흡수, ThunderKittens 가 빠른 두 번째. Triton 은 추상 레벨이 높은 만큼 새 명령의 흡수가 느리다 (spec 정의 → IR 추가 → MLIR pass → backend lowering 전체가 가야 함).

DSL 의 운명을 결정하는 또 하나 — 인프라가 누구 위에 깔리는가. PyTorch, JAX, vLLM, SGLang 같은 layer 가 어떤 DSL 을 default 로 받는지가 산업적 자리.

• Triton — torch.compile 의 Inductor 가 Triton 으로 lowering. 사실상 PyTorch 의 default backend. 가장 넓은 채택.
• CuTe / CUTLASS — cuBLAS, cuDNN, FlashAttention 의 reference 인프라. NVIDIA 가 직접 maintain. 가장 깊은 채택.
• Mosaic-GPU — Google 의 JAX/Pallas 인프라. TPU porting 의 base. JAX 생태계 의존적.
• ThunderKittens — HazyResearch / Together / Stanford. attention 의 reference 구현. research-leading 채택.
• Helion — 강의 시점에 신생. 채택 곡선 미상 (확인 필요).

이 채택 분포가 결정하는 사실 — 새 op 가 등장하면 Triton / CUTLASS 위에 먼저 구현된다. 그 이후 다른 DSL 로 흘러간다. 그래서 "어떤 DSL 이 가장 빨리 새 모델을 따라가는가" 의 답은 인프라 곡선과 1:1.

• Triton — 새 칩 (B200, MI300) 의 명령을 더 빠르게 흡수, autotune 의 cost 감소, inline_asm 의 escape hatch 강화.
• Mosaic-GPU — TPU 에서 GPU 로의 단일 코드 경로 강화, JAX/Pallas 통합의 깊이.
• CuTe DSL — Blackwell 명령의 정식 abstraction, FP4 / sparsity 의 first-class.
• ThunderKittens — HipKittens (AMD), ThunderMittens (Apple) 의 cross-arch 확장. § L075 참조.
• Helion — Triton 의 frontend 추상의 더 위. PyTorch 의 native fusion 표현으로의 후보.
• 새 후보들 — Pallas (JAX), MLX (Apple), 그리고 LLM-as-DSL 같은 generative 접근 (BackendBench 의 후속).

손에 새기기 — 실습 시퀀스

1. 같은 GEMM 5 가지로 — Triton, ThunderKittens, CuTe, Helion, raw CUDA. 코드 길이 + peak% 직접 측정. 4096³ BF16.
2. 새 op 추가의 비용 — softmax + dropout fused 를 각 DSL 에서 새로 짠다. 짜는 시간 + 디버깅 시간 직접 측정.
3. Triton 의 escape hatch — tl.inline_asm_elementwise 또는 PTX 직접 삽입. 어디까지 내려갈 수 있는지 확인.
4. 새 칩 명령 흡수 — H100 의 TMA / WGMMA 를 Triton vs ThunderKittens 로 호출. 코드 양 비교.
5. fault injection — 의도적 wrong indexing 을 각 DSL 에서. 에러 메시지 / 디버깅 워크플로 비교.
6. autotune cost — 같은 GEMM 의 Triton autotune sweep 시간 vs ThunderKittens 의 template 수동 sweep 시간.

• Helion 의 정확한 위치 — Triton 의 frontend 인지 별도 IR 인지. repo 와 docs 직접 확인.
• Mosaic-GPU 의 GPU peak% — TPU 외 GPU 위 성능. NVIDIA 위 성숙도.
• Triton 의 inline_asm 실제 사용 — production 사례에서 PTX inline 이 어떻게 쓰이는지.
• CuTe DSL 의 학습 자료 — CUTLASS docs 외에 좋은 tutorial. (확인 필요)
• LLM 으로 DSL 코드 짜기의 분포 — § L076 의 BackendBench 결과가 DSL 별로 어떻게 다른지.
• 새 후보 — Pallas / MLX / etc — 강의 시점 이후 등장한 DSL 의 위치.
• multi-arch portability 의 진짜 비용 — Triton 이 NVIDIA / AMD / Apple 모두 지원한다고 하지만 실제 peak% 격차.