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

tensorrt_llm/_torch/attention_backend/interface.py

@@ -411,10 +411,14 @@ def create_cuda_graph_metadata(self,
         return cuda_graph_metadata
 
     def prepare_for_spec_dec(self, *fields) -> None:
-        assert len(self._saved_tensors) == 0
+        assert len(self._saved_tensors) == 0, (
+            f"prepare_for_spec_dec called while fields "
+            f"{list(self._saved_tensors)} are still saved; a previous "
+            f"forward likely raised between prepare_for_spec_dec and "
+            f"restore_from_spec_dec")
         for f in fields:
             v = getattr(self, f)
-            assert isinstance(v, torch.Tensor)
+            assert isinstance(v, torch.Tensor), f"{f} is not a torch.Tensor"
             self._saved_tensors[f] = v
             setattr(self, f, v.clone())
 
@@ -423,6 +427,11 @@ def restore_from_spec_dec(self) -> None:
             setattr(self, f, v)
         self._saved_tensors.clear()
 
+    @property
+    def has_spec_dec_saved_state(self) -> bool:
+        """True when prepare_for_spec_dec state has not been restored yet."""
+        return bool(self._saved_tensors)
+
     def update_spec_dec_param(
             self,
             batch_size,

tensorrt_llm/_torch/speculative/dflash.py

@@ -274,13 +274,17 @@ def _prepare_kv_for_draft_forward(
 
             if batch_size > num_contexts:
                 attn_metadata.kv_lens_cuda[num_contexts:batch_size] += 1
+            self._kv_rewind_pending = True
 
             attn_metadata.update_for_spec_dec()
 
     def _apply_kv_rewind_after_draft(self, attn_metadata, spec_metadata):
         """Apply the deferred kv_lens rewind after the draft forward."""
+        self._kv_rewind_pending = False
         is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing()
         if is_warmup:
+            # kv_lens_cuda was saved by prepare_for_spec_dec in this mode and
+            # is restored wholesale, so no rewind is needed.
             return
 
         if hasattr(self, "_kv_rewind_amount") and hasattr(attn_metadata, "kv_lens_cuda"):
@@ -367,7 +371,19 @@ def _store_prefill_context(
                 self._ctx_v_buf[slot, :, cur:end] = chunk_v.permute(1, 0, 2, 3)
             offset += slen
 
-    def forward(
+    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(
         self,
         input_ids,
         position_ids,

tensorrt_llm/_torch/speculative/draft_target.py

@@ -153,7 +153,7 @@ def _update_kv_for_chained_draft_step(
 
         attn_metadata.update_for_spec_dec()
 
-    def forward(
+    def _forward_impl(
         self,
         input_ids,
         position_ids,

tensorrt_llm/_torch/speculative/eagle3.py

@@ -625,6 +625,14 @@ def __init__(self,
         # MTP Eagle: lazily-resolved flag for Mamba hybrid cache support
         self._is_mamba_hybrid_cache = None
 
+        # Worker-side saved spec-dec params; initialized so that a failure
+        # inside _prepare_attn_metadata_for_spec_dec (e.g. an OOM in the
+        # clone calls) cannot turn the cleanup restore into AttributeError.
+        self._saved_packed_mask = None
+        self._saved_position_offsets = None
+        self._saved_position_offsets_cpp = None
+        self._saved_generation_lengths = None
+
     @property
     def max_draft_len(self) -> int:
         return self.spec_config.max_draft_len
@@ -680,15 +688,15 @@ def _restore_attn_metadata_from_spec_dec(self, attn_metadata):
     # Skip torch.compile for now since current Torch is not compatible with Triton 3.4
     # @torch.compile(options={"max-autotune": True})
 
-    def forward(self,
-                input_ids,
-                position_ids,
-                hidden_states,
-                logits,
-                attn_metadata,
-                spec_metadata,
-                draft_model,
-                resource_manager=None):
+    def _forward_impl(self,
+                      input_ids,
+                      position_ids,
+                      hidden_states,
+                      logits,
+                      attn_metadata,
+                      spec_metadata,
+                      draft_model,
+                      resource_manager=None):
 
         runtime_draft_len = spec_metadata.runtime_draft_len
         # skip the draft forward if the runtime draft length is 0

tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py

@@ -375,7 +375,7 @@ def _ensure_spec_tree_manager(self, resource_manager):
                 )
 
     @nvtx_range("eagle3_dyn.forward")
-    def forward(
+    def _forward_impl(
         self,
         input_ids,
         position_ids,
@@ -387,11 +387,11 @@ def forward(
         resource_manager=None,
     ):
         """Override to add accepted_draft_tokens_indices to output."""
-        # Initialize spec_tree_manager before super().forward() which calls
+        # Initialize spec_tree_manager before super()._forward_impl() which calls
         # _forward_draft_loop needing spec_tree_manager.
         if resource_manager is not None:
             self._ensure_spec_tree_manager(resource_manager)
-        output = super().forward(
+        output = super()._forward_impl(
             input_ids,
             position_ids,
             hidden_states,

tensorrt_llm/_torch/speculative/interface.py

@@ -28,6 +28,7 @@
 from tensorrt_llm.logger import logger
 
 from ..._utils import get_sm_version, prefer_pinned
+from ..attention_backend.interface import AttentionMetadata
 from ..attention_backend.trtllm import (AttentionBackend, TrtllmAttention,
                                         TrtllmAttentionMetadata)
 from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE
@@ -925,6 +926,54 @@ def __init__(self, use_separate_draft_kv_cache: bool = False):
         self._force_accept_rng_pool: Optional[torch.Tensor] = None
         self._force_accept_rng_counter: Optional[torch.Tensor] = None
 
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        if "forward" in cls.__dict__:
+            raise TypeError(
+                f"{cls.__name__} must not override SpecWorkerBase.forward; "
+                f"implement _forward_impl instead. SpecWorkerBase.forward "
+                f"guarantees spec-dec attn-metadata cleanup when a forward "
+                f"fails (https://nvbugs/6442074).")
+
+    def forward(self, *args, **kwargs):
+        """Run _forward_impl with guaranteed spec-dec metadata cleanup.
+
+        Tolerated forward failures (e.g. an OOM during the max-shape general
+        warmup, or an error-budget-tolerated serving exception) must not leak
+        the attn-metadata state saved by prepare_for_spec_dec: a stale save
+        fails every subsequent forward at the pairing assert.
+        https://nvbugs/6442074
+        """
+        attn_metadata = kwargs.get("attn_metadata")
+        spec_metadata = kwargs.get("spec_metadata")
+        if attn_metadata is None or spec_metadata is None:
+            for a in args:
+                if attn_metadata is None and isinstance(a, AttentionMetadata):
+                    attn_metadata = a
+                elif spec_metadata is None and isinstance(a, SpecMetadata):
+                    spec_metadata = a
+        try:
+            return self._forward_impl(*args, **kwargs)
+        finally:
+            self._ensure_spec_dec_state_restored(attn_metadata, spec_metadata)
+
+    @abstractmethod
+    def _forward_impl(self, *args, **kwargs):
+        """Worker-specific forward logic, called by SpecWorkerBase.forward."""
+
+    def _ensure_spec_dec_state_restored(self, attn_metadata, spec_metadata):
+        """Restore attn-metadata spec-dec state if a failure skipped it.
+
+        No-op on the success path: workers restore at their preferred point
+        and this sees no saved state. Subclasses with extra transient state
+        (e.g. the deferred kv_lens rewind in PARD/DFlash) extend this.
+        """
+        if attn_metadata is not None and attn_metadata.has_spec_dec_saved_state:
+            logger.warning(
+                "Spec-dec worker forward failed between prepare_for_spec_dec "
+                "and restore_from_spec_dec; restoring attn metadata state.")
+            self._restore_attn_metadata_from_spec_dec(attn_metadata)
+
     @property
     @abstractmethod
     def max_draft_len(self) -> int:

tensorrt_llm/_torch/speculative/mtp.py

@@ -291,7 +291,7 @@ def __init__(self,
     def max_draft_len(self) -> int:
         return self.spec_config.max_draft_len
 
-    def forward(
+    def _forward_impl(
         self,
         input_ids,
         position_ids,

tensorrt_llm/_torch/speculative/pard.py

@@ -144,6 +144,7 @@ def _prepare_kv_for_draft_forward(
 
             if batch_size > num_contexts:
                 attn_metadata.kv_lens_cuda[num_contexts:batch_size] += 1
+            self._kv_rewind_pending = True
 
             attn_metadata.update_for_spec_dec()
 
@@ -155,8 +156,11 @@ def _apply_kv_rewind_after_draft(self, attn_metadata, spec_metadata):
         by prepare_for_spec_dec) to avoid cumulative shrinkage.  Applied during
         capture and normal inference.
         """
+        self._kv_rewind_pending = False
         is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing()
         if is_warmup:
+            # kv_lens_cuda was saved by prepare_for_spec_dec in this mode and
+            # is restored wholesale, so no rewind is needed.
             return
 
         if hasattr(self, "_kv_rewind_amount") and hasattr(attn_metadata, "kv_lens_cuda"):
@@ -165,7 +169,19 @@ def _apply_kv_rewind_after_draft(self, attn_metadata, spec_metadata):
             attn_metadata.kv_lens_cuda[nc:bs] -= self._kv_rewind_amount
             attn_metadata.kv_lens_cuda[nc:bs].clamp_(min=0)
 
-    def forward(
+    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(
         self,
         input_ids,
         position_ids,

tensorrt_llm/_torch/speculative/sa_worker.py

@@ -119,7 +119,7 @@ def __init__(self, spec_config: "SADecodingConfig", model_config=None):
     def max_draft_len(self) -> int:
         return self._max_draft_len
 
-    def forward(
+    def _forward_impl(
         self,
         input_ids: torch.Tensor,
         position_ids: torch.Tensor,

tests/integration/test_lists/waives.txt

@@ -284,7 +284,6 @@ full:H100/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_
 full:H100/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar-eagle3_one_model=False] SKIP (https://nvbugs/6422334)
 full:H100/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=False] SKIP (https://nvbugs/6422343)
 full:H100/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=True] SKIP (https://nvbugs/6422343)
-full:H100/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v2_kv_cache-one_model] SKIP (https://nvbugs/6442074)
 full:H100/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp1-CUTLASS] SKIP (https://nvbugs/6273850)
 full:H100/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6422318)
 full:H100/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6422318)

tests/unittest/_torch/speculative/test_force_accepted_tokens.py

@@ -44,7 +44,7 @@
 
 
 class _StubSpecWorker(SpecWorkerBase):
-    """Concrete ``SpecWorkerBase`` that stubs out the only abstract API.
+    """Concrete ``SpecWorkerBase`` that stubs out the abstract API.
 
     Used purely to drive ``_apply_force_accepted_tokens`` in isolation.
     """
@@ -53,6 +53,9 @@ class _StubSpecWorker(SpecWorkerBase):
     def max_draft_len(self) -> int:
         return 8
 
+    def _forward_impl(self, *args, **kwargs):
+        raise NotImplementedError
+
 
 def _make_worker(value: Optional[float] = None) -> _StubSpecWorker:
     worker = _StubSpecWorker()