[Spec] Windowed draft-decode attention for built-in EAGLE / MTP drafts (#32673)

python/sglang/kernels/ops/speculative/cache_locs.py

@@ -77,16 +77,12 @@ def generate_draft_decode_kv_indices(
     iter_upper: tl.constexpr,
     num_tokens_upper: tl.constexpr,
     page_size: tl.constexpr,
+    window_size: tl.constexpr = 0,
+    sink_size: tl.constexpr = 0,
     NUM_STEPS: tl.constexpr = 0,
 ):
-    # Optional token-block parallelism (NUM_STEPS > 0): the first grid axis
+    # window_size > 0 restricts the draft (not the target) to sink_size prefix
-    # packs (draft step, token block) as ``step + NUM_STEPS * block``,
+    # tokens + the most-recent window_size; window_size == 0 is the identity.
-    # spreading the per-request index copy below over many programs instead
-    # of one program crawling the whole context serially (which bottlenecks
-    # long-context spec decode, where this kernel runs every iteration).
-    # NUM_STEPS == 0 (default) is the historical one-program-per-step kernel:
-    # the same 128-wide copy loop, in the same order, with the token-block
-    # branches folded away at compile time.
     BLOCK_SIZE: tl.constexpr = 128 if NUM_STEPS == 0 else 512
     pid0 = tl.program_id(axis=0)
     bid = tl.program_id(axis=1)
@@ -108,45 +104,52 @@ def generate_draft_decode_kv_indices(
     kv_indptr += kv_indptr_stride * iters
     iters += 1
 
-    if NUM_STEPS == 0:
     load_offset = tl.arange(0, bs_upper)
-        seq_lens = tl.load(
+    seq_lens = tl.load(paged_kernel_lens + load_offset, mask=load_offset < bid, other=0)
-            paged_kernel_lens + load_offset, mask=load_offset < bid, other=0
-        )
     seq_len = tl.load(paged_kernel_lens + bid)
-        cum_seq_len = tl.sum(seq_lens)
+    if window_size > 0:
+        cap = window_size + sink_size
+        seq_lens = tl.minimum(seq_lens, cap)
+        seq_len_w = tl.minimum(seq_len, cap)
+        s_eff = tl.minimum(sink_size, seq_len)
+        recent_start = seq_len - (seq_len_w - s_eff)
     else:
-        seq_len = tl.load(paged_kernel_lens + bid)
+        seq_len_w = seq_len
-        num_loop = tl.cdiv(seq_len, BLOCK_SIZE)
+        s_eff = 0
-        # Blocks with no copy work exit before the O(bs) prefix-sum below;
+        recent_start = 0
-        # block 0 always continues (it owns the extension and kv_indptr).
-        if blk >= num_loop and blk > 0:
-            return
-        load_offset = tl.arange(0, bs_upper)
-        seq_lens = tl.load(
-            paged_kernel_lens + load_offset, mask=load_offset < bid, other=0
-        )
     cum_seq_len = tl.sum(seq_lens)
 
     # Update kv_indices
-    kv_offset = cum_seq_len * topk + bid * iters * topk + topk_id * (seq_len + iters)
+    kv_offset = cum_seq_len * topk + bid * iters * topk + topk_id * (seq_len_w + iters)
     kv_ptr = kv_indices + kv_offset
     token_pool_ptr = req_to_token + tl.load(req_pool_indices + bid) * pool_len
 
+    num_loop = tl.cdiv(seq_len_w, BLOCK_SIZE)
+    if NUM_STEPS != 0 and blk >= num_loop and blk > 0:
+        return
     if NUM_STEPS == 0:
-        kv_offset = tl.arange(0, BLOCK_SIZE)
+        copy_offset = tl.arange(0, BLOCK_SIZE)
-        num_loop = tl.cdiv(seq_len, BLOCK_SIZE)
         for _ in range(num_loop):
-            mask = kv_offset < seq_len
+            mask = copy_offset < seq_len_w
-            data = tl.load(token_pool_ptr + kv_offset, mask=mask)
+            src = tl.where(
-            tl.store(kv_ptr + kv_offset, data, mask=mask)
+                copy_offset < s_eff,
-            kv_offset += BLOCK_SIZE
+                copy_offset,
+                recent_start + copy_offset - s_eff,
+            )
+            data = tl.load(token_pool_ptr + src, mask=mask)
+            tl.store(kv_ptr + copy_offset, data, mask=mask)
+            copy_offset += BLOCK_SIZE
     else:
         for i in range(blk, num_loop, num_blk):
-            tok_off = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            copy_offset = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-            mask = tok_off < seq_len
+            mask = copy_offset < seq_len_w
-            data = tl.load(token_pool_ptr + tok_off, mask=mask)
+            src = tl.where(
-            tl.store(kv_ptr + tok_off, data, mask=mask)
+                copy_offset < s_eff,
+                copy_offset,
+                recent_start + copy_offset - s_eff,
+            )
+            data = tl.load(token_pool_ptr + src, mask=mask)
+            tl.store(kv_ptr + copy_offset, data, mask=mask)
 
     # Extension entries and kv_indptr belong to token block 0 alone; other
     # blocks neither compute nor store them.
@@ -178,18 +181,19 @@ def generate_draft_decode_kv_indices(
             )
 
         tl.store(
-            kv_ptr + seq_len + extend_offset,
+            kv_ptr + seq_len_w + extend_offset,
             extend_data,
             mask=extend_offset < iters,
         )
 
         # Update kv_indptr
         bs_offset = tl.arange(0, num_tokens_upper)
-
         zid = bid * topk + topk_id
         if zid == 0:
             zid = num_seqs * topk
         pos_vals = tl.load(positions + bs_offset, mask=bs_offset < zid, other=0)
+        if window_size > 0:
+            pos_vals = tl.minimum(pos_vals, window_size + sink_size)
         base = tl.sum(pos_vals)
         tl.store(kv_indptr + zid, base + zid * iters)
 

python/sglang/srt/arg_groups/fields/spec.py

@@ -160,7 +160,11 @@ class Spec(msgspec.Struct):
     ] = None
     speculative_draft_window_size: A[
         Optional[int],
-        "Sliding window size for the draft model. Honored by Llama EAGLE-3 (`LlamaForCausalLMEagle3`) and DFLASH only; other EAGLE-3 backends (e.g. MLA-based drafters) silently ignore it. For Llama EAGLE-3, the drafter only attends to the most recent N keys (verifier hidden states + its own outputs); the verifier is unaffected. For DFLASH, the draft worker keeps a recent target-token window in its local KV cache (paged backends may retain up to one extra page on the left for alignment). Default is full attention/context.",
+        "Sliding window size for the draft model. Honored by Llama EAGLE-3 (`LlamaForCausalLMEagle3`), DFLASH, and the built-in EAGLE/MTP draft-decode path on the Triton and FlashInfer draft backends; other EAGLE-3 backends (e.g. MLA-based drafters) silently ignore it. For Llama EAGLE-3, the drafter only attends to the most recent N keys (verifier hidden states + its own outputs); the verifier is unaffected. For DFLASH, the draft worker keeps a recent target-token window in its local KV cache (paged backends may retain up to one extra page on the left for alignment). For the built-in EAGLE/MTP draft, each draft-decode step attends to a --speculative-draft-sink-size sink plus the most recent N tokens, leaving the target verify pass unchanged; it is ignored (with a warning) if the draft model has a native sliding window of its own. Default is full attention/context.",
+    ] = None
+    speculative_draft_sink_size: A[
+        Optional[int],
+        "Number of leading 'attention sink' tokens the draft always attends to, in addition to the --speculative-draft-window-size recent window (StreamingLLM-style). Honored only by the built-in EAGLE/MTP draft-decode path on the Triton and FlashInfer draft backends; the Llama EAGLE-3 and DFLASH windows ignore it. 0/unset => pure recent window. Requires --speculative-draft-window-size.",
     ] = None
     speculative_moe_runner_backend: A[
         Optional[str],

python/sglang/srt/arg_groups/speculative_hook.py

@@ -161,8 +161,11 @@ def handle_speculative_decoding(server_args: ServerArgs) -> None:
         ),
     )
 
