flashinfer.gemm.mm_bf16

flashinfer.gemm.mm_bf16(a: Tensor, b: Tensor, bias: Tensor | None = None, pdl: bool = False, out: Tensor | None = None, out_dtype: dtype = torch.bfloat16, backend: Literal['cudnn', 'cutlass', 'tgv', 'cublaslt', 'tinygemm', 'cutile', 'cute-dsl', 'auto'] = 'cudnn') → Tensor

MM BF16

Parameters:

  • a (torch.Tensor) – Input tensor, shape (m, k), bf16 in row-major layout.

  • b (torch.Tensor) – Weight tensor, shape (k, n), bf16 in column-major layout.

  • bias (Optional[torch.Tensor]) – Optional bias tensor, shape (n,). Enabled for TGV, TinyGEMM, and CuTeDSL backends. Defaults to None.

  • pdl (bool) – Whether to use Programmatic Dependent Launch. Enabled for TGV, TinyGEMM, and CuTeDSL backends. Defaults to False.

  • out (Optional[torch.Tensor]) – Out tensor, shape (m, n), bf16, fp16, or fp32. FP16 and FP32 output are enabled for CUTLASS and cuDNN backends; TinyGEMM and CuTeDSL require bf16 output.

  • out_dtype (torch.dtype) – Output dtype, bf16, fp16, or fp32. Enabled for CUTLASS, cuDNN, and cuBLASLt backends. TinyGEMM and CuTeDSL require torch.bfloat16. Defaults to torch.bfloat16.

  • backend (Literal["cudnn", "cutlass", "tgv", "cublaslt", "tinygemm", "cutile", "cute-dsl", "auto"]) –

    The backend to use for the operation. Defaults to "cudnn". "cudnn" uses the cuDNN backend. "cutlass" uses the CUTLASS backend. "tgv" uses the TGV backend. "cublaslt" uses the cuBLASLt backend with heuristic algorithm search. "tinygemm" uses the TinyGEMM backend for small-M BF16 GEMM. "cutile" uses the cuTile (cuda.tile Python) backend. Pure-Python

    persistent-scheduled GEMM with per-shape exhaustive autotune; ignores bias / pdl. Requires SM >= 90.

    "cute-dsl" uses standalone Blackwell low-M direct and split-K kernels for M <= 32. It is never auto-selected; serving frameworks must select it explicitly. Without autotuning, a measured shape heuristic chooses between the algorithms. With autotuning, both tactic spaces are profiled when bias is disabled; bias uses split-K only. "auto" allows selecting the best tactic from all available backends when autotune is enabled.

Returns:

Out tensor, shape (m, n), bf16, fp16, or fp32 in row-major layout.

Return type:

torch.Tensor

Examples

>>> import torch
>>> import flashinfer
>>> # Using the TGV backend
>>> a = torch.randn([48, 64], device="cuda", dtype=torch.bfloat16)
>>> b = torch.randn([80, 64], device="cuda", dtype=torch.bfloat16).transpose(-2, -1)
>>> bias = torch.randn([80], device="cuda", dtype=torch.bfloat16)
>>> out = flashinfer.mm_bf16(a, b, bias=bias, pdl=True, backend="tgv")
>>> out.shape
torch.Size([48, 80])
>>> out.dtype
torch.bfloat16
>>> # Using the CUTLASS backend
>>> fp16_out = torch.empty([48, 80], device="cuda", dtype=torch.float16)
>>> out = flashinfer.mm_bf16(a, b, out=fp16_out, out_dtype=torch.float16, backend="cutlass")
>>> out.shape
torch.Size([48, 80])
>>> out.dtype
torch.float16
>>> # Using the cuDNN backend
>>> out = flashinfer.mm_bf16(a, b, backend="cudnn")
>>> out.shape
torch.Size([48, 80])
>>> out.dtype
torch.bfloat16
>>> # Using the cuBLASLt backend
>>> out = flashinfer.mm_bf16(a, b, backend="cublaslt")
>>> out.shape
torch.Size([48, 80])
>>> out.dtype
torch.bfloat16