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, 07:00 UTC
[BugFix] Allow qk_nope_head_dim=192 in FlashInfer MLA backend checks (#37475)
Signed-off-by: Kaihang Jiang <kaihangj@nvidia.com>
Kaihang Jiang committed
Mar 20, 2026 at 18:14 UTC
e5ed6c6c134ffac42fde5299943dbe3af1821be2
2 files changed
+8
-8
vllm/v1/attention/backends/mla/flashinfer_mla.py
+4
-4
@@ -77,17 +77,17 @@ class FlashInferMLABackend(MLACommonBackend):
77
use_sparse: bool,
78
device_capability: DeviceCapability,
79
) -> str | None:
80
- # FlashInfer MLA kernel requires qk_nope_head_dim in [64, 128]
80
+ # FlashInfer MLA kernel requires qk_nope_head_dim in [64, 128, 192]
81
from vllm.config import get_current_vllm_config
82
83
vllm_config = get_current_vllm_config()
84
if vllm_config.model_config is not None:
85
hf_text_config = vllm_config.model_config.hf_text_config
86
qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
87
- if qk_nope_head_dim not in [64, 128]:
87
+ if qk_nope_head_dim not in [64, 128, 192]:
88
return (
89
- f"FlashInfer MLA kernel requires qk_nope_head_dim in [64, 128], "
90
- f"but got {qk_nope_head_dim}"
89
+ "FlashInfer MLA kernel requires qk_nope_head_dim "
90
+ f"in [64, 128, 192], but got {qk_nope_head_dim}"
91
)
92
return None
93
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+4
-4
@@ -113,17 +113,17 @@ class FlashInferMLASparseBackend(AttentionBackend):
113
use_sparse: bool,
114
device_capability: DeviceCapability,
115
) -> str | None:
116
- # FlashInfer MLA sparse kernel requires qk_nope_head_dim == 128
116
+ # FlashInfer MLA sparse kernel requires qk_nope_head_dim in [128, 192]
117
from vllm.config import get_current_vllm_config
118
119
vllm_config = get_current_vllm_config()
120
if vllm_config.model_config is not None:
121
hf_text_config = vllm_config.model_config.hf_text_config
122
qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
123
- if qk_nope_head_dim != 128:
123
+ if qk_nope_head_dim not in [128, 192]:
124
return (
125
- f"FlashInfer MLA Sparse kernel requires qk_nope_head_dim == 128, "
126
- f"but got {qk_nope_head_dim}"
125
+ "FlashInfer MLA Sparse kernel requires qk_nope_head_dim "
126
+ f"in [128, 192], but got {qk_nope_head_dim}"
127
)
128
# Check for index_topk which indicates sparse model
129
if not hasattr(hf_text_config, "index_topk"):