@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3ba9907a1d

[Kimi-K3] Overlap low-M TP8 KDA projections (#54697)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Yongye Zhu committed Sep 1, 2026 at 20:35 UTC 3ba9907a1db20c085be5c2d3cfbbbd6ec4a182dc
7 files changed +1092 -12
benchmarks/kernels/benchmark_kimi_k3_kda_projection.py new
+221
@@ -0,0 +1,221 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Compare Kimi-K3 TP8 KDA sequential and multistream projections on B300."""
4 +
5 +from __future__ import annotations
6 +
7 +import argparse
8 +import statistics
9 +from collections.abc import Callable
10 +
11 +import torch
12 +
13 +from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
14 + shape_dynamic_skinny_gemm,
15 +)
16 +from vllm.models.kimi_k3.nvidia.low_latency_gemm import (
17 + KDA_PROJECTION_OVERLAP_MAX_TOKENS,
18 + run_kda_projection_overlap,
19 + try_low_latency_gemm,
20 +)
21 +
22 +ProjectionOutput = tuple[torch.Tensor, torch.Tensor, torch.Tensor]
23 +
24 +
25 +def _capture(
26 + fn: Callable[[], list[ProjectionOutput]],
27 +) -> tuple[torch.cuda.CUDAGraph, list[ProjectionOutput]]:
28 + capture_stream = torch.cuda.Stream()
29 + capture_stream.wait_stream(torch.cuda.current_stream())
30 + with torch.cuda.stream(capture_stream):
31 + fn()
32 + fn()
33 + capture_stream.synchronize()
34 +
35 + graph = torch.cuda.CUDAGraph()
36 + with torch.cuda.graph(graph, stream=capture_stream):
37 + outputs = fn()
38 + torch.cuda.current_stream().wait_stream(capture_stream)
39 + torch.accelerator.synchronize()
40 + return graph, outputs
41 +
42 +
43 +def _benchmark_graph(
44 + graph: torch.cuda.CUDAGraph,
45 + num_weight_sets: int,
46 + repeats: int,
47 +) -> float:
48 + start = torch.cuda.Event(enable_timing=True)
49 + end = torch.cuda.Event(enable_timing=True)
50 + start.record()
51 + for _ in range(repeats):
52 + graph.replay()
53 + end.record()
54 + end.synchronize()
55 + return start.elapsed_time(end) * 1000 / repeats / num_weight_sets
56 +
57 +
58 +def _validate(
59 + output: ProjectionOutput,
60 + hidden_states: torch.Tensor,
61 + packed_weight: torch.Tensor,
62 + f_b_weight: torch.Tensor,
63 +) -> None:
64 + expected_qkvg = torch.nn.functional.linear(
65 + hidden_states.float(), packed_weight[:6144].float()
66 + )
67 + expected_fab = torch.nn.functional.linear(
68 + hidden_states.float(), packed_weight[6144:6284].float()
69 + )
70 + expected_g1 = torch.nn.functional.linear(
71 + expected_fab[:, :128].to(torch.bfloat16).float(),
72 + f_b_weight.float(),
73 + )
74 + for actual, expected in zip(
75 + output,
76 + (expected_qkvg, expected_g1, expected_fab[:, 128:]),
77 + ):
78 + cosine = torch.nn.functional.cosine_similarity(
79 + actual.float().flatten(), expected.flatten(), dim=0
80 + ).item()
81 + if cosine <= 0.999:
82 + raise AssertionError(f"projection cosine similarity is {cosine}")
83 +
84 +
85 +def _sequential_projection(
86 + hidden_states: torch.Tensor,
87 + packed_weight: torch.Tensor,
88 + f_b_weight: torch.Tensor,
89 +) -> ProjectionOutput:
90 + projected = try_low_latency_gemm(hidden_states, packed_weight)
91 + if projected is None:
92 + projected = torch.mm(hidden_states, packed_weight.t())
93 + f_a = projected[:, 6144:6272]
94 + g1 = torch.mm(f_a, f_b_weight.t())
95 + return projected[:, :6144], g1, projected[:, 6272:6284]
96 +
97 +
98 +def main() -> None:
99 + parser = argparse.ArgumentParser(description=__doc__)
100 + parser.add_argument("--tokens", type=int, nargs="+", default=range(1, 17))
101 + parser.add_argument("--weight-sets", type=int, default=4)
102 + parser.add_argument("--repeats", type=int, default=300)
103 + parser.add_argument("--trials", type=int, default=5)
104 + args = parser.parse_args()
105 +
106 + if torch.cuda.get_device_capability() != (10, 3):
107 + raise SystemExit("This benchmark requires an SM103 GPU")
108 + if not shape_dynamic_skinny_gemm.is_available():
109 + raise SystemExit("This benchmark requires CuTe DSL")
110 + required_ops = ("dsv3_fused_a_gemm",)
111 + if not all(hasattr(torch.ops._C, name) for name in required_ops):
112 + raise SystemExit("This benchmark requires the KDA projection GEMMs")
113 + if not all(1 <= num_tokens <= 16 for num_tokens in args.tokens):
114 + raise SystemExit("--tokens must be in [1, 16]")
115 +
116 + torch.manual_seed(42)
117 + packed_weights = [
118 + torch.randn(6288, 7168, dtype=torch.bfloat16, device="cuda")
119 + for _ in range(args.weight_sets)
120 + ]
121 + f_b_weights = [
122 + torch.randn(1536, 128, dtype=torch.bfloat16, device="cuda")
123 + for _ in range(args.weight_sets)
124 + ]
125 +
126 + from flashinfer.autotuner import autotune
127 + from flashinfer.gemm import mm_bf16
128 +
129 + with torch.inference_mode(), autotune(tune_mode=True):
130 + mm_bf16(
131 + torch.empty(
132 + (KDA_PROJECTION_OVERLAP_MAX_TOKENS, 7168),
133 + dtype=torch.bfloat16,
134 + device="cuda",
135 + ),
136 + packed_weights[0][:6144].t(),
137 + pdl=True,
138 + backend="cute-dsl",
139 + )
140 +
141 + print("M sequential (us) multistream (us) fastest")
142 + for num_tokens in args.tokens:
143 + hidden_states = [
144 + torch.randn(
145 + num_tokens,
146 + 7168,
147 + dtype=torch.bfloat16,
148 + device="cuda",
149 + )
150 + for _ in range(args.weight_sets)
151 + ]
152 +
153 + def sequential(
154 + hidden_states: list[torch.Tensor] = hidden_states,
155 + ) -> list[ProjectionOutput]:
156 + return [
157 + _sequential_projection(x, weight, f_b_weight)
158 + for x, weight, f_b_weight in zip(
159 + hidden_states, packed_weights, f_b_weights
160 + )
161 + ]
162 +
163 + candidates: list[tuple[str, Callable[[], list[ProjectionOutput]]]] = [
164 + ("sequential", sequential)
165 + ]
166 + if num_tokens <= KDA_PROJECTION_OVERLAP_MAX_TOKENS:
167 + aux_stream = torch.cuda.Stream()
168 + events = [
169 + (torch.cuda.Event(), torch.cuda.Event())
170 + for _ in range(args.weight_sets)
171 + ]
172 +
173 + def overlap(
174 + hidden_states: list[torch.Tensor] = hidden_states,
175 + aux_stream: torch.cuda.Stream = aux_stream,
176 + events: list[tuple[torch.cuda.Event, torch.cuda.Event]] = events,
177 + ) -> list[ProjectionOutput]:
178 + return [
179 + run_kda_projection_overlap(x, weight, f_b_weight, aux_stream, pair)
180 + for x, weight, f_b_weight, pair in zip(
181 + hidden_states, packed_weights, f_b_weights, events
182 + )
183 + ]
184 +
185 + candidates.append(("multistream", overlap))
186 +
187 + graphs = {}
188 + for name, candidate in candidates:
189 + graph, outputs = _capture(candidate)
190 + graph.replay()
191 + torch.accelerator.synchronize()
192 + _validate(outputs[0], hidden_states[0], packed_weights[0], f_b_weights[0])
193 + graphs[name] = graph
194 +
195 + samples = {name: [] for name in graphs}
196 + graph_items = list(graphs.items())
197 + for trial in range(args.trials):
198 + ordered_graphs = graph_items if trial % 2 == 0 else graph_items[::-1]
199 + for name, graph in ordered_graphs:
200 + samples[name].append(
201 + _benchmark_graph(graph, args.weight_sets, args.repeats)
202 + )
203 + results = {name: statistics.median(values) for name, values in samples.items()}
204 +
205 + sequential_us = results["sequential"]
206 + multistream_us = results.get("multistream")
207 + multistream_text = (
208 + f"{multistream_us:18.3f}"
209 + if multistream_us is not None
210 + else " -"
211 + )
212 + fastest = (
213 + "multistream"
214 + if multistream_us is not None and multistream_us < sequential_us
215 + else "sequential"
216 + )
217 + print(f"{num_tokens:2d} {sequential_us:15.3f}{multistream_text} {fastest}")
218 +
219 +
220 +if __name__ == "__main__":
221 + main()
tests/kernels/test_bf16_skinny_gemm.py
+118
@@ -253,6 +253,65 @@ def test_residual_cute_configs_match_measured_table() -> None:
253 assert actual == EXPECTED_RESIDUAL_CUTE_CONFIGS
254
255
256 +def test_kda_overlap_configs_match_measured_table() -> None:
257 + assert k3_gemm.KDA_PROJECTION_OVERLAP_MAX_TOKENS == 14
258 + assert k3_gemm.KDA_SKINNY_N_MAX_TOKENS == 14
259 + assert k3_gemm.KDA_SKINNY_K_MAX_TOKENS == 14
260 + assert _config_tuple(k3_gemm.KDA_M1_QKVG_CONFIG) == (64, 4, 2, 8)
261 + assert _config_tuple(k3_gemm.KDA_M1_FAB_CONFIG) == (224, 1, 2, 8)
262 + assert {
263 + num_tokens: _config_tuple(config)
264 + for num_tokens, config in k3_gemm.KDA_QKVG_CONFIGS.items()
265 + } == {
266 + 1: (64, 4, 2, 8),
267 + 2: (64, 3, 2, 8),
268 + }
269 +
270 +
271 +@pytest.mark.parametrize("tp_size", [1, 2, 4, 16])
272 +def test_kda_projection_overlap_is_tp8_only(tp_size: int) -> None:
273 + from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
274 +
275 + kda = KimiK3DeltaAttention.__new__(KimiK3DeltaAttention)
276 + nn.Module.__init__(kda)
277 + kda.tp_size = tp_size
278 + kda._projection_overlap_max_tokens = 0
279 +
280 + assert not k3_gemm._enable_kda_projection_overlap(kda)
281 + assert kda._projection_overlap_max_tokens == 0
282 +
283 +
284 +def test_kda_qkvg_autotune_enables_full_overlap(
285 + monkeypatch: pytest.MonkeyPatch,
286 +) -> None:
287 + import flashinfer.gemm
288 +
289 + from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
290 +
291 + kda = KimiK3DeltaAttention.__new__(KimiK3DeltaAttention)
292 + nn.Module.__init__(kda)
293 + kda._projection_overlap_max_tokens = max(k3_gemm.KDA_QKVG_CONFIGS)
294 + kda.in_proj_qkvgfab = nn.Module()
295 + kda.in_proj_qkvgfab.weight = nn.Parameter(
296 + torch.empty(6144, 8, dtype=torch.bfloat16),
297 + requires_grad=False,
298 + )
299 + calls = []
300 +
301 + def fake_mm_bf16(a, b, *, pdl, backend):
302 + calls.append((a.shape, b.shape, pdl, backend))
303 + return torch.empty(a.shape[0], b.shape[1], dtype=a.dtype)
304 +
305 + monkeypatch.setattr(flashinfer.gemm, "mm_bf16", fake_mm_bf16)
306 +
307 + k3_gemm.autotune_kda_qkvg(kda)
308 +
309 + assert calls == [(torch.Size([14, 8]), torch.Size([8, 6144]), True, "cute-dsl")]
310 + assert (
311 + kda._projection_overlap_max_tokens == k3_gemm.KDA_PROJECTION_OVERLAP_MAX_TOKENS
312 + )
313 +
314 +
315 def test_glm52_projection_plans_are_separate() -> None:
316 qkv_a = glm52_gemm.GLM52_QKV_A_PROJECTION
317 q_b = glm52_gemm.GLM52_Q_B_PROJECTION
@@ -812,6 +871,65 @@ def test_cute_selected_shapes(
871 assert cosine > 0.999
872
873
874 +@pytest.mark.parametrize("num_tokens", [1, 2, 3, 4, 8, 9, 10, 11, 12, 13, 14])
875 +def test_kda_projection_overlap_cuda_graph(num_tokens: int) -> None:
876 + """The TP8 split must preserve projection results across both streams."""
877 + _require_sm103_and_cute()
878 + _require_sm103_and_dsv3()
879 + torch.manual_seed(43)
880 + hidden_states = torch.randn(num_tokens, 7168, dtype=torch.bfloat16, device="cuda")
881 + packed_weight = torch.randn(6288, 7168, dtype=torch.bfloat16, device="cuda")
882 + f_b_weight = torch.randn(1536, 128, dtype=torch.bfloat16, device="cuda")
883 + aux_stream = torch.cuda.Stream()
884 + events = (torch.cuda.Event(), torch.cuda.Event())
885 +
886 + k3_gemm.run_kda_projection_overlap(
887 + hidden_states,
888 + packed_weight,
889 + f_b_weight,
890 + aux_stream,
891 + events,
892 + )
893 + torch.accelerator.synchronize()
894 +
895 + capture_stream = torch.cuda.Stream()
896 + capture_stream.wait_stream(torch.cuda.current_stream())
897 + graph = torch.cuda.CUDAGraph()
898 + with torch.cuda.graph(graph, stream=capture_stream):
899 + projected_qkvg, g1, beta = k3_gemm.run_kda_projection_overlap(
900 + hidden_states,
901 + packed_weight,
902 + f_b_weight,
903 + aux_stream,
904 + events,
905 + )
906 + torch.cuda.current_stream().wait_stream(capture_stream)
907 + graph.replay()
908 + torch.accelerator.synchronize()
909 +
910 + assert beta.stride() == ((140 if num_tokens == 1 else 144), 1)
911 +
912 + expected_qkvg = torch.nn.functional.linear(
913 + hidden_states.float(), packed_weight[:6144].float()
914 + )
915 + expected_fab = torch.nn.functional.linear(
916 + hidden_states.float(), packed_weight[6144:6284].float()
917 + )
918 + expected_g1 = torch.nn.functional.linear(
919 + expected_fab[:, :128].to(torch.bfloat16).float(),
920 + f_b_weight.float(),
921 + )
922 + for actual, expected in (
923 + (projected_qkvg, expected_qkvg),
924 + (g1, expected_g1),
925 + (beta, expected_fab[:, 128:]),
926 + ):
927 + cosine = torch.nn.functional.cosine_similarity(
928 + actual.float().flatten(), expected.flatten(), dim=0
929 + ).item()
930 + assert cosine > 0.999
931 +
932 +
933 def _dsv3_probe_tokens(tokens: frozenset[int]) -> set[int]:
934 """Extremes, plus both sides of the kernel's num_tokens<=8 tile_n branch."""
935 if not tokens:
vllm/model_executor/warmup/kernel_warmup.py
+7
@@ -114,6 +114,12 @@ def _warmup_kimi_k3_gemm_rs_ar() -> None:
114 logger.info_once("Warmed up %d Kimi-K3 GEMM-RS/AR variants.", compiled)
115
116
117 +def _autotune_kimi_k3_kda_qkvg(model: torch.nn.Module) -> None:
118 + module = sys.modules.get("vllm.models.kimi_k3.nvidia.low_latency_gemm")
119 + if module is not None:
120 + module.autotune_kda_qkvg(model)
121 +
122 +
123 def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
124 from vllm.model_executor.warmup.minimax_m3_msa_warmup import (
125 minimax_m3_msa_warmup,
@@ -343,6 +349,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
349 ):
350 _run_flashinfer_autotune_dummy_runs(runner)
351 replayssm_autotune_warmup(runner)
352 + _autotune_kimi_k3_kda_qkvg(runner.get_model())
353 finally:
354 set_autotune_process_group(None)
355
vllm/models/kimi_k3/nvidia/kda.py
+43 -12
@@ -399,6 +399,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
399 config: KimiLinearConfig,
400 vllm_config: VllmConfig,
401 prefix: str = "",
402 + aux_stream: torch.cuda.Stream | None = None,
403 run_gemm_rs_ar: bool = False,
404 ) -> None:
405 super().__init__(config, vllm_config, prefix)
@@ -421,6 +422,11 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
422 assert kda_config.get("use_full_rank_gate", False), (
423 "KimiK3DeltaAttention requires a full-rank gate"
424 )
425 + self._projection_aux_stream = aux_stream
426 + self._projection_events = (
427 + (torch.cuda.Event(), torch.cuda.Event()) if aux_stream is not None else None
428 + )
429 + self._projection_overlap_max_tokens = 0
430
431 # Keep f_a before the narrow beta shard, then pad each TP-local row
432 # to select the aligned BF16 GEMM path.
@@ -605,19 +611,44 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
611 positions: torch.Tensor,
612 ) -> torch.Tensor:
613 num_tokens = hidden_states.size(0)
608 - projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]
609 - split_sizes = [
610 - 3 * self.local_projection_size,
611 - self.local_projection_size,
612 - self.head_dim,
613 - self.local_num_heads,
614 - ]
615 - if self.in_proj_padding:
616 - split_sizes.append(self.in_proj_padding)
617 - projected = projected_qkvgfab.split(split_sizes, dim=-1)
618 - mixed_qkv, g_proj_states, f_a, beta = projected[:4]
614 + projection_events = self._projection_events
615 + projection_aux_stream = self._projection_aux_stream
616 + if (
617 + 0 < num_tokens <= self._projection_overlap_max_tokens
618 + and hidden_states.stride() == (self.hidden_size, 1)
619 + and projection_events is not None
620 + and projection_aux_stream is not None
621 + and torch.cuda.is_current_stream_capturing()
622 + ):
623 + from vllm.models.kimi_k3.nvidia.low_latency_gemm import (
624 + run_kda_projection_overlap,
625 + )
626 +
627 + projected_qkvg, g1, beta = run_kda_projection_overlap(
628 + hidden_states,
629 + self.in_proj_qkvgfab.weight,
630 + self.f_b_proj.weight,
631 + projection_aux_stream,
632 + projection_events,
633 + )
634 + mixed_qkv, g_proj_states = projected_qkvg.split(
635 + [3 * self.local_projection_size, self.local_projection_size],
636 + dim=-1,
637 + )
638 + else:
639 + projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]
640 + split_sizes = [
641 + 3 * self.local_projection_size,
642 + self.local_projection_size,
643 + self.head_dim,
644 + self.local_num_heads,
645 + ]
646 + if self.in_proj_padding:
647 + split_sizes.append(self.in_proj_padding)
648 + projected = projected_qkvgfab.split(split_sizes, dim=-1)
649 + mixed_qkv, g_proj_states, f_a, beta = projected[:4]
650 + g1 = self.f_b_proj(f_a)[0]
651
620 - g1 = self.f_b_proj(f_a)[0]
652 beta = beta.unsqueeze(0)
653 g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim)
654 g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)
vllm/models/kimi_k3/nvidia/low_latency_gemm.py
+195
@@ -35,12 +35,32 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
35 UnquantizedEmbeddingMethod,
36 )
37 from vllm.platforms import current_platform
38 +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
39
40 Backend = Literal["cute", "dsv3_fused_a"]
41 # A resolved per-token-count call: the backend plus its CuTe config (None for
42 # dsv3, which needs no config).
43 ResolvedCall = tuple[Backend, SkinnyGemmConfig | None]
44
45 +# TP8 KDA projection split, measured together under CUDA graph capture on B300.
46 +# Q/K/V/G stay on the graph's main stream while F_A/beta and F_B run on the
47 +# model's auxiliary stream.
48 +KDA_M1_QKVG_CONFIG = SkinnyGemmConfig(1, 64, 4, 2, 8)
49 +KDA_M1_FAB_CONFIG = SkinnyGemmConfig(1, 224, 1, 2, 8)
50 +KDA_QKVG_CONFIGS = {
51 + 1: KDA_M1_QKVG_CONFIG,
52 + 2: SkinnyGemmConfig(2, 64, 3, 2, 8),
53 +}
54 +# The captured end-to-end projection sweep wins through M=14 and regresses at
55 +# M=15 and M=16, where the original packed projection remains selected.
56 +KDA_PROJECTION_OVERLAP_MAX_TOKENS = 14
57 +KDA_SKINNY_N_MAX_TOKENS = KDA_PROJECTION_OVERLAP_MAX_TOKENS
58 +KDA_SKINNY_K_MAX_TOKENS = KDA_PROJECTION_OVERLAP_MAX_TOKENS
59 +_KDA_QKVG_SIZE = 4 * 1536
60 +_KDA_FAB_SIZE = 128 + 12
61 +_KDA_PACKED_SIZE = 6288
62 +_KDA_TP_SIZE = 8
63 +
64
65 @dataclass(frozen=True, slots=True)
66 class ProjectionSpec:
@@ -683,6 +703,137 @@ def _run_plan(
703 return output
704
705
706 +def run_kda_projection_overlap(
707 + hidden_states: torch.Tensor,
708 + packed_weight: torch.Tensor,
709 + f_b_weight: torch.Tensor,
710 + aux_stream: torch.cuda.Stream,
711 + events: tuple[torch.cuda.Event, torch.cuda.Event],
712 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
713 + """Run the TP8 KDA decode projection branches concurrently.
714 +
715 + Args:
716 + hidden_states: Packed BF16 input with shape ``[M, 7168]``.
717 + packed_weight: Existing Q/K/V/G/F_A/beta/pad weight with shape
718 + ``[6288, 7168]``.
719 + f_b_weight: F_B weight with shape ``[1536, 128]``.
720 + aux_stream: Stream for the F_A/beta then F_B branch.
721 + events: Start and completion events for the stream fork and join.
722 +
723 + Returns:
724 + The Q/K/V/G projection, F_B output, and beta projection.
725 + """
726 + num_tokens = hidden_states.shape[0]
727 + qkvg_weight = packed_weight[:_KDA_QKVG_SIZE]
728 +
729 + def run_qkvg() -> torch.Tensor:
730 + if config := KDA_QKVG_CONFIGS.get(num_tokens):
731 + return shape_dynamic_skinny_gemm(
732 + hidden_states,
733 + qkvg_weight,
734 + config,
735 + None,
736 + )
737 + if num_tokens <= KDA_PROJECTION_OVERLAP_MAX_TOKENS:
738 + from flashinfer.gemm import mm_bf16
739 +
740 + return mm_bf16(
741 + hidden_states,
742 + qkvg_weight.t(),
743 + pdl=True,
744 + backend="cute-dsl",
745 + )
746 + return torch.mm(hidden_states, qkvg_weight.t())
747 +
748 + def run_fab_fb() -> tuple[torch.Tensor, torch.Tensor]:
749 + if num_tokens == 1:
750 + projected_fab = shape_dynamic_skinny_gemm(
751 + hidden_states,
752 + packed_weight[_KDA_QKVG_SIZE : _KDA_QKVG_SIZE + _KDA_FAB_SIZE],
753 + KDA_M1_FAB_CONFIG,
754 + None,
755 + )
756 + f_a, beta = projected_fab.split([128, 12], dim=-1)
757 + f_a = f_a.as_strided((1, 128), (128, 1))
758 + g1 = torch.empty(
759 + (1, 1536), dtype=hidden_states.dtype, device=hidden_states.device
760 + )
761 + ops.dsv3_fused_a_gemm(g1, f_a, f_b_weight.t(), enable_pdl=True)
762 + else:
763 + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.kda_skinny_gemm import (
764 + kda_skinny_gemm,
765 + )
766 +
767 + if num_tokens <= KDA_SKINNY_N_MAX_TOKENS:
768 + projected_fab = kda_skinny_gemm.run_n(
769 + hidden_states,
770 + packed_weight[_KDA_QKVG_SIZE:],
771 + )
772 + else:
773 + projected_fab = torch.mm(
774 + hidden_states,
775 + packed_weight[_KDA_QKVG_SIZE:].t(),
776 + )
777 + beta = projected_fab[:, 128:140]
778 + if num_tokens <= KDA_SKINNY_K_MAX_TOKENS:
779 + g1 = kda_skinny_gemm.run_k(projected_fab, f_b_weight)
780 + else:
781 + g1 = torch.mm(projected_fab[:, :128], f_b_weight.t())
782 + return g1, beta
783 +
784 + projected_qkvg, (g1, beta) = maybe_execute_in_parallel(
785 + run_qkvg,
786 + run_fab_fb,
787 + events[0],
788 + events[1],
789 + aux_stream,
790 + )
791 + return projected_qkvg, g1, beta
792 +
793 +
794 +def autotune_kda_qkvg(model: nn.Module) -> None:
795 + """Autotune the FlashInfer QKVG GEMM before CUDA graph capture."""
796 + from flashinfer.gemm import mm_bf16
797 +
798 + from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
799 +
800 + children: list[KimiK3DeltaAttention] = []
801 + weights_by_shape: dict[
802 + tuple[int, int, int, torch.dtype, torch.device], torch.Tensor
803 + ] = {}
804 + for child in model.modules():
805 + if not isinstance(child, KimiK3DeltaAttention):
806 + continue
807 + if child._projection_overlap_max_tokens <= 0:
808 + continue
809 + children.append(child)
810 + qkvg_weight = child.in_proj_qkvgfab.weight[:_KDA_QKVG_SIZE]
811 + shape = (
812 + KDA_PROJECTION_OVERLAP_MAX_TOKENS,
813 + qkvg_weight.shape[0],
814 + qkvg_weight.shape[1],
815 + qkvg_weight.dtype,
816 + qkvg_weight.device,
817 + )
818 + weights_by_shape.setdefault(shape, qkvg_weight)
819 +
820 + for shape, qkvg_weight in weights_by_shape.items():
821 + num_tokens = shape[0]
822 + hidden_states = torch.empty(
823 + (num_tokens, qkvg_weight.shape[1]),
824 + dtype=qkvg_weight.dtype,
825 + device=qkvg_weight.device,
826 + )
827 + mm_bf16(
828 + hidden_states,
829 + qkvg_weight.t(),
830 + pdl=True,
831 + backend="cute-dsl",
832 + )
833 + for child in children:
834 + child._projection_overlap_max_tokens = KDA_PROJECTION_OVERLAP_MAX_TOKENS
835 +
836 +
837 def _run_residual_plan(
838 residual_plan: dict[int, SkinnyGemmConfig],
839 x: torch.Tensor,
@@ -777,6 +928,40 @@ class KimiK3LowLatencyEmbeddingMethod(
928 pass
929
930
931 +def _enable_kda_projection_overlap(module: nn.Module) -> bool:
932 + from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
933 +
934 + if envs.VLLM_BATCH_INVARIANT:
935 + return False
936 +
937 + enabled = False
938 + for child in module.modules():
939 + if not isinstance(child, KimiK3DeltaAttention):
940 + continue
941 + if child.tp_size != _KDA_TP_SIZE:
942 + continue
943 + packed_weight = child.in_proj_qkvgfab.weight
944 + f_b_weight = child.f_b_proj.weight
945 + if (
946 + type(child.in_proj_qkvgfab.quant_method) is not KimiK3LowLatencyLinearMethod
947 + or type(child.f_b_proj.quant_method) is not KimiK3LowLatencyLinearMethod
948 + or child._projection_aux_stream is None
949 + or child._projection_events is None
950 + or packed_weight.shape != (_KDA_PACKED_SIZE, 7168)
951 + or f_b_weight.shape != (1536, 128)
952 + or packed_weight.dtype != torch.bfloat16
953 + or f_b_weight.dtype != torch.bfloat16
954 + or not packed_weight.is_cuda
955 + or not f_b_weight.is_cuda
956 + or not packed_weight.is_contiguous()
957 + or not f_b_weight.is_contiguous()
958 + ):
959 + continue
960 + child._projection_overlap_max_tokens = max(KDA_QKVG_CONFIGS)
961 + enabled = True
962 + return enabled
963 +
964 +
965 def enable_kimi_k3_low_latency_gemm(
966 module: nn.Module,
967 dtype: torch.dtype,
@@ -825,6 +1010,16 @@ def enable_kimi_k3_low_latency_gemm(
1010 # TP8 deployment does not compile TP4 configs and vice versa.
1011 warmup_configs.update(config for _, config in spec.cute_configs)
1012 residual_warmup_configs.update(config for _, config in spec.residual_configs)
1013 + if _enable_kda_projection_overlap(module):
1014 + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.kda_skinny_gemm import (
1015 + kda_skinny_gemm,
1016 + )
1017 +
1018 + warmup_configs.update((*KDA_QKVG_CONFIGS.values(), KDA_M1_FAB_CONFIG))
1019 + kda_skinny_gemm.request_warmup(
1020 + set(range(2, KDA_SKINNY_N_MAX_TOKENS + 1)),
1021 + set(range(2, KDA_SKINNY_K_MAX_TOKENS + 1)),
1022 + )
1023
1024 if shape_dynamic_skinny_gemm.is_available():
1025 if warmup_configs:
vllm/models/kimi_k3/nvidia/model.py
+1
@@ -895,6 +895,7 @@ class KimiDecoderLayer(nn.Module):
895 config,
896 vllm_config,
897 prefix=f"{prefix}.self_attn",
898 + aux_stream=aux_stream,
899 run_gemm_rs_ar=run_gemm_rs_ar,
900 )
901 self._self_attn_writes_output = False
vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py new
+507
@@ -0,0 +1,507 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Adapted from
4 +# https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh
5 +"""Kimi-K3 TP8 skinny GEMMs for the KDA F_A/beta and F_B projections."""
6 +
7 +from __future__ import annotations
8 +
9 +from collections.abc import Callable
10 +from functools import partial
11 +
12 +import cuda.bindings.driver as cuda
13 +import cutlass
14 +import cutlass.cute as cute
15 +import torch
16 +from cutlass import BFloat16, Float32
17 +from cutlass._mlir.dialects import llvm
18 +from cutlass.cute.runtime import make_fake_stream, make_fake_tensor
19 +from cutlass.cutlass_dsl import T, dsl_user_op
20 +
21 +_FAB_N = 144
22 +_FAB_K = 7168
23 +_FAB_VECTOR_WIDTH = 16
24 +_FAB_BLOCK_SIZE = _FAB_K // _FAB_VECTOR_WIDTH
25 +_FAB_NUM_WARPS = _FAB_BLOCK_SIZE // cute.arch.WARP_SIZE
26 +
27 +_FB_N = 1536
28 +_FB_K = 128
29 +_FB_VECTOR_WIDTH = 8
30 +_FB_K_LANES = _FB_K // _FB_VECTOR_WIDTH
31 +
32 +SKINNY_N_SPLIT_N = {
33 + 2: 1,
34 + 3: 2,
35 + 4: 3,
36 + 5: 3,
37 + 6: 3,
38 + 7: 3,
39 + 8: 3,
40 + 9: 2,
41 + 10: 2,
42 + 11: 2,
43 + 12: 2,
44 + 13: 2,
45 + 14: 1,
46 + 15: 1,
47 + 16: 1,
48 +}
49 +
50 +SKINNY_K_SPLIT_N = {
51 + 2: 8,
52 + 3: 16,
53 + 4: 16,
54 + 5: 24,
55 + 6: 24,
56 + 7: 24,
57 + 8: 24,
58 + 9: 12,
59 + 10: 16,
60 + 11: 16,
61 + 12: 12,
62 + 13: 8,
63 + 14: 8,
64 + 15: 8,
65 + 16: 8,
66 +}
67 +
68 +
69 +@dsl_user_op
70 +def _fma_f32_bf16(
71 + a: BFloat16,
72 + b: BFloat16,
73 + acc: Float32,
74 + *,
75 + loc=None,
76 + ip=None,
77 +) -> Float32:
78 + a_bits = llvm.bitcast(T.i16(), a.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
79 + b_bits = llvm.bitcast(T.i16(), b.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
80 + result = llvm.inline_asm(
81 + T.f32(),
82 + [a_bits, b_bits, acc.ir_value(loc=loc, ip=ip)],
83 + "fma.rn.f32.bf16 $0, $1, $2, $3;",
84 + "=f,h,h,f",
85 + has_side_effects=False,
86 + is_align_stack=False,
87 + asm_dialect=llvm.AsmDialect.AD_ATT,
88 + loc=loc,
89 + ip=ip,
90 + )
91 + return Float32(result)
92 +
93 +
94 +class _KdaSkinnyNGemm:
95 + def __init__(self, num_rows: int, split_n: int) -> None:
96 + if _FAB_N % split_n:
97 + raise ValueError("split_n must divide the FAB output size")
98 + if num_rows * split_n > _FAB_BLOCK_SIZE:
99 + raise ValueError("the FAB output tile must fit in one block")
100 + self.num_rows = num_rows
101 + self.split_n = split_n
102 +
103 + @cute.jit
104 + def __call__(
105 + self,
106 + a: cute.Tensor,
107 + weight: cute.Tensor,
108 + output: cute.Tensor,
109 + stream: cuda.CUstream,
110 + ) -> None:
111 + copy_a = cute.make_copy_atom(
112 + cute.nvgpu.CopyG2ROp(),
113 + BFloat16,
114 + num_bits_per_copy=_FAB_VECTOR_WIDTH * BFloat16.width,
115 + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
116 + )
117 + copy_weight = cute.make_copy_atom(
118 + cute.nvgpu.CopyG2ROp(),
119 + BFloat16,
120 + num_bits_per_copy=_FAB_VECTOR_WIDTH * BFloat16.width,
121 + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING,
122 + )
123 + self.kernel(a, weight, output, copy_a, copy_weight).launch(
124 + grid=[_FAB_N // self.split_n, 1, 1],
125 + block=[_FAB_BLOCK_SIZE, 1, 1],
126 + smem=_FAB_NUM_WARPS * self.num_rows * self.split_n * 4,
127 + stream=stream,
128 + use_pdl=True,
129 + min_blocks_per_mp=1,
130 + )
131 +
132 + @cute.kernel
133 + def kernel(
134 + self,
135 + a: cute.Tensor,
136 + weight: cute.Tensor,
137 + output: cute.Tensor,
138 + copy_a: cute.CopyAtom,
139 + copy_weight: cute.CopyAtom,
140 + ) -> None:
141 + tidx, _, _ = cute.arch.thread_idx()
142 + block_idx, _, _ = cute.arch.block_idx()
143 + warp_idx = cute.arch.warp_idx()
144 + num_rows: cutlass.Constexpr = self.num_rows
145 + split_n: cutlass.Constexpr = self.split_n
146 +
147 + a_vectors = cute.logical_divide(a, (None, _FAB_VECTOR_WIDTH))
148 + weight_vectors = cute.logical_divide(weight, (None, _FAB_VECTOR_WIDTH))
149 + a_tiles = cute.logical_divide(
150 + a_vectors,
151 + (None, (None, _FAB_BLOCK_SIZE)),
152 + )
153 + weight_tiles = cute.logical_divide(
154 + weight_vectors,
155 + (None, (None, _FAB_BLOCK_SIZE)),
156 + )
157 + a_tile = a_tiles[None, (None, (tidx, None))]
158 + n_base = block_idx * split_n
159 +
160 + a_regs = cute.make_rmem_tensor(
161 + cute.make_layout(
162 + (num_rows, _FAB_VECTOR_WIDTH),
163 + stride=(_FAB_VECTOR_WIDTH, 1),
164 + ),
165 + BFloat16,
166 + )
167 + weight_regs = cute.make_rmem_tensor(
168 + cute.make_layout(
169 + (split_n, _FAB_VECTOR_WIDTH),
170 + stride=(_FAB_VECTOR_WIDTH, 1),
171 + ),
172 + BFloat16,
173 + )
174 + acc = cute.make_rmem_tensor(
175 + cute.make_layout((num_rows, split_n), stride=(split_n, 1)),
176 + Float32,
177 + )
178 + acc.fill(0.0)
179 +
180 + for ni in cutlass.range_constexpr(split_n):
181 + weight_tile = weight_tiles[
182 + n_base + ni,
183 + (None, (tidx, None)),
184 + ]
185 + cute.copy(copy_weight, weight_tile[None, 0], weight_regs[ni, None])
186 +
187 + cute.arch.griddepcontrol_wait()
188 +
189 + for mi in cutlass.range_constexpr(num_rows):
190 + cute.copy(copy_a, a_tile[mi, None, 0], a_regs[mi, None])
191 + for mi in cutlass.range_constexpr(num_rows):
192 + for ni in cutlass.range_constexpr(split_n):
193 + for vi in cutlass.range_constexpr(_FAB_VECTOR_WIDTH):
194 + acc[mi, ni] = _fma_f32_bf16(
195 + a_regs[mi, vi],
196 + weight_regs[ni, vi],
197 + acc[mi, ni],
198 + )
199 + acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni])
200 +
201 + partials_layout = cute.make_layout(
202 + (_FAB_NUM_WARPS, num_rows, split_n),
203 + stride=(num_rows * split_n, split_n, 1),
204 + )
205 + smem = cutlass.utils.SmemAllocator()
206 + partials = smem.allocate_tensor(
207 + Float32,
208 + partials_layout,
209 + byte_alignment=16,
210 + )
211 + with cute.arch.elect_one():
212 + for mi in cutlass.range_constexpr(num_rows):
213 + for ni in cutlass.range_constexpr(split_n):
214 + partials[warp_idx, mi, ni] = acc[mi, ni]
215 +
216 + cute.arch.griddepcontrol_launch_dependents()
217 + cute.arch.sync_threads()
218 +
219 + if tidx < num_rows * split_n:
220 + mi = tidx // split_n
221 + ni = tidx % split_n
222 + total = (
223 + partials[None, mi, ni]
224 + .load()
225 + .reduce(
226 + cute.ReductionOp.ADD,
227 + init_val=Float32(0.0),
228 + reduction_profile=0,
229 + )
230 + )
231 + output[mi, n_base + ni] = Float32(total).to(BFloat16)
232 +
233 +
234 +class _KdaSkinnyKGemm:
235 + def __init__(self, num_rows: int, split_n: int) -> None:
236 + if _FB_N % split_n:
237 + raise ValueError("split_n must divide the FB output size")
238 + if split_n * _FB_K_LANES > 1024:
239 + raise ValueError("the FB output tile exceeds the block-size limit")
240 + self.num_rows = num_rows
241 + self.split_n = split_n
242 +
243 + @cute.jit
244 + def __call__(
245 + self,
246 + a: cute.Tensor,
247 + weight: cute.Tensor,
248 + output: cute.Tensor,
249 + stream: cuda.CUstream,
250 + ) -> None:
251 + copy_atom = cute.make_copy_atom(
252 + cute.nvgpu.CopyG2ROp(),
253 + BFloat16,
254 + num_bits_per_copy=_FB_VECTOR_WIDTH * BFloat16.width,
255 + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
256 + )
257 + self.kernel(a, weight, output, copy_atom).launch(
258 + grid=[_FB_N // self.split_n, 1, 1],
259 + block=[self.split_n * _FB_K_LANES, 1, 1],
260 + stream=stream,
261 + use_pdl=True,
262 + min_blocks_per_mp=1,
263 + )
264 +
265 + @cute.kernel
266 + def kernel(
267 + self,
268 + a: cute.Tensor,
269 + weight: cute.Tensor,
270 + output: cute.Tensor,
271 + copy_atom: cute.CopyAtom,
272 + ) -> None:
273 + tidx, _, _ = cute.arch.thread_idx()
274 + block_idx, _, _ = cute.arch.block_idx()
275 + num_rows: cutlass.Constexpr = self.num_rows
276 + split_n: cutlass.Constexpr = self.split_n
277 + k_lane = tidx % _FB_K_LANES
278 + n_idx = block_idx * split_n + tidx // _FB_K_LANES
279 +
280 + a_vectors = cute.logical_divide(a, (None, _FB_VECTOR_WIDTH))
281 + weight_vectors = cute.logical_divide(weight, (None, _FB_VECTOR_WIDTH))
282 + a_tiles = cute.logical_divide(
283 + a_vectors,
284 + (None, (None, _FB_K_LANES)),
285 + )
286 + weight_tiles = cute.logical_divide(
287 + weight_vectors,
288 + (None, (None, _FB_K_LANES)),
289 + )
290 + a_tile = a_tiles[None, (None, (k_lane, None))]
291 + weight_tile = weight_tiles[n_idx, (None, (k_lane, None))]
292 + weight_regs = cute.make_rmem_tensor(_FB_VECTOR_WIDTH, BFloat16)
293 + input_regs = cute.make_rmem_tensor(
294 + cute.make_layout(
295 + (num_rows, _FB_VECTOR_WIDTH),
296 + stride=(_FB_VECTOR_WIDTH, 1),
297 + ),
298 + BFloat16,
299 + )
300 + cute.copy(
301 + copy_atom,
302 + weight_tile[None, 0],
303 + weight_regs,
304 + )
305 +
306 + cute.arch.griddepcontrol_wait()
307 +
308 + for mi in cutlass.range_constexpr(num_rows):
309 + cute.copy(
310 + copy_atom,
311 + a_tile[mi, None, 0],
312 + input_regs[mi, None],
313 + )
314 + for mi in cutlass.range_constexpr(num_rows):
315 + acc = Float32(0.0)
316 + for vi in cutlass.range_constexpr(_FB_VECTOR_WIDTH):
317 + acc = _fma_f32_bf16(
318 + input_regs[mi, vi],
319 + weight_regs[vi],
320 + acc,
321 + )
322 + for offset in (8, 4, 2, 1):
323 + acc += cute.arch.shuffle_sync_bfly(acc, offset)
324 + output[mi, n_idx] = acc.to(BFloat16)
325 +
326 + cute.arch.griddepcontrol_launch_dependents()
327 +
328 +
329 +class KdaSkinnyGemm:
330 + def __init__(self) -> None:
331 + self._compiled_n: dict[tuple[int, int], Callable[..., None]] = {}
332 + self._compiled_k: dict[tuple[int, int], Callable[..., None]] = {}
333 + self._warmup_n: set[int] = set()
334 + self._warmup_k: set[int] = set()
335 + self._warmup_registered = False
336 +
337 + def _compile_n(self, num_rows: int) -> None:
338 + device = torch.accelerator.current_device_index()
339 + key = (device, num_rows)
340 + if key in self._compiled_n:
341 + return
342 + split_n = SKINNY_N_SPLIT_N[num_rows]
343 + self._compiled_n[key] = cute.compile(
344 + _KdaSkinnyNGemm(num_rows, split_n),
345 + make_fake_tensor(
346 + BFloat16,
347 + (num_rows, _FAB_K),
348 + stride=(_FAB_K, 1),
349 + assumed_align=32,
350 + ),
351 + make_fake_tensor(
352 + BFloat16,
353 + (_FAB_N, _FAB_K),
354 + stride=(_FAB_K, 1),
355 + assumed_align=32,
356 + ),
357 + make_fake_tensor(
358 + BFloat16,
359 + (num_rows, _FAB_N),
360 + stride=(_FAB_N, 1),
361 + assumed_align=32,
362 + ),
363 + make_fake_stream(use_tvm_ffi_env_stream=True),
364 + options="--enable-tvm-ffi --ptxas-options -maxrregcount=255",
365 + )
366 +
367 + def _compile_k(self, num_rows: int) -> None:
368 + device = torch.accelerator.current_device_index()
369 + key = (device, num_rows)
370 + if key in self._compiled_k:
371 + return
372 + split_n = SKINNY_K_SPLIT_N[num_rows]
373 + self._compiled_k[key] = cute.compile(
374 + _KdaSkinnyKGemm(num_rows, split_n),
375 + make_fake_tensor(
376 + BFloat16,
377 + (num_rows, _FB_K),
378 + stride=(_FAB_N, 1),
379 + assumed_align=32,
380 + ),
381 + make_fake_tensor(
382 + BFloat16,
383 + (_FB_N, _FB_K),
384 + stride=(_FB_K, 1),
385 + assumed_align=32,
386 + ),
387 + make_fake_tensor(
388 + BFloat16,
389 + (num_rows, _FB_N),
390 + stride=(_FB_N, 1),
391 + assumed_align=32,
392 + ),
393 + make_fake_stream(use_tvm_ffi_env_stream=True),
394 + options="--enable-tvm-ffi --ptxas-options -maxrregcount=255",
395 + )
396 +
397 + def request_warmup(self, n_tokens: set[int], k_tokens: set[int]) -> None:
398 + if not n_tokens <= SKINNY_N_SPLIT_N.keys():
399 + raise ValueError("unsupported skinny-N warmup token count")
400 + if not k_tokens <= SKINNY_K_SPLIT_N.keys():
401 + raise ValueError("unsupported skinny-K warmup token count")
402 + self._warmup_n.update(n_tokens)
403 + self._warmup_k.update(k_tokens)
404 + if self._warmup_registered:
405 + return
406 + from vllm.model_executor.warmup.cutedsl_warmup import (
407 + register_cutedsl_warmup_provider,
408 + )
409 +
410 + register_cutedsl_warmup_provider(self)
411 + self._warmup_registered = True
412 +
413 + def get_cutedsl_warmup_compile_units(self):
414 + from vllm.model_executor.warmup.cutedsl_warmup import CuTeDSLCompileUnit
415 +
416 + n_units = tuple(
417 + CuTeDSLCompileUnit(
418 + name="Kimi-K3 KDA skinny-N GEMM",
419 + key=("kimi-k3-kda-skinny-n", num_rows),
420 + compile=partial(self._compile_n, num_rows),
421 + )
422 + for num_rows in sorted(self._warmup_n)
423 + )
424 + k_units = tuple(
425 + CuTeDSLCompileUnit(
426 + name="Kimi-K3 KDA skinny-K GEMM",
427 + key=("kimi-k3-kda-skinny-k", num_rows),
428 + compile=partial(self._compile_k, num_rows),
429 + )
430 + for num_rows in sorted(self._warmup_k)
431 + )
432 + return n_units + k_units
433 +
434 + @staticmethod
435 + def _validate(
436 + a: torch.Tensor,
437 + weight: torch.Tensor,
438 + *,
439 + a_shape: tuple[int, int],
440 + weight_shape: tuple[int, int],
441 + ) -> None:
442 + if a.shape != a_shape or weight.shape != weight_shape:
443 + raise ValueError(
444 + f"expected input {a_shape} and weight {weight_shape}, got "
445 + f"{tuple(a.shape)} and {tuple(weight.shape)}"
446 + )
447 + if a.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16:
448 + raise ValueError("KDA skinny GEMMs require BF16 input and weight")
449 + if not a.is_cuda or not weight.is_cuda or a.device != weight.device:
450 + raise ValueError("KDA skinny GEMM operands must share a CUDA device")
451 + if not a.is_contiguous() or not weight.is_contiguous():
452 + raise ValueError("KDA skinny GEMM operands must be contiguous")
453 + if a.data_ptr() % 32 or weight.data_ptr() % 32:
454 + raise ValueError("KDA skinny GEMM operands must be 32-byte aligned")
455 +
456 + def run_n(self, a: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
457 + num_rows = a.shape[0]
458 + if num_rows not in SKINNY_N_SPLIT_N:
459 + raise ValueError("skinny-N requires M in [2, 16]")
460 + self._validate(
461 + a,
462 + weight,
463 + a_shape=(num_rows, _FAB_K),
464 + weight_shape=(_FAB_N, _FAB_K),
465 + )
466 + output = torch.empty(
467 + (num_rows, _FAB_N),
468 + dtype=a.dtype,
469 + device=a.device,
470 + )
471 + with torch.accelerator.device_index(a.device.index):
472 + self._compile_n(num_rows)
473 + key = (a.device.index, num_rows)
474 + self._compiled_n[key](
475 + a,
476 + weight,
477 + output,
478 + )
479 + return output
480 +
481 + def run_k(self, packed_fab: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
482 + num_rows = packed_fab.shape[0]
483 + if num_rows not in SKINNY_K_SPLIT_N:
484 + raise ValueError("skinny-K requires M in [2, 16]")
485 + self._validate(
486 + packed_fab,
487 + weight,
488 + a_shape=(num_rows, _FAB_N),
489 + weight_shape=(_FB_N, _FB_K),
490 + )
491 + output = torch.empty(
492 + (num_rows, _FB_N),
493 + dtype=packed_fab.dtype,
494 + device=packed_fab.device,
495 + )
496 + with torch.accelerator.device_index(packed_fab.device.index):
497 + self._compile_k(num_rows)
498 + key = (packed_fab.device.index, num_rows)
499 + self._compiled_k[key](
500 + packed_fab[:, :_FB_K],
501 + weight,
502 + output,
503 + )
504 + return output
505 +
506 +
507 +kda_skinny_gemm = KdaSkinnyGemm()