flashinfer.quantization.kernels.nvfp4_quantize.nvfp4_quantize_cute_dsl

flashinfer.quantization.kernels.nvfp4_quantize.nvfp4_quantize_cute_dsl(input: ~torch.Tensor, global_scale: float | ~torch.Tensor, sf_layout: int = 0, enable_pdl: bool | None = None, nvfp4_4over6: ~flashinfer.quantization.nvfp4_quantization_utils.NVFP44Over6Config | None = <nvfp4_4over6 unset: read FLASHINFER_NVFP4_4OVER6*>) → Tuple[Tensor, Tensor]

Quantize input tensor to NVFP4 format using the CuTe-DSL kernel.

GPU implementation matching flashinfer.quantization.nvfp4_quantize():

  • E4M3 scale factors (FP8)

  • E2M1 output format (4-bit, 2 values per byte)

  • Supports 128x4, 8x4, and linear scale-factor layouts

  • sf_vec_size = 16

The kernel is compiled separately for host-scalar and device-tensor global scales and handles varying M (batch size) at runtime without recompilation.

Parameters:

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

  • global_scale (float or torch.Tensor) – Host-side NVFP4 global scale factor. A single-element tensor is also accepted for backward compatibility.

  • sf_layout (int) – Scale-factor layout (0=128x4, 1=8x4, 2=linear).

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

  • 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:

(fp4_tensor, scale_tensor) where fp4_tensor is the quantized tensor of shape [M, K/2] with dtype uint8 and scale_tensor holds the E4M3 scale factors (uint8) reshaped to [padded_rows, K/16].

Return type:

Tuple[torch.Tensor, torch.Tensor]