[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
-    # packs (draft step, token block) as ``step + NUM_STEPS * block``,
-    # 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.
+    # window_size > 0 restricts the draft (not the target) to sink_size prefix
+    # tokens + the most-recent window_size; window_size == 0 is the identity.
     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(
-            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)
+    load_offset = tl.arange(0, bs_upper)
+    seq_lens = tl.load(paged_kernel_lens + load_offset, mask=load_offset < bid, other=0)
+    seq_len = tl.load(paged_kernel_lens + bid)
+    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)
-        num_loop = tl.cdiv(seq_len, BLOCK_SIZE)
-        # Blocks with no copy work exit before the O(bs) prefix-sum below;
-        # 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)
+        seq_len_w = seq_len
+        s_eff = 0
+        recent_start = 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)
-        num_loop = tl.cdiv(seq_len, BLOCK_SIZE)
+        copy_offset = tl.arange(0, BLOCK_SIZE)
         for _ in range(num_loop):
-            mask = kv_offset < seq_len
-            data = tl.load(token_pool_ptr + kv_offset, mask=mask)
-            tl.store(kv_ptr + kv_offset, data, mask=mask)
-            kv_offset += BLOCK_SIZE
+            mask = copy_offset < seq_len_w
+            src = tl.where(
+                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)
+            copy_offset += BLOCK_SIZE
     else:
         for i in range(blk, num_loop, num_blk):
-            tok_off = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-            mask = tok_off < seq_len
-            data = tl.load(token_pool_ptr + tok_off, mask=mask)
-            tl.store(kv_ptr + tok_off, data, mask=mask)
+            copy_offset = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = copy_offset < seq_len_w
+            src = tl.where(
+                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.
-    # Consumed by DFLASH (compact draft KV cache) and Llama EAGLE-3 (drafter attention SWA).
+    # Validate --speculative-draft-window-size / --speculative-draft-sink-size once,
+    # 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