[TRTLLM-12807][feat] Wrap FP4 MLA as FMHA lib - #15375
Draft
yuxianq wants to merge 11 commits into
Draft
Conversation
1 task
yuxianq
force-pushed
the
feat/trtllm-12807-fp4-mla-fmha-lib
branch
3 times, most recently
from
June 20, 2026 00:03
e57b660 to
5c5863a
Compare
Allocate a standalone BF16 tensor indexed by sequence slot for recent MLA tokens when FP4 KV cache is active. Add Triton kernels to store context and generation latent cache values into the high-precision pool before attention runs. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
Add shared FP4 MLA KV-cache helpers, V-scale storage, Triton no-dequant decode, and FlashInfer integration for NVFP4 MLA cache. Fold the offset, layout, memory, workspace, tail-scale, and scatter overflow fixes into the implementation, with focused unit coverage. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
Add CuTile and CuTe DSL implementations for the FP4 MLA paged decode path and route them through the existing backend selector. Extend the opt-in FP4 MLA decode benchmark coverage to report the selected backend and estimated memory bandwidth. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
Add a standalone benchmark comparing the FP4 MLA decode backends against FlashInfer BF16 and trtllm-gen baselines. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
Disable compile wrapping for the copy helper and force FLA utility device selection to CUDA for this debug path. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> [None][fix] Avoid CuTe FP4 MLA large KV descriptor overflow Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> Support MPT. Update kernels. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> Support MPT. Update kernels. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> Support MPT. Update kernels. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> Support MPT. Update kernels. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> Support MPT. Update kernels. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
Reuse FP4 MLA V-packed storage across layers Handle partial FP4 MLA page-stat groups Tune FP4 MLA PV block size by batch Optimize FP4 MLA MTP decode path Add experimental FP4 MLA MTP fused QKPV path Optimize FP4 MLA MTP final-page path Fix cutile bug. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> Triton kernel prepack V. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
yuxianq
force-pushed
the
feat/trtllm-12807-fp4-mla-fmha-lib
branch
from
June 20, 2026 00:20
5c5863a to
3549726
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
@coderabbitai summary
Description
Wrap FP4 MLA attention behind the FMHA library abstraction for the PyTorch backend. This adds FMHA interface/registry plumbing, the FP4 MLA FMHA implementation, FlashInfer/TRTLLM integration points, KV cache support updates, and speculative decoding/MTP wiring needed by the new path.
Rebased onto
mainafter #15204 merged.Test Coverage
tests/unittest/_torch/attention/test_fp4_mla.py.python3 -m compileall -q tensorrt_llm/_torch/attention_backend/fmha/fp4_mla.py tensorrt_llm/_torch/pyexecutor/resource_manager.pygit diff --checkENABLE_TILE=0 TRTLLM_FP4_MLA_ATTENTION_BACKEND=triton trtllm-eval --model /home/scratch.trt_llm_data_ci/llm-models/DeepSeek-V3-Lite/nvfp4_moe_only/ --config tmp/extra_llm_api_options.yaml gsm8kpassed, average accuracy60.46.nv-triton==9.9.99.dev20260603+git69a952bc.cudatile:ENABLE_TILE=1 TRTLLM_FP4_MLA_ATTENTION_BACKEND=cutile trtllm-eval --model /home/scratch.trt_llm_data_ci/llm-models/DeepSeek-V3-Lite/nvfp4_moe_only/ --config tmp/extra_llm_api_options.yaml gsm8kpassed, average accuracy61.07.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.