From 813bcd1e6968350e37dd65516af81aea5fa046bf Mon Sep 17 00:00:00 2001 From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> Date: Tue, 23 Jun 2026 15:30:35 -0700 Subject: [PATCH] [nvbugs/6336682][fix] AutoDeploy: pass cache_seq_interface in piecewise CG capture for spec dec The piecewise CUDA-graph capture path's piecewise_named_args_fn returned only cm.named_args, omitting cache_seq_interface. When the EagleWrapper top-level model is invoked with no cache_seq_interface, it dispatches to the prefill-only fallback which calls the FX-graph target_model with just inputs_embeds/position_ids; the post-transform target GraphModule expects the full set of cached SSM/conv state placeholders, raising TypeError: forward() missing 458 required positional arguments. Mirror _get_args_kwargs (which already handles the spec_config branch) so piecewise capture invokes the wrapper through the same KV-cache dispatch the runtime uses (ad_executor.py:1011). Also remove the now-redundant nvbugs/6336682 waiver. Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> --- .../transform/library/compile_model.py | 15 +++++++++------ tests/integration/test_lists/waives.txt | 2 -- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/compile_model.py b/tensorrt_llm/_torch/auto_deploy/transform/library/compile_model.py index 653455fb52f0..7558d86ed367 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/compile_model.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/compile_model.py @@ -123,12 +123,15 @@ def _apply_to_full_model( cm.info.reset() spec_config = cm._spec_config - def _get_args_kwargs(bs: int) -> ArgsKwargs: + def _named_args() -> dict: if spec_config is not None: - cm.info.set_capture_batch(batch_size=bs, max_draft_len=spec_config.max_draft_len) - return (), {**cm.named_args, "cache_seq_interface": cm} - cm.info.set_capture_batch(batch_size=bs) - return (), cm.named_args + return {**cm.named_args, "cache_seq_interface": cm} + return cm.named_args + + def _get_args_kwargs(bs: int) -> ArgsKwargs: + max_draft_len = spec_config.max_draft_len if spec_config is not None else 0 + cm.info.set_capture_batch(batch_size=bs, max_draft_len=max_draft_len) + return (), _named_args() resource_input_names = list(cm.resource_names) if spec_config is not None and "cache_seq_interface" not in resource_input_names: @@ -137,7 +140,7 @@ def _get_args_kwargs(bs: int) -> ArgsKwargs: config_overrides = {} if self.config.piecewise_enabled: extra_kwargs["piecewise_seq_info"] = cm.info - extra_kwargs["piecewise_named_args_fn"] = lambda: cm.named_args + extra_kwargs["piecewise_named_args_fn"] = _named_args max_seq = cm.info.max_seq_len max_batch = cm.info.max_batch_size diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index afb767029b1d..8faf8183ed09 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -6,7 +6,6 @@ accuracy/test_llm_api.py::TestLlama3_1_8BInstruct::test_guided_decoding_4gpus[xg accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[bf16-4-attn_dp_off-trtllm] SKIP (https://nvbugs/6367792) accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[fp8-4-attn_dp_off-trtllm] SKIP (https://nvbugs/6367792) accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[nvfp4-4-attn_dp_off-trtllm] SKIP (https://nvbugs/6367792) -accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[fp8_ws4_80gb-trtllm] SKIP (https://nvbugs/6336682) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp] SKIP (https://nvbugs/6281818) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp_trtllm] SKIP (https://nvbugs/6281818) accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_chunked_prefill[latency] SKIP (https://nvbugs/6276981) @@ -157,7 +156,6 @@ full:B300/accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[ep8] SKIP (h full:B300/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_genpp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6322073) full:B300/unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" SKIP (https://nvbugs/6165866) full:DGX_B200/unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" SKIP (https://nvbugs/6165866) -full:DGX_H100/accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[fp8_ws4_80gb-trtllm] SKIP (https://nvbugs/6336682) full:GB200/accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy SKIP (https://nvbugs/6276923) full:GB200/accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy_contention_opt SKIP (https://nvbugs/6276923) full:GB200/accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy_mode_b_overlap SKIP (https://nvbugs/6276923)