diff --git a/fastvideo/attention/backends/flash_attn.py b/fastvideo/attention/backends/flash_attn.py index 1f5fbc949..30a3f67a2 100644 --- a/fastvideo/attention/backends/flash_attn.py +++ b/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 diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index fa29e5304..98e78d5af 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/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( diff --git a/fastvideo/tests/attention/test_flash_attn_h3_packed_varlen.py b/fastvideo/tests/attention/test_flash_attn_h3_packed_varlen.py new file mode 100644 index 000000000..634ccb753 --- /dev/null +++ b/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)) diff --git a/fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py b/fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py index 176974108..8a543ea88 100644 --- a/fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py +++ b/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"), [