[viraatc]

Description

Port to main of #21496 (open against feat/k3_26.08).

On a disaggregated generation server, a request that arrives from the context server gets zero placeholder drafts in _prepare_disagg_gen_transmission_complete. That happens when the context server sent no drafts, or when its drafter returned zeros for the context step. The executor requests the one-hot draft-probability row only for generation requests that enter with no drafts. So these placeholders are verified against a draft-probability row that nothing has written.

With use_rejection_sampling and a request whose sampling distribution gives token 0 a nonzero probability (temperature > 0 without top-p), the first generation step can accept the placeholders. The output can then lock into a run of token 0. Greedy requests, strict acceptance, and top-p requests that exclude token 0 are not affected.

The fix sets py_needs_onehot_draft_probs at the handoff when the installed drafts are all placeholders and rejection sampling is on. Rejection sampling stays exact with a one-hot proposal row; only the acceptance of that first step changes. Real context drafts, including ones zero-padded to this server's draft width, keep their rows. use_rejection_sampling is read only on the placeholder path, so executors without a spec config are untouched.

Test Coverage

• New tests/unittest/_torch/pyexecutor/test_disagg_gen_placeholder_drafts.py, collected by the unittest/_torch/pyexecutor entry of l0_dgx_b300:
  • Placeholder drafts (none, empty, zeros) request the one-hot row only with rejection sampling.
  • Real drafts keep their rows, full width or zero-padded.
  • Without speculative decoding there are no drafts.
• The same fix and test on feat/k3_26.08 ([None][fix] Disaggregated generation: one-hot draft probabilities for placeholder drafts at the handoff #21496): 3 failed / 7 passed without the fix, 10 passed with it, in a GB300 container.
• On feat/k3_26.08 it removed the token-0 runs end to end: 0 of 120 probe requests vs 63 of 120 without it (DSpark drafter, rejection sampling on, temperature 1.0, no top-p).
• Run on a GB300 node in the public main nightly nvcr.io/nvidia/tensorrt-llm/release:1.4.0rc3.dev202611030007:
  • As shipped, the three placeholder cases with rejection sampling fail (3 failed, 6 passed).
  • With this PR's hunk applied, all 9 pass.
  • The existing handoff tests in test_py_executor.py still pass with the hunk applied (-k disagg_generation).
• ruff check, ruff format and yapf 0.43.0 are clean on the changed lines.

PR Checklist

• Commits signed off (DCO).
• One concern: the handoff path only.

[viraatc]

Closing: the same fix is #21496 on feat/k3_26.08, which reaches main with the feat/k3_26.08 merge-back (#21067).