diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index dff5f6880..d749efefb 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -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." ) diff --git a/python/sglang/srt/debug_utils/pr_fix_toggle.py b/python/sglang/srt/debug_utils/pr_fix_toggle.py index 2f602fdc9..814a995ad 100644 --- a/python/sglang/srt/debug_utils/pr_fix_toggle.py +++ b/python/sglang/srt/debug_utils/pr_fix_toggle.py @@ -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, } diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index fc9f19f72..eda6cae0e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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): diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index 0dae5b27d..eb21eda02 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -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 diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 7072103eb..3661197a4 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -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, diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index ba8d6265a..e90cff6d9 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -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 diff --git a/test/registered/spec/eagle/test_spec_eagle_page.py b/test/registered/spec/eagle/test_spec_eagle_page.py index ece61fbc9..e7aeebe2c 100644 --- a/test/registered/spec/eagle/test_spec_eagle_page.py +++ b/test/registered/spec/eagle/test_spec_eagle_page.py @@ -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() diff --git a/test/registered/spec/eagle/test_spec_eagle_topk_page.py b/test/registered/spec/eagle/test_spec_eagle_topk_page.py new file mode 100644 index 000000000..1db2686c4 --- /dev/null +++ b/test/registered/spec/eagle/test_spec_eagle_topk_page.py @@ -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()