Use FI tinygemm for faster BF16 GEMMs

[askliar]

Purpose

Route small-M BF16 linears (M ≤ 8) through FlashInfer's tinygemm_bf16 kernel to speed up decode-phase QKV/O, MLP up/down, shared MoE, and latent MoE projections. Uses TRT-LLM's warp-specialized TMA + HMMA design tuned for latency-bound GEMVs. Requires SM90+ and FlashInfer; falls back to F.linear otherwise.

Changes

• Add flashinfer_tinygemm_bf16 to the FlashInfer lazy-import registry.
• Register torch.ops.vllm.tinygemm_bf16 via direct_register_custom_op (real + fake impls) for torch.compile fullgraph compatibility.
• Add _tinygemm_unquantized_gemm with small-M BF16 predicate and multi-dim input handling; wire it into dispatch_unquantized_gemm() behind an import-time _TINYGEMM_AVAILABLE flag.

Test Plan

• lm_eval accuracy parity on a BF16 model.
• Decode TPOT benchmark at BS=1 on H100.
• torch.compile fullgraph decode (no graph breaks).
• Fallback path on non-SM90 / no-FlashInfer.

[gemini-code-assist]

Code Review

This pull request introduces support for B12X NVFP4 dense linear and fused MoE kernels for Blackwell (SM120/SM121) GPUs and optimizes unquantized GEMM operations for small batch sizes on SM90+ hardware using FlashInfer's tinygemm kernel. A critical correctness issue was identified in the default_unquantized_gemm implementation where multi-dimensional input tensors were not correctly handled, potentially leading to shape mismatches or incorrect kernel triggering. A code suggestion was provided to flatten the tokens and correctly calculate the number of tokens for the threshold check and output tensor allocation.

[gemini-code-assist]

The current implementation of default_unquantized_gemm using tinygemm_bf16 has a correctness issue when the input tensor x is multi-dimensional (e.g., [batch, seq, hidden]).

1. Incorrect Tiny Detection: x.shape[0] <= 8 only checks the first dimension. For a 3D tensor like [1, 1024, hidden], this would incorrectly trigger the tiny GEMM path for 1024 tokens because shape[0] is 1.
2. Shape Mismatch: The output tensor out is allocated as [x.shape[0], weight.shape[0]], which ignores intermediate dimensions (like seq). This will lead to a shape mismatch or crash when returning to the caller.
3. Kernel Application: The tinygemm kernel should be applied to a flattened view of the tokens to correctly handle multi-dimensional inputs.

I suggest calculating the total number of tokens M and using it for both the threshold check and the tensor views.

Suggested change

	if (
	_tinygemm_bf16_available()
	and x.shape[0] <= 8
	and x.dtype == torch.bfloat16
	and weight.dtype == torch.bfloat16
	and weight.shape[0] % 16 == 0
	and x.is_contiguous()
	and weight.is_contiguous()
	and (bias is None or bias.dtype == torch.bfloat16)
	):
	from vllm.utils.flashinfer import flashinfer_tinygemm_bf16
	out = torch.empty(
	x.shape[0], weight.shape[0], dtype=torch.bfloat16, device=x.device
	)
	flashinfer_tinygemm_bf16(x, weight, out, bias=bias)
	return out
	num_tokens = x.numel() // x.shape[-1]
	if (
	_tinygemm_bf16_available()
	and num_tokens <= 8
	and x.dtype == torch.bfloat16
	and weight.dtype == torch.bfloat16
	and weight.shape[0] % 16 == 0
	and x.is_contiguous()
	and weight.is_contiguous()
	and (bias is None or bias.dtype == torch.bfloat16)
	):
	from vllm.utils.flashinfer import flashinfer_tinygemm_bf16
	out = x.new_empty((*x.shape[:-1], weight.shape[0]))
	flashinfer_tinygemm_bf16(x.view(num_tokens, -1), weight,
	out.view(num_tokens, -1), bias=bias)
	return out

[Himan-D]

Hi @WoosukKwon @youkaichao @robertgshaw2 @mgoin @tlrmchlsmth @houseroad @hmellor @yewentao256 @ProExpertProg @pavanimajety - this PR is ready for review!

Summary

• Uses FlashInfer's tinygemm_bf16 kernel for small-M BF16 GEMVs (M ≤ 8)
• Speeds up decode-phase QKV/O, MLP, MoE projections
• Requires SM90+ and FlashInfer; falls back to F.linear otherwise
• Includes comprehensive unit tests

Please review when you have time. Thanks!

[mgoin]

@Himan-D do you know if this tinygemm is PDL enabled? We are interested in a general gemm PDL replacement for low latency

[mgoin]

