flashinfer.quantization.mxfp4_quantize

flashinfer.quantization.mxfp4_quantize(a: Tensor, backend: str = 'auto', enable_pdl: bool | None = None, sfLayout: SfLayout = SfLayout.layout_128x4, *, batch_invariant: bool = False) → Tuple[Tensor, Tensor]

Quantize input tensor to MXFP4 format.

Parameters:

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

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

  • enable_pdl (bool, optional) – Whether to enable Programmatic Dependent Launch. Only used when backend == "cute-dsl". Auto-detected from device capability when None.

  • sfLayout (SfLayout) – Scale-factor layout. Defaults to SfLayout.layout_128x4. Supported layouts are layout_128x4, layout_8x4, and layout_linear.

  • batch_invariant (bool) – True runs the batch-invariant MXFP4 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 MX chain carried to e2m1, fingerprinted as flashinfer.batch_invariant.fingerprint("mxfp4_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/2] with dtype uint8 (FLOAT4_E2M1X2) and sf is the UE8M0 scale-factor tensor (uint8) 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. Use at your own risk for production workloads.