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

[chienchunhung]

Summary

Fixes NVBUG 6487039.

This PR fixes the model-independent, non-overlap disaggregated attention-DP dummy lifecycle. It generalizes the transactional dummy handling introduced for DeepSeek-V4 in #16279 without enabling generic overlap headroom.

The affected GB300 Kimi E2E test remains waived. This PR fixes the dummy-request assertion and leak, but not the separate Python-transceiver liveness and performance failures also exposed by that workload.

Failure mechanism

The scheduler admits only requests in its configured scheduling-state window. The legacy active-request count could nevertheless treat transfer, wait, or terminal requests outside that window as schedulable. A rank holding only those requests could omit a required pad dummy while peer attention-DP ranks allocated tentative dummies.

If the TP-wide queue decision rejected that iteration, the 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 dummies.

A sticky context/generation role introduced another draining hazard: transfer-only requests could cause a context server to allocate a full-length context dummy even when no rank intended to run a forward pass.

Fix

• Use each scheduler's own state range and readiness rules when counting schedulable real requests, including the ENCODER_INIT boundary for encoder-decoder models.
• On the generic non-PP, non-overlap disaggregated ADP path, derive a fresh none/generation/context forward intent every iteration and combine it with one scalar TP MAX reduction. Transfer-only ranks add no dummy when no peer has real forward work; otherwise they add a dummy matching the peer's compute role.
• Treat dummy allocation as a transaction: check its full KV requirement, tolerate rank-local allocation failure, and roll back only the tentative dummy when the TP-wide queue decision rejects the iteration.
• Preserve the established PP fallback and the previously validated DeepSeek-V4 one-model MTP overlap behavior.

This PR does not expand sequence-slot capacity for default non-PP configurations. Normal non-PP configurations remain at max_batch_size; PP remains at pp_size * max_batch_size. The existing 2 * max_batch_size headroom remains opt-in only for the validated DeepSeek-V4 one-model MTP overlap configuration.

Verification

Conclusion: The NVBUG 6487039 dummy-request lifecycle defect and its reported assertion/recovery signature are fixed and regression-validated. The complete GB300 Kimi E2E test is not verified passing end-to-end and remains waived because it still encounters independent KV-transfer/Python-transceiver liveness failures.

• Exact failure-signature evidence: The affected Kimi selector ran temporarily unwaived for 90 minutes in pipeline #50411. The active-request assertion, executor restart error, and HangDetector/MPI-abort signature did not recur. The run instead failed later with persistent KV-transfer timeouts. This verifies the scoped lifecycle fix, not permanent test unwaiving.
• Current-head regression validation: Targeted multi-GPU CI passed on commit 6408929bc in pipeline #52299:
  • DGX_B200-4_GPUs-PyTorch-3: 8 passed, 4 skipped by pre-existing NVBUG waivers; pytest exit code 0.
  • DGX_B200-8_GPUs-PyTorch-1: 12 passed, 2 skipped by pre-existing NVBUG waivers; pytest exit code 0.

Scope and follow-up

This PR intentionally does not:

• unwaive the GB300 Kimi E2E test;
• change KV-transfer timeouts, transceiver behavior, or transceiver configuration;
• solve generic overlap dummy ownership, concurrent overlap lifecycles, or the small-constant overlap headroom bound;
• generalize the PP dummy path beyond its existing fallback.

Generic overlap lifecycle and headroom are follow-up scope in #17278. The broader Python-transceiver and workload-liveness issues remain tracked in #16918 and #16920.

[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 follows the repository format and clearly describes the generalized ADP dummy lifecycle fix.
Description check	✅ Passed	The description explains the failure, fix, scope, limitations, and targeted test coverage in sufficient detail.

✨ 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.

[tensorrt-cicd]

PR_Github #63137 [ run ] completed with state FAILURE. Commit: 8a8286c
/LLM/main/L0_MergeRequest_PR pipeline #51225 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 --stage-list "A10-Build_Docs, B300-PyTorch-2"

[tensorrt-cicd]

PR_Github #63244 [ run ] triggered by Bot. Commit: 8a8286c Link to invocation

[tensorrt-cicd]

PR_Github #63244 [ run ] completed with state DISABLED
Pipeline is freezed and top-1 instance is under maintenance. For urgent request, contact Yiteng Niu

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast --stage-list "A10-Build_Docs, B300-PyTorch-2"

[tensorrt-cicd]

PR_Github #63557 [ run ] triggered by Bot. Commit: 8a8286c Link to invocation

[tensorrt-cicd]

PR_Github #63557 [ run ] completed with state SUCCESS. Commit: 8a8286c
/LLM/main/L0_MergeRequest_PR pipeline #51523 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #63852 [ run ] triggered by Bot. Commit: 36b2c48 Link to invocation

[BowenFu]

Scheduler eligibility now drives ADP counting, mixed-rank rollback has focused coverage, overlap capacity reaches each slot-indexed consumer, and the waiver change is removed. Approved; please merge after full CI passes.

[tensorrt-cicd]

PR_Github #63852 [ run ] completed with state SUCCESS. Commit: 36b2c48
/LLM/main/L0_MergeRequest_PR pipeline #51793 completed with status: 'FAILURE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ 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

[Shixiaowei02]

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?

[thorjohnsen]

LGTM

[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