[DSv4 Perf] Fix redundant memory allocation and copy for dsv4 pp buffer, 448 MiB GPU memory saved (#50312)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 30, 2026 at 11:38 UTC
904fae8be12f1592b8e489fc0f1004fa49460a89
3 files changed
+24
-28
vllm/models/deepseek_v4/amd/model.py
+8
-10
@@ -573,13 +573,11 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
573
requires_grad=False,
574
)
575
self.hc_head_op = HCHeadOp()
576
- # Pre-hc_head residual stream buffer for the MTP draft. Stable
577
- # address (outside the cudagraph pool) so the copy_ in forward()
578
- # refreshes it correctly across captured shapes.
579
- # refreshes it correctly across captured shapes. Only allocated on
580
- # the last PP rank — that's where MTP target hidden states are
581
- # produced.
582
- if get_pp_group().is_last_rank:
576
+ spec_config = vllm_config.speculative_config
577
+ needs_mtp_hidden_states = spec_config is not None and (
578
+ spec_config.use_eagle() or spec_config.uses_draft_model()
579
+ )
580
+ if get_pp_group().is_last_rank and needs_mtp_hidden_states:
581
self._mtp_hidden_buffer = torch.empty(
582
vllm_config.scheduler_config.max_num_batched_tokens,
583
self.hc_dim,
@@ -679,9 +677,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
677
if not get_pp_group().is_last_rank:
678
return IntermediateTensors({"hidden_states": hidden_states})
679
682
- # Stash pre-hc_head residual for the MTP draft (captured copy_).
683
- num_tokens = hidden_states.shape[0]
684
- self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
680
+ if self._mtp_hidden_buffer is not None:
681
+ num_tokens = hidden_states.shape[0]
682
+ self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
683
684
hidden_states = self.hc_head_op(
685
hidden_states,
vllm/models/deepseek_v4/nvidia/model.py
+8
-10
@@ -1039,13 +1039,11 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1039
torch.empty(1, dtype=torch.float32),
1040
requires_grad=False,
1041
)
1042
- # Pre-hc_head residual stream buffer for the MTP draft. Stable
1043
- # address (outside the cudagraph pool) so the copy_ in forward()
1044
- # refreshes it correctly across captured shapes.
1045
- # refreshes it correctly across captured shapes. Only allocated on
1046
- # the last PP rank — that's where MTP target hidden states are
1047
- # produced.
1048
- if get_pp_group().is_last_rank:
1042
+ spec_config = vllm_config.speculative_config
1043
+ needs_mtp_hidden_states = spec_config is not None and (
1044
+ spec_config.use_eagle() or spec_config.uses_draft_model()
1045
+ )
1046
+ if get_pp_group().is_last_rank and needs_mtp_hidden_states:
1047
self._mtp_hidden_buffer = torch.empty(
1048
vllm_config.scheduler_config.max_num_batched_tokens,
1049
self.hc_dim,
@@ -1130,9 +1128,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1128
if not get_pp_group().is_last_rank:
1129
return IntermediateTensors({"hidden_states": hidden_states})
1130
1133
- # Stash pre-hc_head residual for the MTP draft (captured copy_).
1134
- num_tokens = hidden_states.shape[0]
1135
- self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
1131
+ if self._mtp_hidden_buffer is not None:
1132
+ num_tokens = hidden_states.shape[0]
1133
+ self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
1134
1135
hidden_states = hc_head_fused_kernel_tilelang(
1136
hidden_states,
vllm/models/deepseek_v4/xpu/model.py
+8
-8
@@ -1057,11 +1057,11 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1057
requires_grad=False,
1058
)
1059
self.hc_head_op = HCHeadOp()
1060
- # Pre-hc_head residual stream buffer for the MTP draft. Stable
1061
- # address so the copy_ in forward() refreshes it correctly across
1062
- # captured shapes. Only allocated on the last PP rank — that's
1063
- # where MTP target hidden states are produced.
1064
- if get_pp_group().is_last_rank:
1060
+ spec_config = vllm_config.speculative_config
1061
+ needs_mtp_hidden_states = spec_config is not None and (
1062
+ spec_config.use_eagle() or spec_config.uses_draft_model()
1063
+ )
1064
+ if get_pp_group().is_last_rank and needs_mtp_hidden_states:
1065
self._mtp_hidden_buffer = torch.empty(
1066
vllm_config.scheduler_config.max_num_batched_tokens,
1067
self.hc_dim,
@@ -1139,9 +1139,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1139
if not get_pp_group().is_last_rank:
1140
return IntermediateTensors({"hidden_states": hidden_states})
1141
1142
- # Stash pre-hc_head residual for the MTP draft (captured copy_).
1143
- num_tokens = hidden_states.shape[0]
1144
- self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
1142
+ if self._mtp_hidden_buffer is not None:
1143
+ num_tokens = hidden_states.shape[0]
1144
+ self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
1145
1146
hidden_states = self.hc_head_op(
1147
hidden_states,