flashinfer.activation.silu_and_mul_scaled_nvfp4_experts_quantize

flashinfer.activation.silu_and_mul_scaled_nvfp4_experts_quantize(a, mask, a_global_sf)

使用掩码将批处理输入张量转换为 NVFP4 格式，并进行 Silu 和乘法运算和量化。 :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 :param sf_vec_size: 缩放因子向量大小。默认为 16。 :type sf_vec_size: int, optional

返回值:

包含一个元组

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

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

返回值类型:

Tuple[torch.Tensor, torch.Tensor]