nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[feat]: exact smooth-K / Hadamard Q/K preconditioning and BF16 blocks for dense FP4 attention QAD

Browse Source
This commit is contained in:
Aryan Kumar
2026-11-01 01:51:08 -07:00
parent 90be006f7e
commit 3d4f02673d
4 changed files with 202 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/attention/backends/fp4_dense_qat.py
+49 -13
View File
@@ -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.<i>`` 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"
]
fastvideo/attention/utils/nvfp4_qk_precondition.py
+67
View File
@@ -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"]
fastvideo/tests/train/models/test_pika_h3_flex_qad.py
+71
View File
@@ -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)
fastvideo/train/models/minimax_h3/flex.py
+15 -4
View File
@@ -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