flashinfer.fp4_quantization.fp4_quantize

flashinfer.fp4_quantization.fp4_quantize(input: Tensor, global_scale: Tensor | None = None, sf_vec_size: int = 16, sf_use_ue8m0: bool = False, is_sf_swizzled_layout: bool = True, is_sf_8x4_layout: bool = False, enable_pdl: bool | None = None) → Tuple[Tensor, Tensor]

将输入张量量化为 FP4 格式。

此函数实现 FP4 量化，将输入张量转换为压缩的 FP4 格式，并带有相关的缩放因子。它支持各种输入数据类型和缩放因子布局。

参数:

• input (torch.Tensor) – 输入张量，形状为 [M, K]，数据类型为 fp16/bf16/fp8_quantized。

• global_scale (torch.Tensor, 可选) – 全局缩放因子，形状为 [1]，数据类型为 float32。

• sf_vec_size (int, 可选) – 缩放因子向量大小。默认为 16。

• sf_use_ue8m0 (bool, 可选) – 是否为缩放因子使用 UE8M0 格式。默认为 False。

• is_sf_swizzled_layout (bool, 可选) – 是否为缩放因子使用交错布局。默认为 True。

• is_sf_8x4_layout (bool, 可选) – 是否为缩放因子使用 8x4 布局或 128x4 布局。默认为 False。

• enable_pdl (Optional[bool], 可选) – 是否启用 PDL（程序依赖启动）。如果为 None，则根据设备功能自动检测。默认为 None。

返回值:

包含一个元组

• 量化后的张量，形状为 [M, K/2]，数据类型为 FLOAT4_E2M1X2

• 缩放因子张量，其形状由布局和 sf_vec_size 决定

返回值类型:

Tuple[torch.Tensor, torch.Tensor]

抛出：

NotImplementedError – 如果请求了以下任何功能但尚未实现：- 启用 BFloat16 时输入为 BFloat16 - 启用 FP8 时输入为 FP8 - sf_vec_size 不是 16 或 32