[None][perf] Avoid Index-K cache materialization for MSA#16856
Conversation
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
|
/bot run |
WalkthroughChangesMiniMax-M3 index-K layout handling
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant MSA_backend
participant KVCacheManagerV2
participant MsaIndexer
participant KV_cache
MSA_backend->>KVCacheManagerV2: get_index_k_buffer(layer_idx, kv_layout="HND")
KVCacheManagerV2->>KV_cache: create zero-copy HND view
KV_cache-->>MSA_backend: strided paged index-K
MSA_backend->>KVCacheManagerV2: write_kv_slots(..., layout="HND")
MSA_backend->>MsaIndexer: select_blocks(idx_k_paged)
MsaIndexer-->>MSA_backend: selected blocks
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ 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.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py (1)
1887-1926: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winHND view is only byte-safe when
num_heads == 1; document or assert this.The NHD/HND view-swap trick relies on
TensorWrapperbuilding fresh, default-contiguous strides for the requested logical shape rather than performing a real transpose. That aliases identical bytes only becausepage_shapecollapses to[T, 1, D]vs[1, T, D]whennum_heads == 1— both reduce to the samet*D+dbyte offset. Fornum_heads > 1, requestingkv_layout="HND"would silently reinterpret physically NHD-laid-out bytes as if they were HND, corrupting data without any error. All current callers (MiniMaxM3KVCacheManagerV2.get_index_k_buffer) hardcodenum_heads=1, so this is latent today, but the base method is public/generic and doesn't guard against future multi-head callers.Consider adding an explicit guard and documenting the precondition:
🛡️ Proposed guard
if kv_layout not in ("NHD", "HND"): raise ValueError(f"Unsupported kv_layout: {kv_layout}") + if kv_layout == "HND" and num_heads != 1: + raise NotImplementedError( + "kv_layout='HND' is only supported for num_heads == 1: the " + "NHD/HND views only alias identical bytes when the head " + "dimension is a singleton." + )Also applies to: 1981-1997
🤖 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 `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py` around lines 1887 - 1926, Update get_index_k_buffer to reject kv_layout="HND" when num_heads is greater than one, since the byte-level view is only valid for a single head; raise a clear ValueError before constructing the tensor view. Document this HND precondition in the method docstring while preserving NHD behavior and the existing num_heads=1 callers.
🤖 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.
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py`:
- Around line 1887-1926: Update get_index_k_buffer to reject kv_layout="HND"
when num_heads is greater than one, since the byte-level view is only valid for
a single head; raise a clear ValueError before constructing the tensor view.
Document this HND precondition in the method docstring while preserving NHD
behavior and the existing num_heads=1 callers.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 125b3a2d-5c93-40f0-8b43-65e91261453b
📒 Files selected for processing (7)
tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.pytests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py
|
PR_Github #61666 [ run ] triggered by Bot. Commit: |
|
PR_Github #61666 [ run ] completed with state
|
2 participants
@coderabbitai summary
Description
MiniMax-M3 MSA stores its one-head Index-K side cache through KV cache manager
V2. The cache accessor previously exposed only NHD layout, so every sparse
attention layer converted it to HND with
permute(...).contiguous()beforeinvoking the SM100 FMHA kernel.
This change:
kv_layoutcontract to the KV cache manager V2 accessor,with NHD retained as the default for existing callers;
fmha_sm100;For the one-head Index-K cache, NHD and HND have identical physical bytes
within each page. The change is therefore a metadata/layout-contract change
and avoids materializing a second cache tensor. Existing callers keep NHD as
the default, and no C++ or kernel source changes.
Performance
The original matched GB200 1P1D c64 qualification used 8K input / 1K output
requests. Nsight Systems isolated the following CTX and GEN effects:
_forward_step: -3.72%Each measured phase removed 456 Index-K layout copies. In the matched combined
E2E run, Pareto X (
1000 / median TPOT) improved 34.17% and Pareto Y (totaltoken throughput / 8 GPUs) improved 32.08%. Median TPOT improved 25.47%;
median TTFT regressed 2.91%. The focused phase traces above provide the causal
attribution.
Test Coverage
test_kv_cache_v2_extra_buffers.py.test_minimax_m3_msa_backend.py.FMHA proxy results were bit-identical, and CUDA Graph replay passed.
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.