[None][fix] Fix gen-only async kvtransfer hang - #17107
Conversation
WalkthroughThe executor now receives scheduled-batch idle state for disaggregated transfer checks. Integration coverage adds separate worker environments, optional GPU sharing, and a concurrent Mamba test with its YAML configuration and QA registration. ChangesDisaggregated transfer and Mamba concurrency flow
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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_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 |
|
PR_Github #63444 [ run ] triggered by Bot. Commit: |
|
PR_Github #63444 [ run ] completed with state
|
| @@ -3533,9 +3534,12 @@ def _check_disagg_transfer_progress_when_idle( | |||
| self._check_disagg_gen_cache_transfer_status(1) | |||
| return | |||
|
|
|||
| any_need_check = self._sync_disagg_ctx_status_entry(local_need_check) | |||
| local_need_ctx_check = is_idle or (uses_async_gen_transfer | |||
| and local_needs_progress) | |||
| any_need_check = self._sync_disagg_ctx_status_entry( | |||
| local_need_ctx_check) | |||
| if any_need_check > 0: | |||
| if local_need_check and not all_gen_first: | |||
| if local_need_ctx_check and not all_gen_first: | |||
| logger.warning( | |||
| "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache" | |||
| ) | |||
There was a problem hiding this comment.
I have one concern about the multi-rank synchronous-transfer path.
With overlap disabled, ranks can be at different points: one rank may still be blocked in request_and_receive_sync(), while another rank has no transfer to process and reaches this allreduce. Because every rank must enter collectives in the same order, this could stall or misalign the ranks.
#15737 avoided these progress collectives for synchronous GEN transfer. Could we preserve that behavior for GEN workers while separately allowing CTX workers to poll their asynchronous sends, using a worker-role decision that is consistent across ranks?
2b7ef51 to
30fdacf
Compare
reasonsolo
changed the title
|
/bot run --disable-fail-fast |
|
PR_Github #63664 [ run ] triggered by Bot. Commit: |
|
PR_Github #63664 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63717 [ run ] triggered by Bot. Commit: |
|
PR_Github #63717 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63762 [ run ] triggered by Bot. Commit: |
|
PR_Github #63762 [ run ] completed with state
|
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
30fdacf to
734773c
Compare
|
/bot run --disable-fail-fast |
|
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. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/disaggregated/test_disaggregated.py`:
- Line 1108: Update test_disaggregated_mamba_bs1_concurrency2 to set
TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP to "1" for the regression scenario
instead of removing it, while preserving the two concurrent completion requests.
Ensure the test explicitly exercises the disabled-overlap branch.
🪄 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: a9b34bd7-eb65-4ea9-909e-3ebe85dcf495
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/py_executor.pytests/integration/defs/disaggregated/test_configs/disagg_config_mamba_bs1_concurrency2.yamltests/integration/defs/disaggregated/test_disaggregated.pytests/integration/test_lists/qa/llm_function_core.txt
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/_torch/pyexecutor/py_executor.py
|
PR_Github #63880 [ run ] triggered by Bot. Commit: |
|
PR_Github #63880 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64471 [ run ] triggered by Bot. Commit: |
|
PR_Github #64471 [ run ] completed with state
|
8 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
PyExecutorto poll completed context-side KV transfers during idle iterations.mamba_bs1_concurrency2QA entry.QA Engineer Review
test_disaggregated_mamba_bs1_concurrency2.tests/integration/test_lists/qa/llm_function_core.txt.