[https://nvbugs/6402500][fix] Only snapshot py_last_draft_tokens = py_draft_tokens for requests actually…

[trtllm-agent]

Summary

• Root cause: _prepare_draft_requests recorded dummy draft tokens as py_last_draft_tokens for DISAGG_GENERATION_INIT requests, leaking stale state into the first draft invocation once transfer completed.
• Fix: Only snapshot py_last_draft_tokens = py_draft_tokens for requests actually in GENERATION_IN_PROGRESS; leave py_last_draft_tokens unchanged (None) while a request is still in DISAGG_GENERATION_INIT.
• Automated fix generated by repair-bot

Test plan

• Verify fix on the same GPU type as the original failure
• Check for regressions in related tests

Links

• Bug: https://nvbugs/6402500

Summary by CodeRabbit

• Bug Fixes
  • Fixed an issue with draft-token handling in disaggregated generation requests, helping prevent stale state from affecting the first draft after a request transitions into generation.
• Tests
  • Removed a test waiver so one disaggregated serving accuracy test now runs instead of being skipped.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 4741b861-0de8-4865-b24d-70f3eef816ba

📥 Commits

Reviewing files that changed from the base of the PR and between d9ca6f1 and 38e216a.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt

---

📝 Walkthrough

Walkthrough

Modifies _prepare_draft_requests in py_executor.py to snapshot py_last_draft_tokens only for requests already in GENERATION_IN_PROGRESS state, preventing stale draft-token carryover from DISAGG_GENERATION_INIT. Also removes a related test waiver entry from waives.txt.

Changes

Draft Token Fix and Test Waiver Removal

Layer / File(s)	Summary
Conditional draft-token snapshot tensorrt_llm/_torch/pyexecutor/py_executor.py	py_last_draft_tokens is assigned only when request state is GENERATION_IN_PROGRESS, avoiding stale snapshots from DISAGG_GENERATION_INIT.
Test waiver removal tests/integration/test_lists/waives.txt	Removes the skip entry for the eagle3 disaggregated serving test, re-enabling that test.

Estimated code review effort: 2 (Simple) | ~10 minutes

Related PRs: None identified.

Suggested labels: bug, disaggregated-serving

Suggested reviewers: None identified.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is specific, uses the required bug/fix prefix, and matches the main change in the PR.
Description check	✅ Passed	The description covers the issue, fix, test plan, and bug link, though it does not include the template's PR checklist section.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[ziyixiong-nv]

LGTM

[ziyixiong-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57330 [ run ] triggered by Bot. Commit: a35b7e9 Link to invocation

[tensorrt-cicd]

PR_Github #57330 [ run ] completed with state ABORTED. Commit: a35b7e9
/LLM/main/L0_MergeRequest_PR pipeline #46087 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[trtllm-agent]

NVBug 6402500 is closed as Duplicate. The linked bug appears resolved elsewhere or for a reason that does not prove this PR is redundant. This PR should be judged on its own merits; repair-bot is not auto-closing it.

[brnguyen2]

Approving — the comments below are optional touch-ups, not blockers.

The diff is waiver-removal only — the _prepare_draft_requests gating this describes already landed in main via #16121 (filed under a different bug). Please retitle/reword to [chore] remove stale waivers and say the fix is already in, otherwise a future reader will look for a code change that isn't here.

Also: [waives.txt:252](https://github.com/NVIDIA/TensorRT-LLM/pull/15888/files#diff-621bd2af82a3b97c7a5948368d36c14582ffdf73361deb7f42d2ff395b22167eR252)-253 still waives the same two test_guided_decoding_with_eagle3[...eagle3_one_model=False] cases on H100 under nvbugs/6422334. If the same fix unblocked GB200, those look stale too — worth checking in the same pass.