[perf]: use packed-varlen FA4 for MiniMax H3

fastvideo/attention/backends/flash_attn.py

@@ -10,6 +10,14 @@ from fastvideo.attention.utils.flash_attn_default import (
     flash_attn_func_compilable,
 )
 
+if fa_version == "4":
+    # The FA4 varlen wrapper is already a compile-safe custom op. Keep the
+    # import conditional so FA2/FA3 environments do not need flash_attn.cute.
+    from fastvideo.attention.utils.flash_attn_cute import (
+        flash_attn_varlen_func as flash_attn_varlen_func_compilable, )
+else:
+    flash_attn_varlen_func_compilable = None
+
 from fastvideo.attention.backends.abstract import (
     AttentionBackend,
     AttentionImpl,
@@ -38,6 +46,15 @@ except ImportError:
     _FA4_FP4_AVAILABLE = False
 
 _FA4_QUANT_OPS: tuple | None = None
+_FA4_PACKED_VARLEN_RECEIPT_LOGGED = False
+
+
+def _log_fa4_packed_varlen_receipt() -> None:
+    """Emit one grep-friendly route receipt per worker process."""
+    global _FA4_PACKED_VARLEN_RECEIPT_LOGGED
+    if not _FA4_PACKED_VARLEN_RECEIPT_LOGGED:
+        logger.info("MiniMax-H3 dense attention: FA4 packed-varlen inference route enabled")
+        _FA4_PACKED_VARLEN_RECEIPT_LOGGED = True
 
 
 def _import_fa4_quant_ops() -> tuple:
@@ -228,6 +245,12 @@ class FlashAttentionImpl(AttentionImpl):
     ) -> None:
         self.causal = causal
         self.softmax_scale = softmax_scale
+        # MiniMax-H3's dense DiT explicitly enables this faster FA4 entry
+        # point. It remains off for every other model and for the H3 text
+        # refiner; grad-enabled calls stay on the established fixed path.
+        self.fa4_packed_varlen = bool(extra_impl_args.get("fa4_packed_varlen", False))
+        if self.fa4_packed_varlen and fa_version == "4":
+            _log_fa4_packed_varlen_receipt()
         # An explicit ``nvfp4_fa4`` impl arg wins over the process-wide
         # FASTVIDEO_NVFP4_FA4 env opt-in, so precision-sensitive layers (e.g.
         # the FP32-pinned H3 VAE attention) can force-disable FP4 Q/K
@@ -331,6 +354,29 @@ class FlashAttentionImpl(AttentionImpl):
         elif self.nvfp4_fa4:
             output = self._forward_nvfp4(query, key, value)
 
+        elif (self.fa4_packed_varlen and fa_version == "4" and not torch.is_grad_enabled() and query.shape[0] == 1
+              and query.shape[1] == key.shape[1] == value.shape[1]):
+            # FA4's packed-varlen entry point is materially faster for H3's
+            # long, single-document self-attention. Flatten only the batch
+            # dimension and describe that one sequence with CUDA int32
+            # cumulative lengths; the existing custom-op wrapper keeps this
+            # route traceable under torch.compile(fullgraph=True).
+            assert flash_attn_varlen_func_compilable is not None
+            sequence_length = query.shape[1]
+            cu_seqlens = torch.arange(2, dtype=torch.int32, device=query.device) * sequence_length
+            output = flash_attn_varlen_func_compilable(
+                query.squeeze(0),
+                key.squeeze(0),
+                value.squeeze(0),
+                cu_seqlens,
+                cu_seqlens,
+                sequence_length,
+                sequence_length,
+                dropout_p=0.0,
+                softmax_scale=self.softmax_scale,
+                causal=self.causal,
+            ).unsqueeze(0)
+
         else:
             # Route through the compilable wrapper so dynamo sees a
             # registered op (no graph break) for FA2/FA3; identical

fastvideo/models/dits/minimax_h3.py

@@ -144,6 +144,7 @@ class MiniMaxH3Attention(nn.Module):
         quant_config: QuantizationConfig | None,
         prefix: str,
         fuse_qknorm_rope: bool = False,
+        fa4_packed_varlen: bool = False,
     ) -> None:
         super().__init__()
         self.num_attention_heads = num_attention_heads
@@ -196,6 +197,7 @@ class MiniMaxH3Attention(nn.Module):
             causal=False,
             supported_attention_backends=supported_attention_backends,
             prefix=prefix,
+            fa4_packed_varlen=fa4_packed_varlen,
         )
         self.to_gate_compress: ReplicatedLinear | None = None
         # None = unchecked; the first forward tests the loaded weight once and
@@ -465,6 +467,10 @@ class MiniMaxH3TransformerBlock(nn.Module):
             quant_config,
             prefix=f"{prefix}.attn",
             fuse_qknorm_rope=fuse_qknorm_rope,
