From b951befb74d499d096af9ae8656610b6805ab2eb 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/3] [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 130bf38ad7dc..083d92e794e1 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) @@ -282,9 +270,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) @@ -292,16 +277,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) @@ -435,7 +419,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) @@ -448,7 +435,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) @@ -480,4 +466,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 0acfec7678996b96b4c2834ce61f72c8faa48984 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/3] [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 efa31fa2cc1e..6c653ff969a0 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -428,97 +428,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 43f79e1ac182957a09cb38aec2d725ad40c8f0a0 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/3] [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 6c653ff969a0..6c5ed826f55c 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -274,13 +274,17 @@ def _prepare_kv_for_draft_forward( if batch_size > num_contexts: attn_metadata.kv_lens_cuda[num_contexts:batch_size] += 1 + self._kv_rewind_pending = True attn_metadata.update_for_spec_dec() def _apply_kv_rewind_after_draft(self, attn_metadata, spec_metadata): """Apply the deferred kv_lens rewind after the draft forward.""" + self._kv_rewind_pending = False is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() if is_warmup: + # kv_lens_cuda was saved by prepare_for_spec_dec in this mode and + # is restored wholesale, so no rewind is needed. return if hasattr(self, "_kv_rewind_amount") and hasattr(attn_metadata, "kv_lens_cuda"): @@ -367,7 +371,19 @@ def _store_prefill_context( self._ctx_v_buf[slot, :, cur:end] = chunk_v.permute(1, 0, 2, 3) offset += slen - def forward( + def _ensure_spec_dec_state_restored(self, attn_metadata, spec_metadata): + # Restore first (in warmup mode kv_lens_cuda was saved and comes back + # wholesale), then apply any pending rewind for the other modes so a + # failed draft forward does not leave kv_lens_cuda incremented. + super()._ensure_spec_dec_state_restored(attn_metadata, spec_metadata) + if ( + getattr(self, "_kv_rewind_pending", False) + and attn_metadata is not None + and spec_metadata is not None + ): + self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) + + def _forward_impl( self, input_ids, position_ids, @@ -428,123 +444,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 083d92e794e1..686c3e1bca98 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) @@ -270,6 +282,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) @@ -277,15 +291,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) @@ -419,10 +434,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) @@ -435,6 +447,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) @@ -466,8 +479,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()