[Spec] Naming cleanup: contiguous draft-loc kernel + accepted->accept (#27599)

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,

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(
+        self.move_accept_tokens_to_target_kvcache(batch, accept_index, accept_lens - 1)
-            batch, accept_index, accept_lens - 1
+        predict = self._compact_accept_to_front(predict, accept_index, bs)
-        )
-        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 = 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(
+        accept_out_cache_loc = torch.zeros(size, dtype=torch.int64, device=self.device)
-            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.

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]
+            current_token = draft_tokens[curr]
-            if vocab_size and curr_token_id >= vocab_size:
+            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:

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,

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)

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