[Performance] Optimize GLM-5.3-Flash Triton KPool and gated norm

tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_glm5next_pool_key_indexer_triton.py

@@ -16,7 +16,7 @@
 OUTPUT_WIDTH = INDEX_TOPK + POOL_SIZE - 1
 
 
-def _assert_selection(result, query, cache, weights, ends, pool_lens, table, positions):
+def _assert_selection(result, query, cache, weights, ends, pool_lens, table, positions, pack_tail=False):
     """CPU reference: score each query head before applying the head weights."""
     result, query, cache, weights = (x.cpu() for x in (result, query, cache, weights))
     start = 0
@@ -33,7 +33,8 @@ def _assert_selection(result, query, cache, weights, ends, pool_lens, table, pos
             expected = torch.full((OUTPUT_WIDTH,), -1, dtype=torch.int32)
             expected[: history.numel()] = history.to(torch.int32)
             tail = torch.arange((pos + 1) // POOL_SIZE * POOL_SIZE, pos + 1, dtype=torch.int32)
-            expected[INDEX_TOPK : INDEX_TOPK + tail.numel()] = tail
+            tail_column = min((pos + 1) // POOL_SIZE * POOL_SIZE, INDEX_TOPK) if pack_tail else INDEX_TOPK
+            expected[tail_column : tail_column + tail.numel()] = tail
             actual = result[row, 0]
             # Top-k ordering may differ for tied scores; token membership,
             # multiplicity, fixed tail columns, and every padding lane must match.
@@ -44,6 +45,8 @@ def _assert_selection(result, query, cache, weights, ends, pool_lens, table, pos
                 atol=0,
             )
             torch.testing.assert_close(actual[INDEX_TOPK:], expected[INDEX_TOPK:], rtol=0, atol=0)
+            if pack_tail:
+                torch.testing.assert_close(actual[tail_column : tail_column + tail.numel()], tail, rtol=0, atol=0)
         start = end
 
 
@@ -53,7 +56,7 @@ def _assert_selection(result, query, cache, weights, ends, pool_lens, table, pos
 def test_pool_selection_real_shape_paging_and_causal_tail(max_pool_seq_len, use_graph, monkeypatch):
     generator = torch.Generator().manual_seed(19)
     # Three requests exercise non-power-of-two bucketization. Two extra rows
-    # model graph padding; the caller is responsible for ignoring their output.
+    # model graph padding, which the fused expansion clears.
     ends = torch.tensor([3, 5, 8], dtype=torch.int32)
     num_tokens = 10
     pages = (max_pool_seq_len + CACHE_BLOCK_SIZE - 1) // CACHE_BLOCK_SIZE
@@ -69,6 +72,11 @@ def test_pool_selection_real_shape_paging_and_causal_tail(max_pool_seq_len, use_
     device_lens, device_positions = pool_lens.npu(), positions.npu()
     args = (query, cache, weights, ends.npu(), device_lens, table.npu(), device_positions)
     kwargs = dict(index_topk=INDEX_TOPK, index_kpool=POOL_SIZE, max_pool_seq_len=max_pool_seq_len)
+    if use_graph:
+        # Exercise the model's packed-tail path and a strided SFA output view.
+        backing_output = torch.full((2 * num_tokens, 2080), 77, dtype=torch.int32, device="npu")
+        output_buffer = backing_output[::2]
+        kwargs.update(output_buffer=output_buffer, pack_tail=True)
     # Force several token chunks without a large scratch allocation. This also
     # checks that request lookup uses the batch-global token offset.
     monkeypatch.setattr(indexer, "TRITON_SCORES_CHUNK_BYTES", max(1, max_pool_seq_len) * 4 * 3)
@@ -93,7 +101,12 @@ def test_pool_selection_real_shape_paging_and_causal_tail(max_pool_seq_len, use_
             result = indexer.glm5_next_lightning_indexer_triton(*args, **kwargs)
         assert result.shape == (num_tokens, 1, OUTPUT_WIDTH)
         assert result.dtype == torch.int32
-        _assert_selection(result, query, cache, weights, ends, pool_lens, table, positions)
+        _assert_selection(result, query, cache, weights, ends, pool_lens, table, positions, pack_tail=use_graph)
+        assert (result[int(ends[-1]) :] == -1).all()
+        if use_graph:
+            assert result.data_ptr() == output_buffer.data_ptr()
+            assert (output_buffer[:, OUTPUT_WIDTH:] == -1).all()
+            assert (backing_output[1::2] == 77).all()
         torch.testing.assert_close(cache.cpu(), backing[:, :CACHE_BLOCK_SIZE], rtol=0, atol=0)
 
 

tests/ut/models/test_glm5next_kpool_model_backend.py

@@ -8,6 +8,7 @@
 import pytest
 import torch
 from torch import nn
+from vllm.config.compilation import CUDAGraphMode
 
 import vllm_ascend.attention.indexer_kpool as backend_module
 from vllm_ascend.attention.indexer_kpool import (
@@ -98,8 +99,8 @@ def _tail_metadata() -> AscendIndexerKPoolTailMetadata:
     )
 
 
-@pytest.mark.parametrize("compute_topk", [False, True])
-def test_triton_indexer_updates_both_caches_and_masks_padding(monkeypatch, compute_topk):
+@pytest.mark.parametrize("compute_topk,allow_cache_packing", [(False, True), (True, True), (True, False)])
+def test_triton_indexer_updates_both_caches_and_masks_padding(monkeypatch, compute_topk, allow_cache_packing):
     metadata = _indexer_metadata(num_tokens=10)
     tail_metadata = _tail_metadata()
     tail_metadata.slot_mapping = torch.arange(10)
@@ -117,6 +118,7 @@ def compress(state, cache, k, gate, ape, positions, query_ends, seq_lens, state_
         cache[0, 0].fill_(11)
 
     select = MagicMock(return_value=torch.full((10, 1, 7), -1, dtype=torch.int32))
+    output_buffer = torch.empty(10, 16, dtype=torch.int32)
     monkeypatch.setattr(kpool_module, "glm5_next_kpool_tail_compress_and_write_cache_triton", compress)
     monkeypatch.setattr(kpool_module, "glm5_next_lightning_indexer_triton", select)
     result = SparseAttnIndexerKpool(4, 2)(
@@ -133,13 +135,18 @@ def compress(state, cache, k, gate, ape, positions, query_ends, seq_lens, state_
         index_kpool=4,
         max_pool_seq_len=1,
         compute_topk=compute_topk,
+        output_buffer=output_buffer,
+        allow_cache_packing=allow_cache_packing,
     )
     torch.testing.assert_close(tail_cache[0, 0], torch.full_like(tail_cache[0, 0], 7))
     torch.testing.assert_close(indexer_cache[0, 0], torch.full_like(indexer_cache[0, 0], 11))
     if compute_topk:
         assert result.shape == (10, 1, 7)
         assert (result[8:] == -1).all()
         select.assert_called_once()
+        assert select.call_args.kwargs["output_buffer"] is output_buffer
+        assert select.call_args.kwargs["pack_tail"] is True
+        assert select.call_args.kwargs["allow_cache_packing"] is allow_cache_packing
     else:
         assert result is None
         select.assert_not_called()
@@ -211,8 +218,10 @@ def forward(self, *args, **kwargs):
         return None
 
 
+@pytest.mark.parametrize("graph_mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL])
 def test_backend_uses_normalized_q_c_and_separate_tail_metadata(
     monkeypatch,
+    graph_mode,
 ) -> None:
     backend = Glm5NextKPoolIndexerBackend.__new__(Glm5NextKPoolIndexerBackend)
     nn.Module.__init__(backend)
@@ -233,7 +242,7 @@ def test_backend_uses_normalized_q_c_and_separate_tail_metadata(
         prefix="indexer.tail",
         kv_cache=torch.zeros(2, 2, 4, 2, dtype=torch.float32),
     )
-    backend.topk_indices_buffer = None
+    backend.topk_indices_buffer = torch.empty(64, 16, dtype=torch.int32)
     backend.softmax_scale = 0.5
     backend._wk_weight_f32 = None
     backend.indexer_op = _RecordingKPool()
@@ -243,20 +252,22 @@ def test_backend_uses_normalized_q_c_and_separate_tail_metadata(
         "get_forward_context",
         lambda: SimpleNamespace(
             attn_metadata={"indexer.tail": tail_metadata},
-            cudagraph_runtime_mode=None,
+            cudagraph_runtime_mode=graph_mode,
             virtual_engine=0,
         ),
     )
 
-    normalized_q_c = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
-    hidden = torch.ones(2, 3)
-    metadata = _indexer_metadata(num_tokens=2)
-    metadata.num_actual_tokens = 2
-    metadata.cum_query_lens = torch.tensor([1, 2], dtype=torch.int32)
-    metadata.raw_seq_lens = torch.tensor([1, 1], dtype=torch.int32)
-    metadata.seq_lens = torch.tensor([0, 0], dtype=torch.int32)
-    metadata.positions = torch.tensor([0, 0])
-    metadata.slot_mapping = torch.tensor([-1, -1])
+    # A 64-row graph bucket does not imply reuse for eight one-token requests.
+    normalized_q_c = torch.arange(128, dtype=torch.float32).reshape(64, 2)
+    hidden = torch.ones(64, 3)
+    metadata = _indexer_metadata()
+    metadata.cum_query_lens = torch.arange(1, 9, dtype=torch.int32)
+    metadata.raw_seq_lens = torch.ones(8, dtype=torch.int32)
+    metadata.seq_lens = torch.zeros(8, dtype=torch.int32)
+    metadata.seq_lens_cpu = metadata.seq_lens.clone()
+    metadata.block_table = torch.zeros(8, 1, dtype=torch.int32)
+    metadata.positions = torch.zeros(64, dtype=torch.int64)
+    metadata.slot_mapping = torch.full((64,), -1, dtype=torch.int64)
 
     result = backend.forward(
         hidden,
@@ -268,18 +279,26 @@ def test_backend_uses_normalized_q_c_and_separate_tail_metadata(
 
     assert result is not None
     assert backend.indexer_op.args is not None
-    torch.testing.assert_close(backend.indexer_op.args[1], normalized_q_c.repeat(1, 2).view(2, 2, 2))
+    expected_rows = 64 if graph_mode == CUDAGraphMode.FULL else 8
+    assert result.shape[0] == expected_rows
+    torch.testing.assert_close(
+        backend.indexer_op.args[1], normalized_q_c[:expected_rows].repeat(1, 2).view(expected_rows, 2, 2)
+    )
     expected_k = torch.nn.functional.layer_norm(
-        torch.nn.functional.linear(hidden + 3, backend.wk_weights_proj.weight)[:, :2],
+        torch.nn.functional.linear(hidden[:expected_rows] + 3, backend.wk_weights_proj.weight)[:, :2],
         (2,),
         backend.k_norm.weight,
         backend.k_norm.bias,
         backend.k_norm.eps,
     )
     torch.testing.assert_close(backend.indexer_op.args[0], expected_k)
     assert backend.indexer_op.args[0].dtype == torch.float32
-    expected_weights = torch.nn.functional.linear(hidden, backend.wk_weights_proj.weight[2:]) * (0.5 * 2**-0.5)
+    expected_weights = torch.nn.functional.linear(hidden[:expected_rows], backend.wk_weights_proj.weight[2:]) * (
+        0.5 * 2**-0.5
+    )
     torch.testing.assert_close(backend.indexer_op.args[2], expected_weights)
     assert backend.indexer_op.args[7] is tail_metadata
     assert backend.indexer_op.kwargs is not None
     assert backend.indexer_op.kwargs["compute_topk"] is True
+    assert backend.indexer_op.kwargs["output_buffer"] is backend.topk_indices_buffer
+    assert backend.indexer_op.kwargs["allow_cache_packing"] is (graph_mode != CUDAGraphMode.FULL)

vllm_ascend/attention/indexer_kpool.py

@@ -519,19 +519,9 @@ def forward(
                 else 0
             ),
             compute_topk=compute_topk,
+            output_buffer=self.topk_indices_buffer,
+            # FULL graphs retain padded token rows and replay the captured
+            # dispatch. Keep paged reads independent of the live query count.
+            allow_cache_packing=context.cudagraph_runtime_mode != CUDAGraphMode.FULL,
         )
-        if result is None or self.topk_indices_buffer is None:
-            return result
-
-        if num_tokens > self.topk_indices_buffer.shape[0]:
-            raise RuntimeError(
-                f"GLM KPool output exceeds the top-k buffer rows: {num_tokens} > {self.topk_indices_buffer.shape[0]}."
-            )
-        output = self.topk_indices_buffer[:num_tokens]
-        output.fill_(-1)
-        if result.shape[-1] > output.shape[-1]:
-            raise RuntimeError(
-                f"GLM KPool output exceeds the top-k buffer width: {result.shape[-1]} > {output.shape[-1]}."
-            )
-        output[:, : result.shape[-1]].copy_(result[:, 0])
         return result

vllm_ascend/models/glm5next/sparse_attn_indexer_kpool.py

@@ -79,6 +79,8 @@ def forward(
         index_kpool: int,
         max_pool_seq_len: int,
         compute_topk: bool,
+        output_buffer: torch.Tensor | None = None,
+        allow_cache_packing: bool = True,
     ) -> torch.Tensor | None:
         num_tokens = k.shape[0]
         if index_kpool <= 0 or self.topk_tokens % index_kpool:
@@ -138,10 +140,9 @@ def forward(
             index_topk=self.topk_tokens,
             index_kpool=index_kpool,
             max_pool_seq_len=max_pool_seq_len,
+            output_buffer=output_buffer,
+            pack_tail=True,
+            allow_cache_packing=allow_cache_packing,
         )
-        # A2/A3 SFA requires a contiguous valid prefix; the reference indexer
-        # puts the running tail at the fixed top-k column for short requests.
-        append_causal_tail(indices[:, 0], positions, self.topk_tokens, index_kpool)
-        valid = torch.arange(num_tokens, device=k.device) < indexer_metadata.cum_query_lens[-1]
-        indices.masked_fill_(~valid[:, None, None], -1)
+        # Expansion also packs the causal tail and clears every padded row.
         return indices

vllm_ascend/ops/triton/docs/glm5_next_kpool_tail_compress.md

@@ -61,20 +61,9 @@ This operator cannot reconstruct an uninitialized historical tail itself.
 
 C=R guarantees retention for monotonic non-MTP prefill/decode: each incomplete
 pool needs at most R-1 earlier rows, and all complete pools in a long query read
-the current input before seeding. MTP target verification writes its committed
+the current input before seeding. During MTP, the target model writes its committed
 first row plus up to `num_speculative_tokens` candidate rows before acceptance
 is known. `C=R+num_speculative_tokens` prevents those candidates from wrapping
 onto the committed open pool. Rejection then needs only the normal logical
 length rollback: replacement tokens overwrite the rejected positions while the
 historical rows required to recompute a boundary pool remain intact.
-
-## Verification
-
-Independent tests should compare pooled outputs with an uncompressed FP32
-reference for every chunk-boundary residue modulo R, long multi-request
-prefills, decode and C>R. Use nonadjacent block IDs and padded block/plane/token
-strides with guard values to detect accidental compact-layout addressing.
-Verify final tail slot contents, incomplete-pool writes, empty batches,
-negative sentinels, positive stale graph padding and repeated graph replay.
-Kernel and model results must be recorded on the target NPU; CPU reference or
-syntax checks alone do not establish device correctness or performance.