flashinfer.quantization

Quantization-related kernels for FP4, FP8, and packbits utilities.

Types and Enums

SfLayout(value)

Layout of scale factors for quantization.

NVFP4 4over6 recipe types. These are defined in flashinfer.quantization.nvfp4_quantization_utils and re-exported here; only the canonical flashinfer.quantization spelling is documented, because sphinx -W treats a second description of the same object as an error.

NVFP44Over6Config([e4m3_max, err_mode, ...])

NVFP4 4over6 configuration shared by Python drivers and kernels.

NVFP44Over6ErrMode(value)

Error metric that picks between the two 4over6 block-scale candidates.

Packbits Utilities

packbits(x[, bitorder])

Pack the elements of a binary-valued array into bits in a uint8 array.

segment_packbits(x, indptr[, bitorder])

Pack a batch of binary-valued segments into bits in a uint8 array.

NVFP4 4over6 Recipes

Standard NVFP4 gives each 16-element block a single E4M3 block scale derived from the block’s absolute maximum: amax / 6, because 6 is the largest magnitude an E2M1 element can hold. “4over6” is a two-candidate block-scale search — the quantizer also forms the 1.5x tighter amax / 4 scale, quantizes the block with both, dequantizes both, and keeps whichever reconstructs the block with less error. The amax / 4 candidate buys resolution for the bulk of the block by clipping its largest magnitudes, so it wins on blocks with no dominant outlier. The extra work is paid entirely inside the quantizer: the result is ordinary NVFP4 and needs no change downstream.

4over6 requires fp16/bf16 input and E4M3 (not UE8M0) block scales.

Two knobs matter in practice:

e4m3_max=256

Treat 256, not the full-range 448, as the top of the E4M3 block-scale range. This is what keeps the amax / 4 candidate representable: 256 * 1.5 = 384 still fits E4M3, whereas a 448-based global scale would need 672 and saturate for blocks near the tensor amax. It also changes the global scale — see the warning below.

err_mode="MSE"

Score the two candidates by summed squared error instead of summed absolute error. MSE punishes the few large residuals that amax / 4 clipping creates much harder than MAE does, so it picks amax / 4 less often. Prefer it when a handful of outlier activations dominate downstream accuracy.

err_use_fast_math=True evaluates that error in fp16 rather than exactly in fp32. It is cheaper, and on near-ties it can select the other candidate, so treat it as a distinct recipe rather than a pure speed knob.

Warning

The global scale must be built from the same recipe as the quantizer call. e4m3_max appears both in make_nvfp4_global_scale() and in the kernel’s candidate search; a mismatch silently rescales the whole tensor instead of raising.

import torch
from flashinfer.quantization import (
    NVFP44Over6Config,
    make_nvfp4_global_scale,
    nvfp4_quantize,
)

x = torch.randn(1024, 2048, device="cuda", dtype=torch.bfloat16)
recipe = NVFP44Over6Config(e4m3_max=256, err_mode="MSE")

# Same recipe object on both calls - see the warning above.
global_sf = make_nvfp4_global_scale(
    x, per_token_activation=False, nvfp4_4over6_config=recipe
)
x_q, sf = nvfp4_quantize(x, global_sf, nvfp4_4over6=recipe)

Precedence

Every NVFP4 entry point that used to read the environment now takes an nvfp4_4over6= keyword typed Optional[NVFP44Over6Config]. There are three cases, but only one public type:

nvfp4_4over6=

Meaning

omitted (the default)

Read FLASHINFER_NVFP4_4OVER6; when it is "1", the other three FLASHINFER_NVFP4_4OVER6_* variables supply the recipe. Read on every call. Byte-for-byte the behaviour that predates this parameter, plus a FutureWarning when the environment turns 4over6 on.

None

4over6 off. FLASHINFER_NVFP4_4OVER6=1 cannot turn it back on. This is the same meaning None has always had inside the kernel drivers.

NVFP44Over6Config(...)

On with exactly this recipe. The environment is ignored, and there is no per-field merge: a field left at its dataclass default keeps that default rather than picking up the environment’s value.

The four FLASHINFER_NVFP4_4OVER6* variables are documented in the repository’s CLAUDE.md. They remain supported as a compatibility shim but are deprecated: being process-wide, they cannot express two models served in one process under different recipes. When the shim is retired the default of nvfp4_4over6= becomes None and no signature changes.

Note

make_nvfp4_global_scale() and nvfp4_e4m3_max() take the resolved recipe as nvfp4_4over6_config= and never read the environment (None, their default, has always meant standard NVFP4). When the quantize call leaves nvfp4_4over6 unset, pass resolve_nvfp4_4over6() to the scale helper so both sides agree.

Resolution happens in exactly one place, resolve_nvfp4_4over6(). Below it, every kernel driver sees the two-state Optional[NVFP44Over6Config] in which None means off, and nvfp4_4over6_code() packs that resolved value into the int64 the custom ops take. nvfp4_4over6_from_code() inverts the packing, which is how logs and repro tooling recover the recipe a captured call actually ran with.

resolve_nvfp4_4over6([setting])

Resolve a public nvfp4_4over6= argument to the recipe the kernels use.

