Helion을 vLLM에 통합해 양자화 GEMM을 자동 최적화하다

배경

대규모 언어 모델 추론에서 양자화된 선형 계층은 처리량과 지연시간에 직접적인 영향을 줍니다. vLLM은 FP8, INT8, INT4, NVFP4 같은 형식에 대해 CUTLASS, DeepGEMM, FlashInfer 등의 최적화 구현을 활용하고 있습니다. 그러나 가장 효율적인 행렬 곱셈 방식은 모델, 배치와 입력 형상, 양자화 방식, GPU에 따라 달라집니다. 범용적으로 최적화된 커널이 모든 배포 환경에서 최선이라고 보기는 어렵습니다.

PyTorch가 소개한 이번 작업은 Helion을 vLLM의 linear backend에 통합하는 방식입니다. Helion은 PyTorch 기반의 하드웨어 비종속 커널 DSL로, 타일 프로그래밍 모델을 사용합니다. 개발자가 형상과 하드웨어별 커널을 일일이 작성하는 대신, 알고리즘과 타일링, 메모리 배치, 스케줄링의 선택지를 튜닝 가능한 매개변수로 노출합니다. 이후 AOT autotuner가 주어진 워크로드에 적합한 조합을 탐색합니다.

핵심 내용

• 여러 GEMM 알고리즘을 하나로 통합한다. 하나의 구현이 Standard GEMM뿐 아니라 Split-K와 Swap-AB도 다룰 수 있습니다. Split-K는 K 차원을 여러 스레드 블록으로 나눠 M 또는 N이 작을 때 병렬성을 높이고, Swap-AB는 전치 기반 계산으로 작은 M 형상에서 GPU 활용도를 개선하는 방식입니다.
• 수동 디스패치 규칙을 줄인다. 기존 방식에서는 형상별로 여러 커널을 비교한 뒤 어떤 변형을 사용할지 휴리스틱을 직접 작성해야 했습니다. Helion에서는 알고리즘 선택도 튜닝 공간에 포함되므로, 자동 튜너가 각 형상에 맞는 변형과 저수준 설정을 함께 선택합니다.
• Hopper의 양자화 GEMM에 집중한다. 현재 구현은 Helion의 Triton backend를 사용하며 FP8 dynamic quantization, W8A8 INT8, Block-FP8을 대상으로 합니다. Blackwell에서는 CuteDSL backend의 초기 GEMM 결과도 경쟁력 있는 수준으로 보고됐지만, 폭넓은 지원은 해당 백엔드의 성숙도에 달려 있습니다.
• 혼합 디스패치를 적용한다. Hopper에서 형상별 튜닝과 hybrid dispatch를 결합한 Helion linear backend는 평가된 모델에서 vLLM의 기본 CUTLASS와 DeepGEMM backend보다 높은 성능을 보였습니다. 일부 워크로드에서는 엔드투엔드 처리량이 10% 넘게 증가했습니다. 다만 이 결과는 테스트한 하드웨어와 모델, 형상 범위에 한정됩니다.

성능을 얻기 위한 비용

자동 튜닝은 최적화 비용을 없애기보다 체계화합니다. 세밀한 AOT 탐색에는 여전히 수 시간이 걸릴 수 있습니다. 또한 vLLM 시작 중 CUDA Graph capture가 Helion의 JIT 컴파일을 유발해 콜드 스타트 지연을 늘릴 수 있습니다. 컴파일 결과를 캐시하면 웜 스타트에서 이 부담을 크게 줄일 수 있습니다. CUDA Graph 범위 밖에서는 CPU 측 디스패치와 커널 실행 비용이 추가되어 GPU 성능 향상의 일부를 상쇄할 가능성도 있습니다.

유지보수 역시 중요한 문제입니다. 인기 모델별로 튜닝된 설정을 상위 프로젝트에 많이 포함하면 설정 파일이 커지고, 일반적인 단위 테스트와 CI로 모든 조합을 검증하기 어려워집니다. 즉 성능, 사용 편의성, 유지보수성 사이에는 구조적인 절충이 있습니다. 특정 워크로드에 더 높은 성능을 맞추려면 더 긴 튜닝 시간이나 더 많은 설정 관리가 필요합니다.

의미와 영향

이번 통합은 추론 인프라에서 커널 개발의 방식을 바꿀 수 있는 사례입니다. 개발자는 여러 알고리즘을 하나의 고수준 구현으로 표현하고, 자동 튜너가 하드웨어와 워크로드에 맞는 구성을 찾게 할 수 있습니다. 이에 따라 서비스 운영팀이 복잡한 GPU 커널을 처음부터 작성하지 않고도 자체 모델에 맞춘 최적화를 시도할 여지가 커집니다. 다만 요청 형상이 계속 바뀌거나 빠른 기동이 중요한 서비스에서는 기존의 사전 최적화 backend가 더 단순할 수 있습니다. Helion의 실제 가치는 GPU 성능 이득이 컴파일, 디스패치, 유지보수 비용을 충분히 상쇄하는지에 달려 있습니다.

PyTorch Blog