[feat]: NVFP4 FA4 deploy path applies exact Q/K preconditioning and per-block BF16 attention opt-out

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)

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)

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"]

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"),

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(

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

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}")