vLLM의 Helion: PyTorch 네이티브 DSL이 CUDA 커널을 대체하고 Qwen3 모델의 FP8 추론을 가속화

PyTorch 팀이 오늘 Red Hat과의 협업으로 Qwen3 모델의 FP8 추론(inferencing)을 위해 vLLM에 아홉 가지 Helion 커널을 통합한 내용을 담은 블로그 포스트를 발표했습니다. 저자는 Red Hat의 Sean Chen과 Meta Platforms(PyTorch 팀)의 Yanan Cao입니다.

Helion이란 무엇이며 개발자들이 주목해야 하는 이유는?

Helion은 「타일이 있는 PyTorch」로 설명되는 도메인 특화 언어(DSL, domain-specific language)입니다. 직접 CUDA 코드를 작성하지 않고 타일 프로그래밍(tile-programming) 개념으로 확장된 PyTorch 구문을 사용하여 고성능 GPU 커널을 작성할 수 있습니다. 목표는 하드웨어 이식성입니다. 동일한 커널이 수동 포팅 없이 NVIDIA H100, B200, 미래 아키텍처에서 작동해야 합니다.

Helion은 고급 사전 컴파일(ahead-of-time, AOT) 자동 튜닝 인프라를 제공합니다. 각 대상 GPU 아키텍처와 입력 차원에 대한 최적의 타일 크기와 메모리 레이아웃을 자동으로 탐색합니다. 커널 개발은 빠릅니다. 저자들은 아홉 커널 대부분이 하루 작업 이내에 구현 및 검증되었다고 밝힙니다.

vLLM에 통합된 커널들

아홉 커널은 FP8 추론의 포워드 패스(forward pass) 거의 전체를 커버합니다.

1. dynamic_per_token_scaled_fp8_quant
2. rms_norm_dynamic_per_token_quant
3. silu_and_mul_dynamic_per_token_quant — 새로운 퓨전 커널
4. fused_qk_norm_rope
5. per_token_group_fp8_quant
6. rms_norm_per_block_quant
7. silu_and_mul_per_block_quant
8. scaled_mm
9. scaled_mm_blockwise

통합으로 토큰별 양자화에서 포워드 패스의 커널 호출 수가 11개 이상에서 10개 퓨전 연산으로 줄어듭니다.

개별 커널 수준의 성능

비-GEMM 커널은 이전 구현(torch.ops._C) 대비 H100에서 일관되고 상당한 속도 향상을 보입니다.

• rms_norm_per_block_quant: +67% (torch.ops._C 대비 2.055×)
• silu_and_mul_per_block_quant: +73% (torch.ops._C 대비 2.269×)
• rms_norm_dynamic_per_token_quant: torch.ops._C 대비 1.80×
• fused_qk_norm_rope: torch.compile 대비 1.38×

GEMM 커널(scaled_mm)은 혼합 결과를 보입니다. **H100에서 1.08×**이지만 **B200에서는 0.739×**입니다. 저자들은 Blackwell 아키텍처에 대한 Triton 코드 생성의 한계가 원인이며, 해당 문제를 해결할 CuteDSL 백엔드가 개발 중임을 솔직하게 밝힙니다.

전체 모델 수준의 속도 향상

테스트는 H100 및 B200에서 Qwen3 모델(1.7B, 8B, 32B 파라미터)로 수행되었습니다.

• 토큰별 양자화(H100): Qwen3-1.7B에서 전체 처리량 약 1.05×
• 투기적 디코딩, Qwen3-8B, H100: 전체 처리량 최대 1.09× (배치 크기 16에서 TTFT 1.15× 향상)
• 그룹별 양자화(H100 및 B200): 모든 워크로드에서 약 1.05×

투기적 토큰 3개로 Qwen3-8B를 사용한 투기적 디코딩에서 TTFT(첫 번째 토큰까지의 시간) 1.33× 향상이 측정되었습니다.

실용적인 한계

저자들은 절충안에 대해 투명하게 밝힙니다. scaled_mm과 같은 대형 커널의 자동 튜닝은 세 Qwen3 변형에 대한 모든 입력 형태(168개의 고유 차원)를 커버하는 데 약 하루가 필요합니다. Helion의 런타임 디스패치는 커널 호출당 수십 마이크로초의 지연을 도입하므로, 최적 성능을 위해 CUDA 그래프 캡처가 필수입니다.

Helion 그룹별 양자화 커널은 아직 B200에서 DeepGEMM이 요구하는 UE8M0 형식과 호환되지 않으며, 이는 또 다른 미결 사항입니다.

향후 전망

이 프로젝트는 아홉 커널 모두의 구현과 함께 GitHub 이슈 vllm-project/vllm#32962로 공개되었습니다. 목표는 추가 아키텍처로 Helion 지원을 확장하고 CuteDSL 백엔드를 통해 Blackwell에서의 GEMM 성능을 해결하는 것입니다. 이 문제에 대한 Red Hat과 Meta PyTorch 팀의 협업은 PyTorch 생태계가 추론을 위한 고성능 커널의 하드웨어 이식성 문제를 진지하게 다루고 있다는 신호입니다.