[TRTLLM-14779][fix] Clear capture-only sampling override from cached CUDA graph metadata - #17033
Conversation
…ph metadata The advanced-sampling CUDA graph capture pass sets _force_non_greedy_for_capture=True on spec_metadata so that parameter-less warmup requests scan as non-greedy and the advanced-sampling branch is the one recorded into the graph. create_cuda_graph_metadata() shallow-copies the live SpecMetadata, so every graph captured during that pass caches a copy that inherited the flag. Those copies are reseated as the live spec_metadata on each later replay, but the capture pass clears the flag on the base object only. The flag therefore stays set on the cached copies forever, and _scan_one_model_sampling() takes the capture branch at serving time, replacing EVERY request's sampling parameters with the synthetic capture values (temperature 0.7, top_k 50, top_p 0.9). Symptom: with one-engine speculative decoding and CUDA graphs enabled (the default), a client asking for temperature=1.0 / top_p=1.0 is served at 0.7 / 50 / 0.9 on every row for the whole run. The narrowed nucleus can make a repetition loop effectively inescapable, producing very long degenerate generations that terminate only at the length limit. Only one-engine speculative decoding runs the capture pass, so non-spec decoding is unaffected. Clear the flag on the cached copies in the capture pass's finally, after every graph for the pass has been captured, so the flag is still in effect for the captures that need it. Clearing it at copy time instead would be wrong: the pass-2 populate would then scan the parameter-less warmup requests as greedy and bake the argmax fast path (and absent top-k/top-p kernels) into the graph keyed as the advanced-sampling variant. Also assert in _prepare_inputs that the override is never live outside warmup. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…e lifetime Covers the leak fixed in the previous commit and the constraint that rules out the obvious alternative fix: 1. A graph metadata shallow copy inherits _force_non_greedy_for_capture, and clearing the flag on the base object does not reach the copies. 2. clear_capture_only_spec_state() clears exactly the flagged copies, leaves greedy-pass copies and non-spec entries alone, and is idempotent. 3. After the capture pass tears down, a scan on the cached copy honors the request's own sampling params. This fails without the fix: the copy still carries the flag, so the scan returns the synthetic 0.7 / 50 / 0.9. 4. While the flag is set, the scan still substitutes the synthetic params and reports a non-greedy batch. This pins the flag as load-bearing during capture, so the clear cannot be moved to copy time without breaking the advanced-sampling graph it exists to record. The tests construct a real base SpecMetadata (its __post_init__ is a no-op and no field exercised here is a tensor) and call clear_capture_only_spec_state unbound on a stand-in holding only graph_metadata, so no GPU, no runner construction and no model forward is needed. Style follows the neighboring test_group_all_greedy_sync.py. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
WalkthroughAdds cleanup for capture-only CUDA graph sampling metadata, asserts the override is not active during serving, and introduces speculative decoding tests covering metadata copies, teardown, serving scans, and warmup behavior. ChangesCapture override lifecycle
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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: 1
🧹 Nitpick comments (1)
tests/unittest/_torch/speculative/test_capture_override_leak.py (1)
38-143: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd annotations to the new helpers and tests.
The new functions are unannotated, contrary to the Python review requirement. As per coding guidelines, “Annotate every function.”
🤖 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/_torch/speculative/test_capture_override_leak.py` around lines 38 - 143, Add explicit return-type annotations to the new helper functions _base_meta, _graph_copy, _clear, _request, and _scan, and annotate each test function with its return type as well. Use the existing project types where they are available and an appropriate type for helper collections or test-only returns, without changing behavior.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.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 2000-2004: Gate the cleanup around clear_capture_only_spec_state
in model_engine.py: only clear cached override flags when
cuda_graph_runner.is_warmup_only is false, preserving them through the
warmup-only pass and clearing them after final capture. In
tests/unittest/_torch/speculative/test_capture_override_leak.py, add a lifecycle
test covering warmup-only capture followed by real capture, asserting cached
metadata remains forced during warmup and is cleared after the final capture.
---
Nitpick comments:
In `@tests/unittest/_torch/speculative/test_capture_override_leak.py`:
- Around line 38-143: Add explicit return-type annotations to the new helper
functions _base_meta, _graph_copy, _clear, _request, and _scan, and annotate
each test function with its return type as well. Use the existing project types
where they are available and an appropriate type for helper collections or
test-only returns, without changing 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: 9cfc2b7c-54c8-4cfe-98a1-7ad99a8de3c1
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/model_engine.pytests/integration/test_lists/test-db/l0_h100.ymltests/unittest/_torch/speculative/test_capture_override_leak.py
|
/bot run --disable-fail-fast |
|
PR_Github #62633 [ run ] triggered by Bot. Commit: |
|
PR_Github #62633 [ run ] completed with state
|
Signed-off-by: Xiao Wang <xiaow@nvidia.com> Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #62808 [ run ] triggered by Bot. Commit: |
|
Adjudicating the open CodeRabbit thread on The premise doesn't hold. CodeRabbit's argument is that " self.graph_metadata[key] = {
"attn_metadata": attn_metadata,
"spec_metadata": initial_inputs.get("spec_metadata", None),
}unconditionally at the top of every The two-pass structure is real ( The suggested fix inverts the safety direction. Gating on For what it's worth, the rest of the change reads correctly to me. Clearing in the Not approving yet only because I'd rather not be the first approver on a spec-dec + CUDA-graph change; the diagnosis and the quantified result (GPQA 84.85% → 92.42%, runaways 15/198 → 0/594) are convincing. |
|
PR_Github #62808 [ run ] completed with state
|
|
/bot run --disable-fail-fast --reuse-test |
|
PR_Github #62858 [ run ] triggered by Bot. Commit: |
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
Approving. The default path is provably untouched: spec_metadata is None whenever spec decode is off, the new assert lives under if spec_metadata is not None, and the cleanup hook is only reachable from the force_non_greedy=True capture pass gated by use_one_engine().
Re-confirming my earlier read of the remaining CodeRabbit thread at e0f5aaba7, since it's the only thing still open. capture() assigns self.graph_metadata[key] unconditionally at the top, before if self.is_warmup_only: return output — so the second pass overwrites it with a fresh copy taken while the flag is live, rather than reusing the warmup-pass object. And the "advanced key captured in pass 1" case can't bite either: self.graphs[key] = graph is only reached after that early return, so a warmup-only key has no runnable graph and replay() can never select it. The suggested if not is_warmup_only: gate would be strictly worse — it leaves the flag set for a key captured during estimation but skipped in the final pass, which is the original leak.
Not blocking: the new tests are object-level on SpecMetadata (inheritance, clearing, post-clear sampling) rather than end-to-end through warmup()/capture()/replay(), so they'd miss a future regression in the call site rather than the flag itself.
|
PR_Github #62858 [ run ] completed with state
|
|
/bot run --disable-fail-fast --reuse-test |
|
PR_Github #62958 [ run ] triggered by Bot. Commit: |
|
PR_Github #62958 [ run ] completed with state |
7 participants
Description
The bug
The advanced-sampling CUDA graph capture pass sets
_force_non_greedy_for_capture=Trueon the liveSpecMetadataso that parameter-less warmup requests scan as non-greedy, and the advanced-sampling branch (not the argmax fast path) is the one recorded into the graph.create_cuda_graph_metadata()shallow-copies the liveSpecMetadata, so every graph captured during that pass caches a copy that inherited the flag. Those copies are reseated as the livespec_metadataon each later replay, but the capture pass clears the flag on the base object only. The flag therefore stays set on the cached copies forever, and_scan_one_model_sampling()takes the capture branch at serving time, replacing every request's sampling parameters with the synthetic capture values —temperature 0.7,top_k 50,top_p 0.9.A client asking for
temperature=1.0/top_p=1.0is silently served at0.7 / 50 / 0.9, on every row, for the whole lifetime of the server.Preconditions and blast radius
Only one-engine speculative decoding (
spec_dec_mode.use_one_engine()— MTP, EAGLE3-one-model, etc.) runs the capture pass, and only when CUDA graphs are enabled. Sincecuda_graph_configdefaults to enabled, that is the default serving posture for those modes; no opt-in is required to hit it. Non-speculative decoding never runs the pass and is unaffected.Symptoms vary by workload, which is why this went unnoticed:
temperature 0.7sharpens the distribution andtop_p 0.9can collapse the nucleus to a single token, so a repetition loop becomes inescapable and the generation only stops at the length captemperature=0is sampled at0.7 / 50 / 0.9whenever its batch contains at least one non-greedy request, i.e. non-deterministic "greedy" output depending on batch compositionThe fix
Clear the flag on the cached copies in the capture pass's
finally, after every graph for the pass has been captured, so the flag is still in effect for the captures that need it.Clearing it at copy time instead would be wrong: the pass-2 populate would then scan the parameter-less warmup requests as greedy and bake the argmax fast path — with the top-k/top-p kernels absent — into the graph keyed as the advanced-sampling variant. A unit test pins this constraint so the cleanup is not "simplified" later.
Also asserts in
_prepare_inputsthat the override is never live outside warmup. Kept fatal on purpose: the failure mode is silent output corruption, and a warning would be scrolled past.Measured impact
GLM-5.2 NVFP4 on GB200, GPQA Diamond (198 questions), one-engine MTP
k=1with rejection sampling, CUDA graphs on, reasoning enabled with a 400k-token output budget, client requestingtemperature=1.0/top_p=1.0. "Runaways" are generations that hit the length cap while looping.Pooled over three full runs on the fixed build: 0 / 594 runaways, versus 15 / 198 before (Fisher exact p ≈ 2e-13). Every endpoint — accuracy, mean, p90, max — lands on top of the speculative-decoding-off reference, so the apparent accuracy and token-inflation penalty of enabling MTP on this configuration was entirely this bug.
Test Coverage
tests/unittest/_torch/speculative/test_capture_override_leak.py(new, CPU-only, no GPU or model forward required; registered inl0_h100.yml):test_graph_copy_inherits_flag_and_base_teardown_does_not_reach_it— pins the shallow-copy semantics the bug rests on.test_clear_capture_only_spec_state_clears_every_cached_copy— clears exactly the flagged copies, leaves greedy-pass copies and non-spec entries alone, idempotent.test_serving_scan_honors_client_params_after_capture_teardown— fails without the fix, where the scan returns0.7 / 50 / 0.9instead of the request's own parameters.test_override_stays_live_while_the_flag_is_set— anti-regression for the rejected "clear at copy time" variant.related: #14745
Dev Engineer Review
QA Engineer Review
test_graph_copy_inherits_flag_and_base_teardown_does_not_reach_it()test_clear_capture_only_spec_state_clears_every_cached_copy()test_serving_scan_honors_client_params_after_capture_teardown()test_override_stays_live_while_the_flag_is_set()tests/integration/test_lists/test-db/l0_h100.yml.