| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Per-request scheduling of MLA chunked-context prefill. |
| 4 | |
| 5 | These cover the metadata contract `_compute_prefill_context` relies on: every |
| 6 | context row is gathered exactly once, no chunk exceeds the workspace, and a |
| 7 | chunk never covers a prefill without context (which is why the partial no |
| 8 | longer needs an empty-span masking pass). |
| 9 | """ |
| 10 | |
| 11 | from types import SimpleNamespace |
| 12 | |
| 13 | import pytest |
| 14 | import torch |
| 15 | |
| 16 | import vllm.utils.gpu_sync_debug as gsd |
| 17 | from vllm.model_executor.layers.attention.mla_attention import ( |
| 18 | build_mla_chunked_context_metadata, |
| 19 | init_mla_context_partial, |
| 20 | reorg_kvcache, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.attention.sparse_mla_attention import ( |
| 23 | SparseMLACommonMetadataBuilder, |
| 24 | ) |
| 25 | |
| 26 | BLOCK_SIZE = 16 |
| 27 | |
| 28 | |
| 29 | def build_chunked_context( |
| 30 | context_lens: list[int], |
| 31 | query_lens: list[int], |
| 32 | workspace_size: int, |
| 33 | block_size: int = BLOCK_SIZE, |
| 34 | dcp_world_size: int = 1, |
| 35 | dcp_local_block_size: int = 1, |
| 36 | ): |
| 37 | query_start_loc = torch.zeros(len(query_lens) + 1, dtype=torch.int32) |
| 38 | query_start_loc[1:] = torch.tensor(query_lens, dtype=torch.int32).cumsum(0) |
| 39 | workspace_rows = workspace_size + workspace_size // dcp_world_size |
| 40 | return build_mla_chunked_context_metadata( |
| 41 | context_lens_cpu=torch.tensor(context_lens, dtype=torch.int32), |
| 42 | prefill_query_start_loc_cpu=query_start_loc, |
| 43 | chunked_prefill_workspace=torch.empty((workspace_rows, 1)), |
| 44 | chunked_prefill_workspace_size=workspace_size, |
| 45 | block_size=block_size, |
| 46 | align_chunk_to_block=True, |
| 47 | device=torch.device("cpu"), |
| 48 | dcp_world_size=dcp_world_size, |
| 49 | dcp_local_block_size=dcp_local_block_size, |
| 50 | dcp_virtual_block_size=dcp_local_block_size * dcp_world_size, |
| 51 | ) |
| 52 | |
| 53 | |
| 54 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 55 | def test_sparse_context_lengths_do_not_force_gpu_sync(monkeypatch): |
| 56 | """Sparse MLA must pin computed context lengths before the H2D copy.""" |
| 57 | builder = SimpleNamespace( |
| 58 | _context_lens_cpu=torch.empty(2, dtype=torch.int32, pin_memory=True), |
| 59 | chunked_prefill_workspace=torch.empty((2048, 1)), |
| 60 | chunked_prefill_workspace_size=1024, |
| 61 | kv_cache_spec=SimpleNamespace(block_size=BLOCK_SIZE), |
| 62 | device=torch.device("cuda"), |
| 63 | dcp_world_size=1, |
| 64 | dcp_local_block_size=1, |
| 65 | dcp_virtual_block_size=1, |
| 66 | dcp_manager=None, |
| 67 | ) |
| 68 | common_metadata = SimpleNamespace( |
| 69 | seq_lens_cpu_upper_bound=torch.tensor([17, 2052, 324], dtype=torch.int32), |
| 70 | query_start_loc_cpu=torch.tensor([0, 1, 5, 9], dtype=torch.int32), |
| 71 | ) |
| 72 | monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") |
| 73 | monkeypatch.setattr(gsd, "_sync_check_enabled", True) |
| 74 | gsd._install_copy_checkers() |
| 75 | metadata = gsd.with_gpu_sync_check( |
| 76 | SparseMLACommonMetadataBuilder._build_chunked_context_fields |
| 77 | )( |
| 78 | builder, |
| 79 | common_metadata, |
| 80 | num_decodes=1, |
| 81 | num_prefills=2, |
| 82 | prefill_query_lens_cpu=torch.tensor([4, 4], dtype=torch.int32), |
| 83 | ) |
| 84 | assert metadata is not None |
| 85 | assert metadata.context_lens.device.type == "cuda" |
| 86 | assert metadata.context_lens.tolist() == [2048, 320] |
| 87 | |
| 88 | |
| 89 | @pytest.mark.parametrize( |
| 90 | "context_lens,workspace_size", |
| 91 | [ |
| 92 | # Heterogeneous batch: the long request must not shrink the chunk the |
| 93 | # short ones share. |
| 94 | ([960, 16, 16, 16, 16, 16, 16, 16], 1024), |
| 95 | # A single request that does not fit is split on its own. |
| 96 | ([2048, 320], 1024), |
| 97 | # Every request needs its own chunk. |
| 98 | ([1024, 1024, 1024], 1024), |
| 99 | # Unaligned tails. |
| 100 | ([37, 1000, 5], 1024), |
| 101 | # A zero-context request between active requests must not become an |
| 102 | # empty TRTLLM ragged-attention row. |
| 103 | ([4, 0, 2], 1024), |
| 104 | ([1], 1024), |
| 105 | ], |
| 106 | ) |
| 107 | def test_chunks_gather_every_context_row_exactly_once(context_lens, workspace_size): |
| 108 | """Chunks tile the batch's context without gaps, overlap, or overflow. |
| 109 | |
| 110 | This is the invariant the whole schedule rests on: attention over a chunk |
| 111 | only sees the rows that chunk gathered, so a row counted twice or dropped |
| 112 | silently corrupts the context partial. |
| 113 | """ |
| 114 | query_lens = [4] * len(context_lens) |
| 115 | metadata = build_chunked_context(context_lens, query_lens, workspace_size) |
| 116 | assert metadata is not None |
| 117 | assert metadata.context_lens.tolist() == context_lens |
| 118 | |
| 119 | gathered: dict[int, list[tuple[int, int]]] = {} |
| 120 | previous_request_start = -1 |
| 121 | for chunk in metadata.chunks: |
| 122 | assert chunk.all_rows_active |
| 123 | assert chunk.num_context_tokens <= workspace_size |
| 124 | assert chunk.num_context_tokens == int(chunk.cu_seq_lens[-1]) |
| 125 | assert chunk.token_to_seq.shape[0] == chunk.num_context_tokens |
| 126 | # Chunks are emitted in request order, which is what lets the partial |
| 127 | # treat only a chunk's first request as a continuation. |
| 128 | assert chunk.request_slice.start >= previous_request_start |
| 129 | previous_request_start = chunk.request_slice.start |
| 130 | |
| 131 | starts = chunk.starts.tolist() |
| 132 | seq_lens = chunk.seq_lens.tolist() |
| 133 | assert len(starts) == len(seq_lens) == chunk.num_requests |
| 134 | for offset, (start, length) in enumerate(zip(starts, seq_lens)): |
| 135 | assert length > 0, "a chunk must not cover an empty context span" |
| 136 | gathered.setdefault(chunk.request_slice.start + offset, []).append( |
| 137 | (start, length) |
| 138 | ) |
| 139 | |
| 140 | expected = {i: length for i, length in enumerate(context_lens) if length > 0} |
| 141 | assert gathered.keys() == expected.keys() |
| 142 | for request, spans in gathered.items(): |
| 143 | cursor = 0 |
| 144 | for start, length in spans: |
| 145 | assert start == cursor |
| 146 | cursor += length |
| 147 | assert cursor == expected[request] |
| 148 | |
| 149 | |
| 150 | def test_continuation_is_confined_to_a_chunks_first_request(): |
| 151 | """Only a chunk's first request may continue an earlier chunk. |
| 152 | |
| 153 | `accumulate_mla_context_chunk` merges the continuation token slice and |
| 154 | writes the rest, so a chunk that continues a request other than its first, |
| 155 | or reports the wrong boundary, folds the partial into the wrong tokens. |
| 156 | """ |
| 157 | context_lens = [2048, 32, 32] |
| 158 | query_lens = [3, 5, 7] |
| 159 | metadata = build_chunked_context(context_lens, query_lens, 1024) |
| 160 | assert metadata is not None |
| 161 | query_start_loc = [0, 3, 8, 15] |
| 162 | |
| 163 | assert [chunk.is_continuation for chunk in metadata.chunks] == [ |
| 164 | False, |
| 165 | True, |
| 166 | False, |
| 167 | ] |
| 168 | for chunk in metadata.chunks: |
| 169 | starts = chunk.starts.tolist() |
| 170 | assert chunk.is_continuation == (starts[0] > 0) |
| 171 | # Requests after the first always start at the beginning of their |
| 172 | # context, so they can only ever be initialized, never merged. |
| 173 | assert all(start == 0 for start in starts[1:]) |
| 174 | request_start = chunk.request_slice.start |
| 175 | request_end = chunk.request_slice.stop |
| 176 | assert chunk.continuation_token_end == query_start_loc[request_start + 1] |
| 177 | assert chunk.token_slice == slice( |
| 178 | query_start_loc[request_start], query_start_loc[request_end] |
| 179 | ) |
| 180 | assert chunk.query_start_loc.tolist() == [ |
| 181 | offset - chunk.token_slice.start |
| 182 | for offset in query_start_loc[request_start : request_end + 1] |
| 183 | ] |
| 184 | |
| 185 | |
| 186 | def test_tail_splitting_minimizes_chunks(): |
| 187 | """A tail request may fill one chunk and continue at the next chunk's head. |
| 188 | |
| 189 | Without tail splitting these contexts need three chunks. Splitting on the |
| 190 | block boundary reduces that to the workspace lower bound of two. |
| 191 | """ |
| 192 | metadata = build_chunked_context([768, 512, 512], [3, 5, 7], 1024) |
| 193 | assert metadata is not None |
| 194 | |
| 195 | assert len(metadata.chunks) == 2 |
| 196 | first, second = metadata.chunks |
| 197 | assert first.starts.tolist() == [0, 0] |
| 198 | assert first.seq_lens.tolist() == [768, 256] |
| 199 | assert not first.is_continuation |
| 200 | |
| 201 | assert second.starts.tolist() == [256, 0] |
| 202 | assert second.seq_lens.tolist() == [256, 512] |
| 203 | assert second.is_continuation |
| 204 | |
| 205 | |
| 206 | def test_prefills_without_context_are_skipped_and_neutralized(): |
| 207 | """Context-free prefills get neutral partials instead of chunk work. |
| 208 | |
| 209 | If the full-query partial leaves either an internal or trailing gap |
| 210 | uninitialized, the final merge combines the suffix with undefined scratch. |
| 211 | """ |
| 212 | metadata = build_chunked_context([64, 0, 64, 0], [4, 4, 4, 4], 1024) |
| 213 | assert metadata is not None |
| 214 | |
| 215 | covered = { |
| 216 | request |
| 217 | for chunk in metadata.chunks |
| 218 | for request in range(chunk.request_slice.start, chunk.request_slice.stop) |
| 219 | } |
| 220 | assert covered == {0, 2} |
| 221 | assert metadata.empty_token_slices == [slice(4, 8), slice(12, 16)] |
| 222 | |
| 223 | output, output_lse = init_mla_context_partial( |
| 224 | metadata, |
| 225 | attn_output=torch.empty(4, 2, 3), |
| 226 | attn_softmax_lse=torch.empty(2, 4), |
| 227 | num_tokens=16, |
| 228 | ) |
| 229 | assert output.shape == (16, 2, 3) |
| 230 | assert output_lse.shape == (2, 16) |
| 231 | for token_slice in metadata.empty_token_slices: |
| 232 | assert not torch.count_nonzero(output[token_slice]) |
| 233 | assert torch.isneginf(output_lse[:, token_slice]).all() |
| 234 | |
| 235 | |
| 236 | def test_no_context_needs_no_chunks(): |
| 237 | assert build_chunked_context([0, 0], [4, 4], 1024) is None |
| 238 | |
| 239 | |
| 240 | def test_dcp_chunks_fit_the_per_rank_row_budget(): |
| 241 | """Under DCP each rank gathers its own shard, so the budget is 1/world. |
| 242 | |
| 243 | The local starts and lengths must tile the rank's padded context the same |
| 244 | way the global ones tile the real context, since `reorg_kvcache` unpads |
| 245 | using them. |
| 246 | """ |
| 247 | dcp_world_size, interleave, workspace_size = 2, 64, 1024 |
| 248 | context_lens = [3000, 200, 200] |
| 249 | metadata = build_chunked_context( |
| 250 | context_lens, |
| 251 | [4] * len(context_lens), |
| 252 | workspace_size, |
| 253 | block_size=128, |
| 254 | dcp_world_size=dcp_world_size, |
| 255 | dcp_local_block_size=interleave, |
| 256 | ) |
| 257 | assert metadata is not None |
| 258 | assert metadata.context_lens.tolist() == context_lens |
| 259 | virtual_block_size = interleave * dcp_world_size |
| 260 | |
| 261 | local_cursor: dict[int, int] = {} |
| 262 | for chunk in metadata.chunks: |
| 263 | assert chunk.num_local_context_tokens <= workspace_size // dcp_world_size |
| 264 | assert chunk.local_starts == chunk.starts.tolist() |
| 265 | assert chunk.padded_local_cu_seq_lens.tolist() == [0] + list( |
| 266 | torch.tensor(chunk.padded_local_seq_lens).cumsum(0) |
| 267 | ) |
| 268 | assert ( |
| 269 | chunk.padded_local_token_to_seq.shape[0] == chunk.num_local_context_tokens |
| 270 | ) |
| 271 | for offset, (start, length) in enumerate( |
| 272 | zip(chunk.local_starts, chunk.padded_local_seq_lens) |
| 273 | ): |
| 274 | request = chunk.request_slice.start + offset |
| 275 | assert start == local_cursor.get(request, 0) |
| 276 | local_cursor[request] = start + length |
| 277 | |
| 278 | for request, length in enumerate(context_lens): |
| 279 | padded_local = -(-length // virtual_block_size) * interleave |
| 280 | assert local_cursor[request] == padded_local |
| 281 | |
| 282 | |
| 283 | def test_dcp_reorg_uses_each_chunks_local_starts(): |
| 284 | """Reorganization drops DCP padding at each request's own chunk offset.""" |
| 285 | dcp_world_size, interleave, workspace_size = 2, 64, 1024 |
| 286 | metadata = build_chunked_context( |
| 287 | [3000, 200, 200], |
| 288 | [4, 4, 4], |
| 289 | workspace_size, |
| 290 | block_size=128, |
| 291 | dcp_world_size=dcp_world_size, |
| 292 | dcp_local_block_size=interleave, |
| 293 | ) |
| 294 | assert metadata is not None |
| 295 | |
| 296 | for chunk in metadata.chunks: |
| 297 | assert chunk.padded_local_seq_lens is not None |
| 298 | assert chunk.local_context_lens_allranks is not None |
| 299 | assert chunk.local_starts is not None |
| 300 | |
| 301 | toks = chunk.num_local_context_tokens |
| 302 | rank_buffers = [torch.full((toks, 1, 1), -1) for _ in range(dcp_world_size)] |
| 303 | expected = [] |
| 304 | src_token_idx = 0 |
| 305 | for request, (padded_len, local_lens, local_start) in enumerate( |
| 306 | zip( |
| 307 | chunk.padded_local_seq_lens, |
| 308 | chunk.local_context_lens_allranks, |
| 309 | chunk.local_starts, |
| 310 | ) |
| 311 | ): |
| 312 | for rank, local_len in enumerate(local_lens): |
| 313 | actual_len = min(max(0, local_len - local_start), padded_len) |
| 314 | values = ( |
| 315 | request * 100_000 |
| 316 | + rank * 10_000 |
| 317 | + torch.arange(local_start, local_start + actual_len) |
| 318 | ).view(-1, 1, 1) |
| 319 | rank_buffers[rank][src_token_idx : src_token_idx + actual_len] = values |
| 320 | expected.append(values) |
| 321 | src_token_idx += padded_len |
| 322 | |
| 323 | allgathered = torch.cat(rank_buffers) |
| 324 | reorganized, _ = reorg_kvcache( |
| 325 | allgathered, |
| 326 | allgathered, |
| 327 | padded_local_chunk_seq_lens_lst=chunk.padded_local_seq_lens, |
| 328 | local_context_lens_allranks=chunk.local_context_lens_allranks, |
| 329 | local_starts=chunk.local_starts, |
| 330 | sum_seq_len=chunk.num_context_tokens, |
| 331 | max_seq_len=chunk.max_seq_len, |
| 332 | toks=toks, |
| 333 | ) |
| 334 | |
| 335 | torch.testing.assert_close(reorganized, torch.cat(expected)) |
| 336 | |
| 337 | |
| 338 | def test_chunked_context_metadata_owns_context_lengths(): |
| 339 | """Reusing CPU scratch must not overwrite an earlier batch's context lengths.""" |
| 340 | context_lens = torch.tensor([16, 32], dtype=torch.int32) |
| 341 | metadata = build_mla_chunked_context_metadata( |
| 342 | context_lens_cpu=context_lens, |
| 343 | prefill_query_start_loc_cpu=torch.tensor([0, 4, 8], dtype=torch.int32), |
| 344 | chunked_prefill_workspace=torch.empty((64, 1)), |
| 345 | chunked_prefill_workspace_size=64, |
| 346 | block_size=BLOCK_SIZE, |
| 347 | align_chunk_to_block=True, |
| 348 | device=torch.device("cpu"), |
| 349 | dcp_world_size=1, |
| 350 | dcp_local_block_size=1, |
| 351 | dcp_virtual_block_size=1, |
| 352 | ) |
| 353 | context_lens.zero_() |
| 354 | assert metadata is not None |
| 355 | assert metadata.context_lens.tolist() == [16, 32] |