flashinfer.gemm.bmm_mxfp8

flashinfer.gemm.bmm_mxfp8(A: Tensor, B: Tensor, A_scale: Tensor, B_scale: Tensor, dtype: dtype, out: Tensor | None = None, backend: Literal['cudnn', 'cutlass', 'auto'] = 'auto') → Tensor

BMM MXFP8

Parameters:

  • A (torch.Tensor) – Input tensor, shape (b, m, k), fp8 e4m3 or fp8 e5m2.

  • B (torch.Tensor) – Mat2 tensor, shape (b, k, n), must be column major, fp8 e4m3 or fp8 e5m2. Quantize the contiguous [b, n, k] weight (so the 32-element scale blocks run along k, the reduction dim) and pass the transpose of the quantized tensor, e.g. B = mxfp8_quantize(weight)[0].transpose(-2, -1) (do NOT call .contiguous() on the transpose).

  • A_scale (torch.Tensor) – Scale tensor for A, uint8 (fp8 e8m0 format), in the F8_128x4 swizzled layout produced by mxfp8_quantize(..., is_sf_swizzled_layout=True).

  • B_scale (torch.Tensor) – Scale tensor for B, uint8 (fp8 e8m0 format), in the F8_128x4 swizzled layout, as returned by quantizing the [b, n, k] weight (see B).

  • dtype (torch.dtype) – out dtype, bf16 or fp16.

  • out (Optional[torch.Tensor]) – Out tensor, shape (b, m, n), bf16 or fp16, defaults to None.

  • backend (Literal["cudnn", "cutlass", "auto"]) – The backend to use for the operation. Defaults to "auto". On SM120/121 GPUs, "auto" selects the CUTLASS backend; scales must be 1D swizzled (SfLayout.layout_128x4). Pass B in the standard shape [b, k, n] (column-major); the CUTLASS path transposes internally. Both the cuDNN and CUTLASS backends read the scale tensors in the F8_128x4 swizzled layout; linear-layout scales are not supported. Both layouts are flat 1D buffers, so a linear scale whose length happens to match the padded swizzled length cannot be detected at runtime. Ensure the scale was produced with the swizzled layout rather than relying on a warning.

Returns:

out – Out tensor, shape (b, m, n), bf16 or fp16.

Return type:

torch.Tensor