[None][fix] Fix context-only async kvtransfer hang - #17107
Open
reasonsolo wants to merge 8 commits into
Open
Conversation
reasonsolo
requested review from
SimengLiu-nv,
Tabrizian,
chienchunhung,
eopXD,
joyang-nv,
longlee0622,
ruodil and
yingguo-trt
July 31, 2026 04:41
Contributor
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}
Contributor
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
reasonsolo
force-pushed
the
fix/python-mamba-cache-headroom
branch
from
July 31, 2026 04:48
5cfad04 to
f9fa76c
Compare
reasonsolo
removed request for
a team,
SimengLiu-nv,
eopXD,
joyang-nv,
longlee0622,
ruodil and
yingguo-trt
July 31, 2026 04:51
VALLIS-NERIA
approved these changes
Jul 31, 2026
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62982 [ run ] triggered by Bot. Commit: |
bo-nv
approved these changes
Jul 31, 2026
Collaborator
|
PR_Github #62982 [ run ] completed with state
|
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63015 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63046 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #63046 [ run ] completed with state
|
This was referenced Jul 31, 2026
chienchunhung
left a comment
chienchunhung
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for the PR.
The current unit test validates that the lazy-padding change exposes a second Mamba cache row, but it does not reproduce the deadlock described in the PR or exercise the sync-transfer completion path. Could you add or reinstate an end-to-end BS=1, C=2 regression test and assert request 1’s state after request_and_receive_sync() returns? This is very important because sync receive is expected to set DISAGG_GENERATION_TRANS_COMPLETE itself; if not, we might need to address that bigger issue first.
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
tests/unittest/disaggregated/test_mamba_transfer.py (1)
363-378: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert successful KV-cache allocation in the hybrid path
Changed tests:
- Modified:
test_mamba_disagg_attention_dp_dummy_with_batch_size_one- Added:
test_mamba_disagg_padding_falls_back_when_headroom_is_busy- Removed: none
The fallback test covers the
Nonereturn before KV-cache allocation. The attention-DP test covers Mamba insertion but ignores the KV-cache result. Assert thatmgr.add_dummy_requests(...) is not None, or mockKVCacheManager.add_dummy_requests().The module is listed in
tests/integration/test_lists/test-db/l0_a10.yml. QA lists target integration tests and do not require this unittest. Coverage verdict: insufficient for the full success/failure contract.🤖 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/disaggregated/test_mamba_transfer.py` around lines 363 - 378, Update the successful hybrid-path test around add_dummy_requests to assert that KV-cache allocation succeeds by verifying mgr.add_dummy_requests(...) returns a non-None result, while retaining the existing fallback test’s None assertion for busy headroom. Locate the affected attention-DP test and ensure both success and failure allocation outcomes are covered.Source: Path instructions
tests/unittest/disaggregated/test_py_cache_transceiver_mp.py (1)
902-907: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a diagnostic message to the completion assertion.
request_and_receive_syncsetsrequest.statetoDISAGG_TRANS_ERRORon failure, notDISAGG_GENERATION_TRANS_COMPLETE. If that path triggers, the bareassertraisesAssertionErrorwith no state information. This test already relies onprint(..., flush=True)throughout for multiprocess debugging, so include the actual state in the assertion message to speed up triage of CI failures.♻️ Suggested diff
def _request_and_receive_sync_complete( transceiver: "KvCacheTransceiverV2", request: LlmRequest, ) -> None: transceiver.request_and_receive_sync(request) - assert request.state == LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE + assert request.state == LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE, ( + f"Expected DISAGG_GENERATION_TRANS_COMPLETE, got {request.state!r}" + )🤖 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/disaggregated/test_py_cache_transceiver_mp.py` around lines 902 - 907, Update the completion assertion in _request_and_receive_sync_complete to include request.state in its failure message, while preserving the existing DISAGG_GENERATION_TRANS_COMPLETE expectation and diagnostic print/flush conventions.
🤖 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 `@tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py`:
- Around line 1145-1146: Update MambaCacheManager.add_dummy_requests to return
the boolean result from _impl.add_dummy_requests instead of discarding it, so
the caller’s success check can proceed to KVCacheManager.add_dummy_requests when
appropriate. Add explicit return annotations to this forwarding method and the
corresponding forwarding methods involved in the same delegation path,
preserving the delegated result semantics.
- Around line 593-606: Update add_dummy_requests to preflight the total
free-block requirements for all new ordinary dummies and any new
CUDA_GRAPH_DUMMY_REQUEST_ID sentinel before mutating the pool, rejecting the
entire request list when capacity is insufficient. Alternatively, process the
sentinel allocation first, but ensure no partial state remains when allocation
fails; add a regression test covering an ordinary dummy before a new sentinel
with one free block.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 3493-3507: Update the worker-role decision around
context_only_worker so participation in _sync_disagg_gen_status_entry is
rank-uniform before any synchronous-mode collective. Derive or exchange a
rank-consistent context-only/participation value, while preserving the rule that
synchronous GEN workers do not enter these collectives; use that decision for
the subsequent local_need_gen_check flow.
---
Nitpick comments:
In `@tests/unittest/disaggregated/test_mamba_transfer.py`:
- Around line 363-378: Update the successful hybrid-path test around
add_dummy_requests to assert that KV-cache allocation succeeds by verifying
mgr.add_dummy_requests(...) returns a non-None result, while retaining the
existing fallback test’s None assertion for busy headroom. Locate the affected
attention-DP test and ensure both success and failure allocation outcomes are
covered.
In `@tests/unittest/disaggregated/test_py_cache_transceiver_mp.py`:
- Around line 902-907: Update the completion assertion in
_request_and_receive_sync_complete to include request.state in its failure
message, while preserving the existing DISAGG_GENERATION_TRANS_COMPLETE
expectation and diagnostic print/flush 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: 558b9144-e05d-417a-8db3-571dd3b5ec44
📒 Files selected for processing (8)
tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pytensorrt_llm/_torch/pyexecutor/py_executor.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/_torch/executor/test_py_executor.pytests/unittest/disaggregated/test_mamba_transfer.pytests/unittest/disaggregated/test_py_cache_transceiver_mp.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/unittest/_torch/executor/test_mamba_cache_manager.py
Comment on lines
+593
to
+606
| def add_dummy_requests(self, request_ids: List[int], **kwargs) -> bool: | ||
| # Sentinels alias to the shared _padding_slot; non-sentinel | ||
| # dummies (warmup, attention-DP idle padding) get their own | ||
| # slot and are freed individually. | ||
| if not request_ids: | ||
| return | ||
| return True | ||
| needs_padding_slot = (self._padding_slot is None and any( | ||
| self._is_padding_sentinel(request_id) | ||
| and request_id not in self.mamba_cache_index | ||
| for request_id in request_ids)) | ||
| if needs_padding_slot and not self.mamba_cache_free_blocks: | ||
| # Let CUDAGraphRunner fall back to eager execution when the lazy | ||
| # padding headroom is temporarily occupied by a real request. | ||
| return False |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
rg -n -C 8 --glob '*.py' \
'add_dummy_requests\s*\(|CUDA_GRAPH_DUMMY_REQUEST_ID|ATTENTION_DP_DUMMY_REQUEST_ID' \
tensorrt_llm tests || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
file="tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py"
printf '%s\n' '--- target method ---'
sed -n '540,635p' "$file"
printf '%s\n' '--- related symbols and call sites ---'
rg -n -C 5 \
'def (_is_padding_sentinel|add_dummy_requests)|mamba_cache_free_blocks|mamba_cache_index|_padding_slot|ATTENTION_DP_DUMMY_REQUEST_ID|CUDA_GRAPH_DUMMY_REQUEST_ID' \
"$file" tensorrt_llm/_torch/auto_deploy tensorrt_llm/_torch/pyexecutor tests/unittest/_torch \
--glob '*.py' | head -n 500
printf '%s\n' '--- candidate tests ---'
rg -l \
'mamba_cache_manager|BaseMambaCacheManager|add_dummy_requests' \
tests/unittest/_torch --glob '*.py' | head -n 100Repository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
printf '%s\n' '--- all direct callers of the Mamba add_dummy_requests path ---'
rg -n -C 4 \
'add_dummy_requests\(' \
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py \
tensorrt_llm/_torch/pyexecutor/py_executor.py \
tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py \
tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py \
tests/unittest/_torch --glob '*.py' | head -n 700
printf '%s\n' '--- Mamba-specific tests and fixtures ---'
rg -l \
'MambaCacheManager|MixedMambaHybridCacheManager|mamba_cache_index|padding_slot|ATTENTION_DP_DUMMY_REQUEST_ID' \
tests --glob '*.py' | head -n 100
printf '%s\n' '--- deterministic reproduction of the allocation order ---'
python3 - <<'PY'
CUDA_GRAPH_DUMMY_REQUEST_ID = (1 << 64) - 1
ATTENTION_DP_DUMMY_REQUEST_ID = -1
max_draft_len = 0
request_ids = [123, CUDA_GRAPH_DUMMY_REQUEST_ID]
free_blocks = [7]
mamba_cache_index = {}
padding_slot = None
dummy_request_ids = set()
def is_padding_sentinel(request_id):
return (CUDA_GRAPH_DUMMY_REQUEST_ID - max_draft_len
<= request_id <= CUDA_GRAPH_DUMMY_REQUEST_ID)
needs_padding_slot = (
padding_slot is None and any(
is_padding_sentinel(request_id)
and request_id not in mamba_cache_index
for request_id in request_ids
)
)
precheck_passes = not (needs_padding_slot and not free_blocks)
try:
if precheck_passes:
dummy_request_ids.update(request_ids)
for request_id in request_ids:
if request_id in mamba_cache_index:
continue
if is_padding_sentinel(request_id):
if padding_slot is None:
if len(free_blocks) == 0:
raise RuntimeError("run out of mamba cache blocks")
padding_slot = free_blocks.pop()
block = padding_slot
elif (request_id == ATTENTION_DP_DUMMY_REQUEST_ID
and False):
block = None
else:
if len(free_blocks) == 0:
raise RuntimeError("run out of mamba cache blocks")
block = free_blocks.pop()
mamba_cache_index[request_id] = block
except RuntimeError as error:
print("precheck_passes:", precheck_passes)
print("raised:", str(error))
print("free_blocks:", free_blocks)
print("mamba_cache_index:", mamba_cache_index)
print("dummy_request_ids:", sorted(dummy_request_ids))
print("padding_slot:", padding_slot)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
python3 - <<'PY'
CUDA_GRAPH_DUMMY_REQUEST_ID = (1 << 64) - 1
ATTENTION_DP_DUMMY_REQUEST_ID = -1
request_ids = [123, CUDA_GRAPH_DUMMY_REQUEST_ID]
max_draft_len = 0
free_blocks = [0]
mamba_cache_index = {}
padding_slot = None
dummy_request_ids = set()
def is_padding_sentinel(request_id):
return (CUDA_GRAPH_DUMMY_REQUEST_ID - max_draft_len
<= request_id <= CUDA_GRAPH_DUMMY_REQUEST_ID)
needs_padding_slot = (
padding_slot is None and any(
is_padding_sentinel(request_id)
and request_id not in mamba_cache_index
for request_id in request_ids
)
)
assert needs_padding_slot
assert free_blocks, "The precheck must pass with one free block."
try:
dummy_request_ids.update(request_ids)
for request_id in request_ids:
if request_id in mamba_cache_index:
continue
if is_padding_sentinel(request_id):
if padding_slot is None:
if not free_blocks:
raise RuntimeError("run out of mamba cache blocks")
padding_slot = free_blocks.pop()
block = padding_slot
elif (request_id == ATTENTION_DP_DUMMY_REQUEST_ID
and False):
block = None
else:
if not free_blocks:
raise RuntimeError("run out of mamba cache blocks")
block = free_blocks.pop()
mamba_cache_index[request_id] = block
except RuntimeError as error:
print({
"precheck_passes": True,
"raised": str(error),
"free_blocks": free_blocks,
"mamba_cache_index": mamba_cache_index,
"dummy_request_ids": sorted(dummy_request_ids),
"padding_slot": padding_slot,
})
else:
raise AssertionError("Expected the mixed list to exhaust the pool.")
PY
printf '%s\n' '--- attention-DP caller context ---'
sed -n '5748,5822p' tensorrt_llm/_torch/pyexecutor/py_executor.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 3523
Preflight all new dummy allocations before mutating the pool.
When an ordinary dummy precedes a new CUDA_GRAPH_DUMMY_REQUEST_ID and one free block remains, the precheck passes. The ordinary dummy consumes the block, then the sentinel raises RuntimeError and leaves partial state. Count required blocks before mutation or allocate the sentinel first. Add a mixed-list regression test.
🤖 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 593 -
606, Update add_dummy_requests to preflight the total free-block requirements
for all new ordinary dummies and any new CUDA_GRAPH_DUMMY_REQUEST_ID sentinel
before mutating the pool, rejecting the entire request list when capacity is
insufficient. Alternatively, process the sentinel allocation first, but ensure
no partial state remains when allocation fails; add a regression test covering
an ordinary dummy before a new sentinel with one free block.
reasonsolo
force-pushed
the
fix/python-mamba-cache-headroom
branch
from
August 3, 2026 03:34
658ac8f to
6f44b21
Compare
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
reasonsolo
force-pushed
the
fix/python-mamba-cache-headroom
branch
from
August 3, 2026 04:52
1c44af0 to
e456a8d
Compare
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
reasonsolo
requested review from
BowenFu,
QiJune,
StanleySun639,
jieli-matrix,
yiqingy0 and
yuanjingx87
August 3, 2026 06:33
reasonsolo
changed the title
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
With
max_batch_size=1, the Python Mamba cache manager allocatesBS + 1recurrent-state slots, but previously reserved the extra slot permanently for CUDA graph padding. That left only one row for real requests. In disaggregated serving, request 1 could retain that row during KV-cache transfer and prevent request 2 from being scheduled.Align the Python manager with the C++ behavior by allocating the shared padding row lazily. If all
BS + 1rows are temporarily occupied by real requests when padding is needed, fall back from CUDA graph execution without mutating dummy-request bookkeeping.The BS=1, concurrency=2 regression also exposed a context-side progress bug when
TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1: the context worker skipped idle transfer-status checks even though only synchronous generation receivers need to avoid those collectives. Make the gate role-aware so context-only workers reap completed sends and release their cache rows while generation workers retain the synchronous safety behavior.No public API or configuration changes.
Test Coverage
BS + 1real rows, CUDA graph fallback while headroom is occupied, lazy padding allocation, and attention-DP dummy allocation.DISAGG_GENERATION_TRANS_COMPLETE.BS=1, concurrency 2, synchronous transfer on both roles, and one NIXL progress thread.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.