[https://nvbugs/6487039][fix] Generalize ADP dummy lifecycle

[chienchunhung]

Summary

Fixes NVBUG 6487039.

Attention-DP ranks must enter each collective forward with compatible batches, so a rank with no real work pads its batch with a synthetic request when a peer will run compute. This PR fixes the dummy lifecycle in the non-PP disaggregated path:

• dummy allocation is transactional across the fleet-wide queue decision;
• the generic non-overlap path derives fresh forward intent from the scheduler's eligibility contract; and
• V1 KV-cache managers track pre-prepared dummy ownership locally, avoiding duplicate sequence registration.

The scope is deliberate. Pipeline parallelism keeps its existing fallback. Generic overlap forward-intent is not generalized here; the validated model-specific scheduler-aware padding is preserved. Sequence-slot headroom follows the topology-based disaggregated attention-DP overlap policy already on main via #17282.

Why the failure surfaced

This was a latent lifecycle bug exposed by the high-concurrency Kimi disaggregated workload. The transactional ADP handling introduced by #16279 was scoped to DeepSeek-V4, so Kimi continued to use the legacy path. At concurrency 4096, a fleet-wide scheduling rejection can leave a tentative dummy alive long enough for repeated iterations to expose the broken active-request invariant.

Failure mechanism

The scheduler admits only requests within its configured state window and readiness rules. The legacy active-request count could still treat transfer, wait, or terminal requests outside that window as schedulable. A rank holding only those requests could omit a required dummy while peer ADP ranks allocated tentative dummies.

If the TP-wide queue decision rejected the iteration, tentative peer dummies were not rolled back. They accumulated until the executor hit:

assert self.expected_num_active_requests >= len(self.active_requests)

The subsequent threads can only be started once error and HangDetector/MPI abort were recovery consequences of the leaked dummy requests. A sticky context/generation role also allowed transfer-only requests to retain a stale dummy shape while a cell was draining.

Targeted validation then exposed a second ownership bug in the V1 cache path: add_dummy_requests pre-registered a context dummy's sequence, and normal context resource preparation attempted to register the same sequence again, triggering the C++ emplaceDone assertion.

Fix

• Use the scheduler's own state range and readiness rules when counting schedulable real requests. This also respects the ENCODER_INIT boundary for encoder-decoder models.
• On the generic non-PP, non-overlap disaggregated ADP path, derive a fresh none/generation/context intent every iteration and combine it with one scalar TP MAX reduction. A transfer-only rank adds no dummy when no peer has real forward work; otherwise it pads with the peer's compute role.
• Treat dummy allocation as a transaction: check the full KV requirement, tolerate rank-local allocation failure, and roll back only the tentative dummy when the TP-wide queue decision rejects the iteration.
• Track V1 pre-prepared dummy IDs per KV-cache manager. Target and draft managers retain independent ownership, normal context preparation skips only the already-owned registration, and ownership is cleared after the sequence is freed so the fixed dummy ID can be reused.
• Preserve the PP fallback and the model-specific overlap forward-intent behavior already validated on main.

This PR makes no independent sequence-slot sizing change. On the rebased base, normal non-overlap configurations use max_batch_size, PP uses pp_size * max_batch_size, and non-PP disaggregated attention-DP overlap uses the 2 * max_batch_size headroom supplied by #17282.

Verification

The exact NVBUG workload passed on the pre-rebase PR head d0d14ef2417877b12a89de62b1fd89acf421a2fe. Current-head CI is pending for rebased head 7c9c56733a5afab08fd5a2253c0a789800492a8d.

• Targeted pipeline #52581 ran with test reuse disabled.
• Exact stage: GB300-20_GPUs-5_Nodes-PyTorch-Disagg-PerfSanity-CTX1-NODE1-GPU4-GEN1-NODE4-GPU16-Post-Merge-2.
• Exact selector: perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL].
• The selector executed rather than being skipped and passed in 1,997.36 seconds.
• The benchmark completed 20,480/20,480 requests successfully with zero failed requests in 1,513.46 seconds.
• Slurm job 2901721 completed with exit code 0.
• The logs contain none of the original active-request assertion, duplicate emplaceDone assertion, executor restart error, HangDetector/MPI-abort chain, or KV-transfer/Python-transceiver timeout errors.
• Pre-commit passed on the same head.
• Rebase validation on 7c9c5673: Python syntax compilation and git diff --check passed; all formatting/lint pre-commit hooks passed; both test-list validators passed with Python 3.13 (2,063 unique entries). Local pytest remains unavailable because this host lacks torch.

