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 6, 18:15 UTC
[Feat][Bugfix] Enable additional dimension for Flashinfer MLA and fix routing dtype (#36931)
Signed-off-by: Dimitrios Bariamis <12195802+dbari@users.noreply.github.com> Co-authored-by: Dimitrios Bariamis <12195802+dbari@users.noreply.github.com>
Dimitrios Bariamis committed
Mar 14, 2026 at 00:41 UTC
367cf5cd3eb234e0e191a6d7883bc66f54f42f79
2 files changed
+18
-5
vllm/model_executor/models/deepseek_v2.py
+15
-2
@@ -47,7 +47,11 @@ from vllm.logger import init_logger
47
from vllm.model_executor.layers.activation import SiluAndMul
48
from vllm.model_executor.layers.attention import Attention
49
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
50
-from vllm.model_executor.layers.fused_moe import GateLinear, SharedFusedMoE
50
+from vllm.model_executor.layers.fused_moe import (
51
+ GateLinear,
52
+ RoutingMethodType,
53
+ SharedFusedMoE,
54
+)
55
from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
56
from vllm.model_executor.layers.linear import (
57
ColumnParallelLinear,
@@ -333,8 +337,12 @@ class DeepseekV2MoE(nn.Module):
337
# NOTE(rob): this is a hack until we finish off the PR for
338
# merging TRTLLM kernels into the MK framework. Then we can
339
# query the MonolithicMK for the expected router logits.
340
+ # NOTE(dbari): Use BF16 if routing is not Deepseek, e.g. Mistral Large 3
341
self.gate.set_out_dtype(
337
- torch.float32 if self.experts.quant_method.is_monolithic else torch.bfloat16
342
+ torch.float32
343
+ if self.experts.quant_method.is_monolithic
344
+ and self.experts.routing_method_type == RoutingMethodType.DeepSeekV3
345
+ else torch.bfloat16
346
)
347
348
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -1197,6 +1205,11 @@ class DeepseekV2Model(nn.Module):
1205
if inputs_embeds is not None:
1206
hidden_states = inputs_embeds
1207
else:
1208
+ if input_ids is None:
1209
+ raise ValueError(
1210
+ "Either input_ids or inputs_embeds must be provided "
1211
+ "to DeepseekV2Model.forward"
1212
+ )
1213
hidden_states = self.embed_input_ids(input_ids)
1214
residual = None
1215
else:
vllm/v1/attention/backends/mla/flashinfer_mla.py
+3
-3
@@ -75,16 +75,16 @@ class FlashInferMLABackend(MLACommonBackend):
75
use_sparse: bool,
76
device_capability: DeviceCapability,
77
) -> str | None:
78
- # FlashInfer MLA kernel requires qk_nope_head_dim == 128
78
+ # FlashInfer MLA kernel requires qk_nope_head_dim in [64, 128]
79
from vllm.config import get_current_vllm_config
80
81
vllm_config = get_current_vllm_config()
82
if vllm_config.model_config is not None:
83
hf_text_config = vllm_config.model_config.hf_text_config
84
qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
85
- if qk_nope_head_dim != 128:
85
+ if qk_nope_head_dim not in [64, 128]:
86
return (
87
- f"FlashInfer MLA kernel requires qk_nope_head_dim == 128, "
87
+ f"FlashInfer MLA kernel requires qk_nope_head_dim in [64, 128], "
88
f"but got {qk_nope_head_dim}"
89
)
90
return None