Skip to content

[None][feat] Add FlashInfer MLA attention backend support - #13428

Merged
Tracin merged 15 commits into
NVIDIA:mainfrom
Tracin:flashinfer_mla
May 25, 2026
Merged

[None][feat] Add FlashInfer MLA attention backend support#13428
Tracin merged 15 commits into
NVIDIA:mainfrom
Tracin:flashinfer_mla

Conversation

@Tracin

@Tracin Tracin commented Apr 24, 2026

Copy link
Copy Markdown
Collaborator

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.

Drop redundant contiguous/reshape copies in flashinfer MLA path.

[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.

Summary by CodeRabbit

Release Notes

  • New Features

    • Added Multi-head Latent Attention (MLA) support to the FlashInfer attention backend, enabling optimized attention computation for MLA-based models.
  • Tests

    • Added new bfloat16 accuracy tests with FlashInfer backend for enhanced model evaluation coverage and validation.

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.

@Tracin
Tracin requested review from a team as code owners April 24, 2026 10:06
@Tracin
Tracin requested a review from QiJune April 24, 2026 10:06
@Tracin

Tracin commented Apr 24, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #45385 [ run ] triggered by Bot. Commit: 91e5197 Link to invocation

@coderabbitai

coderabbitai Bot commented Apr 24, 2026

Copy link
Copy Markdown
Contributor
📝 Walkthrough

Walkthrough

This 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

Cohort / File(s) Summary
FlashInfer MLA Backend Implementation
tensorrt_llm/_torch/attention_backend/flashinfer.py
Introduces RaggedPlanParams and MLADecodePlanParams dataclasses for MLA planning. Adds support_mla() class method and mla_params initialization. Extends forward() with latent_cache parameter. Implements context path (ragged prefill with planned decode) and generation path (latent append plus paged decode). Manages stable device buffers for CUDA graph compatibility and explicit synchronization for plan() execution outside stream capture.
MLA Module RoPE Specialization
tensorrt_llm/_torch/modules/attention.py
Refines extract_extra_attrs metadata validation with unified isinstance check. Specializes RoPE application in MLA.forward_impl() for FlashInferAttentionMetadata: context and generation phases now slice position_ids before RoPE application and reconstruct latent_cache_* from RoPE'd k_pe. Updates forward_absorption_generation to compute num_seqs from attn_metadata.num_generations for FlashInfer metadata.
MLA FlashInfer Unit Tests
tests/unittest/_torch/attention/test_attention_mla.py
Adds apply_mla_rope utility for MLA-style RoPE computation. Introduces test_attention_mla_flashinfer parametrized test (SM100 GPU with BF16 constraint). Implements backend-specific RoPE pre-rotation for context and generation phases, bypassing prior fused RoPE paths for FlashInfer execution. Retains non-FlashInfer generation path with prior cuSeqlen allocation logic.
Integration Test Cases
tests/integration/defs/accuracy/test_llm_api_pytorch.py, tests/integration/test_lists/test-db/l0_b200.yml, tests/integration/test_lists/test-db/l0_b300.yml
Adds test_bfloat16_flashinfer() method to accuracy test suite with BF16-only variant running GSM8K evaluation. Registers test in l0_b200 and l0_b300 integration test lists. Uses low-memory skip marker for gating.

Sequence Diagram

sequenceDiagram
    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
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~55 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly summarizes the main change: adding FlashInfer MLA attention backend support.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Description check ✅ Passed PR description includes detailed technical implementation details, design decisions, test coverage information, and a completed PR checklist confirming guideline compliance.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between c4b8e8e and 91e5197.

📒 Files selected for processing (6)
  • tensorrt_llm/_torch/attention_backend/flashinfer.py
  • tensorrt_llm/_torch/modules/attention.py
  • tests/integration/defs/accuracy/test_llm_api_pytorch.py
  • tests/integration/test_lists/test-db/l0_b200.yml
  • tests/integration/test_lists/test-db/l0_b300.yml
  • tests/unittest/_torch/attention/test_attention_mla.py

Comment thread tensorrt_llm/_torch/attention_backend/flashinfer.py Outdated
Comment thread tensorrt_llm/_torch/attention_backend/flashinfer.py
Comment thread tests/integration/test_lists/test-db/l0_b200.yml Outdated
Comment thread tests/unittest/_torch/attention/test_attention_mla.py
Comment thread tests/unittest/_torch/attention/test_attention_mla.py Outdated
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #45385 [ run ] completed with state FAILURE. Commit: 91e5197
/LLM/main/L0_MergeRequest_PR pipeline #35626 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

@Tracin

Tracin commented Apr 27, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #45628 [ run ] triggered by Bot. Commit: d253811 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #45628 [ run ] completed with state SUCCESS. Commit: d253811
/LLM/main/L0_MergeRequest_PR pipeline #35841 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

Tracin added 2 commits May 11, 2026 20:09
  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>
@Tracin
Tracin requested a review from a team as a code owner May 12, 2026 03:14
@Tracin

Tracin commented May 12, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #47852 [ run ] triggered by Bot. Commit: dee75ec Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #47852 [ run ] completed with state SUCCESS. Commit: dee75ec
/LLM/main/L0_MergeRequest_PR pipeline #37723 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
@Tracin

Tracin commented May 12, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #47950 [ run ] triggered by Bot. Commit: a00d56b Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #47950 [ run ] completed with state SUCCESS. Commit: a00d56b
/LLM/main/L0_MergeRequest_PR pipeline #37792 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
@Tracin
Tracin requested a review from a team as a code owner May 13, 2026 02:17
@Tracin

Tracin commented May 13, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #48068 [ run ] triggered by Bot. Commit: 92bda1c Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #48068 [ run ] completed with state SUCCESS. Commit: 92bda1c
/LLM/main/L0_MergeRequest_PR pipeline #37901 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Comment thread tensorrt_llm/_torch/attention_backend/flashinfer.py
Comment thread tensorrt_llm/_torch/attention_backend/flashinfer.py
Comment thread tensorrt_llm/_torch/attention_backend/flashinfer.py Outdated
Tracin added 4 commits May 19, 2026 20:07
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>
@Tracin

Tracin commented May 20, 2026

Copy link
Copy Markdown
Collaborator Author

Hi @yuxianq, new comments are also addressed. Please take a look. Thanks!

Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
@Tracin

Tracin commented May 20, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #49340 [ run ] triggered by Bot. Commit: 0d71ba3 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #49340 [ run ] completed with state SUCCESS. Commit: 0d71ba3
/LLM/main/L0_MergeRequest_PR pipeline #38997 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@Tracin

Tracin commented May 22, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

Signed-off-by: Qi Zhang (qizh) <10434017+Tracin@users.noreply.github.com>
@Tracin

Tracin commented May 22, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #49949 [ run ] triggered by Bot. Commit: 491421e Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #49949 [ run ] completed with state SUCCESS. Commit: 491421e
/LLM/main/L0_MergeRequest_PR pipeline #39520 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@Tracin

Tracin commented May 23, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #50048 [ run ] triggered by Bot. Commit: 491421e Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #50048 [ run ] completed with state SUCCESS. Commit: 491421e
/LLM/main/L0_MergeRequest_PR pipeline #39609 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@Tracin

Tracin commented May 24, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #50073 [ run ] triggered by Bot. Commit: 491421e Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #50073 [ run ] completed with state SUCCESS. Commit: 491421e
/LLM/main/L0_MergeRequest_PR pipeline #39629 completed with status: 'SUCCESS'

CI Report

Link to invocation

@Tracin
Tracin merged commit e45a8e3 into NVIDIA:main May 25, 2026
7 checks passed
KleinBlueC pushed a commit to KleinBlueC/TensorRT-LLM that referenced this pull request May 26, 2026
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
bmarimuthu-nv pushed a commit to nv-auto-deploy/TensorRT-LLM that referenced this pull request May 28, 2026
Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com>
@Tracin
Tracin deleted the flashinfer_mla branch June 9, 2026 09:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants