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,