[vllm] vLLM DeepSeek V4 모델의 불필요한 메모리 할당 및 복사 최적화 분석

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

들어가며

최근 vLLM 프로젝트에서 DeepSeek V4 모델의 성능을 개선하기 위한 Pull Request(PR)가 병합되었습니다. 이 PR의 핵심은 모델의 마지막 파이프라인(PP) 랭크에서 MTP(Multi-Token Prediction) 기능이 활성화되지 않았을 때 발생하는 불필요한 메모리 할당과 데이터 복사를 제거하는 것입니다. MTP는 여러 토큰을 한 번에 예측하여 추론 속도를 높이는 기술이지만, 이 기능이 사용되지 않을 때는 관련 버퍼 할당 및 복사 작업이 오버헤드로 작용할 수 있습니다. 본 글에서는 해당 PR의 코드 변경 사항을 상세히 분석하고, 이러한 최적화가 왜 성능 향상으로 이어지는지, 그리고 이를 통해 얻을 수 있는 일반적인 교훈은 무엇인지 살펴보겠습니다.

코드 변경 분석

이번 PR은 vllm/models/deepseek_v4/amd/model.py, vllm/models/deepseek_v4/nvidia/model.py, vllm/models/deepseek_v4/xpu/model.py 세 파일에서 유사한 변경을 적용했습니다. 변경 사항은 크게 __init__ 메서드에서의 버퍼 할당 조건 수정과 forward 메서드에서의 버퍼 복사 조건 수정으로 나눌 수 있습니다.

1. __init__ 메서드: MTP 미사용 시 버퍼 할당 제거

기존 코드에서는 파이프라인의 마지막 랭크(get_pp_group().is_last_rank)이고 MTP 기능이 활성화되지 않은 경우에도 _mtp_hidden_buffer라는 이름의 버퍼를 torch.empty를 사용하여 할당했습니다. 이는 MTP 기능이 실제로 사용되지 않더라도 메모리를 미리 확보해두는 동작이었습니다.

Before:

--- a/vllm/models/deepseek_v4/amd/model.py
+++ b/vllm/models/deepseek_v4/amd/model.py
@@ -573,13 +573,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             requires_grad=False,
         )
         self.hc_head_op = HCHeadOp()
-        # Pre-hc_head residual stream buffer for the MTP draft. Stable
-        # address (outside the cudagraph pool) so the copy_ in forward()
-        # refreshes it correctly across captured shapes.
-        # refreshes it correctly across captured shapes. Only allocated on
-        # the last PP rank — that's where MTP target hidden states are
-        # produced.
-        if get_pp_group().is_last_rank:
+        spec_config = vllm_config.speculative_config
+        needs_mtp_hidden_states = spec_config is not None and (
+            spec_config.use_eagle() or spec_config.uses_draft_model()
+        )
+        if get_pp_group().is_last_rank and needs_mtp_hidden_states:
             self._mtp_hidden_buffer = torch.empty(
                 vllm_config.scheduler_config.max_num_batched_tokens,
                 self.hc_dim,

수정된 코드에서는 _mtp_hidden_buffer를 할당하는 조건을 강화했습니다. 이제는 마지막 PP 랭크이면서 동시에 speculative_config가 존재하고, use_eagle() 또는 uses_draft_model() 메서드가 참을 반환할 때만 버퍼를 할당합니다. 이는 MTP 기능이 실제로 필요한 경우에만 메모리를 할당하도록 하여 불필요한 메모리 사용을 방지합니다.

After:

--- a/vllm/models/deepseek_v4/amd/model.py
+++ b/vllm/models/deepseek_v4/amd/model.py
@@ -573,13 +573,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             requires_grad=False,
         )
         self.hc_head_op = HCHeadOp()
