From 1d48a069cb9e021acf9190c38739c7862e7d59cc Mon Sep 17 00:00:00 2001 From: dongfengy <99041270+dongfengy@users.noreply.github.com> Date: Mon, 13 Jul 2026 22:18:41 +0000 Subject: [PATCH 1/6] [https://nvbugs/6442074][fix] Make Eagle3 one-model spec-dec metadata save/restore exception-safe TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v2_kv_cache-one_model] fails on H100 with "AssertionError: Sampling failed". The real chain: the final general warmup (max-shape memory-pool prepop) OOMs on H100-80GB and is tolerated by model_engine, but the aborted forward had already run Eagle3OneModelWorker prepare of the attn metadata without try/finally, so the restore never ran. Every subsequent forward then trips the bare "assert len(self._saved_tensors) == 0" in prepare_for_spec_dec, whose empty str() masks the error, cascading into "Sampling failed". - 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, so a failure inside prepare must also reach the restore). - interface.py: add messages to the bare asserts so a recurrence names the unrestored fields instead of logging an empty error. - waives.txt: unwaive the test on H100. Verified on 4x H100-80GB (viking-prod-233) at main f81b9d2787: pristine TOT reproduces the exact QA failure; with this fix the same warmup OOM occurs but serving survives - GSM8K 89.92 (ref 90.30), GPQADiamond 64.14 (ref 65.0), test PASSED. Co-Authored-By: Claude Fable 5 Signed-off-by: dongfengy <99041270+dongfengy@users.noreply.github.com> --- .../_torch/attention_backend/interface.py | 8 ++- tensorrt_llm/_torch/speculative/eagle3.py | 67 +++++++++++-------- tests/integration/test_lists/waives.txt | 34 ++++------ 3 files changed, 56 insertions(+), 53 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index fc6cfdd50222..4ad8cb273666 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()) diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 3978040724b1..67f5a6866d9a 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -737,35 +737,44 @@ def forward(self, max_draft_len=runtime_draft_len, ) - # Save the old attn_metadata and spec_metadata - self._prepare_attn_metadata_for_spec_dec(attn_metadata) - - # Prepare inputs for the 1st draft model forward - position_ids = position_ids.squeeze(0) - inputs = self.prepare_1st_drafter_inputs( - input_ids=input_ids, - position_ids=position_ids, - hidden_states=hidden_states, - accepted_tokens=accepted_tokens, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - draft_model=draft_model) - - # Predict draft tokens. ``original_all_rank_num_tokens`` is saved here - # so the post-loop restore (below) can put attn_metadata back into a - # state the target model expects. - original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens - - # Get the draft KV cache manager if using separate layouts - draft_kv_cache_manager = self.get_draft_kv_cache_manager( - resource_manager) - - next_draft_tokens = self._forward_draft_loop( - inputs, attn_metadata, spec_metadata, draft_model, - draft_kv_cache_manager, num_contexts, num_gens, batch_size, - num_accepted_tokens, original_all_rank_num_tokens, resource_manager) - # restore attn_metadata to support cuda graph - self._restore_attn_metadata_from_spec_dec(attn_metadata) + # Save the old attn_metadata and spec_metadata. The save/restore + # pair must be exception-safe: a failure between them (e.g. an OOM + # during the max-shape general warmup, which model_engine tolerates + # and skips) would otherwise leave stale saved state on the + # persistent attn_metadata and fail every subsequent forward at + # the pairing assert in prepare_for_spec_dec. + # https://nvbugs/6442074 + try: + self._prepare_attn_metadata_for_spec_dec(attn_metadata) + + # Prepare inputs for the 1st draft model forward + position_ids = position_ids.squeeze(0) + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model) + + # Predict draft tokens. ``original_all_rank_num_tokens`` is saved + # here so the post-loop restore (below) can put attn_metadata + # back into a state the target model expects. + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + + # Get the draft KV cache manager if using separate layouts + draft_kv_cache_manager = self.get_draft_kv_cache_manager( + resource_manager) + + next_draft_tokens = self._forward_draft_loop( + inputs, attn_metadata, spec_metadata, draft_model, + draft_kv_cache_manager, num_contexts, num_gens, batch_size, + num_accepted_tokens, original_all_rank_num_tokens, + resource_manager) + finally: + # restore attn_metadata to support cuda graph + self._restore_attn_metadata_from_spec_dec(attn_metadata) # restore all_rank_num_tokens for attention DP if original_all_rank_num_tokens is not None: attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 12a0a3dbf464..3678b67b2f7e 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -22,8 +22,9 @@ accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[u accuracy/test_disaggregated_serving.py::TestQwen3_30B_A3B::test_mixed_ctx_gen_model[ctxpp2gentp2] SKIP (https://nvbugs/5748664) accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_nvfp4] SKIP (https://nvbugs/6336747) accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[qwen3vl_2b_instruct] SKIP (https://nvbugs/6422294) +accuracy/test_llm_api.py::TestLlama3_1_8BInstruct::test_guided_decoding_4gpus[xgrammar] SKIP (https://nvbugs/5346443) +accuracy/test_llm_api_autodeploy.py::TestGemma4MoE::test_bf16 SKIP (https://nvbugs/6445322) accuracy/test_llm_api_autodeploy.py::TestMiniMaxM2::test_finegrained_fp8 SKIP (https://nvbugs/6396422) -accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[nvfp4_ws8_80gb-trtllm] SKIP (https://nvbugs/6450341) accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_nvfp4[8] SKIP (https://nvbugs/6412108) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp] SKIP (https://nvbugs/6428101) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp_trtllm] SKIP (https://nvbugs/6426868) @@ -39,7 +40,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=0- accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6412102) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6388129) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6426847) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6445456) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6402058) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6278337) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6428057) @@ -49,12 +49,9 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mt accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6388153) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6428094) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6428096) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6445456) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6198774) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6445456) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=0-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6198785) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=0-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6198785) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6445456) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6198774) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_online_eplb[mtp_nextn=2-moe_backend=WIDEEP] SKIP (https://nvbugs/6313993) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[pp4-mtp_nextn=0-fp8kv=False-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False-sampler_async_worker=False] SKIP (https://nvbugs/6388153) @@ -63,34 +60,24 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpu accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp2pp2-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6427411) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp2pp2-mtp_nextn=0-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-sampler_async_worker=False] SKIP (https://nvbugs/6427411) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp2pp2-mtp_nextn=0-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6427411) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp4-mtp_nextn=0-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=True-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6445456) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp4-mtp_nextn=0-fp8kv=False-attention_dp=True-cuda_graph=False-overlap_scheduler=False-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6445456) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp4-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6445456) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp4-mtp_nextn=0-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6445456) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=2-fp8kv=False-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6388363) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTEDSL-mtp_nextn=0-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6428087) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTEDSL-mtp_nextn=0-tp2pp2-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6427411) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-ep4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6445456) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6445456) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/5945081) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6384625) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6427411) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6428063) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6384625) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6384625) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp2pp2-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6445472) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6445456) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6272673) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-pp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6422432) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6245394) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6384625) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-tp2pp2-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6384625) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_auto_dtype SKIP (https://nvbugs/6450333) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_guided_decoding_4gpus[one_model] SKIP (https://nvbugs/5596343) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v1_kv_cache-True-True-triton-auto] SKIP (https://nvbugs/6026676) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-dp4-cutlass-auto] SKIP (https://nvbugs/6388142) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-ep4-cutlass-auto] SKIP (https://nvbugs/5596343) -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-ep4-trtllm-fp8] SKIP (https://nvbugs/6450338) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-tp4-cutlass-auto] SKIP (https://nvbugs/5596343) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-tp4-cutlass-fp8] SKIP (https://nvbugs/5651865) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-dp4-cutlass-auto] SKIP (https://nvbugs/5596343) @@ -184,6 +171,7 @@ examples/test_granite.py::test_llm_granite[granite-3.0-2b-instruct-bfloat16] SKI examples/test_medusa.py::test_llm_medusa_with_qaunt_base_model_1gpu[fp8-use_cpp_session-medusa-vicuna-7b-v1.3-4-heads-float16-bs1] SKIP (https://nvbugs/5802248) examples/test_medusa.py::test_llm_medusa_with_qaunt_base_model_1gpu[fp8-use_py_session-medusa-vicuna-7b-v1.3-4-heads-float16-bs1] SKIP (https://nvbugs/5333849) examples/test_multimodal.py::test_llm_fp8_multimodal_general[fp8-fp8-scienceqa-Llama-3.2-11B-Vision-Instruct-pp:1-tp:1-bfloat16-bs:1-cpp_e2e:False] SKIP (https://nvbugs/5222697) +examples/test_multimodal.py::test_llm_multimodal_general[Llama-3.2-11B-Vision-pp:1-tp:1-bfloat16-bs:1-cpp_e2e:False-nb:1] SKIP (https://nvbugs/5333818) examples/test_multimodal.py::test_llm_multimodal_general[Llama-3.2-11B-Vision-pp:1-tp:1-bfloat16-bs:8-cpp_e2e:False-nb:1] SKIP (https://nvbugs/5333818) examples/test_multimodal.py::test_llm_multimodal_general[Llama-3.2-11B-Vision-pp:1-tp:2-bfloat16-bs:1-cpp_e2e:False-nb:1] SKIP (https://nvbugs/5333818) examples/test_nemotron.py::test_llm_nemotron_3_8b_1gpu[bfloat16-fp8] SKIP (https://nvbugs/4961624) @@ -287,9 +275,6 @@ full:H100/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_ full:H100/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=True] SKIP (https://nvbugs/6313072) full:H100/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[llguidance-eagle3_one_model=False] SKIP (https://nvbugs/6422334) 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) @@ -297,16 +282,15 @@ full:H100/accuracy/test_llm_api_pytorch_multimodal.py::TestQwen3_5_27B_VL::test_ full:H100/accuracy/test_llm_api_pytorch_multimodal.py::TestQwen3_5_35B_A3B_VL::test_auto_dtype SKIP (https://nvbugs/6442073) full:H100/disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_tllm_gen_helix[DeepSeek-V3-Lite-bf16-short_prompt] SKIP (https://nvbugs/6414762) full:H100/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959) -full:H100/disaggregated/test_disaggregated.py::test_disaggregated_mixed_stress_test[req10k-conc512-qwen3_32b_fp8_mixed_stress] SKIP (https://nvbugs/6440089) full:H100/disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp4-TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6344107) full:H100/disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-qwen3_32b_fp8_stress] SKIP (https://nvbugs/6312828) -full:H100/test_e2e.py::test_multi_nodes_eval[Qwen3/Qwen3-235B-A22B-tp16-mmlu] SKIP (https://nvbugs/6423926) full:H100/test_e2e.py::test_openai_chat_guided_decoding[openai/gpt-oss-120b] SKIP (https://nvbugs/6384375) full:H100/test_e2e.py::test_ptp_quickstart_advanced_deepseek_multi_nodes[DeepSeek-V3] SKIP (https://nvbugs/6418410) -full:H100/test_e2e.py::test_qwen_e2e_cpprunner_large_new_tokens[DeepSeek-R1-Distill-Qwen-1.5B-DeepSeek-R1-Distill-Qwen-1.5B] SKIP (https://nvbugs/6414760) full:H100_PCIe/unittest/llmapi/test_llm_pytorch.py::test_llama_7b_multi_lora_evict_and_reload_lora_gpu_cache SKIP (https://nvbugs/5682551) full:H20/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=False] SKIP (https://nvbugs/6345827) full:H20/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=True] SKIP (https://nvbugs/6345827) +full:H20/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[llguidance-eagle3_one_model=False] SKIP (https://nvbugs/6422334) +full:H20/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar-eagle3_one_model=False] SKIP (https://nvbugs/6422334) full:H20/accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[nvidia_Llama-3.1-8B-Instruct-FP8-True] SKIP (https://nvbugs/6422351) full:H20/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:H20/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) @@ -439,7 +423,10 @@ test_e2e.py::test_multi_nodes_eval[DeepSeek-R1/DeepSeek-R1-0528-FP4-tp16-mmlu] S test_e2e.py::test_multi_nodes_eval[Kimi-K2-Thinking-NVFP4-tp16-mmlu] SKIP (https://nvbugs/6276983) test_e2e.py::test_multi_nodes_eval[MiniMax-M2-tp16-mmlu] SKIP (https://nvbugs/6373532) test_e2e.py::test_multi_nodes_eval[MiniMax-M3-tp16-mmlu] SKIP (https://nvbugs/6373561) +test_e2e.py::test_openai_chat_example[trt] SKIP (https://nvbugs/5477444) +test_e2e.py::test_openai_completions_example[trt] SKIP (https://nvbugs/5701450) test_e2e.py::test_ptp_quickstart_advanced_deepseek_r1_w4afp8_8gpus[DeepSeek-R1-W4AFP8-DeepSeek-R1/DeepSeek-R1-W4AFP8] SKIP (https://nvbugs/5836830) +test_e2e.py::test_trtllm_bench_iteration_log[TRT-streaming-meta-llama/Llama-3.1-8B-llama-3.1-model/Meta-Llama-3.1-8B] SKIP (https://nvbugs/5448523) unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py::test_prefill_varlen[varlen_hd512_overlap] SKIP (https://nvbugs/6426860) unittest/_torch/misc/test_share_tensor.py::TestShareTensor::test_share_tensor_different_dtypes SKIP (https://nvbugs/6418021) unittest/_torch/modeling -k "modeling_qwen" SKIP (https://nvbugs/6433376) @@ -452,7 +439,6 @@ unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py::test_llm_partial_update_weights_nvfp4[auto-Qwen3/Qwen3-8B] SKIP (https://nvbugs/6372690) unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py::test_llm_partial_update_weights_nvfp4[fp8-Qwen3/Qwen3-30B-A3B] SKIP (https://nvbugs/6372690) unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py::test_llm_partial_update_weights_nvfp4[fp8-Qwen3/Qwen3-8B] SKIP (https://nvbugs/6372690) -unittest/_torch/speculative/test_eagle3.py::test_llama_eagle3[True-TRTLLM-True-False-True-True-True-False-False-False] SKIP (https://nvbugs/6451425) unittest/_torch/thop/parallel/test_fp8_rowwise_linear.py::test_fp8_rowwise_linear[dtype1] SKIP (https://nvbugs/6301807) unittest/_torch/thop/serial/test_moe.py::TestMoeFp4::test_no_autotune[use_score_as_input-RoutingDSv3-swiglu-1024-1024-1] SKIP (https://nvbugs/5908070) unittest/_torch/thop/serial/test_moe.py::TestMoeFp4::test_no_autotune[use_score_as_input-RoutingRenormalize_qwen_next-swiglu-1024-1024-150] SKIP (https://nvbugs/5908070) @@ -484,4 +470,8 @@ unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_tinyllama_logits_processor_t unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[None] SKIP (https://nvbugs/6162504) unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[cuda_graph_config0] SKIP (https://nvbugs/6162504) unittest/llmapi/test_memory_profiling.py::test_profile_kvcache SKIP (https://nvbugs/5580781) +unittest/tools/test_layer_wise_benchmarks.py::test_deepseek_r1_ctx_dep[1] SKIP (https://nvbugs/6162541) +unittest/tools/test_layer_wise_benchmarks.py::test_nemotron_gen_dep[1] SKIP (https://nvbugs/6162541) +unittest/tools/test_layer_wise_benchmarks.py::test_performance_alignment[1] SKIP (https://nvbugs/6127669) +unittest/tools/test_layer_wise_benchmarks.py::test_qwen3_next_gen_tep[1] SKIP (https://nvbugs/6153575) verl/test_verl_cases.py::test_trtllm_abort SKIP (https://nvbugs/6272653) From 66bab0da85ba3fc366320ef01e204bfa2089ca3c Mon Sep 17 00:00:00 2001 From: dongfengy <99041270+dongfengy@users.noreply.github.com> Date: Mon, 13 Jul 2026 22:47:17 +0000 Subject: [PATCH 2/6] [https://nvbugs/6442074][fix] Apply the same spec-dec metadata exception safety to MTP, PARD, DFlash and DraftTarget workers The one-model workers in mtp.py, pard.py, dflash.py and draft_target.py have the same latent bug fixed in eagle3.py: prepare/restore of the attention metadata is not exception-safe, so any tolerated failure in the draft path (e.g. a warmup OOM) poisons attn_metadata._saved_tensors and every later forward dies at the pairing assert. - mtp.py: move change_attn_metadata (which saves state and then mutates more of it) inside the try; restore in finally. - pard.py: prepare and the deferred kv_lens rewind are now covered; the rewind runs in the finally after the restore so a draft failure does not leave kv_lens_cuda incremented for the next request. - dflash.py / draft_target.py: same try/finally treatment. MTPEagleWorker subclasses Eagle3OneModelWorker and is already covered by the eagle3.py fix. Co-Authored-By: Claude Fable 5 Signed-off-by: dongfengy <99041270+dongfengy@users.noreply.github.com> --- tensorrt_llm/_torch/speculative/dflash.py | 192 ++++++++++-------- .../_torch/speculative/draft_target.py | 184 +++++++++-------- tensorrt_llm/_torch/speculative/mtp.py | 132 ++++++------ tensorrt_llm/_torch/speculative/pard.py | 174 +++++++++------- 4 files changed, 373 insertions(+), 309 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index fece18aacc27..d53fb9ca3b49 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -427,97 +427,123 @@ def forward( state_indices=attn_metadata.mamba_metadata.state_indices, ) - self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) - self._prepare_kv_for_draft_forward( - attn_metadata, num_accepted_tokens, num_contexts, batch_size - ) - - # Collapse mrope [3, 1, N] to 1D by taking the first (temporal) dimension. - # The draft model uses standard 1D RoPE, so only scalar positions are needed. - if position_ids.ndim == 3: - position_ids = position_ids[0, 0] - else: - position_ids = position_ids.squeeze(0) - - # Get total tokens processed by target model (for hidden state extraction) - total_target_tokens = input_ids.shape[0] - - # Capture prefill (context) hidden states for future gen steps. - # This gives the draft model the full prompt context, not just gen tokens. - if num_contexts > 0: - self._store_prefill_context( - draft_model, spec_metadata, attn_metadata, position_ids, total_target_tokens + # Save the attn_metadata state that DFlash mutates during drafting. The + # prepare/restore pair must be exception-safe: a failure between them + # (e.g. an OOM during the max-shape general warmup, which model_engine + # tolerates and skips) would otherwise leave stale saved state on the + # persistent attn_metadata and fail every subsequent forward at the + # pairing assert in prepare_for_spec_dec. + # https://nvbugs/6442074 + # kv_prepared gates the deferred kv_lens rewind in the finally block: + # if we fail before _prepare_kv_for_draft_forward runs, applying a + # stale _kv_rewind_amount from a previous step would corrupt + # kv_lens_cuda. + kv_prepared = False + try: + self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) + self._prepare_kv_for_draft_forward( + attn_metadata, num_accepted_tokens, num_contexts, batch_size ) - # Rebuild batch_to_slot after prefill assigns new slots - if self._ctx_buf_inited and spec_metadata.request_ids: - num_seqs = len(spec_metadata.request_ids) - mapping = [self._req_to_slot.get(rid, 0) for rid in spec_metadata.request_ids] - self._batch_to_slot[:num_seqs].copy_( - torch.tensor(mapping, dtype=torch.long, device="cuda") - ) + kv_prepared = True - inputs = self.prepare_1st_drafter_inputs( - input_ids=input_ids, - position_ids=position_ids, - hidden_states=hidden_states, - accepted_tokens=accepted_tokens, - num_accepted_tokens=num_accepted_tokens, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - draft_model=draft_model, - total_target_tokens=total_target_tokens, - ) - - draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + # Collapse mrope [3, 1, N] to 1D by taking the first (temporal) dimension. + # The draft model uses standard 1D RoPE, so only scalar positions are needed. + if position_ids.ndim == 3: + position_ids = position_ids[0, 0] + else: + position_ids = position_ids.squeeze(0) - if num_gens > 0: - with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): - hidden_states_out = draft_model.dflash_forward( - noise_embedding=inputs["noise_embedding"], - query_positions=inputs["query_positions"], - num_ctx_per_req=inputs["num_ctx_per_req"], - ctx_k_cache=inputs["ctx_k_cache"], - ctx_v_cache=inputs["ctx_v_cache"], - ctx_cache_batch_idx=inputs["ctx_cache_batch_idx"], - ) + # Get total tokens processed by target model (for hidden state extraction) + total_target_tokens = input_ids.shape[0] - # Gather K logits per gen request from mask positions (1..K). - # hidden_states_out is flat: [num_gens * block_size, hidden_dim] - block_size = self._resolved_block_size - request_bases = torch.arange(num_gens, dtype=torch.long, device="cuda") * block_size - offsets = torch.arange(K, dtype=torch.long, device="cuda") - # Masks are at positions 1..K in each request's block_size output - gen_gather_ids = (request_bases.unsqueeze(1) + 1 + offsets.unsqueeze(0)).flatten() - gen_gather_ids = gen_gather_ids.clamp(max=hidden_states_out.shape[0] - 1) - - gen_logits = draft_model.logits_processor( - hidden_states_out[gen_gather_ids], draft_model.lm_head, attn_metadata, True + # Capture prefill (context) hidden states for future gen steps. + # This gives the draft model the full prompt context, not just gen tokens. + if num_contexts > 0: + self._store_prefill_context( + draft_model, spec_metadata, attn_metadata, position_ids, total_target_tokens ) + # Rebuild batch_to_slot after prefill assigns new slots + if self._ctx_buf_inited and spec_metadata.request_ids: + num_seqs = len(spec_metadata.request_ids) + mapping = [self._req_to_slot.get(rid, 0) for rid in spec_metadata.request_ids] + self._batch_to_slot[:num_seqs].copy_( + torch.tensor(mapping, dtype=torch.long, device="cuda") + ) + + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + total_target_tokens=total_target_tokens, + ) - vocab_size = gen_logits.shape[-1] - gen_logits = gen_logits.reshape(num_gens, K, vocab_size) - - d2t = getattr(draft_model.model, "d2t", None) - gen_draft_tokens = torch.argmax(gen_logits, dim=-1, keepdim=False).long() - - if d2t is not None: - gen_draft_tokens = d2t[gen_draft_tokens] + gen_draft_tokens - - gen_draft_tokens = gen_draft_tokens.type(torch.int32) - - else: - gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda") + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + + if num_gens > 0: + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + hidden_states_out = draft_model.dflash_forward( + noise_embedding=inputs["noise_embedding"], + query_positions=inputs["query_positions"], + num_ctx_per_req=inputs["num_ctx_per_req"], + ctx_k_cache=inputs["ctx_k_cache"], + ctx_v_cache=inputs["ctx_v_cache"], + ctx_cache_batch_idx=inputs["ctx_cache_batch_idx"], + ) + + # Gather K logits per gen request from mask positions (1..K). + # hidden_states_out is flat: [num_gens * block_size, hidden_dim] + block_size = self._resolved_block_size + request_bases = ( + torch.arange(num_gens, dtype=torch.long, device="cuda") * block_size + ) + offsets = torch.arange(K, dtype=torch.long, device="cuda") + # Masks are at positions 1..K in each request's block_size output + gen_gather_ids = ( + request_bases.unsqueeze(1) + 1 + offsets.unsqueeze(0) + ).flatten() + gen_gather_ids = gen_gather_ids.clamp(max=hidden_states_out.shape[0] - 1) + + gen_logits = draft_model.logits_processor( + hidden_states_out[gen_gather_ids], draft_model.lm_head, attn_metadata, True + ) + + vocab_size = gen_logits.shape[-1] + gen_logits = gen_logits.reshape(num_gens, K, vocab_size) + + d2t = getattr(draft_model.model, "d2t", None) + gen_draft_tokens = torch.argmax(gen_logits, dim=-1, keepdim=False).long() + + if d2t is not None: + gen_draft_tokens = d2t[gen_draft_tokens] + gen_draft_tokens + + gen_draft_tokens = gen_draft_tokens.type(torch.int32) - if num_contexts > 0 and num_gens > 0: - ctx_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") - next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) - elif num_contexts > 0: - next_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") - else: - next_draft_tokens = gen_draft_tokens + else: + gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda") - self._restore_attn_metadata_from_spec_dec(attn_metadata) - self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) + if num_contexts > 0 and num_gens > 0: + ctx_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") + next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) + elif num_contexts > 0: + next_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") + else: + next_draft_tokens = gen_draft_tokens + finally: + # Restore before the deferred rewind: during CUDA-graph warmup + # kv_lens_cuda was saved by prepare_for_spec_dec, so restore + # reinstates the original tensor and _apply_kv_rewind_after_draft + # early-returns. Outside warmup kv_lens_cuda is not saved, so the + # rewind must still run even when drafting failed; otherwise the + # in-place +1 from _prepare_kv_for_draft_forward would persist and + # later forwards would see wrong KV lengths. + self._restore_attn_metadata_from_spec_dec(attn_metadata) + if kv_prepared: + self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) next_new_tokens = self._prepare_next_new_tokens( accepted_tokens, diff --git a/tensorrt_llm/_torch/speculative/draft_target.py b/tensorrt_llm/_torch/speculative/draft_target.py index fe6dc7d1e4c7..12accbe9477d 100644 --- a/tensorrt_llm/_torch/speculative/draft_target.py +++ b/tensorrt_llm/_torch/speculative/draft_target.py @@ -197,102 +197,108 @@ def forward( logits, attn_metadata, spec_metadata ) - # Prepare attention metadata for speculative decoding and save state for restore - self._prepare_attn_metadata_for_draft_target(attn_metadata, spec_metadata) - - # Prepare inputs for the first draft forward - position_ids = position_ids.squeeze(0) - inputs = self.prepare_1st_drafter_inputs( - input_ids=input_ids, - position_ids=position_ids, - accepted_tokens=accepted_tokens, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - ) + # Prepare attention metadata for speculative decoding and save state for restore. + # The save/restore pair must be exception-safe: a failure between them (e.g. an + # OOM during the max-shape general warmup, which model_engine tolerates and + # skips) would otherwise leave stale saved state on the persistent attn_metadata + # and fail every subsequent forward at the pairing assert in + # prepare_for_spec_dec. https://nvbugs/6442074 + try: + self._prepare_attn_metadata_for_draft_target(attn_metadata, spec_metadata) + + # Prepare inputs for the first draft forward + position_ids = position_ids.squeeze(0) + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + ) - next_draft_tokens = [] - original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens - - # Get the draft KV cache manager if using separate layouts - draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - - with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): - for i in range(runtime_draft_len): - if i == 0: - start_ids_gen = ( - spec_metadata.batch_indices_cuda[:num_gens] * (runtime_draft_len + 1) - ).long() - gather_ids_gen = ( - start_ids_gen - + num_accepted_tokens[num_contexts:] - - 1 - + attn_metadata.num_ctx_tokens - ) - gather_ids = torch.concat( - [spec_metadata.gather_ids[:num_contexts], gather_ids_gen], dim=0 - ) - else: - gather_ids = spec_metadata.batch_indices_cuda[:batch_size] + next_draft_tokens = [] + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens - if self.guided_decoder is not None: - new_tokens = inputs["input_ids"][gather_ids] - self.guided_decoder.add_draft_batch( - new_tokens, num_accepted_tokens, draft_step=i - ) + # Get the draft KV cache manager if using separate layouts + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - if original_all_rank_num_tokens is not None: + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + for i in range(runtime_draft_len): if i == 0: - attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens - elif spec_metadata.all_rank_num_seqs is not None: - attn_metadata.all_rank_num_tokens = spec_metadata.all_rank_num_seqs - - hidden_states = draft_model.model(**inputs) - if isinstance(hidden_states, tuple): - hidden_states = hidden_states[0] - - # Disable spec-dec mode for chained draft steps - attn_metadata.use_spec_decoding = False - - logits = draft_model.logits_processor( - hidden_states[gather_ids], draft_model.lm_head, attn_metadata, True - ) - if self.guided_decoder is not None: - d2t = getattr(draft_model.model, "d2t", None) - self.guided_decoder.execute_draft_batch(logits, d2t, draft_step=i) - - new_draft_token = self.draft_decoder(logits, draft_model) - next_draft_tokens.append(new_draft_token) - - # Update inputs and metadata for next draft step - position_ids = inputs["position_ids"][gather_ids] + 1 - if i == 0: - attn_metadata._seq_lens[:batch_size].fill_(1) - attn_metadata._seq_lens_cuda[:batch_size].fill_(1) - attn_metadata.on_update() - if inputs["attn_metadata"].kv_cache_manager is not None: - attn_metadata.host_request_types[: attn_metadata.num_contexts].fill_(1) - attn_metadata.num_contexts = 0 - self._update_kv_after_first_draft_step( - attn_metadata, - num_accepted_tokens, - num_contexts, - batch_size, - runtime_draft_len, + start_ids_gen = ( + spec_metadata.batch_indices_cuda[:num_gens] * (runtime_draft_len + 1) + ).long() + gather_ids_gen = ( + start_ids_gen + + num_accepted_tokens[num_contexts:] + - 1 + + attn_metadata.num_ctx_tokens + ) + gather_ids = torch.concat( + [spec_metadata.gather_ids[:num_contexts], gather_ids_gen], dim=0 + ) + else: + gather_ids = spec_metadata.batch_indices_cuda[:batch_size] + + if self.guided_decoder is not None: + new_tokens = inputs["input_ids"][gather_ids] + self.guided_decoder.add_draft_batch( + new_tokens, num_accepted_tokens, draft_step=i + ) + + if original_all_rank_num_tokens is not None: + if i == 0: + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + elif spec_metadata.all_rank_num_seqs is not None: + attn_metadata.all_rank_num_tokens = spec_metadata.all_rank_num_seqs + + hidden_states = draft_model.model(**inputs) + if isinstance(hidden_states, tuple): + hidden_states = hidden_states[0] + + # Disable spec-dec mode for chained draft steps + attn_metadata.use_spec_decoding = False + + logits = draft_model.logits_processor( + hidden_states[gather_ids], draft_model.lm_head, attn_metadata, True ) - else: - self._update_kv_for_chained_draft_step(attn_metadata, batch_size) - - inputs = { - "input_ids": new_draft_token, - "position_ids": position_ids, - "attn_metadata": attn_metadata, - "spec_metadata": spec_metadata, - } + if self.guided_decoder is not None: + d2t = getattr(draft_model.model, "d2t", None) + self.guided_decoder.execute_draft_batch(logits, d2t, draft_step=i) - next_draft_tokens = torch.stack(next_draft_tokens, dim=1) + new_draft_token = self.draft_decoder(logits, draft_model) + next_draft_tokens.append(new_draft_token) - # Restore attention metadata to original state - self._restore_attn_metadata_from_spec_dec(attn_metadata) + # Update inputs and metadata for next draft step + position_ids = inputs["position_ids"][gather_ids] + 1 + if i == 0: + attn_metadata._seq_lens[:batch_size].fill_(1) + attn_metadata._seq_lens_cuda[:batch_size].fill_(1) + attn_metadata.on_update() + if inputs["attn_metadata"].kv_cache_manager is not None: + attn_metadata.host_request_types[: attn_metadata.num_contexts].fill_(1) + attn_metadata.num_contexts = 0 + self._update_kv_after_first_draft_step( + attn_metadata, + num_accepted_tokens, + num_contexts, + batch_size, + runtime_draft_len, + ) + else: + self._update_kv_for_chained_draft_step(attn_metadata, batch_size) + + inputs = { + "input_ids": new_draft_token, + "position_ids": position_ids, + "attn_metadata": attn_metadata, + "spec_metadata": spec_metadata, + } + + next_draft_tokens = torch.stack(next_draft_tokens, dim=1) + finally: + # Restore attention metadata to original state + 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 diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 6afed0fc02e9..5f8dcdb9070b 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -437,68 +437,80 @@ def forward( spec_metadata=spec_metadata, attn_metadata=attn_metadata) - # update attn metadata - if attn_metadata is not None: - self.change_attn_metadata(num_accepted_tokens, attn_metadata, - spec_metadata) - - # Run MTP layers to predict draft tokens - next_draft_tokens = [] - last_tokens_idx = torch.cumsum( - attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 - - draft_kv_cache_manager = self.get_draft_kv_cache_manager( - resource_manager) - - with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): - for i, mtp_layer in enumerate( - draft_model.mtp_layers[:runtime_draft_len]): - if self.guided_decoder is not None: - new_tokens = draft_inputs['input_ids'][last_tokens_idx] - self.guided_decoder.add_draft_batch(new_tokens, - num_accepted_tokens, - draft_step=i) - - hidden_states = mtp_layer(embed_tokens=draft_model.embed_tokens, - **draft_inputs) - - logits = mtp_layer.shared_head(hidden_states, - draft_model.lm_head, - attn_metadata).float() - if self.guided_decoder is not None: - self.guided_decoder.execute_draft_batch(logits, + # update attn metadata. change_attn_metadata saves the old attn + # metadata state and then mutates more of it, so the whole call must + # be inside the protected region. The save/restore pair must be + # exception-safe: a failure between them (e.g. an OOM during the + # max-shape general warmup, which model_engine tolerates and skips) + # would otherwise leave stale saved state on the persistent + # attn_metadata and fail every subsequent forward at the pairing + # assert in prepare_for_spec_dec. + # https://nvbugs/6442074 + try: + if attn_metadata is not None: + self.change_attn_metadata(num_accepted_tokens, attn_metadata, + spec_metadata) + + # Run MTP layers to predict draft tokens + next_draft_tokens = [] + last_tokens_idx = torch.cumsum( + attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 + + draft_kv_cache_manager = self.get_draft_kv_cache_manager( + resource_manager) + + with self.draft_kv_cache_context(attn_metadata, + draft_kv_cache_manager): + for i, mtp_layer in enumerate( + draft_model.mtp_layers[:runtime_draft_len]): + if self.guided_decoder is not None: + new_tokens = draft_inputs['input_ids'][last_tokens_idx] + self.guided_decoder.add_draft_batch(new_tokens, + num_accepted_tokens, draft_step=i) - new_draft_token = self.draft_sampler(logits) - next_draft_tokens.append(new_draft_token) - # shift input_ids and hidden_states - input_ids = draft_inputs["input_ids"] - input_ids[:-1] = input_ids[1:].clone() - input_ids[last_tokens_idx] = new_draft_token - draft_hidden_states = draft_inputs["hidden_states"] - draft_hidden_states[:-1] = draft_hidden_states[1:].clone() - draft_hidden_states[last_tokens_idx] = hidden_states[ - last_tokens_idx, :] - draft_inputs = { - "input_ids": input_ids, - "position_ids": draft_inputs["position_ids"], - "hidden_states": draft_hidden_states, - "attn_metadata": draft_inputs["attn_metadata"], - } - next_draft_tokens = torch.stack(next_draft_tokens, dim=1) - - # Override with SA draft tokens after all MTP layers have run, - # so that MTP layers never see SA tokens in their inputs. - if self.sa_enhancer is not None: - num_contexts = attn_metadata.num_contexts - gen_draft_tokens = next_draft_tokens[num_contexts:] - gen_draft_tokens = self.sa_enhancer.maybe_override_all_draft_tokens( - gen_draft_tokens) - next_draft_tokens[num_contexts:] = gen_draft_tokens - - # restore attn metadata - if attn_metadata is not None: - self._restore_attn_metadata_from_spec_dec(attn_metadata) + hidden_states = mtp_layer( + embed_tokens=draft_model.embed_tokens, **draft_inputs) + + logits = mtp_layer.shared_head(hidden_states, + draft_model.lm_head, + attn_metadata).float() + if self.guided_decoder is not None: + self.guided_decoder.execute_draft_batch(logits, + draft_step=i) + + new_draft_token = self.draft_sampler(logits) + next_draft_tokens.append(new_draft_token) + # shift input_ids and hidden_states + input_ids = draft_inputs["input_ids"] + input_ids[:-1] = input_ids[1:].clone() + input_ids[last_tokens_idx] = new_draft_token + draft_hidden_states = draft_inputs["hidden_states"] + draft_hidden_states[:-1] = draft_hidden_states[1:].clone() + draft_hidden_states[last_tokens_idx] = hidden_states[ + last_tokens_idx, :] + draft_inputs = { + "input_ids": input_ids, + "position_ids": draft_inputs["position_ids"], + "hidden_states": draft_hidden_states, + "attn_metadata": draft_inputs["attn_metadata"], + } + next_draft_tokens = torch.stack(next_draft_tokens, dim=1) + + # Override with SA draft tokens after all MTP layers have run, + # so that MTP layers never see SA tokens in their inputs. + if self.sa_enhancer is not None: + num_contexts = attn_metadata.num_contexts + gen_draft_tokens = next_draft_tokens[num_contexts:] + gen_draft_tokens = self.sa_enhancer.maybe_override_all_draft_tokens( + gen_draft_tokens) + next_draft_tokens[num_contexts:] = gen_draft_tokens + finally: + # restore attn metadata. restore_from_spec_dec is safe on empty + # or partially saved state, so this is correct even if + # change_attn_metadata itself failed mid-save. + if attn_metadata is not None: + self._restore_attn_metadata_from_spec_dec(attn_metadata) # prepare next new tokens to support overlap scheduler next_new_tokens = self._prepare_next_new_tokens( diff --git a/tensorrt_llm/_torch/speculative/pard.py b/tensorrt_llm/_torch/speculative/pard.py index 34a4509314cd..2622d155f117 100644 --- a/tensorrt_llm/_torch/speculative/pard.py +++ b/tensorrt_llm/_torch/speculative/pard.py @@ -236,100 +236,120 @@ def forward( max_draft_len=K, ) - self._prepare_attn_metadata_for_pard(attn_metadata, spec_metadata) - self._prepare_kv_for_draft_forward( - attn_metadata, num_accepted_tokens, num_contexts, batch_size - ) - - position_ids = position_ids.squeeze(0) - inputs = self.prepare_1st_drafter_inputs( - input_ids=input_ids, - position_ids=position_ids, - hidden_states=hidden_states, - accepted_tokens=accepted_tokens, - num_accepted_tokens=num_accepted_tokens, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - draft_model=draft_model, - ) - - draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + # Save attn_metadata state and adjust kv_lens for the draft forward. + # The save/restore pair must be exception-safe: a failure between them + # (e.g. an OOM during the max-shape general warmup, which model_engine + # tolerates and skips) would otherwise leave stale saved state on the + # persistent attn_metadata and fail every subsequent forward at the + # pairing assert in prepare_for_spec_dec. + # https://nvbugs/6442074 + # + # Drop any rewind state left over from a previous call first, so a + # failure before _prepare_kv_for_draft_forward runs cannot make the + # finally block below re-apply a stale rewind. + if hasattr(self, "_kv_rewind_amount"): + del self._kv_rewind_amount + try: + self._prepare_attn_metadata_for_pard(attn_metadata, spec_metadata) + self._prepare_kv_for_draft_forward( + attn_metadata, num_accepted_tokens, num_contexts, batch_size + ) - if num_gens > 0: - with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): - hidden_states_out = draft_model.model(**inputs) + position_ids = position_ids.squeeze(0) + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) - # Gather K logits per gen request starting at the bonus position. - # Layout: [acc_0..acc_M, masks] (2K total). Positions M..M+K-1 - # produce K unique draft predictions (bonus + K-1 masks). - gen_start_idx = attn_metadata.num_ctx_tokens + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - request_bases = ( - torch.arange(num_gens, dtype=torch.long, device="cuda") * (2 * K) - + gen_start_idx - ) + if num_gens > 0: + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + hidden_states_out = draft_model.model(**inputs) - gen_num_accepted = num_accepted_tokens[num_contexts:batch_size].long() - base_offsets = gen_num_accepted - 1 # M = bonus position - offsets = torch.arange(K, dtype=torch.long, device="cuda") + # Gather K logits per gen request starting at the bonus position. + # Layout: [acc_0..acc_M, masks] (2K total). Positions M..M+K-1 + # produce K unique draft predictions (bonus + K-1 masks). + gen_start_idx = attn_metadata.num_ctx_tokens - gen_gather_ids = ( - request_bases.unsqueeze(1) + base_offsets.unsqueeze(1) + offsets.unsqueeze(0) - ).flatten() - gen_gather_ids = gen_gather_ids.clamp(max=hidden_states_out.shape[0] - 1) + request_bases = ( + torch.arange(num_gens, dtype=torch.long, device="cuda") * (2 * K) + + gen_start_idx + ) - gen_logits = draft_model.logits_processor( - hidden_states_out[gen_gather_ids], draft_model.lm_head, attn_metadata, True - ) + gen_num_accepted = num_accepted_tokens[num_contexts:batch_size].long() + base_offsets = gen_num_accepted - 1 # M = bonus position + offsets = torch.arange(K, dtype=torch.long, device="cuda") - vocab_size = gen_logits.shape[-1] - gen_logits = gen_logits.reshape(num_gens, K, vocab_size) + gen_gather_ids = ( + request_bases.unsqueeze(1) + + base_offsets.unsqueeze(1) + + offsets.unsqueeze(0) + ).flatten() + gen_gather_ids = gen_gather_ids.clamp(max=hidden_states_out.shape[0] - 1) - # Use torch.argmax directly to avoid cute_argmax stride issues - d2t = getattr(draft_model.model, "d2t", None) - gen_draft_tokens = torch.argmax(gen_logits, dim=-1, keepdim=False).long() + gen_logits = draft_model.logits_processor( + hidden_states_out[gen_gather_ids], draft_model.lm_head, attn_metadata, True + ) - if d2t is not None: - gen_draft_tokens = d2t[gen_draft_tokens] + gen_draft_tokens + vocab_size = gen_logits.shape[-1] + gen_logits = gen_logits.reshape(num_gens, K, vocab_size) - gen_draft_tokens = gen_draft_tokens.type(torch.int32) + # Use torch.argmax directly to avoid cute_argmax stride issues + d2t = getattr(draft_model.model, "d2t", None) + gen_draft_tokens = torch.argmax(gen_logits, dim=-1, keepdim=False).long() - if self.sa_enhancer is not None and sa_manager is not None: - gen_draft_tokens = self.sa_enhancer.maybe_override_all_draft_tokens( - gen_draft_tokens - ) + if d2t is not None: + gen_draft_tokens = d2t[gen_draft_tokens] + gen_draft_tokens - # Pad from (num_gens, K) to (num_gens, 2K-1). - if K > 1: - pad = torch.zeros((num_gens, K - 1), dtype=torch.int32, device="cuda") - gen_draft_tokens = torch.cat([gen_draft_tokens, pad], dim=1) + gen_draft_tokens = gen_draft_tokens.type(torch.int32) - elif num_contexts > 0 and self.use_separate_draft_kv_cache: - # Pure context batch: populate the draft KV cache so it's - # ready when generation starts. - with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): - draft_model.model(**inputs) - gen_draft_tokens = torch.empty((0, 2 * K - 1), dtype=torch.int32, device="cuda") + if self.sa_enhancer is not None and sa_manager is not None: + gen_draft_tokens = self.sa_enhancer.maybe_override_all_draft_tokens( + gen_draft_tokens + ) - else: - gen_draft_tokens = torch.empty((0, 2 * K - 1), dtype=torch.int32, device="cuda") + # Pad from (num_gens, K) to (num_gens, 2K-1). + if K > 1: + pad = torch.zeros((num_gens, K - 1), dtype=torch.int32, device="cuda") + gen_draft_tokens = torch.cat([gen_draft_tokens, pad], dim=1) - if num_contexts > 0 and num_gens > 0: - ctx_draft_tokens = torch.zeros( - (num_contexts, 2 * K - 1), dtype=torch.int32, device="cuda" - ) - next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) - elif num_contexts > 0: - next_draft_tokens = torch.zeros( - (num_contexts, 2 * K - 1), dtype=torch.int32, device="cuda" - ) - else: - next_draft_tokens = gen_draft_tokens + elif num_contexts > 0 and self.use_separate_draft_kv_cache: + # Pure context batch: populate the draft KV cache so it's + # ready when generation starts. + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + draft_model.model(**inputs) + gen_draft_tokens = torch.empty((0, 2 * K - 1), dtype=torch.int32, device="cuda") - self._restore_attn_metadata_from_spec_dec(attn_metadata) + else: + gen_draft_tokens = torch.empty((0, 2 * K - 1), dtype=torch.int32, device="cuda") - # Deferred kv_lens rewind (must happen after restore so it persists). - self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) + if num_contexts > 0 and num_gens > 0: + ctx_draft_tokens = torch.zeros( + (num_contexts, 2 * K - 1), dtype=torch.int32, device="cuda" + ) + next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) + elif num_contexts > 0: + next_draft_tokens = torch.zeros( + (num_contexts, 2 * K - 1), dtype=torch.int32, device="cuda" + ) + else: + next_draft_tokens = gen_draft_tokens + finally: + self._restore_attn_metadata_from_spec_dec(attn_metadata) + + # Deferred kv_lens rewind (must happen after restore so it persists). + # Kept inside the finally, ordered after the restore: if the draft + # forward fails after _prepare_kv_for_draft_forward incremented + # kv_lens_cuda, the rewind must still be applied or the next + # iteration sees wrong KV lengths. + self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) next_new_tokens = self._prepare_next_new_tokens( accepted_tokens, From 78592234e19849161468a5223439e366c990c2a2 Mon Sep 17 00:00:00 2001 From: dongfengy <99041270+dongfengy@users.noreply.github.com> Date: Tue, 14 Jul 2026 00:20:13 +0000 Subject: [PATCH 3/6] [https://nvbugs/6442074][refactor] Centralize spec-dec metadata cleanup in SpecWorkerBase Replace the per-worker try/finally blocks with a template method on SpecWorkerBase: forward() extracts attn/spec metadata, runs the worker-specific _forward_impl(), and in a finally block restores any unrestored prepare_for_spec_dec state (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. - speculative/interface.py: template forward + _ensure_spec_dec_state_restored. - attention_backend/interface.py: has_spec_dec_saved_state property. - eagle3/mtp/pard/dflash/draft_target/sa_worker/eagle3_dynamic_tree: forward renamed to _forward_impl; bodies return to straight-line prepare/restore (per-worker try/finally removed). - pard/dflash: _kv_rewind_pending flag so a failed draft forward still rewinds kv_lens_cuda exactly once; consumed by _apply_kv_rewind_after_draft. - eagle3: initialize the worker-side _saved_* attrs so a failure inside the prepare override cannot turn the cleanup restore into AttributeError. - test_force_accepted_tokens.py: stub worker implements _forward_impl. Verified on 4x H100-80GB (viking-prod-233) at main f81b9d2787, A/B: - workers unguarded, no template method: FAILED in 215s with the nvbugs/6442074 signature (warmup OOM then stale-save assert). - template method only: the base finally visibly restores state when the max-shape warmup OOMs, and the test PASSES - GSM8K 90.75 (ref 90.30), GPQADiamond 66.16 (ref 65.0). Co-Authored-By: Claude Fable 5 Signed-off-by: dongfengy <99041270+dongfengy@users.noreply.github.com> Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../_torch/attention_backend/interface.py | 5 + tensorrt_llm/_torch/speculative/dflash.py | 210 +++++++++--------- .../_torch/speculative/draft_target.py | 186 ++++++++-------- tensorrt_llm/_torch/speculative/eagle3.py | 93 ++++---- .../_torch/speculative/eagle3_dynamic_tree.py | 6 +- tensorrt_llm/_torch/speculative/interface.py | 49 ++++ tensorrt_llm/_torch/speculative/mtp.py | 134 +++++------ tensorrt_llm/_torch/speculative/pard.py | 192 ++++++++-------- tensorrt_llm/_torch/speculative/sa_worker.py | 2 +- tests/integration/test_lists/waives.txt | 33 ++- .../speculative/test_force_accepted_tokens.py | 5 +- 11 files changed, 474 insertions(+), 441 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 4ad8cb273666..02f77bbdf582 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -427,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 d53fb9ca3b49..31cdb67e3f5c 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"): @@ -366,7 +370,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, @@ -427,123 +443,97 @@ def forward( state_indices=attn_metadata.mamba_metadata.state_indices, ) - # Save the attn_metadata state that DFlash mutates during drafting. The - # prepare/restore pair must be exception-safe: a failure between them - # (e.g. an OOM during the max-shape general warmup, which model_engine - # tolerates and skips) would otherwise leave stale saved state on the - # persistent attn_metadata and fail every subsequent forward at the - # pairing assert in prepare_for_spec_dec. - # https://nvbugs/6442074 - # kv_prepared gates the deferred kv_lens rewind in the finally block: - # if we fail before _prepare_kv_for_draft_forward runs, applying a - # stale _kv_rewind_amount from a previous step would corrupt - # kv_lens_cuda. - kv_prepared = False - try: - self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) - self._prepare_kv_for_draft_forward( - attn_metadata, num_accepted_tokens, num_contexts, batch_size + self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) + self._prepare_kv_for_draft_forward( + attn_metadata, num_accepted_tokens, num_contexts, batch_size + ) + + # Collapse mrope [3, 1, N] to 1D by taking the first (temporal) dimension. + # The draft model uses standard 1D RoPE, so only scalar positions are needed. + if position_ids.ndim == 3: + position_ids = position_ids[0, 0] + else: + position_ids = position_ids.squeeze(0) + + # Get total tokens processed by target model (for hidden state extraction) + total_target_tokens = input_ids.shape[0] + + # Capture prefill (context) hidden states for future gen steps. + # This gives the draft model the full prompt context, not just gen tokens. + if num_contexts > 0: + self._store_prefill_context( + draft_model, spec_metadata, attn_metadata, position_ids, total_target_tokens ) - kv_prepared = True + # Rebuild batch_to_slot after prefill assigns new slots + if self._ctx_buf_inited and spec_metadata.request_ids: + num_seqs = len(spec_metadata.request_ids) + mapping = [self._req_to_slot.get(rid, 0) for rid in spec_metadata.request_ids] + self._batch_to_slot[:num_seqs].copy_( + torch.tensor(mapping, dtype=torch.long, device="cuda") + ) - # Collapse mrope [3, 1, N] to 1D by taking the first (temporal) dimension. - # The draft model uses standard 1D RoPE, so only scalar positions are needed. - if position_ids.ndim == 3: - position_ids = position_ids[0, 0] - else: - position_ids = position_ids.squeeze(0) + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + total_target_tokens=total_target_tokens, + ) - # Get total tokens processed by target model (for hidden state extraction) - total_target_tokens = input_ids.shape[0] + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - # Capture prefill (context) hidden states for future gen steps. - # This gives the draft model the full prompt context, not just gen tokens. - if num_contexts > 0: - self._store_prefill_context( - draft_model, spec_metadata, attn_metadata, position_ids, total_target_tokens + if num_gens > 0: + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + hidden_states_out = draft_model.dflash_forward( + noise_embedding=inputs["noise_embedding"], + query_positions=inputs["query_positions"], + num_ctx_per_req=inputs["num_ctx_per_req"], + ctx_k_cache=inputs["ctx_k_cache"], + ctx_v_cache=inputs["ctx_v_cache"], + ctx_cache_batch_idx=inputs["ctx_cache_batch_idx"], ) - # Rebuild batch_to_slot after prefill assigns new slots - if self._ctx_buf_inited and spec_metadata.request_ids: - num_seqs = len(spec_metadata.request_ids) - mapping = [self._req_to_slot.get(rid, 0) for rid in spec_metadata.request_ids] - self._batch_to_slot[:num_seqs].copy_( - torch.tensor(mapping, dtype=torch.long, device="cuda") - ) - - inputs = self.prepare_1st_drafter_inputs( - input_ids=input_ids, - position_ids=position_ids, - hidden_states=hidden_states, - accepted_tokens=accepted_tokens, - num_accepted_tokens=num_accepted_tokens, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - draft_model=draft_model, - total_target_tokens=total_target_tokens, - ) - draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - - if num_gens > 0: - with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): - hidden_states_out = draft_model.dflash_forward( - noise_embedding=inputs["noise_embedding"], - query_positions=inputs["query_positions"], - num_ctx_per_req=inputs["num_ctx_per_req"], - ctx_k_cache=inputs["ctx_k_cache"], - ctx_v_cache=inputs["ctx_v_cache"], - ctx_cache_batch_idx=inputs["ctx_cache_batch_idx"], - ) - - # Gather K logits per gen request from mask positions (1..K). - # hidden_states_out is flat: [num_gens * block_size, hidden_dim] - block_size = self._resolved_block_size - request_bases = ( - torch.arange(num_gens, dtype=torch.long, device="cuda") * block_size - ) - offsets = torch.arange(K, dtype=torch.long, device="cuda") - # Masks are at positions 1..K in each request's block_size output - gen_gather_ids = ( - request_bases.unsqueeze(1) + 1 + offsets.unsqueeze(0) - ).flatten() - gen_gather_ids = gen_gather_ids.clamp(max=hidden_states_out.shape[0] - 1) - - gen_logits = draft_model.logits_processor( - hidden_states_out[gen_gather_ids], draft_model.lm_head, attn_metadata, True - ) - - vocab_size = gen_logits.shape[-1] - gen_logits = gen_logits.reshape(num_gens, K, vocab_size) - - d2t = getattr(draft_model.model, "d2t", None) - gen_draft_tokens = torch.argmax(gen_logits, dim=-1, keepdim=False).long() - - if d2t is not None: - gen_draft_tokens = d2t[gen_draft_tokens] + gen_draft_tokens - - gen_draft_tokens = gen_draft_tokens.type(torch.int32) + # Gather K logits per gen request from mask positions (1..K). + # hidden_states_out is flat: [num_gens * block_size, hidden_dim] + block_size = self._resolved_block_size + request_bases = torch.arange(num_gens, dtype=torch.long, device="cuda") * block_size + offsets = torch.arange(K, dtype=torch.long, device="cuda") + # Masks are at positions 1..K in each request's block_size output + gen_gather_ids = (request_bases.unsqueeze(1) + 1 + offsets.unsqueeze(0)).flatten() + gen_gather_ids = gen_gather_ids.clamp(max=hidden_states_out.shape[0] - 1) + + gen_logits = draft_model.logits_processor( + hidden_states_out[gen_gather_ids], draft_model.lm_head, attn_metadata, True + ) - else: - gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda") + vocab_size = gen_logits.shape[-1] + gen_logits = gen_logits.reshape(num_gens, K, vocab_size) - if num_contexts > 0 and num_gens > 0: - ctx_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") - next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) - elif num_contexts > 0: - next_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") - else: - next_draft_tokens = gen_draft_tokens - finally: - # Restore before the deferred rewind: during CUDA-graph warmup - # kv_lens_cuda was saved by prepare_for_spec_dec, so restore - # reinstates the original tensor and _apply_kv_rewind_after_draft - # early-returns. Outside warmup kv_lens_cuda is not saved, so the - # rewind must still run even when drafting failed; otherwise the - # in-place +1 from _prepare_kv_for_draft_forward would persist and - # later forwards would see wrong KV lengths. - self._restore_attn_metadata_from_spec_dec(attn_metadata) - if kv_prepared: - self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) + d2t = getattr(draft_model.model, "d2t", None) + gen_draft_tokens = torch.argmax(gen_logits, dim=-1, keepdim=False).long() + + if d2t is not None: + gen_draft_tokens = d2t[gen_draft_tokens] + gen_draft_tokens + + gen_draft_tokens = gen_draft_tokens.type(torch.int32) + + else: + gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda") + + if num_contexts > 0 and num_gens > 0: + ctx_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") + next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) + elif num_contexts > 0: + next_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") + else: + next_draft_tokens = gen_draft_tokens + + self._restore_attn_metadata_from_spec_dec(attn_metadata) + self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) next_new_tokens = self._prepare_next_new_tokens( accepted_tokens, diff --git a/tensorrt_llm/_torch/speculative/draft_target.py b/tensorrt_llm/_torch/speculative/draft_target.py index 12accbe9477d..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, @@ -197,108 +197,102 @@ def forward( logits, attn_metadata, spec_metadata ) - # Prepare attention metadata for speculative decoding and save state for restore. - # The save/restore pair must be exception-safe: a failure between them (e.g. an - # OOM during the max-shape general warmup, which model_engine tolerates and - # skips) would otherwise leave stale saved state on the persistent attn_metadata - # and fail every subsequent forward at the pairing assert in - # prepare_for_spec_dec. https://nvbugs/6442074 - try: - self._prepare_attn_metadata_for_draft_target(attn_metadata, spec_metadata) - - # Prepare inputs for the first draft forward - position_ids = position_ids.squeeze(0) - inputs = self.prepare_1st_drafter_inputs( - input_ids=input_ids, - position_ids=position_ids, - accepted_tokens=accepted_tokens, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - ) + # Prepare attention metadata for speculative decoding and save state for restore + self._prepare_attn_metadata_for_draft_target(attn_metadata, spec_metadata) + + # Prepare inputs for the first draft forward + position_ids = position_ids.squeeze(0) + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + ) - next_draft_tokens = [] - original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + next_draft_tokens = [] + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + + # Get the draft KV cache manager if using separate layouts + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + for i in range(runtime_draft_len): + if i == 0: + start_ids_gen = ( + spec_metadata.batch_indices_cuda[:num_gens] * (runtime_draft_len + 1) + ).long() + gather_ids_gen = ( + start_ids_gen + + num_accepted_tokens[num_contexts:] + - 1 + + attn_metadata.num_ctx_tokens + ) + gather_ids = torch.concat( + [spec_metadata.gather_ids[:num_contexts], gather_ids_gen], dim=0 + ) + else: + gather_ids = spec_metadata.batch_indices_cuda[:batch_size] - # Get the draft KV cache manager if using separate layouts - draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + if self.guided_decoder is not None: + new_tokens = inputs["input_ids"][gather_ids] + self.guided_decoder.add_draft_batch( + new_tokens, num_accepted_tokens, draft_step=i + ) - with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): - for i in range(runtime_draft_len): + if original_all_rank_num_tokens is not None: if i == 0: - start_ids_gen = ( - spec_metadata.batch_indices_cuda[:num_gens] * (runtime_draft_len + 1) - ).long() - gather_ids_gen = ( - start_ids_gen - + num_accepted_tokens[num_contexts:] - - 1 - + attn_metadata.num_ctx_tokens - ) - gather_ids = torch.concat( - [spec_metadata.gather_ids[:num_contexts], gather_ids_gen], dim=0 - ) - else: - gather_ids = spec_metadata.batch_indices_cuda[:batch_size] - - if self.guided_decoder is not None: - new_tokens = inputs["input_ids"][gather_ids] - self.guided_decoder.add_draft_batch( - new_tokens, num_accepted_tokens, draft_step=i - ) - - if original_all_rank_num_tokens is not None: - if i == 0: - attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens - elif spec_metadata.all_rank_num_seqs is not None: - attn_metadata.all_rank_num_tokens = spec_metadata.all_rank_num_seqs - - hidden_states = draft_model.model(**inputs) - if isinstance(hidden_states, tuple): - hidden_states = hidden_states[0] - - # Disable spec-dec mode for chained draft steps - attn_metadata.use_spec_decoding = False - - logits = draft_model.logits_processor( - hidden_states[gather_ids], draft_model.lm_head, attn_metadata, True + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + elif spec_metadata.all_rank_num_seqs is not None: + attn_metadata.all_rank_num_tokens = spec_metadata.all_rank_num_seqs + + hidden_states = draft_model.model(**inputs) + if isinstance(hidden_states, tuple): + hidden_states = hidden_states[0] + + # Disable spec-dec mode for chained draft steps + attn_metadata.use_spec_decoding = False + + logits = draft_model.logits_processor( + hidden_states[gather_ids], draft_model.lm_head, attn_metadata, True + ) + if self.guided_decoder is not None: + d2t = getattr(draft_model.model, "d2t", None) + self.guided_decoder.execute_draft_batch(logits, d2t, draft_step=i) + + new_draft_token = self.draft_decoder(logits, draft_model) + next_draft_tokens.append(new_draft_token) + + # Update inputs and metadata for next draft step + position_ids = inputs["position_ids"][gather_ids] + 1 + if i == 0: + attn_metadata._seq_lens[:batch_size].fill_(1) + attn_metadata._seq_lens_cuda[:batch_size].fill_(1) + attn_metadata.on_update() + if inputs["attn_metadata"].kv_cache_manager is not None: + attn_metadata.host_request_types[: attn_metadata.num_contexts].fill_(1) + attn_metadata.num_contexts = 0 + self._update_kv_after_first_draft_step( + attn_metadata, + num_accepted_tokens, + num_contexts, + batch_size, + runtime_draft_len, ) - if self.guided_decoder is not None: - d2t = getattr(draft_model.model, "d2t", None) - self.guided_decoder.execute_draft_batch(logits, d2t, draft_step=i) + else: + self._update_kv_for_chained_draft_step(attn_metadata, batch_size) - new_draft_token = self.draft_decoder(logits, draft_model) - next_draft_tokens.append(new_draft_token) + inputs = { + "input_ids": new_draft_token, + "position_ids": position_ids, + "attn_metadata": attn_metadata, + "spec_metadata": spec_metadata, + } - # Update inputs and metadata for next draft step - position_ids = inputs["position_ids"][gather_ids] + 1 - if i == 0: - attn_metadata._seq_lens[:batch_size].fill_(1) - attn_metadata._seq_lens_cuda[:batch_size].fill_(1) - attn_metadata.on_update() - if inputs["attn_metadata"].kv_cache_manager is not None: - attn_metadata.host_request_types[: attn_metadata.num_contexts].fill_(1) - attn_metadata.num_contexts = 0 - self._update_kv_after_first_draft_step( - attn_metadata, - num_accepted_tokens, - num_contexts, - batch_size, - runtime_draft_len, - ) - else: - self._update_kv_for_chained_draft_step(attn_metadata, batch_size) - - inputs = { - "input_ids": new_draft_token, - "position_ids": position_ids, - "attn_metadata": attn_metadata, - "spec_metadata": spec_metadata, - } - - next_draft_tokens = torch.stack(next_draft_tokens, dim=1) - finally: - # Restore attention metadata to original state - self._restore_attn_metadata_from_spec_dec(attn_metadata) + next_draft_tokens = torch.stack(next_draft_tokens, dim=1) + + # Restore attention metadata to original state + 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 diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 67f5a6866d9a..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 @@ -737,44 +745,35 @@ def forward(self, max_draft_len=runtime_draft_len, ) - # Save the old attn_metadata and spec_metadata. The save/restore - # pair must be exception-safe: a failure between them (e.g. an OOM - # during the max-shape general warmup, which model_engine tolerates - # and skips) would otherwise leave stale saved state on the - # persistent attn_metadata and fail every subsequent forward at - # the pairing assert in prepare_for_spec_dec. - # https://nvbugs/6442074 - try: - self._prepare_attn_metadata_for_spec_dec(attn_metadata) - - # Prepare inputs for the 1st draft model forward - position_ids = position_ids.squeeze(0) - inputs = self.prepare_1st_drafter_inputs( - input_ids=input_ids, - position_ids=position_ids, - hidden_states=hidden_states, - accepted_tokens=accepted_tokens, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - draft_model=draft_model) - - # Predict draft tokens. ``original_all_rank_num_tokens`` is saved - # here so the post-loop restore (below) can put attn_metadata - # back into a state the target model expects. - original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens - - # Get the draft KV cache manager if using separate layouts - draft_kv_cache_manager = self.get_draft_kv_cache_manager( - resource_manager) - - next_draft_tokens = self._forward_draft_loop( - inputs, attn_metadata, spec_metadata, draft_model, - draft_kv_cache_manager, num_contexts, num_gens, batch_size, - num_accepted_tokens, original_all_rank_num_tokens, - resource_manager) - finally: - # restore attn_metadata to support cuda graph - self._restore_attn_metadata_from_spec_dec(attn_metadata) + # Save the old attn_metadata and spec_metadata + self._prepare_attn_metadata_for_spec_dec(attn_metadata) + + # Prepare inputs for the 1st draft model forward + position_ids = position_ids.squeeze(0) + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model) + + # Predict draft tokens. ``original_all_rank_num_tokens`` is saved here + # so the post-loop restore (below) can put attn_metadata back into a + # state the target model expects. + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + + # Get the draft KV cache manager if using separate layouts + draft_kv_cache_manager = self.get_draft_kv_cache_manager( + resource_manager) + + next_draft_tokens = self._forward_draft_loop( + inputs, attn_metadata, spec_metadata, draft_model, + draft_kv_cache_manager, num_contexts, num_gens, batch_size, + num_accepted_tokens, original_all_rank_num_tokens, resource_manager) + # restore attn_metadata to support cuda graph + self._restore_attn_metadata_from_spec_dec(attn_metadata) # restore all_rank_num_tokens for attention DP if original_all_rank_num_tokens is not None: attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens 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 5f8dcdb9070b..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, @@ -437,80 +437,68 @@ def forward( spec_metadata=spec_metadata, attn_metadata=attn_metadata) - # update attn metadata. change_attn_metadata saves the old attn - # metadata state and then mutates more of it, so the whole call must - # be inside the protected region. The save/restore pair must be - # exception-safe: a failure between them (e.g. an OOM during the - # max-shape general warmup, which model_engine tolerates and skips) - # would otherwise leave stale saved state on the persistent - # attn_metadata and fail every subsequent forward at the pairing - # assert in prepare_for_spec_dec. - # https://nvbugs/6442074 - try: - if attn_metadata is not None: - self.change_attn_metadata(num_accepted_tokens, attn_metadata, - spec_metadata) - - # Run MTP layers to predict draft tokens - next_draft_tokens = [] - last_tokens_idx = torch.cumsum( - attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 - - draft_kv_cache_manager = self.get_draft_kv_cache_manager( - resource_manager) - - with self.draft_kv_cache_context(attn_metadata, - draft_kv_cache_manager): - for i, mtp_layer in enumerate( - draft_model.mtp_layers[:runtime_draft_len]): - if self.guided_decoder is not None: - new_tokens = draft_inputs['input_ids'][last_tokens_idx] - self.guided_decoder.add_draft_batch(new_tokens, - num_accepted_tokens, + # update attn metadata + if attn_metadata is not None: + self.change_attn_metadata(num_accepted_tokens, attn_metadata, + spec_metadata) + + # Run MTP layers to predict draft tokens + next_draft_tokens = [] + last_tokens_idx = torch.cumsum( + attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 + + draft_kv_cache_manager = self.get_draft_kv_cache_manager( + resource_manager) + + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + for i, mtp_layer in enumerate( + draft_model.mtp_layers[:runtime_draft_len]): + if self.guided_decoder is not None: + new_tokens = draft_inputs['input_ids'][last_tokens_idx] + self.guided_decoder.add_draft_batch(new_tokens, + num_accepted_tokens, + draft_step=i) + + hidden_states = mtp_layer(embed_tokens=draft_model.embed_tokens, + **draft_inputs) + + logits = mtp_layer.shared_head(hidden_states, + draft_model.lm_head, + attn_metadata).float() + if self.guided_decoder is not None: + self.guided_decoder.execute_draft_batch(logits, draft_step=i) - hidden_states = mtp_layer( - embed_tokens=draft_model.embed_tokens, **draft_inputs) - - logits = mtp_layer.shared_head(hidden_states, - draft_model.lm_head, - attn_metadata).float() - if self.guided_decoder is not None: - self.guided_decoder.execute_draft_batch(logits, - draft_step=i) - - new_draft_token = self.draft_sampler(logits) - next_draft_tokens.append(new_draft_token) - # shift input_ids and hidden_states - input_ids = draft_inputs["input_ids"] - input_ids[:-1] = input_ids[1:].clone() - input_ids[last_tokens_idx] = new_draft_token - draft_hidden_states = draft_inputs["hidden_states"] - draft_hidden_states[:-1] = draft_hidden_states[1:].clone() - draft_hidden_states[last_tokens_idx] = hidden_states[ - last_tokens_idx, :] - draft_inputs = { - "input_ids": input_ids, - "position_ids": draft_inputs["position_ids"], - "hidden_states": draft_hidden_states, - "attn_metadata": draft_inputs["attn_metadata"], - } - next_draft_tokens = torch.stack(next_draft_tokens, dim=1) - - # Override with SA draft tokens after all MTP layers have run, - # so that MTP layers never see SA tokens in their inputs. - if self.sa_enhancer is not None: - num_contexts = attn_metadata.num_contexts - gen_draft_tokens = next_draft_tokens[num_contexts:] - gen_draft_tokens = self.sa_enhancer.maybe_override_all_draft_tokens( - gen_draft_tokens) - next_draft_tokens[num_contexts:] = gen_draft_tokens - finally: - # restore attn metadata. restore_from_spec_dec is safe on empty - # or partially saved state, so this is correct even if - # change_attn_metadata itself failed mid-save. - if attn_metadata is not None: - self._restore_attn_metadata_from_spec_dec(attn_metadata) + new_draft_token = self.draft_sampler(logits) + next_draft_tokens.append(new_draft_token) + # shift input_ids and hidden_states + input_ids = draft_inputs["input_ids"] + input_ids[:-1] = input_ids[1:].clone() + input_ids[last_tokens_idx] = new_draft_token + draft_hidden_states = draft_inputs["hidden_states"] + draft_hidden_states[:-1] = draft_hidden_states[1:].clone() + draft_hidden_states[last_tokens_idx] = hidden_states[ + last_tokens_idx, :] + draft_inputs = { + "input_ids": input_ids, + "position_ids": draft_inputs["position_ids"], + "hidden_states": draft_hidden_states, + "attn_metadata": draft_inputs["attn_metadata"], + } + next_draft_tokens = torch.stack(next_draft_tokens, dim=1) + + # Override with SA draft tokens after all MTP layers have run, + # so that MTP layers never see SA tokens in their inputs. + if self.sa_enhancer is not None: + num_contexts = attn_metadata.num_contexts + gen_draft_tokens = next_draft_tokens[num_contexts:] + gen_draft_tokens = self.sa_enhancer.maybe_override_all_draft_tokens( + gen_draft_tokens) + next_draft_tokens[num_contexts:] = gen_draft_tokens + + # restore attn metadata + if attn_metadata is not None: + self._restore_attn_metadata_from_spec_dec(attn_metadata) # prepare next new tokens to support overlap scheduler next_new_tokens = self._prepare_next_new_tokens( diff --git a/tensorrt_llm/_torch/speculative/pard.py b/tensorrt_llm/_torch/speculative/pard.py index 2622d155f117..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, @@ -236,120 +252,100 @@ def forward( max_draft_len=K, ) - # Save attn_metadata state and adjust kv_lens for the draft forward. - # The save/restore pair must be exception-safe: a failure between them - # (e.g. an OOM during the max-shape general warmup, which model_engine - # tolerates and skips) would otherwise leave stale saved state on the - # persistent attn_metadata and fail every subsequent forward at the - # pairing assert in prepare_for_spec_dec. - # https://nvbugs/6442074 - # - # Drop any rewind state left over from a previous call first, so a - # failure before _prepare_kv_for_draft_forward runs cannot make the - # finally block below re-apply a stale rewind. - if hasattr(self, "_kv_rewind_amount"): - del self._kv_rewind_amount - try: - self._prepare_attn_metadata_for_pard(attn_metadata, spec_metadata) - self._prepare_kv_for_draft_forward( - attn_metadata, num_accepted_tokens, num_contexts, batch_size - ) + self._prepare_attn_metadata_for_pard(attn_metadata, spec_metadata) + self._prepare_kv_for_draft_forward( + attn_metadata, num_accepted_tokens, num_contexts, batch_size + ) - position_ids = position_ids.squeeze(0) - inputs = self.prepare_1st_drafter_inputs( - input_ids=input_ids, - position_ids=position_ids, - hidden_states=hidden_states, - accepted_tokens=accepted_tokens, - num_accepted_tokens=num_accepted_tokens, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - draft_model=draft_model, - ) + position_ids = position_ids.squeeze(0) + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) - draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - if num_gens > 0: - with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): - hidden_states_out = draft_model.model(**inputs) + if num_gens > 0: + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + hidden_states_out = draft_model.model(**inputs) - # Gather K logits per gen request starting at the bonus position. - # Layout: [acc_0..acc_M, masks] (2K total). Positions M..M+K-1 - # produce K unique draft predictions (bonus + K-1 masks). - gen_start_idx = attn_metadata.num_ctx_tokens + # Gather K logits per gen request starting at the bonus position. + # Layout: [acc_0..acc_M, masks] (2K total). Positions M..M+K-1 + # produce K unique draft predictions (bonus + K-1 masks). + gen_start_idx = attn_metadata.num_ctx_tokens - request_bases = ( - torch.arange(num_gens, dtype=torch.long, device="cuda") * (2 * K) - + gen_start_idx - ) + request_bases = ( + torch.arange(num_gens, dtype=torch.long, device="cuda") * (2 * K) + + gen_start_idx + ) - gen_num_accepted = num_accepted_tokens[num_contexts:batch_size].long() - base_offsets = gen_num_accepted - 1 # M = bonus position - offsets = torch.arange(K, dtype=torch.long, device="cuda") + gen_num_accepted = num_accepted_tokens[num_contexts:batch_size].long() + base_offsets = gen_num_accepted - 1 # M = bonus position + offsets = torch.arange(K, dtype=torch.long, device="cuda") - gen_gather_ids = ( - request_bases.unsqueeze(1) - + base_offsets.unsqueeze(1) - + offsets.unsqueeze(0) - ).flatten() - gen_gather_ids = gen_gather_ids.clamp(max=hidden_states_out.shape[0] - 1) + gen_gather_ids = ( + request_bases.unsqueeze(1) + base_offsets.unsqueeze(1) + offsets.unsqueeze(0) + ).flatten() + gen_gather_ids = gen_gather_ids.clamp(max=hidden_states_out.shape[0] - 1) - gen_logits = draft_model.logits_processor( - hidden_states_out[gen_gather_ids], draft_model.lm_head, attn_metadata, True - ) + gen_logits = draft_model.logits_processor( + hidden_states_out[gen_gather_ids], draft_model.lm_head, attn_metadata, True + ) - vocab_size = gen_logits.shape[-1] - gen_logits = gen_logits.reshape(num_gens, K, vocab_size) + vocab_size = gen_logits.shape[-1] + gen_logits = gen_logits.reshape(num_gens, K, vocab_size) - # Use torch.argmax directly to avoid cute_argmax stride issues - d2t = getattr(draft_model.model, "d2t", None) - gen_draft_tokens = torch.argmax(gen_logits, dim=-1, keepdim=False).long() + # Use torch.argmax directly to avoid cute_argmax stride issues + d2t = getattr(draft_model.model, "d2t", None) + gen_draft_tokens = torch.argmax(gen_logits, dim=-1, keepdim=False).long() - if d2t is not None: - gen_draft_tokens = d2t[gen_draft_tokens] + gen_draft_tokens + if d2t is not None: + gen_draft_tokens = d2t[gen_draft_tokens] + gen_draft_tokens - gen_draft_tokens = gen_draft_tokens.type(torch.int32) + gen_draft_tokens = gen_draft_tokens.type(torch.int32) - if self.sa_enhancer is not None and sa_manager is not None: - gen_draft_tokens = self.sa_enhancer.maybe_override_all_draft_tokens( - gen_draft_tokens - ) + if self.sa_enhancer is not None and sa_manager is not None: + gen_draft_tokens = self.sa_enhancer.maybe_override_all_draft_tokens( + gen_draft_tokens + ) - # Pad from (num_gens, K) to (num_gens, 2K-1). - if K > 1: - pad = torch.zeros((num_gens, K - 1), dtype=torch.int32, device="cuda") - gen_draft_tokens = torch.cat([gen_draft_tokens, pad], dim=1) + # Pad from (num_gens, K) to (num_gens, 2K-1). + if K > 1: + pad = torch.zeros((num_gens, K - 1), dtype=torch.int32, device="cuda") + gen_draft_tokens = torch.cat([gen_draft_tokens, pad], dim=1) - elif num_contexts > 0 and self.use_separate_draft_kv_cache: - # Pure context batch: populate the draft KV cache so it's - # ready when generation starts. - with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): - draft_model.model(**inputs) - gen_draft_tokens = torch.empty((0, 2 * K - 1), dtype=torch.int32, device="cuda") + elif num_contexts > 0 and self.use_separate_draft_kv_cache: + # Pure context batch: populate the draft KV cache so it's + # ready when generation starts. + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + draft_model.model(**inputs) + gen_draft_tokens = torch.empty((0, 2 * K - 1), dtype=torch.int32, device="cuda") - else: - gen_draft_tokens = torch.empty((0, 2 * K - 1), dtype=torch.int32, device="cuda") + else: + gen_draft_tokens = torch.empty((0, 2 * K - 1), dtype=torch.int32, device="cuda") - if num_contexts > 0 and num_gens > 0: - ctx_draft_tokens = torch.zeros( - (num_contexts, 2 * K - 1), dtype=torch.int32, device="cuda" - ) - next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) - elif num_contexts > 0: - next_draft_tokens = torch.zeros( - (num_contexts, 2 * K - 1), dtype=torch.int32, device="cuda" - ) - else: - next_draft_tokens = gen_draft_tokens - finally: - self._restore_attn_metadata_from_spec_dec(attn_metadata) - - # Deferred kv_lens rewind (must happen after restore so it persists). - # Kept inside the finally, ordered after the restore: if the draft - # forward fails after _prepare_kv_for_draft_forward incremented - # kv_lens_cuda, the rewind must still be applied or the next - # iteration sees wrong KV lengths. - self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) + if num_contexts > 0 and num_gens > 0: + ctx_draft_tokens = torch.zeros( + (num_contexts, 2 * K - 1), dtype=torch.int32, device="cuda" + ) + next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) + elif num_contexts > 0: + next_draft_tokens = torch.zeros( + (num_contexts, 2 * K - 1), dtype=torch.int32, device="cuda" + ) + else: + next_draft_tokens = gen_draft_tokens + + self._restore_attn_metadata_from_spec_dec(attn_metadata) + + # Deferred kv_lens rewind (must happen after restore so it persists). + self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) next_new_tokens = self._prepare_next_new_tokens( accepted_tokens, 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 3678b67b2f7e..7398141badd1 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -22,9 +22,8 @@ accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[u accuracy/test_disaggregated_serving.py::TestQwen3_30B_A3B::test_mixed_ctx_gen_model[ctxpp2gentp2] SKIP (https://nvbugs/5748664) accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_nvfp4] SKIP (https://nvbugs/6336747) accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[qwen3vl_2b_instruct] SKIP (https://nvbugs/6422294) -accuracy/test_llm_api.py::TestLlama3_1_8BInstruct::test_guided_decoding_4gpus[xgrammar] SKIP (https://nvbugs/5346443) -accuracy/test_llm_api_autodeploy.py::TestGemma4MoE::test_bf16 SKIP (https://nvbugs/6445322) accuracy/test_llm_api_autodeploy.py::TestMiniMaxM2::test_finegrained_fp8 SKIP (https://nvbugs/6396422) +accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[nvfp4_ws8_80gb-trtllm] SKIP (https://nvbugs/6450341) accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_nvfp4[8] SKIP (https://nvbugs/6412108) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp] SKIP (https://nvbugs/6428101) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp_trtllm] SKIP (https://nvbugs/6426868) @@ -40,6 +39,7 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=0- accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6412102) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6388129) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6426847) +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6445456) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6402058) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6278337) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6428057) @@ -49,9 +49,12 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mt accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6388153) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6428094) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6428096) +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6445456) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp2pp2-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6198774) +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6445456) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=0-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6198785) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=0-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6198785) +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6445456) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6198774) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_online_eplb[mtp_nextn=2-moe_backend=WIDEEP] SKIP (https://nvbugs/6313993) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[pp4-mtp_nextn=0-fp8kv=False-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False-sampler_async_worker=False] SKIP (https://nvbugs/6388153) @@ -60,24 +63,34 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpu accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp2pp2-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6427411) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp2pp2-mtp_nextn=0-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-sampler_async_worker=False] SKIP (https://nvbugs/6427411) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp2pp2-mtp_nextn=0-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6427411) +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp4-mtp_nextn=0-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=True-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6445456) +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp4-mtp_nextn=0-fp8kv=False-attention_dp=True-cuda_graph=False-overlap_scheduler=False-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6445456) +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp4-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6445456) +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp4-mtp_nextn=0-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6445456) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=2-fp8kv=False-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6388363) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTEDSL-mtp_nextn=0-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6428087) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTEDSL-mtp_nextn=0-tp2pp2-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6427411) +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-ep4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6445456) +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6445456) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/5945081) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6384625) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6427411) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6428063) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6384625) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6384625) +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp2pp2-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6445472) +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6445456) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6272673) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-pp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6422432) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6245394) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6384625) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-tp2pp2-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6384625) +accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_auto_dtype SKIP (https://nvbugs/6450333) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_guided_decoding_4gpus[one_model] SKIP (https://nvbugs/5596343) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v1_kv_cache-True-True-triton-auto] SKIP (https://nvbugs/6026676) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-dp4-cutlass-auto] SKIP (https://nvbugs/6388142) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-ep4-cutlass-auto] SKIP (https://nvbugs/5596343) +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-ep4-trtllm-fp8] SKIP (https://nvbugs/6450338) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-tp4-cutlass-auto] SKIP (https://nvbugs/5596343) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-tp4-cutlass-fp8] SKIP (https://nvbugs/5651865) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-dp4-cutlass-auto] SKIP (https://nvbugs/5596343) @@ -171,7 +184,6 @@ examples/test_granite.py::test_llm_granite[granite-3.0-2b-instruct-bfloat16] SKI examples/test_medusa.py::test_llm_medusa_with_qaunt_base_model_1gpu[fp8-use_cpp_session-medusa-vicuna-7b-v1.3-4-heads-float16-bs1] SKIP (https://nvbugs/5802248) examples/test_medusa.py::test_llm_medusa_with_qaunt_base_model_1gpu[fp8-use_py_session-medusa-vicuna-7b-v1.3-4-heads-float16-bs1] SKIP (https://nvbugs/5333849) examples/test_multimodal.py::test_llm_fp8_multimodal_general[fp8-fp8-scienceqa-Llama-3.2-11B-Vision-Instruct-pp:1-tp:1-bfloat16-bs:1-cpp_e2e:False] SKIP (https://nvbugs/5222697) -examples/test_multimodal.py::test_llm_multimodal_general[Llama-3.2-11B-Vision-pp:1-tp:1-bfloat16-bs:1-cpp_e2e:False-nb:1] SKIP (https://nvbugs/5333818) examples/test_multimodal.py::test_llm_multimodal_general[Llama-3.2-11B-Vision-pp:1-tp:1-bfloat16-bs:8-cpp_e2e:False-nb:1] SKIP (https://nvbugs/5333818) examples/test_multimodal.py::test_llm_multimodal_general[Llama-3.2-11B-Vision-pp:1-tp:2-bfloat16-bs:1-cpp_e2e:False-nb:1] SKIP (https://nvbugs/5333818) examples/test_nemotron.py::test_llm_nemotron_3_8b_1gpu[bfloat16-fp8] SKIP (https://nvbugs/4961624) @@ -275,6 +287,8 @@ full:H100/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_ full:H100/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=True] SKIP (https://nvbugs/6313072) full:H100/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[llguidance-eagle3_one_model=False] SKIP (https://nvbugs/6422334) 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::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) @@ -282,15 +296,16 @@ full:H100/accuracy/test_llm_api_pytorch_multimodal.py::TestQwen3_5_27B_VL::test_ full:H100/accuracy/test_llm_api_pytorch_multimodal.py::TestQwen3_5_35B_A3B_VL::test_auto_dtype SKIP (https://nvbugs/6442073) full:H100/disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_tllm_gen_helix[DeepSeek-V3-Lite-bf16-short_prompt] SKIP (https://nvbugs/6414762) full:H100/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959) +full:H100/disaggregated/test_disaggregated.py::test_disaggregated_mixed_stress_test[req10k-conc512-qwen3_32b_fp8_mixed_stress] SKIP (https://nvbugs/6440089) full:H100/disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp4-TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6344107) full:H100/disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-qwen3_32b_fp8_stress] SKIP (https://nvbugs/6312828) +full:H100/test_e2e.py::test_multi_nodes_eval[Qwen3/Qwen3-235B-A22B-tp16-mmlu] SKIP (https://nvbugs/6423926) full:H100/test_e2e.py::test_openai_chat_guided_decoding[openai/gpt-oss-120b] SKIP (https://nvbugs/6384375) full:H100/test_e2e.py::test_ptp_quickstart_advanced_deepseek_multi_nodes[DeepSeek-V3] SKIP (https://nvbugs/6418410) +full:H100/test_e2e.py::test_qwen_e2e_cpprunner_large_new_tokens[DeepSeek-R1-Distill-Qwen-1.5B-DeepSeek-R1-Distill-Qwen-1.5B] SKIP (https://nvbugs/6414760) full:H100_PCIe/unittest/llmapi/test_llm_pytorch.py::test_llama_7b_multi_lora_evict_and_reload_lora_gpu_cache SKIP (https://nvbugs/5682551) full:H20/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=False] SKIP (https://nvbugs/6345827) full:H20/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=True] SKIP (https://nvbugs/6345827) -full:H20/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[llguidance-eagle3_one_model=False] SKIP (https://nvbugs/6422334) -full:H20/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar-eagle3_one_model=False] SKIP (https://nvbugs/6422334) full:H20/accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[nvidia_Llama-3.1-8B-Instruct-FP8-True] SKIP (https://nvbugs/6422351) full:H20/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:H20/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) @@ -423,10 +438,7 @@ test_e2e.py::test_multi_nodes_eval[DeepSeek-R1/DeepSeek-R1-0528-FP4-tp16-mmlu] S test_e2e.py::test_multi_nodes_eval[Kimi-K2-Thinking-NVFP4-tp16-mmlu] SKIP (https://nvbugs/6276983) test_e2e.py::test_multi_nodes_eval[MiniMax-M2-tp16-mmlu] SKIP (https://nvbugs/6373532) test_e2e.py::test_multi_nodes_eval[MiniMax-M3-tp16-mmlu] SKIP (https://nvbugs/6373561) -test_e2e.py::test_openai_chat_example[trt] SKIP (https://nvbugs/5477444) -test_e2e.py::test_openai_completions_example[trt] SKIP (https://nvbugs/5701450) test_e2e.py::test_ptp_quickstart_advanced_deepseek_r1_w4afp8_8gpus[DeepSeek-R1-W4AFP8-DeepSeek-R1/DeepSeek-R1-W4AFP8] SKIP (https://nvbugs/5836830) -test_e2e.py::test_trtllm_bench_iteration_log[TRT-streaming-meta-llama/Llama-3.1-8B-llama-3.1-model/Meta-Llama-3.1-8B] SKIP (https://nvbugs/5448523) unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py::test_prefill_varlen[varlen_hd512_overlap] SKIP (https://nvbugs/6426860) unittest/_torch/misc/test_share_tensor.py::TestShareTensor::test_share_tensor_different_dtypes SKIP (https://nvbugs/6418021) unittest/_torch/modeling -k "modeling_qwen" SKIP (https://nvbugs/6433376) @@ -439,6 +451,7 @@ unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py::test_llm_partial_update_weights_nvfp4[auto-Qwen3/Qwen3-8B] SKIP (https://nvbugs/6372690) unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py::test_llm_partial_update_weights_nvfp4[fp8-Qwen3/Qwen3-30B-A3B] SKIP (https://nvbugs/6372690) unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py::test_llm_partial_update_weights_nvfp4[fp8-Qwen3/Qwen3-8B] SKIP (https://nvbugs/6372690) +unittest/_torch/speculative/test_eagle3.py::test_llama_eagle3[True-TRTLLM-True-False-True-True-True-False-False-False] SKIP (https://nvbugs/6451425) unittest/_torch/thop/parallel/test_fp8_rowwise_linear.py::test_fp8_rowwise_linear[dtype1] SKIP (https://nvbugs/6301807) unittest/_torch/thop/serial/test_moe.py::TestMoeFp4::test_no_autotune[use_score_as_input-RoutingDSv3-swiglu-1024-1024-1] SKIP (https://nvbugs/5908070) unittest/_torch/thop/serial/test_moe.py::TestMoeFp4::test_no_autotune[use_score_as_input-RoutingRenormalize_qwen_next-swiglu-1024-1024-150] SKIP (https://nvbugs/5908070) @@ -470,8 +483,4 @@ unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_tinyllama_logits_processor_t unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[None] SKIP (https://nvbugs/6162504) unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[cuda_graph_config0] SKIP (https://nvbugs/6162504) unittest/llmapi/test_memory_profiling.py::test_profile_kvcache SKIP (https://nvbugs/5580781) -unittest/tools/test_layer_wise_benchmarks.py::test_deepseek_r1_ctx_dep[1] SKIP (https://nvbugs/6162541) -unittest/tools/test_layer_wise_benchmarks.py::test_nemotron_gen_dep[1] SKIP (https://nvbugs/6162541) -unittest/tools/test_layer_wise_benchmarks.py::test_performance_alignment[1] SKIP (https://nvbugs/6127669) -unittest/tools/test_layer_wise_benchmarks.py::test_qwen3_next_gen_tep[1] SKIP (https://nvbugs/6153575) verl/test_verl_cases.py::test_trtllm_abort SKIP (https://nvbugs/6272653) 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() From 4e8053ceb96a12a1c5f628d3c3e9a96df9800dd9 Mon Sep 17 00:00:00 2001 From: dongfengy <99041270+dongfengy@users.noreply.github.com> Date: Tue, 14 Jul 2026 17:40:24 +0000 Subject: [PATCH 4/6] [https://nvbugs/6442074][chore] Address review: DFlash ctx_len rollback, constructor-initialized rewind state, cleanup regression test - dflash.py: snapshot _ctx_len every forward and roll it back in _ensure_spec_dec_state_restored when the forward fails, so a tolerated failure cannot leave stale per-slot context lengths (the warmup-only restore now reuses the same snapshot). - dflash.py/pard.py: initialize the deferred kv_lens rewind state in the constructor and gate the rewind on _kv_rewind_amount instead of hasattr. - test_force_accepted_tokens.py: annotate the stub hook and add a regression test that a failure between prepare_for_spec_dec and restore leaves no saved state and does not trip the pairing assert afterwards. Co-Authored-By: Claude Fable 5 Signed-off-by: dongfengy <99041270+dongfengy@users.noreply.github.com> --- tensorrt_llm/_torch/speculative/dflash.py | 33 +++++++++++++++---- tensorrt_llm/_torch/speculative/pard.py | 8 ++++- .../speculative/test_force_accepted_tokens.py | 28 +++++++++++++++- 3 files changed, 61 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 31cdb67e3f5c..72ef2d1ef1a5 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] @@ -287,7 +297,7 @@ def _apply_kv_rewind_after_draft(self, attn_metadata, spec_metadata): # 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 @@ -381,6 +391,15 @@ def _ensure_spec_dec_state_restored(self, attn_metadata, spec_metadata): 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, @@ -415,10 +434,11 @@ def _forward_impl( 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() + self._saved_ctx_len = self._ctx_len.clone() + self._ctx_len_restore_pending = True self._execute_guided_decoder_if_present(logits) @@ -543,9 +563,10 @@ def _forward_impl( 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/pard.py b/tensorrt_llm/_torch/speculative/pard.py index 3d19ff4ebc54..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}" ) @@ -163,7 +169,7 @@ def _apply_kv_rewind_after_draft(self, attn_metadata, spec_metadata): # 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 diff --git a/tests/unittest/_torch/speculative/test_force_accepted_tokens.py b/tests/unittest/_torch/speculative/test_force_accepted_tokens.py index e3ac4bffaa8c..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, @@ -53,7 +54,7 @@ class _StubSpecWorker(SpecWorkerBase): def max_draft_len(self) -> int: return 8 - def _forward_impl(self, *args, **kwargs): + def _forward_impl(self, *args: object, **kwargs: object) -> None: raise NotImplementedError @@ -366,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 From 9df1b8b56b097dd1e9407c21c38a562100f9154e Mon Sep 17 00:00:00 2001 From: dongfengy <99041270+dongfengy@users.noreply.github.com> Date: Tue, 14 Jul 2026 19:11:10 +0000 Subject: [PATCH 5/6] [https://nvbugs/6442074][chore] Do not snapshot DFlash ctx_len while capturing a CUDA graph The rollback snapshot added for failed forwards was also taken during CUDA-graph capture, allocating the clone inside the graph memory pool; its replay-time writes could alias pool blocks reused by later captures. Skip the snapshot under capture: captured ops do not mutate _ctx_len until replay, and a capture-time failure aborts the graph itself, so rollback is only meaningful for eager and warmup forwards. Validated on viking-prod-235 (DGX H100): test_dflash_qwen3_5_4b[True] 3/3 passed with this guard. Co-Authored-By: Claude Fable 5 Signed-off-by: dongfengy <99041270+dongfengy@users.noreply.github.com> --- tensorrt_llm/_torch/speculative/dflash.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 72ef2d1ef1a5..a9dca0a4addc 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -437,8 +437,15 @@ def _forward_impl( # 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() - self._saved_ctx_len = self._ctx_len.clone() - self._ctx_len_restore_pending = True + 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) From f69a318c25c57b06b31681122fcd72acacc0da34 Mon Sep 17 00:00:00 2001 From: dongfengy <99041270+dongfengy@users.noreply.github.com> Date: Wed, 15 Jul 2026 00:38:25 +0000 Subject: [PATCH 6/6] [https://nvbugs/6442074][chore] Drop f-prefix from plain assert message fragments Review feedback on PR 16382: the fragments without placeholders do not need to be f-strings. Co-Authored-By: Claude Fable 5 Signed-off-by: dongfengy <99041270+dongfengy@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/interface.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 02f77bbdf582..fec6ba4e6b92 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -412,10 +412,10 @@ def create_cuda_graph_metadata(self, def prepare_for_spec_dec(self, *fields) -> None: assert len(self._saved_tensors) == 0, ( - f"prepare_for_spec_dec called while fields " + "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") + "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), f"{f} is not a torch.Tensor"