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 9, 18:15 UTC
fix: Add glm4_moe_lite to MLA detection (#32614)
Signed-off-by: marksverdhei <marksverdhei@hotmail.com> Signed-off-by: Markus / Mark <46672778+marksverdhei@users.noreply.github.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Markus / Mark committed
Jan 23, 2026 at 21:38 UTC
586a57ad7ede5c5132a70d021fc48fe05b10f5c3
4 files changed
+70
-10
vllm/model_executor/layers/attention/mla_attention.py
+23
-4
@@ -440,18 +440,34 @@ M = TypeVar("M", bound=MLACommonMetadata)
440
A = TypeVar("A", bound=AttentionMetadata)
441
442
443
+def is_deepseek_r1_mla_compatible(vllm_config: VllmConfig) -> bool:
444
+ # Check if model has DeepSeek R1 compatible MLA dimensions:
445
+ # qk_nope_head_dim = 128, qk_rope_head_dim = 64, v_head_dim = 128
446
+ # which results in query/key head dim = 192.
447
+ if vllm_config.model_config is None:
448
+ return False
449
+ hf_text_config = vllm_config.model_config.hf_text_config
450
+ qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
451
+ qk_rope_head_dim = getattr(hf_text_config, "qk_rope_head_dim", 1)
452
+ v_head_dim = getattr(hf_text_config, "v_head_dim", 1)
453
+ return qk_nope_head_dim == 128 and qk_rope_head_dim == 64 and v_head_dim == 128
454
+
455
+
456
def use_flashinfer_prefill() -> bool:
457
# For blackwell default to flashinfer prefill if it's available since
458
# it is faster than FA2.
459
from vllm.config import get_current_vllm_config
460
461
vllm_config = get_current_vllm_config()
449
- return (
462
+ if not (
463
not vllm_config.attention_config.disable_flashinfer_prefill
464
and flashinfer_available
465
and not vllm_config.attention_config.use_cudnn_prefill
466
and current_platform.is_device_capability_family(100)
454
- )
467
+ ):
468
+ return False
469
+
470
+ return is_deepseek_r1_mla_compatible(vllm_config)
471
472
473
def use_cudnn_prefill() -> bool:
@@ -471,11 +487,14 @@ def use_trtllm_ragged_deepseek_prefill() -> bool:
487
from vllm.config import get_current_vllm_config
488
489
vllm_config = get_current_vllm_config()
474
- return (
490
+ if not (
491
flashinfer_available
492
and vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill
493
and current_platform.is_device_capability_family(100)
478
- )
494
+ ):
495
+ return False
496
+
497
+ return is_deepseek_r1_mla_compatible(vllm_config)
498
499
500
@dataclass
vllm/platforms/cuda.py
+19
-6
@@ -180,19 +180,34 @@ class CudaPlatformBase(Platform):
180
use_cutlass_mla = False
181
use_flashinfer_mla = False
182
183
+ from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
184
+
185
if vllm_config.attention_config.backend is None:
186
# Default case
185
- if cls.is_device_capability_family(100) and not use_sparse:
186
- # Blackwell => Force FlashInfer MLA (unless sparse, i.e. DSv3.2).
187
+ hf_text_config = model_config.hf_text_config
188
+ qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
189
+ if (
190
+ cls.is_device_capability_family(100)
191
+ and not use_sparse
192
+ and qk_nope_head_dim == 128
193
+ ):
194
+ # Blackwell => Force FlashInfer MLA (unless sparse, i.e. DSv3.2)
195
+ # and only if qk_nope_head_dim == 128 (kernel constraint)
196
use_flashinfer_mla = True
197
# Set the backend in AttentionConfig so it's used during
198
# backend selection
199
vllm_config.attention_config.backend = (
200
AttentionBackendEnum.FLASHINFER_MLA
201
)
193
- else:
194
- # Not Blackwell
202
+ elif cls.is_device_capability_family(100) and not use_sparse:
203
+ # Fall back to CUTLASS_MLA as 2nd priority on Blackwell
204
+ use_cutlass_mla = True
205
+ elif is_flashmla_dense_supported()[0]:
206
+ # Non-Blackwell with FlashMLA support
207
use_flashmla = True
208
+ else:
209
+ # Fallback: will use Triton MLA or other compatible backend
210
+ pass
211
else:
212
# Forced case
213
backend = vllm_config.attention_config.backend
@@ -200,8 +215,6 @@ class CudaPlatformBase(Platform):
215
use_cutlass_mla = backend == AttentionBackendEnum.CUTLASS_MLA
216
use_flashinfer_mla = backend == AttentionBackendEnum.FLASHINFER_MLA
217
203
- from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
204
-
218
if (
219
use_flashmla
220
and is_flashmla_dense_supported()[0]
vllm/transformers_utils/model_arch_config_convertor.py
+2
@@ -189,6 +189,8 @@ class ModelArchConfigConvertorBase:
189
"deepseek_v3",
190
"deepseek_v32",
191
"deepseek_mtp",
192
+ "glm4_moe_lite",
193
+ "glm4_moe_lite_mtp",
194
"kimi_k2",
195
"kimi_linear",
196
"longcat_flash",
vllm/v1/attention/backends/mla/flashinfer_mla.py
+26
@@ -63,6 +63,32 @@ class FlashInferMLABackend(MLACommonBackend):
63
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
64
return capability.major == 10
65
66
+ @classmethod
67
+ def supports_combination(
68
+ cls,
69
+ head_size: int,
70
+ dtype: torch.dtype,
71
+ kv_cache_dtype: CacheDType | None,
72
+ block_size: int,
73
+ use_mla: bool,
74
+ has_sink: bool,
75
+ use_sparse: bool,
76
+ device_capability: DeviceCapability,
77
+ ) -> str | None:
78
+ # FlashInfer MLA kernel requires qk_nope_head_dim == 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:
86
+ return (
87
+ f"FlashInfer MLA kernel requires qk_nope_head_dim == 128, "
88
+ f"but got {qk_nope_head_dim}"
89
+ )
90
+ return None
91
+
92
@classmethod
93
def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
94
return "HND"