[vllm] vLLM, DeepSeek V4 모델 성능 최적화: AITER MXFP4 BF16 백엔드 개선

2026년 6월 26일수정: 2026년 6월 26일

PR 링크: vllm-project/vllm#46122 상태: Merged | 변경: +28 / -15

들어가며

대규모 언어 모델(LLM)의 발전 속도가 빨라짐에 따라, 모델의 추론 성능을 극대화하는 것은 매우 중요한 과제가 되었습니다. 특히 Mixture-of-Experts (MoE) 아키텍처는 파라미터 수를 늘리면서도 계산량을 효율적으로 관리할 수 있어 주목받고 있습니다. 하지만 MoE 모델은 복잡한 라우팅 및 전문가 선택 로직으로 인해 최적화가 까다롭습니다.

vLLM은 LLM 추론을 위한 고성능 라이브러리로, 다양한 모델 아키텍처와 최적화 기법을 지원합니다. 이번 PR은 vLLM에서 DeepSeek V4 모델의 성능을 향상시키기 위해 AITER (Accelerated Inference Toolkit for ROCm)의 MXFP4 BF16 백엔드를 최적화하는 내용을 다룹니다. 특히 DeepSeek V4 모델의 가중치 준비 과정을 개선하여 추론 속도를 높이는 데 초점을 맞추고 있습니다.

이 글에서는 해당 PR의 코드 변경 사항을 상세히 분석하고, 이러한 변경이 왜 성능 향상으로 이어지는지, 그리고 기술적으로 어떤 의미를 가지는지 살펴보겠습니다.

코드 분석

이번 PR의 핵심 변경 사항은 vllm/model_executor/layers/fused_moe/oracle/mxfp4.py 파일 내 convert_weight_to_mxfp4_moe_kernel_format 함수의 AITER_MXFP4_BF16 백엔드 처리 로직에 집중되어 있습니다.

`vllm/model_executor/layers/fused_moe/oracle/mxfp4.py`

이 파일은 MXFP4 (Mixed Precision FP4) 형식의 가중치를 AITER 백엔드에서 사용할 수 있는 커널 형식으로 변환하는 역할을 합니다. DeepSeek V4와 같은 모델은 MoE 레이어에서 특정 가중치 형식과 처리가 필요하며, 이 PR은 해당 요구사항을 충족하도록 수정되었습니다.

1. AITER MXFP4 BF16 백엔드 로직 개선

가장 중요한 변경은 mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 조건부 블록 내에서 발생했습니다. 이 블록은 DeepSeek V4 모델의 원래 가중치를 AITER가 요구하는 형식으로 변환하는 로직을 포함합니다.

Before:

@@ -1429,38 +1429,51 @@ def convert_weight_to_mxfp4_moe_kernel_format(
         )
 
     elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
-        from vllm._aiter_ops import rocm_aiter_ops  # noqa: F401
+        # Initially introduced for DeepSeekV4
 
         if w13_bias is not None:
             w13_bias = w13_bias.data.to(torch.float32)
         if w2_bias is not None:
             w2_bias = w2_bias.data.to(torch.float32)
 
-        e, n, k = w13_weight.shape
+        import os
 
-        # No de-interleave: standard _load_w13 already produces
-        # [gate_all, up_all] layout.  Use aiter-native shuffle functions
-        # (matching aiter/ops/flydsl/test_flydsl_moe_a4w4.py pattern).
+        from aiter.ops.shuffle import shuffle_scale as _shuf_s
         from aiter.ops.shuffle import shuffle_weight as _shuf_w
-        from aiter.utility.fp4_utils import e8m0_shuffle as _e8m0_shuf
 
-        # w13 (gate+up, stage1): shuffle_weight with layout (16,16)
+        # TODO: Remove this once AITER is fixed
+        # Necessary for AITER side from crashing
+        os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0"
+
         w13_weight = torch.nn.Parameter(
-            _shuf_w(w13_weight.data.view(torch.float4_e2m1fn_x2), (16, 16)),
+            _shuf_w(
+                w13_weight.data.view(torch.float4_e2m1fn_x2),
+                is_guinterleave=True,
+                gate_up=True,
+            ),
             requires_grad=False,
         )
