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

tests/kernels/test_compressor_kv_cache.py

@@ -342,6 +342,12 @@ def test_v41_compressor_metadata_maps_tokens_to_their_ring():
     assert metadata.query_start_loc is query_start_loc
     assert metadata.token_to_req_indices.tolist() == [0, 0, 0, 1, 1]
 
+    # Dummy batches (profiling, CUDA graph capture) carry all-zero block
+    # tables. Block 0 is the shared null page every group overlays, so the ring
+    # must not write it: those tokens stay at PAD like every other cache's.
+    common.block_table_tensor.zero_()
+    assert builder.build(0, common).slot_mapping.tolist() == [-1] * 8
+
 
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA stream coverage")
 @pytest.mark.parametrize(

tests/v1/attention/test_deepseek_v4_swa_visible.py

@@ -57,10 +57,13 @@ def ref_left_right(
     return lefts, rights
 
 
-def ref_swa_bounds(pos: int, window: int, left: int, right: int) -> tuple[int, int]:
-    """Reference [start, end) window bounds for one query token."""
+def ref_swa_bounds(
+    pos: int, window: int, left: int, right: int, replay_start: int = 0
+) -> tuple[int, int]:
+    """Reference [start, end) window bounds for one query token. Under SWA
+    bounded replay no window KV exists below ``replay_start``."""
     left_add = max(left - (window - 1), 0)
-    start = max(pos - (window - 1) - left_add, 0)
+    start = max(pos - (window - 1) - left_add, 0, replay_start)
     return start, pos + right + 1
 
 
@@ -97,20 +100,24 @@ def ref_swa_slot_rows(
     window: int,
     max_image_tokens: int,
     width: int,
+    replay_starts: list[int] | None = None,
 ) -> tuple[list[list[int]], list[int]]:
     """Reference paged slot-id rows and lens for every token in the batch."""
     block_table_cpu = block_table.cpu()
     lefts, rights = ref_left_right(
         seq_lens, query_lens, spans_per_req, max_image_tokens
     )
+    replay_starts = replay_starts or [0] * len(seq_lens)
     rows: list[list[int]] = []
     lens: list[int] = []
     token = 0
     for req, (seq_len, query_len) in enumerate(zip(seq_lens, query_lens)):
         prefix_len = seq_len - query_len
         for i in range(query_len):
             pos = prefix_len + i
-            start, end = ref_swa_bounds(pos, window, lefts[token], rights[token])
+            start, end = ref_swa_bounds(
+                pos, window, lefts[token], rights[token], replay_starts[req]
+            )
             row = []
             for p in range(start, end):
                 blk = int(block_table_cpu[req, p // BLOCK_SIZE])
@@ -139,6 +146,7 @@ def run_swa_kernel(
     swa_indices = torch.zeros(num_tokens, 1, width, dtype=torch.int32, device=device)
     swa_lens = torch.zeros(num_tokens, dtype=torch.int32, device=device)
     is_valid = slot_mapping >= 0
+    replay_start_t = torch.zeros(len(seq_lens), dtype=torch.int32, device=device)
 
     if with_image:
         lefts, rights = ref_left_right(
@@ -164,6 +172,7 @@ def run_swa_kernel(
         block_table,
         block_table.stride(0),
         BLOCK_SIZE,
+        replay_start_t,
         token_offset=0,
         HAS_IMAGE=with_image,
         TRITON_BLOCK_SIZE=1024,
@@ -289,18 +298,25 @@ def combine_case(
     query_lens: list[int],
     spans: list[list[tuple[int, int]]],
     with_image: bool,
+    replay_starts: list[int] | None = None,
+    combine_fn=combine_topk_swa_indices,
 ):
     """Run combine_topk_swa_indices and return (indices, lens, expected)."""
     device = torch.device("cuda")
     num_reqs = len(seq_lens)
+    replay_starts = replay_starts or [0] * num_reqs
     query_start_loc = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device)
     query_start_loc[1:] = torch.tensor(
         query_lens, dtype=torch.int32, device=device
     ).cumsum(0)
     num_tokens = int(query_start_loc[-1])
     seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device)
+    # The builder's gather covers only the context above the replay start.
     gather_lens = torch.tensor(
-        [q + min(s - q, WINDOW - 1) for s, q in zip(seq_lens, query_lens)],
+        [
+            q + min(max(s - q - r, 0), WINDOW - 1)
+            for s, q, r in zip(seq_lens, query_lens, replay_starts)
+        ],
         dtype=torch.int32,
         device=device,
     )
@@ -319,7 +335,7 @@ def combine_case(
     else:
         left_t = right_t = None
 
-    combined_indices, combined_lens = combine_topk_swa_indices(
+    combined_indices, combined_lens = combine_fn(
         topk_indices,
         query_start_loc,
         seq_lens_t,
@@ -354,6 +370,8 @@ def combine_case(
             pos = prefix_len + i
             topk_len = min((pos + 1) // compress_ratio, topk)
             start, end = ref_swa_bounds(pos, WINDOW, lefts[token], rights[token])
+            # The window never reaches below the gathered buffer.
+            start = max(start, gather_start)
             swa_len = end - start
             row = [-1] * combined_topk
             for j in range(topk_len):
@@ -388,6 +406,32 @@ def test_combine_topk_swa_with_image_spans(cfg):
     assert indices.cpu().tolist() == rows
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+@pytest.mark.parametrize("cfg", COMBINE_CASES)
+@pytest.mark.parametrize("with_image", [False, True])
+def test_v41_combine_topk_swa_stops_at_replay_start(cfg, with_image):
+    """SWA bounded replay: the gathered buffer starts at replay_start, so the
+    window (plain or widened by an image span) never indexes below it."""
+    from vllm.models.deepseek_v41.common.ops.cache_utils import (
+        combine_topk_swa_indices as combine_v41,
+    )
+
+    # Request 0 replays [16, 40): the windows of its first rows and of the
+    # span starting at 20 would otherwise reach below 16.
+    indices, lens, rows, exp_lens = combine_case(
+        cfg["compress_ratio"],
+        cfg["topk"],
+        seq_lens=[40, 12],
+        query_lens=[24, 12],
+        spans=[[(20, 27)], []],
+        with_image=with_image,
+        replay_starts=[16, 0],
+        combine_fn=combine_v41,
+    )
+    assert lens.cpu().tolist() == exp_lens
+    assert indices.cpu().tolist() == rows
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 @pytest.mark.parametrize("cfg", COMBINE_CASES)
 def test_combine_topk_swa_without_image_unchanged(cfg):
@@ -494,11 +538,17 @@ def build_metadata(
     seq_lens: list[int],
     query_lens: list[int],
     mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None,
+    replay_starts: list[int] | None = None,
 ):
     device = torch.device("cuda")
     query_start_loc, seq_lens_t, _, slot_mapping, block_table = make_batch(
         seq_lens, query_lens, device
     )
+    replay_start = (
+        None
+        if replay_starts is None
+        else torch.tensor(replay_starts, dtype=torch.int32, device=device)
+    )
     return builder.build(
         0,
         CommonAttentionMetadata(
@@ -514,6 +564,7 @@ def build_metadata(
             slot_mapping=slot_mapping,
             causal=True,
             mm_req_doc_ranges=mm_req_doc_ranges,
+            replay_start=replay_start,
         ),
     )
 
@@ -578,3 +629,98 @@ def test_builder_text_model_unchanged():
     )
     assert md.prefill_swa_lens.cpu().tolist() == lens
     assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows
+
+
+# SWA bounded replay: every prefill index path clamps the window at replay_start.
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_builder_replay_start_bounds_prefill_window_and_gather():
+    """Paged-direct prefill indices stop at replay_start and the FlashMLA
+    gather only covers the context the request may see."""
+    seq_lens = [40, 12, 30]
+    query_lens = [24, 12, 10]
+    replay_starts = [12, 0, 0]
+    builder = make_builder(vision=False)
+    md = build_metadata(builder, seq_lens, query_lens, None, replay_starts)
+    assert md.replay_start is not None
+    assert md.replay_start.cpu().tolist() == replay_starts
+
+    _, _, _, _, block_table = make_batch(seq_lens, query_lens, torch.device("cuda"))
+    rows, lens = ref_swa_slot_rows(
+        seq_lens,
+        query_lens,
+        [[], [], []],
+        block_table,
+        WINDOW,
+        MAX_IMG,
+        WINDOW,
+        replay_starts=replay_starts,
+    )
+    assert md.prefill_swa_lens.cpu().tolist() == lens
+    assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows
+    # Request 0: the first replayed row (pos 16) sees [12, 16]; its last row
+    # (pos 39) is a full window above the replay start and sees all of it.
+    assert md.prefill_swa_lens[0].item() == 16 - 12 + 1
+    assert md.prefill_swa_lens[23].item() == WINDOW
+    # gather_len = query_len + min(prefix_len - replay_start, WINDOW - 1).
+    assert md.prefill_gather_lens.cpu().tolist() == [
+        24 + min(16 - 12, WINDOW - 1),
+        12,
+        10 + min(20, WINDOW - 1),
+    ]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_v41_flashinfer_mixed_sparse_indices_respect_replay_start():
+    from vllm.models.deepseek_v41.common.ops.cache_utils import (
+        build_flashinfer_mixed_sparse_indices as build_v41,
+    )
+
+    device = torch.device("cuda")
+    # 1 decode token (req 0) + two prefill requests; req 1 resumes at 16.
+    seq_lens = [20, 40, 9]
+    query_lens = [1, 24, 9]
+    replay_starts = [0, 16, 0]
+    query_start_loc = torch.tensor([0, 1, 25, 34], dtype=torch.int32, device=device)
+    seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device)
+    token_to_req = torch.tensor(
+        [0] + [1] * 24 + [2] * 9, dtype=torch.int32, device=device
+    )
+    block_table = torch.arange(3, dtype=torch.int32, device=device).view(3, 1)
+    kwargs = dict(
+        decode_swa_indices=torch.zeros((1, WINDOW), dtype=torch.int32, device=device),
+        decode_compressed_indices=None,
+        decode_compressed_topk_lens=None,
+        prefill_topk_indices=torch.empty((33, 0), dtype=torch.int32, device=device),
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens_t,
+        token_to_req_indices=token_to_req,
+        swa_block_table=block_table,
+        swa_block_size=BLOCK_SIZE,
+        compressed_block_table=None,
+        compressed_block_size=BLOCK_SIZE,
+        window_size=WINDOW,
+        compress_ratio=1,
+        topk=0,
+    )
+    plain, _ = build_v41(
+        replay_start=torch.zeros(3, dtype=torch.int32, device=device), **kwargs
+    )
+    bounded, _ = build_v41(
+        replay_start=torch.tensor(replay_starts, dtype=torch.int32, device=device),
+        **kwargs,
+    )
+    rows, _ = ref_swa_slot_rows(
+        seq_lens,
+        query_lens,
+        [[], [], []],
+        block_table,
+        WINDOW,
+        MAX_IMG,
+        WINDOW,
+        replay_starts=replay_starts,
+    )
+    # Decode row is copied through; prefill rows follow the bounded window.
+    assert bounded[0].cpu().tolist() == plain[0].cpu().tolist()
+    assert bounded[1:].cpu().tolist() == rows[1:]

tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py

@@ -14,6 +14,7 @@
 from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.attention.backends.mla.compressor_utils import (
     CompressedSlotMappingKernel,
+    get_compressed_slot_mapping,
 )
 from vllm.v1.attention.backends.mla.indexer import (
     BuildPrefillChunkMetadataKernel,
@@ -203,6 +204,28 @@ def test_compressed_slot_mapping_warmup_includes_index_kpool():
     assert {(key.compress_ratio, key.block_size) for key in keys} == {(32, 2)}
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_compressed_slot_mapping_inherits_padded_token_slots():
+    """A token whose own slot is padded (SWA bounded replay) closes no
+    compressed state either."""
+    device = torch.device("cuda")
+    query_start_loc = torch.tensor([0, 8], dtype=torch.int32, device=device)
+    seq_lens = torch.tensor([8], dtype=torch.int32, device=device)
+    block_table = torch.tensor([[3]], dtype=torch.int32, device=device)
+    slot_mapping = torch.arange(8, dtype=torch.int64, device=device)
+    slot_mapping[:4] = -1
+    compressed = get_compressed_slot_mapping(
+        8,
+        slot_mapping,
+        query_start_loc,
+        seq_lens,
+        block_table,
+        block_size=4,
+        compress_ratio=2,
+    )
+    assert compressed.tolist() == [-1, -1, -1, -1, -1, 3 * 4 + 2, -1, 3 * 4 + 3]
+
+
 def test_index_conversion_warmup_uses_physical_block_stride():
     config = SimpleNamespace(
         cache_config=SimpleNamespace(block_size=64),

tests/v1/core/test_contiguous_kv_packing.py

@@ -651,3 +651,67 @@ def test_ring_beside_paged_groups_keeps_the_kv_block_as_scheduler_block(self):
             scheduler_block_size=128,
         )
         assert len(manager.coordinator.single_type_managers) == len(groups)
+
+
+class TestSWABoundedReplayGrouping:
+    def test_non_cacheable_swa_leaves_hits_to_the_paged_group(self):
+        """Under SWA bounded replay only the paged MLA group hashes (hash
+        block = its block) and the coordinator probes one spec group."""
+        config = _shared_layout_config()
+        config.cache_config.enable_prefix_caching = True
+        specs: dict[str, KVCacheSpec] = {
+            "layers.2.attn": MLAAttentionSpec(
+                block_size=128,
+                num_kv_heads=1,
+                head_size=584,
+                dtype=torch.uint8,
+                tokens_per_state=2,
+                alignment=576,
+            ),
+        }
+        swa = SlidingWindowMLASpec(
+            block_size=32,
+            num_kv_heads=1,
+            head_size=584,
+            dtype=torch.uint8,
+            sliding_window=128,
+            alignment=576,
+            bounded_replay=True,
+        )
+        for layer in range(4):
+            specs[f"layers.{layer}.attn.swa_cache"] = swa
+
+        groups = get_kv_cache_groups(config, specs)
+        # The worker reads the replay window off the (packed) group specs to
+        # arm the window clamp; the wrapper must forward it like cacheability.
+        assert sorted(
+            (g.kv_cache_spec.prefix_cacheable, g.kv_cache_spec.prefix_replay_tokens)
+            for g in groups
+        ) == [(False, 128)] * 4 + [(True, 0)]
+        kv_cache_config = get_kv_cache_config_from_groups(
+            config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups)
+        )
+        manager = KVCacheManager(
+            generate_scheduler_kv_cache_config([kv_cache_config]),
+            max_model_len=8192,
+            enable_caching=True,
+            hash_block_size=128,
+            scheduler_block_size=128,
+        )
+        assert len(manager.coordinator.attention_groups) == 1
+        assert isinstance(
+            manager.coordinator.attention_groups[0].spec, MLAAttentionSpec
+        )
+
+        # Allocating and caching a request goes through every group's manager;
+        # the 32-token SWA group must skip hashing rather than assert on the
+        # 128-token hash block.
+        from tests.v1.core.utils import create_requests
+
+        request = create_requests(num_requests=1, num_tokens=256, block_size=128)[0]
+        assert manager.allocate_slots(request, 256) is not None
+        request.num_computed_tokens = 256
+        manager.cache_blocks(request, 256)
+        for single in manager.coordinator.single_type_managers:
+            cached = single.num_cached_block.get(request.request_id, 0)
+            assert cached == (2 if single.kv_cache_spec.prefix_cacheable else 0)