[https://nvbugs/6442074][fix] Make one-model spec-dec attn-metadata save/restore exception-safe#16382
[https://nvbugs/6442074][fix] Make one-model spec-dec attn-metadata save/restore exception-safe#16382dongfengy wants to merge 6 commits into
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,
PerkzZheng,
YihuiLu512,
brnguyen2,
fredricz-20070104,
lori-ren,
mlefeb01,
niukuo,
yufeiwu-nv and
yunruis
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughSpeculative decoding workers now route execution through a shared forward wrapper that restores attention metadata after failures. DFlash and PARD defer KV rewinds and recover them during cleanup, while worker entrypoints and test stubs use ChangesSpeculative decoding state safety
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant SpecWorkerBase
participant DraftWorker
participant AttentionMetadata
participant KVState
SpecWorkerBase->>DraftWorker: invoke _forward_impl
DraftWorker->>AttentionMetadata: prepare speculative-decoding state
DraftWorker->>KVState: defer KV rewind
DraftWorker-->>SpecWorkerBase: complete or raise
SpecWorkerBase->>AttentionMetadata: restore saved state
SpecWorkerBase->>DraftWorker: apply pending KV rewind
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: 1
🧹 Nitpick comments (4)
tensorrt_llm/_torch/speculative/pard.py (1)
118-184: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winInitialize
_kv_rewind_pending/_kv_rewind_amount/_kv_rewind_nc/_kv_rewind_bsin__init__.The deferred-rewind state machine itself is correct (traced success, mid-failure, warmup, and
K==0paths — all consistent), but these four fields are only ever created dynamically inside_prepare_kv_for_draft_forward, and every read goes throughgetattr(...)/hasattr(...)guards rather than a real default. Declaring them in__init__(e.g.self._kv_rewind_pending = False) removes the reliance on dynamic attribute creation and makes the safety-net invariant explicit at construction time, consistent with the surrounding code's intent.As per coding guidelines, "initialize externally visible class members in the constructor."
♻️ Proposed constructor init
def __init__( self, spec_config: "PARDDecodingConfig", mapping: Mapping, use_separate_draft_kv_cache: bool = False, ): super().__init__(use_separate_draft_kv_cache) self.spec_config = spec_config self.mapping = mapping self.sa_enhancer: Optional[SADraftEnhancer] = None if getattr(spec_config, "sa_config", None) is not None: self.sa_enhancer = SADraftEnhancer(spec_config.sa_config.threshold) + # Deferred kv_lens_cuda rewind state (see _prepare_kv_for_draft_forward / + # _apply_kv_rewind_after_draft / _ensure_spec_dec_state_restored). + self._kv_rewind_pending = False + self._kv_rewind_amount = None + self._kv_rewind_nc = None + self._kv_rewind_bs = None logger.info( f"PARDWorker initialized with use_separate_draft_kv_cache={use_separate_draft_kv_cache}" )🤖 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/pard.py` around lines 118 - 184, Initialize _kv_rewind_pending, _kv_rewind_amount, _kv_rewind_nc, and _kv_rewind_bs in the class __init__ with safe defaults. Keep the existing assignments in _prepare_kv_for_draft_forward and update _apply_kv_rewind_after_draft and _ensure_spec_dec_state_restored to rely on these constructor-initialized fields rather than hasattr/getattr guards where appropriate.Source: Coding guidelines
tensorrt_llm/_torch/speculative/interface.py (1)
938-975: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a unit test exercising the failure-cleanup path.
This
forward()/_ensure_spec_dec_state_restored()pair is the core fix for the reported nvbugs regression, but the only spec-worker test visible in this context (test_force_accepted_tokens.py's_StubSpecWorker) only stubs_forward_implto raiseNotImplementedErrorand doesn't exercise the cleanup path. A test that makes a stub's_forward_implcallprepare_for_spec_decthen raise, and assertsattn_metadata.has_spec_dec_saved_stateisFalseafterward (and that a subsequent call doesn't trip the pairing assert), would directly regression-test this fix.🤖 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 - 975, The spec-worker tests lack coverage for cleanup after _forward_impl fails following prepare_for_spec_dec. Add a unit test using a stub worker whose _forward_impl prepares spec-dec state and then raises, assert has_spec_dec_saved_state is false afterward, and invoke forward again to confirm the restored state avoids the pairing assertion.tests/unittest/_torch/speculative/test_force_accepted_tokens.py (2)
56-57: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the new test hook.
Add parameter and return annotations to
_forward_impl, for example:- def _forward_impl(self, *args, **kwargs): + def _forward_impl(self, *args: object, **kwargs: object) -> None:As per coding guidelines, every Python function must be annotated.
🤖 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 56 - 57, Update the _forward_impl test hook with parameter and return type annotations, including an appropriate annotation for variadic arguments and keyword arguments, while preserving its existing NotImplementedError behavior.Source: Coding guidelines
47-58: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd a regression test for centralized failure cleanup.
This change only makes the stub satisfy the new subclass contract; it does not verify that
SpecWorkerBase.forwardrestores attention metadata when_forward_implraises. Extendtests/unittest/_torch/speculative/test_force_accepted_tokens.pywith a failing-hook case that asserts the saved state is cleared/restored after the exception.As per path instructions, test coverage must explicitly identify whether the changed behavior is sufficiently covered and provide concrete follow-up.
🤖 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 47 - 58, Add a regression test in the existing SpecWorkerBase test coverage that configures the attention metadata, makes the stub’s _forward_impl raise, and calls forward while asserting the exception propagates. Verify afterward that the saved attention state is cleared and the original metadata is restored, covering centralized failure cleanup.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 373-385: Extend the DFlash speculative state rollback to cover
persistent context lengths mutated by prepare_1st_drafter_inputs. Snapshot the
affected _ctx_len values before mutation, track the snapshot as pending, and
restore it from _ensure_spec_dec_state_restored alongside KV state when a
forward fails. Clear the pending context snapshot only after the complete
forward succeeds, preserving current successful-forward behavior.
---
Nitpick comments:
In `@tensorrt_llm/_torch/speculative/interface.py`:
- Around line 938-975: The spec-worker tests lack coverage for cleanup after
_forward_impl fails following prepare_for_spec_dec. Add a unit test using a stub
worker whose _forward_impl prepares spec-dec state and then raises, assert
has_spec_dec_saved_state is false afterward, and invoke forward again to confirm
the restored state avoids the pairing assertion.
In `@tensorrt_llm/_torch/speculative/pard.py`:
- Around line 118-184: Initialize _kv_rewind_pending, _kv_rewind_amount,
_kv_rewind_nc, and _kv_rewind_bs in the class __init__ with safe defaults. Keep
the existing assignments in _prepare_kv_for_draft_forward and update
_apply_kv_rewind_after_draft and _ensure_spec_dec_state_restored to rely on
these constructor-initialized fields rather than hasattr/getattr guards where
appropriate.
In `@tests/unittest/_torch/speculative/test_force_accepted_tokens.py`:
- Around line 56-57: Update the _forward_impl test hook with parameter and
return type annotations, including an appropriate annotation for variadic
arguments and keyword arguments, while preserving its existing
NotImplementedError behavior.
- Around line 47-58: Add a regression test in the existing SpecWorkerBase test
coverage that configures the attention metadata, makes the stub’s _forward_impl
raise, and calls forward while asserting the exception propagates. Verify
afterward that the saved attention state is cleared and the original metadata is
restored, covering centralized failure cleanup.
🪄 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: 55dd2371-718a-4614-a9e0-3d85ce86e4bd
📒 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
|
PR_Github #59246 [ run ] triggered by Bot. Commit: |
…ck, constructor-initialized rewind state, cleanup regression test - dflash.py: snapshot _ctx_len every forward and roll it back in _ensure_spec_dec_state_restored when the forward fails, so a tolerated failure cannot leave stale per-slot context lengths (the warmup-only restore now reuses the same snapshot). - dflash.py/pard.py: initialize the deferred kv_lens rewind state in the constructor and gate the rewind on _kv_rewind_amount instead of hasattr. - test_force_accepted_tokens.py: annotate the stub hook and add a regression test that a failure between prepare_for_spec_dec and restore leaves no saved state and does not trip the pairing assert afterwards. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: dongfengy <99041270+dongfengy@users.noreply.github.com>
|
/bot kill |
|
/bot run --disable-fail-fast |
|
PR_Github #59247 [ kill ] triggered by Bot. Commit: |
|
PR_Github #59248 [ run ] triggered by Bot. Commit: |
|
PR_Github #59247 [ kill ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #59258 [ run ] triggered by Bot. Commit: |
|
PR_Github #59259 [ kill ] triggered by Bot. Commit: |
|
PR_Github #59258 [ run ] completed with state |
|
PR_Github #59248 [ run ] completed with state |
|
PR_Github #59259 [ kill ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #59262 [ run ] triggered by Bot. Commit: |
| fails every subsequent forward at the pairing assert. | ||
| https://nvbugs/6442074 | ||
| """ | ||
| attn_metadata = kwargs.get("attn_metadata") |
There was a problem hiding this comment.
Could also do something like this:
with prepare_for_spec_dec(attn_metata) as attn_metadata:
# Run forward implwhere the prepare_for_spec_dec context manager restores everything on exit. It's basically the same thing but would be a bit more idiomatic. You could avoid adding attributes to self because you can save stuff directly in the context manager as well.
There was a problem hiding this comment.
Hi Mike, thanks for the good suggestion! Some of the states are in subclasses, looks like putting them in a context manager requires additional work, and it's possible this will result in the code being a little bit more complicated.
|
PR_Github #59262 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59296 [ run ] triggered by Bot. Commit: |
…ge fragments Review feedback on PR 16382: the fragments without placeholders do not need to be f-strings. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: dongfengy <99041270+dongfengy@users.noreply.github.com>
|
/bot kill |
|
PR_Github #59312 [ kill ] triggered by Bot. Commit: |
|
PR_Github #59296 [ run ] completed with state |
|
PR_Github #59312 [ kill ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #59315 [ run ] triggered by Bot. Commit: |
|
PR_Github #59315 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59377 [ run ] triggered by Bot. Commit: |
|
PR_Github #59377 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59426 [ run ] triggered by Bot. Commit: |
|
PR_Github #59426 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59439 [ run ] triggered by Bot. Commit: |
|
PR_Github #59439 [ run ] completed with state |
4 participants
Description
The max-shape general warmup can OOM on smaller GPUs — that is expected, and
model_enginedeliberately catches and skips it. The real bug: the one-model spec-dec workers (Eagle3 and siblings) are not exception-safe. The aborted forward has already saved attention metadata viaprepare_for_spec_decand never restores it, so dirty state stays behind on the persistent metadata, and every later forward dies on the stale-state assert — surfacing as the QA signatureAssertionError: Sampling failed(nvbugs/6442074).Fix:
SpecWorkerBase.forwardbecomes a template method: workers implement_forward_impl, and the basefinallyrestores any unrestored spec-dec metadata (plus the pending kv_lens rewind for PARD/DFlash).__init_subclass__rejectsforwardoverrides, so a future worker cannot reintroduce the leak.prepare_for_spec_decget messages (they used to produce an empty error string, which is why the log lineEncountered an error in forward function:was blank).Verification
A/B on the same 4x H100-80GB node at main: pristine TOT reproduces the QA failure (warmup OOM → stale-state assert, FAILED in ~9 min); with the fix the same warmup OOM occurs and is survived — GSM8K 90.2–90.8 (ref 90.3), GPQADiamond 63.1–66.2 (ref 65.0), PASSED. Control runs: workers unguarded → FAILED; template method only → PASSED.
test_force_accepted_tokens.py: 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