[None] [chore] Mamba cache in separate file #6796
Conversation
📝 WalkthroughWalkthroughMambaCacheManager and MambaHybridCacheManager were moved from resource_manager.py into a new module mamba_cache_manager.py; imports were updated in _util.py and CODEOWNERS entries were updated to reference the new path. No runtime behavior changes beyond relocation. Changes
Sequence Diagram(s)sequenceDiagram
participant Scheduler as Scheduler/Caller
participant Hybrid as MambaHybridCacheManager
participant Mamba as MambaCacheManager
participant KV as KVCacheManager
Scheduler->>Hybrid: prepare_resources(scheduled_batch)
Hybrid->>Mamba: prepare_resources(scheduled_batch)
Hybrid->>KV: prepare_resources(scheduled_batch)
Hybrid-->>Scheduler: done
Scheduler->>Hybrid: free_resources(request)
Hybrid->>Mamba: free_resources(request)
Hybrid->>KV: free_resources(request)
Hybrid-->>Scheduler: done
Scheduler->>Hybrid: shutdown()
Hybrid->>Mamba: shutdown()
Hybrid->>KV: shutdown()
Hybrid-->>Scheduler: done
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Suggested reviewers
📜 Recent review detailsConfiguration used: .coderabbit.yaml 📒 Files selected for processing (2)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
Status, Documentation and Community
|
Signed-off-by: Tomer Asida <[email protected]>
…de ownership Signed-off-by: Tomer Asida <[email protected]>
tomeras91
force-pushed
the
mamba-cache-separate-file
branch
from
25829fb to
91b7bba
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (3)
tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py (3)
28-44: Add class docstring for MambaCacheManager (guidelines prefer Google-style docstrings).A brief summary of purpose and args improves maintainability and Sphinx rendering.
class MambaCacheManager(BaseResourceManager): + """GPU-resident cache for Mamba conv and SSM states. + + Args: + d_state: SSM state size. + d_conv: Convolution kernel length. + num_heads: Number of attention heads (pre-TP). + n_groups: Number of groups for state partitioning. + head_dim: Per-head dimension. + num_layers: Total transformer layers (pre-PP; will be sharded). + max_batch_size: Maximum batch size per rank. + mapping: Parallelism mapping. + dtype: Torch dtype for conv states. + ssm_cache_dtype: Torch dtype for SSM states. + layer_mask: Optional mask to select subset of layers locally. + """
167-199: Add class docstring for MambaHybridCacheManager.Documenting the hybrid behavior clarifies dual-manager responsibilities.
class MambaHybridCacheManager(KVCacheManager, MambaCacheManager): + """Hybrid cache manager combining Mamba conv/SSM states with KV cache. + + Initializes both MambaCacheManager and KVCacheManager and delegates + prepare/free/shutdown to both. + """
77-88: Guard against invalid d_conv values (conv_states uses d_conv - 1).If d_conv could be 1, this shape would be invalid. A defensive assert avoids silent mis-shapes.
# mamba conv states + assert d_conv >= 2, "d_conv must be >= 2 since conv state uses d_conv - 1" self.conv_states = torch.empty(
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
.github/CODEOWNERS(1 hunks)tensorrt_llm/_torch/pyexecutor/_util.py(1 hunks)tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py(1 hunks)tensorrt_llm/_torch/pyexecutor/resource_manager.py(0 hunks)
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/pyexecutor/resource_manager.py
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.py: Python code should conform to Python 3.8+.
Indent Python code with 4 spaces. Do not use tabs.
Always maintain the namespace when importing in Python, even if only one class or function from a module is used.
Python filenames should use snake_case (e.g., some_file.py).
Python classes should use PascalCase (e.g., class SomeClass).
Python functions and methods should use snake_case (e.g., def my_awesome_function():).
Python local variables should use snake_case. Prefix k for variable names that start with a number (e.g., k_99th_percentile).
Python global variables should use upper snake_case and prefix G (e.g., G_MY_GLOBAL).
Python constants should use upper snake_case (e.g., MY_CONSTANT).
Avoid shadowing variables declared in an outer scope in Python.
Initialize all externally visible members of a Python class in the constructor.
For interfaces that may be used outside a Python file, prefer docstrings over comments.
Comments in Python should be reserved for code within a function, or interfaces that are local to a file.
Use Google style docstrings for Python classes and functions, which can be parsed by Sphinx.
Attributes and variables in Python can be documented inline; attribute docstrings will be rendered under the class docstring.
Avoid using reflection in Python when functionality can be easily achieved without it.
When using try-except blocks in Python, limit the except to the smallest set of errors possible.
When using try-except blocks to handle multiple possible variable types in Python, keep the body of the try as small as possible, using the else block to implement the logic.
Files:
tensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
**/*.{cpp,h,hpp,cc,cxx,cu,py}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
All TensorRT-LLM Open Source Software code should contain an NVIDIA copyright header that includes the current year. This includes .cpp, .h, .cu, .py, and any other source files which are compiled or interpreted.
Files:
tensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
🧠 Learnings (1)
📚 Learning: 2025-08-06T13:58:07.506Z
Learnt from: galagam
PR: NVIDIA/TensorRT-LLM#6487
File: tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_trtllm_bench.py:1-12
Timestamp: 2025-08-06T13:58:07.506Z
Learning: In TensorRT-LLM, test files (files under tests/ directories) do not require NVIDIA copyright headers, unlike production source code files. Test files typically start directly with imports, docstrings, or code.
Applied to files:
.github/CODEOWNERS
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
197-197: Undefined name DecodingBaseConfig
(F821)
201-201: Line too long (130 > 120)
(E501)
🔇 Additional comments (2)
.github/CODEOWNERS (1)
110-110: CODEOWNERS entry for mamba_cache_manager.py is correct and old entry removed
- No occurrences of
/tensorrt_llm/_torch/pyexecutor/resource_manager.pyremain in.github/CODEOWNERS.- The new
/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pyentry on line 110 correctly assigns ownership to the Nemotron teams.tensorrt_llm/_torch/pyexecutor/_util.py (1)
28-33: Verification complete – no stale imports of MambaHybridCacheManager remainAll occurrences of MambaHybridCacheManager are either in its definition file or in
_util.pywith the updated import path. No other modules import it fromresource_manager. Safe to merge.
…che_manager.py - it's redundant. They don't need to specifically own this file, and it is already included under tensorrt_llm/_torch/pyexecutor Signed-off-by: Tomer Asida <[email protected]>
Signed-off-by: Tomer Asida <[email protected]>
|
/bot run |
|
PR_Github #14817 [ run ] triggered by Bot |
|
PR_Github #14817 [ run ] completed with state |
|
/bot run |
|
/bot run |
|
PR_Github #14834 [ run ] triggered by Bot |
|
PR_Github #14834 [ run ] completed with state |
|
/bot run |
Signed-off-by: Tomer Asida <[email protected]>
|
/bot run |
|
PR_Github #14986 [ run ] triggered by Bot |
|
PR_Github #14986 [ run ] completed with state |
|
/bot run |
|
PR_Github #15145 [ run ] triggered by Bot |
|
/bot run |
|
PR_Github #15304 [ run ] triggered by Bot |
|
PR_Github #15304 [ run ] completed with state |
Signed-off-by: Tomer Asida <[email protected]>
Signed-off-by: Tomer Asida <[email protected]> Signed-off-by: Wangshanshan <[email protected]>
Signed-off-by: Tomer Asida <[email protected]> Signed-off-by: Wangshanshan <[email protected]>
Signed-off-by: Tomer Asida <[email protected]>
Signed-off-by: Tomer Asida <[email protected]> Signed-off-by: Wangshanshan <[email protected]>
Signed-off-by: Tomer Asida <[email protected]>
Signed-off-by: Tomer Asida <[email protected]> Signed-off-by: Wangshanshan <[email protected]>
5 participants
Summary by CodeRabbit
Refactor
Chores
Description
prior to this PR, nemotron-devs were code owners of
resource_manager.py, because it containedMambaCacheManagerandMambaHybridCacheManagerwhich are unique and needed for mamba and hybrid attention mamba models (specifically, Nemotron-H architecture). This resulted in many false positives of for reviewing PRs which aren't really relevant for Nemotron models. This PR moves these classes to a separate filemamba_cache_manager.py, and replaces code ownership ofresource_manager.pywith this new file.