biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[BugFix] Allow qk_nope_head_dim=192 in FlashInfer MLA backend checks (#37475)

Browse Source 
Signed-off-by: Kaihang Jiang <kaihangj@nvidia.com>
This commit is contained in:
Kaihang Jiang
2026-03-20 18:14:55 -04:00
committed by GitHub
parent b3d0b37908
commit e5ed6c6c13
2 changed files with 8 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
vllm/v1/attention/backends/mla/flashinfer_mla.py
View File

@@ -77,17 +77,17 @@ class FlashInferMLABackend(MLACommonBackend):
use_sparse: bool,
device_capability: DeviceCapability,
) -> str | None:
# FlashInfer MLA kernel requires qk_nope_head_dim in [64, 128]
# FlashInfer MLA kernel requires qk_nope_head_dim in [64, 128, 192]
from vllm.config import get_current_vllm_config
vllm_config = get_current_vllm_config()
if vllm_config.model_config is not None:
hf_text_config = vllm_config.model_config.hf_text_config
qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
if qk_nope_head_dim not in [64, 128]:
if qk_nope_head_dim not in [64, 128, 192]:
return (
f"FlashInfer MLA kernel requires qk_nope_head_dim in [64, 128], "
f"but got {qk_nope_head_dim}"
"FlashInfer MLA kernel requires qk_nope_head_dim "
f"in [64, 128, 192], but got {qk_nope_head_dim}"
)
return None

8
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
View File

@@ -113,17 +113,17 @@ class FlashInferMLASparseBackend(AttentionBackend):
use_sparse: bool,
device_capability: DeviceCapability,
) -> str | None:
# FlashInfer MLA sparse kernel requires qk_nope_head_dim == 128
# FlashInfer MLA sparse kernel requires qk_nope_head_dim in [128, 192]
from vllm.config import get_current_vllm_config
vllm_config = get_current_vllm_config()
if vllm_config.model_config is not None:
hf_text_config = vllm_config.model_config.hf_text_config
qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
if qk_nope_head_dim != 128:
if qk_nope_head_dim not in [128, 192]:
return (
f"FlashInfer MLA Sparse kernel requires qk_nope_head_dim == 128, "
f"but got {qk_nope_head_dim}"
"FlashInfer MLA Sparse kernel requires qk_nope_head_dim "
f"in [128, 192], but got {qk_nope_head_dim}"
)
# Check for index_topk which indicates sparse model
if not hasattr(hf_text_config, "index_topk"):