[Spec] Rename accepted_indices -> accept_indices; drop _token_id suffix per Rule 5 (#25038)

This commit is contained in:
Liangsheng Yin
2026-05-11 22:29:02 -07:00
committed by GitHub
parent ecf5d844f5
commit 1f985c50db
8 changed files with 51 additions and 51 deletions
+3 -3
View File
@@ -346,14 +346,14 @@ def add_output_logprobs_for_spec_v1(
top_logprobs_nums = batch.top_logprobs_nums top_logprobs_nums = batch.top_logprobs_nums
token_ids_logprobs = batch.token_ids_logprobs token_ids_logprobs = batch.token_ids_logprobs
accepted_indices = res.accepted_indices accept_indices = res.accept_indices
assert len(accepted_indices) == len(logits_output.next_token_logits) assert len(accept_indices) == len(logits_output.next_token_logits)
temperatures = batch.sampling_info.temperatures temperatures = batch.sampling_info.temperatures
num_draft_tokens = batch.spec_info.draft_token_num num_draft_tokens = batch.spec_info.draft_token_num
# acceptance indices are the indices in a "flattened" batch. # acceptance indices are the indices in a "flattened" batch.
# dividing it to num_draft_tokens will yield the actual batch index. # dividing it to num_draft_tokens will yield the actual batch index.
temperatures = temperatures[accepted_indices // num_draft_tokens] temperatures = temperatures[accept_indices // num_draft_tokens]
if envs.SGLANG_RETURN_ORIGINAL_LOGPROB.get(): if envs.SGLANG_RETURN_ORIGINAL_LOGPROB.get():
logprobs = torch.nn.functional.log_softmax( logprobs = torch.nn.functional.log_softmax(
logits_output.next_token_logits, dim=-1 logits_output.next_token_logits, dim=-1
@@ -406,7 +406,7 @@ class SchedulerOutputProcessorMixin:
dp_cooperation_info=batch.dp_cooperation_info, dp_cooperation_info=batch.dp_cooperation_info,
) )
def _resolve_spec_overlap_token_ids( def _resolve_spec_overlap_tokens(
self: Scheduler, result: GenerationBatchResult, batch: ScheduleBatch self: Scheduler, result: GenerationBatchResult, batch: ScheduleBatch
) -> List[List[int]]: ) -> List[List[int]]:
"""Resolve the padding next token ids for speculative decoding with overlap.""" """Resolve the padding next token ids for speculative decoding with overlap."""
@@ -487,7 +487,7 @@ class SchedulerOutputProcessorMixin:
if batch.spec_algorithm.is_none() or batch.is_spec_v2: if batch.spec_algorithm.is_none() or batch.is_spec_v2:
if batch.is_spec_v2: if batch.is_spec_v2:
next_token_ids = self._resolve_spec_overlap_token_ids(result, batch) next_token_ids = self._resolve_spec_overlap_tokens(result, batch)
elif isinstance(next_token_ids, list): elif isinstance(next_token_ids, list):
pass # MLX path: already a list[int], skip torch round-trip pass # MLX path: already a list[int], skip torch round-trip
else: else:
@@ -738,7 +738,7 @@ class DFlashWorker:
added_vocab_start = int(shard.added_vocab_start_index) added_vocab_start = int(shard.added_vocab_start_index)
num_tokens = int(hidden_states.shape[0]) num_tokens = int(hidden_states.shape[0])
out_token_ids = torch.empty( out_tokens = torch.empty(
(num_tokens,), dtype=torch.long, device=hidden_states.device (num_tokens,), dtype=torch.long, device=hidden_states.device
) )
@@ -753,13 +753,13 @@ class DFlashWorker:
hs = _cast_hs(hidden_states[start:end]) hs = _cast_hs(hidden_states[start:end])
if num_org > 0: if num_org > 0:
base_logits = torch.matmul(hs, weight[:num_org].T) base_logits = torch.matmul(hs, weight[:num_org].T)
out_token_ids[start:end] = ( out_tokens[start:end] = (
torch.argmax(base_logits, dim=-1).to(torch.long) torch.argmax(base_logits, dim=-1).to(torch.long)
+ org_vocab_start + org_vocab_start
) )
else: else:
out_token_ids[start:end] = 0 out_tokens[start:end] = 0
return out_token_ids return out_tokens
for start in range(0, num_tokens, int(chunk_size)): for start in range(0, num_tokens, int(chunk_size)):
end = min(num_tokens, start + int(chunk_size)) end = min(num_tokens, start + int(chunk_size))
@@ -812,7 +812,7 @@ class DFlashWorker:
) )
if tp_size == 1: if tp_size == 1:
out_token_ids[start:end] = global_ids.to(torch.long) out_tokens[start:end] = global_ids.to(torch.long)
continue continue
# Gather per-rank maxima and associated global ids, then select the global max. # Gather per-rank maxima and associated global ids, then select the global max.
@@ -869,9 +869,9 @@ class DFlashWorker:
rank_index[0].copy_(best_rank) rank_index[0].copy_(best_rank)
selected_ids = self._draft_greedy_selected_ids_buf[:, :chunk_len] selected_ids = self._draft_greedy_selected_ids_buf[:, :chunk_len]
torch.gather(gathered_ids, 0, rank_index, out=selected_ids) torch.gather(gathered_ids, 0, rank_index, out=selected_ids)
out_token_ids[start:end].copy_(selected_ids.view(-1)) out_tokens[start:end].copy_(selected_ids.view(-1))
return out_token_ids return out_tokens
def _append_target_hidden_to_draft_kv( def _append_target_hidden_to_draft_kv(
self, self,
+4 -4
View File
@@ -574,7 +574,7 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
logits_output=logits_output, logits_output=logits_output,
accept_tokens=accept_tokens, accept_tokens=accept_tokens,
num_correct_drafts_per_req_cpu=num_correct_drafts_list, num_correct_drafts_per_req_cpu=num_correct_drafts_list,
accepted_indices=accept_index, accept_indices=accept_index,
) )
else: else:
if page_size == 1 or self.topk == 1: if page_size == 1 or self.topk == 1:
@@ -651,7 +651,7 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
logits_output=logits_output, logits_output=logits_output,
accept_tokens=accept_tokens, accept_tokens=accept_tokens,
num_correct_drafts_per_req_cpu=num_correct_drafts_list, num_correct_drafts_per_req_cpu=num_correct_drafts_list,
accepted_indices=accept_index, accept_indices=accept_index,
) )
@@ -972,7 +972,7 @@ class EagleVerifyOutput:
# Accepted token length per sequence in a batch in CPU (full set). # Accepted token length per sequence in a batch in CPU (full set).
num_correct_drafts_per_req_cpu: List[int] num_correct_drafts_per_req_cpu: List[int]
# Accepted indices from logits_output.next_token_logits # Accepted indices from logits_output.next_token_logits
accepted_indices: torch.Tensor accept_indices: torch.Tensor
@classmethod @classmethod
def create_idle( def create_idle(
@@ -988,7 +988,7 @@ class EagleVerifyOutput:
logits_output=logits_output, logits_output=logits_output,
accept_tokens=torch.empty(0, dtype=torch.long, device=device), accept_tokens=torch.empty(0, dtype=torch.long, device=device),
num_correct_drafts_per_req_cpu=[], num_correct_drafts_per_req_cpu=[],
accepted_indices=torch.full( accept_indices=torch.full(
(0, spec_steps + 1), -1, dtype=torch.int32, device=device (0, spec_steps + 1), -1, dtype=torch.int32, device=device
), ),
) )
@@ -971,9 +971,9 @@ class EAGLEWorker(TpModelWorker):
# Post process based on verified outputs. # Post process based on verified outputs.
# Pick indices that we care (accepted) # Pick indices that we care (accepted)
logits_output.next_token_logits = logits_output.next_token_logits[ logits_output.next_token_logits = logits_output.next_token_logits[
res.accepted_indices res.accept_indices
] ]
logits_output.hidden_states = logits_output.hidden_states[res.accepted_indices] logits_output.hidden_states = logits_output.hidden_states[res.accept_indices]
if ( if (
self.target_worker.model_runner.hybrid_gdn_config is not None self.target_worker.model_runner.hybrid_gdn_config is not None
@@ -1033,16 +1033,16 @@ class EAGLEWorker(TpModelWorker):
) )
# If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask # If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask
# res.accepted_indices.shape[0] > 0 skips DP attn idle batch # res.accept_indices.shape[0] > 0 skips DP attn idle batch
if spec_info.topk > 1 and res.accepted_indices.shape[0] > 0: if spec_info.topk > 1 and res.accept_indices.shape[0] > 0:
# accepted_indices=[0,2,3,4,5,7,9,10,11], num_accept_tokens=[4, 3, 2], cumulative_num_accept_tokens=[4, 7, 9] # accept_indices=[0,2,3,4,5,7,9,10,11], num_accept_tokens=[4, 3, 2], cumulative_num_accept_tokens=[4, 7, 9]
# first_token_indices_per_req=prepend(0, accepted_indices[cumulative_num_accept_tokens[:-1]]) = [0, 5, 10] # first_token_indices_per_req=prepend(0, accept_indices[cumulative_num_accept_tokens[:-1]]) = [0, 5, 10]
# last_token_indices_per_req=accepted_indices[cumulative_num_accept_tokens - 1] = [4, 9, 11] (last token ID of each req) # last_token_indices_per_req=accept_indices[cumulative_num_accept_tokens - 1] = [4, 9, 11] (last token ID of each req)
# last_correct_step_indices = [4,4,1]; those are the per-req spec-decoding step offsets that contain the correct mamba caches # last_correct_step_indices = [4,4,1]; those are the per-req spec-decoding step offsets that contain the correct mamba caches
# equivalent: last_correct_step_indices = last_token_indices_per_req - first_token_indices_per_req; # equivalent: last_correct_step_indices = last_token_indices_per_req - first_token_indices_per_req;
# `accepted_indices_offset` equals `first_token_indices_per_req` because the first accepted slot of each req is its "current token" at logical position i * draft_token_num. # `accepted_indices_offset` equals `first_token_indices_per_req` because the first accepted slot of each req is its "current token" at logical position i * draft_token_num.
last_correct_step_indices = ( last_correct_step_indices = (
res.accepted_indices[cumulative_num_accept_tokens - 1] res.accept_indices[cumulative_num_accept_tokens - 1]
- accepted_indices_offset - accepted_indices_offset
) )
else: else:
@@ -1062,7 +1062,7 @@ class EAGLEWorker(TpModelWorker):
to_track_ith = torch.clamp(tracking_point - seq_lens_pre_verify - 1, min=0) to_track_ith = torch.clamp(tracking_point - seq_lens_pre_verify - 1, min=0)
mamba_steps_to_track = torch.where( mamba_steps_to_track = torch.where(
to_track_mask, to_track_mask,
res.accepted_indices[to_track_ith + accepted_indices_start] res.accept_indices[to_track_ith + accepted_indices_start]
- accepted_indices_offset, - accepted_indices_offset,
-1, -1,
) )
@@ -757,9 +757,9 @@ class FrozenKVMTPWorker(TpModelWorker):
) )
logits_output.next_token_logits = logits_output.next_token_logits[ logits_output.next_token_logits = logits_output.next_token_logits[
res.accepted_indices res.accept_indices
] ]
logits_output.hidden_states = logits_output.hidden_states[res.accepted_indices] logits_output.hidden_states = logits_output.hidden_states[res.accept_indices]
if ( if (
self.target_worker.model_runner.hybrid_gdn_config is not None self.target_worker.model_runner.hybrid_gdn_config is not None
@@ -556,9 +556,9 @@ class MultiLayerEagleWorker(TpModelWorker):
# Post process based on verified outputs. # Post process based on verified outputs.
# Pick indices that we care (accepted) # Pick indices that we care (accepted)
logits_output.next_token_logits = logits_output.next_token_logits[ logits_output.next_token_logits = logits_output.next_token_logits[
res.accepted_indices res.accept_indices
] ]
logits_output.hidden_states = logits_output.hidden_states[res.accepted_indices] logits_output.hidden_states = logits_output.hidden_states[res.accept_indices]
if self.target_worker.model_runner.hybrid_gdn_config is not None: if self.target_worker.model_runner.hybrid_gdn_config is not None:
num_accept_tokens = ( num_accept_tokens = (
@@ -571,11 +571,11 @@ class MultiLayerEagleWorker(TpModelWorker):
) )
# If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask # If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask
# res.accepted_indices.shape[0] > 0 skips DP attn idle batch # res.accept_indices.shape[0] > 0 skips DP attn idle batch
if spec_info.topk > 1 and res.accepted_indices.shape[0] > 0: if spec_info.topk > 1 and res.accept_indices.shape[0] > 0:
# accepted_indices=[0,2,3,4,5,7,9,10,11], num_accept_tokens=[4, 3, 2], cumulative_num_accept_tokens=[4, 7, 9] # accept_indices=[0,2,3,4,5,7,9,10,11], num_accept_tokens=[4, 3, 2], cumulative_num_accept_tokens=[4, 7, 9]
# first_token_indices_per_req=prepend(0, accepted_indices[cumulative_num_accept_tokens[:-1]]) = [0, 5, 10] # first_token_indices_per_req=prepend(0, accept_indices[cumulative_num_accept_tokens[:-1]]) = [0, 5, 10]
# last_token_indices_per_req=accepted_indices[cumulative_num_accept_tokens - 1] = [4, 9, 11] (last token ID of each req) # last_token_indices_per_req=accept_indices[cumulative_num_accept_tokens - 1] = [4, 9, 11] (last token ID of each req)
# last_correct_step_indices = [4,4,1]; those are the per-req spec-decoding step offsets that contain the correct mamba caches # last_correct_step_indices = [4,4,1]; those are the per-req spec-decoding step offsets that contain the correct mamba caches
cumulative_num_accept_tokens = torch.cumsum(num_accept_tokens, dim=0) cumulative_num_accept_tokens = torch.cumsum(num_accept_tokens, dim=0)
req_start_positions = torch.cat( req_start_positions = torch.cat(
@@ -588,8 +588,8 @@ class MultiLayerEagleWorker(TpModelWorker):
cumulative_num_accept_tokens[:-1], cumulative_num_accept_tokens[:-1],
] ]
) )
first_token_indices_per_req = res.accepted_indices[req_start_positions] first_token_indices_per_req = res.accept_indices[req_start_positions]
last_token_indices_per_req = res.accepted_indices[ last_token_indices_per_req = res.accept_indices[
cumulative_num_accept_tokens - 1 cumulative_num_accept_tokens - 1
] ]
last_correct_step_indices = ( last_correct_step_indices = (
+15 -15
View File
@@ -157,7 +157,7 @@ class NgramVerifyInput(SpecInput):
batch: ScheduleBatch, batch: ScheduleBatch,
logits_output: torch.Tensor, logits_output: torch.Tensor,
): ):
accept_index_cpu = self.accepted_indices.tolist() accept_index_cpu = self.accept_indices.tolist()
predict_cpu = self.predict.tolist() predict_cpu = self.predict.tolist()
has_finished = False has_finished = False
think_end_id = batch.model_config.think_end_id think_end_id = batch.model_config.think_end_id
@@ -176,7 +176,7 @@ class NgramVerifyInput(SpecInput):
if req.finished(): if req.finished():
has_finished = True has_finished = True
# set all tokens after finished token to -1 and break # set all tokens after finished token to -1 and break
self.accepted_indices[i, j + 1 :] = -1 self.accept_indices[i, j + 1 :] = -1
break break
else: else:
if req.grammar is not None: if req.grammar is not None:
@@ -185,7 +185,7 @@ class NgramVerifyInput(SpecInput):
except ValueError as e: except ValueError as e:
logger.info( logger.info(
f"{i=}, {req=}\n" f"{i=}, {req=}\n"
f"{self.accepted_indices=}\n" f"{self.accept_indices=}\n"
f"{self.predict=}\n" f"{self.predict=}\n"
) )
raise e raise e
@@ -197,17 +197,17 @@ class NgramVerifyInput(SpecInput):
req.update_spec_correct_drafts_histogram(num_correct_drafts_this_req) req.update_spec_correct_drafts_histogram(num_correct_drafts_this_req)
if has_finished: if has_finished:
self.num_correct_drafts = (self.accepted_indices != -1).sum(dim=1) - 1 self.num_correct_drafts = (self.accept_indices != -1).sum(dim=1) - 1
self.accepted_indices = self.accepted_indices[self.accepted_indices != -1] self.accept_indices = self.accept_indices[self.accept_indices != -1]
logits_output.next_token_logits = logits_output.next_token_logits[ logits_output.next_token_logits = logits_output.next_token_logits[
self.accepted_indices self.accept_indices
] ]
if logits_output.hidden_states: if logits_output.hidden_states:
logits_output.hidden_states = logits_output.hidden_states[ logits_output.hidden_states = logits_output.hidden_states[
self.accepted_indices self.accept_indices
] ]
self.accept_tokens = self.predict[self.accepted_indices] self.accept_tokens = self.predict[self.accept_indices]
def _free_cache( def _free_cache(
self, self,
@@ -220,16 +220,16 @@ class NgramVerifyInput(SpecInput):
if page_size == 1: if page_size == 1:
# TODO: boolean array index leads to a device sync. Remove it. # TODO: boolean array index leads to a device sync. Remove it.
evict_mask = torch.full_like(self.draft_token, True, dtype=torch.bool) evict_mask = torch.full_like(self.draft_token, True, dtype=torch.bool)
evict_mask[self.accepted_indices] = False evict_mask[self.accept_indices] = False
batch.token_to_kv_pool_allocator.free(batch.out_cache_loc[evict_mask]) batch.token_to_kv_pool_allocator.free(batch.out_cache_loc[evict_mask])
batch.out_cache_loc = batch.out_cache_loc[self.accepted_indices] batch.out_cache_loc = batch.out_cache_loc[self.accept_indices]
else: else:
# Shift the accepted tokens to the beginning. # Shift the accepted tokens to the beginning.
# Only evict the last part # Only evict the last part
src_cache_loc, tgt_cache_loc, to_free_num_slots = get_src_tgt_cache_loc( src_cache_loc, tgt_cache_loc, to_free_num_slots = get_src_tgt_cache_loc(
batch.seq_lens, batch.seq_lens,
batch.out_cache_loc, batch.out_cache_loc,
self.accepted_indices, self.accept_indices,
self.num_correct_drafts, self.num_correct_drafts,
self.draft_token_num, self.draft_token_num,
page_size, page_size,
@@ -297,7 +297,7 @@ class NgramVerifyInput(SpecInput):
predict_shape = list(logits_output.next_token_logits.shape)[:-1] predict_shape = list(logits_output.next_token_logits.shape)[:-1]
predict_shape[-1] += 1 predict_shape[-1] += 1
self.predict = torch.empty(predict_shape, dtype=torch.int32, device=self.device) self.predict = torch.empty(predict_shape, dtype=torch.int32, device=self.device)
self.accepted_indices = torch.full( self.accept_indices = torch.full(
(bs, self.draft_token_num), -1, dtype=torch.int32, device=self.device (bs, self.draft_token_num), -1, dtype=torch.int32, device=self.device
) )
self.num_correct_drafts = torch.empty( self.num_correct_drafts = torch.empty(
@@ -306,7 +306,7 @@ class NgramVerifyInput(SpecInput):
verify_tree_greedy( verify_tree_greedy(
predicts=self.predict, # mutable predicts=self.predict, # mutable
accept_index=self.accepted_indices, # mutable accept_index=self.accept_indices, # mutable
accept_token_num=self.num_correct_drafts, # mutable accept_token_num=self.num_correct_drafts, # mutable
candidates=candidates, candidates=candidates,
# kwarg LHS retained as `retrive_*` to match sgl_kernel op schema. # kwarg LHS retained as `retrive_*` to match sgl_kernel op schema.
@@ -327,7 +327,7 @@ class NgramVerifyInput(SpecInput):
predict_shape = list(logits_output.next_token_logits.shape)[:-1] predict_shape = list(logits_output.next_token_logits.shape)[:-1]
predict_shape[-1] += 1 predict_shape[-1] += 1
self.predict = torch.empty(predict_shape, dtype=torch.int32, device=self.device) self.predict = torch.empty(predict_shape, dtype=torch.int32, device=self.device)
self.accepted_indices = torch.full( self.accept_indices = torch.full(
(bs, self.draft_token_num), -1, dtype=torch.int32, device=self.device (bs, self.draft_token_num), -1, dtype=torch.int32, device=self.device
) )
self.num_correct_drafts = torch.empty( self.num_correct_drafts = torch.empty(
@@ -371,7 +371,7 @@ class NgramVerifyInput(SpecInput):
) )
tree_speculative_sampling_target_only( tree_speculative_sampling_target_only(
predicts=self.predict, # mutable predicts=self.predict, # mutable
accept_index=self.accepted_indices, # mutable accept_index=self.accept_indices, # mutable
accept_token_num=self.num_correct_drafts, # mutable accept_token_num=self.num_correct_drafts, # mutable
candidates=candidates.to(torch.int64), candidates=candidates.to(torch.int64),
# kwarg LHS retained as `retrive_*` to match sgl_kernel op schema. # kwarg LHS retained as `retrive_*` to match sgl_kernel op schema.