@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c810937573

[Kernel] Extend CuTe DSL skinny GEMM to GLM-5.2 (#49791)

Signed-off-by: Peiyuan Zhou <peiyuanzhou1994@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

xiaozhoupy committed Aug 4, 2026 at 07:49 UTC c8109375733e6b788e73c7cc1ed2004234834392
7 files changed +619 -51
csrc/libtorch_stable/dsv3_fused_a_gemm.cu
+42 -24
@@ -391,7 +391,7 @@ struct MmaComputer {
391 static constexpr int n_iter_cnt =
392 (tile_n + 7) /
393 8; // Possible to have non-1 n_iter_cnt for ab_swap m16 case.
394 - static_assert(m_iter_cnt == 1);
394 + static_assert(m_iter_cnt == 1 || m_iter_cnt == 2);
395 static_assert(n_iter_cnt == 1 || n_iter_cnt == 2);
396
397 __device__ MmaComputer(bf16_t* gmem_c_local_, bf16_t* smem_a_,
@@ -417,12 +417,15 @@ struct MmaComputer {
417 __device__ void prepare() {
418 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
419 #pragma unroll
420 - for (int i = 0; i < k_phase_cnt; i++) {
421 - int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256;
422 - int m_idx = linear_idx % tile_m;
423 - int k_idx = linear_idx / tile_m + warp_k_offset_in_tile_k;
424 - k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
425 - a_smem_offsets[0][i] = m_idx * tile_k + k_idx;
420 + for (int m = 0; m < m_iter_cnt; m++) {
421 + #pragma unroll
422 + for (int i = 0; i < k_phase_cnt; i++) {
423 + int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256;
424 + int m_idx = linear_idx % 16 + m * 16;
425 + int k_idx = linear_idx / 16 + warp_k_offset_in_tile_k;
426 + k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
427 + a_smem_offsets[m][i] = m_idx * tile_k + k_idx;
428 + }
429 }
430 #pragma unroll
431 for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
@@ -446,11 +449,14 @@ struct MmaComputer {
449 wait_barrier(smem_barrier + 0 + stage_idx * 2, phase_bit);
450
451 #pragma unroll
449 - for (int i = 0; i < k_phase_cnt; i++) {
450 - int smem_offset = a_smem_offsets[0][i];
451 - bf16_t* smem_ptr_this_iter =
452 - smem_a + stage_idx * tile_m * tile_k + smem_offset;
453 - ldsm_x4(smem_ptr_this_iter, reinterpret_cast<uint32_t*>(a_reg[0][i]));
452 + for (int m = 0; m < m_iter_cnt; m++) {
453 + #pragma unroll
454 + for (int i = 0; i < k_phase_cnt; i++) {
455 + int smem_offset = a_smem_offsets[m][i];
456 + bf16_t* smem_ptr_this_iter =
457 + smem_a + stage_idx * tile_m * tile_k + smem_offset;
458 + ldsm_x4(smem_ptr_this_iter, reinterpret_cast<uint32_t*>(a_reg[m][i]));
459 + }
460 }
461
462 #pragma unroll
@@ -469,9 +475,12 @@ struct MmaComputer {
475 for (int k_iter_idx = 0; k_iter_idx < k_phase_cnt; k_iter_idx++) {
476 #pragma unroll
477 for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
472 - hmma_16_8_16_f32acc_bf16ab(
473 - acc_reg[0][n_iter_idx], a_reg[0][k_iter_idx],
474 - b_reg[n_iter_idx][k_iter_idx], acc_reg[0][n_iter_idx]);
478 + #pragma unroll
479 + for (int m = 0; m < m_iter_cnt; m++) {
480 + hmma_16_8_16_f32acc_bf16ab(
481 + acc_reg[m][n_iter_idx], a_reg[m][k_iter_idx],
482 + b_reg[n_iter_idx][k_iter_idx], acc_reg[m][n_iter_idx]);
483 + }
484 }
485 }
486 ::arrive_barrier(smem_barrier + 1 + stage_idx * 2);
@@ -486,14 +495,14 @@ struct MmaComputer {
495 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
496 asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
497 // reorganize the acc_reg
489 - constexpr int thread_m = 2;
498 + constexpr int thread_m = 2 * m_iter_cnt;
499 constexpr int thread_n = 2 * n_iter_cnt;
500 constexpr int cta_mma_n = n_iter_cnt * 8;
501 float acc_reg_reorg[thread_m][thread_n];
502
503 for (int i = 0; i < thread_m; i++) {
504 for (int j = 0; j < thread_n; j++) {
496 - acc_reg_reorg[i][j] = acc_reg[0][j / 2][(j % 2) + (i * 2)];
505 + acc_reg_reorg[i][j] = acc_reg[i / 2][j / 2][(j % 2) + (i % 2) * 2];
506 }
507 }
508
@@ -514,7 +523,8 @@ struct MmaComputer {
523 for (int m_idx_thread = 0; m_idx_thread < thread_m; m_idx_thread++) {
524 #pragma unroll
525 for (int n_idx_thread = 0; n_idx_thread < thread_n; n_idx_thread++) {
517 - int m_idx = (lane_idx / 4) + m_idx_thread * 8;
526 + int m_idx =
527 + (lane_idx / 4) + (m_idx_thread % 2) * 8 + (m_idx_thread / 2) * 16;
528 int n_idx =
529 ((lane_idx % 4) * 2) + (n_idx_thread % 2) + (n_idx_thread / 2) * 8;
530 smem_c[cosize_smem_c * warp_idx + smem_c_index_func(m_idx, n_idx)] =
@@ -587,7 +597,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
597 static_assert(
598 tile_k == 128 || tile_k == 256 || tile_k == 512 ||
599 tile_k == 1024); // tile_k must be larger than 64 since 4 warp splitK.
590 - static_assert(tile_m == 16);
600 + static_assert(tile_m == 16 || tile_m == 32);
601 constexpr int g2s_vec_bytes = 16;
602 constexpr int a_elem_bytes = 2;
603 constexpr int b_elem_bytes = 2;
@@ -647,7 +657,8 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
657 #endif
658 }
659
650 -template <typename T, int kHdIn, int kHdOut, int kTileN, int kTileK = 256>
660 +template <typename T, int kHdIn, int kHdOut, int kTileN, int kTileK = 256,
661 + int kTileM = 16>
662 void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
663 cudaStream_t const stream, bool enable_pdl) {
664 constexpr int gemm_m = kHdOut;
@@ -655,7 +666,7 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
666 constexpr int gemm_k = kHdIn;
667 constexpr int batch_size = 1;
668 std::swap(mat_a, mat_b);
658 - constexpr int tile_m = 16;
669 + constexpr int tile_m = kTileM;
670 constexpr int tile_n = kTileN;
671 constexpr int tile_k = kTileK;
672 constexpr int max_stage_cnt =
@@ -695,15 +706,15 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
706 output, mat_a, mat_b, gemm_n);
707 }
708
698 -template <typename T, int kHdIn, int kHdOut, int kTileK = 256>
709 +template <typename T, int kHdIn, int kHdOut, int kTileK = 256, int kTileM = 16>
710 void invokeFusedAGemmForTokens(T* output, T const* mat_a, T const* mat_b,
711 int num_tokens, cudaStream_t const stream,
712 bool enable_pdl) {
713 if (num_tokens <= 8) {
703 - invokeFusedAGemm<T, kHdIn, kHdOut, 8, kTileK>(
714 + invokeFusedAGemm<T, kHdIn, kHdOut, 8, kTileK, kTileM>(
715 output, mat_a, mat_b, num_tokens, stream, enable_pdl);
716 } else {
706 - invokeFusedAGemm<T, kHdIn, kHdOut, 16, kTileK>(
717 + invokeFusedAGemm<T, kHdIn, kHdOut, 16, kTileK, kTileM>(
718 output, mat_a, mat_b, num_tokens, stream, enable_pdl);
719 }
720 }
@@ -782,6 +793,13 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
793 DISPATCH_DSV3_SHAPE(7168, 3216)
794 DISPATCH_DSV3_SHAPE(7168, 4224)
795
796 + if (hd_in == 6144 && hd_out == 2624) {
797 + invokeFusedAGemmForTokens<__nv_bfloat16, 6144, 2624, 256, 32>(
798 + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
799 + return;
800 + }
801 + DISPATCH_DSV3_SHAPE(2048, 2048)
802 +
803 #ifdef VLLM_K3_BENCH_SHAPES
804 // The selector routes these shapes to CuTe or the default GEMM, so they are
805 // never reached in production. They are compiled only for offline
tests/kernels/test_bf16_skinny_gemm.py
+255 -6
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Tests for the Kimi-K3 SM103 decode GEMM selector (shape-only dispatch)."""
3 +"""Tests for BF16 skinny GEMMs and the Kimi-K3 SM103 selector."""
4
5 from pathlib import Path
6 from types import SimpleNamespace
@@ -11,7 +11,10 @@ import regex as re
11 import torch
12 from torch import nn
13
14 -from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import SkinnyGemmConfig
14 +from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
15 + SkinnyGemmConfig,
16 +)
17 +from vllm.models.deepseek_v32.nvidia import glm52_low_latency_gemm as glm52_gemm
18 from vllm.models.kimi_k3.nvidia import low_latency_gemm as k3_gemm
19 from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS
20
@@ -60,6 +63,12 @@ RESIDUAL_CUTE_CASES = [
63 for num_tokens, _ in spec.residual_configs
64 ]
65
66 +GLM_CUTE_CASES = [
67 + (spec, config)
68 + for spec in glm52_gemm.GLM52_PROJECTIONS.values()
69 + for _, config in spec.cute_configs
70 +]
71 +
72 EXPECTED_CUTE_CONFIGS = {
73 (3072, 7168, 1): (224, 3, 4, 8),
74 (3072, 7168, 2): (128, 3, 2, 8),
@@ -164,13 +173,21 @@ def test_every_dsv3_routed_shape_is_instantiated() -> None:
173 )
174 } | {
175 (int(hd_in), int(hd_out))
167 - for hd_in, hd_out in re.findall(r"hd_in == (\d+) && hd_out == (\d+)", explicit)
176 + for hd_in, hd_out in re.findall(
177 + r"hd_in == (\d+) && hd_out == (\d+)",
178 + production_macros + explicit,
179 + )
180 }
181 assert compiled, "failed to parse the dispatch list"
182
183 + specs = [
184 + *KIMI_K3_PROJECTIONS.values(),
185 + glm52_gemm.GLM52_QKV_A_PROJECTION,
186 + glm52_gemm.GLM52_Q_B_PROJECTION,
187 + ]
188 missing = sorted(
189 (spec.n, spec.k)
173 - for spec in KIMI_K3_PROJECTIONS.values()
190 + for spec in specs
191 if spec.dsv3_tokens and (spec.k, spec.n) not in compiled
192 )
193 assert not missing, (
@@ -190,12 +207,17 @@ def test_packed_row_major_rejects_single_row_slice() -> None:
207
208
209 def test_cute_configs_match_measured_table() -> None:
193 - actual = {
194 - (spec.n, spec.k, num_tokens): _config_tuple(config)
210 + configs = [
211 + (spec.n, spec.k, num_tokens, config)
212 for spec in k3_gemm.KIMI_K3_PROJECTIONS.values()
213 for num_tokens, config in spec.cute_configs
214 + ]
215 + actual = {
216 + (n, k, num_tokens): _config_tuple(config)
217 + for n, k, num_tokens, config in configs
218 }
219 assert actual == EXPECTED_CUTE_CONFIGS
220 + assert all(config.static_k is None for *_, config in configs)
221
222
223 def test_residual_cute_configs_match_measured_table() -> None:
@@ -207,6 +229,134 @@ def test_residual_cute_configs_match_measured_table() -> None:
229 assert actual == EXPECTED_RESIDUAL_CUTE_CONFIGS
230
231
232 +def test_glm52_projection_plans_are_separate() -> None:
233 + qkv_a = glm52_gemm.GLM52_QKV_A_PROJECTION
234 + q_b = glm52_gemm.GLM52_Q_B_PROJECTION
235 +
236 + qkv_a_plan = qkv_a.build_plan()
237 + q_b_plan = q_b.build_plan()
238 +
239 + assert (qkv_a.n, qkv_a.k, set(qkv_a_plan)) == (
240 + 2624,
241 + 6144,
242 + set(range(1, 17)),
243 + )
244 + assert (q_b.n, q_b.k, set(q_b_plan)) == (
245 + 2048,
246 + 2048,
247 + set(range(1, 17)),
248 + )
249 + assert {
250 + num_tokens
251 + for num_tokens, (backend, _) in qkv_a_plan.items()
252 + if backend == "cute"
253 + } == {1, 2}
254 + assert {
255 + num_tokens
256 + for num_tokens, (backend, _) in qkv_a_plan.items()
257 + if backend == "dsv3_fused_a"
258 + } == set(range(3, 17))
259 + assert {
260 + num_tokens for num_tokens, (backend, _) in q_b_plan.items() if backend == "cute"
261 + } == {1, 2}
262 + assert {
263 + num_tokens
264 + for num_tokens, (backend, _) in q_b_plan.items()
265 + if backend == "dsv3_fused_a"
266 + } == set(range(3, 17))
267 +
268 + eh = glm52_gemm.GLM52_EH_PROJECTION
269 + eh_plan = eh.build_plan()
270 + assert (eh.n, eh.k) == (6144, 12288)
271 + # The MTP eh_proj has no dsv3 winners; M >= 4 falls back to cuBLAS.
272 + assert set(eh_plan) == {1, 2, 3}
273 + assert all(backend == "cute" for backend, _ in eh_plan.values())
274 +
275 +
276 +def test_glm52_layout_rejects_nonpacked_single_row_view() -> None:
277 + single_row = torch.empty(1, 144)[:, :128]
278 + multiple_rows = torch.empty(2, 144)[:, :128]
279 +
280 + assert single_row.stride() == multiple_rows.stride() == (144, 1)
281 + assert not glm52_gemm._is_supported_row_major(single_row)
282 + assert not glm52_gemm._is_supported_row_major(multiple_rows)
283 +
284 +
285 +def test_glm52_installer_maps_only_selected_unquantized_shapes(
286 + monkeypatch: pytest.MonkeyPatch,
287 +) -> None:
288 + class FakeLinearBase(nn.Module):
289 + def __init__(
290 + self,
291 + n: int,
292 + k: int,
293 + quant_method: object,
294 + ) -> None:
295 + super().__init__()
296 + self.weight = nn.Parameter(
297 + torch.empty(
298 + n,
299 + k,
300 + dtype=torch.bfloat16,
301 + device="meta",
302 + )
303 + )
304 + self.quant_method = quant_method
305 +
306 + qkv_a = glm52_gemm.GLM52_QKV_A_PROJECTION
307 + q_b = glm52_gemm.GLM52_Q_B_PROJECTION
308 + root = nn.Module()
309 + root.attn = nn.Module()
310 + root.attn.fused_qkv_a_proj = FakeLinearBase(
311 + qkv_a.n, qkv_a.k, glm52_gemm.UnquantizedLinearMethod()
312 + )
313 + root.attn.q_b_proj = FakeLinearBase(
314 + q_b.n, q_b.k, glm52_gemm.UnquantizedLinearMethod()
315 + )
316 + root.same_shape_other_name = FakeLinearBase(
317 + qkv_a.n, qkv_a.k, glm52_gemm.UnquantizedLinearMethod()
318 + )
319 + root.quantized = nn.Module()
320 + quantized_method = object()
321 + root.quantized.q_b_proj = FakeLinearBase(q_b.n, q_b.k, quantized_method)
322 + root.wrong_shape = nn.Module()
323 + root.wrong_shape.fused_qkv_a_proj = FakeLinearBase(
324 + qkv_a.n + 1, qkv_a.k, glm52_gemm.UnquantizedLinearMethod()
325 + )
326 + monkeypatch.setattr(glm52_gemm, "LinearBase", FakeLinearBase)
327 + monkeypatch.setattr(glm52_gemm, "_is_sm103", lambda: True)
328 + monkeypatch.setattr(
329 + glm52_gemm.shape_dynamic_skinny_gemm,
330 + "is_available",
331 + lambda: False,
332 + )
333 +
334 + glm52_gemm.enable_glm52_low_latency_gemm(
335 + root,
336 + torch.bfloat16,
337 + )
338 +
339 + assert isinstance(
340 + root.attn.fused_qkv_a_proj.quant_method,
341 + glm52_gemm.GLM52LowLatencyLinearMethod,
342 + )
343 + assert isinstance(
344 + root.attn.q_b_proj.quant_method,
345 + glm52_gemm.GLM52LowLatencyLinearMethod,
346 + )
347 + assert root.attn.fused_qkv_a_proj.quant_method._plan == qkv_a.build_plan()
348 + assert root.attn.q_b_proj.quant_method._plan == q_b.build_plan()
349 + assert isinstance(
350 + root.same_shape_other_name.quant_method,
351 + glm52_gemm.GLM52LowLatencyLinearMethod,
352 + )
353 + assert root.quantized.q_b_proj.quant_method is quantized_method
354 + assert (
355 + type(root.wrong_shape.fused_qkv_a_proj.quant_method)
356 + is glm52_gemm.UnquantizedLinearMethod
357 + )
358 +
359 +
360 @pytest.mark.parametrize("key", EXPECTED_SELECTIONS)
361 def test_sm103_selector_table(key: tuple[int, int]) -> None:
362 n, k = key
@@ -361,6 +511,76 @@ def _require_sm103_and_cute() -> None:
511 pytest.skip("CuTe DSL is not available")
512
513
514 +@pytest.mark.parametrize("spec,config", GLM_CUTE_CASES)
515 +def test_glm_cute_selected_shapes(
516 + spec: glm52_gemm.GLM52ProjectionSpec,
517 + config: SkinnyGemmConfig,
518 +) -> None:
519 + _require_sm103_and_cute()
520 + torch.manual_seed(42)
521 + x = torch.randn(
522 + config.num_rows,
523 + spec.k,
524 + dtype=torch.bfloat16,
525 + device="cuda",
526 + )
527 + weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda")
528 + plan = spec.build_plan()
529 +
530 + output = glm52_gemm.run_glm52_plan(plan, x, weight)
531 +
532 + assert output is not None
533 + reference = x.float() @ weight.float().t()
534 + torch.testing.assert_close(output.float(), reference, rtol=2e-2, atol=2e-1)
535 +
536 +
537 +def test_glm52_q_b_nonpacked_single_row_falls_back() -> None:
538 + _require_sm103_and_cute()
539 + spec = glm52_gemm.GLM52_Q_B_PROJECTION
540 + storage = torch.randn(
541 + 1,
542 + glm52_gemm.GLM52_QKV_A_PROJECTION.n,
543 + dtype=torch.bfloat16,
544 + device="cuda",
545 + )
546 + x = storage[:, : spec.k]
547 + weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda")
548 +
549 + assert x.stride() == (glm52_gemm.GLM52_QKV_A_PROJECTION.n, 1)
550 + assert not glm52_gemm._runtime_ok(x, weight)
551 + output = glm52_gemm.run_glm52_plan(spec.build_plan(), x, weight)
552 +
553 + assert output is None
554 +
555 +
556 +@pytest.mark.parametrize("spec,config", GLM_CUTE_CASES)
557 +def test_glm_cute_selected_shapes_cuda_graph_capture(
558 + spec: glm52_gemm.GLM52ProjectionSpec,
559 + config: SkinnyGemmConfig,
560 +) -> None:
561 + _require_sm103_and_cute()
562 + x = torch.randn(
563 + config.num_rows,
564 + spec.k,
565 + dtype=torch.bfloat16,
566 + device="cuda",
567 + )
568 + weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda")
569 + plan = spec.build_plan()
570 + glm52_gemm.run_glm52_plan(plan, x, weight)
571 + torch.accelerator.synchronize()
572 +
573 + graph = torch.cuda.CUDAGraph()
574 + with torch.cuda.graph(graph):
575 + output = glm52_gemm.run_glm52_plan(plan, x, weight)
576 + graph.replay()
577 + torch.accelerator.synchronize()
578 +
579 + assert output is not None
580 + reference = x.float() @ weight.float().t()
581 + torch.testing.assert_close(output.float(), reference, rtol=2e-2, atol=2e-1)
582 +
583 +
584 @pytest.mark.parametrize("n,k,num_tokens", CUTE_CASES)
585 def test_cute_selected_shapes(n: int, k: int, num_tokens: int) -> None:
586 _require_sm103_and_cute()
@@ -393,6 +613,15 @@ DSV3_CASES = sorted(
613 for num_tokens in _dsv3_probe_tokens(spec.dsv3_tokens)
614 )
615
616 +GLM_DSV3_CASES = [
617 + (num_tokens, spec)
618 + for spec in (
619 + glm52_gemm.GLM52_QKV_A_PROJECTION,
620 + glm52_gemm.GLM52_Q_B_PROJECTION,
621 + )
622 + for num_tokens in sorted(_dsv3_probe_tokens(spec.dsv3_tokens))
623 +]
624 +
625
626 @pytest.mark.parametrize("num_tokens,n,k", DSV3_CASES)
627 def test_dsv3_selected_shapes(num_tokens: int, n: int, k: int) -> None:
@@ -413,6 +642,26 @@ def test_dsv3_selected_shapes(num_tokens: int, n: int, k: int) -> None:
642 assert cosine > 0.999
643
644
645 +@pytest.mark.parametrize("num_tokens,spec", GLM_DSV3_CASES)
646 +def test_glm_dsv3_selected_shapes(
647 + num_tokens: int,
648 + spec: glm52_gemm.GLM52ProjectionSpec,
649 +) -> None:
650 + _require_sm103_and_dsv3()
651 + torch.manual_seed(42)
652 + x = torch.randn(num_tokens, spec.k, dtype=torch.bfloat16, device="cuda")
653 + weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda")
654 +
655 + output = glm52_gemm.run_glm52_plan(spec.build_plan(), x, weight)
656 +
657 + assert output is not None
658 + reference = torch.nn.functional.linear(x, weight)
659 + cosine = torch.nn.functional.cosine_similarity(
660 + output.float().flatten(), reference.float().flatten(), dim=0
661 + ).item()
662 + assert cosine > 0.999
663 +
664 +
665 def test_nonpacked_single_token_dsv3_falls_back() -> None:
666 _require_sm103_and_dsv3()
667 n, k = 1536, 128
vllm/model_executor/kernels/linear/cute_dsl/_skinny_gemm.py
+92 -19
@@ -8,6 +8,8 @@ import cutlass.cute as cute
8 from cuda.bindings.driver import CUstream
9 from cutlass import const_expr
10
11 +_PREFETCH_TILES = 2
12 +
13
14 class CuteSkinnyGemm:
15 """Shape-dynamic low-latency GEMM for small token counts.
@@ -15,8 +17,9 @@ class CuteSkinnyGemm:
17 Computes ``C[M, N] = A[M, K] @ B[N, K].T + residual`` with BF16 or FP16
18 inputs, FP32 accumulators, and an output matching the input dtype. The
19 residual term is optional and is added before the output conversion. N and
18 - K are runtime values. The tiny M dimension is fully unrolled alongside a
19 - small set of tuning parameters.
20 + K are runtime values unless the configuration specializes K at compile
21 + time. The tiny M dimension is fully unrolled alongside a small set of
22 + tuning parameters.
23 """
24
25 def __init__(
@@ -30,9 +33,15 @@ class CuteSkinnyGemm:
33 k_unroll: int = 1,
34 has_residual: bool = False,
35 use_pdl: bool = False,
36 + static_k: int | None = None,
37 ) -> None:
38 if block_size % cute.arch.WARP_SIZE != 0:
39 raise ValueError("block_size must be a multiple of the warp size")
40 + k_tile_size = block_size * vector_width
41 + if static_k is not None and (
42 + static_k % k_tile_size or static_k < _PREFETCH_TILES * k_tile_size
43 + ):
44 + raise ValueError("static K must contain at least two complete tiles")
45 self.element_type = element_type
46 self.num_rows = num_rows
47 self.block_size = block_size
@@ -41,6 +50,7 @@ class CuteSkinnyGemm:
50 self.k_unroll = k_unroll
51 self.has_residual = has_residual
52 self.use_pdl = use_pdl
53 + self.static_k = static_k
54 self.num_warps = block_size // cute.arch.WARP_SIZE
55
56 @cute.jit
@@ -66,7 +76,22 @@ class CuteSkinnyGemm:
76 num_bits_per_copy=self.vector_width * self.element_type.width,
77 load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING,
78 )
69 - self.kernel(gA, gB, gResidual, gC, k, copy_a, copy_b).launch(
79 + copy_b_prefetch = cute.make_copy_atom(
80 + cute.nvgpu.CopyG2ROp(),
81 + self.element_type,
82 + num_bits_per_copy=self.vector_width * self.element_type.width,
83 + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
84 + )
85 + self.kernel(
86 + gA,
87 + gB,
88 + gResidual,
89 + gC,
90 + k,
91 + copy_a,
92 + copy_b,
93 + copy_b_prefetch,
94 + ).launch(
95 grid=[cute.ceil_div(n, self.outputs_per_block), 1, 1],
96 block=[self.block_size, 1, 1],
97 smem=self.num_rows * self.outputs_per_block * self.num_warps * 4,
@@ -85,6 +110,7 @@ class CuteSkinnyGemm:
110 k_extent: cutlass.Int32,
111 copy_a: cute.CopyAtom,
112 copy_b: cute.CopyAtom,
113 + copy_b_prefetch: cute.CopyAtom,
114 ) -> None:
115 tidx, _, _ = cute.arch.thread_idx()
116 block_idx, _, _ = cute.arch.block_idx()
@@ -102,12 +128,8 @@ class CuteSkinnyGemm:
128 acc = cute.make_rmem_tensor(acc_layout, cutlass.Float32)
129 acc.fill(0.0)
130
105 - if const_expr(self.use_pdl):
106 - cute.arch.griddepcontrol_wait()
107 -
131 n_base = block_idx * outputs_per_block
132 k_tile_size: cutlass.Constexpr = block_size * vector_width
110 - num_k_tiles = k_extent // k_tile_size
133
134 gA_vec = cute.logical_divide(gA, (None, vector_width))
135 gB_vec = cute.logical_divide(gB, (None, vector_width))
@@ -127,21 +149,72 @@ class CuteSkinnyGemm:
149 self.element_type,
150 )
151
130 - for k_tile in cutlass.range(num_k_tiles, unroll=self.k_unroll):
131 - for mi in cutlass.range_constexpr(num_rows):
132 - cute.copy(copy_a, tA[mi, None, k_tile], a_regs[mi, None])
152 + if const_expr(self.static_k is not None):
153 + num_k_tiles: cutlass.Constexpr = self.static_k // k_tile_size
154 + prefetched_b = cute.make_rmem_tensor(
155 + cute.make_layout(
156 + (_PREFETCH_TILES, outputs_per_block, vector_width),
157 + stride=(outputs_per_block * vector_width, vector_width, 1),
158 + ),
159 + self.element_type,
160 + )
161 + for k_tile in cutlass.range_constexpr(_PREFETCH_TILES):
162 + for ni in cutlass.range_constexpr(outputs_per_block):
163 + tB_static = tB_all[n_base + ni, (None, (tidx, None))]
164 + cute.copy(
165 + copy_b_prefetch,
166 + tB_static[None, k_tile],
167 + prefetched_b[k_tile, ni, None],
168 + )
169
134 - for ni in cutlass.range_constexpr(outputs_per_block):
135 - n_idx = n_base + ni
136 - tB = tB_all[n_idx, (None, (tidx, None))]
137 - cute.copy(copy_b, tB[None, k_tile], b_regs[ni, None])
170 + if const_expr(self.use_pdl):
171 + cute.arch.griddepcontrol_wait()
172
139 - for vi in cutlass.range_constexpr(vector_width):
140 - for mi in cutlass.range_constexpr(num_rows):
173 + for k_tile in cutlass.range_constexpr(num_k_tiles):
174 + if const_expr(k_tile < _PREFETCH_TILES):
175 + b_f32 = prefetched_b[k_tile, None, None].load().to(cutlass.Float32)
176 + else:
177 for ni in cutlass.range_constexpr(outputs_per_block):
142 - acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to(
143 - cutlass.Float32
144 - ) * b_regs[ni, vi].to(cutlass.Float32)
178 + tB_static = tB_all[n_base + ni, (None, (tidx, None))]
179 + cute.copy(
180 + copy_b,
181 + tB_static[None, k_tile],
182 + b_regs[ni, None],
183 + )
184 + b_f32 = b_regs.load().to(cutlass.Float32)
185 +
186 + for mi in cutlass.range_constexpr(num_rows):
187 + cute.copy(copy_a, tA[mi, None, k_tile], a_regs[mi, None])
188 + a_f32 = a_regs.load().to(cutlass.Float32)
189 +
190 + for vi in cutlass.range_constexpr(vector_width):
191 + for mi in cutlass.range_constexpr(num_rows):
192 + for ni in cutlass.range_constexpr(outputs_per_block):
193 + acc[mi, ni] += a_f32[mi, vi] * b_f32[ni, vi]
194 + else:
195 + if const_expr(self.use_pdl):
196 + cute.arch.griddepcontrol_wait()
197 +
198 + num_k_tiles = k_extent // k_tile_size
199 + for k_tile in cutlass.range(num_k_tiles, unroll=self.k_unroll):
200 + for mi in cutlass.range_constexpr(num_rows):
201 + cute.copy(copy_a, tA[mi, None, k_tile], a_regs[mi, None])
202 +
203 + for ni in cutlass.range_constexpr(outputs_per_block):
204 + n_idx = n_base + ni
205 + tB_dynamic = tB_all[n_idx, (None, (tidx, None))]
206 + cute.copy(
207 + copy_b,
208 + tB_dynamic[None, k_tile],
209 + b_regs[ni, None],
210 + )
211 +
212 + for vi in cutlass.range_constexpr(vector_width):
213 + for mi in cutlass.range_constexpr(num_rows):
214 + for ni in cutlass.range_constexpr(outputs_per_block):
215 + acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to(
216 + cutlass.Float32
217 + ) * b_regs[ni, vi].to(cutlass.Float32)
218
219 for mi in cutlass.range_constexpr(num_rows):
220 for ni in cutlass.range_constexpr(outputs_per_block):
vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm.py
+9 -1
@@ -23,6 +23,7 @@ class SkinnyGemmConfig:
23 outputs_per_block: int
24 k_unroll: int = 1
25 vector_width: int = 8
26 + static_k: int | None = None
27
28
29 class ShapeDynamicSkinnyGemm:
@@ -112,7 +113,11 @@ class ShapeDynamicSkinnyGemm:
113
114 element_type = self._cutlass_dtype(dtype)
115 n = cute.sym_int(divisibility=config.outputs_per_block)
115 - k = cute.sym_int(divisibility=config.block_size * config.vector_width)
116 + k = (
117 + config.static_k
118 + if config.static_k is not None
119 + else cute.sym_int(divisibility=config.block_size * config.vector_width)
120 + )
121 a = make_fake_tensor(
122 element_type,
123 (config.num_rows, k),
@@ -134,6 +139,7 @@ class ShapeDynamicSkinnyGemm:
139 k_unroll=config.k_unroll,
140 has_residual=has_residual,
141 use_pdl=self._use_pdl(),
142 + static_k=config.static_k,
143 )
144 self._compiled[(dtype, config, has_residual)] = cute.compile(
145 kernel,
@@ -239,6 +245,8 @@ class ShapeDynamicSkinnyGemm:
245 raise ValueError(
246 "K must be divisible by block_size * vector_width for this config"
247 )
248 + if config.static_k is not None and a.shape[1] != config.static_k:
249 + raise ValueError("input K must match config static_k")
250 has_residual = residual is not None
251 cache_key = (a.dtype, config, has_residual)
252 if cache_key not in self._compiled:
vllm/models/deepseek_v32/nvidia/glm52_low_latency_gemm.py new
+199
@@ -0,0 +1,199 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""GLM-5.2 decode GEMM selection for unquantized BF16 on SM103."""
4 +
5 +from __future__ import annotations
6 +
7 +from dataclasses import dataclass
8 +from typing import Literal
9 +
10 +import torch
11 +from torch import nn
12 +
13 +import vllm.envs as envs
14 +from vllm import _custom_ops as ops
15 +from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
16 + SkinnyGemmConfig,
17 + shape_dynamic_skinny_gemm,
18 +)
19 +from vllm.model_executor.layers.linear import (
20 + LinearBase,
21 + UnquantizedLinearMethod,
22 +)
23 +from vllm.platforms import current_platform
24 +
25 +Backend = Literal["cute", "dsv3_fused_a"]
26 +ResolvedCall = tuple[Backend, SkinnyGemmConfig | None]
27 +
28 +
29 +@dataclass(frozen=True, slots=True)
30 +class GLM52ProjectionSpec:
31 + n: int
32 + k: int
33 + cute_configs: tuple[tuple[int, SkinnyGemmConfig], ...]
34 + dsv3_tokens: frozenset[int] = frozenset()
35 +
36 + def build_plan(self) -> dict[int, ResolvedCall]:
37 + plan: dict[int, ResolvedCall] = {
38 + num_tokens: ("cute", config) for num_tokens, config in self.cute_configs
39 + }
40 + plan.update(
41 + (num_tokens, ("dsv3_fused_a", None)) for num_tokens in self.dsv3_tokens
42 + )
43 + return plan
44 +
45 +
46 +GLM52_QKV_A_PROJECTION = GLM52ProjectionSpec(
47 + n=2624,
48 + k=6144,
49 + cute_configs=(
50 + (1, SkinnyGemmConfig(1, 128, 4, static_k=6144)),
51 + (2, SkinnyGemmConfig(2, 128, 2)),
52 + ),
53 + dsv3_tokens=frozenset(range(3, 17)),
54 +)
55 +
56 +GLM52_Q_B_PROJECTION = GLM52ProjectionSpec(
57 + n=2048,
58 + k=2048,
59 + cute_configs=(
60 + (1, SkinnyGemmConfig(1, 128, 4, static_k=2048)),
61 + (2, SkinnyGemmConfig(2, 64, 2, k_unroll=2)),
62 + ),
63 + dsv3_tokens=frozenset(range(3, 17)),
64 +)
65 +
66 +# The MTP eh_proj is a plain nn.Linear, so it gets its plan through
67 +# build_glm52_plan rather than a quant_method swap. cuBLAS wins from M=4.
68 +GLM52_EH_PROJECTION = GLM52ProjectionSpec(
69 + n=6144,
70 + k=12288,
71 + cute_configs=(
72 + (1, SkinnyGemmConfig(1, 256, 2, vector_width=4, static_k=12288)),
73 + (2, SkinnyGemmConfig(2, 64, 2)),
74 + (3, SkinnyGemmConfig(3, 64, 2)),
75 + ),
76 +)
77 +
78 +GLM52_PROJECTIONS = {
79 + (spec.n, spec.k): spec
80 + for spec in (
81 + GLM52_QKV_A_PROJECTION,
82 + GLM52_Q_B_PROJECTION,
83 + GLM52_EH_PROJECTION,
84 + )
85 +}
86 +
87 +
88 +def _is_sm103() -> bool:
89 + return current_platform.is_device_capability((10, 3))
90 +
91 +
92 +def _is_supported_row_major(tensor: torch.Tensor) -> bool:
93 + return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1)
94 +
95 +
96 +def _runtime_ok(x: torch.Tensor, weight: torch.Tensor) -> bool:
97 + return (
98 + not envs.VLLM_BATCH_INVARIANT
99 + and _is_supported_row_major(x)
100 + and _is_supported_row_major(weight)
101 + and x.dtype == torch.bfloat16
102 + and weight.dtype == torch.bfloat16
103 + and x.is_cuda
104 + and weight.is_cuda
105 + and x.device == weight.device
106 + and x.shape[1] == weight.shape[1]
107 + )
108 +
109 +
110 +def run_glm52_plan(
111 + plan: dict[int, ResolvedCall] | None,
112 + x: torch.Tensor,
113 + weight: torch.Tensor,
114 +) -> torch.Tensor | None:
115 + if plan is None or not _runtime_ok(x, weight):
116 + return None
117 + entry = plan.get(x.shape[0])
118 + if entry is None:
119 + return None
120 +
121 + backend, config = entry
122 + if backend == "cute":
123 + if not shape_dynamic_skinny_gemm.is_available():
124 + return None
125 + return shape_dynamic_skinny_gemm(x, weight, config)
126 +
127 + if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"):
128 + return None
129 + output = torch.empty(
130 + (x.shape[0], weight.shape[0]),
131 + dtype=x.dtype,
132 + device=x.device,
133 + )
134 + ops.dsv3_fused_a_gemm(output, x, weight.t(), enable_pdl=True)
135 + return output
136 +
137 +
138 +def _request_warmup(dtype: torch.dtype, configs: set[SkinnyGemmConfig]) -> None:
139 + if configs and shape_dynamic_skinny_gemm.is_available():
140 + shape_dynamic_skinny_gemm.request_warmup_configs(dtype, configs)
141 +
142 +
143 +def build_glm52_plan(
144 + weight: torch.Tensor | None, dtype: torch.dtype
145 +) -> dict[int, ResolvedCall] | None:
146 + """Plan for a weight the walk below cannot reach (a plain ``nn.Linear``)."""
147 + if dtype != torch.bfloat16 or not _is_sm103():
148 + return None
149 + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16:
150 + return None
151 + spec = GLM52_PROJECTIONS.get(tuple(weight.shape))
152 + if spec is None:
153 + return None
154 + _request_warmup(dtype, {config for _, config in spec.cute_configs})
155 + return spec.build_plan()
156 +
157 +
158 +class GLM52LowLatencyLinearMethod(UnquantizedLinearMethod):
159 + def __init__(self, plan: dict[int, ResolvedCall]) -> None:
160 + super().__init__()
161 + self._plan = plan
162 +
163 + def apply(
164 + self,
165 + layer: nn.Module,
166 + x: torch.Tensor,
167 + bias: torch.Tensor | None = None,
168 + ) -> torch.Tensor:
169 + if bias is None:
170 + output = run_glm52_plan(self._plan, x, layer.weight)
171 + if output is not None:
172 + return output
173 + return super().apply(layer, x, bias)
174 +
175 +
176 +def enable_glm52_low_latency_gemm(
177 + module: nn.Module,
178 + dtype: torch.dtype,
179 +) -> None:
180 + if dtype != torch.bfloat16 or not _is_sm103():
181 + return
182 +
183 + warmup_configs: set[SkinnyGemmConfig] = set()
184 + for child in module.modules():
185 + if (
186 + not isinstance(child, LinearBase)
187 + or type(child.quant_method) is not UnquantizedLinearMethod
188 + ):
189 + continue
190 + weight = getattr(child, "weight", None)
191 + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16:
192 + continue
193 + spec = GLM52_PROJECTIONS.get(tuple(weight.shape))
194 + if spec is None:
195 + continue
196 + child.quant_method = GLM52LowLatencyLinearMethod(spec.build_plan())
197 + warmup_configs.update(config for _, config in spec.cute_configs)
198 +
199 + _request_warmup(dtype, warmup_configs)
vllm/models/deepseek_v32/nvidia/model.py
+7
@@ -42,6 +42,8 @@ from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_
42 from vllm.models.deepseek_v32.attention import DeepseekV32Attention
43 from vllm.sequence import IntermediateTensors
44
45 +from .glm52_low_latency_gemm import enable_glm52_low_latency_gemm
46 +
47
48 def _all_gather_sp_states(
49 hidden_states: torch.Tensor,
@@ -405,6 +407,11 @@ class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM):
407
408 model_cls = DeepseekV32Model
409
410 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
411 + super().__init__(vllm_config=vllm_config, prefix=prefix)
412 + if self.config.model_type == "glm_moe_dsa":
413 + enable_glm52_low_latency_gemm(self, vllm_config.model_config.dtype)
414 +
415 def set_moe_parameters(self):
416 # Same as the base, but keyed on the MoE block type rather than the
417 # decoder-layer type (DeepseekV32DecoderLayer is a plain nn.Module).
vllm/models/deepseek_v32/nvidia/mtp.py
+15 -1
@@ -42,6 +42,11 @@ from vllm.models.deepseek_v32.common.kernels import fused_eh_norm
42 from vllm.platforms import current_platform
43 from vllm.sequence import IntermediateTensors
44
45 +from .glm52_low_latency_gemm import (
46 + build_glm52_plan,
47 + enable_glm52_low_latency_gemm,
48 + run_glm52_plan,
49 +)
50 from .model import DeepseekV32DecoderLayer
51
52
@@ -56,6 +61,11 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
61 self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
62 self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
63 self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
64 + self._eh_plan = (
65 + build_glm52_plan(self.eh_proj.weight, vllm_config.model_config.dtype)
66 + if config.model_type == "glm_moe_dsa"
67 + else None
68 + )
69
70 topk_indices_buffer = torch.empty(
71 vllm_config.scheduler_config.max_num_batched_tokens,
@@ -91,7 +101,9 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
101 self.hnorm.weight,
102 self.enorm.variance_epsilon,
103 )
94 - hidden_states = self.eh_proj(eh_input)
104 + hidden_states = run_glm52_plan(self._eh_plan, eh_input, self.eh_proj.weight)
105 + if hidden_states is None:
106 + hidden_states = self.eh_proj(eh_input)
107 hidden_states, residual = self.mtp_block(
108 positions=positions, hidden_states=hidden_states, residual=None
109 )
@@ -199,6 +211,8 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
211 self.model = DeepseekV32MultiTokenPredictor(
212 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
213 )
214 + if self.config.model_type == "glm_moe_dsa":
215 + enable_glm52_low_latency_gemm(self, vllm_config.model_config.dtype)
216 self.set_moe_parameters()
217
218 def set_moe_parameters(self):