diff --git a/fastvideo/attention/backends/fp4_dense_qat.py b/fastvideo/attention/backends/fp4_dense_qat.py index d577cd637..b64ae4cf9 100644 --- a/fastvideo/attention/backends/fp4_dense_qat.py +++ b/fastvideo/attention/backends/fp4_dense_qat.py @@ -13,15 +13,18 @@ cut back. """ from __future__ import annotations +import re import types from typing import Any import torch from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics, load_qat_vsa_modules +from fastvideo.attention.utils.nvfp4_qk_precondition import precondition_qk TILE = 128 _ORIGINAL = "_fp4_dense_original_forward" +_BLOCK = re.compile(r"(?:^|\.)transformer_blocks\.(\d+)\.") def dense_block_lists(seq_len: int, batch: int, heads: int, reference: Any, @@ -37,8 +40,13 @@ def dense_block_lists(seq_len: int, batch: int, heads: int, reference: Any, class FP4DenseAttention: """Dense attention ``forward`` replacement: ``(q, k, v, metadata) -> [B, L, H, D]`` through the FP4 kernel.""" - def __init__(self, numerics: FP4AttentionNumerics, modules: tuple[Any, Any] | None = None) -> None: + def __init__(self, + numerics: FP4AttentionNumerics, + modules: tuple[Any, Any] | None = None, + qk_transform: tuple[bool, bool] = (False, False)) -> None: self.numerics = numerics + # (smooth K, Hadamard-rotate Q/K) applied before quantization, exactly as the deployed FA4 NVFP4 path does. + self.smooth_k_mean, self.hadamard = qk_transform self._reference, self._train = modules or load_qat_vsa_modules() self._lists: dict[tuple[int, int, int, str], tuple[torch.Tensor, ...]] = {} @@ -52,6 +60,7 @@ class FP4DenseAttention: if getattr(metadata, "packed_sequence", None) is not None: raise NotImplementedError("dense FP4 attention emulator does not support packed sequences") batch, seq_len, heads, _ = query.shape + query, key = precondition_qk(query, key, smooth_k=self.smooth_k_mean, hadamard=self.hadamard) pad = -seq_len % TILE q, k, v = (torch.nn.functional.pad(t.to(torch.bfloat16).transpose(1, 2), (0, 0, 0, pad)) for t in (query, key, value)) @@ -80,31 +89,58 @@ def _dense_impl_types() -> tuple[type, ...]: return (FlashAttentionImpl, SDPAImpl) -def dense_attention_impls(transformer: torch.nn.Module) -> list[Any]: - """Every dense (FlashAttention or SDPA) impl reached through the transformer's distributed attention layers.""" +def named_dense_attention_impls(transformer: torch.nn.Module) -> list[tuple[str, Any]]: + """``(module name, impl)`` for every dense (FlashAttention or SDPA) impl behind a distributed attention layer.""" types_ = _dense_impl_types() impls = [] - for module in transformer.modules(): + for name, module in transformer.named_modules(): impl = getattr(getattr(module, "distributed_attention", None), "attn_impl", None) if isinstance(impl, types_): - impls.append(impl) + impls.append((name, impl)) return impls +def dense_attention_impls(transformer: torch.nn.Module) -> list[Any]: + """Every dense (FlashAttention or SDPA) impl reached through the transformer's distributed attention layers.""" + return [impl for _, impl in named_dense_attention_impls(transformer)] + + +def block_index(module_name: str) -> int | None: + """``transformer_blocks.`` index of a module name (``None`` outside the block stack).""" + match = _BLOCK.search(module_name + ".") + return int(match.group(1)) if match else None + + def install_fp4_dense_attention(transformer: torch.nn.Module, numerics: FP4AttentionNumerics | None, - modules: tuple[Any, Any] | None = None) -> int: - """Route every dense attention layer through the FP4 emulator (``None`` restores the original forward).""" - override = FP4DenseAttention(numerics, modules) if numerics is not None else None - impls = dense_attention_impls(transformer) - for impl in impls: + modules: tuple[Any, Any] | None = None, + *, + qk_transform: tuple[bool, bool] = (False, False), + bf16_blocks: tuple[int, ...] = ()) -> int: + """Route dense attention layers through the FP4 emulator; returns how many now run FP4. + + ``numerics=None`` restores every original forward. ``bf16_blocks`` keeps the attention of those + ``transformer_blocks`` indices on its original (BF16) forward, matching a deployment that opts them out of NVFP4. + """ + override = FP4DenseAttention(numerics, modules, qk_transform) if numerics is not None else None + named = named_dense_attention_impls(transformer) + skip = frozenset(int(b) for b in bf16_blocks) + missing = skip - {block_index(name) for name, _ in named} + if missing: + raise ValueError(f"bf16_blocks {sorted(missing)} match no transformer_blocks attention layer") + installed = 0 + for name, impl in named: if _ORIGINAL not in impl.__dict__: impl.__dict__[_ORIGINAL] = impl.forward - if override is None: + if override is None or block_index(name) in skip: impl.forward = impl.__dict__[_ORIGINAL] else: impl.forward = types.MethodType(lambda self, q, k, v, meta, _o=override: _o(q, k, v, meta), impl) - return len(impls) + installed += 1 + return installed if override is not None else len(named) -__all__ = ["FP4DenseAttention", "dense_block_lists", "dense_attention_impls", "install_fp4_dense_attention"] +__all__ = [ + "FP4DenseAttention", "block_index", "dense_block_lists", "dense_attention_impls", "install_fp4_dense_attention", + "named_dense_attention_impls" +] diff --git a/fastvideo/attention/utils/nvfp4_qk_precondition.py b/fastvideo/attention/utils/nvfp4_qk_precondition.py new file mode 100644 index 000000000..8766fa2ff --- /dev/null +++ b/fastvideo/attention/utils/nvfp4_qk_precondition.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Exact Q/K preconditioning ahead of NVFP4 (per-16 block) Q/K quantization. + +Both transforms leave the attention output unchanged in exact arithmetic and only redistribute what the FP4 blocks +have to represent: + +* ``smooth_k``: subtract each head's token-mean of K. ``q . k_mean`` is the same for every key of a query row, and + softmax is shift-invariant per row. +* ``hadamard``: rotate Q and K by the same orthonormal head-dim Hadamard matrix ``H`` (``Q H (K H)^T = Q K^T``), + which spreads outlier channels across the 16-wide scale blocks. + +The QAD emulator and the deployed FlashAttention-4 NVFP4 path apply the same function, so training and inference +quantize identical tensors. +""" +from __future__ import annotations + +import math +from functools import lru_cache + +import torch + +QK_TRANSFORMS = { + "none": (False, False), + "smooth_k": (True, False), + "hadamard": (False, True), + "smooth_k+hadamard": (True, True), +} + + +def parse_qk_transform(name: str | None) -> tuple[bool, bool]: + """``"smooth_k+hadamard"`` -> ``(smooth_k, hadamard)``; ``None``/``"none"`` -> no transform.""" + key = (name or "none").strip().lower() + if key not in QK_TRANSFORMS: + raise ValueError(f"unknown NVFP4 Q/K transform {name!r}; expected one of {', '.join(QK_TRANSFORMS)}") + return QK_TRANSFORMS[key] + + +@lru_cache(maxsize=8) +def _hadamard(n: int, device: torch.device) -> torch.Tensor: + if n <= 0 or n & (n - 1): + raise ValueError(f"Hadamard rotation needs a power-of-two head dim, got {n}") + h = torch.ones(1, 1) + while h.shape[0] < n: + h = torch.cat([torch.cat([h, h], 1), torch.cat([h, -h], 1)], 0) + return (h / math.sqrt(n)).to(device) + + +def hadamard_matrix(n: int, device: torch.device | str) -> torch.Tensor: + """Orthonormal Sylvester Hadamard matrix of size ``n`` (fp32).""" + return _hadamard(n, torch.device(device)) + + +def precondition_qk(query: torch.Tensor, key: torch.Tensor, *, smooth_k: bool, + hadamard: bool) -> tuple[torch.Tensor, torch.Tensor]: + """Transform ``[B, L, H, D]`` Q/K (computed in fp32, returned in the input dtype).""" + if not (smooth_k or hadamard): + return query, key + q, k = query.float(), key.float() + if smooth_k: + k = k - k.mean(dim=1, keepdim=True) + if hadamard: + rotation = hadamard_matrix(q.shape[-1], q.device) + q, k = q @ rotation, k @ rotation + return q.to(query.dtype), k.to(key.dtype) + + +__all__ = ["QK_TRANSFORMS", "hadamard_matrix", "parse_qk_transform", "precondition_qk"] diff --git a/fastvideo/tests/train/models/test_pika_h3_flex_qad.py b/fastvideo/tests/train/models/test_pika_h3_flex_qad.py index 5ea28e0f7..28065b988 100644 --- a/fastvideo/tests/train/models/test_pika_h3_flex_qad.py +++ b/fastvideo/tests/train/models/test_pika_h3_flex_qad.py @@ -248,6 +248,77 @@ def test_install_fp4_dense_attention_covers_sdpa_fallback(): assert all(impl.forward == original for impl, original in zip(impls, originals, strict=True)) +@pytest.mark.parametrize("name", ["smooth_k", "hadamard", "smooth_k+hadamard"]) +def test_precondition_qk_leaves_attention_unchanged(name): + from fastvideo.attention.utils.nvfp4_qk_precondition import parse_qk_transform, precondition_qk + torch.manual_seed(0) + q, k, v = (torch.randn(2, 37, 3, 16, dtype=torch.float32) for _ in range(3)) + k = k + 4.0 # shared K bias that smoothing removes + smooth, hadamard = parse_qk_transform(name) + tq, tk = precondition_qk(q, k, smooth_k=smooth, hadamard=hadamard) + + def attend(a, b): + return torch.nn.functional.scaled_dot_product_attention(*(t.transpose(1, 2) for t in (a, b, v))).transpose(1, 2) + + torch.testing.assert_close(attend(tq, tk), attend(q, k), atol=1e-5, rtol=1e-5) + if smooth: + assert float(tk.mean(dim=1).abs().max()) < 1e-4 + + +def test_parse_qk_transform_rejects_unknown_names(): + from fastvideo.attention.utils.nvfp4_qk_precondition import parse_qk_transform + assert parse_qk_transform(None) == (False, False) + with pytest.raises(ValueError, match="unknown NVFP4 Q/K transform"): + parse_qk_transform("rotate") + + +def test_fp4_dense_attention_with_qk_transform_matches_dense_attention(): + from fastvideo.attention.backends.fp4_dense_qat import FP4DenseAttention + from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics + reference = _reference_module() + kernel = _ExactKernel(reference) + attention = FP4DenseAttention(FP4AttentionNumerics(quantize=True, quantize_pv=False), (reference, kernel), + qk_transform=(True, True)) + torch.manual_seed(0) + q, k, v = (torch.randn(1, 200, 2, 16, dtype=torch.bfloat16) for _ in range(3)) + out = attention(q, k, v, None) + expected = torch.nn.functional.scaled_dot_product_attention(*(t.float().transpose(1, 2) for t in (q, k, v))) + torch.testing.assert_close(out.float(), expected.transpose(1, 2), atol=3e-2, rtol=3e-2) + + +def test_install_fp4_dense_attention_keeps_bf16_blocks_on_original_forward(): + from fastvideo.attention.backends.fp4_dense_qat import block_index, install_fp4_dense_attention + from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics + from fastvideo.attention.backends.sdpa import SDPAImpl + + class Block(torch.nn.Module): + + def __init__(self): + super().__init__() + self.attn = torch.nn.Module() + self.attn.distributed_attention = SimpleNamespace( + attn_impl=SDPAImpl(1, 8, causal=False, softmax_scale=8**-0.5)) + + model = torch.nn.Module() + model.transformer_blocks = torch.nn.ModuleList([Block(), Block(), Block()]) + impls = [block.attn.distributed_attention.attn_impl for block in model.transformer_blocks] + originals = [impl.forward for impl in impls] + reference = _reference_module() + kernel = _ExactKernel(reference) + numerics = FP4AttentionNumerics(quantize_pv=False) + assert block_index("transformer_blocks.12._checkpoint_wrapped_module.attn") == 12 + assert block_index("token_refiner.blocks.0.attn") is None + assert install_fp4_dense_attention(model, numerics, (reference, kernel), bf16_blocks=(2, )) == 2 + assert impls[2].forward == originals[2] and impls[0].forward != originals[0] + q = torch.randn(1, 130, 1, 8, dtype=torch.bfloat16) + impls[2].forward(q, q, q, None) + assert kernel.calls == [] + impls[1].forward(q, q, q, None) + assert len(kernel.calls) == 1 + with pytest.raises(ValueError, match="match no transformer_blocks"): + install_fp4_dense_attention(model, numerics, (reference, kernel), bf16_blocks=(7, )) + + def _write_trunk(trunk: Path, weights: dict[str, torch.Tensor]) -> None: trunk.mkdir(parents=True) names = sorted(weights) diff --git a/fastvideo/train/models/minimax_h3/flex.py b/fastvideo/train/models/minimax_h3/flex.py index 50f577ab5..d57b958fe 100644 --- a/fastvideo/train/models/minimax_h3/flex.py +++ b/fastvideo/train/models/minimax_h3/flex.py @@ -173,12 +173,18 @@ class MiniMaxH3FlexRef2VAModel(MiniMaxH3OmniRefPDDModel): len(self.qad_linears), self.fp4_numerics or "BF16 dense", self.num_trainable / 1e9) def _install_fp4_attention(self, config: dict[str, Any]) -> None: - """Dense FP4 attention as deployed by FlashAttention-4 NVFP4: Q/K in NVFP4, P and V in BF16 by default.""" + """Dense FP4 attention as deployed by FlashAttention-4 NVFP4: Q/K in NVFP4, P and V in BF16 by default. + + ``qk_transform`` (``none``/``smooth_k``/``hadamard``/``smooth_k+hadamard``) preconditions Q/K exactly before + quantization; ``bf16_blocks`` keeps those blocks' attention in BF16. Deployment must use the same settings. + """ self.fp4_numerics = None + self.fp4_install: dict[str, Any] = {} if not config.get("enabled", False): return from fastvideo.attention.backends.fp4_dense_qat import install_fp4_dense_attention from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics + from fastvideo.attention.utils.nvfp4_qk_precondition import parse_qk_transform floor = config.get("first_block_max_floor") self.fp4_numerics = FP4AttentionNumerics(quantize=True, @@ -186,10 +192,15 @@ class MiniMaxH3FlexRef2VAModel(MiniMaxH3OmniRefPDDModel): two_level_p=bool(config.get("two_level_p", False)), smooth_k=bool(config.get("smooth_k", False)), first_block_max_floor=None if floor is None else float(floor)) - installed = install_fp4_dense_attention(self.transformer, self.fp4_numerics) + self.fp4_install = { + "qk_transform": parse_qk_transform(config.get("qk_transform")), + "bf16_blocks": tuple(int(b) for b in config.get("bf16_blocks") or ()), + } + installed = install_fp4_dense_attention(self.transformer, self.fp4_numerics, **self.fp4_install) if installed == 0: raise ValueError("qad.fp4_attention needs dense (FlashAttention or SDPA) attention layers; none were found") - logger.info("Dense FP4 attention emulator installed on %d attention layers", installed) + logger.info("Dense FP4 attention emulator installed on %d attention layers (qk_transform=%s, bf16_blocks=%s)", + installed, config.get("qk_transform") or "none", list(self.fp4_install["bf16_blocks"])) @contextlib.contextmanager def quantization(self, *, linears: bool = True, attention: bool = True) -> Iterator[None]: @@ -206,7 +217,7 @@ class MiniMaxH3FlexRef2VAModel(MiniMaxH3OmniRefPDDModel): finally: set_nvfp4_qad_enabled(self.qad_linears, True) if fp4 is not None: - install_fp4_dense_attention(self.transformer, fp4) + install_fp4_dense_attention(self.transformer, fp4, **self.fp4_install) # ------------------------------------------------------------------ Flex ladder @property