Add per-request retraction count (#11177)

This commit is contained in:
Scott Lee
2025-10-27 23:22:34 -07:00
committed by GitHub
parent 81a632ace6
commit bf8f7a944f
7 changed files with 69 additions and 0 deletions
@@ -730,6 +730,7 @@ class SchedulerOutputProcessorMixin:
cached_tokens = []
spec_verify_ct = []
spec_accepted_tokens = []
retraction_counts = []
output_hidden_states = None
if return_logprob:
@@ -831,6 +832,8 @@ class SchedulerOutputProcessorMixin:
completion_tokens.append(len(output_ids_))
cached_tokens.append(req.cached_tokens)
retraction_counts.append(req.retraction_count)
if not self.spec_algorithm.is_none():
spec_verify_ct.append(req.spec_verify_ct)
spec_accepted_tokens.append(req.spec_accepted_tokens)
@@ -953,6 +956,7 @@ class SchedulerOutputProcessorMixin:
http_worker_ipcs=http_worker_ipcs,
placeholder_tokens_idx=None,
placeholder_tokens_val=None,
retraction_counts=retraction_counts,
)
)
@@ -964,6 +968,7 @@ class SchedulerOutputProcessorMixin:
embeddings = []
prompt_tokens = []
cached_tokens = []
retraction_counts = []
for req in reqs:
if req.finished():
rids.append(req.rid)
@@ -972,6 +977,7 @@ class SchedulerOutputProcessorMixin:
embeddings.append(req.embedding)
prompt_tokens.append(len(req.origin_input_ids))
cached_tokens.append(req.cached_tokens)
retraction_counts.append(req.retraction_count)
self.send_to_detokenizer.send_output(
BatchEmbeddingOutput(
finished_reasons,
@@ -982,5 +988,6 @@ class SchedulerOutputProcessorMixin:
http_worker_ipcs=http_worker_ipcs,
placeholder_tokens_idx=None,
placeholder_tokens_val=None,
retraction_counts=retraction_counts,
)
)