[None][feat] Reserve one more slots for attention_dp in mixed mamba cache manager

[Wanli-Jiang]

Feature

• Will fix out-of-mamba-block error with ADP + max_batch_size=1 + disagg scenario.

Summary by CodeRabbit

• New Features

  • Enhanced attention-DP (distributed pipeline) support with dedicated padding slot reservation in cache management.

• Tests

  • Added test coverage for attention-DP dummy request slot reservation and Mamba cache manager behavior in distributed pipeline scenarios.

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces a dedicated constant ATTENTION_DP_DUMMY_REQUEST_ID for attention-DP padding dummy requests and updates mamba cache management to reserve and isolate a separate slot for such requests, preventing them from competing with real request resource pools. The executor is wired to use the constant, and tests validate the reservation behavior across both standard and disaggregated scenarios.

Changes

Attention-DP dummy slot management

Layer / File(s)	Summary
Attention-DP dummy request ID constant definition tensorrt_llm/_torch/pyexecutor/llm_request.py	Module-level constant ATTENTION_DP_DUMMY_REQUEST_ID = 0 with documentation clarifying its reserved role for padding dummy requests and differentiation from other ID ranges.
Mamba cache manager slot reservation and lifecycle management tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py	Import constant, allocate _attention_dp_dummy_slot when mapping.enable_attention_dp is enabled, route the dummy ID to that reserved slot in add_dummy_requests, exclude it from the free pool in free_resources. Add _get_mamba_hybrid_pool_size(max_batch_size, mapping) helper to compute mamba pool size with attention-DP headroom, and update MixedMambaHybridCacheManager to use it.
Executor integration with dummy request constant tensorrt_llm/_torch/pyexecutor/py_executor.py	Import ATTENTION_DP_DUMMY_REQUEST_ID and refactor _pad_attention_dp_dummy_request() to construct dummy allocations using the constant, passing the ID to KV cache and SPEC resource managers.
Mamba cache manager regression tests for Attention-DP slot reservation tests/unittest/_torch/executor/test_mamba_cache_manager.py	Update _make_mgr helper to accept enable_attention_dp and compute pool size via _get_mamba_hybrid_pool_size with configured Mapping. Add test test_attention_dp_dummy_has_reserved_slot_with_batch_size_one to verify reserved slot behavior with batch size one.
Disaggregated mamba transfer tests with Attention-DP support tests/unittest/disaggregated/test_mamba_transfer.py	Refactor _create_managers to accept max_batch_size and enable_attention_dp, threading them into Mapping and KvCacheConfig. Add test test_mamba_disagg_attention_dp_dummy_with_batch_size_one to verify dummy request handling in disaggregated scenarios.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description lacks the required template sections (Description, Test Coverage, PR Checklist properly filled). Only a brief feature statement and linked issue mention are provided.	Expand the description to include: detailed explanation of the issue and solution, clear list of relevant tests covering the changes, and verification of checklist items per the repository template.
Docstring Coverage	⚠️ Warning	Docstring coverage is 40.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately reflects the main change: reserving additional cache slots for attention_dp in the mamba cache manager, addressing the core feature.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

tests/unittest/disaggregated/test_mamba_transfer.py (1)

350-357: 💤 Low value

Optional: assert reserved-slot stability after free_resources to match the unit test.

This test verifies dummy routing but, unlike test_attention_dp_dummy_has_reserved_slot_with_batch_size_one in test_mamba_cache_manager.py, it doesn't confirm the reserved slot stays out of the free pool after the dummy is freed — the exact lifecycle the disagg + ADP + max_batch_size=1 fix targets.

💚 Optional lifecycle assertion

         assert mgr._impl.mamba_cache_index[100] != mgr._impl._attention_dp_dummy_slot
+
+        mgr._impl.free_resources(SimpleNamespace(py_request_id=ATTENTION_DP_DUMMY_REQUEST_ID))
+        assert mgr._impl._attention_dp_dummy_slot not in mgr._impl.mamba_cache_free_blocks
     finally:
         mgr.shutdown()

This also requires importing SimpleNamespace (from types import SimpleNamespace).

🤖 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 350 - 357,
After adding the dummy request, also verify the reserved dummy slot remains
excluded from the free pool after freeing the dummy: call
mgr.free_resources([ATTENTION_DP_DUMMY_REQUEST_ID]) (or the appropriate free
method on mgr), then assert that mgr._impl.mamba_cache_free_blocks does not
contain mgr._impl._attention_dp_dummy_slot and that mgr._impl.mamba_cache_index
either no longer maps the request id or still maps the id to the reserved slot
as per the other test; also add "from types import SimpleNamespace" to imports
if needed for parity with test_mamba_cache_manager.py.

🤖 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.

Nitpick comments:
In `@tests/unittest/disaggregated/test_mamba_transfer.py`:
- Around line 350-357: After adding the dummy request, also verify the reserved
dummy slot remains excluded from the free pool after freeing the dummy: call
mgr.free_resources([ATTENTION_DP_DUMMY_REQUEST_ID]) (or the appropriate free
method on mgr), then assert that mgr._impl.mamba_cache_free_blocks does not
contain mgr._impl._attention_dp_dummy_slot and that mgr._impl.mamba_cache_index
either no longer maps the request id or still maps the id to the reserved slot
as per the other test; also add "from types import SimpleNamespace" to imports
if needed for parity with test_mamba_cache_manager.py.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: f43be582-9614-457c-8576-0a3c6e0af29d

📥 Commits

Reviewing files that changed from the base of the PR and between 06e3a77 and e8da248.

📒 Files selected for processing (5)

• tensorrt_llm/_torch/pyexecutor/llm_request.py
• tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tests/unittest/_torch/executor/test_mamba_cache_manager.py
• tests/unittest/disaggregated/test_mamba_transfer.py

[Wanli-Jiang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #51542 [ run ] triggered by Bot. Commit: d2e0805 Link to invocation

[tensorrt-cicd]

PR_Github #51542 [ run ] completed with state SUCCESS. Commit: d2e0805
/LLM/main/L0_MergeRequest_PR pipeline #40938 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[Wanli-Jiang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #51690 [ run ] triggered by Bot. Commit: d2e0805 Link to invocation

[tensorrt-cicd]

PR_Github #51690 [ run ] completed with state SUCCESS. Commit: d2e0805
/LLM/main/L0_MergeRequest_PR pipeline #41070 completed with status: 'SUCCESS'

CI Report

Link to invocation

[HuiGao-NV]

LGTM