[vllm] vLLM, DeepSeek-V3.2/GLM-5.2 MTP 경로 최적화: All-Reduce 융합 및 로컬 Argmax 도입

2026년 8월 15일수정: 2026년 8월 15일

PR 링크: vllm-project/vllm#49793 상태: Merged | 변경: +33 / -5

들어가며

최근 대규모 언어 모델(LLM)의 발전 속도는 눈부십니다. 특히 추론(inference) 단계에서의 성능 최적화는 모델의 실제 활용성을 결정짓는 중요한 요소입니다. vLLM은 이러한 추론 성능 향상을 위해 지속적으로 다양한 최적화 기법을 도입하고 있습니다. 이번 글에서는 vLLM 프로젝트의 GitHub Pull Request(PR)에서 제안된 DeepSeek-V3.2 및 GLM-5.2 모델의 Multi-Token Prediction (MTP) 경로에 대한 두 가지 주요 성능 최적화를 심층적으로 분석하고자 합니다.

이 PR은 특히 MTP 경로에서 발생하는 두 가지 병목 현상을 해결하는 데 중점을 둡니다. 첫째, 시퀀스 병렬화(sequence parallelism)를 사용하지 않는 경로에서 발생하는 후반부 All-Reduce 연산을 최종 RMSNorm 연산과 융합(fuse)하여 통신 오버헤드를 줄입니다. 둘째, 전체 어휘(vocabulary)에 대한 All-Gather 연산을 건너뛰고 각 텐서 병렬 처리(tensor model parallelism) 단위 내에서 지역적(local)으로 최대값을 찾는 Argmax 연산을 사용하여, 제안된(draft) 토큰을 더 효율적으로 생성합니다. 이 두 가지 개선 사항이 어떻게 실제 코드에서 구현되었고, 어떤 성능 향상을 가져왔는지 자세히 살펴보겠습니다.

코드 분석

이번 PR은 주로 vllm/models/deepseek_v32/nvidia/mtp.py 파일을 수정하여 두 가지 주요 최적화를 적용했습니다.

1. 후반부 All-Reduce와 최종 RMSNorm 융합

기존의 MTP(Multi-Token Prediction) 경로에서는 시퀀스 병렬화(sequence parallelism)를 사용하지 않는 경우, MoE(Mixture of Experts) 레이어의 출력이 최종적으로 tensor_model_parallel_all_reduce를 통해 모든 장치에 걸쳐 동기화되었습니다. 이후 별도의 RMSNorm 연산이 수행되었습니다.

Before:

