flashinfer.grouped_mm.moe_gemm_fp8_nt_groupwise

flashinfer.grouped_mm.moe_gemm_fp8_nt_groupwise(a: Tensor, b: Tensor, a_scale: Tensor, b_scale: Tensor, m_indptr: Tensor, scale_granularity_mnk: Tuple[int, int, int] = (1, 128, 128), scale_major_mode: Literal['MN'] = 'MN', backend: Literal['cute'] = 'cute', out: Tensor | None = None, out_dtype: dtype | None = None) → Tensor

Perform grouped GEMM with FP8 inputs in zero-padding mode using groupwise float32 scaling. Currently only supported on NVIDIA RTX PRO 6000 Blackwell (SM120) architecture.

Zero-padding mode accepts token-packed input a (no per-expert pre-padding along M) with 4-row per-expert padding on the scale tensor a_scale. The group descriptor is a CSR cumsum m_indptr. This mode is optimized for decoding with small per-expert M (down to m_per_expert = 1) where DeepGEMM-style contiguous padding would waste memory and compute.

Parameters:

  • a (torch.Tensor) – Row-major input tensor shape (cum_m, k), data type is torch.float8_e4m3fn. Token-packed across experts; cum_m is the cumulative sum of segment lengths.

  • b (torch.Tensor) – Column-major input tensor shape (num_experts, n, k), data type is torch.float8_e4m3fn.

  • a_scale (torch.Tensor) – Float32 scale tensor for a in zero-padding MN-major layout: contiguous shape (k_blocks, m_padded) where k_blocks = ceil(k / 128) and m_padded = (cum_m + num_experts * 3) // 4 * 4. Expert i’s scales start at column (m_indptr[i] + 3 * i) // 4 * 4; padding columns are ignored by the kernel. The data pointer must be 16-byte aligned.

  • b_scale (torch.Tensor) – Float32 scale tensor for b, contiguous shape (num_experts, k_blocks, n_blocks) with n_blocks = ceil(n / 128) (one scale per (128, 128) block).

  • m_indptr (torch.Tensor) – The indptr of the segment lengths, shape (num_experts + 1,), data type is torch.int32. m_indptr[0] = 0, m_indptr[num_experts] = cum_m.

  • scale_granularity_mnk (Tuple[int, int, int]) – The granularity of the scale tensors. The FP8 float-scale kernel contract is fixed to (1, 128, 128) (per-token A-scale, (128, 128)-block B-scale); anything else raises ValueError.

  • scale_major_mode (Literal["MN"]) – The layout mode of scale tensors. Currently only "MN" is supported. Defaults to "MN".

  • backend (Literal["cute"]) – Backend selector. Currently only "cute" is implemented.

  • out (Optional[torch.Tensor]) – The output tensor, shape (cum_m, n). If not specified, an output tensor will be created.

  • out_dtype (Optional[torch.dtype]) – The data type of the output tensor. Currently only torch.bfloat16 is supported.

Returns:

out – The output tensor, shape (cum_m, n).

Return type:

torch.Tensor

Notes

  • Unlike flashinfer.grouped_mm.moe_gemm_mxfp8_nt_groupwise() (int32-packed UE8M0 scales, per-token along both M and N), the FP8 entry consumes plain float32 scales: per-token along M for a and one scale per (128, 128) block for b.

  • A row-major (cum_m, k_blocks) per-token A-scale must be transposed and re-packed into the padded (k_blocks, m_padded) layout described above (zero-fill the padding columns; copy expert i’s transposed scales to its 4-row-aligned start column).