[None][fix] fix disagg overlap slot headroom without MTP - #17282
[None][fix] fix disagg overlap slot headroom without MTP#17282reasonsolo wants to merge 3 commits into
Conversation
|
/bot run --disable-fail-fast |
WalkthroughSequence-slot overlap headroom now applies to disaggregated attention-DP deployments. Eligibility uses topology, cache-transceiver configuration, pipeline parallelism, and overlap scheduling. Unit and integration tests cover the updated behavior. ChangesDisaggregated attention-DP overlap headroom
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant ModelEngine
participant HeadroomUtility
participant SequenceSlotSizing
ModelEngine->>HeadroomUtility: evaluate attention-DP, disaggregation, pipeline, and overlap settings
HeadroomUtility-->>ModelEngine: return headroom eligibility
ModelEngine->>SequenceSlotSizing: calculate maximum sequence slots
SequenceSlotSizing-->>ModelEngine: return slot capacity
🚥 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 (1)
tests/unittest/_torch/executor/test_seq_slot_sizing.py (1)
50-52: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd annotations to the modified test function.
test_dsv4_overlap_headroom_gatehas untyped parameters and no return annotation. Add precise parameter annotations and-> None.Proposed change
def test_dsv4_overlap_headroom_gate( - model_type, pp_size, disable_overlap, expected -) : + model_type: str, + pp_size: int, + disable_overlap: bool, + expected: bool, +) -> None:As per coding guidelines, annotate every function.
🤖 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 50 - 52, Update the test_dsv4_overlap_headroom_gate function signature with precise annotations for model_type, pp_size, disable_overlap, and expected, and add a -> None return annotation. Preserve the existing parameterization and test behavior.Source: Coding guidelines
🤖 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 42-56: Complete the coverage follow-up for
test_dsv4_overlap_headroom_gate by running pytest tests/unittest/ and reporting
the results; preserve the existing parameterized cases and test behavior.
---
Nitpick comments:
In `@tests/unittest/_torch/executor/test_seq_slot_sizing.py`:
- Around line 50-52: Update the test_dsv4_overlap_headroom_gate function
signature with precise annotations for model_type, pp_size, disable_overlap, and
expected, and add a -> None return annotation. Preserve the existing
parameterization and test behavior.
🪄 Autofix
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: 5a15a08d-d92f-4ad4-b6eb-603d9dd513b8
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/model_engine.pytests/unittest/_torch/executor/test_seq_slot_sizing.py
|
PR_Github #63941 [ run ] triggered by Bot. Commit: |
Signed-off-by: Lizhi Zhou <lizhiz@oci-aga-slurm-1-login-01.cm.cluster>
d55d013 to
e2f7e71
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
/bot run --disable-fail-fast |
|
PR_Github #63947 [ run ] triggered by Bot. Commit: |
|
PR_Github #63941 [ run ] completed with state |
|
PR_Github #63947 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64021 [ run ] triggered by Bot. Commit: |
|
PR_Github #64021 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64041 [ run ] triggered by Bot. Commit: |
|
PR_Github #64041 [ run ] completed with state
|
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
reasonsolo
changed the title
reasonsolo
changed the title
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #64285 [ run ] triggered by Bot. Commit: |
ZhanruiSunCh
left a comment
ZhanruiSunCh
left a comment
There was a problem hiding this comment.
LGTM for infra part.
| def should_enable_disagg_adp_overlap_headroom( | ||
| mapping: Mapping, | ||
| cache_transceiver_config: Optional[CacheTransceiverConfig], | ||
| disable_overlap_scheduler: bool) -> bool: | ||
| """Gate extra sequence slots to non-PP disaggregated attention-DP.""" | ||
| is_disagg = (cache_transceiver_config is not None | ||
| and cache_transceiver_config.backend is not None) | ||
| return (mapping.enable_attention_dp and is_disagg and not mapping.has_pp() |
There was a problem hiding this comment.
Agreed that the MTP conjunct in #16279 was over-narrow — draft length has nothing to do with slot lifetime. But swapping it for enable_attention_dp and is_disagg narrows in a new direction, and two configurations that #16279 deliberately protected lose the headroom here:
- DeepSeek-V4 + MTP-one-model + non-PP + overlap, running aggregated (no
cache_transceiver_config) - the same, disaggregated but with
enable_attention_dp=False
The old predicate required neither disagg nor ADP (model_type == "deepseek_v4" and not mapping.has_pp() + MTP + overlap), so both go from 2 * max_batch_size back to max_batch_size.
That matters because the race isn't disagg-specific. #16279 describes it purely as a V2-scheduler/overlap lifetime issue, and the code agrees: V2 excludes GENERATION_TO_COMPLETE from its schedulable range (scheduler_v2.py:151,357), overlap moves the completing request into that state (py_executor.py:6272), the backfilled request takes a slot in generic prepare_resources (py_executor.py:4485, seq_slot_manager.py:17), and only afterwards does _process_previous_batch free the old one (py_executor.py:4651). No transceiver anywhere on that path. Manager selection keys off use_kv_cache_manager_v2 (_util.py:81), not disaggregation, so an aggregated DSV4/MTP/V2 deployment can hit exactly the same ValueError("No free slots").
The mirror image is that the new gate doesn't test for V2 either, so V1 disagg ADP deployments now pay 2x for a race V1 doesn't have.
Could the gate key off the thing the mechanism actually depends on — V2 + overlap + non-PP — rather than topology? kv_cache_config.use_kv_cache_manager_v2 is available at the same point in __init__ as cache_transceiver_config, so it doesn't reintroduce the load-order problem that made you move this ahead of the model load. Failing that, OR-ing the old DSV4/MTP condition back in would at least keep this from being a narrowing.
Not free either: 2x propagates into TorchSampler state, the legacy TRTLLMSampler decoder state (sampler.py:4876,4886,4900), the guided decoder, and SpecMetadata. Greedy non-spec is negligible (~0.3 MiB at B=256), but beam search adds ~12 * B * beam_width * max_seq_len bytes (sampler.py:2284) and spec rejection sampling adds 4 * B * max_draft_len * vocab_size for draft_probs (interface.py:576,590) — both in the hundreds of MiB per rank at B=256. Worth a line in the description, since disagg ADP + overlap is a very common serving shape and this is now on by default for all of it.
|
PR_Github #64285 [ run ] completed with state
|
10 participants
Dev Engineer Review
2 * max_batch_sizesequence-slot headroom only when attention-DP, disaggregation, non-pipeline execution, and overlap scheduling are enabled._enable_disagg_adp_overlap_headroomflag.QA Engineer Review
tests/integration/test_lists/waives.txt.Description
With overlap scheduling, requests completed in the previous iteration still hold their sequence slots when the next iteration calls prepare_resources(). Meanwhile, the scheduler has already removed those terminal requests from its capacity budget and may schedule replacement requests or ADP dummy requests.
This creates a transient requirement of up to 2 * max_batch_size sequence slots.
DeepSeek-V4 already had this overlap headroom, but it was incorrectly restricted to one-model MTP configurations. With draft length 0, the slot pool remained at max_batch_size, allowing dummy allocation to fail when the pool was saturated.
Once some rank event loops exited with NoFreeSlotsError, the remaining ranks blocked in distributed collectives, presenting externally as a server hang.
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.