flashinfer.gemm.nvfp4_quantize_smooth

flashinfer.gemm.nvfp4_quantize_smooth(x: Tensor, pre_quant_scale: Tensor, global_scale: Tensor, enable_pdl: bool | None = None, backend: Literal['cutlass', 'cute-dsl', 'auto'] = 'auto') → Tuple[Tensor, Tensor]

Smooth + NVFP4 quantize: (xq, sf) = nvfp4-quantize(x * pre_quant_scale).

The SM100/SM103 CUTLASS backend applies the SVDQuant per-input-channel smoothing scale and NVFP4-quantizes in one pass. The SM120/SM121 CuTe DSL backend also applies smoothing inside the NVFP4 quantizer, avoiding a BF16 intermediate. Both use ue4m3 block scales, the 128x4 swizzled layout, and SF vector size 16.

Parameters:

  • x (torch.Tensor) – Input activation, shape (m, n) bf16.

  • pre_quant_scale (torch.Tensor) – Per-input-channel smoothing scale, shape (n,) bf16.

  • global_scale (torch.Tensor) – Global scale, float32 device scalar: (448 * 6) / (x * pre_quant_scale).abs().max().

  • enable_pdl (Optional[bool]) – Whether to launch with Programmatic Dependent Launch. Defaults to the device default.

  • backend (Literal["cutlass", "cute-dsl", "auto"]) – "cutlass" selects fused smoothing and quantization on SM100/SM103; "cute-dsl" selects fused smoothing plus CuTe DSL quantization on SM120/SM121; "auto" (default) selects by compute capability.

Returns:

  • xq (torch.Tensor) – Quantized tensor, shape (m, n // 2) uint8 (packed e2m1).

  • sf (torch.Tensor) – Block scales, uint8 (ue4m3), 128x4 swizzled layout, 1-D of size ceil(m / 128) * 128 * ceil(n / 16 / 4) * 4.