[None][fix] Seq-slot pool overlap headroom with consistent slot-indexed buffer sizing - #16279
Conversation
|
/bot run |
|
PR_Github #58778 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThe changes adjust overlap-aware sequence capacity calculations, exclude ineligible requests from scheduling counts, and make attention-DP dummy insertion tolerate unavailable KV-cache slots. Tests cover slot sizing, generation-completion states, disaggregated mode, and failed dummy allocation. ChangesExecutor scheduling updates
Estimated code review effort: 3 (Moderate) | ~20 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
tests/unittest/_torch/executor/test_py_executor.py (1)
1383-1404: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider also asserting the resulting
active_requestsstate.The test confirms a dummy-add was requested (
len(stub.add_dummy_calls) == 1) but not that the dummy actually lands instub.active_requests, which is the behavior that keeps ADP collectives aligned. Strengthening the assertion would catch a regression where the call happens but the returned dummy request is never appended._run_pad(stub) assert len(stub.add_dummy_calls) == 1 + assert len(stub.active_requests) == 2🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/executor/test_py_executor.py` around lines 1383 - 1404, The test test_pad_dummy_added_for_rank_whose_requests_are_all_to_complete should also verify that the dummy request is present in stub.active_requests after _run_pad. Add an assertion for the resulting active request collection, preserving the existing add_dummy_calls assertion.Source: Path instructions
tensorrt_llm/_torch/pyexecutor/_util.py (1)
2088-2097: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate headroom-calc logic — extract a shared helper.
The
num_micro_batchesderivation is copy-pasted verbatim increate_py_executor_instanceandcreate_torch_sampler_args, with only a comment ("Must match create_py_executor_instance's slot-pool sizing.") enforcing consistency. If one copy is edited later without the other, the sampler'smax_num_sequenceswill silently diverge from the slot-pool size and reintroduce the "No free slots" failure this PR fixes.♻️ Proposed refactor: shared helper
+def _compute_slot_pool_micro_batches(mapping: Mapping, + disable_overlap_scheduler: bool) -> int: + """Micro-batch multiplier for slot-pool / sampler sequence-capacity sizing. + + Under overlap, prev-iter finished requests still hold slots when the + next iteration's prepare_resources runs, so headroom for two in-flight + iterations is needed. PP is structurally overlap-style regardless of + disable_overlap_scheduler, so it always uses pp_size as the multiplier. + """ + if mapping.has_pp(): + return mapping.pp_size + return 1 if disable_overlap_scheduler else 2 + + def create_py_executor_instance( ... - if mapping.has_pp(): - num_micro_batches = mapping.pp_size - else: - num_micro_batches = 1 if llm_args.disable_overlap_scheduler else 2 - max_num_sequences = max_batch_size * num_micro_batches + num_micro_batches = _compute_slot_pool_micro_batches( + mapping, llm_args.disable_overlap_scheduler) + max_num_sequences = max_batch_size * num_micro_batchesdef create_torch_sampler_args(...): - # Must match create_py_executor_instance's slot-pool sizing. - if mapping.has_pp(): - num_micro_batches = mapping.pp_size - else: - num_micro_batches = 1 if disable_overlap_scheduler else 2 - max_num_sequences = max_batch_size * num_micro_batches + num_micro_batches = _compute_slot_pool_micro_batches( + mapping, disable_overlap_scheduler) + max_num_sequences = max_batch_size * num_micro_batchesAlso applies to: 2454-2459
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/pyexecutor/_util.py` around lines 2088 - 2097, Extract the duplicated num_micro_batches and max_num_sequences calculation into a shared helper in _util.py, then call it from both create_py_executor_instance and create_torch_sampler_args. Remove the copy-pasted derivation and the comment-only consistency dependency, ensuring both callers use identical slot-pool sizing behavior for pipeline parallelism and overlap scheduling.tests/unittest/_torch/executor/test_seq_slot_sizing.py (1)
1-45: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest coverage: sufficient for
create_torch_sampler_args, but tied to duplicated logic.Coverage of the four
pp_size/disable_overlap/expected_factorcombinations is sufficient and correctly matches the implementation in_util.py. However, this test only exercisescreate_torch_sampler_args; the identical (but separately duplicated) logic increate_py_executor_instanceformax_num_sequencessizing has no equivalent direct test. If the duplicated logic in_util.pyis extracted into a shared helper (see review on that file), consider testing that helper directly so both call sites stay covered by one test.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/executor/test_seq_slot_sizing.py` around lines 1 - 45, Extend coverage beyond create_torch_sampler_args by extracting the shared max_num_sequences sizing logic used by create_torch_sampler_args and create_py_executor_instance into a common helper, then test that helper with the existing pp_size/disable_overlap/expected_factor cases. Update both call sites to reuse the helper so the same behavior remains covered consistently.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/executor/test_seq_slot_sizing.py`:
- Around line 30-31: Run ruff format on the test_seq_slot_sizing.py file and
apply its formatting changes, particularly to
test_sampler_max_num_sequences_includes_overlap_headroom, so the file passes the
CI ruff-format check.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 2088-2097: Extract the duplicated num_micro_batches and
max_num_sequences calculation into a shared helper in _util.py, then call it
from both create_py_executor_instance and create_torch_sampler_args. Remove the
copy-pasted derivation and the comment-only consistency dependency, ensuring
both callers use identical slot-pool sizing behavior for pipeline parallelism
and overlap scheduling.
In `@tests/unittest/_torch/executor/test_py_executor.py`:
- Around line 1383-1404: The test
test_pad_dummy_added_for_rank_whose_requests_are_all_to_complete should also
verify that the dummy request is present in stub.active_requests after _run_pad.
Add an assertion for the resulting active request collection, preserving the
existing add_dummy_calls assertion.
In `@tests/unittest/_torch/executor/test_seq_slot_sizing.py`:
- Around line 1-45: Extend coverage beyond create_torch_sampler_args by
extracting the shared max_num_sequences sizing logic used by
create_torch_sampler_args and create_py_executor_instance into a common helper,
then test that helper with the existing pp_size/disable_overlap/expected_factor
cases. Update both call sites to reuse the helper so the same behavior remains
covered consistently.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c681afc1-5027-4071-a43b-6a047838c13c
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/py_executor.pytests/unittest/_torch/executor/test_py_executor.pytests/unittest/_torch/executor/test_seq_slot_sizing.py
|
/bot run |
|
PR_Github #58778 [ run ] completed with state
|
|
PR_Github #58779 [ run ] triggered by Bot. Commit: |
|
Validation passed. A/B run completed: GB300 DSV4-Pro disagg e2e
|
|
PR_Github #58779 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58785 [ run ] triggered by Bot. Commit: |
|
PR_Github #58785 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58854 [ run ] triggered by Bot. Commit: |
|
PR_Github #58854 [ run ] completed with state
|
…location failure Adopt the fix from PR NVIDIA#16161 (closed in favor of this PR): 1. _count_schedulable_active_requests (disaggregated mode): count only requests inside the scheduler's schedulability window [CONTEXT_INIT, GENERATION_TO_COMPLETE), mirroring no_schedule_until_state / no_schedule_after_state. States outside the window (awaiting KV transfer, DISAGG_CONTEXT_WAIT_SCHEDULER, GENERATION_TO_COMPLETE) cannot produce forward work; counting them suppresses the ADP pad dummy while the rank schedules batch=0, which turns can_queue False fleet-wide and stalls every rank (surfaces as the expected_num_active_requests assert or leaked pad dummies). Non-disaggregated behavior is unchanged. 2. _pad_attention_dp_dummy_request: add_dummy_requests returns None when no cache resources are free for even a 1-token dummy; skip padding for the iteration instead of crashing on an unchecked [0]. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>
…ed buffer sizing
Under the overlap scheduler a finishing request holds its sequence slot
until the next iteration's _process_previous_batch, while the V2
scheduler has already dropped it from the request budget
(no_schedule_after_state=GENERATION_TO_COMPLETE) and backfilled its
seat. Transient slot demand therefore reaches 2 * max_batch_size and
SeqSlotManager raises ValueError("No free slots") (DSV4-Pro disagg GEN
dep32/batch2/mtp3 dies at the first request burst).
Size the pool with overlap headroom via a single-sourced
compute_max_num_sequences (max_batch_size * num_micro_batches: pp_size
with PP, else 2 under overlap), and size every buffer indexed by
py_seq_slot from the same number so slot ids stay in range:
- sampler state (create_torch_sampler_args) and the legacy TRTLLMSampler
decoder state
- guided-decoder per-slot buffers (py_executor_creator; AutoDeploy's
guided decoder aligned to its own pool, which keeps pp_size headroom)
- SpecMetadata.draft_probs for rejection sampling (new num_seq_slots
field, threaded from model_engine)
- eagle3 dynamic-tree DynamicTreeSlotStorage (new max_num_seq_slots,
threaded via get_spec_resource_manager); batch-position work buffers
keep num_trees sizing
V1 scheduler capacity stays max_batch_size * pp_size (V1 retains
GENERATION_TO_COMPLETE in capacity and handles overlap via
two_step_lookahead); the V2 request budget stays max_batch_size.
Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>
4c1b4f8 to
865e9c0
Compare
|
/bot run --post-merge --disable-fail-fast --stage-list "GB300-4_GPUs-PyTorch-Post-Merge-1, GB300-4_GPUs-PyTorch-Post-Merge-2, GB300-4_GPUs-PyTorch-Post-Merge-3, A30-PyTorch-1, A30-PyTorch-2, DGX_H100-4_GPUs-PyTorch-Others-1, DGX_H100-4_GPUs-PyTorch-Others-2, DGX_H100-4_GPUs-PyTorch-DeepSeek-1, DGX_H100-4_GPUs-PyTorch-Post-Merge-1" |
|
PR_Github #59026 [ run ] triggered by Bot. Commit: |
|
Hi @qiaoxj07, Based on the offline discussion, I pushed |
|
Thanks @Tabrizian — I’ve reworked the patch to keep PR #16279 focused on the merge-critical DSv4 path. The Qwen, Mamba, AutoDeploy, generic Eagle/dynamic-tree, PP, and other non-DSv4 behavior has been restored and can be handled in follow-up PRs. The remaining changes are now gated at two deliberately different boundaries:
One thing I'd like to call out: restricting this to DSv4 is more involved than adding a single model check: model identity is established after loading, while |
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
|
By the way, it does not seem like we have |
|
/bot run --post-merge --disable-fail-fast --stage-list "GB300-4_GPUs-PyTorch-Post-Merge-1, GB300-4_GPUs-PyTorch-Post-Merge-2, GB300-4_GPUs-PyTorch-Post-Merge-3, GB200-4_GPUs-PyTorch-1, GB200-4_GPUs-PyTorch-3, DGX_H100-4_GPUs-PyTorch-Others-1, DGX_H100-4_GPUs-PyTorch-Others-2, DGX_H100-4_GPUs-PyTorch-DeepSeek-1" |
|
PR_Github #59051 [ run ] triggered by Bot. Commit: |
chienchunhung
left a comment
chienchunhung
left a comment
There was a problem hiding this comment.
PR looks good to me for DSv4-only scope. A couple of follow-up items should:
- Generalize scheduler eligibility beyond DSv4
- Generalize safe dummy allocation
- Generalize sequence-slot capacity
- Add E2E coverage for "DSv4+ disagg + ADP + MTP Eagle one-model + overlap scheduler"
|
PR_Github #59026 [ run ] completed with state |
|
PR_Github #59051 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #59079 [ run ] triggered by Bot. Commit: |
|
PR_Github #59079 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59136 [ run ] triggered by Bot. Commit: |
|
PR_Github #59136 [ run ] completed with state
|
|
/bot run |
|
PR_Github #59180 [ run ] triggered by Bot. Commit: |
|
PR_Github #59180 [ run ] completed with state |
10 participants
Background
Under the overlap scheduler, a request that reaches
GENERATION_TO_COMPLETE(GtC) vacates its V2-scheduler budget seat one iteration before its sequence slot is freed (_process_previous_batchruns after the next iteration'sprepare_resources). The scheduler backfills the vacated seat with a new request that needs a second concurrent slot, so transient slot demand can reach2 x max_batch_sizewhile theSeqSlotManagerpool holds onlymax_batch_size * pp_size→ValueError("No free slots"). On attention-DP the same over-accumulation family also surfaces as the_pad_attention_dp_dummy_requestassert, and a rank whose only requests sit outside the scheduler window schedules an empty batch without inserting a pad dummy, desyncing collectives;add_dummy_requestscan also legitimately returnNoneunder resource pressure, which the pad path indexed unconditionally.Summary (2 commits)
Commit 1 — ADP pad correctness (adopts closed #16161 by @lingjiew, with thanks):
_count_schedulable_active_requests(disagg path) counts only requests inside the scheduler window[CONTEXT_INIT, GENERATION_TO_COMPLETE), so ranks holding only GtC / WAIT_SCHEDULER requests insert a pad dummy and keep collectives aligned. Non-disaggregated behavior is unchanged._pad_attention_dp_dummy_requesttoleratesadd_dummy_requestsreturningNone(skip padding for one iteration instead of crashing on[0]).Commit 2 — pool headroom with single-sourced sizing:
compute_max_num_sequences()is the one sizing function:max_batch_size * pp_sizewith PP, else2xunder overlap /1xwithout.py_seq_slot: theSeqSlotManagerpool,TorchSamplerstate, legacyTRTLLMSamplerdecoder state, the guided decoder (py_executor_creator),SpecMetadata.draft_probs(newnum_seq_slotsfield, Eagle3/MTP one-model rejection sampling), the Eagle3 dynamic-tree slot storage, and the AutoDeploy guided decoder (aligned to AD's own pool).2048 vs 4096size mismatch (sampler was resized without its consumers) and the Nemotron ADP4_MTP mamba dummy-mask assert (previous revision changed non-disagg dummy accounting; this revision does not).Tests
test_seq_slot_sizing.py: sizing helper + sampler-args consistency (8 cases).test_py_executor.py: pad dummy added for GtC-only / WAIT_SCHEDULER-only disagg ranks; allocation-failure skip (from [https://nvbugs/6438685][fix] Prevent leaked ADP pad dummy from crashing the executor #16161). Existing non-disagg pad regression tests unchanged.Validation (clean
main@ a201a43 image, hot-patched with exactly this diff)test_py_executor/test_benchmark_disagg/test_kv_cache_v2_scheduler/ sizing tests (325 passed). Sampler-suite and eagle3-unittest shards (the previously failing CI stages) running at time of writing; CI on this revision is the arbiter.TestNemotronV3Ultra::test_nvfp4_4gpus_block_reuse[ADP4_MTP]config (NVFP4 550B, tp4/ep4/attention-DP/MTP3/overlap/block-reuse, 96 concurrent staggered requests) completes successfully with this patch.No free slotsbenchmark shape (DSV4 dep32/batch2/mtp3) is a flaky race — an unpatched pass exists — so single benchmark passes are not treated as proof; repeated-run disagg validation is tracked as follow-up. The mechanism itself (GtC budget-vacate vs late slot free) is verified line-by-line, and an equivalent hot-patch previously completed the full 80-minute benchmark on the failing image.Impact
max_batch_sizeto2 x max_batch_sizeentries under overlap without PP; no change with PP or overlap disabled. No scheduler admission behavior change.