| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """CPU tests for the decode top-k row->token mapping (no GPU required). |
| 4 | |
| 5 | On non-uniform decode batches (``requires_padding`` -- mixed plain-decode and |
| 6 | spec-verify requests, or variable MTP verify lens; taken on Hopper where the |
| 7 | varlen/flatten logits path is unavailable), the pool-topk rows follow the |
| 8 | PADDED ``[batch_size, next_n]`` grid: row ``(b, t)`` is flat decode token |
| 9 | ``offset_b + t``. The former inline ``dec_seq = positions[:n] + 1`` |
| 10 | (``n = batch_size * next_n``) indexes the flat per-token layout with padded |
| 11 | coordinates, so rows after the first non-uniform request read another |
| 12 | request's positions and rows past the decode region read prefill tokens -- |
| 13 | ``expand_pools_and_append_tail`` then anchors the tail at a foreign length. |
| 14 | |
| 15 | These tests pin the defect with the exact production arithmetic and verify |
| 16 | the layout-aware replacement (``_decode_topk_seq_lens``), including the tail |
| 17 | expansion consequences via the pure-torch expand/append pair that the fused |
| 18 | kernel is documented to replicate. |
| 19 | """ |
| 20 | |
| 21 | import torch |
| 22 | |
| 23 | # Bootstrap the glm5next package before importing kpool_compress directly: its |
| 24 | # import runs glm5next/__init__, which pulls model -> attention -> the indexer |
| 25 | # dispatcher (sparse_indexer) -> back into the nvidia/amd sparse_indexer |
| 26 | # modules. Production always enters via attention.py first. |
| 27 | # isort: off |
| 28 | import vllm.models.glm5next # noqa: F401 |
| 29 | from vllm.models.glm5next.nvidia.ops.kpool_compress import ( # noqa: E402 |
| 30 | append_tail_to_topk, |
| 31 | expand_pools_to_tokens, |
| 32 | ) |
| 33 | # isort: on |
| 34 | |
| 35 | from vllm.models.glm5next.common.sparse_indexer import ( |
| 36 | _decode_topk_seq_lens, |
| 37 | _fill_short_decode_causal_indices, |
| 38 | ) |
| 39 | from vllm.platforms import current_platform |
| 40 | |
| 41 | if current_platform.is_rocm(): |
| 42 | from vllm.models.glm5next.amd import sparse_indexer as indexer_mod |
| 43 | else: |
| 44 | from vllm.models.glm5next.nvidia import sparse_indexer as indexer_mod |
| 45 | |
| 46 | KPOOL = 4 |
| 47 | TOPK_TOKENS = 16 |
| 48 | SELECT_K = TOPK_TOKENS // KPOOL |
| 49 | |
| 50 | |
| 51 | def test_kpool_ops_dispatch_matches_platform(): |
| 52 | expected_backend = ".amd." if current_platform.is_rocm() else ".nvidia." |
| 53 | assert expected_backend in indexer_mod.kpool_ops.__name__ |
| 54 | |
| 55 | |
| 56 | def test_short_decode_fills_exact_causal_rows(): |
| 57 | topk = torch.full((3, 8), 99, dtype=torch.int32) |
| 58 | positions = torch.tensor([0, 3, 7], dtype=torch.int64) |
| 59 | |
| 60 | assert _fill_short_decode_causal_indices(topk, positions, 3, 8, 8) |
| 61 | assert topk.tolist() == [ |
| 62 | [0, -1, -1, -1, -1, -1, -1, -1], |
| 63 | [0, 1, 2, 3, -1, -1, -1, -1], |
| 64 | [0, 1, 2, 3, 4, 5, 6, 7], |
| 65 | ] |
| 66 | |
| 67 | |
| 68 | def test_short_decode_leaves_buffer_unchanged_for_sparse_context(): |
| 69 | topk = torch.full((2, 8), 99, dtype=torch.int32) |
| 70 | before = topk.clone() |
| 71 | |
| 72 | assert not _fill_short_decode_causal_indices(topk, torch.tensor([7, 8]), 2, 9, 8) |
| 73 | assert torch.equal(topk, before) |
| 74 | |
| 75 | |
| 76 | def make_non_uniform_batch(): |
| 77 | """3 requests: plain decode (1 token), MTP verify (4), adaptive verify (3). |
| 78 | |
| 79 | Flat decode positions (production layout: decode tokens first, then any |
| 80 | prefill tokens of the same batch): |
| 81 | req0: [30] (context len 30) |
| 82 | req1: [100..103] (verify at context len 100) |
| 83 | req2: [7, 8, 9] (verify at context len 7) |
| 84 | Followed by 4 prefill tokens at positions 555..558 so the flat tensor is |
| 85 | at least ``n = 3 * 4`` long, as it is in a real mixed batch. |
| 86 | """ |
| 87 | decode_lens = torch.tensor([1, 4, 3], dtype=torch.int64) |
| 88 | per_req_positions = [[30], [100, 101, 102, 103], [7, 8, 9]] |
| 89 | flat_decode = [p for req in per_req_positions for p in req] |
| 90 | positions = torch.tensor(flat_decode + [555, 556, 557, 558], dtype=torch.int64) |
| 91 | return decode_lens, per_req_positions, positions |
| 92 | |
| 93 | |
| 94 | def expected_row_seq_lens(per_req_positions, batch_size, next_n): |
| 95 | """Ground truth: row (b, t) -> pos + 1 for real rows, 0 for pad rows.""" |
| 96 | out = torch.zeros(batch_size * next_n, dtype=torch.int32) |
| 97 | for b, req in enumerate(per_req_positions): |
| 98 | for t, pos in enumerate(req): |
| 99 | out[b * next_n + t] = pos + 1 |
| 100 | return out |
| 101 | |
| 102 | |
| 103 | def test_legacy_flat_layout_misaligns_and_bleeds(): |
| 104 | """The bug: ``positions[:n] + 1`` with padded-row coordinates reads other |
| 105 | requests' positions and prefill positions.""" |
| 106 | decode_lens, per_req, positions = make_non_uniform_batch() |
| 107 | batch_size = decode_lens.shape[0] |
| 108 | next_n = int(decode_lens.max()) |
| 109 | n = batch_size * next_n |
| 110 | assert n == 12 |
| 111 | |
| 112 | legacy = positions[:n].to(torch.int32) + 1 |
| 113 | expected = expected_row_seq_lens(per_req, batch_size, next_n) |
| 114 | |
| 115 | # Row (1, 0): req1's first verify token (pos 100, seq 101) reads flat |
| 116 | # index 4 -> req1's LAST verify token (pos 103). Its true tail token is |
| 117 | # dropped and the tail anchors 3 tokens late. |
| 118 | assert int(legacy[4]) == 104 and int(expected[4]) == 101 |
| 119 | |
| 120 | # Rows (2, *): req2 starts at flat offset 5, but padded coordinates point |
| 121 | # at flat indices 8..11 -- the batch's PREFILL positions. |
| 122 | for t in range(3): |
| 123 | assert int(legacy[2 * next_n + t]) >= 556, ( |
| 124 | "legacy row (2, t) should read prefill positions" |
| 125 | ) |
| 126 | assert not torch.equal(legacy, expected) |
| 127 | |
| 128 | |
| 129 | def test_helper_uniform_layout_matches_flat_slice(): |
| 130 | """Uniform batches keep the flat shortcut (zero behavior/perf change).""" |
| 131 | per_req = [[200, 201, 202, 203], [50, 51, 52, 53]] |
| 132 | positions = torch.tensor([p for r in per_req for p in r], dtype=torch.int64) |
| 133 | decode_lens = torch.tensor([4, 4], dtype=torch.int64) |
| 134 | out = _decode_topk_seq_lens(positions, decode_lens, 8, 2, 4, requires_padding=False) |
| 135 | assert torch.equal(out, positions[:8].to(torch.int32) + 1) |
| 136 | |
| 137 | |
| 138 | def test_helper_padded_layout_per_row(): |
| 139 | """Non-uniform batches map every padded row to its own token's position; |
| 140 | pad rows get 0 (empty tail).""" |
| 141 | decode_lens, per_req, positions = make_non_uniform_batch() |
| 142 | out = _decode_topk_seq_lens(positions, decode_lens, 8, 3, 4, requires_padding=True) |
| 143 | expected = expected_row_seq_lens(per_req, 3, 4) |
| 144 | assert torch.equal(out, expected) |
| 145 | # Pad rows (0, 1..3) and (2, 3) collapse to 0 -> no tail appended. |
| 146 | assert out[1] == 0 and out[2] == 0 and out[3] == 0 and out[11] == 0 |
| 147 | |
| 148 | |
| 149 | def expand_tail_region(dec_seq): |
| 150 | """Run the production pure-torch expand + append pair (the fused kernel |
| 151 | replicates it exactly on the identity path) and return the tail columns |
| 152 | [TOPK_TOKENS, TOPK_TOKENS + KPOOL - 1).""" |
| 153 | rows = dec_seq.shape[0] |
| 154 | pool_ids = torch.arange(SELECT_K, dtype=torch.int64).expand(rows, SELECT_K) |
| 155 | valid = torch.ones_like(pool_ids, dtype=torch.bool) |
| 156 | expanded = expand_pools_to_tokens(pool_ids, valid, TOPK_TOKENS, KPOOL) |
| 157 | seq_lens = dec_seq.to(torch.int32) |
| 158 | pool_lens = (seq_lens // KPOOL).to(torch.int32) |
| 159 | out = append_tail_to_topk(expanded, seq_lens, pool_lens, KPOOL) |
| 160 | return out[:, TOPK_TOKENS:] |
| 161 | |
| 162 | |
| 163 | def test_tail_expansion_legacy_vs_fixed(): |
| 164 | """End-to-end tail consequence: the fixed mapping appends exactly the |
| 165 | request's trailing incomplete pool; the legacy flat slice drops real tail |
| 166 | tokens and emits indices far past the request's own sequence.""" |
| 167 | decode_lens, per_req, positions = make_non_uniform_batch() |
| 168 | n = 12 |
| 169 | fixed = _decode_topk_seq_lens( |
| 170 | positions, decode_lens, 8, 3, 4, requires_padding=True |
| 171 | ) |
| 172 | legacy = positions[:n].to(torch.int32) + 1 |
| 173 | |
| 174 | fixed_tail = expand_tail_region(fixed) |
| 175 | legacy_tail = expand_tail_region(legacy) |
| 176 | |
| 177 | # Fixed: row (1, 0) (seq 101) keeps its single tail token 100; row (2, 2) |
| 178 | # (seq 10) keeps its full trailing pool [8, 9]. |
| 179 | assert fixed_tail[4, 0].item() == 100 |
| 180 | assert fixed_tail[10, :2].tolist() == [8, 9] |
| 181 | |
| 182 | # Legacy: row (1, 0) loses its tail token (anchored at 104, count 0)... |
| 183 | assert legacy_tail[4, 0].item() == -1 |
| 184 | # ...and row (2, 2) reads prefill position 557 -> tail indices 556/557, |
| 185 | # way past req2's 10-token sequence -> out-of-bounds block-table reads. |
| 186 | assert legacy_tail[10, :2].tolist() == [556, 557] |
| 187 | |
| 188 | assert not torch.equal(legacy_tail, fixed_tail) |