@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1dc2d854c1

[Feat][Model] Support encoder-side SWA-bounded replay for DeepSeek-V4.1-Flash (#56227)

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>

Yifan Qiao committed Sep 18, 2026 at 15:01 UTC 1dc2d854c12099f1f08233c886368a8c7baf430f
30 files changed +1067 -38
tests/models/test_deepseek_v41_replay_start.py new
+112
@@ -0,0 +1,112 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""SWA bounded replay in DeepseekV41ModelState: the batch's replay starts reach
4 +the sliding-window builders, and the replayed tokens' slots are padded in the
5 +prefix-cacheable groups only."""
6 +
7 +from dataclasses import replace
8 +from types import SimpleNamespace
9 +from unittest.mock import MagicMock
10 +
11 +import numpy as np
12 +import pytest
13 +import torch
14 +
15 +from vllm.config import CUDAGraphMode
16 +from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState
17 +from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadataBuilder
18 +from vllm.v1.attention.backends.utils import PAD_SLOT_ID
19 +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
20 +from vllm.v1.worker.gpu.model_states.default import DefaultModelState
21 +
22 +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
23 +
24 +DEVICE = torch.device("cuda")
25 +WINDOW = 4
26 +# Group 0 is prefix-cacheable (compressed KV), group 1 is the replayed window.
27 +KV_CACHE_CONFIG = SimpleNamespace(
28 + kv_cache_groups=[
29 + SimpleNamespace(
30 + kv_cache_spec=SimpleNamespace(prefix_cacheable=True, prefix_replay_tokens=0)
31 + ),
32 + SimpleNamespace(
33 + kv_cache_spec=SimpleNamespace(
34 + prefix_cacheable=False, prefix_replay_tokens=WINDOW
35 + )
36 + ),
37 + ]
38 +)
39 +
40 +
41 +@pytest.fixture
42 +def state(monkeypatch):
43 + cfg = MagicMock()
44 + cfg.model_config.enable_prompt_embeds = False
45 + cfg.model_config.uses_mrope = False
46 + cfg.model_config.is_multimodal_model = False
47 + cfg.scheduler_config.max_num_seqs = 8
48 + cfg.scheduler_config.max_num_batched_tokens = 64
49 + seen: dict = {}
50 +
51 + def prepare_attn(self, input_batch, cg_mode, block_tables, slot_mappings, *a, **kw):
52 + seen["slot_mappings"] = slot_mappings.clone()
53 + seen["extra"] = kw["model_specific_attn_metadata"]
54 + return {}
55 +
56 + monkeypatch.setattr(DefaultModelState, "prepare_attn", prepare_attn)
57 + state = DeepseekV41ModelState(
58 + cfg, SimpleNamespace(token_lookback_depth=0), None, DEVICE
59 + )
60 + state.seen = seen # type: ignore[attr-defined]
61 + return state
62 +
63 +
64 +def _batch(query_lens, seq_lens, idx_mapping, is_prefilling):
65 + num_reqs, num_tokens = len(query_lens), sum(query_lens)
66 + batch = InputBatch.make_dummy(num_reqs, num_tokens, InputBuffers(8, 64, DEVICE))
67 + return replace(
68 + batch,
69 + query_start_loc=torch.tensor(
70 + [0, *np.cumsum(query_lens)], dtype=torch.int32, device=DEVICE
71 + ),
72 + idx_mapping_np=np.array(idx_mapping, dtype=np.intp),
73 + is_prefilling_np=np.array(is_prefilling, dtype=np.bool_),
74 + positions=torch.cat(
75 + [torch.arange(s - q, s) for q, s in zip(query_lens, seq_lens)]
76 + ).to(DEVICE),
77 + )
78 +
79 +
80 +def _prepare(state, batch):
81 + slot_mappings = torch.arange(2 * batch.num_tokens, device=DEVICE).view(2, -1)
82 + state.prepare_attn(
83 + batch, CUDAGraphMode.NONE, (), slot_mappings, [], KV_CACHE_CONFIG
84 + )
85 + swa_builder = MagicMock(spec=DeepseekSparseSWAMetadataBuilder)
86 + extra = state.seen["extra"]
87 + replay_start = extra.get_extra_attn_kwargs(swa_builder, batch.num_reqs)[
88 + "replay_start"
89 + ]
90 + assert extra.get_extra_attn_kwargs(MagicMock(), batch.num_reqs) == {}
91 + return replay_start.tolist(), state.seen["slot_mappings"]
92 +
93 +
94 +def test_replayed_tokens_write_only_the_window_group(state):
95 + # Request state 3 resumes a hit at 16: positions 16..19 are replayed.
96 + state.add_request(3, SimpleNamespace(replay_start=16))
97 + state.add_request(1, SimpleNamespace(replay_start=0))
98 + batch = _batch([6, 2], [22, 2], idx_mapping=[3, 1], is_prefilling=[True, True])
99 + replay_start, slots = _prepare(state, batch)
100 + assert replay_start == [16, 0]
101 + padded = [PAD_SLOT_ID] * WINDOW + [4, 5, 6, 7]
102 + assert slots[0].tolist() == padded # compressed KV keeps its cached rows
103 + assert slots[1].tolist() == list(range(8, 16)) # the window is rebuilt
104 +
105 +
106 +def test_only_prefills_carry_a_replay_start(state):
107 + state.add_request(0, SimpleNamespace(replay_start=16))
108 + # The same request, decoding: its rows sit above the hit.
109 + batch = _batch([1], [40], idx_mapping=[0], is_prefilling=[False])
110 + replay_start, slots = _prepare(state, batch)
111 + assert replay_start == [0]
112 + assert slots.tolist() == [[0], [1]]
tests/v1/attention/test_deepseek_v4_swa_visible.py
+174 -10
@@ -11,6 +11,8 @@ 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
@@ -57,10 +59,13 @@ def ref_left_right(
59 return lefts, rights
60
61
60 -def ref_swa_bounds(pos: int, window: int, left: int, right: int) -> tuple[int, int]:
61 - """Reference [start, end) window bounds for one query token."""
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)
63 - start = max(pos - (window - 1) - left_add, 0)
68 + start = max(pos - (window - 1) - left_add, 0, replay_start)
69 return start, pos + right + 1
70
71
@@ -97,12 +102,14 @@ def ref_swa_slot_rows(
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
@@ -110,7 +117,9 @@ def ref_swa_slot_rows(
117 prefix_len = seq_len - query_len
118 for i in range(query_len):
119 pos = prefix_len + i
113 - start, end = ref_swa_bounds(pos, window, lefts[token], rights[token])
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])
@@ -139,6 +148,7 @@ def run_swa_kernel(
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(
@@ -164,6 +174,7 @@ def run_swa_kernel(
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,
@@ -289,18 +300,25 @@ def combine_case(
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(
303 - [q + min(s - q, WINDOW - 1) for s, q in zip(seq_lens, query_lens)],
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 )
@@ -319,7 +337,14 @@ def combine_case(
337 else:
338 left_t = right_t = None
339
322 - combined_indices, combined_lens = combine_topk_swa_indices(
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,
@@ -329,9 +354,7 @@ def combine_case(
354 topk,
355 M,
356 N,
332 - left_visible=left_t,
333 - right_visible=right_t,
334 - max_image_tokens=MAX_IMG,
357 + **image_kwargs,
358 )
359
360 # Reference rows.
@@ -342,7 +365,7 @@ def combine_case(
365 MAX_IMG,
366 )
367 topk_cpu = topk_indices.cpu()
345 - width = WINDOW + MAX_IMG
368 + width = WINDOW + (MAX_IMG if has_image else 0)
369 combined_topk = (topk + width + 127) // 128 * 128
370 rows = []
371 lens = []
@@ -354,6 +377,8 @@ def combine_case(
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):
@@ -388,6 +413,31 @@ def test_combine_topk_swa_with_image_spans(cfg):
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):
@@ -472,6 +522,7 @@ def make_builder(
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:
@@ -498,6 +549,7 @@ def make_builder(
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,
@@ -512,11 +564,17 @@ def build_metadata(
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(
@@ -533,6 +591,7 @@ def build_metadata(
591 causal=True,
592 mm_req_doc_ranges=mm_req_doc_ranges,
593 ),
594 + replay_start=replay_start,
595 )
596
597
@@ -653,3 +712,108 @@ def test_v41_image_prefill_uses_causal_swa(compress_ratio, query_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:]
tests/v1/attention/test_dspark_noncausal_sparse_mla.py
+1
@@ -669,6 +669,7 @@ def test_dsv41_flashinfer_dspark_window_matches_reference(
669 ),
670 block_table=torch.arange(48, device=device, dtype=torch.int32).view(3, -1),
671 block_size=128,
672 + replay_start=torch.zeros(3, device=device, dtype=torch.int32),
673 flashinfer_sparse_index_cache={},
674 max_decode_query_len=5,
675 )
tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
+23
@@ -14,6 +14,7 @@ from vllm.models.deepseek_v41.sparse_mla import (
14 from vllm.v1.attention.backend import CommonAttentionMetadata
15 from vllm.v1.attention.backends.mla.compressor_utils import (
16 CompressedSlotMappingKernel,
17 + get_compressed_slot_mapping,
18 )
19 from vllm.v1.attention.backends.mla.indexer import (
20 BuildPrefillChunkMetadataKernel,
@@ -235,6 +236,28 @@ def test_compressed_slot_mapping_warmup_includes_index_kpool():
236 assert {(key.compress_ratio, key.block_size) for key in keys} == {(32, 2)}
237
238
239 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
240 +def test_compressed_slot_mapping_inherits_padded_token_slots():
241 + """A token whose own slot is padded (SWA bounded replay) closes no
242 + compressed state either."""
243 + device = torch.device("cuda")
244 + query_start_loc = torch.tensor([0, 8], dtype=torch.int32, device=device)
245 + seq_lens = torch.tensor([8], dtype=torch.int32, device=device)
246 + block_table = torch.tensor([[3]], dtype=torch.int32, device=device)
247 + slot_mapping = torch.arange(8, dtype=torch.int64, device=device)
248 + slot_mapping[:4] = -1
249 + compressed = get_compressed_slot_mapping(
250 + 8,
251 + slot_mapping,
252 + query_start_loc,
253 + seq_lens,
254 + block_table,
255 + block_size=4,
256 + compress_ratio=2,
257 + )
258 + assert compressed.tolist() == [-1, -1, -1, -1, -1, 3 * 4 + 2, -1, 3 * 4 + 3]
259 +
260 +
261 def test_index_conversion_warmup_uses_physical_block_stride():
262 config = SimpleNamespace(
263 cache_config=SimpleNamespace(block_size=64),
tests/v1/core/test_contiguous_kv_packing.py
+64
@@ -651,3 +651,67 @@ class TestCompressorRingGroup:
651 scheduler_block_size=128,
652 )
653 assert len(manager.coordinator.single_type_managers) == len(groups)
654 +
655 +
656 +class TestSWABoundedReplayGrouping:
657 + def test_non_cacheable_swa_leaves_hits_to_the_paged_group(self):
658 + """Under SWA bounded replay only the paged MLA group hashes (hash
659 + block = its block) and the coordinator probes one spec group."""
660 + config = _shared_layout_config()
661 + config.cache_config.enable_prefix_caching = True
662 + specs: dict[str, KVCacheSpec] = {
663 + "layers.2.attn": MLAAttentionSpec(
664 + block_size=128,
665 + num_kv_heads=1,
666 + head_size=584,
667 + dtype=torch.uint8,
668 + tokens_per_state=2,
669 + alignment=576,
670 + ),
671 + }
672 + swa = SlidingWindowMLASpec(
673 + block_size=32,
674 + num_kv_heads=1,
675 + head_size=584,
676 + dtype=torch.uint8,
677 + sliding_window=128,
678 + alignment=576,
679 + bounded_replay=True,
680 + )
681 + for layer in range(4):
682 + specs[f"layers.{layer}.attn.swa_cache"] = swa
683 +
684 + groups = get_kv_cache_groups(config, specs)
685 + # The worker reads the replay window off the (packed) group specs to
686 + # arm the window clamp; the wrapper must forward it like cacheability.
687 + assert sorted(
688 + (g.kv_cache_spec.prefix_cacheable, g.kv_cache_spec.prefix_replay_tokens)
689 + for g in groups
690 + ) == [(False, 128)] * 4 + [(True, 0)]
691 + kv_cache_config = get_kv_cache_config_from_groups(
692 + config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups)
693 + )
694 + manager = KVCacheManager(
695 + generate_scheduler_kv_cache_config([kv_cache_config]),
696 + max_model_len=8192,
697 + enable_caching=True,
698 + hash_block_size=128,
699 + scheduler_block_size=128,
700 + )
701 + assert len(manager.coordinator.attention_groups) == 1
702 + assert isinstance(
703 + manager.coordinator.attention_groups[0].spec, MLAAttentionSpec
704 + )
705 +
706 + # Allocating and caching a request goes through every group's manager;
707 + # the 32-token SWA group must skip hashing rather than assert on the
708 + # 128-token hash block.
709 + from tests.v1.core.utils import create_requests
710 +
711 + request = create_requests(num_requests=1, num_tokens=256, block_size=128)[0]
712 + assert manager.allocate_slots(request, 256) is not None
713 + request.num_computed_tokens = 256
714 + manager.cache_blocks(request, 256)
715 + for single in manager.coordinator.single_type_managers:
716 + cached = single.num_cached_block.get(request.request_id, 0)
717 + assert cached == (2 if single.kv_cache_spec.prefix_cacheable else 0)
tests/v1/core/test_prefix_caching.py
+3 -1
@@ -6,7 +6,7 @@ import copy
6 from collections.abc import Callable
7 from dataclasses import replace
8 from math import lcm
9 -from types import SimpleNamespace
9 +from types import MethodType, SimpleNamespace
10 from unittest.mock import MagicMock
11
12 import pytest
@@ -533,7 +533,9 @@ def test_nixl_hisparse_full_block_import_keeps_a_writable_tail(num_tokens):
533 kv_cache_manager=manager,
534 failed_recving_kv_req_ids=set(),
535 finished_recving_kv_req_ids={request.request_id},
536 + prefix_replay_tokens=0,
537 )
538 + scheduler._mark_prefix_replay = MethodType(Scheduler._mark_prefix_replay, scheduler)
539 Scheduler._update_waiting_for_remote_kv(scheduler, request)
540 assert request.num_tokens - request.num_computed_tokens == 1
541 assert manager.allocate_slots(request, num_new_tokens=1) is not None
tests/v1/core/test_prefix_replay.py new
+279
@@ -0,0 +1,279 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""SWA bounded replay: after a prefix hit the scheduler recomputes the tail of
4 +the hit to rebuild the non-cacheable sliding-window group, keeps the hit's
5 +blocks, and hands the worker the replay start."""
6 +
7 +import pytest
8 +import torch
9 +
10 +from tests.v1.kv_connector.unit.utils import MockKVConfig, create_model_runner_output
11 +from vllm.v1.core.sched.scheduler import Scheduler
12 +from vllm.v1.kv_cache_interface import (
13 + FullAttentionSpec,
14 + KVCacheConfig,
15 + KVCacheGroupSpec,
16 + SlidingWindowMLASpec,
17 +)
18 +from vllm.v1.outputs import ModelRunnerOutput
19 +from vllm.v1.request import RequestStatus
20 +from vllm.v1.structured_output import StructuredOutputManager
21 +
22 +from .utils import create_requests, create_scheduler
23 +
24 +BLOCK_SIZE = 16
25 +WINDOW = 32
26 +NUM_PROMPT_TOKENS = 100
27 +# 100 tokens -> 6 full blocks cached -> 96-token hit.
28 +HIT_TOKENS = NUM_PROMPT_TOKENS // BLOCK_SIZE * BLOCK_SIZE
29 +FULL, SWA = 0, 1
30 +
31 +
32 +def _replay_scheduler(
33 + *,
34 + long_prefill_token_threshold: int = 0,
35 + use_kv_connector: MockKVConfig | None = None,
36 +) -> Scheduler:
37 + """A hybrid layout: one prefix-cacheable full-attention group and one
38 + replayed sliding-window group."""
39 + base = create_scheduler(
40 + block_size=BLOCK_SIZE,
41 + enable_prefix_caching=True,
42 + long_prefill_token_threshold=long_prefill_token_threshold,
43 + use_kv_connector=use_kv_connector,
44 + )
45 + vllm_config = base.vllm_config
46 + kv_cache_config = KVCacheConfig(
47 + num_blocks=10000,
48 + kv_cache_tensors=[],
49 + kv_cache_groups=[
50 + KVCacheGroupSpec(
51 + ["full"],
52 + FullAttentionSpec(
53 + block_size=BLOCK_SIZE,
54 + num_kv_heads=1,
55 + head_size=1,
56 + dtype=torch.float32,
57 + ),
58 + ),
59 + KVCacheGroupSpec(
60 + ["swa"],
61 + SlidingWindowMLASpec(
62 + block_size=BLOCK_SIZE,
63 + num_kv_heads=1,
64 + head_size=64,
65 + dtype=torch.bfloat16,
66 + sliding_window=WINDOW,
67 + bounded_replay=True,
68 + ),
69 + ),
70 + ],
71 + )
72 + scheduler = Scheduler(
73 + vllm_config=vllm_config,
74 + kv_cache_config=kv_cache_config,
75 + block_size=BLOCK_SIZE,
76 + log_stats=True,
77 + structured_output_manager=StructuredOutputManager(vllm_config),
78 + )
79 + scheduler.use_v2_model_runner = True
80 + assert scheduler.prefix_replay_tokens == WINDOW
81 + return scheduler
82 +
83 +
84 +def _prefill(scheduler: Scheduler, request) -> None:
85 + scheduler.add_request(request)
86 + out = scheduler.schedule()
87 + assert out.num_scheduled_tokens[request.request_id] == request.num_prompt_tokens
88 + scheduler.update_from_output(out, _step_output(out, [request]))
89 +
90 +
91 +def _step_output(out, requests) -> ModelRunnerOutput:
92 + """Model output for every scheduled request: a token once its prefill is
93 + done (non-zero, so outputs never extend the all-zero shared prompt in the
94 + cache), nothing while it is still prefilling. schedule() has already
95 + advanced num_computed_tokens past this step's chunk."""
96 + scheduled = [r for r in requests if r.request_id in out.num_scheduled_tokens]
97 + return ModelRunnerOutput(
98 + req_ids=[r.request_id for r in scheduled],
99 + req_id_to_index={r.request_id: i for i, r in enumerate(scheduled)},
100 + sampled_token_ids=[
101 + [1000] if r.num_computed_tokens >= r.num_prompt_tokens else []
102 + for r in scheduled
103 + ],
104 + logprobs=None,
105 + prompt_logprobs_dict={},
106 + pooler_output=None,
107 + )
108 +
109 +
110 +def _new_req_data(out, request):
111 + return next(r for r in out.scheduled_new_reqs if r.req_id == request.request_id)
112 +
113 +
114 +def test_hit_replays_window_without_reallocating():
115 + scheduler = _replay_scheduler()
116 + expected_replay = WINDOW
117 + first, second = create_requests(
118 + num_requests=2,
119 + num_tokens=NUM_PROMPT_TOKENS,
120 + same_prompt=True,
121 + block_size=BLOCK_SIZE,
122 + )
123 + _prefill(scheduler, first)
124 + manager = scheduler.kv_cache_manager
125 +
126 + scheduler.add_request(second)
127 + out = scheduler.schedule()
128 +
129 + replay_start = HIT_TOKENS - expected_replay
130 + new_req = _new_req_data(out, second)
131 + assert new_req.num_computed_tokens == replay_start
132 + assert new_req.replay_start == replay_start
133 + assert (
134 + out.num_scheduled_tokens[second.request_id] == NUM_PROMPT_TOKENS - replay_start
135 + )
136 +
137 + first_blocks = manager.get_blocks(first.request_id).blocks
138 + second_blocks = manager.get_blocks(second.request_id).blocks
139 + # The full-attention hit is shared (replayed tokens included); only the
140 + # tail block is new.
141 + num_hit_blocks = HIT_TOKENS // BLOCK_SIZE
142 + assert second_blocks[FULL][:num_hit_blocks] == first_blocks[FULL][:num_hit_blocks]
143 + assert all(b.ref_cnt == 2 for b in second_blocks[FULL][:num_hit_blocks])
144 + assert len(second_blocks[FULL]) == num_hit_blocks + 1
145 + # The window group never hits: blocks below the window are null, the
146 + # window (replayed and new tokens) gets fresh blocks.
147 + num_skipped_blocks = (HIT_TOKENS - WINDOW + 1) // BLOCK_SIZE
148 + assert all(b.is_null for b in second_blocks[SWA][:num_skipped_blocks])
149 + num_window_blocks = len(second_blocks[SWA]) - num_skipped_blocks
150 + assert num_window_blocks == -(-NUM_PROMPT_TOKENS // BLOCK_SIZE) - num_skipped_blocks
151 + assert not any(b.is_null for b in second_blocks[SWA][num_skipped_blocks:])
152 +
153 + scheduler.update_from_output(out, _step_output(out, [first, second]))
154 + assert second.num_computed_tokens == NUM_PROMPT_TOKENS
155 +
156 +
157 +def test_replay_only_chunks_make_progress():
158 + """A chunk cap smaller than the replay window schedules replay-only
159 + chunks; they advance the request instead of stalling it."""
160 + scheduler = _replay_scheduler(long_prefill_token_threshold=BLOCK_SIZE)
161 + first, second = create_requests(
162 + num_requests=2,
163 + num_tokens=NUM_PROMPT_TOKENS,
164 + same_prompt=True,
165 + block_size=BLOCK_SIZE,
166 + )
167 + scheduler.add_request(first)
168 + while first.num_computed_tokens < NUM_PROMPT_TOKENS:
169 + out = scheduler.schedule()
170 + scheduler.update_from_output(out, _step_output(out, [first]))
171 +
172 + scheduler.add_request(second)
173 +
174 + def full_blocks():
175 + return scheduler.kv_cache_manager.get_blocks(second.request_id).blocks[FULL]
176 +
177 + scheduled: list[int] = []
178 + while second.num_computed_tokens < NUM_PROMPT_TOKENS:
179 + out = scheduler.schedule()
180 + if not scheduled:
181 + new_req = _new_req_data(out, second)
182 + assert new_req.num_computed_tokens == HIT_TOKENS - WINDOW
183 + assert new_req.replay_start == HIT_TOKENS - WINDOW
184 + # A replay-only chunk adopts the hit and allocates nothing new.
185 + assert len(full_blocks()) == HIT_TOKENS // BLOCK_SIZE
186 + scheduled.append(out.num_scheduled_tokens[second.request_id])
187 + scheduler.update_from_output(out, _step_output(out, [first, second]))
188 + # 32 replayed + 4 new tokens in 16-token chunks: two replay-only chunks.
189 + assert scheduled == [BLOCK_SIZE, BLOCK_SIZE, NUM_PROMPT_TOKENS - HIT_TOKENS]
190 + assert len(full_blocks()) == -(-NUM_PROMPT_TOKENS // BLOCK_SIZE)
191 + assert second.status == RequestStatus.RUNNING
192 +
193 +
194 +def test_async_remote_kv_hit_replays_after_load():
195 + """A KV-connector hit loaded asynchronously carries no window state
196 + either; the replay is applied when the load completes."""
197 + matched = 64
198 + scheduler = _replay_scheduler(
199 + use_kv_connector=MockKVConfig(matched_tokens=matched, is_async=True)
200 + )
201 + request = create_requests(
202 + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE
203 + )[0]
204 + scheduler.add_request(request)
205 + out = scheduler.schedule()
206 + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
207 + assert request.request_id not in out.num_scheduled_tokens
208 + scheduler.update_from_output(
209 + out, create_model_runner_output([], finished_recving={request.request_id})
210 + )
211 +
212 + out = scheduler.schedule()
213 + new_req = _new_req_data(out, request)
214 + assert new_req.num_computed_tokens == matched - WINDOW
215 + assert new_req.replay_start == matched - WINDOW
216 + assert (
217 + out.num_scheduled_tokens[request.request_id]
218 + == NUM_PROMPT_TOKENS - matched + WINDOW
219 + )
220 +
221 +
222 +def test_remote_kv_hit_is_taken_in_whole_blocks():
223 + """A hit ending one token short of a block boundary would put the replay
224 + window's first token in a block the sliding-window group retires, so a
225 + connector hit is cut back to whole blocks."""
226 + matched = 4 * BLOCK_SIZE - 1
227 + scheduler = _replay_scheduler(
228 + use_kv_connector=MockKVConfig(matched_tokens=matched, is_async=True)
229 + )
230 + request = create_requests(
231 + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE
232 + )[0]
233 + scheduler.add_request(request)
234 + out = scheduler.schedule()
235 + scheduler.update_from_output(
236 + out, create_model_runner_output([], finished_recving={request.request_id})
237 + )
238 + out = scheduler.schedule()
239 + new_req = _new_req_data(out, request)
240 + hit = 3 * BLOCK_SIZE
241 + assert new_req.replay_start == hit - WINDOW
242 + assert new_req.num_computed_tokens == hit - WINDOW
243 + swa_manager = scheduler.kv_cache_manager.coordinator.single_type_managers[SWA]
244 + swa_blocks = scheduler.kv_cache_manager.get_blocks(request.request_id).blocks[SWA]
245 + assert swa_blocks[new_req.replay_start // BLOCK_SIZE] is not swa_manager._null_block
246 +
247 +
248 +@pytest.mark.parametrize("via_connector", [False, True])
249 +def test_hit_no_longer_than_window_is_ignored(via_connector):
250 + """Such a hit would be recomputed in full anyway. It is not adopted (an
251 + async load of it is not even started), so a request has computed tokens
252 + iff it replays."""
253 + if via_connector:
254 + scheduler = _replay_scheduler(
255 + use_kv_connector=MockKVConfig(matched_tokens=WINDOW, is_async=True)
256 + )
257 + request = create_requests(
258 + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE
259 + )[0]
260 + else:
261 + scheduler = _replay_scheduler()
262 + first = create_requests(
263 + num_requests=1, num_tokens=WINDOW, same_prompt=True, block_size=BLOCK_SIZE
264 + )[0]
265 + _prefill(scheduler, first)
266 + request = create_requests(
267 + num_requests=1,
268 + num_tokens=NUM_PROMPT_TOKENS,
269 + same_prompt=True,
270 + block_size=BLOCK_SIZE,
271 + req_ids=["long"],
272 + )[0]
273 + scheduler.add_request(request)
274 + out = scheduler.schedule()
275 + new_req = _new_req_data(out, request)
276 + assert request.status == RequestStatus.RUNNING
277 + assert new_req.num_computed_tokens == 0
278 + assert new_req.replay_start == 0
279 + assert out.num_scheduled_tokens[request.request_id] == NUM_PROMPT_TOKENS
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+1
@@ -1505,6 +1505,7 @@ def test_failed_load_rezeroes_unwritten_skipped_blocks():
1505 scheduler.kv_cache_manager.cache_blocks = MagicMock()
1506 scheduler.failed_recving_kv_req_ids = {"req-1"}
1507 scheduler.finished_recving_kv_req_ids = {"req-1"}
1508 + scheduler.prefix_replay_tokens = 0
1509
1510 request = MagicMock()
1511 request.request_id = "req-1"
vllm/config/cache.py
+7 -1
@@ -227,6 +227,11 @@ class CacheConfig:
227 necessary for implementing this optimization in some models (e.g. Gemma3n)
228 """
229
230 + swa_bounded_replay: bool = True
231 + """Keep the sliding-window KV of models that support it (DeepSeek-V4.1)
232 + out of prefix caching and rebuild it after a prefix hit by recomputing the
233 + hit's last window. Requires model runner V2."""
234 +
235 kv_cache_memory_bytes: int | None = None
236 """Size of KV Cache per GPU in bytes. By default, this is set to None
237 and vllm can automatically infer the kv cache size based on
@@ -281,8 +286,9 @@ class CacheConfig:
286 "effective_attention_block_size",
287 "kv_cache_size_tokens",
288 "kv_cache_max_concurrency",
284 - # WIP feature toggle not impacting compiled graph shape
289 + # Feature toggles not impacting compiled graph shape
290 "kv_sharing_fast_prefill",
291 + "swa_bounded_replay",
292 }
293
294 from vllm.config.utils import get_hash_factors, hash_factors
vllm/engine/arg_utils.py
+5
@@ -781,6 +781,7 @@ class EngineArgs:
781 stream_interval: int = SchedulerConfig.stream_interval
782
783 kv_sharing_fast_prefill: bool = CacheConfig.kv_sharing_fast_prefill
784 + swa_bounded_replay: bool = CacheConfig.swa_bounded_replay
785 optimization_level: OptimizationLevel = VllmConfig.optimization_level
786 performance_mode: PerformanceMode = VllmConfig.performance_mode
787
@@ -1310,6 +1311,9 @@ class EngineArgs:
1311 cache_group.add_argument(
1312 "--kv-sharing-fast-prefill", **cache_kwargs["kv_sharing_fast_prefill"]
1313 )
1314 + cache_group.add_argument(
1315 + "--swa-bounded-replay", **cache_kwargs["swa_bounded_replay"]
1316 + )
1317 cache_group.add_argument(
1318 "--mamba-cache-dtype", **cache_kwargs["mamba_cache_dtype"]
1319 )
@@ -2118,6 +2122,7 @@ class EngineArgs:
2122 prefix_cache_retention_interval=self.prefix_cache_retention_interval,
2123 kv_cache_dtype_skip_layers=self.kv_cache_dtype_skip_layers,
2124 kv_sharing_fast_prefill=self.kv_sharing_fast_prefill,
2125 + swa_bounded_replay=self.swa_bounded_replay,
2126 mamba_cache_dtype=self.mamba_cache_dtype,
2127 mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype,
2128 mamba_block_size=self.mamba_block_size,
vllm/models/deepseek_v4/amd/rocm.py
+2
@@ -495,11 +495,13 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild
495 common_prefix_len: int,
496 common_attn_metadata: CommonAttentionMetadata,
497 fast_build: bool = False,
498 + replay_start: torch.Tensor | None = None,
499 ) -> DeepseekV4ROCMAiterSparseSWAMetadata:
500 base = super().build(
501 common_prefix_len=common_prefix_len,
502 common_attn_metadata=common_attn_metadata,
503 fast_build=fast_build,
504 + replay_start=replay_start,
505 )
506
507 ragged_indices = None
vllm/models/deepseek_v4/sparse_mla.py
+1
@@ -195,6 +195,7 @@ class DeepseekV4SparseMLAMetadataBuilder(
195 if self.compress_ratio > 1:
196 slot_mapping = get_compressed_slot_mapping(
197 cm.num_actual_tokens,
198 + cm.slot_mapping,
199 cm.query_start_loc,
200 cm.seq_lens,
201 cm.block_table_tensor.clamp_(min=0),
vllm/models/deepseek_v41/amd/rocm.py
+2
@@ -425,11 +425,13 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu
425 common_prefix_len: int,
426 common_attn_metadata: CommonAttentionMetadata,
427 fast_build: bool = False,
428 + replay_start: torch.Tensor | None = None,
429 ) -> DeepseekV4ROCMAiterSparseSWAMetadata:
430 base = super().build(
431 common_prefix_len=common_prefix_len,
432 common_attn_metadata=common_attn_metadata,
433 fast_build=fast_build,
434 + replay_start=replay_start,
435 )
436
437 ragged_indices = None
vllm/models/deepseek_v41/attention.py
+20
@@ -512,6 +512,25 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
512 "decodes only on SM100."
513 )
514
515 + swa_bounded_replay = cache_config.swa_bounded_replay
516 + if swa_bounded_replay and not vllm_config.use_v2_model_runner:
517 + logger.warning_once(
518 + "SWA bounded replay needs model runner V2 (only it skips the "
519 + "paged-KV writes of replayed tokens); the sliding-window cache "
520 + "takes part in prefix caching instead."
521 + )
522 + swa_bounded_replay = False
523 + if (
524 + swa_bounded_replay
525 + and vllm_config.parallel_config.prefill_context_parallel_size > 1
526 + ):
527 + logger.warning_once(
528 + "SWA bounded replay is off under prefill context parallelism "
529 + "(the replayed tokens' slot padding knows the rank-local batch "
530 + "only); the sliding-window cache takes part in prefix caching "
531 + "instead."
532 + )
533 + swa_bounded_replay = False
534 self.swa_cache_layer = DeepseekV4SWACache(
535 head_dim=self.head_dim,
536 window_size=self.window_size,
@@ -522,6 +541,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
541 block_size=32,
542 packed_bytes_per_token=self.swa_bytes_per_token,
543 packed_page_alignment=self.kv_page_alignment,
544 + bounded_replay=swa_bounded_replay,
545 )
546
547 # The attention layer itself was already registered with the
vllm/models/deepseek_v41/common/ops/cache_utils.py
+13
@@ -982,6 +982,10 @@ class CombineTopkSwaIndicesKernel(
982 # turns the store index below negative).
983 topk_len = 0
984 swa_start = tl.maximum(pos - (WINDOW_SIZE - 1), 0)
985 + # gather_len already excludes context below the request's replay
986 + # start (SWA bounded replay), so the window cannot start before
987 + # the gathered buffer does.
988 + swa_start = tl.maximum(swa_start, gather_start)
989 swa_len = pos - swa_start + 1
990
991 offset = tl.arange(0, PADDED_TOP_K)
@@ -1128,6 +1132,8 @@ def build_flashinfer_mixed_sparse_indices(
1132 decode_is_valid_token: torch.Tensor | None = None,
1133 swa_block_span: int | None = None,
1134 compressed_block_span: int | None = None,
1135 + *,
1136 + replay_start: torch.Tensor,
1137 ) -> tuple[torch.Tensor, torch.Tensor]:
1138 """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches.
1139
@@ -1137,6 +1143,9 @@ def build_flashinfer_mixed_sparse_indices(
1143 per token). Decode tokens read precomputed SWA/compressed indices; prefill
1144 tokens derive their SWA window from the position and translate local
1145 compressed indices to global slots via the block tables.
1146 +
1147 + ``replay_start`` ([num_reqs], SWA bounded replay) lower-bounds every
1148 + prefill token's window: positions below it hold no window KV.
1149 """
1150 assert decode_swa_indices.dtype == torch.int32
1151 assert decode_swa_indices.dim() == 2
@@ -1227,6 +1236,7 @@ def build_flashinfer_mixed_sparse_indices(
1236 query_start_loc,
1237 seq_lens,
1238 token_to_req_indices,
1239 + replay_start,
1240 swa_block_table,
1241 swa_block_table.stride(0),
1242 swa_block_size,
@@ -1295,6 +1305,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
1305 query_start_loc_ptr,
1306 seq_lens_ptr,
1307 token_to_req_indices_ptr,
1308 + replay_start_ptr,
1309 swa_block_table_ptr,
1310 swa_block_table_stride,
1311 swa_block_size,
@@ -1394,6 +1405,8 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
1405 token_idx_in_query = token_idx - query_start
1406 pos = start_pos + token_idx_in_query
1407 swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1), 0)
1408 + # SWA bounded replay: no window KV exists below the request's replay start.
1409 + swa_start_pos = tl.maximum(swa_start_pos, tl.load(replay_start_ptr + req_idx))
1410 swa_len = pos - swa_start_pos + 1
1411 if COMPRESS_RATIO > 0:
1412 topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py
+6 -1
@@ -200,8 +200,11 @@ class DeepseekSparseSWAFlashInferMetadataBuilder(DeepseekV41SparseSWAMetadataBui
200 common_prefix_len: int,
201 common_attn_metadata: CommonAttentionMetadata,
202 fast_build: bool = False,
203 + replay_start: torch.Tensor | None = None,
204 ) -> "DeepseekSparseSWAMetadata":
204 - metadata = super().build(common_prefix_len, common_attn_metadata, fast_build)
205 + metadata = super().build(
206 + common_prefix_len, common_attn_metadata, fast_build, replay_start
207 + )
208 num_tokens = metadata.num_decode_tokens
209 if not common_attn_metadata.causal and num_tokens > 0:
210 assert metadata.decode_swa_lens is not None
@@ -375,6 +378,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
378 assert swa_metadata.token_to_req_indices is not None
379 assert swa_metadata.decode_swa_indices is not None
380 assert swa_metadata.block_table is not None
381 + assert swa_metadata.replay_start is not None
382
383 decode_swa_indices = swa_metadata.decode_swa_indices.reshape(
384 num_decode_tokens, swa_metadata.decode_swa_width
@@ -454,6 +458,7 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
458 decode_is_valid_token=decode_is_valid_token,
459 swa_block_span=swa_block_span,
460 compressed_block_span=compressed_block_span,
461 + replay_start=swa_metadata.replay_start[:num_reqs],
462 )
463 if swa_only:
464 swa_metadata.flashinfer_sparse_index_cache["swa_only"] = (
vllm/models/deepseek_v41/nvidia/model_state.py
+143 -2
@@ -2,15 +2,23 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from typing import Any
4
5 +import numpy as np
6 import torch
7 import torch.nn as nn
8
8 -from vllm.config import VllmConfig
9 +from vllm.config import CUDAGraphMode, VllmConfig
10 from vllm.triton_utils import tl, triton
11 +from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadataBuilder
12 +from vllm.v1.attention.backends.utils import PAD_SLOT_ID
13 +from vllm.v1.core.sched.output import NewRequestData
14 +from vllm.v1.kv_cache_interface import KVCacheConfig
15 +from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool
16 from vllm.v1.worker.gpu.input_batch import InputBatch
17 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
18 from vllm.v1.worker.gpu.model_states.default import DefaultModelState
19 +from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
20 from vllm.v1.worker.gpu.states import RequestState
21 +from vllm.v1.worker.utils import AttentionGroup
22
23
24 @triton.jit
@@ -41,14 +49,69 @@ def _gather_lookback_kernel(
49 tl.store(lookback_ptr + batch_idx * DEPTH + offs, ids, mask=offs < DEPTH)
50
51
52 +@triton.jit
53 +def _pad_replayed_slots_kernel(
54 + slot_mappings_ptr, # [num_groups, num_tokens_padded]
55 + group_stride,
56 + cacheable_groups_ptr, # [NUM_GROUPS] indices of the prefix-cacheable groups
57 + query_start_loc_ptr, # [num_reqs + 1]
58 + positions_ptr, # [num_tokens]
59 + replay_start_ptr, # [num_reqs]
60 + window,
61 + pad_slot_id,
62 + NUM_GROUPS: tl.constexpr,
63 + BLOCK: tl.constexpr,
64 +):
65 + req = tl.program_id(0)
66 + start = tl.load(replay_start_ptr + req)
67 + if start == 0:
68 + return
69 + begin = tl.load(query_start_loc_ptr + req)
70 + end = tl.load(query_start_loc_ptr + req + 1)
71 + for tok in range(begin, end, BLOCK):
72 + offs = tok + tl.arange(0, BLOCK)
73 + pos = tl.load(positions_ptr + offs, mask=offs < end, other=0)
74 + replayed = (offs < end) & (pos >= start) & (pos < start + window)
75 + for i in tl.static_range(NUM_GROUPS):
76 + group = tl.load(cacheable_groups_ptr + i)
77 + tl.store(
78 + slot_mappings_ptr + group * group_stride + offs,
79 + pad_slot_id,
80 + mask=replayed,
81 + )
82 +
83 +
84 +class ReplayAttnMetadata(ModelSpecificAttnMetadata):
85 + """Hands the batch's replay starts to the sliding-window builders."""
86 +
87 + def __init__(self, replay_start: torch.Tensor) -> None:
88 + self.replay_start = replay_start
89 +
90 + def get_extra_attn_kwargs(
91 + self, attn_metadata_builder: Any, num_reqs: int
92 + ) -> dict[str, Any]:
93 + if isinstance(attn_metadata_builder, DeepseekSparseSWAMetadataBuilder):
94 + return {"replay_start": self.replay_start[:num_reqs]}
95 + return {}
96 +
97 +
98 class DeepseekV41ModelState(DefaultModelState):
45 - """DefaultModelState plus the engram lookback window.
99 + """DefaultModelState plus the engram lookback window and SWA bounded replay.
100
101 The engram n-gram hash needs the ids of the ``depth`` tokens preceding
102 each request's chunk start (see ``common/engram.py``). The runner keeps
103 the full token history on device, so the window is gathered there every
104 step: exact for prompt and generated tokens alike, whatever instance
105 produced their KV.
106 +
107 + SWA bounded replay (``CacheConfig.swa_bounded_replay``) keeps the
108 + sliding-window KV out of prefix caching and rebuilds it after a prefix hit
109 + by recomputing the hit's last window; the scheduler tells each request the
110 + position from which it holds window KV (``NewRequestData.replay_start``).
111 + ``prepare_attn`` gathers those per batch, pads the replayed tokens' slots
112 + in the prefix-cacheable groups so the cached KV stays as is, and hands the
113 + starts to the sliding-window metadata builders, whose kernels read no
114 + window KV below them.
115 """
116
117 def __init__(
@@ -67,6 +130,20 @@ class DeepseekV41ModelState(DefaultModelState):
130 (self.max_num_reqs, depth), -1, dtype=torch.int32, device=device
131 )
132
133 + # Per request state index; batches gather from it (see prepare_attn).
134 + self._replay_start_np = np.zeros(self.max_num_reqs, dtype=np.int32)
135 + self._replay_start = torch.zeros(
136 + self.max_num_reqs, dtype=torch.int32, device=device
137 + )
138 + self._replay_start_staging = UvaBufferPool(self.max_num_reqs, torch.int32)
139 + # The replay window and the indices of the prefix-cacheable groups, from
140 + # the KV cache config on first use; None until then.
141 + self._replay: tuple[int, torch.Tensor] | None = None
142 +
143 + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
144 + super().add_request(req_index, new_req_data)
145 + self._replay_start_np[req_index] = new_req_data.replay_start
146 +
147 def prepare_inputs(
148 self, input_batch: InputBatch, req_states: RequestState
149 ) -> dict[str, torch.Tensor | None]:
@@ -96,3 +173,67 @@ class DeepseekV41ModelState(DefaultModelState):
173 self.lookback_token_ids.fill_(-1)
174 model_inputs["lookback_token_ids"] = self.lookback_token_ids
175 return model_inputs
176 +
177 + def prepare_attn(
178 + self,
179 + input_batch: InputBatch,
180 + cudagraph_mode: CUDAGraphMode,
181 + block_tables: tuple[torch.Tensor, ...],
182 + slot_mappings: torch.Tensor,
183 + attn_groups: list[list[AttentionGroup]],
184 + kv_cache_config: KVCacheConfig,
185 + for_capture: bool = False,
186 + ubatch_idx: int = 0,
187 + model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
188 + ) -> dict[str, Any]:
189 + if self._replay is None:
190 + specs = [group.kv_cache_spec for group in kv_cache_config.kv_cache_groups]
191 + self._replay = (
192 + max(spec.prefix_replay_tokens for spec in specs),
193 + torch.tensor(
194 + [i for i, spec in enumerate(specs) if spec.prefix_cacheable],
195 + dtype=torch.int32,
196 + device=self.device,
197 + ),
198 + )
199 + window, cacheable_groups = self._replay
200 + if window:
201 + num_reqs = input_batch.num_reqs
202 + # Decode rows sit above the hit, so only prefills carry a replay
203 + # start; dummy batches (captures, profiling) carry none.
204 + replay_start_np = np.where(
205 + input_batch.is_prefilling_np[:num_reqs],
206 + self._replay_start_np[input_batch.idx_mapping_np[:num_reqs]],
207 + 0,
208 + ).astype(np.int32)
209 + replay_start = self._replay_start_staging.copy_to_gpu(
210 + replay_start_np, out=self._replay_start[:num_reqs]
211 + )
212 + if replay_start_np.any():
213 + # The replayed tokens rebuild window KV only: their slots in the
214 + # prefix-cacheable groups are padded so the cached KV stays as is.
215 + _pad_replayed_slots_kernel[(num_reqs,)](
216 + slot_mappings,
217 + slot_mappings.stride(0),
218 + cacheable_groups,
219 + input_batch.query_start_loc,
220 + input_batch.positions,
221 + replay_start,
222 + window,
223 + PAD_SLOT_ID,
224 + NUM_GROUPS=cacheable_groups.numel(),
225 + BLOCK=1024,
226 + )
227 + assert model_specific_attn_metadata is None
228 + model_specific_attn_metadata = ReplayAttnMetadata(replay_start)
229 + return super().prepare_attn(
230 + input_batch,
231 + cudagraph_mode,
232 + block_tables,
233 + slot_mappings,
234 + attn_groups,
235 + kv_cache_config,
236 + for_capture=for_capture,
237 + ubatch_idx=ubatch_idx,
238 + model_specific_attn_metadata=model_specific_attn_metadata,
239 + )
vllm/models/deepseek_v41/sparse_mla.py
+1
@@ -195,6 +195,7 @@ class DeepseekV4SparseMLAMetadataBuilder(
195 if self.compress_ratio > 1:
196 slot_mapping = get_compressed_slot_mapping(
197 cm.num_actual_tokens,
198 + cm.slot_mapping,
199 cm.query_start_loc,
200 cm.seq_lens,
201 cm.block_table_tensor.clamp_(min=0),
vllm/v1/attention/backends/mla/compressor_utils.py
+26 -5
@@ -42,6 +42,8 @@ class CompressedSlotMappingKernel(
42 @triton.jit(do_not_specialize=["block_table_stride"])
43 def kernel(
44 # [num_tokens]
45 + compressed_slot_mapping_ptr,
46 + # [num_tokens] the tokens' own slots
47 slot_mapping_ptr,
48 # [num_reqs + 1]
49 query_start_loc_ptr,
@@ -69,7 +71,12 @@ class CompressedSlotMappingKernel(
71 mask = offset < query_len
72
73 pos = start_pos + i + tl.arange(0, TRITON_BLOCK_SIZE)
72 - is_valid = (pos + 1) % COMPRESS_RATIO == 0
74 + # A replayed token has a PAD slot (SWA bounded replay: its KV is
75 + # cached already); its compressed KV is cached too, so PAD here.
76 + slot = tl.load(
77 + slot_mapping_ptr + query_start + offset, mask=mask, other=PAD_ID
78 + )
79 + is_valid = ((pos + 1) % COMPRESS_RATIO == 0) & (slot != PAD_ID)
80 pos_after_compress = pos // COMPRESS_RATIO
81
82 block_ids = pos_after_compress // block_size
@@ -81,7 +88,9 @@ class CompressedSlotMappingKernel(
88
89 # NOTE
90 slot_ids = tl.where(is_valid, slot_ids, PAD_ID)
84 - tl.store(slot_mapping_ptr + query_start + offset, slot_ids, mask=mask)
91 + tl.store(
92 + compressed_slot_mapping_ptr + query_start + offset, slot_ids, mask=mask
93 + )
94
95 def dispatch( # type: ignore[override]
96 self,
@@ -121,6 +130,7 @@ class CompressedSlotMappingKernel(
130 def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
131 int32_ptr = TritonWarmupTensor(torch.int32)
132 return dict(
133 + compressed_slot_mapping=TritonWarmupTensor(torch.int64),
134 slot_mapping=TritonWarmupTensor(torch.int64),
135 query_start_loc=int32_ptr,
136 seq_lens=int32_ptr,
@@ -132,6 +142,7 @@ class CompressedSlotMappingKernel(
142 @kernel_launcher
143 def __call__(
144 self,
145 + compressed_slot_mapping: torch.Tensor,
146 slot_mapping: torch.Tensor,
147 query_start_loc: torch.Tensor,
148 seq_lens: torch.Tensor,
@@ -149,6 +160,7 @@ class CompressedSlotMappingKernel(
160
161 def get_compressed_slot_mapping(
162 num_tokens: int,
163 + slot_mapping: torch.Tensor,
164 query_start_loc: torch.Tensor,
165 seq_lens: torch.Tensor,
166 block_table: torch.Tensor,
@@ -156,19 +168,28 @@ def get_compressed_slot_mapping(
168 compress_ratio: int,
169 out: torch.Tensor | None = None,
170 ) -> torch.Tensor:
171 + """Slot mapping for writing the compressed states of ``num_tokens`` tokens.
172 +
173 + Every ``compress_ratio`` tokens share one compressed state, written by the
174 + last of them: that token maps to the state's slot, the others to PAD. A
175 + token whose own ``slot_mapping`` entry is PAD maps to PAD too: SWA bounded
176 + replay recomputes tokens whose KV is cached already, and their compressed
177 + states must not be rewritten either.
178 + """
179 if out is not None:
180 # Guard: for padded / invalid sequences.
181 # Negative positions produce bogus block indices that lead to illegal memory
182 # accesses inside the block_table load.
183 # NOTE: Fill -1 to the whole tensor, not just the first `num_tokens`.
184 out.fill_(-1)
165 - slot_mapping = out[:num_tokens]
185 + compressed_slot_mapping = out[:num_tokens]
186 else:
167 - slot_mapping = torch.full(
187 + compressed_slot_mapping = torch.full(
188 (num_tokens,), -1, dtype=torch.int64, device=query_start_loc.device
189 )
190
191 _COMPRESSED_SLOT_MAPPING_KERNEL(
192 + compressed_slot_mapping,
193 slot_mapping,
194 query_start_loc,
195 seq_lens,
@@ -176,7 +197,7 @@ def get_compressed_slot_mapping(
197 block_size,
198 compress_ratio,
199 )
179 - return slot_mapping
200 + return compressed_slot_mapping
201
202
203 _COMPRESSED_SLOT_MAPPING_KERNEL = CompressedSlotMappingKernel()
vllm/v1/attention/backends/mla/indexer.py
+10 -1
@@ -795,6 +795,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
795 self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
796 self.pcp_world_size = parallel_config.prefill_context_parallel_size
797 self.use_pcp = self.pcp_world_size > 1
798 + self.pcp_rank = get_pcp_group().rank_in_group if self.use_pcp else 0
799 self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
800 # NOTE(Chen):an estimated max size of flattened_kv. Need to double check.
801 self.max_prefill_buffer_size = get_max_prefill_buffer_size(self.vllm_config)
@@ -1217,10 +1218,18 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1218 factor = self.kv_cache_spec.block_size // kernel_block_size
1219 indexer_block_table = (block_table[:, ::factor] // factor).contiguous()
1220 padded_num_tokens = num_tokens
1220 - if self.pcp_world_size > 1:
1221 + local_slot_mapping = slot_mapping
1222 + if self.use_pcp:
1223 + # The gathered layout holds each rank's local tokens, padded, in
1224 + # rank order, so this rank's segment lines up with query_start_loc.
1225 padded_num_tokens = slot_mapping.shape[0] // self.pcp_world_size
1226 + local_slot_mapping = slot_mapping[
1227 + self.pcp_rank * padded_num_tokens : (self.pcp_rank + 1)
1228 + * padded_num_tokens
1229 + ]
1230 compressed_slot_mapping = get_compressed_slot_mapping(
1231 num_tokens,
1232 + local_slot_mapping,
1233 query_start_loc,
1234 seq_lens,
1235 indexer_block_table,
vllm/v1/attention/backends/mla/sparse_swa.py
+50 -5
@@ -81,9 +81,12 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
81 block_size: int = 64,
82 packed_bytes_per_token: int = 584,
83 packed_page_alignment: int = 576,
84 + bounded_replay: bool = False,
85 ):
86 super().__init__()
87 self.backend_cls = backend_cls or DeepseekSparseSWABackend
88 + # DeepseekV4.1 SWA bounded replay.
89 + self.bounded_replay = bounded_replay
90 self.kv_cache = torch.tensor([])
91 self.head_dim = head_dim
92 self.window_size = window_size
@@ -117,6 +120,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
120 "nvfp4_ds_mla",
121 )
122 return SlidingWindowMLASpec(
123 + bounded_replay=self.bounded_replay,
124 block_size=self.block_size,
125 num_kv_heads=1,
126 head_size=self.head_dim,
@@ -198,6 +202,10 @@ class DeepseekSparseSWAMetadata:
202 # None when the model is text-only or the batch has no image spans.
203 prefill_left_visible: torch.Tensor | None = None
204 prefill_right_visible: torch.Tensor | None = None
205 + # SWA bounded replay: [num_reqs] lower bound of the positions a request's
206 + # window attention may read (SWA bounded replay);
207 + # zeros when nothing replays.
208 + replay_start: torch.Tensor | None = None
209
210 # Number of decode/prefill requests/tokens (batch is reordered: decodes first)
211 num_decodes: int = 0
@@ -256,10 +264,12 @@ class DeepseekSparseSWAMetadata:
264 has_compressed = compress_ratio > 1
265
266 # query_len <= max_num_batched_tokens and
259 - # gather_len = query_len + min(prefix_len, window_size - 1), so the
260 - # worst-case gathered width is bounded by
261 - # max_num_batched_tokens + window_size - 1. The compressed prefix pool
262 - # is bounded by ceil(max_model_len / compress_ratio).
267 + # gather_len <= query_len + min(prefix_len, window_size - 1) (a replay
268 + # boundary only shrinks it), so the worst-case gathered width is
269 + # bounded by max_num_batched_tokens + window_size - 1. The compressed
270 + # prefix pool is bounded by ceil(max_model_len / compress_ratio). The
271 + # CPU gather lengths below are this upper bound, for chunk planning
272 + # only; the kernels use the exact per-request values.
273 max_workspace_area = prefill_chunk_size * (
274 (cdiv(self.prefill_max_model_len, compress_ratio) if has_compressed else 0)
275 + self.prefill_window_size
@@ -331,6 +341,7 @@ class ComputePrefillMetadataKernel(
341 # Inputs
342 seq_lens_ptr,
343 query_start_loc_ptr,
344 + replay_start_ptr,
345 num_prefills,
346 num_decodes,
347 window_size,
@@ -352,7 +363,10 @@ class ComputePrefillMetadataKernel(
363
364 query_len = qsl_end - qsl_start
365 prefix_len = seq_len - query_len
355 - gather_len = query_len + tl.minimum(prefix_len, window_size - 1)
366 + # Context below replay_start has no window KV (SWA bounded replay).
367 + replay_start = tl.load(replay_start_ptr + num_decodes + safe_offset, mask=mask)
368 + visible_prefix = tl.maximum(prefix_len - replay_start, 0)
369 + gather_len = query_len + tl.minimum(visible_prefix, window_size - 1)
370
371 tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask)
372
@@ -384,6 +398,7 @@ class ComputePrefillMetadataKernel(
398 prefill_gather_lens=int32_ptr,
399 seq_lens=int32_ptr,
400 query_start_loc=int32_ptr,
401 + replay_start=int32_ptr,
402 num_prefills=compile_key.block_size,
403 num_decodes=0,
404 window_size=1,
@@ -395,6 +410,7 @@ class ComputePrefillMetadataKernel(
410 prefill_gather_lens: torch.Tensor,
411 seq_lens: torch.Tensor,
412 query_start_loc: torch.Tensor,
413 + replay_start: torch.Tensor,
414 num_prefills: int,
415 num_decodes: int,
416 window_size: int,
@@ -550,12 +566,19 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
566 )
567 self.decode_swa_indices_noncausal: torch.Tensor | None = None
568 self._max_tokens = max_tokens
569 + # replay_start when no group replays: no lower bound.
570 + self.no_replay_start = torch.zeros(
571 + self.vllm_config.scheduler_config.max_num_seqs + 1,
572 + dtype=torch.int32,
573 + device=self.device,
574 + )
575
576 def build(
577 self,
578 common_prefix_len: int,
579 common_attn_metadata: CommonAttentionMetadata,
580 fast_build: bool = False,
581 + replay_start: torch.Tensor | None = None,
582 ) -> DeepseekSparseSWAMetadata:
583 """Build SWA metadata for mixed decode/prefill batches.
584
@@ -564,6 +587,10 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
587 separate window_topk_idxs for each portion.
588
589 For prefill, we use chunked prefill to align with the indexer's chunking.
590 +
591 + ``replay_start`` ([num_reqs], SWA bounded replay): the position from
592 + which each request holds window KV; the model state passes it for every
593 + batch of a replaying group, graph captures pass nothing.
594 """
595 seq_lens = common_attn_metadata.seq_lens
596 seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
@@ -587,6 +614,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
614
615 is_valid_token = self.is_valid_token[: slot_mapping.shape[0]]
616 is_valid_token.copy_(slot_mapping >= 0)
617 + if replay_start is None:
618 + # Graph captures build without it; their dummy batches replay nothing.
619 + replay_start = self.no_replay_start
620
621 non_causal = not common_attn_metadata.causal
622 decode_swa_width = (
@@ -609,6 +639,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
639 device=self.device,
640 )
641 decode_swa_indices = self.decode_swa_indices_noncausal
642 + # No replay_start: decode rows sit above the hit, so a
643 + # request's replayed window never reaches them.
644 _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL(
645 decode_swa_indices,
646 self.decode_swa_lens,
@@ -637,6 +669,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
669 is_valid_token,
670 block_table,
671 self.block_size,
672 + replay_start,
673 num_tokens=num_decode_tokens,
674 token_offset=0,
675 )
@@ -683,6 +716,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
716 is_valid_token,
717 block_table,
718 self.block_size,
719 + replay_start,
720 num_tokens=num_prefill_tokens,
721 token_offset=num_decode_tokens,
722 has_image=has_image,
@@ -696,6 +730,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
730 seq_lens_cpu,
731 query_start_loc,
732 query_start_loc_cpu,
733 + replay_start,
734 )
735
736 # Per-layer-type tile-scheduler plan holders. Empty FlashMLASchedMeta
@@ -727,6 +762,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
762 ),
763 prefill_left_visible=prefill_left_visible,
764 prefill_right_visible=prefill_right_visible,
765 + replay_start=replay_start,
766 block_size=self.block_size,
767 num_decodes=num_decodes,
768 num_prefills=num_prefills,
@@ -803,6 +839,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
839 assert metadata.is_valid_token is not None
840 assert metadata.decode_swa_indices is not None
841 assert metadata.decode_swa_lens is not None
842 + assert metadata.replay_start is not None
843
844 _COMPUTE_SWA_INDICES_AND_LENS_KERNEL(
845 metadata.decode_swa_indices,
@@ -817,6 +854,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
854 metadata.is_valid_token,
855 metadata.block_table,
856 self.block_size,
857 + metadata.replay_start,
858 num_tokens=metadata.num_decode_tokens,
859 token_offset=0,
860 )
@@ -872,6 +910,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
910 seq_lens_cpu: torch.Tensor | None,
911 query_start_loc: torch.Tensor,
912 query_start_loc_cpu: torch.Tensor,
913 + replay_start: torch.Tensor,
914 ) -> dict[str, torch.Tensor | int | None]:
915 """Pre-compute DeepseekV4 prefill metadata during the metadata build phase.
916
@@ -893,6 +932,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
932 pfx_gather_lens,
933 seq_lens,
934 query_start_loc,
935 + replay_start,
936 num_prefills,
937 num_decodes,
938 self.window_size,
@@ -982,6 +1022,7 @@ def _compute_swa_indices_and_lens_kernel(
1022 block_table_ptr,
1023 block_table_stride,
1024 block_size,
1025 + replay_start_ptr,
1026 token_offset,
1027 HAS_IMAGE: tl.constexpr,
1028 TRITON_BLOCK_SIZE: tl.constexpr,
@@ -1022,6 +1063,8 @@ def _compute_swa_indices_and_lens_kernel(
1063 right = 0
1064 left_add = tl.maximum(left - (window_size - 1), 0)
1065 start_pos = tl.maximum(pos - (window_size - 1) - left_add, 0)
1066 + # SWA bounded replay: no window KV exists below the request's replay start.
1067 + start_pos = tl.maximum(start_pos, tl.load(replay_start_ptr + req_idx))
1068 end_pos = pos + right + 1
1069
1070 swa_len = end_pos - start_pos
@@ -1111,6 +1154,7 @@ class ComputeSWAIndicesAndLensKernel(
1154 is_valid_token=TritonWarmupTensor(torch.bool),
1155 block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
1156 block_size=compile_key.block_size,
1157 + replay_start=int32_ptr,
1158 num_tokens=1,
1159 token_offset=0,
1160 has_image=bool(compile_key.has_image),
@@ -1131,6 +1175,7 @@ class ComputeSWAIndicesAndLensKernel(
1175 is_valid_token: torch.Tensor,
1176 block_table: torch.Tensor,
1177 block_size: int,
1178 + replay_start: torch.Tensor,
1179 *,
1180 num_tokens: int,
1181 token_offset: int,
vllm/v1/core/kv_cache_manager.py
+9 -4
@@ -465,12 +465,17 @@ class KVCacheManager:
465 A list of new allocated blocks.
466
467 """
468 - # When loading KV data asynchronously, we may have zero new tokens to
469 - # compute while still allocating slots for externally computed tokens.
470 - if num_new_tokens == 0 and num_external_computed_tokens == 0:
468 + # A step may need no slots of its own while still adopting computed
469 + # tokens: an async KV load, or a chunk that ends inside the replayed
470 + # range of a hit (SWA bounded replay).
471 + if (
472 + num_new_tokens == 0
473 + and num_external_computed_tokens == 0
474 + and num_new_computed_tokens == 0
475 + ):
476 raise ValueError(
477 "num_new_tokens must be greater than 0 when there are no "
473 - "external computed tokens"
478 + "computed tokens to adopt"
479 )
480
481 if new_computed_blocks is not None:
vllm/v1/core/sched/output.py
+3
@@ -48,6 +48,8 @@ class NewRequestData:
48
49 # Only used for v2 model runner.
50 prefill_token_ids: list[int] | None = None
51 + # DeepSeek-V4.1 only: SWA bounded replay; see Request.replay_start.
52 + replay_start: int = 0
53
54 @classmethod
55 def from_request(
@@ -73,6 +75,7 @@ class NewRequestData:
75 prompt_embeds=request.prompt_embeds,
76 prompt_is_token_ids=request.prompt_is_token_ids,
77 prefill_token_ids=prefill_token_ids,
78 + replay_start=request.replay_start,
79 )
80
81 @property
vllm/v1/core/sched/scheduler.py
+73 -6
@@ -271,6 +271,18 @@ class Scheduler(SchedulerInterface):
271 self.use_eagle_block_drop = False
272 self.num_spec_tokens = vllm_config.num_speculative_tokens
273 self.num_lookahead_tokens = vllm_config.num_lookahead_tokens
274 + # DSV41 SWA bounded replay: groups that declare a replay window are rebuilt
275 + # after a prefix hit by recomputing its trailing tokens. One window
276 + # for all such groups, so the rewind matches every group's allocation.
277 + replay_windows = {
278 + group.kv_cache_spec.prefix_replay_tokens
279 + for group in kv_cache_config.kv_cache_groups
280 + if group.kv_cache_spec.prefix_replay_tokens > 0
281 + }
282 + assert len(replay_windows) <= 1, (
283 + f"Prefix replay windows should agree: {sorted(replay_windows)}"
284 + )
285 + self.prefix_replay_tokens = replay_windows.pop() if replay_windows else 0
286 self.num_prefill_lookahead = vllm_config.num_prefill_lookahead_tokens
287 self.dynamic_sd_lookup: list[int] | None = None
288 if speculative_config is not None:
@@ -946,6 +958,17 @@ class Scheduler(SchedulerInterface):
958 step_skipped_waiting.prepend_request(request)
959 continue
960
961 + if self.prefix_replay_tokens:
962 + # SWA bounded replay recomputes the hit's last
963 + # window, from the block holding its first
964 + # token; the sliding-window groups retire
965 + # whole blocks below the window of the hit's
966 + # next token. Hits end on a block boundary, or
967 + # a hit ending one token short of one would
968 + # retire the block the replay starts in.
969 + ext_tokens -= ext_tokens % self.block_size
970 + load_kv_async = load_kv_async and ext_tokens > 0
971 +
972 if partial_tail and ext_tokens > partial_tail:
973 # Remote strictly exceeds the full local hit: drop the
974 # sub-block tail so no CoW is needed, and let the load
@@ -988,6 +1011,17 @@ class Scheduler(SchedulerInterface):
1011 num_new_local_computed_tokens + num_external_computed_tokens
1012 )
1013 assert num_computed_tokens <= request.num_tokens
1014 + if 0 < num_computed_tokens <= self.prefix_replay_tokens:
1015 + # SWA bounded replay: a hit no longer than the replayed
1016 + # window would be recomputed in full and save nothing.
1017 + new_computed_blocks = (
1018 + self.kv_cache_manager.empty_kv_cache_blocks
1019 + )
1020 + num_new_local_computed_tokens = 0
1021 + num_external_computed_tokens = 0
1022 + connector_prefix_cache_hits = 0
1023 + num_computed_tokens = 0
1024 + load_kv_async = False
1025
1026 # Skip request with pending mm encoding prefetches
1027 if self._ec_transfer_pending(request, num_computed_tokens):
@@ -1022,6 +1056,11 @@ class Scheduler(SchedulerInterface):
1056 new_encoder_compute_budget = encoder_compute_budget
1057 pad_spec_decode = False
1058
1059 + # SWA bounded replay: recompute the tail of the hit without
1060 + # rewriting its cached KV. An async load replays once the KV
1061 + # has arrived (_update_waiting_for_remote_kv).
1062 + num_replay_tokens = 0
1063 +
1064 if load_kv_async:
1065 # KVTransfer: loading remote KV, do not allocate for new work.
1066 assert num_external_computed_tokens > 0
@@ -1032,6 +1071,12 @@ class Scheduler(SchedulerInterface):
1071 break
1072 else:
1073 request_token_budget = min(token_budget, input_budget - draft_slots)
1074 + if did_prefix_cache_lookup:
1075 + # Fresh admission: local and/or sync external hit.
1076 + num_replay_tokens = self._mark_prefix_replay(
1077 + request, num_computed_tokens
1078 + )
1079 + num_computed_tokens -= num_replay_tokens
1080 # Number of tokens to be scheduled.
1081 # We use `request.num_tokens` instead of
1082 # `request.num_prompt_tokens` to consider the resumed
@@ -1159,9 +1204,12 @@ class Scheduler(SchedulerInterface):
1204 + self._spec_decode_step_blocks()
1205 )
1206
1207 + # Replayed tokens are already counted in the adopted hit; a
1208 + # chunk that ends inside the replayed range needs no new slots.
1209 + num_tokens_past_hit = max(num_new_tokens - num_replay_tokens, 0)
1210 new_blocks = self.kv_cache_manager.allocate_slots(
1211 request,
1164 - num_new_tokens,
1212 + num_tokens_past_hit,
1213 num_new_computed_tokens=num_new_local_computed_tokens,
1214 new_computed_blocks=new_computed_blocks,
1215 num_lookahead_tokens=effective_lookahead_tokens,
@@ -2963,6 +3011,20 @@ class Scheduler(SchedulerInterface):
3011 self._request_remaining_blocks(req) for req in self._inflight_prefills
3012 )
3013
3014 + def _mark_prefix_replay(self, request: Request, num_hit_tokens: int) -> int:
3015 + """Record where a prefix hit's replay starts and return the number of
3016 + hit tokens to recompute; the caller rewinds the computed count by it.
3017 +
3018 + The replayed tokens are the hit's last window (see
3019 + ``Request.replay_start``): the worker rebuilds their sliding-window KV
3020 + and leaves their cached KV alone.
3021 + """
3022 + if not self.prefix_replay_tokens:
3023 + return 0
3024 + num_replay_tokens = min(self.prefix_replay_tokens, num_hit_tokens)
3025 + request.replay_start = num_hit_tokens - num_replay_tokens
3026 + return num_replay_tokens
3027 +
3028 def _update_waiting_for_remote_kv(self, request: Request) -> None:
3029 """KV Connector: update request state after async recv is finished.
3030
@@ -2998,11 +3060,16 @@ class Scheduler(SchedulerInterface):
3060 # This will cache the blocks iff caching is enabled.
3061 self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens)
3062
3001 - # on a full prompt hit, we need to re-compute the last token
3002 - # in order to be able to sample the next token
3003 - if request.num_computed_tokens == request.num_tokens:
3004 - request.num_computed_tokens = request.num_tokens - 1
3005 -
3063 + # SWA bounded replay recomputes the tail of the hit, which covers the
3064 + # last token; otherwise a full prompt hit re-computes that token so the
3065 + # next one can be sampled.
3066 + num_replay_tokens = self._mark_prefix_replay(
3067 + request, request.num_computed_tokens
3068 + )
3069 + if num_replay_tokens > 0:
3070 + request.num_computed_tokens -= num_replay_tokens
3071 + elif request.num_computed_tokens == request.num_tokens:
3072 + request.num_computed_tokens = request.num_tokens - 1
3073 self.finished_recving_kv_req_ids.remove(request.request_id)
3074
3075 def _try_promote_blocked_waiting_request(self, request: Request) -> bool:
vllm/v1/core/single_type_kv_cache_manager.py
+2
@@ -492,6 +492,8 @@ class SingleTypeKVCacheManager(ABC):
492 can resume at, from ``get_replay_boundaries``.
493
494 """
495 + if not self.kv_cache_spec.prefix_cacheable:
496 + return
497 num_cached_blocks = self.num_cached_block.get(request.request_id, 0)
498 num_full_blocks = num_tokens // self.block_size
499
vllm/v1/kv_cache_interface.py
+23 -1
@@ -164,6 +164,11 @@ class KVCacheSpec:
164 """Whether this spec's group participates in prefix caching."""
165 return True
166
167 + @property
168 + def prefix_replay_tokens(self) -> int:
169 + """DeepSeek-V4.1 only: bounded replay. Currently only for DSV41 SWA."""
170 + return 0
171 +
172 @property
173 def num_heads(self) -> int:
174 raise NotImplementedError
@@ -908,6 +913,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
913 # DeepseekV4-only: see MLAAttentionSpec.model_version.
914 alignment: int | None = None # Default to None for no padding.
915 model_version: str | None = None
916 + bounded_replay: bool = False
917
918 # MLA stores a single latent vector per state; there is no separate V.
919 head_size_v: int = 0
@@ -919,6 +925,14 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
925 super().__post_init__()
926 _apply_alignment_padding(self)
927
928 + @property
929 + def prefix_cacheable(self) -> bool:
930 + return not self.bounded_replay
931 +
932 + @property
933 + def prefix_replay_tokens(self) -> int:
934 + return self.sliding_window if self.bounded_replay else 0
935 +
936 @classmethod
937 def merge(cls, specs: list[Self]) -> Self:
938 assert all(isinstance(spec, SlidingWindowMLASpec) for spec in specs), (
@@ -930,16 +944,18 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
944 model_version_set = set(spec.model_version for spec in specs)
945 sliding_window_set = set(spec.sliding_window for spec in specs)
946 extra_retained_set = set(spec.extra_retained_tokens for spec in specs)
947 + bounded_replay_set = set(spec.bounded_replay for spec in specs)
948 assert (
949 len(cache_dtype_str_set) == 1
950 and len(tokens_per_state_set) == 1
951 and len(model_version_set) == 1
952 and len(sliding_window_set) == 1
953 and len(extra_retained_set) == 1
954 + and len(bounded_replay_set) == 1
955 ), (
956 "All attention layers in the same KV cache group must use the same "
957 "quantization method, tokens per state, model version, sliding "
942 - "window size, and retained token count."
958 + "window size, retained token count, and replay policy."
959 )
960 return cls(
961 block_size=specs[0].block_size,
@@ -954,6 +970,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
970 cache_dtype_str=cache_dtype_str_set.pop(),
971 tokens_per_state=tokens_per_state_set.pop(),
972 model_version=model_version_set.pop(),
973 + bounded_replay=bounded_replay_set.pop(),
974 )
975
976 def is_uniform_with_collection(
@@ -962,6 +979,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
979 return all(
980 isinstance(spec, SlidingWindowMLASpec)
981 and spec.sliding_window == self.sliding_window
982 + and spec.bounded_replay == self.bounded_replay
983 for spec in kv_cache_specs.values()
984 )
985
@@ -1195,6 +1213,10 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
1213 def prefix_cacheable(self) -> bool:
1214 return all(spec.prefix_cacheable for spec in self.kv_cache_specs.values())
1215
1216 + @property
1217 + def prefix_replay_tokens(self) -> int:
1218 + return max(spec.prefix_replay_tokens for spec in self.kv_cache_specs.values())
1219 +
1220 @property
1221 def first_spec(self) -> KVCacheSpec:
1222 """Return the first spec in the group."""
vllm/v1/request.py
+3
@@ -201,6 +201,9 @@ class Request:
201 # in the (sparse) prefix cache; 0 means none. Set at admission for
202 # hybrid/Mamba models when a shared prefix is detected (Marconi-style).
203 self.shared_prefix_boundary = 0
204 + # DeepSeek-V4.1 only: SWA bounded replay. The request holds no
205 + # sliding-window KV below this position; 0 when nothing replays.
206 + self.replay_start = 0
207
208 # The number of NaNs in logits. A value greater than 0
209 # indicates that the output is corrupted
vllm/v1/worker/gpu/model_states/default.py
+6 -1
@@ -16,7 +16,10 @@ from vllm.v1.worker.gpu.attn_utils import (
16 from vllm.v1.worker.gpu.input_batch import InputBatch
17 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
18 from vllm.v1.worker.gpu.mm.rope import get_rope_state
19 -from vllm.v1.worker.gpu.model_states.interface import ModelState
19 +from vllm.v1.worker.gpu.model_states.interface import (
20 + ModelSpecificAttnMetadata,
21 + ModelState,
22 +)
23 from vllm.v1.worker.gpu.model_states.mm_pruning import maybe_create_mm_pruner
24 from vllm.v1.worker.gpu.model_states.prompt_embeds import PromptEmbedsState
25 from vllm.v1.worker.gpu.states import RequestState
@@ -173,6 +176,7 @@ class DefaultModelState(ModelState):
176 kv_cache_config: KVCacheConfig,
177 for_capture: bool = False,
178 ubatch_idx: int = 0,
179 + model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
180 ) -> dict[str, Any]:
181 if cudagraph_mode == CUDAGraphMode.FULL:
182 # Use padded sizes - padding is handled by model_runner.prepare_attn.
@@ -229,5 +233,6 @@ class DefaultModelState(ModelState):
233 req_idx=input_batch.idx_mapping_np,
234 ubatch_idx=ubatch_idx,
235 fast_prefill=input_batch.fast_prefill,
236 + model_specific_attn_metadata=model_specific_attn_metadata,
237 )
238 return attn_metadata
vllm/v1/worker/gpu/model_states/encoder_only.py
+3
@@ -19,6 +19,7 @@ from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec, KVCacheConfig
19 from vllm.v1.worker.gpu.input_batch import InputBatch
20 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
21 from vllm.v1.worker.gpu.model_states.default import DefaultModelState
22 +from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
23 from vllm.v1.worker.gpu.states import RequestState
24 from vllm.v1.worker.utils import AttentionGroup
25
@@ -168,6 +169,7 @@ class EncoderOnlyModelState(DefaultModelState):
169 kv_cache_config: KVCacheConfig,
170 for_capture: bool = False,
171 ubatch_idx: int = 0,
172 + model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
173 ) -> dict[str, Any]:
174 assert ubatch_idx == 0, "DBO is not supported"
175 attn_metadata = super().prepare_attn(
@@ -178,6 +180,7 @@ class EncoderOnlyModelState(DefaultModelState):
180 attn_groups,
181 kv_cache_config,
182 for_capture,
183 + model_specific_attn_metadata=model_specific_attn_metadata,
184 )
185 attn_metadata.update(
186 self._build_encoder_attn_metadata(input_batch, cudagraph_mode, for_capture)
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+2
@@ -239,8 +239,10 @@ class MambaHybridModelState(DefaultModelState):
239 kv_cache_config: KVCacheConfig,
240 for_capture: bool = False,
241 ubatch_idx: int = 0,
242 + model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
243 ) -> dict[str, Any]:
244 assert ubatch_idx == 0, "DBO is not supported"
245 + assert model_specific_attn_metadata is None
246 if cudagraph_mode == CUDAGraphMode.FULL:
247 num_reqs = input_batch.num_reqs_after_padding
248 num_tokens = input_batch.num_tokens_after_padding