+            # The packed multimodal document is a single long self-attention
+            # sequence. FA4's varlen API is substantially faster for this
+            # shape; the backend keeps grad/training and non-FA4 calls fixed.
+            fa4_packed_varlen=True,
         )
         self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps)
         self.ff = MiniMaxH3FeedForward(

fastvideo/tests/attention/test_flash_attn_h3_packed_varlen.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Routing guards for MiniMax-H3's packed-varlen FA4 inference path."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+
+flash_attn_module = pytest.importorskip(
+    "fastvideo.attention.backends.flash_attn",
+    reason="no usable flash-attention package installed",
+    exc_type=ImportError,
+)
+FlashAttentionImpl = flash_attn_module.FlashAttentionImpl
+
+
+def _build_impl(*, fa4_packed_varlen: bool) -> FlashAttentionImpl:
+    return FlashAttentionImpl(
+        num_heads=2,
+        head_size=8,
+        causal=False,
+        softmax_scale=0.125,
+        num_kv_heads=2,
+        fa4_packed_varlen=fa4_packed_varlen,
+    )
+
+
+def _qkv(batch: int = 1) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    query = torch.randn(batch, 7, 2, 8)
+    return query, torch.randn_like(query), torch.randn_like(query)
+
+
+def test_h3_fa4_no_grad_uses_single_sequence_varlen(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls = []
+
+    def varlen(query, key, value, cu_q, cu_k, max_q, max_k, **kwargs):
+        calls.append((query.shape, key.shape, value.shape, cu_q.tolist(), cu_k.tolist(), max_q, max_k, kwargs))
+        return value + 1
+
+    def unexpected_fixed(*args, **kwargs):
+        raise AssertionError("H3 FA4 inference took the fixed-length path")
+
+    monkeypatch.setattr(flash_attn_module, "fa_version", "4")
+    monkeypatch.setattr(flash_attn_module, "flash_attn_varlen_func_compilable", varlen)
+    monkeypatch.setattr(flash_attn_module, "flash_attn_func_compilable", unexpected_fixed)
+    query, key, value = _qkv()
+
+    with torch.no_grad():
+        output = _build_impl(fa4_packed_varlen=True)._forward_impl(query, key, value, None)
+
+    torch.testing.assert_close(output, value + 1)
+    assert calls == [
+        (
+            torch.Size([7, 2, 8]),
+            torch.Size([7, 2, 8]),
+            torch.Size([7, 2, 8]),
+            [0, 7],
+            [0, 7],
+            7,
+            7,
+            {"dropout_p": 0.0, "softmax_scale": 0.125, "causal": False},
+        )
+    ]
+
+
+@pytest.mark.parametrize(
+    ("fa4_packed_varlen", "version", "batch", "requires_grad"),
+    [
+        (False, "4", 1, False),
+        (True, "3", 1, False),
+        (True, "4", 2, False),
+        (True, "4", 1, True),
+    ],
+)
+def test_h3_varlen_guard_preserves_fixed_path(
+    monkeypatch: pytest.MonkeyPatch,
+    fa4_packed_varlen: bool,
+    version: str,
+    batch: int,
+    requires_grad: bool,
+) -> None:
+    def unexpected_varlen(*args, **kwargs):
+        raise AssertionError("packed-varlen route escaped its H3 FA4 inference guard")
+
+    def fixed(query, key, value, **kwargs):
+        del key, value, kwargs
+        return query + 2
+
+    monkeypatch.setattr(flash_attn_module, "fa_version", version)
+    monkeypatch.setattr(flash_attn_module, "flash_attn_varlen_func_compilable", unexpected_varlen)
+    monkeypatch.setattr(flash_attn_module, "flash_attn_func_compilable", fixed)
+    query, key, value = _qkv(batch)
+    query.requires_grad_(requires_grad)
+
+    context = torch.enable_grad() if requires_grad else torch.no_grad()
+    with context:
+        output = _build_impl(fa4_packed_varlen=fa4_packed_varlen)._forward_impl(query, key, value, None)
+
+    torch.testing.assert_close(output, query + 2)
+    if requires_grad:
+        output.sum().backward()
+        torch.testing.assert_close(query.grad, torch.ones_like(query))

fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py

@@ -9,6 +9,42 @@ import torch
 from fastvideo.platforms import AttentionBackendEnum
 
 
+def test_only_packed_minimax_h3_blocks_enable_fa4_varlen(monkeypatch: pytest.MonkeyPatch) -> None:
+    import fastvideo.models.dits.minimax_h3 as h3
+
+    routed: list[bool] = []
+
+    class RecordingAttention(torch.nn.Module):
+
+        def __init__(self, *args, fa4_packed_varlen: bool = False, **kwargs) -> None:
+            super().__init__()
+            del args, kwargs
+            routed.append(fa4_packed_varlen)
+
+    monkeypatch.setattr(h3, "MiniMaxH3Attention", RecordingAttention)
+    common = dict(
+        hidden_size=8,
+        num_attention_heads=1,
+        attention_head_dim=8,
+        ffn_dim=16,
+        norm_eps=1e-5,
+        qk_norm_eps=1e-5,
+        supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, ),
+        quant_config=None,
+    )
+    h3.MiniMaxH3TransformerBlock(
+        **common,
+        time_embed_dim=4,
+        prefix="minimax_h3.transformer_blocks.0",
+    )
+    h3.MiniMaxH3TokenRefinerBlock(
+        **common,
+        prefix="minimax_h3.token_refiner.refiner_blocks.0",
+    )
+
+    assert routed == [True, False]
+
+
 @pytest.mark.parametrize(
     ("raw", "expected"),
     [