--- a/vllm/models/deepseek_v32/nvidia/mtp.py
+++ b/vllm/models/deepseek_v32/nvidia/mtp.py
@@ -119,9 +119,6 @@ def forward(
         hidden_states, residual = self.mtp_block(
             positions=positions, hidden_states=hidden_states, residual=None
         )
-        if not is_sequence_parallel:
-            # Without sequence parallelism, the MoE output is left un-reduced.
-            hidden_states = tensor_model_parallel_all_reduce(hidden_states)
         # Recycle the POST-final-norm hidden into the next draft step. The
         # residual-add is fused into the final RMSNorm so it is computed
         # exactly once, and the result is returned for both tuple positions:
@@ -132,9 +129,15 @@ def forward(
         # is understood by both the V2 speculator (isinstance-tuple check) and
         # the legacy proposer (model_returns_tuple is True for the
         # DeepSeekMTPModel architecture).
-        hidden_states, _ = self.shared_head.norm(hidden_states, residual)
         if is_sequence_parallel:
+            hidden_states, _ = self.shared_head.norm(hidden_states, residual)
             hidden_states = sp_all_gather(hidden_states)[: positions.shape[0]]
+        else:
+            # The MoE output is left un-reduced; fuse its all-reduce into the
+            # final norm, as the main model does at layer boundaries.
+            hidden_states, _ = fused_allreduce_rms_norm(
+                hidden_states, residual, self.shared_head.norm
+            )
         return hidden_states, hidden_states

After:

--- a/vllm/models/deepseek_v32/nvidia/mtp.py
+++ b/vllm/models/deepseek_v32/nvidia/mtp.py
@@ -119,9 +119,6 @@ def forward(
         hidden_states, residual = self.mtp_block(
             positions=positions, hidden_states=hidden_states, residual=None
         )
-        if not is_sequence_parallel:
-            # Without sequence parallelism, the MoE output is left un-reduced.
-            hidden_states = tensor_model_parallel_all_reduce(hidden_states)
         # Recycle the POST-final-norm hidden into the next draft step. The
         # residual-add is fused into the final RMSNorm so it is computed
         # exactly once, and the result is returned for both tuple positions:
@@ -132,9 +129,15 @@ def forward(
         # is understood by both the V2 speculator (isinstance-tuple check) and
         # the legacy proposer (model_returns_tuple is True for the
         # DeepSeekMTPModel architecture).
-        hidden_states, _ = self.shared_head.norm(hidden_states, residual)
         if is_sequence_parallel:
+            hidden_states, _ = self.shared_head.norm(hidden_states, residual)
             hidden_states = sp_all_gather(hidden_states)[: positions.shape[0]]
+        else:
+            # The MoE output is left un-reduced; fuse its all-reduce into the
+            # final norm, as the main model does at layer boundaries.
+            hidden_states, _ = fused_allreduce_rms_norm(
+                hidden_states, residual, self.shared_head.norm
+            )
         return hidden_states, hidden_states

vllm/models/deepseek_v32/nvidia/mtp.py 파일의 forward 메서드에서, is_sequence_parallel이 False일 경우, 기존에는 tensor_model_parallel_all_reduce를 호출한 뒤 self.shared_head.norm을 적용했습니다. 하지만 이 PR에서는 이 두 연산을 fused_allreduce_rms_norm 함수를 사용하여 하나로 융합했습니다. 이 함수는 내부적으로 All-Reduce 연산과 RMSNorm 연산을 결합하여 실행합니다. 이는 통신과 연산의 오버헤드를 줄여 성능을 향상시킵니다. 특히, 메인 모델의 레이어 경계에서 이미 유사한 융합이 이루어지고 있다는 점을 고려할 때, MTP 경로에서도 일관된 최적화를 적용한 것입니다.

2. 제안된 토큰 생성을 위한 로컬 Argmax 도입

Speculative Decoding에서 제안된(draft) 토큰을 생성할 때, 일반적으로 전체 어휘(vocabulary)에 대한 로그 확률(logits)을 계산한 후, 가장 확률이 높은 토큰을 선택합니다. 이 과정에서 전체 어휘에 대한 정보를 취합하기 위해 compute_logits 함수 내에서 all_gather 연산이 필요했습니다. 하지만 이 PR에서는 get_top_tokens라는 새로운 메서드를 도입하여 이 부분을 최적화했습니다.

Before (Conceptual): compute_logits에서 전체 어휘에 대한 all_gather 후 Argmax 수행

After: get_top_tokens 메서드에서 각 텐서 병렬 처리(TP) 장치 내에서 지역적(local)으로 Argmax 수행

--- a/vllm/models/deepseek_v32/nvidia/mtp.py
+++ b/vllm/models/deepseek_v32/nvidia/mtp.py
@@ -212,6 +215,23 @@ def compute_logits(
         # second RMSNorm.        return self.logits_processor(mtp_layer.shared_head.head, hidden_states)
 
+    def get_top_tokens(
+        self, 
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        """Greedy draft token ids via per-rank argmax over the vocab shard.
+
+        Saves the full-vocab all-gather ``compute_logits`` does; same tokens.
+        Name is fixed by the protocol the proposer probes for
+        (``use_local_argmax_reduction``). 
+        """
+        current_step_idx = spec_step_idx % self.num_mtp_layers
+        mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
+        return self.logits_processor.get_top_tokens(
+            mtp_layer.shared_head.head, hidden_states
+        )
+
 
 class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -262,6 +282,14 @@ def compute_logits(
     ) -> torch.Tensor | None:
         return self.model.compute_logits(hidden_states, spec_step_idx)
 
+    def get_top_tokens(
+        self,
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        """See ``DeepseekV32MultiTokenPredictor.get_top_tokens``."""
+        return self.model.get_top_tokens(hidden_states, spec_step_idx)
+
     def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
         spec_layer_weight_names = [
             "embed_tokens",

DeepseekV32MTP 클래스 내에 get_top_tokens 메서드가 추가되었습니다. 이 메서드는 compute_logits와 유사하게 MTP 레이어의 헤드(mtp_layer.shared_head.head)를 사용하지만, logits_processor.get_top_tokens를 호출하여 각 TP 장치 내에서 독립적으로 가장 높은 확률을 가진 토큰 ID를 반환합니다. 이는 전체 어휘에 대한 all_gather 연산을 생략함으로써 통신 비용을 크게 절감합니다. Proposer 측에서 use_local_argmax_reduction 플래그를 통해 이 기능을 사용할 수 있는지 확인하므로, 모델 측에서도 이에 대응하는 기능을 제공하는 것입니다.

왜 이게 좋은가?

이 PR에서 제안된 두 가지 최적화는 vLLM의 MTP 경로 성능을 실질적으로 향상시킵니다.

성능 향상

벤치마크 결과에 따르면, 특히 동시성(concurrency)이 높은 환경(64)에서 상당한 성능 향상이 관찰되었습니다. 배치 크기 1에서는 큰 변화가 없었지만, 이는 해당 시나리오에서 all_gather 연산이 차지하는 비중이 매우 작기 때문입니다. 그러나 동시성 64 환경에서는 다음과 같은 결과가 나타났습니다:

Baseline: 2923.2 output tok/s
+both (최적화 적용): 3321.1 output tok/s (+13.6%)

이는 단순히 추측(speculation)의 성공률 증가 때문이 아니라, 실제 처리 속도가 향상되었음을 의미합니다. 특히, acceptance length가 약간 감소했음에도 불구하고 처리량이 증가했다는 점은 주목할 만합니다. 이는 get_top_tokens가 제거한 all_gather 연산이 동시성 환경에서 상당한 병목이었음을 시사합니다. 이 연산은 동시성 64 환경에서 각 토큰당 약 19MB의 데이터를 전송해야 했으며, 이는 5번 반복될 수 있었습니다. 로컬 Argmax는 이 통신 비용을 제거하여 처리량을 약 17.8% (throughput normalized by acceptance length) 향상시켰습니다.

또한, 후반부 All-Reduce와 RMSNorm 융합은 통신과 연산의 결합을 통해 잠재적인 레이턴시 감소 효과를 제공합니다. 비록 벤치마크에서 이 융합 단독의 효과는 명확히 분리되지 않았지만, 메인 모델 레이어에서의 유사한 융합이 성능 향상에 기여했음을 고려할 때 MTP 경로에서도 긍정적인 영향을 기대할 수 있습니다.

정확도 유지

성능 향상과 더불어, 정확도 측면에서도 큰 변화가 없다는 점이 중요합니다. GSM8K 데이터셋을 사용한 5-shot 평가 결과, Baseline과 최적화 적용 후의 결과가 표준 편차 범위 내에서 거의 동일했습니다. 이는 성능 최적화가 모델의 추론 정확성을 저해하지 않음을 보여줍니다.

일반적인 교훈

통신과 연산의 융합 (Fusion): 서로 다른 연산들이 특정 조건 하에서 함께 수행될 수 있다면, 이를 융합하여 중간 결과 저장을 줄이고 통신 및 연산 오버헤드를 감소시킬 수 있습니다. 특히 GPU와 같은 병렬 처리 환경에서는 이러한 융합이 성능에 큰 영향을 미칩니다.
지역적 연산 활용 (Local Computation): 전체 데이터를 동기화하여 연산하는 대신, 각 처리 단위(예: TP 장치) 내에서 독립적으로 연산을 수행하고 결과를 취합하는 방식은 통신 비용을 크게 절감할 수 있습니다. 이는 특히 분산 환경에서 중요한 최적화 기법입니다.
병목 지점 식별 및 제거: 벤치마크 결과를 통해 성능 병목이 발생하는 지점을 정확히 파악하고, 해당 지점을 제거하거나 완화하는 데 집중하는 것이 효과적입니다. 동시성이 높아질수록 통신 관련 병목이 두드러지는 경향이 있습니다.
정확도와 성능의 균형: 성능 최적화를 적용할 때는 항상 정확도에 미치는 영향을 함께 고려해야 합니다. 이 PR은 성능을 크게 향상시키면서도 정확도를 유지하는 성공적인 사례를 보여줍니다.

리뷰 피드백 반영

리뷰 과정에서 몇 가지 중요한 논의가 있었습니다. [WoosukKwon]님은 이 PR이 여전히 필요한지 질문했고, [zhou9402]님은 이전 PR에서 일부 기능이 이미 메인으로 병합되었기 때문에 이 PR이 더 단순해졌다고 설명했습니다. 이는 PR의 유지보수성과 관련성을 명확히 하는 과정이었습니다.

또한, [TheEpicDolphin]님은 Speculator의 on_prefill_begin 및 on_prefill_end 호출 누락을 지적했습니다. [zhou9402]님은 이 함수들이 현재는 no-op(No Operation)이지만, 유연성과 완전성을 위해 호출하는 것이 좋다고 동의하며 수정했습니다. 이는 코드의 견고성(robustness)과 향후 확장성을 고려한 중요한 개선 사항입니다.

결론

vLLM의 이번 PR은 DeepSeek-V3.2 및 GLM-5.2 모델의 MTP 경로에서 발생하는 통신 및 연산 병목을 효과적으로 해결했습니다. 후반부 All-Reduce와 RMSNorm의 융합, 그리고 로컬 Argmax를 통한 전체 어휘 All-Gather 연산의 제거는 특히 동시성이 높은 환경에서 상당한 성능 향상을 가져왔습니다. 이러한 최적화는 LLM 추론 성능을 한 단계 끌어올리는 데 기여하며, vLLM이 지속적으로 발전하고 있음을 보여주는 좋은 예시입니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [onnxruntime] ONNX Runtime CUDA ArgMax/ArgMin 연산 최적화: 대규모 행렬 병렬 처리 개선
현재글 : [vllm] vLLM, DeepSeek-V3.2/GLM-5.2 MTP 경로 최적화: All-Reduce 융합 및 로컬 Argmax 도입
다음글 없음