flashinfer.fp4_quantization.scaled_fp4_grouped_quantize

flashinfer.fp4_quantization.scaled_fp4_grouped_quantize(a, mask, a_global_sf)

使用掩码将批处理输入张量量化为 NVFP4 格式。 :param a: 形状为 [B, M, K] 且 dtype 为 fp16/bf16 的输入张量。 :type a: torch.Tensor :param a_global_sf: 形状为 [1] 且 dtype 为 float32 的全局缩放因子。 :type a_global_sf: torch.Tensor :param mask: 应用于量化之前的掩码张量。 :type mask: torch.Tensor

返回值:

包含一个元组

• 量化后的张量，形状为 [B, M, K/2]，dtype 为 FLOAT4_E2M1X2

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

返回值类型:

Tuple[torch.Tensor, torch.Tensor]