[Perf][GLM5.2] Fuse sparse MLA Q concatenation with head padding (#53878)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Chauncey committed
Aug 27, 2026 at 23:54 UTC
d21c5b50a1d3db5950ab905fcf78b713c81e15e1
3 files changed
+37
-6
tests/kernels/test_concat_mla_q.py
+11
@@ -34,6 +34,17 @@ def test_concat_mla_q_contiguous(num_tokens, num_heads, nope_dim, rope_dim, dtyp
34
torch.testing.assert_close(q_out, ref, atol=0, rtol=0)
35
36
37
+def test_concat_mla_q_padded_output_heads():
38
+ """The fused concat can write directly into a head-padded workspace."""
39
+ ql_nope = torch.randn(4, 8, 512, dtype=torch.bfloat16, device="cuda")
40
+ q_pe = torch.randn(4, 8, 64, dtype=torch.bfloat16, device="cuda")
41
+ q_out = torch.empty(4, 64, 576, dtype=torch.bfloat16, device="cuda")
42
+
43
+ ops.concat_mla_q(ql_nope, q_pe, q_out)
44
+
45
+ torch.testing.assert_close(q_out[:, :8], torch.cat((ql_nope, q_pe), dim=-1))
46
+
47
+
48
@pytest.mark.parametrize("num_tokens", [t for t in NUM_TOKENS if t > 1])
49
@pytest.mark.parametrize("num_heads", NUM_HEADS)
50
@pytest.mark.parametrize("nope_dim", NOPE_DIM)
tests/v1/attention/test_sparse_mla_backends.py
+4
-2
@@ -824,18 +824,20 @@ def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens():
824
825
captured = {}
826
827
- def _stub_kernel(q, kv, indices, lengths):
827
+ def _stub_kernel(q, kv, indices, lengths, actual_num_heads):
828
captured["indices"] = indices
829
+ captured["actual_num_heads"] = actual_num_heads
830
return torch.zeros(q.shape[0], q.shape[1], 512, dtype=q.dtype, device=q.device)
831
832
stub_impl = SimpleNamespace(_bf16_flash_mla_kernel=_stub_kernel)
833
834
out = FlashMLASparseImpl._forward_bf16_kv(
834
- stub_impl, q, kv_cache, topk_indices, attn_metadata
835
+ stub_impl, q, kv_cache, topk_indices, attn_metadata, q.shape[1]
836
)
837
838
assert out.shape[0] == num_mqa_tokens
839
assert captured["indices"].shape[0] == num_mqa_tokens
840
+ assert captured["actual_num_heads"] == q.shape[1]
841
reference = _triton_convert_reference_impl(
842
attn_metadata.req_id_per_token[:num_mqa_tokens],
843
attn_metadata.block_table,
vllm/v1/attention/backends/mla/flashmla_sparse.py
+22
-4
@@ -587,7 +587,14 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
587
588
vllm_config = get_current_vllm_config()
589
max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
590
- q_concat_shape = (max_tokens, num_heads, head_size)
590
+ q_concat_heads = num_heads
591
+ if not is_quantized_kv_cache(kv_cache_dtype):
592
+ q_concat_heads = (
593
+ (num_heads + self.prefill_padding - 1)
594
+ // self.prefill_padding
595
+ * self.prefill_padding
596
+ )
597
+ q_concat_shape = (max_tokens, q_concat_heads, head_size)
598
if is_quantized_kv_cache(kv_cache_dtype):
599
assert kv_cache_dtype == "fp8_ds_mla", (
600
"FlashMLA Sparse Attention backend fp8 only supports "
@@ -626,6 +633,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
633
kv_c_and_k_pe_cache: torch.Tensor,
634
topk_indices: torch.Tensor,
635
attn_metadata: FlashMLASparseMetadata,
636
+ actual_num_heads: int,
637
) -> tuple[torch.Tensor, torch.Tensor]:
638
# Convert per-request indices to global slots (decode) or workspace
639
# offsets (prefill). req_id_per_token covers the whole batch; slice it
@@ -648,6 +656,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
656
kv_rows,
657
topk_indices,
658
topk_length,
659
+ actual_num_heads,
660
)
661
662
def _forward_fp8_kv_separate_prefill_decode(
@@ -880,6 +889,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
889
kv_c_and_k_pe_cache: torch.Tensor,
890
topk_indices: torch.Tensor,
891
topk_length: torch.Tensor | None = None,
892
+ actual_num_heads: int | None = None,
893
) -> tuple[torch.Tensor, torch.Tensor]:
894
num_tokens = q.shape[0]
895
kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(
@@ -889,13 +899,14 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
899
# NOTE(Chen): kernel requires num_local_head to be a multiple of
900
# 64 on hopper and 128 on blackwell. Pad from q's head count, not
901
# self.num_heads: under DCP the heads are all-gathered before this.
892
- actual_num_heads = q.shape[1]
902
+ if actual_num_heads is None:
903
+ actual_num_heads = q.shape[1]
904
padded_num_heads = (
905
(actual_num_heads + self.prefill_padding - 1)
906
// self.prefill_padding
907
* self.prefill_padding
908
)
898
- if actual_num_heads < padded_num_heads:
909
+ if q.shape[1] < padded_num_heads:
910
logger.warning_once(
911
f"Padding num_heads from {actual_num_heads} to "
912
f"{padded_num_heads} for BF16 sparse prefill kernel"
@@ -928,10 +939,13 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
939
# MQA 576/512 approach for both prefill and decode
940
941
# Concatenate q if it's a tuple (ql_nope, q_pe)
942
+ actual_num_heads = self.num_heads
943
if isinstance(q, tuple):
944
ql_nope, q_pe = q
945
q = self.q_concat_buffer[: ql_nope.shape[0]]
946
ops.concat_mla_q(ql_nope, q_pe, q)
947
+ else:
948
+ actual_num_heads = q.shape[1]
949
950
num_actual_toks = q.shape[0]
951
@@ -945,7 +959,11 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
959
960
if not use_fp8_cache:
961
attn_out, bf16_lse = self._forward_bf16_kv(
948
- q, kv_c_and_k_pe_cache, topk_indices, attn_metadata
962
+ q,
963
+ kv_c_and_k_pe_cache,
964
+ topk_indices,
965
+ attn_metadata,
966
+ actual_num_heads,
967
)
968
if self.need_to_return_lse_for_decode:
969
lse = bf16_lse