[None][fix] reserve draft KV at disagg transition for gen-only-dspark case - #16990
Conversation
e5630ec to
912580d
Compare
|
/bot run --disable-fail-fast |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (4)
WalkthroughThe changes update speculative-decoding draft reservation and attention-DP dummy sizing. Capacity checks now include engine, sequence, and KV-cache limits. Unit tests cover disaggregated generation transitions and disabled-dsv4 dummy allocation. ChangesCapacity accounting
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 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: 2
🤖 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_kv_cache_v2_capacity_only.py`:
- Around line 197-198: Update the expected value in the test assertion for
manager._required_gen_capacity(request, 128) from 137 to 133, while preserving
the existing _effective_draft_len assertion.
In `@tests/unittest/_torch/executor/test_py_executor.py`:
- Around line 1507-1517: Add test-list entries for
test_disabled_dsv4_gate_checks_full_generation_capacity to the appropriate
tests/integration/test_lists/test-db and tests/integration/test_lists/qa files,
using the existing format and suite placement conventions.
🪄 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: 310c0807-d825-4df4-b337-900e7d0b216f
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/_torch/pyexecutor/py_executor.pytests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.pytests/unittest/_torch/executor/test_py_executor.py
|
PR_Github #62461 [ run ] triggered by Bot. Commit: |
|
PR_Github #62461 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62671 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #62695 [ run ] triggered by Bot. Commit: |
|
PR_Github #62671 [ run ] completed with state |
|
[by Codex] @thorjohnsen Friendly review reminder: this PR is awaiting your review. Thanks! |
|
/bot run --disable-fail-fast |
|
PR_Github #62755 [ run ] triggered by Bot. Commit: |
|
PR_Github #62695 [ run ] completed with state |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py (1)
602-626: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winKeep
get_max_resource_count()aligned with real-request capacity.
PyExecutoruses this method only as a fallback forget_kv_cache_capacity(). Scheduler admission uses live cache checks. After_padding_slotis claimed, the method still overreports the documented maximum real-request count by one. Return_max_resource_count - 1when_padding_slotis set.🤖 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/mamba_cache_manager.py` around lines 602 - 626, Update get_max_resource_count() to return _max_resource_count - 1 once _padding_slot has been claimed, while retaining _max_resource_count when it remains unset. Keep the result aligned with the maximum real-request capacity and preserve the existing lazy padding-slot behavior.
🧹 Nitpick comments (1)
tests/integration/defs/accuracy/test_disaggregated_serving.py (1)
2216-2284: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueTest coverage summary for
tests/integration/defs/accuracy/test_disaggregated_serving.py.
- Test function added:
TestNemotronNano9BV2::test_sync_transfer_bs1_concurrency2(Line 2216-2284). It launches a 1-context/1-generation disaggregated deployment withmax_batch_size=1, overlap scheduler enabled,TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1(synchronous cache transfer), NIXL Python-runtime cache transceiver, and float16 Mamba SSM cache, then submits two concurrent prompts and asserts each returns exactly one output.- Listed in test list files: confirmed present in
tests/integration/test_lists/test-db/l0_dgx_b200.ymlasaccuracy/test_disaggregated_serving.py::TestNemotronNano9BV2::test_sync_transfer_bs1_concurrency2.- Coverage verdict: sufficient for its stated purpose (regression coverage for the back-to-back batch-size-1 Mamba cache-holding scenario under synchronous transfer), though it is a structural/crash-regression check rather than an output-quality check: the assertion (
len(output.outputs) == 1) does not verify generated text content. Since the test lives inaccuracy/test_disaggregated_serving.py, consider a lightweight content sanity check (e.g., non-empty text) if intended to also guard against silent corruption of recurrent state, not just crashes.As per path instructions, "Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM. Always produce a test coverage summary, even if no issues are found," and confirming listing via "Whether each changed test is listed in the appropriate test list files under tests/integration/test_lists/ (test-db/ for CI, qa/ for manual QA)."
🤖 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/integration/defs/accuracy/test_disaggregated_serving.py` around lines 2216 - 2284, Strengthen TestNemotronNano9BV2::test_sync_transfer_bs1_concurrency2 with a lightweight output-content assertion in addition to checking one output per response. Verify each generated output contains non-empty text while preserving the existing concurrency, configuration, and response-count checks.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.
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py`:
- Around line 602-626: Update get_max_resource_count() to return
_max_resource_count - 1 once _padding_slot has been claimed, while retaining
_max_resource_count when it remains unset. Keep the result aligned with the
maximum real-request capacity and preserve the existing lazy padding-slot
behavior.
---
Nitpick comments:
In `@tests/integration/defs/accuracy/test_disaggregated_serving.py`:
- Around line 2216-2284: Strengthen
TestNemotronNano9BV2::test_sync_transfer_bs1_concurrency2 with a lightweight
output-content assertion in addition to checking one output per response. Verify
each generated output contains non-empty text while preserving the existing
concurrency, configuration, and response-count checks.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a74738ea-d7cd-4524-887f-24d6c4b53cbc
📒 Files selected for processing (5)
tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pytests/integration/defs/accuracy/test_disaggregated_serving.pytests/integration/test_lists/test-db/l0_dgx_b200.ymltests/unittest/_torch/executor/test_mamba_cache_manager.pytests/unittest/disaggregated/test_mamba_transfer.py
6bb5556 to
5141693
Compare
|
PR_Github #62936 [ 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>
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
5141693 to
d463494
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #62955 [ run ] triggered by Bot. Commit: |
|
PR_Github #62955 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63040 [ run ] triggered by Bot. Commit: |
|
PR_Github #63040 [ run ] completed with state |
… case (#16990) Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
7 participants
Dev Engineer Review
KVCacheManagerV2._effective_draft_lennow reserves transferred context draft tokens and falls back to configured draft capacity for speculative decoding._pad_attention_dp_dummy_requestnow checks engine, sequence, and KV-cache capacity before allocating a context dummy request.QA Engineer Review
test_disagg_gen_reserves_target_draft_and_required_capacity.test_disagg_gen_disables_speculative_decoding.test_disagg_gen_prefers_context_phase_draft_tokens.test_disabled_dsv4_gate_checks_full_generation_capacity._StubADPExecutor.__init__with capacity-related parameters.test-db/coverage.test_disabled_dsv4_gate_checks_full_generation_capacitydefinition causes the earlier test to be overwritten. Rename or remove one definition.Description
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.