flashinfer.quantization.nvfp4_quantize

flashinfer.quantization.nvfp4_quantize(a, a_global_sf, sfLayout=SfLayout.layout_128x4, do_shuffle=False, sf_vec_size=16, enable_pdl=None, backend: str = 'auto', per_token_activation: bool = False, expanded_idx_to_permuted_idx: ~torch.Tensor | None = None, out_scale: ~torch.Tensor | None = None, nvfp4_4over6: ~flashinfer.quantization.nvfp4_quantization_utils.NVFP44Over6Config | None = <nvfp4_4over6 unset: read FLASHINFER_NVFP4_4OVER6*>, *, batch_invariant: bool = False)

Quantize input tensor to NVFP4 format.

Parameters:

  • a (torch.Tensor) – Input tensor of shape [M, K] with dtype fp16/bf16/float8_e4m3fn.

  • a_global_sf (float or torch.Tensor) – Global scale factor. The CuTe-DSL backend accepts a host-side float or, for backward compatibility, a single-element tensor.

  • sfLayout (SfLayout) – Scale-factor layout. Defaults to SfLayout.layout_128x4.

  • do_shuffle (bool) – Whether to shuffle the scale factors. Only the TRT-LLM backend needs to shuffle the tensor-B scale factors. Defaults to False.

  • sf_vec_size (int) – Scale-factor vector size. Defaults to 16.

  • enable_pdl (bool, optional) – Whether to enable Programmatic Dependent Launch. Auto-detected from device capability when None.

  • backend (str) –

    Backend to use for quantization:

    • "auto" (default): "cuda", or the batch-invariant kernel under batch_invariant=True (which requires "auto").

    • "cuda": stable CUDA kernel on SM100+.

    • "cute-dsl": CuTe-DSL kernel (SM100+, experimental); supports all sfLayout values (layout_128x4 / layout_8x4 / layout_linear) and input dtypes fp16/bf16/float8_e4m3fn, but only sf_vec_size == 16.

    • "cake": generated Cake per-token kernel (SM100/SM103, experimental, explicit opt-in); per_token_activation=True with layout_128x4, no shuffle and fp16/bf16 input only, outputs bitwise equal to "cute-dsl".

  • per_token_activation (bool) – Whether to use per-token NVFP4 activation scaling. In this mode a_global_sf is the inverse base scale multiplier (typically 1 / (448 * 6)) and the function also returns per-token FP32 scales.

  • expanded_idx_to_permuted_idx (torch.Tensor, optional) – Optional row-remapping buffer for per-token activation quantization.

  • out_scale (torch.Tensor, optional) – Scalar the returned per-token scales are multiplied by. Only for per_token_activation=True with backend="cute-dsl" or backend="cake". Does not change the quantized values.

  • nvfp4_4over6 (NVFP44Over6Config or None) –

    NVFP4 “4over6” scale-candidate search. Requires fp16 / bf16 input.

    • omitted (the default): read the legacy FLASHINFER_NVFP4_4OVER6* environment variables.

    • None: 4over6 off. The environment is ignored.

    • NVFP44Over6Config: on with exactly that recipe. The environment is ignored.

  • batch_invariant (bool) – True runs the batch-invariant NVFP4 quantization class in place of the backends (requires backend="auto"): a row’s bits depend on its values and the global scale alone, not on the batch it is quantized in, its position, the configuration chosen or the SM count, and the result is deterministic. Per-tensor (the class nvfp4.quant: a_global_sf is an input; flashinfer.batch_invariant.fingerprint("nvfp4_quantize")) and per-token (nvfp4.quant.pertoken: the row’s scale is derived from its amax; fingerprint("nvfp4_quantize_per_token", K=)), all three scale-factor layouts and do_shuffle. A 4over6 recipe (explicit, or enabled through the FLASHINFER_NVFP4_4OVER6 environment when nvfp4_4over6 is omitted), sf_vec_size != 16, an fp8 input, expanded_idx_to_permuted_idx and out_scale raise flashinfer.batch_invariant.BatchInvariantUnsupportedError instead of falling back. See flashinfer/batch_invariant/README.md.

Returns:

(x_q, sf) where x_q has shape [M, K/2] with dtype FLOAT4_E2M1X2 and sf is the scale-factor tensor (shape depends on the layout and sf_vec_size). When per_token_activation=True, a third tensor containing per-token FP32 scales is also returned.

Return type:

Tuple[torch.Tensor, torch.Tensor]

Warning

The "cute-dsl" backend is experimental and not part of the stable API. It may change or be removed in future versions without notice.

Notes

Build the global scale from the same recipe: flashinfer.make_nvfp4_global_scale(..., nvfp4_4over6_config=<same value>). With per_token_activation=True the pairing is checked when nvfp4_4over6 is explicit and a_global_sf is readable without a device synchronization.