@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a369a7becc

[Perf][Attention] Reduce GLM sparse MLA preparation overhead (#57458)

Signed-off-by: Summer Yang <girasoleyang@gmail.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>

Summer Yang committed Sep 21, 2026 at 17:24 UTC a369a7becc20c6c4dda601aedef51b8cc8387971
7 files changed +248 -94
tests/kernels/attention/test_flashinfer_mla_decode.py
+103 -49
@@ -130,69 +130,123 @@ def test_flashinfer_mla_decode(dtype: torch.dtype, bs: int, block_size: int):
130 torch.testing.assert_close(out_ans, out_ref, atol=1e-2, rtol=1e-2)
131
132
133 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])
134 +@pytest.mark.parametrize("mode", ["tp", "dcp", "hisparse"])
135 @requires_sm10x
134 -def test_flashinfer_trtllm_sparse_mla_decode_without_rope():
135 - """The native sparse MLA path supports a zero-width rotary tail."""
136 +def test_flashinfer_trtllm_sparse_mla_decode_without_rope(dtype, mode, monkeypatch):
137 + """TP skips preparation; empty DCP contributions and HiSparse outputs vanish."""
138 + import vllm.v1.attention.backends.mla.flashinfer_mla_sparse as backend
139 + from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
140 +
141 torch.set_default_device("cuda")
142 torch.manual_seed(42)
138 -
139 - batch_size = 2
140 - block_size = 64
141 - num_blocks = 4
142 - sparse_topk = 128
143 - valid_lens = torch.tensor([17, 73], dtype=torch.int32)
144 -
145 - query = torch.randn(
146 - batch_size,
147 - 1,
148 - NUM_HEADS,
149 - KV_LORA_RANK,
150 - dtype=torch.bfloat16,
151 - )
152 - kv_cache = torch.randn(
153 - num_blocks,
154 - block_size,
155 - KV_LORA_RANK,
156 - dtype=torch.bfloat16,
157 - )
158 -
143 + batch_size, num_heads, block_size, num_blocks, sparse_topk = 3, 8, 64, 4, 2176
144 + valid_lens = [1 if mode == "tp" else 0, 17, 73]
145 + counts = torch.tensor(valid_lens, dtype=torch.int32)
146 + query = torch.randn(batch_size, 1, num_heads, KV_LORA_RANK).to(dtype)
147 + kv_cache = torch.randn(num_blocks, block_size, KV_LORA_RANK).to(dtype)
148 num_slots = num_blocks * block_size
160 - slot_tables = torch.stack(
161 - [torch.randperm(num_slots)[:sparse_topk] for _ in range(batch_size)]
162 - ).to(torch.int32)
163 - for row, valid_len in zip(slot_tables, valid_lens.tolist()):
164 - row[valid_len:] = -1
165 -
166 - workspace_buffer = torch.empty(
167 - FLASHINFER_WORKSPACE_BUFFER_SIZE,
168 - dtype=torch.int8,
149 + indices = torch.full((batch_size, sparse_topk), -1, dtype=torch.int32)
150 + for token, length in enumerate(valid_lens):
151 + indices[token, :length] = torch.randperm(num_slots)[:length].int()
152 + empty = counts == 0
153 + reference_indices = indices.clone()
154 + reference_indices[:, 0].masked_fill_(empty, 0)
155 + workspace = torch.empty(FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.int8)
156 +
157 + impl = object.__new__(backend.FlashInferMLASparseImpl)
158 + impl.qk_nope_head_dim = QK_NOPE_HEAD_DIM
159 + impl.kv_lora_rank = KV_LORA_RANK
160 + impl.qk_rope_head_dim = 0
161 + impl.index_group = (
162 + object.__new__(HiSparseMLAIndexGroup) if mode == "hisparse" else None
163 )
170 - out = trtllm_batch_decode_with_kv_cache_mla(
164 + impl._workspace_buffer = workspace
165 + impl.bmm1_scale = QK_NOPE_HEAD_DIM**-0.5
166 + impl.bmm2_scale = 1.0
167 + impl.is_nope_mla = True
168 + impl.need_to_return_lse_for_decode = mode == "dcp"
169 + if mode == "tp":
170 +
171 + def unexpected_helper(*args, **kwargs):
172 + pytest.fail("TP must not launch empty-query preparation")
173 +
174 + monkeypatch.setattr(
175 + backend, "prepare_sparse_mla_safe_lengths", unexpected_helper
176 + )
177 +
178 + out, lse = impl._run_mqa_kernel(query.squeeze(1), kv_cache, indices, counts)
179 + reference = trtllm_batch_decode_with_kv_cache_mla(
180 query=query,
181 kv_cache=kv_cache.unsqueeze(1),
173 - workspace_buffer=workspace_buffer,
182 + workspace_buffer=workspace,
183 qk_nope_head_dim=QK_NOPE_HEAD_DIM,
184 kv_lora_rank=KV_LORA_RANK,
185 qk_rope_head_dim=0,
177 - block_tables=slot_tables.unsqueeze(1),
178 - seq_lens=valid_lens,
186 + block_tables=reference_indices.unsqueeze(1),
187 + seq_lens=counts,
188 max_seq_len=sparse_topk,
189 sparse_mla_top_k=sparse_topk,
181 - sparse_mla_top_k_lens=valid_lens,
190 + sparse_mla_top_k_lens=counts.clamp(min=1),
191 bmm1_scale=QK_NOPE_HEAD_DIM**-0.5,
192 bmm2_scale=1.0,
184 - ).squeeze(1)
185 -
186 - flat_cache = kv_cache.view(num_slots, KV_LORA_RANK).float()
187 - refs = []
188 - for batch_idx, valid_len in enumerate(valid_lens.tolist()):
189 - selected_kv = flat_cache[slot_tables[batch_idx, :valid_len].long()]
190 - scores = torch.einsum("hd,kd->hk", query[batch_idx, 0].float(), selected_kv)
191 - probs = torch.softmax(scores * QK_NOPE_HEAD_DIM**-0.5, dim=-1)
192 - refs.append(torch.einsum("hk,kd->hd", probs, selected_kv))
193 - ref = torch.stack(refs).to(torch.bfloat16)
194 -
195 - torch.testing.assert_close(out, ref, atol=2e-2, rtol=2e-2)
193 + return_lse=mode == "dcp",
194 + )
195 + if mode == "dcp":
196 + reference, reference_lse = reference
197 + reference_lse = reference_lse.reshape(batch_size, num_heads)
198 + reference_lse.masked_fill_(empty[:, None], float("-inf"))
199 + torch.testing.assert_close(lse, reference_lse, atol=0, rtol=0)
200 +
201 + from vllm.v1.attention.ops.dcp import CPTritonContext, correct_attn_out
202 +
203 + # The peer owns only our empty selections. Check our contribution after
204 + # the production merge, without zeroing the local attention output.
205 + peer_lse = torch.full_like(lse, float("-inf")).masked_fill(empty[:, None], 0)
206 + out, _ = correct_attn_out(
207 + out, torch.stack((lse, peer_lse)), cp_rank=0, ctx=CPTritonContext()
208 + )
209 + else:
210 + assert lse is None
211 + reference = reference.squeeze(1).masked_fill(empty[:, None, None], 0)
212 + torch.testing.assert_close(out, reference, atol=0, rtol=0)
213 +
214 + if dtype == torch.bfloat16:
215 + flat_cache = kv_cache.view(num_slots, KV_LORA_RANK).float()
216 + refs = []
217 + for token, length in enumerate(valid_lens):
218 + selected_kv = flat_cache[reference_indices[token, :length].long()]
219 + scores = torch.einsum("hd,kd->hk", query[token, 0].float(), selected_kv)
220 + probs = torch.softmax(scores * QK_NOPE_HEAD_DIM**-0.5, dim=-1)
221 + refs.append(torch.einsum("hk,kd->hd", probs, selected_kv))
222 + ref = torch.stack(refs).to(torch.bfloat16)
223 + torch.testing.assert_close(out, ref, atol=2e-2, rtol=2e-2)
224 +
225 +
226 +@requires_sm10x
227 +def test_sparse_mla_safe_lengths():
228 + from vllm.v1.attention.backends.mla.sparse_utils import (
229 + prepare_sparse_mla_safe_lengths,
230 + )
231 +
232 + device = "cuda"
233 + num_tokens = 259
234 + # Include a partial tile and expose accidental writes into adjacent storage.
235 + counts_storage = torch.full((num_tokens, 2), 17, dtype=torch.int32, device=device)
236 + counts = counts_storage[:, 0]
237 + counts[::3] = 0
238 + original_counts = counts_storage.clone()
239 + empty = counts == 0
240 + indices_storage = torch.full(
241 + (num_tokens, 256), 42, dtype=torch.int32, device=device
242 + )
243 + indices = indices_storage[:, ::2]
244 + expected_indices = indices_storage.clone()
245 + expected_indices[:, 0].masked_fill_(empty, 0)
246 + safe = prepare_sparse_mla_safe_lengths(indices, counts)
247 + torch.testing.assert_close(safe, counts.clamp(min=1), rtol=0, atol=0)
248 + torch.testing.assert_close(indices_storage, expected_indices, rtol=0, atol=0)
249 + torch.testing.assert_close(counts_storage, original_counts, rtol=0, atol=0)
250
251
252 @requires_sm90
tests/kernels/test_concat_mla_q.py
+29
@@ -148,3 +148,32 @@ def test_concat_mla_q_values_preserved(num_tokens):
148 assert torch.equal(out_bits[..., :nope_dim], ql_nope_bits)
149
150 assert torch.equal(out_bits[..., nope_dim:], q_pe_bits)
151 +
152 +
153 +@pytest.mark.parametrize("rope_dim,transposed", [(0, False), (0, True), (64, False)])
154 +@pytest.mark.parametrize("enabled", [False, True])
155 +def test_concat_mla_q_fp8_nope_and_rope(rope_dim, transposed, enabled):
156 + """FP8 query packing preserves values with empty RoPE and strided BMM output."""
157 + from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
158 + from vllm.model_executor.layers.attention.mla_attention import _DecodeConcatQuantFP8
159 + from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
160 +
161 + torch.manual_seed(42)
162 + shape = (8, 17, 512) if transposed else (17, 8, 512)
163 + ql = torch.randn(shape, dtype=torch.bfloat16, device="cuda") * 128
164 + if transposed:
165 + ql = ql.transpose(0, 1)
166 + qp = torch.randn(17, 8, rope_dim, dtype=torch.bfloat16, device="cuda") * 128
167 + scale = torch.tensor(0.3, device="cuda", dtype=torch.float32)
168 + joined = torch.cat((ql, qp), dim=-1)
169 + expected = (joined.float() * scale.reciprocal()).clamp(-448, 448)
170 + expected = expected.to(torch.float8_e4m3fn)
171 + custom_ops = ["none"] + (["+mla_decode_concat_quant_fp8"] if enabled else [])
172 + config = VllmConfig(
173 + compilation_config=CompilationConfig(mode=0, custom_ops=custom_ops)
174 + )
175 + with set_current_vllm_config(config):
176 + quant = _DecodeConcatQuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
177 + actual = quant(ql, qp, scale)
178 + torch.testing.assert_close(actual.float(), expected.float(), atol=0, rtol=0)
179 + assert actual.is_contiguous()
tests/v1/attention/test_indexer_dcp_localize.py
+3 -4
@@ -700,9 +700,8 @@ def test_dcp_filter_compacts_valid_slots_for_sparse_kernel(
700 @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
701 @pytest.mark.parametrize("interleave", [1, 2])
702 @pytest.mark.parametrize("dcp_rank", [0, 1])
703 -# 384 is not a power of two, so it exercises the multi-tile atomic allocator
704 -# rather than the single-tile path 1024 takes.
705 -@pytest.mark.parametrize("num_topk", [1024, 384])
703 +# Include padded single-tile compaction and the multi-tile atomic fallback.
704 +@pytest.mark.parametrize("num_topk", [1024, 2176, 4224])
705 def test_dcp_filter_compaction_matches_reference(
706 interleave: int, dcp_rank: int, num_topk: int
707 ):
@@ -729,7 +728,7 @@ def test_dcp_filter_compaction_matches_reference(
728 for r in range(num_rows):
729 # Ids are distinct, so a row holds at most `seq` valid entries.
730 hi = min(num_topk * 3 // 4, seq)
732 - n_valid = int(torch.randint(num_topk // 4, hi, (1,)).item())
731 + n_valid = int(torch.randint(hi // 3, hi, (1,)).item())
732 perm = torch.randperm(seq, device=device)[:n_valid].to(torch.int32)
733 token_indices[r, :n_valid] = perm
734
tests/v1/attention/test_sparse_mla_backends.py
+18 -7
@@ -1663,10 +1663,12 @@ def test_split_indexer_prefill_chunks_single_request_overflow():
1663 assert out == expected
1664
1665
1666 -# 384 is not a power of two, so it counts via the tiled atomic accumulation
1667 -# rather than the single-tile path 128 takes.
1668 -@pytest.mark.parametrize("num_topk_tokens", [128, 384])
1669 -def test_triton_convert_returns_valid_counts(num_topk_tokens: int):
1666 +# Power-of-two, GLM's padded tile, and atomic fallback, with reused buffers.
1667 +@pytest.mark.parametrize(
1668 + "num_topk_tokens,reuse_buffers",
1669 + [(128, False), (2176, False), (2176, True), (4224, False), (4224, True)],
1670 +)
1671 +def test_triton_convert_returns_valid_counts(num_topk_tokens: int, reuse_buffers: bool):
1672 """Test that return_valid_counts correctly counts non-negative indices."""
1673 device = torch.device(DEVICE_TYPE)
1674 num_tokens = 8
@@ -1679,8 +1681,7 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int):
1681 num_requests * max_blocks_per_req, dtype=torch.int32, device=device
1682 ).view(num_requests, max_blocks_per_req)
1683
1682 - # Create token indices with varying numbers of valid entries: half the row,
1683 - # a quarter of it, the whole row, then a single valid entry -- twice over.
1684 + # Include partly valid, fully valid, and empty selections for each request.
1685 token_indices = torch.full(
1686 (num_tokens, num_topk_tokens), -1, dtype=torch.int32, device=device
1687 )
@@ -1688,7 +1689,7 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int):
1689 num_topk_tokens // 2,
1690 num_topk_tokens // 4,
1691 num_topk_tokens,
1691 - 1,
1692 + 0,
1693 ] * 2
1694 expected_valid = []
1695 for i in range(num_tokens):
@@ -1702,6 +1703,15 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int):
1703 expected_valid, dtype=torch.int32, device=device
1704 )
1705
1706 + buffers = {}
1707 + if reuse_buffers:
1708 + buffers = dict(
1709 + out=torch.full_like(token_indices, -99),
1710 + valid_counts_out=torch.full(
1711 + (num_tokens,), 99, dtype=torch.int32, device=device
1712 + ),
1713 + )
1714 +
1715 # Test with return_valid_counts=True
1716 result, valid_counts = triton_convert_req_index_to_global_index(
1717 req_id,
@@ -1710,6 +1720,7 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int):
1720 BLOCK_SIZE=block_size,
1721 NUM_TOPK_TOKENS=num_topk_tokens,
1722 return_valid_counts=True,
1723 + **buffers,
1724 )
1725
1726 torch.testing.assert_close(valid_counts, expected_valid_tensor, rtol=0, atol=0)
vllm/model_executor/layers/attention/mla_attention.py
+4 -1
@@ -1553,7 +1553,10 @@ class _DecodeConcatQuantFP8(QuantFP8):
1553 scale: torch.Tensor,
1554 scale_ub: torch.Tensor | None = None,
1555 ) -> torch.Tensor:
1556 - decode_q0 = torch.cat((decode_ql_nope, decode_q_pe), dim=-1)
1556 + if decode_q_pe.shape[-1] == 0 and decode_ql_nope.is_contiguous():
1557 + decode_q0 = decode_ql_nope
1558 + else:
1559 + decode_q0 = torch.cat((decode_ql_nope, decode_q_pe), dim=-1)
1560 decode_q_flat = decode_q0.reshape(decode_q0.shape[0], -1)
1561 decode_q, _ = quant_fn(self, decode_q_flat, scale, scale_ub)
1562 return decode_q.view(decode_q0.shape)
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+19 -15
@@ -30,6 +30,7 @@ from vllm.v1.attention.backend import (
30 from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
31 from vllm.v1.attention.backends.mla.sparse_utils import (
32 flat_kv_row_view,
33 + prepare_sparse_mla_safe_lengths,
34 triton_convert_req_index_to_global_index,
35 triton_filter_and_convert_dcp_index,
36 )
@@ -697,18 +698,16 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
698 sparse_topk_capacity = topk_indices.shape[1]
699
700 extra_kwargs: dict[str, torch.Tensor] = {}
700 - empty_rows: torch.Tensor | None = None
701 + needs_empty_query_guard = self.need_to_return_lse_for_decode or isinstance(
702 + self.index_group, HiSparseMLAIndexGroup
703 + )
704 if self.is_nope_mla:
702 - # The native no-rope kernel takes the active top-k length per query
703 - # token (``seq_lens`` here is already the compacted per-token valid
704 - # count, int32) and rejects zero-length rows. Point empty rows at a
705 - # single valid dummy slot with length 1 and zero their output after
706 - # the launch. ``triton_convert_req_index_to_global_index`` packs the
707 - # valid indices into a contiguous prefix, which is what the kernel
708 - # requires of the page table.
709 - empty_rows = seq_lens == 0
710 - topk_indices[:, 0] = topk_indices[:, 0].masked_fill(empty_rows, 0)
711 - extra_kwargs["sparse_mla_top_k_lens"] = seq_lens.clamp(min=1)
705 + # Resident TP queries have nonempty selections. Preserve empty-query
706 + # handling for DCP's local selections and host-backed HiSparse.
707 + topk_lens = seq_lens
708 + if needs_empty_query_guard:
709 + topk_lens = prepare_sparse_mla_safe_lengths(topk_indices, seq_lens)
710 + extra_kwargs["sparse_mla_top_k_lens"] = topk_lens
711
712 kernel_out = trtllm_batch_decode_with_kv_cache_mla(
713 query=query,
@@ -737,12 +736,17 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
736 out = o.view(-1, o.shape[-2], o.shape[-1])
737 if lse is not None:
738 lse = self._normalize_lse(lse, out.shape[0], out.shape[1])
740 - if empty_rows is None and lse is not None:
739 + if self.is_nope_mla and needs_empty_query_guard:
740 + empty_queries = seq_lens == 0
741 + if lse is not None:
742 + # DCP combine already suppresses outputs with zero LSE weight.
743 + lse.masked_fill_(empty_queries[:, None], float("-inf"))
744 + else:
745 + out.masked_fill_(empty_queries[:, None, None], 0.0)
746 + elif lse is not None:
747 empty_rows = (topk_indices == -1).all(dim=-1)
742 - if empty_rows is not None:
748 out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0)
744 - if lse is not None:
745 - lse.masked_fill_(empty_rows.view(-1, 1), float("-inf"))
749 + lse.masked_fill_(empty_rows.view(-1, 1), float("-inf"))
750 return out, lse
751
752 @staticmethod
vllm/v1/attention/backends/mla/sparse_utils.py
+72 -18
@@ -86,7 +86,7 @@ class ConvertReqIndexToGlobalIndexKernel(
86
87 @staticmethod
88 @triton.jit(do_not_specialize=["max_num_blocks_per_req", "workspace_rank_stride"])
89 - def kernel(
89 + def sparse_mla_index_remap_kernel(
90 req_id_ptr, # int32 [num_tokens]
91 block_table_ptr, # int32 [num_requests, max_num_blocks_per_req]
92 token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]
@@ -100,15 +100,16 @@ class ConvertReqIndexToGlobalIndexKernel(
100 BLOCK_SIZE: tl.constexpr,
101 BLOCK_STRIDE_ROWS: tl.constexpr,
102 BLOCK_N: tl.constexpr, # tile width along columns
103 + NUM_TOPK_TOKENS: tl.constexpr,
104 HAS_PREFILL: tl.constexpr,
105 COUNT_VALID: tl.constexpr, # whether to count valid indices
105 - # BLOCK_N == NUM_TOPK_TOKENS: one program owns the row, so the valid count
106 + # BLOCK_N >= NUM_TOPK_TOKENS: one program owns the query, so its count
107 # is an in-register reduction and needs no atomic.
108 SINGLE_TILE: tl.constexpr,
109 # When set, scatter valid slots to a contiguous prefix [0, valid_count) using
110 # valid_count_ptr as an atomic slot allocator (DCP filtering leaves interior
111 # -1 gaps; the trtllm-gen sparse kernel reads the first valid_count entries).
111 - # Requires COUNT_VALID and an out buffer pre-filled with -1. Order within the
112 + # Requires COUNT_VALID; multi-tile output must be pre-filled with -1. The
113 # prefix is unspecified (only the selected set matters).
114 COMPACT_TO_FRONT: tl.constexpr,
115 # DCP de-interleave: with DCP_SIZE == 1 these are an exact no-op
@@ -136,7 +137,7 @@ class ConvertReqIndexToGlobalIndexKernel(
137
138 # Load token indices for this tile
139 ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1
139 - tok = tl.load(ti_ptr) # int32
140 + tok = tl.load(ti_ptr, mask=indice_id < NUM_TOPK_TOKENS, other=-1)
141
142 # Only token == -1 should propagate as -1
143 is_invalid_tok = tok < 0
@@ -197,10 +198,19 @@ class ConvertReqIndexToGlobalIndexKernel(
198 dest = base + local_offset
199 out_ptr_dest = out_ptr + token_id * out_stride0 + dest * out_stride1
200 tl.store(out_ptr_dest, out_val, mask=is_valid == 1)
201 + if SINGLE_TILE:
202 + # This program owns both the compacted prefix and its padding.
203 + # The stores are disjoint, so no separate initialization is needed.
204 + tl.store(
205 + out_ptr + token_id * out_stride0 + indice_id * out_stride1,
206 + -1,
207 + mask=(indice_id >= tile_valid_count)
208 + & (indice_id < NUM_TOPK_TOKENS),
209 + )
210 else:
211 # Store results in place (input column == output column).
212 out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1
203 - tl.store(out_ptr_ij, out_val)
213 + tl.store(out_ptr_ij, out_val, mask=indice_id < NUM_TOPK_TOKENS)
214
215 # Accumulate the tile's valid count into the row total; a single tile's
216 # reduction *is* the total.
@@ -211,6 +221,9 @@ class ConvertReqIndexToGlobalIndexKernel(
221 else:
222 tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
223
224 + # Keep the warmup wrapper interface while exposing a descriptive CUDA symbol.
225 + kernel = sparse_mla_index_remap_kernel
226 +
227 def dispatch( # type: ignore[override]
228 self,
229 *,
@@ -421,18 +434,19 @@ def _remap_tiling(
434 Counting the valid slots per row is the only reason the column tiles have to
435 talk to each other, so when counting give one program the whole row: the
436 count becomes an in-register reduction plus a plain store, needing neither
424 - atomics nor a zero-initialized counter. The row is one ``tl.arange``, so this
425 - needs a power-of-two width; other top-k sizes stay tiled and atomic.
437 + atomics nor a zero-initialized counter. Pad modest non-power-of-two widths
438 + (including GLM's 2176 entries) to one tile; larger widths stay tiled and atomic.
439
440 Returns:
441 (single_tile, block_n, tiles_per_row, num_warps)
442
443 """
431 - single_tile = (
432 - count_valid and triton.next_power_of_2(NUM_TOPK_TOKENS) == NUM_TOPK_TOKENS
444 + padded_width = triton.next_power_of_2(NUM_TOPK_TOKENS)
445 + single_tile = count_valid and (
446 + padded_width == NUM_TOPK_TOKENS or padded_width <= 4096
447 )
448 if single_tile:
435 - return True, NUM_TOPK_TOKENS, 1, 8
449 + return True, padded_width, 1, 8
450 return False, BLOCK_N, NUM_TOPK_TOKENS // BLOCK_N, 4
451
452
@@ -504,14 +518,11 @@ def triton_convert_req_index_to_global_index(
518 req_id_c = req_id.contiguous()
519 block_table_c = block_table.contiguous()
520 token_indices_c = token_indices.contiguous()
507 - # When return_valid_counts, the kernel scatters valid entries to a
508 - # contiguous prefix [0, valid_count) and leaves the tail unwritten, so
509 - # pre-fill -1 there. flash_mla_sparse_fwd then bounds attention to
510 - # [:topk_length] == exactly the valid set (no dropped tokens).
521 + # Only multi-tile compaction needs separate padding initialization.
522 if out is None:
523 out = (
524 torch.full_like(token_indices_c, -1)
514 - if return_valid_counts
525 + if return_valid_counts and not single_tile
526 else torch.empty_like(token_indices_c)
527 )
528 else:
@@ -519,7 +530,7 @@ def triton_convert_req_index_to_global_index(
530 assert out.device == token_indices_c.device
531 assert out.shape == token_indices_c.shape
532 assert out.is_contiguous()
522 - if return_valid_counts:
533 + if return_valid_counts and not single_tile:
534 out.fill_(-1)
535
536 valid_counts: torch.Tensor | None = None
@@ -636,13 +647,13 @@ def triton_filter_and_convert_dcp_index(
647 token_indices_c = token_indices.contiguous()
648
649 # The compaction uses the valid-count buffer as a slot allocator, so it
639 - # requires counting. Pre-fill out with -1 so the unwritten tail stays -1.
650 + # requires counting. Only the multi-tile path needs pre-filled padding.
651 count_valid = return_valid_counts or compact_valid_to_front
652
653 # The compaction builds on the counting, so it shares the tiling.
654 single_tile, _, _, _ = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, count_valid)
655
645 - if compact_valid_to_front:
656 + if compact_valid_to_front and not single_tile:
657 out = torch.full_like(token_indices_c, -1)
658 else:
659 out = torch.empty_like(token_indices_c)
@@ -685,3 +696,46 @@ def triton_filter_and_convert_dcp_index(
696
697
698 _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL = ConvertReqIndexToGlobalIndexKernel()
699 +
700 +
701 +@triton.jit(do_not_specialize=["num_tokens"])
702 +def sparse_mla_prepare_safe_lengths_kernel(
703 + indices_ptr,
704 + counts_ptr,
705 + safe_lengths_ptr,
706 + num_tokens,
707 + index_stride0: tl.constexpr,
708 + count_stride: tl.constexpr,
709 + BLOCK: tl.constexpr,
710 +):
711 + token = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
712 + count = tl.load(counts_ptr + token * count_stride, token < num_tokens, other=1)
713 + tl.store(safe_lengths_ptr + token, tl.maximum(count, 1), token < num_tokens)
714 + # TRTLLM's NoPE sparse kernel requires at least one valid KV-cache entry.
715 + tl.store(
716 + indices_ptr + token * index_stride0, 0, (token < num_tokens) & (count == 0)
717 + )
718 +
719 +
720 +def prepare_sparse_mla_safe_lengths(
721 + physical_indices: torch.Tensor, valid_counts: torch.Tensor
722 +) -> torch.Tensor:
723 + """Install dummy slots for empty queries and return nonzero kernel lengths.
724 +
725 + Preserve the raw counts so empty outputs and LSE can be neutralized later.
726 + """
727 + num_tokens = valid_counts.numel()
728 + safe_lengths = torch.empty(
729 + (num_tokens,), dtype=valid_counts.dtype, device=valid_counts.device
730 + )
731 + if num_tokens:
732 + sparse_mla_prepare_safe_lengths_kernel[(triton.cdiv(num_tokens, 256),)](
733 + physical_indices,
734 + valid_counts,
735 + safe_lengths,
736 + num_tokens,
737 + physical_indices.stride(0),
738 + valid_counts.stride(0),
739 + BLOCK=256,
740 + )
741 + return safe_lengths