From 1d4ee060c264dab218586054c990f2c455fbceaf Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Mon, 1 Jun 2026 15:37:20 -0700 Subject: [PATCH] Support spec v2 tree drafting (eagle topk>1) with page_size==1 (#26866) --- .../sglang/srt/arg_groups/speculative_hook.py | 6 +- .../sglang/srt/speculative/eagle_info_v2.py | 5 +- .../sglang/srt/speculative/eagle_worker_v2.py | 57 ++++++++++++++++++- .../multi_layer_eagle_worker_v2.py | 3 +- .../spec/eagle/test_spec_eagle_stress.py | 16 +++++- .../spec/eagle/test_spec_eagle_topk.py | 15 ++++- 6 files changed, 92 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 1c38fcb40..ae423fc53 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -278,10 +278,14 @@ 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 and not server_args.disable_overlap_schedule ): + # Spec v2 tree drafting supports topk > 1 with page_size == 1. The + # page_size > 1 + topk > 1 draft KV allocation (partial-page duplication) + # is not yet ported to v2, so fall back to v1 only for that case. server_args.disable_overlap_schedule = True - spec_v1_reason = "spec v2 currently only supports topk = 1" + spec_v1_reason = "spec v2 topk > 1 currently requires page_size == 1" elif ( not envs.SGLANG_ENABLE_SPEC_V2.get() and not server_args.disable_overlap_schedule diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index 88e0bd03b..2ee5eab12 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -506,7 +506,7 @@ def fill_bonus_tokens( accept_tokens, accept_lens, bonus_tokens_ptr, - num_draft_tokens: tl.constexpr, + accept_stride: tl.constexpr, ): # NOTE: we cannot fuse any in-place operations of `accept_lens` inside this kernel # because this kernel reads accept_lens @@ -514,7 +514,8 @@ def fill_bonus_tokens( # `accept_lens` includes the bonus token; the last accepted slot is at -1. accept_len = tl.load(accept_lens + pid) - bonus_token_idx = num_draft_tokens * pid + accept_len - 1 + # accept_stride = per-req width of accept_tokens (= accept_index.shape[1]). + bonus_token_idx = accept_stride * pid + accept_len - 1 bonus_token = tl.load(accept_tokens + bonus_token_idx) tl.store(bonus_tokens_ptr + pid, bonus_token) diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 13c4dc52c..6a841ed53 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -1202,11 +1202,13 @@ class EAGLEWorkerV2(BaseSpecWorker): if not batch.forward_mode.is_idle(): accept_tokens = predict[accept_index] bonus_tokens = torch.empty_like(accept_lens, dtype=torch.int32) + # stride = accept_tokens per-req width = accept_index.shape[1] + # (spec_steps + 1); NOT num_draft_tokens, wrong for topk > 1 trees. fill_bonus_tokens[(bs,)]( accept_tokens, accept_lens, bonus_tokens, - self.speculative_num_draft_tokens, + accept_index.shape[1], ) else: bonus_tokens = torch.empty((0,), device=self.device, dtype=torch.int32) @@ -1216,6 +1218,13 @@ class EAGLEWorkerV2(BaseSpecWorker): batch, logits_output, predict, accept_index, self.speculative_num_steps ) + if not batch.forward_mode.is_idle() and self.topk > 1: + # topk == 1 needs nothing here: the accepted path is already the front + # chain, so the whole compaction is an identity transform. + predict = self._finalize_accepted_tree_path( + batch, accept_index, accept_lens, predict, logits_output, bs + ) + next_draft_input = EagleDraftInput(bonus_tokens=bonus_tokens) # verify_forward_batch transitively holds verify-time GPU tensors @@ -1297,6 +1306,30 @@ class EAGLEWorkerV2(BaseSpecWorker): model=self.target_worker.model_runner.model, ) + def _finalize_accepted_tree_path( + self, + batch: ScheduleBatch, + accept_index: torch.Tensor, + accept_lens: torch.Tensor, + predict: torch.Tensor, + logits_output, + bs: int, + ) -> torch.Tensor: + """Tree drafting (topk > 1): move the accepted path -- KV slots, predict, + hidden_states -- to the contiguous front of each per-req block, which the + downstream chain-layout code (draft-extend select_index, committed-KV reads) + assumes. Returns compacted predict; mutates logits_output.hidden_states + (moved only when present).""" + self.move_accepted_tokens_to_target_kvcache( + batch, accept_index, accept_lens - 1 + ) + predict = self._compact_accepted_to_front(predict, accept_index, bs) + if logits_output.hidden_states is not None: + logits_output.hidden_states = self._compact_accepted_to_front( + logits_output.hidden_states, accept_index, bs + ) + return predict + def move_accepted_tokens_to_target_kvcache( self, batch: ScheduleBatch, @@ -1313,7 +1346,9 @@ class EAGLEWorkerV2(BaseSpecWorker): seq_lens is advanced by ``num_correct_drafts + 1`` to cover the bonus slot. """ bs = len(batch.seq_lens) - size = bs * self.speculative_num_draft_tokens + # accept_index element count, NOT bs * num_draft_tokens: for topk > 1 the + # tree exceeds the accepted chain, over-reading accept_index (illegal memory). + size = bs * accept_index.shape[1] # fill_accepted_out_cache_loc reads out_cache_loc[accept_index]; -1 sentinel ok. maybe_detect_oob( @@ -1350,6 +1385,24 @@ class EAGLEWorkerV2(BaseSpecWorker): tgt_cache_loc, accepted_out_cache_loc ) + def _compact_accepted_to_front( + self, x: torch.Tensor, accept_index: torch.Tensor, bs: int + ) -> torch.Tensor: + """Gather the accepted tree path to the front of each per-req block. + + ``x`` is node-indexed over the whole tree (``[bs * num_draft_tokens, ...]``), + ``accept_index`` is ``[bs, spec_steps + 1]`` global node indices (-1 padded). + Padded entries clamp to node 0 but land past accept_lens (never read); + trailing unaccepted slots stay and are freed as overshoot. + """ + nd = self.speculative_num_draft_tokens + s1 = accept_index.shape[1] # spec_steps + 1 + safe = accept_index.to(torch.int64).clamp(min=0).reshape(-1) + gathered = x[safe] + out = x.clone() + out.view(bs, nd, *x.shape[1:])[:, :s1] = gathered.view(bs, s1, *x.shape[1:]) + return out + def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput): success, message = self._draft_worker.draft_runner.update_weights_from_disk( recv_req.model_path, diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 725142669..4b74c65aa 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -790,11 +790,12 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker): if not batch.forward_mode.is_idle(): accept_tokens = predict[accept_index] bonus_tokens = torch.empty_like(accept_lens, dtype=torch.int32) + # stride = accept_tokens per-req width = accept_index.shape[1]. fill_bonus_tokens[(bs,)]( accept_tokens, accept_lens, bonus_tokens, - self.speculative_num_draft_tokens, + accept_index.shape[1], ) else: bonus_tokens = torch.empty((0,), device=self.device, dtype=torch.int32) diff --git a/test/registered/spec/eagle/test_spec_eagle_stress.py b/test/registered/spec/eagle/test_spec_eagle_stress.py index eff30e1db..af4730377 100644 --- a/test/registered/spec/eagle/test_spec_eagle_stress.py +++ b/test/registered/spec/eagle/test_spec_eagle_stress.py @@ -20,7 +20,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=600, stage="base-b", runner_config="1-gpu-large") +register_cuda_ci(est_time=780, stage="base-b", runner_config="1-gpu-large") class TestEagle3Perf(Eagle3Base, SpecPerfKit): @@ -40,6 +40,20 @@ class TestEagleLlama2Retract(EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit): ) +class TestEagle3Topk16V2Retract(Eagle3Base, SpecAccuracyKit, SpecFeatureKit): + """EAGLE3 topk=16 tree on spec v2 under retract; must not leak KV. Stresses + the accepted-path KV move (move_accepted_tokens_to_target_kvcache).""" + + spec_topk = 16 + spec_tokens = 64 + disable_overlap = False + cuda_graph_max_bs = 5 + max_running_requests = 64 + gsm8k_accept_len_thres = 2.4 + extra_args = ("--max-total-tokens", 4500) # small KV to trigger retract + env_overrides = ((envs.SGLANG_TEST_RETRACT, True),) + + class TestEagleLlama2AbortAll(EagleLlama2Base, AbortAllMixin): abort_all_max_new_tokens = 4000 diff --git a/test/registered/spec/eagle/test_spec_eagle_topk.py b/test/registered/spec/eagle/test_spec_eagle_topk.py index 2b3914bba..369a5bb7f 100644 --- a/test/registered/spec/eagle/test_spec_eagle_topk.py +++ b/test/registered/spec/eagle/test_spec_eagle_topk.py @@ -1,7 +1,9 @@ """topk > 1 tree drafting (EAGLE3 topk16 + EAGLE/Llama-2 topk8). -topk > 1 always routes to spec v1; flashinfer is pinned (topk > 1 can't use fa3). -Runs on the cheap (5090) runner -- functional sanity only, no perf/stress. +topk > 1 routes to spec v1, except page_size==1 which can also stay on spec v2 +(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. +(topk > 1 on fa3 is covered on the Hopper runner in test_spec_eagle_fa3.py.) """ import unittest @@ -16,7 +18,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=840, stage="base-b", runner_config="1-gpu-small") +register_cuda_ci(est_time=1180, stage="base-b", runner_config="1-gpu-small") class TestEagle3Topk16(Eagle3Base, SpecCorrectnessKit, SpecAccuracyKit, SpecLogprobKit): @@ -31,6 +33,13 @@ class TestEagle3Topk16(Eagle3Base, SpecCorrectnessKit, SpecAccuracyKit, SpecLogp gsm8k_accept_len_thres = 2.4 # EAGLE3 topk16 gsm8k accept ~2.48 +class TestEagle3Topk16SpecV2(TestEagle3Topk16, SpecFeatureKit): + """EAGLE3 topk=16 tree on spec v2 (overlap, page1): guards the v2 tree path's + accepted-path compaction, validated by logprob_spec_v2_match.""" + + disable_overlap = False + + class TestEagleLlama2Suite( EagleLlama2Base, SpecCorrectnessKit,