[DSv4] Remove sparse-MLA q-head padding for FlashInfer >=0.6.14 (#48047)
Signed-off-by: Jeff Ma <jeffjma@umich.edu> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Jeff (Junze) Ma committed
Jul 30, 2026 at 20:30 UTC
b49eaf205a7546267ef4dd1ec6563ba01d940b5b
1 file changed
+16
-19
vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
+16
-19
@@ -60,6 +60,20 @@ def _packed_block_span(pool: torch.Tensor) -> int:
60
return block_stride // token_stride
61
62
63
+# Sparse MLA h_q counts accepted natively (flashinfer>=0.6.14, #3545).
64
+_SPARSE_MLA_SUPPORTED_Q_HEADS = (8, 16, 32, 64, 128)
65
+
66
+
67
+def _pad_to_supported_q_heads(num_heads: int) -> int:
68
+ for supported in _SPARSE_MLA_SUPPORTED_Q_HEADS:
69
+ if num_heads <= supported:
70
+ return supported
71
+ raise ValueError(
72
+ f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads "
73
+ "(sparse MLA kernel requires h_q in {8, 16, 32, 64, 128})."
74
+ )
75
+
76
+
77
class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLABackend):
78
"""FlashInfer backend using the DSv4 sparse metadata/cache layout.
79
@@ -164,13 +178,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
178
179
@classmethod
180
def get_padded_num_q_heads(cls, num_heads: int) -> int:
167
- # FP8 decode kernel only supports h_q = 64 or 128.
168
- if num_heads > 128:
169
- raise ValueError(
170
- f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads "
171
- "(FP8 decode kernel requires h_q in {64, 128})."
172
- )
173
- return 64 if num_heads <= 64 else 128
181
+ return _pad_to_supported_q_heads(num_heads)
182
183
def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
184
return deep_gemm_fp8_o_proj(
@@ -545,18 +553,7 @@ class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):
553
554
@classmethod
555
def get_padded_num_q_heads(cls, num_heads: int) -> int:
548
- if num_heads <= 16:
549
- return 16
550
- if num_heads <= 32:
551
- return 32
552
- if num_heads <= 64:
553
- return 64
554
- if num_heads <= 128:
555
- return 128
556
- raise ValueError(
557
- f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads "
558
- "(SM120 kernel requires h_q in {16, 32, 64, 128})."
559
- )
556
+ return _pad_to_supported_q_heads(num_heads)
557
558
def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
559
return deep_gemm_fp8_o_proj(