[None][fix] Fix gen-only async kvtransfer hang - #17107
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
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 Sequence Diagram(s)sequenceDiagram
participant Executor
participant TransferProgressHelper
participant ContextTransfer
Executor->>TransferProgressHelper: pass scheduled-batch idle state
TransferProgressHelper->>ContextTransfer: poll context-transfer progress during idle iteration
ContextTransfer-->>TransferProgressHelper: report completed transfers
TransferProgressHelper-->>Executor: return transfer-progress result
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 #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
|
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
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/unittest/_torch/executor/test_py_executor.py`:
- Around line 692-699: Register the three changed PyExecutor tests, including
test_sync_benchmark_polls_context_transfer_when_idle, in the appropriate test
list under tests/integration/test_lists/test-db/ or
tests/integration/test_lists/qa/. Then run pytest tests/unittest/ and record its
result; leave coverage follow-up unchanged.
🪄 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: 99fdc52a-4be3-4423-a6bf-79e3a4b6219a
📒 Files selected for processing (2)
tests/unittest/_torch/executor/test_benchmark_disagg.pytests/unittest/_torch/executor/test_py_executor.py
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #64542 [ run ] triggered by Bot. Commit: |
|
PR_Github #64542 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64593 [ run ] triggered by Bot. Commit: |
9 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.