A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 10:30 UTC
[Bugfix] Fix DSV4 MTP after ROCm mHC integration (#42930)
Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
Mohammad Miadh Angkad committed
May 19, 2026 at 01:02 UTC
67f58ce23f469e118688a50687ef0fbb14a1c028
2 files changed
+17
-12
vllm/model_executor/models/deepseek_v4.py
+12
-8
@@ -1261,10 +1261,12 @@ class DeepseekV4DecoderLayer(nn.Module):
1261
x: torch.Tensor,
1262
positions: torch.Tensor,
1263
input_ids: torch.Tensor | None,
1264
- post_mix: torch.Tensor | None,
1265
- res_mix: torch.Tensor | None,
1266
- residual: torch.Tensor | None,
1267
- ) -> torch.Tensor:
1264
+ post_mix: torch.Tensor | None = None,
1265
+ res_mix: torch.Tensor | None = None,
1266
+ residual: torch.Tensor | None = None,
1267
+ ) -> tuple[
1268
+ torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
1269
+ ]:
1270
residual = x
1271
x, post, comb = self.hc_pre(
1272
x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
@@ -1288,10 +1290,12 @@ class DeepseekV4DecoderLayer(nn.Module):
1290
x: torch.Tensor,
1291
positions: torch.Tensor,
1292
input_ids: torch.Tensor | None,
1291
- post_mix: torch.Tensor | None,
1292
- res_mix: torch.Tensor | None,
1293
- residual: torch.Tensor | None,
1294
- ) -> torch.Tensor:
1293
+ post_mix: torch.Tensor | None = None,
1294
+ res_mix: torch.Tensor | None = None,
1295
+ residual: torch.Tensor | None = None,
1296
+ ) -> tuple[
1297
+ torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
1298
+ ]:
1299
if current_platform.is_rocm():
1300
return self._forward_rocm(
1301
x, positions, input_ids, post_mix, res_mix, residual
vllm/model_executor/models/deepseek_v4_mtp.py
+5
-4
@@ -146,9 +146,10 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
146
hidden_states, residual, post_mix, res_mix = self.mtp_block(
147
positions=positions, x=hidden_states, input_ids=None
148
)
149
- hidden_states = self.mtp_block.hc_post(
150
- hidden_states, residual, post_mix, res_mix
151
- )
149
+ if current_platform.is_cuda():
150
+ hidden_states = self.mtp_block.hc_post(
151
+ hidden_states, residual, post_mix, res_mix
152
+ )
153
# Return the flat pre-hc_head residual so it can be re-fed as the
154
# next spec step's `previous_hidden_states` when
155
# num_speculative_tokens > 1. hc_head is deferred to compute_logits.
@@ -235,7 +236,7 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
236
hidden_states = hidden_states.view(
237
-1, mtp_layer.hc_mult, mtp_layer.config.hidden_size
238
)
238
- hidden_states = self.hc_head_op(
239
+ hidden_states = mtp_layer.hc_head_op(
240
hidden_states,
241
mtp_layer.hc_head_fn,
242
mtp_layer.hc_head_scale,