[None][fix] Fix context-only async kvtransfer hang - #17107
Conversation
WalkthroughThe Mamba cache manager now allocates the shared CUDA-graph padding slot on demand and reports allocation failure. Disaggregated transfer handling now polls context-only transfers correctly. Unit and integration tests cover cache capacity, dummy requests, and synchronous transfers. ChangesMamba cache and transfer flow
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 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_mamba_cache_manager.py (1)
58-86: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd annotations to the changed test functions.
Annotate all changed test functions and fixture parameters. Use
enable_attention_dp: booland-> Nonefor the unit tests. Add a precise type formonkeypatchand-> Nonefor the integration test.
tests/unittest/_torch/executor/test_mamba_cache_manager.py#L58-L86: annotateenable_attention_dpand both test return types.tests/unittest/_torch/executor/test_mamba_cache_manager.py#L153-L157: add the test return type.tests/integration/defs/accuracy/test_disaggregated_serving.py#L2170-L2170: annotatemonkeypatchand the test return type.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_mamba_cache_manager.py` around lines 58 - 86, Annotate the changed test functions: in tests/unittest/_torch/executor/test_mamba_cache_manager.py lines 58-86, add enable_attention_dp: bool and -> None to both tests; at lines 153-157, add -> None to the test. In tests/integration/defs/accuracy/test_disaggregated_serving.py line 2170, add the precise type for monkeypatch and -> None to the integration test.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/integration/defs/accuracy/test_disaggregated_serving.py`:
- Around line 2178-2179: Update the test setup around the existing monkeypatch
environment configuration before launch_disaggregated_llm() to remove inherited
UCX_NET_DEVICES and set TRTLLM_NIXL_NUM_THREADS to "1"; apply the same
deterministic settings to both server-role setup blocks, including the
corresponding lines around 2211-2219.
---
Nitpick comments:
In `@tests/unittest/_torch/executor/test_mamba_cache_manager.py`:
- Around line 58-86: Annotate the changed test functions: in
tests/unittest/_torch/executor/test_mamba_cache_manager.py lines 58-86, add
enable_attention_dp: bool and -> None to both tests; at lines 153-157, add ->
None to the test. In
tests/integration/defs/accuracy/test_disaggregated_serving.py line 2170, add the
precise type for monkeypatch and -> None to the integration test.
🪄 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: 80a49c2d-a76f-46fa-ad4a-c0caa76f6988
📒 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
5cfad04 to
f9fa76c
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #62982 [ run ] triggered by Bot. Commit: |
|
PR_Github #62982 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
1c44af0 to
e456a8d
Compare
reasonsolo
changed the title
73906bc to
dd62818
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #63401 [ run ] triggered by Bot. Commit: |
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
dd62818 to
2b7ef51
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #63403 [ run ] triggered by Bot. Commit: |
|
PR_Github #63401 [ run ] completed with state |
|
PR_Github #63403 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63444 [ run ] triggered by Bot. Commit: |
|
PR_Github #63444 [ run ] completed with state
|
6 participants
Description
When no requests are scheduled, the PyExecutor forward pass is skipped. Previously, this could also skip checking completed context-side KV transfers, leaving the context request incomplete and its cache resources unreleased. With
max_batch_size=1 and concurrency 2, the second request then remains blocked, causing a deadlock.
This change detects idle iterations from the scheduled batch and polls context transfer status regardless of TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP. The regression test runs context and generation workers on one GPU with batch size
1 and two concurrent Mamba requests.
Test Coverage
PR Checklist
Dev Engineer Review
Falseand allows eager-mode fallback.QA Engineer Review
DISAGG_GENERATION_TRANS_COMPLETE.TestNemotronNano9BV2::test_sync_transfer_bs1_concurrency2.tests/integration/test_lists/test-db/l0_dgx_b200.yml.test-db/orqa/files.Verdict: needs follow-up.