[https://nvbugs/6442074][fix] Make one-model spec-dec attn-metadata save/restore exception-safe#16381
Closed
dongfengy wants to merge 3 commits into
Closed
Conversation
… save/restore exception-safe TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v2_kv_cache-one_model] fails on H100 with "AssertionError: Sampling failed". The real chain: the final general warmup (max-shape memory-pool prepop) OOMs on H100-80GB and is tolerated by model_engine, but the aborted forward had already run Eagle3OneModelWorker prepare of the attn metadata without try/finally, so the restore never ran. Every subsequent forward then trips the bare "assert len(self._saved_tensors) == 0" in prepare_for_spec_dec, whose empty str() masks the error, cascading into "Sampling failed". - eagle3.py: wrap the draft path in try/finally with the prepare INSIDE the try (the Eagle3 prepare override allocates clones after the base save, so a failure inside prepare must also reach the restore). - interface.py: add messages to the bare asserts so a recurrence names the unrestored fields instead of logging an empty error. - waives.txt: unwaive the test on H100. Verified on 4x H100-80GB (viking-prod-233) at main f81b9d2: pristine TOT reproduces the exact QA failure; with this fix the same warmup OOM occurs but serving survives - GSM8K 89.92 (ref 90.30), GPQADiamond 64.14 (ref 65.0), test PASSED. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: dongfengy <99041270+dongfengy@users.noreply.github.com>
…ion safety to MTP, PARD, DFlash and DraftTarget workers The one-model workers in mtp.py, pard.py, dflash.py and draft_target.py have the same latent bug fixed in eagle3.py: prepare/restore of the attention metadata is not exception-safe, so any tolerated failure in the draft path (e.g. a warmup OOM) poisons attn_metadata._saved_tensors and every later forward dies at the pairing assert. - mtp.py: move change_attn_metadata (which saves state and then mutates more of it) inside the try; restore in finally. - pard.py: prepare and the deferred kv_lens rewind are now covered; the rewind runs in the finally after the restore so a draft failure does not leave kv_lens_cuda incremented for the next request. - dflash.py / draft_target.py: same try/finally treatment. MTPEagleWorker subclasses Eagle3OneModelWorker and is already covered by the eagle3.py fix. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: dongfengy <99041270+dongfengy@users.noreply.github.com>
…up in SpecWorkerBase Replace the per-worker try/finally blocks with a template method on SpecWorkerBase: forward() extracts attn/spec metadata, runs the worker-specific _forward_impl(), and in a finally block restores any unrestored prepare_for_spec_dec state (plus the pending kv_lens rewind for PARD/DFlash). __init_subclass__ rejects subclasses that override forward directly, so a future worker cannot reintroduce the leak. - speculative/interface.py: template forward + _ensure_spec_dec_state_restored. - attention_backend/interface.py: has_spec_dec_saved_state property. - eagle3/mtp/pard/dflash/draft_target/sa_worker/eagle3_dynamic_tree: forward renamed to _forward_impl; bodies return to straight-line prepare/restore (per-worker try/finally removed). - pard/dflash: _kv_rewind_pending flag so a failed draft forward still rewinds kv_lens_cuda exactly once; consumed by _apply_kv_rewind_after_draft. - eagle3: initialize the worker-side _saved_* attrs so a failure inside the prepare override cannot turn the cleanup restore into AttributeError. - test_force_accepted_tokens.py: stub worker implements _forward_impl. Verified on 4x H100-80GB (viking-prod-233) at main f81b9d2, A/B: - workers unguarded, no template method: FAILED in 215s with the nvbugs/6442074 signature (warmup OOM then stale-save assert). - template method only: the base finally visibly restores state when the max-shape warmup OOMs, and the test PASSES - GSM8K 90.75 (ref 90.30), GPQADiamond 66.16 (ref 65.0). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: dongfengy <99041270+dongfengy@users.noreply.github.com> Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
dongfengy
requested review from
BowenFu,
brnguyen2,
cascade812,
dpitman-nvda,
pengbowang-nv,
ruodil,
tongyuantongyu,
yingguo-trt,
yiqingy0 and
yunruis
July 14, 2026 17:02
3 tasks
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #59242 [ run ] triggered by Bot. Commit: |
2 tasks
Collaborator
Author
|
/bot kill |
Collaborator
Author
|
Closing in favor of #16382 (same commits rebased onto current main TOT). |
Collaborator
|
PR_Github #59244 [ ] completed with state |
Contributor
📝 WalkthroughWalkthroughSpeculative-decoding workers now implement ChangesSpeculative decoding forward lifecycle
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant SpecWorkerBase.forward
participant Worker._forward_impl
participant AttentionMetadata
participant DFlashWorker_or_PARDWorker
SpecWorkerBase.forward->>Worker._forward_impl: run speculative decoding
Worker._forward_impl->>DFlashWorker_or_PARDWorker: adjust KV lengths
Worker._forward_impl-->>SpecWorkerBase.forward: complete or fail
SpecWorkerBase.forward->>AttentionMetadata: restore saved state
AttentionMetadata->>DFlashWorker_or_PARDWorker: apply pending KV rewind
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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
tensorrt_llm/_torch/speculative/interface.py (1)
938-963: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider a dedicated unit test for the new exception-safety guarantee.
This wrapper is the core fix for nvbugs/6442074 (tolerated OOM leaving spec-dec state populated). A focused unit test — subclass
SpecWorkerBasewith an_forward_implthat raises after callingattn_metadata.prepare_for_spec_dec(...), then assertattn_metadata.has_spec_dec_saved_state is Falseafter the exception — would directly verify the guarantee this PR provides, independent of the H100 integration test.🤖 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/_torch/speculative/interface.py` around lines 938 - 963, Add a focused unit test for SpecWorkerBase.forward that uses a test subclass whose _forward_impl calls attn_metadata.prepare_for_spec_dec(...) and then raises an exception. Assert the exception propagates and attn_metadata.has_spec_dec_saved_state is False afterward, verifying cleanup independently of integration tests.tensorrt_llm/_torch/speculative/dflash.py (1)
277-288: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
DFlashWorkerandPARDWorkerduplicate the entire deferred-KV-rewind cleanup mechanism verbatim. Both independently implement_kv_rewind_pendingbookkeeping plus an identical_ensure_spec_dec_state_restoredoverride (same code, same docstring). Any future fix or edge case discovered in one will need to be manually ported to the other, risking divergence — a natural candidate to hoist intoSpecWorkerBase(e.g. a small mixin or a generic "deferred rewind" helper keyed by a callback) now, while the two implementations are still identical.
tensorrt_llm/_torch/speculative/dflash.py#L277-L288: extract the_kv_rewind_pending/_kv_rewind_amount/_kv_rewind_nc/_kv_rewind_bsbookkeeping in_prepare_kv_for_draft_forward/_apply_kv_rewind_after_draftinto a shared base-class helper instead of maintaining a second copy inpard.py.tensorrt_llm/_torch/speculative/dflash.py#L374-L384: replace this override with a call into the shared helper fromSpecWorkerBaserather than duplicatingpard.py's identical override.tensorrt_llm/_torch/speculative/pard.py#L147-L163: extract this bookkeeping (identical todflash.py's) into the same shared helper.tensorrt_llm/_torch/speculative/pard.py#L172-L182: replace this override with a call into the shared helper instead of duplicatingdflash.py's identical override.🤖 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/_torch/speculative/dflash.py` around lines 277 - 288, The deferred KV-rewind bookkeeping and restoration logic is duplicated between DFlashWorker and PARDWorker; centralize it in SpecWorkerBase through a shared helper. Update tensorrt_llm/_torch/speculative/dflash.py lines 277-288 and tensorrt_llm/_torch/speculative/pard.py lines 147-163 to use that helper for _kv_rewind_pending, _kv_rewind_amount, _kv_rewind_nc, and _kv_rewind_bs. Replace the duplicate _ensure_spec_dec_state_restored overrides in tensorrt_llm/_torch/speculative/dflash.py lines 374-384 and tensorrt_llm/_torch/speculative/pard.py lines 172-182 with calls to the shared base implementation.tests/unittest/_torch/speculative/test_force_accepted_tokens.py (1)
46-57: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd regression coverage for the cleanup wrapper.
This change only makes
_StubSpecWorkerconcrete; it does not exerciseSpecWorkerBase.forward()invoking_forward_impl()and restoring attention state after an exception. Add or verify a focused test intests/unittest/_torch/speculative/test_force_accepted_tokens.py(or a dedicated cleanup test) that raises from_forward_impland asserts saved state is cleared and restored.As per path instructions, this test change needs an explicit coverage assessment and concrete follow-up when coverage is insufficient.
🤖 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_force_accepted_tokens.py` around lines 46 - 57, The new _StubSpecWorker only enables instantiation and does not cover SpecWorkerBase.forward() cleanup. Add a focused regression test that makes _forward_impl raise, invokes forward(), and asserts the saved attention state is both restored and cleared afterward; include the required explicit coverage assessment and a concrete follow-up if coverage remains insufficient.Source: Path instructions
🤖 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/speculative/dflash.py`:
- Around line 374-386: Update _ensure_spec_dec_state_restored to restore the
warmup _ctx_len snapshot during cleanup, including when _store_prefill_context
or prepare_1st_drafter_inputs fails. Reuse the existing warmup snapshot/state
symbols and perform restoration in the finally-driven hook alongside the
existing superclass restoration and pending KV rewind handling.
---
Nitpick comments:
In `@tensorrt_llm/_torch/speculative/dflash.py`:
- Around line 277-288: The deferred KV-rewind bookkeeping and restoration logic
is duplicated between DFlashWorker and PARDWorker; centralize it in
SpecWorkerBase through a shared helper. Update
tensorrt_llm/_torch/speculative/dflash.py lines 277-288 and
tensorrt_llm/_torch/speculative/pard.py lines 147-163 to use that helper for
_kv_rewind_pending, _kv_rewind_amount, _kv_rewind_nc, and _kv_rewind_bs. Replace
the duplicate _ensure_spec_dec_state_restored overrides in
tensorrt_llm/_torch/speculative/dflash.py lines 374-384 and
tensorrt_llm/_torch/speculative/pard.py lines 172-182 with calls to the shared
base implementation.
In `@tensorrt_llm/_torch/speculative/interface.py`:
- Around line 938-963: Add a focused unit test for SpecWorkerBase.forward that
uses a test subclass whose _forward_impl calls
attn_metadata.prepare_for_spec_dec(...) and then raises an exception. Assert the
exception propagates and attn_metadata.has_spec_dec_saved_state is False
afterward, verifying cleanup independently of integration tests.
In `@tests/unittest/_torch/speculative/test_force_accepted_tokens.py`:
- Around line 46-57: The new _StubSpecWorker only enables instantiation and does
not cover SpecWorkerBase.forward() cleanup. Add a focused regression test that
makes _forward_impl raise, invokes forward(), and asserts the saved attention
state is both restored and cleared afterward; include the required explicit
coverage assessment and a concrete follow-up if coverage remains insufficient.
🪄 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: 7f883c59-129a-445d-9633-e583faed47df
📒 Files selected for processing (11)
tensorrt_llm/_torch/attention_backend/interface.pytensorrt_llm/_torch/speculative/dflash.pytensorrt_llm/_torch/speculative/draft_target.pytensorrt_llm/_torch/speculative/eagle3.pytensorrt_llm/_torch/speculative/eagle3_dynamic_tree.pytensorrt_llm/_torch/speculative/interface.pytensorrt_llm/_torch/speculative/mtp.pytensorrt_llm/_torch/speculative/pard.pytensorrt_llm/_torch/speculative/sa_worker.pytests/integration/test_lists/waives.txttests/unittest/_torch/speculative/test_force_accepted_tokens.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
Comment on lines
+374
to
+386
| def _ensure_spec_dec_state_restored(self, attn_metadata, spec_metadata): | ||
| # Restore first (in warmup mode kv_lens_cuda was saved and comes back | ||
| # wholesale), then apply any pending rewind for the other modes so a | ||
| # failed draft forward does not leave kv_lens_cuda incremented. | ||
| super()._ensure_spec_dec_state_restored(attn_metadata, spec_metadata) | ||
| if ( | ||
| getattr(self, "_kv_rewind_pending", False) | ||
| and attn_metadata is not None | ||
| and spec_metadata is not None | ||
| ): | ||
| self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) | ||
|
|
||
| def _forward_impl( |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
path = Path('tensorrt_llm/_torch/speculative/dflash.py')
lines = path.read_text().splitlines()
for start, end in [(330, 620)]:
print(f'লINES {start}-{end}')
for i in range(start, min(end, len(lines)) + 1):
print(f"{i:4d}: {lines[i-1]}")
PY
rg -n "_ctx_len|_kv_rewind_pending|_ensure_spec_dec_state_restored|_store_prefill_context|prepare_1st_drafter_inputs" tensorrt_llm/_torch/speculative/dflash.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 16244
🏁 Script executed:
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
path = Path('tensorrt_llm/_torch/speculative/dflash.py')
lines = path.read_text().splitlines()
for start, end in [(330, 620)]:
print(f'LINES {start}-{end}')
for i in range(start, min(end, len(lines)) + 1):
print(f"{i:4d}: {lines[i-1]}")
PY
rg -n "_ctx_len|_kv_rewind_pending|_ensure_spec_dec_state_restored|_store_prefill_context|prepare_1st_drafter_inputs" tensorrt_llm/_torch/speculative/dflash.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 16244
🏁 Script executed:
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
path = Path('tensorrt_llm/_torch/speculative/dflash.py')
lines = path.read_text().splitlines()
for start, end in [(400, 590)]:
print(f'LINES {start}-{end}')
for i in range(start, min(end, len(lines)) + 1):
print(f"{i:4d}: {lines[i-1]}")
PY
printf '\nSEARCH\n'
rg -n "finally|except|saved_ctx_len|_ctx_len.copy_|is_warmup" tensorrt_llm/_torch/speculative/dflash.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 10010
🏁 Script executed:
set -euo pipefail
rg -n "_ensure_spec_dec_state_restored|_forward_impl\(" tensorrt_llm/_torch -g '*.py'Repository: NVIDIA/TensorRT-LLM
Length of output: 2189
🏁 Script executed:
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
path = Path('tensorrt_llm/_torch/speculative/interface.py')
lines = path.read_text().splitlines()
for start, end in [(944, 970)]:
print(f'LINES {start}-{end}')
for i in range(start, min(end, len(lines)) + 1):
print(f"{i:4d}: {lines[i-1]}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 1556
Restore warmup _ctx_len in _ensure_spec_dec_state_restored
tensorrt_llm/_torch/speculative/dflash.py:419-549 snapshots _ctx_len for warmup but only restores it on the success path. Because SpecWorkerBase.forward always calls _ensure_spec_dec_state_restored(...) in finally, a failure in _store_prefill_context or prepare_1st_drafter_inputs still leaves the warmup context lengths behind and can corrupt later DFlash context/RoPE positions. Fold the warmup restore into that cleanup hook.
🤖 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/_torch/speculative/dflash.py` around lines 374 - 386, Update
_ensure_spec_dec_state_restored to restore the warmup _ctx_len snapshot during
cleanup, including when _store_prefill_context or prepare_1st_drafter_inputs
fails. Reuse the existing warmup snapshot/state symbols and perform restoration
in the finally-driven hook alongside the existing superclass restoration and
pending KV rewind handling.
Collaborator
|
PR_Github #59242 [ run ] completed with state
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Note
Supersedes #16325. The original PR's head branch is owned by the repair-bot automation, whose periodic auto-rebases invalidated in-flight CI runs and re-triggered reviewer assignment; this PR carries the identical, verified commits on a developer-owned branch. See #16325 for prior review history (both CodeRabbit threads there were addressed and resolved).
Summary
f81b9d2787on 4x H100-80GB): the final general warmup — the max-shape memory-pool pre-population that runs right after CUDA-graph capture — hits atorch.OutOfMemoryErroron H100-80GB, whichmodel_engine._general_warmup_impldeliberately catches and skips. That aborted forward had already runEagle3OneModelWorker's_prepare_attn_metadata_for_spec_decwith no try/finally, so the restore never ran andattn_metadata._saved_tensorswas left populated on the persistent metadata object. Every subsequent forward then trips the bareassert len(self._saved_tensors) == 0inAttentionMetadata.prepare_for_spec_dec; a bareAssertionErrorhas an emptystr(), so the log showsEncountered an error in forward function:with no message, the executor swallows it,sample_statebecomesNone, and the loop dies with the secondaryAssertionError: Sampling failed(plus a tertiarythreads can only be started onceduring restart) — the signature QA reported. Not a regression: thev2_kv_cacheparametrization was added 2026-07-10 by [None][test] KV cache manager v2: add V2 + VSWA multi-GPU test coverage #16114 and this was its first QA run; KVCacheManagerV2 shifts available-token accounting enough to push the max-shape warmup into OOM on 80GB where V1 stays under.eagle3.py: wrap the draft path in try/finally with the prepare inside the try — the Eagle3 prepare override allocates clones after the base save populates_saved_tensors, so a failure inside prepare itself must also reach the restore (restore is safe on empty/partial state). CoversMTPEagleWorkervia inheritance.mtp.py:change_attn_metadata(which saves state and then mutates more of it) moved inside the try; restore in finally.pard.py: stale_kv_rewind_amountcleared up front; the deferred kv_lens rewind runs in the finally after the restore, so a draft failure cannot leavekv_lens_cudaincremented.dflash.py: same try/finally; akv_preparedflag gates the deferred rewind so a failure before_prepare_kv_for_draft_forwardcannot apply a stale rewind.draft_target.py: same try/finally.attention_backend/interface.py: messages added to the bare asserts so a recurrence names the unrestored fields instead of logging an empty error.waives.txt: unwaive the test on H100.Verification
A/B on 4x H100-80GB (dlcluster viking-prod-233, dgxh100 — same node class as the QA failure), both runs at main
f81b9d2787:OOM during general warmup with 2048 tokens, 2048 generation tokens. Skipping.followed by all 4 ranks asserting on the first real batch. FAILED in 563s.89.92(ref 90.30), GPQADiamond64.14(ref 65.0). PASSED in 628s.Follow-up in this PR: structural guarantee (template method)
Per review feedback that per-worker try/finally is not future-proof, the third commit centralizes the invariant in
SpecWorkerBase:forward()is now a base-class template method: it runs the worker-specific_forward_impl()and restores any unrestoredprepare_for_spec_decstate in afinally(with a warning log), plus the pending kv_lens rewind for PARD/DFlash.__init_subclass__rejects subclasses that overrideforwarddirectly, so a future worker cannot reintroduce the leak.forward->_forward_implrenames vs pre-fix code).finallyvisibly restores state when the max-shape warmup OOMs and the test PASSES (GSM8K 90.75/90.22, GPQADiamond 66.16/63.13 across two runs).test_force_accepted_tokens.pyunit tests: 21 passed.Test plan
tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v2_kv_cache-one_model]PASSED (unwaived by this PR)Links
This PR originally carried an automated repair-bot fix; its head was replaced by the verified fix above (root cause corrected: the trigger is the tolerated warmup OOM, not a transient serving failure; prepare moved inside the try; sibling workers covered).
Summary by CodeRabbit
Reliability
Testing