flashinfer.fp4_quantization.nvfp4_quantize

flashinfer.fp4_quantization.nvfp4_quantize(a, a_global_sf, sfLayout=SfLayout.layout_128x4, do_shuffle=False, sf_vec_size=16, enable_pdl=None)

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

参数:

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

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

• sfLayout (SfLayout, 可选) – 缩放因子布局。默认为 SfLayout.layout_128x4。

• do_shuffle (bool, 可选) – 是否打乱缩放因子。默认为 False。只有 TRTLLM 后端才需要打乱张量 B 的缩放因子。

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

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

返回值:

包含一个元组

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

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

返回值类型:

Tuple[torch.Tensor, torch.Tensor]