[None][feat] Add FlashInfer MLA attention backend support - #13428
Conversation
|
/bot run |
|
PR_Github #45385 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis changeset adds MLA (Multi-head Latent Attention) support to the FlashInfer attention backend through dedicated planning dataclasses and state management. Core logic extends to handle both context and generation phases with phase-specific RoPE handling and output tensor allocation. Integration and unit tests validate the new functionality across GPU targets. Changes
Sequence DiagramsequenceDiagram
participant Test as Test/Application
participant FA as FlashInferAttention
participant MLAModule as MLA Module
participant FFI as FlashInfer Library
participant CUDA as CUDA/GPU
Test->>FA: forward(q, k, v, attn_metadata, latent_cache, phase)
alt Context Phase (Prefill)
FA->>FA: plan_ragged() for KV cache layout
FA->>CUDA: allocate stable device buffers
FA->>CUDA: sync (explicit plan execution)
FA->>MLAModule: extract RoPE positions from attn_metadata
MLAModule->>MLAModule: apply_mla_rope(k_pe, positions)
MLAModule->>FA: return RoPE'd k_pe
FA->>FFI: append_paged_mla_kv_cache(RoPE'd k, v, latent_cache)
FFI->>CUDA: update MLA latent cache
FA->>FFI: forward_ragged_mla(q, latent_cache)
FFI->>CUDA: compute attention with v_head_dim output
FA->>Test: return context output (v_head_dim per head)
else Generation Phase (Decode)
MLAModule->>MLAModule: apply_mla_rope(q_pe, k_pe, positions)
MLAModule->>FA: return RoPE'd tensors
FA->>FFI: append_paged_mla_kv_cache(RoPE'd k, v, latent_cache)
FFI->>CUDA: update MLA latent cache
FA->>FA: execute cached BatchMLAPagedAttentionWrapper plan
FA->>FFI: forward_paged_mla(q, planned_decode_wrapper)
FFI->>CUDA: compute attention with kv_lora_rank output
FA->>Test: return generation output (kv_lora_rank per head)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~55 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/_torch/attention_backend/flashinfer.py`:
- Around line 158-185: plan_ragged currently calls
BatchPrefillWithRaggedKVCacheWrapper.plan() during forward execution which can
violate CUDA graph capture; remove the call from plan_ragged so it only
initializes/returns self._ragged_prefill_wrapper, and move the actual planning
call into prepare (mirroring MLA decode): in prepare, ensure you create
self._ragged_prefill_wrapper if None and invoke
self._ragged_prefill_wrapper.plan(...) with qo_indptr, kv_indptr and the same
fields from plan_params (num_heads, num_kv_heads, head_dim, head_dim_vo,
use_fp16_qk_reduction=False, causal=True, q_data_type, kv_data_type, sm_scale)
so planning happens outside the forward path.
- Around line 927-956: This path runs ragged prefill only on the new chunk's
q_ctx/k_ctx/v_ctx and ignores any cached paged MLA KV produced by
append_paged_mla_kv_cache(), causing incorrect attention when a cached prefix
exists; update the logic around metadata.plan_ragged and wrapper.run to detect a
cached paged MLA KV (e.g., check a flag like metadata.has_paged_mla_kv_cache or
metadata.paged_mla_kv_present) and if present, short-circuit this branch (either
skip ragged prefill and use the paged-KV-aware path or fall back to a full
prefill that reads the cached KV) so that q_ctx/k_ctx/v_ctx are combined with
the cached prefix rather than used in isolation. Ensure the guard is placed
before creating RaggedPlanParams/wrapper and before wrapper.run so the cached
prefix is never ignored.
In `@tests/integration/test_lists/test-db/l0_b200.yml`:
- Line 39: The QA scheduled functional lists are missing the integration test
node; add the test identifier
accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer
to the QA functional list named llm_function_full.txt (and also append it to
llm_function_core.txt if you want core single-node FlashInfer coverage),
ensuring the entry follows the existing README format for functional test cases
so it propagates from the pre-merge list into the QA scheduled lists.
In `@tests/unittest/_torch/attention/test_attention_mla.py`:
- Around line 776-785: gen_metadata_kwargs currently sets max_num_tokens to
sum(context_sequence_lengths) which under-allocates for multi-token generation;
change max_num_tokens to reflect the tokens in the current decode batch (i.e.,
use generation_seq_len_q * len(context_sequence_lengths) or otherwise sum the
per-request generation lengths) so FlashInferAttentionMetadata is sized for the
generation_seq_len_q across the decode batch; update the assignment to the
max_num_tokens key in gen_metadata_kwargs accordingly (referencing
gen_metadata_kwargs, max_num_tokens, context_sequence_lengths, and
generation_seq_len_q).
- Around line 795-798: The assignment to gen_metadata_kwargs['enable_flash_mla']
is incorrectly indented causing a Flake8 E123; locate the conditional that
checks backend_name == "TRTLLM" and reformat the assignment so the key and its
value expression (torch.cuda.get_device_capability() == (9, 0)) are on the same
logical line or properly aligned (e.g., gen_metadata_kwargs['enable_flash_mla']
= torch.cuda.get_device_capability() == (9, 0)) to satisfy linting while
preserving the existing condition and semantics.
🪄 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: a7c2d5b9-66d6-49b5-aeac-4a5c2ec68275
📒 Files selected for processing (6)
tensorrt_llm/_torch/attention_backend/flashinfer.pytensorrt_llm/_torch/modules/attention.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_b300.ymltests/unittest/_torch/attention/test_attention_mla.py
|
PR_Github #45385 [ run ] completed with state
|
|
/bot run |
|
PR_Github #45628 [ run ] triggered by Bot. Commit: |
|
PR_Github #45628 [ run ] completed with state
|
Implement FlashInfer-based Multi-head Latent Attention (MLA) for
DeepSeek-style models, enabling use of BatchPrefillWithRaggedKVCacheWrapper
for context and BatchMLAPagedAttentionWrapper for generation.
Key design decisions:
- KV cache: zero-copy ckv/kpe views split from the existing paged KV pool
buffer (kv_lora_rank + qk_rope_head_dim), avoiding extra allocation
- Context phase: ragged prefill with expanded Q/K/V after appending latent
to paged MLA caches via append_paged_mla_kv_cache
- Generation phase: paged MLA decode with q_nope/q_pe split from fused_q;
latent cache append handled before each decode step
- RoPE applied externally in MLA.forward before latent_cache construction;
FlashInfer kernels receive pre-RoPE'd inputs
- Plans called unconditionally each forward pass (no stale caching)
- backend=auto on BatchMLAPagedAttentionWrapper for FA3 on Hopper
- mla_rope_generation stub copies RoPE'd q_pe into fused_q to satisfy
the forward_absorption_generation call chain
Also fix metadata type assertion in attention.py to allow
FlashInferAttentionMetadata for MLA, and replace hasattr duck-typing
with isinstance check for kv_lens_cuda_runtime.
Validated with 84/84 unit tests and end-to-end DeepSeek-V3-Lite NVFP4
inference with --attention_backend FLASHINFER.
Support CUDA graph for flashinfer MLA.
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
Drop redundant contiguous/reshape copies in flashinfer MLA path.
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
[None][test] Fix FlashInfer MLA unit test and add B200/B300 smoke coverage
- test_attention_mla_flashinfer: restrict to SM100 and set max_num_tokens
to sum(context_sequence_lengths) so FlashInfer's batch_indices/positions
buffers are sized for the full batch (was max(...), causing OOB copy).
- Add TestDeepSeekV3Lite::test_bfloat16_flashinfer accuracy smoke test and
register it in l0_b200.yml and l0_b300.yml.
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
|
/bot run |
|
PR_Github #47852 [ run ] triggered by Bot. Commit: |
|
PR_Github #47852 [ run ] completed with state
|
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
|
/bot run |
|
PR_Github #47950 [ run ] triggered by Bot. Commit: |
|
PR_Github #47950 [ run ] completed with state
|
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
|
/bot run |
|
PR_Github #48068 [ run ] triggered by Bot. Commit: |
|
PR_Github #48068 [ run ] completed with state
|
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
|
Hi @yuxianq, new comments are also addressed. Please take a look. Thanks! |
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
|
/bot run |
|
PR_Github #49340 [ run ] triggered by Bot. Commit: |
|
PR_Github #49340 [ run ] completed with state
|
|
/bot run |
Signed-off-by: Qi Zhang (qizh) <10434017+Tracin@users.noreply.github.com>
|
/bot run |
|
PR_Github #49949 [ run ] triggered by Bot. Commit: |
|
PR_Github #49949 [ run ] completed with state
|
|
/bot run |
|
PR_Github #50048 [ run ] triggered by Bot. Commit: |
|
PR_Github #50048 [ run ] completed with state
|
|
/bot run |
|
PR_Github #50073 [ run ] triggered by Bot. Commit: |
|
PR_Github #50073 [ run ] completed with state |
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
Implement FlashInfer-based Multi-head Latent Attention (MLA) for
DeepSeek-style models, enabling use of BatchPrefillWithRaggedKVCacheWrapper
for context and BatchMLAPagedAttentionWrapper for generation.
Key design decisions:
Also fix metadata type assertion in attention.py to allow
FlashInferAttentionMetadata for MLA, and replace hasattr duck-typing
with isinstance check for kv_lens_cuda_runtime.
Validated with 84/84 unit tests and end-to-end DeepSeek-V3-Lite NVFP4
inference with --attention_backend FLASHINFER.
Support CUDA graph for flashinfer MLA.
Drop redundant contiguous/reshape copies in flashinfer MLA path.
[None][test] Fix FlashInfer MLA unit test and add B200/B300 smoke coverage
Summary by CodeRabbit
Release Notes
New Features
Tests
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.