[https://nvbugs/6442074][fix] Make one-model spec-dec attn-metadata save/restore exception-safe

[trtllm-agent]

Summary

• Root cause (reproduced at main f81b9d2787 on 4x H100-80GB): the final general warmup — the max-shape memory-pool pre-population that runs right after CUDA-graph capture — hits a torch.OutOfMemoryError on H100-80GB, which model_engine._general_warmup_impl deliberately catches and skips. That aborted forward had already run Eagle3OneModelWorker's _prepare_attn_metadata_for_spec_dec with no try/finally, so the restore never ran and attn_metadata._saved_tensors was left populated on the persistent metadata object. Every subsequent forward then trips the bare assert len(self._saved_tensors) == 0 in AttentionMetadata.prepare_for_spec_dec; a bare AssertionError has an empty str(), so the log shows Encountered an error in forward function: with no message, the executor swallows it, sample_state becomes None, and the loop dies with the secondary AssertionError: Sampling failed (plus a tertiary threads can only be started once during restart) — the signature QA reported. Not a regression: the v2_kv_cache parametrization 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.
• Fix: make the spec-dec attn-metadata save/restore exception-safe in all one-model workers:
  • 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). Covers MTPEagleWorker via 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_amount cleared up front; the deferred kv_lens rewind runs in the finally after the restore, so a draft failure cannot leave kv_lens_cuda incremented.
  • dflash.py: same try/finally; a kv_prepared flag gates the deferred rewind so a failure before _prepare_kv_for_draft_forward cannot 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:

• Pristine TOT: reproduces the exact QA failure — 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.
• With this fix: the same warmup OOM occurs and is survived — GSM8K 89.92 (ref 90.30), GPQADiamond 64.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 unrestored prepare_for_spec_dec state in a finally (with a warning log), 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.
• The five per-worker try/finally blocks from the second commit are removed; worker bodies return to their original straight-line form (pure forward -> _forward_impl renames vs pre-fix code).
• A/B on the same 4x H100-80GB node: workers unguarded with no template method -> FAILED in 215s with the nvbugs/6442074 signature; template method only -> the base finally visibly 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.py unit tests: 21 passed.

Test plan

• Reproduce the original failure at main TOT on the same GPU type (H100-80GB x4)
• Verify the fix on the same node with the OOM trigger present in the log
• 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

• Bug: https://nvbugs/6442074

---

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

• Bug Fixes

  • Improved speculative decoding reliability by ensuring attention state is restored after draft processing, including when errors occur.
  • Prevented stale attention metadata from affecting subsequent CUDA-graph executions.

• Tests

  • Removed an outdated waiver for an H100 Eagle3 accuracy test, allowing it to run normally.

[coderabbitai]

📝 Walkthrough

Walkthrough

Eagle3, MTP, and PARD draft execution now restores speculative-decoding attention metadata when failures occur. The related H100 integration waiver is removed.

Changes

Speculative decoding metadata restoration

Layer / File(s)	Summary
Guarded draft execution and coverage update tensorrt_llm/_torch/speculative/eagle3.py, tensorrt_llm/_torch/speculative/mtp.py, tensorrt_llm/_torch/speculative/pard.py, tests/integration/test_lists/waives.txt	Eagle3, MTP, and PARD move attention-metadata restoration into finally blocks around draft execution, while the H100 Eagle3 reuse waiver is removed.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: cascade812

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title matches the bug fix and accurately summarizes the exception-safe attn-metadata restore change.
Description check	✅ Passed	The description explains the root cause, fix, verification, and test plan, though it doesn't follow the template headings exactly.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tensorrt_llm/_torch/speculative/eagle3.py (1)

757-783: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Restore all mutated attention state in finally.

Line 771 restores saved tensors, but Lines 774-776 and 783 run only on success. If the draft loop raises after changing all_rank_num_tokens or use_spec_decoding, the worker remains in draft state. Capture the original value before the try and move both restorations into finally.

Proposed fix

+        original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens
         try:
             ...
-            original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens
             ...
         finally:
             self._restore_attn_metadata_from_spec_dec(attn_metadata)
-        if original_all_rank_num_tokens is not None:
-            attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens
+            if original_all_rank_num_tokens is not None:
+                attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens
+            attn_metadata.use_spec_decoding = True
             ...
-        attn_metadata.use_spec_decoding = True

