[https://nvbugs/6094224][fix] Fix mamba disagg issues when conc > mbs#13274
Conversation
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughAdded Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py`:
- Around line 579-582: The C++ path in filter_ctx_requests_by_capacity
short-circuits when self._use_cpp is true, bypassing the new capacity guard;
either implement the same capacity-filtering logic in the C++ binding class
(e.g., add a filter_ctx_requests_by_capacity method to CppMambaCacheManager and
expose it via the bindings) so PyExecutor._schedule will call into the C++
check, or detect the unsupported combination and fail fast by raising a clear
error when _use_cpp is true and disaggregated scheduling is requested (check
TRTLLM_USE_CPP_MAMBA environment flag or the manager type) so you don't return
unfiltered context_requests. Ensure you reference and update
filter_ctx_requests_by_capacity, CppMambaCacheManager (and its bindings), and
the call site in PyExecutor._schedule accordingly.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 2947-2959: The code currently overwrites num_fitting (set from
scheduler_output.num_fitting_requests) with the filtered context count after
calling kv_cache_manager.filter_ctx_requests_by_capacity, which loses the
original "did anything fit?" signal; revert that overwrite so num_fitting
remains scheduler_output.num_fitting_requests, still apply the filtered
scheduled_context_requests to ScheduledRequests.reset_context_requests, and if
you need the filtered count keep it in a separate local (e.g.,
filtered_ctx_count) rather than assigning it back to num_fitting; refer to the
variables/methods num_fitting, scheduler_output.num_fitting_requests,
scheduled_context_requests, kv_cache_manager.filter_ctx_requests_by_capacity,
and ScheduledRequests.reset_context_requests to locate and fix the code.
🪄 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: Pro Plus
Run ID: aecbc7f5-b3da-4393-9ada-71446a7b2ab1
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pytensorrt_llm/_torch/pyexecutor/py_executor.py
|
PR_Github #44710 [ run ] triggered by Bot. Commit: |
|
PR_Github #44710 [ run ] completed with state
|
|
/bot run |
|
PR_Github #44805 [ run ] triggered by Bot. Commit: |
|
PR_Github #44805 [ run ] completed with state
|
|
/bot run |
|
PR_Github #44812 [ run ] triggered by Bot. Commit: |
|
PR_Github #44812 [ run ] completed with state
|
|
/bot run |
|
PR_Github #44828 [ run ] triggered by Bot. Commit: |
|
PR_Github #44828 [ run ] completed with state
|
|
/bot run |
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #45257 [ run ] completed with state |
|
@Tabrizian @Wanli-Jiang Please help review, thanks! |
|
better to add a test to cover the corner case (conc>max batch size). |
|
/bot skip --comment "The new commit just adds a new test, and the test passed on nsc cluster" |
|
PR_Github #45715 [ skip ] triggered by Bot. Commit: |
|
PR_Github #45715 [ skip ] completed with state |
|
@bo-nv Is it possible to understand why scheduler is not handling this case correctly? Perhaps CapacityScheduler needs to be updated to make sure we have enough space in MambaCacheManager. |
Signed-off-by: Bo Deng <deemod@nvidia.com>
|
/bot skip --comment "Just add a comment" |
|
/bot run |
|
PR_Github #46042 [ run ] triggered by Bot. Commit: |
|
PR_Github #46042 [ run ] completed with state
|
|
/bot run |
|
PR_Github #46080 [ run ] triggered by Bot. Commit: |
|
/bot skip --comment "Just add a comment" |
|
PR_Github #46106 [ skip ] triggered by Bot. Commit: |
|
PR_Github #46106 [ skip ] completed with state |
…NVIDIA#13274) Signed-off-by: Bo Deng <deemod@nvidia.com> Co-authored-by: xinhe-nv <200704525+xinhe-nv@users.noreply.github.com>
5 participants
Summary by CodeRabbit
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)
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.