[None][fix] Use HND mapping for MiniMax-M3 MSA KV cache - #17374
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (5)
🚧 Files skipped from review as they are similar to previous changes (5)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. WalkthroughMiniMax M3 cache management now selects NHD or HND layouts from the attention implementation. Mapper kinds, buffer access, KV-transfer test setup, and verification follow the selected layout. ChangesMiniMax M3 KV layout handling
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to This PR narrows MiniMax-M3 MSA KV-cache transfers to the correct HND mapping while preserving existing Triton behavior; no actionable merge-blocking risk remains after normal checks and review. Sequence Diagram(s)sequenceDiagram
participant ContextManager
participant GenerationManager
participant KVTransferTest
ContextManager->>GenerationManager: transfer layout-specific KV buffers
GenerationManager->>KVTransferTest: expose transferred buffers
KVTransferTest->>KVTransferTest: compare HND or NHD slices
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Description checkExplanation The description explains the issue, solution, test coverage, and known GPU test limitation. The checklist is included and marked as reviewed. The API-change checklist item does not state whether an api-compatible label was added, but the description is otherwise complete.
✨ 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: 1
🤖 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/unittest/disaggregated/test_minimax_m3_kv_transfer.py`:
- Around line 101-109: Add type annotations to all newly changed functions:
annotate the parametrized arguments of test_minimax_disagg_role_mapper_kinds and
the parameters of its nested fake_base_init, then annotate
update_before_transfer as bool in
tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py at lines 101-109 and
660 respectively.
🪄 Autofix
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: f3f6b46b-52f1-4280-9113-7c8a73bd351d
📒 Files selected for processing (4)
tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytests/unittest/disaggregated/test_minimax_m3_kv_transfer.pytests/unittest/disaggregated/test_pool_matching.py
|
/bot run |
|
PR_Github #64698 [ run ] triggered by Bot. Commit: |
|
PR_Github #64698 [ run ] completed with state
|
|
/bot run |
|
Removed the "ci: full pre-merge approved" label because @peihu-nv could not be verified as an active member of NVIDIA/trt-llm-ci-approvers. Ask a member of that team to apply it. |
|
PR_Github #64747 [ run ] triggered by Bot. Commit: |
|
/bot run --only-multi-gpu-test |
|
PR_Github #64748 [ run ] triggered by Bot. Commit: |
|
PR_Github #64747 [ run ] completed with state |
|
PR_Github #64748 [ run ] completed with state
|
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
7245457 to
df50617
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
/bot run --add-multi-gpu-test |
|
PR_Github #66852 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #68921 [ run ] triggered by Bot. Commit: |
|
PR_Github #68921 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #69231 [ run ] triggered by Bot. Commit: |
|
PR_Github #69231 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #70757 [ run ] triggered by Bot. Commit: |
|
PR_Github #70757 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #70796 [ run ] triggered by Bot. Commit: |
|
PR_Github #70796 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #70817 [ run ] triggered by Bot. Commit: |
|
PR_Github #70817 [ run ] completed with state
|
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
/bot run --disable-fail-fast |
|
PR_Github #71028 [ run ] triggered by Bot. Commit: |
|
PR_Github #71028 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #71066 [ run ] triggered by Bot. Commit: |
|
PR_Github #71066 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #71193 [ run ] triggered by Bot. Commit: |
|
PR_Github #71193 [ run ] completed with state |
6 participants
Dev Engineer Review
INDEXEDmapper.REPLICATED.QA Engineer Review
Test functions added or modified:
test_minimax_m3_msa_hnd_head_mismatch_transfertest_minimax_disagg_role_mapper_kindstests/integration/test_lists/test-db/l0_h100.ymlincludesunittest/disaggregated/test_minimax_m3_kv_transfer.py, which covers the KV-transfer test module. No separateqa/entry covers these tests.Verdict: sufficient. GPU CI is still required to validate the TP1-to-TP4 transfer path.
Description
MiniMax-M3 MSA stores its main paged K/V cache in HND layout, but disaggregated KV transfer declared the cache as NHD. This causes head-mismatched context and generation parallelism to use the wrong transfer mapping.
Select the
INDEXEDmapper and HND cache views for MSA while preserving the existing NHD behavior for the Triton backend. Add focused TP1-to-TP4 transfer coverage for the MSA path. MiniMax-M3 MSA uses FP8 rather than NVFP4 KV cache, so NVFP4 scale-pool handling remains explicitly scoped to the Triton/NHD path.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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.