🤖 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/eagle3.py` around lines 757 - 783, Move
capture of the original all_rank_num_tokens value before the try block, then
update the finally block in the surrounding draft-token flow to restore both
attn_metadata.all_rank_num_tokens and attn_metadata.use_spec_decoding on every
exit, including exceptions. Remove the success-only restoration and ensure
use_spec_decoding is restored to its pre-loop value rather than unconditionally
set afterward.

🤖 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/mtp.py`:
- Around line 440-448: Move the attn_metadata update in the MTP forward flow
into the existing try/finally block so change_attn_metadata and its internal
_prepare_attn_metadata_for_spec_dec updates are protected. Ensure the finally
cleanup always clears _saved_tensors when either metadata preparation or the
draft loop raises, while preserving the current update behavior when no
exception occurs.

In `@tensorrt_llm/_torch/speculative/pard.py`:
- Around line 243-339: Move the deferred _apply_kv_rewind_after_draft call into
the existing finally block so it runs after _restore_attn_metadata_from_spec_dec
whether draft execution succeeds or raises. Track whether
_prepare_kv_for_draft_forward completed successfully if needed, and apply the
rewind only when KV preparation incremented the metadata; remove the
post-finally call while preserving restore-before-rewind ordering.

---

Outside diff comments:
In `@tensorrt_llm/_torch/speculative/eagle3.py`:
- Around line 757-783: Move capture of the original all_rank_num_tokens value
before the try block, then update the finally block in the surrounding
draft-token flow to restore both attn_metadata.all_rank_num_tokens and
attn_metadata.use_spec_decoding on every exit, including exceptions. Remove the
success-only restoration and ensure use_spec_decoding is restored to its
pre-loop value rather than unconditionally set afterward.

🪄 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: 4d120b8c-9036-406f-a7ca-9ea57204f518

📥 Commits

Reviewing files that changed from the base of the PR and between 5603308 and f598b6a.

📒 Files selected for processing (4)

• tensorrt_llm/_torch/speculative/eagle3.py
• tensorrt_llm/_torch/speculative/mtp.py
• tensorrt_llm/_torch/speculative/pard.py
• tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt

[dongfengy]

Re the outside-diff review comment on eagle3.py ("Restore all mutated attention state in finally" — all_rank_num_tokens / use_spec_decoding): fact-checked against head adc97c0, both premises do not hold, so no change is made:

• all_rank_num_tokens: the draft-loop mutation of attn_metadata.all_rank_num_tokens happens inside Eagle3DraftModel.forward, which already restores it in its own try/finally (modeling_speculative.py:475-477). Additionally, under attention DP the field is unconditionally reassigned from a fresh allgather during every forward's input prep (model_engine.py _prepare_tp_inputs), and with enable_attention_dp=False it is None throughout. The success-path assignment in the worker is redundant defense, not the last line of it.
• use_spec_decoding: recomputed unconditionally at the start of every engine forward via attn_metadata.update_spec_dec_param(...) (model_engine.py → trtllm.py update_spec_dec_param, which sets self.use_spec_decoding = self.is_spec_decoding_enabled) before any reader; CUDA-graph replay executes no host code and never reads the flag. A stale False after a tolerated failure is dead state that self-heals before the next consumer.

Unlike _saved_tensors — which has a pairing assert and is fatal if stale (the subject of this PR) — neither field has a cross-iteration reader before being recomputed.

[dongfengy]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59053 [ run ] triggered by Bot. Commit: adc97c0 Link to invocation

[tensorrt-cicd]

PR_Github #59053 [ run ] completed with state SUCCESS. Commit: adc97c0
/LLM/main/L0_MergeRequest_PR pipeline #47579 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[dongfengy]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59116 [ run ] triggered by Bot. Commit: adc97c0 Link to invocation

[dongfengy]

/bot kill

[dongfengy]

/bot run --disable-fail-fast

[dongfengy]

/bot kill

[dongfengy]

Closing in favor of #16381, which carries the identical (content-verified) commits on a developer-owned branch. The repair-bot automation that owns this branch periodically force-rebases it, which unfortunately invalidates in-flight CI runs and re-triggers reviewer assignment, so continuing here would keep racing the automation. Review history in this PR remains valid: both inline CodeRabbit findings were addressed structurally and the threads resolved; see #16381 for the current CI. Thanks repair-bot for the initial draft and the base-freshness diligence!

[tensorrt-cicd]

PR_Github #59241 [ ] completed with state FAILURE. Commit: 43f79e1
Not allowed on merged PR

Link to invocation

[tensorrt-cicd]

PR_Github #59240 [ ] completed with state FAILURE. Commit: 43f79e1
Not allowed on merged PR

Link to invocation

[tensorrt-cicd]

PR_Github #59116 [ run ] completed with state ABORTED. Commit: adc97c0

Link to invocation