From 47bb2539f84d21a4e05e5b648d864485b80ae8ad Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Sun, 1 Nov 2026 02:22:35 -0800 Subject: [PATCH] [feat]: NVFP4 FA4 deploy path applies exact Q/K preconditioning and per-block BF16 attention opt-out --- .../inference/basic/basic_pika_h3_flex.py | 32 ++++- fastvideo/attention/backends/flash_attn.py | 12 +- .../attention/utils/nvfp4_qk_precondition.py | 14 +- fastvideo/envs.py | 13 ++ fastvideo/models/dits/minimax_h3.py | 13 ++ .../test_nvfp4_qk_transform_deploy.py | 129 ++++++++++++++++++ .../check_fp4_attention_parity.py | 25 ++-- 7 files changed, 226 insertions(+), 12 deletions(-) create mode 100644 fastvideo/tests/attention/test_nvfp4_qk_transform_deploy.py diff --git a/examples/inference/basic/basic_pika_h3_flex.py b/examples/inference/basic/basic_pika_h3_flex.py index ef59480e0..3b3d39400 100644 --- a/examples/inference/basic/basic_pika_h3_flex.py +++ b/examples/inference/basic/basic_pika_h3_flex.py @@ -17,6 +17,12 @@ parallel VAE tiles) before the generator starts. ``--nvfp4-decoder`` and ``--nvfp4-dit h3_dit`` loads a packed NVFP4 DiT export (e.g. FastVideo/Pika-H3-Flex-NVFP4-QAD). They change weights, not the six-forward sampler. Kernel/VSA work is separate. +``--nvfp4-attention`` turns on FlashAttention-4 NVFP4 attention (Q/K in NVFP4, P/V in BF16; Blackwell and FA4 only) +with the named exact Q/K preconditioning, which must match the ``qk_transform`` the checkpoint was QAD-trained with +(``smooth_k+hadamard`` for Pika-H3-Flex-NVFP4-QAD). ``--nvfp4-attention-bf16-blocks 45,49`` keeps those +``transformer_blocks`` on BF16 attention, as in training. ``none`` quantizes Q/K without preconditioning. +Without ``--nvfp4-attention`` attention stays BF16. + ``--resolution`` (480p/768p), ``--aspect`` (one of Pika's eleven buckets), and ``--duration`` (5-15 s) pick the bucket and the ``17n + 5`` frame count; ``--height``/``--width``/``--num-frames`` override them with explicit values. @@ -48,7 +54,7 @@ import os from pathlib import Path from typing import Any -from fastvideo import VideoGenerator +from fastvideo import VideoGenerator, envs from fastvideo.api import ( CompileConfig, ComponentConfig, @@ -63,6 +69,7 @@ from fastvideo.api import ( QuantizationConfig, SamplingConfig, ) +from fastvideo.attention.utils.nvfp4_qk_precondition import QK_TRANSFORMS, parse_block_indices from fastvideo.models.vision_utils import load_image from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference from fastvideo.pipelines.basic.pika_h3_flex.buckets import ( @@ -135,8 +142,20 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--nvfp4-dit", choices=("h3_dit", "h3_dit_ffn"), help="load a packed NVFP4 DiT export with this layer profile (h3_dit: attention + FFN linears)") + parser.add_argument("--nvfp4-attention", + choices=tuple(QK_TRANSFORMS), + help="FA4 NVFP4 Q/K attention with this exact Q/K preconditioning (default: BF16 attention)") + parser.add_argument("--nvfp4-attention-bf16-blocks", + default="", + help="comma-separated transformer_blocks indices that keep BF16 attention, e.g. 45,49") parser.add_argument("--output", default="outputs/pika_h3_flex") args = parser.parse_args() + if args.nvfp4_attention_bf16_blocks and not args.nvfp4_attention: + parser.error("--nvfp4-attention-bf16-blocks needs --nvfp4-attention") + try: + parse_block_indices(args.nvfp4_attention_bf16_blocks) + except ValueError as error: + parser.error(f"--nvfp4-attention-bf16-blocks: {error}") if args.task in ("i2va", "fl2va") and not args.image: parser.error(f"--task {args.task} needs --image") @@ -210,8 +229,19 @@ def _pipeline(args: argparse.Namespace) -> PipelineSelection: ) +def _apply_nvfp4_attention(args: argparse.Namespace) -> None: + """Workers inherit the environment, so set it before VideoGenerator starts.""" + if not args.nvfp4_attention: + return + envs.FASTVIDEO_NVFP4_FA4.set(True) + envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.set(args.nvfp4_attention) + envs.FASTVIDEO_NVFP4_FA4_BF16_BLOCKS.set(args.nvfp4_attention_bf16_blocks) + envs.setdefault_external("CUTE_DSL_ENABLE_TVM_FFI", "1") + + def main() -> None: args = parse_args() + _apply_nvfp4_attention(args) if args.lossless_accel: for key, value in LOSSLESS_ACCEL_ENV.items(): os.environ.setdefault(key, value) diff --git a/fastvideo/attention/backends/flash_attn.py b/fastvideo/attention/backends/flash_attn.py index 8a9a5e29e..919ad3593 100644 --- a/fastvideo/attention/backends/flash_attn.py +++ b/fastvideo/attention/backends/flash_attn.py @@ -27,6 +27,7 @@ from fastvideo.attention.backends.abstract import ( PackedSequenceMetadata, ) from fastvideo.attention.backends.attn_qat_infer import _resolve_fa4_pv_mode +from fastvideo.attention.utils.nvfp4_qk_precondition import parse_qk_transform, precondition_qk from fastvideo.logger import init_logger logger = init_logger(__name__) @@ -279,6 +280,12 @@ class FlashAttentionImpl(AttentionImpl): if nvfp4_fa4 is None: nvfp4_fa4 = envs.FASTVIDEO_NVFP4_FA4.get() self.nvfp4_fa4 = bool(nvfp4_fa4) + # Exact Q/K preconditioning ahead of NVFP4 quantization, shared with the QAD emulator + # (fastvideo.attention.utils.nvfp4_qk_precondition). Same arg-over-env precedence as nvfp4_fa4. + qk_transform = extra_impl_args.get("nvfp4_qk_transform") + if qk_transform is None: + qk_transform = envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.get() + self.nvfp4_smooth_k, self.nvfp4_hadamard = parse_qk_transform(qk_transform) # PV-mode knob for the FA4-FP4 path. Resolved + validated # unconditionally (shared helper) so a typo fails at construction, not # at the first forward on a Blackwell box, and so the ATTN_QAT_INFER @@ -291,7 +298,8 @@ class FlashAttentionImpl(AttentionImpl): assert cap in [(10, 0), (10, 3)], (f"NVFP4 FA4 requires Blackwell (sm100a/sm103a), got sm{cap[0]}{cap[1]}") assert _FA4_FP4_AVAILABLE, ("NVFP4 FA4 requires flash-attention-fp4 (flash_attn.cute). " "Install via instructions in docs/inference/optimizations.md") - logger.info("NVFP4 FA4 enabled for FlashAttentionImpl (quant_qk, pv_mode=%s)", self.fa4_pv_mode) + logger.info("NVFP4 FA4 enabled for FlashAttentionImpl (quant_qk, pv_mode=%s, smooth_k=%s, hadamard=%s)", + self.fa4_pv_mode, self.nvfp4_smooth_k, self.nvfp4_hadamard) def forward( self, @@ -522,6 +530,8 @@ class FlashAttentionImpl(AttentionImpl): orig_seqlen_q = query.shape[1] orig_seqlen_k = key.shape[1] + query, key = precondition_qk(query, key, smooth_k=self.nvfp4_smooth_k, hadamard=self.nvfp4_hadamard) + # Quantize Q/K to FP4 (internally pads to multiple of 128 for SF layout) q_fp4, q_sf = _nvfp4_quantize_for_fa4(query) k_fp4, k_sf = _nvfp4_quantize_for_fa4(key) diff --git a/fastvideo/attention/utils/nvfp4_qk_precondition.py b/fastvideo/attention/utils/nvfp4_qk_precondition.py index 8766fa2ff..e130b266e 100644 --- a/fastvideo/attention/utils/nvfp4_qk_precondition.py +++ b/fastvideo/attention/utils/nvfp4_qk_precondition.py @@ -35,6 +35,18 @@ def parse_qk_transform(name: str | None) -> tuple[bool, bool]: return QK_TRANSFORMS[key] +def parse_block_indices(raw: str | None) -> frozenset[int]: + """``"45, 49"`` -> ``frozenset({45, 49})``; empty/``None`` -> empty set.""" + items = [item.strip() for item in (raw or "").split(",") if item.strip()] + try: + indices = frozenset(int(item) for item in items) + except ValueError as error: + raise ValueError(f"expected a comma-separated list of block indices, got {raw!r}") from error + if any(index < 0 for index in indices): + raise ValueError(f"block indices must be non-negative, got {raw!r}") + return indices + + @lru_cache(maxsize=8) def _hadamard(n: int, device: torch.device) -> torch.Tensor: if n <= 0 or n & (n - 1): @@ -64,4 +76,4 @@ def precondition_qk(query: torch.Tensor, key: torch.Tensor, *, smooth_k: bool, return q.to(query.dtype), k.to(key.dtype) -__all__ = ["QK_TRANSFORMS", "hadamard_matrix", "parse_qk_transform", "precondition_qk"] +__all__ = ["QK_TRANSFORMS", "hadamard_matrix", "parse_block_indices", "parse_qk_transform", "precondition_qk"] diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 0173cfc3d..b93a124d1 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -368,6 +368,19 @@ FASTVIDEO_NVFP4_FA4 = EnvBool( category="attention", doc="FlashAttention-4 quantizes Q and K to NVFP4. An explicit nvfp4_fa4 attention implementation argument " "takes precedence.") +FASTVIDEO_NVFP4_FA4_QK_TRANSFORM = EnvChoice( + "none", + choices=("none", "smooth_k", "hadamard", "smooth_k+hadamard"), + category="attention", + doc="Exact Q/K preconditioning applied right before NVFP4 quantization on the FlashAttention-4 NVFP4 path: " + "smooth_k subtracts each head's token-mean of K, hadamard rotates Q and K by the head-dim Hadamard matrix. " + "Must match the qk_transform the model was QAD-trained with. An explicit nvfp4_qk_transform attention " + "implementation argument takes precedence.") +FASTVIDEO_NVFP4_FA4_BF16_BLOCKS = EnvStr( + "", + category="attention", + doc="Comma-separated MiniMax-H3 transformer_blocks indices whose attention stays in BF16 while " + "FASTVIDEO_NVFP4_FA4 quantizes the rest, for example 45,49.") FASTVIDEO_FA4_PV_MODE = EnvChoice( "bf16", choices=("bf16", "fp8"), diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index faa70ce60..8662d058c 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -19,6 +19,7 @@ from fastvideo.attention.backends.abstract import AttentionMetadata, layer_idx_f from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadata from fastvideo.attention.layer import DistributedAttention_VSA from fastvideo.attention.selector import get_attn_backend +from fastvideo.attention.utils.nvfp4_qk_precondition import parse_block_indices from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config from fastvideo.distributed.communication_op import ( sequence_model_parallel_all_gather_with_unpad, @@ -207,6 +208,7 @@ class MiniMaxH3Attention(nn.Module): fa4_packed_varlen: bool = False, exact_rope: bool = False, vsa_fp4: dict[str, Any] | None = None, + bf16_attention: bool = False, ) -> None: super().__init__() self.num_attention_heads = num_attention_heads @@ -254,6 +256,8 @@ class MiniMaxH3Attention(nn.Module): supported_attention_backends=supported_attention_backends) use_vsa = resolved_backend.get_name() == "VIDEO_SPARSE_ATTN_H3" attention_cls = DistributedAttention_VSA if use_vsa else DistributedAttention + # Blocks listed in FASTVIDEO_NVFP4_FA4_BF16_BLOCKS opt out of the process-wide NVFP4 FA4 Q/K quantization. + precision_args: dict[str, Any] = {"nvfp4_fa4": False} if bf16_attention else {} self.distributed_attention = attention_cls( num_heads=num_attention_heads, head_size=attention_head_dim, @@ -261,6 +265,7 @@ class MiniMaxH3Attention(nn.Module): supported_attention_backends=supported_attention_backends, prefix=prefix, fa4_packed_varlen=fa4_packed_varlen, + **precision_args, ) # Opt-in inference route: VSA-H3 selection on the block-sparse FP4 # kernel (see minimax_h3_vsa_fp4); grad and compile keep the generic path. @@ -658,6 +663,7 @@ class MiniMaxH3TransformerBlock(nn.Module): fa4_packed_varlen: bool = False, exact_kernels: frozenset[str] = frozenset(), vsa_fp4: dict[str, Any] | None = None, + bf16_attention: bool = False, ) -> None: super().__init__() self.norm1 = nn.RMSNorm(hidden_size, eps=norm_eps) @@ -673,6 +679,7 @@ class MiniMaxH3TransformerBlock(nn.Module): fa4_packed_varlen=fa4_packed_varlen, exact_rope="rope" in exact_kernels, vsa_fp4=vsa_fp4, + bf16_attention=bf16_attention, ) self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps) self.ff = MiniMaxH3FeedForward( @@ -916,6 +923,11 @@ class MiniMaxH3Transformer3DModel(BaseDiT): config.quant_config, prefix=f"{config.prefix}.token_refiner", ) + bf16_blocks = parse_block_indices(envs.FASTVIDEO_NVFP4_FA4_BF16_BLOCKS.get()) + unknown_blocks = sorted(index for index in bf16_blocks if index >= arch.num_layers) + if unknown_blocks: + raise ValueError(f"FASTVIDEO_NVFP4_FA4_BF16_BLOCKS {unknown_blocks} exceed the {arch.num_layers} " + "transformer_blocks") self.transformer_blocks = nn.ModuleList([ MiniMaxH3TransformerBlock( arch.hidden_size, @@ -935,6 +947,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT): fa4_packed_varlen=envs.FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN.get(), exact_kernels=self.enabled_exact_kernels, vsa_fp4=config.vsa_fp4, + bf16_attention=index in bf16_blocks, ) for index in range(arch.num_layers) ]) self.norm_out = MiniMaxH3AdaLayerNormOut( diff --git a/fastvideo/tests/attention/test_nvfp4_qk_transform_deploy.py b/fastvideo/tests/attention/test_nvfp4_qk_transform_deploy.py new file mode 100644 index 000000000..36982f804 --- /dev/null +++ b/fastvideo/tests/attention/test_nvfp4_qk_transform_deploy.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Deployed NVFP4 FA4 Q/K preconditioning and per-block BF16 opt-out. + +The transform must be applied to Q/K right before NVFP4 quantization with the same function the QAD emulator uses, +and MiniMax-H3 blocks listed in ``FASTVIDEO_NVFP4_FA4_BF16_BLOCKS`` must build BF16 attention impls. +""" +from unittest.mock import MagicMock, patch + +import pytest + +torch = pytest.importorskip("torch") + +import fastvideo.envs as envs # noqa: E402 +from fastvideo.attention.utils.nvfp4_qk_precondition import ( # noqa: E402 + parse_block_indices, parse_qk_transform, precondition_qk, +) + +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 _impl(**extra) -> "FlashAttentionImpl": + return FlashAttentionImpl(num_heads=2, head_size=64, causal=False, softmax_scale=0.125, num_kv_heads=2, **extra) + + +@pytest.fixture +def fa4_runtime_available(): + with ( + patch.object(flash_attn_module, "_FA4_FP4_AVAILABLE", True), + patch("torch.cuda.get_device_capability", return_value=(10, 0)), + ): + yield + + +def test_block_list_parsing() -> None: + assert parse_block_indices("") == frozenset() + assert parse_block_indices(None) == frozenset() + assert parse_block_indices(" 45, 49,45 ") == frozenset({45, 49}) + for bad in ("a,1", "-1", "4.5"): + with pytest.raises(ValueError): + parse_block_indices(bad) + + +def test_transform_env_default_and_validation(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.override(None)) + assert envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.get() == "none" + assert _impl().nvfp4_smooth_k is False and _impl().nvfp4_hadamard is False + env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.override("Smooth_K+Hadamard")) + impl = _impl() + assert (impl.nvfp4_smooth_k, impl.nvfp4_hadamard) == (True, True) + env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.override("bogus")) + with pytest.raises(envs.EnvVarError): + envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.get() + + +def test_impl_arg_wins_over_transform_env(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4_QK_TRANSFORM.override("hadamard")) + impl = _impl(nvfp4_qk_transform="smooth_k") + assert (impl.nvfp4_smooth_k, impl.nvfp4_hadamard) == (True, False) + with pytest.raises(ValueError): + _impl(nvfp4_qk_transform="bogus") + + +@pytest.mark.parametrize("name", ["none", "smooth_k", "hadamard", "smooth_k+hadamard"]) +def test_transform_applied_before_quantization(name, env_overrides, fa4_runtime_available) -> None: + env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4.override(None)) + impl = _impl(nvfp4_fa4=True, nvfp4_qk_transform=name) + torch.manual_seed(0) + q = torch.randn(1, 20, 2, 64, dtype=torch.bfloat16) + 3.0 + k = torch.randn(1, 20, 2, 64, dtype=torch.bfloat16) + 3.0 + v = torch.randn(1, 20, 2, 64, dtype=torch.bfloat16) + seen = [] + + def fake_quant(tensor): + seen.append(tensor.clone()) + return tensor, tensor + + fp4 = MagicMock(side_effect=lambda *a, **kw: a[2]) + with ( + patch.object(flash_attn_module, "_nvfp4_quantize_for_fa4", side_effect=fake_quant), + patch.object(flash_attn_module, "flash_attn_fp4_func", fp4), + ): + impl._forward_nvfp4(q, k, v) + + smooth_k, hadamard = parse_qk_transform(name) + expected_q, expected_k = precondition_qk(q, k, smooth_k=smooth_k, hadamard=hadamard) + assert torch.equal(seen[0], expected_q) and torch.equal(seen[1], expected_k) + assert (name == "none") == (torch.equal(seen[0], q) and torch.equal(seen[1], k)) + + +def test_smooth_k_removes_mean_and_preserves_attention() -> None: + torch.manual_seed(1) + q, k, v = (torch.randn(1, 32, 2, 64) for _ in range(3)) + k = k + 5.0 + qt, kt = precondition_qk(q, k, smooth_k=True, hadamard=True) + ref = torch.nn.functional.scaled_dot_product_attention(*(t.transpose(1, 2) for t in (q, k, v))) + out = torch.nn.functional.scaled_dot_product_attention(*(t.transpose(1, 2) for t in (qt, kt, v))) + torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4) + + +def _build_h3_attention(bf16_attention: bool): + h3 = pytest.importorskip("fastvideo.models.dits.minimax_h3") + backend = MagicMock() + backend.get_name.return_value = "FLASH_ATTN" + distributed = MagicMock() + with ( + patch.object(h3, "get_attn_backend", return_value=backend), + patch.object(h3, "get_compute_dtype", return_value=torch.bfloat16), + patch.object(h3, "ReplicatedLinear", side_effect=lambda *a, **kw: torch.nn.Identity()), + patch.object(h3, "DistributedAttention", distributed), + ): + h3.MiniMaxH3Attention(64, 2, 32, 1e-6, (), None, prefix="transformer_blocks.3.attn", + bf16_attention=bf16_attention) + return distributed.call_args.kwargs + + +def test_h3_attention_bf16_opt_out_wiring() -> None: + assert _build_h3_attention(True)["nvfp4_fa4"] is False + assert "nvfp4_fa4" not in _build_h3_attention(False) + + +def test_bf16_block_impl_ignores_env_opt_in(env_overrides, fa4_runtime_available) -> None: + env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4.override(True)) + assert _impl(nvfp4_fa4=_build_h3_attention(True)["nvfp4_fa4"]).nvfp4_fa4 is False + assert _impl(**{k: v for k, v in _build_h3_attention(False).items() if k == "nvfp4_fa4"}).nvfp4_fa4 is True # env wins diff --git a/scripts/distill/pika_h3_flex_qad/check_fp4_attention_parity.py b/scripts/distill/pika_h3_flex_qad/check_fp4_attention_parity.py index 87a358756..134fe5f6d 100644 --- a/scripts/distill/pika_h3_flex_qad/check_fp4_attention_parity.py +++ b/scripts/distill/pika_h3_flex_qad/check_fp4_attention_parity.py @@ -1,11 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 """Parity of the dense FP4 attention emulator (QAD training) with FlashAttention-4 NVFP4 (deployment). -Both should compute Q/K in NVFP4 (per-16 E4M3 block scales, unit global scale) with P and V in BF16. Reports the -relative L2 distance of the emulator to the FA4 NVFP4 kernel next to the FP4-vs-BF16 distance: the emulator is a +With ``--qk-transform`` both sides precondition Q/K first (the emulator through ``qk_transform=``, the deployed +``FlashAttentionImpl._forward_nvfp4`` through the ``nvfp4_qk_transform`` impl arg). Both should compute Q/K in NVFP4 +(per-16 E4M3 block scales, unit global scale) with P and V in BF16. Reports the relative L2 distance of the emulator to the FA4 NVFP4 kernel next to the FP4-vs-BF16 distance: the emulator is a faithful training target when the first is a small fraction of the second. Needs a Blackwell GPU with FA4 FP4. - python scripts/distill/pika_h3_flex_qad/check_fp4_attention_parity.py --seq-len 8269 --heads 14 + python scripts/distill/pika_h3_flex_qad/check_fp4_attention_parity.py --seq-len 8269 --heads 14 \\ + --qk-transform smooth_k+hadamard """ from __future__ import annotations @@ -23,28 +25,33 @@ def main() -> None: parser.add_argument("--seq-len", type=int, default=8269, help="not a multiple of 128 on purpose") parser.add_argument("--heads", type=int, default=14) parser.add_argument("--dim", type=int, default=128) + parser.add_argument("--qk-transform", default="none", help="none, smooth_k, hadamard or smooth_k+hadamard") + parser.add_argument("--k-bias", type=float, default=4.0, help="per-channel K offset that smooth_k removes") parser.add_argument("--outlier-scale", type=float, default=8.0, help="scale of a few outlier K channels") args = parser.parse_args() from fastvideo.attention.backends import flash_attn as fa from fastvideo.attention.backends.fp4_dense_qat import FP4DenseAttention from fastvideo.attention.backends.fp4_vsa_qat import FP4AttentionNumerics + from fastvideo.attention.utils.nvfp4_qk_precondition import parse_qk_transform torch.manual_seed(0) shape = (1, args.seq_len, args.heads, args.dim) q, k, v = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(3)) k[..., :4] *= args.outlier_scale + k += args.k_bias * torch.randn(1, 1, args.heads, args.dim, device="cuda", dtype=torch.bfloat16) + transform = parse_qk_transform(args.qk_transform) scale = args.dim**-0.5 - q_fp4, q_sf = fa._nvfp4_quantize_for_fa4(q) - k_fp4, k_sf = fa._nvfp4_quantize_for_fa4(k) - fa4_fp4 = fa.flash_attn_fp4_func(q_fp4[:, :args.seq_len], k_fp4[:, :args.seq_len], v, q_sf, k_sf, - softmax_scale=scale, causal=False) - fa4_fp4 = fa4_fp4[0] if isinstance(fa4_fp4, tuple) else fa4_fp4 + impl = fa.FlashAttentionImpl(num_heads=args.heads, head_size=args.dim, causal=False, softmax_scale=scale, + num_kv_heads=args.heads, nvfp4_fa4=True, nvfp4_qk_transform=args.qk_transform) + fa4_fp4 = impl._forward_nvfp4(q, k, v) bf16 = torch.nn.functional.scaled_dot_product_attention(*(t.transpose(1, 2) for t in (q, k, v)), scale=scale).transpose(1, 2) - emulator = FP4DenseAttention(FP4AttentionNumerics(quantize=True, quantize_pv=False))(q, k, v, None) + numerics = FP4AttentionNumerics(quantize=True, quantize_pv=False) + emulator = FP4DenseAttention(numerics, qk_transform=transform)(q, k, v, None) + print(f"qk_transform={args.qk_transform} k_bias={args.k_bias} outlier_scale={args.outlier_scale}") fp4_vs_bf16 = rel(fa4_fp4, bf16) emu_vs_fa4 = rel(emulator, fa4_fp4) print(f"FA4 NVFP4 vs BF16 attention: rel L2 {fp4_vs_bf16:.5f}")