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; cuBLASLt supports bias with BF16 output. Defaults to None.

  • pdl (bool) – Whether to use Programmatic Dependent Launch. Enabled for TGV, TinyGEMM, and CuTeDSL backends. The CuTeDSL M > 32 fallback ignores PDL because cuBLASLt does not expose it. 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 and an optional fused BF16 bias epilogue for BF16 output. "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 kernels for M <= 32 (direct, cluster Split-K and warp Split-K) and cuBLASLt above that. It is never auto-selected; serving frameworks must select it explicitly. Without autotuning, M > 32 runs cuBLASLt; below, the direct kernel runs where its shape heuristic applies, otherwise the warp Split-K kernel whenever it is eligible (N % 16 == 0, K % 128 == 0 with at most 64 K tiles), and cluster Split-K otherwise. With autotuning, one call profiles the low-M kernels on the M <= 32 buckets and the cuBLASLt fallback on the larger ones, so a single large-M warm-up tunes both ranges; with bias the direct kernel is excluded. "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