minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
avalliappan-nvidia
2026-09-22 08:17:57 +08:00
committed by GitHub
parent 9fdb71732a
commit 61d0cf2074
8 changed files with 560 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/speculative/cache_locs.py
+43 -39
View File
@@ -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
+5 -1
View File
@@ -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
+24 -4
View File
@@ -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
+6
View File
@@ -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
+6
View File
@@ -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
+36
View File
@@ -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
+439
View File
@@ -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
+1
View File
@@ -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,
)