diff --git a/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_glm5next_kpool_triton.py b/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_glm5next_kpool_triton.py new file mode 100644 index 000000000000..5131b04ece3d --- /dev/null +++ b/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_glm5next_kpool_triton.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch_npu # noqa: F401 + +from vllm_ascend.ops.triton.glm5_next_kpool_state_compress import glm5_next_kpool_state_compress_and_write_cache_triton + +HEAD_DIM = 128 +BF16_RTOL = 1e-2 +BF16_ATOL = 1e-2 + + +@pytest.mark.parametrize("pool,capacity", [(4, 4), (4, 16), (8, 8)]) +@pytest.mark.parametrize("use_graph", [False, True]) +@torch.inference_mode() +def test_paged_state_long_prefill_padding_and_rollback(pool, capacity, use_graph): + generator = torch.Generator().manual_seed(13) + dim = HEAD_DIM + # Physical pages deliberately have padding and request IDs are reordered. + state_backing = torch.full((32, capacity + 2, 2 * dim), -7.0, device="npu") + cache_backing = torch.full((3, 18, 1, dim), -7.0, dtype=torch.bfloat16, device="npu") + state, cache = state_backing[:, :capacity], cache_backing[:, :16] + expected_state_backing, expected_cache_backing = state_backing.cpu(), cache_backing.cpu() + expected_state, expected_cache = expected_state_backing[:, :capacity], expected_cache_backing[:, :16] + ape = torch.randn(pool, dim, generator=generator) * 0.1 + state_table = torch.tensor([[5, 9, 3, 12, 7, 16, 2, 20], [11, 8, 1, 18, 4, 22, 15, 23]], dtype=torch.int32) + cache_blocks = [1, 2] + history: list[dict[int, tuple[torch.Tensor, torch.Tensor]]] = [{}, {}] + graph, captured_args = None, None + graph_capacity = 32 + + def run(starts, lengths, invalidate=False): + nonlocal graph, captured_args + positions = torch.cat([torch.arange(s, s + n) for s, n in zip(starts, lengths)]) + num_tokens = positions.numel() + keys = torch.randn(num_tokens, dim, generator=generator) + gates = torch.randn(num_tokens, dim, generator=generator) * 0.1 + ends = torch.tensor(lengths, dtype=torch.int32).cumsum(0).to(torch.int32) + seq_lens = torch.tensor([s + n for s, n in zip(starts, lengths)], dtype=torch.int32) + state_slots, indexer_slots = [], [] + cursor = 0 + for req, (start, length) in enumerate(zip(starts, lengths)): + for local in range(length): + row, pos = cursor + local, start + local + history[req][pos] = (keys[row], gates[row]) + state_slots.append(int(state_table[req, pos // capacity]) * capacity + pos % capacity) + slot = cache_blocks[req] * 16 + pos // pool if (pos + 1) % pool == 0 else -1 + indexer_slots.append(slot) + if slot >= 0: + window = [history[req][p] for p in range(pos - pool + 1, pos + 1)] + pooled = ( + torch.softmax(torch.stack([g for _, g in window]) + ape, dim=0) + * torch.stack([k for k, _ in window]) + ).sum(0) + expected_cache[cache_blocks[req], pos // pool, 0] = pooled.bfloat16() + cursor += length + if invalidate: + state_slots[0] = -1 + state_slots[1] = state.shape[0] * capacity + indexer_slots[0] = cache.shape[0] * cache.shape[1] + # Each logical position has its own scheduler-provided physical slot. + cursor = 0 + for req, (start, length) in enumerate(zip(starts, lengths)): + for local in range(length): + row, pos = cursor + local, start + local + if 0 <= state_slots[row] < state.shape[0] * capacity: + expected_state[state_slots[row] // capacity, pos % capacity] = torch.cat((keys[row], gates[row])) + cursor += length + # Graph-capacity rows have no owning request, even if a stale slot is positive. + padded_keys = torch.nn.functional.pad(keys, (0, 0, 0, graph_capacity - num_tokens)).npu() + padded_gates = torch.nn.functional.pad(gates, (0, 0, 0, graph_capacity - num_tokens)).npu() + args = ( + state, + cache, + padded_keys, + padded_gates, + ape.npu(), + torch.cat((positions, torch.zeros(graph_capacity - num_tokens))).long().npu(), + ends.npu(), + seq_lens.npu(), + torch.tensor(state_slots + [0] + [-1] * (graph_capacity - num_tokens - 1), device="npu"), + state_table.npu(), + torch.tensor(indexer_slots + [0] + [-1] * (graph_capacity - num_tokens - 1), device="npu"), + pool, + ) + if use_graph: + if graph is None: + captured_args = args + glm5_next_kpool_state_compress_and_write_cache_triton(*args) + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph): + glm5_next_kpool_state_compress_and_write_cache_triton(*captured_args) + else: + for target, value in zip(captured_args[2:-1], args[2:-1]): + target.copy_(value) + graph.replay() + else: + glm5_next_kpool_state_compress_and_write_cache_triton(*args) + torch.testing.assert_close(state_backing.cpu(), expected_state_backing, rtol=0, atol=0) + torch.testing.assert_close(cache_backing.cpu(), expected_cache_backing, rtol=BF16_RTOL, atol=BF16_ATOL) + + run([0, 0], [2, 3], invalidate=True) + run([0, 3], [2, 13]) # Rewrite invalid req0 slots; req1 reads history before a long prefill. + run([2, 16], [3, 3]) # Verification writes future candidates into separate paged slots. + run([3, 17], [1, 3]) # Reject candidates and overwrite their positions. + state_table[1, : 16 // capacity] = -1 # Evicted old pages must not be consulted. + run([4, 20], [pool, pool]) + + +@pytest.mark.parametrize("empty", ["tokens", "requests"]) +def test_empty_compression_preserves_caches(empty): + state = torch.ones(1, 4, 2 * HEAD_DIM, device="npu") + cache = torch.ones(1, 16, 1, HEAD_DIM, dtype=torch.bfloat16, device="npu") + count = 0 if empty == "tokens" else 1 + keys = torch.zeros(count, HEAD_DIM, device="npu") + ends = torch.tensor([count] if empty == "tokens" else [], dtype=torch.int32, device="npu") + glm5_next_kpool_state_compress_and_write_cache_triton( + state, + cache, + keys, + keys, + torch.zeros(4, HEAD_DIM, device="npu"), + torch.zeros(count, dtype=torch.int64, device="npu"), + ends, + ends, + torch.full((count,), -1, device="npu"), + torch.zeros(1, 1, dtype=torch.int32, device="npu"), + torch.full((count,), -1, device="npu"), + 4, + ) + torch.testing.assert_close(state.cpu(), torch.ones_like(state, device="cpu"), rtol=0, atol=0) + torch.testing.assert_close(cache.cpu(), torch.ones_like(cache, device="cpu"), rtol=0, atol=0) + + +@pytest.mark.parametrize("block_size,pages", [(2, 1), (4, 0)]) +def test_invalid_state_layout_raises(block_size, pages): + state = torch.zeros(1, block_size, 2 * HEAD_DIM, device="npu") + cache = torch.zeros(1, 16, 1, HEAD_DIM, dtype=torch.bfloat16, device="npu") + keys = torch.zeros(1, HEAD_DIM, device="npu") + ends = torch.ones(1, dtype=torch.int32, device="npu") + slots = torch.zeros(1, dtype=torch.int64, device="npu") + with pytest.raises(ValueError, match="nonempty state page table and block size >= pool size"): + glm5_next_kpool_state_compress_and_write_cache_triton( + state, + cache, + keys, + keys, + torch.zeros(4, HEAD_DIM, device="npu"), + slots, + ends, + ends, + slots, + torch.zeros(1, pages, dtype=torch.int32, device="npu"), + slots, + 4, + ) diff --git a/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_glm5next_pool_key_indexer_triton.py b/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_glm5next_pool_key_indexer_triton.py new file mode 100644 index 000000000000..4c6eef3f980e --- /dev/null +++ b/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_glm5next_pool_key_indexer_triton.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch_npu # noqa: F401 + +from vllm_ascend.ops.triton import glm5_next_lightning_indexer as indexer + +# GLM-5.3-Flash indexer parameters, including its actual top-k width. +HEAD_DIM = 128 +NUM_HEADS = 32 +POOL_SIZE = 4 +INDEX_TOPK = 2048 +CACHE_BLOCK_SIZE = 16 +OUTPUT_WIDTH = INDEX_TOPK + POOL_SIZE - 1 + + +def _assert_selection(result, query, cache, weights, ends, pool_lens, table, positions): + """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 + for req, end in enumerate(ends.tolist()): + for row in range(start, end): + pos = int(positions[row]) + count = min((pos + 1) // POOL_SIZE, int(pool_lens[req])) + ids = torch.arange(count) + keys = cache[table[req, ids // CACHE_BLOCK_SIZE].long(), ids % CACHE_BLOCK_SIZE, 0].float() + per_head_scores = query[row].float() @ keys.T + scores = (per_head_scores * weights[row].float()[:, None]).sum(0) + selected = torch.topk(scores, min(INDEX_TOPK // POOL_SIZE, count)).indices + history = (selected[:, None] * POOL_SIZE + torch.arange(POOL_SIZE)).flatten() + 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 + actual = result[row, 0] + # Top-k ordering may differ for tied scores; token membership, + # multiplicity, fixed tail columns, and every padding lane must match. + torch.testing.assert_close( + actual[:INDEX_TOPK].sort().values, + expected[:INDEX_TOPK].sort().values, + rtol=0, + atol=0, + ) + torch.testing.assert_close(actual[INDEX_TOPK:], expected[INDEX_TOPK:], rtol=0, atol=0) + start = end + + +@pytest.mark.parametrize("max_pool_seq_len", [0, 4, 512, 2050]) +@pytest.mark.parametrize("use_graph", [False, True]) +@torch.inference_mode() +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. + ends = torch.tensor([3, 5, 8], dtype=torch.int32) + num_tokens = 10 + pages = (max_pool_seq_len + CACHE_BLOCK_SIZE - 1) // CACHE_BLOCK_SIZE + num_blocks = max(1, 3 * pages) + table = torch.randperm(num_blocks, generator=generator)[: 3 * pages].reshape(3, pages).to(torch.int32) + backing = torch.randn(num_blocks, CACHE_BLOCK_SIZE + 2, 1, HEAD_DIM, generator=generator).bfloat16() + cache = backing.npu()[:, :CACHE_BLOCK_SIZE] + query = torch.randn(num_tokens, NUM_HEADS, HEAD_DIM, generator=generator).bfloat16().npu() + weights = torch.randn(num_tokens, NUM_HEADS, generator=generator).bfloat16().npu() + pool_lens = torch.tensor([0, min(4, max_pool_seq_len), max_pool_seq_len], dtype=torch.int32) + last_pos = max(2, max_pool_seq_len * POOL_SIZE + 2) + positions = torch.tensor([0, 1, 2, 1, 2, last_pos - 2, last_pos - 1, last_pos, 0, 0]) + 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) + # 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) + + if use_graph: + indexer.glm5_next_lightning_indexer_triton(*args, **kwargs) + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph): + result = indexer.glm5_next_lightning_indexer_triton(*args, **kwargs) + for step in range(2): + if step: + query.copy_(torch.randn(num_tokens, NUM_HEADS, HEAD_DIM, generator=generator).bfloat16().npu()) + weights.copy_(torch.randn(num_tokens, NUM_HEADS, generator=generator).bfloat16().npu()) + pool_lens[2] //= 2 + positions[3:5] = torch.tensor([15, 16]) + device_lens.copy_(pool_lens) + device_positions.copy_(positions) + if use_graph: + graph.replay() + else: + 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) + torch.testing.assert_close(cache.cpu(), backing[:, :CACHE_BLOCK_SIZE], rtol=0, atol=0) + + +def test_empty_query(): + result = indexer.glm5_next_lightning_indexer_triton( + torch.empty(0, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="npu"), + torch.empty(1, CACHE_BLOCK_SIZE, 1, HEAD_DIM, dtype=torch.bfloat16, device="npu"), + torch.empty(0, NUM_HEADS, dtype=torch.bfloat16, device="npu"), + torch.empty(0, dtype=torch.int32, device="npu"), + torch.empty(0, dtype=torch.int32, device="npu"), + torch.empty(0, 0, dtype=torch.int32, device="npu"), + torch.empty(0, dtype=torch.int64, device="npu"), + index_topk=INDEX_TOPK, + index_kpool=POOL_SIZE, + max_pool_seq_len=0, + ) + assert result.shape == (0, 1, OUTPUT_WIDTH) + assert result.dtype == torch.int32 diff --git a/tests/ut/models/test_glm5next_indexer_backend.py b/tests/ut/models/test_glm5next_indexer_backend.py new file mode 100644 index 000000000000..568495c7b4bc --- /dev/null +++ b/tests/ut/models/test_glm5next_indexer_backend.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch +from torch import nn + +from vllm_ascend.ops.mla import IndexerWrapper + + +class _FakeKPoolBackend(nn.Module): + def __init__(self, source, rope_dim): + super().__init__() + del rope_dim + self.k_cache = source.k_cache + self.head_dim = source.head_dim + self.enable_sparse_li_c8 = False + self.num_cache_tensors = 1 + + def process_weights_after_loading(self): + return None + + +class _FakeIndexer(nn.Module): + def __init__(self): + super().__init__() + self.n_head = 2 + self.head_dim = 4 + self.topk_tokens = 8 + self.q_lora_rank = 3 + self.wq_b = nn.Linear(3, 8, bias=False) + self.wk_weights_proj = nn.Linear(5, 6, bias=False) + self.k_norm = nn.LayerNorm(4) + self.softmax_scale = 0.5 + self.index_kpool_compress_ape = nn.Parameter(torch.zeros(4, 4)) + self.index_kpool_compress_gate = nn.Parameter(torch.zeros(4, 5)) + self.k_cache = SimpleNamespace(prefix="indexer.k_cache") + + def get_ascend_indexer_backend_cls(self): + return _FakeKPoolBackend + + +def test_wrapper_selects_model_backend_and_preserves_weight_names() -> None: + wrapper = IndexerWrapper(_FakeIndexer(), qk_rope_head_dim=0) + + assert isinstance(wrapper.impl, _FakeKPoolBackend) + names = set(dict(wrapper.named_parameters())) + assert { + "wq_b.weight", + "wk_weights_proj.weight", + "k_norm.weight", + "k_norm.bias", + "index_kpool_compress_ape", + "index_kpool_compress_gate", + }.issubset(names) + assert not any(name.startswith("impl.") for name in names) diff --git a/vllm_ascend/ops/mla.py b/vllm_ascend/ops/mla.py index 39c1ce8db05a..662d0bbdc7c8 100644 --- a/vllm_ascend/ops/mla.py +++ b/vllm_ascend/ops/mla.py @@ -32,10 +32,7 @@ from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.backend import AttentionMetadata # type: ignore -from vllm_ascend.attention.indexer import ( - AscendSFAIndexerBackend, - AscendSFAIndexerMetadata, -) +from vllm_ascend.attention.indexer import AscendSFAIndexerBackend class IndexerWrapper(nn.Module): @@ -60,7 +57,16 @@ def __init__(self, vllm_indexer: nn.Module, qk_rope_head_dim: int) -> None: self.wk_weights_proj = vllm_indexer.wk_weights_proj self.k_norm = vllm_indexer.k_norm self.softmax_scale = vllm_indexer.softmax_scale - self.impl = AscendSFAIndexerBackend(vllm_indexer, qk_rope_head_dim) + # Preserve checkpoint-visible direct Parameters for every indexer + # family. Registering them here keeps paths at ``...indexer.`` + # rather than adding an implementation segment. + if isinstance(vllm_indexer, nn.Module): + for name, parameter in vllm_indexer.named_parameters(recurse=False): + self.register_parameter(name, parameter) + + backend_factory = getattr(type(vllm_indexer), "get_ascend_indexer_backend_cls", None) + backend_cls = backend_factory(vllm_indexer) if backend_factory is not None else AscendSFAIndexerBackend + self.impl = backend_cls(vllm_indexer, qk_rope_head_dim) # Interface consumed by the SFA impl - delegated to the backend impl. @property @@ -89,7 +95,7 @@ def forward( cos: torch.Tensor, sin: torch.Tensor, k_hidden_states: torch.Tensor, - indexer_metadata: AscendSFAIndexerMetadata, + indexer_metadata: AttentionMetadata, compute_topk: bool = True, ) -> torch.Tensor | None: return self.impl(hidden_states, q_c, cos, sin, k_hidden_states, indexer_metadata, compute_topk) diff --git a/vllm_ascend/ops/triton/docs/glm5_next_kpool_state_compress.md b/vllm_ascend/ops/triton/docs/glm5_next_kpool_state_compress.md new file mode 100644 index 000000000000..ed294865b412 --- /dev/null +++ b/vllm_ascend/ops/triton/docs/glm5_next_kpool_state_compress.md @@ -0,0 +1,57 @@ +# glm5_next_kpool_state_compress + +## Description + +- **Function**: `glm5_next_kpool_state_compress_and_write_cache_triton` compresses completed KeyPool windows into the paged indexer cache, then saves the current keys and gates in the paged compressor state. GLM-5.3-Flash uses FP32 keys, gates, positional bias, and state, with BF16 compressed keys; its pool size is 4 and head dimension is 128. +- **Formula**: For a pool ending at token position `p`, let `K[j, d]` and `G[j, d]` denote the key and gate at position `p - P + 1 + j`. With positional bias `A[j, d]`, compute `W[:, d] = softmax(G[:, d] + A[:, d])` and `C[d] = sum_j(W[j, d] * K[j, d])`. Accumulate in FP32 and cast `C` to the cache dtype on store. State rows contain the concatenation `[K, G]`. +- **Algorithm flow**: + 1. Each program handles one token and a tile of dimensions. Locate its request using cumulative query ends. A valid indexer slot identifies a token completing a pool. + 2. Read window entries in the current query directly from `k` and `gate_score`; read older entries through `state_block_table`. Apply the stable softmax and write the compressed key to `indexer_slot_mapping`. + 3. Launch a second kernel on the same stream to save current keys and gates to `state_slot_mapping`. This ordering ensures compression reads the preceding step's historical state before current state writes. +- **Supported modes**: Eager execution and fixed-shape NPU graph capture/replay on Atlas A2/A3 with Triton-Ascend. See Test Cases for validation scope. Ascend 950: N/A (not validated by this change). + +## Parameters + +`T` is the number of input rows including graph padding, `R` the request count, `D` the head dimension, `P` the pool size, `Bs` the state block size, and `Bi` the indexer block size. All tensor parameters are on the same NPU. + +| Parameter | Input/Output/Attribute | Description | Data type | Data format | +| --- | --- | --- | --- | --- | +| `state_cache` | Input/Output | Historical and updated keys/gates, shape `[Ns, Bs, 2 * D]` | FP32 | ND; block, token, and dimension strides supported | +| `indexer_cache` | Output | Completed compressed keys, shape `[Ni, Bi, 1, D]`; other slots remain unchanged | BF16 | ND; block, token, and dimension strides supported | +| `k` | Input | Current normalized keys, `[T, D]` | FP32 | ND | +| `gate_score` | Input | Current compression gate scores, `[T, D]` | FP32 | ND | +| `ape` | Input | Positional bias within a pool, `[P, D]` | FP32 | ND | +| `positions` | Input | Absolute token positions, `[T]` | int64 | ND | +| `cum_query_lens` | Input | Cumulative exclusive ends of current queries, `[R]`, without a leading zero | int32 | ND | +| `seq_lens` | Input | Total sequence lengths including the current query, `[R]` | int32 | ND | +| `state_slot_mapping` | Input | Flattened physical state slots, `[T]`; `-1` skips a state write | int64 | ND | +| `state_block_table` | Input | Request logical state page to physical block mapping, `[R, Ms]` | int32 | ND | +| `indexer_slot_mapping` | Input | Flattened physical compressed-cache slots, `[T]`; `-1` skips compression output | int64 | ND | +| `index_kpool` | Attribute | Number of keys per pool, `P`; GLM-5.3-Flash uses 4 | Python int | Scalar | +| Return value | Output | N/A; updates both caches in place | None | N/A | + +## Constraints + +- Inference only. The model configuration uses `D = 128`, `P = 4`. Tests also cover `P = 8` as a boundary variant. `P` and `D` must be positive; `Bs >= P`, `Ms > 0`, and both caches must provide storage for every valid slot. An empty page table or a state block smaller than the pool raises `ValueError`. +- Queries are packed in request order. Within each request, positions are contiguous and end at `seq_lens[r] - 1`. Cumulative ends are nondecreasing and the last end is at most `T`. An empty token batch or empty request list returns without modifying caches. +- Only pool-completing rows may have valid indexer slots. All historical positions needed to complete their windows must already be present in the state cache and mapped by the page table. Evicted pages outside these windows may contain invalid entries. +- Valid state slots are in `[0, Ns * Bs)` and valid indexer slots in `[0, Ni * Bi)`. Negative and out-of-range slots are skipped. Rows after the last query end do not write either cache, even if their padded slots contain stale positive values. Valid destination slots must be unique within a launch; duplicate writes are not ordered. +- Cache tensors may have noncontiguous storage as described above. The wrapper makes the other tensor inputs contiguous when needed. Avoid layout conversion in a captured hot path by preparing contiguous inputs beforehand. +- Capture/replay requires stable shapes, strides, pool size, request capacity, and tensor addresses. Metadata and tensor contents can change between replays. Compression and state writes must execute in order on the same stream. + +## Origin and Differences + +- **Origin**: Developed for the GLM-5.3-Flash KeyPool compression and paged state-cache sequence. +- **Differences**: Fuses request lookup, current/history gathering, softmax, and compressed-cache writes in Triton. Uses a separate ordered state-write kernel instead of a chain of host-dispatched indexing operations. Uses the existing scheduler-provided slots and the shared `vllm.utils.math_utils.next_power_of_2` helper. + +## Test Cases + +The single-operator test uses the model's FP32 state/key/gate/bias and BF16 compressed-cache dtypes with `D = 128`, `P = 4`, and state block sizes 4 and 16. A `P = 8` variant checks a wider window. A CPU reference builds each complete window from logical token history independently of device page lookup. + +Cases cover prefill spanning multiple pools, decode/verification updates, rollback, historical tails, reordered physical pages, eviction of unneeded pages, noncontiguous cache storage, negative/out-of-range slots, graph padding with stale positive slots, empty inputs, and invalid state layouts. Eager and graph cases compare the entire backing allocations, including untouched padding. Graph cases update keys, gates, positions, lengths, slots, and page-table contents between replays. + +State copies require bit-exact equality (`rtol = 0`, `atol = 0`). BF16 compressed keys use one fixed tolerance across cases (`rtol = 1e-2`, `atol = 1e-2`). Validated on Atlas A3 with PyTorch 2.10.0, torch-npu 2.10.0.post4, and Triton-Ascend 3.2.0. Atlas A2 and Ascend 950 were not exercised in this validation. These are accuracy tests, not throughput measurements. + +```bash +pytest -sv tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_glm5next_kpool_triton.py +``` diff --git a/vllm_ascend/ops/triton/docs/glm5_next_lightning_indexer.md b/vllm_ascend/ops/triton/docs/glm5_next_lightning_indexer.md new file mode 100644 index 000000000000..9b631cd87abf --- /dev/null +++ b/vllm_ascend/ops/triton/docs/glm5_next_lightning_indexer.md @@ -0,0 +1,56 @@ +# glm5_next_lightning_indexer + +## Description + +- **Function**: `glm5_next_lightning_indexer_triton` selects compressed KeyPools for each query, expands the selected pools into original token indices, and appends the visible incomplete pool as a causal tail. GLM-5.3-Flash uses 32 query heads, head dimension 128, pool size 4, and `index_topk = 2048`, producing 2051 output columns. +- **Formula**: For query token `t`, head `h`, dimension `d`, and pool `j`, `qbar[t, d] = sum_h(weights[t, h] * query[t, h, d])` and `score[t, j] = sum_d(qbar[t, d] * cache[j, d])`. Only pools before `min((positions[t] + 1) // P, indexer_seq_lens[r])` are visible. Select up to `index_topk // P` pools by descending score, then expand pool `j` to `[j * P, ..., j * P + P - 1]`. Tail positions range from `((positions[t] + 1) // P) * P` through `positions[t]`. +- **Algorithm flow**: + 1. Compute the head-weighted query in FP32. Split the token batch into chunks targeting a 256 MiB score-buffer budget. + 2. A Triton kernel maps tokens to requests, gathers paged compressed keys, and writes FP32 scores in pool tiles. Invisible scores remain negative infinity. + 3. Apply `torch.topk`, expand pool IDs, pad missing history with `-1`, and write the incomplete tail at the fixed `index_topk` column. The result contains logical token indices, not physical cache slots. +- **Supported modes**: Eager execution and fixed-shape NPU graph capture/replay on Atlas A2/A3 with Triton-Ascend. See Test Cases for validation scope. Ascend 950: N/A (not validated by this change). + +## Parameters + +`T` includes graph padding, `H` is the query head count, `D` the head dimension, `R` the request count, `B` the compressed-cache block size, and `P = index_kpool`. All tensor parameters are on the same NPU. + +| Parameter | Input/Output/Attribute | Description | Data type | Data format | +| --- | --- | --- | --- | --- | +| `query` | Input | Query vectors, `[T, H, D]`; model shape `[T, 32, 128]` | BF16 | ND | +| `indexer_cache` | Input | Compressed keys, `[N, B, 1, D]` | BF16 | ND; block, token, and dimension strides supported | +| `weights` | Input | Per-head query weights with model scaling already applied, `[T, H]` | BF16 | ND | +| `cum_query_lens` | Input | Cumulative exclusive query ends, `[R]`, without a leading zero | int32 | Contiguous ND | +| `indexer_seq_lens` | Input | Number of available complete pools per request, `[R]`, not raw token lengths | int32 | Contiguous ND | +| `indexer_block_table` | Input | Logical compressed-cache page to physical block mapping, `[R, M]` | int32 | ND; request and page strides supported | +| `positions` | Input | Absolute query token positions, `[T]` | int64 | Contiguous ND | +| `index_topk` | Attribute | Maximum number of history tokens selected through complete pools; model value 2048 | Python int | Scalar, keyword-only | +| `index_kpool` | Attribute | Number of original tokens in each complete pool; model value 4 | Python int | Scalar, keyword-only | +| `max_pool_seq_len` | Attribute | Upper bound on complete pool count and width of the score buffer | Python int | Scalar, keyword-only | +| Return value | Output | Logical token indices, `[T, 1, index_topk + P - 1]`; unused columns are `-1` | int32 | ND | + +## Constraints + +- Inference only. `D` must be a power of two; this path is intended and tested for `D = 128`. `H`, `P`, and `B` must be positive. `index_topk` must be a positive multiple of `P`. +- `max_pool_seq_len >= 0`, `0 <= indexer_seq_lens[r] <= max_pool_seq_len`, and `max_pool_seq_len <= M * B`. If scoring is needed, the cache and request list must be nonempty. The caller must provide valid physical blocks for visible pools; clamping an invalid physical block is not a substitute for valid cache metadata. +- Queries are packed in request order, cumulative ends are nondecreasing, and the last end does not exceed `T`. Positions are nonnegative. Empty queries return shape `[0, 1, index_topk + P - 1]`. A zero maximum pool count returns only the causal tail, with the history region filled with `-1`. +- The first `index_topk` columns hold selected history. Tail tokens always start at column `index_topk`, even when fewer history tokens are available. Thus valid entries need not form a contiguous prefix. Callers requiring a contiguous prefix must compact the result separately. +- Rows beyond the final query end are graph padding and their values are unspecified. The caller must mask or ignore them. Equal-score pools may be returned in any top-k order. +- Capture/replay requires stable shapes, addresses, strides, and scalar attributes (including `max_pool_seq_len`). Query values, weights, positions, pool lengths, and page-table contents may change in the existing buffers. The kernel skips invisible pool sub-tiles at runtime. +- The scratch budget controls token chunking; at least one score row is allocated. Its size is `max_pool_seq_len * sizeof(float32)`, so a single extremely long row can exceed the budget. No constant context-length cutoff is imposed by the wrapper. + +## Origin and Differences + +- **Origin**: Developed for GLM-5.3-Flash pooled-key selection. +- **Differences**: Combines request lookup, paged-cache gathering, and pool scoring in Triton, while retaining the device `torch.topk` operation and tensor operations for index expansion. The head-weighted query is computed once per token chunk. This replaces repeated host-dispatched indexing operations without changing cache allocation or pool-to-token semantics. Reuses `vllm.utils.math_utils.next_power_of_2` for request-capacity rounding. + +## Test Cases + +The test uses the model's actual `[T, 32, 128]` BF16 queries, BF16 head weights/cache, pool size 4, and top-k 2048. Pool capacities 0, 4, 512, and 2050 cover tail-only output, insufficient history, exactly the selection width, and multiple 2048-pool tiles with a partial final tile. Three requests exercise non-power-of-two request counts, distinct lengths, randomized physical page mappings, and noncontiguous cache blocks. + +An independent CPU reference scores each head against the logical keys before weighting and summing the scores. Selected token membership and multiplicity, all history padding, and the fixed tail columns must match exactly (`rtol = 0`, `atol = 0`); history is sorted only for comparison because equal scores do not define a unique ordering. Test inputs use a fixed local generator for reproducibility. + +Both eager and graph cases force several token chunks with a reduced scratch budget. They change queries, head weights, positions, and visible pool lengths between calls/replays, and verify the cache remains unchanged. Output shape/dtype and the empty-query path are checked separately. Validated on Atlas A3 with PyTorch 2.10.0, torch-npu 2.10.0.post4, and Triton-Ascend 3.2.0. Atlas A2 and Ascend 950 were not exercised in this validation. These are accuracy tests, not throughput measurements. + +```bash +pytest -sv tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_glm5next_pool_key_indexer_triton.py +``` diff --git a/vllm_ascend/ops/triton/glm5_next_kpool_state_compress.py b/vllm_ascend/ops/triton/glm5_next_kpool_state_compress.py new file mode 100644 index 000000000000..6698ace1cb64 --- /dev/null +++ b/vllm_ascend/ops/triton/glm5_next_kpool_state_compress.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton fast path for the GLM5 Next indexer pre-compress sequence. + +Two ordered launches per sparse layer perform the following: + +1. For tokens that complete a pool (valid indexer slot), gather the + ``index_kpool`` window of states ending at the token's position. Window + entries covered by the current query chunk are read straight from the + ``k``/``gate_score`` inputs; older entries come from the paged state cache. +2. Compress the window with ``softmax(gate_score + ape)`` over the pool axis + and write the BF16 vector into the paged indexer cache. +3. A second launch writes each valid token to its scheduler-provided state slot. + This ordering keeps long prefills from overwriting historical rows while + another program is still reading them to complete the first pool. + +Doing this in torch lowers to an aclnnIndex/SearchSorted/where small-op flood +per layer, so the sequence stays in Triton. +""" + +from __future__ import annotations + +import torch +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import next_power_of_2 + +TRITON_MAX_BLOCK_D = 128 + + +# Keep batch-varying inputs unspecialized to avoid recompiling per step. +# REQ_POW2 stays constexpr for tl.arange; warm up its power-of-two variants. +@triton.jit(do_not_specialize=["num_reqs", "num_tokens"]) +def _glm5_next_kpool_state_compress_kernel( + state_cache_ptr, + indexer_cache_ptr, + k_ptr, + gate_score_ptr, + ape_ptr, + positions_ptr, + cum_query_lens_ptr, + seq_lens_ptr, + state_slot_mapping_ptr, + state_block_table_ptr, + indexer_slot_mapping_ptr, + num_reqs, + num_tokens, + k_stride_t: tl.constexpr, + gate_score_stride_t: tl.constexpr, + ape_stride_p: tl.constexpr, + state_cache_stride_block: tl.constexpr, + state_cache_stride_offset: tl.constexpr, + state_cache_stride_d: tl.constexpr, + indexer_cache_stride_block: tl.constexpr, + indexer_cache_stride_offset: tl.constexpr, + indexer_cache_stride_d: tl.constexpr, + state_block_table_stride_req: tl.constexpr, + state_block_table_stride_page: tl.constexpr, + state_num_slots: tl.constexpr, + indexer_num_slots: tl.constexpr, + state_num_blocks: tl.constexpr, + state_max_pages: tl.constexpr, + state_block_size: tl.constexpr, + indexer_block_size: tl.constexpr, + REQ_POW2: tl.constexpr, + HEAD_DIM: tl.constexpr, + POOL_SIZE: tl.constexpr, + BLOCK_P: tl.constexpr, + BLOCK_D: tl.constexpr, +): + token_idx = tl.program_id(0) + dim_offsets = tl.program_id(1) * BLOCK_D + tl.arange(0, BLOCK_D) + dim_mask = dim_offsets < HEAD_DIM + + # Request bucketize: requests are packed contiguously in the batch. + req_offsets = tl.arange(0, REQ_POW2) + query_ends = tl.load( + cum_query_lens_ptr + req_offsets, + mask=req_offsets < num_reqs, + other=2147483647, + ) + req_id = tl.sum(tl.where(token_idx >= query_ends, 1, 0)) + # Full ACL graphs keep padded rows beyond the last request; keep their + # pointer arithmetic in bounds even though their stores are masked. + req_id = tl.minimum(req_id, num_reqs - 1) + + # 2) Pool window gather for pool-completing tokens. + indexer_slot = tl.load(indexer_slot_mapping_ptr + token_idx).to(tl.int64) + last_query_end = tl.load(cum_query_lens_ptr + num_reqs - 1) + indexer_valid = (indexer_slot >= 0) & (indexer_slot < indexer_num_slots) & (token_idx < last_query_end) + + pos = tl.load(positions_ptr + token_idx).to(tl.int32) + query_end = tl.load(cum_query_lens_ptr + req_id) + prev_query_end = tl.load(cum_query_lens_ptr + req_id - 1, mask=req_id > 0, other=0) + seq_len = tl.load(seq_lens_ptr + req_id) + request_query_start = seq_len - (query_end - prev_query_end) + + pool_offsets = tl.arange(0, BLOCK_P) + pool_mask = pool_offsets < POOL_SIZE + # Column j holds the state at position pos - (POOL_SIZE - 1 - j). + pool_pos = pos - (POOL_SIZE - 1 - pool_offsets) + eff_pos = tl.maximum(pool_pos, 0) + in_window = pool_mask & (pool_pos >= request_query_start) + + # In-window rows live in this launch's k/gate_score inputs; the matching + # input row is the batch row of the token at that position. + src_row = prev_query_end + eff_pos - request_query_start + src_row = tl.minimum(tl.maximum(src_row, 0), num_tokens - 1) + window_mask = in_window[:, None] & dim_mask[None, :] + pool_k_in = tl.load( + k_ptr + src_row[:, None] * k_stride_t + dim_offsets[None, :], + mask=window_mask, + other=0.0, + ).to(tl.float32) + pool_g_in = tl.load( + gate_score_ptr + src_row[:, None] * gate_score_stride_t + dim_offsets[None, :], + mask=window_mask, + other=0.0, + ).to(tl.float32) + + # Older rows come from the paged state cache (written by earlier steps). + page = eff_pos // state_block_size + history_valid = indexer_valid & pool_mask & (pool_pos >= 0) & (~in_window) & (page < state_max_pages) + page_offset = eff_pos % state_block_size + physical = tl.load( + state_block_table_ptr + req_id * state_block_table_stride_req + page * state_block_table_stride_page, + mask=history_valid, + other=-1, + ).to(tl.int64) + history_valid = history_valid & (physical >= 0) & (physical < state_num_blocks) + physical = tl.where(history_valid, physical, 0) + hist_addr = physical[:, None] * state_cache_stride_block + page_offset[:, None] * state_cache_stride_offset + hist_mask = history_valid[:, None] & dim_mask[None, :] + pool_k_hist = tl.load( + state_cache_ptr + hist_addr + dim_offsets[None, :] * state_cache_stride_d, + mask=hist_mask, + other=0.0, + ).to(tl.float32) + pool_g_hist = tl.load( + state_cache_ptr + hist_addr + (HEAD_DIM + dim_offsets[None, :]) * state_cache_stride_d, + mask=hist_mask, + other=0.0, + ).to(tl.float32) + + pool_k = tl.where(in_window[:, None], pool_k_in, pool_k_hist) + pool_g = tl.where(in_window[:, None], pool_g_in, pool_g_hist) + + # 3) softmax(gate + ape) over the pool axis, weighted sum of K. + ape = tl.load( + ape_ptr + pool_offsets[:, None] * ape_stride_p + dim_offsets[None, :], + mask=pool_mask[:, None] & dim_mask[None, :], + other=0.0, + ).to(tl.float32) + scores = tl.where(pool_mask[:, None], pool_g + ape, float("-inf")) + score_max = tl.max(scores, axis=0) + weights = tl.exp(scores - score_max[None, :]) + weights = weights / tl.sum(weights, axis=0)[None, :] + compressed = tl.sum(weights * pool_k, axis=0) + + safe_indexer_slot = tl.where(indexer_valid, indexer_slot, 0) + indexer_block = safe_indexer_slot // indexer_block_size + indexer_offset = safe_indexer_slot % indexer_block_size + tl.store( + indexer_cache_ptr + + indexer_block * indexer_cache_stride_block + + indexer_offset * indexer_cache_stride_offset + + dim_offsets * indexer_cache_stride_d, + compressed, + mask=dim_mask & indexer_valid, + ) + + +@triton.jit(do_not_specialize=["num_reqs"]) +def _store_kpool_state_kernel( + state, + k, + gate, + query_ends, + state_slots, + num_reqs, + state_stride_b: tl.constexpr, + state_stride_t: tl.constexpr, + state_stride_d: tl.constexpr, + k_stride_t: tl.constexpr, + gate_stride_t: tl.constexpr, + NUM_BLOCKS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, +): + row = tl.program_id(0) + dims = tl.program_id(1) * BLOCK_D + tl.arange(0, BLOCK_D) + end = tl.load(query_ends + num_reqs - 1) + slot = tl.load(state_slots + row).to(tl.int64) + valid = (row < end) & (slot >= 0) & (slot < NUM_BLOCKS * BLOCK_SIZE) + safe_slot = tl.where(valid, slot, 0) + block, offset = safe_slot // BLOCK_SIZE, safe_slot % BLOCK_SIZE + mask = valid & (dims < HEAD_DIM) + key = tl.load(k + row * k_stride_t + dims, mask=mask, other=0) + score = tl.load(gate + row * gate_stride_t + dims, mask=mask, other=0) + addr = state + block * state_stride_b + offset * state_stride_t + tl.store(addr + dims * state_stride_d, key, mask=mask) + tl.store(addr + (HEAD_DIM + dims) * state_stride_d, score, mask=mask) + + +def glm5_next_kpool_state_compress_and_write_cache_triton( + state_cache: torch.Tensor, + indexer_cache: torch.Tensor, + k: torch.Tensor, + gate_score: torch.Tensor, + ape: torch.Tensor, + positions: torch.Tensor, + cum_query_lens: torch.Tensor, + seq_lens: torch.Tensor, + state_slot_mapping: torch.Tensor, + state_block_table: torch.Tensor, + indexer_slot_mapping: torch.Tensor, + index_kpool: int, +) -> None: + """Compress pools, then write states using the original paged slot mapping.""" + num_tokens, head_dim = k.shape + if num_tokens == 0 or cum_query_lens.numel() == 0: + return + if state_block_table.shape[1] == 0 or state_cache.shape[1] < index_kpool: + raise ValueError("KPool requires a nonempty state page table and block size >= pool size.") + + if not k.is_contiguous(): + k = k.contiguous() + if not gate_score.is_contiguous(): + gate_score = gate_score.contiguous() + if not ape.is_contiguous(): + ape = ape.contiguous() + if not positions.is_contiguous(): + positions = positions.contiguous() + if not cum_query_lens.is_contiguous(): + cum_query_lens = cum_query_lens.contiguous() + if not seq_lens.is_contiguous(): + seq_lens = seq_lens.contiguous() + if not state_slot_mapping.is_contiguous(): + state_slot_mapping = state_slot_mapping.contiguous() + if not state_block_table.is_contiguous(): + state_block_table = state_block_table.contiguous() + if not indexer_slot_mapping.is_contiguous(): + indexer_slot_mapping = indexer_slot_mapping.contiguous() + + block_p = next_power_of_2(index_kpool) + block_d = min(next_power_of_2(head_dim), TRITON_MAX_BLOCK_D) + num_reqs = cum_query_lens.shape[0] + _glm5_next_kpool_state_compress_kernel[(num_tokens, triton.cdiv(head_dim, block_d))]( + state_cache, + indexer_cache, + k, + gate_score, + ape, + positions, + cum_query_lens, + seq_lens, + state_slot_mapping, + state_block_table, + indexer_slot_mapping, + num_reqs, + num_tokens, + k.stride(0), + gate_score.stride(0), + ape.stride(0), + state_cache.stride(0), + state_cache.stride(1), + state_cache.stride(2), + indexer_cache.stride(0), + indexer_cache.stride(1), + indexer_cache.stride(3), + state_block_table.stride(0), + state_block_table.stride(1), + state_cache.shape[0] * state_cache.shape[1], + indexer_cache.shape[0] * indexer_cache.shape[1], + state_cache.shape[0], + state_block_table.shape[1], + state_cache.shape[1], + indexer_cache.shape[1], + next_power_of_2(max(1, num_reqs)), + head_dim, + index_kpool, + block_p, + block_d, + ) + _store_kpool_state_kernel[(num_tokens, triton.cdiv(head_dim, block_d))]( + state_cache, + k, + gate_score, + cum_query_lens, + state_slot_mapping, + num_reqs, + state_cache.stride(0), + state_cache.stride(1), + state_cache.stride(2), + k.stride(0), + gate_score.stride(0), + state_cache.shape[0], + state_cache.shape[1], + head_dim, + block_d, + ) diff --git a/vllm_ascend/ops/triton/glm5_next_lightning_indexer.py b/vllm_ascend/ops/triton/glm5_next_lightning_indexer.py new file mode 100644 index 000000000000..7acee89fc68f --- /dev/null +++ b/vllm_ascend/ops/triton/glm5_next_lightning_indexer.py @@ -0,0 +1,234 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton fast path for the narrow GLM5 Next KPool lightning indexer. + +The heavy part is scoring every compressed pool against the token's +head-weighted query: a paged-cache gather plus a 128-dim matvec. Doing that in +torch lowers to an aclnnIndex/SearchSorted small-op flood, so it stays in one +triton kernel that writes the raw pool scores to a scratch buffer. The top-k +selection itself is a single fused aclnn ``topk`` call, followed by a few +element-wise ops for pool->token expansion and causal tail append. + +There is no hard length limit: the kernel tiles pools dynamically, and the +wrapper chunks the token dimension so the fp32 scores scratch stays under +``TRITON_SCORES_CHUNK_BYTES`` even for long-context prefill (e.g. 128K input). + +Note: greedy/sort-based in-kernel top-k is deliberately avoided — reductions +over wide vectors scalarize on the Ascend triton backend, which is slower than +the torch fallback by two orders of magnitude. +""" + +from __future__ import annotations + +import torch +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import next_power_of_2 + +# Fix the pool tile width to avoid recompiling as sequence lengths grow. +TRITON_POOL_CHUNK_SIZE = 2048 +# Sub-tile of pools whose K rows are loaded as one coalesced 2D block. +# 128 pools x 128 dims x fp32 = 64KB, keeping the UB footprint small. +TRITON_POOL_SUB_TILE_SIZE = 128 +# Chunk tokens to limit the FP32 score buffer to this budget where possible. +TRITON_SCORES_CHUNK_BYTES = 256 * 1024 * 1024 + + +# Keep batch-varying inputs unspecialized to avoid recompiling per step. +# REQ_POW2 stays constexpr for tl.arange; warm up its power-of-two variants. +@triton.jit(do_not_specialize=["token_offset", "max_pool_seq_len", "num_reqs", "num_cache_blocks"]) +def _glm5_next_lightning_indexer_score_kernel( + qbar_ptr, + indexer_cache_ptr, + cum_query_lens_ptr, + indexer_seq_lens_ptr, + indexer_block_table_ptr, + positions_ptr, + scores_ptr, + token_offset, + max_pool_seq_len, + num_reqs, + num_cache_blocks, + cache_stride_block: tl.constexpr, + cache_stride_offset: tl.constexpr, + cache_stride_d: tl.constexpr, + block_table_stride_req: tl.constexpr, + block_table_stride_page: tl.constexpr, + pool_block_size: tl.constexpr, + REQ_POW2: tl.constexpr, + HEAD_DIM: tl.constexpr, + INDEX_KPOOL: tl.constexpr, + BLOCK_POOL: tl.constexpr, + SUB_POOL: tl.constexpr, +): + # qbar/scores rows are chunk-local; positions/query-boundary lookups use + # the batch-global token index. + local_token_idx = tl.program_id(0) + token_idx = local_token_idx + token_offset + chunk = tl.program_id(1) + + req_offsets = tl.arange(0, REQ_POW2) + query_ends = tl.load(cum_query_lens_ptr + req_offsets, mask=req_offsets < num_reqs, other=2147483647) + req_id = tl.sum(tl.where(token_idx >= query_ends, 1, 0)) + # Full ACL graphs keep padded rows beyond the last request; keep their + # pointer arithmetic in bounds even though their outputs are unused. + req_id = tl.minimum(req_id, num_reqs - 1) + + pos = tl.load(positions_ptr + token_idx).to(tl.int32) + request_pool_len = tl.load(indexer_seq_lens_ptr + req_id).to(tl.int32) + causal_pool_len = (pos + 1) // INDEX_KPOOL + visible_pool_len = tl.minimum(causal_pool_len, request_pool_len) + + dim_offsets = tl.arange(0, HEAD_DIM) + qbar = tl.load(qbar_ptr + local_token_idx * HEAD_DIM + dim_offsets) + + chunk_start = chunk * BLOCK_POOL + # Dynamic trip count: requests shorter than the static max pool length + # skip their out-of-range sub-tiles even inside captured graphs. Cells + # beyond ``visible_pool_len`` keep the -inf the wrapper initialized. + chunk_visible = tl.maximum(tl.minimum(visible_pool_len, chunk_start + BLOCK_POOL) - chunk_start, 0) + num_subs = tl.cdiv(chunk_visible, SUB_POOL) + for sub in tl.range(num_subs): + pool_offsets = chunk_start + sub * SUB_POOL + tl.arange(0, SUB_POOL) + in_range = pool_offsets < max_pool_seq_len + valid_pool = in_range & (pool_offsets < visible_pool_len) + logical_pages = pool_offsets // pool_block_size + page_offsets = pool_offsets % pool_block_size + physical_blocks = tl.load( + indexer_block_table_ptr + req_id * block_table_stride_req + logical_pages * block_table_stride_page, + mask=in_range, + other=0, + ).to(tl.int64) + # Clamp both sides: padded/stale block-table entries must never form + # an out-of-range cache address, even though their loads are masked. + physical_blocks = tl.minimum(tl.maximum(physical_blocks, 0), num_cache_blocks - 1) + k_addrs = ( + physical_blocks[:, None] * cache_stride_block + + page_offsets[:, None] * cache_stride_offset + + dim_offsets[None, :] * cache_stride_d + ) + k_tile = tl.load(indexer_cache_ptr + k_addrs, mask=valid_pool[:, None], other=0.0).to(tl.float32) + scores = tl.sum(k_tile * qbar[None, :], axis=1) + scores = tl.where(valid_pool, scores, float("-inf")) + tl.store(scores_ptr + local_token_idx * max_pool_seq_len + pool_offsets, scores, mask=in_range) + + +def glm5_next_lightning_indexer_triton( + query: torch.Tensor, + indexer_cache: torch.Tensor, + weights: torch.Tensor, + cum_query_lens: torch.Tensor, + indexer_seq_lens: torch.Tensor, + indexer_block_table: torch.Tensor, + positions: torch.Tensor, + *, + index_topk: int, + index_kpool: int, + max_pool_seq_len: int, +) -> torch.Tensor: + pool_topk = index_topk // index_kpool + output_width = index_topk + index_kpool - 1 + num_tokens = query.shape[0] + if num_tokens == 0: + return torch.empty( + (0, 1, output_width), + dtype=torch.int32, + device=query.device, + ) + + output = torch.empty( + (num_tokens, 1, output_width), + dtype=torch.int32, + device=query.device, + ) + if max_pool_seq_len == 0: + output.fill_(-1) + tail_offsets = torch.arange(index_kpool - 1, device=query.device) + tail_start = (positions + 1) // index_kpool * index_kpool + tail = torch.where( + tail_offsets[None, :] < (positions + 1 - tail_start)[:, None], + tail_start[:, None] + tail_offsets[None, :], + -1, + ) + output[:, 0, index_topk:] = tail.to(torch.int32) + return output + block_pool = TRITON_POOL_CHUNK_SIZE + num_chunks = (max_pool_seq_len + block_pool - 1) // block_pool + topk = min(pool_topk, max_pool_seq_len) + token_offsets = torch.arange(index_kpool, device=query.device) + tail_offsets = torch.arange(index_kpool - 1, device=query.device) + + # Chunk the token dimension so the fp32 scores scratch stays bounded; + # long-context prefill would otherwise need num_tokens x max_pool_seq_len + # x 4 bytes (over 1GB at 128K context with a full prefill batch). + token_chunk = max(1, TRITON_SCORES_CHUNK_BYTES // (max_pool_seq_len * 4)) + for token_start in range(0, num_tokens, token_chunk): + token_end = min(token_start + token_chunk, num_tokens) + rows = token_end - token_start + # Head-weighted query, computed once here instead of per chunk program. + qbar = ( + (query[token_start:token_end].float() * weights[token_start:token_end].float().unsqueeze(-1)) + .sum(dim=1) + .contiguous() + ) + # -inf init: the kernel skips sub-tiles beyond a request's visible pools, + # and those cells must stay excluded from the top-k. + scores = torch.full( + (rows, max_pool_seq_len), + float("-inf"), + dtype=torch.float32, + device=query.device, + ) + _glm5_next_lightning_indexer_score_kernel[(rows, num_chunks)]( + qbar, + indexer_cache, + cum_query_lens, + indexer_seq_lens, + indexer_block_table, + positions, + scores, + token_start, + max_pool_seq_len, + cum_query_lens.shape[0], + indexer_cache.shape[0], + indexer_cache.stride(0), + indexer_cache.stride(1), + indexer_cache.stride(3), + indexer_block_table.stride(0), + indexer_block_table.stride(1), + indexer_cache.shape[1], + next_power_of_2(max(1, cum_query_lens.shape[0])), + query.shape[2], + index_kpool, + block_pool, + TRITON_POOL_SUB_TILE_SIZE, + ) + + topk_vals, pool_ids = torch.topk(scores, topk, dim=1) + pool_ids = torch.where( + topk_vals == float("-inf"), + torch.full_like(pool_ids, -1), + pool_ids, + ) + history = pool_ids.unsqueeze(-1) * index_kpool + token_offsets + history = torch.where( + pool_ids.unsqueeze(-1) >= 0, + history, + torch.full_like(history, -1), + ).reshape(rows, topk * index_kpool) + if topk < pool_topk: + history = torch.nn.functional.pad( + history, + (0, (pool_topk - topk) * index_kpool), + value=-1, + ) + + pos = positions[token_start:token_end] + tail_start = (pos + 1) // index_kpool * index_kpool + tail_count = pos + 1 - tail_start + tail = torch.where( + tail_offsets[None, :] < tail_count[:, None], + tail_start[:, None] + tail_offsets[None, :], + -1, + ) + output[token_start:token_end, 0] = torch.cat([history, tail], dim=1).to(torch.int32) + return output