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 4, 17:30 UTC
[ROCm] Cast score correction bias tensor during model construction for DeepSeek/Kimi-K2 (#39999)
Signed-off-by: Hemanth Acharya <heachary@amd.com>
Hemanth Acharya committed
Apr 24, 2026 at 05:32 UTC
fa4b70555bc1163c2ee8fa4b193be9957f05dda6
4 files changed
+19
-2
vllm/_aiter_ops.py
+2
@@ -1782,6 +1782,8 @@ class rocm_aiter_ops:
1782
need_renorm: bool,
1783
routed_scaling_factor: float = 1.0,
1784
) -> None:
1785
+ if correction_bias.dtype != gating_output.dtype:
1786
+ correction_bias = correction_bias.to(gating_output.dtype)
1787
torch.ops.vllm.rocm_aiter_biased_grouped_topk(
1788
gating_output,
1789
correction_bias,
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+1
-1
@@ -152,7 +152,7 @@ def rocm_aiter_grouped_topk(
152
if e_score_correction_bias is not None:
153
rocm_aiter_ops.biased_grouped_topk(
154
gating_output,
155
- e_score_correction_bias.to(gating_output.dtype),
155
+ e_score_correction_bias,
156
topk_weights,
157
topk_ids,
158
num_expert_group,
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+1
-1
@@ -136,7 +136,7 @@ def fused_topk_bias(
136
)
137
rocm_aiter_ops.biased_grouped_topk(
138
gating_output,
139
- e_score_correction_bias.to(gating_output.dtype),
139
+ e_score_correction_bias,
140
topk_weights,
141
topk_ids,
142
num_expert_group=num_expert_group,
vllm/model_executor/models/deepseek_v2.py
+15
@@ -349,6 +349,21 @@ class DeepseekV2MoE(nn.Module):
349
else torch.bfloat16
350
)
351
352
+ # Pre-cast the bias to match the gate output dtype so the
353
+ # conversion is not repeated on every forward pass. All
354
+ # downstream references (FusedMoE, router) share the same
355
+ # nn.Parameter object, so mutating .data propagates everywhere.
356
+ # Weight loading uses copy_(), which handles the dtype conversion.
357
+ # Only needed on ROCm where the aiter biased_grouped_topk kernel
358
+ # requires the bias dtype to match the gating output dtype.
359
+ if (
360
+ self.is_rocm_aiter_moe_enabled
361
+ and self.gate.e_score_correction_bias is not None
362
+ ):
363
+ self.gate.e_score_correction_bias.data = (
364
+ self.gate.e_score_correction_bias.data.to(self.gate.out_dtype)
365
+ )
366
+
367
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
368
num_tokens, hidden_dim = hidden_states.shape
369
hidden_states = hidden_states.view(-1, hidden_dim)