From 3fe6bc390bdcfc45873e796f5c1351f2001a91ae Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Mon, 8 Jun 2026 15:04:58 -0700 Subject: [PATCH] [Spec] Naming cleanup: contiguous draft-loc kernel + `accepted`->`accept` (#27599) --- .../sglang/srt/speculative/eagle_info_v2.py | 6 ++-- .../sglang/srt/speculative/eagle_worker_v2.py | 36 +++++++++---------- python/sglang/srt/speculative/spec_utils.py | 6 ++-- .../srt/speculative/triton_ops/cache_locs.py | 2 +- .../srt/speculative/triton_ops/eagle.py | 6 ++-- .../spec/eagle/test_spec_eagle_stress.py | 2 +- 6 files changed, 27 insertions(+), 31 deletions(-) diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index e458cab29..bbe7ed9c9 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -37,7 +37,7 @@ from sglang.srt.speculative.spec_utils import ( generate_simulated_accept_index, ) from sglang.srt.speculative.triton_ops.cache_locs import ( - assign_draft_cache_locs_page_size_1 as assign_draft_cache_locs_page_size_1, + assign_draft_cache_locs_contiguous as assign_draft_cache_locs_contiguous, ) from sglang.srt.speculative.triton_ops.cache_locs import ( assign_extend_cache_locs as assign_extend_cache_locs, @@ -46,7 +46,7 @@ from sglang.srt.speculative.triton_ops.cache_locs import ( assign_extend_cache_locs_func as assign_extend_cache_locs_func, ) from sglang.srt.speculative.triton_ops.eagle import ( - fill_accepted_out_cache_loc as fill_accepted_out_cache_loc, + fill_accept_out_cache_loc as fill_accept_out_cache_loc, ) from sglang.srt.speculative.triton_ops.eagle import ( fill_bonus_tokens as fill_bonus_tokens, @@ -241,7 +241,7 @@ class EagleDraftInputV2Mixin: device=batch.device, ) # FIXME(lsyin): align with the default code path - assign_draft_cache_locs_page_size_1[(bs,)]( + assign_draft_cache_locs_contiguous[(bs,)]( batch.req_pool_indices, req_to_token_pool.req_to_token, batch.seq_lens, diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 8468c17b1..6bf5d6182 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -53,7 +53,7 @@ from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import ( from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput from sglang.srt.speculative.eagle_info_v2 import ( assign_extend_cache_locs, - fill_accepted_out_cache_loc, + fill_accept_out_cache_loc, fill_bonus_tokens, ) from sglang.srt.speculative.eagle_utils import ( @@ -1300,7 +1300,7 @@ class EAGLEWorkerV2(BaseSpecWorker): 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( + predict = self._finalize_accept_tree_path( batch, accept_index, accept_lens, predict, logits_output, bs ) @@ -1333,7 +1333,7 @@ class EAGLEWorkerV2(BaseSpecWorker): """Update mamba state for hybrid GDN models after verification.""" # `accept_lens` already includes the bonus token (drafts + 1 per req). if not batch.forward_mode.is_idle() and accept_index.numel() > 0: - accepted_indices_offset = torch.arange( + accept_indices_offset = torch.arange( 0, bs * self.speculative_num_draft_tokens, step=self.speculative_num_draft_tokens, @@ -1345,7 +1345,7 @@ class EAGLEWorkerV2(BaseSpecWorker): # mamba state to commit; reduces to accept_lens - 1 for topk == 1. last_correct_step_indices = ( accept_index[req_idx, (accept_lens - 1).to(torch.int64)] - - accepted_indices_offset + - accept_indices_offset ) if batch.mamba_track_indices is not None: @@ -1365,7 +1365,7 @@ class EAGLEWorkerV2(BaseSpecWorker): tracking_point - seq_lens_pre_verify - 1, min=0 ).to(torch.int64) candidate_track_steps = ( - accept_index[req_idx, to_track_ith] - accepted_indices_offset + accept_index[req_idx, to_track_ith] - accept_indices_offset ) mamba_steps_to_track = torch.where( to_track_mask, @@ -1382,7 +1382,7 @@ class EAGLEWorkerV2(BaseSpecWorker): model=self.target_worker.model_runner.model, ) - def _finalize_accepted_tree_path( + def _finalize_accept_tree_path( self, batch: ScheduleBatch, accept_index: torch.Tensor, @@ -1396,17 +1396,15 @@ class EAGLEWorkerV2(BaseSpecWorker): 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) + self.move_accept_tokens_to_target_kvcache(batch, accept_index, accept_lens - 1) + predict = self._compact_accept_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 = self._compact_accept_to_front( logits_output.hidden_states, accept_index, bs ) return predict - def move_accepted_tokens_to_target_kvcache( + def move_accept_tokens_to_target_kvcache( self, batch: ScheduleBatch, accept_index: torch.Tensor, @@ -1426,7 +1424,7 @@ class EAGLEWorkerV2(BaseSpecWorker): # 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. + # fill_accept_out_cache_loc reads out_cache_loc[accept_index]; -1 sentinel ok. maybe_detect_oob( accept_index, -1, @@ -1439,9 +1437,7 @@ class EAGLEWorkerV2(BaseSpecWorker): dtype=torch.int64, device=self.device, ) - accepted_out_cache_loc = torch.zeros( - size, dtype=torch.int64, device=self.device - ) + accept_out_cache_loc = torch.zeros(size, dtype=torch.int64, device=self.device) assign_extend_cache_locs[(bs,)]( batch.req_pool_indices, self.req_to_token_pool.req_to_token, @@ -1451,17 +1447,17 @@ class EAGLEWorkerV2(BaseSpecWorker): self.req_to_token_pool.req_to_token.shape[1], next_power_of_2(bs), ) - fill_accepted_out_cache_loc[(size,)]( + fill_accept_out_cache_loc[(size,)]( accept_index, batch.out_cache_loc, - accepted_out_cache_loc, + accept_out_cache_loc, next_power_of_2(size), ) self.token_to_kv_pool_allocator.get_kvcache().move_kv_cache( - tgt_cache_loc, accepted_out_cache_loc + tgt_cache_loc, accept_out_cache_loc ) - def _compact_accepted_to_front( + def _compact_accept_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. diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index abdee622b..48d07334a 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -335,13 +335,13 @@ def traverse_tree( is_accepted = True else: parent_bitmask = allocate_token_bitmask[parent_pos] - curr_token_id = draft_tokens[curr] - if vocab_size and curr_token_id >= vocab_size: + current_token = draft_tokens[curr] + if vocab_size and current_token >= vocab_size: is_accepted = False else: # 32 boolean bitmask values are packed into 32-bit integers is_accepted = ( - parent_bitmask[curr_token_id // 32] & (1 << (curr_token_id % 32)) + parent_bitmask[current_token // 32] & (1 << (current_token % 32)) ) != 0 if is_accepted: diff --git a/python/sglang/srt/speculative/triton_ops/cache_locs.py b/python/sglang/srt/speculative/triton_ops/cache_locs.py index e8a6a754d..e16a15ad5 100644 --- a/python/sglang/srt/speculative/triton_ops/cache_locs.py +++ b/python/sglang/srt/speculative/triton_ops/cache_locs.py @@ -94,7 +94,7 @@ def assign_req_to_token_pool_func( @triton.jit -def assign_draft_cache_locs_page_size_1( +def assign_draft_cache_locs_contiguous( req_pool_indices, req_to_token, seq_lens, diff --git a/python/sglang/srt/speculative/triton_ops/eagle.py b/python/sglang/srt/speculative/triton_ops/eagle.py index bd80d9a6e..de9b2ec4c 100644 --- a/python/sglang/srt/speculative/triton_ops/eagle.py +++ b/python/sglang/srt/speculative/triton_ops/eagle.py @@ -22,10 +22,10 @@ def fill_bonus_tokens( @triton.jit -def fill_accepted_out_cache_loc( +def fill_accept_out_cache_loc( accept_index, out_cache_loc, - accepted_out_cache_loc, + accept_out_cache_loc, size_upper: tl.constexpr, ): pid = tl.program_id(axis=0) @@ -36,4 +36,4 @@ def fill_accepted_out_cache_loc( src = tl.load(accept_index + pid) if src > -1: value = tl.load(out_cache_loc + src) - tl.store(accepted_out_cache_loc + dst, value) + tl.store(accept_out_cache_loc + dst, value) diff --git a/test/registered/spec/eagle/test_spec_eagle_stress.py b/test/registered/spec/eagle/test_spec_eagle_stress.py index 40e54f585..353d8a83b 100644 --- a/test/registered/spec/eagle/test_spec_eagle_stress.py +++ b/test/registered/spec/eagle/test_spec_eagle_stress.py @@ -43,7 +43,7 @@ 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).""" + the accepted-path KV move (move_accept_tokens_to_target_kvcache).""" spec_topk = 16 spec_tokens = 64