flashinfer.gemm.mm_bf16

flashinfer.gemm.mm_bf16(a: Tensor, b: Tensor, bias: Tensor | None = None, pdl: bool | None = False, out: Tensor | None = None, out_dtype: dtype = torch.bfloat16, backend: Literal['cudnn', 'cutlass', 'tgv', 'cublaslt', 'tinygemm', 'cutile', 'cute-dsl', 'auto'] = 'cudnn', *, batch_invariant: bool = False) → 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 (Optional[bool]) – Whether to use Programmatic Dependent Launch. Enabled for TGV, TinyGEMM, and CuTeDSL backends. None leaves it to the autotuner: under autotuning, each of these backends is profiled both without and with PDL and the faster variant is cached; without a tuned entry they run without PDL. These entries use the op name "bf16_gemm_pdl_auto", while skip_ops={"bf16_gemm"} still skips such calls. PDL kernels may read b and bias before waiting for the preceding kernel, so pass pdl=False if a kernel that may still be running writes them. Defaults to False so autotuning cannot silently enable this input-dependency assumption.

  • 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 (direct, cluster Split-K and warp Split-K) and supports M <= 32 only; use "auto" or "cublaslt" for larger M. Without autotuning, 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; requires CuTe DSL >= 4.7), and cluster Split-K otherwise. With autotuning, one call profiles the available low-M kernels on every M <= 32 bucket; under "auto" a single large-M warm-up also tunes these buckets. The direct kernel is profiled only for M <= 3 or M * N <= 4096, and never with bias. Direct and warp Split-K reuse compiled kernels across M for matching tactics, while each M bucket still tunes independently. "auto" allows selecting the best tactic from all available backends when autotune is enabled.

  • batch_invariant (bool) – True runs the batch-invariant bf16 GEMM class in place of the backends (requires backend="auto" and bf16 output): a row of the output is bit-identical whatever the number of rows, the row’s position in the batch or the configuration chosen, and the result is deterministic. Under autotuning the class’s top configurations for the shape are profiled; all of them compute the same bits. b may be the row-major (k, n) tensor or the transpose of a row-major (n, k) weight. A bf16 bias of shape (n,) is added once, in fp32, to the fully reduced accumulator before the single rounding to bf16 (a biased call is a different specification: flashinfer.batch_invariant.fingerprint("mm_bf16", N=, K=, bias=True)). Unsupported inputs raise flashinfer.batch_invariant.BatchInvariantUnsupportedError instead of falling back. See flashinfer/batch_invariant/README.md. Defaults to False.

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