From 904fae8be12f1592b8e489fc0f1004fa49460a89 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Thu, 30 Jul 2026 11:38:51 -0400 Subject: [PATCH] [DSv4 Perf] Fix redundant memory allocation and copy for dsv4 pp buffer, 448 MiB GPU memory saved (#50312) Signed-off-by: yewentao256 --- vllm/models/deepseek_v4/amd/model.py | 18 ++++++++---------- vllm/models/deepseek_v4/nvidia/model.py | 18 ++++++++---------- vllm/models/deepseek_v4/xpu/model.py | 16 ++++++++-------- 3 files changed, 24 insertions(+), 28 deletions(-) diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index b382f7fd6f..cdd338d636 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -573,13 +573,11 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): 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, @@ -679,9 +677,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): 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, diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 96a3c77074..bbda009372 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1039,13 +1039,11 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): torch.empty(1, dtype=torch.float32), requires_grad=False, ) - # 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, @@ -1130,9 +1128,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): 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 = hc_head_fused_kernel_tilelang( hidden_states, diff --git a/vllm/models/deepseek_v4/xpu/model.py b/vllm/models/deepseek_v4/xpu/model.py index f70cac1bce..4c6155c842 100644 --- a/vllm/models/deepseek_v4/xpu/model.py +++ b/vllm/models/deepseek_v4/xpu/model.py @@ -1057,11 +1057,11 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): requires_grad=False, ) self.hc_head_op = HCHeadOp() - # Pre-hc_head residual stream buffer for the MTP draft. Stable - # address so the copy_ in forward() 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, @@ -1139,9 +1139,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): 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,