flashinfer.gemm.tinygemm_bf16

flashinfer.gemm.tinygemm_bf16(input: Tensor, weight: Tensor, out: Tensor, bias: Tensor | None = None, use_pdl: bool = False) → None

SM90+ optimized small GEMM: out = input @ weight.T + bias (equivalent to F.linear).

A latency-optimized, warp-specialized GEMM designed for tiny batch sizes (ideally 1-8 rows, where a single TILE_N=8 tile covers the entire batch dimension) using Ampere-style HMMA instructions. Uses TMA for async bulk data loads and mma.sync.aligned.m16n8k16 tensor-core instructions with BF16 input/weight/bias/output and FP32 internal accumulation. The warp-specialized design (384 threads: 4 compute + 8 DMA warps) with 16 pipeline stages and 4x stage unroll trades off peak throughput in favor of minimal latency. Adapted from the TensorRT-LLM tinygemm2 kernel.

Parameters:

  • input (torch.Tensor) – Input activations of shape (batch_size, input_features). Must be bfloat16, contiguous. input_features must be a multiple of 64.

  • weight (torch.Tensor) – Weight matrix of shape (output_features, input_features). Must be bfloat16, contiguous (row-major). output_features must be a multiple of 16.

  • out (torch.Tensor) – Pre-allocated output tensor of shape (batch_size, output_features). Must be bfloat16, contiguous. Mutated in place.

  • bias (Optional[torch.Tensor]) – Optional bias vector of shape (output_features,). Must be bfloat16, contiguous. If None, zero bias is used.

  • use_pdl (bool) – Enable Programmatic Dependent Launch (stream serialization). When True, the kernel uses cudaGridDependencySynchronize() to overlap DMA with the preceding kernel’s compute. Only enable when ALL preceding stream operations also use PDL, otherwise the kernel hangs. The kernel may read weight and bias before that dependency wait, so the preceding kernel must not write them; only input may be its output. Defaults to False.

Notes

Requires SM90+ (Hopper or newer). Raises ValueError if tensor dimensions, dtypes, or alignment constraints are violated.

On SM100/SM103 (B200/B300 class) devices the bias path dispatches to tinygemm2_sm100 — generated variants of the same kernel with bit-identical outputs and lower latency (see csrc/tinygemm2_sm100.cu). Set FLASHINFER_DISABLE_TINYGEMM2_SM100=1 to force the reference implementation everywhere.

On the same devices, shapes with large weights (up to 64 batch rows) and narrow outputs over long inputs dispatch to tinygemm2_tcgen05 instead: generated tcgen05 programs whose outputs are bit-identical to the reference’s (flashinfer.jit.cake_tinygemm2_tcgen05). Set FLASHINFER_DISABLE_TINYGEMM2_TCGEN05=1 to keep the kernels above.

On SM90 (H100/H200 class) devices, the shapes the measured windows name dispatch to tinygemm2_wgmma: generated wgmma programs whose outputs are bit-identical to the reference’s (flashinfer.jit.cake_tinygemm2_wgmma). Set FLASHINFER_DISABLE_TINYGEMM2_WGMMA=1 to keep the reference kernel.