diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index ae423fc53..1c38fcb40 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -278,14 +278,10 @@ 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 topk > 1 currently requires page_size == 1" + spec_v1_reason = "spec v2 currently only supports topk = 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 2ee5eab12..88e0bd03b 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, - accept_stride: tl.constexpr, + num_draft_tokens: tl.constexpr, ): # NOTE: we cannot fuse any in-place operations of `accept_lens` inside this kernel # because this kernel reads accept_lens @@ -514,8 +514,7 @@ def fill_bonus_tokens( # `accept_lens` includes the bonus token; the last accepted slot is at -1. accept_len = tl.load(accept_lens + pid) - # accept_stride = per-req width of accept_tokens (= accept_index.shape[1]). - bonus_token_idx = accept_stride * pid + accept_len - 1 + bonus_token_idx = num_draft_tokens * 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 6e606ac37..4f367a256 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -1232,13 +1232,11 @@ 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, - accept_index.shape[1], + self.speculative_num_draft_tokens, ) else: bonus_tokens = torch.empty((0,), device=self.device, dtype=torch.int32) @@ -1248,13 +1246,6 @@ 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 @@ -1336,30 +1327,6 @@ 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, @@ -1376,9 +1343,7 @@ class EAGLEWorkerV2(BaseSpecWorker): seq_lens is advanced by ``num_correct_drafts + 1`` to cover the bonus slot. """ bs = len(batch.seq_lens) - # 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] + size = bs * self.speculative_num_draft_tokens # fill_accepted_out_cache_loc reads out_cache_loc[accept_index]; -1 sentinel ok. maybe_detect_oob( @@ -1415,24 +1380,6 @@ 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 4b74c65aa..725142669 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -790,12 +790,11 @@ 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, - accept_index.shape[1], + self.speculative_num_draft_tokens, ) 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 af4730377..eff30e1db 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=780, stage="base-b", runner_config="1-gpu-large") +register_cuda_ci(est_time=600, stage="base-b", runner_config="1-gpu-large") class TestEagle3Perf(Eagle3Base, SpecPerfKit): @@ -40,20 +40,6 @@ 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 369a5bb7f..2b3914bba 100644 --- a/test/registered/spec/eagle/test_spec_eagle_topk.py +++ b/test/registered/spec/eagle/test_spec_eagle_topk.py @@ -1,9 +1,7 @@ """topk > 1 tree drafting (EAGLE3 topk16 + EAGLE/Llama-2 topk8). -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.) +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. """ import unittest @@ -18,7 +16,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=1180, stage="base-b", runner_config="1-gpu-small") +register_cuda_ci(est_time=840, stage="base-b", runner_config="1-gpu-small") class TestEagle3Topk16(Eagle3Base, SpecCorrectnessKit, SpecAccuracyKit, SpecLogprobKit): @@ -33,13 +31,6 @@ 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,