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 = 'cuda', 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*>)

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:

    • "cuda": stable CUDA kernel (default).

    • "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.

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.