[vllm] vLLM, MXFP4 양자화 MoE 모델을 위한 CUTLASS 기반 SM100 커널 추가로 성능 향상

2026년 4월 18일수정: 2026년 4월 18일

PR 링크: vllm-project/vllm#37463 상태: Merged | 변경: +None / -None

들어가며

최근 대규모 언어 모델(LLM) 분야에서는 모델의 크기를 줄이면서도 성능을 유지하기 위한 양자화(Quantization) 기술이 핵심적인 역할을 하고 있습니다. 특히 Mixture-of-Experts (MoE) 아키텍처는 모델의 파라미터 수를 크게 늘리면서도 추론 시에는 일부 전문가(expert)만 활성화하여 효율성을 높이는 방식으로 주목받고 있습니다. vLLM은 이러한 최신 기술들을 빠르게 도입하여 LLM 추론 성능을 극대화하는 라이브러리로 알려져 있습니다.

이번 PR(#37128)은 vLLM에 MXFP4 양자화 방식을 사용하는 MoE 모델을 위한 새로운 커널을 추가하여 성능을 향상시키는 것을 목표로 합니다. 기존에는 NVFP4와 같은 다른 양자화 방식이나 다른 하드웨어 아키텍처에 대한 지원은 있었지만, MXFP4와 SM100(NVIDIA Blackwell 아키텍처) 조합에 대한 최적화된 커널은 부족했습니다. 이 PR은 CUTLASS 라이브러리를 활용하여 이러한 격차를 해소하고, 특정 하드웨어 및 양자화 방식에서 더 나은 추론 성능을 제공하고자 합니다.

코드 변경사항 분석

이번 PR의 핵심은 MXFP4 양자화된 MoE 모델을 SM100 GPU에서 효율적으로 실행하기 위한 새로운 CUDA 커널들을 추가하는 것입니다. 주요 변경사항은 다음과 같습니다.

1. 빌드 시스템 및 CMake 설정 업데이트

새로운 커널 파일들을 컴파일하고 링크하기 위해 CMakeLists.txt 파일이 수정되었습니다. 또한, CI/CD 파이프라인에서 새로운 커널에 대한 테스트가 포함되도록 .buildkite/test_areas/kernels.yaml 파일도 업데이트되었습니다.

CMakeLists.txt 변경:

--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -944,7 +944,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
       "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
       "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
-      "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
+      "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
+      "csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
+      "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
     set_gencode_flags_for_srcs(
       SRCS "${SRCS}"
       CUDA_ARCHS "${FP4_ARCHS}")

.buildkite/test_areas/kernels.yaml 변경:

--- a/.buildkite/test_areas/kernels.yaml
+++ b/.buildkite/test_areas/kernels.yaml
@@ -141,6 +141,7 @@ steps:
     - pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py
     - pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
     - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
+    - pytest -v -s tests/kernels/moe/test_mxfp4_moe.py
     - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
     - pytest -v -s tests/kernels/moe/test_flashinfer.py
     - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py

2. Ops 헤더 파일 업데이트 (`ops.h`)

새로운 MXFP4 관련 커널 함수들을 선언하기 위해 ops.h 파일에 함수 프로토타입이 추가되었습니다.

--- a/csrc/libtorch_stable/ops.h
+++ b/csrc/libtorch_stable/ops.h
@@ -134,4 +134,27 @@ void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out,
                               torch::stable::Tensor& input,
                               torch::stable::Tensor& input_global_scale);
 
+void mxfp4_experts_quant(
+    torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
+    torch::stable::Tensor const& input,
+    torch::stable::Tensor const& input_offset_by_experts,
+    torch::stable::Tensor const& output_scale_offset_by_experts,
+    int64_t n_experts);
+
+void silu_and_mul_mxfp4_experts_quant(
+    torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
+    torch::stable::Tensor const& input,
+    torch::stable::Tensor const& input_offset_by_experts,
+    torch::stable::Tensor const& output_scale_offset_by_experts,
+    int64_t n_experts);
+
+void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
+                            const torch::stable::Tensor& a,
+                            const torch::stable::Tensor& b,
+                            const torch::stable::Tensor& a_blockscale,
+                            const torch::stable::Tensor& b_blockscales,
+                            const torch::stable::Tensor& problem_sizes,
+                            const torch::stable::Tensor& expert_offsets,
+                            const torch::stable::Tensor& sf_offsets);
+    
 #endif

3. 새로운 MXFP4 커널 구현 (`mxfp4_blockwise_moe_kernel.cu`, `mxfp4_experts_quant.cu`)

이 PR의 핵심은 두 개의 새로운 CUDA 커널 파일입니다:

csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu: MXFP4 양자화된 가중치와 활성화를 위한 그룹화된 GEMM(General Matrix Multiply) 연산을 CUTLASS 라이브러리를 사용하여 구현합니다. 이 커널은 SM100 아키텍처에 최적화되어 있으며, mx_float4_t 타입과 E8M0 스케일 형식을 사용합니다. 기존 NVFP4 경로와는 다른 스케일링 방식(block-32 scales)을 사용합니다.
csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu: 활성화 양자화 커널로, SiLU 활성화 함수와 Mul 연산을 융합(fused)하여 성능을 향상시킵니다.

이 파일들은 CUTLASS의 Collective Operations API를 활용하여 GPU 하드웨어의 특성을 최대한 활용하도록 설계되었습니다. 특히, TMA(Tensor Memory Access)를 사용하여 효율적인 데이터 로딩 및 처리를 수행합니다.

4. 압축 텐서 모델 체크포인트와의 통합

