[feat] Pika H3 Flex NVFP4 QAD inference: packed NVFP4 DiT, FP4 attention and --weights presets
This commit is contained in:
Aryan Kumar
@@ -766,7 +766,7 @@
|
||||
"stage": "inference",
|
||||
"task": "Text to video (with audio)",
|
||||
"label": "Pika H3 Flex",
|
||||
"summary": "Pika's post-trained MiniMax-H3: synchronized video and stereo audio from text in six forwards, at 480p or 768p, 5-15 s, with optional NVFP4 decoder and text encoder on Blackwell.",
|
||||
"summary": "Pika's post-trained MiniMax-H3: synchronized video and stereo audio from text in six forwards, at 480p or 768p, 5-15 s, in BF16 or as FastVideo's NVFP4 QAD release on Blackwell, with optional FP4 attention.",
|
||||
"model": "Pika-Labs/Pika-H3-Flex",
|
||||
"source": "examples/inference/basic/basic_pika_h3_flex.py",
|
||||
"command": "FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN FASTVIDEO_FA4=1 python examples/inference/basic/basic_pika_h3_flex.py --task t2va --offload resident --prompt 'A paper boat drifts down a rainy city street at night; rain patters on the water.'",
|
||||
@@ -783,6 +783,14 @@
|
||||
"expected_artifact": "MP4 with synchronized stereo audio at outputs/pika_h3_flex/pika_h3_flex_t2va.mp4",
|
||||
"modes": ["T2VA"],
|
||||
"knobs": [
|
||||
{
|
||||
"key": "weights",
|
||||
"label": "Weights",
|
||||
"hint": "NVFP4 QAD needs Blackwell",
|
||||
"flag": "--weights",
|
||||
"options": [{"value": "bf16", "label": "BF16 (Pika release)"}, {"value": "nvfp4-qad", "label": "NVFP4 QAD"}, {"value": "nvfp4-qad-fp4-attention", "label": "NVFP4 QAD + FP4 attention"}],
|
||||
"default": "bf16"
|
||||
},
|
||||
{
|
||||
"key": "resolution",
|
||||
"label": "Resolution",
|
||||
@@ -855,6 +863,7 @@
|
||||
}
|
||||
],
|
||||
"limitations": [
|
||||
"NVFP4 QAD release (--weights), one GB200, 768p 16:9, 5 s, warm: 16.8 s with BF16 attention and 15.2 s with FP4 attention; at 768p, 15 s, FP4 attention is 16% faster than BF16 attention (85.3 s vs 101.8 s).",
|
||||
"Measured on one GB200 at 768p 16:9, 5 s, warm: official pika-h3-pro CLI 59.7 s; FastVideo BF16 26.2 s; with the NVFP4 quality-36L decoder 18.0 s; with the fast-8L decoder 17.4 s. The 36L decoder measured 40.5 dB PSNR against the BF16 decoder.",
|
||||
"Peak resident memory was about 129 GiB for the all-BF16 run and 104.6 GiB with the NVFP4 text encoder (25.49 s) at 768p / 5 s on GB200; the official CLI's sampling peak was 71.7 GiB, a different measurement. The NVFP4 text encoder's first request in a fresh process pays a one-time cost of about 183 s.",
|
||||
"--offload resident keeps every component on the GPU (104.6-108.5 GiB peak with the NVFP4 text encoder, about 129 GiB all-BF16, on GB200). --offload lazy loads components on demand for smaller devices; on GB200 it took 48.06 s for the same T2VA request.",
|
||||
@@ -887,6 +896,14 @@
|
||||
"expected_artifact": "MP4 with synchronized stereo audio at outputs/pika_h3_flex/pika_h3_flex_<task>.mp4",
|
||||
"modes": ["I2VA", "L2VA", "FL2VA"],
|
||||
"knobs": [
|
||||
{
|
||||
"key": "weights",
|
||||
"label": "Weights",
|
||||
"hint": "NVFP4 QAD needs Blackwell",
|
||||
"flag": "--weights",
|
||||
"options": [{"value": "bf16", "label": "BF16 (Pika release)"}, {"value": "nvfp4-qad", "label": "NVFP4 QAD"}, {"value": "nvfp4-qad-fp4-attention", "label": "NVFP4 QAD + FP4 attention"}],
|
||||
"default": "bf16"
|
||||
},
|
||||
{
|
||||
"key": "task",
|
||||
"label": "Task",
|
||||
@@ -1002,6 +1019,14 @@
|
||||
"expected_artifact": "MP4 with synchronized stereo audio at outputs/pika_h3_flex/pika_h3_flex_ref2va.mp4",
|
||||
"modes": ["Ref2VA"],
|
||||
"knobs": [
|
||||
{
|
||||
"key": "weights",
|
||||
"label": "Weights",
|
||||
"hint": "NVFP4 QAD needs Blackwell",
|
||||
"flag": "--weights",
|
||||
"options": [{"value": "bf16", "label": "BF16 (Pika release)"}, {"value": "nvfp4-qad", "label": "NVFP4 QAD"}, {"value": "nvfp4-qad-fp4-attention", "label": "NVFP4 QAD + FP4 attention"}],
|
||||
"default": "bf16"
|
||||
},
|
||||
{
|
||||
"key": "resolution",
|
||||
"label": "Resolution",
|
||||
|
||||
@@ -200,6 +200,9 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
|
||||
| `FASTVIDEO_VSA_PTX_BLACKWELL` | bool | `0` | attention | VIDEO_SPARSE_ATTN_H3 sends tile-64 forwards to the data-center Blackwell (sm_100a / sm_103a) kernel; grad-tracking forwards need its autograd operator, else Triton. fastvideo-kernel reads the same variable with the same rule. Deprecated names: `FASTVIDEO_VSA_SM100A`. |
|
||||
| `FASTVIDEO_VSA_TRITON` | bool | `0` | attention | Force the Triton MiniMax-H3 sparse attention kernel. fastvideo-kernel reads the same variable. |
|
||||
| `FASTVIDEO_NVFP4_FA4` | bool | `0` | attention | FlashAttention-4 quantizes Q and K to NVFP4. An explicit nvfp4_fa4 attention implementation argument takes precedence. |
|
||||
| `FASTVIDEO_NVFP4_FA4_QK_TRANSFORM` | one of none, smooth_k, hadamard, smooth_k+hadamard | `none` | attention | 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_FUSED_QK_QUANT` | bool | `1` | attention | With FASTVIDEO_NVFP4_FA4_QK_TRANSFORM other than none, apply the transform and the NVFP4 quantization in one fused Triton kernel per tensor (same FP4 codes and scale factors as the unfused path). Set it to 0 to use the unfused PyTorch transform and flashinfer quantization. |
|
||||
| `FASTVIDEO_NVFP4_FA4_BF16_BLOCKS` | str | `""` | attention | 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` | one of bf16, fp8 | `bf16` | attention | V dtype on the NVFP4 FlashAttention-4 path (FLASH_ATTN with nvfp4_fa4, ATTN_QAT_INFER on sm_100a/sm_103a): bf16 keeps V in BF16; fp8 casts V to float8 e4m3 without scaling. An explicit fa4_pv_mode attention implementation argument takes precedence. |
|
||||
| `FASTVIDEO_DISABLE_ATTENTION_COMPILE` | bool | `1` | attention | Keep attention forward out of torch.compile graphs (torch.compiler.disable). Set it to 0 to let attention constructed under that setting be traced. Setting it explicitly to true also blocks regional compile. |
|
||||
| `FASTVIDEO_MLX_WINDOW` | int | `0` | attention | MLX FastWan windowed attention size in tokens. 0 uses full attention. |
|
||||
|
||||
@@ -14,7 +14,15 @@ does not expose a step-count flag. ``--lossless-accel`` turns on the existing
|
||||
training-free FastVideo switches (bit-identical kernels, reference memo,
|
||||
parallel VAE tiles) before the generator starts. ``--nvfp4-decoder`` and
|
||||
``--nvfp4-text-encoder`` opt into the existing Blackwell FP4 weight paths;
|
||||
they change weights, not the six-forward sampler. Kernel/VSA work is separate.
|
||||
``--weights nvfp4-qad`` (or ``nvfp4-qad-fp4-attention``) selects FastVideo's NVFP4 QAD release;
|
||||
``--nvfp4-dit h3_dit`` loads any 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;
|
||||
@@ -47,7 +55,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,
|
||||
@@ -59,8 +67,10 @@ from fastvideo.api import (
|
||||
OutputConfig,
|
||||
ParallelismConfig,
|
||||
PipelineSelection,
|
||||
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 (
|
||||
@@ -95,9 +105,31 @@ def _reference(value: str) -> MiniMaxH3Reference:
|
||||
return MiniMaxH3Reference(source=source, media_type=media_type)
|
||||
|
||||
|
||||
WEIGHT_PRESETS = {
|
||||
"bf16": {
|
||||
"model_path": "Pika-Labs/Pika-H3-Flex"
|
||||
},
|
||||
"nvfp4-qad": {
|
||||
"model_path": "FastVideo/Pika-H3-Flex-NVFP4-QAD",
|
||||
"nvfp4_dit": "h3_dit"
|
||||
},
|
||||
"nvfp4-qad-fp4-attention": {
|
||||
"model_path": "FastVideo/Pika-H3-Flex-NVFP4-FP4Attn-QAD",
|
||||
"nvfp4_dit": "h3_dit",
|
||||
"nvfp4_attention": "smooth_k+hadamard",
|
||||
"nvfp4_attention_bf16_blocks": "45,49"
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("--model-path", default="Pika-Labs/Pika-H3-Flex")
|
||||
parser.add_argument("--weights",
|
||||
choices=tuple(WEIGHT_PRESETS),
|
||||
default="bf16",
|
||||
help="bf16: Pika's release; nvfp4-qad: FastVideo's NVFP4 QAD release; "
|
||||
"nvfp4-qad-fp4-attention: NVFP4 QAD with FP4 attention (both need Blackwell)")
|
||||
parser.add_argument("--model-path", help="local model directory or Hub repo (overrides --weights)")
|
||||
parser.add_argument("--task", choices=TASKS, default="t2va")
|
||||
parser.add_argument("--prompt", required=True)
|
||||
parser.add_argument("--image", help="first frame (i2va/fl2va)")
|
||||
@@ -130,8 +162,25 @@ def parse_args() -> argparse.Namespace:
|
||||
help="Blackwell-only NVFP4 H3 decoder (quality-36L or fast-8L decoder-nvfp4.safetensors)")
|
||||
parser.add_argument("--nvfp4-text-encoder",
|
||||
help="Blackwell-only serialized NVFP4 Qwen3-VL text_encoder directory")
|
||||
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",
|
||||
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()
|
||||
for name, value in WEIGHT_PRESETS[args.weights].items():
|
||||
if getattr(args, name) is None:
|
||||
setattr(args, name, value)
|
||||
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")
|
||||
@@ -202,8 +251,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 or "")
|
||||
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)
|
||||
@@ -223,6 +283,8 @@ def main() -> None:
|
||||
parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus),
|
||||
offload=offload,
|
||||
compile=CompileConfig(enabled=args.compile),
|
||||
quantization=(QuantizationConfig(transformer_quant="NVFP4", layer_profile=args.nvfp4_dit)
|
||||
if args.nvfp4_dit else None),
|
||||
),
|
||||
pipeline=_pipeline(args),
|
||||
))
|
||||
|
||||
@@ -28,6 +28,12 @@ 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.attention.utils.nvfp4_fused_qk_quant import (
|
||||
fused_qk_quant_available,
|
||||
fused_qk_quant_supported,
|
||||
nvfp4_fused_qk_quantize_fa4,
|
||||
)
|
||||
from fastvideo.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -284,6 +290,12 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
if nvfp4_fa4 is None:
|
||||
nvfp4_fa4 = envs.FASTVIDEO_NVFP4_FA4.get()
|
||||
self.nvfp4_fa4 = bool(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)
|
||||
self.nvfp4_fused_qk_quant = ((self.nvfp4_smooth_k or self.nvfp4_hadamard)
|
||||
and envs.FASTVIDEO_NVFP4_FA4_FUSED_QK_QUANT.get() and fused_qk_quant_available())
|
||||
# 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
|
||||
@@ -301,7 +313,10 @@ 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, "
|
||||
"fused_qk_quant=%s)", self.fa4_pv_mode, self.nvfp4_smooth_k, self.nvfp4_hadamard,
|
||||
self.nvfp4_fused_qk_quant)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
@@ -537,8 +552,15 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
orig_seqlen_k = key.shape[1]
|
||||
|
||||
# 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)
|
||||
if self.nvfp4_fused_qk_quant and fused_qk_quant_supported(query, key):
|
||||
q_fp4, q_sf, k_fp4, k_sf = nvfp4_fused_qk_quantize_fa4(query,
|
||||
key,
|
||||
smooth_k=self.nvfp4_smooth_k,
|
||||
hadamard=self.nvfp4_hadamard)
|
||||
else:
|
||||
query, key = precondition_qk(query, key, smooth_k=self.nvfp4_smooth_k, hadamard=self.nvfp4_hadamard)
|
||||
q_fp4, q_sf = _nvfp4_quantize_for_fa4(query)
|
||||
k_fp4, k_sf = _nvfp4_quantize_for_fa4(key)
|
||||
|
||||
# fp8 PV: unscaled e4m3 cast (no mSFV/v_descale); output stays BF16.
|
||||
if self.fa4_pv_mode == "fp8":
|
||||
|
||||
@@ -0,0 +1,401 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Fused Q/K preconditioning + NVFP4 quantization for the FlashAttention-4 NVFP4 path.
|
||||
|
||||
One persistent Triton kernel per tensor replaces ``precondition_qk`` followed by ``_nvfp4_quantize_for_fa4`` (fp32
|
||||
casts, mean subtraction, an fp32 Hadamard matmul, the BF16 round, ``F.pad``, flashinfer's ``nvfp4_quantize`` and the
|
||||
scale-factor permute). Each tile of BLOCK_M tokens x one head:
|
||||
|
||||
1. (K, ``smooth_k``) subtracts the per-(batch, head, channel) token mean, computed by
|
||||
:func:`reference_token_mean`: the reference's own ``k.float().mean(dim=1)`` reduction, bit for bit, without
|
||||
materializing the fp32 copy.
|
||||
2. (``hadamard``) multiplies by ``hadamard_matrix(D)`` the way the reference's cuBLAS GEMM does. When torch lets
|
||||
cuBLAS use TF32 (``TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=1`` in NVIDIA's PyTorch containers, so in training and
|
||||
deployment on the cluster), that is a TF32 tensor-core GEMM: operands rounded to nearest-even TF32, exact
|
||||
products, fp32 accumulation. The kernel issues the same TF32 MMA, which reproduces cuBLAS bit for bit. With TF32
|
||||
off, the fp32 operand is split exactly into three BF16 terms multiplied by the +/-1 matrix and scaled by fp32
|
||||
``1/sqrt(D)``; only fp32 summation order then differs from cuBLAS's fp32 GEMM.
|
||||
3. Rounds to BF16, as the reference does before quantizing.
|
||||
4. Quantizes per 16 channels exactly like flashinfer's default (fast-math) ``nvfp4_quantize`` with a unit global
|
||||
scale: ``sf = e4m3(amax * rcp.approx(6))``, ``code = cvt.rn.satfinite.e2m1(x * rcp.approx(sf))``.
|
||||
5. Writes the packed FP4 data as ``(batch, seqlen_padded, nheads, headdim // 2)`` and the scale factors straight into
|
||||
FlashAttention-4's MMA layout (a permuted view of the contiguous ``(batch, nheads, rest_m, rest_k, 32, 4, 4)``
|
||||
buffer, as the reference returns), with zero codes and zero scales in the rows padding seqlen to a multiple of
|
||||
128.
|
||||
|
||||
The ``fastvideo::nvfp4_fused_qk_quantize_fa4`` custom op (with a stride-exact fake) keeps
|
||||
``torch.compile(fullgraph=True)`` tracing.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from functools import lru_cache
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.attention.utils.nvfp4_qk_precondition import hadamard_matrix
|
||||
|
||||
try:
|
||||
import triton
|
||||
import triton.language as tl
|
||||
_TRITON_AVAILABLE = True
|
||||
except ImportError: # pragma: no cover - exercised only on hosts without Triton
|
||||
triton = None
|
||||
tl = None
|
||||
_TRITON_AVAILABLE = False
|
||||
|
||||
TILE_M = 128 # seqlen padding and scale-factor tile rows (flashinfer layout_128x4)
|
||||
SF_VEC = 16 # NVFP4 channels per scale factor
|
||||
SUPPORTED_HEAD_DIMS = (64, 128)
|
||||
|
||||
|
||||
def _flag_set(value: str | None) -> bool:
|
||||
return (value or "").strip().lower() not in ("", "0", "false", "no", "off")
|
||||
|
||||
|
||||
def fused_qk_quant_available() -> bool:
|
||||
"""Process-level gate: Triton is importable and flashinfer's NVFP4 quantization runs its default fast-math
|
||||
recipe (the one reproduced here). Read once at attention construction, outside compiled code."""
|
||||
non_default_recipe = (_flag_set(os.environ.get("FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH"))
|
||||
or _flag_set(os.environ.get("TRTLLM_DISABLE_FP4_QUANT_FAST_MATH"))
|
||||
or _flag_set(os.environ.get("FLASHINFER_NVFP4_4OVER6")))
|
||||
return _TRITON_AVAILABLE and not non_default_recipe
|
||||
|
||||
|
||||
def fused_qk_quant_supported(query: torch.Tensor, key: torch.Tensor) -> bool:
|
||||
"""Tensor-level gate: BF16 CUDA ``[B, L, H, D]`` Q/K with a supported, channel-contiguous head dim."""
|
||||
return (query.is_cuda and key.is_cuda and query.dtype == torch.bfloat16 and key.dtype == torch.bfloat16
|
||||
and query.dim() == 4 and key.dim() == 4 and query.shape[0] == key.shape[0]
|
||||
and query.shape[-1] == key.shape[-1] and query.shape[-1] in SUPPORTED_HEAD_DIMS and query.stride(-1) == 1
|
||||
and key.stride(-1) == 1 and _token_mean_reproducible(key))
|
||||
|
||||
|
||||
def _padded(seqlen: int) -> int:
|
||||
return (seqlen + TILE_M - 1) // TILE_M * TILE_M
|
||||
|
||||
|
||||
def _output_shapes(tensor_4d: torch.Tensor) -> tuple[tuple[int, ...], tuple[int, ...]]:
|
||||
batch, seqlen, nheads, headdim = tensor_4d.shape
|
||||
seqlen_padded = _padded(seqlen)
|
||||
fp4_shape = (batch, seqlen_padded, nheads, headdim // 2)
|
||||
sf_canonical = (batch, nheads, seqlen_padded // TILE_M, headdim // SF_VEC // 4, 32, 4, 4)
|
||||
return fp4_shape, sf_canonical
|
||||
|
||||
|
||||
# FA4's MMA view of the contiguous canonical scale buffer: (32, 4, rest_m, 4, rest_k, nheads, batch), stride[3] = 1.
|
||||
_SF_MMA_PERMUTE = (4, 5, 2, 6, 3, 1, 0)
|
||||
|
||||
if _TRITON_AVAILABLE:
|
||||
|
||||
@triton.jit
|
||||
def _rcp_approx_ftz(x):
|
||||
return tl.inline_asm_elementwise("rcp.approx.ftz.f32 $0, $1;",
|
||||
"=f,f", [x],
|
||||
dtype=tl.float32,
|
||||
is_pure=True,
|
||||
pack=1)
|
||||
|
||||
@triton.jit
|
||||
def _e2m1x2(lo, hi):
|
||||
# cvt.rn.satfinite.e2m1x2.f32 d, a, b puts a in the high nibble: element 2j -> low nibble (flashinfer order).
|
||||
return tl.inline_asm_elementwise(
|
||||
"""
|
||||
{
|
||||
.reg .b8 b0, b1, b2, b3;
|
||||
cvt.rn.satfinite.e2m1x2.f32 b0, $5, $1;
|
||||
cvt.rn.satfinite.e2m1x2.f32 b1, $6, $2;
|
||||
cvt.rn.satfinite.e2m1x2.f32 b2, $7, $3;
|
||||
cvt.rn.satfinite.e2m1x2.f32 b3, $8, $4;
|
||||
mov.b32 $0, {b0, b1, b2, b3};
|
||||
}
|
||||
""",
|
||||
"=r,f,f,f,f,f,f,f,f",
|
||||
[lo, hi],
|
||||
dtype=tl.uint8,
|
||||
is_pure=True,
|
||||
pack=4,
|
||||
)
|
||||
|
||||
@triton.jit
|
||||
def _precondition_tile(x, xf_mean, rotation, hadamard_scale, SUBTRACT_MEAN: tl.constexpr, HADAMARD: tl.constexpr,
|
||||
TF32: tl.constexpr):
|
||||
"""Steps 1-2 on one [BLOCK_M, D] BF16 tile; returns fp32."""
|
||||
if HADAMARD and TF32:
|
||||
# Reproduce cuBLAS's TF32 GEMM bit for bit: operands rounded to nearest-even TF32, exact products with
|
||||
# the TF32-rounded +/-1/sqrt(D) entries (rotation, fp32), fp32 tensor-core accumulation.
|
||||
xf = x.to(tl.float32)
|
||||
if SUBTRACT_MEAN:
|
||||
xf = xf - xf_mean
|
||||
bits = xf.to(tl.int32, bitcast=True)
|
||||
bits = (bits + 0xFFF + ((bits >> 13) & 1)) & -8192
|
||||
xf = bits.to(tl.float32, bitcast=True)
|
||||
xf = tl.dot(xf, rotation, input_precision="tf32", out_dtype=tl.float32)
|
||||
elif HADAMARD:
|
||||
# fp32 GEMM: split the fp32 operand exactly into three BF16 terms, multiply each by the exact +/-1
|
||||
# matrix (rotation, BF16) on tensor cores, then scale by fp32 1/sqrt(D).
|
||||
if SUBTRACT_MEAN:
|
||||
xf = x.to(tl.float32) - xf_mean
|
||||
hi = xf.to(tl.bfloat16)
|
||||
rest = xf - hi.to(tl.float32)
|
||||
mid = rest.to(tl.bfloat16)
|
||||
lo = (rest - mid.to(tl.float32)).to(tl.bfloat16)
|
||||
acc = tl.dot(hi, rotation, out_dtype=tl.float32)
|
||||
acc = tl.dot(mid, rotation, acc, out_dtype=tl.float32)
|
||||
acc = tl.dot(lo, rotation, acc, out_dtype=tl.float32)
|
||||
else:
|
||||
acc = tl.dot(x, rotation, out_dtype=tl.float32)
|
||||
xf = acc * hadamard_scale
|
||||
elif SUBTRACT_MEAN:
|
||||
xf = x.to(tl.float32) - xf_mean
|
||||
else:
|
||||
xf = x.to(tl.float32)
|
||||
return xf
|
||||
|
||||
@triton.jit
|
||||
def _fused_precondition_nvfp4_kernel(
|
||||
x_ptr,
|
||||
mean_ptr,
|
||||
rotation_ptr,
|
||||
fp4_ptr,
|
||||
sf_ptr,
|
||||
seqlen,
|
||||
seqlen_padded,
|
||||
nheads,
|
||||
num_tiles,
|
||||
stride_b,
|
||||
stride_l,
|
||||
stride_h,
|
||||
hadamard_scale,
|
||||
D: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
SUBTRACT_MEAN: tl.constexpr,
|
||||
HADAMARD: tl.constexpr,
|
||||
TF32: tl.constexpr,
|
||||
NUM_STAGES: tl.constexpr,
|
||||
):
|
||||
cols = tl.arange(0, D)
|
||||
if HADAMARD: # noqa: SIM108 (Triton resolves constexpr branches; keep the statement form)
|
||||
rotation = tl.load(rotation_ptr + cols[:, None] * D + cols[None, :])
|
||||
else:
|
||||
rotation = None
|
||||
tiles_m = seqlen_padded // BLOCK_M
|
||||
r = tl.arange(0, BLOCK_M)
|
||||
kk = tl.arange(0, D // 16)
|
||||
half_cols = tl.arange(0, D // 2)
|
||||
# Persistent: each program walks tiles (BLOCK_M tokens x one head) so the rotation is loaded once.
|
||||
for tile in tl.range(tl.program_id(0), num_tiles, tl.num_programs(0), num_stages=NUM_STAGES):
|
||||
bh = tile // tiles_m
|
||||
pid_m = tile % tiles_m
|
||||
b = bh // nheads
|
||||
h = bh % nheads
|
||||
rows = pid_m * BLOCK_M + r
|
||||
valid = (rows < seqlen)[:, None]
|
||||
src = (x_ptr + b.to(tl.int64) * stride_b + h.to(tl.int64) * stride_h +
|
||||
rows.to(tl.int64)[:, None] * stride_l + cols[None, :])
|
||||
x = tl.load(src, mask=valid, other=0.0)
|
||||
if SUBTRACT_MEAN: # noqa: SIM108
|
||||
xf_mean = tl.load(mean_ptr + bh * D + cols)[None, :]
|
||||
else:
|
||||
xf_mean = None
|
||||
xf = _precondition_tile(x, xf_mean, rotation, hadamard_scale, SUBTRACT_MEAN, HADAMARD, TF32)
|
||||
|
||||
# Step 3: round to BF16. Rows padding seqlen to a multiple of 128 quantize as zeros.
|
||||
y = tl.where(valid, xf, 0.0).to(tl.bfloat16).to(tl.float32)
|
||||
|
||||
# Step 4: flashinfer's fast-math NVFP4 recipe with a unit global scale.
|
||||
blocks = tl.reshape(y, [BLOCK_M, D // 16, 16])
|
||||
vec_max = tl.max(tl.abs(blocks), axis=2)
|
||||
sf_value = vec_max * _rcp_approx_ftz(tl.full([BLOCK_M, D // 16], 6.0, tl.float32))
|
||||
sf_e4m3 = sf_value.to(tl.float8e4nv)
|
||||
out_scale = tl.where(vec_max != 0.0, _rcp_approx_ftz(sf_e4m3.to(tl.float32)), 0.0)
|
||||
scaled = tl.reshape(blocks * out_scale[:, :, None], [BLOCK_M, D // 2, 2])
|
||||
even, odd = tl.split(scaled)
|
||||
codes = _e2m1x2(even, odd)
|
||||
|
||||
out_rows = (b * seqlen_padded + rows).to(tl.int64)
|
||||
tl.store(fp4_ptr + (out_rows[:, None] * nheads + h) * (D // 2) + half_cols[None, :], codes)
|
||||
|
||||
# Scale bytes go straight to the canonical (batch, nheads, rest_m, rest_k, 32, 4, 4) buffer. A 128-row
|
||||
# tile owns D // 64 contiguous 512-byte atoms; inside an atom, row r and scale column c land at
|
||||
# (r % 32) * 16 + (r // 32) * 4 + c (flashinfer layout_128x4).
|
||||
row128 = rows % 128
|
||||
atom_base = ((b * nheads + h) * (seqlen_padded // 128) + rows // 128).to(tl.int64) * (D // 16 * 128)
|
||||
sf_offsets = (atom_base[:, None] + (kk[None, :] // 4) * 512 + (row128[:, None] % 32) * 16 +
|
||||
(row128[:, None] // 32) * 4 + kk[None, :] % 4)
|
||||
tl.store(sf_ptr + sf_offsets, sf_e4m3.to(tl.uint8, bitcast=True))
|
||||
|
||||
|
||||
def _round_tf32(t: torch.Tensor) -> torch.Tensor:
|
||||
"""fp32 -> nearest-even TF32 (kept in fp32), the operand rounding cuBLAS TF32 GEMMs apply."""
|
||||
bits = t.contiguous().view(torch.int32)
|
||||
return ((bits + 0xFFF + ((bits >> 13) & 1)) & -8192).view(torch.float32)
|
||||
|
||||
|
||||
@lru_cache(maxsize=8)
|
||||
def _hadamard_operand(n: int, device: torch.device, tf32: bool) -> torch.Tensor:
|
||||
"""Kernel operand for the reference ``x @ hadamard_matrix(n)``: TF32-rounded entries in fp32 under TF32, else
|
||||
the exact +/-1 signs in BF16 (the kernel scales by fp32 ``1/sqrt(n)`` afterwards)."""
|
||||
rotation = hadamard_matrix(n, "cpu")
|
||||
if tf32:
|
||||
return _round_tf32(rotation).to(device)
|
||||
return torch.sign(rotation).to(device=device, dtype=torch.bfloat16).contiguous()
|
||||
|
||||
|
||||
@lru_cache(maxsize=8)
|
||||
def _hadamard_scale(n: int) -> float:
|
||||
"""The reference matrix entry ``1/sqrt(n)`` in fp32 (applied after the exact +/-1 product when TF32 is off)."""
|
||||
return float(hadamard_matrix(n, "cpu")[0, 0].item())
|
||||
|
||||
|
||||
# (block_m, num_warps, num_stages, ctas_per_sm): persistent-grid launch per mode, tuned on GB200 at L = 58k, H = 56.
|
||||
LAUNCH_CONFIGS = {"plain": (64, 4, 1, 8), "tf32": (64, 4, 3, 2), "fp32": (64, 4, 3, 2)}
|
||||
|
||||
|
||||
def _launch_config(hadamard: bool, tf32: bool) -> tuple[int, int, int, int]:
|
||||
return LAUNCH_CONFIGS["plain" if not hadamard else ("tf32" if tf32 else "fp32")]
|
||||
|
||||
|
||||
@lru_cache(maxsize=8)
|
||||
def _num_sms(device: torch.device) -> int:
|
||||
return torch.cuda.get_device_properties(device).multi_processor_count
|
||||
|
||||
|
||||
_INT32_MAX = 2**31 - 1
|
||||
|
||||
|
||||
def _fits_32bit_indexing(batch: int, seqlen: int, rows: int, row_bytes: int, row_numel: int) -> bool:
|
||||
max_offset = (batch - 1) * seqlen * row_bytes + (rows - 1) * row_bytes + (row_bytes - 4)
|
||||
return max_offset <= _INT32_MAX and batch * rows * row_numel <= _INT32_MAX
|
||||
|
||||
|
||||
def _token_mean_reproducible(tensor_4d: torch.Tensor) -> bool:
|
||||
"""Batched inputs whose fp32 copy needs 64-bit indexing split differently in torch; those take the unfused path."""
|
||||
batch, seqlen, nheads, headdim = tensor_4d.shape
|
||||
return batch == 1 or _fits_32bit_indexing(batch, seqlen, seqlen, nheads * headdim * 4, nheads * headdim)
|
||||
|
||||
|
||||
def _fp32_reduction_chunks(tensor_4d: torch.Tensor) -> list[tuple[int, int]]:
|
||||
"""Token ranges torch reduces separately in ``k.float().mean(dim=1)``.
|
||||
|
||||
A CUDA reduction whose operand needs 64-bit indexing is split (TensorIterator ``with_32bit_indexing``) by halving
|
||||
the largest dimension, here the tokens of the contiguous fp32 copy, until every piece's byte offsets fit in int32;
|
||||
the pieces' partial sums are added left to right. 768p/15 s H3 (113k tokens x 56 heads) splits in two.
|
||||
"""
|
||||
batch, seqlen, nheads, headdim = tensor_4d.shape
|
||||
row_bytes = nheads * headdim * 4
|
||||
if not _token_mean_reproducible(tensor_4d):
|
||||
raise ValueError(f"the split fp32 token mean of a batched {tuple(tensor_4d.shape)} input is not reproduced")
|
||||
|
||||
def split(start: int, rows: int) -> list[tuple[int, int]]:
|
||||
if _fits_32bit_indexing(batch, seqlen, rows, row_bytes, nheads * headdim):
|
||||
return [(start, rows)]
|
||||
first = rows // 2
|
||||
return split(start, first) + split(start + first, rows - first)
|
||||
|
||||
return split(0, seqlen)
|
||||
|
||||
|
||||
def reference_token_mean(tensor_4d: torch.Tensor) -> torch.Tensor:
|
||||
"""``tensor_4d.float().mean(dim=1)`` bit for bit (``[B, H, D]`` fp32) without materializing the fp32 copy.
|
||||
|
||||
Reading BF16 with an fp32 accumulator runs the same reduction as the fp32 copy; when that copy needs 64-bit
|
||||
indexing, torch reduces token ranges separately, adds the partial sums and scales by ``fp32(H * D) / numel``,
|
||||
which is reproduced here. Exact for contiguous (or non-dense) inputs, whose fp32 copy is contiguous.
|
||||
"""
|
||||
chunks = _fp32_reduction_chunks(tensor_4d)
|
||||
if len(chunks) == 1:
|
||||
return tensor_4d.mean(dim=1, dtype=torch.float32).contiguous()
|
||||
total = None
|
||||
for start, rows in chunks:
|
||||
partial = tensor_4d[:, start:start + rows].sum(dim=1, dtype=torch.float32)
|
||||
total = partial if total is None else total + partial
|
||||
batch, seqlen, nheads, headdim = tensor_4d.shape
|
||||
factor = (torch.tensor(batch * nheads * headdim, dtype=torch.float32) /
|
||||
torch.tensor(batch * seqlen * nheads * headdim, dtype=torch.float32)).item()
|
||||
return (total * factor).contiguous()
|
||||
|
||||
|
||||
def _precondition_quantize(tensor_4d: torch.Tensor, *, subtract_mean: bool, hadamard: bool,
|
||||
tf32: bool) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
batch, seqlen, nheads, headdim = tensor_4d.shape
|
||||
fp4_shape, sf_canonical = _output_shapes(tensor_4d)
|
||||
seqlen_padded = fp4_shape[1]
|
||||
fp4 = torch.empty(fp4_shape, device=tensor_4d.device, dtype=torch.uint8)
|
||||
sf = torch.empty(sf_canonical, device=tensor_4d.device, dtype=torch.uint8)
|
||||
# The reference's own reduction (``k.float().mean(dim=1)``): reading BF16 with an fp32 accumulator gives the
|
||||
# same bits without materializing the fp32 copy.
|
||||
mean = rotation = fp4 # placeholder pointers for the operands a mode does not read
|
||||
if subtract_mean:
|
||||
mean = reference_token_mean(tensor_4d).view(batch * nheads, headdim)
|
||||
if hadamard:
|
||||
rotation = _hadamard_operand(headdim, tensor_4d.device, tf32)
|
||||
block_m, num_warps, num_stages, ctas_per_sm = _launch_config(hadamard, tf32)
|
||||
num_tiles = seqlen_padded // block_m * batch * nheads
|
||||
grid = (min(num_tiles, _num_sms(tensor_4d.device) * ctas_per_sm), )
|
||||
_fused_precondition_nvfp4_kernel[grid](
|
||||
tensor_4d,
|
||||
mean,
|
||||
rotation,
|
||||
fp4,
|
||||
sf,
|
||||
seqlen,
|
||||
seqlen_padded,
|
||||
nheads,
|
||||
num_tiles,
|
||||
tensor_4d.stride(0),
|
||||
tensor_4d.stride(1),
|
||||
tensor_4d.stride(2),
|
||||
_hadamard_scale(headdim) if hadamard else 1.0,
|
||||
D=headdim,
|
||||
BLOCK_M=block_m,
|
||||
SUBTRACT_MEAN=subtract_mean,
|
||||
HADAMARD=hadamard,
|
||||
TF32=tf32 and hadamard,
|
||||
NUM_STAGES=num_stages,
|
||||
num_warps=num_warps,
|
||||
)
|
||||
return fp4.view(torch.float4_e2m1fn_x2), sf.permute(*_SF_MMA_PERMUTE)
|
||||
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo::nvfp4_fused_qk_quantize_fa4",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def _nvfp4_fused_qk_quantize_fa4_op(query: torch.Tensor, key: torch.Tensor, smooth_k: bool,
|
||||
hadamard: bool) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
# Mirror the reference ``x_fp32 @ H`` exactly: cuBLAS runs it in TF32 whenever torch allows TF32 matmuls.
|
||||
tf32 = bool(torch.backends.cuda.matmul.allow_tf32)
|
||||
q_fp4, q_sf = _precondition_quantize(query, subtract_mean=False, hadamard=hadamard, tf32=tf32)
|
||||
k_fp4, k_sf = _precondition_quantize(key, subtract_mean=smooth_k, hadamard=hadamard, tf32=tf32)
|
||||
return q_fp4, q_sf, k_fp4, k_sf
|
||||
|
||||
|
||||
def _fake_outputs(tensor_4d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
fp4_shape, sf_canonical = _output_shapes(tensor_4d)
|
||||
fp4 = tensor_4d.new_empty(fp4_shape, dtype=torch.float4_e2m1fn_x2)
|
||||
# Reproduce the real strides: sf is a permuted view of a contiguous canonical buffer.
|
||||
sf = tensor_4d.new_empty(sf_canonical, dtype=torch.uint8).permute(*_SF_MMA_PERMUTE)
|
||||
return fp4, sf
|
||||
|
||||
|
||||
@torch.library.register_fake("fastvideo::nvfp4_fused_qk_quantize_fa4")
|
||||
def _nvfp4_fused_qk_quantize_fa4_fake(query: torch.Tensor, key: torch.Tensor, smooth_k: bool,
|
||||
hadamard: bool) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
del smooth_k, hadamard
|
||||
return (*_fake_outputs(query), *_fake_outputs(key))
|
||||
|
||||
|
||||
def nvfp4_fused_qk_quantize_fa4(query: torch.Tensor, key: torch.Tensor, *, smooth_k: bool,
|
||||
hadamard: bool) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""``precondition_qk`` + ``_nvfp4_quantize_for_fa4`` on Q and K in fused kernels.
|
||||
|
||||
Returns ``(q_fp4, q_sf, k_fp4, k_sf)`` with the reference layout contract (seqlen padded to 128; slice
|
||||
``[:, :seqlen]`` before FA4). Check :func:`fused_qk_quant_supported` first.
|
||||
"""
|
||||
return torch.ops.fastvideo.nvfp4_fused_qk_quantize_fa4(query, key, smooth_k, hadamard)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"SUPPORTED_HEAD_DIMS", "fused_qk_quant_available", "fused_qk_quant_supported", "nvfp4_fused_qk_quantize_fa4",
|
||||
"reference_token_mean"
|
||||
]
|
||||
@@ -0,0 +1,78 @@
|
||||
# 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.
|
||||
|
||||
A checkpoint trained with one of these transforms must be served with the same one.
|
||||
"""
|
||||
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]
|
||||
|
||||
|
||||
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):
|
||||
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_block_indices", "parse_qk_transform", "precondition_qk"]
|
||||
@@ -378,6 +378,25 @@ 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_FUSED_QK_QUANT = EnvBool(
|
||||
True,
|
||||
category="attention",
|
||||
doc="With FASTVIDEO_NVFP4_FA4_QK_TRANSFORM other than none, apply the transform and the NVFP4 quantization in one "
|
||||
"fused Triton kernel per tensor (same FP4 codes and scale factors as the unfused path). Set it to 0 to use the "
|
||||
"unfused PyTorch transform and flashinfer quantization.")
|
||||
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"),
|
||||
|
||||
@@ -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,
|
||||
@@ -226,6 +227,7 @@ class MiniMaxH3Attention(nn.Module):
|
||||
exact_rope: bool = False,
|
||||
mk_qkv: bool = False,
|
||||
vsa_fp4: dict[str, Any] | None = None,
|
||||
bf16_attention: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.num_attention_heads = num_attention_heads
|
||||
@@ -274,6 +276,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,
|
||||
@@ -282,6 +286,7 @@ class MiniMaxH3Attention(nn.Module):
|
||||
prefix=prefix,
|
||||
fa4_packed_varlen=fa4_packed_varlen,
|
||||
cudnn_sdpa=cudnn_sdpa,
|
||||
**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.
|
||||
@@ -698,6 +703,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
|
||||
exact_kernels: frozenset[str] = frozenset(),
|
||||
vsa_fp4: dict[str, Any] | None = None,
|
||||
mk_tier: bool = False,
|
||||
bf16_attention: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.norm1 = nn.RMSNorm(hidden_size, eps=norm_eps)
|
||||
@@ -715,6 +721,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
|
||||
exact_rope="rope" in exact_kernels,
|
||||
vsa_fp4=vsa_fp4,
|
||||
mk_qkv=mk_tier,
|
||||
bf16_attention=bf16_attention,
|
||||
)
|
||||
self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps)
|
||||
self.ff = MiniMaxH3FeedForward(
|
||||
@@ -977,6 +984,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,
|
||||
@@ -998,6 +1010,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
|
||||
exact_kernels=self.enabled_exact_kernels,
|
||||
vsa_fp4=config.vsa_fp4,
|
||||
mk_tier=self.use_mk_kernels,
|
||||
bf16_attention=index in bf16_blocks,
|
||||
) for index in range(arch.num_layers)
|
||||
])
|
||||
self.norm_out = MiniMaxH3AdaLayerNormOut(
|
||||
|
||||
@@ -234,6 +234,33 @@ def load_adapter(source: str | os.PathLike[str], *, task: str, revision: str | N
|
||||
)
|
||||
|
||||
|
||||
PIKA_MERGED_MARKER = "pika_h3_flex_merged.json"
|
||||
|
||||
|
||||
def merged_adapter(model_dir: str | os.PathLike[str], task: str) -> dict[str, Any] | None:
|
||||
"""The marker entry when ``model_dir``'s trunk for ``task`` already contains the Pika adapter, else None.
|
||||
|
||||
Exported checkpoints (e.g. NVFP4 QAD releases) merge the adapter into the trunk before quantizing it, so
|
||||
applying the released adapter again would adapt the model twice. Such a checkpoint ships
|
||||
``pika_h3_flex_merged.json``: ``{"format_version": 1, "merged": {"<task>": {"trunk": ..., "release_id": ...,
|
||||
"adapter_sha256": ...}}}``. A marker that does not cover ``task`` is an error: that trunk is absent or unadapted.
|
||||
"""
|
||||
path = Path(model_dir) / PIKA_MERGED_MARKER
|
||||
if not path.is_file():
|
||||
return None
|
||||
marker = _load_manifest(path)
|
||||
if _as_int(marker.get("format_version"), "format_version") != PIKA_ADAPTER_FORMAT_VERSION:
|
||||
raise PikaAdapterError(f"Unsupported {PIKA_MERGED_MARKER} format {marker.get('format_version')!r}.")
|
||||
merged = marker.get("merged")
|
||||
if not isinstance(merged, Mapping) or task not in merged:
|
||||
raise PikaAdapterError(f"{str(path)!r} declares merged Pika adapters for {sorted(merged or {})}, not "
|
||||
f"{task!r}; this checkpoint does not carry the {task} model.")
|
||||
entry = merged[task]
|
||||
if not isinstance(entry, Mapping) or entry.get("trunk") != expected_trunk(task):
|
||||
raise PikaAdapterError(f"{str(path)!r}: the merged {task} adapter must sit on trunk {expected_trunk(task)!r}.")
|
||||
return dict(entry)
|
||||
|
||||
|
||||
def adapter_hf_patterns(task: str | None = None) -> list[str]:
|
||||
"""The adapter files the pipeline must include in a partial Hub snapshot."""
|
||||
tasks = tuple(TASK_TO_TRUNK) if task is None else (task, )
|
||||
@@ -251,6 +278,7 @@ __all__ = [
|
||||
"PIKA_ADAPTER_RANK",
|
||||
"PIKA_BASE_REPO",
|
||||
"PIKA_MANIFEST_FILENAME",
|
||||
"PIKA_MERGED_MARKER",
|
||||
"PIKA_SAMPLING_CONTRACT",
|
||||
"PIKA_TARGET_MODULES",
|
||||
"PikaAdapter",
|
||||
@@ -259,6 +287,7 @@ __all__ = [
|
||||
"adapter_hf_patterns",
|
||||
"expected_trunk",
|
||||
"load_adapter",
|
||||
"merged_adapter",
|
||||
"task_for_pipeline",
|
||||
"validate_adapter_manifest",
|
||||
]
|
||||
|
||||
@@ -14,7 +14,8 @@ from fastvideo.configs.pipelines.pika_h3_flex import (
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3BasePipeline
|
||||
from fastvideo.pipelines.basic.pika_h3_flex.adapter import adapter_hf_patterns, load_adapter, task_for_pipeline
|
||||
from fastvideo.pipelines.basic.pika_h3_flex.adapter import (PIKA_MERGED_MARKER, adapter_hf_patterns, load_adapter,
|
||||
merged_adapter, task_for_pipeline)
|
||||
from fastvideo.pipelines.basic.pika_h3_flex.stages.denoising import PikaH3FlexDenoisingStage
|
||||
from fastvideo.pipelines.basic.pika_h3_flex.stages.input_preparation import PikaH3FlexInputPreparationStage
|
||||
|
||||
@@ -36,7 +37,7 @@ class PikaH3FlexBasePipeline(MiniMaxH3BasePipeline):
|
||||
# The concrete class picks the adapter, so a T2VA download does not
|
||||
# fetch the 2.7 GB Ref2VA adapter (and vice versa).
|
||||
task = "ref2va" if issubclass(cls, PikaH3FlexRefPipeline) else "fl2va"
|
||||
return [*patterns, *adapter_hf_patterns(task)]
|
||||
return [*patterns, *adapter_hf_patterns(task), PIKA_MERGED_MARKER]
|
||||
|
||||
def _resolve_pika_adapter(self) -> None:
|
||||
"""Point ``lora_path`` at the released adapter before LoRA layers attach.
|
||||
@@ -47,6 +48,13 @@ class PikaH3FlexBasePipeline(MiniMaxH3BasePipeline):
|
||||
"""
|
||||
task = task_for_pipeline(self._ref2va)
|
||||
explicit = self.fastvideo_args.lora_path
|
||||
merged = merged_adapter(self.model_path, task) if explicit is None else None
|
||||
if merged is not None:
|
||||
# The trunk already carries the adapter (an exported, e.g. NVFP4 QAD, checkpoint): attach nothing.
|
||||
logger.info("Pika H3 Flex: the %s trunk of %s already contains adapter %s; not applying it again", task,
|
||||
self.model_path, merged.get("release_id"))
|
||||
self._pika_adapter = None
|
||||
return
|
||||
source = explicit if explicit is not None else self.model_path
|
||||
revision = None if explicit is not None else self.fastvideo_args.revision
|
||||
adapter = load_adapter(source, task=task, revision=revision)
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Fused NVFP4 Q/K preconditioning + quantization vs the reference path.
|
||||
|
||||
The fused kernel (``fastvideo.attention.utils.nvfp4_fused_qk_quant``) must produce the same FP4 codes, the same
|
||||
scale-factor bytes in FlashAttention-4's layout, and therefore the same FA4 NVFP4 attention output as
|
||||
``precondition_qk`` followed by ``_nvfp4_quantize_for_fa4``. GPU tests need Blackwell (sm_100a/sm_103a), flashinfer
|
||||
and the FA4 NVFP4 kernel; the selection tests run anywhere.
|
||||
"""
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
|
||||
import fastvideo.envs as envs # noqa: E402
|
||||
from fastvideo.attention.utils.nvfp4_qk_precondition import parse_qk_transform, precondition_qk # noqa: E402
|
||||
|
||||
flash_attn_module = pytest.importorskip(
|
||||
"fastvideo.attention.backends.flash_attn",
|
||||
reason="no usable flash-attention package installed",
|
||||
exc_type=ImportError,
|
||||
)
|
||||
fused = pytest.importorskip("fastvideo.attention.utils.nvfp4_fused_qk_quant")
|
||||
|
||||
TRANSFORMS = ["none", "smooth_k", "hadamard", "smooth_k+hadamard"]
|
||||
|
||||
|
||||
def _blackwell_fa4_fp4() -> bool:
|
||||
if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in [(10, 0), (10, 3)]:
|
||||
return False
|
||||
try:
|
||||
import flashinfer.quantization # noqa: F401
|
||||
except ImportError:
|
||||
return False
|
||||
return flash_attn_module._FA4_FP4_AVAILABLE and fused.fused_qk_quant_available()
|
||||
|
||||
|
||||
requires_blackwell = pytest.mark.skipif(not _blackwell_fa4_fp4(),
|
||||
reason="needs Blackwell, flashinfer, Triton and FA4 NVFP4")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tf32_matmul(request):
|
||||
previous = torch.backends.cuda.matmul.allow_tf32
|
||||
torch.backends.cuda.matmul.allow_tf32 = request.param
|
||||
yield request.param
|
||||
torch.backends.cuda.matmul.allow_tf32 = previous
|
||||
|
||||
|
||||
def _inputs(batch, seqlen, heads, dim, *, strided=False, seed=0):
|
||||
gen = torch.Generator(device="cuda").manual_seed(seed)
|
||||
if strided: # Q/K/V sliced out of one packed buffer, as a fused QKV projection produces them
|
||||
qkv = torch.randn(batch, seqlen, 3, heads, dim, device="cuda", dtype=torch.bfloat16, generator=gen)
|
||||
q, k, v = qkv.unbind(2)
|
||||
else:
|
||||
q, k, v = (torch.randn(batch, seqlen, heads, dim, device="cuda", dtype=torch.bfloat16, generator=gen)
|
||||
for _ in range(3))
|
||||
# A per-channel K offset (what smooth_k removes) and a few outlier channels (what the Hadamard spreads).
|
||||
k = k.clone()
|
||||
k[..., :4] *= 8.0
|
||||
k += 4.0 * torch.randn(batch, 1, heads, dim, device="cuda", dtype=torch.bfloat16, generator=gen)
|
||||
if strided:
|
||||
buffer = torch.empty(batch, seqlen, 2, heads, dim, device="cuda", dtype=torch.bfloat16)
|
||||
buffer[:, :, 1] = k
|
||||
k = buffer[:, :, 1]
|
||||
return q, k, v
|
||||
|
||||
|
||||
def _reference(q, k, transform):
|
||||
smooth_k, hadamard = parse_qk_transform(transform)
|
||||
q, k = precondition_qk(q, k, smooth_k=smooth_k, hadamard=hadamard)
|
||||
return (*flash_attn_module._nvfp4_quantize_for_fa4(q), *flash_attn_module._nvfp4_quantize_for_fa4(k))
|
||||
|
||||
|
||||
def _fused(q, k, transform):
|
||||
smooth_k, hadamard = parse_qk_transform(transform)
|
||||
return fused.nvfp4_fused_qk_quantize_fa4(q, k, smooth_k=smooth_k, hadamard=hadamard)
|
||||
|
||||
|
||||
def _codes(fp4):
|
||||
data = fp4.view(torch.uint8)
|
||||
return torch.stack([data & 0xF, data >> 4], dim=-1)
|
||||
|
||||
|
||||
def _attention(q4, q_sf, k4, k_sf, v):
|
||||
out = flash_attn_module.flash_attn_fp4_func(q4[:, :v.shape[1]], k4[:, :v.shape[1]], v, q_sf, k_sf,
|
||||
softmax_scale=v.shape[-1]**-0.5, causal=False)
|
||||
return out[0] if isinstance(out, tuple) else out
|
||||
|
||||
|
||||
SHAPES = [
|
||||
pytest.param((1, 1000, 4, 128, False), id="odd_seqlen"),
|
||||
pytest.param((2, 300, 2, 64, False), id="batch2_d64"),
|
||||
pytest.param((1, 777, 3, 128, True), id="strided"),
|
||||
]
|
||||
|
||||
|
||||
@requires_blackwell
|
||||
@pytest.mark.parametrize("tf32_matmul", [True], indirect=True)
|
||||
@pytest.mark.parametrize("shape", SHAPES)
|
||||
@pytest.mark.parametrize("transform", TRANSFORMS)
|
||||
def test_fused_matches_reference_bit_exactly_under_tf32(transform, shape, tf32_matmul) -> None:
|
||||
"""Under TF32 matmuls (NVIDIA PyTorch containers) the fused path is bit-identical to the reference."""
|
||||
batch, seqlen, heads, dim, strided = shape
|
||||
q, k, v = _inputs(batch, seqlen, heads, dim, strided=strided)
|
||||
ref, out = _reference(q, k, transform), _fused(q, k, transform)
|
||||
for r, o in zip(ref, out, strict=True):
|
||||
assert r.shape == o.shape and r.dtype == o.dtype and r.stride() == o.stride()
|
||||
assert torch.equal(r.view(torch.uint8), o.view(torch.uint8))
|
||||
for fp4, length in ((out[0], seqlen), (out[2], seqlen)):
|
||||
assert not _codes(fp4)[:, length:].any(), "rows padding seqlen to 128 must hold zero codes"
|
||||
assert torch.equal(_attention(*out, v), _attention(*ref, v))
|
||||
|
||||
|
||||
@requires_blackwell
|
||||
@pytest.mark.parametrize("tf32_matmul", [False], indirect=True)
|
||||
@pytest.mark.parametrize("transform", ["hadamard", "smooth_k+hadamard"])
|
||||
def test_fused_matches_reference_under_fp32(transform, tf32_matmul) -> None:
|
||||
"""With TF32 off only fp32 summation order differs from cuBLAS's fp32 GEMM: codes/scales agree to >= 99.99%."""
|
||||
q, k, v = _inputs(1, 4000, 8, 128, seed=1)
|
||||
ref, out = _reference(q, k, transform), _fused(q, k, transform)
|
||||
for r, o in ((ref[0], out[0]), (ref[2], out[2])):
|
||||
assert (_codes(r) == _codes(o)).float().mean().item() >= 0.9999
|
||||
for r, o in ((ref[1], out[1]), (ref[3], out[3])):
|
||||
assert r.stride() == o.stride()
|
||||
assert (r == o).float().mean().item() >= 0.9999
|
||||
|
||||
|
||||
@requires_blackwell
|
||||
def test_forward_nvfp4_uses_fused_path_and_compiles_fullgraph() -> None:
|
||||
impl = flash_attn_module.FlashAttentionImpl(num_heads=4, head_size=128, causal=False, softmax_scale=128**-0.5,
|
||||
num_kv_heads=4, nvfp4_fa4=True,
|
||||
nvfp4_qk_transform="smooth_k+hadamard")
|
||||
assert impl.nvfp4_fused_qk_quant
|
||||
q, k, v = _inputs(1, 1000, 4, 128, seed=2)
|
||||
with patch.object(flash_attn_module, "precondition_qk", side_effect=AssertionError("unfused path taken")):
|
||||
eager = impl._forward_nvfp4(q, k, v)
|
||||
compiled = torch.compile(impl._forward_nvfp4, fullgraph=True)(q, k, v)
|
||||
assert torch.equal(eager, compiled)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA GPU (about 4 GiB)")
|
||||
@pytest.mark.parametrize("seqlen", [74898, 74899])
|
||||
def test_token_mean_matches_reference_across_the_64bit_index_split(seqlen) -> None:
|
||||
"""At 56 x 128 channels the fp32 copy of K needs 64-bit indexing from 74,899 tokens on (768p/15 s H3 has 113k),
|
||||
where torch reduces it in two halves; the fused path's mean must keep the reference's bits on both sides."""
|
||||
gen = torch.Generator(device="cuda").manual_seed(3)
|
||||
k = (torch.randn(1, seqlen, 56, 128, device="cuda", generator=gen) * 2 + 1).to(torch.bfloat16)
|
||||
assert len(fused._fp32_reduction_chunks(k)) == (1 if seqlen == 74898 else 2)
|
||||
assert torch.equal(fused.reference_token_mean(k), k.float().mean(dim=1))
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("shape", "reproducible"), [
|
||||
((1, 113000, 56, 128), True),
|
||||
((2, 1000, 56, 128), True),
|
||||
((2, 40000, 56, 128), False),
|
||||
((2, 80000, 56, 128), False),
|
||||
])
|
||||
def test_batched_inputs_that_split_the_token_mean_take_the_unfused_path(shape, reproducible) -> None:
|
||||
k = torch.empty(shape, dtype=torch.bfloat16, device="meta")
|
||||
assert fused._token_mean_reproducible(k) is reproducible
|
||||
if not reproducible:
|
||||
with pytest.raises(ValueError, match="not reproduced"):
|
||||
fused._fp32_reduction_chunks(k)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA GPU")
|
||||
def test_batched_token_mean_matches_reference() -> None:
|
||||
gen = torch.Generator(device="cuda").manual_seed(5)
|
||||
k = (torch.randn(2, 4096, 56, 128, device="cuda", generator=gen) * 2 + 1).to(torch.bfloat16)
|
||||
assert torch.equal(fused.reference_token_mean(k), k.float().mean(dim=1))
|
||||
|
||||
@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 _impl(**extra):
|
||||
return flash_attn_module.FlashAttentionImpl(num_heads=2, head_size=64, causal=False, softmax_scale=0.125,
|
||||
num_kv_heads=2, nvfp4_fa4=True, **extra)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("transform", TRANSFORMS)
|
||||
def test_fused_path_selected_only_with_a_transform(transform, env_overrides, fa4_runtime_available) -> None:
|
||||
env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4_FUSED_QK_QUANT.override(None))
|
||||
with patch.object(flash_attn_module, "fused_qk_quant_available", return_value=True):
|
||||
assert _impl(nvfp4_qk_transform=transform).nvfp4_fused_qk_quant == (transform != "none")
|
||||
|
||||
|
||||
def test_fused_path_env_opt_out(env_overrides, fa4_runtime_available) -> None:
|
||||
env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4_FUSED_QK_QUANT.override(False))
|
||||
with patch.object(flash_attn_module, "fused_qk_quant_available", return_value=True):
|
||||
assert _impl(nvfp4_qk_transform="smooth_k+hadamard").nvfp4_fused_qk_quant is False
|
||||
|
||||
|
||||
def test_fused_path_off_for_non_default_flashinfer_recipe(env_overrides) -> None:
|
||||
env_overrides.enter_context(envs.override_external("FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH", "1"))
|
||||
assert fused.fused_qk_quant_available() is False
|
||||
@@ -0,0 +1,128 @@
|
||||
# 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, 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
|
||||
@@ -93,6 +93,8 @@ EXTERNAL_ALLOWLIST = {
|
||||
"PYTORCH_MPS_*": "PyTorch MPS memory settings.",
|
||||
"USER": "Login name of the current user.",
|
||||
"PYTHONPATH": "Python module search path.",
|
||||
"FLASHINFER_*": "flashinfer settings; the fused NVFP4 Q/K quantizer checks the quantization recipe switches.",
|
||||
"TRTLLM_DISABLE_FP4_QUANT_FAST_MATH": "TensorRT-LLM FP4 quantization switch that flashinfer also reads.",
|
||||
"FASTVIDEO_VSA_CUTEDSL": "Read by fastvideo-kernel, which is outside this policy.",
|
||||
"FASTVIDEO_VSA_TRITON": "Read by fastvideo-kernel, which is outside this policy.",
|
||||
"FASTVIDEO_KERNEL_VSA_FORCE_TRITON": "Read by fastvideo-kernel, which is outside this policy.",
|
||||
@@ -142,6 +144,7 @@ EXTERNAL_WRITE_ALLOWLIST = {
|
||||
"RAY_USAGE_STATS_ENABLED": "Turns off Ray usage statistics unless the user turned them on.",
|
||||
"CUTE_DSL_ENABLE_TVM_FFI": "CUTLASS DSL switch that the NVFP4 FlashAttention-4 path needs.",
|
||||
"CUBLAS_WORKSPACE_CONFIG": "cuBLAS workspace size; golden-gate tests set it before cuBLAS initializes.",
|
||||
"FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH": "Tests pin the flashinfer recipe gate of the fused NVFP4 Q/K quantizer.",
|
||||
}
|
||||
EXTERNAL_WRITE_HELPERS = {"set_external", "setdefault_external", "unset_external", "override_external"}
|
||||
|
||||
|
||||
@@ -93,3 +93,24 @@ def test_load_adapter_requires_the_manifest(tmp_path: Path) -> None:
|
||||
(task_dir / "manifest.json").unlink()
|
||||
with pytest.raises(adapter.PikaAdapterError, match="manifest"):
|
||||
adapter.load_adapter(tmp_path, task="fl2va")
|
||||
|
||||
|
||||
def _write_marker(root: Path, merged: dict) -> None:
|
||||
(root / adapter.PIKA_MERGED_MARKER).write_text(json.dumps({"format_version": 1, "merged": merged}),
|
||||
encoding="utf-8")
|
||||
|
||||
|
||||
def test_merged_marker_skips_the_released_adapter(tmp_path):
|
||||
assert adapter.merged_adapter(tmp_path, "ref2va") is None
|
||||
_write_marker(tmp_path, {"ref2va": {"trunk": "transformer_ref", "release_id": "pika-flex-ref2va-v1"}})
|
||||
assert adapter.merged_adapter(tmp_path, "ref2va")["release_id"] == "pika-flex-ref2va-v1"
|
||||
|
||||
|
||||
def test_merged_marker_must_cover_the_task_and_trunk(tmp_path):
|
||||
_write_marker(tmp_path, {"ref2va": {"trunk": "transformer_ref"}})
|
||||
with pytest.raises(adapter.PikaAdapterError, match="does not carry the fl2va model"):
|
||||
adapter.merged_adapter(tmp_path, "fl2va")
|
||||
_write_marker(tmp_path, {"ref2va": {"trunk": "transformer"}})
|
||||
with pytest.raises(adapter.PikaAdapterError, match="must sit on trunk"):
|
||||
adapter.merged_adapter(tmp_path, "ref2va")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user