Full merge-gate pipeline #52601 on pre-rebase head d0d14ef completed with unrelated Gemma4 timeout, GB200 performance-threshold, Jenkins, and Slurm flakes; it did not expose an ADP dummy-lifecycle regression. A fresh full run was started for rebased head 7c9c5673 from the trigger comment: PR_Github #65097 launched L0_MergeRequest_PR #52898, whose live console checked out the exact rebased SHA.

The exact Kimi test is already unwaived on main by #16920; this PR makes no waiver or timeout change.

Scope and related work

This PR intentionally does not:

• change KV-transfer timeouts, Python-transceiver behavior, or transceiver configuration;
• generalize cross-rank forward-intent consensus for overlap configurations;
• generalize the PP dummy path beyond its existing fallback; or
• modify test waivers.

#17278, already merged and included in this PR's base, makes overlap pad-dummy surplus fail-soft instead of crashing at the assertion. #17282, also included in the rebased base, supplies topology-based sequence-slot headroom for disaggregated attention-DP overlap. This PR preserves both protections while fixing the transactional leak and scheduler/ownership invariants within the scope above. Earlier Python-transceiver diagnostics are documented in #16918, and #16920 owns the test unwaiving.

[chienchunhung]

/bot run --disable-fail-fast --stage-list "GB300-20_GPUs-5_Nodes-PyTorch-Disagg-PerfSanity-CTX1-NODE1-GPU4-GEN1-NODE4-GPU16-Post-Merge-2"

[tensorrt-cicd]

PR_Github #62036 [ run ] triggered by Bot. Commit: 153a3df Link to invocation

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

The change generalizes ADP dummy handling, makes scheduler state ranges explicit, and propagates runtime sequence-slot capacity into model loading, MRoPE caches, speculative metadata, guided decoding, and executor padding behavior.

Changes

Runtime capacity and ADP scheduling

