[kernel] Integrate FlashInfer BF16 CuTeDSL Low Latency GEMM (#50572)
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
Jiahan Chang (Cyrus) committed
Aug 28, 2026 at 06:29 UTC
7c877062acab618beb377ed6aae959faccc727c0
7 files changed
+349
-23
tests/kernels/test_flashinfer_bf16_gemm.py
new
+49
@@ -0,0 +1,49 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+import torch.nn.functional as F
7
+
8
+from vllm.model_executor.layers import utils as layer_utils
9
+from vllm.utils.flashinfer import is_flashinfer_cutedsl_bf16_gemm_supported
10
+
11
+
12
+@pytest.fixture(scope="module", autouse=True)
13
+def require_flashinfer_bf16_cutedsl() -> None:
14
+ if not torch.accelerator.is_available():
15
+ pytest.skip("CUDA is required")
16
+ if not is_flashinfer_cutedsl_bf16_gemm_supported():
17
+ pytest.skip("FlashInfer BF16 cute-dsl backend is unavailable")
18
+
19
+
20
+@pytest.mark.parametrize("n,k", [(1024, 1024), (2048, 3072)])
21
+@pytest.mark.parametrize(
22
+ "m,use_bias,pdl",
23
+ [
24
+ (1, False, False),
25
+ (8, True, False),
26
+ (16, False, True),
27
+ (32, True, True),
28
+ ],
29
+)
30
+def test_flashinfer_bf16_cutedsl_correctness(
31
+ m: int,
32
+ n: int,
33
+ k: int,
34
+ use_bias: bool,
35
+ pdl: bool,
36
+) -> None:
37
+ torch.manual_seed(0)
38
+ x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 0.1
39
+ weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * 0.1
40
+ bias = (
41
+ torch.randn(n, device="cuda", dtype=torch.bfloat16) * 0.1 if use_bias else None
42
+ )
43
+
44
+ actual = layer_utils.cuda_flashinfer_bf16_gemm_impl(
45
+ x, weight, bias, pdl, "flashinfer_cutedsl"
46
+ )
47
+ expected = F.linear(x, weight, bias)
48
+
49
+ torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-1)
vllm/config/kernel.py
+6
-2
@@ -267,12 +267,16 @@ class KernelConfig:
267
"""
268
269
linear_backend: LinearBackend = "auto"
270
- """Backend for quantized linear layer GEMM kernels. Available options:
270
+ """Backend for linear layer GEMM kernels. Available options:
271
+
272
+ Layer types without an implementation from the requested backend use
273
+ automatic selection.
274
275
- "auto": Automatically select the best backend based on model and hardware
276
- "cutlass": Use CUTLASS-based kernels
277
- "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
275
- - "flashinfer_cutedsl": Use FlashInfer with CuTe-DSL kernels (NVFP4, MXFP8)
278
+ - "flashinfer_cutedsl": Use FlashInfer with CuTe-DSL kernels
279
+ (BF16, NVFP4, MXFP8)
280
- "flashinfer_trtllm": Use FlashInfer with TensorRT-LLM kernels
281
- "flashinfer_cudnn": Use FlashInfer with cuDNN kernels
282
- "flashinfer_b12x": Use FlashInfer b12x CuteDSL NVFP4 GEMM (SM120+)
vllm/model_executor/kernels/linear/__init__.py
+11
-6
@@ -237,10 +237,7 @@ def _get_linear_backend() -> str:
237
return "auto"
238
239
240
-# Mapping from linear_backend name to the set of kernel classes it covers.
241
-# When a user sets --linear-backend <name>, only kernels in the corresponding
242
-# set are considered candidates. If none can implement the layer config,
243
-# an error is raised to respect the user's explicit intent.
240
+# Kernel classes covered by each --linear-backend value.
241
_LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
242
"b12x": {
243
B12xFp8BlockScaledMMKernel,
@@ -345,9 +342,17 @@ def _filter_kernels_by_backend(
342
backend: str,
343
kernels: list[type],
344
) -> list[type]:
348
- """Filter a kernel priority list to only those matching the backend."""
345
+ """Prefer matching kernels, falling back to auto by layer type."""
346
backend_kernels = _LINEAR_BACKEND_KERNEL_MAP.get(backend, set())
350
- return [k for k in kernels if k in backend_kernels]
347
+ filtered = [kernel for kernel in kernels if kernel in backend_kernels]
348
+ if not filtered:
349
+ logger.warning_once(
350
+ "--linear-backend=%s has no kernel for this linear layer type; "
351
+ "using automatic selection for that layer type.",
352
+ backend,
353
+ )
354
+ return kernels
355
+ return filtered
356
357
358
def _resolve_backend_kernels(
vllm/model_executor/layers/linear.py
+9
-2
@@ -10,7 +10,7 @@ from torch.nn.parameter import Parameter
10
from typing_extensions import TypeIs
11
12
import vllm.envs as envs
13
-from vllm.config import get_current_vllm_config
13
+from vllm.config import get_current_vllm_config, get_current_vllm_config_or_none
14
from vllm.distributed import (
15
divide,
16
get_tensor_model_parallel_rank,
@@ -166,6 +166,13 @@ class LinearMethodBase(QuantizeMethodBase):
166
class UnquantizedLinearMethod(LinearMethodBase):
167
"""Linear method without quantization."""
168
169
+ def __init__(self) -> None:
170
+ config = get_current_vllm_config_or_none()
171
+ linear_backend = (
172
+ config.kernel_config.linear_backend if config is not None else "auto"
173
+ )
174
+ self._gemm_impl = dispatch_unquantized_gemm(linear_backend)
175
+
176
def create_weights(
177
self,
178
layer: torch.nn.Module,
@@ -215,7 +222,7 @@ class UnquantizedLinearMethod(LinearMethodBase):
222
) -> torch.Tensor:
223
if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike():
224
return linear_batch_invariant(x, layer.weight, bias)
218
- return dispatch_unquantized_gemm()(layer, x, layer.weight, bias)
225
+ return self._gemm_impl(layer, x, layer.weight, bias)
226
227
228
class LinearBase(PluggableLayer):
vllm/model_executor/layers/utils.py
+172
-2
@@ -2,7 +2,9 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Utility methods for model layers."""
4
5
+import functools
6
from collections.abc import Callable
7
+from dataclasses import dataclass
8
9
import torch
10
@@ -11,6 +13,10 @@ from vllm import envs
13
from vllm._aiter_ops import rocm_aiter_ops
14
from vllm.logger import init_logger
15
from vllm.platforms import CpuArchEnum, current_platform
16
+from vllm.utils.flashinfer import (
17
+ flashinfer_bf16_mm,
18
+ is_flashinfer_cutedsl_bf16_gemm_supported,
19
+)
20
from vllm.utils.platform_utils import num_compute_units
21
from vllm.utils.torch_utils import direct_register_custom_op
22
@@ -92,6 +98,144 @@ def default_unquantized_gemm(
98
return torch.nn.functional.linear(x, weight, bias)
99
100
101
+_FlashInferBf16RuntimeCheck = Callable[
102
+ [torch.Tensor, torch.Tensor, torch.Tensor | None], bool
103
+]
104
+
105
+
106
+@dataclass(frozen=True)
107
+class _FlashInferBf16Backend:
108
+ flashinfer_backend: str
109
+ is_supported: Callable[[], bool]
110
+ can_implement: _FlashInferBf16RuntimeCheck
111
+
112
+
113
+def _can_use_flashinfer_cutedsl_bf16(
114
+ x: torch.Tensor,
115
+ weight: torch.Tensor,
116
+ bias: torch.Tensor | None,
117
+) -> bool:
118
+ if not (
119
+ current_platform.is_cuda() and current_platform.is_device_capability_family(100)
120
+ ):
121
+ return False
122
+ if x.ndim < 1 or weight.ndim != 2:
123
+ return False
124
+ if (
125
+ not x.is_cuda
126
+ or not weight.is_cuda
127
+ or x.device != weight.device
128
+ or x.dtype != torch.bfloat16
129
+ or weight.dtype != torch.bfloat16
130
+ or not x.is_contiguous()
131
+ or not weight.is_contiguous()
132
+ ):
133
+ return False
134
+
135
+ k = x.shape[-1]
136
+ n = weight.shape[0]
137
+ if (
138
+ k <= 0
139
+ or n <= 0
140
+ or weight.shape[1] != k
141
+ or k % 128 != 0
142
+ or x.data_ptr() % 32 != 0
143
+ or weight.data_ptr() % 32 != 0
144
+ ):
145
+ return False
146
+
147
+ m = x.numel() // k
148
+ if not 1 <= m <= 32:
149
+ return False
150
+ return bias is None or (
151
+ bias.is_cuda
152
+ and bias.device == x.device
153
+ and bias.dtype == torch.bfloat16
154
+ and bias.ndim == 1
155
+ and bias.shape[0] == n
156
+ and bias.is_contiguous()
157
+ )
158
+
159
+
160
+_FLASHINFER_BF16_BACKENDS = {
161
+ "flashinfer_cutedsl": _FlashInferBf16Backend(
162
+ flashinfer_backend="cute-dsl",
163
+ is_supported=is_flashinfer_cutedsl_bf16_gemm_supported,
164
+ can_implement=_can_use_flashinfer_cutedsl_bf16,
165
+ ),
166
+}
167
+
168
+
169
+def _get_flashinfer_bf16_backend(vllm_backend: str) -> _FlashInferBf16Backend:
170
+ backend_spec = _FLASHINFER_BF16_BACKENDS.get(vllm_backend)
171
+ if backend_spec is None:
172
+ supported = ", ".join(sorted(_FLASHINFER_BF16_BACKENDS))
173
+ raise ValueError(
174
+ f"Unsupported vLLM FlashInfer BF16 backend {vllm_backend!r}; "
175
+ f"supported backends: {supported}"
176
+ )
177
+ return backend_spec
178
+
179
+
180
+def cuda_flashinfer_bf16_gemm_impl(
181
+ x: torch.Tensor,
182
+ weight: torch.Tensor,
183
+ bias: torch.Tensor | None,
184
+ pdl: bool,
185
+ vllm_backend: str,
186
+) -> torch.Tensor:
187
+ backend_spec = _get_flashinfer_bf16_backend(vllm_backend)
188
+ if not backend_spec.can_implement(x, weight, bias):
189
+ return torch.nn.functional.linear(x, weight, bias)
190
+
191
+ k = x.shape[-1]
192
+ n = weight.shape[0]
193
+ x_2d = x.view(-1, k)
194
+ out_2d = flashinfer_bf16_mm(
195
+ x_2d,
196
+ weight.t(),
197
+ bias,
198
+ pdl,
199
+ backend_spec.flashinfer_backend,
200
+ )
201
+ return out_2d.view(*x.shape[:-1], n)
202
+
203
+
204
+def cuda_flashinfer_bf16_gemm_fake(
205
+ x: torch.Tensor,
206
+ weight: torch.Tensor,
207
+ bias: torch.Tensor | None,
208
+ pdl: bool,
209
+ vllm_backend: str,
210
+) -> torch.Tensor:
211
+ return x.new_empty((*x.shape[:-1], weight.shape[0]))
212
+
213
+
214
+def cuda_flashinfer_bf16_gemm(
215
+ layer: torch.nn.Module,
216
+ x: torch.Tensor,
217
+ weight: torch.Tensor,
218
+ bias: torch.Tensor | None = None,
219
+ *,
220
+ vllm_backend: str,
221
+ pdl: bool,
222
+) -> torch.Tensor:
223
+ return torch.ops.vllm.cuda_flashinfer_bf16_gemm(
224
+ x,
225
+ weight,
226
+ bias,
227
+ pdl,
228
+ vllm_backend,
229
+ )
230
+
231
+
232
+direct_register_custom_op(
233
+ op_name="cuda_flashinfer_bf16_gemm",
234
+ op_func=cuda_flashinfer_bf16_gemm_impl,
235
+ fake_impl=cuda_flashinfer_bf16_gemm_fake,
236
+)
237
+
238
+
239
def use_aiter_triton_gemm(n, m, k, dtype):
240
if (
241
not rocm_aiter_ops.is_triton_gemm_enabled()
@@ -391,10 +535,36 @@ def cpu_unquantized_gemm(
535
return layer.cpu_linear(x, weight, bias)
536
537
394
-def dispatch_unquantized_gemm() -> Callable[..., torch.Tensor]:
538
+def dispatch_unquantized_gemm(
539
+ linear_backend: str = "auto",
540
+) -> Callable[..., torch.Tensor]:
541
if current_platform.is_rocm():
542
return rocm_unquantized_gemm
543
elif current_platform.is_cpu():
544
return cpu_unquantized_gemm
399
- else:
545
+ elif not current_platform.is_cuda():
546
return default_unquantized_gemm
547
+
548
+ backend_spec = _FLASHINFER_BF16_BACKENDS.get(linear_backend)
549
+ if backend_spec is None:
550
+ return default_unquantized_gemm
551
+
552
+ if not backend_spec.is_supported():
553
+ logger.warning_once(
554
+ "--linear-backend=%s requested FlashInfer mm_bf16 backend %r, "
555
+ "but it is unavailable on the current hardware or environment; "
556
+ "using automatic selection for unquantized linear layers.",
557
+ linear_backend,
558
+ backend_spec.flashinfer_backend,
559
+ )
560
+ return default_unquantized_gemm
561
+
562
+ logger.info_once(
563
+ "Using FlashInfer %s for eligible unquantized BF16 GEMMs.",
564
+ backend_spec.flashinfer_backend,
565
+ )
566
+ return functools.partial(
567
+ cuda_flashinfer_bf16_gemm,
568
+ vllm_backend=linear_backend,
569
+ pdl=current_platform.is_arch_support_pdl(),
570
+ )
vllm/model_executor/warmup/kernel_warmup.py
+26
-11
@@ -247,6 +247,31 @@ def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None:
247
return None
248
249
250
+_FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS = 32
251
+
252
+
253
+def _flashinfer_autotune_token_counts(runner: "GPUModelRunner") -> tuple[int, ...]:
254
+ max_tokens = runner.scheduler_config.max_num_batched_tokens
255
+ linear_backend = runner.vllm_config.kernel_config.linear_backend
256
+ if (
257
+ linear_backend == "flashinfer_cutedsl"
258
+ and max_tokens > _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS
259
+ ):
260
+ return max_tokens, _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS
261
+ return (max_tokens,)
262
+
263
+
264
+def _run_flashinfer_autotune_dummy_runs(runner: "GPUModelRunner") -> None:
265
+ for num_tokens in _flashinfer_autotune_token_counts(runner):
266
+ logger.info("Running FlashInfer autotune with %d tokens.", num_tokens)
267
+ runner._dummy_run(
268
+ num_tokens=num_tokens,
269
+ skip_eplb=True,
270
+ is_profile=True,
271
+ randomize_inputs=True,
272
+ )
273
+
274
+
275
def flashinfer_autotune(runner: "GPUModelRunner") -> None:
276
"""
277
Autotune FlashInfer operations.
@@ -284,20 +309,10 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
309
logger.info_once("Using FlashInfer autotune cache file: %s", cache_path)
310
311
# We skip EPLB here since we don't want to record dummy metrics.
287
- # When autotuning with number of tokens m, flashinfer will autotune
288
- # operations for all number of tokens up to m, so we only need to
289
- # run with the max number of tokens.
312
# Randomize inputs to avoid every token pick the same experts,
313
# which lead to some EP ranks receiving no tokens and skipping their
314
# MoE kernel entirely, and cause hang due to all-reduce collective
315
# during synchronized autotuning.
294
- dummy_run_kwargs = dict(
295
- num_tokens=runner.scheduler_config.max_num_batched_tokens,
296
- skip_eplb=True,
297
- is_profile=True,
298
- randomize_inputs=True,
299
- )
300
-
316
# Read cached autotune results and broadcast to all ranks.
317
cached_results: bytes | None = None
318
if is_leader and cache_path.exists():
@@ -316,7 +331,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
331
torch.inference_mode(),
332
fi_utils.autotune(tune_mode=True, **autotune_kwargs),
333
):
319
- runner._dummy_run(**dummy_run_kwargs)
334
+ _run_flashinfer_autotune_dummy_runs(runner)
335
finally:
336
set_autotune_process_group(None)
337
vllm/utils/flashinfer.py
+76
@@ -24,6 +24,26 @@ from vllm.utils.math_utils import cdiv
24
25
logger = init_logger(__name__)
26
27
+
28
+def flashinfer_bf16_mm(
29
+ a: torch.Tensor,
30
+ b: torch.Tensor,
31
+ bias: torch.Tensor | None,
32
+ pdl: bool,
33
+ backend: str,
34
+) -> torch.Tensor:
35
+ from flashinfer import mm_bf16
36
+
37
+ return mm_bf16(
38
+ a,
39
+ b,
40
+ bias=bias,
41
+ pdl=pdl,
42
+ out_dtype=torch.bfloat16,
43
+ backend=backend,
44
+ )
45
+
46
+
47
# This is the storage path for the cubins, it can be replaced
48
# with a local path for testing.
49
# Referenced from https://github.com/flashinfer-ai/flashinfer/blob/0c9a92c3d9a7e043ab6f3f7b2273269caf6ab044/flashinfer/jit/cubin_loader.py#L35 # noqa: E501
@@ -63,6 +83,58 @@ def has_flashinfer() -> bool:
83
return True
84
85
86
+@functools.cache
87
+def has_flashinfer_bf16_gemm() -> bool:
88
+ """Return whether FlashInfer exposes the BF16 dense GEMM API."""
89
+ if not has_flashinfer():
90
+ return False
91
+ mod = _get_submodule("flashinfer")
92
+ return mod is not None and callable(getattr(mod, "mm_bf16", None))
93
+
94
+
95
+@functools.cache
96
+def is_flashinfer_bf16_gemm_supported(
97
+ backend: str,
98
+ compute_capability: int | None = None,
99
+) -> bool:
100
+ """Return whether an exact FlashInfer BF16 backend is available."""
101
+ if not current_platform.is_cuda() or not has_flashinfer_bf16_gemm():
102
+ return False
103
+
104
+ mod = _get_submodule("flashinfer")
105
+ mm_bf16 = getattr(mod, "mm_bf16", None) if mod is not None else None
106
+ backend_supported = getattr(mm_bf16, "is_backend_supported", None)
107
+ if not callable(backend_supported):
108
+ return False
109
+
110
+ if compute_capability is None:
111
+ device_capability = current_platform.get_device_capability()
112
+ if device_capability is None:
113
+ return False
114
+ compute_capability = device_capability.to_int()
115
+
116
+ try:
117
+ return bool(backend_supported(backend, compute_capability))
118
+ except (AttributeError, ImportError, RuntimeError, TypeError, ValueError):
119
+ return False
120
+
121
+
122
+@functools.cache
123
+def is_flashinfer_cutedsl_bf16_gemm_supported() -> bool:
124
+ """Return whether the CuTeDSL BF16 dense GEMM backend is available."""
125
+ if not is_flashinfer_bf16_gemm_supported("cute-dsl"):
126
+ return False
127
+ try:
128
+ from flashinfer.cute_dsl.utils import is_cute_dsl_available
129
+ from flashinfer.utils import is_sm100a_supported
130
+ except (ImportError, ModuleNotFoundError):
131
+ return False
132
+ try:
133
+ return is_cute_dsl_available() and is_sm100a_supported(torch.device("cuda"))
134
+ except (RuntimeError, TypeError, ValueError):
135
+ return False
136
+
137
+
138
def _missing(*_: Any, **__: Any) -> NoReturn:
139
"""Placeholder for unavailable FlashInfer backend."""
140
raise RuntimeError(
@@ -1079,6 +1151,10 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool:
1151
1152
__all__ = [
1153
"has_flashinfer",
1154
+ "flashinfer_bf16_mm",
1155
+ "has_flashinfer_bf16_gemm",
1156
+ "is_flashinfer_bf16_gemm_supported",
1157
+ "is_flashinfer_cutedsl_bf16_gemm_supported",
1158
"flashinfer_trtllm_fp8_block_scale_moe",
1159
"flashinfer_cutlass_fused_moe",
1160
"flashinfer_cutedsl_grouped_gemm_nt_masked",