[None][feat] Add Prometheus metrics for prompt cache, speculative decoding, perplexity, and batch occupancy#12636
Conversation
f673921 to
68e4ece
Compare
📝 WalkthroughWalkthroughThe changes extend speculative decoding tracking and metrics collection across the inference pipeline. Per-position draft/accept counters are added to request tracking, integrated into response handling during execution, and exposed through Prometheus metrics with perplexity histograms. Changes
Sequence DiagramsequenceDiagram
participant Client
participant Executor as Executor<br/>(py_executor)
participant Request as Request<br/>(llm_request)
participant Result as Result<br/>(result)
participant Collector as MetricsCollector
Client->>Request: Create LlmRequest with<br/>per_pos_drafted/accepted
Request->>Request: Initialize state arrays<br/>[0] * MAX_SPEC_DECODE_POSITIONS
Executor->>Request: Process iteration<br/>with draft tokens
Note over Executor: Read num_draft_tokens<br/>& py_num_accepted_draft_tokens
Executor->>Request: Update py_per_pos_drafted[pos]<br/>for pos in draft_len
Executor->>Request: Update py_per_pos_accepted[pos]<br/>for pos in accepted
Executor->>Result: Create response with<br/>updated arrays
Result->>Result: record_stats(): Extract<br/>per_pos_drafted/accepted
Result->>Result: Compute prefill & generation<br/>perplexity from logprobs
Result->>Collector: Log metrics dict with<br/>per-position counters
Collector->>Collector: log_request_metrics_dict():<br/>Emit Prometheus counters
Note over Collector: SPEC_DEC_DRAFTED_PER_POS<br/>SPEC_DEC_ACCEPTED_PER_POS<br/>PREFILL_PERPLEXITY<br/>GENERATION_PERPLEXITY
Collector->>Client: Metrics recorded
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 3489-3495: The code incorrectly infers draft length by counting
non-zero values in py_draft_tokens (sum(1 for t in py_draft if t != 0)), which
misclassifies legitimate token ID 0 as padding; replace that inference with the
existing explicit bookkeeping by calling the imported get_draft_token_length (or
otherwise using the authoritative num_draft_tokens bookkeeping) to compute
draft_len instead of treating 0 as sentinel—update the branch that reads
py_draft_tokens so it uses get_draft_token_length(py_draft) or assigns the known
num_draft_tokens value, touching symbols py_draft_tokens, draft_len,
num_draft_tokens, and get_draft_token_length.
In `@tensorrt_llm/executor/result.py`:
- Around line 581-583: The current code only sets
metrics_stats[MetricNames.PROMPT_CACHE_CACHED_TOKENS] when self.cached_tokens >
0, which drops zero-values; change the logic so the metric key is always emitted
(e.g., always assign metrics_stats[MetricNames.PROMPT_CACHE_CACHED_TOKENS] =
self.cached_tokens or metrics_stats[MetricNames.PROMPT_CACHE_CACHED_TOKENS] =
int(self.cached_tokens or 0)) so that cache-miss (0) is recorded; update the
block around self.cached_tokens and metrics_stats to unconditionally add the key
instead of gating on > 0.
- Around line 621-623: The broad except Exception handlers around the prefill
perplexity computation should be narrowed to only catch predictable errors;
update the except blocks that currently read like the one logging "Failed to
compute prefill perplexity" (and the similar handler at the other location) to
except (ValueError, TypeError): so they still log with exc_info=True but no
longer mask unrelated exceptions raised by float(entry), math.exp(), or
type-mismatch arithmetic in the prefill perplexity calculation.
In `@tensorrt_llm/metrics/collector.py`:
- Around line 279-284: The current truthiness check on
MetricNames.PROMPT_CACHE_CACHED_TOKENS causes zero values to be skipped; change
the branch to detect presence (or non-None) instead of truthiness so zeros are
processed. Specifically, read cached_tokens from metrics_dict (use
metrics_dict.get(MetricNames.PROMPT_CACHE_CACHED_TOKENS) or check
MetricNames.PROMPT_CACHE_CACHED_TOKENS in metrics_dict) and if the key exists
(cached_tokens is not None) call
self._log_counter(self.counter_tokens_cached_prompt, self.labels, cached_tokens)
and self._log_histogram(self.histogram_tokens_cached_prompt, cached_tokens) so
that histogram updates include zero-cache requests.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 12948839-87e5-419a-8226-171f9cc2989a
📒 Files selected for processing (6)
tensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/executor/result.pytensorrt_llm/metrics/collector.pytensorrt_llm/metrics/enums.pytensorrt_llm/serve/openai_server.py
f6082e6 to
4868671
Compare
svc-trtllm-gh-bot
added
the
Community want to contribute
|
@vedularaghu, thank you for your contribution! |
dede667 to
c781447
Compare
Appreciate your feedback @karljang. I have removed the metrics which are merged along with the PR and updated my PR. |
JunyiXu-nv
left a comment
JunyiXu-nv
left a comment
There was a problem hiding this comment.
Overall LGTM. Could you add some tests to guard this new metrics? There are some references in tests/unittest/metrics/test_collector.py and tests/unittest/llmapi/apps/_test_openai_prometheus.py
Appreciate your review, @JunyiXu-nv. I have added tests to guard the new metrics. |
|
/bot run |
|
PR_Github #48635 [ run ] triggered by Bot. Commit: |
|
PR_Github #48635 [ run ] completed with state
|
|
@schetlur-nv thank you for running the CI bot. I'm unable to open the CI Agent Failure Analysis link. I have pushed some pre-commit fixes. Do you mind sharing the failures from the link? |
|
Hi @vedularaghu , the CI failed because of pre-commit check failure. You can run the pre-commit hooks locally and fix the issues. Here's the reference about setting pre-commit: https://github.com/NVIDIA/TensorRT-LLM/blob/main/CONTRIBUTING.md#coding-guidelines |
|
@JunyiXu-nv thanks for the pointer. Turns out the failing step was actually bandit -r tensorrt_llm that scripts/release_check.py runs after the pre-commit hooks (the hooks themselves all passed) — bandit's |
|
/bot run --disable-fail-fast |
|
PR_Github #49062 [ run ] triggered by Bot. Commit: |
|
PR_Github #49062 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49226 [ run ] triggered by Bot. Commit: |
|
PR_Github #49226 [ run ] completed with state
|
b934e39 to
bb262e5
Compare
|
PR_Github #52263 [ run ] triggered by Bot. Commit: |
|
PR_Github #52263 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #52358 [ run ] triggered by Bot. Commit: |
|
PR_Github #52358 [ run ] completed with state
|
|
Hi @karljang, the only failure in pipeline #41656 is accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16 (DGX_H100-PyTorch-6), which is a Qwen3 model-accuracy test unrelated to this PR (this PR only adds Prometheus metrics — no changes to Qwen3, accuracy tests, or model code). It looks like a known main-side flake/regression of the type that's been getting waived in post-merge. Could you re-trigger, or confirm whether it can be waived so this can move forward? Thanks! |
|
/bot run --disable-fail-fast |
|
PR_Github #52455 [ run ] triggered by Bot. Commit: |
|
PR_Github #52455 [ run ] completed with state
|
|
@karljang I rebased it, if you don't mind could you please re-trigger? |
|
/bot run --disable-fail-fast |
|
PR_Github #52557 [ run ] triggered by Bot. Commit: |
|
PR_Github #52557 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #52660 [ run ] triggered by Bot. Commit: |
|
PR_Github #52660 [ run ] completed with state
|
|
/bot run |
|
PR_Github #52846 [ run ] triggered by Bot. Commit: |
|
@vedularaghu Thanks for checking it and rebasing the PR 👍 When the failures are unrelated to the PR, rebasing resets the CI/test history for the PR head, so we lose the ability to rely on the previous test records and have to rerun the full CI set again. For unrelated/flaky/infra failures, it is usually better to keep the same commit and ask for a CI rerun instead. |
|
@karljang that makes sense. I'll keep that in mind from next time. Thanks for re-running the bot. |
|
PR_Github #52846 [ run ] completed with state
|
|
@karljang the failures are unrelated to the PR |
|
/bot run --disable-fail-fast |
|
PR_Github #52953 [ run ] triggered by Bot. Commit: |
|
PR_Github #52953 [ run ] completed with state |
|
@karljang looks like we can merge it finally. Really appreciate your help with all the re-runs. |
|
Thanks for your contribution~! |
7 participants
Description
Add Prometheus metrics to TRT-LLM's serving layer for finer-grained observability of prompt caching, per-position speculative-decoding acceptance, request perplexity, prefill batch occupancy, and request validation errors. All new metrics are exposed via the existing
/prometheus/metricsendpoint and the background iteration stats collector — no existing routes or APIs change.This PR is now rebased on top of
mainafter #12545 was merged, and duplicates with #12545 have been removed:MetricNames.PROMPT_TOKENS(now upstream) and thetrtllm_prompt_tokens_totalcounter creation (already provided ascounter_prompt_tokensby [None][feat] Add production-level Prometheus metrics (iteration stats, config info, token counters, phase histograms) #12545).trtllm_server_queue_length,trtllm_num_active_requests, andtrtllm_prefill_num_context_requestsgauges — [None][feat] Add production-level Prometheus metrics (iteration stats, config info, token counters, phase histograms) #12545 already provides equivalent metrics from the same iteration-stat fields astrtllm_num_requests_waiting,trtllm_num_requests_running, andtrtllm_num_context_requests.metrics_stats[MetricNames.PROMPT_TOKENS] = …line inrecord_stats([None][feat] Add production-level Prometheus metrics (iteration stats, config info, token counters, phase histograms) #12545 already records it atsequence_index == 0).New per-request metrics (via
log_request_metrics_dict)trtllm_prompt_cached_tokens_total/trtllm_prompt_cached_tokens: KV cache hit tracking (counter + per-request histogram).trtllm_spec_decode_drafted_tokens_total/trtllm_spec_decode_accepted_tokens_total: per-position speculative decoding acceptance counters, labeled bytoken_position. Complements [None][feat] Add production-level Prometheus metrics (iteration stats, config info, token counters, phase histograms) #12545's iteration-level aggregate totals (trtllm_spec_decode_num_draft_tokens_total/trtllm_spec_decode_num_accepted_tokens_total) by exposing the per-draft-position distribution so consumers can compute per-position acceptance rates.trtllm_prefill_perplexity/trtllm_generation_perplexity: per-request perplexity histograms computed from logprobs.trtllm_request_error_total: error counter labeled by HTTP status code.New iteration-level metrics (via
log_iteration_stats)trtllm_prefill_batch_occupancy: fraction of max active slots occupied by context (prefill) requests (numContextRequests / maxNumActiveRequests).trtllm_prefill_batch_tokens: histogram of total context tokens per iteration. Distribution view alongside [None][feat] Add production-level Prometheus metrics (iteration stats, config info, token counters, phase histograms) #12545'strtllm_total_context_tokensgauge.Implementation details
py_per_pos_drafted/py_per_pos_acceptedarrays onLlmRequest(sized by newMAX_SPEC_DECODE_POSITIONS = 16constant), accumulated during_handle_responsesinpy_executor.pyand propagated throughLlmResponse.result.per_pos_drafted/acceptedtoGenerationResultBase.record_stats.prompt_logprobs(prefill) and eithercumulative_logprobor per-tokenlogprobs(generation), withmath.isfiniteguards and narrowedexcept (ValueError, TypeError)handling.RequestPerfMetrics.speculative_decodingaggregate totals (collapsed into a single-element per-position array) when per-position arrays are unavailable (TRT backend).sequence_index == 0 and output.finish_reasonto avoid double-counting acrossn > 1candidates — matching the pattern [None][feat] Add production-level Prometheus metrics (iteration stats, config info, token counters, phase histograms) #12545 uses forPROMPT_TOKENS. Generation perplexity is recorded per candidate.@staticmethod create_error_responseremain unchanged. The new metrics integrate into the existing_iteration_stats_collector_loopbackground task and theRequestValidationErrorexception handler.Files changed
tensorrt_llm/metrics/enums.pyMetricNamesenum values (PROMPT_CACHE_CACHED_TOKENS,SPEC_DEC_ACCEPTED_PER_POS,SPEC_DEC_DRAFTED_PER_POS,PREFILL_PERPLEXITY,GENERATION_PERPLEXITY)tensorrt_llm/metrics/collector.pylog_request_errormethodtensorrt_llm/_torch/pyexecutor/llm_request.pyMAX_SPEC_DECODE_POSITIONSconstant and per-position drafted/accepted counterstensorrt_llm/_torch/pyexecutor/py_executor.pyLlmResponse.resulttensorrt_llm/executor/result.pyrecord_stats(with the existing #12545PROMPT_TOKENSrecording untouched)tensorrt_llm/serve/openai_server.pyRequestValidationErrorexception handlerTotal: +244 / −5 across 6 files.
Test Coverage
/prometheus/metricsendpoint.PR Checklist
collector.py)mainafter [None][feat] Add production-level Prometheus metrics (iteration stats, config info, token counters, phase histograms) #12545 merged; overlapping metrics removed