PR 설명에 따르면, 이 변경사항은 압축된 텐서 형식의 모델 체크포인트와도 통합됩니다. 이는 MXFP4로 양자화된 모델을 vLLM에서 직접 로드하고 실행할 수 있음을 의미합니다. compressed_tensors 관련 모듈의 수정이 있었을 것으로 예상됩니다 (구체적인 diff는 제공되지 않았으나, compressed_tensors_moe.py 파일 언급이 리뷰 댓글에 있습니다).

왜 이게 좋은가?

이 PR은 여러 측면에서 vLLM의 성능과 기능성을 향상시킵니다.

MXFP4 양자화 지원 강화: MXFP4는 모델의 메모리 사용량을 크게 줄이면서도 비교적 높은 정확도를 유지할 수 있는 양자화 방식입니다. 이 PR을 통해 vLLM은 MXFP4로 양자화된 MoE 모델을 더 효율적으로 지원하게 되었습니다.
SM100 아키텍처 최적화: 새로운 커널은 SM100(NVIDIA Blackwell) 아키텍처의 특성을 고려하여 설계되었습니다. 이는 해당 아키텍처를 사용하는 GPU에서 더 높은 처리량과 낮은 지연 시간을 달성할 수 있게 합니다.
CUTLASS 활용: CUTLASS는 NVIDIA GPU를 위한 고성능 딥러닝 기본 연산(특히 GEMM)을 위한 템플릿 라이브러리입니다. CUTLASS를 활용함으로써 vLLM은 복잡한 GPU 프로그래밍을 직접 수행하는 대신, 검증되고 최적화된 라이브러리를 사용하여 개발 생산성과 성능을 동시에 확보할 수 있습니다.
성능 향상 (실제 결과): PR 설명에 포함된 테스트 결과는 이 최적화의 효과를 명확히 보여줍니다.
- MXFP4 W4A4 커널 사용 시:
  - Total latency: 24.276 s
  - Questions per second: 54.333
  - Output tokens per second: 6218.226
- Marlin w4a16 커널 사용 시 (기존):
  - Total latency: 28.203 s
  - Questions per second: 46.768
  - Output tokens per second: 5552.307
MXFP4 W4A4 커널을 사용했을 때, 지연 시간(latency)이 약 14% 감소했고, 초당 처리량(QPS, Tokens/sec)은 약 16% ~ 18% 증가했습니다. 이는 상당한 성능 향상입니다.
일반적인 교훈:
- 하드웨어별 최적화의 중요성: 특정 GPU 아키텍처(SM100)에 맞춰 커널을 최적화하면 상당한 성능 향상을 얻을 수 있습니다.
- 최신 양자화 기술 도입: MXFP4와 같은 새로운 양자화 기법을 지원함으로써 모델 압축 및 효율성 증대에 기여할 수 있습니다.
- 고성능 라이브러리 활용: CUTLASS와 같은 검증된 라이브러리를 활용하여 복잡한 연산을 효율적으로 구현하는 것은 좋은 전략입니다.
- MoE 아키텍처 지원: MoE 모델의 복잡성을 효율적으로 처리하기 위한 전용 커널 개발은 필수적입니다.

리뷰 댓글 분석 및 반영

리뷰 댓글들은 이 PR의 맥락과 향후 작업 방향을 이해하는 데 도움을 줍니다.

yzong-rh의 코멘트: 이 PR이 zyongye의 MXFP4 리팩토링 작업과 충돌할 수 있음을 지적합니다. 이는 코드 변경이 기존 작업과 어떻게 상호작용하는지, 잠재적 충돌을 어떻게 해결해야 하는지에 대한 중요한 고려사항입니다. 후속 작업에서 이러한 충돌이 해결되었거나, 이 PR이 해당 리팩토링을 포함하고 있을 가능성이 있습니다.
geraldstanje의 질문: RTX 6000 Pro (SM120)와 같은 다른 GPU에서도 작동하는지에 대한 질문입니다. mgoin은 이 PR이 SM100에만 해당하며, SM120 지원은 후속 작업으로 가능할 것이라고 답변합니다. 이는 이 PR의 범위를 명확히 하고, 향후 확장 가능성을 제시합니다.
dsikka의 코멘트: compressed_tensors_moe.py 파일의 업데이트 필요성을 언급합니다. mgoin은 후속 작업으로 미룰 수 있다고 답하며, 이는 기능 추가와 코드 정리를 분리하여 진행하는 일반적인 개발 방식입니다.
mgoin의 디버그 메시지 제거: compressed_tensors_w4a16_mxfp4.py 파일에서 불필요한 디버그 메시지를 제거한 것에 대한 논의입니다. 이는 코드의 가독성과 유지보수성을 위한 결정으로 보입니다.
zyongye의 dtype 질문: MoE 설정에서 a dtype을 명시하지 않았을 때의 의미에 대한 질문입니다. mgoin은 가중치와 활성화의 dtype이 동일함을 확인시켜 줍니다. 이는 코드의 동작 방식을 명확히 하는 데 도움이 됩니다.

종합적으로, 이 PR은 SM100 GPU에서 MXFP4 양자화 MoE 모델의 성능을 크게 향상시키는 중요한 개선 사항입니다. CUTLASS를 활용한 최적화된 커널 추가는 vLLM이 최신 하드웨어 및 모델 아키텍처를 지원하는 능력을 강화하며, 실제 성능 측정 결과는 이러한 개선의 효과를 입증합니다. 향후 SM120 지원 등 추가적인 개선이 기대됩니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [vllm] vLLM 멀티모달 스케줄러 오버헤드 최적화: Python List 캐싱으로 27% 성능 향상
현재글 : [vllm] vLLM, MXFP4 양자화 MoE 모델을 위한 CUTLASS 기반 SM100 커널 추가로 성능 향상
다음글 없음

댓글

관련 포스트

PR Analysis 의 다른글