[TRTLLM-13685][chore] Move MLA module to separate file - #15648
Conversation
|
/bot run |
|
PR_Github #55972 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThe PR moves MLA implementation into ChangesMLA module split and import rewiring
Sequence Diagram(s)sequenceDiagram
participant F as MLA.forward
participant P as trtllm.mla_dsa_proj
participant A as trtllm.mla_dsa_attn_inplace
participant I as forward_impl
participant C as forward_context
participant G as forward_absorption_generation
alt DSA path
F->>P: project hidden_states and position_ids
F->>A: dispatch q, compressed_kv, k_pe, latent_cache, indexer_intermediates
else Non-DSA path
F->>I: run pure-PyTorch MLA execution
I->>C: select a context strategy
I->>G: run absorption generation
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 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/models/modeling_speculative.py`:
- Around line 16-23: The remaining MLA import is still tied to the legacy
attention module path, so update the affected import usage to come from the
dedicated MLA module instead. In the speculative model code and any related test
setup such as test_model_loader_mx, reference MLA via the symbol in modules.mla
rather than modules.attention, unless you are intentionally keeping a legacy
compatibility coverage path. Keep the other module imports unchanged and make
sure all MLA references point to the new location.
In `@tensorrt_llm/_torch/modules/mla.py`:
- Around line 1914-1917: The FlashMLA padding path in mla.py creates extra heads
with uninitialized memory even though the comment says they should be zeros. In
the q_concat padding block, replace the q_padded allocation with a
zero-initialized tensor or explicitly zero the padded tail before copying
q_concat so the extra heads remain zeroed before dispatch to the sparse kernel.
Use the q_concat/q_padded logic in the padded tensor construction to locate the
fix.
- Around line 1291-1300: The chunked-prefill scratch tensors in mla.py are
hardcoded to allocate on the default CUDA device, which can mismatch the rank’s
actual tensor device in multi-GPU runs. Update the allocations in the relevant
MLA forward/chunked-prefill path to use q.device instead of a fixed cuda device,
and keep both softmax_stats_tensor and temp_softmax_stats_tensor on the same
device as q.
- Line 123: The unpacking in _extract_mla_extra_attrs within mla.py binds
metadata only to discard it, which triggers the Ruff RUF059 warning. Rename the
unused metadata binding to _metadata in the mla_layer extraction so the unused
variable is clearly marked and lint passes.
- Around line 1018-1041: The mixed context/generation path in the MLA forward
flow is using the full-batch position_ids and stale latent cache values after
external RoPE. Update the logic around apply_rope in the context and generation
slices so q_ctx/q_gen and k_pe_ctx/k_pe_gen receive matching sliced
position_ids, and refresh latent_cache_ctx/latent_cache_gen after RoPE the same
way the non-DSA path does. Use the existing forward_context_dsa, apply_rope,
latent_cache_ctx, and latent_cache_gen blocks as the place to keep cache writes
consistent when self.apply_rotary_emb is enabled.
🪄 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: 152ce3bf-9546-4f13-9ded-271ea999e266
📒 Files selected for processing (12)
AGENTS.mdtensorrt_llm/_torch/custom_ops/__init__.pytensorrt_llm/_torch/models/modeling_deepseekv3.pytensorrt_llm/_torch/models/modeling_speculative.pytensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/modules/attention.pytensorrt_llm/_torch/modules/mla.pytensorrt_llm/tools/layer_wise_benchmarks/mark_utils.pytests/unittest/_torch/attention/sparse/test_short_seq_mha.pytests/unittest/_torch/attention/sparse/test_sparse_mla_forward.pytests/unittest/_torch/attention/test_attention_mla.pytests/unittest/_torch/modules/test_mla_helix.py
|
/bot run --disable-fail-fast |
|
PR_Github #55974 [ run ] triggered by Bot. Commit: |
|
PR_Github #55972 [ run ] completed with state |
|
PR_Github #55974 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #56080 [ run ] triggered by Bot. Commit: |
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #56914 [ run ] triggered by Bot. Commit: |
|
PR_Github #56914 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #56976 [ run ] triggered by Bot. Commit: |
|
PR_Github #56976 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #57047 [ run ] triggered by Bot. Commit: |
|
PR_Github #57047 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #57101 [ run ] triggered by Bot. Commit: |
|
PR_Github #57101 [ run ] completed with state
|
|
/bot skip --comment "only DGX_H100-PyTorch-4/test_unittests.py::test_unittests_v2[unittest/disaggregated/test_cache_transceiver_harness.py] fails, this failed test also appear in other PR's CI and it is unrelated to this PR, skip CI" |
|
PR_Github #57125 Bot args parsing error: Failed to parse bot args |
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
/bot skip --comment "only DGX_H100-PyTorch-4/test_unittests.py::test_unittests_v2[unittest/disaggregated/test_cache_transceiver_harness.py] fails, this failed test also appear in other PR's CI and it is unrelated to this PR, skip CI" |
|
PR_Github #57138 Bot args parsing error: Failed to parse bot args |
|
/bot skip --comment "this single failed test also appears in other PR's CI and it is unrelated to this PR, skip CI" |
|
PR_Github #57143 Bot args parsing error: Failed to parse bot args |
|
/bot skip --comment "only DGX_H100-PyTorch-4/test_unittests.py::test_unittests_v2[unittest/disaggregated/test_cache_transceiver_harness.py] fails, this failed test also appear in CI of other PR and it is unrelated to this PR, skip CI" |
|
PR_Github #57144 [ skip ] triggered by Bot. Commit: |
|
PR_Github #57144 [ skip ] completed with state |
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Tests
Description
Move the
MLAmodule and MLA-specific custom ops/helpers out oftensorrt_llm/_torch/modules/attention.pyintotensorrt_llm/_torch/modules/mla.py.Shared attention utilities, including the standard
Attentionmodule and Helix CP helpers, remain inattention.pyand are imported bymla.pywhere needed. Model, test, benchmark, and documentation imports are updated to use the new MLA module path.Test Coverage
pre-commit run --files AGENTS.md tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md tensorrt_llm/_torch/custom_ops/__init__.py tensorrt_llm/_torch/models/modeling_deepseekv3.py tensorrt_llm/_torch/models/modeling_speculative.py tensorrt_llm/_torch/modules/attention.py tensorrt_llm/_torch/modules/mla.py tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py tests/unittest/_torch/attention/test_attention_mla.py tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py tests/unittest/_torch/attention/sparse/test_short_seq_mha.py tests/unittest/_torch/modules/test_mla_helix.pypython3 -m compileall -q tensorrt_llm/_torch/modules/attention.py tensorrt_llm/_torch/modules/mla.py tensorrt_llm/_torch/models/modeling_deepseekv3.py tensorrt_llm/_torch/models/modeling_speculative.py tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py tests/unittest/_torch/attention/test_attention_mla.py tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py tests/unittest/_torch/attention/sparse/test_short_seq_mha.py tests/unittest/_torch/modules/test_mla_helix.pyTargeted pytest could not run locally because this environment is missing TRT-LLM native extensions:
PYTHONPATH=$PWD:operator trtllm::deepseek_v4_q_norm does not existPYTHONPATH=$PWD TRT_LLM_NO_LIB_INIT=1:ModuleNotFoundError: No module named 'kv_cache_manager_v2.rawref._rawref'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.