diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index fc6cfdd50222..fec6ba4e6b92 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, ( + "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, diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index fece18aacc27..a9dca0a4addc 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/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, 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..01938943020d 100644 --- a/tensorrt_llm/_torch/speculative/pard.py +++ b/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, 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 12a0a3dbf464..7398141badd1 100644 --- a/tests/integration/test_lists/waives.txt +++ b/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) diff --git a/tests/unittest/_torch/speculative/test_force_accepted_tokens.py b/tests/unittest/_torch/speculative/test_force_accepted_tokens.py index ac8228562bf8..fbfaca610386 100644 --- a/tests/unittest/_torch/speculative/test_force_accepted_tokens.py +++ b/tests/unittest/_torch/speculative/test_force_accepted_tokens.py @@ -35,6 +35,7 @@ import pytest import torch +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata from tensorrt_llm._torch.speculative.interface import ( FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR, SpecWorkerBase, @@ -44,7 +45,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 +54,9 @@ class _StubSpecWorker(SpecWorkerBase): def max_draft_len(self) -> int: return 8 + def _forward_impl(self, *args: object, **kwargs: object) -> None: + raise NotImplementedError + def _make_worker(value: Optional[float] = None) -> _StubSpecWorker: worker = _StubSpecWorker() @@ -363,3 +367,28 @@ def test_cuda_graph_replay_advances_rng_state(): if __name__ == "__main__": pytest.main([__file__, "-v"]) + + +class _FailingSpecWorker(_StubSpecWorker): + """Stub whose forward saves spec-dec metadata state and then fails.""" + + def _forward_impl(self, *args: object, **kwargs: object) -> None: + attn_metadata = kwargs["attn_metadata"] + attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda") + raise RuntimeError("simulated draft failure") + + +def test_forward_restores_spec_dec_state_on_failure() -> None: + """A failure between prepare_for_spec_dec and restore must not leak saved + attn-metadata state; SpecWorkerBase.forward restores it in its cleanup + (https://nvbugs/6442074).""" + _require_cuda() + attn_metadata = AttentionMetadata(max_num_requests=2, max_num_tokens=16) + attn_metadata.seq_lens = torch.ones(2, dtype=torch.int32) + worker = _FailingSpecWorker() + for _ in range(2): + # The second iteration would trip the pairing assert inside + # prepare_for_spec_dec if the first failure had leaked saved state. + with pytest.raises(RuntimeError, match="simulated draft failure"): + worker(attn_metadata=attn_metadata, spec_metadata=None) + assert not attn_metadata.has_spec_dec_saved_state