[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, (
+            "prepare_for_spec_dec called while fields "
+            f"{list(self._saved_tensors)} are still saved; a previous "
+            "forward likely raised between prepare_for_spec_dec and "
+            "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

@@ -172,6 +172,16 @@ def __init__(
         # graph compatible.
         self._ctx_buf_inited = False
         self._ctx_len = None
+        # Snapshot for rolling back in-place _ctx_len updates when a forward
+        # fails (or after warmup). See _ensure_spec_dec_state_restored.
+        self._saved_ctx_len = None
+        self._ctx_len_restore_pending = False
+        # 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
         self._batch_to_slot = None
         self._max_ctx = 0
         self._ctx_k_buf = None  # [max_batch, L, max_ctx+block, nkv, hd]
@@ -274,16 +284,20 @@ 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"):
+        if self._kv_rewind_amount is not None and hasattr(attn_metadata, "kv_lens_cuda"):
             nc = self._kv_rewind_nc
             bs = self._kv_rewind_bs
             attn_metadata.kv_lens_cuda[nc:bs] -= self._kv_rewind_amount
@@ -366,7 +380,28 @@ 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)
+        if (
+            getattr(self, "_ctx_len_restore_pending", False)
+            and self._ctx_len is not None
+            and self._saved_ctx_len is not None
+        ):
+            # A failed forward must not keep this iteration's in-place
+            # _ctx_len updates: roll back to the pre-forward snapshot.
+            self._ctx_len.copy_(self._saved_ctx_len)
+            self._ctx_len_restore_pending = False
+
+    def _forward_impl(
         self,
         input_ids,
         position_ids,
@@ -399,10 +434,18 @@ def forward(
         self._lazy_init_ctx_buffers(draft_model, spec_metadata, attn_metadata)
         spec_metadata._dflash_worker = self
 
-        # Save context lengths before warmup to prevent accumulation
+        # Save context lengths so both warmup and a failed forward can roll
+        # back the in-place _ctx_len updates made during drafting.
         is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing()
-        if is_warmup:
-            saved_ctx_len = self._ctx_len.clone()
+        if not torch.cuda.is_current_stream_capturing():
+            # Never allocate the snapshot while capturing a CUDA graph: the
+            # clone would live in the graph memory pool, and its replay-time
+            # writes could alias blocks reused by later captures. Rollback is
+            # only meaningful for eager/warmup forwards anyway; a failure
+            # during capture aborts the graph itself, and captured ops do not
+            # mutate _ctx_len until replay.
+            self._saved_ctx_len = self._ctx_len.clone()
+            self._ctx_len_restore_pending = True
 
         self._execute_guided_decoder_if_present(logits)
 
@@ -527,9 +570,10 @@ def forward(
             num_accepted_tokens,
         )
 
-        # Restore context lengths after warmup
+        # Restore context lengths after warmup; real runs keep the updates.
         if is_warmup:
-            self._ctx_len.copy_(saved_ctx_len)
+            self._ctx_len.copy_(self._saved_ctx_len)
+        self._ctx_len_restore_pending = False
 
         return {
             "logits": raw_logits,

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

@@ -91,6 +91,12 @@ def __init__(
         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}"
         )
@@ -144,6 +150,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,17 +162,32 @@ 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"):
+        if self._kv_rewind_amount is not None and hasattr(attn_metadata, "kv_lens_cuda"):
             nc = self._kv_rewind_nc
             bs = self._kv_rewind_bs
             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

@@ -289,7 +289,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)