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, 03:45 UTC
[ROCm][DeepSeek] Enable V3.2 TP4 AITER MLA (#41835)
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
akii96 committed
May 7, 2026 at 23:10 UTC
c936548ce6b0f6ef2f71968315c13213898ecd31
2 files changed
+12
-10
vllm/model_executor/models/deepseek_v2.py
+11
-9
@@ -299,6 +299,15 @@ class DeepseekV2MoE(nn.Module):
299
self.is_fusion_moe_shared_experts_enabled = (
300
rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
301
)
302
+ if (
303
+ self.is_rocm_aiter_moe_enabled
304
+ and self.gate.e_score_correction_bias is not None
305
+ ):
306
+ # AITER biased_grouped_topk requires the correction bias dtype to
307
+ # match the router logits. Keep DeepSeek's correction bias in fp32
308
+ # by requesting fp32 router logits for this routing path.
309
+ self.gate.set_out_dtype(torch.float32)
310
+
311
if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled:
312
self.shared_experts = None
313
else:
@@ -338,22 +347,15 @@ class DeepseekV2MoE(nn.Module):
347
n_shared_experts=config.n_shared_experts
348
if self.is_fusion_moe_shared_experts_enabled
349
else None,
350
+ router_logits_dtype=self.gate.out_dtype,
351
)
352
343
- # Pre-cast the bias to match the gate output dtype so the
344
- # conversion is not repeated on every forward pass. All
345
- # downstream references (FusedMoE, router) share the same
346
- # nn.Parameter object, so mutating .data propagates everywhere.
347
- # Weight loading uses copy_(), which handles the dtype conversion.
348
- # Only needed on ROCm where the aiter biased_grouped_topk kernel
349
- # requires the bias dtype to match the gating output dtype.
353
if (
354
self.is_rocm_aiter_moe_enabled
355
and self.gate.e_score_correction_bias is not None
356
):
354
- gate_out_dtype = self.gate.out_dtype or self.gate.weight.dtype
357
self.gate.e_score_correction_bias.data = (
356
- self.gate.e_score_correction_bias.data.to(gate_out_dtype)
358
+ self.gate.e_score_correction_bias.data.to(self.gate.out_dtype)
359
)
360
361
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+1
-1
@@ -396,7 +396,7 @@ class AiterMLAHelper:
396
"""
397
398
_AITER_MIN_MLA_HEADS: Final = 16
399
- _AITER_UNSUPPORTED_HEADS = [32]
399
+ _AITER_UNSUPPORTED_HEADS: ClassVar[tuple[int, ...]] = ()
400
401
@staticmethod
402
def check_num_heads_validity(num_heads: int):