diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index fc6cfdd50222..02f77bbdf582 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/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, diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index efa31fa2cc1e..6c5ed826f55c 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/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, diff --git a/tensorrt_llm/_torch/speculative/draft_target.py b/tensorrt_llm/_torch/speculative/draft_target.py index fe6dc7d1e4c7..dcd2ac0d2d03 100644 --- a/tensorrt_llm/_torch/speculative/draft_target.py +++ b/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, diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 3978040724b1..f2320f48f427 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/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 diff --git a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py index 7d9e836f5b9a..00dff45ec03a 100644 --- a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py +++ b/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, diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 6244873efc43..acd16b5ba583 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/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: diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 6afed0fc02e9..acd9579997d3 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/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, diff --git a/tensorrt_llm/_torch/speculative/pard.py b/tensorrt_llm/_torch/speculative/pard.py index 34a4509314cd..3d19ff4ebc54 100644 --- a/tensorrt_llm/_torch/speculative/pard.py +++ b/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, diff --git a/tensorrt_llm/_torch/speculative/sa_worker.py b/tensorrt_llm/_torch/speculative/sa_worker.py index 90cd617cef2a..07023fe6d661 100644 --- a/tensorrt_llm/_torch/speculative/sa_worker.py +++ b/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, diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 130bf38ad7dc..686c3e1bca98 100644 --- a/tests/integration/test_lists/waives.txt +++ b/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) diff --git a/tests/unittest/_torch/speculative/test_force_accepted_tokens.py b/tests/unittest/_torch/speculative/test_force_accepted_tokens.py index ac8228562bf8..e3ac4bffaa8c 100644 --- a/tests/unittest/_torch/speculative/test_force_accepted_tokens.py +++ b/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()