[https://nvbugs/6438685][fix] Prevent leaked ADP pad dummy from crashing the executor

[lingjiew]

Description

Resolves NVBug 6438685.

With attention DP enabled, a pad dummy request can leak across iterations and crash the executor with:

AssertionError
  File "tensorrt_llm/_torch/pyexecutor/py_executor.py", in _pad_attention_dp_dummy_request
    assert self.expected_num_active_requests >= len(self.active_requests)

Observed as a deterministic crash in disaggregated serving (wide-EP generation servers, GB300) whenever a burst of KV-cache transfers coincides with requests finishing: 21/21 benchmark cells died within the first minutes of a high-concurrency agentic workload.

Root cause

Three interacting behaviors:

1. Pad predicate wider than the scheduler's. _count_schedulable_active_requests counts requests at state GENERATION_TO_COMPLETE as schedulable, but MicroBatchScheduler refuses them (no_schedule_after_state=GENERATION_TO_COMPLETE). A rank holding only such requests therefore gets no pad dummy yet schedules batch = 0.
2. Fleet-wide skip leaks dummies. can_queue is "no zero in the all-gathered batch sizes", so that one rank turns can_queue False on every rank; _forward_step and _update_request_states are both skipped. The only same-iteration dummy removal lives in _update_request_states_tp, so pad dummies inserted on other ranks (e.g. ranks whose requests are all awaiting KV transfer) survive into the next iteration.
3. Hard cap makes the leak fatal. Next iteration, gather_all_rank_states counts the leaked dummy; with expected_num_active_requests capped at pp_size * max_batch_size, a rank at cap + 1 leaked dummy trips the assert above. (The cap itself is correct and untouched.)

Fix

_count_schedulable_active_requests: align the pad predicate with the scheduler's FULL schedulability window [CONTEXT_INIT, GENERATION_TO_COMPLETE) (disaggregated path), mirroring both no_schedule_until_state and no_schedule_after_state. The lower bound also covers DISAGG_CONTEXT_WAIT_SCHEDULER (gen-first mode on the context server) — the left-boundary mirror of the same mismatch. A rank that cannot actually schedule anything now receives its pad dummy and stays inside the collectives instead of forcing a fleet-wide skipped iteration — the leak-producing skip no longer occurs, so no dummy survives into the next round.

_pad_attention_dp_dummy_request: guard the add_dummy_requests() return — both V1 and V2 managers may return None when the rank has no free cache resources (possible while non-schedulable requests still hold theirs). On failure, log a warning and skip padding for the iteration instead of crashing on an unchecked [0]; the resource-holding requests complete within an iteration or two and padding succeeds on the retry.

Review note: an earlier revision also carried a defensive cleanup in _fetch_new_requests that terminated any dummy surviving a skipped iteration. Per review feedback it was removed — such a net would also silently absorb future leak bugs the predicate fix does not cover, hiding them from the assert. The root-cause fix stands alone.

Test Coverage

• New regression tests in tests/unittest/_torch/executor/test_py_executor.py:
  • test_pad_dummy_added_when_only_to_complete_requests_disagg — in disaggregated mode a rank holding only GENERATION_TO_COMPLETE requests now receives a pad dummy. The existing test_pad_dummy_counts_generation_to_complete (non-disagg path, transceiver None) is unchanged and still passes, since the predicate change is scoped to the disaggregated path.
  • test_pad_dummy_added_when_only_wait_scheduler_requests_disagg — left-boundary mirror: a rank holding only DISAGG_CONTEXT_WAIT_SCHEDULER requests receives its pad dummy.
  • test_pad_dummy_allocation_failure_skips_padding — add_dummy_requests() -> None degrades to a skipped iteration instead of a crash.
  • ADP mock request factories now provide the state_value attribute (with real disagg state ints) the window comparison reads.
• Validated on GB300 disaggregated serving (Qwen3.5 agentic benchmark, 42 configurations × 3600 s each, including the high-concurrency cache-pressure warmup that previously crashed 21/21 cells deterministically): no assert, no leaked-dummy warnings after the predicate fix, throughput/latency in line with expectations.

Notes / known limitations

• The state >= GENERATION_TO_COMPLETE bound also excludes GENERATION_COMPLETE, DISAGG_CONTEXT_TRANS_IN_PROGRESS, DISAGG_CONTEXT_COMPLETE, and DISAGG_GENERATION_WAIT_TOKENS — all likewise refused by the scheduler, so each exclusion is in the safe direction and extends the same fix to ctx-server / wait-tokens variants of the stall.
• Non-disaggregated ADP (transceiver None) keeps its existing counting semantics; a rank whose requests are all on their final token (GENERATION_TO_COMPLETE via the overlap path) can in principle still cause a fleet-wide skipped iteration there. This is deliberately out of scope for this PR (the current focus is the disaggregated path, and the non-disagg semantics are pinned by an existing unit test); it can be revisited if the stall is ever observed in aggregated serving.

PR Checklist

• The PR title follows the TensorRT-LLM convention ([None][fix] ...)
• Commit is signed off (DCO)
• Regression unit tests included

[coderabbitai]

📝 Walkthrough

Walkthrough

The executor’s disaggregated scheduling count now excludes requests in or beyond GENERATION_TO_COMPLETE. Test helpers populate request state values, and a regression test verifies _run_pad adds a dummy for an otherwise empty schedulable batch.

Changes

Disaggregated schedulability

Layer / File(s)	Summary
Filter to-complete requests tensorrt_llm/_torch/pyexecutor/py_executor.py	_count_schedulable_active_requests excludes disaggregated requests whose state reaches GENERATION_TO_COMPLETE, alongside existing transfer-state exclusions.
Cover request-state padding behavior tests/unittest/_torch/executor/test_benchmark_disagg.py, tests/unittest/_torch/executor/test_py_executor.py	Test request stubs now expose state_value, and a regression test verifies one pad dummy is added when only to-complete requests are active.

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

Suggested reviewers: suyoggupta, Superjomn, pengbowang-nv, hyukn, liji-nv, longlee0622, brb-nv, JunyiXu-nv, juney-nvidia

Sequence Diagram(s)

sequenceDiagram
  participant PyExecutor_run_pad
  participant PyExecutor_count_schedulable
  participant active_requests
  participant scheduler
  PyExecutor_run_pad->>PyExecutor_count_schedulable: count schedulable active requests
  PyExecutor_count_schedulable->>active_requests: inspect request state_value
  active_requests-->>PyExecutor_count_schedulable: return request states
  PyExecutor_count_schedulable-->>PyExecutor_run_pad: return filtered count
  PyExecutor_run_pad->>scheduler: queue pad dummy when needed

Loading

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title matches the main change: fixing leaked ADP pad dummy crashes in the executor.
Description check	✅ Passed	The description includes the issue, root cause, fix, test coverage, and checklist items required by the template.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[chienchunhung]

Thanks for catching and fixing this issue, while preserving the cap logic!

Also I think this PR potentially resolves https://nvbugspro.nvidia.com/bug/6438685. Can you include the bug info in PR title & description? Thanks!

[chienchunhung]

By the way, this return predicate makes a rank containing only GENERATION_TO_COMPLETE requests enter the existing dummy-allocation path. Since both V1 and V2 add_dummy_requests() may return None, could you confirm that disaggregated ADP reserves sufficient dummy capacity while terminal requests still retain their resources?

[lingjiew]

It can't be decided that disaggregated ADP reserves sufficient dummy capacity for now. This will be another change to reserve the dummy capacity. Added a warning here but not totally fix it in this PR.

[lancelly]

It seems that this PR is a subset of #16279, could you please take a look at it? @lingjiew @chienchunhung

[lingjiew]

It seems that this PR is a subset of #16279, could you please take a look at it? @lingjiew @chienchunhung
Yes. They are same fix. I'll close this PR since #16729 already fix it.

[chienchunhung]

It seems that this PR is a subset of #16279, could you please take a look at it? @lingjiew @chienchunhung

Sure; will take a look today.