From e7b25f0c4a61b4d99bdaee6d8f04e7862ee1d5d2 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Wed, 8 Jul 2026 05:57:45 -0700 Subject: [PATCH] [nvbugs/6422334][fix] Skip py_last_draft_tokens snapshot for DISAGG_GENERATION_INIT Two-model Eagle3 disagg with disable_overlap_scheduler=True crashes with File "tensorrt_llm/_torch/speculative/model_drafter.py", line 240 assert num_draft_tokens == 0 AssertionError when guided decoding is enabled. The AssertionError kills the drafter thread on the generation server; every subsequent request fails with "Event loop terminated with error", so pytest's HTTP client can no longer connect and the test times out at 3600 s. Root cause: while a request waits in DISAGG_GENERATION_INIT for KV cache transfer, PyExecutor._prepare_draft_requests runs each iteration and unconditionally snapshotted py_last_draft_tokens = py_draft_tokens for both GENERATION_IN_PROGRESS and DISAGG_GENERATION_INIT. Because it also refills py_draft_tokens with [0]*max_total_draft_tokens right after, py_last_draft_tokens accumulates a dummy [0, 0, 0]. When the request finally transitions to GENERATION_IN_PROGRESS, ModelDrafter enters the first-time-context branch and asserts num_draft_tokens == 0, but the stale snapshot makes _initialize_draft_tokens return 3 -> assert fires. Only snapshot py_last_draft_tokens for requests that have actually been through the drafter (GENERATION_IN_PROGRESS). Also remove the two H20 waivers that this fix makes obsolete. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 4 +++- tests/integration/test_lists/waives.txt | 2 -- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index f03feff26193..b0b27309e558 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3966,7 +3966,9 @@ def _prepare_draft_requests(self): LlmRequestState.DISAGG_GENERATION_INIT): continue - req.py_last_draft_tokens = req.py_draft_tokens + # Skip DISAGG_GENERATION_INIT: snapshotting the dummy py_draft_tokens leaks stale state into the first real draft. + if req.state == LlmRequestState.GENERATION_IN_PROGRESS: + req.py_last_draft_tokens = req.py_draft_tokens if self.max_total_draft_tokens > 0 and self.use_spec_decode and not req.py_disable_speculative_decoding: req.py_draft_tokens = [0] * self.max_total_draft_tokens diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index d67988ac1721..01e95380cdbd 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -270,8 +270,6 @@ full:H100/test_e2e.py::test_ptp_quickstart_advanced_deepseek_multi_nodes[DeepSee 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)