We should we should not import tinygemm on import of this file. It should be delayed to the first call of the unquantized gemm function itself

[mgoin]

Can we do the capability check only for arches that we have validated? I am wary of enabling this on all NVIDIA GPUs >= Hopper without benchmarks on each platform.

[askliar]

I will do microbenchmarking across multiple platforms to make sure we are honest here.

[mgoin]

It seems very suboptimal to be always allocating and filling bias on a low latency path. I would avoid this

[benchislett]

https://github.com/flashinfer-ai/flashinfer/blob/9e3d8b9d1d11af893f7e7389baa25192db34bd8e/flashinfer/gemm/routergemm.py#L313

Seems non-optional. Not sure what a better solution looks like, but agree that filling the bias every time is a pain

[benchislett]

maybe we can use a pattern like the workspace buffer for FlashInfer attention?

[askliar]

Yes, indeed. I will talk to FlashInfer kernel author to see if that is something that can be changed.

[mgoin]

I would like to see a microbenchmark script to make sure this gemm is good across a range of weight shapes, currently this will apply to most weights small or large

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

Hi @askliar, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[askliar]

@mgoin tinygemm has PDL enabled. Check here: https://github.com/askliar/flashinfer/blob/main/csrc/tinygemm2.cu

[benchislett]

Is M <= 8 an arbitrary heuristic or is this the recommended choice? Is this backed by benchmarks?

[mergify]

Hi @askliar, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[vadiklyutiy]

What do you think about bypassing all bf16 gemms to FalshInfer?
We do it for fp4 and fp8.

FlashInfer has autotuning and can choose the best kernel including this tinygemm.

Also I found early some perf gaps for medium size batches (from 257-271) where flashinfer can provide improvement. See #35467 and #27173

(But need 2 small modifications of flashinfer flashinfer-ai/flashinfer#2958 and flashinfer-ai/flashinfer#2914 - both ready to be merged).

[Himan-D]

Thanks for the review feedback!

A few updates on the current implementation:

1. Delayed import - Already handled via which uses and is called lazily on first use, not at module import time.

2. Architecture validation - Currently only enables on SM90+ (Hopper) as specified in the PR. We check which is the correct threshold.

3. Bias handling - You're right that the FlashInfer kernel appears to require bias. We've kept the current approach of handling None bias in the predicate check to skip tinygemm when bias is None. As @askliar mentioned, we can follow up with FlashInfer about making bias optional.

4. Multi-dimensional inputs - Fixed in commit - we now use to correctly count total tokens across all dimensions.

5. Microbenchmark - Added benchmark script in commit at

Would you like us to make any additional changes?

[Himan-D]

Thanks for the review! Here are the updates:

1. Delayed import - Already handled via lazy check with functools.cache
2. Architecture - Only enables on SM90+ (Hopper)
3. Bias handling - We skip tinygemm when bias is None (predicate check)
4. Multi-dimensional inputs - Fixed with num_tokens calculation
5. Microbenchmark - Added benchmark script

Let me know if you have more feedback!

[Himan-D]

Great question!

Yes, the FlashInfer tinygemm_bf16 kernel supports PDL (Programmatic Dependent Launch) via the parameter.

Looking at the FlashInfer source (flashinfer/gemm/routergemm.py):

When , it uses to overlap DMA with compute from preceding kernels - ideal for low-latency decode scenarios.

Currently our implementation doesn't enable PDL - we could add this as an option. Would you like us to:

1. Add PDL support to the tinygemm path?
2. Make it configurable via environment variable?

This would be particularly useful for PD (Prefill-Decode) disaggregation where we want to overlap the GEMM with previous kernels.

[Himan-D]

Hey @mgoin - to follow up on your question about PDL:

Yes the FlashInfer tinygemm_bf16 supports PDL (Programmatic Dependent Launch) via the use_pdl parameter. Currently we have it disabled (use_pdl=False).

For PD workloads, we could add PDL support - would you like us to:

1. Add use_pdl parameter to default_unquantized_gemm?
2. Make it configurable via env var like VLLM_TINYGEMM_USE_PDL?

This would help overlap GEMM compute with previous kernel DMA for lower latency.

[Himan-D]

Good question! Yes the FlashInfer tinygemm supports PDL via use_pdl parameter. Currently use_pdl=False. For PD workloads we could add PDL support - want us to add it as configurable option?

[Himan-D]

Following up on PDL question - yes FlashInfer tinygemm supports PDL (use_pdl param). Currently use_pdl=False. Could add as configurable option for PD workloads.

[benchislett]

We should use PDL.

[ProExpertProg]

I think if we can just call flashinfer.gemm generally and rely on FI-builtin autotuning to select the fastest one, that would be the best approach?