flashinfer.quantization.mxfp8_quantize

flashinfer.quantization.mxfp8_quantize(input: Tensor, is_sf_swizzled_layout: bool = True, alignment: int = 32, enable_pdl: bool | None = None, backend: Literal['auto', 'cuda', 'cute-dsl'] = 'auto', sf_swizzle_layout: SfLayout | None = None, *, batch_invariant: bool = False) → Tuple[Tensor, Tensor]

Quantize input tensor to MxFP8 format.

Implements MxFP8 quantization that converts input tensors to a compressed MxFP8 format with associated scale factors. Supports various input data types and scale-factor layouts.

Parameters:

  • input (torch.Tensor) – Input tensor of shape [M, K] with dtype fp16/bf16/fp8_quantized. The "cute-dsl" backend additionally supports fp32 input; the "cuda" backend does not.

  • is_sf_swizzled_layout (bool) – Whether to use the swizzled layout for scale factors. Defaults to True.

  • alignment (int) – sfVecSize. Defaults to 32.

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

  • backend ({"auto", "cuda", "cute-dsl"}) –

    Backend to use:

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

    • "cuda": stable JIT-compiled CUDA kernel.

    • "cute-dsl": CuTe-DSL kernel (SM100+, experimental).

  • sf_swizzle_layout (SfLayout, optional) – Swizzle layout for scale factors; when supplied this overrides is_sf_swizzled_layout.

  • batch_invariant (bool) – True runs the batch-invariant MXFP8 quantization class in place of the backends (requires backend="auto"): a row’s bits depend on its values alone, not on the batch it is quantized in, its position, the configuration chosen or the SM count, and the result is deterministic. The dataflow is TransformerEngine’s MXFP8 chain, fingerprinted as flashinfer.batch_invariant.fingerprint("mxfp8_quantize"). Unsupported inputs 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] with dtype FLOAT8_E4M3 and sf is the scale-factor tensor whose shape depends on the chosen layout and sf_vec_size (fixed at 32 here).

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.