| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for DeepSeek-V4 in-image bidirectional SWA visibility (vision variant). |
| 4 | |
| 5 | The reference semantics are a torch port of the official |
| 6 | ``get_image_visible`` / ``get_window_topk_idxs_visible``: inside an image span |
| 7 | [span_start, span_end] (inclusive), a token at ``pos`` sees |
| 8 | ``min(pos - span_start, max_image_tokens - 1)`` extra tokens to the left and |
| 9 | ``min(span_end - pos, max_image_tokens)`` to the right; the window then starts |
| 10 | at ``max(pos - (window - 1) - max(left - (window - 1), 0), 0)`` and ends at |
| 11 | ``pos + right`` (inclusive). |
| 12 | """ |
| 13 | |
| 14 | import inspect |
| 15 | |
| 16 | import pytest |
| 17 | import torch |
| 18 | from typing_extensions import TypedDict |
| 19 | |
| 20 | from tests.v1.attention.utils import create_vllm_config |
| 21 | from vllm.models.deepseek_v4.common.ops.cache_utils import ( |
| 22 | build_flashinfer_mixed_sparse_indices, |
| 23 | combine_topk_swa_indices, |
| 24 | ) |
| 25 | from vllm.v1.attention.backend import CommonAttentionMetadata |
| 26 | from vllm.v1.attention.backends.mla.sparse_swa import ( |
| 27 | DeepseekSparseSWAMetadataBuilder, |
| 28 | _compute_image_visibility_kernel, |
| 29 | _compute_swa_indices_and_lens_kernel, |
| 30 | ) |
| 31 | from vllm.v1.kv_cache_interface import SlidingWindowMLASpec |
| 32 | |
| 33 | WINDOW = 8 |
| 34 | MAX_IMG = 6 |
| 35 | WIDTH = WINDOW + MAX_IMG |
| 36 | BLOCK_SIZE = 64 |
| 37 | |
| 38 | |
| 39 | def ref_left_right( |
| 40 | seq_lens: list[int], |
| 41 | query_lens: list[int], |
| 42 | spans_per_req: list[list[tuple[int, int]]], |
| 43 | max_image_tokens: int, |
| 44 | ) -> tuple[list[int], list[int]]: |
| 45 | """Per-token (left, right) for the flattened decode-first token stream.""" |
| 46 | lefts: list[int] = [] |
| 47 | rights: list[int] = [] |
| 48 | for seq_len, query_len, spans in zip(seq_lens, query_lens, spans_per_req): |
| 49 | prefix_len = seq_len - query_len |
| 50 | for i in range(query_len): |
| 51 | pos = prefix_len + i |
| 52 | left = right = 0 |
| 53 | for span_start, span_end in spans: |
| 54 | if span_start <= pos <= span_end: |
| 55 | left = min(pos - span_start, max_image_tokens - 1) |
| 56 | right = min(span_end - pos, max_image_tokens) |
| 57 | lefts.append(left) |
| 58 | rights.append(right) |
| 59 | return lefts, rights |
| 60 | |
| 61 | |
| 62 | def ref_swa_bounds( |
| 63 | pos: int, window: int, left: int, right: int, replay_start: int = 0 |
| 64 | ) -> tuple[int, int]: |
| 65 | """Reference [start, end) window bounds for one query token. Under SWA |
| 66 | bounded replay no window KV exists below ``replay_start``.""" |
| 67 | left_add = max(left - (window - 1), 0) |
| 68 | start = max(pos - (window - 1) - left_add, 0, replay_start) |
| 69 | return start, pos + right + 1 |
| 70 | |
| 71 | |
| 72 | def make_batch( |
| 73 | seq_lens: list[int], |
| 74 | query_lens: list[int], |
| 75 | device: torch.device, |
| 76 | ): |
| 77 | """Build the kernel inputs for a decode-first batch.""" |
| 78 | num_reqs = len(seq_lens) |
| 79 | query_start_loc = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device) |
| 80 | query_start_loc[1:] = torch.tensor( |
| 81 | query_lens, dtype=torch.int32, device=device |
| 82 | ).cumsum(0) |
| 83 | num_tokens = int(query_start_loc[-1]) |
| 84 | seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device) |
| 85 | token_to_req = torch.repeat_interleave( |
| 86 | torch.arange(num_reqs, dtype=torch.int32, device=device), |
| 87 | torch.tensor(query_lens, dtype=torch.int32, device=device), |
| 88 | ) |
| 89 | max_blocks = (max(seq_lens) + BLOCK_SIZE - 1) // BLOCK_SIZE |
| 90 | block_table = torch.arange( |
| 91 | num_reqs * max_blocks, dtype=torch.int32, device=device |
| 92 | ).view(num_reqs, max_blocks) |
| 93 | slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) |
| 94 | return query_start_loc, seq_lens_t, token_to_req, slot_mapping, block_table |
| 95 | |
| 96 | |
| 97 | def ref_swa_slot_rows( |
| 98 | seq_lens: list[int], |
| 99 | query_lens: list[int], |
| 100 | spans_per_req: list[list[tuple[int, int]]], |
| 101 | block_table: torch.Tensor, |
| 102 | window: int, |
| 103 | max_image_tokens: int, |
| 104 | width: int, |
| 105 | replay_starts: list[int] | None = None, |
| 106 | ) -> tuple[list[list[int]], list[int]]: |
| 107 | """Reference paged slot-id rows and lens for every token in the batch.""" |
| 108 | block_table_cpu = block_table.cpu() |
| 109 | lefts, rights = ref_left_right( |
| 110 | seq_lens, query_lens, spans_per_req, max_image_tokens |
| 111 | ) |
| 112 | replay_starts = replay_starts or [0] * len(seq_lens) |
| 113 | rows: list[list[int]] = [] |
| 114 | lens: list[int] = [] |
| 115 | token = 0 |
| 116 | for req, (seq_len, query_len) in enumerate(zip(seq_lens, query_lens)): |
| 117 | prefix_len = seq_len - query_len |
| 118 | for i in range(query_len): |
| 119 | pos = prefix_len + i |
| 120 | start, end = ref_swa_bounds( |
| 121 | pos, window, lefts[token], rights[token], replay_starts[req] |
| 122 | ) |
| 123 | row = [] |
| 124 | for p in range(start, end): |
| 125 | blk = int(block_table_cpu[req, p // BLOCK_SIZE]) |
| 126 | row.append(blk * BLOCK_SIZE + p % BLOCK_SIZE) |
| 127 | lens.append(len(row)) |
| 128 | row.extend([-1] * (width - len(row))) |
| 129 | rows.append(row) |
| 130 | token += 1 |
| 131 | return rows, lens |
| 132 | |
| 133 | |
| 134 | def run_swa_kernel( |
| 135 | seq_lens: list[int], |
| 136 | query_lens: list[int], |
| 137 | spans_per_req: list[list[tuple[int, int]]], |
| 138 | window: int = WINDOW, |
| 139 | max_image_tokens: int = MAX_IMG, |
| 140 | with_image: bool = True, |
| 141 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 142 | device = torch.device("cuda") |
| 143 | query_start_loc, seq_lens_t, token_to_req, slot_mapping, block_table = make_batch( |
| 144 | seq_lens, query_lens, device |
| 145 | ) |
| 146 | num_tokens = int(query_start_loc[-1]) |
| 147 | width = window + (max_image_tokens if with_image else 0) |
| 148 | swa_indices = torch.zeros(num_tokens, 1, width, dtype=torch.int32, device=device) |
| 149 | swa_lens = torch.zeros(num_tokens, dtype=torch.int32, device=device) |
| 150 | is_valid = slot_mapping >= 0 |
| 151 | replay_start_t = torch.zeros(len(seq_lens), dtype=torch.int32, device=device) |
| 152 | |
| 153 | if with_image: |
| 154 | lefts, rights = ref_left_right( |
| 155 | seq_lens, query_lens, spans_per_req, max_image_tokens |
| 156 | ) |
| 157 | left_t = torch.tensor(lefts, dtype=torch.int32, device=device) |
| 158 | right_t = torch.tensor(rights, dtype=torch.int32, device=device) |
| 159 | else: |
| 160 | left_t = right_t = swa_lens # unused dummies |
| 161 | |
| 162 | _compute_swa_indices_and_lens_kernel[(num_tokens,)]( |
| 163 | swa_indices, |
| 164 | swa_indices.stride(0), |
| 165 | swa_lens, |
| 166 | window, |
| 167 | width, |
| 168 | left_t, |
| 169 | right_t, |
| 170 | query_start_loc, |
| 171 | seq_lens_t, |
| 172 | token_to_req, |
| 173 | is_valid, |
| 174 | block_table, |
| 175 | block_table.stride(0), |
| 176 | BLOCK_SIZE, |
| 177 | replay_start_t, |
| 178 | token_offset=0, |
| 179 | HAS_IMAGE=with_image, |
| 180 | TRITON_BLOCK_SIZE=1024, |
| 181 | ) |
| 182 | return swa_indices[:, 0], swa_lens |
| 183 | |
| 184 | |
| 185 | # seq_lens, query_lens, spans per request (positions are prompt-absolute, |
| 186 | # inclusive on both ends, matching the reference's sentinel-bracketed spans). |
| 187 | class _Case(TypedDict): |
| 188 | seq_lens: list[int] |
| 189 | query_lens: list[int] |
| 190 | spans: list[list[tuple[int, int]]] |
| 191 | |
| 192 | |
| 193 | CASES: list[_Case] = [ |
| 194 | # two image spans in one request, one pure-text request |
| 195 | { |
| 196 | "seq_lens": [30, 12], |
| 197 | "query_lens": [30, 12], |
| 198 | "spans": [[(4, 12), (20, 24)], []], |
| 199 | }, |
| 200 | # span larger than max_image_tokens, straddling the window boundary |
| 201 | { |
| 202 | "seq_lens": [26], |
| 203 | "query_lens": [26], |
| 204 | "spans": [[(2, 17)]], |
| 205 | }, |
| 206 | # chunked prefill: span fully inside the second chunk's query |
| 207 | { |
| 208 | "seq_lens": [40, 9], |
| 209 | "query_lens": [16, 9], |
| 210 | "spans": [[(30, 38)], []], |
| 211 | }, |
| 212 | # span ending exactly at the prompt end; tiny request |
| 213 | { |
| 214 | "seq_lens": [7, 15], |
| 215 | "query_lens": [7, 15], |
| 216 | "spans": [[(0, 6)], [(10, 14)]], |
| 217 | }, |
| 218 | ] |
| 219 | |
| 220 | |
| 221 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 222 | @pytest.mark.parametrize("case", CASES) |
| 223 | def test_swa_indices_kernel_with_image_spans(case): |
| 224 | rows, lens = ref_swa_slot_rows( |
| 225 | case["seq_lens"], |
| 226 | case["query_lens"], |
| 227 | case["spans"], |
| 228 | make_batch(case["seq_lens"], case["query_lens"], torch.device("cuda"))[4], |
| 229 | WINDOW, |
| 230 | MAX_IMG, |
| 231 | WIDTH, |
| 232 | ) |
| 233 | indices, actual_lens = run_swa_kernel( |
| 234 | case["seq_lens"], case["query_lens"], case["spans"], with_image=True |
| 235 | ) |
| 236 | assert actual_lens.cpu().tolist() == lens |
| 237 | assert indices.cpu().tolist() == rows |
| 238 | |
| 239 | |
| 240 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 241 | @pytest.mark.parametrize("case", CASES) |
| 242 | def test_swa_indices_kernel_without_image_unchanged(case): |
| 243 | """HAS_IMAGE=False must reproduce the plain causal sliding window.""" |
| 244 | rows, lens = ref_swa_slot_rows( |
| 245 | case["seq_lens"], |
| 246 | case["query_lens"], |
| 247 | [[] for _ in case["seq_lens"]], |
| 248 | make_batch(case["seq_lens"], case["query_lens"], torch.device("cuda"))[4], |
| 249 | WINDOW, |
| 250 | MAX_IMG, |
| 251 | WINDOW, |
| 252 | ) |
| 253 | indices, actual_lens = run_swa_kernel( |
| 254 | case["seq_lens"], case["query_lens"], case["spans"], with_image=False |
| 255 | ) |
| 256 | assert indices.shape[-1] == WINDOW |
| 257 | assert actual_lens.cpu().tolist() == lens |
| 258 | assert indices.cpu().tolist() == rows |
| 259 | |
| 260 | |
| 261 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 262 | def test_image_visibility_kernel(): |
| 263 | """The builder's visibility kernel must match get_image_visible.""" |
| 264 | device = torch.device("cuda") |
| 265 | seq_lens = [30, 12] |
| 266 | query_lens = [30, 12] |
| 267 | spans = [[(4, 12), (20, 24)], []] |
| 268 | query_start_loc, seq_lens_t, token_to_req, _, _ = make_batch( |
| 269 | seq_lens, query_lens, device |
| 270 | ) |
| 271 | num_tokens = int(query_start_loc[-1]) |
| 272 | |
| 273 | # CSR span layout: request -> contiguous [start, end) span list. |
| 274 | indptr = [0, 2, 2] |
| 275 | starts = [4, 20] |
| 276 | ends = [12, 24] |
| 277 | left = torch.zeros(num_tokens, dtype=torch.int32, device=device) |
| 278 | right = torch.zeros(num_tokens, dtype=torch.int32, device=device) |
| 279 | _compute_image_visibility_kernel[(num_tokens,)]( |
| 280 | left, |
| 281 | right, |
| 282 | torch.tensor(indptr, dtype=torch.int32, device=device), |
| 283 | torch.tensor(starts, dtype=torch.int32, device=device), |
| 284 | torch.tensor(ends, dtype=torch.int32, device=device), |
| 285 | query_start_loc, |
| 286 | seq_lens_t, |
| 287 | token_to_req, |
| 288 | MAX_IMG, |
| 289 | token_offset=0, |
| 290 | ) |
| 291 | ref_left, ref_right = ref_left_right(seq_lens, query_lens, spans, MAX_IMG) |
| 292 | assert left.cpu().tolist() == ref_left |
| 293 | assert right.cpu().tolist() == ref_right |
| 294 | |
| 295 | |
| 296 | def combine_case( |
| 297 | compress_ratio: int, |
| 298 | topk: int, |
| 299 | seq_lens: list[int], |
| 300 | query_lens: list[int], |
| 301 | spans: list[list[tuple[int, int]]], |
| 302 | with_image: bool, |
| 303 | replay_starts: list[int] | None = None, |
| 304 | combine_fn=combine_topk_swa_indices, |
| 305 | ): |
| 306 | """Run combine_topk_swa_indices and return (indices, lens, expected).""" |
| 307 | device = torch.device("cuda") |
| 308 | num_reqs = len(seq_lens) |
| 309 | replay_starts = replay_starts or [0] * num_reqs |
| 310 | query_start_loc = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device) |
| 311 | query_start_loc[1:] = torch.tensor( |
| 312 | query_lens, dtype=torch.int32, device=device |
| 313 | ).cumsum(0) |
| 314 | num_tokens = int(query_start_loc[-1]) |
| 315 | seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device) |
| 316 | # The builder's gather covers only the context above the replay start. |
| 317 | gather_lens = torch.tensor( |
| 318 | [ |
| 319 | q + min(max(s - q - r, 0), WINDOW - 1) |
| 320 | for s, q, r in zip(seq_lens, query_lens, replay_starts) |
| 321 | ], |
| 322 | dtype=torch.int32, |
| 323 | device=device, |
| 324 | ) |
| 325 | N = (max(seq_lens) + compress_ratio - 1) // compress_ratio |
| 326 | M = N + int(gather_lens.max()) + 8 |
| 327 | gen = torch.Generator(device="cpu").manual_seed(0) |
| 328 | topk_indices = torch.randint( |
| 329 | 0, 4096, (num_tokens, max(topk, 1)), generator=gen, dtype=torch.int32 |
| 330 | ).to(device) |
| 331 | topk_indices = topk_indices[:, : max(topk, 1)] |
| 332 | |
| 333 | if with_image: |
| 334 | lefts, rights = ref_left_right(seq_lens, query_lens, spans, MAX_IMG) |
| 335 | left_t = torch.tensor(lefts, dtype=torch.int32, device=device) |
| 336 | right_t = torch.tensor(rights, dtype=torch.int32, device=device) |
| 337 | else: |
| 338 | left_t = right_t = None |
| 339 | |
| 340 | # V4.1's combine has no image windows (causal SWA only). |
| 341 | has_image = "left_visible" in inspect.signature(combine_fn).parameters |
| 342 | image_kwargs = ( |
| 343 | dict(left_visible=left_t, right_visible=right_t, max_image_tokens=MAX_IMG) |
| 344 | if has_image |
| 345 | else {} |
| 346 | ) |
| 347 | combined_indices, combined_lens = combine_fn( |
| 348 | topk_indices, |
| 349 | query_start_loc, |
| 350 | seq_lens_t, |
| 351 | gather_lens, |
| 352 | WINDOW, |
| 353 | compress_ratio, |
| 354 | topk, |
| 355 | M, |
| 356 | N, |
| 357 | **image_kwargs, |
| 358 | ) |
| 359 | |
| 360 | # Reference rows. |
| 361 | lefts, rights = ref_left_right( |
| 362 | seq_lens, |
| 363 | query_lens, |
| 364 | spans if with_image else [[] for _ in seq_lens], |
| 365 | MAX_IMG, |
| 366 | ) |
| 367 | topk_cpu = topk_indices.cpu() |
| 368 | width = WINDOW + (MAX_IMG if has_image else 0) |
| 369 | combined_topk = (topk + width + 127) // 128 * 128 |
| 370 | rows = [] |
| 371 | lens = [] |
| 372 | token = 0 |
| 373 | for b, (seq_len, query_len) in enumerate(zip(seq_lens, query_lens)): |
| 374 | prefix_len = seq_len - query_len |
| 375 | gather_start = seq_len - int(gather_lens[b]) |
| 376 | for i in range(query_len): |
| 377 | pos = prefix_len + i |
| 378 | topk_len = min((pos + 1) // compress_ratio, topk) |
| 379 | start, end = ref_swa_bounds(pos, WINDOW, lefts[token], rights[token]) |
| 380 | # The window never reaches below the gathered buffer. |
| 381 | start = max(start, gather_start) |
| 382 | swa_len = end - start |
| 383 | row = [-1] * combined_topk |
| 384 | for j in range(topk_len): |
| 385 | row[j] = int(topk_cpu[token, j]) + M * b |
| 386 | for j in range(swa_len): |
| 387 | row[topk_len + j] = M * b + N + start + j - gather_start |
| 388 | rows.append(row) |
| 389 | lens.append(topk_len + swa_len) |
| 390 | token += 1 |
| 391 | return combined_indices, combined_lens, rows, lens |
| 392 | |
| 393 | |
| 394 | COMBINE_CASES = [ |
| 395 | dict(compress_ratio=1, topk=0), # SWA-only layer |
| 396 | dict(compress_ratio=4, topk=16), # C4A layer |
| 397 | ] |
| 398 | |
| 399 | |
| 400 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 401 | @pytest.mark.parametrize("cfg", COMBINE_CASES) |
| 402 | def test_combine_topk_swa_with_image_spans(cfg): |
| 403 | case = CASES[0] |
| 404 | indices, lens, rows, exp_lens = combine_case( |
| 405 | cfg["compress_ratio"], |
| 406 | cfg["topk"], |
| 407 | case["seq_lens"], |
| 408 | case["query_lens"], |
| 409 | case["spans"], |
| 410 | with_image=True, |
| 411 | ) |
| 412 | assert lens.cpu().tolist() == exp_lens |
| 413 | assert indices.cpu().tolist() == rows |
| 414 | |
| 415 | |
| 416 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 417 | @pytest.mark.parametrize("cfg", COMBINE_CASES) |
| 418 | def test_v41_combine_topk_swa_stops_at_replay_start(cfg): |
| 419 | """SWA bounded replay: the gathered buffer starts at replay_start, so the |
| 420 | window never indexes below it.""" |
| 421 | from vllm.models.deepseek_v41.common.ops.cache_utils import ( |
| 422 | combine_topk_swa_indices as combine_v41, |
| 423 | ) |
| 424 | |
| 425 | # Request 0 replays [16, 40): the windows of its first rows would |
| 426 | # otherwise reach below 16. |
| 427 | indices, lens, rows, exp_lens = combine_case( |
| 428 | cfg["compress_ratio"], |
| 429 | cfg["topk"], |
| 430 | seq_lens=[40, 12], |
| 431 | query_lens=[24, 12], |
| 432 | spans=[[], []], |
| 433 | with_image=False, |
| 434 | replay_starts=[16, 0], |
| 435 | combine_fn=combine_v41, |
| 436 | ) |
| 437 | assert lens.cpu().tolist() == exp_lens |
| 438 | assert indices.cpu().tolist() == rows |
| 439 | |
| 440 | |
| 441 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 442 | @pytest.mark.parametrize("cfg", COMBINE_CASES) |
| 443 | def test_combine_topk_swa_without_image_unchanged(cfg): |
| 444 | """left_visible=None must reproduce the plain causal combined indices.""" |
| 445 | case = CASES[0] |
| 446 | indices, lens, rows, exp_lens = combine_case( |
| 447 | cfg["compress_ratio"], |
| 448 | cfg["topk"], |
| 449 | case["seq_lens"], |
| 450 | case["query_lens"], |
| 451 | case["spans"], |
| 452 | with_image=False, |
| 453 | ) |
| 454 | assert lens.cpu().tolist() == exp_lens |
| 455 | assert indices.cpu().tolist() == rows |
| 456 | |
| 457 | |
| 458 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 459 | def test_flashinfer_mixed_sparse_indices_with_image_spans(): |
| 460 | device = torch.device("cuda") |
| 461 | # 1 decode token (req 0) + two prefill requests (reqs 1, 2). |
| 462 | seq_lens = [20, 30, 9] |
| 463 | query_lens = [1, 30, 9] |
| 464 | spans = [[], [(4, 12)], [(0, 6)]] |
| 465 | query_start_loc = torch.tensor([0, 1, 31, 40], dtype=torch.int32, device=device) |
| 466 | seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device) |
| 467 | token_to_req = torch.tensor( |
| 468 | [0] + [1] * 30 + [2] * 9, dtype=torch.int32, device=device |
| 469 | ) |
| 470 | block_table = torch.arange(3, dtype=torch.int32, device=device).view(3, 1) |
| 471 | decode_swa = torch.tensor( |
| 472 | [[100, 101, 102, 103, 104, 105, 106, 107]], dtype=torch.int32, device=device |
| 473 | ) |
| 474 | prefill_topk = torch.zeros(39, 0, dtype=torch.int32, device=device) |
| 475 | |
| 476 | lefts, rights = ref_left_right(seq_lens, query_lens, spans, MAX_IMG) |
| 477 | left_t = torch.tensor(lefts, dtype=torch.int32, device=device) |
| 478 | right_t = torch.tensor(rights, dtype=torch.int32, device=device) |
| 479 | |
| 480 | sparse_indices, sparse_lens = build_flashinfer_mixed_sparse_indices( |
| 481 | decode_swa_indices=decode_swa, |
| 482 | decode_compressed_indices=None, |
| 483 | decode_compressed_topk_lens=None, |
| 484 | prefill_topk_indices=prefill_topk, |
| 485 | query_start_loc=query_start_loc, |
| 486 | seq_lens=seq_lens_t, |
| 487 | token_to_req_indices=token_to_req, |
| 488 | swa_block_table=block_table, |
| 489 | swa_block_size=BLOCK_SIZE, |
| 490 | compressed_block_table=None, |
| 491 | compressed_block_size=BLOCK_SIZE, |
| 492 | window_size=WINDOW, |
| 493 | compress_ratio=1, |
| 494 | topk=0, |
| 495 | prefill_left_visible=left_t, |
| 496 | prefill_right_visible=right_t, |
| 497 | max_image_tokens=MAX_IMG, |
| 498 | ) |
| 499 | |
| 500 | swa_total = WINDOW + MAX_IMG |
| 501 | assert sparse_indices.shape == (40, swa_total) |
| 502 | assert sparse_lens.cpu().tolist() == [swa_total] * 40 |
| 503 | # Decode row: slots copied, image-extension columns padded with -1. |
| 504 | decode_row = sparse_indices[0].cpu().tolist() |
| 505 | assert decode_row[:WINDOW] == list(range(100, 108)) |
| 506 | assert decode_row[WINDOW:] == [-1] * MAX_IMG |
| 507 | # Prefill rows: paged slot ids over the widened window. |
| 508 | exp_rows, _ = ref_swa_slot_rows( |
| 509 | seq_lens, query_lens, spans, block_table, WINDOW, MAX_IMG, swa_total |
| 510 | ) |
| 511 | actual = sparse_indices[1:].cpu().tolist() |
| 512 | assert actual == exp_rows[1:] |
| 513 | |
| 514 | |
| 515 | def make_builder( |
| 516 | vision: bool, |
| 517 | window: int = WINDOW, |
| 518 | max_image_tokens: int = MAX_IMG, |
| 519 | max_num_batched_tokens: int = 64, |
| 520 | model_name: str = "meta-llama/Meta-Llama-3-8B", |
| 521 | builder_cls: type[ |
| 522 | DeepseekSparseSWAMetadataBuilder |
| 523 | ] = DeepseekSparseSWAMetadataBuilder, |
| 524 | mm_prefix_clamp_sliding_window: bool = True, |
| 525 | bounded_replay: bool = False, |
| 526 | ) -> DeepseekSparseSWAMetadataBuilder: |
| 527 | overrides: dict = {"sliding_window": window} |
| 528 | if vision: |
| 529 | # Emulate the V4 vision config, whose mm_prefix_clamp_sliding_window |
| 530 | # gates the in-kernel SWA widening. V4.1 (causal image tokens) leaves |
| 531 | # it off. |
| 532 | overrides.update( |
| 533 | vision_n_layers=2, |
| 534 | vision_max_n_token=max_image_tokens, |
| 535 | mm_prefix_clamp_sliding_window=mm_prefix_clamp_sliding_window, |
| 536 | ) |
| 537 | vllm_config = create_vllm_config( |
| 538 | model_name=model_name, |
| 539 | max_model_len=4096, |
| 540 | max_num_batched_tokens=max_num_batched_tokens, |
| 541 | max_num_seqs=8, |
| 542 | hf_config_override=overrides, |
| 543 | ) |
| 544 | spec = SlidingWindowMLASpec( |
| 545 | block_size=BLOCK_SIZE, |
| 546 | num_kv_heads=1, |
| 547 | head_size=512, |
| 548 | dtype=torch.bfloat16, |
| 549 | sliding_window=window, |
| 550 | cache_dtype_str="auto", |
| 551 | model_version="deepseek_v4", |
| 552 | bounded_replay=bounded_replay, |
| 553 | ) |
| 554 | return builder_cls( |
| 555 | kv_cache_spec=spec, |
| 556 | layer_names=["layer0"], |
| 557 | vllm_config=vllm_config, |
| 558 | device=torch.device("cuda"), |
| 559 | ) |
| 560 | |
| 561 | |
| 562 | def build_metadata( |
| 563 | builder: DeepseekSparseSWAMetadataBuilder, |
| 564 | seq_lens: list[int], |
| 565 | query_lens: list[int], |
| 566 | mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None, |
| 567 | replay_starts: list[int] | None = None, |
| 568 | ): |
| 569 | device = torch.device("cuda") |
| 570 | query_start_loc, seq_lens_t, _, slot_mapping, block_table = make_batch( |
| 571 | seq_lens, query_lens, device |
| 572 | ) |
| 573 | replay_start = ( |
| 574 | None |
| 575 | if replay_starts is None |
| 576 | else torch.tensor(replay_starts, dtype=torch.int32, device=device) |
| 577 | ) |
| 578 | return builder.build( |
| 579 | 0, |
| 580 | CommonAttentionMetadata( |
| 581 | query_start_loc=query_start_loc, |
| 582 | query_start_loc_cpu=query_start_loc.cpu(), |
| 583 | seq_lens=seq_lens_t, |
| 584 | seq_lens_cpu_upper_bound=seq_lens_t.cpu(), |
| 585 | num_reqs=len(seq_lens), |
| 586 | num_actual_tokens=int(query_start_loc[-1]), |
| 587 | max_query_len=max(query_lens), |
| 588 | max_seq_len=max(seq_lens), |
| 589 | block_table_tensor=block_table, |
| 590 | slot_mapping=slot_mapping, |
| 591 | causal=True, |
| 592 | mm_req_doc_ranges=mm_req_doc_ranges, |
| 593 | ), |
| 594 | replay_start=replay_start, |
| 595 | ) |
| 596 | |
| 597 | |
| 598 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 599 | def test_builder_in_image_visibility(): |
| 600 | seq_lens = [30, 12] |
| 601 | query_lens = [30, 12] |
| 602 | spans = [[(4, 12), (20, 24)], []] |
| 603 | builder = make_builder(vision=True) |
| 604 | md = build_metadata(builder, seq_lens, query_lens, {0: spans[0], 1: spans[1]}) |
| 605 | assert md.num_prefills == 2 |
| 606 | assert md.num_decode_tokens == 0 |
| 607 | assert md.prefill_swa_indices.shape[-1] == WIDTH |
| 608 | assert md.prefill_left_visible is not None |
| 609 | assert md.prefill_right_visible is not None |
| 610 | |
| 611 | ref_left, ref_right = ref_left_right(seq_lens, query_lens, spans, MAX_IMG) |
| 612 | assert md.prefill_left_visible.cpu().tolist() == ref_left |
| 613 | assert md.prefill_right_visible.cpu().tolist() == ref_right |
| 614 | |
| 615 | _, _, _, _, block_table = make_batch(seq_lens, query_lens, torch.device("cuda")) |
| 616 | rows, lens = ref_swa_slot_rows( |
| 617 | seq_lens, query_lens, spans, block_table, WINDOW, MAX_IMG, WIDTH |
| 618 | ) |
| 619 | assert md.prefill_swa_lens.cpu().tolist() == lens |
| 620 | assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows |
| 621 | |
| 622 | |
| 623 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 624 | def test_builder_no_image_spans_fast_path(): |
| 625 | """Vision model, image-free batch: no visibility tensors, plain window.""" |
| 626 | seq_lens = [30, 12] |
| 627 | query_lens = [30, 12] |
| 628 | builder = make_builder(vision=True) |
| 629 | md = build_metadata(builder, seq_lens, query_lens, {0: [], 1: []}) |
| 630 | assert md.prefill_left_visible is None |
| 631 | assert md.prefill_right_visible is None |
| 632 | |
| 633 | _, _, _, _, block_table = make_batch(seq_lens, query_lens, torch.device("cuda")) |
| 634 | rows, lens = ref_swa_slot_rows( |
| 635 | seq_lens, query_lens, [[], []], block_table, WINDOW, MAX_IMG, WIDTH |
| 636 | ) |
| 637 | assert md.prefill_swa_lens.cpu().tolist() == lens |
| 638 | assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows |
| 639 | |
| 640 | |
| 641 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 642 | def test_builder_text_model_unchanged(): |
| 643 | """Text-only model: buffers stay window-sized and spans are ignored.""" |
| 644 | seq_lens = [30, 12] |
| 645 | query_lens = [30, 12] |
| 646 | builder = make_builder(vision=False) |
| 647 | assert builder.max_image_tokens == 0 |
| 648 | md = build_metadata(builder, seq_lens, query_lens, None) |
| 649 | assert md.prefill_swa_indices.shape[-1] == WINDOW |
| 650 | assert md.prefill_left_visible is None |
| 651 | |
| 652 | _, _, _, _, block_table = make_batch(seq_lens, query_lens, torch.device("cuda")) |
| 653 | rows, lens = ref_swa_slot_rows( |
| 654 | seq_lens, query_lens, [[], []], block_table, WINDOW, MAX_IMG, WINDOW |
| 655 | ) |
| 656 | assert md.prefill_swa_lens.cpu().tolist() == lens |
| 657 | assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows |
| 658 | |
| 659 | |
| 660 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 661 | @pytest.mark.parametrize("compress_ratio", [0, 1, 2]) |
| 662 | @pytest.mark.parametrize("query_len", [100, 4000]) |
| 663 | def test_v41_image_prefill_uses_causal_swa(compress_ratio, query_len): |
| 664 | """V4.1 image tokens use the same causal SWA as text, including in chunks.""" |
| 665 | from vllm.models.deepseek_v41.common.ops.cache_utils import ( |
| 666 | combine_topk_swa_indices as combine_v41, |
| 667 | ) |
| 668 | from vllm.models.deepseek_v41.sparse_mla import DeepseekV41SparseSWAMetadataBuilder |
| 669 | |
| 670 | seq_len, window, max_image_tokens = 4000, 128, 2048 |
| 671 | spans = [(1900, 3947)] |
| 672 | builder = make_builder( |
| 673 | True, |
| 674 | window, |
| 675 | max_image_tokens, |
| 676 | query_len, |
| 677 | builder_cls=DeepseekV41SparseSWAMetadataBuilder, |
| 678 | mm_prefix_clamp_sliding_window=False, |
| 679 | ) |
| 680 | assert builder.max_image_tokens == 0 |
| 681 | md = build_metadata(builder, [seq_len], [query_len], {0: spans}) |
| 682 | assert md.prefill_left_visible is None |
| 683 | assert md.prefill_right_visible is None |
| 684 | assert md.prefill_swa_indices.shape[-1] == window |
| 685 | plan = md.get_prefill_chunk_plan( |
| 686 | compress_ratio, prefill_chunk_size=4, has_compressed=compress_ratio > 0 |
| 687 | ) |
| 688 | [(start, end, n, m)] = plan |
| 689 | gather_len = int(md.prefill_gather_lens[0]) |
| 690 | assert (start, end, m - n) == (0, 1, gather_len) |
| 691 | assert gather_len == query_len + min(seq_len - query_len, window - 1) |
| 692 | indices, lens = combine_v41( |
| 693 | torch.empty(query_len, 0, dtype=torch.int32, device="cuda"), |
| 694 | md.query_start_loc, |
| 695 | md.prefill_seq_lens, |
| 696 | md.prefill_gather_lens, |
| 697 | window, |
| 698 | compress_ratio, |
| 699 | 0, |
| 700 | m, |
| 701 | n, |
| 702 | ) |
| 703 | indices, lens = indices.cpu(), lens.cpu().tolist() |
| 704 | paged_indices = md.prefill_swa_indices.cpu() |
| 705 | paged_lens = md.prefill_swa_lens.cpu().tolist() |
| 706 | for token, pos in enumerate(range(seq_len - query_len, seq_len)): |
| 707 | lo, hi = max(0, pos - window + 1), pos + 1 |
| 708 | row = indices[token, : lens[token]] |
| 709 | assert lens[token] == hi - lo |
| 710 | assert torch.all((row >= n) & (row < m)) |
| 711 | absolute_keys = row - n + seq_len - gather_len |
| 712 | assert absolute_keys.tolist() == list(range(lo, hi)) |
| 713 | assert paged_lens[token] == hi - lo |
| 714 | assert paged_indices[token, 0, : hi - lo].tolist() == list(range(lo, hi)) |
| 715 | |
| 716 | |
| 717 | # SWA bounded replay: every prefill index path clamps the window at replay_start. |
| 718 | |
| 719 | |
| 720 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 721 | def test_builder_replay_start_bounds_prefill_window_and_gather(): |
| 722 | """Paged-direct prefill indices stop at replay_start and the FlashMLA |
| 723 | gather only covers the context the request may see.""" |
| 724 | seq_lens = [40, 12, 30] |
| 725 | query_lens = [24, 12, 10] |
| 726 | replay_starts = [12, 0, 0] |
| 727 | builder = make_builder(vision=False) |
| 728 | md = build_metadata(builder, seq_lens, query_lens, None, replay_starts) |
| 729 | assert md.replay_start is not None |
| 730 | assert md.replay_start.cpu().tolist() == replay_starts |
| 731 | |
| 732 | _, _, _, _, block_table = make_batch(seq_lens, query_lens, torch.device("cuda")) |
| 733 | rows, lens = ref_swa_slot_rows( |
| 734 | seq_lens, |
| 735 | query_lens, |
| 736 | [[], [], []], |
| 737 | block_table, |
| 738 | WINDOW, |
| 739 | MAX_IMG, |
| 740 | WINDOW, |
| 741 | replay_starts=replay_starts, |
| 742 | ) |
| 743 | assert md.prefill_swa_lens.cpu().tolist() == lens |
| 744 | assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows |
| 745 | # Request 0: the first replayed row (pos 16) sees [12, 16]; its last row |
| 746 | # (pos 39) is a full window above the replay start and sees all of it. |
| 747 | assert md.prefill_swa_lens[0].item() == 16 - 12 + 1 |
| 748 | assert md.prefill_swa_lens[23].item() == WINDOW |
| 749 | # gather_len = query_len + min(prefix_len - replay_start, WINDOW - 1). |
| 750 | assert md.prefill_gather_lens.cpu().tolist() == [ |
| 751 | 24 + min(16 - 12, WINDOW - 1), |
| 752 | 12, |
| 753 | 10 + min(20, WINDOW - 1), |
| 754 | ] |
| 755 | |
| 756 | |
| 757 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 758 | def test_builder_replaying_group_builds_without_replay_start(): |
| 759 | """Graph captures build without model-specific kwargs: a replaying group's |
| 760 | metadata then bounds nothing.""" |
| 761 | builder = make_builder(vision=False, bounded_replay=True) |
| 762 | md = build_metadata(builder, [40, 12], [24, 12], None) |
| 763 | assert not md.replay_start.any() |
| 764 | assert md.prefill_swa_lens[0].item() == WINDOW |
| 765 | |
| 766 | |
| 767 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 768 | def test_v41_flashinfer_mixed_sparse_indices_respect_replay_start(): |
| 769 | from vllm.models.deepseek_v41.common.ops.cache_utils import ( |
| 770 | build_flashinfer_mixed_sparse_indices as build_v41, |
| 771 | ) |
| 772 | |
| 773 | device = torch.device("cuda") |
| 774 | # 1 decode token (req 0) + two prefill requests; req 1 resumes at 16. |
| 775 | seq_lens = [20, 40, 9] |
| 776 | query_lens = [1, 24, 9] |
| 777 | replay_starts = [0, 16, 0] |
| 778 | query_start_loc = torch.tensor([0, 1, 25, 34], dtype=torch.int32, device=device) |
| 779 | seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device) |
| 780 | token_to_req = torch.tensor( |
| 781 | [0] + [1] * 24 + [2] * 9, dtype=torch.int32, device=device |
| 782 | ) |
| 783 | block_table = torch.arange(3, dtype=torch.int32, device=device).view(3, 1) |
| 784 | kwargs = dict( |
| 785 | decode_swa_indices=torch.zeros((1, WINDOW), dtype=torch.int32, device=device), |
| 786 | decode_compressed_indices=None, |
| 787 | decode_compressed_topk_lens=None, |
| 788 | prefill_topk_indices=torch.empty((33, 0), dtype=torch.int32, device=device), |
| 789 | query_start_loc=query_start_loc, |
| 790 | seq_lens=seq_lens_t, |
| 791 | token_to_req_indices=token_to_req, |
| 792 | swa_block_table=block_table, |
| 793 | swa_block_size=BLOCK_SIZE, |
| 794 | compressed_block_table=None, |
| 795 | compressed_block_size=BLOCK_SIZE, |
| 796 | window_size=WINDOW, |
| 797 | compress_ratio=1, |
| 798 | topk=0, |
| 799 | ) |
| 800 | plain, _ = build_v41( |
| 801 | replay_start=torch.zeros(3, dtype=torch.int32, device=device), **kwargs |
| 802 | ) |
| 803 | bounded, _ = build_v41( |
| 804 | replay_start=torch.tensor(replay_starts, dtype=torch.int32, device=device), |
| 805 | **kwargs, |
| 806 | ) |
| 807 | rows, _ = ref_swa_slot_rows( |
| 808 | seq_lens, |
| 809 | query_lens, |
| 810 | [[], [], []], |
| 811 | block_table, |
| 812 | WINDOW, |
| 813 | MAX_IMG, |
| 814 | WINDOW, |
| 815 | replay_starts=replay_starts, |
| 816 | ) |
| 817 | # Decode row is copied through; prefill rows follow the bounded window. |
| 818 | assert bounded[0].cpu().tolist() == plain[0].cpu().tolist() |
| 819 | assert bounded[1:].cpu().tolist() == rows[1:] |