[Feature] Support for request-level speculative decoding metrics monitoring. (#5518)

* support spec metrics monitor per request

* fix bug

* remove debug log

* fix ut bugs
This commit is contained in:
GoldPancake
2025-12-12 12:22:18 +08:00
committed by GitHub
parent 3c1f7b85a4
commit 909059c60a
11 changed files with 154 additions and 47 deletions
@@ -261,6 +261,7 @@ class OpenAIServingCompletion:
aggregated_token_ids = [[] for _ in range(num_choices)]
aggregated_prompt_logprobs_tensors = [None] * num_choices
completion_batched_token_ids = [[] for _ in range(num_choices)]
aggregated_speculate_metrics = [None] * num_choices
current_waiting_time = 0
while num_choices > 0:
if self.engine_client.check_model_weight_status():
@@ -315,12 +316,18 @@ class OpenAIServingCompletion:
)
output_tokens[rid] += len(data["outputs"]["token_ids"])
completion_batched_token_ids[rid].extend(data["outputs"]["token_ids"])
output_speculate_metrics = data["metrics"].get("speculate_metrics", None)
if output_speculate_metrics is not None:
aggregated_speculate_metrics[rid] = output_speculate_metrics
if data.get("finished", False):
data["output_token_ids"] = output_tokens[rid]
data["outputs"]["top_logprobs"] = aggregated_top_logprobs[rid]
data["outputs"]["draft_top_logprobs"] = aggregated_draft_top_logprobs[rid]
data["outputs"]["token_ids"] = aggregated_token_ids[rid]
data["prompt_logprobs_tensors"] = aggregated_prompt_logprobs_tensors[rid]
data["speculate_metrics"] = aggregated_speculate_metrics[rid]
valid_results[rid] = data
num_choices -= 1
break
@@ -512,6 +519,7 @@ class OpenAIServingCompletion:
output_tokens[idx] += output.get("num_image_tokens")
num_image_tokens[idx] += output.get("num_image_tokens")
reasoning_tokens[idx] += output.get("reasoning_token_num", 0)
output_speculate_metrics = res["metrics"].get("speculate_metrics", None)
delta_message = CompletionResponseStreamChoice(
index=idx,
text=output["text"],
@@ -524,6 +532,7 @@ class OpenAIServingCompletion:
logprobs=logprobs_res,
prompt_logprobs=clamp_prompt_logprobs(prompt_logprobs_res),
draft_logprobs=draft_logprobs_res,
speculate_metrics=output_speculate_metrics,
)
if not res["finished"] and "delta_message" in output:
delta_message_output = output["delta_message"]
@@ -686,6 +695,7 @@ class OpenAIServingCompletion:
draft_logprobs=aggregated_draft_logprobs,
prompt_logprobs=clamp_prompt_logprobs(prompt_logprobs_res),
finish_reason=finish_reason,
speculate_metrics=final_res["metrics"].get("speculate_metrics", None),
)
choices.append(choice_data)