-        shuffled_w13_scale = _e8m0_shuf(
-            w13_weight_scale.view(-1, w13_weight_scale.shape[-1])
+        shuffled_w13_scale = _shuf_s(
+            w13_weight_scale.reshape(-1, w13_weight_scale.shape[-1]),
+            num_experts,
+            True,
+            True,
         )
 
-        # w2 (down-proj, stage2): same shuffle as w13 for a4w4 fp4x2
-        # (tuning script uses shuffle_weight((16,16)) + e8m0_shuffle for both)
         w2_weight = torch.nn.Parameter(
-            _shuf_w(w2_weight.data.view(torch.float4_e2m1fn_x2), (16, 16)),
+            _shuf_w(
+                w2_weight.data.view(torch.float4_e2m1fn_x2),
+                is_guinterleave=False,
+                gate_up=False,
+            ),
             requires_grad=False,
         )
-        shuffled_w2_scale = _e8m0_shuf(
-            w2_weight_scale.view(-1, w2_weight_scale.shape[-1])
+        # use_gu_interleave
+        shuffled_w2_scale = _shuf_s(
+            w2_weight_scale.reshape(-1, w2_weight_scale.shape[-1]),
+            num_experts,
+            True,
+            False,
         )
 
         return (

After:

from vllm._aiter_ops import rocm_aiter_ops 제거: 이 줄은 실제 사용되지 않아 제거되었습니다. (주석으로 noqa: F401이 있었지만, 명시적으로 제거하는 것이 좋습니다.)
주석 변경: Initially introduced for DeepSeekV4라는 주석이 추가되어 이 로직이 특정 모델에 맞춰져 있음을 명확히 했습니다. 또한, 리뷰어 Rohan138의 피드백에 따라 AITER의 근본적인 문제를 지적하는 주석(TODO: Remove this once AITER is fixed)과 환경 변수 설정(os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0")의 필요성을 설명하는 주석이 추가되었습니다. 이는 AITER 라이브러리 자체의 제약으로 인해 특정 환경 변수 설정을 강제해야 함을 나타냅니다.
aiter.ops.shuffle 사용 변경:
- shuffle_weight 함수 호출 시 is_guinterleave=True 및 gate_up 인자가 명시적으로 추가되었습니다. 이는 가중치 데이터의 인터리브(interleave) 여부와 게이트(gate) 및 업 프로젝션(up-projection) 가중치인지 여부를 정확히 지정합니다. DeepSeek V4의 가중치 레이아웃에 맞추기 위한 중요한 변경입니다.
- e8m0_shuffle 함수가 제거되고 aiter.ops.shuffle.shuffle_scale 함수가 사용되었습니다. shuffle_scale 함수는 num_experts, is_guinterleave, gate_up 등의 인자를 받아 스케일(scale) 데이터를 올바르게 셔플링합니다. 이는 이전의 e8m0_shuffle보다 더 명확하고 유연한 스케일 처리 방식을 제공합니다.
reshape 사용: w13_weight_scale.view(-1, w13_weight_scale.shape[-1]) 대신 w13_weight_scale.reshape(-1, w13_weight_scale.shape[-1])가 사용되었습니다. reshape는 view와 유사하지만, 데이터가 연속적이지 않은 경우에도 작동할 수 있어 더 안전한 선택일 수 있습니다. 여기서는 큰 차이가 없을 수 있으나, 일관성을 위해 사용된 것으로 보입니다.

이러한 변경은 DeepSeek V4 모델의 가중치 형식이 AITER의 MXFP4 BF16 커널과 정확히 일치하도록 보장하며, 특히 가중치 데이터의 인터리빙 및 스케일링 방식을 최적화하여 GPU에서 효율적으로 처리될 수 있도록 합니다.

왜 이게 좋은가?

성능 향상

PR 설명에 포함된 테스트 결과는 이러한 최적화가 실제 성능 향상으로 이어진다는 것을 명확히 보여줍니다.

성능 비교 (moeopt vs baseline):

C	output tok/s	total tok/s	mean TPOT	p99 TPOT	mean TTFT	p99 TTFT	mean E2E
1	+3.10%	+3.10%	+2.94%	+2.94%	+14.24%	+51.77%	+3.01%
16	+9.06%	+9.06%	+8.37%	+8.51%	+5.68%	-16.01%	+8.31%
32	+9.32%	+9.32%	+8.57%	+8.80%	+6.79%	+20.23%	+8.52%
64	+8.81%	+8.81%	+8.24%	+8.41%	+2.34%	+6.63%	+8.10%

표에서 볼 수 있듯이, moeopt (최적화된 버전)는 baseline (최적화 이전 버전)에 비해 출력 토큰/초 (output tok/s) 및 총 토큰/초 (total tok/s)에서 평균 3.10%에서 9.32%까지 향상되었습니다. 또한, 평균 응답 시간 (mean TTFT)은 최대 14.24% 감소했으며, 평균 처리 시간 (mean TPOT)도 최대 8.80% 단축되었습니다. 이는 GPU에서의 가중치 로딩 및 연산이 더 효율적으로 이루어지고 있음을 시사합니다.

특히 mean TTFT (Time To First Token)의 경우, 배치 크기(C)가 작을 때는 증가했지만, 배치 크기가 커질수록 감소하는 경향을 보입니다. 이는 배치 처리 시의 오버헤드가 줄어들고 GPU 활용률이 높아졌기 때문으로 해석할 수 있습니다. p99 TTFT의 경우, 배치 크기 16에서 오히려 감소한 것은 주목할 만하며, 이는 특정 워크로드에서 더 안정적인 성능을 제공할 수 있음을 나타냅니다.

LM eval 결과에서도 미미하지만 정확도 향상(gsm8k exact_match 0.9484 -> 0.9500)이 관찰되었습니다. 이는 성능 최적화가 모델의 정확성에 부정적인 영향을 미치지 않았음을 보여줍니다.

일반적 교훈

모델별 특화 최적화의 중요성: DeepSeek V4와 같이 특정 아키텍처나 가중치 형식을 가진 모델은 일반적인 최적화 기법만으로는 최고의 성능을 내기 어렵습니다. 모델의 세부적인 특성을 이해하고 이에 맞춰 라이브러리 내부 로직을 조정하는 것이 중요합니다.
하드웨어 및 라이브러리 의존성: AITER와 같은 하드웨어 가속 라이브러리는 특정 하드웨어(ROCm) 및 소프트웨어 스택에 최적화되어 있습니다. 이러한 라이브러리의 내부 동작 방식(예: 가중치 셔플링, 데이터 형식)을 정확히 이해하고 활용하는 것이 성능의 핵심입니다.
정확한 데이터 형식 및 레이아웃: MoE 모델에서 가중치는 종종 복잡한 형식(예: FP4, BF16)과 레이아웃(예: interleaved, non-interleaved)을 가집니다. 이러한 데이터가 GPU 커널에서 효율적으로 처리되려면, 라이브러리 레벨에서 해당 형식과 레이아웃을 정확하게 변환하고 준비하는 과정이 필수적입니다.
지속적인 성능 측정 및 검증: PR 설명에 포함된 LM eval 및 성능 측정 결과는 변경 사항이 의도한 대로 작동하는지, 그리고 회귀(regression)는 없는지를 확인하는 데 매우 중요합니다. 특히 복잡한 모델이나 라이브러리에서는 다양한 시나리오에서의 테스트가 필수적입니다.
오픈소스 협업 및 피드백: 리뷰어 Rohan138의 피드백은 gpt-oss-120b와 같은 다른 모델에 대한 회귀 테스트의 중요성을 강조합니다. 또한, AITER의 근본적인 이슈에 대한 주석은 오픈소스 프로젝트에서 개발자들이 서로의 작업에 대한 이해를 높이고, 장기적인 개선 방향을 설정하는 데 도움을 줍니다.

결론

이번 vLLM PR은 DeepSeek V4 모델의 추론 성능을 향상시키기 위해 AITER MXFP4 BF16 백엔드의 가중치 준비 로직을 정교하게 최적화했습니다. aiter.ops.shuffle 함수의 올바른 사용, is_guinterleave 및 gate_up 인자의 명시적 지정, 그리고 스케일 데이터 처리 방식의 개선을 통해 DeepSeek V4 모델의 가중치가 GPU에서 더욱 효율적으로 처리될 수 있도록 만들었습니다. 테스트 결과는 이러한 변경이 실제 추론 속도 향상으로 이어짐을 입증했으며, 동시에 모델의 정확성에는 부정적인 영향을 미치지 않았습니다.

이 PR은 LLM 추론 최적화 분야에서 모델별 특화, 하드웨어 라이브러리 활용, 그리고 정확한 데이터 처리의 중요성을 다시 한번 강조합니다. 앞으로도 vLLM과 같은 라이브러리들이 다양한 최신 모델들을 효율적으로 지원하기 위한 지속적인 노력과 개선이 기대됩니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [vllm] vLLM ROCm 환경에서 Shared-Expert Fusion을 통한 MoE 추론 성능 최적화
현재글 : [vllm] vLLM, DeepSeek V4 모델 성능 최적화: AITER MXFP4 BF16 백엔드 개선
다음글 [sglang] SGLang, CUDA 그래프 재실행 시 호스트-디바이스 동기화 제거로 성능 향상