[None][fix] CppMambaHybridCacheManager: handle ranks with zero local mamba layers

[VALLIS-NERIA]

Summary

Under PP sharding, a rank can end up with only full-attention layers and no mamba layers. The previous CppMambaHybridCacheManager.__init__ unconditionally allocated SSM/conv state and set up linear-attention metadata, leaving these ranks in an invalid state.

This PR adds an early-exit path after super().__init__() when local_num_mamba_layers == 0:

• Forward num_layers (not mamba_num_layers + num_layers) and the union of mamba+attention layer masks to the parent KVCacheManager.
• Skip all mamba-only setup (ssm/conv_state_shape, linear_attention_metadata, MTP intermediate states, replay buffers, layer offsets, etc.).
• Initialize self.requests = [] before the early return so the guards below are safe.

Guard prepare_resources / update_mamba_states / _setup_state_indices with the same condition so they no-op on these ranks (the parent's prepare_resources still runs).

Test plan

• New unit test test_cpp_hybrid_zero_local_mamba_layers in tests/unittest/_torch/executor/test_mamba_cache_manager.py:
  • Constructs CppMambaHybridCacheManager with the real parent KVCacheManager (no mocks) on world_size=1 with an all-attention layer mask — same Python branch as a PP rank with zero local mamba layers.
  • Verifies early-exit invariants: local_num_mamba_layers == 0, mamba_pp_layers == [], requests == [], pp_layers == [0,1,2,3], parent self.impl / blocks_per_window initialized, is_linear_attention is False, no mamba-only attributes set.
  • Calls the three guarded methods on the zero-mamba manager to confirm they no-op without raising on the uninitialized mamba state.
• Existing test_mamba_cache_manager.py suite still passes (10/10).

Summary by CodeRabbit

• Refactor

  • Optimized Mamba cache manager initialization to skip unnecessary state setup when ranks have no local Mamba layers, improving efficiency in distributed configurations.

• Tests

  • Added test coverage for Mamba cache manager edge cases in distributed tensor-parallel scenarios.

[VALLIS-NERIA]

/bot run

[coderabbitai]

📝 Walkthrough

Walkthrough

CppMambaHybridCacheManager now detects when the current rank has zero local mamba layers during __init__, normalizes the layer mask and PP-sharding upfront, and short-circuits initialization and state-handling methods to skip unneeded mamba machinery on those ranks. Test coverage validates the early-exit path.

Changes

Zero local mamba layers optimization

Layer / File(s)	Summary
Early-exit initialization logic tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py	__init__ normalizes the combined layer_mask, computes pp_layers before KV backend initialization, derives mamba_pp_layers and local_num_mamba_layers, and returns early after base KVCacheManager construction when local_num_mamba_layers == 0, skipping mamba state setup.
Attention window configuration tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py	After kv_cache_config.max_attention_window is set, per-layer attention window types and per_layer_kv_heads are derived using the already-normalized layer_mask, including zeroing KV heads on mamba layers.
State-handling method short-circuits tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py	prepare_resources, update_mamba_states, and _setup_state_indices return immediately when local_num_mamba_layers == 0.
Test coverage for zero-layers edge case tests/unittest/_torch/executor/test_mamba_cache_manager.py	Module docstring expanded to document CppMambaHybridCacheManager PP-sharding edge cases; new imports for cache types and configuration; _build_zero_mamba_hybrid() helper constructs a manager with zero local layers; test_cpp_hybrid_zero_local_mamba_layers() verifies early-exit invariants, absence of mamba-only attributes, and no-op behavior of mamba methods.

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main fix: handling the edge case of zero local mamba layers in CppMambaHybridCacheManager, which is the core change across the modified files.
Description check	✅ Passed	The description covers the problem, solution approach, and test coverage. It explains the early-exit path, what is skipped, how methods are guarded, and provides a comprehensive test plan with verification steps.
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

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

• Generate code and open pull requests
• Plan features and break down work
• Investigate incidents and troubleshoot customer tickets together
• Automate recurring tasks and respond to alerts with triggers
• Summarize progress and report instantly

Built for teams:

• Shared memory across your entire org—no repeating context
• Per-thread sandboxes to safely plan and execute work
• Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get started

---

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

[coderabbitai]

Actionable comments posted: 2

🤖 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 1000-1015: The constructor reads and copies layer_mask
unconditionally which breaks when layer_mask is None; before calling .copy() or
indexing, compute total_layers = len(mamba_layer_mask) and if layer_mask is None
set full_attention_layer_mask = [False] * total_layers else set
full_attention_layer_mask = layer_mask.copy(), then build the combined
layer_mask = [mamba_layer_mask[i] or full_attention_layer_mask[i] for i in
range(total_layers)]; update the code paths around the existing variables
(layer_mask, mamba_layer_mask, full_attention_layer_mask) to use this guarded
logic.
- Around line 1029-1053: Seed the publicly visible mamba-related fields before
the zero-layer early return so accessor calls don't raise AttributeError: assign
self._use_replay_state_update = use_replay_state_update and initialize the
backing attribute(s) used by get_mamba_ssm_cache_dtype (e.g.
self._mamba_ssm_cache_dtype or whatever private field that method reads) as well
as any other public/externally accessed mamba fields that are normally set
later; do this immediately after setting self.local_num_mamba_layers and
self.requests and before the if self.local_num_mamba_layers == 0: early return,
then keep the existing super().__init__ call and return.

🪄 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: 2ae55170-1756-4283-a131-19a7bf8a466f

📥 Commits

Reviewing files that changed from the base of the PR and between 9547230 and ca348ad.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
• tests/unittest/_torch/executor/test_mamba_cache_manager.py

[tensorrt-cicd]

PR_Github #47737 [ run ] triggered by Bot. Commit: ac30ded Link to invocation

[tensorrt-cicd]

PR_Github #47737 [ run ] completed with state SUCCESS. Commit: ac30ded
/LLM/main/L0_MergeRequest_PR pipeline #37632 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

[VALLIS-NERIA]

Addressed CodeRabbit feedback:

• API Usage Error #1015: Handle layer_mask=None (treat as all-False full-attention mask) and validate length against mamba_layer_mask.
• New Model Support #1053: Seed self._use_replay_state_update and self.ssm_state_dtype before the zero-mamba early-return so get_mamba_ssm_cache_dtype() and the use_replay_state_update property don't AttributeError on ranks with no local mamba layers.

[VALLIS-NERIA]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #47860 [ run ] triggered by Bot. Commit: 5aad39a Link to invocation

[Wanli-Jiang]

Verified that it can run well for new ckpt, thanks!

[tensorrt-cicd]

PR_Github #47860 [ run ] completed with state SUCCESS. Commit: 5aad39a
/LLM/main/L0_MergeRequest_PR pipeline #37724 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

[VALLIS-NERIA]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #47978 [ run ] triggered by Bot. Commit: 5aad39a Link to invocation

[tensorrt-cicd]

PR_Github #47978 [ run ] completed with state SUCCESS. Commit: 5aad39a
/LLM/main/L0_MergeRequest_PR pipeline #37818 completed with status: 'SUCCESS'

CI Report

Link to invocation