-    # Validate --speculative-draft-window-size once, regardless of algorithm.
+    # Validate --speculative-draft-window-size / --speculative-draft-sink-size once,
-    # Consumed by DFLASH (compact draft KV cache) and Llama EAGLE-3 (drafter attention SWA).
+    # regardless of algorithm. Consumed by DFLASH (compact draft KV cache), Llama
+    # EAGLE-3 (drafter attention SWA), and the built-in MTP/NEXTN + EAGLE draft-decode
+    # path on the Triton and FlashInfer draft attention backends (StreamingLLM sink +
+    # recent window).
     if cfg.speculative_draft_window_size is not None:
         window_size = int(cfg.speculative_draft_window_size)
         if window_size <= 0:
@@ -174,13 +177,30 @@ def handle_speculative_decoding(server_args: ServerArgs) -> None:
             "handle_speculative_decoding",
             speculative_draft_window_size=window_size,
         )
-        if cfg.speculative_algorithm not in ("EAGLE3", "DFLASH"):
+        if cfg.speculative_algorithm not in ("EAGLE", "EAGLE3", "DFLASH"):
             logger.warning(
                 "--speculative-draft-window-size has no effect with "
-                "speculative_algorithm=%s (honored by Llama EAGLE-3 and DFLASH only).",
+                "speculative_algorithm=%s (honored by DFLASH, Llama EAGLE-3, and the "
+                "EAGLE/MTP/NEXTN draft-decode path on the Triton/FlashInfer draft backends).",
                 cfg.speculative_algorithm,
             )
 
