@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0762f2afeb

[Perf][Feat] Add generic cuteDSL LL BF16 router (GEMM) (#42562)

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>

Roberto L. Castro committed Jul 14, 2026 at 08:01 UTC 0762f2afeb74f790e8c5ebe2b95a012eb38d499e
8 files changed +1812 -8
.buildkite/test_areas/kernels.yaml
+5
@@ -232,6 +232,10 @@ steps:
232 - vllm/v1/attention/backends/mla/flashinfer_mla.py
233 - vllm/v1/attention/selector.py
234 - vllm/platforms/cuda.py
235 + - vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py
236 + - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py
237 + - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py
238 + - tests/kernels/test_ll_bf16_gemm.py
239 - tests/kernels/test_top_k_per_row.py
240 commands:
241 - nvidia-smi
@@ -260,6 +264,7 @@ steps:
264 - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
265 - pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py
266 - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
267 + - pytest -v -s tests/kernels/test_ll_bf16_gemm.py
268 # e2e
269 - pytest -v -s tests/models/quantization/test_nvfp4.py
270
tests/kernels/test_ll_bf16_gemm.py new
+553
@@ -0,0 +1,553 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for cuteDSL low-latency router GEMM (dot-product + split-K)."""
4 +
5 +import pytest
6 +import torch
7 +import torch.nn.functional as F
8 +
9 +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
10 +
11 +
12 +@pytest.fixture(autouse=True, scope="module")
13 +def _require_sm90_and_cutedsl():
14 + if torch.cuda.get_device_capability()[0] < 9:
15 + pytest.skip("Requires SM90+ (Hopper/Blackwell)")
16 + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import (
17 + is_available,
18 + )
19 +
20 + if not is_available():
21 + pytest.skip("cuteDSL (CUTLASS Python) not installed")
22 +
23 +
24 +# ===== Helpers =====
25 +
26 +
27 +def _ref(a, b):
28 + return torch.mm(a.float(), b.float().T)
29 +
30 +
31 +def _assert_close(out, ref, *, min_cos_sim=0.99, context=""):
32 + assert out.device.type == "cuda", f"{context}: not on CUDA"
33 + assert torch.isfinite(out).all(), f"{context}: NaN/Inf"
34 + cos = F.cosine_similarity(
35 + out.reshape(-1).float(), ref.reshape(-1).float(), dim=0
36 + ).item()
37 + assert cos > min_cos_sim, (
38 + f"{context}: cos_sim {cos:.6f} < {min_cos_sim} "
39 + f"(abs_err={(out.float() - ref.float()).abs().max().item():.2e})"
40 + )
41 +
42 +
43 +def _can_precompile(a, b):
44 + return (
45 + a.dim() == 2
46 + and b.dim() == 2
47 + and a.dtype == torch.bfloat16
48 + and b.dtype == torch.bfloat16
49 + and a.device.type == "cuda"
50 + and b.device.type == "cuda"
51 + and a.device == b.device
52 + and a.shape[1] == b.shape[1]
53 + and a.is_contiguous()
54 + and b.is_contiguous()
55 + )
56 +
57 +
58 +def _gemm(a, b):
59 + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import (
60 + ll_bf16_gemm,
61 + ll_bf16_gemm_kernel,
62 + )
63 +
64 + if _can_precompile(a, b):
65 + compile_key = ll_bf16_gemm_kernel.dispatch(
66 + M=a.shape[0], K=a.shape[1], N=b.shape[0]
67 + )
68 + ll_bf16_gemm_kernel.compile(compile_key)
69 + return ll_bf16_gemm(a, b)
70 +
71 +
72 +# ===== Shapes =====
73 +
74 +SHAPES = [
75 + (256, 7168, "DSV3"),
76 + (256, 14400, "DSV4-Flash"),
77 + (128, 5120, "DeepSeek-V2"),
78 + (8, 4096, "Mixtral-8x7B"),
79 + (64, 2880, "non-tile-aligned-K"),
80 + (256, 2048, "split-K-boundary"),
81 +]
82 +
83 +SHAPES_SPLITK = [(n, k, d) for n, k, d in SHAPES if k >= 2048]
84 +
85 +
86 +# =================================================================
87 +# Dot-product kernel (M<=4 or K<2048)
88 +# =================================================================
89 +
90 +
91 +@pytest.mark.parametrize("M", [1, 2, 3, 4])
92 +@pytest.mark.parametrize("N,K,desc", SHAPES, ids=[s[2] for s in SHAPES])
93 +def test_dotprod(M, N, K, desc):
94 + torch.manual_seed(42)
95 + a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda")
96 + b = torch.randn(N, K, dtype=torch.bfloat16, device="cuda")
97 + out = _gemm(a, b)
98 + assert out.dtype == torch.float32
99 + assert out.shape == (M, N)
100 + _assert_close(out, _ref(a, b), context=f"dotprod {M}x{N}x{K}")
101 +
102 +
103 +@pytest.mark.parametrize("M", [1, 4, 8, 16])
104 +def test_dotprod_small_K(M):
105 + """K<2048 forces dot-product regardless of M."""
106 + torch.manual_seed(42)
107 + a = torch.randn(M, 1024, dtype=torch.bfloat16, device="cuda")
108 + b = torch.randn(64, 1024, dtype=torch.bfloat16, device="cuda")
109 + _assert_close(_gemm(a, b), _ref(a, b), context=f"small_K M={M}")
110 +
111 +
112 +@pytest.mark.parametrize("K", [16, 32, 64, 128, 256, 512, 1024, 1536])
113 +def test_dotprod_K_sweep(K):
114 + torch.manual_seed(42)
115 + a = torch.randn(4, K, dtype=torch.bfloat16, device="cuda")
116 + b = torch.randn(32, K, dtype=torch.bfloat16, device="cuda")
117 + _assert_close(_gemm(a, b), _ref(a, b), context=f"K={K}")
118 +
119 +
120 +# =================================================================
121 +# Split-K kernel (M>4 and K>=2048)
122 +# =================================================================
123 +
124 +
125 +@pytest.mark.parametrize("M", [5, 6, 8, 12, 16])
126 +@pytest.mark.parametrize("N,K,desc", SHAPES_SPLITK, ids=[s[2] for s in SHAPES_SPLITK])
127 +def test_splitk(M, N, K, desc):
128 + torch.manual_seed(42)
129 + a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda")
130 + b = torch.randn(N, K, dtype=torch.bfloat16, device="cuda")
131 + out = _gemm(a, b)
132 + assert out.dtype == torch.float32
133 + assert out.shape == (M, N)
134 + _assert_close(out, _ref(a, b), context=f"splitk {M}x{N}x{K}")
135 +
136 +
137 +@pytest.mark.parametrize("K", [2048, 2304, 2880, 3072, 4096, 5120, 7168, 14400])
138 +def test_splitk_K_sweep(K):
139 + """Includes non-tile-aligned K (2880) and uneven split (2304)."""
140 + torch.manual_seed(42)
141 + a = torch.randn(8, K, dtype=torch.bfloat16, device="cuda")
142 + b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda")
143 + _assert_close(_gemm(a, b), _ref(a, b), context=f"splitk K={K}")
144 +
145 +
146 +# =================================================================
147 +# Dispatch boundary (M=4/5, K=2032/2048)
148 +# =================================================================
149 +
150 +
151 +@pytest.mark.parametrize("M", [4, 5])
152 +@pytest.mark.parametrize("K", [2032, 2048])
153 +def test_dispatch_boundary(M, K):
154 + torch.manual_seed(42)
155 + a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda")
156 + b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda")
157 + path = "splitk" if M > 4 and K >= 2048 else "dotprod"
158 + _assert_close(_gemm(a, b), _ref(a, b), context=f"M={M} K={K} ({path})")
159 +
160 +
161 +# =================================================================
162 +# Arbitrary N
163 +# =================================================================
164 +
165 +
166 +@pytest.mark.parametrize("N", [1, 3, 7, 16, 17, 64, 128, 256, 384])
167 +def test_arbitrary_N_dotprod(N):
168 + torch.manual_seed(42)
169 + a = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda")
170 + b = torch.randn(N, 2048, dtype=torch.bfloat16, device="cuda")
171 + out = _gemm(a, b)
172 + assert out.shape == (4, N)
173 + _assert_close(out, _ref(a, b), context=f"dotprod N={N}")
174 +
175 +
176 +@pytest.mark.parametrize("N", [1, 8, 16, 17, 64, 128, 256])
177 +def test_arbitrary_N_splitk(N):
178 + torch.manual_seed(42)
179 + a = torch.randn(8, 4096, dtype=torch.bfloat16, device="cuda")
180 + b = torch.randn(N, 4096, dtype=torch.bfloat16, device="cuda")
181 + out = _gemm(a, b)
182 + assert out.shape == (8, N)
183 + _assert_close(out, _ref(a, b), context=f"splitk N={N}")
184 +
185 +
186 +# =================================================================
187 +# Single token (M=1, decode path)
188 +# =================================================================
189 +
190 +
191 +@pytest.mark.parametrize(
192 + "N,K",
193 + [(256, 7168), (256, 14400), (8, 4096), (384, 7168)],
194 + ids=["DSV3", "DSV4-Flash", "Mixtral", "DSV4-Pro"],
195 +)
196 +def test_single_token(N, K):
197 + torch.manual_seed(42)
198 + a = torch.randn(1, K, dtype=torch.bfloat16, device="cuda")
199 + b = torch.randn(N, K, dtype=torch.bfloat16, device="cuda")
200 + out = _gemm(a, b)
201 + assert out.shape == (1, N)
202 + _assert_close(out, _ref(a, b), context=f"M=1 {N}x{K}")
203 +
204 +
205 +# =================================================================
206 +# Numerical robustness
207 +# =================================================================
208 +
209 +
210 +@pytest.mark.parametrize("M,K", [(4, 2048), (8, 4096)], ids=["dotprod", "splitk"])
211 +def test_large_values(M, K):
212 + torch.manual_seed(42)
213 + a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") * 100
214 + b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda") * 100
215 + out = _gemm(a, b)
216 + assert torch.isfinite(out).all()
217 + _assert_close(out, _ref(a, b), context=f"large M={M}")
218 +
219 +
220 +def test_near_zero():
221 + torch.manual_seed(42)
222 + a = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") * 1e-4
223 + b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda") * 1e-4
224 + out = _gemm(a, b)
225 + assert torch.isfinite(out).all()
226 + assert out.abs().max() < 1.0
227 +
228 +
229 +def test_zeros():
230 + a = torch.zeros(4, 2048, dtype=torch.bfloat16, device="cuda")
231 + b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda")
232 + assert (_gemm(a, b) == 0).all()
233 +
234 +
235 +def test_ones():
236 + a = torch.ones(1, 2048, dtype=torch.bfloat16, device="cuda")
237 + b = torch.randn(32, 2048, dtype=torch.bfloat16, device="cuda")
238 + _assert_close(_gemm(a, b), _ref(a, b), context="ones")
239 +
240 +
241 +# =================================================================
242 +# Output dtype
243 +# =================================================================
244 +
245 +
246 +@pytest.mark.parametrize("M,K", [(4, 2048), (8, 4096)], ids=["dotprod", "splitk"])
247 +def test_output_fp32(M, K):
248 + torch.manual_seed(42)
249 + a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda")
250 + b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda")
251 + assert _gemm(a, b).dtype == torch.float32
252 +
253 +
254 +# =================================================================
255 +# Determinism
256 +# =================================================================
257 +
258 +
259 +@pytest.mark.parametrize(
260 + "M,K", [(1, 4096), (4, 4096), (5, 4096), (8, 4096), (16, 4096)]
261 +)
262 +def test_deterministic(M, K):
263 + torch.manual_seed(42)
264 + a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda")
265 + b = torch.randn(128, K, dtype=torch.bfloat16, device="cuda")
266 + torch.testing.assert_close(_gemm(a, b), _gemm(a, b), atol=0, rtol=0)
267 +
268 +
269 +# =================================================================
270 +# Cross-kernel consistency
271 +# =================================================================
272 +
273 +
274 +@pytest.mark.parametrize("K", [2048, 4096, 7168])
275 +def test_dotprod_vs_splitk(K):
276 + """First 4 rows from split-K (M=5) match dot-product (M=4)."""
277 + torch.manual_seed(42)
278 + a = torch.randn(5, K, dtype=torch.bfloat16, device="cuda")
279 + b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda")
280 + _assert_close(
281 + _gemm(a, b)[:4],
282 + _gemm(a[:4], b),
283 + min_cos_sim=0.999,
284 + context=f"cross-kernel K={K}",
285 + )
286 +
287 +
288 +# =================================================================
289 +# CUDA graph
290 +# =================================================================
291 +
292 +
293 +@pytest.mark.parametrize("M,K", [(4, 2048), (8, 4096)], ids=["dotprod", "splitk"])
294 +def test_cudagraph(M, K):
295 + torch.manual_seed(42)
296 + a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda")
297 + b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda")
298 + _gemm(a, b)
299 + torch.accelerator.synchronize()
300 +
301 + g = torch.cuda.CUDAGraph()
302 + with torch.cuda.graph(g):
303 + out = _gemm(a, b)
304 + for _ in range(5):
305 + g.replay()
306 + torch.accelerator.synchronize()
307 + _assert_close(out, _ref(a, b), context=f"cudagraph M={M}")
308 +
309 +
310 +def test_cudagraph_20x_replay():
311 + torch.manual_seed(42)
312 + a = torch.randn(4, 4096, dtype=torch.bfloat16, device="cuda")
313 + b = torch.randn(256, 4096, dtype=torch.bfloat16, device="cuda")
314 + _gemm(a, b)
315 + torch.accelerator.synchronize()
316 +
317 + g = torch.cuda.CUDAGraph()
318 + with torch.cuda.graph(g):
319 + out = _gemm(a, b)
320 + results = []
321 + for _ in range(20):
322 + g.replay()
323 + torch.accelerator.synchronize()
324 + results.append(out.clone())
325 + for i in range(1, len(results)):
326 + torch.testing.assert_close(
327 + results[0], results[i], atol=0, rtol=0, msg=f"Replay {i} differs"
328 + )
329 +
330 +
331 +def test_cudagraph_input_update():
332 + torch.manual_seed(42)
333 + a = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda")
334 + b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda")
335 + _gemm(a, b)
336 + torch.accelerator.synchronize()
337 +
338 + g = torch.cuda.CUDAGraph()
339 + with torch.cuda.graph(g):
340 + out = _gemm(a, b)
341 + a.copy_(torch.randn_like(a))
342 + g.replay()
343 + torch.accelerator.synchronize()
344 + _assert_close(out, _ref(a, b), context="cudagraph input update")
345 +
346 +
347 +# =================================================================
348 +# GateLinear dispatch integration
349 +# =================================================================
350 +
351 +
352 +def _make_gate_linear(monkeypatch, *, params_dtype, out_dtype=torch.float32):
353 + monkeypatch.setattr(
354 + "vllm.model_executor.layers.linear.get_tensor_model_parallel_rank",
355 + lambda: 0,
356 + )
357 + monkeypatch.setattr(
358 + "vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size",
359 + lambda: 1,
360 + )
361 + monkeypatch.setattr(
362 + "vllm.model_executor.parameter.get_tensor_model_parallel_rank",
363 + lambda: 0,
364 + )
365 + monkeypatch.setattr(
366 + "vllm.model_executor.parameter.get_tensor_model_parallel_world_size",
367 + lambda: 1,
368 + )
369 +
370 + from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
371 +
372 + monkeypatch.setattr(
373 + "vllm.model_executor.kernels.linear.cute_dsl.ll_bf16.is_available",
374 + lambda: True,
375 + )
376 + return GateLinear(
377 + input_size=2048,
378 + output_size=64,
379 + bias=False,
380 + out_dtype=out_dtype,
381 + params_dtype=params_dtype,
382 + ).cuda()
383 +
384 +
385 +def test_gate_linear_uses_ll_bf16_for_bf16_fast_path(monkeypatch):
386 + gate = _make_gate_linear(monkeypatch, params_dtype=torch.bfloat16)
387 + x = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda")
388 + calls = []
389 +
390 + def fake_ll_bf16_gemm(hidden_states, router_weight):
391 + calls.append((hidden_states, router_weight))
392 + return torch.full(
393 + (hidden_states.shape[0], router_weight.shape[0]),
394 + 1.0,
395 + dtype=torch.float32,
396 + device=hidden_states.device,
397 + )
398 +
399 + monkeypatch.setattr(
400 + "vllm.model_executor.kernels.linear.cute_dsl.ll_bf16.ll_bf16_gemm",
401 + fake_ll_bf16_gemm,
402 + )
403 + out, bias = gate(x)
404 + assert bias is None
405 + assert out.shape == (4, 64)
406 + assert out.dtype == torch.float32
407 + assert len(calls) == 1
408 + assert calls[0][0] is x
409 + assert calls[0][1] is gate.weight
410 +
411 +
412 +def test_gate_linear_fp32_weight_falls_back(monkeypatch):
413 + gate = _make_gate_linear(monkeypatch, params_dtype=torch.float32)
414 + assert not gate.allow_ll_bf16_gemm
415 + x = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda")
416 +
417 + def fail_ll_bf16_gemm(hidden_states, router_weight):
418 + raise AssertionError("ll_bf16_gemm should not run for fp32 weights")
419 +
420 + monkeypatch.setattr(
421 + "vllm.model_executor.kernels.linear.cute_dsl.ll_bf16.ll_bf16_gemm",
422 + fail_ll_bf16_gemm,
423 + )
424 + out, _ = gate(x)
425 + assert out.shape == (4, 64)
426 + assert out.dtype == torch.float32
427 +
428 +
429 +def test_gate_linear_non_bf16_activation_falls_back(monkeypatch):
430 + gate = _make_gate_linear(monkeypatch, params_dtype=torch.bfloat16)
431 + x = torch.randn(4, 2048, dtype=torch.float16, device="cuda")
432 +
433 + def fail_ll_bf16_gemm(hidden_states, router_weight):
434 + raise AssertionError("ll_bf16_gemm should not run for non-bf16 activations")
435 +
436 + monkeypatch.setattr(
437 + "vllm.model_executor.kernels.linear.cute_dsl.ll_bf16.ll_bf16_gemm",
438 + fail_ll_bf16_gemm,
439 + )
440 + out, _ = gate(x)
441 + assert out.shape == (4, 64)
442 + assert out.dtype == torch.float32
443 +
444 +
445 +def test_gate_linear_set_out_dtype_enables_ll_bf16(monkeypatch):
446 + gate = _make_gate_linear(monkeypatch, params_dtype=torch.bfloat16, out_dtype=None)
447 + assert not gate.allow_ll_bf16_gemm
448 + gate.set_out_dtype(torch.float32)
449 + assert gate.allow_ll_bf16_gemm
450 +
451 +
452 +def test_gate_linear_non_fp32_out_dtype_disables_ll_bf16(monkeypatch):
453 + gate = _make_gate_linear(monkeypatch, params_dtype=torch.bfloat16, out_dtype=None)
454 + gate.set_out_dtype(torch.bfloat16)
455 + assert not gate.allow_ll_bf16_gemm
456 +
457 +
458 +def test_gate_linear_m_gt_16_falls_back(monkeypatch):
459 + gate = _make_gate_linear(monkeypatch, params_dtype=torch.bfloat16)
460 + x = torch.randn(17, 2048, dtype=torch.bfloat16, device="cuda")
461 +
462 + def fail_ll_bf16_gemm(hidden_states, router_weight):
463 + raise AssertionError("ll_bf16_gemm should not run for M > 16")
464 +
465 + monkeypatch.setattr(
466 + "vllm.model_executor.kernels.linear.cute_dsl.ll_bf16.ll_bf16_gemm",
467 + fail_ll_bf16_gemm,
468 + )
469 + out, _ = gate(x)
470 + assert out.shape == (17, 64)
471 + assert out.dtype == torch.float32
472 +
473 +
474 +# =================================================================
475 +# Negative tests — invalid inputs
476 +# =================================================================
477 +
478 +
479 +@pytest.mark.parametrize(
480 + "M,K,N,dtype",
481 + [
482 + pytest.param(4, 2048, 64, torch.float32, id="fp32_input"),
483 + pytest.param(4, 2048, 64, torch.float16, id="fp16_input"),
484 + pytest.param(8, 4096, 64, torch.float32, id="fp32_splitk_path"),
485 + ],
486 +)
487 +def test_invalid_dtype(M, K, N, dtype):
488 + a = torch.randn(M, K, device="cuda", dtype=dtype)
489 + b = torch.randn(N, K, device="cuda", dtype=dtype)
490 + with pytest.raises(ValueError, match="dtype=bfloat16"):
491 + _gemm(a, b)
492 +
493 +
494 +def test_invalid_device_cpu():
495 + a = torch.randn(4, 2048, dtype=torch.bfloat16)
496 + b = torch.randn(64, 2048, dtype=torch.bfloat16)
497 + with pytest.raises(ValueError, match="device_type=cuda"):
498 + _gemm(a, b)
499 +
500 +
501 +def test_invalid_1d_input():
502 + a = torch.randn(2048, dtype=torch.bfloat16, device="cuda")
503 + b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda")
504 + with pytest.raises(ValueError, match="2D tensors"):
505 + _gemm(a, b)
506 +
507 +
508 +def test_mismatched_K():
509 + a = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda")
510 + b = torch.randn(64, 1024, dtype=torch.bfloat16, device="cuda")
511 + with pytest.raises(ValueError, match="matching K dimensions"):
512 + _gemm(a, b)
513 +
514 +
515 +def test_invalid_K_divisibility():
516 + a = torch.randn(4, 2049, dtype=torch.bfloat16, device="cuda")
517 + b = torch.randn(64, 2049, dtype=torch.bfloat16, device="cuda")
518 +
519 + with pytest.raises(ValueError, match="K to be divisible by 8"):
520 + _gemm(a, b)
521 +
522 +
523 +def test_non_contiguous_input():
524 + a = torch.randn(2048, 4, dtype=torch.bfloat16, device="cuda").T
525 + b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda")
526 + assert not a.is_contiguous()
527 + with pytest.raises(ValueError, match="contiguous row-major"):
528 + _gemm(a, b)
529 +
530 +
531 +def test_invalid_output_dtype():
532 + a = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda")
533 + b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda")
534 + with pytest.raises(ValueError, match="output_dtype=torch.float32"):
535 + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm
536 +
537 + ll_bf16_gemm(a, b, output_dtype=torch.bfloat16)
538 +
539 +
540 +def test_cache_miss_compiles_dotprod():
541 + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import LLBf16Gemm
542 +
543 + torch.manual_seed(42)
544 + a = torch.randn(3, 64, dtype=torch.bfloat16, device="cuda")
545 + b = torch.randn(17, 64, dtype=torch.bfloat16, device="cuda")
546 + kernel = LLBf16Gemm()
547 + out = kernel(a, b)
548 + assert out.shape == (3, 17)
549 + _assert_close(out, _ref(a, b), context="cache miss dotprod")
550 +
551 +
552 +if __name__ == "__main__":
553 + pytest.main([__file__, "-v"])
vllm/model_executor/kernels/linear/cute_dsl/__init__.py
vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py new
+313
@@ -0,0 +1,313 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import cutlass
5 +import cutlass.cute as cute
6 +from cuda.bindings.driver import CUstream
7 +from cutlass import const_expr
8 +
9 +
10 +class LLBf16Dotprod:
11 + """BF16 router GEMM kernel based on CTA-local dot products.
12 +
13 + This kernel computes C[M, N] = A[M, K] @ B[N, K]^T for bf16 inputs
14 + and fp32 output. It launches one CTA per output column, distributes K
15 + across CTA threads with vectorized loads, accumulates one fp32 dot product
16 + per token, and reduces through warp shuffles plus shared memory.
17 +
18 + :param k: Compile-time K dimension specialized into the generated kernel.
19 + :type k: int
20 + :param bs: Threads per CTA and K-stripe width used by the reduction.
21 + :type bs: int
22 + :param main_vec_width: bf16 elements loaded per thread in the main loop.
23 + :type main_vec_width: int
24 + :param tail_vec_width: bf16 elements loaded per thread in the vector tail.
25 + :type tail_vec_width: int
26 + :param use_pdl: Whether to launch with Programmatic Dependent Launch.
27 + :type use_pdl: bool
28 +
29 + :note: Supported A/B data types:
30 + - BFloat16/BFloat16
31 + :note: Supported accumulator data types:
32 + - Float32
33 + :note: Supported C data types:
34 + - Float32
35 + :note: Constraints:
36 + - K must preserve 16-byte row alignment for contiguous bf16 inputs.
37 +
38 + :note: K is handled as vectorized main/tail loops plus scalar remainder.
39 +
40 + :compile-key: ``(M, K, bs)`` selects the token count, hidden size,
41 + and CTA thread/K-stripe width specialization.
42 + """
43 +
44 + def __init__(
45 + self,
46 + k: int,
47 + bs: int = 128,
48 + main_vec_width: int = 8,
49 + tail_vec_width: int = 4,
50 + use_pdl: bool = False,
51 + ):
52 + """Initialize the dot-product kernel configuration.
53 +
54 + This configuration fixes the CTA thread count, reduction warp count,
55 + bf16 vector widths, and K-loop decomposition used by the generated
56 + kernel.
57 +
58 + :param k: Hidden size K used to specialize the K-loop decomposition.
59 + :type k: int
60 + :param bs: Threads per CTA and K-stripe width for the reduction.
61 + :type bs: int
62 + :param main_vec_width: BF16 elements loaded per thread in the main loop.
63 + :type main_vec_width: int
64 + :param tail_vec_width: BF16 elements loaded per thread in the vector tail.
65 + :type tail_vec_width: int
66 + :param use_pdl: Whether to launch with Programmatic Dependent Launch.
67 + :type use_pdl: bool
68 + """
69 + self.bs = bs
70 + self.main_vec_width = main_vec_width
71 + self.tail_vec_width = tail_vec_width
72 + self.use_pdl = use_pdl
73 + self.num_warps = bs // cute.arch.WARP_SIZE
74 + self._init_k_tiles(k)
75 +
76 + def _vectorized_elems(self, k_extent: int, vec_width: int) -> int:
77 + vector_tile = vec_width * self.bs
78 + return (k_extent // vector_tile) * vector_tile
79 +
80 + def _init_k_tiles(self, k: int) -> None:
81 + """Split K into vector loops, scalar rounds, and ragged tail."""
82 + self.k_main_elems = self._vectorized_elems(k, self.main_vec_width)
83 + self.k_after_main = k - self.k_main_elems
84 + self.k_tail_elems = self._vectorized_elems(
85 + self.k_after_main, self.tail_vec_width
86 + )
87 + self.k_done_all = self.k_main_elems + self.k_tail_elems
88 + self.scalar_rem = k - self.k_done_all
89 + self.ks_full = self.scalar_rem // self.bs
90 + self.ks_part = self.scalar_rem % self.bs
91 + self.k_scalar_full = self.ks_full * self.bs
92 + self.k_part_offset = self.k_done_all + self.k_scalar_full
93 + self.main_tiles = self.k_main_elems // (self.main_vec_width * self.bs)
94 + self.tail_tiles = self.k_tail_elems // (self.tail_vec_width * self.bs)
95 +
96 + @cute.jit
97 + def _vector_dotprod(
98 + self,
99 + acc: cute.Tensor,
100 + tA: cute.Tensor,
101 + tB: cute.Tensor,
102 + M: cutlass.Constexpr,
103 + num_tiles: cutlass.Constexpr,
104 + align_bytes: cutlass.Constexpr,
105 + ):
106 + for tile in cutlass.range_constexpr(num_tiles):
107 + bt = tB[None, tile]
108 + br = cute.make_rmem_tensor_like(bt)
109 + cute.autovec_copy(bt, br)
110 + br_f32 = br.load().to(cutlass.Float32)
111 +
112 + for m in cutlass.range_constexpr(M):
113 + at = tA[m, None, tile]
114 + ar = cute.make_rmem_tensor_like(at)
115 + cute.autovec_copy(at, ar)
116 + vec_width: cutlass.Constexpr = cute.size(ar)
117 + for v in cutlass.range_constexpr(vec_width):
118 + acc[m] = acc[m] + ar[v].to(cutlass.Float32) * br_f32[v]
119 +
120 + def _make_thread_vector_slice(
121 + self,
122 + gA_vec: cute.Tensor,
123 + gB_vec: cute.Tensor,
124 + tidx: cutlass.Int32,
125 + n_idx: cutlass.Int32,
126 + bs: cutlass.Constexpr,
127 + ):
128 + # (M/N, K_TILE, K_LANE, K_VEC); tidx selects K_LANE.
129 + tA = cute.logical_divide(gA_vec, (None, (None, bs)))
130 + tB = cute.logical_divide(gB_vec, (None, (None, bs)))
131 + return tA[None, (None, (tidx, None))], tB[n_idx, (None, (tidx, None))]
132 +
133 + def _make_k_slice(
134 + self,
135 + gX: cute.Tensor,
136 + k_offset: cutlass.Constexpr,
137 + k_extent: cutlass.Constexpr,
138 + ):
139 + k_layout_extent: cutlass.Constexpr = (
140 + 1 if const_expr(k_extent == 0) else k_extent
141 + )
142 +
143 + if const_expr(k_offset == 0):
144 + return cute.local_tile(
145 + gX, (cute.size(gX, mode=[0]), k_layout_extent), (0, 0)
146 + )
147 + return cute.local_tile(
148 + cute.domain_offset((0, k_offset), gX),
149 + (cute.size(gX, mode=[0]), k_layout_extent),
150 + (0, 0),
151 + )
152 +
153 + @cute.jit
154 + def __call__(
155 + self,
156 + gA: cute.Tensor,
157 + gB: cute.Tensor,
158 + gC: cute.Tensor,
159 + M: cutlass.Constexpr,
160 + K_dim: cutlass.Constexpr,
161 + N_dim: cutlass.Int32,
162 + stream: CUstream,
163 + ):
164 + """Execute the dot-product GEMM operation in steps:
165 + - Launch one CTA per output column ``n`` with ``bs`` threads.
166 + - Keep one FP32 accumulator per token ``m`` in each thread.
167 + - Traverse K with vectorized 128-bit, vectorized 64-bit, scalar, and
168 + ragged-tail loops from the precomputed K decomposition.
169 + - Reduce each token accumulator first within the warp, then across
170 + warps through shared memory, and store ``C[:, n]``.
171 +
172 + :param gA: Input tensor A with shape ``[M, K]``.
173 + :type gA: cute.Tensor
174 + :param gB: Input tensor B with shape ``[N, K]``.
175 + :type gB: cute.Tensor
176 + :param gC: Output tensor C with shape ``[M, N]``.
177 + :type gC: cute.Tensor
178 + :param M: Token count selected by the compile key.
179 + :type M: cutlass.Constexpr
180 + :param K_dim: Hidden size selected by the compile key.
181 + :type K_dim: cutlass.Constexpr
182 + :param N_dim: Output column count used for the launch grid.
183 + :type N_dim: cutlass.Int32
184 + :param stream: CUDA stream for asynchronous execution.
185 + :type stream: CUstream
186 + """
187 + self.kernel(
188 + gA,
189 + gB,
190 + gC,
191 + M,
192 + self.main_vec_width,
193 + self.tail_vec_width,
194 + self.bs,
195 + self.num_warps,
196 + self.k_main_elems,
197 + self.k_tail_elems,
198 + self.k_done_all,
199 + self.ks_full,
200 + self.ks_part,
201 + self.k_scalar_full,
202 + self.k_part_offset,
203 + self.main_tiles,
204 + self.tail_tiles,
205 + ).launch(
206 + grid=[N_dim, 1, 1],
207 + block=[self.bs, 1, 1],
208 + smem=M * 4 * self.num_warps,
209 + stream=stream,
210 + use_pdl=self.use_pdl,
211 + min_blocks_per_mp=1,
212 + )
213 +
214 + @cute.kernel
215 + def kernel(
216 + self,
217 + gA: cute.Tensor,
218 + gB: cute.Tensor,
219 + gC: cute.Tensor,
220 + M: cutlass.Constexpr,
221 + main_vec_width: cutlass.Constexpr,
222 + tail_vec_width: cutlass.Constexpr,
223 + bs: cutlass.Constexpr,
224 + num_warps: cutlass.Constexpr,
225 + k_main_elems: cutlass.Constexpr,
226 + k_tail_elems: cutlass.Constexpr,
227 + k_done_all: cutlass.Constexpr,
228 + ks_full: cutlass.Constexpr,
229 + ks_part: cutlass.Constexpr,
230 + k_scalar_full: cutlass.Constexpr,
231 + k_part_offset: cutlass.Constexpr,
232 + main_tiles: cutlass.Constexpr,
233 + tail_tiles: cutlass.Constexpr,
234 + ):
235 + tidx, _, _ = cute.arch.thread_idx()
236 + n_idx, _, _ = cute.arch.block_idx()
237 + wid = cute.arch.warp_idx()
238 +
239 + # One FP32 accumulator per token.
240 + acc = cute.make_rmem_tensor((M,), cutlass.Float32)
241 + acc.fill(0.0)
242 +
243 + if const_expr(self.use_pdl):
244 + cute.arch.griddepcontrol_wait()
245 +
246 + # 128-bit vectorized main loop
247 + if const_expr(k_main_elems > 0):
248 + gA_main = self._make_k_slice(gA, 0, k_main_elems)
249 + gB_main = self._make_k_slice(gB, 0, k_main_elems)
250 + gA_vec = cute.logical_divide(gA_main, (None, main_vec_width))
251 + gB_vec = cute.logical_divide(gB_main, (None, main_vec_width))
252 + tA, tB = self._make_thread_vector_slice(gA_vec, gB_vec, tidx, n_idx, bs)
253 + self._vector_dotprod(acc, tA, tB, M, main_tiles, 16)
254 +
255 + # 64-bit vectorized tail (K remainder after main loop)
256 + if const_expr(k_tail_elems > 0):
257 + gA_tail = self._make_k_slice(gA, k_main_elems, k_tail_elems)
258 + gB_tail = self._make_k_slice(gB, k_main_elems, k_tail_elems)
259 + gA_tail_vec = cute.logical_divide(gA_tail, (None, tail_vec_width))
260 + gB_tail_vec = cute.logical_divide(gB_tail, (None, tail_vec_width))
261 + tA_t, tB_t = self._make_thread_vector_slice(
262 + gA_tail_vec, gB_tail_vec, tidx, n_idx, bs
263 + )
264 + self._vector_dotprod(acc, tA_t, tB_t, M, tail_tiles, 8)
265 +
266 + # Full scalar rounds use CuTe width-1 tiles; KS_PART is the ragged tail.
267 + if const_expr(ks_full > 0):
268 + gA_scalar = self._make_k_slice(gA, k_done_all, k_scalar_full)
269 + gB_scalar = self._make_k_slice(gB, k_done_all, k_scalar_full)
270 + gA_scalar_vec = cute.logical_divide(gA_scalar, (None, 1))
271 + gB_scalar_vec = cute.logical_divide(gB_scalar, (None, 1))
272 + tA_s, tB_s = self._make_thread_vector_slice(
273 + gA_scalar_vec, gB_scalar_vec, tidx, n_idx, bs
274 + )
275 + self._vector_dotprod(acc, tA_s, tB_s, M, ks_full, 2)
276 +
277 + # Only threads below KS_PART load the ragged tail.
278 + if const_expr(ks_part > 0):
279 + gA_part = self._make_k_slice(gA, k_part_offset, ks_part)
280 + gB_part = self._make_k_slice(gB, k_part_offset, ks_part)
281 + if tidx < ks_part:
282 + bv2 = gB_part[n_idx, tidx].to(cutlass.Float32)
283 + for m in cutlass.range_constexpr(M):
284 + acc[m] = acc[m] + gA_part[m, tidx].to(cutlass.Float32) * bv2
285 +
286 + # Intra-warp shuffle reduction
287 + for m in cutlass.range_constexpr(M):
288 + acc[m] = cute.arch.warp_reduction_sum(acc[m])
289 +
290 + # Cross-warp reduction via shared memory
291 + smem_red_layout = cute.make_layout((M, num_warps), stride=(num_warps, 1))
292 + smem = cutlass.utils.SmemAllocator()
293 + sm = smem.allocate_tensor(cutlass.Float32, smem_red_layout, byte_alignment=16)
294 + with cute.arch.elect_one():
295 + for m in cutlass.range_constexpr(M):
296 + sm[m, wid] = acc[m]
297 +
298 + # Final reduction and output
299 + cute.arch.sync_threads()
300 + if tidx == 0:
301 + for m in cutlass.range_constexpr(M):
302 + partials = sm[m, None].load()
303 + gC[m, n_idx] = partials.reduce(
304 + cute.ReductionOp.ADD,
305 + init_val=cutlass.Float32(0.0),
306 + reduction_profile=0,
307 + )
308 + if const_expr(self.use_pdl):
309 + cute.arch.griddepcontrol_launch_dependents()
310 +
311 +
312 +def make_host_bf16(k_val: int, bs: int = 128):
313 + return LLBf16Dotprod(k=k_val, bs=bs)
vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py new
+585
@@ -0,0 +1,585 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import math
5 +
6 +import cutlass
7 +import cutlass.cute as cute
8 +from cuda.bindings.driver import CUstream
9 +from cutlass import const_expr
10 +from cutlass._mlir import ir as _ir
11 +from cutlass._mlir.dialects import llvm as _llvm
12 +from cutlass.cutlass_dsl import dsl_user_op
13 +from cutlass.pipeline import sm90 as pipeline
14 +
15 +
16 +# Map a local smem address to the same peer CTA DSMEM offset.
17 +@dsl_user_op
18 +def set_block_rank(smem_ptr, peer_rank, *, loc=None, ip=None):
19 + dsmem_ptr = cute.arch.map_dsmem_ptr(smem_ptr, peer_rank, loc=loc, ip=ip)
20 + return cutlass.Int32(dsmem_ptr.toint(loc=loc, ip=ip))
21 +
22 +
23 +# Plain DSMEM store; mbarrier helpers do not model this reduction.
24 +@dsl_user_op
25 +def st_shared_remote_f32(remote_addr, val, *, loc=None, ip=None):
26 + i32 = _ir.IntegerType.get_signless(32)
27 + addr_ir = remote_addr.ir_value(loc=loc, ip=ip)
28 + val_ir = val.ir_value(loc=loc, ip=ip)
29 + _llvm.inline_asm(
30 + i32,
31 + [addr_ir, val_ir],
32 + "st.shared::cluster.f32 [$0], $1; mov.u32 $2, 0;",
33 + "r,f,=r",
34 + has_side_effects=True, # keep the inline store ordered
35 + loc=loc,
36 + ip=ip,
37 + )
38 +
39 +
40 +class LLBf16SplitK:
41 + """BF16 router GEMM kernel based on clustered split-K MMA.
42 +
43 + This kernel computes C[M, N] = A[M, K] @ B[N, K]^T for bf16 inputs
44 + and fp32 output. It partitions K across a CTA cluster, uses DMA warps to
45 + stage A/B tiles with cp.async, uses MMA warps to accumulate fp32 partials,
46 + and reduces split-K partials through DSMEM before storing C.
47 +
48 + :param ab_dtype: Element type for A and B operands.
49 + :param acc_dtype: Accumulator type used by MMA and reductions.
50 + :param out_dtype: Output element type. The public wrapper uses fp32.
51 + :param tile_n: CTA tile size in N. M is fixed to 16 for router batches.
52 + :type tile_n: int
53 + :param tile_k: K tile size staged through shared memory.
54 + :type tile_k: int
55 + :param num_stages: Number of cp.async pipeline stages.
56 + :type num_stages: int
57 + :param num_dma_warps: Producer warps that issue GMEM to SMEM copies.
58 + :type num_dma_warps: int
59 + :param split_k: Number of CTAs in the cluster-level split-K reduction.
60 + :type split_k: int
61 + :param use_pdl: Whether to launch with Programmatic Dependent Launch.
62 + :type use_pdl: bool
63 +
64 + :note: Supported A/B data types:
65 + - BFloat16/BFloat16
66 + :note: Supported accumulator data types:
67 + - Float32
68 + :note: Supported C data types:
69 + - Float32
70 + :note: Constraints:
71 + - K must preserve 16-byte row alignment for contiguous bf16 inputs.
72 +
73 + :compile-key: ``(split_k, num_stages)`` selects the cluster split
74 + count and cp.async pipeline depth specialization.
75 + """
76 +
77 + def __init__(
78 + self,
79 + ab_dtype=cutlass.BFloat16,
80 + acc_dtype=cutlass.Float32,
81 + out_dtype=cutlass.Float32,
82 + tile_n: int = 16,
83 + tile_k: int = 256,
84 + num_stages: int = 2,
85 + num_dma_warps: int = 4,
86 + split_k: int = 8,
87 + use_pdl: bool = False,
88 + ):
89 + """Initialize the split-K kernel configuration.
90 +
91 + This configuration fixes the CTA tile shape, cp.async pipeline depth,
92 + producer/consumer warp split, and CTA-cluster split count used by the
93 + DSMEM reduction.
94 +
95 + :param ab_dtype: Element type for A and B operands.
96 + :param acc_dtype: Accumulator type used by MMA and reductions.
97 + :param out_dtype: Output element type.
98 + :param tile_n: CTA tile size in N. M is fixed to 16.
99 + :type tile_n: int
100 + :param tile_k: K tile size staged through shared memory.
101 + :type tile_k: int
102 + :param num_stages: Number of cp.async pipeline stages.
103 + :type num_stages: int
104 + :param num_dma_warps: Producer warps that issue GMEM to SMEM copies.
105 + :type num_dma_warps: int
106 + :param split_k: CTAs in the cluster-level split-K reduction.
107 + :type split_k: int
108 + :param use_pdl: Whether to launch with Programmatic Dependent Launch.
109 + :type use_pdl: bool
110 + """
111 + self.ab_dtype = ab_dtype
112 + self.acc_dtype = acc_dtype
113 + self.out_dtype = out_dtype
114 + self.tile_m = 16
115 + self.tile_n = tile_n
116 + self.tile_k = tile_k
117 + self.copy_bits = 128
118 + self.num_stages = num_stages
119 + self.split_k = split_k
120 + self.use_pdl = use_pdl
121 + self.mma_shape = (16, 8, 16) # mma.sync.aligned.m16n8k16
122 + self.atom_layout = (1, 1, 1) # one MMA atom per warp
123 + self.num_dma_warps = num_dma_warps
124 + self.num_mma_warps = 4
125 + self.num_dma_threads = self.num_dma_warps * cute.arch.WARP_SIZE
126 + self.num_mma_threads = self.num_mma_warps * cute.arch.WARP_SIZE
127 + self.num_threads = self.num_dma_threads + self.num_mma_threads
128 + self.num_epilogue_elems = self.tile_m * self.tile_n
129 + self.epilogue_elems_per_thread = self.num_epilogue_elems // self.num_mma_threads
130 +
131 + def _make_smem_layout_AB(self, dtype, copy_bits, smem_tiler):
132 + """Build the staged swizzled SMEM layout for A or B tiles."""
133 + major_size = min(smem_tiler[1], 64)
134 + # Match swizzle span to contiguous K bytes, capped by CuTe 3-bit swizzle.
135 + swizzle_bits = int(math.log2(major_size * dtype.width // copy_bits))
136 + swizzle_bits = min(swizzle_bits, 3)
137 + # Tile the swizzled atom across (M_or_N, K, stages).
138 + layout_atom_outer = cute.make_layout((8, major_size), stride=(major_size, 1))
139 + layout_atom = cute.make_composed_layout(
140 + cute.make_swizzle(swizzle_bits, 3, 3), 0, layout_atom_outer
141 + )
142 + return cute.tile_to_shape(layout_atom, smem_tiler, (0, 1, 2))
143 +
144 + def _make_gmem_tiled_copy(self, atom_copy, dtype, copy_bits, num_threads):
145 + """Build the per-thread cp.async vector-copy layout."""
146 + # Lay threads across K so each lane issues one vector copy.
147 + copy_elems = copy_bits // dtype.width
148 + k_threads = cute.size(self.tile_k) // copy_elems # threads along K
149 + thread_layout = cute.make_layout(
150 + (num_threads // k_threads, k_threads), stride=(k_threads, 1)
151 + )
152 + value_layout = cute.make_layout((1, copy_elems))
153 + return cute.make_tiled_copy_tv(atom_copy, thread_layout, value_layout)
154 +
155 + @cute.jit
156 + def _fill_pred(self, pred_flat, coord_tensor, k_tile, dim_limit, K_total):
157 + # Predicate one K tile and keep the stage-broadcast view in sync.
158 + coord_ktile = coord_tensor[None, None, 0, k_tile]
159 + num_vec = pred_flat.shape[0]
160 + num_mn = pred_flat.shape[1]
161 + for v in cutlass.range_constexpr(num_vec):
162 + # pred_flat is (K_VEC, M/N) for one K_TILE.
163 + for j in cutlass.range_constexpr(num_mn):
164 + pred_flat[v, j] = cute.elem_less(
165 + coord_ktile[(0, v), j], (dim_limit, K_total)
166 + )
167 +
168 + def _make_pred(self, tXcX, k_tile, dim_limit, K_total):
169 + # pred_flat is (K_VEC, M/N); pred is (K_VEC, M/N, STAGE).
170 + num_vec = tXcX.shape[0][1]
171 + num_mn = cute.size(tXcX, mode=[1])
172 + pred_flat = cute.make_rmem_tensor(
173 + cute.make_layout((num_vec, num_mn), stride=(num_mn, 1)),
174 + cutlass.Boolean,
175 + )
176 + pred = cute.make_tensor(
177 + pred_flat.iterator,
178 + cute.make_layout(
179 + (num_vec, num_mn, cute.size(tXcX, mode=[2])),
180 + stride=(num_mn, 1, 0),
181 + ),
182 + )
183 + return pred_flat, pred
184 +
185 + @cute.jit
186 + def __call__(
187 + self,
188 + mA: cute.Tensor,
189 + mB: cute.Tensor,
190 + mC: cute.Tensor,
191 + stream: CUstream,
192 + scale: float = 1.0,
193 + ):
194 + """Execute the split-K GEMM operation in steps:
195 + - Build swizzled staged SMEM layouts for A ``[16, tile_k, stages]``
196 + and B ``[tile_n, tile_k, stages]``.
197 + - Build 128-bit cp.async GMEM-to-SMEM tiled copies for DMA warps.
198 + - Build an m16n8k16 BF16 MMA tiled across the CTA N tile.
199 + - Launch grid ``ceil(M/16) x ceil(N/tile_n) x split_k`` with one
200 + cluster along split-K, so each cluster rank owns K tiles
201 + ``rank, rank + split_k, ...``.
202 + - Reduce MMA-warp partials within the CTA, exchange split-K partials
203 + through DSMEM, then reduce cluster partials and store C.
204 +
205 + :param mA: Input tensor A with shape ``[M, K]``.
206 + :type mA: cute.Tensor
207 + :param mB: Input tensor B with shape ``[N, K]``.
208 + :type mB: cute.Tensor
209 + :param mC: Output tensor C with shape ``[M, N]``.
210 + :type mC: cute.Tensor
211 + :param stream: CUDA stream for asynchronous execution.
212 + :type stream: CUstream
213 + :param scale: Epilogue scale applied before storing C.
214 + :type scale: float
215 + """
216 + bM, bN, bK = self.tile_m, self.tile_n, self.tile_k
217 + copy_bits: cutlass.Constexpr = self.copy_bits
218 + sA_layout = self._make_smem_layout_AB(
219 + mA.element_type, copy_bits, (bM, bK, self.num_stages)
220 + )
221 + sB_layout = self._make_smem_layout_AB(
222 + mB.element_type, copy_bits, (bN, bK, self.num_stages)
223 + )
224 +
225 + @cute.struct
226 + class SharedStorage:
227 + a: cute.struct.Align[
228 + cute.struct.MemRange[mA.element_type, cute.cosize(sA_layout)], 16
229 + ]
230 + b: cute.struct.Align[
231 + cute.struct.MemRange[mB.element_type, cute.cosize(sB_layout)], 16
232 + ]
233 + mbar: cute.struct.Align[
234 + cute.struct.MemRange[cutlass.Int64, self.num_stages * 2], 8
235 + ]
236 +
237 + atom_g2s = cute.make_copy_atom(
238 + cute.nvgpu.cpasync.CopyG2SOp(
239 + cache_mode=cute.nvgpu.cpasync.LoadCacheMode.GLOBAL
240 + ),
241 + mA.element_type,
242 + num_bits_per_copy=copy_bits,
243 + ) # cp.async GMEM -> SMEM, bypassing L1
244 + tiled_copy_A = self._make_gmem_tiled_copy(
245 + atom_g2s, mA.element_type, copy_bits, self.num_dma_threads
246 + )
247 + tiled_copy_B = self._make_gmem_tiled_copy(
248 + atom_g2s, mB.element_type, copy_bits, self.num_dma_threads
249 + )
250 + op = cute.nvgpu.warp.MmaF16BF16Op(self.ab_dtype, self.acc_dtype, self.mma_shape)
251 + # Repeat the m16n8k16 atom along N to cover the CTA output tile.
252 + perm_mnk = (
253 + self.atom_layout[0] * self.mma_shape[0],
254 + self.atom_layout[1] * self.mma_shape[1] * (self.tile_n // 8),
255 + self.atom_layout[2] * self.mma_shape[2],
256 + )
257 + tiled_mma = cute.make_tiled_mma(
258 + op, cute.make_layout(self.atom_layout), permutation_mnk=perm_mnk
259 + )
260 + tiler_mn = (bM, bN)
261 + grid_m, grid_n = cute.ceil_div(mC.shape, tiler_mn)
262 + self.kernel(
263 + mA,
264 + mB,
265 + mC,
266 + scale,
267 + sA_layout,
268 + sB_layout,
269 + tiled_copy_A,
270 + tiled_copy_B,
271 + tiled_mma,
272 + SharedStorage,
273 + ).launch(
274 + grid=[
275 + cute.size(grid_m),
276 + cute.size(grid_n),
277 + self.split_k,
278 + ],
279 + block=[self.num_threads, 1, 1],
280 + cluster=[
281 + 1,
282 + 1,
283 + self.split_k,
284 + ], # split-K CTAs form one cluster
285 + stream=stream,
286 + use_pdl=self.use_pdl,
287 + )
288 +
289 + @cute.kernel
290 + def kernel(
291 + self,
292 + mA,
293 + mB,
294 + mC,
295 + scale: cutlass.Float32,
296 + sA_layout: cute.ComposedLayout,
297 + sB_layout: cute.ComposedLayout,
298 + tiled_copy_A: cute.TiledCopy,
299 + tiled_copy_B: cute.TiledCopy,
300 + tiled_mma: cute.TiledMma,
301 + shared_storage: cutlass.Constexpr,
302 + ):
303 + bM, bN, bK = self.tile_m, self.tile_n, self.tile_k
304 + num_stages = self.num_stages
305 + tidx, _, _ = cute.arch.thread_idx()
306 + bid_m, bid_n, bid_z = cute.arch.block_idx()
307 + warp_idx = cute.arch.warp_idx()
308 + lane_id = cute.arch.lane_idx()
309 + num_dma_warps: cutlass.Constexpr = self.num_dma_warps
310 + is_dma_warp = warp_idx < num_dma_warps
311 + dma_tidx = tidx
312 + mma_tidx = tidx - self.num_dma_threads
313 + mma_warp_idx = warp_idx - num_dma_warps
314 +
315 + cta_tiler = (bM, bN, bK)
316 + coord = (bid_m, bid_n, None) # all K tiles
317 + # CTA-local tiles.
318 + gA = cute.local_tile(
319 + mA, tiler=cta_tiler, coord=coord, proj=(1, None, 1)
320 + ) # skip N
321 + gB = cute.local_tile(
322 + mB, tiler=cta_tiler, coord=coord, proj=(None, 1, 1)
323 + ) # skip M
324 + gC = cute.local_tile(
325 + mC, tiler=cta_tiler, coord=coord, proj=(1, 1, None)
326 + ) # skip K
327 +
328 + mcA = cute.make_identity_tensor(mA.layout.shape)
329 + mcB = cute.make_identity_tensor(mB.layout.shape)
330 + mcC = cute.make_identity_tensor(mC.layout.shape)
331 + # Coordinate modes: cA=(M,K,k_tile), cB=(N,K,k_tile), cC=(M,N).
332 + cA = cute.local_tile(mcA, tiler=cta_tiler, coord=coord, proj=(1, None, 1))
333 + cB = cute.local_tile(mcB, tiler=cta_tiler, coord=coord, proj=(None, 1, 1))
334 + cC = cute.local_tile(mcC, tiler=cta_tiler, coord=coord, proj=(1, 1, None))
335 +
336 + # 128-bit cp.async copies require 16-byte aligned GMEM views.
337 + gA = cute.make_tensor(gA.iterator.align(16), gA.layout)
338 + gB = cute.make_tensor(gB.iterator.align(16), gB.layout)
339 +
340 + smem = cutlass.utils.SmemAllocator()
341 + storage_ptr = smem.allocate(shared_storage.size_in_bytes(), byte_alignment=16) # type: ignore[attr-defined]
342 + storage = shared_storage(storage_ptr) # type: ignore[call-arg]
343 + sA = storage.a.get_tensor(sA_layout)
344 + sB = storage.b.get_tensor(sB_layout)
345 +
346 + # Pipeline cp.async producers into MMA consumers.
347 + producer_group = pipeline.CooperativeGroup(
348 + pipeline.Agent.Thread, self.num_dma_threads
349 + )
350 + consumer_group = pipeline.CooperativeGroup(
351 + pipeline.Agent.Thread, self.num_mma_threads
352 + )
353 + mainloop_pipeline = pipeline.PipelineCpAsync.create(
354 + barrier_storage=storage.mbar.data_ptr(),
355 + num_stages=num_stages,
356 + producer_group=producer_group,
357 + consumer_group=consumer_group,
358 + )
359 +
360 + # Round-robin split-K: split z handles tiles z, z+split_k, ...
361 + K_total = cute.size(mA, mode=[1])
362 + k_tile_count = cute.size(gA, mode=[2])
363 + k_start = bid_z
364 + num_k_tiles = cute.ceil_div(k_tile_count - k_start, self.split_k)
365 +
366 + if is_dma_warp:
367 + # DMA warps trade registers for copy throughput.
368 + cute.arch.setmaxregister_decrease(40)
369 + thr_A = tiled_copy_A.get_slice(dma_tidx)
370 + thr_B = tiled_copy_B.get_slice(dma_tidx)
371 + tAgA = thr_A.partition_S(gA)
372 + tAsA = thr_A.partition_D(sA)
373 + tBgB = thr_B.partition_S(gB)
374 + tBsB = thr_B.partition_D(sB)
375 + tAcA = thr_A.partition_S(cA)
376 + tBcB = thr_B.partition_S(cB)
377 +
378 + # Build M/K and N/K predicates, broadcast across K-tile copies.
379 + tApA_flat, tApA = self._make_pred(tAcA, k_start, mA.shape[0], K_total)
380 + tBpB_flat, tBpB = self._make_pred(tBcB, k_start, mB.shape[0], K_total)
381 + self._fill_pred(tApA_flat, tAcA, k_start, mA.shape[0], K_total)
382 + self._fill_pred(tBpB_flat, tBcB, k_start, mB.shape[0], K_total)
383 +
384 + producer_state = pipeline.make_pipeline_state(
385 + pipeline.PipelineUserType.Producer, num_stages
386 + )
387 +
388 + # Prime the first pipeline stage.
389 + mainloop_pipeline.producer_acquire(producer_state)
390 + cute.copy(
391 + tiled_copy_B,
392 + tBgB[None, None, None, k_start],
393 + tBsB[None, None, None, producer_state.index],
394 + pred=tBpB,
395 + )
396 + if const_expr(self.use_pdl):
397 + cute.arch.griddepcontrol_wait()
398 + cute.copy(
399 + tiled_copy_A,
400 + tAgA[None, None, None, k_start],
401 + tAsA[None, None, None, producer_state.index],
402 + pred=tApA,
403 + )
404 + mainloop_pipeline.producer_commit(producer_state)
405 + producer_state.advance()
406 +
407 + for k_tile in cutlass.range(
408 + k_start + self.split_k, k_tile_count, self.split_k, unroll=1
409 + ):
410 + self._fill_pred(tApA_flat, tAcA, k_tile, mA.shape[0], K_total)
411 + self._fill_pred(tBpB_flat, tBcB, k_tile, mB.shape[0], K_total)
412 + mainloop_pipeline.producer_acquire(producer_state)
413 + cute.copy(
414 + tiled_copy_A,
415 + tAgA[None, None, None, k_tile],
416 + tAsA[None, None, None, producer_state.index],
417 + pred=tApA,
418 + )
419 + cute.copy(
420 + tiled_copy_B,
421 + tBgB[None, None, None, k_tile],
422 + tBsB[None, None, None, producer_state.index],
423 + pred=tBpB,
424 + )
425 + mainloop_pipeline.producer_commit(producer_state)
426 + producer_state.advance()
427 +
428 + mainloop_pipeline.producer_tail(producer_state)
429 +
430 + else:
431 + # MMA warps with k-phase interleaving
432 + cute.arch.setmaxregister_increase(232) # large MMA fragments
433 +
434 + num_mma_warps: cutlass.Constexpr = self.num_mma_warps
435 +
436 + thr_mma = tiled_mma.get_slice(lane_id)
437 + tCsA = thr_mma.partition_A(sA)
438 + tCsB = thr_mma.partition_B(sB)
439 + tCgC = thr_mma.partition_C(gC)
440 +
441 + tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0])
442 + tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0])
443 + tCrC = tiled_mma.make_fragment_C(tCgC)
444 + tCrC.fill(0.0)
445 +
446 + # ldmatrix moves SMEM fragments into MMA registers.
447 + atom_s2r_A = cute.make_copy_atom(
448 + cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), mA.element_type
449 + ) # non-transposed, x4
450 + atom_s2r_B = cute.make_copy_atom(
451 + cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), mB.element_type
452 + )
453 +
454 + # SMEM -> register copy path.
455 + tiled_s2r_A = cute.make_tiled_copy_A(atom_s2r_A, tiled_mma)
456 + tiled_s2r_B = cute.make_tiled_copy_B(atom_s2r_B, tiled_mma)
457 + thr_s2r_A = tiled_s2r_A.get_slice(lane_id)
458 + thr_s2r_B = tiled_s2r_B.get_slice(lane_id)
459 + tCsA_v = thr_s2r_A.partition_S(sA) # views, not copies
460 + tCrA_v = thr_s2r_A.retile(tCrA)
461 + tCsB_v = thr_s2r_B.partition_S(sB)
462 + tCrB_v = thr_s2r_B.retile(tCrB)
463 + # Split the MMA K-fragments across the MMA warps.
464 + tCsA_warp_v = cute.logical_divide(tCsA_v, (None, None, num_mma_warps, None))
465 + tCsB_warp_v = cute.logical_divide(tCsB_v, (None, None, num_mma_warps, None))
466 +
467 + num_k_blocks = cute.size(tCrA, mode=[2])
468 + k_blocks_per_warp: cutlass.Constexpr = num_k_blocks // num_mma_warps
469 +
470 + consumer_state = pipeline.make_pipeline_state(
471 + pipeline.PipelineUserType.Consumer, num_stages
472 + )
473 +
474 + # Shape-dynamic split-K count, so this stays a runtime range.
475 + for _ in cutlass.range(num_k_tiles, unroll_full=True):
476 + mainloop_pipeline.consumer_wait(consumer_state)
477 + for ki in cutlass.range_constexpr(k_blocks_per_warp):
478 + cute.copy(
479 + tiled_s2r_A,
480 + tCsA_warp_v[
481 + None, None, (mma_warp_idx, ki), consumer_state.index
482 + ],
483 + tCrA_v[None, None, 0],
484 + ) # ldmatrix
485 + cute.copy(
486 + tiled_s2r_B,
487 + tCsB_warp_v[
488 + None, None, (mma_warp_idx, ki), consumer_state.index
489 + ],
490 + tCrB_v[None, None, 0],
491 + )
492 + cute.gemm(
493 + tiled_mma, tCrC, tCrA[None, None, 0], tCrB[None, None, 0], tCrC
494 + ) # mma.sync
495 + mainloop_pipeline.consumer_release(consumer_state)
496 + consumer_state.advance()
497 +
498 + # Cluster reduction epilogue.
499 + # Reduce per-warp accumulators within this CTA.
500 + num_elems: cutlass.Constexpr = self.num_epilogue_elems
501 + elems_per_thread: cutlass.Constexpr = self.epilogue_elems_per_thread
502 + # Map MMA threads to linear CTA output elements.
503 + epilogue_thread_layout = cute.make_layout(
504 + (elems_per_thread, self.num_mma_threads),
505 + stride=(self.num_mma_threads, 1),
506 + )
507 + epilogue_slots = cute.make_tensor(0, epilogue_thread_layout)
508 + epilogue_slot_coords = cute.make_identity_tensor((bN, bM))
509 + # Layout: (mma_warp, linear MN element).
510 + smem_red = cute.make_tensor(
511 + cute.arch.alloc_smem(
512 + cutlass.Float32, num_elems * num_mma_warps, alignment=16
513 + ),
514 + cute.make_layout((num_mma_warps, num_elems), stride=(num_elems, 1)),
515 + )
516 + smem_warp = cute.make_tensor(
517 + cute.domain_offset((mma_warp_idx, 0), smem_red).iterator,
518 + cute.make_layout((bM, bN), stride=(bN, 1)),
519 + )
520 + tCsC_partial = thr_mma.partition_C(smem_warp)
521 + cute.autovec_copy(tCrC, tCsC_partial)
522 + cute.arch.sync_threads()
523 +
524 + # Layout: (split-K rank, linear MN element).
525 + partials = cute.make_tensor(
526 + cute.arch.alloc_smem(
527 + cutlass.Float32, num_elems * self.split_k, alignment=16
528 + ),
529 + cute.make_layout((self.split_k, num_elems), stride=(num_elems, 1)),
530 + )
531 + cta_rank = cute.arch.block_idx_in_cluster()
532 +
533 + for ei in cutlass.range_constexpr(elems_per_thread):
534 + elem_idx = epilogue_slots[ei, mma_tidx]
535 + local_coord = cute.select(epilogue_slot_coords[elem_idx], mode=[1, 0])
536 + total = cutlass.Float32(0.0)
537 + if cute.elem_less(local_coord, gC.shape):
538 + total = (
539 + smem_red[None, elem_idx]
540 + .load()
541 + .reduce(
542 + cute.ReductionOp.ADD,
543 + init_val=cutlass.Float32(0.0),
544 + reduction_profile=0,
545 + )
546 + )
547 + total = total * scale
548 + partials[cta_rank, elem_idx] = total
549 +
550 + cute.arch.sync_threads()
551 +
552 + # Broadcast this CTA's partials to peer DSMEM.
553 + for ei in cutlass.range_constexpr(elems_per_thread):
554 + elem_idx = epilogue_slots[ei, mma_tidx]
555 + my_slot = cute.domain_offset((cta_rank, elem_idx), partials).iterator
556 + my_val = partials[cta_rank, elem_idx]
557 + for peer in cutlass.range_constexpr(self.split_k):
558 + remote = set_block_rank(my_slot, cutlass.Int32(peer))
559 + st_shared_remote_f32(remote, my_val)
560 +
561 + cute.arch.cluster_arrive()
562 + cute.arch.cluster_wait() # peer DSMEM stores are now visible
563 +
564 + if const_expr(self.use_pdl) and mma_tidx == 0:
565 + cute.arch.griddepcontrol_launch_dependents()
566 + cute.arch.sync_threads()
567 +
568 + # Reduce split-K partials and write global output.
569 + for ei in cutlass.range_constexpr(elems_per_thread):
570 + elem_idx = epilogue_slots[ei, mma_tidx]
571 + local_coord = cute.select(epilogue_slot_coords[elem_idx], mode=[1, 0])
572 + global_coord = cC[local_coord]
573 + if cute.elem_less(global_coord, mC.shape):
574 + acc = (
575 + partials[None, elem_idx]
576 + .load()
577 + .reduce(
578 + cute.ReductionOp.ADD,
579 + init_val=cutlass.Float32(0.0),
580 + reduction_profile=0,
581 + )
582 + )
583 + gC[local_coord] = acc
584 +
585 + cute.arch.sync_threads()
vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py new
+282
@@ -0,0 +1,282 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from __future__ import annotations
5 +
6 +import logging
7 +import math
8 +from collections.abc import Iterable
9 +from dataclasses import dataclass
10 +from typing import Any, Literal
11 +
12 +import torch
13 +
14 +logger = logging.getLogger(__name__)
15 +
16 +_cutedsl_available: bool | None = None
17 +
18 +
19 +def is_available() -> bool:
20 + global _cutedsl_available
21 + if _cutedsl_available is not None:
22 + return _cutedsl_available
23 + try:
24 + import cutlass # noqa: F401
25 + import cutlass.cute # noqa: F401
26 +
27 + _cutedsl_available = True
28 + except ImportError:
29 + _cutedsl_available = False
30 + logger.info("cuteDSL (CUTLASS Python) not available, ll_bf16_gemm disabled")
31 + return _cutedsl_available
32 +
33 +
34 +_DEFAULT_DOTPROD_BS = 128
35 +_DEFAULT_DOTPROD_MAX_M = 4
36 +_DEFAULT_SPLITK_CONFIG = (6, 4)
37 +_TUNED_DOTPROD_MAX_M: dict[tuple[int, int], int] = {
38 + (7168, 256): 6,
39 +}
40 +_TUNED_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = {
41 + (7168, 384): {
42 + 5: (4, 4),
43 + **{M: (5, 4) for M in range(6, 17)},
44 + },
45 +}
46 +
47 +
48 +_cute_ctx = None
49 +
50 +
51 +def _cute():
52 + global _cute_ctx
53 + if _cute_ctx is not None:
54 + return _cute_ctx
55 + import cutlass.cute as cute
56 + from cuda.bindings.driver import CUstream
57 +
58 + _cute_ctx = (cute, CUstream)
59 + return _cute_ctx
60 +
61 +
62 +def _stream():
63 + _, CUstream = _cute()
64 + from vllm.utils.torch_utils import current_stream
65 +
66 + return CUstream(current_stream().cuda_stream)
67 +
68 +
69 +def _use_pdl() -> bool:
70 + from vllm.platforms import current_platform
71 +
72 + return current_platform.is_arch_support_pdl()
73 +
74 +
75 +class LLBf16Gemm:
76 + @dataclass(frozen=True, slots=True)
77 + class CompileKey:
78 + backend: Literal["dotprod", "splitk"]
79 + M: int = 0
80 + K: int = 0
81 + bs: int = 0
82 + split_k: int = 0
83 + num_stages: int = 0
84 +
85 + def __init__(self) -> None:
86 + # Dot-prod: keyed on (M, K, bs), because M and K are Constexpr.
87 + self._compiled_cache: dict[tuple[int, int, int], Any] = {}
88 + # Split-K: keyed on (split_k, num_stages), fully shape-dynamic.
89 + self._splitk_cache: dict[tuple[int, int], Any] = {}
90 +
91 + def dispatch(self, *, M: int, K: int, N: int) -> CompileKey:
92 + dotprod_max_m = _TUNED_DOTPROD_MAX_M.get((K, N), _DEFAULT_DOTPROD_MAX_M)
93 + if dotprod_max_m >= M or K < 2048:
94 + return self.CompileKey(backend="dotprod", M=M, K=K, bs=_DEFAULT_DOTPROD_BS)
95 +
96 + split_k, num_stages = _TUNED_CONFIGS.get((K, N), {}).get(
97 + M, _DEFAULT_SPLITK_CONFIG
98 + )
99 + return self.CompileKey(backend="splitk", split_k=split_k, num_stages=num_stages)
100 +
101 + def get_warmup_keys(
102 + self,
103 + *,
104 + shapes: Iterable[tuple[int, int]],
105 + m_values: Iterable[int],
106 + ) -> list[CompileKey]:
107 + return list(
108 + dict.fromkeys(
109 + self.dispatch(M=M, K=K, N=N) for K, N in shapes for M in m_values
110 + )
111 + )
112 +
113 + @staticmethod
114 + def _fake_gemm_tensors(*, M, K, N, divisibility: int):
115 + from cutlass import BFloat16, Float32
116 + from quack.compile_utils import make_fake_tensor
117 +
118 + hidden_states = make_fake_tensor(BFloat16, (M, K), divisibility=divisibility)
119 + router_weight = make_fake_tensor(BFloat16, (N, K), divisibility=divisibility)
120 + output = make_fake_tensor(Float32, (M, N), divisibility=1)
121 + return hidden_states, router_weight, output
122 +
123 + def _compile_splitk(self, compile_key: CompileKey) -> None:
124 + cute, _ = _cute()
125 + from ._ll_bf16_splitk import LLBf16SplitK
126 +
127 + hidden_states, router_weight, output = self._fake_gemm_tensors(
128 + M=cute.sym_int(),
129 + K=cute.sym_int(),
130 + N=cute.sym_int(),
131 + divisibility=8,
132 + )
133 + gemm = LLBf16SplitK(
134 + tile_n=16,
135 + tile_k=256,
136 + num_stages=compile_key.num_stages,
137 + num_dma_warps=4,
138 + split_k=compile_key.split_k,
139 + use_pdl=_use_pdl(),
140 + )
141 + compiled = cute.compile(
142 + gemm,
143 + hidden_states,
144 + router_weight,
145 + output,
146 + _stream(),
147 + options="--enable-tvm-ffi",
148 + )
149 + self._splitk_cache[(compile_key.split_k, compile_key.num_stages)] = compiled
150 + logger.debug(
151 + "Compiled ll_bf16_splitk: sk=%d ns=%d",
152 + compile_key.split_k,
153 + compile_key.num_stages,
154 + )
155 +
156 + def _compile_dotprod(self, compile_key: CompileKey) -> None:
157 + cute, _ = _cute()
158 + from ._ll_bf16_dotprod import LLBf16Dotprod
159 +
160 + N = cute.sym_int()
161 + stride_divisibility = math.gcd(8, compile_key.K)
162 + hidden_states, router_weight, output = self._fake_gemm_tensors(
163 + M=compile_key.M,
164 + K=compile_key.K,
165 + N=N,
166 + divisibility=stride_divisibility,
167 + )
168 + gemm = LLBf16Dotprod(k=compile_key.K, bs=compile_key.bs, use_pdl=_use_pdl())
169 + compiled = cute.compile(
170 + gemm,
171 + hidden_states,
172 + router_weight,
173 + output,
174 + compile_key.M,
175 + compile_key.K,
176 + 1, # runtime N placeholder for fake-tensor compile
177 + _stream(),
178 + options="--enable-tvm-ffi --ptxas-options -maxrregcount=64",
179 + )
180 + self._compiled_cache[(compile_key.M, compile_key.K, compile_key.bs)] = compiled
181 + logger.debug(
182 + "Compiled ll_bf16_dotprod: M=%d, K=%d, bs=%d",
183 + compile_key.M,
184 + compile_key.K,
185 + compile_key.bs,
186 + )
187 +
188 + def compile(self, compile_key: CompileKey) -> None:
189 + if compile_key.backend == "splitk":
190 + splitk_cache_key = (compile_key.split_k, compile_key.num_stages)
191 + if splitk_cache_key not in self._splitk_cache:
192 + self._compile_splitk(compile_key)
193 + return
194 +
195 + dotprod_cache_key = (compile_key.M, compile_key.K, compile_key.bs)
196 + if dotprod_cache_key not in self._compiled_cache:
197 + self._compile_dotprod(compile_key)
198 +
199 + def warmup(
200 + self,
201 + *,
202 + shapes: Iterable[tuple[int, int]],
203 + m_values: Iterable[int],
204 + ) -> None:
205 + for compile_key in self.get_warmup_keys(shapes=shapes, m_values=m_values):
206 + self.compile(compile_key)
207 +
208 + @staticmethod
209 + def _validate_inputs(
210 + hidden_states: torch.Tensor,
211 + router_weight: torch.Tensor,
212 + output_dtype: torch.dtype,
213 + ) -> None:
214 + if hidden_states.dim() != 2 or router_weight.dim() != 2:
215 + raise ValueError("hidden_states and router_weight must be 2D tensors")
216 + if (
217 + hidden_states.dtype != torch.bfloat16
218 + or router_weight.dtype != torch.bfloat16
219 + ):
220 + raise ValueError("hidden_states and router_weight must have dtype=bfloat16")
221 + if hidden_states.device.type != "cuda" or router_weight.device.type != "cuda":
222 + raise ValueError(
223 + "hidden_states and router_weight must have device_type=cuda"
224 + )
225 + if hidden_states.device != router_weight.device:
226 + raise ValueError(
227 + "hidden_states and router_weight must be on the same CUDA device"
228 + )
229 + if output_dtype != torch.float32:
230 + raise ValueError("ll_bf16_gemm only supports output_dtype=torch.float32")
231 + if hidden_states.shape[1] != router_weight.shape[1]:
232 + raise ValueError(
233 + "hidden_states and router_weight must have matching K dimensions"
234 + )
235 + # Kernels use vectorized bf16 loads and require 16-byte row alignment.
236 + if hidden_states.shape[1] % 8 != 0:
237 + raise ValueError("ll_bf16_gemm requires K to be divisible by 8")
238 + if not hidden_states.is_contiguous() or not router_weight.is_contiguous():
239 + raise ValueError(
240 + "hidden_states and router_weight must be contiguous row-major inputs"
241 + )
242 +
243 + def __call__(
244 + self,
245 + hidden_states: torch.Tensor, # [M, K] bf16
246 + router_weight: torch.Tensor, # [N, K] bf16
247 + output_dtype: torch.dtype = torch.float32,
248 + ) -> torch.Tensor: # [M, N] fp32
249 + self._validate_inputs(hidden_states, router_weight, output_dtype)
250 +
251 + M, K = hidden_states.shape
252 + N = router_weight.shape[0]
253 + compile_key = self.dispatch(M=M, K=K, N=N)
254 + if compile_key.backend == "splitk":
255 + splitk_cache_key = (compile_key.split_k, compile_key.num_stages)
256 + if splitk_cache_key not in self._splitk_cache:
257 + self.compile(compile_key)
258 + kernel = self._splitk_cache[splitk_cache_key]
259 + else:
260 + dotprod_cache_key = (compile_key.M, compile_key.K, compile_key.bs)
261 + if dotprod_cache_key not in self._compiled_cache:
262 + self.compile(compile_key)
263 + kernel = self._compiled_cache[dotprod_cache_key]
264 +
265 + stream = _stream()
266 + output = torch.empty(M, N, dtype=output_dtype, device=hidden_states.device)
267 + if compile_key.backend == "splitk":
268 + kernel(hidden_states, router_weight, output, stream, 1.0)
269 + else:
270 + kernel(hidden_states, router_weight, output, N, stream)
271 + return output
272 +
273 +
274 +ll_bf16_gemm_kernel = LLBf16Gemm()
275 +
276 +
277 +def ll_bf16_gemm(
278 + hidden_states: torch.Tensor,
279 + router_weight: torch.Tensor,
280 + output_dtype: torch.dtype = torch.float32,
281 +) -> torch.Tensor:
282 + return ll_bf16_gemm_kernel(hidden_states, router_weight, output_dtype)
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+46 -8
@@ -14,11 +14,13 @@ from vllm.utils.torch_utils import direct_register_custom_op
14 class GateLinear(ReplicatedLinear):
15 """MoE gate linear layer with multi-tier GEMM dispatch:
16
17 - 1. DSV3 specialized kernel (SM90+, M<=16, H=7168 E=256/384, H=6144 E=256)
18 - 2. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, M<=32,
17 + 1. cuteDSL ll_bf16_gemm (SM90+, M<=16, bf16 in, fp32 out,
18 + K divisible by 8)
19 + 2. DSV3 specialized kernel (SM90+, M<=16, H=7168 E=256/384, H=6144 E=256)
20 + 3. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, M<=32,
21 (H, E) in {(3072, 256), (6144, 128), (6144, 256)})
20 - 3. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
21 - 4. F.linear via ReplicatedLinear (ultimate fallback)
22 + 4. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
23 + 5. F.linear via ReplicatedLinear (ultimate fallback)
24
25 The ``out_dtype`` attribute is mutable and can be set after init
26 (e.g. when the required dtype depends on the expert quantization
@@ -99,6 +101,21 @@ class GateLinear(ReplicatedLinear):
101 and self.out_dtype == torch.float32
102 )
103
104 + # cuteDSL ll_bf16_gemm eligibility. Any dims supported, but SM90+ required bc:
105 + # 1. PDL support. Both dot-product and split-K kernels.
106 + # 2. Thread Block Clusters. Split-K kernel for cross-CTA reduction.
107 + self.allow_ll_bf16_gemm = False
108 + if can_use_specialized_kernels:
109 + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import (
110 + is_available,
111 + )
112 +
113 + self.allow_ll_bf16_gemm = (
114 + self.weight.dtype == torch.bfloat16
115 + and self.out_dtype == torch.float32
116 + and is_available()
117 + )
118 +
119 def set_out_dtype(self, out_dtype: torch.dtype) -> None:
120 """Set output dtype for the router logits after init.
121
@@ -116,10 +133,31 @@ class GateLinear(ReplicatedLinear):
133 ):
134 self.allow_cublas_router_gemm = self.weight.dtype == torch.bfloat16
135
136 + # out_dtype may start as None -> recompute eligibility here
137 + if self.allow_specialized_router_gemm:
138 + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import (
139 + is_available,
140 + )
141 +
142 + self.allow_ll_bf16_gemm = (
143 + self.weight.dtype == torch.bfloat16
144 + and out_dtype == torch.float32
145 + and is_available()
146 + )
147 +
148 def forward(
149 self, x: torch.Tensor
150 ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
122 - # Tier 1: DSV3 specialized kernel
151 + # Tier 1: cuteDSL ll_bf16_gemm (SM90+, any dims)
152 + if self.allow_ll_bf16_gemm and x.shape[0] <= 16 and x.dtype == torch.bfloat16:
153 + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import (
154 + ll_bf16_gemm,
155 + )
156 +
157 + output = ll_bf16_gemm(x, self.weight)
158 + return output, None
159 +
160 + # Tier 2: DSV3 specialized kernel (fallback for when cuteDSL unavailable)
161 if self.allow_dsv3_router_gemm and x.shape[0] <= self._dsv3_max_batch:
162 output = ops.dsv3_router_gemm(
163 hidden_states=x,
@@ -128,7 +166,7 @@ class GateLinear(ReplicatedLinear):
166 )
167 return output, None
168
131 - # Tier 2: fp32 specialized kernel (H=3072, E=256, M<=32)
169 + # Tier 3: fp32 specialized kernel (H=3072, E=256, M<=32)
170 # Dispatch is wrapped in a custom op so that torch.compile/CUDA-graph
171 # capture does not freeze the runtime num_tokens branch.
172 if self.allow_fp32_router_gemm and x.dtype in (
@@ -138,12 +176,12 @@ class GateLinear(ReplicatedLinear):
176 output = torch.ops.vllm.fp32_router_gemm_dispatch(x, self.weight)
177 return output, None
178
141 - # Tier 3: cuBLAS bf16→fp32
179 + # Tier 4: cuBLAS bf16→fp32
180 if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16:
181 output = torch.mm(x, self.weight.T, out_dtype=torch.float32)
182 return output, None
183
146 - # Tier 4: F.linear (ReplicatedLinear)
184 + # Tier 5: F.linear (ReplicatedLinear)
185 if self.out_dtype is not None and x.dtype != self.weight.dtype:
186 x = x.to(self.weight.dtype)
187 output, output_bias = super().forward(x)
vllm/model_executor/warmup/kernel_warmup.py
+28
@@ -42,6 +42,31 @@ if TYPE_CHECKING:
42
43 logger = init_logger(__name__)
44
45 +_LL_BF16_WARMUP_MODEL_SHAPES: tuple[tuple[int, int], ...] = (
46 + (7168, 256), # DSV3
47 + (7168, 384), # DSV4-Pro
48 + (14400, 256), # DSV4-Flash
49 +)
50 +_LL_BF16_WARMUP_M_RANGE = range(1, 17)
51 +
52 +
53 +def _warmup_ll_bf16_router_gemm() -> None:
54 + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import (
55 + is_available as is_ll_bf16_gemm_available,
56 + )
57 + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import (
58 + ll_bf16_gemm_kernel,
59 + )
60 +
61 + if not is_ll_bf16_gemm_available():
62 + return
63 +
64 + logger.info("Warming up ll_bf16 router GEMM kernels.")
65 + ll_bf16_gemm_kernel.warmup(
66 + shapes=_LL_BF16_WARMUP_MODEL_SHAPES,
67 + m_values=_LL_BF16_WARMUP_M_RANGE,
68 + )
69 +
70
71 def kernel_warmup(worker: "Worker"):
72 from vllm.model_executor.warmup.minimax_m3_msa_warmup import (
@@ -94,6 +119,9 @@ def kernel_warmup(worker: "Worker"):
119 elif has_flashinfer() and current_platform.has_device_capability(90):
120 flashinfer_autotune(worker.model_runner)
121
122 + if current_platform.has_device_capability(90):
123 + _warmup_ll_bf16_router_gemm()
124 +
125 # FlashInfer attention warmup
126 # Only warmup if the model has FlashInfer attention groups
127 # and is not a pooling model