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'] = 'cutlass') → Tuple[Tensor, Tensor]

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

Applies the SVDQuant per-input-channel smoothing scale and NVFP4-quantizes in one pass over the input; the result is byte-identical to quantizing x * pre_quant_scale with the stock NVFP4 quantizer (ue4m3 block scales, 128x4 swizzled layout, 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"]) – Only the CUDA backend exists.

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.