nvfp4_4over6_code(config)

Pack a resolved recipe into the int64 the custom ops / TVM-FFI take.

nvfp4_4over6_from_code(code)

Inverse of nvfp4_4over6_code(), for logs, repro tooling and tests.

make_nvfp4_global_scale(input_tensor, ...[, ...])

Build the NVFP4 global scale implied by a resolved 4over6 recipe.

nvfp4_e4m3_max([nvfp4_4over6_config])

E4M3 block-scale clamp implied by a resolved recipe.

FP4 Quantization

Core kernels for NVFP4 / MXFP4 (de)quantization and the scale-factor layout helpers used by the FP4 GEMM/MoE pipelines.

fp4_quantize(input[, global_scale, ...])

Quantize input tensor to FP4 format.

nvfp4_quantize(a, a_global_sf[, sfLayout, ...])

Quantize input tensor to NVFP4 format.

nvfp4_batched_quantize(a, a_global_sf[, ...])

Quantize batched input tensor to NVFP4 format.

mxfp4_quantize(a[, backend, enable_pdl, ...])

Quantize input tensor to MXFP4 format.

mxfp4_dequantize(a_fp4, a_sf[, sfLayout])

Dequantize MXFP4 packed weights back to float32.

mxfp4_dequantize_host(weight, scale[, ...])

Host-side MXFP4 dequantization.

block_scale_interleave(unswizzled_sf)

Swizzle a block-scale tensor for FP4 layouts.

e2m1_and_ufp8sf_scale_to_float(e2m1_tensor, ...)

Dequantize an E2M1 tensor with UFP8 scales back to float32.

scaled_fp4_grouped_quantize(a, mask, a_global_sf)

Quantize a batched input tensor to NVFP4 with a per-row mask.

silu_and_mul_nvfp4_quantize(input, global_scale)

Apply SwiGLU and NVFP4 quantization in one CuTe-DSL kernel.

shuffle_matrix_a(input_tensor, epilogue_tile_m)

PyTorch equivalent of TRT-LLM-gen shuffleMatrixA.

shuffle_matrix_sf_a(input_tensor, ...[, ...])

CUDA implementation of TRT-LLM-gen shuffleMatrixSfA for linear-layout SF.

Note

flashinfer.quantization.nvfp4_block_scale_interleave is an alias for block_scale_interleave() (same Python object). Use either name; we document the canonical block_scale_interleave to avoid Sphinx duplicate object description warnings under -W.

FP4 KV Cache Quantization

GPU-accelerated quantization / dequantization for KV-cache data using the linear (non-swizzled) block-scale layout.

nvfp4_kv_quantize(input, global_scale)

GPU quantization to the NVFP4 KV-cache format with linear block-scale layout.

nvfp4_kv_dequantize(fp4_data, block_scales, ...)

GPU dequantization of an NVFP4 KV cache with linear block-scale layout.

nvfp4_kv_dequantize_paged(paged_kv_cache, ...)

Dequantize a paged NVFP4 KV cache into caller-owned contiguous outputs.

nvfp4_quantize_paged_kv_cache(k_cache, v_cache)

Quantize a paged KV cache to NVFP4 for the trtllm-gen MHA kernel.

FP8 Quantization

mxfp8_quantize(input[, ...])

Quantize input tensor to MxFP8 format.

mxfp8_grouped_quantize(a, mask[, backend])

Quantize grouped inputs to MXFP8 with UE8M0 block scales.

mxfp8_dequantize_host(input, scale_tensor[, ...])

Host-side dequantization of an MxFP8 tensor back to float32.

per_token_group_quant_8bit(x, group_size[, ...])

Per-token group 8-bit quantization (FP8 or INT8).

Fused gated-activation MXFP8 quantization

silu_and_mul_mxfp8_quantize(gated_input, *)

Apply SwiGLU and emit RCEIL MXFP8 rowwise and/or colwise outputs.

silu_and_mul_mxfp8_quantize_backward(...[, ...])

Apply the SwiGLU backward transform and emit RCEIL MXFP8 outputs.

Note

mxfp8_grouped_quantize uses a cuTile backend and requires SM100+ and cuda.tile (a requirements.txt dependency). K must be divisible by 32 and is padded internally to 128-column tiles.

CuTe-DSL Quantization Kernels (experimental)

The CuTe-DSL backends are conditionally available when the nvidia-cutlass-dsl package is installed. At runtime they are also re-exported as flashinfer.quantization.{nvfp4,mxfp4,mxfp8}_quantize_cute_dsl when available; documenting them here via their canonical submodule path keeps the docs build from depending on the CuTe-DSL stack being importable.

nvfp4_quantize_cute_dsl(input, global_scale)

Quantize input tensor to NVFP4 format using the CuTe-DSL kernel.

nvfp4_quantize_per_token_cute_dsl(input, ...)

Per-token NVFP4 activation quantization using the CuTe-DSL kernel.

mxfp4_quantize_cute_dsl(input[, sf_layout, ...])

Quantize input tensor to MXFP4 format using the CuTe-DSL kernel.

mxfp8_quantize_cute_dsl(input[, ...])

Quantize input tensor to MXFP8 format using the CuTe-DSL kernel.