-        # Pre-hc_head residual stream buffer for the MTP draft. Stable
-        # address (outside the cudagraph pool) so the copy_ in forward()
-        # refreshes it correctly across captured shapes.
-        # refreshes it correctly across captured shapes. Only allocated on
-        # the last PP rank — that's where MTP target hidden states are
-        # produced.
-        if get_pp_group().is_last_rank:
+        spec_config = vllm_config.speculative_config
+        needs_mtp_hidden_states = spec_config is not None and (
+            spec_config.use_eagle() or spec_config.uses_draft_model()
+        )
+        if get_pp_group().is_last_rank and needs_mtp_hidden_states:
             self._mtp_hidden_buffer = torch.empty(
                 vllm_config.scheduler_config.max_num_batched_tokens,
                 self.hc_dim,

2. forward 메서드: MTP 미사용 시 버퍼 복사 제거

forward 메서드에서는 이전 단계의 hidden_states를 _mtp_hidden_buffer로 복사하는 작업이 수행되었습니다. 이 복사 작업 역시 MTP 기능이 비활성화된 경우에는 불필요한 연산이었습니다.

Before:

--- a/vllm/models/deepseek_v4/amd/model.py
+++ b/vllm/models/deepseek_v4/amd/model.py
@@ -679,9 +677,9 @@ def forward(
         if not get_pp_group().is_last_rank:
             return IntermediateTensors({"hidden_states": hidden_states})
 
-        # Stash pre-hc_head residual for the MTP draft (captured copy_).
-        num_tokens = hidden_states.shape[0]
-        self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
+        if self._mtp_hidden_buffer is not None:
+            num_tokens = hidden_states.shape[0]
+            self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
 
         hidden_states = self.hc_head_op(
             hidden_states,

수정된 코드에서는 _mtp_hidden_buffer가 실제로 할당되었는지(is not None) 확인하는 조건을 추가했습니다. 이 조건이 참일 때만 copy_ 연산을 수행하도록 하여, MTP 기능이 비활성화되어 버퍼가 할당되지 않은 경우에는 이 복사 작업이 건너뛰게 됩니다.

After:

--- a/vllm/models/deepseek_v4/amd/model.py
+++ b/vllm/models/deepseek_v4/amd/model.py
@@ -679,9 +677,9 @@ def forward(
         if not get_pp_group().is_last_rank:
             return IntermediateTensors({"hidden_states": hidden_states})
 
-        # Stash pre-hc_head residual for the MTP draft (captured copy_).
-        num_tokens = hidden_states.shape[0]
-        self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
+        if self._mtp_hidden_buffer is not None:
+            num_tokens = hidden_states.shape[0]
+            self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
 
         hidden_states = self.hc_head_op(
             hidden_states,

왜 이게 좋은가?

이 PR의 가장 큰 장점은 불필요한 메모리 할당 및 데이터 복사를 제거함으로써 성능을 향상시켰다는 점입니다. 특히 MTP 기능이 사용되지 않는 시나리오에서는 이러한 변경이 직접적인 성능 개선으로 이어집니다.

PR 설명에 포함된 테스트 스크립트 결과는 이러한 개선 효과를 명확히 보여줍니다. MTP 버퍼 크기는 최대 448 MiB에 달하며, 이는 MAX_NUM_TOKENS가 8192일 때 HC_MULT * HIDDEN_SIZE의 곱으로 계산됩니다. 이 버퍼에 데이터를 복사하는 데 걸리는 시간 또한 상당합니다. 예를 들어, 8192 토큰의 경우 142.941 us가 소요됩니다.

tokens  copy size       copy latency
1       0.055 MiB       5.320 us
8       0.438 MiB       5.243 us
32      1.750 MiB       4.644 us
128     7.000 MiB       4.589 us
512     28.000 MiB      6.103 us
2048    112.000 MiB     37.346 us
8192    448.000 MiB     142.941 us

이 PR을 통해 MTP 기능이 비활성화된 경우, 위에서 언급된 448 MiB의 메모리 할당과 해당 버퍼로의 데이터 복사가 완전히 생략됩니다. 이는 다음과 같은 이점을 제공합니다:

1. 메모리 절약: 특히 대규모 모델이나 많은 GPU를 사용하는 분산 환경에서 상당한 GPU 메모리를 절약할 수 있습니다. PR 설명에 따르면 최대 448 MiB의 GPU 메모리를 절약할 수 있다고 합니다.
2. 계산 시간 단축: 불필요한 torch.empty 호출과 copy_ 연산이 제거되므로, 해당 연산에 소요되던 시간이 절약됩니다. 이는 전체 추론 지연 시간을 줄이는 데 기여합니다.
3. 코드 간결성 및 유지보수성 향상: 조건부 로직을 통해 코드의 복잡성을 줄이고, 실제 필요한 경우에만 리소스를 사용하도록 하여 코드의 의도를 명확히 합니다.

일반적인 교훈:

• 조건부 리소스 관리: 모든 기능을 항상 활성화하거나 리소스를 할당하는 대신, 기능의 활성화 여부에 따라 리소스를 동적으로 할당하고 해제하는 것이 중요합니다. 이는 특히 옵션이 많은 라이브러리나 프레임워크에서 유용합니다.
• 측정의 중요성: PR 설명에 포함된 성능 측정 스크립트는 최적화의 효과를 정량적으로 입증하는 데 결정적인 역할을 했습니다. 실제 사용 사례를 기반으로 성능을 측정하고 그 결과를 공유하는 것은 코드 변경의 가치를 설득하는 데 필수적입니다.
• 파이프라인 병렬 처리에서의 주의점: 분산 환경, 특히 파이프라인 병렬 처리에서는 각 랭크의 역할과 책임이 명확해야 합니다. 특정 랭크에서만 필요한 연산이나 할당을 다른 랭크로 확장하지 않도록 주의해야 합니다.

결론

이번 vLLM PR은 DeepSeek V4 모델에서 MTP 기능이 비활성화되었을 때 발생하는 불필요한 메모리 할당 및 복사 작업을 제거함으로써, GPU 메모리 사용량을 줄이고 추론 성능을 향상시키는 효과적인 최적화를 이루었습니다. 이는 코드의 효율성을 높이는 좋은 사례이며, 조건부 리소스 관리와 성능 측정의 중요성을 다시 한번 강조합니다. 이러한 세심한 최적화는 vLLM이 고성능 언어 모델 추론을 위한 강력한 도구로 자리매김하는 데 기여할 것입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.empty.html
• https://pytorch.org/docs/stable/generated/torch.Tensor.copy_.html

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