| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """The ROCm MXFP4 indexer reads the paged cache in place, so each prefill |
| 4 | chunk is split back into its requests' rows, which one launch finds through |
| 5 | the chunk's query_start_loc. A chunk is whole requests or one query slice of a |
| 6 | long request, and both have to come back with the right rows and the right |
| 7 | block-table rows. |
| 8 | """ |
| 9 | |
| 10 | import pytest |
| 11 | import torch |
| 12 | |
| 13 | from vllm.platforms import current_platform |
| 14 | |
| 15 | if not current_platform.is_rocm(): |
| 16 | pytest.skip("ROCm-only", allow_module_level=True) |
| 17 | |
| 18 | from vllm.v1.attention.backends.mla.indexer import ( |
| 19 | DeepseekV32IndexerPrefillChunkMetadata, |
| 20 | ) |
| 21 | from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import ( |
| 22 | native_decode, |
| 23 | plan_gather_launches, |
| 24 | plan_prefill_chunks, |
| 25 | split_prefill_chunks, |
| 26 | ) |
| 27 | from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import rocm_mxfp4_consumer_rows |
| 28 | |
| 29 | |
| 30 | def _chunk(block_table, token_start, token_end, row_ends): |
| 31 | ends = torch.tensor(row_ends, dtype=torch.int32) |
| 32 | return DeepseekV32IndexerPrefillChunkMetadata( |
| 33 | block_table=block_table, |
| 34 | cu_seqlen_ks=torch.full_like(ends, 7), |
| 35 | cu_seqlen_ke=ends + 7, |
| 36 | cu_seq_lens=torch.zeros(1, dtype=torch.int32), |
| 37 | token_to_seq=torch.zeros(1, dtype=torch.int32), |
| 38 | total_seq_lens=1, |
| 39 | token_start=token_start, |
| 40 | token_end=token_end, |
| 41 | num_reqs=block_table.shape[0], |
| 42 | ) |
| 43 | |
| 44 | |
| 45 | # A decode (1 token), then prefills of 3, 3, 5 and 6 query tokens; the last |
| 46 | # one is sliced into two chunks. |
| 47 | QUERY_START_LOC = [0, 1, 4, 7, 12, 18] |
| 48 | SEQ_LENS = torch.tensor([40, 10, 12, 21, 64]) |
| 49 | BLOCK_TABLE = torch.arange(20, dtype=torch.int32).view(5, 4) |
| 50 | CHUNKS = [ |
| 51 | _chunk(BLOCK_TABLE[1:4], 1, 12, list(range(11))), |
| 52 | _chunk(BLOCK_TABLE[4:5], 12, 15, [29, 30, 30]), |
| 53 | _chunk(BLOCK_TABLE[4:5], 15, 18, [31, 31, 32]), |
| 54 | ] |
| 55 | |
| 56 | |
| 57 | def _plans(min_gather_width, block=8): |
| 58 | return plan_prefill_chunks( |
| 59 | CHUNKS, |
| 60 | QUERY_START_LOC, |
| 61 | SEQ_LENS, |
| 62 | SEQ_LENS.int() // 2, |
| 63 | 2, |
| 64 | block, |
| 65 | min_gather_width, |
| 66 | torch.tensor(QUERY_START_LOC, dtype=torch.int32), |
| 67 | ) |
| 68 | |
| 69 | |
| 70 | def test_chunks_split_into_requests(): |
| 71 | plans = _plans(20.0) |
| 72 | |
| 73 | assert [p.requests for p in plans] == [ |
| 74 | [(0, 3, 0), (3, 6, 1), (6, 11, 2)], |
| 75 | [(0, 3, 0)], |
| 76 | [(0, 3, 0)], |
| 77 | ] |
| 78 | # every chunk launches once, with offsets local to the chunk |
| 79 | assert [p.query_start_loc.tolist() for p in plans] == [ |
| 80 | [0, 3, 6, 11], |
| 81 | [0, 3], |
| 82 | [0, 3], |
| 83 | ] |
| 84 | # the logits are as wide as the longest compressed context in the chunk |
| 85 | assert [p.width for p in plans] == [10, 32, 32] |
| 86 | assert [p.context_lens.tolist() for p in plans] == [[5, 6, 10], [32], [32]] |
| 87 | torch.testing.assert_close(plans[2].row_ends, torch.tensor([31, 31, 32]).int()) |
| 88 | torch.testing.assert_close(plans[2].block_ends, torch.tensor([4, 4, 4]).int()) |
| 89 | assert [p.use_gather for p in plans] == [False, True, True] |
| 90 | assert [p.first_request for p in plans] == [1, 4, 4] |
| 91 | # no candidate blocks, as in the model-neutral builder |
| 92 | for plan in _plans(None, block=0): |
| 93 | assert plan.block_ends is None and not plan.use_gather |
| 94 | |
| 95 | |
| 96 | def test_launches_follow_the_logits_budget(monkeypatch): |
| 97 | """Each launch's fp32 logits fit in VLLM_SPARSE_INDEXER_MAX_LOGITS_MB. A dense |
| 98 | chunk's logits are [rows, longest context in the chunk]: K is read in place, |
| 99 | not gathered, so requests share a launch until their rows no longer fit, and |
| 100 | a request that does not fit alone is split on its query rows. A consumer's |
| 101 | logits are [rows, pool].""" |
| 102 | budget = 4 * 1000 * 30 |
| 103 | contexts, queries = torch.tensor([1000, 600, 200]), torch.tensor([10, 10, 10]) |
| 104 | assert split_prefill_chunks(contexts, queries, budget, 2) == [ |
| 105 | (slice(2, 5), slice(0, 30)) |
| 106 | ] |
| 107 | # the third request would take the chunk past the budget |
| 108 | queries = torch.tensor([10, 10, 11]) |
| 109 | assert split_prefill_chunks(contexts, queries, budget) == [ |
| 110 | (slice(0, 2), slice(0, 20)), |
| 111 | (slice(2, 3), slice(0, 11)), |
| 112 | ] |
| 113 | assert split_prefill_chunks(torch.tensor([1000]), torch.tensor([100]), budget) == [ |
| 114 | (slice(0, 1), slice(lo, min(lo + 30, 100))) for lo in range(0, 100, 30) |
| 115 | ] |
| 116 | # the default 512 MiB holds 128 rows of a 1M context, or 8192 of a 16K pool |
| 117 | assert split_prefill_chunks( |
| 118 | torch.tensor([1 << 20]), torch.tensor([200]), 512 << 20 |
| 119 | ) == [(slice(0, 1), slice(0, 128)), (slice(0, 1), slice(128, 200))] |
| 120 | monkeypatch.setenv("VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", "512") |
| 121 | assert rocm_mxfp4_consumer_rows(16384) == 8192 |
| 122 | |
| 123 | |
| 124 | def test_gather_launches_rejoin_a_requests_rows(): |
| 125 | """The consumers' logits are [rows, pool] whatever the context, so a |
| 126 | request's rows go back together across the chunks that cut them, and a |
| 127 | launch never mixes requests.""" |
| 128 | # the first chunk is too narrow to gather, the last request's slices rejoin |
| 129 | (joined,) = plan_gather_launches(CHUNKS, _plans(20.0), 100) |
| 130 | assert (joined.token_start, joined.token_end) == (12, 18) |
| 131 | # every chunk gathers: the first chunk's three requests launch apart |
| 132 | launches = plan_gather_launches(CHUNKS, _plans(0.0), 100) |
| 133 | assert [(g.token_start, g.token_end) for g in launches] == [ |
| 134 | (1, 4), |
| 135 | (4, 7), |
| 136 | (7, 12), |
| 137 | (12, 18), |
| 138 | ] |
| 139 | assert plan_gather_launches(CHUNKS, _plans(None), 100) == [] |
| 140 | |
| 141 | |
| 142 | def test_decode_launches_native_on_uniform_steps(): |
| 143 | """A decode step goes to the kernel as next_n-row sequences only when its |
| 144 | shape says so: each request has next_n rows, and cudagraph padding (query |
| 145 | length 0) comes only after them. A FULL graph captured on a uniform batch |
| 146 | then replays the same launch on a padded one.""" |
| 147 | rows = torch.zeros(8, dtype=torch.int32) |
| 148 | row_block_table = torch.arange(32, dtype=torch.int32).view(8, 4) |
| 149 | context_lens = torch.tensor([9, 4, 7, 0, 5], dtype=torch.int32) |
| 150 | |
| 151 | native = native_decode(rows, row_block_table, [2, 2, 2, 0], 2, context_lens) |
| 152 | assert native is not None and native.next_n == 2 |
| 153 | assert native.context_lens.tolist() == [9, 4, 7, 0] |
| 154 | # every request's first flattened row, as a view |
| 155 | assert native.block_table.tolist() == row_block_table[::2].tolist() |
| 156 | assert native.block_table.data_ptr() == row_block_table.data_ptr() |
| 157 | |
| 158 | assert native_decode(rows, row_block_table, [2, 0, 2, 2], 2, context_lens) is None |
| 159 | # a ragged step, rows padded past the requests, and no speculation |
| 160 | assert native_decode(rows[:5], row_block_table, [2, 1, 2], 2, context_lens) is None |
| 161 | assert native_decode(rows, row_block_table, [2, 2, 2], 2, context_lens) is None |
| 162 | assert native_decode(rows[:4], row_block_table, [1] * 4, 1, context_lens) is None |