[perf] Fused Q/K preconditioning + NVFP4 quantize for FA4 NVFP4 attention (bit-identical) (#2052)
Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com>
This commit is contained in:
Aryan Kumar
co-authored by
Aryan Kumar
Aryan Kumar
parent
7066408fb3
commit
4b4df1a94e
@@ -202,6 +202,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
|
||||
| `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. |
|
||||
|
||||
@@ -29,6 +29,11 @@ from fastvideo.attention.backends.abstract import (
|
||||
)
|
||||
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__)
|
||||
@@ -291,6 +296,8 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
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
|
||||
@@ -308,8 +315,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, smooth_k=%s, hadamard=%s)",
|
||||
self.fa4_pv_mode, self.nvfp4_smooth_k, self.nvfp4_hadamard)
|
||||
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,
|
||||
@@ -544,11 +553,16 @@ 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)
|
||||
if self.nvfp4_fused_qk_quant and fused_qk_quant_supported(query, key, hadamard=self.nvfp4_hadamard):
|
||||
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,407 @@
|
||||
# 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, *, hadamard: bool = False) -> bool:
|
||||
"""Tensor-level gate: BF16 CUDA ``[B, L, H, D]`` Q/K with a supported, channel-contiguous head dim.
|
||||
|
||||
The Hadamard rotation is bit-exact only against cuBLAS's TF32 GEMM, so with TF32 matmuls off it takes the
|
||||
reference path.
|
||||
"""
|
||||
if hadamard and not torch.backends.cuda.matmul.allow_tf32:
|
||||
return False
|
||||
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"
|
||||
]
|
||||
@@ -392,6 +392,12 @@ FASTVIDEO_NVFP4_FA4_QK_TRANSFORM = EnvChoice(
|
||||
"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",
|
||||
|
||||
@@ -0,0 +1,212 @@
|
||||
# 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
|
||||
@pytest.mark.parametrize("tf32_matmul", [False], indirect=True)
|
||||
def test_hadamard_without_tf32_takes_the_reference_path(tf32_matmul) -> None:
|
||||
q, k, _ = _inputs(1, 256, 4, 128, seed=4)
|
||||
assert not fused.fused_qk_quant_supported(q, k, hadamard=True)
|
||||
assert fused.fused_qk_quant_supported(q, k, hadamard=False)
|
||||
|
||||
|
||||
@requires_blackwell
|
||||
@pytest.mark.parametrize("tf32_matmul", [True], indirect=True)
|
||||
def test_forward_nvfp4_uses_fused_path_and_compiles_fullgraph(tf32_matmul) -> 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
|
||||
@@ -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"}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user