[feat] Add FlashInfer attention backend support for MiniMax H3

[klhhhhh]

Dependency

This PR depends on the following PRs and should be reviewed and merged after both of them:

1. #1799 — Add opt-in FlashInfer attention backend for Wan
2. #1827 — Add FlashInfer cuDNN batched prefill attention

Dependency order:

#1799: Base FlashInfer backend
  └── #1827: FlashInfer cuDNN prefill implementation
        └── This PR: MiniMax H3 integration

PR #1799 introduces the shared opt-in FLASHINFER attention backend using
single_prefill_with_kv_cache.

PR #1827 extends that backend with the optional
cudnn_batch_prefill_with_kv_cache implementation.

This PR enables MiniMax H3 to use the shared FlashInfer backend introduced by
those PRs. It does not duplicate the FlashInfer kernel implementation.

Summary

This PR adds opt-in FlashInfer attention support for MiniMax H3.

MiniMax H3 already uses FastVideo's shared DistributedAttention abstraction.
Its attention tensors follow the BSHD contract expected by the existing
FlashInfer backend:

query: [batch, sequence, heads, head_dim]
key:   [batch, sequence, heads, head_dim]
value: [batch, sequence, heads, head_dim]

H3 also uses a head dimension of 128 and dense non-causal self-attention, which
is compatible with both FlashInfer implementations provided by the dependency
PRs.

The integration applies to:

• The 50 main MiniMax H3 Transformer blocks
• The two MiniMax H3 text-refiner blocks
• H3's Q/K RMSNorm path
• H3's partial rotary embedding path
• Sequence-parallel attention through DistributedAttention

Usage

Use the FlashInfer single-request implementation:

FASTVIDEO_ATTENTION_BACKEND=FLASHINFER \
FASTVIDEO_FLASHINFER_PREFILL_BACKEND=single \
<minimax-h3-inference-command>

Use the FlashInfer cuDNN implementation:

FASTVIDEO_ATTENTION_BACKEND=FLASHINFER \
FASTVIDEO_FLASHINFER_PREFILL_BACKEND=cudnn \
<minimax-h3-inference-command>

The integration remains opt-in and does not change the default MiniMax H3
attention backend.

Implementation

Declare FlashInfer support for MiniMax H3

This PR adds AttentionBackendEnum.FLASHINFER to:

MiniMaxH3ArchConfig._supported_attention_backends

The supported backend list becomes:

_supported_attention_backends = (
    AttentionBackendEnum.TORCH_SDPA,
    AttentionBackendEnum.FLASH_ATTN,
    AttentionBackendEnum.FLASHINFER,
    AttentionBackendEnum.ATTN_QAT_INFER,
    AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3,
)

Because MiniMaxH3Attention already constructs DistributedAttention, no
separate FlashInfer call is added to the H3 model forward.

The resulting dispatch path is:

MiniMaxH3Attention
  → DistributedAttention
  → attention backend selector
  → FlashInferImpl
  → single_prefill_with_kv_cache
     or cudnn_batch_prefill_with_kv_cache

Preserve H3-specific attention behavior

MiniMax H3 rotates only 96 of its 128 attention-head channels.

The H3 model continues to apply its Q/K RMSNorm and partial RoPE before calling
DistributedAttention. FlashInfer therefore receives the same normalized and
rotated Q/K tensors as the existing Torch SDPA and FlashAttention paths.

No H3 packing, RoPE, projection, or output logic is changed by this PR.

Files changed

fastvideo/configs/models/dits/minimax_h3.py

Adds AttentionBackendEnum.FLASHINFER to the backends supported by
MiniMaxH3ArchConfig.

fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py

Adds a regression test verifying that MiniMax H3 continues to declare
FlashInfer support.

fastvideo/tests/transformers/test_minimax_h3_flashinfer.py

Adds MiniMax H3-specific FlashInfer integration tests covering:

• Routing from MiniMaxH3Attention to DistributedAttention
• Selection of the FLASHINFER backend
• The real H3 attention execution path
• H3 Q/K RMSNorm
• H3 partial rotary embedding
• BF16 inputs
• Head dimension 128
• FlashInfer single prefill versus Torch SDPA
• FlashInfer cuDNN prefill versus Torch SDPA
• Final attention output projection

Tests

Run the model-routing and CUDA parity tests with:

CUDA_VISIBLE_DEVICES=0 python -m pytest \
  fastvideo/tests/transformers/test_minimax_h3_flashinfer.py \
  fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py \
  fastvideo/tests/attention/test_flashinfer_backend.py \
  -vs

The CUDA parity test constructs real MiniMaxH3Attention layers and compares
the complete attention output against Torch SDPA for both:

FASTVIDEO_FLASHINFER_PREFILL_BACKEND=single
FASTVIDEO_FLASHINFER_PREFILL_BACKEND=cudnn

This is a model-level attention integration test rather than a direct
FlashInfer kernel-only test.

Scope

This PR covers MiniMax H3 dense-attention routing and model-level numerical
parity.

It does not yet include:

• Full MiniMax H3 video-and-audio generation parity
• End-to-end H3 performance results
• Automatic selection between FlashInfer single and cuDNN implementations
• FlashInfer training or backward support
• Replacement of the H3 Video Sparse Attention backend

Performance considerations

MiniMax H3 packs text, conditioning, audio, and video tokens into one long
attention sequence. In common inference configurations, this packed document
uses an attention batch size of one.

As a result, the FlashInfer single-request implementation may be more suitable
for typical H3 inference than the cuDNN batched implementation. Both paths
remain explicitly selectable so that they can be evaluated on real H3
workloads before choosing any model-specific default.

Follow-ups

• Run end-to-end MiniMax H3 generation parity against Torch SDPA and
  FlashAttention.
• Benchmark DiT latency, total generation latency, and peak GPU memory.
• Compare FlashInfer single and cuDNN on representative H3 sequence lengths.
• Validate sequence-parallel execution with multiple GPUs.
• Determine whether H3 should automatically prefer the single-request
  implementation for batch-size-one packed attention.

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed

This rule is failing.

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[klhhhhh]

@claude