[TRTLLM-14708][fix] Populate speculative-decoding request perf metrics in the PyTorch one-engine flow - #17082
[TRTLLM-14708][fix] Populate speculative-decoding request perf metrics in the PyTorch one-engine flow#17082brnguyen2 wants to merge 5 commits into
Conversation
|
/bot run |
|
PR_Github #62851 [ run ] triggered by Bot. Commit: |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe change propagates speculative-decoding token totals through PyTorch responses, backfills request performance metrics, and reports lm-eval acceptance rates. Tests cover metric persistence, evaluation controls, windowed generation, partial scoring, and multimodal output handling. ChangesEvaluation and speculative metrics
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant LlmResult
participant GenerationResultBase
participant LmEval
PyExecutor->>LlmResult: Attach speculative-decoding totals
LlmResult->>GenerationResultBase: Pass response totals
GenerationResultBase->>LmEval: Expose request performance metrics
LmEval->>LmEval: Aggregate acceptance rate
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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: 2
🧹 Nitpick comments (3)
tensorrt_llm/executor/result.py (1)
288-289: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate the
perf_metricsparameter.The helper has a
-> Nonereturn annotation but the parameter is untyped;tllm.RequestPerfMetricsis the concrete type flowing in from_handle_sequence.♻️ Proposed annotation
- def _maybe_fill_spec_dec_perf_metrics(self, perf_metrics) -> None: + def _maybe_fill_spec_dec_perf_metrics( + self, perf_metrics: tllm.RequestPerfMetrics) -> None:As per coding guidelines: "Annotate every function, use
Nonefor non-returning functions, avoidAnyand unnecessary type ignores".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/executor/result.py` around lines 288 - 289, Update the _maybe_fill_spec_dec_perf_metrics method signature to annotate perf_metrics as tllm.RequestPerfMetrics, preserving its existing -> None return annotation and implementation.Source: Coding guidelines
tensorrt_llm/evaluate/lm_eval.py (1)
138-142: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider narrowing the swallowed exception type.
The catch-all is deliberate (partial scoring must never fail the eval) and that intent is documented, but it also hides genuine bugs in the tracker itself — and the repo standard is to catch specific types. The realistic failure modes here are attribute/key/index/type/value errors from exotic task, filter, or metric shapes.
♻️ Proposed narrowing
- except Exception as e: + except (AttributeError, IndexError, KeyError, TypeError, + ValueError) as e: logger.info( f"Partial scoring disabled for this run ({type(e).__name__}: {e})" ) self.disabled = TrueIf the blanket catch is intentional for third-party filter code, a short
# noqa-style rationale noting the deliberate deviation would help future readers.As per coding guidelines: "Catch specific exceptions instead of using broad or bare exception handling such as
except:".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/evaluate/lm_eval.py` around lines 138 - 142, In the partial-scoring handling around the existing `self.disabled` assignment, replace the broad `except Exception` with a tuple of the expected `AttributeError`, `KeyError`, `IndexError`, `TypeError`, and `ValueError` exceptions. Preserve the informational log and disabling behavior, and add a brief rationale comment only if a broader catch remains necessary for third-party filter code.Source: Coding guidelines
tests/unittest/others/test_lm_eval.py (1)
1493-1497: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winGuard the optional
lm_eval/datasetsimports. Wrap these runtime imports inpytest.importorskip(...)so the E2E block skips cleanly when either package is missing;ConfigurableTaskalready accepts acustom_datasetcallable here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/others/test_lm_eval.py` around lines 1493 - 1497, Update _toy_task so its optional runtime dependencies are loaded through pytest.importorskip for both datasets and lm_eval, allowing the E2E tests to skip cleanly when either package is unavailable. Preserve the existing ConfigurableTask custom_dataset setup and behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/evaluate/lm_eval.py`:
- Around line 64-70: Update the comment above SPEC_STATS_ENV_VAR to reflect that
acceptance-rate reporting is now supported, removing the outdated deferral
rationale and noting that speculative-decoding summaries include both acceptance
length and acceptance rate when performance metrics are enabled by
_get_sampling_params.
In `@tests/unittest/executor/test_spec_dec_perf_metrics.py`:
- Around line 1-18: Mark the test module containing the GenerationResultBase
speculative-decoding metrics tests with pytest.mark.cpu_only so it is selected
by the CPU-only test stages. Apply the marker at module scope, using the
existing pytest import and preserving the current test logic.
---
Nitpick comments:
In `@tensorrt_llm/evaluate/lm_eval.py`:
- Around line 138-142: In the partial-scoring handling around the existing
`self.disabled` assignment, replace the broad `except Exception` with a tuple of
the expected `AttributeError`, `KeyError`, `IndexError`, `TypeError`, and
`ValueError` exceptions. Preserve the informational log and disabling behavior,
and add a brief rationale comment only if a broader catch remains necessary for
third-party filter code.
In `@tensorrt_llm/executor/result.py`:
- Around line 288-289: Update the _maybe_fill_spec_dec_perf_metrics method
signature to annotate perf_metrics as tllm.RequestPerfMetrics, preserving its
existing -> None return annotation and implementation.
In `@tests/unittest/others/test_lm_eval.py`:
- Around line 1493-1497: Update _toy_task so its optional runtime dependencies
are loaded through pytest.importorskip for both datasets and lm_eval, allowing
the E2E tests to skip cleanly when either package is unavailable. Preserve the
existing ConfigurableTask custom_dataset setup and behavior.
🪄 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: Enterprise
Run ID: f8258d1d-ba63-4621-becb-f4ee03c5322c
📒 Files selected for processing (7)
tensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/evaluate/lm_eval.pytensorrt_llm/executor/result.pytests/integration/test_lists/test-db/l0_a10.ymltests/unittest/executor/test_spec_dec_perf_metrics.pytests/unittest/others/test_lm_eval.py
|
PR_Github #62851 [ run ] completed with state
|
0354935 to
5d06271
Compare
|
/bot run |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/executor/result.py (1)
288-289: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate
perf_metrics.The helper is otherwise fully typed;
perf_metricsis atllm.RequestPerfMetrics.♻️ Proposed change
- def _maybe_fill_spec_dec_perf_metrics(self, perf_metrics) -> None: + def _maybe_fill_spec_dec_perf_metrics( + self, perf_metrics: tllm.RequestPerfMetrics) -> None:As per coding guidelines, "Annotate every function, use
Nonefor non-returning functions".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/executor/result.py` around lines 288 - 289, Update the _maybe_fill_spec_dec_perf_metrics method signature to annotate perf_metrics as tllm.RequestPerfMetrics, while preserving its existing None return annotation and implementation.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tensorrt_llm/executor/result.py`:
- Around line 288-289: Update the _maybe_fill_spec_dec_perf_metrics method
signature to annotate perf_metrics as tllm.RequestPerfMetrics, while preserving
its existing None return annotation and implementation.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4fa6fc08-2d56-4d07-aef6-2d73eec748e0
📒 Files selected for processing (7)
tensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/evaluate/lm_eval.pytensorrt_llm/executor/result.pytests/integration/test_lists/test-db/l0_a10.ymltests/unittest/executor/test_spec_dec_perf_metrics.pytests/unittest/others/test_lm_eval.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/integration/test_lists/test-db/l0_a10.yml
|
PR_Github #62880 [ run ] triggered by Bot. Commit: |
|
Reviewed Two things. The new test file never runs in CI. The two paths don't accumulate the denominator the same way. C++ does Holding my approval mainly because this is stacked and will be rebased once #16752 lands — the diff I'd be approving isn't the one that merges. Happy to re-review promptly after the retarget. |
|
PR_Github #62880 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62925 [ run ] triggered by Bot. Commit: |
|
PR_Github #62925 [ run ] completed with state
|
…s in the PyTorch one-engine flow With SamplingParams(return_perf_metrics=True), the request_perf_metrics.speculative_decoding section arrived zeroed in the PyTorch flow even when drafting ran: its only writer is C++ LlmRequest::updateNumTokensPerIteration, which the TRT engine runtime and TRTLLMSampler call but the TorchSampler path never does. Accumulate exact per-request draft/accepted totals in the PyTorch executor, attach them to the response, and backfill the serialized section client-side; the runtime-populated paths are left untouched. Also restore the acceptance-rate (AR) line to the TLLM_EVAL_SPEC_STATS lm-eval summary, deferred by the parent change until these counters were populated; enabling the stats now also opts requests into perf metrics. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
5d06271 to
486171b
Compare
|
PR_Github #62994 [ run ] triggered by Bot. Commit: |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/executor/test_spec_dec_perf_metrics.py`:
- Line 1: Add the repository-standard NVIDIA copyright header at the top of the
new test module, before its module docstring, while leaving the existing test
content unchanged.
- Around line 39-46: Update the no-op tests around test_noop_without_totals and
test_noop_on_nonpositive_drafted to assert all three speculative-decoding
metrics: total_accepted_draft_tokens, total_draft_tokens, and acceptance_rate.
Add the missing assertions in every no-op branch while preserving the expected
zero/default values established by the implementation.
🪄 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: Enterprise
Run ID: d2a9089d-5e5c-4b85-b36d-60e7617847f5
📒 Files selected for processing (6)
tensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/evaluate/lm_eval.pytensorrt_llm/executor/result.pytests/unittest/executor/test_spec_dec_perf_metrics.pytests/unittest/others/test_lm_eval.py
🚧 Files skipped from review as they are similar to previous changes (5)
- tensorrt_llm/executor/result.py
- tensorrt_llm/evaluate/lm_eval.py
- tensorrt_llm/_torch/pyexecutor/llm_request.py
- tensorrt_llm/_torch/pyexecutor/py_executor.py
- tests/unittest/others/test_lm_eval.py
|
PR_Github #62994 [ run ] completed with state
|
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
C++ LlmRequest::updateNumTokensPerIteration accumulates totalDraftTokens += min(getNumDraftTokens(), maxDraftPathLen) so that tree-based drafting counts at most one path per step in the acceptance-rate denominator. The PyTorch accumulation added by this PR used the unclamped draft token count, which agrees for linear chains but would over-count for tree drafters. Clamp to max_draft_len (the per-path depth; max_total_draft_tokens is the tree size) to keep both paths on the same definition. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
The SPEC_STATS_ENV_VAR comment still described AR reporting as deferred to a follow-up; this PR is that follow-up. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
@BowenFu Both points addressed, thanks — the branch is now rebased onto main post-#16752, so the diff you'd approve is the one that merges. CI wiring ( Denominator clamp ( Also refreshed the now-stale "AR reporting is deferred" comment above |
|
/bot run |
|
PR_Github #63030 [ run ] triggered by Bot. Commit: |
|
PR_Github #63030 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63043 [ run ] triggered by Bot. Commit: |
|
PR_Github #63043 [ run ] completed with state
|
Address review: py_draft_tokens is not a valid denominator at response-handling time. In the one-model flow (SpecSamplerBase), update_requests has already replaced it with the NEXT step's buffer, and the request's prefill iteration carries dummy draft tokens; in the drafter flows, pad_draft_tokens_for_cuda_graph has padded it to the static max. All of these inflated the drafted count and biased the acceptance rate (and the per-position stats) low. _handle_responses also re-read the same stale counters for every active request each iteration, double-counting steps for requests that sat unscheduled. Samplers now write the paired denominator at the same place they write py_num_accepted_draft_tokens: - SpecSamplerBase.sample_async snapshots per-request draft counts into the sample state (before dummy-token insertion; preferring the pre-padding count the drafter recorded), and update_requests copies them onto the request as py_num_draft_tokens_verified. - TorchSampler pairs against the pre-padding count recorded by pad_draft_tokens_for_cuda_graph (py_draft_tokens_effective_len); py_rewind_len keeps using the padded length, since padding occupies KV cache. Accumulation moves from _handle_responses into the executor right after Sampler.update_requests (_accumulate_spec_dec_stats), consuming py_num_draft_tokens_verified exactly once per verified step. This fixes the cumulative totals and the per-position arrays alike. New tests in tests/unittest/executor/test_spec_dec_stats_pairing.py pin the pairing: verified-count-vs-replaced-buffer, consume-once semantics, prefill steps counting nothing, the tree-drafting clamp, and the drafter's pre-padding recording. Wired into l0_a10.yml rather than the CPU-only list: the tests are GPU-free but import py_executor, which no cpu_only-stage test currently does. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
|
PR_Github #63142 [ run ] triggered by Bot. Commit: |
|
PR_Github #63142 [ run ] completed with state
|
|
/bot run |
7 participants
Description
Update (2026-07-31): #16752 has merged and this branch has been rebased onto post-merge
main, so the stacked structure described here originally is gone. All 4 commits belong to this PR (the original fix plus three review-response commits); the full diff is the change to review. This is the follow-up promised in #16752: "AR reporting should return in a follow-up together with a fix populating those counters."Problem
With
SamplingParams(return_perf_metrics=True), per-requestoutput.outputs[0].request_perf_metrics.speculative_decoding(
total_accepted_draft_tokens,total_draft_tokens,acceptance_rate)comes back zeroed in the PyTorch flow even when drafting ran, while
output.avg_decoded_tokens_per_iterflows correctly.Root cause
The only writer of
RequestPerfMetrics.speculativeDecodingis C++LlmRequest::updateNumTokensPerIteration, which is called by the TRT engineruntime (
trtGptModelInflightBatching.cpp) and by theTRTLLMSampler(
sampler.py). The PyTorch flow withTorchSamplernever calls it, socreateSerializedResultcopies a zeroed section into every response.Generality: this is PyTorch-flow generic (all spec modes, one- and
two-model), not specific to any drafter. Acceptance counts exist per step
on the Python side (
LlmRequest.py_num_accepted_draft_tokens) but were onlyaggregated into ITERATION stats and the capped per-position arrays
(
py_per_pos_*), never into the request's cumulative perf metrics.Fix (Python-only, no cpp rebuild)
llm_request.py:LlmRequestgains exact cumulativepy_total_draft_tokens/py_total_accepted_draft_tokenscounters (unlikethe per-pos arrays, not capped at
MAX_SPEC_DECODE_POSITIONS).py_executor.py(_handle_responses): accumulate them alongside theexisting per-pos stats and attach
(accepted, drafted)to the emittedresponse as
LlmResult.spec_dec_totals. (The C++Resultis alreadyserialized by
create_serialized_resultat that point, so the sectioncannot be written runtime-side without cpp changes.)
executor/result.py:GenerationResultBase._maybe_fill_spec_dec_perf_metricsbackfills
RequestPerfMetrics.speculative_decoding(totals + acceptancerate) from those totals when the section arrives empty. TRT-engine /
TRTLLMSamplerpaths that already populate it runtime-side are detected(
total_draft_tokens > 0) and left untouched.This also fixes the existing
record_statsfallback (executor/result.py)that reads the same section for
SPEC_DEC_*Prometheus metrics when theper-pos arrays are absent.
Eval-side follow-up (the stacked part)
_log_spec_statsinevaluate/lm_eval.py(added by #16752) now also logs acorpus acceptance-rate line —
sum(accepted)/sum(drafted)over per-requestrequest_perf_metrics.speculative_decoding, gated ontotal_draft_tokens > 0— alongside the existing AL line.Design decision: #16752 deliberately does not request perf metrics (zero
overhead), but AR needs
return_perf_metrics=True. This PR opts requestsinto perf metrics when
TLLM_EVAL_SPEC_STATS=1is set, i.e. the (small)per-request perf-metrics overhead is bundled with explicitly asking for
spec-dec stats. The alternative (log AR only when the caller independently
enabled perf metrics) would make the AR line appear only with an extra,
undocumented knob.
Risks
spec_dec_totalsaccumulates in the same loop as the per-pos stats, usingthe same
draft_len/py_num_accepted_draft_tokensper-iteration values;any bias in those (e.g. final-iteration draft handling) is shared with the
already-validated per-pos stats.
request_perf_metrics(i.e.
return_perf_metrics=True), matching C++ behavior.Test Coverage
tests/unittest/executor/test_spec_dec_perf_metrics.py:fills a zeroed section from totals (30/40 → AR 0.75), no-op without totals,
no-op when the section is already populated runtime-side, no-op on
non-positive drafted counts, and pickle round-trip (responses cross the
worker/proxy IPC boundary pickled).
tests/unittest/others/test_lm_eval.py: new cases for the AR line(corpus summation, skip-when-no-drafting, AR-without-AL) and for the
TLLM_EVAL_SPEC_STATS=1→return_perf_metricsopt-in; existing AL testsupdated so fake outputs carry explicit (or absent) perf-metrics counters.
(draft_len 2, TP16, 48 prompts): the AR line now prints
(
acceptance 661/5392 = 12.3%) alongside an unchanged AL line(tokens/step mean ~1.28); the pre-fix baseline printed the AL line only.
PR Checklist
pre-commit run --files <touched>clean.Dev Engineer Review
The implementation correctly adds cumulative speculative-decoding counters to
LlmRequest, transfers them through PyTorch executor responses, and backfills missingRequestPerfMetrics.speculative_decodingvalues. Existing populated metrics and TRT metrics remain unchanged.The drafted-token count is clamped to
max_draft_len, matching the C++ denominator for tree-based drafting.lm_evalreports corpus-level acceptance rate whenTLLM_EVAL_SPEC_STATS=1. The acceptance-rate calculation and pickle round-trip behavior are covered.The CPU test is included in
tests/integration/test_lists/test-db/l0_cpu_x86.yml. No configuration or test-list format issues are reported.QA Engineer Review
tests/unittest/executor/test_spec_dec_perf_metrics.pyadds:test_fills_zeroed_section_from_totalstest_noop_without_totalstest_noop_when_section_already_populatedtest_noop_on_nonpositive_draftedtest_survives_pickle_roundtripThe test file is covered by
tests/integration/test_lists/test-db/l0_cpu_x86.yml.tests/unittest/others/test_lm_eval.pyadds or updates speculative-statistics tests for acceptance length, acceptance rate, missing-token handling,return_perf_metricsgating, request scheduling, partial scoring, multimodal scoring, and end-to-end execution. This file is covered bytests/integration/test_lists/test-db/l0_a10.yml.Verdict: sufficient.