flashinfer.quantization.silu_and_mul_mxfp8_experts_quantize

flashinfer.quantization.silu_and_mul_mxfp8_experts_quantize(a: Tensor, mask: Tensor, *, enable_pdl: bool | None = None) → Tuple[Tensor, Tensor]

Fused SiLU + mul + MXFP8 quantization with a per-expert row mask.

Quantizes silu(a[..., :K]) * a[..., K:], rounded to the input dtype, to MXFP8 (E4M3 values, one UE8M0 scale per 32 values). Requires SM100+. The outputs are the (A, SFA) operand of the second flashinfer.gemm.grouped_gemm_nt_masked() of a MoE.

Parameters:

  • a (torch.Tensor) – Contiguous, 32-byte aligned float16 or bfloat16 tensor of shape [B, M, 2 * K]; each row holds the gate values, then the up values. K must be divisible by 32.

  • mask (torch.Tensor) – Contiguous int32 CUDA tensor of shape [B], read on the device. Only the first min(mask[i], M) rows of expert i are read and written.

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

Returns:

(x_q, sf) in the layout of mxfp8_grouped_quantize(): x_q has logical shape [M, padded_K, B] with dtype float8_e4m3fn and sf has logical shape [32, 4, padded_M // 128, 4, padded_K // 128, B] with dtype uint8, where padded_K and padded_M round K and M up to multiples of 128. Rows past the valid ones (and their scale factors) are unspecified.

NaN and infinities quantize to 448, and the other values of a block with an infinity to zero.

Return type:

Tuple[torch.Tensor, torch.Tensor]