Layer / File(s)	Summary
Sequence-slot capacity propagation tensorrt_llm/_torch/pyexecutor/_util.py, tensorrt_llm/_torch/pyexecutor/model_engine.py, tensorrt_llm/_torch/pyexecutor/model_loader.py, tensorrt_llm/_torch/models/*, tensorrt_llm/_torch/pyexecutor/py_executor_creator.py, tensorrt_llm/_torch/speculative/*, tests/unittest/_torch/executor/test_seq_slot_sizing.py, tests/unittest/_torch/executor/test_model_loader_gms.py, tests/unittest/_torch/executor/test_pytorch_model_engine.py, tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py, tests/unittest/_torch/speculative/test_rejection_buffers_guard.py	Runtime sequence-slot sizing replaces DeepSeek-V4-specific overlap-headroom wiring and drives model configuration, MRoPE cache allocation, speculative metadata, guided decoding, and dummy slot indexing.
Scheduler state-range contract tensorrt_llm/_torch/pyexecutor/_util.py, tensorrt_llm/_torch/pyexecutor/scheduler/*, tests/unittest/_torch/executor/test_dual_pool_kv_cache.py, tests/unittest/_torch/executor/test_py_executor.py	Schedulers expose admissible request-state ranges, reject decoder requests waiting for encoder output, and forward configured bounds to the C++ micro-batch scheduler.
Generic ADP dummy handling tensorrt_llm/_torch/pyexecutor/model_engine.py, tensorrt_llm/_torch/pyexecutor/py_executor.py, tests/unittest/_torch/executor/test_benchmark_disagg.py, tests/unittest/_torch/executor/test_py_executor.py	ADP dummy-fix gating uses a generic mapping-based flag for request counting, dummy allocation finalization, and attention-DP padding, with expanded state-window and mixed-rank coverage.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant ModelEngine
  participant ModelLoader
  participant ModelConfig
  participant PyExecutor
  participant RequestScheduler
  ModelEngine->>ModelLoader: pass max_num_seq_slots
  ModelLoader->>ModelConfig: attach runtime slot capacity
  ModelEngine->>PyExecutor: provide generic ADP dummy-fix flag
  PyExecutor->>RequestScheduler: check request schedulability
  RequestScheduler-->>PyExecutor: return state-range eligibility
  PyExecutor->>PyExecutor: allocate or finalize ADP padding dummy

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#16190: Both touch PyTorchModelEngine initialization and related runtime behavior.
• NVIDIA/TensorRT-LLM#16990: Both modify the ADP dummy-padding path in py_executor.py.

Suggested labels: api-compatible

Suggested reviewers: bowenfu

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 26.42% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
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 identifies the bug, change type, and main change to generalize the ADP dummy lifecycle.
Description check	✅ Passed	The description clearly explains the problem, fix, scope, and extensive validation, with relevant test coverage provided.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[tensorrt-cicd]

PR_Github #62036 [ run ] completed with state SUCCESS. Commit: 153a3df
/LLM/main/L0_MergeRequest_PR pipeline #50222 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[BowenFu]

Flagging one thing before this goes further: the gate widening isn't inert for non-DeepSeek-V4 models in the disagg + attention-DP case, which is the case it targets.

Comparing against main, dropping the model_type == "deepseek_v4" term changes behavior in three places for any non-PP model that has a kv_cache_transceiver:

1. _count_schedulable_active_requests — the old non-DSv4 branch counted every active request except generation-init / transmission-in-progress. The new one counts only [CONTEXT_INIT, GENERATION_TO_COMPLETE), which excludes DISAGG_CONTEXT_WAIT_SCHEDULER and terminal GENERATION_TO_COMPLETE. A rank holding only those requests now reports zero schedulable and requests dummy padding where it previously didn't. Since this feeds needs_dummy per rank, ranks can now disagree about whether to pad.

2. _pad_attention_dp_dummy_request — the old non-DSv4 branch allocated the dummy unconditionally. The new path can return without padding (live-dummy check, _has_adp_dummy_kv_capacity precheck, OutOfPagesError), leaving that rank with an empty scheduled batch.

3. _finalize_adp_dummy_allocation — on can_queue=False it now removes the dummy and frees KV/spec resources for all non-PP models, where non-DSv4 previously returned immediately.

That combination is what #16279 deliberately scoped to DSv4, and the PR description acknowledges as much.

The test change makes this harder to catch rather than easier: test_dsv4_adp_dummy_fix_gate loses its ("deepseek_v3", 1, False) case, so the one assertion pinning non-DSv4 to the established path is deleted rather than replaced. If the widening is intended, could you add a non-DSv4 disagg + ADP case showing the two paths agree — plus something covering the multi-rank pad/no-pad agreement in (1)?

Minor, unrelated: this removes the same gb300_kimi-k25-thinking-fp4...NIXL e2e waiver line as #16920, so whichever merges second will need a rebase.

[chienchunhung]

/bot run --disable-fail-fast

[Shixiaowei02]

The dummy path is generalized but the paired 2x seq-slot headroom stays DSv4-only, worth generalizing both, or gating on the overlap scheduler. Could we also land the code fix alone and leave the waiver removal to a follow-up?

[chienchunhung]

Flagging one thing before this goes further: the gate widening isn't inert for non-DeepSeek-V4 models in the disagg + attention-DP case, which is the case it targets.

Comparing against main, dropping the model_type == "deepseek_v4" term changes behavior in three places for any non-PP model that has a kv_cache_transceiver:

1. _count_schedulable_active_requests — the old non-DSv4 branch counted every active request except generation-init / transmission-in-progress. The new one counts only [CONTEXT_INIT, GENERATION_TO_COMPLETE), which excludes DISAGG_CONTEXT_WAIT_SCHEDULER and terminal GENERATION_TO_COMPLETE. A rank holding only those requests now reports zero schedulable and requests dummy padding where it previously didn't. Since this feeds needs_dummy per rank, ranks can now disagree about whether to pad.
2. _pad_attention_dp_dummy_request — the old non-DSv4 branch allocated the dummy unconditionally. The new path can return without padding (live-dummy check, _has_adp_dummy_kv_capacity precheck, OutOfPagesError), leaving that rank with an empty scheduled batch.
3. _finalize_adp_dummy_allocation — on can_queue=False it now removes the dummy and frees KV/spec resources for all non-PP models, where non-DSv4 previously returned immediately.

That combination is what #16279 deliberately scoped to DSv4, and the PR description acknowledges as much.

The test change makes this harder to catch rather than easier: test_dsv4_adp_dummy_fix_gate loses its ("deepseek_v3", 1, False) case, so the one assertion pinning non-DSv4 to the established path is deleted rather than replaced. If the widening is intended, could you add a non-DSv4 disagg + ADP case showing the two paths agree — plus something covering the multi-rank pad/no-pad agreement in (1)?

Minor, unrelated: this removes the same gb300_kimi-k25-thinking-fp4...NIXL e2e waiver line as #16920, so whichever merges second will need a rebase.

Thanks.

I added test_non_dsv4_disagg_adp_mixed_rank_states_stay_queueable to cover the generic non-PP path with one busy rank and one terminal-only rank. It verifies that only the terminal rank pads, both ranks expose batch size 1 to _can_queue, the TP-wide decision stays queueable, and the dummy is committed.

Existing allocation-failure and rollback tests cover the other two behavior changes.

[tburt-nv]

The whole unittest/_torch/executor directory is covered in the test-db.

[chienchunhung]

The dummy path is generalized but the paired 2x seq-slot headroom stays DSv4-only, worth generalizing both, or gating on the overlap scheduler. Could we also land the code fix alone and leave the waiver removal to a follow-up?

IIUC these are separate lifecycle conditions rather than paired requirements.

The dummy fix is needed when rank-local padding or allocation disagrees, including this Kimi configuration where overlap is disabled. The 2× slot pool addresses overlap backfill before previous-iteration slot release; this case uses MTP0 with disable_overlap_scheduler: true, so additional headroom would not participate.

Gating dummy handling on overlap would therefore disable the reported fix, while generalizing slot headroom would require broader validation of all py_seq_slot-indexed consumers. I'd keep the independent gates and moved the waiver removal out of this PR because #16920 already owns the same unwaive with its timeout change.

[brnguyen2]

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

The fix matches the failure mechanism and the generalization is the right direction — delegating the ADP count to the scheduler's own state contract removes the duplicated window that caused the desync. Two things to settle before merge:

1. Blast radius: the gate widens from one model to every non-PP disagg ADP deployment, and the 2× slot sizing now applies to every non-PP overlap config (the default). Targeted B200 stages passed; please confirm pipeline #51225 lands green before merging, since that's the first full-coverage run of the generalized path.
2. The description is thorough and honest about what this does and doesn't fix (test stays waived) — appreciated.

Remaining comments are inline; nothing blocking.

[Shixiaowei02]

The extra slots are on by default for every non-pipeline setup now, but the bound worked out on this thread is a small constant, not double. The biggest buffer is allocated before the cache pool is sized from free memory, so this shrinks the cache for every default speculative setup.

[chienchunhung]

/bot run --disable-fail-fast --stage-list "DGX_B200-4_GPUs-PyTorch-3, DGX_B200-8_GPUs-PyTorch-1"

[tensorrt-cicd]

PR_Github #64417 [ run ] triggered by Bot. Commit: 6408929 Link to invocation

[chienchunhung]

Hi @chienchunhung , thanks for the effort. Please ensure that multi-GPU testing is fully validated.

The new schedulable window excludes the in-transfer states, so a draining context server now builds a full-length context dummy every iteration. Fail-soft, but new work for every non-DeepSeek model. Can the role be re-evaluated while draining?

Hi @Shixiaowei02 Thanks for raising this good point.

The generic non-overlap ADP path now derives forward intent from scheduler-eligible real requests on every iteration and combines a scalar context/generation/none intent across ADP ranks. A rank that is only draining transfers therefore adds no dummy when no peer has compute, or adds a correctly shaped dummy matching the peer’s context or generation work.

The additional collective is limited to non-PP, non-overlap disaggregated ADP. The existing overlap path remains unchanged and does not incur this collective.

I added coverage for transfer-only draining, context-role re-evaluation, and preservation of the overlap behavior. I will validate the resulting head with targeted multi-GPU CI.

The extra slots are on by default for every non-pipeline setup now, but the bound worked out on this thread is a small constant, not double. The biggest buffer is allocated before the cache pool is sized from free memory, so this shrinks the cache for every default speculative setup.

The generic 2 * max_batch_size expansion has been removed. Normal non-PP configurations retain max_batch_size, while the extra headroom remains opt-in only for the previously validated DeepSeek-V4 one-model MTP overlap path. PP retains its existing pp_size * max_batch_size sizing.

The broader overlap bound and small-constant solution remain follow-up scope in #17278.

[tensorrt-cicd]

PR_Github #64417 [ run ] completed with state SUCCESS. Commit: 6408929
/LLM/main/L0_MergeRequest_PR pipeline #52299 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast --disable-reuse-test --stage-list "GB300-20_GPUs-5_Nodes-PyTorch-Disagg-PerfSanity-CTX1-NODE1-GPU4-GEN1-NODE4-GPU16-Post-Merge-2"

[tensorrt-cicd]

PR_Github #64501 [ run ] triggered by Bot. Commit: 6408929 Link to invocation

[tensorrt-cicd]

PR_Github #64501 [ run ] completed with state FAILURE. Commit: 6408929
/LLM/main/L0_MergeRequest_PR pipeline #52371 (Partly Tested) 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

[chienchunhung]

/bot run --disable-fail-fast --disable-reuse-test --stage-list "DGX_B200-PyTorch-5, GB300-20_GPUs-5_Nodes-PyTorch-Disagg-PerfSanity-CTX1-NODE1-GPU4-GEN1-NODE4-GPU16-Post-Merge-2"

[tensorrt-cicd]

PR_Github #64728 [ run ] triggered by Bot. Commit: d0d14ef Link to invocation

[tensorrt-cicd]

PR_Github #64729 [ run ] triggered by Bot. Commit: d0d14ef Link to invocation

[tensorrt-cicd]

PR_Github/16921-d0d14ef #64728 was force-killed by a newer pipeline run.
L0 job information not available (job may not have been triggered yet).

Link to superseding invocation

[tensorrt-cicd]

PR_Github #64729 [ run ] completed with state SUCCESS. Commit: d0d14ef
/LLM/main/L0_MergeRequest_PR pipeline #52581 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #64752 [ run ] triggered by Bot. Commit: d0d14ef Link to invocation

[tensorrt-cicd]

PR_Github #64752 [ run ] completed with state FAILURE. Commit: d0d14ef
/LLM/main/L0_MergeRequest_PR pipeline #52601 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

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65097 [ run ] triggered by Bot. Commit: 7c9c567 Link to invocation