[Feat][Model] DeepSeek-V4.1 decoder-side SWA bounded replay

tests/v1/cudagraph/test_breakable_cudagraph.py

@@ -165,6 +165,44 @@ def test_nested_capture_raises(cuda_capture_stream):
         pass
 
 
+def test_eager_break_may_capture_its_own_graphs(cuda_capture_stream):
+    """An eager break runs between outer segments, so it may capture (and on
+    replay, replay) graphs of its own, e.g. for a differently shaped sub-batch."""
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+
+    x = torch.zeros(4, device="cuda")
+    y = torch.zeros(4, device="cuda")
+    outer = BreakableCUDAGraphCapture()
+    inner = BreakableCUDAGraphCapture()
+
+    def eager_step():
+        if inner.num_graphs == 0:
+            # Capture time: the outer capture is current but paused.
+            assert BreakableCUDAGraphCapture.current() is outer
+            with inner:
+                y.add_(x)
+            assert BreakableCUDAGraphCapture.current() is outer
+        else:
+            assert BreakableCUDAGraphCapture.current() is None
+            inner.replay()
+
+    with outer:
+        x.add_(1.0)
+        outer.add_eager(eager_step)
+        x.add_(1.0)
+    assert BreakableCUDAGraphCapture.current() is None
+    assert inner.num_graphs == 1
+
+    outer.replay()
+    torch.accelerator.synchronize()
+    assert x.tolist() == [2.0] * 4
+    assert y.tolist() == [1.0] * 4
+    outer.replay()
+    torch.accelerator.synchronize()
+    assert x.tolist() == [4.0] * 4
+    assert y.tolist() == [4.0] * 4
+
+
 def test_active_state_isolated_across_threads(cuda_capture_stream):
     """Verify the thread-local 'active capture' slot is per-thread.
 

tests/v1/spec_decode/test_dspark_context_rows.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""A sliding-window DSpark drafter only inserts context KV for the last
+``window`` scheduled tokens of each request."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool
+from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="UVA buffers")
+
+
+def _rows(window, query_lens):
+    speculator = SimpleNamespace(
+        context_window=window, _context_row_pool=UvaBufferPool(512, torch.int64)
+    )
+    batch = SimpleNamespace(
+        num_reqs=len(query_lens),
+        query_start_loc_np=np.concatenate([[0], np.cumsum(query_lens)]).astype(
+            np.int32
+        ),
+    )
+    return DSparkSpeculator._context_rows(speculator, batch)
+
+
+def test_context_rows_keep_each_request_tail():
+    rows = _rows(128, [1, 300, 100])
+    assert rows.tolist() == [0, *range(301 - 128, 301), *range(301, 401)]
+
+
+@pytest.mark.parametrize(
+    "window, query_lens",
+    [(128, [6, 6, 100]), (None, [1, 300])],
+    ids=["all fit the window", "no window"],
+)
+def test_context_rows_none_when_nothing_to_skip(window, query_lens):
+    assert _rows(window, query_lens) is None

vllm/compilation/breakable_cudagraph.py

@@ -156,12 +156,22 @@ def __init__(self, pool: Any | None = None) -> None:
         self._num_eager_breaks: int = 0
         self._current_graph: torch.cuda.CUDAGraph | None = None
         self._capturing: bool = False
+        self._outer: BreakableCUDAGraphCapture | None = None
+
+    @property
+    def capturing(self) -> bool:
+        """Whether a graph segment is being captured right now."""
+        return self._capturing
 
     # --- context manager protocol ----------------------------------------
 
     def __enter__(self) -> BreakableCUDAGraphCapture:
-        if getattr(BreakableCUDAGraphCapture._tls, "active", None) is not None:
+        outer = BreakableCUDAGraphCapture.current()
+        if outer is not None and outer._capturing:
             raise RuntimeError("Nested BreakableCUDAGraphCapture is not supported.")
+        # An eager break of an outer capture may capture graphs of its own
+        # (see add_eager); the outer capture becomes current again on exit.
+        self._outer = outer
         BreakableCUDAGraphCapture._tls.active = self
         self._begin_segment()
         return self
@@ -170,7 +180,8 @@ def __exit__(self, exc_type, exc, tb) -> None:
         try:
             self._end_segment()
         finally:
-            BreakableCUDAGraphCapture._tls.active = None
+            BreakableCUDAGraphCapture._tls.active = self._outer
+            self._outer = None
 
     # --- segment management ----------------------------------------------
 

vllm/config/cache.py

@@ -228,7 +228,8 @@ class CacheConfig:
     swa_bounded_replay: bool = True
     """Keep the sliding-window KV of models that support it (DeepSeek-V4.1)
     out of prefix caching and rebuild it after a prefix hit by recomputing the
-    hit's last window. Requires model runner V2."""
+    hit's last window; layers past the last KV-source layer then also prefill
+    only each request's trailing window. Requires model runner V2."""
 
     kv_cache_memory_bytes: int | None = None
     """Size of KV Cache per GPU in bytes. By default, this is set to None