diff --git a/python/sglang/kernels/ops/speculative/cache_locs.py b/python/sglang/kernels/ops/speculative/cache_locs.py index af887c303..98697a020 100644 --- a/python/sglang/kernels/ops/speculative/cache_locs.py +++ b/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) diff --git a/python/sglang/srt/arg_groups/fields/spec.py b/python/sglang/srt/arg_groups/fields/spec.py index f9dde3e74..5ae41ce08 100644 --- a/python/sglang/srt/arg_groups/fields/spec.py +++ b/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], diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 76964243d..5afef4f83 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index eb4946be2..b8018cb47 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 6cbfc87f0..6bbf0ec71 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/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: diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index a547702ba..0a836db76 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/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 diff --git a/test/registered/kernel/spec/test_draft_decode_window.py b/test/registered/kernel/spec/test_draft_decode_window.py new file mode 100644 index 000000000..0617d295f --- /dev/null +++ b/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() diff --git a/test/registered/unit/spec/test_spec_registry.py b/test/registered/unit/spec/test_spec_registry.py index 685c6ede1..88d10973d 100644 --- a/test/registered/unit/spec/test_spec_registry.py +++ b/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, )