flashinfer.gemm.mm_mxfp8

flashinfer.gemm.mm_mxfp8(a: Tensor, b: Tensor, a_descale: Tensor, b_descale: Tensor, out: Tensor | None = None, out_dtype: dtype = torch.bfloat16, use_8x4_sf_layout: bool = False, backend: Literal['cutlass', 'cute-dsl', 'trtllm', 'cudnn', 'auto'] = 'auto') → Tensor

MM MXFP8 (block size 32)

Parameters:

  • a (torch.Tensor) – Input A tensor, shape (m, k), mxfp8 e4m3.

  • b (torch.Tensor) – Input B tensor, shape (k, n), should be column major, mxfp8 e4m3.

  • a_descale (torch.Tensor) – Block scale tensor for A, uint8 (fp8 e8m0), 1D swizzled layout: shape (M_padded * K_padded,) where M_padded = round_up(m, 8 if 8x4 layout else 128) and K_padded = round_up(k // 32, 4). Produced by mxfp8_quantize(..., is_sf_swizzled_layout=True). The 8x4 layout (use_8x4_sf_layout=True) is only consumed by the trtllm backend. 2D linear scales are not supported by any backend and raise ValueError.

  • b_descale (torch.Tensor) – Block scale tensor for B, uint8 (fp8 e8m0), 1D swizzled 128x4 layout: shape (N_padded * K_padded,) where N_padded = round_up(n, 128) and K_padded = round_up(k // 32, 4), flattened from the (N_padded, K_padded) grid. For the trtllm backend, quantize with the linear layout and shuffle with shuffle_matrix_sf_a instead (it emits the swizzled+shuffled layout trtllm expects).

  • out (Optional[torch.Tensor]) – Out tensor, shape (m, n), bf16 or fp16. If provided, the result is written into it (supported by the CUTLASS and cuDNN backends). Defaults to None.

  • out_dtype (torch.dtype) – Output dtype, bf16 or fp16. Defaults to torch.bfloat16.

  • use_8x4_sf_layout (bool) – Whether the scale tensors for a are in 8x4 layout (vs 128x4).

  • backend (Literal["cutlass", "cute-dsl", "trtllm", "cudnn", "auto"]) –

    The backend to use for the operation. Defaults to "auto". "auto" selects the CUTLASS backend when available and otherwise falls back to the cuDNN backend. - The "cute-dsl" backend currently requires swizzled 1D scales

    (mxfp8_quantize(..., is_sf_swizzled_layout=True)).

    • The "trtllm" requires b to be quantized with 128x4 swizzle layout and shuffled. a can be quantized with either 128x4 or 8x4 layout (controlled by use_8x4_sf_layout).

    • The "cutlass" backend only supports 1D swizzled scales (SfLayout.layout_128x4); the kernel has no linear-scale path. Passing 2D linear scales raises ValueError. Use mxfp8_quantize(..., is_sf_swizzled_layout=True).

    • The "cudnn" backend consumes block scales in the F8_128x4 swizzled layout (use_8x4_sf_layout=False) and is supported on SM100/103/110/120/121.

Returns:

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

Return type:

torch.Tensor

Examples

>>> import torch
>>> from flashinfer import mxfp8_quantize, mm_mxfp8
>>> m, n, k = 512, 256, 128
>>> # Create input tensors - note: weight is [n, k] for typical NN layers
>>> a = torch.randn([m, k], device="cuda", dtype=torch.bfloat16)
>>> weight = torch.randn([n, k], device="cuda", dtype=torch.bfloat16)
>>>
>>> # Quantize input [m, k] - scales are 1D swizzled for (M, K/32) layout
>>> a_mx, a_sf = mxfp8_quantize(input=a, is_sf_swizzled_layout=True)
>>> # Quantize weight [n, k] - scales are 1D swizzled for (N, K/32) layout
>>> w_mx, w_sf = mxfp8_quantize(input=weight, is_sf_swizzled_layout=True)
>>> # Pass weight.T as [k, n] and 1D swizzled scales directly
>>> out = mm_mxfp8(a_mx, w_mx.t(), a_sf, w_sf, out_dtype=torch.bfloat16)
>>> out.shape
torch.Size([512, 256])