+    if cfg.speculative_draft_sink_size is not None:
+        sink_size = int(cfg.speculative_draft_sink_size)
+        if sink_size < 0:
+            raise ValueError(
+                f"--speculative-draft-sink-size must be non-negative, got {sink_size}."
+            )
+        if cfg.speculative_draft_window_size is None:
+            raise ValueError(
+                "--speculative-draft-sink-size requires --speculative-draft-window-size."
+            )
+        declare_resolution(
+            server_args,
+            "handle_speculative_decoding",
+            speculative_draft_sink_size=sink_size,
+        )
+
     algo = None
     if cfg.speculative_algorithm is not None:
         from sglang.srt.speculative.spec_info import SpeculativeAlgorithm

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -50,6 +50,7 @@ from sglang.srt.speculative.spec_utils import (
     draft_kv_indices_buffer_width,
     draft_kv_indices_used_len,
     generate_draft_decode_kv_indices,
+    resolve_draft_decode_window,
 )
 from sglang.srt.utils import (
     get_cuda_graph_max_batch_size,
@@ -2357,6 +2358,9 @@ class FlashInferMultiStepDraftBackend:
         # Cached variables for generate_draft_decode_kv_indices
         self.pool_len = model_runner.req_to_token_pool.req_to_token.shape[1]
         self.req_to_token_pool = model_runner.req_to_token_pool
+        self.draft_window_size, self.draft_sink_size = resolve_draft_decode_window(
+            model_runner
+        )
 
     def common_template(
         self,
@@ -2395,6 +2399,8 @@ class FlashInferMultiStepDraftBackend:
             next_power_of_2(self.speculative_num_steps),
             next_power_of_2(bs),
             self.page_size,
+            self.draft_window_size,
+            self.draft_sink_size,
         )
 
         assert forward_batch.spec_info is not None

python/sglang/srt/layers/attention/triton_backend.py

@@ -44,6 +44,7 @@ from sglang.srt.speculative.spec_utils import (
     draft_kv_indices_buffer_width,
     draft_kv_indices_used_len,
     generate_draft_decode_kv_indices,
+    resolve_draft_decode_window,
 )
 from sglang.srt.utils import (
     get_bool_env_var,
@@ -2343,6 +2344,9 @@ class TritonMultiStepDraftBackend:
         self.req_to_token_pool = model_runner.req_to_token_pool
         self.pool_len = model_runner.req_to_token_pool.req_to_token.shape[1]
         self.page_size = get_schedule().page_size
+        self.draft_window_size, self.draft_sink_size = resolve_draft_decode_window(
+            model_runner
+        )
 
     def common_template(
         self,
@@ -2377,6 +2381,8 @@ class TritonMultiStepDraftBackend:
             next_power_of_2(self.speculative_num_steps),
             next_power_of_2(bs),
             self.page_size,
+            self.draft_window_size,
+            self.draft_sink_size,
         )
 
         if call_fn is None:

python/sglang/srt/speculative/spec_utils.py

@@ -31,6 +31,7 @@ from sglang.kernels.ops.speculative.cache_locs import (
 from sglang.kernels.ops.speculative.eagle import (
     fill_accept_out_cache_loc_func as fill_accept_out_cache_loc_func,
 )
+from sglang.srt.arg_groups.overrides import resolving_view
 from sglang.srt.configs.hybrid_arch import mambaish_config
 from sglang.srt.constrained.base_grammar_backend import GrammarMask
 from sglang.srt.distributed.parallel_state import (
@@ -238,6 +239,41 @@ def draft_kv_indices_used_len(
     return seq_lens_sum * topk + bs * num_steps
 
 
+def resolve_draft_decode_window(model_runner) -> Tuple[int, int]:
+    """Resolve (window_size, sink_size) for generate_draft_decode_kv_indices.
+
+    Returns (0, 0) -- full draft attention, the pristine read plan -- when
+    --speculative-draft-window-size is unset, and also when the draft model
+    already has a sliding window of its own: the index builder emits one KV list
+    shared by every draft layer, so it cannot express a per-layer window, and the
+    draft model's own window is authoritative -- whether it comes from the
+    checkpoint or, for LlamaForCausalLMEagle3, from this same flag.
+    """
+    # Read through the resolving view: handle_speculative_decoding declares both
+    # fields rather than assigning them, so the raw field holds the unvalidated input.
+    cfg = resolving_view(model_runner.server_args)
+    window_size = int(cfg.speculative_draft_window_size or 0)
+    if window_size <= 0:
+        return 0, 0
+    # The runner's resolved window, not the raw config field: config keys
+    # (sliding_window / window_size) are overloaded across model families, while
+    # this is the same value the attention backends key their own SWA paths on.
+    native_window = getattr(model_runner, "sliding_window_size", None)
+    if native_window is not None and native_window > 0:
+        # An equal window is the one that was asked for, applied per layer instead
+        # of here (LlamaForCausalLMEagle3 routes this flag into its own window).
+        if native_window != window_size:
+            logger.warning(
+                "Ignoring --speculative-draft-window-size=%d: this draft model has a "
+                "sliding window of %d, which the attention backend applies per layer. "
+                "Draft-decode windowing stays off.",
+                window_size,
+                native_window,
+            )
+        return 0, 0
+    return window_size, int(cfg.speculative_draft_sink_size or 0)
+
+
 def record_stream_each(tensors, stream):
     """Call record_stream(stream) on each cuda tensor in `tensors`, skipping
     non-tensor / non-cuda entries. Tells the caching allocator that the

test/registered/kernel/spec/test_draft_decode_window.py

@@ -0,0 +1,439 @@
+"""Unit test: windowed draft-decode KV index builder (StreamingLLM sink + window).
+
+Covers the ``window_size`` / ``sink_size`` additions to
+``generate_draft_decode_kv_indices`` (the shared kv-index kernel used by the
+built-in MTP/NEXTN + EAGLE draft-decode path on BOTH the Triton and FlashInfer
+draft attention backends).
+
+Properties checked:
+
+  1. IDENTITY (off-by-default is a no-op) -- ``window_size == 0`` reproduces the
+     full-KV read plan. Both the CSR offsets (``kv_indptr``) and the gathered
+     slots (``kv_indices``) are checked against an independent closed-form oracle
+     derived only from the input ``seq_lens`` (NOT from the kernel's packing
+     logic), so an offset bug is not mirrored by the oracle. This is the
+     losslessness/regression guarantee.
+  2. KV-INDPTR lengths -- with a window each draft-decode step keeps
+     ``min(seq_len, sink + window)`` base tokens + ``it + 1`` never-windowed tree
+     tokens; compared to a clamp+cumsum oracle. Parametrized past the
+     256K/512K/1M thresholds (incl. the perf-smoke W4032/S64 params).
+  3. CONTENT -- the actual gathered slots are, in order,
+     ``[first sink] + [most-recent window] + [draft tree]`` (the StreamingLLM
+     layout), read back at buffer offset 0 (num_seqs=1, topk=1).
+  3b. TREE (``topk > 1``) -- with a tree draft every ``(request, topk)`` slot gets
+     its own copy of the windowed base list plus its own branch of tree tokens.
+     The windowed per-slot write offset is ``topk_id * (kept + it + 1)``, so a
+     windowing bug there would silently overlap neighbouring branches while the
+     lengths still look right. Slices are located by the *oracle* ``kv_indptr``
+     (independent closed form) and their content checked, which ties the offset
+     and the CSR bounds together. Covered for ``page_size == 1`` and for the
+     ``page_size > 1 and topk > 1`` paged-tree branch, where the tree tokens are
+     read from the UNCAPPED ``seq_len`` but written at the capped offset.
+
+  4. RESOLUTION -- ``resolve_draft_decode_window`` maps server args to the
+     kernel's (window, sink), and returns (0, 0) for a draft model that already
+     has a sliding window of its own (that per-layer window wins over the flag,
+     since this index builder emits one KV list for all draft layers). It warns
+     only when it substitutes a different window than the one requested.
+  5. BACKEND WIRING -- both ``TritonMultiStepDraftBackend`` and
+     ``FlashInferMultiStepDraftBackend`` resolve the pair in ``__init__`` and
+     forward it to the kernel, so windowing is honored identically on both
+     backends (CPU-only guard).
+
+The kernel is imported normally from ``sglang`` (no by-file-path staging), so it
+tracks the installed tree.
+"""
+
+import inspect
+import unittest
+from types import SimpleNamespace
+from unittest import mock
+
+import torch
+
+from sglang.srt.speculative import spec_utils
+from sglang.srt.speculative.spec_utils import (
+    generate_draft_decode_kv_indices,
+    resolve_draft_decode_window,
+)
+from sglang.srt.utils import next_power_of_2
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+_HAS_CUDA = torch.cuda.is_available()
+
+
+# --- Harness: mirror {Triton,FlashInfer}MultiStepDraftBackend.common_template ---
+def _make_inputs(seq_lens, num_steps, topk, page_size=1, pool_len=None):
+    dev = "cuda"
+    num_seqs = len(seq_lens)
+    bs = num_seqs * topk
+    seq_lens_t = torch.tensor(seq_lens, dtype=torch.int64, device=dev)
+    max_seq = int(seq_lens_t.max().item())
+    if pool_len is None:
+        # Room for every branch's tree tokens: topk*num_steps slots (page_size==1)
+        # or up to topk*num_new_pages*page_size past the prefix's last page.
+        pool_len = max_seq + topk * (num_steps + page_size) + 16
+    # Unique physical slot ids per (req, position) so gather checks are unambiguous.
+    req_to_token = (
+        torch.arange(num_seqs * pool_len, dtype=torch.int64, device=dev).reshape(
+            num_seqs, pool_len
+        )
+        + 1000
+    )
+    req_pool_indices = torch.arange(num_seqs, dtype=torch.int64, device=dev)
+    positions = torch.empty(bs, dtype=torch.int64, device=dev)
+    for bid in range(num_seqs):
+        for tk in range(topk):
+            positions[bid * topk + tk] = seq_lens[bid]
+
+    width = bs * (max_seq + num_steps) + 16
+    kv_indices = torch.zeros((num_steps, width), dtype=torch.int64, device=dev)
+    kv_indptr = torch.zeros((num_steps, bs + 1), dtype=torch.int64, device=dev)
+    return dict(
+        req_pool_indices=req_pool_indices,
+        req_to_token=req_to_token,
+        seq_lens=seq_lens_t,
+        positions=positions,
+        kv_indices=kv_indices,
+        kv_indptr=kv_indptr,
+        pool_len=pool_len,
+        page_size=page_size,
+        num_seqs=num_seqs,
+        topk=topk,
+        num_steps=num_steps,
+        bs=bs,
+    )
+
+
+def _run(io, window_size=0, sink_size=0):
+    kv_indices = io["kv_indices"].clone()
+    kv_indptr = io["kv_indptr"].clone()
+    generate_draft_decode_kv_indices[(io["num_steps"], io["num_seqs"], io["topk"])](
+        io["req_pool_indices"],
+        io["req_to_token"],
+        io["seq_lens"],
+        kv_indices,
+        kv_indptr,
+        io["positions"],
+        io["pool_len"],
+        kv_indices.shape[1],
+        kv_indptr.shape[1],
+        next_power_of_2(io["num_seqs"]),
+        next_power_of_2(io["num_steps"]),
+        next_power_of_2(io["bs"]),
+        io["page_size"],
+        window_size,
+        sink_size,
+    )
+    torch.cuda.synchronize()
+    return kv_indices, kv_indptr
+
+
+def _expected_kv_indptr(io, cap=None):
+    """Closed form: kv_indptr[it][zid] = sum(clamp(pos[:zid], cap)) + zid*(it+1)."""
+    pos = io["positions"].to("cpu")
+    if cap is not None:
+        pos = torch.clamp(pos, max=cap)
+    csum = torch.cat([torch.zeros(1, dtype=torch.int64), torch.cumsum(pos, 0)])
+    out = torch.zeros((io["num_steps"], io["bs"] + 1), dtype=torch.int64)
+    for it in range(io["num_steps"]):
+        for zid in range(io["bs"] + 1):
+            out[it, zid] = csum[zid] + zid * (it + 1)
+    return out
+
+
+def _expected_base(r2t, seq_len, window, sink):
+    """StreamingLLM base gather for one request (window==0 => full context)."""
+    if window == 0:
+        return r2t[0:seq_len]
+    cap = window + sink
+    kept = min(seq_len, cap)
+    s_eff = min(sink, seq_len)
+    recent_start = seq_len - (kept - s_eff)
+    return torch.cat([r2t[0:s_eff], r2t[recent_start : recent_start + (kept - s_eff)]])
+
+
+def _expected_tree(r2t, seq_len, topk_id, topk, num_steps, n_tree, page_size):
+    """This branch's tree tokens, always read from the UNCAPPED prefix end.
+
+    Mirrors the kernel's two extend layouts (linear, and the paged per-topk page
+    stride), both of which windowing must leave untouched.
+    """
+    if page_size == 1 or topk == 1:
+        start = seq_len + topk_id * num_steps
+    else:
+        last_page_len = seq_len % page_size
+        num_new_pages_per_topk = (
+            last_page_len + num_steps + page_size - 1
+        ) // page_size
+        start = (
+            (seq_len // page_size) * page_size
+            + topk_id * num_new_pages_per_topk * page_size
+            + last_page_len
+        )
+    return r2t[start : start + n_tree]
+
+
+@unittest.skipUnless(_HAS_CUDA, "draft-decode kv-index kernel requires a CUDA GPU")
+class TestDraftDecodeWindowKernel(unittest.TestCase):
+    def test_window_zero_is_full_kv(self):
+        """window_size==0 => original full-KV read plan (indptr + content)."""
+        for seq_lens, num_steps, topk in [
+            ([128], 4, 1),
+            ([300, 130, 517], 5, 1),
+            ([1024, 777], 4, 2),
+        ]:
+            with self.subTest(seq_lens=seq_lens, num_steps=num_steps, topk=topk):
+                io = _make_inputs(seq_lens, num_steps, topk)
+                self._assert_slot_layout(io, 0, 0)
+
+    def test_windowed_kv_indptr_lengths(self):
+        for seq_lens, window, sink in [
+            ([5000, 2048, 900], 1024, 64),
+            ([5000, 2048, 900], 1024, 0),  # pure recent window
+            ([100, 200], 4096, 64),  # cap exceeds seq_len -> keep all
+            # long-context regression (perf smoke used W4032/S64 past 256K/512K/1M)
+            ([1029306], 4032, 64),
+            ([1048576, 300000, 70000], 4032, 64),
+            ([262145, 131072], 4032, 64),
+        ]:
+            with self.subTest(seq_lens=seq_lens, window=window, sink=sink):
+                io = _make_inputs(seq_lens, num_steps=4, topk=1)
+                _, indptr = _run(io, window_size=window, sink_size=sink)
+                expected = _expected_kv_indptr(io, cap=window + sink)
+                self.assertTrue(torch.equal(indptr.cpu(), expected))
+
+    def test_windowed_content_sink_plus_recent(self):
+        for seq_len, window, sink in [
+            (5000, 1024, 64),
+            (65536, 4032, 64),
+            (262145, 4032, 64),  # just past native 262144
+            (524288, 4032, 64),
+            (1029306, 4032, 64),  # exact perf-smoke seq_len (W4032/S64)
+        ]:
+            with self.subTest(seq_len=seq_len, window=window, sink=sink):
+                io = _make_inputs([seq_len], num_steps=4, topk=1)
+                kv_indices, _ = _run(io, window_size=window, sink_size=sink)
+                r2t = io["req_to_token"][0].cpu()
+                kept = min(seq_len, window + sink)
+                expected_base = _expected_base(r2t, seq_len, window, sink)
+                for step in range(io["num_steps"]):
+                    row = kv_indices[step].cpu()
+                    self.assertTrue(
+                        torch.equal(row[0:kept], expected_base),
+                        f"base gather wrong at step {step} (seq_len={seq_len})",
+                    )
+                    n_tree = step + 1
+                    self.assertTrue(
+                        torch.equal(
+                            row[kept : kept + n_tree], r2t[seq_len : seq_len + n_tree]
+                        ),
+                        f"tree tokens wrong at step {step} (seq_len={seq_len})",
+                    )
+
+    def _assert_slot_layout(self, io, window, sink):
+        """Every (request, topk) slot holds [sink + recent] + its own tree branch.
+
+        Slots are located by the closed-form kv_indptr oracle, so this also pins
+        the per-slot write offset (topk_id * (kept + iters)). ``topk == 1`` is a
+        one-branch tree and ``window == 0`` asks for the whole prefix, so this
+        covers the identity case too.
+        """
+        cap = None if window == 0 else window + sink
+        kv_indices, indptr = _run(io, window_size=window, sink_size=sink)
+        expected_indptr = _expected_kv_indptr(io, cap=cap)
+        self.assertTrue(torch.equal(indptr.cpu(), expected_indptr))
+
+        topk, num_steps = io["topk"], io["num_steps"]
+        seq_lens = io["seq_lens"].tolist()
+        for step in range(num_steps):
+            row = kv_indices[step].cpu()
+            n_tree = step + 1
+            for bid, seq_len in enumerate(seq_lens):
+                r2t = io["req_to_token"][bid].cpu()
+                kept = seq_len if cap is None else min(seq_len, cap)
+                for topk_id in range(topk):
+                    zid = bid * topk + topk_id
+                    start = int(expected_indptr[step, zid])
+                    where = f"step={step} bid={bid} topk_id={topk_id}"
+                    self.assertTrue(
+                        torch.equal(
+                            row[start : start + kept],
+                            _expected_base(r2t, seq_len, window, sink),
+                        ),
+                        f"base gather wrong at {where}",
+                    )
+                    self.assertTrue(
+                        torch.equal(
+                            row[start + kept : start + kept + n_tree],
+                            _expected_tree(
+                                r2t,
+                                seq_len,
+                                topk_id,
+                                topk,
+                                num_steps,
+                                n_tree,
+                                io["page_size"],
+                            ),
+                        ),
+                        f"tree tokens wrong at {where}",
+                    )
+
+    def test_tree_topk_layout(self):
+        """topk > 1 (tree draft): per-branch slices, windowed and unwindowed."""
+        for seq_lens, num_steps, topk, window, sink in [
+            ([600], 4, 2, 128, 16),
+            ([600, 250], 3, 4, 128, 16),
+            ([5000, 900], 4, 2, 1024, 64),
+            ([5000, 900], 4, 2, 1024, 0),  # pure recent window
+            ([300, 130], 3, 2, 4096, 64),  # cap exceeds seq_len -> keep all
+            ([65536, 1029306], 3, 2, 4032, 64),  # long-context tree
+            ([600, 250], 3, 4, 0, 0),  # off-by-default control
+        ]:
+            with self.subTest(seq_lens=seq_lens, topk=topk, window=window, sink=sink):
+                io = _make_inputs(seq_lens, num_steps, topk)
+                self._assert_slot_layout(io, window, sink)
+
+    def test_tree_topk_paged(self):
+        """page_size > 1 AND topk > 1: the paged-tree extend branch.
+
+        The tree tokens are read from the uncapped prefix end with a per-topk page
+        stride and stored at the CAPPED offset, so this is the one place where the
+        windowed and unwindowed lengths must both be respected in one statement.
+        """
+        for seq_lens, num_steps, topk, page_size, window, sink in [
+            ([600], 4, 2, 4, 128, 16),
+            ([601, 255], 3, 2, 8, 128, 16),  # prefix not page-aligned
+            ([604, 256], 3, 4, 4, 128, 0),  # page-aligned prefix, no sink
+            ([5000, 900], 4, 2, 16, 1024, 64),
+            ([601, 255], 3, 2, 8, 0, 0),  # off-by-default control
+        ]:
+            with self.subTest(
+                seq_lens=seq_lens, topk=topk, page_size=page_size, window=window
+            ):
+                io = _make_inputs(seq_lens, num_steps, topk, page_size=page_size)
+                self._assert_slot_layout(io, window, sink)
+
+
+class TestResolveDraftDecodeWindow(unittest.TestCase):
+    """CPU-only: server args -> (window, sink), incl. the opt-out for an SWA draft."""
+
+    @staticmethod
+    def _runner(window=None, sink=None, native_window=None, declared=None):
+        server_args = SimpleNamespace(
+            speculative_draft_window_size=window,
+            speculative_draft_sink_size=sink,
+        )
+        if declared is not None:
+            server_args._resolved_overrides = (
+                ("handle_speculative_decoding", declared),
+            )
+        return SimpleNamespace(
+            server_args=server_args,
+            sliding_window_size=native_window,
+        )
+
+    def test_unset_is_full_attention(self):
+        self.assertEqual(resolve_draft_decode_window(self._runner()), (0, 0))
+
+    def test_window_only(self):
+        self.assertEqual(
+            resolve_draft_decode_window(self._runner(window=4032)), (4032, 0)
+        )
+
+    def test_window_and_sink(self):
+        self.assertEqual(
+            resolve_draft_decode_window(self._runner(window=4032, sink=64)), (4032, 64)
+        )
+
+    def test_declared_values_win_over_the_raw_fields(self):
+        """handle_speculative_decoding declares these fields instead of assigning
+        them, so a raw field read would answer with the unvalidated input."""
+        runner = self._runner(
+            window=None,
+            sink=None,
+            declared={
+                "speculative_draft_window_size": 4032,
+                "speculative_draft_sink_size": 64,
+            },
+        )
+        self.assertEqual(resolve_draft_decode_window(runner), (4032, 64))
+
+    def test_native_sliding_window_draft_opts_out(self):
+        """A draft with its own window keeps it: the flag must not override it."""
+        with mock.patch.object(spec_utils.logger, "warning") as warn:
+            self.assertEqual(
+                resolve_draft_decode_window(
+                    self._runner(window=4032, sink=64, native_window=1024)
+                ),
+                (0, 0),
+            )
+        warn.assert_called_once()
+
+    def test_native_window_equal_to_flag_is_quiet(self):
+        """LlamaForCausalLMEagle3 routes this flag into its own window, so an equal
+        native window is the requested one applied per layer, not a conflict."""
+        with mock.patch.object(spec_utils.logger, "warning") as warn:
+            self.assertEqual(
+                resolve_draft_decode_window(
+                    self._runner(window=4032, sink=64, native_window=4032)
+                ),
+                (0, 0),
+            )
+        warn.assert_not_called()
+
+    def test_no_native_window_still_windows(self):
+        """Guard must not fire for full-attention drafts (0 / None / absent)."""
+        for native in (None, 0):
+            with self.subTest(native_window=native):
+                self.assertEqual(
+                    resolve_draft_decode_window(
+                        self._runner(window=4032, sink=64, native_window=native)
+                    ),
+                    (4032, 64),
+                )
+        runner = self._runner(window=4032, sink=64)
+        del runner.sliding_window_size
+        self.assertEqual(resolve_draft_decode_window(runner), (4032, 64))
+
+
+class TestDraftDecodeWindowBackendWiring(unittest.TestCase):
+    """CPU-only guard: both draft backends must forward window/sink to the kernel.
+
+    The kernel tests above prove correctness once the args reach the kernel; this
+    guards the per-backend plumbing (init resolution + kernel passthrough) so it
+    cannot silently regress on either backend without a GPU-loaded model.
+    """
+
+    def _assert_backend_wires(self, cls):
+        init_src = inspect.getsource(cls.__init__)
+        self.assertIn("resolve_draft_decode_window", init_src)
+        self.assertIn("self.draft_window_size", init_src)
+        self.assertIn("self.draft_sink_size", init_src)
+        tmpl_src = inspect.getsource(cls.common_template)
+        self.assertIn("self.draft_window_size", tmpl_src)
+        self.assertIn("self.draft_sink_size", tmpl_src)
+
+    def test_triton_backend_wires_window_sink(self):
+        try:
+            from sglang.srt.layers.attention.triton_backend import (
+                TritonMultiStepDraftBackend,
+            )
+        except ImportError as e:  # pragma: no cover
+            self.skipTest(f"triton backend import unavailable: {e}")
+        self._assert_backend_wires(TritonMultiStepDraftBackend)
+
+    def test_flashinfer_backend_wires_window_sink(self):
+        try:
+            from sglang.srt.layers.attention.flashinfer_backend import (
+                FlashInferMultiStepDraftBackend,
+            )
+        except ImportError as e:  # pragma: no cover
+            self.skipTest(f"flashinfer backend import unavailable: {e}")
+        self._assert_backend_wires(FlashInferMultiStepDraftBackend)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/spec/test_spec_registry.py

@@ -232,6 +232,7 @@ class TestServerArgsHook(_RegistryIsolated):
             decrypted_draft_config_file=None,
             trust_remote_code=False,
             speculative_draft_window_size=None,
+            speculative_draft_sink_size=None,
             speculative_skip_dp_mlp_sync=False,
             speculative_adaptive=False,
         )