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

Spec v2 tree drafting (topk>1) with page_size>1 (#26972)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-06 12:00:27 -07:00
committed by GitHub
parent 88a7b0fd30
commit 84ca0ffb8c
8 changed files with 228 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/speculative_hook.py
+14 -7
View File
@@ -285,6 +285,8 @@ def _handle_eagle_family(server_args: "ServerArgs") -> None:
"Max running requests is reset to 48 for speculative decoding. You can override this by explicitly setting --max-running-requests."
)
# Spec v2 tree drafting supports topk > 1 with page_size == 1 and page_size > 1
# (the latter via partial-page duplication; backend-gated below).
spec_v1_reason = None
# mamba / linear-attn state models only support topk == 1 on spec v2.
# mamba2_cache_params exists iff the config carries such state; check the
@@ -294,16 +296,15 @@ def _handle_eagle_family(server_args: "ServerArgs") -> None:
if (
server_args.speculative_eagle_topk is not None
and server_args.speculative_eagle_topk > 1
and (server_args.page_size > 1 or is_mamba_state_model)
and is_mamba_state_model
and not server_args.disable_overlap_schedule
):
# Spec v2 topk > 1 only supports page_size == 1 on non-mamba models;
# page_size > 1 (partial-page dup) isn't ported to v2 yet -> fall back to v1.
# Spec v2 topk > 1 is not supported for mamba/linear-attn state models
# (only topk == 1); fall back to v1 for those. page_size > 1 is supported
# on v2 (partial-page duplication), so it no longer forces v1.
server_args.disable_overlap_schedule = True
spec_v1_reason = (
"spec v2 topk > 1 is not supported for mamba/linear-attn models"
if is_mamba_state_model
else "spec v2 topk > 1 currently requires page_size == 1"
)
elif (
not envs.SGLANG_ENABLE_SPEC_V2.get()
@@ -389,13 +390,19 @@ def _handle_eagle_family(server_args: "ServerArgs") -> None:
)
server_args.speculative_num_draft_tokens = server_args.speculative_num_steps + 1
# topk > 1 + page_size > 1 needs the two-pass cascade draft-decode (shared prefix
# pass + per-branch expand pass with prefix-tail dup). Only these backends implement
# it; flashmla / trtllm_mla / cutlass_mla can't express the per-branch tree, so reject.
_PAGE_TREE_SPEC_BACKENDS = ("flashinfer", "fa3", "triton")
if (
server_args.speculative_eagle_topk > 1
and server_args.page_size > 1
and server_args.attention_backend not in ["flashinfer", "fa3"]
and server_args.attention_backend not in _PAGE_TREE_SPEC_BACKENDS
):
raise ValueError(
"speculative_eagle_topk > 1 with page_size > 1 is unstable and produces incorrect results for paged attention backends. This combination is only supported for the 'flashinfer' backend."
f"speculative_eagle_topk > 1 with page_size > 1 is only supported on "
f"{_PAGE_TREE_SPEC_BACKENDS}; got attention_backend="
f"{server_args.attention_backend!r}. Use page_size == 1 or one of those backends."
)
python/sglang/srt/debug_utils/pr_fix_toggle.py
+21
View File
@@ -87,11 +87,32 @@ patches:
"""
_PR_REVERT_YAML_26972 = """
patches:
- target: sglang.srt.model_executor.model_runner_kv_cache_mixin.ModelRunnerKVCacheMixin._init_pools
edits:
- match: |
if (
self.server_args.speculative_algorithm is not None
and self.server_args.page_size > 1
and (self.server_args.speculative_eagle_topk or 1) > 1
):
from sglang.srt.managers.utils import get_alloc_len_per_decode
extra_max_context_len = max(
extra_max_context_len,
2 * get_alloc_len_per_decode(self.server_args),
)
replacement: ""
"""
_PR_FIX_REVERT_YAML: Dict[int, str] = {
25015: _PR_REVERT_YAML_25015,
26329: _PR_REVERT_YAML_26329,
27338: _PR_REVERT_YAML_27338,
27360: _PR_REVERT_YAML_27360,
26972: _PR_REVERT_YAML_26972,
}
python/sglang/srt/managers/scheduler.py
+4 -2
View File
@@ -419,6 +419,10 @@ class Scheduler(
# Init mamba backend
self.init_mamba_backend()
# Must precede init_model_worker: revert targets like _init_pools run during it,
# so patching them afterwards is a no-op.
maybe_revert_pr_fix()
# Launch a model worker and draft model worker if using speculative decoding
self.init_model_worker()
@@ -564,8 +568,6 @@ class Scheduler(
self.init_batch_result_processor()
maybe_revert_pr_fix()
self.is_initializing = False
def init_zbal_on_npu(self):
python/sglang/srt/managers/utils.py
+8 -3
View File
@@ -246,9 +246,14 @@ def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int:
if page_size == 1 or spec_topk == 1:
return max(spec_steps * spec_topk, spec_tokens)
else:
raise NotImplementedError(
"get_alloc_len_per_decode not implemented for page_size > 1 and spec_topk > 1"
)
# page_size > 1 + topk > 1 (spec v2 tree): worst-case page-aligned tree
# footprint. Per topk branch needs ceil((last_page_len + num_steps) / page)
# pages; the partial tail page can be up to page_size - 1, and each branch
# gets its own (duplicated) copy -- so reserve for all topk branches.
num_new_pages_per_topk = (
(page_size - 1) + spec_steps + page_size - 1
) // page_size
return max(num_new_pages_per_topk * page_size * spec_topk, spec_tokens)
@dataclass
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
+15
View File
@@ -300,6 +300,21 @@ class ModelRunnerKVCacheMixin:
if max_spec_draft_tokens is not None:
extra_max_context_len += max_spec_draft_tokens
# page>1 + topk>1 reserves a holey draft footprint (2 * get_alloc_len_per_decode
# = topk * num_new_pages * page) far beyond the default num_draft_tokens
# headroom; widen the row to hold it, else free leaks KV and the holey gather OOBs.
if (
self.server_args.speculative_algorithm is not None
and self.server_args.page_size > 1
and (self.server_args.speculative_eagle_topk or 1) > 1
):
from sglang.srt.managers.utils import get_alloc_len_per_decode
extra_max_context_len = max(
extra_max_context_len,
2 * get_alloc_len_per_decode(self.server_args),
)
if self.server_args.disaggregation_mode == "decode":
from sglang.srt.disaggregation.decode import (
DecodeReqToTokenPool,
python/sglang/srt/speculative/eagle_info_v2.py
+121 -16
View File
@@ -74,6 +74,53 @@ if is_cuda() or is_musa():
)
def duplicate_prefix_tail_to_draft_branches(
token_to_kv_pool,
rows: torch.Tensor,
prefix_base: torch.Tensor,
last_page: torch.Tensor,
num_new_pages: torch.Tensor,
topk: int,
page_size: int,
) -> None:
"""Copy the prefix partial-tail page into each branch's first-page holes (page>1 + topk>1).
The draft-decode expand pass reads each branch's own draft page by block id
(cache_loc // page_size), so branch b>=1's hole slots [0, last_page) must hold the
real prefix tail (branch 0's first page already is it). Mirrors V1 #7725.
"""
if topk <= 1:
return
bs = rows.shape[0]
page_off = torch.arange(page_size, device=rows.device, dtype=torch.int64)
branches = torch.arange(1, topk, device=rows.device, dtype=torch.int64).view(
1, topk - 1, 1
)
# Source: the prefix tail page [prefix_base, prefix_base + page_size), one per branch.
src_pos = (prefix_base.view(bs, 1, 1) + page_off.view(1, 1, page_size)).expand(
bs, topk - 1, page_size
)
# Target: branch b's first page [prefix_base + b*num_new_pages*page, + page_size).
tgt_pos = (
prefix_base.view(bs, 1, 1)
+ branches * (num_new_pages.view(bs, 1, 1) * page_size)
+ page_off.view(1, 1, page_size)
)
# Only [0, last_page) holds real prefix KV; [last_page, page_size) are the branch's
# own draft slots and must not be overwritten.
vmask = (page_off.view(1, 1, page_size) < last_page.view(bs, 1, 1)).expand(
bs, topk - 1, page_size
)
src_slots = torch.gather(rows, 1, src_pos.reshape(bs, -1)).reshape(
bs, topk - 1, page_size
)[vmask]
tgt_slots = torch.gather(rows, 1, tgt_pos.reshape(bs, -1)).reshape(
bs, topk - 1, page_size
)[vmask]
if src_slots.numel() > 0:
token_to_kv_pool.move_kv_cache(tgt_slots, src_slots)
@dataclass
class EagleDraftInputV2Mixin:
def prepare_for_decode(self: EagleDraftInput, batch: ScheduleBatch):
@@ -128,6 +175,21 @@ class EagleDraftInputV2Mixin:
cur_kv_lens_cpu = torch.tensor(cur_kv_lens, dtype=torch.int32, device="cpu")
nxt_kv_lens_cpu = torch.tensor(nxt_kv_lens, dtype=torch.int32, device="cpu")
# Fail fast if the page>1 + topk>1 draft over-allocation
# (2 * get_alloc_len_per_decode) outgrows the req_to_token row: the write below
# would OOB and free would leak KV. The row is widened to hold it in _init_pools
# (PR #26972); fail here with a clear error, not on a later cryptic CUDA assert.
from sglang.srt.server_args import get_global_server_args
if page_size > 1 and (get_global_server_args().speculative_eagle_topk or 1) > 1:
max_alloc_len = int(nxt_kv_lens_cpu.max())
row_width = batch.req_to_token_pool.req_to_token.shape[1]
assert max_alloc_len <= row_width, (
f"spec v2 page>1 topk>1 draft over-allocation ({max_alloc_len}) exceeds "
f"req_to_token row width ({row_width}); page_size={page_size}. Widen the "
f"row to hold committed + 2 * get_alloc_len_per_decode (PR #26972)."
)
# non_blocking H2D: a blocking .to() syncs the schedule stream, which the WAR
# barrier has chained to the prev forward -> host stalls a full forward.
cur_kv_lens_device = cur_kv_lens_cpu.to(device=batch.device, non_blocking=True)
@@ -171,22 +233,65 @@ class EagleDraftInputV2Mixin:
if not batch.forward_mode.is_idle():
bs = len(batch.seq_lens)
# Assign cache locations
batch.out_cache_loc = torch.empty(
(bs * topk * num_steps,),
dtype=torch.int64,
device=batch.device,
)
# FIXME(lsyin): align with the default code path
assign_draft_cache_locs_page_size_1[(bs,)](
batch.req_pool_indices,
req_to_token_pool.req_to_token,
batch.seq_lens,
batch.out_cache_loc,
req_to_token_pool.req_to_token.shape[1],
topk,
num_steps,
)
# Assign cache locations (draft-write targets).
page_size = batch.token_to_kv_pool_allocator.page_size
if page_size == 1 or topk == 1:
batch.out_cache_loc = torch.empty(
(bs * topk * num_steps,),
dtype=torch.int64,
device=batch.device,
)
# FIXME(lsyin): align with the default code path
assign_draft_cache_locs_page_size_1[(bs,)](
batch.req_pool_indices,
req_to_token_pool.req_to_token,
batch.seq_lens,
batch.out_cache_loc,
req_to_token_pool.req_to_token.shape[1],
topk,
num_steps,
)
else:
# page_size > 1 + topk > 1: per-branch page-aligned draft pages.
# Reduce out_cache_loc from the page-aligned tree region down to the
# dense draft slots (skip each branch's duplicated prefix-tail slots
# and trailing padding), matching generate_draft_decode_kv_indices'
# paged read formula: prefix_base + t*num_new_pages*page + last_page + s.
# base is batch.seq_lens (== KV-ready committed prefix at draft time;
# the bonus is the tree root written by verify, not part of [0:seq_lens]).
rows = req_to_token_pool.req_to_token[batch.req_pool_indices.long()]
seq_lens = batch.seq_lens.to(torch.int64)
last_page = seq_lens % page_size
prefix_base = seq_lens - last_page
num_new_pages = (last_page + num_steps + page_size - 1) // page_size
topk_ids = torch.arange(
topk, device=rows.device, dtype=torch.int64
).view(1, topk)
starts = (
prefix_base.view(bs, 1)
+ topk_ids * (num_new_pages.view(bs, 1) * page_size)
+ last_page.view(bs, 1)
)
steps = torch.arange(
num_steps, device=rows.device, dtype=torch.int64
).view(1, 1, num_steps)
pos = (starts.view(bs, topk, 1) + steps).reshape(bs, topk * num_steps)
batch.out_cache_loc = (
torch.gather(rows, 1, pos).reshape(-1).contiguous()
)
# Each branch's page-aligned region starts with `last_page` hole slots
# overlapping the prefix tail page; duplicate the real prefix-tail KV
# into them so whole-page reads stay coherent (see helper docstring).
duplicate_prefix_tail_to_draft_branches(
draft_model_runner.token_to_kv_pool,
rows,
prefix_base,
last_page,
num_new_pages,
topk,
page_size,
)
# Get a forward batch
self.num_tokens_per_req = topk
test/registered/spec/eagle/test_spec_eagle_page.py
+4 -11
View File
@@ -1,7 +1,7 @@
"""page_size > 1 variants (flashinfer).
"""page_size > 1 variants at topk=1 (flashinfer).
EAGLE3 page64 (spec v2) + EAGLE/Llama-2 page4 (topk1 and topk8, spec v1).
Runs on the cheap (5090) runner.
EAGLE3 page64 (spec v2) + EAGLE/Llama-2 page4 (spec v1). topk>1 page variants
live in test_spec_eagle_topk.py. Runs on the cheap (5090) runner.
"""
import unittest
@@ -15,7 +15,7 @@ from sglang.test.kits.spec_server_kits import (
)
from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base, EagleLlama2Base
register_cuda_ci(est_time=540, stage="base-b", runner_config="1-gpu-small")
register_cuda_ci(est_time=360, stage="base-b", runner_config="1-gpu-small")
class TestEagle3Page64(Eagle3Base, SpecAccuracyKit, SpecLogprobKit, SpecFeatureKit):
@@ -35,12 +35,5 @@ class TestEagleLlama2Page4Topk1(EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
class TestEagleLlama2Page4Topk8(EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit):
"""Llama-2 topk>1 tree + page_size=4 (spec v1)."""
page_size = 4
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
if __name__ == "__main__":
unittest.main()
test/registered/spec/eagle/test_spec_eagle_topk_page.py
+41
View File
@@ -0,0 +1,41 @@
"""topk > 1 tree drafting at page_size > 1 (EAGLE3 topk8 + EAGLE/Llama-2 topk8).
page64 stays on spec v2 (overlap), page4 runs on spec v1 (no overlap). flashinfer is
pinned because this runs on the cheap (5090) runner, where fa3 (Hopper-only) isn't
available -- functional sanity only, no perf/stress. (page>1 topk>1 on fa3 is covered
on the Hopper runner in test_spec_eagle_fa3.py.)
"""
import unittest
from sglang.srt.environ import envs
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.spec_server_kits import (
SpecAccuracyKit,
SpecFeatureKit,
)
from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base, EagleLlama2Base
register_cuda_ci(est_time=720, stage="base-b", runner_config="1-gpu-small")
class TestEagle3Page64Topk8(Eagle3Base, SpecAccuracyKit, SpecFeatureKit):
"""EAGLE3 topk=8 tree + page_size=64 (spec v2)."""
page_size = 64
spec_topk = 8
spec_tokens = 32
disable_overlap = False
cuda_graph_max_bs = 5
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
class TestEagleLlama2Page4Topk8(EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit):
"""Llama-2 topk>1 tree + page_size=4 (spec v1)."""
page_size = 4
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
if __name__ == "__main__":
unittest.main()