[Model] Add Kimi K3 support: model files and kernels [1/N] (#50089)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Ziming Huang <zelda.huanghuang@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Isotr0py <mozf@inferact.ai> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Summer Yang <girasoleyang@gmail.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com> Co-authored-by: Bowen Wang <abmfy@icloud.com> Co-authored-by: gnovack <novackgm@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: xiaozhoupy <peiyuanzhou1994@gmail.com> Co-authored-by: Roy Wang <yasong.wang@inferact.ai> Co-authored-by: Jeff (Junze) Ma <93145857+majunze2001@users.noreply.github.com> Co-authored-by: Codex <codex@openai.com>
Jiangyun Zhu committed
Jul 29, 2026 at 14:10 UTC
7c6729b769597541d327f666273cd972b8a5318d
109 files changed
+28009
-773
.buildkite/check-torch-abi.py
+1
@@ -13,6 +13,7 @@ from torch_abi_audit.report import ExtensionReport, PackageReport
13
# Temporary allowlist of extensions not yet on the stable ABI.
14
# Shrink and remove over time.
15
ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = (
16
+ "_flashkda_C.abi3.so",
17
"vllm_flash_attn/_vllm_fa2_C.abi3.so",
18
"vllm_flash_attn/_vllm_fa3_C.abi3.so",
19
"third_party/deep_gemm/_C*.so",
.buildkite/test_areas/kernels.yaml
-11
@@ -214,17 +214,6 @@ steps:
214
commands:
215
- pytest -v -s kernels/mamba
216
217
-- label: Kernels KDA Test
218
- timeout_in_minutes: 25
219
- device: h200_18gb
220
- source_file_dependencies:
221
- - vllm/third_party/flash_linear_attention/ops/kda.py
222
- - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
223
- - vllm/third_party/flash_linear_attention/ops/l2norm.py
224
- - tests/kernels/test_kda.py
225
- commands:
226
- - pytest -v -s kernels/test_kda.py
227
-
217
- label: Kernels DeepGEMM Test (H100)
218
key: kernels-deepgemm-test-h100
219
timeout_in_minutes: 35
.buildkite/test_areas/models_basic.yaml
+4
-2
@@ -69,9 +69,11 @@ steps:
69
- vllm/models/kimi_k3/
70
- csrc/libtorch_stable/kimi_k3/
71
- tests/models/kimi_k3/
72
+ - tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py
73
+ - tests/kernels/test_bf16_skinny_gemm.py
74
commands:
73
- # The native NVIDIA AttnRes kernel requires the SM100 family.
74
- - pytest -v -s models/kimi_k3
75
+ # The native NVIDIA Kimi K3 kernels require the SM100 family.
76
+ - pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py
77
78
- label: Basic Models Test (Other CPU) # 5min
79
key: basic-models-test-other-cpu
.pre-commit-config.yaml
+1
-1
@@ -4,7 +4,7 @@ default_install_hook_types:
4
default_stages:
5
- pre-commit # Run locally
6
- manual # Run in CI
7
-exclude: 'vllm/third_party/.*'
7
+exclude: 'vllm/third_party/.*|vllm/models/kimi_k3/nvidia/ops/third_party/.*|vllm/models/kimi_k3/amd/ops/third_party/.*'
8
repos:
9
- repo: https://github.com/astral-sh/ruff-pre-commit
10
rev: v0.14.0
CMakeLists.txt
+26
-1
@@ -421,8 +421,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
421
"csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
422
"csrc/libtorch_stable/cache_kernels.cu"
423
"csrc/libtorch_stable/cache_kernels_fused.cu"
424
+ "csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu"
425
+ "csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp"
426
"csrc/libtorch_stable/custom_all_reduce.cu"
425
- "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
427
+ "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu"
428
+ "csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu")
429
430
if(VLLM_GPU_LANG STREQUAL "CUDA" AND
431
DEFINED CMAKE_CUDA_COMPILER_VERSION AND
@@ -1079,6 +1082,23 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1082
set(MLA_ARCHS)
1083
endif()
1084
1085
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
1086
+ cuda_archs_loose_intersection(FUSED_KDA_DECODE_ARCHS
1087
+ "9.0a;10.0f;12.0f" "${CUDA_ARCHS}")
1088
+ endif()
1089
+ if(FUSED_KDA_DECODE_ARCHS)
1090
+ set(FUSED_KDA_DECODE_SRC
1091
+ "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu")
1092
+ set_gencode_flags_for_srcs(
1093
+ SRCS "${FUSED_KDA_DECODE_SRC}"
1094
+ CUDA_ARCHS "${FUSED_KDA_DECODE_ARCHS}")
1095
+ set_property(SOURCE ${FUSED_KDA_DECODE_SRC} APPEND PROPERTY
1096
+ COMPILE_OPTIONS "$<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>")
1097
+ list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_KDA_DECODE_SRC}")
1098
+ message(STATUS
1099
+ "Building fused KDA decode for archs: ${FUSED_KDA_DECODE_ARCHS}")
1100
+ endif()
1101
+
1102
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
1103
cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS
1104
"10.0f" "${CUDA_ARCHS}")
@@ -1138,6 +1158,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1158
target_compile_definitions(_C_stable_libtorch PRIVATE
1159
VLLM_ENABLE_COOPERATIVE_TOPK=1)
1160
endif()
1161
+ if(FUSED_KDA_DECODE_ARCHS)
1162
+ target_compile_definitions(_C_stable_libtorch PRIVATE
1163
+ VLLM_ENABLE_FUSED_KDA_DECODE=1)
1164
+ endif()
1165
if(KIMI_K3_ATTN_RES_ARCHS)
1166
target_compile_definitions(_C_stable_libtorch PRIVATE
1167
VLLM_ENABLE_KIMI_K3_ATTN_RES=1)
@@ -1439,6 +1463,7 @@ if (VLLM_GPU_LANG STREQUAL "CUDA")
1463
include(cmake/external_projects/deepgemm.cmake)
1464
include(cmake/external_projects/fmha_sm100.cmake)
1465
include(cmake/external_projects/flashmla.cmake)
1466
+ include(cmake/external_projects/flashkda.cmake)
1467
include(cmake/external_projects/qutlass.cmake)
1468
include(cmake/external_projects/tml_fa4.cmake)
1469
benchmarks/kernels/benchmark_k3_cutedsl_residual.py
new
+367
@@ -0,0 +1,367 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Benchmark the Kimi-K3 latent MoE addmm against CuTe residual GEMM.
4
+
5
+The benchmark covers ``BF16[M, 3584] @ BF16[7168, 3584].T + BF16[M, 7168]``
6
+with FP32 accumulation and BF16 output. Both backends execute through CUDA
7
+Graph replay. Weights and residuals rotate across buffers exceeding L2 so the
8
+comparison models the full latent MoE projection-and-add path.
9
+"""
10
+
11
+from __future__ import annotations
12
+
13
+import argparse
14
+import dataclasses
15
+import importlib.util
16
+import json
17
+import math
18
+import statistics
19
+from collections.abc import Callable, Sequence
20
+from pathlib import Path
21
+from typing import Any
22
+
23
+import cutlass
24
+import cutlass.cute as cute
25
+import torch
26
+from cuda.bindings import driver as cuda
27
+from cuda.bindings.driver import CUstream
28
+from quack.compile_utils import make_fake_tensor
29
+
30
+N = 7168
31
+K = 3584
32
+
33
+
34
+@dataclasses.dataclass(frozen=True, slots=True)
35
+class Config:
36
+ block_size: int
37
+ outputs_per_block: int
38
+ k_unroll: int
39
+ vector_width: int = 8
40
+
41
+
42
+def parse_config(value: str) -> Config:
43
+ try:
44
+ parts = [int(part) for part in value.split(",")]
45
+ except ValueError as error:
46
+ raise argparse.ArgumentTypeError(
47
+ "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]"
48
+ ) from error
49
+ if len(parts) == 3:
50
+ return Config(*parts)
51
+ if len(parts) == 4:
52
+ return Config(*parts)
53
+ raise argparse.ArgumentTypeError(
54
+ "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]"
55
+ )
56
+
57
+
58
+def production_residual_config(m: int) -> Config | None:
59
+ """The measured Latent-MoE residual config for M, from the K3 table."""
60
+ from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS
61
+
62
+ spec = KIMI_K3_PROJECTIONS.get((N, K))
63
+ config = spec.residual_config(m) if spec is not None else None
64
+ if config is None:
65
+ return None
66
+ return Config(
67
+ config.block_size,
68
+ config.outputs_per_block,
69
+ config.k_unroll,
70
+ config.vector_width,
71
+ )
72
+
73
+
74
+def candidate_configs(mode: str, selected: Config | None, m: int) -> list[Config]:
75
+ if mode == "selected":
76
+ if selected is not None:
77
+ return [selected]
78
+ # No explicit --config: fall back to the production table for this M.
79
+ config = production_residual_config(m)
80
+ return [config] if config is not None else []
81
+ if mode == "baseline":
82
+ return [Config(224, 4, 2)]
83
+ return [
84
+ Config(block_size, outputs_per_block, k_unroll, vector_width)
85
+ for vector_width in (4, 8)
86
+ for block_size in (32, 64, 128, 224, 448)
87
+ if block_size % 32 == 0 and K % (block_size * vector_width) == 0
88
+ for outputs_per_block in (1, 2, 4, 7, 8)
89
+ if N % outputs_per_block == 0
90
+ for k_unroll in (1, 2, 4)
91
+ ]
92
+
93
+
94
+def load_kernel_class(path: Path):
95
+ spec = importlib.util.spec_from_file_location("cute_skinny_device", path)
96
+ if spec is None or spec.loader is None:
97
+ raise RuntimeError(f"cannot load CuTe kernel from {path}")
98
+ module = importlib.util.module_from_spec(spec)
99
+ spec.loader.exec_module(module)
100
+ return module.CuteSkinnyGemm
101
+
102
+
103
+def stream() -> CUstream:
104
+ return CUstream(torch.cuda.current_stream().cuda_stream)
105
+
106
+
107
+def compile_kernel(kernel_class, m: int, config: Config, max_registers: int):
108
+ element_type = cutlass.BFloat16
109
+ n = cute.sym_int(divisibility=config.outputs_per_block)
110
+ k = cute.sym_int(divisibility=config.block_size * config.vector_width)
111
+ a = make_fake_tensor(element_type, (m, k), divisibility=config.vector_width)
112
+ b = make_fake_tensor(element_type, (n, k), divisibility=config.vector_width)
113
+ residual = make_fake_tensor(element_type, (m, n), divisibility=1)
114
+ c = make_fake_tensor(element_type, (m, n), divisibility=1)
115
+ kernel = kernel_class(
116
+ element_type=element_type,
117
+ num_rows=m,
118
+ block_size=config.block_size,
119
+ outputs_per_block=config.outputs_per_block,
120
+ vector_width=config.vector_width,
121
+ k_unroll=config.k_unroll,
122
+ has_residual=True,
123
+ use_pdl=True,
124
+ )
125
+ return cute.compile(
126
+ kernel,
127
+ a,
128
+ b,
129
+ residual,
130
+ c,
131
+ stream(),
132
+ options=(
133
+ "--enable-tvm-ffi --keep-cubin "
134
+ f"--ptxas-options -maxrregcount={max_registers} "
135
+ "--ptxas-options -lineinfo"
136
+ ),
137
+ )
138
+
139
+
140
+def resource_usage(compiled) -> dict[str, Any]:
141
+ executor = getattr(compiled, "_default_executor", None)
142
+ context = getattr(executor, "exec_context", None)
143
+ functions = getattr(context, "kernel_functions", None)
144
+ if not functions:
145
+ return {"resource_metrics_available": False}
146
+
147
+ def attribute(name, function) -> int:
148
+ error, value = cuda.cuFuncGetAttribute(name, function)
149
+ if error != cuda.CUresult.CUDA_SUCCESS:
150
+ raise RuntimeError(f"cuFuncGetAttribute failed with {error}")
151
+ return int(value)
152
+
153
+ registers = [
154
+ attribute(cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS, function)
155
+ for function in functions
156
+ ]
157
+ local_bytes = [
158
+ attribute(
159
+ cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES,
160
+ function,
161
+ )
162
+ for function in functions
163
+ ]
164
+ return {
165
+ "resource_metrics_available": True,
166
+ "registers_per_thread": max(registers, default=0),
167
+ "spill_bytes": max(local_bytes, default=0),
168
+ }
169
+
170
+
171
+def rotating_buffer_count(m: int, multiplier: float, limit: int) -> int:
172
+ properties = torch.cuda.get_device_properties(0)
173
+ bytes_per_pair = (N * K + m * N) * 2
174
+ target = math.ceil(multiplier * properties.L2_cache_size)
175
+ return max(2, min(limit, math.ceil(target / bytes_per_pair)))
176
+
177
+
178
+def graph_samples(
179
+ launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], None],
180
+ activation: torch.Tensor,
181
+ weights: Sequence[torch.Tensor],
182
+ residuals: Sequence[torch.Tensor],
183
+ repeats: int,
184
+ replays: int,
185
+) -> tuple[list[float], list[torch.Tensor]]:
186
+ outputs = [torch.empty_like(residual) for residual in residuals]
187
+ for weight, residual, output in zip(weights, residuals, outputs):
188
+ launch(activation, weight, residual, output)
189
+ torch.accelerator.synchronize()
190
+
191
+ graph = torch.cuda.CUDAGraph()
192
+ with torch.cuda.graph(graph):
193
+ for weight, residual, output in zip(weights, residuals, outputs):
194
+ launch(activation, weight, residual, output)
195
+ for _ in range(20):
196
+ graph.replay()
197
+ torch.accelerator.synchronize()
198
+
199
+ samples = []
200
+ for _ in range(repeats):
201
+ start = torch.cuda.Event(enable_timing=True)
202
+ end = torch.cuda.Event(enable_timing=True)
203
+ start.record()
204
+ for _ in range(replays):
205
+ graph.replay()
206
+ end.record()
207
+ end.synchronize()
208
+ samples.append(start.elapsed_time(end) * 1000.0 / (replays * len(weights)))
209
+ return samples, outputs
210
+
211
+
212
+def summarize(samples: Sequence[float]) -> dict[str, Any]:
213
+ ordered = sorted(samples)
214
+
215
+ def percentile(fraction: float) -> float:
216
+ position = fraction * (len(ordered) - 1)
217
+ lower = math.floor(position)
218
+ upper = math.ceil(position)
219
+ if lower == upper:
220
+ return ordered[lower]
221
+ weight = position - lower
222
+ return ordered[lower] * (1.0 - weight) + ordered[upper] * weight
223
+
224
+ mean = statistics.mean(samples)
225
+ return {
226
+ "median_us": statistics.median(samples),
227
+ "p10_us": percentile(0.1),
228
+ "p90_us": percentile(0.9),
229
+ "mean_us": mean,
230
+ "cv_pct": statistics.pstdev(samples) / mean * 100.0,
231
+ "samples_us": list(samples),
232
+ }
233
+
234
+
235
+def correctness(
236
+ output: torch.Tensor,
237
+ activation: torch.Tensor,
238
+ weight: torch.Tensor,
239
+ residual: torch.Tensor,
240
+) -> dict[str, Any]:
241
+ actual = output.float()
242
+ reference = activation.float() @ weight.float().t() + residual.float()
243
+ error = (actual - reference).abs()
244
+ scaled_error = error / (reference.abs() + 1.0)
245
+ cosine = torch.nn.functional.cosine_similarity(
246
+ actual.flatten(), reference.flatten(), dim=0
247
+ ).item()
248
+ return {
249
+ "valid": cosine > 0.999,
250
+ "cosine": cosine,
251
+ "max_abs_error": error.max().item(),
252
+ "max_scaled_error": scaled_error.max().item(),
253
+ }
254
+
255
+
256
+def main() -> None:
257
+ parser = argparse.ArgumentParser(description=__doc__)
258
+ parser.add_argument("--kernel", type=Path, required=True)
259
+ parser.add_argument("--output", type=Path, required=True)
260
+ parser.add_argument(
261
+ "--mode", choices=("baseline", "sweep", "selected"), default="baseline"
262
+ )
263
+ parser.add_argument("--config", type=parse_config)
264
+ parser.add_argument("--m", type=int, action="append")
265
+ parser.add_argument("--config-shard", type=int, default=0)
266
+ parser.add_argument("--num-config-shards", type=int, default=1)
267
+ parser.add_argument("--repeats", type=int, default=21)
268
+ parser.add_argument("--replays", type=int, default=200)
269
+ parser.add_argument("--cache-multiplier", type=float, default=3.0)
270
+ parser.add_argument("--max-buffers", type=int, default=32)
271
+ parser.add_argument("--max-registers", type=int, default=64)
272
+ args = parser.parse_args()
273
+
274
+ token_counts = args.m or list(range(1, 17))
275
+ if any(not 1 <= m <= 16 for m in token_counts):
276
+ raise ValueError("expected 1 <= M <= 16")
277
+ if not 0 <= args.config_shard < args.num_config_shards:
278
+ raise ValueError("config shard must be in [0, num_config_shards)")
279
+ torch.accelerator.set_device_index(0)
280
+ if torch.cuda.get_device_capability() != (10, 3):
281
+ raise RuntimeError("this benchmark requires SM103")
282
+
283
+ kernel_class = load_kernel_class(args.kernel)
284
+ properties = torch.cuda.get_device_properties(0)
285
+ metadata = {
286
+ "device": properties.name,
287
+ "compute_capability": list(torch.cuda.get_device_capability()),
288
+ "torch_version": torch.__version__,
289
+ "cuda_version": torch.version.cuda,
290
+ }
291
+ args.output.parent.mkdir(parents=True, exist_ok=True)
292
+ with args.output.open("w", encoding="utf-8") as output_file:
293
+ for m in token_counts:
294
+ configs = candidate_configs(args.mode, args.config, m)
295
+ torch.manual_seed(20260722 + m)
296
+ count = rotating_buffer_count(m, args.cache_multiplier, args.max_buffers)
297
+ activation = torch.randn((m, K), device="cuda", dtype=torch.bfloat16)
298
+ weights = [
299
+ torch.randn((N, K), device="cuda", dtype=torch.bfloat16)
300
+ for _ in range(count)
301
+ ]
302
+ residuals = [
303
+ torch.randn((m, N), device="cuda", dtype=torch.bfloat16)
304
+ for _ in range(count)
305
+ ]
306
+ candidates: list[tuple[str, Config | None]] = [("cublas_addmm", None)]
307
+ candidates.extend(
308
+ ("cute_residual", config)
309
+ for index, config in enumerate(configs)
310
+ if index % args.num_config_shards == args.config_shard
311
+ )
312
+ for backend, config in candidates:
313
+ row: dict[str, Any] = {
314
+ "m": m,
315
+ "n": N,
316
+ "k": K,
317
+ "backend": backend,
318
+ "mode": args.mode,
319
+ "config": dataclasses.asdict(config) if config else {},
320
+ "num_buffers": count,
321
+ "cache_multiplier": args.cache_multiplier,
322
+ **metadata,
323
+ }
324
+ try:
325
+ if backend == "cublas_addmm":
326
+ launch = lambda a, b, residual, c: torch.addmm(
327
+ residual, a, b.t(), out=c
328
+ )
329
+ else:
330
+ if config is None:
331
+ raise AssertionError("missing CuTe config")
332
+ compiled = compile_kernel(
333
+ kernel_class, m, config, args.max_registers
334
+ )
335
+ launch = lambda a, b, residual, c, fn=compiled: fn(
336
+ a, b, residual, c, stream()
337
+ )
338
+ row.update(resource_usage(compiled))
339
+ samples, outputs = graph_samples(
340
+ launch,
341
+ activation,
342
+ weights,
343
+ residuals,
344
+ args.repeats,
345
+ args.replays,
346
+ )
347
+ row.update(
348
+ correctness(outputs[0], activation, weights[0], residuals[0])
349
+ )
350
+ row.update(summarize(samples))
351
+ except Exception as error: # noqa: BLE001
352
+ row.update(
353
+ {
354
+ "valid": False,
355
+ "error": f"{type(error).__name__}: {error}",
356
+ }
357
+ )
358
+ output_file.write(json.dumps(row, sort_keys=True) + "\n")
359
+ output_file.flush()
360
+ print(json.dumps(row, sort_keys=True), flush=True)
361
+
362
+ del activation, weights, residuals
363
+ torch.accelerator.empty_cache()
364
+
365
+
366
+if __name__ == "__main__":
367
+ main()
benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py
new
+806
@@ -0,0 +1,806 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Benchmark the Kimi K3 latent-MoE tail and its up-projection kernels.
4
+
5
+The ``up-projection`` subcommand isolates the TP-local dynamic and static-M
6
+skinny GEMMs. It rotates weights through a working set larger than L2 to model
7
+successive model layers.
8
+
9
+The ``whole-tail`` subcommand measures the distributed operator. Its reference
10
+path includes two AllReduces, RMSNorm, the replicated up-projection, and the
11
+final add. CUDA-event samples report the slowest rank so cross-rank skew is
12
+included.
13
+
14
+Examples:
15
+
16
+.. code-block:: console
17
+
18
+ .venv/bin/python \
19
+ benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py up-projection
20
+
21
+ torchrun --nproc-per-node=8 \
22
+ benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py whole-tail
23
+
24
+For multi-node runs, launch one ``torchrun`` agent per node and use a shared
25
+rendezvous endpoint.
26
+"""
27
+
28
+from __future__ import annotations
29
+
30
+import argparse
31
+import json
32
+import math
33
+import os
34
+import statistics
35
+from collections.abc import Callable, Sequence
36
+from dataclasses import asdict
37
+from pathlib import Path
38
+from typing import Any
39
+
40
+import cutlass
41
+import cutlass.utils as utils
42
+import torch
43
+import torch.distributed as dist
44
+import torch.nn.functional as F
45
+from cuda.bindings import driver as cuda
46
+
47
+from vllm.distributed import get_tp_group
48
+from vllm.distributed.parallel_state import (
49
+ init_distributed_environment,
50
+ initialize_model_parallel,
51
+ set_custom_all_reduce,
52
+)
53
+from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
54
+from vllm.models.kimi_k3.nvidia.ops import latent_moe_tail
55
+from vllm.models.kimi_k3.nvidia.ops.cute_dsl.latent_moe_tail import (
56
+ fused_add_multicast_gemm,
57
+ fused_add_multicast_skinny_gemm,
58
+)
59
+
60
+HIDDEN_SIZE = 7168
61
+LATENT_SIZE = 3584
62
+RMS_EPS = 0.1
63
+MAX_NUM_TOKENS = 16
64
+MMA_TILER_MN = (64, 32)
65
+CLUSTER_SHAPE_MN = (1, 8)
66
+B_PRIME_STAGES = 2
67
+
68
+
69
+def parse_up_projection_config(
70
+ value: str,
71
+) -> fused_add_multicast_skinny_gemm.SkinnyConfig:
72
+ try:
73
+ values = [int(part) for part in value.split(",")]
74
+ except ValueError as error:
75
+ raise argparse.ArgumentTypeError(
76
+ "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]]"
77
+ ) from error
78
+ if len(values) in (3, 4):
79
+ return fused_add_multicast_skinny_gemm.SkinnyConfig(*values)
80
+ if len(values) == 5 and values[4] in (0, 1):
81
+ return fused_add_multicast_skinny_gemm.SkinnyConfig(
82
+ *values[:4],
83
+ prefetch_b_before_pdl=bool(values[4]),
84
+ )
85
+ raise argparse.ArgumentTypeError(
86
+ "config must be BLOCK,OUTPUTS,K_UNROLL"
87
+ "[,VECTOR_WIDTH[,PREFETCH_B]], where PREFETCH_B is 0 or 1"
88
+ )
89
+
90
+
91
+def parse_tail_skinny_config(
92
+ value: str,
93
+) -> tuple[int, fused_add_multicast_skinny_gemm.SkinnyConfig]:
94
+ try:
95
+ values = [int(part) for part in value.split(",")]
96
+ except ValueError as error:
97
+ raise argparse.ArgumentTypeError(
98
+ "config must be M,BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]]"
99
+ ) from error
100
+ if len(values) == 4:
101
+ num_tokens, *config = values
102
+ return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(*config)
103
+ if len(values) == 5:
104
+ num_tokens, *config = values
105
+ return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(*config)
106
+ if len(values) == 6 and values[5] in (0, 1):
107
+ num_tokens, block, outputs, unroll, vector_width, prefetch = values
108
+ return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(
109
+ block,
110
+ outputs,
111
+ unroll,
112
+ vector_width,
113
+ bool(prefetch),
114
+ )
115
+ raise argparse.ArgumentTypeError(
116
+ "config must be M,BLOCK,OUTPUTS,K_UNROLL"
117
+ "[,VECTOR_WIDTH[,PREFETCH_B]], where PREFETCH_B is 0 or 1"
118
+ )
119
+
120
+
121
+def parse_args() -> argparse.Namespace:
122
+ parser = argparse.ArgumentParser(description=__doc__)
123
+ subparsers = parser.add_subparsers(dest="scope", required=True)
124
+
125
+ up_projection = subparsers.add_parser(
126
+ "up-projection",
127
+ help="Benchmark the isolated TP-local up-projection kernels.",
128
+ )
129
+ up_projection.add_argument(
130
+ "--backend",
131
+ choices=("dynamic", "skinny", "both"),
132
+ default="both",
133
+ )
134
+ up_projection.add_argument("--tp-size", type=int, default=16)
135
+ up_projection.add_argument(
136
+ "--num-tokens",
137
+ type=int,
138
+ nargs="+",
139
+ default=[*range(1, 9), 16],
140
+ )
141
+ up_projection.add_argument(
142
+ "--skinny-config",
143
+ type=parse_up_projection_config,
144
+ action="append",
145
+ help="Benchmark a static-M config for every selected token count.",
146
+ )
147
+ up_projection.add_argument("--cache-multiplier", type=float, default=2.0)
148
+ up_projection.add_argument("--max-weights", type=int, default=64)
149
+ up_projection.add_argument("--warmup-replays", type=int, default=10)
150
+ up_projection.add_argument("--samples", type=int, default=31)
151
+ up_projection.add_argument("--output", type=Path)
152
+
153
+ whole_tail = subparsers.add_parser(
154
+ "whole-tail",
155
+ help="Benchmark the distributed latent-MoE tail operator.",
156
+ )
157
+ whole_tail.add_argument(
158
+ "--backend",
159
+ choices=("reference", "fused", "both"),
160
+ default="both",
161
+ )
162
+ whole_tail.add_argument(
163
+ "--num-tokens",
164
+ type=int,
165
+ nargs="+",
166
+ default=[1, 5, 8, 16],
167
+ )
168
+ whole_tail.add_argument("--warmup-replays", type=int, default=20)
169
+ whole_tail.add_argument("--samples", type=int, default=51)
170
+ whole_tail.add_argument(
171
+ "--skinny-max-num-tokens",
172
+ type=int,
173
+ nargs="+",
174
+ help="Override the fused operator's static-M cutoff; use 0 for dynamic-only.",
175
+ )
176
+ whole_tail.add_argument(
177
+ "--skinny-config",
178
+ type=parse_tail_skinny_config,
179
+ action="append",
180
+ help="Override one static-M config for tuning.",
181
+ )
182
+ whole_tail.add_argument("--output", type=Path)
183
+ return parser.parse_args()
184
+
185
+
186
+def percentile(samples: Sequence[float], fraction: float) -> float:
187
+ ordered = sorted(samples)
188
+ position = fraction * (len(ordered) - 1)
189
+ lower = math.floor(position)
190
+ upper = math.ceil(position)
191
+ if lower == upper:
192
+ return ordered[lower]
193
+ upper_weight = position - lower
194
+ return ordered[lower] * (1.0 - upper_weight) + ordered[upper] * upper_weight
195
+
196
+
197
+def summarize(samples_us: Sequence[float]) -> dict[str, Any]:
198
+ mean_us = statistics.mean(samples_us)
199
+ return {
200
+ "median_us": statistics.median(samples_us),
201
+ "p10_us": percentile(samples_us, 0.1),
202
+ "p90_us": percentile(samples_us, 0.9),
203
+ "mean_us": mean_us,
204
+ "cv_pct": statistics.pstdev(samples_us) / mean_us * 100.0,
205
+ "samples_us": list(samples_us),
206
+ }
207
+
208
+
209
+def rotating_weight_count(
210
+ shard_size: int,
211
+ cache_multiplier: float,
212
+ limit: int,
213
+) -> int:
214
+ properties = torch.cuda.get_device_properties(
215
+ torch.accelerator.current_device_index()
216
+ )
217
+ weight_bytes = shard_size * LATENT_SIZE * 2
218
+ target_bytes = math.ceil(properties.L2_cache_size * cache_multiplier)
219
+ return max(2, min(limit, math.ceil(target_bytes / weight_bytes)))
220
+
221
+
222
+def capture_up_projection_graph(
223
+ launches: Sequence[Callable[[], None]],
224
+) -> torch.cuda.CUDAGraph:
225
+ for launch in launches:
226
+ launch()
227
+ torch.accelerator.synchronize()
228
+
229
+ graph = torch.cuda.CUDAGraph()
230
+ with torch.cuda.graph(graph):
231
+ for launch in launches:
232
+ launch()
233
+ torch.accelerator.synchronize()
234
+ return graph
235
+
236
+
237
+def benchmark_up_projection_graph(
238
+ graph: torch.cuda.CUDAGraph,
239
+ *,
240
+ operations_per_replay: int,
241
+ warmup_replays: int,
242
+ samples: int,
243
+) -> dict[str, Any]:
244
+ for _ in range(warmup_replays):
245
+ graph.replay()
246
+ torch.accelerator.synchronize()
247
+
248
+ samples_us = []
249
+ start = torch.cuda.Event(enable_timing=True)
250
+ end = torch.cuda.Event(enable_timing=True)
251
+ for _ in range(samples):
252
+ start.record()
253
+ graph.replay()
254
+ end.record()
255
+ end.synchronize()
256
+ samples_us.append(start.elapsed_time(end) * 1000.0 / operations_per_replay)
257
+ return summarize(samples_us)
258
+
259
+
260
+class DynamicKernel:
261
+ def __init__(
262
+ self,
263
+ shard_size: int,
264
+ mailbox: torch.Tensor,
265
+ shared_shard: torch.Tensor,
266
+ ) -> None:
267
+ self.shard_size = shard_size
268
+ self.mailbox = mailbox
269
+ self.mailbox_c = fused_add_multicast_gemm._as_cute(mailbox)
270
+ compile_latent = torch.empty(
271
+ (1, MAX_NUM_TOKENS, LATENT_SIZE),
272
+ dtype=torch.bfloat16,
273
+ device=mailbox.device,
274
+ )
275
+ compile_weight = torch.empty(
276
+ (1, shard_size, LATENT_SIZE),
277
+ dtype=torch.bfloat16,
278
+ device=mailbox.device,
279
+ )
280
+ cluster_size = math.prod(CLUSTER_SHAPE_MN)
281
+ max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size)
282
+ self.compiled = fused_add_multicast_gemm.compile_kernel(
283
+ (MAX_NUM_TOKENS, shard_size, LATENT_SIZE, 1),
284
+ fused_add_multicast_gemm._as_cute(
285
+ compile_latent,
286
+ dynamic_m=True,
287
+ ),
288
+ fused_add_multicast_gemm._as_cute(compile_weight),
289
+ self.mailbox_c,
290
+ fused_add_multicast_gemm._as_cute(shared_shard),
291
+ HIDDEN_SIZE,
292
+ shard_size,
293
+ MMA_TILER_MN,
294
+ CLUSTER_SHAPE_MN,
295
+ max_active_clusters,
296
+ B_PRIME_STAGES,
297
+ )
298
+
299
+ def launch(
300
+ self,
301
+ latent: torch.Tensor,
302
+ weight: torch.Tensor,
303
+ shared_shard: torch.Tensor,
304
+ ) -> None:
305
+ stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
306
+ self.compiled(
307
+ fused_add_multicast_gemm._as_cute(
308
+ latent.unsqueeze(0),
309
+ dynamic_m=True,
310
+ ),
311
+ fused_add_multicast_gemm._as_cute(weight.unsqueeze(0)),
312
+ self.mailbox_c,
313
+ fused_add_multicast_gemm._as_cute(shared_shard),
314
+ cutlass.Int64(latent.shape[0]),
315
+ cutlass.Int64(self.mailbox.data_ptr()),
316
+ stream,
317
+ )
318
+
319
+
320
+class SkinnyKernel:
321
+ def __init__(
322
+ self,
323
+ num_tokens: int,
324
+ shard_size: int,
325
+ config: fused_add_multicast_skinny_gemm.SkinnyConfig,
326
+ ) -> None:
327
+ self.compiled = fused_add_multicast_skinny_gemm.compile_kernel(
328
+ num_rows=num_tokens,
329
+ latent_dim=LATENT_SIZE,
330
+ hidden_dim=HIDDEN_SIZE,
331
+ shard_dim=shard_size,
332
+ config=config,
333
+ )
334
+
335
+ def launch(
336
+ self,
337
+ latent: torch.Tensor,
338
+ weight: torch.Tensor,
339
+ shared_shard: torch.Tensor,
340
+ mailbox: torch.Tensor,
341
+ ) -> None:
342
+ self.compiled(
343
+ fused_add_multicast_skinny_gemm._as_cute(latent),
344
+ fused_add_multicast_skinny_gemm._as_cute(weight),
345
+ fused_add_multicast_skinny_gemm._as_cute(shared_shard),
346
+ cutlass.Int64(mailbox.data_ptr()),
347
+ cuda.CUstream(torch.cuda.current_stream().cuda_stream),
348
+ )
349
+
350
+
351
+def check_up_projection_output(
352
+ actual: torch.Tensor,
353
+ latent: torch.Tensor,
354
+ weight: torch.Tensor,
355
+ shared_shard: torch.Tensor,
356
+) -> None:
357
+ gemm = F.linear(latent.float(), weight.float()).to(torch.bfloat16)
358
+ expected = (gemm.float() + shared_shard.float()).to(torch.bfloat16)
359
+ torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
360
+
361
+
362
+def make_up_projection_launches(
363
+ launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], None],
364
+ latent: torch.Tensor,
365
+ weights: Sequence[torch.Tensor],
366
+ shared_shard: torch.Tensor,
367
+) -> list[Callable[[], None]]:
368
+ return [
369
+ lambda weight=weight: launch(latent, weight, shared_shard) for weight in weights
370
+ ]
371
+
372
+
373
+def benchmark_up_projection(args: argparse.Namespace) -> None:
374
+ if args.tp_size <= 0 or HIDDEN_SIZE % args.tp_size:
375
+ raise ValueError("TP size must be positive and divide the hidden size")
376
+ if any(not 1 <= num_tokens <= MAX_NUM_TOKENS for num_tokens in args.num_tokens):
377
+ raise ValueError("--num-tokens values must be in [1, 16]")
378
+ if args.cache_multiplier <= 0 or args.max_weights <= 0:
379
+ raise ValueError("cache multiplier and max weights must be positive")
380
+ if args.warmup_replays < 0 or args.samples <= 0:
381
+ raise ValueError("warmup replays must be nonnegative and samples positive")
382
+
383
+ torch.accelerator.set_device_index(0)
384
+ device = torch.device("cuda", 0)
385
+ if torch.cuda.get_device_capability(device)[0] != 10:
386
+ raise RuntimeError("Kimi K3 latent-MoE tail requires SM100")
387
+
388
+ shard_size = HIDDEN_SIZE // args.tp_size
389
+ weight_count = rotating_weight_count(
390
+ shard_size,
391
+ args.cache_multiplier,
392
+ args.max_weights,
393
+ )
394
+ torch.manual_seed(20260726)
395
+ weights = [
396
+ torch.randn(
397
+ (shard_size, LATENT_SIZE),
398
+ dtype=torch.bfloat16,
399
+ device=device,
400
+ )
401
+ / LATENT_SIZE**0.5
402
+ for _ in range(weight_count)
403
+ ]
404
+ mailbox = torch.empty(
405
+ (1, MAX_NUM_TOKENS, HIDDEN_SIZE),
406
+ dtype=torch.bfloat16,
407
+ device=device,
408
+ )
409
+ shared = torch.randn(
410
+ (MAX_NUM_TOKENS, HIDDEN_SIZE),
411
+ dtype=torch.bfloat16,
412
+ device=device,
413
+ )
414
+ shared_shard = shared[:, :shard_size]
415
+ use_dynamic = args.backend in ("dynamic", "both")
416
+ use_skinny = args.backend in ("skinny", "both")
417
+ dynamic_kernel = (
418
+ DynamicKernel(shard_size, mailbox, shared_shard) if use_dynamic else None
419
+ )
420
+
421
+ results = []
422
+ for num_tokens in args.num_tokens:
423
+ latent = torch.randn(
424
+ (num_tokens, LATENT_SIZE),
425
+ dtype=torch.bfloat16,
426
+ device=device,
427
+ )
428
+ result: dict[str, Any] = {"num_tokens": num_tokens}
429
+ if dynamic_kernel is not None:
430
+ launches = make_up_projection_launches(
431
+ dynamic_kernel.launch,
432
+ latent,
433
+ weights,
434
+ shared_shard,
435
+ )
436
+ graph = capture_up_projection_graph(launches)
437
+ result["dynamic"] = benchmark_up_projection_graph(
438
+ graph,
439
+ operations_per_replay=len(launches),
440
+ warmup_replays=args.warmup_replays,
441
+ samples=args.samples,
442
+ )
443
+ check_up_projection_output(
444
+ mailbox[0, :num_tokens, :shard_size],
445
+ latent,
446
+ weights[-1],
447
+ shared_shard[:num_tokens],
448
+ )
449
+ if use_skinny:
450
+ configs = args.skinny_config or [
451
+ fused_add_multicast_skinny_gemm.config_for_m(
452
+ num_tokens,
453
+ shard_size,
454
+ )
455
+ ]
456
+ skinny_results = []
457
+ for config in configs:
458
+ skinny_kernel = SkinnyKernel(num_tokens, shard_size, config)
459
+
460
+ def launch_skinny(
461
+ latent: torch.Tensor,
462
+ weight: torch.Tensor,
463
+ shared_shard: torch.Tensor,
464
+ *,
465
+ skinny_kernel: SkinnyKernel = skinny_kernel,
466
+ num_tokens: int = num_tokens,
467
+ ) -> None:
468
+ skinny_kernel.launch(
469
+ latent,
470
+ weight,
471
+ shared_shard[:num_tokens],
472
+ mailbox,
473
+ )
474
+
475
+ launches = make_up_projection_launches(
476
+ launch_skinny,
477
+ latent,
478
+ weights,
479
+ shared_shard,
480
+ )
481
+ graph = capture_up_projection_graph(launches)
482
+ timing = benchmark_up_projection_graph(
483
+ graph,
484
+ operations_per_replay=len(launches),
485
+ warmup_replays=args.warmup_replays,
486
+ samples=args.samples,
487
+ )
488
+ check_up_projection_output(
489
+ mailbox[0, :num_tokens, :shard_size],
490
+ latent,
491
+ weights[-1],
492
+ shared_shard[:num_tokens],
493
+ )
494
+ skinny_results.append(
495
+ {
496
+ "config": asdict(config),
497
+ **timing,
498
+ }
499
+ )
500
+ result["skinny"] = skinny_results
501
+ results.append(result)
502
+
503
+ properties = torch.cuda.get_device_properties(device)
504
+ report = {
505
+ "scope": "up-projection",
506
+ "device": properties.name,
507
+ "compute_capability": list(torch.cuda.get_device_capability(device)),
508
+ "tp_size": args.tp_size,
509
+ "shard_size": shard_size,
510
+ "weight_count": weight_count,
511
+ "cache_multiplier": args.cache_multiplier,
512
+ "warmup_replays": args.warmup_replays,
513
+ "samples": args.samples,
514
+ "results": results,
515
+ }
516
+ rendered = json.dumps(report, indent=2)
517
+ print(rendered, flush=True)
518
+ if args.output is not None:
519
+ args.output.parent.mkdir(parents=True, exist_ok=True)
520
+ args.output.write_text(rendered + "\n", encoding="utf-8")
521
+
522
+
523
+def capture_tail_graph(
524
+ operation: Callable[[], torch.Tensor],
525
+ cpu_group: dist.ProcessGroup,
526
+) -> tuple[torch.cuda.CUDAGraph, torch.Tensor]:
527
+ for _ in range(3):
528
+ dist.barrier(group=cpu_group)
529
+ output = operation()
530
+ torch.accelerator.synchronize()
531
+
532
+ dist.barrier(group=cpu_group)
533
+ graph = torch.cuda.CUDAGraph()
534
+ with torch.cuda.graph(graph):
535
+ output = operation()
536
+ torch.accelerator.synchronize()
537
+ return graph, output
538
+
539
+
540
+def benchmark_tail_graph(
541
+ graph: torch.cuda.CUDAGraph,
542
+ *,
543
+ warmup_replays: int,
544
+ samples: int,
545
+ device_group: dist.ProcessGroup,
546
+ cpu_group: dist.ProcessGroup,
547
+) -> dict[str, Any]:
548
+ for _ in range(warmup_replays):
549
+ graph.replay()
550
+ torch.accelerator.synchronize()
551
+
552
+ dist.barrier(group=cpu_group)
553
+ starts = [torch.cuda.Event(enable_timing=True) for _ in range(samples + 1)]
554
+ ends = [torch.cuda.Event(enable_timing=True) for _ in range(samples + 1)]
555
+ for start, end in zip(starts, ends):
556
+ start.record()
557
+ graph.replay()
558
+ end.record()
559
+ torch.accelerator.synchronize()
560
+
561
+ samples_us = torch.tensor(
562
+ [start.elapsed_time(end) * 1000.0 for start, end in zip(starts, ends)],
563
+ dtype=torch.float64,
564
+ device=torch.accelerator.current_device_index(),
565
+ )
566
+ dist.all_reduce(samples_us, op=dist.ReduceOp.MAX, group=device_group)
567
+ return summarize(samples_us[1:].tolist())
568
+
569
+
570
+def make_inputs(
571
+ num_tokens: int,
572
+ rank: int,
573
+ device: torch.device,
574
+) -> tuple[torch.Tensor, torch.Tensor]:
575
+ torch.manual_seed(20260726 + 100 * num_tokens + rank)
576
+ routed = torch.randn(
577
+ (num_tokens, LATENT_SIZE),
578
+ dtype=torch.bfloat16,
579
+ device=device,
580
+ ).mul_(0.01)
581
+ shared = torch.randn(
582
+ (num_tokens, HIDDEN_SIZE),
583
+ dtype=torch.bfloat16,
584
+ device=device,
585
+ )
586
+ return routed, shared
587
+
588
+
589
+def make_reference(
590
+ routed: torch.Tensor,
591
+ shared: torch.Tensor,
592
+ rms_weight: torch.Tensor,
593
+ up_weight: torch.Tensor,
594
+ device_group: dist.ProcessGroup,
595
+) -> Callable[[], torch.Tensor]:
596
+ routed_workspace = torch.empty_like(routed)
597
+ shared_workspace = torch.empty_like(shared)
598
+
599
+ def reference() -> torch.Tensor:
600
+ routed_workspace.copy_(routed)
601
+ dist.all_reduce(routed_workspace, group=device_group)
602
+ normalized = F.rms_norm(
603
+ routed_workspace,
604
+ (LATENT_SIZE,),
605
+ rms_weight,
606
+ RMS_EPS,
607
+ )
608
+ projected = F.linear(normalized, up_weight)
609
+ shared_workspace.copy_(shared)
610
+ dist.all_reduce(shared_workspace, group=device_group)
611
+ return projected.add(shared_workspace)
612
+
613
+ return reference
614
+
615
+
616
+def check_fused_output(
617
+ fused_output: torch.Tensor,
618
+ reference: Callable[[], torch.Tensor],
619
+ cpu_group: dist.ProcessGroup,
620
+) -> None:
621
+ dist.barrier(group=cpu_group)
622
+ expected = reference()
623
+ torch.testing.assert_close(fused_output, expected, atol=8e-2, rtol=3e-2)
624
+
625
+
626
+def benchmark_whole_tail(args: argparse.Namespace) -> None:
627
+ if any(not 1 <= num_tokens <= 16 for num_tokens in args.num_tokens):
628
+ raise ValueError("--num-tokens values must be in [1, 16]")
629
+ if args.warmup_replays < 0 or args.samples <= 0:
630
+ raise ValueError("warmup replays must be nonnegative and samples positive")
631
+ if args.skinny_max_num_tokens is not None and any(
632
+ not 0 <= cutoff <= 8 for cutoff in args.skinny_max_num_tokens
633
+ ):
634
+ raise ValueError("--skinny-max-num-tokens must be in [0, 8]")
635
+ skinny_configs = dict(args.skinny_config or ())
636
+ if len(skinny_configs) != len(args.skinny_config or ()):
637
+ raise ValueError("--skinny-config must not repeat an M value")
638
+ if any(not 1 <= num_tokens <= 8 for num_tokens in skinny_configs):
639
+ raise ValueError("--skinny-config M values must be in [1, 8]")
640
+ if not {"RANK", "WORLD_SIZE", "LOCAL_RANK"} <= os.environ.keys():
641
+ raise RuntimeError("launch this benchmark with torchrun")
642
+
643
+ rank = int(os.environ["RANK"])
644
+ world_size = int(os.environ["WORLD_SIZE"])
645
+ local_rank = int(os.environ["LOCAL_RANK"])
646
+ device = torch.device("cuda", local_rank)
647
+ torch.accelerator.set_device_index(device)
648
+ init_distributed_environment()
649
+ if world_size > 8:
650
+ set_custom_all_reduce(False)
651
+ initialize_model_parallel(tensor_model_parallel_size=world_size)
652
+ device_group = get_tp_group().device_group
653
+ cpu_group = dist.new_group(backend="gloo")
654
+
655
+ if torch.cuda.get_device_capability(device)[0] != 10:
656
+ raise RuntimeError("Kimi K3 latent-MoE tail requires SM100")
657
+
658
+ torch.manual_seed(20260726)
659
+ rms_weight = 1 + 0.1 * torch.randn(
660
+ LATENT_SIZE,
661
+ dtype=torch.bfloat16,
662
+ device=device,
663
+ )
664
+ up_weight = (
665
+ torch.randn(
666
+ (HIDDEN_SIZE, LATENT_SIZE),
667
+ dtype=torch.bfloat16,
668
+ device=device,
669
+ )
670
+ / LATENT_SIZE**0.5
671
+ )
672
+
673
+ use_reference = args.backend in ("reference", "both")
674
+ use_fused = args.backend in ("fused", "both")
675
+ fused_ops = []
676
+ if use_fused:
677
+ production_config_for_m = fused_add_multicast_skinny_gemm.config_for_m
678
+
679
+ def config_for_m(
680
+ num_rows: int,
681
+ shard_dim: int = 896,
682
+ ) -> fused_add_multicast_skinny_gemm.SkinnyConfig:
683
+ config = skinny_configs.get(num_rows)
684
+ if config is not None:
685
+ return config
686
+ return production_config_for_m(num_rows, shard_dim)
687
+
688
+ fused_add_multicast_skinny_gemm.config_for_m = config_for_m
689
+ cutoffs = args.skinny_max_num_tokens or [latent_moe_tail._SKINNY_MAX_NUM_TOKENS]
690
+ for cutoff in cutoffs:
691
+ latent_moe_tail._SKINNY_MAX_NUM_TOKENS = cutoff
692
+ latent_moe_tail.KimiK3LatentMoETailOp._instances.clear()
693
+ fused_ops.append(
694
+ (
695
+ cutoff,
696
+ latent_moe_tail.KimiK3LatentMoETailOp.initialize(
697
+ hidden_size=HIDDEN_SIZE,
698
+ latent_size=LATENT_SIZE,
699
+ dtype=torch.bfloat16,
700
+ device=device,
701
+ rms_eps=RMS_EPS,
702
+ ),
703
+ )
704
+ )
705
+ cutedsl_warmup()
706
+
707
+ results = []
708
+ for num_tokens in args.num_tokens:
709
+ routed, shared = make_inputs(num_tokens, rank, device)
710
+ reference = make_reference(
711
+ routed,
712
+ shared,
713
+ rms_weight,
714
+ up_weight,
715
+ device_group,
716
+ )
717
+ result: dict[str, Any] = {"num_tokens": num_tokens}
718
+ if use_reference:
719
+ reference_graph, _ = capture_tail_graph(reference, cpu_group)
720
+ result["reference"] = benchmark_tail_graph(
721
+ reference_graph,
722
+ warmup_replays=args.warmup_replays,
723
+ samples=args.samples,
724
+ device_group=device_group,
725
+ cpu_group=cpu_group,
726
+ )
727
+ for cutoff, fused_op in fused_ops:
728
+
729
+ def fused(
730
+ routed: torch.Tensor = routed,
731
+ shared: torch.Tensor = shared,
732
+ fused_op: latent_moe_tail.KimiK3LatentMoETailOp = fused_op,
733
+ ) -> torch.Tensor:
734
+ return fused_op(routed, shared, rms_weight, up_weight)
735
+
736
+ fused_graph, fused_output = capture_tail_graph(fused, cpu_group)
737
+ fused_key = "fused" if len(fused_ops) == 1 else f"fused_skinny_max_{cutoff}"
738
+ result[fused_key] = benchmark_tail_graph(
739
+ fused_graph,
740
+ warmup_replays=args.warmup_replays,
741
+ samples=args.samples,
742
+ device_group=device_group,
743
+ cpu_group=cpu_group,
744
+ )
745
+ check_fused_output(fused_output, reference, cpu_group)
746
+ if "reference" in result:
747
+ speedup = (
748
+ result["reference"]["median_us"] / result[fused_key]["median_us"]
749
+ )
750
+ if len(fused_ops) == 1:
751
+ result["speedup"] = speedup
752
+ else:
753
+ result[f"{fused_key}_speedup"] = speedup
754
+ results.append(result)
755
+
756
+ properties = torch.cuda.get_device_properties(device)
757
+ report = {
758
+ "scope": "whole-tail",
759
+ "device": properties.name,
760
+ "compute_capability": list(torch.cuda.get_device_capability(device)),
761
+ "world_size": world_size,
762
+ "torch_version": torch.__version__,
763
+ "cuda_version": torch.version.cuda,
764
+ "warmup_replays": args.warmup_replays,
765
+ "samples": args.samples,
766
+ "skinny_max_num_tokens": [cutoff for cutoff, _ in fused_ops],
767
+ "skinny_configs": {
768
+ str(num_tokens): asdict(config)
769
+ for num_tokens, config in skinny_configs.items()
770
+ },
771
+ "timing_scope": {
772
+ "reference": (
773
+ "two input copies, two AllReduces, RMSNorm, full replicated "
774
+ "up-projection GEMM, and final add"
775
+ ),
776
+ "fused": (
777
+ "routed AllReduce/RMSNorm plus shared ReduceScatter, sharded "
778
+ "up-projection/multicast, and Lamport copy"
779
+ ),
780
+ },
781
+ "results": results,
782
+ }
783
+ if rank == 0:
784
+ rendered = json.dumps(report, indent=2)
785
+ print(rendered, flush=True)
786
+ if args.output is not None:
787
+ args.output.parent.mkdir(parents=True, exist_ok=True)
788
+ args.output.write_text(rendered + "\n", encoding="utf-8")
789
+
790
+ dist.barrier(group=cpu_group)
791
+
792
+
793
+def main() -> None:
794
+ args = parse_args()
795
+ if args.scope == "up-projection":
796
+ benchmark_up_projection(args)
797
+ return
798
+
799
+ from vllm.config import VllmConfig, set_current_vllm_config
800
+
801
+ with set_current_vllm_config(VllmConfig()):
802
+ benchmark_whole_tail(args)
803
+
804
+
805
+if __name__ == "__main__":
806
+ main()
benchmarks/kernels/benchmark_kimi_k3_sp_collectives.py
new
+239
@@ -0,0 +1,239 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import argparse
5
+import json
6
+import os
7
+import statistics
8
+from collections.abc import Callable
9
+
10
+import torch
11
+import torch.distributed as dist
12
+
13
+import vllm._custom_ops as ops
14
+from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
15
+
16
+
17
+def parse_args() -> argparse.Namespace:
18
+ parser = argparse.ArgumentParser()
19
+ parser.add_argument("--tokens", type=int, nargs="+", default=[8, 32, 128, 1024])
20
+ parser.add_argument("--hidden-size", type=int, default=7168)
21
+ parser.add_argument("--graph-repeats", type=int, default=20)
22
+ parser.add_argument("--warmup-replays", type=int, default=5)
23
+ parser.add_argument("--samples", type=int, default=15)
24
+ return parser.parse_args()
25
+
26
+
27
+def capture_graph(op: Callable[[], None], repeats: int) -> torch.cuda.CUDAGraph:
28
+ stream = torch.cuda.Stream()
29
+ stream.wait_stream(torch.cuda.current_stream())
30
+ with torch.cuda.stream(stream):
31
+ for _ in range(3):
32
+ op()
33
+ stream.synchronize()
34
+
35
+ graph = torch.cuda.CUDAGraph()
36
+ with torch.cuda.graph(graph, stream=stream):
37
+ for _ in range(repeats):
38
+ op()
39
+ torch.cuda.current_stream().wait_stream(stream)
40
+ return graph
41
+
42
+
43
+def max_rank_graph_time(
44
+ graph: torch.cuda.CUDAGraph,
45
+ repeats: int,
46
+ warmup_replays: int,
47
+ samples: int,
48
+ device_group: dist.ProcessGroup,
49
+ cpu_group: dist.ProcessGroup,
50
+) -> float:
51
+ for _ in range(warmup_replays):
52
+ graph.replay()
53
+ torch.accelerator.synchronize()
54
+
55
+ timings = []
56
+ start = torch.cuda.Event(enable_timing=True)
57
+ end = torch.cuda.Event(enable_timing=True)
58
+ for _ in range(samples):
59
+ dist.barrier(group=cpu_group)
60
+ start.record()
61
+ graph.replay()
62
+ end.record()
63
+ end.synchronize()
64
+ elapsed = torch.tensor(
65
+ start.elapsed_time(end) / repeats,
66
+ dtype=torch.float64,
67
+ device=torch.accelerator.current_device_index(),
68
+ )
69
+ dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group)
70
+ timings.append(elapsed.item())
71
+ return statistics.median(timings)
72
+
73
+
74
+def check_outputs(
75
+ comm: CustomAllreduce,
76
+ local: torch.Tensor,
77
+ reduce_input: torch.Tensor,
78
+ device_group: dist.ProcessGroup,
79
+) -> None:
80
+ expected_gather = torch.empty(
81
+ (local.shape[0] * dist.get_world_size(), local.shape[1]),
82
+ dtype=local.dtype,
83
+ device=local.device,
84
+ )
85
+ dist.all_gather_into_tensor(expected_gather, local, group=device_group)
86
+ gathered = comm.custom_all_gather(local)
87
+ assert gathered is not None
88
+ torch.testing.assert_close(gathered, expected_gather)
89
+
90
+ expected_scatter = torch.empty_like(local)
91
+ dist.reduce_scatter_tensor(
92
+ expected_scatter,
93
+ reduce_input.clone(),
94
+ group=device_group,
95
+ )
96
+ scattered = comm.custom_reduce_scatter(reduce_input)
97
+ assert scattered is not None
98
+ torch.testing.assert_close(scattered, expected_scatter)
99
+
100
+
101
+def benchmark_shape(
102
+ comm: CustomAllreduce,
103
+ global_tokens: int,
104
+ hidden_size: int,
105
+ graph_repeats: int,
106
+ warmup_replays: int,
107
+ samples: int,
108
+ device_group: dist.ProcessGroup,
109
+ cpu_group: dist.ProcessGroup,
110
+) -> dict[str, float | int]:
111
+ world_size = dist.get_world_size()
112
+ rank = dist.get_rank()
113
+ padded_tokens = (global_tokens + world_size - 1) // world_size * world_size
114
+ local_tokens = padded_tokens // world_size
115
+ local = torch.full(
116
+ (local_tokens, hidden_size),
117
+ rank + 1,
118
+ dtype=torch.bfloat16,
119
+ device=torch.accelerator.current_device_index(),
120
+ )
121
+ reduce_input = torch.full(
122
+ (padded_tokens, hidden_size),
123
+ rank + 1,
124
+ dtype=torch.bfloat16,
125
+ device=local.device,
126
+ )
127
+ check_outputs(comm, local, reduce_input, device_group)
128
+
129
+ custom_gather_out = torch.empty(
130
+ (padded_tokens, hidden_size),
131
+ dtype=local.dtype,
132
+ device=local.device,
133
+ )
134
+ custom_scatter_out = torch.empty_like(local)
135
+ nccl_gather_out = torch.empty_like(custom_gather_out)
136
+ nccl_scatter_out = torch.empty_like(local)
137
+
138
+ def custom_ag() -> None:
139
+ ops.mnnvl_lamport_all_gather(
140
+ comm._ptr,
141
+ local,
142
+ custom_gather_out,
143
+ comm.mnnvl_lamport_ag_local_ptr,
144
+ comm.mnnvl_lamport_ag_multicast_ptr,
145
+ comm.mnnvl_lamport_ag_epoch_ptr,
146
+ comm.mnnvl_buffer_size,
147
+ )
148
+
149
+ def custom_rs() -> None:
150
+ ops.mnnvl_lamport_reduce_scatter(
151
+ comm._ptr,
152
+ reduce_input,
153
+ custom_scatter_out,
154
+ comm.mnnvl_lamport_rs_local_ptr,
155
+ comm.mnnvl_lamport_rs_epoch_ptr,
156
+ comm.mnnvl_buffer_size,
157
+ )
158
+
159
+ def nccl_ag() -> None:
160
+ dist.all_gather_into_tensor(nccl_gather_out, local, group=device_group)
161
+
162
+ def nccl_rs() -> None:
163
+ dist.reduce_scatter_tensor(
164
+ nccl_scatter_out,
165
+ reduce_input,
166
+ group=device_group,
167
+ )
168
+
169
+ graphs = {
170
+ "custom_ag_us": capture_graph(custom_ag, graph_repeats),
171
+ "nccl_ag_us": capture_graph(nccl_ag, graph_repeats),
172
+ "custom_rs_us": capture_graph(custom_rs, graph_repeats),
173
+ "nccl_rs_us": capture_graph(nccl_rs, graph_repeats),
174
+ }
175
+ times = {
176
+ name: max_rank_graph_time(
177
+ graph,
178
+ graph_repeats,
179
+ warmup_replays,
180
+ samples,
181
+ device_group,
182
+ cpu_group,
183
+ )
184
+ * 1000
185
+ for name, graph in graphs.items()
186
+ }
187
+ torch.testing.assert_close(custom_gather_out, nccl_gather_out)
188
+ torch.testing.assert_close(custom_scatter_out, nccl_scatter_out)
189
+ return {
190
+ "global_tokens": global_tokens,
191
+ "padded_tokens": padded_tokens,
192
+ "local_bytes": local.nbytes,
193
+ "full_bytes": reduce_input.nbytes,
194
+ **times,
195
+ "ag_speedup": times["nccl_ag_us"] / times["custom_ag_us"],
196
+ "rs_speedup": times["nccl_rs_us"] / times["custom_rs_us"],
197
+ }
198
+
199
+
200
+def main() -> None:
201
+ args = parse_args()
202
+ local_rank = int(os.environ["LOCAL_RANK"])
203
+ torch.accelerator.set_device_index(local_rank)
204
+ dist.init_process_group("nccl")
205
+ device_group = dist.group.WORLD
206
+ cpu_group = dist.new_group(backend="gloo")
207
+
208
+ comm = CustomAllreduce(
209
+ group=cpu_group,
210
+ device=torch.device("cuda", local_rank),
211
+ )
212
+ assert not comm.disabled
213
+ assert comm.world_size == 16
214
+ assert comm.mnnvl_only
215
+ assert comm.mnnvl_multicast_ptr
216
+
217
+ results = [
218
+ benchmark_shape(
219
+ comm,
220
+ tokens,
221
+ args.hidden_size,
222
+ args.graph_repeats,
223
+ args.warmup_replays,
224
+ args.samples,
225
+ device_group,
226
+ cpu_group,
227
+ )
228
+ for tokens in args.tokens
229
+ ]
230
+ if dist.get_rank() == 0:
231
+ print(json.dumps(results, indent=2), flush=True)
232
+
233
+ comm.close()
234
+ dist.destroy_process_group(cpu_group)
235
+ dist.destroy_process_group()
236
+
237
+
238
+if __name__ == "__main__":
239
+ main()
cmake/external_projects/flashkda.cmake
new
+74
@@ -0,0 +1,74 @@
1
+include(FetchContent)
2
+
3
+if(DEFINED ENV{FLASH_KDA_SRC_DIR})
4
+ set(FLASH_KDA_SRC_DIR $ENV{FLASH_KDA_SRC_DIR})
5
+endif()
6
+
7
+if(FLASH_KDA_SRC_DIR)
8
+ FetchContent_Declare(
9
+ flashkda
10
+ SOURCE_DIR ${FLASH_KDA_SRC_DIR}
11
+ )
12
+else()
13
+ FetchContent_Declare(
14
+ flashkda
15
+ GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git
16
+ GIT_TAG a3e42bbbece3bb38f7c426b880315294a336e82f
17
+ GIT_PROGRESS TRUE
18
+ GIT_SUBMODULES cutlass
19
+ )
20
+endif()
21
+
22
+FetchContent_MakeAvailable(flashkda)
23
+message(STATUS "FlashKDA is available at ${flashkda_SOURCE_DIR}")
24
+
25
+set(FLASH_KDA_SUPPORT_ARCHS)
26
+if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0)
27
+ list(APPEND FLASH_KDA_SUPPORT_ARCHS "9.0a")
28
+endif()
29
+if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
30
+ list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0f" "12.0f")
31
+elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
32
+ list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0a" "10.3a" "12.0a")
33
+endif()
34
+
35
+cuda_archs_loose_intersection(
36
+ FLASH_KDA_ARCHS "${FLASH_KDA_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
37
+
38
+if(FLASH_KDA_ARCHS)
39
+ message(STATUS "FlashKDA CUDA architectures: ${FLASH_KDA_ARCHS}")
40
+
41
+ set(FLASH_KDA_SOURCES
42
+ csrc/flashkda_registration.cpp
43
+ ${flashkda_SOURCE_DIR}/csrc/flash_kda.cpp
44
+ ${flashkda_SOURCE_DIR}/csrc/smxx/fwd_launch.cu)
45
+ set(FLASH_KDA_INCLUDES
46
+ ${flashkda_SOURCE_DIR}/csrc
47
+ ${flashkda_SOURCE_DIR}/cutlass/include
48
+ ${flashkda_SOURCE_DIR}/cutlass/examples/common
49
+ ${flashkda_SOURCE_DIR}/cutlass/tools/util/include)
50
+
51
+ set_gencode_flags_for_srcs(
52
+ SRCS "${FLASH_KDA_SOURCES}"
53
+ CUDA_ARCHS "${FLASH_KDA_ARCHS}")
54
+
55
+ define_extension_target(
56
+ _flashkda_C
57
+ DESTINATION vllm
58
+ LANGUAGE ${VLLM_GPU_LANG}
59
+ SOURCES ${FLASH_KDA_SOURCES}
60
+ COMPILE_FLAGS ${VLLM_GPU_FLAGS}
61
+ ARCHITECTURES ${VLLM_GPU_ARCHES}
62
+ INCLUDE_DIRECTORIES ${FLASH_KDA_INCLUDES}
63
+ USE_SABI 3
64
+ WITH_SOABI)
65
+
66
+ target_compile_options(_flashkda_C PRIVATE
67
+ $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API --expt-relaxed-constexpr --expt-extended-lambda --use_fast_math -O3>
68
+ $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>)
69
+else()
70
+ message(STATUS
71
+ "FlashKDA will not compile: CUDA >=12.0 and a supported architecture "
72
+ "(SM90, SM10x, or SM12x) are required")
73
+ add_custom_target(_flashkda_C)
74
+endif()
csrc/custom_all_gather_reduce_scatter.cuh
new
+327
@@ -0,0 +1,327 @@
1
+#pragma once
2
+
3
+#include "custom_collective_common.cuh"
4
+
5
+namespace vllm {
6
+
7
+constexpr int kMnnvlLamportAgThreads = 128;
8
+constexpr int kMnnvlLamportRsThreads = 256;
9
+constexpr int kMnnvlLamportConcurrentPollMaxPacks = 8192;
10
+
11
+using CopyPack = array_t<uint64_t, 2>;
12
+
13
+template <int ngpus>
14
+__global__ void __launch_bounds__(512, 1)
15
+ cross_device_all_gather(RankData* _dp, RankSignals sg, Signal* self_sg,
16
+ CopyPack* __restrict__ result, int rank,
17
+ int size_per_rank) {
18
+ auto dp = *_dp;
19
+ int tid = blockIdx.x * blockDim.x + threadIdx.x;
20
+ int stride = gridDim.x * blockDim.x;
21
+ barrier_at_start<ngpus>(sg, self_sg, rank);
22
+#pragma unroll
23
+ for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
24
+ auto src = reinterpret_cast<const CopyPack*>(dp.ptrs[src_rank]);
25
+ auto dst = result + src_rank * size_per_rank;
26
+ for (int idx = tid; idx < size_per_rank; idx += stride) {
27
+ dst[idx] = src[idx];
28
+ }
29
+ }
30
+ barrier_at_end<ngpus, true>(sg, self_sg, rank);
31
+}
32
+
33
+template <typename T, int ngpus>
34
+__global__ void __launch_bounds__(512, 1)
35
+ cross_device_reduce_scatter(RankData* _dp, RankSignals sg, Signal* self_sg,
36
+ T* __restrict__ result, int rank,
37
+ int size_per_rank) {
38
+ using P = typename packed_t<T>::P;
39
+ using A = typename packed_t<T>::A;
40
+ auto dp = *_dp;
41
+ auto offset = rank * size_per_rank;
42
+ barrier_at_start<ngpus>(sg, self_sg, rank);
43
+ for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size_per_rank;
44
+ idx += gridDim.x * blockDim.x) {
45
+ reinterpret_cast<P*>(result)[idx] =
46
+ packed_reduce<P, ngpus, A>((const P**)&dp.ptrs[0], offset + idx);
47
+ }
48
+ barrier_at_end<ngpus, true>(sg, self_sg, rank);
49
+}
50
+
51
+template <typename P>
52
+union LamportPack {
53
+ P packed;
54
+ uint32_t words[sizeof(P) / sizeof(uint32_t)];
55
+};
56
+
57
+template <typename P>
58
+DINLINE LamportPack<P> load_lamport_pack(const P* ptr) {
59
+ static_assert(sizeof(P) == 16);
60
+ LamportPack<P> value;
61
+#if !defined(USE_ROCM)
62
+ asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];"
63
+ : "=r"(value.words[0]), "=r"(value.words[1]),
64
+ "=r"(value.words[2]), "=r"(value.words[3])
65
+ : "l"(ptr)
66
+ : "memory");
67
+#else
68
+ const volatile uint32_t* src =
69
+ reinterpret_cast<const volatile uint32_t*>(ptr);
70
+ #pragma unroll
71
+ for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
72
+ value.words[i] = src[i];
73
+ }
74
+#endif
75
+ return value;
76
+}
77
+
78
+template <typename P>
79
+DINLINE bool is_lamport_dirty(const LamportPack<P>& value) {
80
+#pragma unroll
81
+ for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
82
+ if (value.words[i] == 0x80000000U) return true;
83
+ }
84
+ return false;
85
+}
86
+
87
+template <typename P>
88
+DINLINE P lamport_sentinel() {
89
+ LamportPack<P> value;
90
+#pragma unroll
91
+ for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
92
+ value.words[i] = 0x80000000U;
93
+ }
94
+ return value.packed;
95
+}
96
+
97
+template <typename P>
98
+DINLINE P sanitize_lamport_payload(P packed) {
99
+ LamportPack<P> value{.packed = packed};
100
+#pragma unroll
101
+ for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) {
102
+ if (value.words[i] == 0x80000000U) value.words[i] = 0;
103
+ }
104
+ return value.packed;
105
+}
106
+
107
+template <typename P>
108
+DINLINE P wait_lamport_payload(const P* ptr) {
109
+ auto value = load_lamport_pack(ptr);
110
+ while (is_lamport_dirty(value)) value = load_lamport_pack(ptr);
111
+ return value.packed;
112
+}
113
+
114
+template <typename P, int ngpus>
115
+DINLINE void wait_lamport_payloads(const P* base, int rank, int rank_stride,
116
+ P local_value, P (&values)[ngpus]) {
117
+ bool ready[ngpus];
118
+#pragma unroll
119
+ for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
120
+ ready[src_rank] = src_rank == rank;
121
+ if (src_rank == rank) values[src_rank] = local_value;
122
+ }
123
+
124
+ int remaining = ngpus - 1;
125
+ while (remaining != 0) {
126
+#pragma unroll
127
+ for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
128
+ if (!ready[src_rank]) {
129
+ auto value = load_lamport_pack(base + src_rank * rank_stride);
130
+ if (!is_lamport_dirty(value)) {
131
+ values[src_rank] = value.packed;
132
+ ready[src_rank] = true;
133
+ --remaining;
134
+ }
135
+ }
136
+ }
137
+ }
138
+}
139
+
140
+template <typename P, typename A, int ngpus>
141
+DINLINE P reduce_lamport_payloads(const P* current_local, const P* packed_input,
142
+ int rank, int size_per_rank, int idx) {
143
+ P source_zero =
144
+ rank == 0 ? packed_input[idx] : wait_lamport_payload(current_local + idx);
145
+ A tmp = upcast(source_zero);
146
+#pragma unroll
147
+ for (int src_rank = 1; src_rank < ngpus; ++src_rank) {
148
+ P value = src_rank == rank
149
+ ? packed_input[rank * size_per_rank + idx]
150
+ : wait_lamport_payload(current_local +
151
+ src_rank * size_per_rank + idx);
152
+ packed_assign_add(tmp, upcast(value));
153
+ }
154
+ return sanitize_lamport_payload(downcast<P>(tmp));
155
+}
156
+
157
+DINLINE void lamport_cta_arrive(uint32_t* counter) {
158
+#if !defined(USE_ROCM)
159
+ if (threadIdx.x < 32) {
160
+ asm volatile("barrier.cta.sync 1, %0;" : : "r"(blockDim.x) : "memory");
161
+ if (threadIdx.x == 0) {
162
+ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
163
+ asm volatile("red.async.release.global.gpu.add.u32 [%0], 1;"
164
+ :
165
+ : "l"(counter)
166
+ : "memory");
167
+ #elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
168
+ asm volatile("red.release.global.gpu.add.u32 [%0], 1;"
169
+ :
170
+ : "l"(counter)
171
+ : "memory");
172
+ #else
173
+ atomicAdd(counter, 1);
174
+ #endif
175
+ }
176
+ } else {
177
+ asm volatile("barrier.cta.arrive 1, %0;" : : "r"(blockDim.x) : "memory");
178
+ }
179
+#else
180
+ __syncthreads();
181
+ if (threadIdx.x == 0) atomicAdd(counter, 1);
182
+#endif
183
+}
184
+
185
+template <typename T, int ngpus>
186
+__global__ void __launch_bounds__(kMnnvlLamportAgThreads, 1)
187
+ mnnvl_lamport_all_gather(RankData* _dp, const T* __restrict__ input,
188
+ T* __restrict__ result,
189
+ T* __restrict__ multicast_buffer,
190
+ uint32_t* __restrict__ epochs, int rank,
191
+ int size_per_rank, int stage_size) {
192
+ using P = typename packed_t<T>::P;
193
+#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
194
+ (__CUDA_ARCH__ >= 900)
195
+ cudaGridDependencySynchronize();
196
+#endif
197
+ auto dp = *_dp;
198
+ int tid = blockIdx.x * blockDim.x + threadIdx.x;
199
+ int stride = gridDim.x * blockDim.x;
200
+ uint32_t epoch = epochs[0];
201
+ int current_stage = epoch % 3;
202
+ int dirty_stage = (epoch + 1) % 3;
203
+ int dirty_size = epochs[2 + dirty_stage];
204
+ auto local_buffer = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[rank]));
205
+ auto current_local = local_buffer + current_stage * stage_size;
206
+ auto dirty_local = local_buffer + dirty_stage * stage_size;
207
+ auto current_multicast =
208
+ reinterpret_cast<P*>(multicast_buffer) + current_stage * stage_size;
209
+ auto packed_input = reinterpret_cast<const P*>(input);
210
+ auto packed_result = reinterpret_cast<P*>(result);
211
+
212
+ int total_size = size_per_rank * ngpus;
213
+ P local_value;
214
+ if (tid < size_per_rank) {
215
+ local_value = packed_input[tid];
216
+ current_multicast[rank * size_per_rank + tid] =
217
+ sanitize_lamport_payload(local_value);
218
+ }
219
+#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
220
+ (__CUDA_ARCH__ >= 900)
221
+ cudaTriggerProgrammaticLaunchCompletion();
222
+#endif
223
+
224
+ lamport_cta_arrive(&epochs[1]);
225
+
226
+ for (int idx = tid; idx < dirty_size; idx += stride) {
227
+ dirty_local[idx] = lamport_sentinel<P>();
228
+ }
229
+
230
+ if (tid < size_per_rank) {
231
+#pragma unroll
232
+ for (int src_rank = 0; src_rank < ngpus; ++src_rank) {
233
+ int output_idx = src_rank * size_per_rank + tid;
234
+ P value = src_rank == rank
235
+ ? local_value
236
+ : wait_lamport_payload(current_local + output_idx);
237
+ packed_result[output_idx] = value;
238
+ }
239
+ }
240
+
241
+ if (tid == 0) {
242
+ while (*reinterpret_cast<volatile uint32_t*>(&epochs[1]) < gridDim.x);
243
+ epochs[2 + current_stage] = total_size;
244
+ epochs[0] = epoch + 1;
245
+ epochs[1] = 0;
246
+ }
247
+}
248
+
249
+template <typename T, int ngpus>
250
+__global__ void __launch_bounds__(kMnnvlLamportRsThreads, 1)
251
+ mnnvl_lamport_reduce_scatter_kernel(RankData* _dp,
252
+ const T* __restrict__ input,
253
+ T* __restrict__ result,
254
+ uint32_t* __restrict__ epochs, int rank,
255
+ int size_per_rank, int stage_size) {
256
+ using P = typename packed_t<T>::P;
257
+ using A = typename packed_t<T>::A;
258
+#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
259
+ (__CUDA_ARCH__ >= 900)
260
+ cudaGridDependencySynchronize();
261
+#endif
262
+ auto dp = *_dp;
263
+ int dst_rank = blockIdx.x % ngpus;
264
+ int tile = blockIdx.x / ngpus;
265
+ int idx = tile * blockDim.x + threadIdx.x;
266
+ int tid = blockIdx.x * blockDim.x + threadIdx.x;
267
+ int stride = gridDim.x * blockDim.x;
268
+ uint32_t epoch = epochs[0];
269
+ int current_stage = epoch % 3;
270
+ int dirty_stage = (epoch + 1) % 3;
271
+ int dirty_size = epochs[2 + dirty_stage];
272
+ auto local_buffer = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[rank]));
273
+ auto current_local = local_buffer + current_stage * stage_size;
274
+ auto dirty_local = local_buffer + dirty_stage * stage_size;
275
+ auto packed_input = reinterpret_cast<const P*>(input);
276
+
277
+ if (idx < size_per_rank && dst_rank != rank) {
278
+ auto dst = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[dst_rank])) +
279
+ current_stage * stage_size + rank * size_per_rank;
280
+ auto src = packed_input + dst_rank * size_per_rank;
281
+ dst[idx] = sanitize_lamport_payload(src[idx]);
282
+ }
283
+#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \
284
+ (__CUDA_ARCH__ >= 900)
285
+ cudaTriggerProgrammaticLaunchCompletion();
286
+#endif
287
+
288
+ lamport_cta_arrive(&epochs[1]);
289
+
290
+ for (int idx = tid; idx < dirty_size; idx += stride) {
291
+ dirty_local[idx] = lamport_sentinel<P>();
292
+ }
293
+
294
+ if (idx < size_per_rank && dst_rank == rank) {
295
+ if constexpr (ngpus == 4) {
296
+ if (size_per_rank > kMnnvlLamportConcurrentPollMaxPacks) {
297
+ reinterpret_cast<P*>(result)[idx] =
298
+ reduce_lamport_payloads<P, A, ngpus>(current_local, packed_input,
299
+ rank, size_per_rank, idx);
300
+ } else {
301
+ P values[ngpus];
302
+ wait_lamport_payloads<P, ngpus>(
303
+ current_local + idx, rank, size_per_rank,
304
+ packed_input[rank * size_per_rank + idx], values);
305
+ A tmp = upcast(values[0]);
306
+#pragma unroll
307
+ for (int src_rank = 1; src_rank < ngpus; ++src_rank) {
308
+ packed_assign_add(tmp, upcast(values[src_rank]));
309
+ }
310
+ reinterpret_cast<P*>(result)[idx] =
311
+ sanitize_lamport_payload(downcast<P>(tmp));
312
+ }
313
+ } else {
314
+ reinterpret_cast<P*>(result)[idx] = reduce_lamport_payloads<P, A, ngpus>(
315
+ current_local, packed_input, rank, size_per_rank, idx);
316
+ }
317
+ }
318
+
319
+ if (tid == 0) {
320
+ while (*reinterpret_cast<volatile uint32_t*>(&epochs[1]) < gridDim.x);
321
+ epochs[2 + current_stage] = size_per_rank * ngpus;
322
+ epochs[0] = epoch + 1;
323
+ epochs[1] = 0;
324
+ }
325
+}
326
+
327
+} // namespace vllm
csrc/custom_all_reduce.cuh
+20
-296
@@ -1,299 +1,8 @@
1
#pragma once
2
3
-#include <cuda.h>
4
-#include <cuda_bf16.h>
5
-#include <cuda_fp16.h>
6
-#include <cuda_runtime.h>
7
-
8
-#if defined(USE_ROCM)
9
-typedef __hip_bfloat16 nv_bfloat16;
10
-#endif
11
-
12
-#include <iostream>
13
-#include <array>
14
-#include <limits>
15
-#include <map>
16
-#include <unordered_map>
17
-#include <vector>
18
-#include <cstdlib>
19
-#include <cstring>
3
+#include "custom_collective_common.cuh"
4
5
namespace vllm {
22
-#define CUDACHECK(cmd) \
23
- do { \
24
- cudaError_t e = cmd; \
25
- if (e != cudaSuccess) { \
26
- printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, \
27
- cudaGetErrorString(e)); \
28
- exit(EXIT_FAILURE); \
29
- } \
30
- } while (0)
31
-
32
-// Maximal number of blocks in allreduce kernel.
33
-constexpr int kMaxBlocks = 36;
34
-
35
-// Default number of blocks in allreduce kernel.
36
-#ifndef USE_ROCM
37
-const int defaultBlockLimit = 36;
38
-CUpointer_attribute rangeStartAddrAttr = CU_POINTER_ATTRIBUTE_RANGE_START_ADDR;
39
-#else
40
-const int defaultBlockLimit = 16;
41
-hipPointer_attribute rangeStartAddrAttr =
42
- HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR;
43
-#endif
44
-
45
-// Counter may overflow, but it's fine since unsigned int overflow is
46
-// well-defined behavior.
47
-using FlagType = uint32_t;
48
-
49
-// Two sets of peer counters are needed for two syncs: starting and ending an
50
-// operation. The reason is that it's possible for peer GPU block to arrive at
51
-// the second sync point while the current GPU block haven't passed the first
52
-// sync point. Thus, peer GPU may write counter+1 while current GPU is busy
53
-// waiting for counter. We use alternating counter array to avoid this
54
-// possibility.
55
-struct Signal {
56
- alignas(128) FlagType start[kMaxBlocks][8];
57
- alignas(128) FlagType end[kMaxBlocks][8];
58
- alignas(128) FlagType _flag[kMaxBlocks]; // incremental flags for each rank
59
-};
60
-
61
-struct __align__(16) RankData {
62
- const void* ptrs[8];
63
-};
64
-
65
-struct __align__(16) RankSignals {
66
- Signal* signals[8];
67
-};
68
-
69
-// like std::array, but aligned
70
-template <typename T, int sz>
71
-struct __align__(alignof(T) * sz) array_t {
72
- T data[sz];
73
- using type = T;
74
- static constexpr int size = sz;
75
-};
76
-
77
-// use packed type to maximize memory efficiency
78
-// goal: generate ld.128 and st.128 instructions
79
-template <typename T>
80
-struct packed_t {
81
- // the (P)acked type for load/store
82
- using P = array_t<T, 16 / sizeof(T)>;
83
- // the (A)ccumulator type for reduction
84
- using A = array_t<float, 16 / sizeof(T)>;
85
-};
86
-
87
-#define DINLINE __device__ __forceinline__
88
-
89
-// scalar cast functions
90
-DINLINE float upcast_s(half val) { return __half2float(val); }
91
-
92
-template <typename T>
93
-DINLINE T downcast_s(float val);
94
-template <>
95
-DINLINE half downcast_s(float val) {
96
- return __float2half(val);
97
-}
98
-
99
-// scalar add functions
100
-// for some reason when compiling with Pytorch, the + operator for half and
101
-// bfloat is disabled so we call the intrinsics directly
102
-DINLINE half& assign_add(half& a, half b) {
103
- a = __hadd(a, b);
104
- return a;
105
-}
106
-DINLINE float& assign_add(float& a, float b) { return a += b; }
107
-
108
-#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
109
-DINLINE float upcast_s(nv_bfloat16 val) { return __bfloat162float(val); }
110
-template <>
111
-DINLINE nv_bfloat16 downcast_s(float val) {
112
- return __float2bfloat16(val);
113
-}
114
-DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) {
115
- a = __hadd(a, b);
116
- return a;
117
-}
118
-#endif
119
-
120
-template <typename T, int N>
121
-DINLINE array_t<T, N>& packed_assign_add(array_t<T, N>& a, array_t<T, N> b) {
122
-#pragma unroll
123
- for (int i = 0; i < N; i++) {
124
- assign_add(a.data[i], b.data[i]);
125
- }
126
- return a;
127
-}
128
-
129
-template <typename T, int N>
130
-DINLINE array_t<float, N> upcast(array_t<T, N> val) {
131
- if constexpr (std::is_same<T, float>::value) {
132
- return val;
133
- } else {
134
- array_t<float, N> out;
135
-#pragma unroll
136
- for (int i = 0; i < N; i++) {
137
- out.data[i] = upcast_s(val.data[i]);
138
- }
139
- return out;
140
- }
141
-}
142
-
143
-template <typename O>
144
-DINLINE O downcast(array_t<float, O::size> val) {
145
- if constexpr (std::is_same<typename O::type, float>::value) {
146
- return val;
147
- } else {
148
- O out;
149
-#pragma unroll
150
- for (int i = 0; i < O::size; i++) {
151
- out.data[i] = downcast_s<typename O::type>(val.data[i]);
152
- }
153
- return out;
154
- }
155
-}
156
-
157
-#if !defined(USE_ROCM)
158
-
159
-static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) {
160
- #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
161
- asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag),
162
- "l"(flag_addr));
163
- #else
164
- asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag),
165
- "l"(flag_addr));
166
- #endif
167
-}
168
-
169
-static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) {
170
- FlagType flag;
171
- #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
172
- asm volatile("ld.acquire.sys.global.u32 %0, [%1];"
173
- : "=r"(flag)
174
- : "l"(flag_addr));
175
- #else
176
- asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;"
177
- : "=r"(flag)
178
- : "l"(flag_addr));
179
- #endif
180
- return flag;
181
-}
182
-
183
-static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) {
184
- asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr));
185
-}
186
-
187
-static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) {
188
- FlagType flag;
189
- asm volatile("ld.volatile.global.u32 %0, [%1];"
190
- : "=r"(flag)
191
- : "l"(flag_addr));
192
- return flag;
193
-}
194
-
195
-// This function is meant to be used as the first synchronization in the all
196
-// reduce kernel. Thus, it doesn't need to make any visibility guarantees for
197
-// prior memory accesses. Note: volatile writes will not be reordered against
198
-// other volatile writes.
199
-template <int ngpus>
200
-DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg,
201
- int rank) {
202
- uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
203
- if (threadIdx.x < ngpus) {
204
- auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank];
205
- auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x];
206
- // Write the expected counter value to peer and wait for correct value
207
- // from peer.
208
- st_flag_volatile(peer_counter_ptr, flag);
209
- while (ld_flag_volatile(self_counter_ptr) != flag);
210
- }
211
- __syncthreads();
212
- // use one thread to update flag
213
- if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
214
-}
215
-
216
-// This function is meant to be used as the second or the final
217
-// synchronization barrier in the all reduce kernel. If it's the final
218
-// synchronization barrier, we don't need to make any visibility guarantees
219
-// for prior memory accesses.
220
-template <int ngpus, bool final_sync = false>
221
-DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) {
222
- __syncthreads();
223
- uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
224
- if (threadIdx.x < ngpus) {
225
- auto peer_counter_ptr = &sg.signals[threadIdx.x]->end[blockIdx.x][rank];
226
- auto self_counter_ptr = &self_sg->end[blockIdx.x][threadIdx.x];
227
- // Write the expected counter value to peer and wait for correct value from
228
- // peer.
229
- if constexpr (!final_sync) {
230
- st_flag_release(peer_counter_ptr, flag);
231
- while (ld_flag_acquire(self_counter_ptr) != flag);
232
- } else {
233
- st_flag_volatile(peer_counter_ptr, flag);
234
- while (ld_flag_volatile(self_counter_ptr) != flag);
235
- }
236
- }
237
- if constexpr (!final_sync) __syncthreads();
238
-
239
- // use one thread to update flag
240
- if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
241
-}
242
-
243
-#else
244
-
245
-template <int ngpus>
246
-DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg,
247
- int rank) {
248
- uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
249
- if (threadIdx.x < ngpus) {
250
- // simultaneously write to the corresponding flag of all ranks.
251
- // Latency = 1 p2p write
252
- __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank],
253
- flag, __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM);
254
- // wait until we got true from all ranks
255
- while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x],
256
- __ATOMIC_RELAXED,
257
- __MEMORY_SCOPE_DEVICE) < flag);
258
- }
259
- __syncthreads();
260
- // use one thread to update flag
261
- if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
262
-}
263
-
264
-template <int ngpus, bool final_sync = false>
265
-DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) {
266
- __syncthreads();
267
- uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
268
- if (threadIdx.x < ngpus) {
269
- // simultaneously write to the corresponding flag of all ranks.
270
- // Latency = 1 p2p write
271
- __scoped_atomic_store_n(&sg.signals[threadIdx.x]->end[blockIdx.x][rank],
272
- flag,
273
- final_sync ? __ATOMIC_RELAXED : __ATOMIC_RELEASE,
274
- __MEMORY_SCOPE_SYSTEM);
275
- // wait until we got true from all ranks
276
- while (
277
- __scoped_atomic_load_n(&self_sg->end[blockIdx.x][threadIdx.x],
278
- final_sync ? __ATOMIC_RELAXED : __ATOMIC_ACQUIRE,
279
- __MEMORY_SCOPE_DEVICE) < flag);
280
- }
281
- if constexpr (!final_sync) __syncthreads();
282
- // use one thread to update flag
283
- if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
284
-}
285
-
286
-#endif
287
-
288
-template <typename P, int ngpus, typename A>
289
-DINLINE P packed_reduce(const P* ptrs[], int idx) {
290
- A tmp = upcast(ptrs[0][idx]);
291
-#pragma unroll
292
- for (int i = 1; i < ngpus; i++) {
293
- packed_assign_add(tmp, upcast(ptrs[i][idx]));
294
- }
295
- return downcast<P>(tmp);
296
-}
6
7
template <typename T, int ngpus>
8
__global__ void __launch_bounds__(512, 1)
@@ -616,6 +325,21 @@ class CustomAllreduce {
325
#undef KL
326
}
327
328
+ void allgather(cudaStream_t stream, void* input, void* output, int size_bytes,
329
+ int threads = 512, int block_limit = defaultBlockLimit);
330
+ template <typename T>
331
+ void mnnvl_lamport_allgather(cudaStream_t stream, T* input, T* output,
332
+ void* local_buffer, void* multicast_buffer,
333
+ uint32_t* epochs, int size_bytes,
334
+ int stage_size_bytes);
335
+ template <typename T>
336
+ void reduce_scatter(cudaStream_t stream, T* input, T* output, int size,
337
+ int threads = 512, int block_limit = defaultBlockLimit);
338
+ template <typename T>
339
+ void mnnvl_lamport_reduce_scatter(cudaStream_t stream, T* input, T* output,
340
+ void* local_buffer, uint32_t* epochs,
341
+ int size, int stage_size_bytes);
342
+
343
~CustomAllreduce() {
344
for (auto [_, ptr] : ipc_handles_) {
345
CUDACHECK(cudaIpcCloseMemHandle(ptr));
@@ -625,8 +349,8 @@ class CustomAllreduce {
349
350
/**
351
* To inspect PTX/SASS, copy paste this header file to compiler explorer and
628
- add a template instantiation:
352
+ * add a template instantiation:
353
* template void vllm::CustomAllreduce::allreduce<half>(cudaStream_t, half *,
630
- half *, int, int, int);
631
-*/
632
-} // namespace vllm
\ No newline at end of file
354
+ * half *, int, int, int);
355
+ */
356
+} // namespace vllm
csrc/custom_collective_common.cuh
new
+332
@@ -0,0 +1,332 @@
1
+#pragma once
2
+
3
+#include <cuda.h>
4
+#include <cuda_bf16.h>
5
+#include <cuda_fp16.h>
6
+#include <cuda_runtime.h>
7
+
8
+#if defined(USE_ROCM)
9
+typedef __hip_bfloat16 nv_bfloat16;
10
+#endif
11
+
12
+#include <iostream>
13
+#include <array>
14
+#include <limits>
15
+#include <map>
16
+#include <unordered_map>
17
+#include <vector>
18
+#include <cstdlib>
19
+#include <cstring>
20
+
21
+namespace vllm {
22
+constexpr int kMaxCustomCollectiveRanks = 16;
23
+
24
+#define CUDACHECK(cmd) \
25
+ do { \
26
+ cudaError_t e = cmd; \
27
+ if (e != cudaSuccess) { \
28
+ printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, \
29
+ cudaGetErrorString(e)); \
30
+ exit(EXIT_FAILURE); \
31
+ } \
32
+ } while (0)
33
+
34
+// Maximal number of blocks in allreduce kernel.
35
+constexpr int kMaxBlocks = 36;
36
+
37
+// Default number of blocks in allreduce kernel.
38
+#ifndef USE_ROCM
39
+inline constexpr int defaultBlockLimit = 36;
40
+inline CUpointer_attribute rangeStartAddrAttr =
41
+ CU_POINTER_ATTRIBUTE_RANGE_START_ADDR;
42
+#else
43
+inline constexpr int defaultBlockLimit = 16;
44
+inline hipPointer_attribute rangeStartAddrAttr =
45
+ HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR;
46
+#endif
47
+
48
+// Counter may overflow, but it's fine since unsigned int overflow is
49
+// well-defined behavior.
50
+using FlagType = uint32_t;
51
+
52
+// Two sets of peer counters are needed for two syncs: starting and ending an
53
+// operation. The reason is that it's possible for peer GPU block to arrive at
54
+// the second sync point while the current GPU block haven't passed the first
55
+// sync point. Thus, peer GPU may write counter+1 while current GPU is busy
56
+// waiting for counter. We use alternating counter array to avoid this
57
+// possibility.
58
+struct Signal {
59
+ alignas(128) FlagType start[kMaxBlocks][kMaxCustomCollectiveRanks];
60
+ alignas(128) FlagType end[kMaxBlocks][kMaxCustomCollectiveRanks];
61
+ alignas(128) FlagType _flag[kMaxBlocks]; // incremental flags for each rank
62
+};
63
+
64
+struct __align__(16) RankData {
65
+ const void* ptrs[kMaxCustomCollectiveRanks];
66
+};
67
+
68
+struct __align__(16) RankSignals {
69
+ Signal* signals[kMaxCustomCollectiveRanks];
70
+};
71
+
72
+// like std::array, but aligned
73
+template <typename T, int sz>
74
+struct __align__(alignof(T) * sz) array_t {
75
+ T data[sz];
76
+ using type = T;
77
+ static constexpr int size = sz;
78
+};
79
+
80
+// use packed type to maximize memory efficiency
81
+// goal: generate ld.128 and st.128 instructions
82
+template <typename T>
83
+struct packed_t {
84
+ // the (P)acked type for load/store
85
+ using P = array_t<T, 16 / sizeof(T)>;
86
+ // the (A)ccumulator type for reduction
87
+ using A = array_t<float, 16 / sizeof(T)>;
88
+};
89
+
90
+#define DINLINE __device__ __forceinline__
91
+
92
+// scalar cast functions
93
+DINLINE float upcast_s(half val) { return __half2float(val); }
94
+
95
+template <typename T>
96
+DINLINE T downcast_s(float val);
97
+template <>
98
+DINLINE half downcast_s(float val) {
99
+ return __float2half(val);
100
+}
101
+
102
+// scalar add functions
103
+// for some reason when compiling with Pytorch, the + operator for half and
104
+// bfloat is disabled so we call the intrinsics directly
105
+DINLINE half& assign_add(half& a, half b) {
106
+ a = __hadd(a, b);
107
+ return a;
108
+}
109
+DINLINE float& assign_add(float& a, float b) { return a += b; }
110
+
111
+#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
112
+DINLINE float upcast_s(nv_bfloat16 val) { return __bfloat162float(val); }
113
+template <>
114
+DINLINE nv_bfloat16 downcast_s(float val) {
115
+ return __float2bfloat16(val);
116
+}
117
+DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) {
118
+ a = __hadd(a, b);
119
+ return a;
120
+}
121
+#endif
122
+
123
+template <typename T, int N>
124
+DINLINE array_t<T, N>& packed_assign_add(array_t<T, N>& a, array_t<T, N> b) {
125
+#pragma unroll
126
+ for (int i = 0; i < N; i++) {
127
+ assign_add(a.data[i], b.data[i]);
128
+ }
129
+ return a;
130
+}
131
+
132
+template <typename T, int N>
133
+DINLINE array_t<float, N> upcast(array_t<T, N> val) {
134
+ if constexpr (std::is_same<T, float>::value) {
135
+ return val;
136
+ } else {
137
+ array_t<float, N> out;
138
+#pragma unroll
139
+ for (int i = 0; i < N; i++) {
140
+ out.data[i] = upcast_s(val.data[i]);
141
+ }
142
+ return out;
143
+ }
144
+}
145
+
146
+template <typename O>
147
+DINLINE O downcast(array_t<float, O::size> val) {
148
+ if constexpr (std::is_same<typename O::type, float>::value) {
149
+ return val;
150
+ } else {
151
+ O out;
152
+#pragma unroll
153
+ for (int i = 0; i < O::size; i++) {
154
+ out.data[i] = downcast_s<typename O::type>(val.data[i]);
155
+ }
156
+ return out;
157
+ }
158
+}
159
+
160
+#if !defined(USE_ROCM)
161
+
162
+static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) {
163
+ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
164
+ asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag),
165
+ "l"(flag_addr));
166
+ #else
167
+ asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag),
168
+ "l"(flag_addr));
169
+ #endif
170
+}
171
+
172
+static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) {
173
+ FlagType flag;
174
+ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
175
+ asm volatile("ld.acquire.sys.global.u32 %0, [%1];"
176
+ : "=r"(flag)
177
+ : "l"(flag_addr));
178
+ #else
179
+ asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;"
180
+ : "=r"(flag)
181
+ : "l"(flag_addr));
182
+ #endif
183
+ return flag;
184
+}
185
+
186
+static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) {
187
+ asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr));
188
+}
189
+
190
+static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) {
191
+ FlagType flag;
192
+ asm volatile("ld.volatile.global.u32 %0, [%1];"
193
+ : "=r"(flag)
194
+ : "l"(flag_addr));
195
+ return flag;
196
+}
197
+
198
+// This function is meant to be used as the first synchronization in the all
199
+// reduce kernel. Thus, it doesn't need to make any visibility guarantees for
200
+// prior memory accesses. Note: volatile writes will not be reordered against
201
+// other volatile writes.
202
+template <int ngpus>
203
+DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg,
204
+ int rank) {
205
+ uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
206
+ if (threadIdx.x < ngpus) {
207
+ auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank];
208
+ auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x];
209
+ // Write the expected counter value to peer and wait for correct value
210
+ // from peer.
211
+ st_flag_volatile(peer_counter_ptr, flag);
212
+ while (ld_flag_volatile(self_counter_ptr) != flag);
213
+ }
214
+ __syncthreads();
215
+ // use one thread to update flag
216
+ if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
217
+}
218
+
219
+template <int ngpus>
220
+DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg,
221
+ int rank) {
222
+ __syncthreads();
223
+ uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
224
+ if (threadIdx.x < ngpus) {
225
+ auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank];
226
+ auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x];
227
+ st_flag_release(peer_counter_ptr, flag);
228
+ while (ld_flag_acquire(self_counter_ptr) != flag);
229
+ }
230
+ __syncthreads();
231
+ if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
232
+}
233
+
234
+// This function is meant to be used as the second or the final
235
+// synchronization barrier in the all reduce kernel. If it's the final
236
+// synchronization barrier, we don't need to make any visibility guarantees
237
+// for prior memory accesses.
238
+template <int ngpus, bool final_sync = false>
239
+DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) {
240
+ __syncthreads();
241
+ uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
242
+ if (threadIdx.x < ngpus) {
243
+ auto peer_counter_ptr = &sg.signals[threadIdx.x]->end[blockIdx.x][rank];
244
+ auto self_counter_ptr = &self_sg->end[blockIdx.x][threadIdx.x];
245
+ // Write the expected counter value to peer and wait for correct value from
246
+ // peer.
247
+ if constexpr (!final_sync) {
248
+ st_flag_release(peer_counter_ptr, flag);
249
+ while (ld_flag_acquire(self_counter_ptr) != flag);
250
+ } else {
251
+ st_flag_volatile(peer_counter_ptr, flag);
252
+ while (ld_flag_volatile(self_counter_ptr) != flag);
253
+ }
254
+ }
255
+ if constexpr (!final_sync) __syncthreads();
256
+
257
+ // use one thread to update flag
258
+ if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
259
+}
260
+
261
+#else
262
+
263
+template <int ngpus>
264
+DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg,
265
+ int rank) {
266
+ uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
267
+ if (threadIdx.x < ngpus) {
268
+ // simultaneously write to the corresponding flag of all ranks.
269
+ // Latency = 1 p2p write
270
+ __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank],
271
+ flag, __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM);
272
+ // wait until we got true from all ranks
273
+ while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x],
274
+ __ATOMIC_RELAXED,
275
+ __MEMORY_SCOPE_DEVICE) < flag);
276
+ }
277
+ __syncthreads();
278
+ // use one thread to update flag
279
+ if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
280
+}
281
+
282
+template <int ngpus>
283
+DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg,
284
+ int rank) {
285
+ __syncthreads();
286
+ uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
287
+ if (threadIdx.x < ngpus) {
288
+ __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank],
289
+ flag, __ATOMIC_RELEASE, __MEMORY_SCOPE_SYSTEM);
290
+ while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x],
291
+ __ATOMIC_ACQUIRE,
292
+ __MEMORY_SCOPE_DEVICE) < flag);
293
+ }
294
+ __syncthreads();
295
+ if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
296
+}
297
+
298
+template <int ngpus, bool final_sync = false>
299
+DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) {
300
+ __syncthreads();
301
+ uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
302
+ if (threadIdx.x < ngpus) {
303
+ // simultaneously write to the corresponding flag of all ranks.
304
+ // Latency = 1 p2p write
305
+ __scoped_atomic_store_n(&sg.signals[threadIdx.x]->end[blockIdx.x][rank],
306
+ flag,
307
+ final_sync ? __ATOMIC_RELAXED : __ATOMIC_RELEASE,
308
+ __MEMORY_SCOPE_SYSTEM);
309
+ // wait until we got true from all ranks
310
+ while (
311
+ __scoped_atomic_load_n(&self_sg->end[blockIdx.x][threadIdx.x],
312
+ final_sync ? __ATOMIC_RELAXED : __ATOMIC_ACQUIRE,
313
+ __MEMORY_SCOPE_DEVICE) < flag);
314
+ }
315
+ if constexpr (!final_sync) __syncthreads();
316
+ // use one thread to update flag
317
+ if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
318
+}
319
+
320
+#endif
321
+
322
+template <typename P, int ngpus, typename A>
323
+DINLINE P packed_reduce(const P* ptrs[], int idx) {
324
+ A tmp = upcast(ptrs[0][idx]);
325
+#pragma unroll
326
+ for (int i = 1; i < ngpus; i++) {
327
+ packed_assign_add(tmp, upcast(ptrs[i][idx]));
328
+ }
329
+ return downcast<P>(tmp);
330
+}
331
+
332
+} // namespace vllm
csrc/flashkda_registration.cpp
new
+17
@@ -0,0 +1,17 @@
1
+#include "core/registration.h"
2
+#include "flash_kda.h"
3
+
4
+TORCH_LIBRARY(_flashkda_C, m) {
5
+ m.def("get_workspace_size(int T_total, int H, int N=1) -> int",
6
+ &get_workspace_size);
7
+ m.def(
8
+ "fwd(Tensor q, Tensor k, Tensor v, Tensor g, Tensor beta, float scale, "
9
+ "Tensor(a!) out, Tensor workspace, Tensor A_log, Tensor dt_bias, "
10
+ "float lower_bound, "
11
+ "Tensor? initial_state=None, Tensor(b!)? final_state=None, "
12
+ "Tensor? cu_seqlens=None) -> ()");
13
+}
14
+
15
+TORCH_LIBRARY_IMPL(_flashkda_C, CUDA, m) { m.impl("fwd", &fwd); }
16
+
17
+REGISTER_EXTENSION(_flashkda_C)
csrc/libtorch_stable/activation_kernels.cu
+108
@@ -464,6 +464,66 @@ __global__ void swigluoai_and_mul_kernel(
464
}
465
}
466
467
+// SITU (Kimi SituGLU) gated activation. Non-interleaved layout:
468
+// input = [gate(d), up(d)] per token.
469
+// gate_out = beta * tanh(gate / beta) * sigmoid(gate)
470
+// up_out = (linear_beta > 0) ? linear_beta * tanh(up / linear_beta) : up
471
+// out = gate_out * up_out
472
+// Compute is done in fp32 and written straight to `out` -- no intermediate
473
+// tensors and no full-tensor fp32 upcast (the pure-torch forward_native
474
+// allocated ~8 fp32 temporaries per call, which blows up MoE profiling).
475
+template <typename scalar_t>
476
+__global__ void situ_and_mul_kernel(
477
+ scalar_t* __restrict__ out, // [..., d]
478
+ const scalar_t* __restrict__ input, // [..., 2, d]
479
+ const int d, const float beta, const float linear_beta) {
480
+ const int64_t row = blockIdx.x;
481
+ const scalar_t* gate_ptr = input + row * 2 * d;
482
+ const scalar_t* up_ptr = gate_ptr + d;
483
+ scalar_t* out_ptr = out + row * d;
484
+ const bool clamp_up = linear_beta > 0.0f;
485
+ const float inv_beta = 1.0f / beta;
486
+ const float inv_linear_beta = clamp_up ? 1.0f / linear_beta : 0.0f;
487
+ for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
488
+ const float g = (float)VLLM_LDG(&gate_ptr[idx]);
489
+ const float u = (float)VLLM_LDG(&up_ptr[idx]);
490
+ const float gate_out = beta * tanhf(g * inv_beta) / (1.0f + expf(-g));
491
+ const float up_out =
492
+ clamp_up ? linear_beta * tanhf(u * inv_linear_beta) : u;
493
+ out_ptr[idx] = (scalar_t)(gate_out * up_out);
494
+ }
495
+}
496
+
497
+template <typename scalar_t>
498
+__global__ void masked_situ_and_mul_kernel(
499
+ scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
500
+ const int* __restrict__ expert_num_tokens, const int max_num_tokens,
501
+ const int d, const float beta, const float linear_beta) {
502
+ const int expert = blockIdx.y;
503
+ const int num_tokens = expert_num_tokens[expert];
504
+ const int idx = blockIdx.x * blockDim.x + threadIdx.x;
505
+ if (idx >= d || num_tokens == 0) {
506
+ return;
507
+ }
508
+
509
+ const bool clamp_up = linear_beta > 0.0f;
510
+ const float inv_beta = 1.0f / beta;
511
+ const float inv_linear_beta = clamp_up ? 1.0f / linear_beta : 0.0f;
512
+ const int64_t expert_row = static_cast<int64_t>(expert) * max_num_tokens;
513
+ for (int token = 0; token < num_tokens; ++token) {
514
+ const int64_t row = expert_row + token;
515
+ const scalar_t* gate_ptr = input + row * 2 * d;
516
+ const scalar_t* up_ptr = gate_ptr + d;
517
+ scalar_t* out_ptr = out + row * d;
518
+ const float g = (float)VLLM_LDG(&gate_ptr[idx]);
519
+ const float u = (float)VLLM_LDG(&up_ptr[idx]);
520
+ const float gate_out = beta * tanhf(g * inv_beta) / (1.0f + expf(-g));
521
+ const float up_out =
522
+ clamp_up ? linear_beta * tanhf(u * inv_linear_beta) : u;
523
+ out_ptr[idx] = (scalar_t)(gate_out * up_out);
524
+ }
525
+}
526
+
527
} // namespace vllm
528
529
#define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \
@@ -553,6 +613,54 @@ void swigluoai_and_mul(torch::stable::Tensor& out, // [..., d]
613
double alpha, double limit) {
614
LAUNCH_SIGLUOAI_AND_MUL(vllm::swigluoai_and_mul, alpha, limit);
615
}
616
+
617
+// Kimi SITU gated activation. `linear_beta <= 0` means "unset" (up passed
618
+// through), matching SituAndMul(linear_beta=None) on the Python side.
619
+void situ_and_mul(torch::stable::Tensor& out, // [..., d]
620
+ torch::stable::Tensor& input, // [..., 2 * d]
621
+ double beta, double linear_beta) {
622
+ int d = input.size(-1) / 2;
623
+ int64_t num_tokens = input.numel() / input.size(-1);
624
+ if (num_tokens == 0) {
625
+ return;
626
+ }
627
+ dim3 grid(num_tokens);
628
+ dim3 block(std::min(d, 1024));
629
+ const torch::stable::accelerator::DeviceGuard device_guard(
630
+ input.get_device_index());
631
+ const cudaStream_t stream = get_current_cuda_stream();
632
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
633
+ input.scalar_type(), "situ_and_mul_kernel", [&] {
634
+ vllm::situ_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>(
635
+ out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),
636
+ d, (float)beta, (float)linear_beta);
637
+ });
638
+}
639
+
640
+void masked_situ_and_mul(torch::stable::Tensor& out, // [E, T, d]
641
+ torch::stable::Tensor& input, // [E, T, 2 * d]
642
+ const torch::stable::Tensor& expert_num_tokens,
643
+ double beta, double linear_beta) {
644
+ int num_experts = input.size(0);
645
+ int max_num_tokens = input.size(1);
646
+ int d = input.size(2) / 2;
647
+ if (num_experts == 0 || max_num_tokens == 0) {
648
+ return;
649
+ }
650
+ constexpr int block_size = 256;
651
+ dim3 grid((d + block_size - 1) / block_size, num_experts);
652
+ dim3 block(block_size);
653
+ const torch::stable::accelerator::DeviceGuard device_guard(
654
+ input.get_device_index());
655
+ const cudaStream_t stream = get_current_cuda_stream();
656
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
657
+ input.scalar_type(), "masked_situ_and_mul_kernel", [&] {
658
+ vllm::masked_situ_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>(
659
+ out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),
660
+ expert_num_tokens.const_data_ptr<int>(), max_num_tokens, d,
661
+ (float)beta, (float)linear_beta);
662
+ });
663
+}
664
namespace vllm {
665
666
// Element-wise activation kernel template.
csrc/libtorch_stable/attention/merge_attn_states.cu
+32
-7
@@ -21,7 +21,10 @@ __global__ void merge_attn_states_kernel(
21
const float* prefix_lse, const scalar_t* suffix_output,
22
const float* suffix_lse, const uint num_tokens, const uint num_heads,
23
const uint head_size, const uint prefix_head_stride,
24
- const uint output_head_stride, const uint prefix_num_tokens,
24
+ const uint output_head_stride, const uint prefix_lse_head_stride,
25
+ const uint prefix_lse_token_stride, const uint suffix_lse_head_stride,
26
+ const uint suffix_lse_token_stride, const uint output_lse_head_stride,
27
+ const uint output_lse_token_stride, const uint prefix_num_tokens,
28
const float* output_scale) {
29
// Inputs always load 128-bit packs (pack_size elements of scalar_t).
30
// Outputs store pack_size elements of output_t, which is smaller for FP8.
@@ -84,15 +87,19 @@ __global__ void merge_attn_states_kernel(
87
}
88
}
89
if (output_lse != nullptr && pack_idx == 0) {
87
- float s_lse = suffix_lse[head_idx * num_tokens + token_idx];
88
- output_lse[head_idx * num_tokens + token_idx] = s_lse;
90
+ float s_lse = suffix_lse[head_idx * suffix_lse_head_stride +
91
+ token_idx * suffix_lse_token_stride];
92
+ output_lse[head_idx * output_lse_head_stride +
93
+ token_idx * output_lse_token_stride] = s_lse;
94
}
95
return;
96
}
97
98
// For tokens within prefix range, merge prefix and suffix
94
- float p_lse = prefix_lse[head_idx * num_tokens + token_idx];
95
- float s_lse = suffix_lse[head_idx * num_tokens + token_idx];
99
+ float p_lse = prefix_lse[head_idx * prefix_lse_head_stride +
100
+ token_idx * prefix_lse_token_stride];
101
+ float s_lse = suffix_lse[head_idx * suffix_lse_head_stride +
102
+ token_idx * suffix_lse_token_stride];
103
p_lse = std::isinf(p_lse) ? -std::numeric_limits<float>::infinity() : p_lse;
104
s_lse = std::isinf(s_lse) ? -std::numeric_limits<float>::infinity() : s_lse;
105
@@ -132,7 +139,8 @@ __global__ void merge_attn_states_kernel(
139
}
140
// We only need to write to output_lse once per head.
141
if (output_lse != nullptr && pack_idx == 0) {
135
- output_lse[head_idx * num_tokens + token_idx] = max_lse;
142
+ output_lse[head_idx * output_lse_head_stride +
143
+ token_idx * output_lse_token_stride] = max_lse;
144
}
145
return;
146
}
@@ -187,7 +195,8 @@ __global__ void merge_attn_states_kernel(
195
// We only need to write to output_lse once per head.
196
if (output_lse != nullptr && pack_idx == 0) {
197
float out_lse = logf(out_se) + max_lse;
190
- output_lse[head_idx * num_tokens + token_idx] = out_lse;
198
+ output_lse[head_idx * output_lse_head_stride +
199
+ token_idx * output_lse_token_stride] = out_lse;
200
}
201
}
202
@@ -221,6 +230,9 @@ __global__ void merge_attn_states_kernel(
230
reinterpret_cast<scalar_t*>(suffix_output.data_ptr()), \
231
reinterpret_cast<float*>(suffix_lse.data_ptr()), num_tokens, \
232
num_heads, head_size, prefix_head_stride, output_head_stride, \
233
+ prefix_lse_head_stride, prefix_lse_token_stride, \
234
+ suffix_lse_head_stride, suffix_lse_token_stride, \
235
+ output_lse_head_stride, output_lse_token_stride, \
236
prefix_num_tokens, output_scale_ptr); \
237
}
238
@@ -259,6 +271,19 @@ void merge_attn_states_launcher(
271
const uint head_size = output.size(2);
272
const uint prefix_head_stride = prefix_output.stride(1);
273
const uint output_head_stride = output.stride(1);
274
+ // lse tensors are [NUM_HEADS, NUM_TOKENS] but may be non-contiguous views
275
+ // (e.g. a transpose of a backend's [NUM_TOKENS, NUM_HEADS] output), so index
276
+ // them by their actual strides rather than assuming a contiguous layout.
277
+ const uint prefix_lse_head_stride = prefix_lse.stride(0);
278
+ const uint prefix_lse_token_stride = prefix_lse.stride(1);
279
+ const uint suffix_lse_head_stride = suffix_lse.stride(0);
280
+ const uint suffix_lse_token_stride = suffix_lse.stride(1);
281
+ uint output_lse_head_stride = 0;
282
+ uint output_lse_token_stride = 0;
283
+ if (output_lse.has_value()) {
284
+ output_lse_head_stride = output_lse.value().stride(0);
285
+ output_lse_token_stride = output_lse.value().stride(1);
286
+ }
287
// Thread mapping is based on input BF16 pack_size
288
const uint pack_size = 16 / sizeof(scalar_t);
289
STD_TORCH_CHECK(head_size % pack_size == 0,
csrc/libtorch_stable/cache_kernels.cu
+96
@@ -443,6 +443,55 @@ __global__ void concat_and_cache_mla_kernel(
443
copy(k_pe, kv_cache, k_pe_stride, block_stride, pe_dim, kv_lora_rank);
444
}
445
446
+// Grouped variant of concat_and_cache_mla: inserts the context K/V for every
447
+// draft layer in a single launch. Grid is (num_tokens, num_layers); each layer
448
+// reads its own cache base pointer from kv_cache_ptrs (same pointer-array
449
+// pattern as copy_blocks_kernel). bf16 only, so it is a raw 16-bit copy with no
450
+// scaling or quantization; scalar_t is uint16_t for portability.
451
+template <typename scalar_t>
452
+__global__ void concat_and_cache_mla_grouped_kernel(
453
+ const scalar_t* __restrict__ kv_c, // [num_layers, num_tokens,
454
+ // kv_lora_rank]
455
+ const scalar_t* __restrict__ k_pe, // [num_layers, num_tokens, pe_dim]
456
+ const int64_t* __restrict__ kv_cache_ptrs, // [num_layers]
457
+ const int64_t* __restrict__ slot_mapping, // [num_layers, num_tokens]
458
+ const int64_t kv_c_layer_stride, const int64_t kv_c_token_stride,
459
+ const int64_t k_pe_layer_stride, const int64_t k_pe_token_stride,
460
+ const int64_t slot_layer_stride, const int64_t block_stride,
461
+ const int64_t entry_stride, const int kv_lora_rank, const int pe_dim,
462
+ const int block_size) {
463
+ const int64_t token_idx = blockIdx.x;
464
+ const int64_t layer_idx = blockIdx.y;
465
+ const int64_t slot_idx =
466
+ slot_mapping[layer_idx * slot_layer_stride + token_idx];
467
+ // NOTE: slot_idx can be -1 if the token is padded
468
+ if (slot_idx < 0) {
469
+ return;
470
+ }
471
+ const int64_t block_idx = slot_idx / block_size;
472
+ const int64_t block_offset = slot_idx % block_size;
473
+
474
+ scalar_t* __restrict__ kv_cache =
475
+ reinterpret_cast<scalar_t*>(kv_cache_ptrs[layer_idx]);
476
+ const scalar_t* __restrict__ kv_c_layer =
477
+ kv_c + layer_idx * kv_c_layer_stride;
478
+ const scalar_t* __restrict__ k_pe_layer =
479
+ k_pe + layer_idx * k_pe_layer_stride;
480
+
481
+ auto copy = [&](const scalar_t* __restrict__ src, int64_t src_token_stride,
482
+ int size, int offset) {
483
+ for (int i = threadIdx.x; i < size; i += blockDim.x) {
484
+ const int64_t src_idx = token_idx * src_token_stride + i;
485
+ const int64_t dst_idx =
486
+ block_idx * block_stride + block_offset * entry_stride + i + offset;
487
+ kv_cache[dst_idx] = src[src_idx];
488
+ }
489
+ };
490
+
491
+ copy(kv_c_layer, kv_c_token_stride, kv_lora_rank, 0);
492
+ copy(k_pe_layer, k_pe_token_stride, pe_dim, kv_lora_rank);
493
+}
494
+
495
template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
496
__global__ void concat_and_cache_ds_mla_kernel(
497
const scalar_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank]
@@ -902,6 +951,53 @@ void concat_and_cache_mla(
951
}
952
}
953
954
+void concat_and_cache_mla_grouped(
955
+ torch::stable::Tensor& kv_c, // [num_layers, num_tokens, kv_lora_rank]
956
+ torch::stable::Tensor& k_pe, // [num_layers, num_tokens, pe_dim]
957
+ torch::stable::Tensor& kv_cache_ptrs, // [num_layers] int64, on device
958
+ torch::stable::Tensor& slot_mapping, // [num_layers, num_tokens] int64
959
+ int64_t block_size, int64_t block_stride, int64_t entry_stride) {
960
+ int num_layers = kv_c.size(0);
961
+ int num_tokens = kv_c.size(1);
962
+ int kv_lora_rank = kv_c.size(2);
963
+ int pe_dim = k_pe.size(2);
964
+
965
+ STD_TORCH_CHECK(
966
+ kv_c.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
967
+ k_pe.scalar_type() == torch::headeronly::ScalarType::BFloat16,
968
+ "concat_and_cache_mla_grouped only supports a bf16 KV cache; got kv_c=",
969
+ kv_c.scalar_type(), ", k_pe=", k_pe.scalar_type());
970
+ STD_TORCH_CHECK(
971
+ kv_cache_ptrs.scalar_type() == torch::headeronly::ScalarType::Long,
972
+ "kv_cache_ptrs must be int64");
973
+
974
+ if (num_tokens == 0 || num_layers == 0) {
975
+ return;
976
+ }
977
+
978
+ const int64_t kv_c_layer_stride = kv_c.stride(0);
979
+ const int64_t kv_c_token_stride = kv_c.stride(1);
980
+ const int64_t k_pe_layer_stride = k_pe.stride(0);
981
+ const int64_t k_pe_token_stride = k_pe.stride(1);
982
+ const int64_t slot_layer_stride = slot_mapping.stride(0);
983
+
984
+ const torch::stable::accelerator::DeviceGuard device_guard(
985
+ kv_c.get_device_index());
986
+ const cudaStream_t stream = get_current_cuda_stream();
987
+
988
+ dim3 grid(num_tokens, num_layers);
989
+ dim3 block(std::min(kv_lora_rank, 512));
990
+ vllm::concat_and_cache_mla_grouped_kernel<uint16_t>
991
+ <<<grid, block, 0, stream>>>(
992
+ reinterpret_cast<const uint16_t*>(kv_c.data_ptr()),
993
+ reinterpret_cast<const uint16_t*>(k_pe.data_ptr()),
994
+ kv_cache_ptrs.const_data_ptr<int64_t>(),
995
+ slot_mapping.const_data_ptr<int64_t>(), kv_c_layer_stride,
996
+ kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride,
997
+ slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim,
998
+ block_size);
999
+}
1000
+
1001
namespace vllm {
1002
1003
template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>
csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu
new
+362
@@ -0,0 +1,362 @@
1
+#include "torch_utils.h"
2
+
3
+#include <torch/csrc/stable/macros.h>
4
+#include <torch/csrc/stable/accelerator.h>
5
+#include <torch/csrc/stable/tensor.h>
6
+#include <torch/headeronly/core/ScalarType.h>
7
+
8
+#include "custom_all_reduce.cuh"
9
+#include "custom_all_gather_reduce_scatter.cuh"
10
+
11
+namespace vllm {
12
+
13
+void CustomAllreduce::allgather(cudaStream_t stream, void* input, void* output,
14
+ int size_bytes, int threads, int block_limit) {
15
+ if (size_bytes % sizeof(CopyPack) != 0)
16
+ throw std::runtime_error(
17
+ "custom allgather requires input byte size to be a multiple of " +
18
+ std::to_string(sizeof(CopyPack)));
19
+
20
+ auto ptrs = buffers_.at(input);
21
+ int size_per_rank = size_bytes / sizeof(CopyPack);
22
+ int total_size = size_per_rank * world_size_;
23
+ int blocks = std::min(block_limit, (total_size + threads - 1) / threads);
24
+
25
+#define AG_CASE(ngpus) \
26
+ case ngpus: \
27
+ cross_device_all_gather<ngpus><<<blocks, threads, 0, stream>>>( \
28
+ ptrs, sg_, self_sg_, reinterpret_cast<CopyPack*>(output), rank_, \
29
+ size_per_rank); \
30
+ break;
31
+
32
+ switch (world_size_) {
33
+ AG_CASE(2)
34
+ AG_CASE(4)
35
+ AG_CASE(6)
36
+ AG_CASE(8)
37
+ default:
38
+ throw std::runtime_error(
39
+ "custom allgather only supports num gpus in (2,4,6,8)");
40
+ }
41
+#undef AG_CASE
42
+}
43
+
44
+template <typename T>
45
+void CustomAllreduce::mnnvl_lamport_allgather(cudaStream_t stream, T* input,
46
+ T* output, void* local_buffer,
47
+ void* multicast_buffer,
48
+ uint32_t* epochs, int size_bytes,
49
+ int stage_size_bytes) {
50
+ if (size_bytes % sizeof(typename packed_t<T>::P) != 0 ||
51
+ stage_size_bytes % sizeof(typename packed_t<T>::P) != 0)
52
+ throw std::runtime_error(
53
+ "MNNVL Lamport allgather requires 16-byte aligned sizes");
54
+
55
+ auto ptrs = buffers_.at(local_buffer);
56
+ int size_per_rank = size_bytes / sizeof(typename packed_t<T>::P);
57
+ int stage_size = stage_size_bytes / sizeof(typename packed_t<T>::P);
58
+ int blocks =
59
+ (size_per_rank + kMnnvlLamportAgThreads - 1) / kMnnvlLamportAgThreads;
60
+
61
+#if !defined(USE_ROCM) && CUDA_VERSION >= 12000
62
+ cudaLaunchAttribute attributes[1]{};
63
+ attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
64
+ attributes[0].val.programmaticStreamSerializationAllowed = 1;
65
+ cudaLaunchConfig_t config{.gridDim = dim3(blocks),
66
+ .blockDim = dim3(kMnnvlLamportAgThreads),
67
+ .dynamicSmemBytes = 0,
68
+ .stream = stream,
69
+ .attrs = attributes,
70
+ .numAttrs = 1};
71
+ #define MNNVL_LAMPORT_AG_LAUNCH(ngpus) \
72
+ CUDACHECK(cudaLaunchKernelEx(&config, &mnnvl_lamport_all_gather<T, ngpus>, \
73
+ ptrs, input, output, \
74
+ reinterpret_cast<T*>(multicast_buffer), \
75
+ epochs, rank_, size_per_rank, stage_size))
76
+#else
77
+ #define MNNVL_LAMPORT_AG_LAUNCH(ngpus) \
78
+ mnnvl_lamport_all_gather<T, ngpus> \
79
+ <<<blocks, kMnnvlLamportAgThreads, 0, stream>>>( \
80
+ ptrs, input, output, reinterpret_cast<T*>(multicast_buffer), \
81
+ epochs, rank_, size_per_rank, stage_size)
82
+#endif
83
+
84
+#define MNNVL_LAMPORT_AG_CASE(ngpus) \
85
+ case ngpus: \
86
+ MNNVL_LAMPORT_AG_LAUNCH(ngpus); \
87
+ break;
88
+
89
+ switch (world_size_) {
90
+ MNNVL_LAMPORT_AG_CASE(2)
91
+ MNNVL_LAMPORT_AG_CASE(4)
92
+ MNNVL_LAMPORT_AG_CASE(6)
93
+ MNNVL_LAMPORT_AG_CASE(8)
94
+ MNNVL_LAMPORT_AG_CASE(16)
95
+ default:
96
+ throw std::runtime_error(
97
+ "MNNVL Lamport allgather only supports num gpus in (2,4,6,8,16)");
98
+ }
99
+#undef MNNVL_LAMPORT_AG_CASE
100
+#undef MNNVL_LAMPORT_AG_LAUNCH
101
+}
102
+
103
+template <typename T>
104
+void CustomAllreduce::reduce_scatter(cudaStream_t stream, T* input, T* output,
105
+ int size, int threads, int block_limit) {
106
+ auto packed_size = packed_t<T>::P::size;
107
+ if (size % (packed_size * world_size_) != 0)
108
+ throw std::runtime_error(
109
+ "custom reduce-scatter requires each output shard byte size to be "
110
+ "a multiple of 16");
111
+
112
+ auto ptrs = buffers_.at(input);
113
+ int size_per_rank = size / packed_size / world_size_;
114
+ int blocks = std::min(block_limit, (size_per_rank + threads - 1) / threads);
115
+
116
+#define RS_CASE(ngpus) \
117
+ case ngpus: \
118
+ cross_device_reduce_scatter<T, ngpus><<<blocks, threads, 0, stream>>>( \
119
+ ptrs, sg_, self_sg_, output, rank_, size_per_rank); \
120
+ break;
121
+
122
+ switch (world_size_) {
123
+ RS_CASE(2)
124
+ RS_CASE(4)
125
+ RS_CASE(6)
126
+ RS_CASE(8)
127
+ default:
128
+ throw std::runtime_error(
129
+ "custom reduce-scatter only supports num gpus in (2,4,6,8)");
130
+ }
131
+#undef RS_CASE
132
+}
133
+
134
+template <typename T>
135
+void CustomAllreduce::mnnvl_lamport_reduce_scatter(cudaStream_t stream,
136
+ T* input, T* output,
137
+ void* local_buffer,
138
+ uint32_t* epochs, int size,
139
+ int stage_size_bytes) {
140
+ auto packed_size = packed_t<T>::P::size;
141
+ if (size % (packed_size * world_size_) != 0 ||
142
+ stage_size_bytes % sizeof(typename packed_t<T>::P) != 0)
143
+ throw std::runtime_error(
144
+ "MNNVL Lamport reduce-scatter requires 16-byte aligned sizes");
145
+
146
+ auto ptrs = buffers_.at(local_buffer);
147
+ int size_per_rank = size / packed_size / world_size_;
148
+ int stage_size = stage_size_bytes / sizeof(typename packed_t<T>::P);
149
+ int blocks_per_rank =
150
+ (size_per_rank + kMnnvlLamportRsThreads - 1) / kMnnvlLamportRsThreads;
151
+ int blocks = blocks_per_rank * world_size_;
152
+
153
+#if !defined(USE_ROCM) && CUDA_VERSION >= 12000
154
+ cudaLaunchAttribute attributes[1]{};
155
+ attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
156
+ attributes[0].val.programmaticStreamSerializationAllowed = 1;
157
+ cudaLaunchConfig_t config{.gridDim = dim3(blocks),
158
+ .blockDim = dim3(kMnnvlLamportRsThreads),
159
+ .dynamicSmemBytes = 0,
160
+ .stream = stream,
161
+ .attrs = attributes,
162
+ .numAttrs = 1};
163
+ #define MNNVL_LAMPORT_RS_LAUNCH(ngpus) \
164
+ CUDACHECK(cudaLaunchKernelEx( \
165
+ &config, &mnnvl_lamport_reduce_scatter_kernel<T, ngpus>, ptrs, input, \
166
+ output, epochs, rank_, size_per_rank, stage_size))
167
+#else
168
+ #define MNNVL_LAMPORT_RS_LAUNCH(ngpus) \
169
+ mnnvl_lamport_reduce_scatter_kernel<T, ngpus> \
170
+ <<<blocks, kMnnvlLamportRsThreads, 0, stream>>>( \
171
+ ptrs, input, output, epochs, rank_, size_per_rank, stage_size)
172
+#endif
173
+
174
+#define MNNVL_LAMPORT_RS_CASE(ngpus) \
175
+ case ngpus: \
176
+ MNNVL_LAMPORT_RS_LAUNCH(ngpus); \
177
+ break;
178
+
179
+ switch (world_size_) {
180
+ MNNVL_LAMPORT_RS_CASE(2)
181
+ MNNVL_LAMPORT_RS_CASE(4)
182
+ MNNVL_LAMPORT_RS_CASE(6)
183
+ MNNVL_LAMPORT_RS_CASE(8)
184
+ MNNVL_LAMPORT_RS_CASE(16)
185
+ default:
186
+ throw std::runtime_error(
187
+ "MNNVL Lamport reduce-scatter only supports num gpus in "
188
+ "(2,4,6,8,16)");
189
+ }
190
+#undef MNNVL_LAMPORT_RS_CASE
191
+#undef MNNVL_LAMPORT_RS_LAUNCH
192
+}
193
+
194
+} // namespace vllm
195
+
196
+using fptr_t = int64_t;
197
+static_assert(sizeof(void*) == sizeof(fptr_t));
198
+
199
+bool _is_weak_contiguous(torch::stable::Tensor& t);
200
+
201
+void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp,
202
+ torch::stable::Tensor& out, fptr_t _reg_buffer,
203
+ int64_t reg_buffer_sz_bytes) {
204
+ auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
205
+ const torch::stable::accelerator::DeviceGuard device_guard(
206
+ inp.get_device_index());
207
+ const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
208
+
209
+ STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
210
+ STD_TORCH_CHECK((inp.numel() * fa->world_size_) == (out.numel()));
211
+ STD_TORCH_CHECK(_is_weak_contiguous(out));
212
+ STD_TORCH_CHECK(_is_weak_contiguous(inp));
213
+ auto input_size = inp.numel() * inp.element_size();
214
+ auto reg_buffer = reinterpret_cast<void*>(_reg_buffer);
215
+ STD_TORCH_CHECK(reg_buffer != nullptr);
216
+ STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes));
217
+ STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size,
218
+ cudaMemcpyDeviceToDevice, stream));
219
+ fa->allgather(stream, reg_buffer, out.mutable_data_ptr(), input_size);
220
+}
221
+
222
+void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp,
223
+ torch::stable::Tensor& out, fptr_t _local_buffer,
224
+ fptr_t _multicast_buffer, fptr_t _epoch_buffer,
225
+ int64_t stage_sz_bytes) {
226
+ auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
227
+ const torch::stable::accelerator::DeviceGuard device_guard(
228
+ inp.get_device_index());
229
+ const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
230
+
231
+ STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
232
+ STD_TORCH_CHECK((inp.numel() * fa->world_size_) == (out.numel()));
233
+ STD_TORCH_CHECK(_is_weak_contiguous(out));
234
+ STD_TORCH_CHECK(_is_weak_contiguous(inp));
235
+ auto input_size = inp.numel() * inp.element_size();
236
+ STD_TORCH_CHECK((input_size * fa->world_size_) <= stage_sz_bytes);
237
+ auto local_buffer = reinterpret_cast<void*>(_local_buffer);
238
+ auto multicast_buffer = reinterpret_cast<void*>(_multicast_buffer);
239
+ auto epochs = reinterpret_cast<uint32_t*>(_epoch_buffer);
240
+ switch (out.scalar_type()) {
241
+ case torch::headeronly::ScalarType::Float: {
242
+ fa->mnnvl_lamport_allgather<float>(
243
+ stream, reinterpret_cast<float*>(inp.mutable_data_ptr()),
244
+ reinterpret_cast<float*>(out.mutable_data_ptr()), local_buffer,
245
+ multicast_buffer, epochs, input_size, stage_sz_bytes);
246
+ break;
247
+ }
248
+ case torch::headeronly::ScalarType::Half: {
249
+ fa->mnnvl_lamport_allgather<half>(
250
+ stream, reinterpret_cast<half*>(inp.mutable_data_ptr()),
251
+ reinterpret_cast<half*>(out.mutable_data_ptr()), local_buffer,
252
+ multicast_buffer, epochs, input_size, stage_sz_bytes);
253
+ break;
254
+ }
255
+#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
256
+ case torch::headeronly::ScalarType::BFloat16: {
257
+ fa->mnnvl_lamport_allgather<nv_bfloat16>(
258
+ stream, reinterpret_cast<nv_bfloat16*>(inp.mutable_data_ptr()),
259
+ reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), local_buffer,
260
+ multicast_buffer, epochs, input_size, stage_sz_bytes);
261
+ break;
262
+ }
263
+#endif
264
+ default:
265
+ throw std::runtime_error(
266
+ "MNNVL Lamport allgather only supports float32, float16 and "
267
+ "bfloat16");
268
+ }
269
+}
270
+
271
+void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
272
+ torch::stable::Tensor& out, fptr_t _reg_buffer,
273
+ int64_t reg_buffer_sz_bytes) {
274
+ auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
275
+ const torch::stable::accelerator::DeviceGuard device_guard(
276
+ inp.get_device_index());
277
+ const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
278
+
279
+ STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
280
+ STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel()));
281
+ STD_TORCH_CHECK(_is_weak_contiguous(out));
282
+ STD_TORCH_CHECK(_is_weak_contiguous(inp));
283
+ auto input_size = inp.numel() * inp.element_size();
284
+ auto reg_buffer = reinterpret_cast<void*>(_reg_buffer);
285
+ STD_TORCH_CHECK(reg_buffer != nullptr);
286
+ STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes));
287
+ STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size,
288
+ cudaMemcpyDeviceToDevice, stream));
289
+ switch (out.scalar_type()) {
290
+ case torch::headeronly::ScalarType::Float: {
291
+ fa->reduce_scatter<float>(
292
+ stream, reinterpret_cast<float*>(reg_buffer),
293
+ reinterpret_cast<float*>(out.mutable_data_ptr()), inp.numel());
294
+ break;
295
+ }
296
+ case torch::headeronly::ScalarType::Half: {
297
+ fa->reduce_scatter<half>(stream, reinterpret_cast<half*>(reg_buffer),
298
+ reinterpret_cast<half*>(out.mutable_data_ptr()),
299
+ inp.numel());
300
+ break;
301
+ }
302
+#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
303
+ case torch::headeronly::ScalarType::BFloat16: {
304
+ fa->reduce_scatter<nv_bfloat16>(
305
+ stream, reinterpret_cast<nv_bfloat16*>(reg_buffer),
306
+ reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), inp.numel());
307
+ break;
308
+ }
309
+#endif
310
+ default:
311
+ throw std::runtime_error(
312
+ "custom reduce-scatter only supports float32, float16 and bfloat16");
313
+ }
314
+}
315
+
316
+void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
317
+ torch::stable::Tensor& out,
318
+ fptr_t _local_buffer, fptr_t _epoch_buffer,
319
+ int64_t stage_sz_bytes) {
320
+ auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
321
+ const torch::stable::accelerator::DeviceGuard device_guard(
322
+ inp.get_device_index());
323
+ const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
324
+
325
+ STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
326
+ STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel()));
327
+ STD_TORCH_CHECK(_is_weak_contiguous(out));
328
+ STD_TORCH_CHECK(_is_weak_contiguous(inp));
329
+ auto input_size = inp.numel() * inp.element_size();
330
+ STD_TORCH_CHECK(input_size <= stage_sz_bytes);
331
+ auto local_buffer = reinterpret_cast<void*>(_local_buffer);
332
+ auto epochs = reinterpret_cast<uint32_t*>(_epoch_buffer);
333
+ switch (out.scalar_type()) {
334
+ case torch::headeronly::ScalarType::Float: {
335
+ fa->mnnvl_lamport_reduce_scatter<float>(
336
+ stream, reinterpret_cast<float*>(inp.mutable_data_ptr()),
337
+ reinterpret_cast<float*>(out.mutable_data_ptr()), local_buffer,
338
+ epochs, inp.numel(), stage_sz_bytes);
339
+ break;
340
+ }
341
+ case torch::headeronly::ScalarType::Half: {
342
+ fa->mnnvl_lamport_reduce_scatter<half>(
343
+ stream, reinterpret_cast<half*>(inp.mutable_data_ptr()),
344
+ reinterpret_cast<half*>(out.mutable_data_ptr()), local_buffer, epochs,
345
+ inp.numel(), stage_sz_bytes);
346
+ break;
347
+ }
348
+#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
349
+ case torch::headeronly::ScalarType::BFloat16: {
350
+ fa->mnnvl_lamport_reduce_scatter<nv_bfloat16>(
351
+ stream, reinterpret_cast<nv_bfloat16*>(inp.mutable_data_ptr()),
352
+ reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), local_buffer,
353
+ epochs, inp.numel(), stage_sz_bytes);
354
+ break;
355
+ }
356
+#endif
357
+ default:
358
+ throw std::runtime_error(
359
+ "MNNVL Lamport reduce-scatter only supports float32, float16 and "
360
+ "bfloat16");
361
+ }
362
+}
csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp
new
+29
@@ -0,0 +1,29 @@
1
+#include "ops.h"
2
+#include "core/registration.h"
3
+
4
+#include <torch/csrc/stable/library.h>
5
+
6
+STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ag_rs) {
7
+ custom_ag_rs.def(
8
+ "custom_all_gather(int fa, Tensor inp, Tensor! out, int reg_buffer, "
9
+ "int reg_buffer_sz_bytes) -> ()");
10
+ custom_ag_rs.def(
11
+ "mnnvl_lamport_all_gather(int fa, Tensor inp, Tensor! out, int "
12
+ "local_buffer, int multicast_buffer, int epoch_buffer, int "
13
+ "stage_sz_bytes) -> ()");
14
+ custom_ag_rs.def(
15
+ "custom_reduce_scatter(int fa, Tensor inp, Tensor! out, int reg_buffer, "
16
+ "int reg_buffer_sz_bytes) -> ()");
17
+ custom_ag_rs.def(
18
+ "mnnvl_lamport_reduce_scatter(int fa, Tensor inp, Tensor! out, int "
19
+ "local_buffer, int epoch_buffer, int stage_sz_bytes) -> ()");
20
+}
21
+
22
+STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ag_rs) {
23
+ custom_ag_rs.impl("custom_all_gather", TORCH_BOX(&custom_all_gather));
24
+ custom_ag_rs.impl("mnnvl_lamport_all_gather",
25
+ TORCH_BOX(&mnnvl_lamport_all_gather));
26
+ custom_ag_rs.impl("custom_reduce_scatter", TORCH_BOX(&custom_reduce_scatter));
27
+ custom_ag_rs.impl("mnnvl_lamport_reduce_scatter",
28
+ TORCH_BOX(&mnnvl_lamport_reduce_scatter));
29
+}
csrc/libtorch_stable/custom_all_reduce.cu
+4
-4
@@ -18,14 +18,14 @@ fptr_t init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs,
18
torch::stable::Tensor& rank_data, int64_t rank,
19
bool fully_connected) {
20
int world_size = fake_ipc_ptrs.size();
21
- if (world_size > 8)
22
- throw std::invalid_argument("world size > 8 is not supported");
21
+ if (world_size > vllm::kMaxCustomCollectiveRanks)
22
+ throw std::invalid_argument("world size > 16 is not supported");
23
if (world_size % 2 != 0)
24
throw std::invalid_argument("Odd num gpus is not supported for now");
25
if (rank < 0 || rank >= world_size)
26
throw std::invalid_argument("invalid rank passed in");
27
28
- vllm::Signal* ipc_ptrs[8];
28
+ vllm::Signal* ipc_ptrs[vllm::kMaxCustomCollectiveRanks];
29
for (int i = 0; i < world_size; i++) {
30
ipc_ptrs[i] = reinterpret_cast<vllm::Signal*>(fake_ipc_ptrs[i]);
31
}
@@ -124,7 +124,7 @@ int64_t meta_size() { return sizeof(vllm::Signal); }
124
void register_buffer(fptr_t _fa, const std::vector<fptr_t>& fake_ipc_ptrs) {
125
auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
126
STD_TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_);
127
- void* ipc_ptrs[8];
127
+ void* ipc_ptrs[vllm::kMaxCustomCollectiveRanks];
128
for (int i = 0; i < fake_ipc_ptrs.size(); i++) {
129
ipc_ptrs[i] = reinterpret_cast<void*>(fake_ipc_ptrs[i]);
130
}
csrc/libtorch_stable/dsv3_fused_a_gemm.cu
+116
-35
@@ -647,17 +647,17 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
647
#endif
648
}
649
650
-template <typename T, int kHdIn, int kHdOut, int kTileN>
650
+template <typename T, int kHdIn, int kHdOut, int kTileN, int kTileK = 256>
651
void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
652
- cudaStream_t const stream) {
653
- constexpr int gemm_m = kHdOut; // 2112
654
- int const gemm_n = num_tokens; // 1-16
655
- constexpr int gemm_k = kHdIn; // 7168
652
+ cudaStream_t const stream, bool enable_pdl) {
653
+ constexpr int gemm_m = kHdOut;
654
+ int const gemm_n = num_tokens;
655
+ constexpr int gemm_k = kHdIn;
656
constexpr int batch_size = 1;
657
std::swap(mat_a, mat_b);
658
constexpr int tile_m = 16;
659
- constexpr int tile_n = kTileN; // 8 or 16
660
- constexpr int tile_k = std::max(256, 1024 / tile_n); // 256
659
+ constexpr int tile_n = kTileN;
660
+ constexpr int tile_k = kTileK;
661
constexpr int max_stage_cnt =
662
1024 * 192 / ((tile_m + tile_n) * tile_k * sizeof(bf16_t));
663
constexpr int k_iter_cnt = gemm_k / tile_k;
@@ -679,7 +679,8 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
679
config.stream = stream;
680
cudaLaunchAttribute attrs[1];
681
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
682
- attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
682
+ attrs[0].val.programmaticStreamSerializationAllowed =
683
+ enable_pdl || getEnvEnablePDL();
684
config.numAttrs = 1;
685
config.attrs = attrs;
686
if (smem_bytes >= (48 * 1024)) {
@@ -694,36 +695,50 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
695
output, mat_a, mat_b, gemm_n);
696
}
697
697
-template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>(
698
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
699
- cudaStream_t);
700
-
701
-template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
702
- __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
703
- cudaStream_t);
698
+template <typename T, int kHdIn, int kHdOut, int kTileK = 256>
699
+void invokeFusedAGemmForTokens(T* output, T const* mat_a, T const* mat_b,
700
+ int num_tokens, cudaStream_t const stream,
701
+ bool enable_pdl) {
702
+ if (num_tokens <= 8) {
703
+ invokeFusedAGemm<T, kHdIn, kHdOut, 8, kTileK>(
704
+ output, mat_a, mat_b, num_tokens, stream, enable_pdl);
705
+ } else {
706
+ invokeFusedAGemm<T, kHdIn, kHdOut, 16, kTileK>(
707
+ output, mat_a, mat_b, num_tokens, stream, enable_pdl);
708
+ }
709
+}
710
711
void dsv3_fused_a_gemm(torch::stable::Tensor& output,
712
torch::stable::Tensor const& mat_a,
707
- torch::stable::Tensor const& mat_b) {
713
+ torch::stable::Tensor const& mat_b, bool enable_pdl) {
714
STD_TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
715
int const num_tokens = mat_a.size(0);
716
int const hd_in = mat_a.size(1);
717
int const hd_out = mat_b.size(1);
718
713
- constexpr int kHdIn = 7168;
714
- constexpr int kHdOut = 2112;
719
STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
720
"required 1 <= mat_a.shape[0] <= 16");
717
- STD_TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168");
718
- STD_TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112");
721
STD_TORCH_CHECK(output.size(0) == num_tokens,
722
"required output.shape[0] == mat_a.shape[0]");
723
STD_TORCH_CHECK(output.size(1) == hd_out,
724
"required output.shape[1] == mat_b.shape[1]");
725
+ STD_TORCH_CHECK(mat_b.size(0) == hd_in,
726
+ "required mat_b.shape[0] == mat_a.shape[1]");
727
724
- STD_TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
725
- STD_TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
726
- STD_TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
728
+ STD_TORCH_CHECK(mat_a.get_device_index() == mat_b.get_device_index() &&
729
+ mat_a.get_device_index() == output.get_device_index(),
730
+ "mat_a, mat_b, and output must be on the same device");
731
+
732
+ // The kernels index global memory with raw pointers and packed strides, so
733
+ // reject any padded or transposed view rather than reading out of bounds.
734
+ STD_TORCH_CHECK(
735
+ mat_a.stride(0) == hd_in && mat_a.stride(1) == 1,
736
+ "mat_a must be a packed row-major [num_tokens, hd_in] tensor");
737
+ STD_TORCH_CHECK(
738
+ output.stride(0) == hd_out && output.stride(1) == 1,
739
+ "output must be a packed row-major [num_tokens, hd_out] tensor");
740
+ STD_TORCH_CHECK(mat_b.stride(0) == 1 && mat_b.stride(1) == hd_in,
741
+ "mat_b must be a packed column-major [hd_in, hd_out] tensor");
742
743
STD_TORCH_CHECK(
744
mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
@@ -738,19 +753,85 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,
753
STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
754
755
auto stream = get_current_cuda_stream(mat_a.get_device_index());
741
- if (num_tokens <= 8) {
742
- invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>(
743
- reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
744
- reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
745
- reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
746
- stream);
747
- } else {
748
- invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>(
749
- reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
750
- reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
751
- reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
752
- stream);
756
+ auto* output_ptr =
757
+ reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr());
758
+ auto const* mat_a_ptr =
759
+ reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
760
+ auto const* mat_b_ptr =
761
+ reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr());
762
+
763
+#define DISPATCH_DSV3_SHAPE(HD_IN, HD_OUT) \
764
+ if (hd_in == HD_IN && hd_out == HD_OUT) { \
765
+ invokeFusedAGemmForTokens<__nv_bfloat16, HD_IN, HD_OUT>( \
766
+ output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); \
767
+ return; \
768
}
769
+
770
+ // Shapes the Kimi-K3 selector routes to dsv3_fused_a (see the dsv3 winners
771
+ // in KIMI_K3_PROJECTIONS) plus the DeepSeek V2/V3 QKV A-projection.
772
+ DISPATCH_DSV3_SHAPE(7168, 1536)
773
+ DISPATCH_DSV3_SHAPE(7168, 2112)
774
+ DISPATCH_DSV3_SHAPE(1536, 2304)
775
+ DISPATCH_DSV3_SHAPE(1536, 4608)
776
+ DISPATCH_DSV3_SHAPE(7168, 3584)
777
+ DISPATCH_DSV3_SHAPE(768, 7168)
778
+ // TP16 dsv3 winners, as (hd_in=K, hd_out=N). TP16 dense down_proj is absent
779
+ // because hd_in=2112 is not a multiple of any supported tile_k.
780
+ DISPATCH_DSV3_SHAPE(1536, 1152)
781
+ DISPATCH_DSV3_SHAPE(7168, 768)
782
+ DISPATCH_DSV3_SHAPE(7168, 3216)
783
+ DISPATCH_DSV3_SHAPE(7168, 4224)
784
+
785
+#ifdef VLLM_K3_BENCH_SHAPES
786
+ // The selector routes these shapes to CuTe or the default GEMM, so they are
787
+ // never reached in production. They are compiled only for offline
788
+ // DSV3-vs-CuTe benchmarking.
789
+ DISPATCH_DSV3_SHAPE(7168, 6288)
790
+ DISPATCH_DSV3_SHAPE(1536, 7168)
791
+ DISPATCH_DSV3_SHAPE(3584, 7168)
792
+ DISPATCH_DSV3_SHAPE(7168, 8448)
793
+ DISPATCH_DSV3_SHAPE(7168, 20480)
794
+ DISPATCH_DSV3_SHAPE(7168, 3072)
795
+ DISPATCH_DSV3_SHAPE(7168, 12448)
796
+ DISPATCH_DSV3_SHAPE(3072, 7168)
797
+ DISPATCH_DSV3_SHAPE(8448, 7168)
798
+ DISPATCH_DSV3_SHAPE(7168, 16896)
799
+ DISPATCH_DSV3_SHAPE(7168, 40960)
800
+#endif
801
+
802
+#undef DISPATCH_DSV3_SHAPE
803
+
804
+ if (hd_in == 128 && hd_out == 1536) {
805
+ invokeFusedAGemmForTokens<__nv_bfloat16, 128, 1536, 128>(
806
+ output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
807
+ return;
808
+ }
809
+ if (hd_in == 128 && hd_out == 3072) {
810
+ invokeFusedAGemmForTokens<__nv_bfloat16, 128, 3072, 128>(
811
+ output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
812
+ return;
813
+ }
814
+ // TP16 KDA f_b_proj and shared_expert down_proj. Neither hd_in is a multiple
815
+ // of 256, so both need the 128 tile_k.
816
+ if (hd_in == 128 && hd_out == 768) {
817
+ invokeFusedAGemmForTokens<__nv_bfloat16, 128, 768, 128>(
818
+ output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
819
+ return;
820
+ }
821
+ if (hd_in == 384 && hd_out == 7168) {
822
+ invokeFusedAGemmForTokens<__nv_bfloat16, 384, 7168, 128>(
823
+ output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
824
+ return;
825
+ }
826
+#ifdef VLLM_K3_BENCH_SHAPES
827
+ if (hd_in == 4224 && hd_out == 7168) {
828
+ invokeFusedAGemmForTokens<__nv_bfloat16, 4224, 7168, 128>(
829
+ output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl);
830
+ return;
831
+ }
832
+#endif
833
+
834
+ STD_TORCH_CHECK(false, "unsupported DSV3 fused-A GEMM shape");
835
}
836
837
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu
new
+1237
@@ -0,0 +1,1237 @@
1
+/*
2
+ * SPDX-License-Identifier: Apache-2.0
3
+ * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+ *
5
+ * Fused Kimi-K3 MLA prefill + decode epilogues with optional RoPE.
6
+ *
7
+ * Prefill: runs after the q_b_proj / kv_b_proj GEMMs, one launch per token
8
+ * slice. Decode: runs after BMM1 (q_nope x W_UK) right before forward_mqa,
9
+ * concatenating mqa_q = [ql_nope | q_pe] and inserting the latent cache
10
+ * (fused_kimi_k3_mla_decode_q_concat_kv_cache_{,fp8_,ds_mla_}insert).
11
+ *
12
+ * Prefill variants:
13
+ *
14
+ * bf16 (fused_kimi_k3_mla_key_concat_kv_cache_insert):
15
+ * - optional in-place q RoPE: rotate q[t, h, 128:192]
16
+ * - full key concat: k_out[t, h] = [k_nope[t, h] | k_pe[t]] (per head)
17
+ * - latent cache insert: cache[slot(t)] = [kv_c_normed[t] | k_pe[t]]
18
+ * (v is used as-is in bf16, so it is not touched here.)
19
+ *
20
+ * fp8 (fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert):
21
+ * - q_fp8[t, h] = quant(q[t, h], q_scale)
22
+ * - k_fp8[t, h] = quant([k_nope[t, h] | k_pe[t]], k_scale)
23
+ * - v_fp8[t, h] = quant(v[t, h], v_scale)
24
+ * - cache[slot(t)] = quant([kv_c_normed[t] | k_pe[t]], k_scale)
25
+ * matching MLA's _q_scale / _k_scale / _v_scale (the cache latent uses a
26
+ * separate cache scale). Per-tensor E4M3.
27
+ *
28
+ * fp8_ds_mla (fused_kimi_k3_mla_key_concat_ds_mla_insert):
29
+ * - full key concat (bf16), and cache insert in DeepSeek's 656-byte
30
+ * block-scaled layout (NoPE fp8 in 4 tiles of 128 with per-tile dynamic
31
+ * scales, RoPE bf16), bit-compatible with concat_and_cache_ds_mla_kernel.
32
+ *
33
+ * Both use Programmatic Dependent Launch (PDL) to overlap the tail of the
34
+ * producing GEMMs on sm_90+, and are structured after
35
+ * `fusedDeepseekV4FullCacheKernel`: one grid, one warp per (token, slot) with
36
+ * `slotsPerToken = num_heads + 1`. Slots [0, H) do the per-head work; the extra
37
+ * slot H does the per-token cache insert. All dims are multiples of 8, so bf16
38
+ * copies move one uint4 (8 elems) per step and fp8 stores pack 8 elems into a
39
+ * uint2.
40
+ */
41
+
42
+#include "torch_utils.h"
43
+
44
+#include <torch/csrc/stable/macros.h>
45
+#include <torch/csrc/stable/accelerator.h>
46
+#include <torch/csrc/stable/tensor.h>
47
+#include <torch/csrc/stable/ops.h>
48
+#include <torch/headeronly/core/ScalarType.h>
49
+#include <torch/csrc/stable/device.h>
50
+
51
+#include "cuda_compat.h"
52
+#include "dispatch_utils.h"
53
+#include "type_convert.cuh"
54
+
55
+#ifndef USE_ROCM
56
+ #include <cuda_fp8.h>
57
+ #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
58
+#else
59
+ #include <hip/hip_fp8.h>
60
+ #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
61
+#endif
62
+#include <cuda_runtime.h>
63
+#include <cfloat>
64
+#include <type_traits>
65
+
66
+#ifdef USE_ROCM
67
+__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) {
68
+ #if defined(__gfx942__)
69
+ __hip_fp8_e4m3_fnuz fp8_val(val);
70
+ #else
71
+ __hip_fp8_e4m3 fp8_val(val);
72
+ #endif
73
+ return reinterpret_cast<uint8_t&>(fp8_val);
74
+}
75
+#endif
76
+
77
+namespace vllm {
78
+namespace kimi_k3_fused_ops {
79
+
80
+namespace {
81
+inline int getSMVersion() {
82
+ auto* props = get_device_prop();
83
+ return props->major * 10 + props->minor;
84
+}
85
+} // namespace
86
+
87
+// ────────────────────────────────────────────────────────────────────────────
88
+// Constants (Kimi-K3 MLA)
89
+// ────────────────────────────────────────────────────────────────────────────
90
+constexpr int kKvLoraRank = 512; // L
91
+constexpr int kQkNopeHeadDim = 128; // P
92
+constexpr int kQkRopeHeadDim = 64; // R
93
+constexpr int kQkHeadDim = kQkNopeHeadDim + kQkRopeHeadDim; // 192
94
+constexpr int kVHeadDim = 128; // V
95
+constexpr int kCacheEntry = kKvLoraRank + kQkRopeHeadDim; // 576
96
+constexpr int kVecElems = 8; // 8 bf16 == one uint4 load / one uint2 fp8 store
97
+
98
+#if defined(USE_ROCM) && defined(__gfx942__)
99
+constexpr float kFp8Max = 224.0f;
100
+#else
101
+constexpr float kFp8Max = 448.0f;
102
+#endif
103
+// Divisor for fp8_ds_mla per-tile dynamic scales (matches cache_kernels.cu).
104
+// fp8_ds_mla 656B entry: [0,512) NoPE fp8 (4 tiles of 128), [512,528) 4 fp32
105
+// tile scales, [528,656) RoPE 64 bf16.
106
+constexpr float kFp8ScaleDivisor = kFp8Max;
107
+
108
+// Copy 8 source elements (one uint4 of bf16/fp16) to `dst`. FP8=false stores a
109
+// uint4 (bf16); FP8=true decodes to fp32, scales by `scale_inv`, saturates to
110
+// ±kFp8Max and packs into a uint2 of E4M3.
111
+template <typename scalar_t, bool FP8, bool APPLY_ROPE = false>
112
+__device__ __forceinline__ void copyChunk8(void* dst, const scalar_t* src,
113
+ float scale_inv,
114
+ const float* cos_sin = nullptr,
115
+ int rope_elem_base = 0) {
116
+ uint4 const v = *reinterpret_cast<const uint4*>(src);
117
+ if constexpr (FP8 || APPLY_ROPE) {
118
+ using Converter = vllm::_typeConvert<scalar_t>;
119
+ auto const* p =
120
+ reinterpret_cast<typename Converter::packed_hip_type const*>(&v);
121
+ float f[kVecElems];
122
+#pragma unroll
123
+ for (int i = 0; i < 4; i++) {
124
+ float2 x = Converter::convert(p[i]);
125
+ f[2 * i] = x.x;
126
+ f[2 * i + 1] = x.y;
127
+ }
128
+ if constexpr (APPLY_ROPE) {
129
+#pragma unroll
130
+ for (int i = 0; i < kVecElems / 2; i++) {
131
+ int const pair_idx = rope_elem_base / 2 + i;
132
+ float const cos = static_cast<float>(cos_sin[pair_idx]);
133
+ float const sin =
134
+ static_cast<float>(cos_sin[pair_idx + kQkRopeHeadDim / 2]);
135
+ float const x = f[2 * i];
136
+ float const y = f[2 * i + 1];
137
+ f[2 * i] = x * cos - y * sin;
138
+ f[2 * i + 1] = x * sin + y * cos;
139
+ }
140
+ }
141
+ if constexpr (!FP8) {
142
+ uint4 out;
143
+ auto* o = reinterpret_cast<typename Converter::packed_hip_type*>(&out);
144
+#pragma unroll
145
+ for (int i = 0; i < kVecElems / 2; i++) {
146
+ o[i] = Converter::convert(make_float2(f[2 * i], f[2 * i + 1]));
147
+ }
148
+ *reinterpret_cast<uint4*>(dst) = out;
149
+ return;
150
+ }
151
+#ifndef USE_ROCM
152
+ uint2 out;
153
+ auto* o2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&out);
154
+ #pragma unroll
155
+ for (int i = 0; i < 4; i++) {
156
+ float2 s = make_float2(f[2 * i] * scale_inv, f[2 * i + 1] * scale_inv);
157
+ s.x = fminf(fmaxf(s.x, -kFp8Max), kFp8Max);
158
+ s.y = fminf(fmaxf(s.y, -kFp8Max), kFp8Max);
159
+ o2[i] = __nv_cvt_float2_to_fp8x2(s, __NV_SATFINITE, __NV_E4M3);
160
+ }
161
+ *reinterpret_cast<uint2*>(dst) = out;
162
+#else
163
+ uint8_t out[kVecElems];
164
+ #pragma unroll
165
+ for (int i = 0; i < kVecElems; i++) {
166
+ float s = fminf(fmaxf(f[i] * scale_inv, -kFp8Max), kFp8Max);
167
+ out[i] = rocm_cvt_float_to_fp8_e4m3(s);
168
+ }
169
+ *reinterpret_cast<uint2*>(dst) = *reinterpret_cast<uint2 const*>(out);
170
+#endif
171
+ } else {
172
+ *reinterpret_cast<uint4*>(dst) = v;
173
+ }
174
+}
175
+
176
+// Concat + store one head's full key: dst[e] = [k_nope | k_pe], e in [0, 192).
177
+// FP8 dst is byte-addressed; bf16 dst is scalar_t-addressed (dst_elem_size).
178
+template <typename scalar_t, bool FP8, bool APPLY_ROPE = false>
179
+__device__ __forceinline__ void writeFullKey(void* dst, const scalar_t* k_nope,
180
+ const scalar_t* k_pe, int laneId,
181
+ int dst_elem_size, float scale_inv,
182
+ const float* cos_sin = nullptr) {
183
+ auto* d = reinterpret_cast<uint8_t*>(dst);
184
+ for (int e = laneId * kVecElems; e < kQkHeadDim; e += 32 * kVecElems) {
185
+ if (e < kQkNopeHeadDim) {
186
+ copyChunk8<scalar_t, FP8>(d + e * dst_elem_size, k_nope + e, scale_inv);
187
+ } else {
188
+ int const rope_e = e - kQkNopeHeadDim;
189
+ copyChunk8<scalar_t, FP8, APPLY_ROPE>(
190
+ d + e * dst_elem_size, k_pe + rope_e, scale_inv, cos_sin, rope_e);
191
+ }
192
+ }
193
+}
194
+
195
+// Store a prefill query, rotating only q[..., 128:192]. For bf16 dst may alias
196
+// q (in-place); fp8 writes the quantized query directly to its output.
197
+template <typename scalar_t, bool FP8, bool APPLY_ROPE = false>
198
+__device__ __forceinline__ void writePrefillQuery(
199
+ void* dst, const scalar_t* q, int laneId, int dst_elem_size,
200
+ float scale_inv, const float* cos_sin = nullptr) {
201
+ auto* d = reinterpret_cast<uint8_t*>(dst);
202
+ if constexpr (FP8) {
203
+ for (int e = laneId * kVecElems; e < kQkHeadDim; e += 32 * kVecElems) {
204
+ if (e < kQkNopeHeadDim) {
205
+ copyChunk8<scalar_t, true>(d + e * dst_elem_size, q + e, scale_inv);
206
+ } else {
207
+ int const rope_e = e - kQkNopeHeadDim;
208
+ copyChunk8<scalar_t, true, APPLY_ROPE>(d + e * dst_elem_size, q + e,
209
+ scale_inv, cos_sin, rope_e);
210
+ }
211
+ }
212
+ } else if constexpr (APPLY_ROPE) {
213
+ for (int e = laneId * kVecElems; e < kQkRopeHeadDim; e += 32 * kVecElems) {
214
+ copyChunk8<scalar_t, false, true>(
215
+ d + (kQkNopeHeadDim + e) * dst_elem_size, q + kQkNopeHeadDim + e,
216
+ scale_inv, cos_sin, e);
217
+ }
218
+ }
219
+}
220
+
221
+// Concat + store a 576-wide latent: dst[e] = [a512 | b64], e in [0, 576). Used
222
+// for the decode query mqa_q = [ql_nope | q_pe] and the plain latent cache
223
+// entry [kv_c | k_pe]. FP8 packs to E4M3 (dst_elem_size 1); bf16 stores uint4.
224
+template <typename scalar_t, bool FP8, bool APPLY_ROPE = false>
225
+__device__ __forceinline__ void writeLatent576(void* dst, const scalar_t* a512,
226
+ const scalar_t* b64, int laneId,
227
+ int dst_elem_size,
228
+ float scale_inv,
229
+ const float* cos_sin = nullptr) {
230
+ auto* d = reinterpret_cast<uint8_t*>(dst);
231
+ for (int e = laneId * kVecElems; e < kCacheEntry; e += 32 * kVecElems) {
232
+ if (e < kKvLoraRank) {
233
+ copyChunk8<scalar_t, FP8>(d + e * dst_elem_size, a512 + e, scale_inv);
234
+ } else {
235
+ int const rope_e = e - kKvLoraRank;
236
+ copyChunk8<scalar_t, FP8, APPLY_ROPE>(d + e * dst_elem_size, b64 + rope_e,
237
+ scale_inv, cos_sin, rope_e);
238
+ }
239
+ }
240
+}
241
+
242
+// Write [kv_c | k_pe] into the fp8_ds_mla 656B entry using one warp: NoPE 512
243
+// as fp8 in 4 tiles of 128 (per-tile dynamic absmax scale, 4 fp32 scales at
244
+// [512,528)), RoPE 64 as bf16 at [528,656). Bit-compatible with
245
+// concat_and_cache_ds_mla_kernel.
246
+template <typename scalar_t, bool APPLY_ROPE = false>
247
+__device__ __forceinline__ void writeDsMlaCache(
248
+ uint8_t* row, const scalar_t* kvc, const scalar_t* pe, int laneId,
249
+ const float* cos_sin = nullptr) {
250
+ constexpr int kElemsPerLane = kKvLoraRank / 32; // 16
251
+ int const tile = laneId >> 3; // 8 lanes per tile
252
+ scalar_t vals[kElemsPerLane];
253
+ *reinterpret_cast<uint4*>(vals) =
254
+ *reinterpret_cast<const uint4*>(kvc + laneId * kElemsPerLane);
255
+ *reinterpret_cast<uint4*>(vals + 8) =
256
+ *reinterpret_cast<const uint4*>(kvc + laneId * kElemsPerLane + 8);
257
+
258
+ float max_abs = 0.0f;
259
+#pragma unroll
260
+ for (int i = 0; i < kElemsPerLane; i++) {
261
+ max_abs = fmaxf(max_abs, fabsf(static_cast<float>(vals[i])));
262
+ }
263
+#pragma unroll
264
+ for (int offset = 4; offset > 0; offset /= 2) {
265
+ max_abs = fmaxf(max_abs, VLLM_SHFL_XOR_SYNC_WIDTH(max_abs, offset, 8));
266
+ }
267
+ float const tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, FLT_MIN);
268
+ if ((laneId & 7) == 0) {
269
+ reinterpret_cast<float*>(row)[kKvLoraRank / 4 + tile] = tile_scale;
270
+ }
271
+ uint8_t res[kElemsPerLane];
272
+#pragma unroll
273
+ for (int i = 0; i < kElemsPerLane; i++) {
274
+ res[i] =
275
+ fp8::scaled_convert<uint8_t, scalar_t, Fp8KVCacheDataType::kFp8E4M3>(
276
+ vals[i], tile_scale);
277
+ }
278
+ *reinterpret_cast<uint4*>(row + laneId * kElemsPerLane) =
279
+ *reinterpret_cast<const uint4*>(res);
280
+ scalar_t* row16 = reinterpret_cast<scalar_t*>(row);
281
+ scalar_t* rope_dst = row16 + kKvLoraRank / 2 + 8 + laneId * 2;
282
+ if constexpr (APPLY_ROPE) {
283
+ using Converter = vllm::_typeConvert<scalar_t>;
284
+ using packed_t = typename Converter::packed_hip_type;
285
+ packed_t const src = *reinterpret_cast<packed_t const*>(pe + laneId * 2);
286
+ float2 const xy = Converter::convert(src);
287
+ float const cos = static_cast<float>(cos_sin[laneId]);
288
+ float const sin = static_cast<float>(cos_sin[laneId + kQkRopeHeadDim / 2]);
289
+ *reinterpret_cast<packed_t*>(rope_dst) = Converter::convert(
290
+ make_float2(xy.x * cos - xy.y * sin, xy.x * sin + xy.y * cos));
291
+ } else {
292
+ *reinterpret_cast<int32_t*>(rope_dst) =
293
+ *reinterpret_cast<const int32_t*>(pe + laneId * 2);
294
+ }
295
+}
296
+
297
+// ────────────────────────────────────────────────────────────────────────────
298
+// bf16 variant: optional q RoPE + full-key concat + latent cache insert
299
+// ────────────────────────────────────────────────────────────────────────────
300
+template <typename scalar_t, bool APPLY_ROPE>
301
+__global__ void fusedKimiK3MLAKeyConcatKVCacheInsertKernel(
302
+ scalar_t* __restrict__ q, int64_t const q_tok_stride,
303
+ int64_t const q_head_stride, const scalar_t* __restrict__ k_nope,
304
+ int64_t const kn_tok_stride, int64_t const kn_head_stride,
305
+ const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride,
306
+ const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride,
307
+ scalar_t* __restrict__ k_out, int64_t const ko_tok_stride,
308
+ int64_t const ko_head_stride, scalar_t* __restrict__ k_cache,
309
+ int64_t const cache_block_stride, int64_t const cache_token_stride,
310
+ const int64_t* __restrict__ slot_mapping,
311
+ const int64_t* __restrict__ position_ids,
312
+ const float* __restrict__ cos_sin_cache, int const num_tokens,
313
+ int const num_heads, int const cache_block_size) {
314
+ int const warpsPerBlock = blockDim.x / 32;
315
+ int const laneId = threadIdx.x % 32;
316
+ int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32;
317
+ int const slotsPerToken = num_heads + 1;
318
+ int const tokenIdx = globalWarpIdx / slotsPerToken;
319
+ int const slotIdx = globalWarpIdx % slotsPerToken;
320
+ if (tokenIdx >= num_tokens) return;
321
+
322
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
323
+ cudaGridDependencySynchronize();
324
+#endif
325
+
326
+ const float* rope_cache = nullptr;
327
+ if constexpr (APPLY_ROPE) {
328
+ rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim;
329
+ }
330
+
331
+ if (slotIdx < num_heads) {
332
+ scalar_t* qh = q + tokenIdx * q_tok_stride + slotIdx * q_head_stride;
333
+ writePrefillQuery<scalar_t, false, APPLY_ROPE>(
334
+ qh, qh, laneId, sizeof(scalar_t), 1.0f, rope_cache);
335
+ writeFullKey<scalar_t, false, APPLY_ROPE>(
336
+ k_out + tokenIdx * ko_tok_stride + slotIdx * ko_head_stride,
337
+ k_nope + tokenIdx * kn_tok_stride + slotIdx * kn_head_stride,
338
+ k_pe + tokenIdx * k_pe_tok_stride, laneId, sizeof(scalar_t), 1.0f,
339
+ rope_cache);
340
+ } else {
341
+ int64_t const slot_id = slot_mapping[tokenIdx];
342
+ if (slot_id >= 0) {
343
+ scalar_t* row = k_cache +
344
+ (slot_id / cache_block_size) * cache_block_stride +
345
+ (slot_id % cache_block_size) * cache_token_stride;
346
+ writeLatent576<scalar_t, false, APPLY_ROPE>(
347
+ row, kv_c + tokenIdx * kv_c_tok_stride,
348
+ k_pe + tokenIdx * k_pe_tok_stride, laneId, sizeof(scalar_t), 1.0f,
349
+ rope_cache);
350
+ }
351
+ }
352
+
353
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
354
+ cudaTriggerProgrammaticLaunchCompletion();
355
+#endif
356
+}
357
+
358
+// ────────────────────────────────────────────────────────────────────────────
359
+// fp8 variant: quant q / k / v + latent cache insert
360
+// ────────────────────────────────────────────────────────────────────────────
361
+template <typename scalar_t, bool APPLY_ROPE>
362
+__global__ void fusedKimiK3MLAQKVQuantKVCacheFp8Kernel(
363
+ const scalar_t* __restrict__ q, int64_t const q_tok_stride,
364
+ int64_t const q_head_stride, const scalar_t* __restrict__ k_nope,
365
+ int64_t const kn_tok_stride, int64_t const kn_head_stride,
366
+ const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride,
367
+ const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride,
368
+ const scalar_t* __restrict__ v, int64_t const v_tok_stride,
369
+ int64_t const v_head_stride, uint8_t* __restrict__ q_fp8,
370
+ int64_t const qo_tok_stride, int64_t const qo_head_stride,
371
+ uint8_t* __restrict__ k_fp8, int64_t const ko_tok_stride,
372
+ int64_t const ko_head_stride, uint8_t* __restrict__ v_fp8,
373
+ int64_t const vo_tok_stride, int64_t const vo_head_stride,
374
+ uint8_t* __restrict__ k_cache, int64_t const cache_block_stride,
375
+ int64_t const cache_token_stride, const int64_t* __restrict__ slot_mapping,
376
+ const float* __restrict__ q_scale_inv,
377
+ const float* __restrict__ k_scale_inv,
378
+ const float* __restrict__ v_scale_inv,
379
+ const float* __restrict__ cache_scale_inv, int const num_tokens,
380
+ int const num_heads, int const cache_block_size,
381
+ const int64_t* __restrict__ position_ids,
382
+ const float* __restrict__ cos_sin_cache) {
383
+ int const warpsPerBlock = blockDim.x / 32;
384
+ int const laneId = threadIdx.x % 32;
385
+ int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32;
386
+ int const slotsPerToken = num_heads + 1;
387
+ int const tokenIdx = globalWarpIdx / slotsPerToken;
388
+ int const slotIdx = globalWarpIdx % slotsPerToken;
389
+ if (tokenIdx >= num_tokens) return;
390
+
391
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
392
+ cudaGridDependencySynchronize();
393
+#endif
394
+
395
+ const float* rope_cache = nullptr;
396
+ if constexpr (APPLY_ROPE) {
397
+ rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim;
398
+ }
399
+
400
+ if (slotIdx < num_heads) {
401
+ int const h = slotIdx;
402
+ // q_fp8[t, h] = quant(q[t, h], q_scale)
403
+ float const qsi = __ldg(q_scale_inv);
404
+ const scalar_t* qh = q + tokenIdx * q_tok_stride + h * q_head_stride;
405
+ uint8_t* qo = q_fp8 + tokenIdx * qo_tok_stride + h * qo_head_stride;
406
+ writePrefillQuery<scalar_t, true, APPLY_ROPE>(qo, qh, laneId, 1, qsi,
407
+ rope_cache);
408
+ // k_fp8[t, h] = quant([k_nope | k_pe], k_scale)
409
+ writeFullKey<scalar_t, true, APPLY_ROPE>(
410
+ k_fp8 + tokenIdx * ko_tok_stride + h * ko_head_stride,
411
+ k_nope + tokenIdx * kn_tok_stride + h * kn_head_stride,
412
+ k_pe + tokenIdx * k_pe_tok_stride, laneId, 1, __ldg(k_scale_inv),
413
+ rope_cache);
414
+ // v_fp8[t, h] = quant(v[t, h], v_scale)
415
+ float const vsi = __ldg(v_scale_inv);
416
+ const scalar_t* vh = v + tokenIdx * v_tok_stride + h * v_head_stride;
417
+ uint8_t* vo = v_fp8 + tokenIdx * vo_tok_stride + h * vo_head_stride;
418
+ for (int e = laneId * kVecElems; e < kVHeadDim; e += 32 * kVecElems) {
419
+ copyChunk8<scalar_t, true>(vo + e, vh + e, vsi);
420
+ }
421
+ } else {
422
+ int64_t const slot_id = slot_mapping[tokenIdx];
423
+ if (slot_id >= 0) {
424
+ // The cache latent uses _k_scale (read back by decode / context); the
425
+ // attention key (k_fp8 above) uses its own k_scale.
426
+ float const ksi = __ldg(cache_scale_inv);
427
+ uint8_t* row = k_cache +
428
+ (slot_id / cache_block_size) * cache_block_stride +
429
+ (slot_id % cache_block_size) * cache_token_stride;
430
+ writeLatent576<scalar_t, true, APPLY_ROPE>(
431
+ row, kv_c + tokenIdx * kv_c_tok_stride,
432
+ k_pe + tokenIdx * k_pe_tok_stride, laneId, 1, ksi, rope_cache);
433
+ }
434
+ }
435
+
436
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
437
+ cudaTriggerProgrammaticLaunchCompletion();
438
+#endif
439
+}
440
+
441
+// ────────────────────────────────────────────────────────────────────────────
442
+// ds_mla variant: concat full key (bf16) + fp8_ds_mla latent cache insert
443
+//
444
+// Cache entry (656 bytes), matching concat_and_cache_ds_mla_kernel:
445
+// [0, 512) NoPE 512 vals as fp8, 4 tiles of 128, each dynamically scaled
446
+// [512, 528) 4 fp32 per-tile scales
447
+// [528, 656) RoPE 64 vals as bf16 (unquantized)
448
+// The cache slot uses one warp: lane L quantizes NoPE elems [L*16, L*16+16)
449
+// (tile = L>>3, absmax-reduced within its 8-lane tile group), then all 32 lanes
450
+// write 2 RoPE bf16 each.
451
+// ────────────────────────────────────────────────────────────────────────────
452
+template <typename scalar_t, bool APPLY_ROPE>
453
+__global__ void fusedKimiK3MLAKeyConcatDsMlaInsertKernel(
454
+ scalar_t* __restrict__ q, int64_t const q_tok_stride,
455
+ int64_t const q_head_stride, const scalar_t* __restrict__ k_nope,
456
+ int64_t const kn_tok_stride, int64_t const kn_head_stride,
457
+ const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride,
458
+ const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride,
459
+ scalar_t* __restrict__ k_out, int64_t const ko_tok_stride,
460
+ int64_t const ko_head_stride, uint8_t* __restrict__ k_cache,
461
+ int64_t const cache_block_stride, int64_t const cache_token_stride,
462
+ const int64_t* __restrict__ slot_mapping, int const num_tokens,
463
+ int const num_heads, int const cache_block_size,
464
+ const int64_t* __restrict__ position_ids,
465
+ const float* __restrict__ cos_sin_cache) {
466
+ int const warpsPerBlock = blockDim.x / 32;
467
+ int const laneId = threadIdx.x % 32;
468
+ int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32;
469
+ int const slotsPerToken = num_heads + 1;
470
+ int const tokenIdx = globalWarpIdx / slotsPerToken;
471
+ int const slotIdx = globalWarpIdx % slotsPerToken;
472
+ if (tokenIdx >= num_tokens) return;
473
+
474
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
475
+ cudaGridDependencySynchronize();
476
+#endif
477
+
478
+ const float* rope_cache = nullptr;
479
+ if constexpr (APPLY_ROPE) {
480
+ rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim;
481
+ }
482
+
483
+ if (slotIdx < num_heads) {
484
+ scalar_t* qh = q + tokenIdx * q_tok_stride + slotIdx * q_head_stride;
485
+ writePrefillQuery<scalar_t, false, APPLY_ROPE>(
486
+ qh, qh, laneId, sizeof(scalar_t), 1.0f, rope_cache);
487
+ // Full key (bf16): k_out[t, h] = [k_nope[t, h] | k_pe[t]].
488
+ writeFullKey<scalar_t, false, APPLY_ROPE>(
489
+ k_out + tokenIdx * ko_tok_stride + slotIdx * ko_head_stride,
490
+ k_nope + tokenIdx * kn_tok_stride + slotIdx * kn_head_stride,
491
+ k_pe + tokenIdx * k_pe_tok_stride, laneId, sizeof(scalar_t), 1.0f,
492
+ rope_cache);
493
+ } else {
494
+ int64_t const slot_id = slot_mapping[tokenIdx];
495
+ if (slot_id >= 0) {
496
+ uint8_t* row = k_cache +
497
+ (slot_id / cache_block_size) * cache_block_stride +
498
+ (slot_id % cache_block_size) * cache_token_stride;
499
+ writeDsMlaCache<scalar_t, APPLY_ROPE>(
500
+ row, kv_c + tokenIdx * kv_c_tok_stride,
501
+ k_pe + tokenIdx * k_pe_tok_stride, laneId, rope_cache);
502
+ }
503
+ }
504
+
505
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
506
+ cudaTriggerProgrammaticLaunchCompletion();
507
+#endif
508
+}
509
+
510
+// ────────────────────────────────────────────────────────────────────────────
511
+// Decode epilogue: concat mqa_q = [ql_nope | q_pe] (576) + latent cache insert,
512
+// run right before forward_mqa. Q_FP8 quantizes mqa_q; KV_FP8 quantizes the
513
+// plain per-tensor cache. (ds_mla cache uses the separate kernel below.)
514
+// ────────────────────────────────────────────────────────────────────────────
515
+template <typename scalar_t, bool Q_FP8, bool KV_FP8, bool APPLY_ROPE>
516
+__global__ void fusedKimiK3MLADecodeQConcatKVCacheKernel(
517
+ const scalar_t* __restrict__ ql_nope, int64_t const qn_tok_stride,
518
+ int64_t const qn_head_stride, const scalar_t* __restrict__ q_pe,
519
+ int64_t const qpe_tok_stride, int64_t const qpe_head_stride,
520
+ const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride,
521
+ const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride,
522
+ void* __restrict__ mqa_q, int64_t const mq_tok_stride,
523
+ int64_t const mq_head_stride, void* __restrict__ k_cache,
524
+ int64_t const cache_block_stride, int64_t const cache_token_stride,
525
+ const int64_t* __restrict__ slot_mapping,
526
+ const float* __restrict__ q_scale_inv,
527
+ const float* __restrict__ cache_scale_inv, int const num_tokens,
528
+ int const num_heads, int const cache_block_size,
529
+ const int64_t* __restrict__ position_ids,
530
+ const float* __restrict__ cos_sin_cache) {
531
+ constexpr int kMqElem = Q_FP8 ? 1 : sizeof(scalar_t);
532
+ constexpr int kCacheElem = KV_FP8 ? 1 : sizeof(scalar_t);
533
+ int const warpsPerBlock = blockDim.x / 32;
534
+ int const laneId = threadIdx.x % 32;
535
+ int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32;
536
+ int const slotsPerToken = num_heads + 1;
537
+ int const tokenIdx = globalWarpIdx / slotsPerToken;
538
+ int const slotIdx = globalWarpIdx % slotsPerToken;
539
+ if (tokenIdx >= num_tokens) return;
540
+
541
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
542
+ cudaGridDependencySynchronize();
543
+#endif
544
+
545
+ const float* rope_cache = nullptr;
546
+ if constexpr (APPLY_ROPE) {
547
+ rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim;
548
+ }
549
+
550
+ if (slotIdx < num_heads) {
551
+ float const qsi = Q_FP8 ? __ldg(q_scale_inv) : 1.0f;
552
+ writeLatent576<scalar_t, Q_FP8, APPLY_ROPE>(
553
+ reinterpret_cast<uint8_t*>(mqa_q) +
554
+ (tokenIdx * mq_tok_stride + slotIdx * mq_head_stride) * kMqElem,
555
+ ql_nope + tokenIdx * qn_tok_stride + slotIdx * qn_head_stride,
556
+ q_pe + tokenIdx * qpe_tok_stride + slotIdx * qpe_head_stride, laneId,
557
+ kMqElem, qsi, rope_cache);
558
+ } else {
559
+ int64_t const slot_id = slot_mapping[tokenIdx];
560
+ if (slot_id >= 0) {
561
+ float const ksi = KV_FP8 ? __ldg(cache_scale_inv) : 1.0f;
562
+ writeLatent576<scalar_t, KV_FP8, APPLY_ROPE>(
563
+ reinterpret_cast<uint8_t*>(k_cache) +
564
+ (slot_id / cache_block_size * cache_block_stride +
565
+ slot_id % cache_block_size * cache_token_stride) *
566
+ kCacheElem,
567
+ kv_c + tokenIdx * kv_c_tok_stride, k_pe + tokenIdx * k_pe_tok_stride,
568
+ laneId, kCacheElem, ksi, rope_cache);
569
+ }
570
+ }
571
+
572
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
573
+ cudaTriggerProgrammaticLaunchCompletion();
574
+#endif
575
+}
576
+
577
+// Decode epilogue for fp8_ds_mla: concat mqa_q (bf16) + ds_mla cache insert.
578
+template <typename scalar_t, bool APPLY_ROPE>
579
+__global__ void fusedKimiK3MLADecodeQConcatDsMlaKernel(
580
+ const scalar_t* __restrict__ ql_nope, int64_t const qn_tok_stride,
581
+ int64_t const qn_head_stride, const scalar_t* __restrict__ q_pe,
582
+ int64_t const qpe_tok_stride, int64_t const qpe_head_stride,
583
+ const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride,
584
+ const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride,
585
+ scalar_t* __restrict__ mqa_q, int64_t const mq_tok_stride,
586
+ int64_t const mq_head_stride, uint8_t* __restrict__ k_cache,
587
+ int64_t const cache_block_stride, int64_t const cache_token_stride,
588
+ const int64_t* __restrict__ slot_mapping, int const num_tokens,
589
+ int const num_heads, int const cache_block_size,
590
+ const int64_t* __restrict__ position_ids,
591
+ const float* __restrict__ cos_sin_cache) {
592
+ int const warpsPerBlock = blockDim.x / 32;
593
+ int const laneId = threadIdx.x % 32;
594
+ int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32;
595
+ int const slotsPerToken = num_heads + 1;
596
+ int const tokenIdx = globalWarpIdx / slotsPerToken;
597
+ int const slotIdx = globalWarpIdx % slotsPerToken;
598
+ if (tokenIdx >= num_tokens) return;
599
+
600
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
601
+ cudaGridDependencySynchronize();
602
+#endif
603
+
604
+ const float* rope_cache = nullptr;
605
+ if constexpr (APPLY_ROPE) {
606
+ rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim;
607
+ }
608
+
609
+ if (slotIdx < num_heads) {
610
+ writeLatent576<scalar_t, false, APPLY_ROPE>(
611
+ mqa_q + tokenIdx * mq_tok_stride + slotIdx * mq_head_stride,
612
+ ql_nope + tokenIdx * qn_tok_stride + slotIdx * qn_head_stride,
613
+ q_pe + tokenIdx * qpe_tok_stride + slotIdx * qpe_head_stride, laneId,
614
+ sizeof(scalar_t), 1.0f, rope_cache);
615
+ } else {
616
+ int64_t const slot_id = slot_mapping[tokenIdx];
617
+ if (slot_id >= 0) {
618
+ uint8_t* row = k_cache +
619
+ (slot_id / cache_block_size) * cache_block_stride +
620
+ (slot_id % cache_block_size) * cache_token_stride;
621
+ writeDsMlaCache<scalar_t, APPLY_ROPE>(
622
+ row, kv_c + tokenIdx * kv_c_tok_stride,
623
+ k_pe + tokenIdx * k_pe_tok_stride, laneId, rope_cache);
624
+ }
625
+ }
626
+
627
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
628
+ cudaTriggerProgrammaticLaunchCompletion();
629
+#endif
630
+}
631
+
632
+// PDL-aware launch of a (token, num_heads + 1)-warp grid.
633
+template <typename KernelT, typename... Args>
634
+static void launchPdl(KernelT kernel, int num_tokens, int num_heads,
635
+ cudaStream_t stream, Args... args) {
636
+ constexpr int kBlockSize = 256;
637
+ constexpr int kWarpsPerBlock = kBlockSize / 32;
638
+ int64_t const total_warps =
639
+ static_cast<int64_t>(num_tokens) * (num_heads + 1);
640
+ int const grid =
641
+ static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
642
+#ifndef USE_ROCM
643
+ static int const sm_version = getSMVersion();
644
+ cudaLaunchConfig_t config;
645
+ config.gridDim = dim3(grid);
646
+ config.blockDim = dim3(kBlockSize);
647
+ config.dynamicSmemBytes = 0;
648
+ config.stream = stream;
649
+ cudaLaunchAttribute attrs[1];
650
+ attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
651
+ attrs[0].val.programmaticStreamSerializationAllowed = 1;
652
+ config.attrs = attrs;
653
+ config.numAttrs = (sm_version >= 90) ? 1 : 0;
654
+ cudaLaunchKernelEx(&config, kernel, args...);
655
+#else
656
+ // clang-format off
657
+ // hipify's CUDA->HIP regex catastrophically backtracks on "> > >"; keep the
658
+ // launch closer as ">>>". clang-format would otherwise re-split it (it does
659
+ // not parse the CUDA launch syntax in this libtorch_stable file).
660
+ kernel<<<grid, kBlockSize, 0, stream>>>(args...);
661
+ // clang-format on
662
+#endif
663
+}
664
+
665
+} // namespace kimi_k3_fused_ops
666
+} // namespace vllm
667
+
668
+// ────────────────────────────────────────────────────────────────────────────
669
+// Torch op wrappers
670
+// ────────────────────────────────────────────────────────────────────────────
671
+namespace {
672
+bool check_rope_inputs(
673
+ std::optional<torch::stable::Tensor> const& position_ids,
674
+ std::optional<torch::stable::Tensor> const& cos_sin_cache,
675
+ torch::stable::Tensor const& /*input*/, int64_t num_tokens) {
676
+ using torch::headeronly::ScalarType;
677
+ STD_TORCH_CHECK(position_ids.has_value() == cos_sin_cache.has_value(),
678
+ "position_ids and cos_sin_cache must be provided together");
679
+ if (!position_ids.has_value()) return false;
680
+
681
+ auto const& positions = position_ids.value();
682
+ auto const& rope_cache = cos_sin_cache.value();
683
+ STD_TORCH_CHECK(positions.device().is_cuda() && positions.dim() == 1 &&
684
+ positions.scalar_type() == ScalarType::Long &&
685
+ positions.size(0) == num_tokens,
686
+ "position_ids must be int64 CUDA with shape [num_tokens]");
687
+ STD_TORCH_CHECK(rope_cache.device().is_cuda() && rope_cache.dim() == 2 &&
688
+ rope_cache.size(1) == 64 && rope_cache.stride(1) == 1 &&
689
+ rope_cache.scalar_type() == ScalarType::Float,
690
+ "cos_sin_cache must have shape [max_position, 64], unit "
691
+ "last-dim stride, and be fp32 (RoPE math runs in fp32)");
692
+ return true;
693
+}
694
+} // namespace
695
+
696
+void fused_kimi_k3_mla_key_concat_kv_cache_insert(
697
+ torch::stable::Tensor& q, // [Tp, H, 192]
698
+ torch::stable::Tensor const& k_nope, // [Tp, H, 128]
699
+ torch::stable::Tensor const& k_pe, // [Tp, 64]
700
+ torch::stable::Tensor const& kv_c_normed, // [Tp, 512]
701
+ torch::stable::Tensor& k_out, // [Tp, H, 192] bf16, written
702
+ torch::stable::Tensor& k_cache, // [nblk, bs, 576] bf16, written
703
+ torch::stable::Tensor const& slot_mapping, // [Tp] int64
704
+ int64_t cache_block_size, std::optional<torch::stable::Tensor> position_ids,
705
+ std::optional<torch::stable::Tensor> cos_sin_cache) {
706
+ using torch::headeronly::ScalarType;
707
+ namespace kk3 = vllm::kimi_k3_fused_ops;
708
+ STD_TORCH_CHECK(
709
+ k_nope.device().is_cuda() && k_nope.dim() == 3 && k_nope.size(2) == 128,
710
+ "k_nope shape [Tp, H, 128] CUDA");
711
+ STD_TORCH_CHECK(q.device().is_cuda() && q.dim() == 3 && q.size(2) == 192,
712
+ "q shape [Tp, H, 192] CUDA");
713
+ // k_pe is a strided view of the fused QKV-LoRA GEMM output; the kernel takes
714
+ // its row stride and reads the 64 cols contiguously, so unit last-dim stride
715
+ // (not full contiguity) is all that is required.
716
+ STD_TORCH_CHECK(k_pe.device().is_cuda() && k_pe.dim() == 2 &&
717
+ k_pe.stride(1) == 1 && k_pe.size(1) == 64,
718
+ "k_pe shape [Tp, 64], unit last-dim stride, CUDA");
719
+ STD_TORCH_CHECK(kv_c_normed.device().is_cuda() &&
720
+ kv_c_normed.is_contiguous() && kv_c_normed.dim() == 2 &&
721
+ kv_c_normed.size(1) == 512,
722
+ "kv_c_normed shape [Tp, 512] contiguous CUDA");
723
+ STD_TORCH_CHECK(k_out.device().is_cuda() && k_out.is_contiguous() &&
724
+ k_out.dim() == 3 && k_out.size(2) == 192,
725
+ "k_out shape [Tp, H, 192] contiguous CUDA");
726
+ STD_TORCH_CHECK(k_cache.device().is_cuda() && k_cache.dim() == 3 &&
727
+ k_cache.size(1) == cache_block_size &&
728
+ k_cache.size(2) == 576 && k_cache.stride(2) == 1,
729
+ "k_cache shape [nblk, block_size, 576] contiguous CUDA");
730
+ STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
731
+ slot_mapping.scalar_type() == ScalarType::Long,
732
+ "slot_mapping must be int64 CUDA");
733
+ ScalarType const dt = k_nope.scalar_type();
734
+ STD_TORCH_CHECK(q.scalar_type() == dt && k_pe.scalar_type() == dt &&
735
+ kv_c_normed.scalar_type() == dt &&
736
+ k_out.scalar_type() == dt && k_cache.scalar_type() == dt,
737
+ "all tensors must share k_nope's (bf16/fp16) dtype");
738
+
739
+ int const num_tokens = static_cast<int>(k_nope.size(0));
740
+ int const num_heads = static_cast<int>(k_nope.size(1));
741
+ STD_TORCH_CHECK(static_cast<int>(k_out.size(1)) == num_heads,
742
+ "k_out head count must match k_nope");
743
+ STD_TORCH_CHECK(q.size(0) == num_tokens && q.size(1) == num_heads,
744
+ "q token/head dimensions must match k_nope");
745
+ bool const apply_rope =
746
+ check_rope_inputs(position_ids, cos_sin_cache, q, num_tokens);
747
+ if (num_tokens == 0) return;
748
+
749
+ const torch::stable::accelerator::DeviceGuard device_guard(
750
+ k_nope.get_device_index());
751
+ const cudaStream_t stream =
752
+ get_current_cuda_stream(k_nope.get_device_index());
753
+
754
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
755
+ dt, "fused_kimi_k3_mla_key_concat_kv_cache_insert", [&] {
756
+ auto launch = [&](auto kernel) {
757
+ kk3::launchPdl(
758
+ kernel, num_tokens, num_heads, stream,
759
+ reinterpret_cast<scalar_t*>(q.mutable_data_ptr()), q.stride(0),
760
+ q.stride(1),
761
+ reinterpret_cast<scalar_t const*>(k_nope.const_data_ptr()),
762
+ k_nope.stride(0), k_nope.stride(1),
763
+ reinterpret_cast<scalar_t const*>(k_pe.const_data_ptr()),
764
+ k_pe.stride(0),
765
+ reinterpret_cast<scalar_t const*>(kv_c_normed.const_data_ptr()),
766
+ kv_c_normed.stride(0),
767
+ reinterpret_cast<scalar_t*>(k_out.mutable_data_ptr()),
768
+ k_out.stride(0), k_out.stride(1),
769
+ reinterpret_cast<scalar_t*>(k_cache.mutable_data_ptr()),
770
+ k_cache.stride(0), k_cache.stride(1),
771
+ slot_mapping.const_data_ptr<int64_t>(),
772
+ apply_rope ? position_ids.value().const_data_ptr<int64_t>()
773
+ : nullptr,
774
+ apply_rope ? reinterpret_cast<float const*>(
775
+ cos_sin_cache.value().const_data_ptr())
776
+ : nullptr,
777
+ num_tokens, num_heads, static_cast<int>(cache_block_size));
778
+ };
779
+ if (apply_rope) {
780
+ launch(
781
+ kk3::fusedKimiK3MLAKeyConcatKVCacheInsertKernel<scalar_t, true>);
782
+ } else {
783
+ launch(
784
+ kk3::fusedKimiK3MLAKeyConcatKVCacheInsertKernel<scalar_t, false>);
785
+ }
786
+ });
787
+}
788
+
789
+void fused_kimi_k3_mla_key_concat_ds_mla_insert(
790
+ torch::stable::Tensor& q, // [Tp, H, 192]
791
+ torch::stable::Tensor const& k_nope, // [Tp, H, 128] bf16
792
+ torch::stable::Tensor const& k_pe, // [Tp, 64] bf16
793
+ torch::stable::Tensor const& kv_c_normed, // [Tp, 512] bf16
794
+ torch::stable::Tensor& k_out, // [Tp, H, 192] bf16, written
795
+ torch::stable::Tensor& k_cache, // [nblk, bs, 656] uint8, written
796
+ torch::stable::Tensor const& slot_mapping, // [Tp] int64
797
+ int64_t cache_block_size, std::optional<torch::stable::Tensor> position_ids,
798
+ std::optional<torch::stable::Tensor> cos_sin_cache) {
799
+ using torch::headeronly::ScalarType;
800
+ namespace kk3 = vllm::kimi_k3_fused_ops;
801
+ ScalarType const dt = k_nope.scalar_type();
802
+ STD_TORCH_CHECK(
803
+ k_nope.device().is_cuda() && k_nope.dim() == 3 && k_nope.size(2) == 128,
804
+ "k_nope shape [Tp, H, 128] CUDA");
805
+ STD_TORCH_CHECK(q.device().is_cuda() && q.scalar_type() == dt &&
806
+ q.dim() == 3 && q.size(2) == 192,
807
+ "q shape [Tp, H, 192]");
808
+ STD_TORCH_CHECK(k_pe.device().is_cuda() && k_pe.dim() == 2 &&
809
+ k_pe.stride(1) == 1 && k_pe.scalar_type() == dt &&
810
+ k_pe.size(1) == 64,
811
+ "k_pe shape [Tp, 64], unit last-dim stride");
812
+ STD_TORCH_CHECK(kv_c_normed.device().is_cuda() &&
813
+ kv_c_normed.is_contiguous() &&
814
+ kv_c_normed.scalar_type() == dt &&
815
+ kv_c_normed.dim() == 2 && kv_c_normed.size(1) == 512,
816
+ "kv_c_normed shape [Tp, 512] contiguous");
817
+ STD_TORCH_CHECK(k_out.device().is_cuda() && k_out.is_contiguous() &&
818
+ k_out.scalar_type() == dt && k_out.dim() == 3 &&
819
+ k_out.size(2) == 192,
820
+ "k_out shape [Tp, H, 192] contiguous");
821
+ // fp8_ds_mla entry is 656 bytes stored as uint8.
822
+ STD_TORCH_CHECK(
823
+ k_cache.device().is_cuda() && k_cache.scalar_type() == ScalarType::Byte &&
824
+ k_cache.dim() == 3 && k_cache.size(1) == cache_block_size &&
825
+ k_cache.size(2) == 656 && k_cache.stride(2) == 1,
826
+ "k_cache shape [nblk, block_size, 656] uint8 contiguous");
827
+ STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
828
+ slot_mapping.scalar_type() == ScalarType::Long,
829
+ "slot_mapping must be int64 CUDA");
830
+
831
+ int const num_tokens = static_cast<int>(k_nope.size(0));
832
+ int const num_heads = static_cast<int>(k_nope.size(1));
833
+ STD_TORCH_CHECK(static_cast<int>(k_out.size(1)) == num_heads,
834
+ "k_out head count must match k_nope");
835
+ STD_TORCH_CHECK(q.size(0) == num_tokens && q.size(1) == num_heads,
836
+ "q token/head dimensions must match k_nope");
837
+ bool const apply_rope =
838
+ check_rope_inputs(position_ids, cos_sin_cache, q, num_tokens);
839
+ if (num_tokens == 0) return;
840
+
841
+ const torch::stable::accelerator::DeviceGuard device_guard(
842
+ k_nope.get_device_index());
843
+ const cudaStream_t stream =
844
+ get_current_cuda_stream(k_nope.get_device_index());
845
+
846
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
847
+ dt, "fused_kimi_k3_mla_key_concat_ds_mla_insert", [&] {
848
+ auto launch = [&](auto kernel) {
849
+ kk3::launchPdl(
850
+ kernel, num_tokens, num_heads, stream,
851
+ reinterpret_cast<scalar_t*>(q.mutable_data_ptr()), q.stride(0),
852
+ q.stride(1),
853
+ reinterpret_cast<scalar_t const*>(k_nope.const_data_ptr()),
854
+ k_nope.stride(0), k_nope.stride(1),
855
+ reinterpret_cast<scalar_t const*>(k_pe.const_data_ptr()),
856
+ k_pe.stride(0),
857
+ reinterpret_cast<scalar_t const*>(kv_c_normed.const_data_ptr()),
858
+ kv_c_normed.stride(0),
859
+ reinterpret_cast<scalar_t*>(k_out.mutable_data_ptr()),
860
+ k_out.stride(0), k_out.stride(1),
861
+ reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr()),
862
+ k_cache.stride(0), k_cache.stride(1),
863
+ slot_mapping.const_data_ptr<int64_t>(), num_tokens, num_heads,
864
+ static_cast<int>(cache_block_size),
865
+ apply_rope ? position_ids.value().const_data_ptr<int64_t>()
866
+ : nullptr,
867
+ apply_rope ? reinterpret_cast<float const*>(
868
+ cos_sin_cache.value().const_data_ptr())
869
+ : nullptr);
870
+ };
871
+ if (apply_rope) {
872
+ launch(kk3::fusedKimiK3MLAKeyConcatDsMlaInsertKernel<scalar_t, true>);
873
+ } else {
874
+ launch(
875
+ kk3::fusedKimiK3MLAKeyConcatDsMlaInsertKernel<scalar_t, false>);
876
+ }
877
+ });
878
+}
879
+
880
+void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(
881
+ torch::stable::Tensor const& q, // [Tp, H, 192] bf16
882
+ torch::stable::Tensor const& k_nope, // [Tp, H, 128] bf16
883
+ torch::stable::Tensor const& k_pe, // [Tp, 64] bf16
884
+ torch::stable::Tensor const& kv_c_normed, // [Tp, 512] bf16
885
+ torch::stable::Tensor const& v, // [Tp, H, 128] bf16
886
+ torch::stable::Tensor& q_fp8, // [Tp, H, 192] fp8, written
887
+ torch::stable::Tensor& k_fp8, // [Tp, H, 192] fp8, written
888
+ torch::stable::Tensor& v_fp8, // [Tp, H, 128] fp8, written
889
+ torch::stable::Tensor& k_cache, // [nblk, bs, 576] fp8, written
890
+ torch::stable::Tensor const& slot_mapping, // [Tp] int64
891
+ torch::stable::Tensor const& q_scale_inv, // scalar fp32 (1 / q scale)
892
+ torch::stable::Tensor const& k_scale_inv, // scalar fp32 (1 / k scale)
893
+ torch::stable::Tensor const& v_scale_inv, // scalar fp32 (1 / v scale)
894
+ torch::stable::Tensor const& cache_scale_inv, // scalar fp32 (1 / kv scale)
895
+ int64_t cache_block_size, std::optional<torch::stable::Tensor> position_ids,
896
+ std::optional<torch::stable::Tensor> cos_sin_cache) {
897
+ using torch::headeronly::ScalarType;
898
+ namespace kk3 = vllm::kimi_k3_fused_ops;
899
+ ScalarType const dt = k_nope.scalar_type();
900
+ auto check_in = [&](torch::stable::Tensor const& t, int d2, char const* n) {
901
+ STD_TORCH_CHECK(t.device().is_cuda() && t.scalar_type() == dt &&
902
+ t.dim() == 3 && t.size(2) == d2,
903
+ n);
904
+ };
905
+ check_in(q, 192, "q shape [Tp, H, 192]");
906
+ check_in(k_nope, 128, "k_nope shape [Tp, H, 128]");
907
+ check_in(v, 128, "v shape [Tp, H, 128]");
908
+ STD_TORCH_CHECK(k_pe.device().is_cuda() && k_pe.dim() == 2 &&
909
+ k_pe.stride(1) == 1 && k_pe.scalar_type() == dt &&
910
+ k_pe.size(1) == 64,
911
+ "k_pe shape [Tp, 64], unit last-dim stride");
912
+ STD_TORCH_CHECK(kv_c_normed.device().is_cuda() &&
913
+ kv_c_normed.is_contiguous() &&
914
+ kv_c_normed.scalar_type() == dt &&
915
+ kv_c_normed.dim() == 2 && kv_c_normed.size(1) == 512,
916
+ "kv_c_normed shape [Tp, 512] contiguous");
917
+ auto check_out = [&](torch::stable::Tensor const& t, int d2, char const* n) {
918
+ STD_TORCH_CHECK(t.device().is_cuda() && t.is_contiguous() &&
919
+ t.scalar_type() == ScalarType::Float8_e4m3fn &&
920
+ t.dim() == 3 && t.size(2) == d2,
921
+ n);
922
+ };
923
+ check_out(q_fp8, 192, "q_fp8 shape [Tp, H, 192] fp8 contiguous");
924
+ check_out(k_fp8, 192, "k_fp8 shape [Tp, H, 192] fp8 contiguous");
925
+ check_out(v_fp8, 128, "v_fp8 shape [Tp, H, 128] fp8 contiguous");
926
+ STD_TORCH_CHECK(k_cache.device().is_cuda() && k_cache.dim() == 3 &&
927
+ k_cache.size(1) == cache_block_size &&
928
+ k_cache.size(2) == 576 && k_cache.stride(2) == 1 &&
929
+ k_cache.scalar_type() == ScalarType::Float8_e4m3fn,
930
+ "k_cache shape [nblk, block_size, 576] fp8 contiguous");
931
+ STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
932
+ slot_mapping.scalar_type() == ScalarType::Long,
933
+ "slot_mapping must be int64 CUDA");
934
+ auto check_scale = [&](torch::stable::Tensor const& s, char const* n) {
935
+ STD_TORCH_CHECK(s.device().is_cuda() &&
936
+ s.scalar_type() == ScalarType::Float && s.size(0) == 1,
937
+ n);
938
+ };
939
+ check_scale(q_scale_inv, "q_scale_inv must be scalar float32 CUDA");
940
+ check_scale(k_scale_inv, "k_scale_inv must be scalar float32 CUDA");
941
+ check_scale(v_scale_inv, "v_scale_inv must be scalar float32 CUDA");
942
+ check_scale(cache_scale_inv, "cache_scale_inv must be scalar float32 CUDA");
943
+
944
+ int const num_tokens = static_cast<int>(k_nope.size(0));
945
+ int const num_heads = static_cast<int>(k_nope.size(1));
946
+ bool const apply_rope =
947
+ check_rope_inputs(position_ids, cos_sin_cache, q, num_tokens);
948
+ if (num_tokens == 0) return;
949
+
950
+ const torch::stable::accelerator::DeviceGuard device_guard(
951
+ k_nope.get_device_index());
952
+ const cudaStream_t stream =
953
+ get_current_cuda_stream(k_nope.get_device_index());
954
+
955
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
956
+ dt, "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert", [&] {
957
+ auto launch = [&](auto kernel) {
958
+ kk3::launchPdl(
959
+ kernel, num_tokens, num_heads, stream,
960
+ reinterpret_cast<scalar_t const*>(q.const_data_ptr()),
961
+ q.stride(0), q.stride(1),
962
+ reinterpret_cast<scalar_t const*>(k_nope.const_data_ptr()),
963
+ k_nope.stride(0), k_nope.stride(1),
964
+ reinterpret_cast<scalar_t const*>(k_pe.const_data_ptr()),
965
+ k_pe.stride(0),
966
+ reinterpret_cast<scalar_t const*>(kv_c_normed.const_data_ptr()),
967
+ kv_c_normed.stride(0),
968
+ reinterpret_cast<scalar_t const*>(v.const_data_ptr()),
969
+ v.stride(0), v.stride(1),
970
+ reinterpret_cast<uint8_t*>(q_fp8.mutable_data_ptr()),
971
+ q_fp8.stride(0), q_fp8.stride(1),
972
+ reinterpret_cast<uint8_t*>(k_fp8.mutable_data_ptr()),
973
+ k_fp8.stride(0), k_fp8.stride(1),
974
+ reinterpret_cast<uint8_t*>(v_fp8.mutable_data_ptr()),
975
+ v_fp8.stride(0), v_fp8.stride(1),
976
+ reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr()),
977
+ k_cache.stride(0), k_cache.stride(1),
978
+ slot_mapping.const_data_ptr<int64_t>(),
979
+ q_scale_inv.const_data_ptr<float>(),
980
+ k_scale_inv.const_data_ptr<float>(),
981
+ v_scale_inv.const_data_ptr<float>(),
982
+ cache_scale_inv.const_data_ptr<float>(), num_tokens, num_heads,
983
+ static_cast<int>(cache_block_size),
984
+ apply_rope ? position_ids.value().const_data_ptr<int64_t>()
985
+ : nullptr,
986
+ apply_rope ? reinterpret_cast<float const*>(
987
+ cos_sin_cache.value().const_data_ptr())
988
+ : nullptr);
989
+ };
990
+ if (apply_rope) {
991
+ launch(kk3::fusedKimiK3MLAQKVQuantKVCacheFp8Kernel<scalar_t, true>);
992
+ } else {
993
+ launch(kk3::fusedKimiK3MLAQKVQuantKVCacheFp8Kernel<scalar_t, false>);
994
+ }
995
+ });
996
+}
997
+
998
+// ────────────────────────────────────────────────────────────────────────────
999
+// Decode epilogue torch ops
1000
+// ────────────────────────────────────────────────────────────────────────────
1001
+namespace {
1002
+// Shared shape checks for the decode ops. Verifies the query/latent inputs and
1003
+// slot_mapping; caller checks mqa_q / k_cache dtypes for its variant.
1004
+void check_decode_inputs(torch::stable::Tensor const& ql_nope,
1005
+ torch::stable::Tensor const& q_pe,
1006
+ torch::stable::Tensor const& kv_c_normed,
1007
+ torch::stable::Tensor const& k_pe,
1008
+ torch::stable::Tensor const& mqa_q,
1009
+ torch::stable::Tensor const& slot_mapping) {
1010
+ using torch::headeronly::ScalarType;
1011
+ auto const dt = ql_nope.scalar_type();
1012
+ STD_TORCH_CHECK(ql_nope.device().is_cuda() && ql_nope.dim() == 3 &&
1013
+ ql_nope.size(2) == 512,
1014
+ "ql_nope shape [B, H, 512] CUDA");
1015
+ STD_TORCH_CHECK(q_pe.device().is_cuda() && q_pe.scalar_type() == dt &&
1016
+ q_pe.dim() == 3 && q_pe.size(2) == 64,
1017
+ "q_pe shape [B, H, 64]");
1018
+ STD_TORCH_CHECK(kv_c_normed.device().is_cuda() &&
1019
+ kv_c_normed.is_contiguous() &&
1020
+ kv_c_normed.scalar_type() == dt &&
1021
+ kv_c_normed.dim() == 2 && kv_c_normed.size(1) == 512,
1022
+ "kv_c_normed shape [B, 512] contiguous");
1023
+ STD_TORCH_CHECK(k_pe.device().is_cuda() && k_pe.dim() == 2 &&
1024
+ k_pe.stride(1) == 1 && k_pe.scalar_type() == dt &&
1025
+ k_pe.size(1) == 64,
1026
+ "k_pe shape [B, 64], unit last-dim stride");
1027
+ STD_TORCH_CHECK(mqa_q.device().is_cuda() && mqa_q.is_contiguous() &&
1028
+ mqa_q.dim() == 3 && mqa_q.size(2) == 576,
1029
+ "mqa_q shape [B, H, 576] contiguous");
1030
+ STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
1031
+ slot_mapping.scalar_type() == ScalarType::Long,
1032
+ "slot_mapping must be int64 CUDA");
1033
+}
1034
+} // namespace
1035
+
1036
+void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(
1037
+ torch::stable::Tensor const& ql_nope, // [B, H, 512] bf16
1038
+ torch::stable::Tensor const& q_pe, // [B, H, 64] bf16
1039
+ torch::stable::Tensor const& kv_c_normed, // [B, 512] bf16
1040
+ torch::stable::Tensor const& k_pe, // [B, 64] bf16
1041
+ torch::stable::Tensor& mqa_q, // [B, H, 576] bf16, written
1042
+ torch::stable::Tensor& k_cache, // [nblk, bs, 576] bf16, written
1043
+ torch::stable::Tensor const& slot_mapping, // [B] int64
1044
+ int64_t cache_block_size, std::optional<torch::stable::Tensor> position_ids,
1045
+ std::optional<torch::stable::Tensor> cos_sin_cache) {
1046
+ using torch::headeronly::ScalarType;
1047
+ namespace kk3 = vllm::kimi_k3_fused_ops;
1048
+ ScalarType const dt = ql_nope.scalar_type();
1049
+ check_decode_inputs(ql_nope, q_pe, kv_c_normed, k_pe, mqa_q, slot_mapping);
1050
+ STD_TORCH_CHECK(mqa_q.scalar_type() == dt && k_cache.scalar_type() == dt,
1051
+ "mqa_q / k_cache must match ql_nope dtype (bf16)");
1052
+ STD_TORCH_CHECK(k_cache.device().is_cuda() && k_cache.dim() == 3 &&
1053
+ k_cache.size(1) == cache_block_size &&
1054
+ k_cache.size(2) == 576 && k_cache.stride(2) == 1,
1055
+ "k_cache shape [nblk, block_size, 576] contiguous");
1056
+
1057
+ int const num_tokens = static_cast<int>(ql_nope.size(0));
1058
+ int const num_heads = static_cast<int>(ql_nope.size(1));
1059
+ bool const apply_rope =
1060
+ check_rope_inputs(position_ids, cos_sin_cache, q_pe, num_tokens);
1061
+ if (num_tokens == 0) return;
1062
+ const torch::stable::accelerator::DeviceGuard device_guard(
1063
+ ql_nope.get_device_index());
1064
+ const cudaStream_t stream =
1065
+ get_current_cuda_stream(ql_nope.get_device_index());
1066
+
1067
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
1068
+ dt, "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", [&] {
1069
+ auto launch = [&](auto kernel) {
1070
+ kk3::launchPdl(
1071
+ kernel, num_tokens, num_heads, stream,
1072
+ reinterpret_cast<scalar_t const*>(ql_nope.const_data_ptr()),
1073
+ ql_nope.stride(0), ql_nope.stride(1),
1074
+ reinterpret_cast<scalar_t const*>(q_pe.const_data_ptr()),
1075
+ q_pe.stride(0), q_pe.stride(1),
1076
+ reinterpret_cast<scalar_t const*>(kv_c_normed.const_data_ptr()),
1077
+ kv_c_normed.stride(0),
1078
+ reinterpret_cast<scalar_t const*>(k_pe.const_data_ptr()),
1079
+ k_pe.stride(0), mqa_q.mutable_data_ptr(), mqa_q.stride(0),
1080
+ mqa_q.stride(1), k_cache.mutable_data_ptr(), k_cache.stride(0),
1081
+ k_cache.stride(1), slot_mapping.const_data_ptr<int64_t>(),
1082
+ nullptr, nullptr, num_tokens, num_heads,
1083
+ static_cast<int>(cache_block_size),
1084
+ apply_rope ? position_ids.value().const_data_ptr<int64_t>()
1085
+ : nullptr,
1086
+ apply_rope ? reinterpret_cast<float const*>(
1087
+ cos_sin_cache.value().const_data_ptr())
1088
+ : nullptr);
1089
+ };
1090
+ if (apply_rope) {
1091
+ launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, false,
1092
+ false, true>);
1093
+ } else {
1094
+ launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, false,
1095
+ false, false>);
1096
+ }
1097
+ });
1098
+}
1099
+
1100
+void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(
1101
+ torch::stable::Tensor const& ql_nope, // [B, H, 512] bf16
1102
+ torch::stable::Tensor const& q_pe, // [B, H, 64] bf16
1103
+ torch::stable::Tensor const& kv_c_normed, // [B, 512] bf16
1104
+ torch::stable::Tensor const& k_pe, // [B, 64] bf16
1105
+ torch::stable::Tensor& mqa_q, // [B, H, 576] fp8, written
1106
+ torch::stable::Tensor& k_cache, // [nblk, bs, 576] fp8, written
1107
+ torch::stable::Tensor const& slot_mapping, // [B] int64
1108
+ torch::stable::Tensor const& q_scale_inv, // scalar fp32 (1 / q scale)
1109
+ torch::stable::Tensor const& cache_scale_inv, // scalar fp32 (1 / kv scale)
1110
+ int64_t cache_block_size, std::optional<torch::stable::Tensor> position_ids,
1111
+ std::optional<torch::stable::Tensor> cos_sin_cache) {
1112
+ using torch::headeronly::ScalarType;
1113
+ namespace kk3 = vllm::kimi_k3_fused_ops;
1114
+ ScalarType const dt = ql_nope.scalar_type();
1115
+ check_decode_inputs(ql_nope, q_pe, kv_c_normed, k_pe, mqa_q, slot_mapping);
1116
+ STD_TORCH_CHECK(mqa_q.scalar_type() == ScalarType::Float8_e4m3fn,
1117
+ "mqa_q must be float8_e4m3fn");
1118
+ STD_TORCH_CHECK(k_cache.device().is_cuda() && k_cache.dim() == 3 &&
1119
+ k_cache.size(1) == cache_block_size &&
1120
+ k_cache.size(2) == 576 && k_cache.stride(2) == 1 &&
1121
+ k_cache.scalar_type() == ScalarType::Float8_e4m3fn,
1122
+ "k_cache shape [nblk, block_size, 576] fp8 contiguous");
1123
+ auto check_scale = [&](torch::stable::Tensor const& s, char const* n) {
1124
+ STD_TORCH_CHECK(s.device().is_cuda() &&
1125
+ s.scalar_type() == ScalarType::Float && s.size(0) == 1,
1126
+ n);
1127
+ };
1128
+ check_scale(q_scale_inv, "q_scale_inv must be scalar float32 CUDA");
1129
+ check_scale(cache_scale_inv, "cache_scale_inv must be scalar float32 CUDA");
1130
+
1131
+ int const num_tokens = static_cast<int>(ql_nope.size(0));
1132
+ int const num_heads = static_cast<int>(ql_nope.size(1));
1133
+ bool const apply_rope =
1134
+ check_rope_inputs(position_ids, cos_sin_cache, q_pe, num_tokens);
1135
+ if (num_tokens == 0) return;
1136
+ const torch::stable::accelerator::DeviceGuard device_guard(
1137
+ ql_nope.get_device_index());
1138
+ const cudaStream_t stream =
1139
+ get_current_cuda_stream(ql_nope.get_device_index());
1140
+
1141
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
1142
+ dt, "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert", [&] {
1143
+ auto launch = [&](auto kernel) {
1144
+ kk3::launchPdl(
1145
+ kernel, num_tokens, num_heads, stream,
1146
+ reinterpret_cast<scalar_t const*>(ql_nope.const_data_ptr()),
1147
+ ql_nope.stride(0), ql_nope.stride(1),
1148
+ reinterpret_cast<scalar_t const*>(q_pe.const_data_ptr()),
1149
+ q_pe.stride(0), q_pe.stride(1),
1150
+ reinterpret_cast<scalar_t const*>(kv_c_normed.const_data_ptr()),
1151
+ kv_c_normed.stride(0),
1152
+ reinterpret_cast<scalar_t const*>(k_pe.const_data_ptr()),
1153
+ k_pe.stride(0), mqa_q.mutable_data_ptr(), mqa_q.stride(0),
1154
+ mqa_q.stride(1), k_cache.mutable_data_ptr(), k_cache.stride(0),
1155
+ k_cache.stride(1), slot_mapping.const_data_ptr<int64_t>(),
1156
+ q_scale_inv.const_data_ptr<float>(),
1157
+ cache_scale_inv.const_data_ptr<float>(), num_tokens, num_heads,
1158
+ static_cast<int>(cache_block_size),
1159
+ apply_rope ? position_ids.value().const_data_ptr<int64_t>()
1160
+ : nullptr,
1161
+ apply_rope ? reinterpret_cast<float const*>(
1162
+ cos_sin_cache.value().const_data_ptr())
1163
+ : nullptr);
1164
+ };
1165
+ if (apply_rope) {
1166
+ launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, true,
1167
+ true, true>);
1168
+ } else {
1169
+ launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, true,
1170
+ true, false>);
1171
+ }
1172
+ });
1173
+}
1174
+
1175
+void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(
1176
+ torch::stable::Tensor const& ql_nope, // [B, H, 512] bf16
1177
+ torch::stable::Tensor const& q_pe, // [B, H, 64] bf16
1178
+ torch::stable::Tensor const& kv_c_normed, // [B, 512] bf16
1179
+ torch::stable::Tensor const& k_pe, // [B, 64] bf16
1180
+ torch::stable::Tensor& mqa_q, // [B, H, 576] bf16, written
1181
+ torch::stable::Tensor& k_cache, // [nblk, bs, 656] uint8, written
1182
+ torch::stable::Tensor const& slot_mapping, // [B] int64
1183
+ int64_t cache_block_size, std::optional<torch::stable::Tensor> position_ids,
1184
+ std::optional<torch::stable::Tensor> cos_sin_cache) {
1185
+ using torch::headeronly::ScalarType;
1186
+ namespace kk3 = vllm::kimi_k3_fused_ops;
1187
+ ScalarType const dt = ql_nope.scalar_type();
1188
+ check_decode_inputs(ql_nope, q_pe, kv_c_normed, k_pe, mqa_q, slot_mapping);
1189
+ STD_TORCH_CHECK(mqa_q.scalar_type() == dt, "mqa_q must be bf16 for ds_mla");
1190
+ STD_TORCH_CHECK(
1191
+ k_cache.device().is_cuda() && k_cache.scalar_type() == ScalarType::Byte &&
1192
+ k_cache.dim() == 3 && k_cache.size(1) == cache_block_size &&
1193
+ k_cache.size(2) == 656 && k_cache.stride(2) == 1,
1194
+ "k_cache shape [nblk, block_size, 656] uint8 contiguous");
1195
+
1196
+ int const num_tokens = static_cast<int>(ql_nope.size(0));
1197
+ int const num_heads = static_cast<int>(ql_nope.size(1));
1198
+ bool const apply_rope =
1199
+ check_rope_inputs(position_ids, cos_sin_cache, q_pe, num_tokens);
1200
+ if (num_tokens == 0) return;
1201
+ const torch::stable::accelerator::DeviceGuard device_guard(
1202
+ ql_nope.get_device_index());
1203
+ const cudaStream_t stream =
1204
+ get_current_cuda_stream(ql_nope.get_device_index());
1205
+
1206
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
1207
+ dt, "fused_kimi_k3_mla_decode_q_concat_ds_mla_insert", [&] {
1208
+ auto launch = [&](auto kernel) {
1209
+ kk3::launchPdl(
1210
+ kernel, num_tokens, num_heads, stream,
1211
+ reinterpret_cast<scalar_t const*>(ql_nope.const_data_ptr()),
1212
+ ql_nope.stride(0), ql_nope.stride(1),
1213
+ reinterpret_cast<scalar_t const*>(q_pe.const_data_ptr()),
1214
+ q_pe.stride(0), q_pe.stride(1),
1215
+ reinterpret_cast<scalar_t const*>(kv_c_normed.const_data_ptr()),
1216
+ kv_c_normed.stride(0),
1217
+ reinterpret_cast<scalar_t const*>(k_pe.const_data_ptr()),
1218
+ k_pe.stride(0),
1219
+ reinterpret_cast<scalar_t*>(mqa_q.mutable_data_ptr()),
1220
+ mqa_q.stride(0), mqa_q.stride(1),
1221
+ reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr()),
1222
+ k_cache.stride(0), k_cache.stride(1),
1223
+ slot_mapping.const_data_ptr<int64_t>(), num_tokens, num_heads,
1224
+ static_cast<int>(cache_block_size),
1225
+ apply_rope ? position_ids.value().const_data_ptr<int64_t>()
1226
+ : nullptr,
1227
+ apply_rope ? reinterpret_cast<float const*>(
1228
+ cos_sin_cache.value().const_data_ptr())
1229
+ : nullptr);
1230
+ };
1231
+ if (apply_rope) {
1232
+ launch(kk3::fusedKimiK3MLADecodeQConcatDsMlaKernel<scalar_t, true>);
1233
+ } else {
1234
+ launch(kk3::fusedKimiK3MLADecodeQConcatDsMlaKernel<scalar_t, false>);
1235
+ }
1236
+ });
1237
+}
csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu
new
+1130
@@ -0,0 +1,1130 @@
1
+/*
2
+ * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
3
+ */
4
+
5
+#include <cstdint>
6
+#include <cuda_bf16.h>
7
+#include <cuda_runtime.h>
8
+#include <array>
9
+#include <optional>
10
+
11
+#include "../torch_utils.h"
12
+#include "../../cuda_compat.h"
13
+
14
+namespace {
15
+
16
+constexpr int kDimK = 128;
17
+constexpr int kDimV = 128;
18
+constexpr int kKernelWidth = 4;
19
+constexpr int kConvStateWidth = kKernelWidth - 1;
20
+constexpr int kThreads = 256;
21
+constexpr int kWarps = kThreads / 32;
22
+constexpr int kChunkV = 32;
23
+constexpr int kNumChunks = kDimV / kChunkV;
24
+constexpr int kRowsPerWarp = kChunkV / kWarps;
25
+
26
+struct KdaDecodeStrides {
27
+ int64_t x_row;
28
+ int64_t beta_row;
29
+ int64_t onorm_row;
30
+ int64_t conv_slot;
31
+ int64_t state_slot;
32
+};
33
+
34
+__device__ __forceinline__ float bf16_load(const __nv_bfloat16* ptr,
35
+ int64_t idx) {
36
+ return __bfloat162float(ptr[idx]);
37
+}
38
+
39
+__device__ __forceinline__ float bf16_load(const float* ptr, int64_t idx) {
40
+ return ptr[idx];
41
+}
42
+
43
+template <int kChannels>
44
+__device__ __forceinline__ float conv_weight_load(const float* ptr, int channel,
45
+ int width) {
46
+ return ptr[width * kChannels + channel];
47
+}
48
+
49
+__device__ __forceinline__ __nv_bfloat16 bf16_store(float value) {
50
+ return __float2bfloat16(value);
51
+}
52
+
53
+template <bool kUseCacheGlobalStore>
54
+__device__ __forceinline__ void store_state_float4(float* ptr, float4 value) {
55
+ if constexpr (kUseCacheGlobalStore) {
56
+ __stcg(reinterpret_cast<float4*>(ptr), value);
57
+ } else {
58
+ *reinterpret_cast<float4*>(ptr) = value;
59
+ }
60
+}
61
+
62
+__device__ __forceinline__ float sigmoid_fast(float x) {
63
+ return 1.0f / (1.0f + __expf(-x));
64
+}
65
+
66
+__device__ __forceinline__ float silu_fast(float x) {
67
+ return x * sigmoid_fast(x);
68
+}
69
+
70
+__device__ __forceinline__ float softplus_fast(float x) {
71
+ return x > 20.0f ? x : log1pf(__expf(x));
72
+}
73
+
74
+__device__ __forceinline__ float warp_reduce_sum(float value) {
75
+#pragma unroll
76
+ for (int offset = 16; offset > 0; offset >>= 1) {
77
+ value += __shfl_xor_sync(0xffffffffu, value, offset);
78
+ }
79
+ return value;
80
+}
81
+
82
+__device__ __forceinline__ void cp_async_cg_16b(float* smem_ptr,
83
+ const float* gmem_ptr) {
84
+ uint32_t smem_addr =
85
+ static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
86
+ asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n"
87
+ :
88
+ : "r"(smem_addr), "l"(gmem_ptr));
89
+}
90
+
91
+__device__ __forceinline__ void cp_async_commit() {
92
+ asm volatile("cp.async.commit_group;\n" ::);
93
+}
94
+
95
+__device__ __forceinline__ void cp_async_wait_all() {
96
+ asm volatile("cp.async.wait_all;\n" ::: "memory");
97
+}
98
+
99
+__device__ __forceinline__ void cp_async_wait_group_1() {
100
+ asm volatile("cp.async.wait_group 1;\n" ::: "memory");
101
+}
102
+
103
+template <int kCopyThreads>
104
+__device__ __forceinline__ void cp_async_state_chunk_for(float* s_state,
105
+ const float* state,
106
+ int slot, int i_hv,
107
+ int HV, int chunk) {
108
+ constexpr int kFloat4PerChunk = kChunkV * kDimK / 4;
109
+ const int tid = threadIdx.x;
110
+ const int stage = chunk & 1;
111
+ const int v_base = chunk * kChunkV;
112
+ for (int linear4 = tid; linear4 < kFloat4PerChunk; linear4 += kCopyThreads) {
113
+ const int elem = linear4 * 4;
114
+ const int row = elem / kDimK;
115
+ const int k = elem - row * kDimK;
116
+ float* dst = s_state + (stage * kChunkV + row) * kDimK + k;
117
+ const float* src =
118
+ state + ((slot * HV + i_hv) * kDimV + v_base + row) * kDimK + k;
119
+ cp_async_cg_16b(dst, src);
120
+ }
121
+ cp_async_commit();
122
+}
123
+
124
+__device__ __forceinline__ void cp_async_state_chunk(float* s_state,
125
+ const float* state,
126
+ int slot, int i_hv, int HV,
127
+ int chunk) {
128
+ cp_async_state_chunk_for<kThreads>(s_state, state, slot, i_hv, HV, chunk);
129
+}
130
+
131
+__device__ __forceinline__ float block_reduce_sum(float value, float* scratch) {
132
+ const int tid = threadIdx.x;
133
+ const int lane = tid & 31;
134
+ const int warp = tid >> 5;
135
+
136
+ float warp_total = warp_reduce_sum(value);
137
+ if (lane == 0) {
138
+ scratch[warp] = warp_total;
139
+ }
140
+ __syncthreads();
141
+
142
+ float block_total = 0.0f;
143
+ if (warp == 0) {
144
+ block_total = lane < kWarps ? scratch[lane] : 0.0f;
145
+ block_total = warp_reduce_sum(block_total);
146
+ if (lane == 0) {
147
+ scratch[0] = block_total;
148
+ }
149
+ }
150
+ __syncthreads();
151
+ return scratch[0];
152
+}
153
+
154
+struct Sum2 {
155
+ float x;
156
+ float y;
157
+};
158
+
159
+__device__ __forceinline__ Sum2 warp_reduce_sum_pair(float x, float y) {
160
+#pragma unroll
161
+ for (int offset = 16; offset > 0; offset >>= 1) {
162
+ x += __shfl_xor_sync(0xffffffffu, x, offset);
163
+ y += __shfl_xor_sync(0xffffffffu, y, offset);
164
+ }
165
+ return {x, y};
166
+}
167
+
168
+template <int kReduceWarps>
169
+__device__ __forceinline__ Sum2 block_reduce_sum2_for(float x, float y,
170
+ float* scratch) {
171
+ const int lane = threadIdx.x & 31;
172
+ const int warp = threadIdx.x >> 5;
173
+
174
+ const float warp_x = warp_reduce_sum(x);
175
+ const float warp_y = warp_reduce_sum(y);
176
+ if (lane == 0) {
177
+ scratch[warp] = warp_x;
178
+ scratch[kReduceWarps + warp] = warp_y;
179
+ }
180
+ __syncthreads();
181
+
182
+ float block_x = 0.0f;
183
+ float block_y = 0.0f;
184
+ if (warp == 0) {
185
+ block_x = lane < kReduceWarps ? scratch[lane] : 0.0f;
186
+ block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f;
187
+ block_x = warp_reduce_sum(block_x);
188
+ block_y = warp_reduce_sum(block_y);
189
+ if (lane == 0) {
190
+ scratch[0] = block_x;
191
+ scratch[1] = block_y;
192
+ }
193
+ }
194
+ __syncthreads();
195
+ return {scratch[0], scratch[1]};
196
+}
197
+
198
+__device__ __forceinline__ Sum2 block_reduce_sum2(float x, float y,
199
+ float* scratch) {
200
+ return block_reduce_sum2_for<kWarps>(x, y, scratch);
201
+}
202
+
203
+template <int kReduceWarps>
204
+__device__ __forceinline__ float block_reduce_sum_active_for(float value,
205
+ float* scratch) {
206
+ const int lane = threadIdx.x & 31;
207
+ const int warp = threadIdx.x >> 5;
208
+
209
+ float warp_total = 0.0f;
210
+ if (warp < kReduceWarps) {
211
+ warp_total = warp_reduce_sum(value);
212
+ }
213
+ if (lane == 0 && warp < kReduceWarps) {
214
+ scratch[warp] = warp_total;
215
+ }
216
+ __syncthreads();
217
+
218
+ float block_total = 0.0f;
219
+ if (warp == 0) {
220
+ block_total = lane < kReduceWarps ? scratch[lane] : 0.0f;
221
+ block_total = warp_reduce_sum(block_total);
222
+ if (lane == 0) {
223
+ scratch[0] = block_total;
224
+ }
225
+ }
226
+ __syncthreads();
227
+ return scratch[0];
228
+}
229
+
230
+template <int kReduceWarps>
231
+__device__ __forceinline__ Sum2 block_reduce_sum2_active_for(float x, float y,
232
+ float* scratch) {
233
+ const int lane = threadIdx.x & 31;
234
+ const int warp = threadIdx.x >> 5;
235
+
236
+ float warp_x = 0.0f;
237
+ float warp_y = 0.0f;
238
+ if (warp < kReduceWarps) {
239
+ warp_x = warp_reduce_sum(x);
240
+ warp_y = warp_reduce_sum(y);
241
+ }
242
+ if (lane == 0 && warp < kReduceWarps) {
243
+ scratch[warp] = warp_x;
244
+ scratch[kReduceWarps + warp] = warp_y;
245
+ }
246
+ __syncthreads();
247
+
248
+ float block_x = 0.0f;
249
+ float block_y = 0.0f;
250
+ if (warp == 0) {
251
+ block_x = lane < kReduceWarps ? scratch[lane] : 0.0f;
252
+ block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f;
253
+ block_x = warp_reduce_sum(block_x);
254
+ block_y = warp_reduce_sum(block_y);
255
+ if (lane == 0) {
256
+ scratch[0] = block_x;
257
+ scratch[1] = block_y;
258
+ }
259
+ }
260
+ __syncthreads();
261
+ return {scratch[0], scratch[1]};
262
+}
263
+
264
+template <bool kApplyOnorm, bool kUseStaticDecodeLayout = false,
265
+ int kFixedHeads = 0, int kFixedValueHeads = 0,
266
+ bool kUseHeadGrid = false, bool kAccumulateOnormSumsq = false,
267
+ bool kUseActiveQkReduction = false, bool kUseCacheGlobalStore = false,
268
+ bool kComputeOutputBeforeStore = false, bool kSkipWarpSync = false,
269
+ bool kPreloadOnormParams = false,
270
+ bool kPrefetchNextStateChunk = false,
271
+ bool kUseActiveOnormReduction = false, bool kUpdateConvState = false,
272
+ bool kUseLowerBound = false, bool kApplyBetaSigmoid = true>
273
+__global__
274
+__launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel(
275
+ const __nv_bfloat16* __restrict__ x_q,
276
+ const __nv_bfloat16* __restrict__ x_k,
277
+ const __nv_bfloat16* __restrict__ x_v, const float* __restrict__ w_q_t,
278
+ const float* __restrict__ w_k_t, const float* __restrict__ w_v_t,
279
+ const float* __restrict__ bias_q, const float* __restrict__ bias_k,
280
+ const float* __restrict__ bias_v, __nv_bfloat16* __restrict__ cs_q,
281
+ __nv_bfloat16* __restrict__ cs_k, __nv_bfloat16* __restrict__ cs_v,
282
+ const float* __restrict__ a_log, const __nv_bfloat16* __restrict__ g,
283
+ const float* __restrict__ dt_bias, const __nv_bfloat16* __restrict__ beta,
284
+ const __nv_bfloat16* __restrict__ onorm_g,
285
+ const float* __restrict__ onorm_weight,
286
+ const int* __restrict__ ssm_state_indices,
287
+ const int* __restrict__ cu_seqlens, float* __restrict__ state,
288
+ __nv_bfloat16* __restrict__ out, int B, int H, int HV, float lower_bound,
289
+ float scale, float onorm_eps, KdaDecodeStrides strides) {
290
+ const int tid = threadIdx.x;
291
+ const int lane = tid & 31;
292
+ const int warp = tid >> 5;
293
+ int i_n;
294
+ int i_hv;
295
+ int i_h;
296
+ int bos;
297
+ int slot;
298
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
299
+ cudaGridDependencySynchronize();
300
+#endif
301
+ if constexpr (kUseStaticDecodeLayout) {
302
+ if constexpr (kUseHeadGrid) {
303
+ i_n = blockIdx.x;
304
+ i_hv = blockIdx.y;
305
+ } else {
306
+ const int nhv = blockIdx.x;
307
+ i_n = nhv / kFixedValueHeads;
308
+ i_hv = nhv - i_n * kFixedValueHeads;
309
+ }
310
+ i_h = i_hv;
311
+ bos = i_n;
312
+ slot = ssm_state_indices == nullptr ? i_n : ssm_state_indices[i_n];
313
+ } else {
314
+ const int nhv = blockIdx.x;
315
+ i_n = nhv / HV;
316
+ i_hv = nhv - i_n * HV;
317
+ const int hv_per_h = HV / H;
318
+ i_h = i_hv / hv_per_h;
319
+
320
+ bos = cu_seqlens == nullptr ? i_n : cu_seqlens[i_n];
321
+ const int eos = cu_seqlens == nullptr ? i_n + 1 : cu_seqlens[i_n + 1];
322
+ if (eos <= bos) {
323
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
324
+ cudaTriggerProgrammaticLaunchCompletion();
325
+#endif
326
+ return;
327
+ }
328
+ slot = ssm_state_indices == nullptr ? i_n : ssm_state_indices[i_n];
329
+ }
330
+
331
+ constexpr int kLocalDim = kFixedHeads * kDimK;
332
+ constexpr int kPackedDim = 3 * kLocalDim;
333
+ const int hk_off = i_h * kDimK;
334
+ const int hv_off = i_hv * kDimV;
335
+ constexpr int hv_count = kFixedValueHeads;
336
+ float* const state_for_slot = state + slot * strides.state_slot;
337
+ const int64_t conv_slot_offset = slot * strides.conv_slot;
338
+ __nv_bfloat16* const cs_q_for_slot = cs_q + conv_slot_offset;
339
+ __nv_bfloat16* const cs_k_for_slot = cs_k + conv_slot_offset;
340
+ __nv_bfloat16* const cs_v_for_slot = cs_v + conv_slot_offset;
341
+
342
+ __shared__ float s_state[2][kChunkV][kDimK];
343
+ __shared__ float s_q[kDimK];
344
+ __shared__ float s_k[kDimK];
345
+ __shared__ float s_decay[kDimK];
346
+ __shared__ float s_v[kDimV];
347
+ __shared__ float s_o[kDimV];
348
+ __shared__ float s_reduce[kThreads];
349
+ __shared__ float s_beta;
350
+ float pre_onorm_gate = 0.0f;
351
+ float pre_onorm_weight = 0.0f;
352
+
353
+ cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, hv_count, 0);
354
+
355
+ if constexpr (kUpdateConvState) {
356
+ if (tid < kDimK) {
357
+ const int k = tid;
358
+ const int hk = hk_off + k;
359
+ const int64_t xq_idx = bos * strides.x_row + i_h * kDimK + k;
360
+ const float exp_a =
361
+ __shfl_sync(0xffffffffu, lane == 0 ? __expf(a_log[i_h]) : 0.0f, 0);
362
+
363
+ float q_acc = bias_q == nullptr ? 0.0f : bf16_load(bias_q, hk);
364
+ float k_acc = bias_k == nullptr ? 0.0f : bf16_load(bias_k, hk);
365
+ __nv_bfloat16 q_shift0 = __float2bfloat16(0.0f);
366
+ __nv_bfloat16 q_shift1 = __float2bfloat16(0.0f);
367
+ __nv_bfloat16 k_shift0 = __float2bfloat16(0.0f);
368
+ __nv_bfloat16 k_shift1 = __float2bfloat16(0.0f);
369
+#pragma unroll
370
+ for (int w = 0; w < kConvStateWidth; ++w) {
371
+ const __nv_bfloat16 q_state = cs_q_for_slot[hk + w * kPackedDim];
372
+ const __nv_bfloat16 k_state = cs_k_for_slot[hk + w * kPackedDim];
373
+ q_acc += __bfloat162float(q_state) *
374
+ conv_weight_load<kLocalDim>(w_q_t, hk, w);
375
+ k_acc += __bfloat162float(k_state) *
376
+ conv_weight_load<kLocalDim>(w_k_t, hk, w);
377
+ if (w == 1) {
378
+ q_shift0 = q_state;
379
+ k_shift0 = k_state;
380
+ } else if (w == 2) {
381
+ q_shift1 = q_state;
382
+ k_shift1 = k_state;
383
+ }
384
+ }
385
+ const __nv_bfloat16 q_new = x_q[xq_idx];
386
+ const __nv_bfloat16 k_new = x_k[xq_idx];
387
+ q_acc += __bfloat162float(q_new) *
388
+ conv_weight_load<kLocalDim>(w_q_t, hk, kKernelWidth - 1);
389
+ k_acc += __bfloat162float(k_new) *
390
+ conv_weight_load<kLocalDim>(w_k_t, hk, kKernelWidth - 1);
391
+
392
+ cs_q_for_slot[hk] = q_shift0;
393
+ cs_q_for_slot[hk + kPackedDim] = q_shift1;
394
+ cs_q_for_slot[hk + 2 * kPackedDim] = q_new;
395
+ cs_k_for_slot[hk] = k_shift0;
396
+ cs_k_for_slot[hk + kPackedDim] = k_shift1;
397
+ cs_k_for_slot[hk + 2 * kPackedDim] = k_new;
398
+
399
+ s_q[k] = silu_fast(q_acc);
400
+ s_k[k] = silu_fast(k_acc);
401
+
402
+ const int64_t gate_idx = bos * kLocalDim + i_hv * kDimK + k;
403
+ const float g_raw = bf16_load(g, gate_idx) + dt_bias[hk];
404
+ if constexpr (kUseLowerBound) {
405
+ s_decay[k] = __expf(lower_bound * sigmoid_fast(exp_a * g_raw));
406
+ } else {
407
+ s_decay[k] = __expf(-exp_a * softplus_fast(g_raw));
408
+ }
409
+ }
410
+ } else {
411
+ if (tid < kDimK) {
412
+ const int k = tid;
413
+ const int hk = hk_off + k;
414
+ const float exp_a =
415
+ __shfl_sync(0xffffffffu, lane == 0 ? __expf(a_log[i_h]) : 0.0f, 0);
416
+
417
+ float q_acc = bias_q == nullptr ? 0.0f : bf16_load(bias_q, hk);
418
+ float k_acc = bias_k == nullptr ? 0.0f : bf16_load(bias_k, hk);
419
+#pragma unroll
420
+ for (int w = 0; w < kConvStateWidth; ++w) {
421
+ const int cs_idx = hk + w * kPackedDim;
422
+ q_acc += bf16_load(cs_q_for_slot, cs_idx) *
423
+ conv_weight_load<kLocalDim>(w_q_t, hk, w);
424
+ k_acc += bf16_load(cs_k_for_slot, cs_idx) *
425
+ conv_weight_load<kLocalDim>(w_k_t, hk, w);
426
+ }
427
+ q_acc += bf16_load(x_q, bos * strides.x_row + i_h * kDimK + k) *
428
+ conv_weight_load<kLocalDim>(w_q_t, hk, kKernelWidth - 1);
429
+ k_acc += bf16_load(x_k, bos * strides.x_row + i_h * kDimK + k) *
430
+ conv_weight_load<kLocalDim>(w_k_t, hk, kKernelWidth - 1);
431
+
432
+ s_q[k] = silu_fast(q_acc);
433
+ s_k[k] = silu_fast(k_acc);
434
+
435
+ const int64_t gate_idx = bos * kLocalDim + i_hv * kDimK + k;
436
+ const float g_raw = bf16_load(g, gate_idx) + dt_bias[hk];
437
+ if constexpr (kUseLowerBound) {
438
+ s_decay[k] = __expf(lower_bound * sigmoid_fast(exp_a * g_raw));
439
+ } else {
440
+ s_decay[k] = __expf(-exp_a * softplus_fast(g_raw));
441
+ }
442
+ }
443
+ }
444
+
445
+ if constexpr (kUpdateConvState) {
446
+ if (tid < kDimV) {
447
+ const int v = tid;
448
+ const int hvv = hv_off + v;
449
+ const int64_t xv_idx = bos * strides.x_row + i_hv * kDimV + v;
450
+
451
+ float v_acc = bias_v == nullptr ? 0.0f : bf16_load(bias_v, hvv);
452
+ __nv_bfloat16 v_shift0 = __float2bfloat16(0.0f);
453
+ __nv_bfloat16 v_shift1 = __float2bfloat16(0.0f);
454
+#pragma unroll
455
+ for (int w = 0; w < kConvStateWidth; ++w) {
456
+ const __nv_bfloat16 v_state = cs_v_for_slot[hvv + w * kPackedDim];
457
+ v_acc += __bfloat162float(v_state) *
458
+ conv_weight_load<kLocalDim>(w_v_t, hvv, w);
459
+ if (w == 1) {
460
+ v_shift0 = v_state;
461
+ } else if (w == 2) {
462
+ v_shift1 = v_state;
463
+ }
464
+ }
465
+ const __nv_bfloat16 v_new = x_v[xv_idx];
466
+ v_acc += __bfloat162float(v_new) *
467
+ conv_weight_load<kLocalDim>(w_v_t, hvv, kKernelWidth - 1);
468
+ cs_v_for_slot[hvv] = v_shift0;
469
+ cs_v_for_slot[hvv + kPackedDim] = v_shift1;
470
+ cs_v_for_slot[hvv + 2 * kPackedDim] = v_new;
471
+ s_v[v] = silu_fast(v_acc);
472
+
473
+ if constexpr (kApplyOnorm && kPreloadOnormParams) {
474
+ const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + v;
475
+ pre_onorm_gate = sigmoid_fast(bf16_load(onorm_g, gate_idx));
476
+ pre_onorm_weight = onorm_weight[v];
477
+ }
478
+ }
479
+ } else {
480
+ if (tid < kDimV) {
481
+ const int v = tid;
482
+ const int hvv = hv_off + v;
483
+
484
+ float v_acc = bias_v == nullptr ? 0.0f : bf16_load(bias_v, hvv);
485
+#pragma unroll
486
+ for (int w = 0; w < kConvStateWidth; ++w) {
487
+ const int cs_idx = hvv + w * kPackedDim;
488
+ v_acc += bf16_load(cs_v_for_slot, cs_idx) *
489
+ conv_weight_load<kLocalDim>(w_v_t, hvv, w);
490
+ }
491
+ v_acc += bf16_load(x_v, bos * strides.x_row + i_hv * kDimV + v) *
492
+ conv_weight_load<kLocalDim>(w_v_t, hvv, kKernelWidth - 1);
493
+ s_v[v] = silu_fast(v_acc);
494
+
495
+ if constexpr (kApplyOnorm && kPreloadOnormParams) {
496
+ const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + v;
497
+ pre_onorm_gate = sigmoid_fast(bf16_load(onorm_g, gate_idx));
498
+ pre_onorm_weight = onorm_weight[v];
499
+ }
500
+ }
501
+ }
502
+
503
+ if (tid == 0) {
504
+ const float beta_raw = bf16_load(beta, bos * strides.beta_row + i_hv);
505
+ if constexpr (kApplyBetaSigmoid) {
506
+ s_beta = sigmoid_fast(beta_raw);
507
+ } else {
508
+ s_beta = beta_raw;
509
+ }
510
+ }
511
+ __syncthreads();
512
+
513
+ if constexpr (kPrefetchNextStateChunk && kNumChunks > 1) {
514
+ cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, hv_count,
515
+ 1);
516
+ }
517
+
518
+ const float q_sq = tid < kDimK ? s_q[tid] * s_q[tid] : 0.0f;
519
+ const float k_sq = tid < kDimK ? s_k[tid] * s_k[tid] : 0.0f;
520
+ Sum2 qk_sum;
521
+ if constexpr (kUseActiveQkReduction) {
522
+ qk_sum = block_reduce_sum2_active_for<kDimK / 32>(q_sq, k_sq, s_reduce);
523
+ } else {
524
+ qk_sum = block_reduce_sum2(q_sq, k_sq, s_reduce);
525
+ }
526
+ if (tid < kDimK) {
527
+ s_q[tid] *= rsqrtf(qk_sum.x + 1.0e-6f) * scale;
528
+ s_k[tid] *= rsqrtf(qk_sum.y + 1.0e-6f);
529
+ }
530
+ __syncthreads();
531
+
532
+ const int k_base = lane * 4;
533
+ const float4 q4 = *reinterpret_cast<const float4*>(s_q + k_base);
534
+ const float4 k4 = *reinterpret_cast<const float4*>(s_k + k_base);
535
+ const float4 decay4 = *reinterpret_cast<const float4*>(s_decay + k_base);
536
+ float r_q[4] = {q4.x, q4.y, q4.z, q4.w};
537
+ float r_k[4] = {k4.x, k4.y, k4.z, k4.w};
538
+ float r_decay[4] = {decay4.x, decay4.y, decay4.z, decay4.w};
539
+ float o_sumsq = 0.0f;
540
+
541
+#pragma unroll
542
+ for (int chunk = 0; chunk < kNumChunks; ++chunk) {
543
+ if constexpr (kPrefetchNextStateChunk && kNumChunks > 1) {
544
+ if (chunk + 1 < kNumChunks) {
545
+ cp_async_wait_group_1();
546
+ } else {
547
+ cp_async_wait_all();
548
+ }
549
+ } else {
550
+ cp_async_wait_all();
551
+ }
552
+ if constexpr (!kSkipWarpSync) {
553
+ __syncwarp();
554
+ }
555
+
556
+ if constexpr (!kPrefetchNextStateChunk) {
557
+ if (chunk + 1 < kNumChunks) {
558
+ cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv,
559
+ hv_count, chunk + 1);
560
+ }
561
+ }
562
+
563
+#pragma unroll
564
+ for (int row = 0; row < kRowsPerWarp; row += 2) {
565
+ const int v_row_a = warp + row * kWarps;
566
+ const int v_row_b = warp + (row + 1) * kWarps;
567
+ const int v0 = chunk * kChunkV + v_row_a;
568
+ const int v1 = chunk * kChunkV + v_row_b;
569
+ float h_a_vals[4];
570
+ float h_b_vals[4];
571
+ float dot_hk_a = 0.0f;
572
+ float dot_hk_b = 0.0f;
573
+
574
+ const float4 raw_h_a = *reinterpret_cast<const float4*>(
575
+ &s_state[chunk & 1][v_row_a][k_base]);
576
+ const float4 raw_h_b = *reinterpret_cast<const float4*>(
577
+ &s_state[chunk & 1][v_row_b][k_base]);
578
+ h_a_vals[0] = raw_h_a.x * r_decay[0];
579
+ h_a_vals[1] = raw_h_a.y * r_decay[1];
580
+ h_a_vals[2] = raw_h_a.z * r_decay[2];
581
+ h_a_vals[3] = raw_h_a.w * r_decay[3];
582
+ h_b_vals[0] = raw_h_b.x * r_decay[0];
583
+ h_b_vals[1] = raw_h_b.y * r_decay[1];
584
+ h_b_vals[2] = raw_h_b.z * r_decay[2];
585
+ h_b_vals[3] = raw_h_b.w * r_decay[3];
586
+ dot_hk_a = h_a_vals[0] * r_k[0] + h_a_vals[1] * r_k[1] +
587
+ h_a_vals[2] * r_k[2] + h_a_vals[3] * r_k[3];
588
+ dot_hk_b = h_b_vals[0] * r_k[0] + h_b_vals[1] * r_k[1] +
589
+ h_b_vals[2] * r_k[2] + h_b_vals[3] * r_k[3];
590
+
591
+ const Sum2 dot_hk = warp_reduce_sum_pair(dot_hk_a, dot_hk_b);
592
+ const float v_new0 = (s_v[v0] - dot_hk.x) * s_beta;
593
+ const float v_new1 = (s_v[v1] - dot_hk.y) * s_beta;
594
+
595
+ float dot_hq_a = 0.0f;
596
+ float dot_hq_b = 0.0f;
597
+ const int state_idx_a = (i_hv * kDimV + v0) * kDimK + k_base;
598
+ const int state_idx_b = (i_hv * kDimV + v1) * kDimK + k_base;
599
+ const float h_a_0 = h_a_vals[0] + r_k[0] * v_new0;
600
+ const float h_a_1 = h_a_vals[1] + r_k[1] * v_new0;
601
+ const float h_a_2 = h_a_vals[2] + r_k[2] * v_new0;
602
+ const float h_a_3 = h_a_vals[3] + r_k[3] * v_new0;
603
+ const float h_b_0 = h_b_vals[0] + r_k[0] * v_new1;
604
+ const float h_b_1 = h_b_vals[1] + r_k[1] * v_new1;
605
+ const float h_b_2 = h_b_vals[2] + r_k[2] * v_new1;
606
+ const float h_b_3 = h_b_vals[3] + r_k[3] * v_new1;
607
+ if constexpr (kComputeOutputBeforeStore) {
608
+ dot_hq_a =
609
+ h_a_0 * r_q[0] + h_a_1 * r_q[1] + h_a_2 * r_q[2] + h_a_3 * r_q[3];
610
+ dot_hq_b =
611
+ h_b_0 * r_q[0] + h_b_1 * r_q[1] + h_b_2 * r_q[2] + h_b_3 * r_q[3];
612
+ store_state_float4<kUseCacheGlobalStore>(
613
+ state_for_slot + state_idx_a,
614
+ make_float4(h_a_0, h_a_1, h_a_2, h_a_3));
615
+ store_state_float4<kUseCacheGlobalStore>(
616
+ state_for_slot + state_idx_b,
617
+ make_float4(h_b_0, h_b_1, h_b_2, h_b_3));
618
+ } else {
619
+ store_state_float4<kUseCacheGlobalStore>(
620
+ state_for_slot + state_idx_a,
621
+ make_float4(h_a_0, h_a_1, h_a_2, h_a_3));
622
+ store_state_float4<kUseCacheGlobalStore>(
623
+ state_for_slot + state_idx_b,
624
+ make_float4(h_b_0, h_b_1, h_b_2, h_b_3));
625
+ dot_hq_a =
626
+ h_a_0 * r_q[0] + h_a_1 * r_q[1] + h_a_2 * r_q[2] + h_a_3 * r_q[3];
627
+ dot_hq_b =
628
+ h_b_0 * r_q[0] + h_b_1 * r_q[1] + h_b_2 * r_q[2] + h_b_3 * r_q[3];
629
+ }
630
+
631
+ const Sum2 dot_hq = warp_reduce_sum_pair(dot_hq_a, dot_hq_b);
632
+ if (lane == 0) {
633
+ s_o[v0] = dot_hq.x;
634
+ s_o[v1] = dot_hq.y;
635
+ if constexpr (kApplyOnorm && kAccumulateOnormSumsq) {
636
+ o_sumsq += dot_hq.x * dot_hq.x + dot_hq.y * dot_hq.y;
637
+ }
638
+ }
639
+ }
640
+
641
+ if constexpr (kPrefetchNextStateChunk) {
642
+ if (chunk + 2 < kNumChunks) {
643
+ cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv,
644
+ hv_count, chunk + 2);
645
+ }
646
+ }
647
+ }
648
+ __syncthreads();
649
+
650
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
651
+ cudaTriggerProgrammaticLaunchCompletion();
652
+#endif
653
+
654
+ if constexpr (kApplyOnorm) {
655
+ if constexpr (kAccumulateOnormSumsq) {
656
+ if (lane == 0) {
657
+ s_reduce[warp] = o_sumsq;
658
+ }
659
+ __syncthreads();
660
+
661
+ float total_sumsq = 0.0f;
662
+ if (warp == 0) {
663
+ total_sumsq = lane < kWarps ? s_reduce[lane] : 0.0f;
664
+ total_sumsq = warp_reduce_sum(total_sumsq);
665
+ if (lane == 0) {
666
+ s_reduce[0] = total_sumsq;
667
+ }
668
+ }
669
+ __syncthreads();
670
+
671
+ if (tid < kDimV) {
672
+ const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + tid;
673
+ const float raw_o = s_o[tid];
674
+ const float rstd =
675
+ rsqrtf(s_reduce[0] / static_cast<float>(kDimV) + onorm_eps);
676
+ float gate;
677
+ float weight;
678
+ if constexpr (kPreloadOnormParams) {
679
+ gate = pre_onorm_gate;
680
+ weight = pre_onorm_weight;
681
+ } else {
682
+ const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + tid;
683
+ gate = sigmoid_fast(bf16_load(onorm_g, gate_idx));
684
+ weight = onorm_weight[tid];
685
+ }
686
+ const float y = raw_o * rstd * weight * gate;
687
+ out[out_idx] = bf16_store(y);
688
+ }
689
+ } else {
690
+ const float raw_o = tid < kDimV ? s_o[tid] : 0.0f;
691
+ const float o_sq = raw_o * raw_o;
692
+ float sumsq;
693
+ if constexpr (kUseActiveOnormReduction || kUseActiveQkReduction) {
694
+ sumsq = block_reduce_sum_active_for<kDimV / 32>(o_sq, s_reduce);
695
+ } else {
696
+ sumsq = block_reduce_sum(o_sq, s_reduce);
697
+ }
698
+
699
+ if (tid < kDimV) {
700
+ const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + tid;
701
+ const float rstd =
702
+ rsqrtf(sumsq / static_cast<float>(kDimV) + onorm_eps);
703
+ float gate;
704
+ float weight;
705
+ if constexpr (kPreloadOnormParams) {
706
+ gate = pre_onorm_gate;
707
+ weight = pre_onorm_weight;
708
+ } else {
709
+ const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + tid;
710
+ gate = sigmoid_fast(bf16_load(onorm_g, gate_idx));
711
+ weight = onorm_weight[tid];
712
+ }
713
+ const float y = raw_o * rstd * weight * gate;
714
+ out[out_idx] = bf16_store(y);
715
+ }
716
+ }
717
+ } else {
718
+ if (tid < kDimV) {
719
+ const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + tid;
720
+ out[out_idx] = bf16_store(s_o[tid]);
721
+ }
722
+ }
723
+}
724
+
725
+template <int kHeads, bool kApplyOnorm, bool kUpdateConvState,
726
+ bool kUseLowerBound, bool kApplyBetaSigmoid>
727
+void launch_kda_decode_many_heads_raw(
728
+ const void* x_q, const void* x_k, const void* x_v, const void* w_q_t,
729
+ const void* w_k_t, const void* w_v_t, const void* bias_q,
730
+ const void* bias_k, const void* bias_v, void* cs_q, void* cs_k, void* cs_v,
731
+ const float* a_log, const void* g, const float* dt_bias, const void* beta,
732
+ const void* onorm_g, const float* onorm_weight,
733
+ const int* ssm_state_indices, const int* cu_seqlens, float* state,
734
+ void* out, int B, int H, int HV, float lower_bound, float scale,
735
+ float onorm_eps, KdaDecodeStrides strides, cudaStream_t stream) {
736
+ auto kernel = &kda_decode_fusion_many_heads_kernel<
737
+ kApplyOnorm, true, kHeads, kHeads, true, false, false, false, false,
738
+ false, true, true, true, kUpdateConvState, kUseLowerBound,
739
+ kApplyBetaSigmoid>;
740
+ cudaLaunchConfig_t config{};
741
+ config.gridDim = dim3(B, kHeads);
742
+ config.blockDim = dim3(kThreads);
743
+ config.dynamicSmemBytes = 0;
744
+ config.stream = stream;
745
+ cudaLaunchAttribute attrs[1];
746
+ attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
747
+ attrs[0].val.programmaticStreamSerializationAllowed = 1;
748
+ config.attrs = attrs;
749
+ config.numAttrs = 1;
750
+ cudaLaunchKernelEx(&config, kernel,
751
+ reinterpret_cast<const __nv_bfloat16*>(x_q),
752
+ reinterpret_cast<const __nv_bfloat16*>(x_k),
753
+ reinterpret_cast<const __nv_bfloat16*>(x_v),
754
+ reinterpret_cast<const float*>(w_q_t),
755
+ reinterpret_cast<const float*>(w_k_t),
756
+ reinterpret_cast<const float*>(w_v_t),
757
+ reinterpret_cast<const float*>(bias_q),
758
+ reinterpret_cast<const float*>(bias_k),
759
+ reinterpret_cast<const float*>(bias_v),
760
+ reinterpret_cast<__nv_bfloat16*>(cs_q),
761
+ reinterpret_cast<__nv_bfloat16*>(cs_k),
762
+ reinterpret_cast<__nv_bfloat16*>(cs_v), a_log,
763
+ reinterpret_cast<const __nv_bfloat16*>(g), dt_bias,
764
+ reinterpret_cast<const __nv_bfloat16*>(beta),
765
+ reinterpret_cast<const __nv_bfloat16*>(onorm_g),
766
+ onorm_weight, ssm_state_indices, cu_seqlens, state,
767
+ reinterpret_cast<__nv_bfloat16*>(out), B, H, HV,
768
+ lower_bound, scale, onorm_eps, strides);
769
+}
770
+
771
+template <bool kApplyOnorm, bool kUseLowerBound, bool kApplyBetaSigmoid>
772
+void launch_kda_decode_many_heads_selected(
773
+ const void* x_q, const void* x_k, const void* x_v, const void* w_q_t,
774
+ const void* w_k_t, const void* w_v_t, const void* bias_q,
775
+ const void* bias_k, const void* bias_v, void* cs_q, void* cs_k, void* cs_v,
776
+ const float* a_log, const void* g, const float* dt_bias, const void* beta,
777
+ const void* onorm_g, const float* onorm_weight,
778
+ const int* ssm_state_indices, const int* cu_seqlens, float* state,
779
+ void* out, int B, int H, int HV, bool update_conv_cache, float lower_bound,
780
+ float scale, float onorm_eps, KdaDecodeStrides strides,
781
+ cudaStream_t stream) {
782
+#define LAUNCH_KDA_DECODE(NUM_HEADS) \
783
+ do { \
784
+ if (update_conv_cache) { \
785
+ launch_kda_decode_many_heads_raw<NUM_HEADS, kApplyOnorm, true, \
786
+ kUseLowerBound, kApplyBetaSigmoid>( \
787
+ x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \
788
+ cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \
789
+ ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \
790
+ scale, onorm_eps, strides, stream); \
791
+ } else { \
792
+ launch_kda_decode_many_heads_raw<NUM_HEADS, kApplyOnorm, false, \
793
+ kUseLowerBound, kApplyBetaSigmoid>( \
794
+ x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \
795
+ cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \
796
+ ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \
797
+ scale, onorm_eps, strides, stream); \
798
+ } \
799
+ } while (false)
800
+
801
+ switch (H) {
802
+ case 12:
803
+ LAUNCH_KDA_DECODE(12);
804
+ break;
805
+ case 24:
806
+ LAUNCH_KDA_DECODE(24);
807
+ break;
808
+ case 48:
809
+ LAUNCH_KDA_DECODE(48);
810
+ break;
811
+ case 96:
812
+ LAUNCH_KDA_DECODE(96);
813
+ break;
814
+ default:
815
+ STD_TORCH_CHECK(false, "Unsupported number of heads: ", H);
816
+ }
817
+
818
+#undef LAUNCH_KDA_DECODE
819
+}
820
+
821
+struct KdaDecodeLaunchParams {
822
+ const void* x_q;
823
+ const void* x_k;
824
+ const void* x_v;
825
+ const void* w_q_t;
826
+ const void* w_k_t;
827
+ const void* w_v_t;
828
+ const void* bias_q;
829
+ const void* bias_k;
830
+ const void* bias_v;
831
+ void* cs_q;
832
+ void* cs_k;
833
+ void* cs_v;
834
+ const float* a_log;
835
+ const void* g;
836
+ const float* dt_bias;
837
+ const void* beta;
838
+ const void* onorm_g;
839
+ const float* onorm_weight;
840
+ const int* ssm_state_indices;
841
+ const int* cu_seqlens;
842
+ float* state;
843
+ void* out;
844
+ int B;
845
+ int H;
846
+ int HV;
847
+ bool update_conv_cache;
848
+ float lower_bound;
849
+ float scale;
850
+ float onorm_eps;
851
+ KdaDecodeStrides strides;
852
+ cudaStream_t stream;
853
+};
854
+
855
+template <bool kApplyOnorm, bool kUseLowerBound, bool kApplyBetaSigmoid>
856
+void launch_kda_decode_selected_backend(const KdaDecodeLaunchParams& p) {
857
+ launch_kda_decode_many_heads_selected<kApplyOnorm, kUseLowerBound,
858
+ kApplyBetaSigmoid>(
859
+ p.x_q, p.x_k, p.x_v, p.w_q_t, p.w_k_t, p.w_v_t, p.bias_q, p.bias_k,
860
+ p.bias_v, p.cs_q, p.cs_k, p.cs_v, p.a_log, p.g, p.dt_bias, p.beta,
861
+ p.onorm_g, p.onorm_weight, p.ssm_state_indices, p.cu_seqlens, p.state,
862
+ p.out, p.B, p.H, p.HV, p.update_conv_cache, p.lower_bound, p.scale,
863
+ p.onorm_eps, p.strides, p.stream);
864
+}
865
+
866
+template <bool kApplyOnorm, bool kUseLowerBound>
867
+void dispatch_kda_decode_beta(const KdaDecodeLaunchParams& p,
868
+ bool apply_beta_sigmoid) {
869
+ if (apply_beta_sigmoid) {
870
+ launch_kda_decode_selected_backend<kApplyOnorm, kUseLowerBound, true>(p);
871
+ } else {
872
+ launch_kda_decode_selected_backend<kApplyOnorm, kUseLowerBound, false>(p);
873
+ }
874
+}
875
+
876
+template <bool kApplyOnorm>
877
+void dispatch_kda_decode_decay(const KdaDecodeLaunchParams& p,
878
+ bool use_lower_bound, bool apply_beta_sigmoid) {
879
+ if (use_lower_bound) {
880
+ dispatch_kda_decode_beta<kApplyOnorm, true>(p, apply_beta_sigmoid);
881
+ } else {
882
+ dispatch_kda_decode_beta<kApplyOnorm, false>(p, apply_beta_sigmoid);
883
+ }
884
+}
885
+
886
+void dispatch_kda_decode_features(const KdaDecodeLaunchParams& p,
887
+ bool apply_onorm, bool use_lower_bound,
888
+ bool apply_beta_sigmoid) {
889
+ if (apply_onorm) {
890
+ dispatch_kda_decode_decay<true>(p, use_lower_bound, apply_beta_sigmoid);
891
+ } else {
892
+ dispatch_kda_decode_decay<false>(p, use_lower_bound, apply_beta_sigmoid);
893
+ }
894
+}
895
+
896
+} // namespace
897
+
898
+extern "C" void launch_kda_decode_many_heads_cuda(
899
+ const void* x_q, const void* x_k, const void* x_v, const void* w_q_t,
900
+ const void* w_k_t, const void* w_v_t, const void* bias_q,
901
+ const void* bias_k, const void* bias_v, void* cs_q, void* cs_k, void* cs_v,
902
+ const float* a_log, const void* g, const float* dt_bias, const void* beta,
903
+ const void* onorm_g, const float* onorm_weight,
904
+ const int* ssm_state_indices, const int* cu_seqlens, float* state,
905
+ void* out, int B, int H, int HV, bool apply_onorm, bool update_conv_cache,
906
+ bool use_lower_bound, bool apply_beta_sigmoid, float lower_bound,
907
+ float scale, float onorm_eps, const int64_t* raw_strides,
908
+ cudaStream_t stream) {
909
+ const KdaDecodeStrides strides{raw_strides[0], raw_strides[1], raw_strides[2],
910
+ raw_strides[3], raw_strides[4]};
911
+ const KdaDecodeLaunchParams params{x_q,
912
+ x_k,
913
+ x_v,
914
+ w_q_t,
915
+ w_k_t,
916
+ w_v_t,
917
+ bias_q,
918
+ bias_k,
919
+ bias_v,
920
+ cs_q,
921
+ cs_k,
922
+ cs_v,
923
+ a_log,
924
+ g,
925
+ dt_bias,
926
+ beta,
927
+ onorm_g,
928
+ onorm_weight,
929
+ ssm_state_indices,
930
+ cu_seqlens,
931
+ state,
932
+ out,
933
+ B,
934
+ H,
935
+ HV,
936
+ update_conv_cache,
937
+ lower_bound,
938
+ scale,
939
+ onorm_eps,
940
+ strides,
941
+ stream};
942
+ dispatch_kda_decode_features(params, apply_onorm, use_lower_bound,
943
+ apply_beta_sigmoid);
944
+}
945
+
946
+void fused_kda_decode(
947
+ torch::stable::Tensor const& x, torch::stable::Tensor const& weight,
948
+ std::optional<torch::stable::Tensor> bias,
949
+ torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g,
950
+ torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log,
951
+ torch::stable::Tensor const& dt_bias,
952
+ torch::stable::Tensor const& state_indices, torch::stable::Tensor& state,
953
+ torch::stable::Tensor& out, std::optional<double> lower_bound,
954
+ std::optional<torch::stable::Tensor> output_gate,
955
+ std::optional<torch::stable::Tensor> norm_weight, double norm_eps) {
956
+ using torch::headeronly::ScalarType;
957
+ constexpr int kHeadDim = 128;
958
+ constexpr int kConvWidth = 4;
959
+
960
+ STD_TORCH_CHECK(x.is_cuda() && x.scalar_type() == ScalarType::BFloat16,
961
+ "x must be a CUDA bfloat16 tensor");
962
+ STD_TORCH_CHECK(weight.is_cuda() && weight.scalar_type() == ScalarType::Float,
963
+ "weight must be a CUDA float32 tensor");
964
+ STD_TORCH_CHECK(
965
+ conv_state.is_cuda() && conv_state.scalar_type() == ScalarType::BFloat16,
966
+ "conv_state must be a CUDA bfloat16 tensor");
967
+ STD_TORCH_CHECK(
968
+ raw_g.is_cuda() && raw_g.scalar_type() == ScalarType::BFloat16,
969
+ "raw_g must be a CUDA bfloat16 tensor");
970
+ STD_TORCH_CHECK(
971
+ raw_beta.is_cuda() && raw_beta.scalar_type() == ScalarType::BFloat16,
972
+ "raw_beta must be a CUDA bfloat16 tensor");
973
+ STD_TORCH_CHECK(a_log.is_cuda() && a_log.scalar_type() == ScalarType::Float,
974
+ "A_log must be a CUDA float32 tensor");
975
+ STD_TORCH_CHECK(
976
+ dt_bias.is_cuda() && dt_bias.scalar_type() == ScalarType::Float,
977
+ "dt_bias must be a CUDA float32 tensor");
978
+ STD_TORCH_CHECK(state.is_cuda() && state.scalar_type() == ScalarType::Float,
979
+ "state must be a CUDA float32 tensor");
980
+ STD_TORCH_CHECK(out.is_cuda() && out.scalar_type() == ScalarType::BFloat16,
981
+ "out must be a CUDA bfloat16 tensor");
982
+ STD_TORCH_CHECK(
983
+ state_indices.is_cuda() && state_indices.scalar_type() == ScalarType::Int,
984
+ "state_indices must be a CUDA int32 tensor");
985
+
986
+ STD_TORCH_CHECK(x.dim() == 2, "x must have shape [B, 3 * H * 128]");
987
+ int const batch_size = static_cast<int>(x.size(0));
988
+ int64_t const qkv_width = x.size(1);
989
+ STD_TORCH_CHECK(qkv_width % (3 * kHeadDim) == 0,
990
+ "x must have shape [B, 3 * H * 128]");
991
+ int64_t const num_heads = qkv_width / (3 * kHeadDim);
992
+ STD_TORCH_CHECK(
993
+ num_heads == 12 || num_heads == 24 || num_heads == 48 || num_heads == 96,
994
+ "H must be 12, 24, 48, or 96, got ", num_heads);
995
+ STD_TORCH_CHECK(batch_size > 0,
996
+ "KDA decode fusion requires at least one row");
997
+ int const dim = num_heads * kHeadDim;
998
+
999
+ STD_TORCH_CHECK(weight.dim() == 3 && weight.is_contiguous() &&
1000
+ weight.size(0) == 3 && weight.size(1) == kConvWidth &&
1001
+ weight.size(2) == dim,
1002
+ "weight must have shape [3, 4, H * 128]");
1003
+ STD_TORCH_CHECK(conv_state.dim() == 3 && conv_state.size(1) == 3 * dim &&
1004
+ conv_state.size(2) == kConvWidth - 1,
1005
+ "conv_state must have shape [slots, 3 * H * 128, 3]");
1006
+ STD_TORCH_CHECK(raw_g.dim() == 4 && raw_g.size(0) == 1 &&
1007
+ raw_g.size(1) == batch_size &&
1008
+ raw_g.size(2) == num_heads && raw_g.size(3) == kHeadDim,
1009
+ "raw_g must have shape [1, B, H, 128]");
1010
+ STD_TORCH_CHECK(raw_beta.dim() == 3 && raw_beta.size(0) == 1 &&
1011
+ raw_beta.size(1) == batch_size &&
1012
+ raw_beta.size(2) == num_heads,
1013
+ "raw_beta must have shape [1, B, H]");
1014
+ STD_TORCH_CHECK(a_log.is_contiguous() && a_log.numel() == num_heads,
1015
+ "A_log must be contiguous with H elements");
1016
+ STD_TORCH_CHECK(dt_bias.is_contiguous() && dt_bias.numel() == dim,
1017
+ "dt_bias must be contiguous with H * 128 elements");
1018
+ STD_TORCH_CHECK(
1019
+ state_indices.is_contiguous() && state_indices.numel() == batch_size,
1020
+ "state_indices must be contiguous with B elements");
1021
+ STD_TORCH_CHECK(state.dim() == 4 && state.size(1) == num_heads &&
1022
+ state.size(2) == kHeadDim && state.size(3) == kHeadDim,
1023
+ "state must have shape [slots, H, 128, 128]");
1024
+ STD_TORCH_CHECK(out.dim() == 4 && out.size(0) == 1 &&
1025
+ out.size(1) == batch_size && out.size(2) == num_heads &&
1026
+ out.size(3) == kHeadDim,
1027
+ "out must have shape [1, B, H, 128]");
1028
+ STD_TORCH_CHECK(x.stride(1) == 1,
1029
+ "x must be contiguous in its channel dimension");
1030
+ STD_TORCH_CHECK(conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) &&
1031
+ conv_state.stride(1) == 1 &&
1032
+ conv_state.stride(2) == 3 * dim,
1033
+ "conv_state must use the SD cache layout");
1034
+ STD_TORCH_CHECK(state.stride(0) >= num_heads * kHeadDim * kHeadDim &&
1035
+ state.stride(1) == kHeadDim * kHeadDim &&
1036
+ state.stride(2) == kHeadDim && state.stride(3) == 1,
1037
+ "state must have contiguous [H, 128, 128] slot contents");
1038
+ STD_TORCH_CHECK(raw_g.is_contiguous(), "raw_g must be contiguous");
1039
+ STD_TORCH_CHECK(raw_beta.stride(2) == 1,
1040
+ "raw_beta must be contiguous in its head dimension");
1041
+ STD_TORCH_CHECK(out.is_contiguous(), "out must be contiguous");
1042
+
1043
+ bool const apply_onorm = output_gate.has_value();
1044
+ STD_TORCH_CHECK(apply_onorm == norm_weight.has_value(),
1045
+ "output_gate and norm_weight must be provided together");
1046
+ void const* output_gate_ptr = nullptr;
1047
+ float const* norm_weight_ptr = nullptr;
1048
+ int64_t output_gate_row_stride = 0;
1049
+ if (apply_onorm) {
1050
+ STD_TORCH_CHECK(output_gate->is_cuda() &&
1051
+ output_gate->scalar_type() == ScalarType::BFloat16,
1052
+ "output_gate must be a CUDA bfloat16 tensor");
1053
+ bool const gate_is_3d =
1054
+ output_gate->dim() == 3 && output_gate->size(0) == batch_size &&
1055
+ output_gate->size(1) == num_heads && output_gate->size(2) == kHeadDim;
1056
+ bool const gate_is_4d =
1057
+ output_gate->dim() == 4 && output_gate->size(0) == 1 &&
1058
+ output_gate->size(1) == batch_size &&
1059
+ output_gate->size(2) == num_heads && output_gate->size(3) == kHeadDim;
1060
+ STD_TORCH_CHECK(gate_is_3d || gate_is_4d,
1061
+ "output_gate must have shape [B, H, 128] or "
1062
+ "[1, B, H, 128]");
1063
+ int const row_dim = gate_is_3d ? 0 : 1;
1064
+ STD_TORCH_CHECK(output_gate->stride(output_gate->dim() - 1) == 1,
1065
+ "output_gate must be contiguous in its last dimension");
1066
+ STD_TORCH_CHECK(output_gate->stride(row_dim + 1) == kHeadDim,
1067
+ "output_gate must have contiguous head rows");
1068
+ STD_TORCH_CHECK(norm_weight->is_cuda() &&
1069
+ norm_weight->scalar_type() == ScalarType::Float,
1070
+ "norm_weight must be a CUDA float32 tensor");
1071
+ STD_TORCH_CHECK(
1072
+ norm_weight->is_contiguous() && norm_weight->numel() == kHeadDim,
1073
+ "norm_weight must be contiguous with 128 elements");
1074
+ STD_TORCH_CHECK(norm_eps >= 0.0, "norm_eps must be non-negative");
1075
+ output_gate_ptr = output_gate->data_ptr();
1076
+ norm_weight_ptr = static_cast<float const*>(norm_weight->data_ptr());
1077
+ output_gate_row_stride = output_gate->stride(row_dim);
1078
+ }
1079
+
1080
+ void const* bias_ptr = nullptr;
1081
+ if (bias.has_value()) {
1082
+ STD_TORCH_CHECK(bias->is_cuda() && bias->scalar_type() == ScalarType::Float,
1083
+ "bias must be a CUDA float32 tensor");
1084
+ STD_TORCH_CHECK(bias->is_contiguous() && bias->numel() == 3 * dim,
1085
+ "bias must be contiguous with 3 * H * 128 elements");
1086
+ bias_ptr = bias->data_ptr();
1087
+ }
1088
+
1089
+ auto const* x_ptr = static_cast<char const*>(x.data_ptr());
1090
+ auto const* weight_ptr = static_cast<char const*>(weight.data_ptr());
1091
+ auto* conv_ptr = static_cast<char*>(conv_state.data_ptr());
1092
+ auto const* bias_bytes = static_cast<char const*>(bias_ptr);
1093
+ int64_t const segment_bytes = dim * sizeof(__nv_bfloat16);
1094
+ int64_t const weight_segment_bytes =
1095
+ dim * kConvWidth * static_cast<int64_t>(sizeof(float));
1096
+ int64_t const conv_segment_bytes =
1097
+ dim * conv_state.stride(1) * sizeof(__nv_bfloat16);
1098
+ int64_t const bias_segment_bytes = dim * sizeof(float);
1099
+ std::array<int64_t, 5> const strides{
1100
+ x.stride(0), raw_beta.stride(1), output_gate_row_stride,
1101
+ conv_state.stride(0), state.stride(0),
1102
+ };
1103
+ bool const use_lower_bound = lower_bound.has_value();
1104
+ float const lower_bound_value =
1105
+ use_lower_bound ? static_cast<float>(*lower_bound) : 0.0f;
1106
+
1107
+ torch::stable::accelerator::DeviceGuard const device_guard(
1108
+ x.get_device_index());
1109
+ cudaStream_t const stream = get_current_cuda_stream(x.get_device_index());
1110
+ launch_kda_decode_many_heads_cuda(
1111
+ x_ptr, x_ptr + segment_bytes, x_ptr + 2 * segment_bytes, weight_ptr,
1112
+ weight_ptr + weight_segment_bytes, weight_ptr + 2 * weight_segment_bytes,
1113
+ bias_bytes,
1114
+ bias_bytes == nullptr ? nullptr : bias_bytes + bias_segment_bytes,
1115
+ bias_bytes == nullptr ? nullptr : bias_bytes + 2 * bias_segment_bytes,
1116
+ conv_ptr, conv_ptr + conv_segment_bytes,
1117
+ conv_ptr + 2 * conv_segment_bytes,
1118
+ static_cast<float const*>(a_log.data_ptr()), raw_g.data_ptr(),
1119
+ static_cast<float const*>(dt_bias.data_ptr()), raw_beta.data_ptr(),
1120
+ output_gate_ptr, norm_weight_ptr,
1121
+ static_cast<int const*>(state_indices.data_ptr()), nullptr,
1122
+ static_cast<float*>(state.data_ptr()), out.data_ptr(), batch_size,
1123
+ num_heads, num_heads, apply_onorm, true, use_lower_bound, true,
1124
+ lower_bound_value, 0.08838834764831845f, static_cast<float>(norm_eps),
1125
+ strides.data(), stream);
1126
+ cudaError_t const error = cudaGetLastError();
1127
+ STD_TORCH_CHECK(
1128
+ error == cudaSuccess,
1129
+ "Kimi K3 KDA decode kernel launch failed: ", cudaGetErrorString(error));
1130
+}
csrc/libtorch_stable/moe/grouped_topk_kernels.cu
+437
-1
@@ -25,6 +25,7 @@
25
#include "libtorch_stable/torch_utils.h"
26
27
#include <cmath>
28
+#include <cstdint>
29
#include <tuple>
30
#include <cuda_fp16.h>
31
#include <cuda_bf16.h>
@@ -448,7 +449,8 @@ enum ScoringFunc {
449
SCORING_SIGMOID = 1 // apply sigmoid
450
};
451
451
-// Efficient sigmoid approximation from TensorRT-LLM
452
+// Adapted from
453
+// https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu
454
__device__ inline float sigmoid_accurate(float x) {
455
return 0.5f * tanhf(0.5f * x) + 0.5f;
456
}
@@ -890,6 +892,434 @@ __global__ void grouped_topk_fused_small_expert_count_kernel(
892
#endif
893
}
894
895
+// Adapted from
896
+// https://github.com/flashinfer-ai/flashinfer/blob/06400d062a2d51564bbe781f6f811d0b75ca593e/include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh
897
+namespace single_group_topk {
898
+namespace detail {
899
+
900
+static constexpr int BlockDim = 256;
901
+static constexpr uint32_t FullWarpMask = 0xffffffffU;
902
+static constexpr float InvalidScore = -INFINITY;
903
+
904
+// TopK-only tuning: use wider workers and keep these tiers on the block path.
905
+template <int MaxNumExperts, int MaxNumTopExperts>
906
+static constexpr bool UseTunedBlockPath =
907
+ MaxNumTopExperts == 16 && (MaxNumExperts == 896 || MaxNumExperts == 1024);
908
+
909
+template <typename T, typename BiasT, ScoringFunc SF>
910
+__device__ __forceinline__ void preprocess_score(T input, BiasT correction_bias,
911
+ float& unbiased_score,
912
+ float& selection_score) {
913
+ unbiased_score = 0.0F;
914
+ selection_score = InvalidScore;
915
+ float const input_float = cuda_cast<float, T>(input);
916
+ float const bias = cuda_cast<float, BiasT>(correction_bias);
917
+ if (!is_finite(input_float) || !is_finite(bias)) {
918
+ return;
919
+ }
920
+
921
+ float const unbiased = apply_scoring<SF>(input_float);
922
+ float const biased = unbiased + bias;
923
+ if constexpr (SF == SCORING_NONE) {
924
+ if (!is_finite(biased)) {
925
+ return;
926
+ }
927
+ }
928
+ unbiased_score = unbiased;
929
+ selection_score = biased == 0.0F ? 0.0F : biased;
930
+}
931
+
932
+template <typename IdxT>
933
+__device__ __forceinline__ void write_outputs(
934
+ cg::thread_block_tile<WARP_SIZE> const& warp, float lane_selection_score,
935
+ float lane_unbiased, int32_t lane_expert, int32_t lane, int32_t token,
936
+ int32_t topk, float* topk_values, IdxT* topk_indices, bool renormalize,
937
+ float routed_scaling_factor) {
938
+ bool const finite_selection =
939
+ lane < topk && lane_selection_score != InvalidScore;
940
+ lane_unbiased = finite_selection ? lane_unbiased : 0.0F;
941
+ unsigned const finite_mask = __ballot_sync(FullWarpMask, finite_selection);
942
+ float const sum = cg::reduce(warp, lane_unbiased, cg::plus<float>{});
943
+
944
+ if (lane < topk) {
945
+ float output = 0.0F;
946
+ if (finite_mask == 0) {
947
+ if (renormalize) {
948
+ output = 1.0F / static_cast<float>(topk);
949
+ }
950
+ } else if (finite_selection) {
951
+ float scale = routed_scaling_factor;
952
+ if (renormalize) {
953
+ scale /= sum + 1e-20F;
954
+ }
955
+ output = lane_unbiased * scale;
956
+ }
957
+
958
+ int64_t const output_index = int64_t{token} * topk + lane;
959
+ topk_values[output_index] = output;
960
+ topk_indices[output_index] = static_cast<IdxT>(lane_expert);
961
+ }
962
+}
963
+
964
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
965
+ int MaxNumExperts, int MaxNumTopExperts>
966
+__global__ void __launch_bounds__(BlockDim)
967
+ single_group_topk_block_kernel(T const* scores, float* topk_values,
968
+ IdxT* topk_indices, BiasT const* bias,
969
+ int64_t num_experts, int64_t topk,
970
+ bool renormalize,
971
+ float routed_scaling_factor,
972
+ bool enable_pdl) {
973
+ static constexpr int NumChunks = (MaxNumExperts + WARP_SIZE - 1) / WARP_SIZE;
974
+ static constexpr int WorkerValuesPerLane =
975
+ UseTunedBlockPath<MaxNumExperts, MaxNumTopExperts> ? 8 : 4;
976
+ static constexpr int ExpertsPerWorkerWarp = WorkerValuesPerLane * WARP_SIZE;
977
+ using LaneOwnedRange =
978
+ reduce_topk::HighExpertLaneOwnedTopKRange<MaxNumExperts,
979
+ MaxNumTopExperts>;
980
+ static constexpr int NumWorkerWarps =
981
+ (MaxNumExperts + ExpertsPerWorkerWarp - 1) / ExpertsPerWorkerWarp;
982
+ static constexpr int NumIntermediate = NumWorkerWarps * MaxNumTopExperts;
983
+ static constexpr int MergeValuesPerLane =
984
+ (NumIntermediate + WARP_SIZE - 1) / WARP_SIZE;
985
+ static constexpr bool LaneOwnedResourcesFit =
986
+ NumWorkerWarps <= BlockDim / WARP_SIZE && MergeValuesPerLane <= 64;
987
+ static constexpr bool UseHierarchicalLaneTopK =
988
+ LaneOwnedRange::kEnabled && LaneOwnedResourcesFit;
989
+
990
+ static_assert(NumChunks <= 64);
991
+ static_assert(MaxNumTopExperts <= WARP_SIZE);
992
+
993
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
994
+ if (enable_pdl) {
995
+ cudaGridDependencySynchronize();
996
+ }
997
+#endif
998
+
999
+ __shared__ float __attribute((aligned(128))) biased_scores[MaxNumExperts];
1000
+ __shared__ float __attribute((aligned(128))) unbiased_scores[MaxNumExperts];
1001
+
1002
+ int32_t const token = static_cast<int32_t>(blockIdx.x);
1003
+ int32_t const lane = static_cast<int32_t>(threadIdx.x) % WARP_SIZE;
1004
+ int32_t const warp_id = static_cast<int32_t>(threadIdx.x) / WARP_SIZE;
1005
+ int32_t const num_experts_i32 = static_cast<int32_t>(num_experts);
1006
+ int32_t const topk_i32 = static_cast<int32_t>(topk);
1007
+ T const* token_scores = scores + int64_t{token} * num_experts;
1008
+
1009
+ for (int32_t expert = static_cast<int32_t>(threadIdx.x);
1010
+ expert < num_experts_i32; expert += BlockDim) {
1011
+ preprocess_score<T, BiasT, SF>(token_scores[expert], bias[expert],
1012
+ unbiased_scores[expert],
1013
+ biased_scores[expert]);
1014
+ }
1015
+ __syncthreads();
1016
+
1017
+ auto warp = cg::tiled_partition<WARP_SIZE>(cg::this_thread_block());
1018
+
1019
+ if constexpr (UseHierarchicalLaneTopK) {
1020
+ __shared__ float
1021
+ __attribute((aligned(128))) intermediate_scores[NumIntermediate];
1022
+ __shared__ int32_t
1023
+ __attribute((aligned(128))) intermediate_indices[NumIntermediate];
1024
+
1025
+ if (warp_id < NumWorkerWarps) {
1026
+ float local_scores[WorkerValuesPerLane];
1027
+ int32_t local_indices[WorkerValuesPerLane];
1028
+#pragma unroll
1029
+ for (int index = 0; index < WorkerValuesPerLane; ++index) {
1030
+ int32_t const expert =
1031
+ warp_id * ExpertsPerWorkerWarp + index * WARP_SIZE + lane;
1032
+ local_scores[index] =
1033
+ expert < num_experts_i32 ? biased_scores[expert] : InvalidScore;
1034
+ local_indices[index] = expert;
1035
+ }
1036
+
1037
+ float lane_score;
1038
+ int32_t lane_expert;
1039
+ reduce_topk::reduceTopKForLane<MaxNumTopExperts>(
1040
+ warp, lane_score, lane_expert, local_scores, local_indices,
1041
+ InvalidScore, lane);
1042
+ if (lane < MaxNumTopExperts) {
1043
+ int32_t const intermediate = warp_id * MaxNumTopExperts + lane;
1044
+ bool const active = lane < topk_i32;
1045
+ intermediate_scores[intermediate] = active ? lane_score : InvalidScore;
1046
+ intermediate_indices[intermediate] =
1047
+ active ? lane_expert : MaxNumExperts;
1048
+ }
1049
+ }
1050
+ __syncthreads();
1051
+
1052
+ if (warp_id != 0) {
1053
+ return;
1054
+ }
1055
+
1056
+ float merge_scores[MergeValuesPerLane];
1057
+ int32_t merge_indices[MergeValuesPerLane];
1058
+#pragma unroll
1059
+ for (int index = 0; index < MergeValuesPerLane; ++index) {
1060
+ int32_t const intermediate = index * WARP_SIZE + lane;
1061
+ bool const active = intermediate < NumIntermediate;
1062
+ merge_scores[index] =
1063
+ active ? intermediate_scores[intermediate] : InvalidScore;
1064
+ merge_indices[index] =
1065
+ active ? intermediate_indices[intermediate] : MaxNumExperts;
1066
+ }
1067
+
1068
+ float lane_score;
1069
+ int32_t lane_expert;
1070
+ reduce_topk::reduceTopKForLane<MaxNumTopExperts>(
1071
+ warp, lane_score, lane_expert, merge_scores, merge_indices,
1072
+ InvalidScore, lane);
1073
+ float const lane_unbiased =
1074
+ lane < topk_i32 && lane_expert >= 0 && lane_expert < num_experts_i32
1075
+ ? unbiased_scores[lane_expert]
1076
+ : 0.0F;
1077
+ write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token,
1078
+ topk_i32, topk_values, topk_indices, renormalize,
1079
+ routed_scaling_factor);
1080
+ } else {
1081
+ if (warp_id != 0) {
1082
+ return;
1083
+ }
1084
+
1085
+ float local_scores[NumChunks];
1086
+ int32_t local_indices[NumChunks];
1087
+#pragma unroll
1088
+ for (int index = 0; index < NumChunks; ++index) {
1089
+ int32_t const expert = index * WARP_SIZE + lane;
1090
+ local_scores[index] =
1091
+ expert < num_experts_i32 ? biased_scores[expert] : InvalidScore;
1092
+ local_indices[index] = expert;
1093
+ }
1094
+
1095
+ float top_scores[MaxNumTopExperts];
1096
+ int32_t top_experts[MaxNumTopExperts];
1097
+ reduce_topk::reduceTopK(warp, top_scores, top_experts, local_scores,
1098
+ local_indices, InvalidScore, topk_i32);
1099
+ float const lane_score = lane < topk_i32 ? top_scores[lane] : InvalidScore;
1100
+ int32_t const lane_expert = lane < topk_i32 ? top_experts[lane] : -1;
1101
+ float const lane_unbiased =
1102
+ lane < topk_i32 && lane_expert >= 0 && lane_expert < num_experts_i32
1103
+ ? unbiased_scores[lane_expert]
1104
+ : 0.0F;
1105
+ write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token,
1106
+ topk_i32, topk_values, topk_indices, renormalize,
1107
+ routed_scaling_factor);
1108
+ }
1109
+
1110
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
1111
+ if (enable_pdl) {
1112
+ cudaTriggerProgrammaticLaunchCompletion();
1113
+ }
1114
+#endif
1115
+}
1116
+
1117
+template <int MaxNumExperts>
1118
+struct WarpTopKLaunchConfig {
1119
+ static constexpr int DefaultBlockDim =
1120
+ MaxNumExperts <= 1024 ? MaxNumExperts : 1024;
1121
+ static constexpr int BlockDim = DefaultBlockDim > 256 ? 256 : DefaultBlockDim;
1122
+ static constexpr int NumWarps = BlockDim / WARP_SIZE;
1123
+ static constexpr int MaxBlockScale =
1124
+ (DefaultBlockDim + BlockDim - 1) / BlockDim;
1125
+ static constexpr int MaxBlocks = 1024 * MaxBlockScale;
1126
+
1127
+ static_assert(BlockDim % WARP_SIZE == 0);
1128
+
1129
+ static uint32_t grid_dim(int64_t num_tokens) {
1130
+ int64_t const token_blocks = (num_tokens + NumWarps - 1) / NumWarps;
1131
+ int64_t const selected =
1132
+ token_blocks < MaxBlocks ? token_blocks : MaxBlocks;
1133
+ return static_cast<uint32_t>(selected > 0 ? selected : 1);
1134
+ }
1135
+};
1136
+
1137
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
1138
+ int MaxNumExperts, int MaxNumTopExperts>
1139
+__global__ void __launch_bounds__(WarpTopKLaunchConfig<MaxNumExperts>::BlockDim)
1140
+ single_group_topk_warp_kernel(T const* scores, float* topk_values,
1141
+ IdxT* topk_indices, BiasT const* bias,
1142
+ int64_t num_tokens, int64_t num_experts,
1143
+ int64_t topk, bool renormalize,
1144
+ float routed_scaling_factor,
1145
+ bool enable_pdl) {
1146
+ static constexpr int NumChunks = (MaxNumExperts + WARP_SIZE - 1) / WARP_SIZE;
1147
+ static constexpr int WarpBlockDim =
1148
+ WarpTopKLaunchConfig<MaxNumExperts>::BlockDim;
1149
+ using LaneOwnedRange =
1150
+ reduce_topk::HighExpertLaneOwnedTopKRange<MaxNumExperts,
1151
+ MaxNumTopExperts>;
1152
+
1153
+ static_assert(NumChunks <= 64);
1154
+ static_assert(MaxNumTopExperts <= WARP_SIZE);
1155
+
1156
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
1157
+ if (enable_pdl) {
1158
+ cudaGridDependencySynchronize();
1159
+ }
1160
+#endif
1161
+
1162
+ int32_t const lane = static_cast<int32_t>(threadIdx.x) % WARP_SIZE;
1163
+ int32_t const warp_id = static_cast<int32_t>(threadIdx.x) / WARP_SIZE;
1164
+ int32_t const global_warp =
1165
+ static_cast<int32_t>(blockIdx.x) * WarpBlockDim / WARP_SIZE + warp_id;
1166
+ int32_t const global_warp_stride =
1167
+ static_cast<int32_t>(gridDim.x) * WarpBlockDim / WARP_SIZE;
1168
+ int32_t const num_experts_i32 = static_cast<int32_t>(num_experts);
1169
+ int32_t const topk_i32 = static_cast<int32_t>(topk);
1170
+ auto warp = cg::tiled_partition<WARP_SIZE>(cg::this_thread_block());
1171
+
1172
+ for (int32_t token = global_warp; token < num_tokens;
1173
+ token += global_warp_stride) {
1174
+ T const* token_scores = scores + int64_t{token} * num_experts;
1175
+ float local_scores[NumChunks];
1176
+ int32_t local_indices[NumChunks];
1177
+#pragma unroll
1178
+ for (int index = 0; index < NumChunks; ++index) {
1179
+ int32_t const expert = index * WARP_SIZE + lane;
1180
+ float unbiased;
1181
+ float selection;
1182
+ if (expert < num_experts_i32) {
1183
+ preprocess_score<T, BiasT, SF>(token_scores[expert], bias[expert],
1184
+ unbiased, selection);
1185
+ } else {
1186
+ selection = InvalidScore;
1187
+ }
1188
+ local_scores[index] = selection;
1189
+ local_indices[index] = expert;
1190
+ }
1191
+
1192
+ float lane_score;
1193
+ int32_t lane_expert;
1194
+ if constexpr (LaneOwnedRange::kEnabled) {
1195
+ reduce_topk::reduceTopKForLane<MaxNumTopExperts>(
1196
+ warp, lane_score, lane_expert, local_scores, local_indices,
1197
+ InvalidScore, lane);
1198
+ } else {
1199
+ float top_scores[MaxNumTopExperts];
1200
+ int32_t top_experts[MaxNumTopExperts];
1201
+ reduce_topk::reduceTopK(warp, top_scores, top_experts, local_scores,
1202
+ local_indices, InvalidScore, topk_i32);
1203
+ lane_score = lane < topk_i32 ? top_scores[lane] : InvalidScore;
1204
+ lane_expert = lane < topk_i32 ? top_experts[lane] : -1;
1205
+ }
1206
+
1207
+ float lane_unbiased = 0.0F;
1208
+ if (lane < topk_i32 && lane_expert >= 0 && lane_expert < num_experts_i32) {
1209
+ lane_unbiased = lane_score - cuda_cast<float, BiasT>(bias[lane_expert]);
1210
+ }
1211
+ write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token,
1212
+ topk_i32, topk_values, topk_indices, renormalize,
1213
+ routed_scaling_factor);
1214
+ }
1215
+
1216
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
1217
+ if (enable_pdl) {
1218
+ cudaTriggerProgrammaticLaunchCompletion();
1219
+ }
1220
+#endif
1221
+}
1222
+
1223
+template <int Experts, int TopK>
1224
+struct Tier {
1225
+ static constexpr int kExperts = Experts;
1226
+ static constexpr int kTopK = TopK;
1227
+};
1228
+
1229
+template <typename... Tiers>
1230
+struct TierList {};
1231
+
1232
+using SigmoidBiasTiers =
1233
+ TierList<Tier<128, 8>, Tier<256, 8>, Tier<384, 8>, Tier<512, 8>,
1234
+ Tier<512, 22>, Tier<768, 16>, Tier<896, 16>, Tier<1024, 16>>;
1235
+
1236
+using PrecomputedSoftmaxBiasTiers =
1237
+ TierList<Tier<128, 4>, Tier<128, 8>, Tier<160, 8>, Tier<256, 8>,
1238
+ Tier<256, 16>, Tier<512, 8>, Tier<512, 16>, Tier<512, 22>,
1239
+ Tier<512, 32>, Tier<576, 8>, Tier<768, 16>, Tier<896, 16>,
1240
+ Tier<1024, 16>>;
1241
+
1242
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
1243
+ int MaxNumExperts, int MaxNumTopExperts>
1244
+void launch(T* scores, float* topk_values, IdxT* topk_indices,
1245
+ BiasT const* bias, int64_t num_tokens, int64_t num_experts,
1246
+ int64_t topk, bool renormalize, double routed_scaling_factor,
1247
+ bool enable_pdl, cudaLaunchConfig_t& config) {
1248
+ config.dynamicSmemBytes = 0;
1249
+ bool const use_block_kernel =
1250
+ UseTunedBlockPath<MaxNumExperts, MaxNumTopExperts> ||
1251
+ MaxNumExperts > 1024 || num_experts >= 1024 ||
1252
+ (num_experts >= 256 && num_tokens <= 1024);
1253
+ if (use_block_kernel) {
1254
+ config.gridDim = static_cast<uint32_t>(num_tokens);
1255
+ config.blockDim = BlockDim;
1256
+ cudaLaunchKernelEx(
1257
+ &config,
1258
+ &single_group_topk_block_kernel<T, BiasT, IdxT, SF, MaxNumExperts,
1259
+ MaxNumTopExperts>,
1260
+ scores, topk_values, topk_indices, bias, num_experts, topk, renormalize,
1261
+ static_cast<float>(routed_scaling_factor), enable_pdl);
1262
+ } else {
1263
+ using WarpConfig = WarpTopKLaunchConfig<MaxNumExperts>;
1264
+ config.gridDim = WarpConfig::grid_dim(num_tokens);
1265
+ config.blockDim = WarpConfig::BlockDim;
1266
+ cudaLaunchKernelEx(
1267
+ &config,
1268
+ &single_group_topk_warp_kernel<T, BiasT, IdxT, SF, MaxNumExperts,
1269
+ MaxNumTopExperts>,
1270
+ scores, topk_values, topk_indices, bias, num_tokens, num_experts, topk,
1271
+ renormalize, static_cast<float>(routed_scaling_factor), enable_pdl);
1272
+ }
1273
+}
1274
+
1275
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
1276
+bool dispatch(TierList<>*, T*, float*, IdxT*, BiasT const*, int64_t, int64_t,
1277
+ int64_t, bool, double, bool, cudaLaunchConfig_t&) {
1278
+ return false;
1279
+}
1280
+
1281
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,
1282
+ typename First, typename... Rest>
1283
+bool dispatch(TierList<First, Rest...>*, T* scores, float* topk_values,
1284
+ IdxT* topk_indices, BiasT const* bias, int64_t num_tokens,
1285
+ int64_t num_experts, int64_t topk, bool renormalize,
1286
+ double routed_scaling_factor, bool enable_pdl,
1287
+ cudaLaunchConfig_t& config) {
1288
+ if (num_experts <= First::kExperts && topk <= First::kTopK) {
1289
+ launch<T, BiasT, IdxT, SF, First::kExperts, First::kTopK>(
1290
+ scores, topk_values, topk_indices, bias, num_tokens, num_experts, topk,
1291
+ renormalize, routed_scaling_factor, enable_pdl, config);
1292
+ return true;
1293
+ }
1294
+ return dispatch<T, BiasT, IdxT, SF>(
1295
+ static_cast<TierList<Rest...>*>(nullptr), scores, topk_values,
1296
+ topk_indices, bias, num_tokens, num_experts, topk, renormalize,
1297
+ routed_scaling_factor, enable_pdl, config);
1298
+}
1299
+
1300
+} // namespace detail
1301
+
1302
+template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
1303
+bool invoke(T* scores, float* topk_values, IdxT* topk_indices,
1304
+ BiasT const* bias, int64_t num_tokens, int64_t num_experts,
1305
+ int64_t topk, bool renormalize, double routed_scaling_factor,
1306
+ bool enable_pdl, cudaLaunchConfig_t& config) {
1307
+ static_assert(SF == SCORING_NONE || SF == SCORING_SIGMOID);
1308
+ if constexpr (SF == SCORING_SIGMOID) {
1309
+ return detail::dispatch<T, BiasT, IdxT, SF>(
1310
+ static_cast<detail::SigmoidBiasTiers*>(nullptr), scores, topk_values,
1311
+ topk_indices, bias, num_tokens, num_experts, topk, renormalize,
1312
+ routed_scaling_factor, enable_pdl, config);
1313
+ } else {
1314
+ return detail::dispatch<T, BiasT, IdxT, SF>(
1315
+ static_cast<detail::PrecomputedSoftmaxBiasTiers*>(nullptr), scores,
1316
+ topk_values, topk_indices, bias, num_tokens, num_experts, topk,
1317
+ renormalize, routed_scaling_factor, enable_pdl, config);
1318
+ }
1319
+}
1320
+
1321
+} // namespace single_group_topk
1322
+
1323
template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>
1324
void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
1325
BiasT const* bias, int64_t const num_tokens,
@@ -905,6 +1335,12 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
1335
attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl;
1336
config.numAttrs = 1;
1337
config.attrs = attrs;
1338
+ if (n_group == 1 && topk_group == 1 &&
1339
+ single_group_topk::invoke<T, BiasT, IdxT, SF>(
1340
+ scores, topk_values, topk_indices, bias, num_tokens, num_experts,
1341
+ topk, renormalize, routed_scaling_factor, enable_pdl, config)) {
1342
+ return;
1343
+ }
1344
1345
// Check if we can use the optimized
1346
// grouped_topk_fused_small_expert_count_kernel
csrc/libtorch_stable/moe/moeTopKFuncs.cuh
+210
-99
@@ -1,6 +1,7 @@
1
/*
2
* Adapted from
3
* https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh
4
+ * https://github.com/flashinfer-ai/flashinfer/blob/06400d062a2d51564bbe781f6f811d0b75ca593e/include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh
5
* Copyright (c) 2026, The vLLM team.
6
* SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION. All rights
7
* reserved. SPDX-License-Identifier: Apache-2.0
@@ -23,6 +24,9 @@
24
#include <cooperative_groups/reduce.h>
25
#include <cub/cub.cuh>
26
27
+#include <cstdint>
28
+#include <type_traits>
29
+
30
namespace vllm {
31
namespace moe {
32
namespace reduce_topk {
@@ -38,11 +42,10 @@ struct TopKRedType {
42
"Top K reduction only implemented for int, float, float16 and bfloat16");
43
44
using TypeCmp = std::conditional_t<sizeof(T) == 4, uint64_t, uint32_t>;
41
- using IdxT = std::conditional_t<sizeof(T) == 4, int32_t, int16_t>;
45
46
static constexpr int kMoveBits = (sizeof(T) == 4) ? 32 : 16;
47
static constexpr int kMaxIdx = 65535;
45
- TypeCmp compValIdx;
48
+ TypeCmp compVal;
49
50
static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) {
51
auto valueBits = cub::Traits<T>::TwiddleIn(
@@ -69,69 +72,175 @@ struct TopKRedType {
72
__host__ __device__ TopKRedType() = default;
73
74
__host__ __device__ TopKRedType(T val, int32_t idx)
72
- : compValIdx(makeCmpVal(val, idx)) {}
75
+ : compVal(makeCmpVal(val, idx)) {}
76
74
- __host__ __device__ operator TypeCmp() const noexcept { return compValIdx; }
77
+ __host__ __device__ operator TypeCmp() const noexcept { return compVal; }
78
79
__device__ inline TypeCmp reduce(
80
cg::thread_block_tile<kWARP_SIZE> const& warp) {
78
- return cg::reduce(warp, compValIdx, cg::greater<TypeCmp>{});
81
+#ifdef __CUDA_ARCH__
82
+ static constexpr bool kHAS_FAST_REDUX = (__CUDA_ARCH__ / 100) >= 10;
83
+#else
84
+ static constexpr bool kHAS_FAST_REDUX = false;
85
+#endif
86
+ if constexpr (!kHAS_FAST_REDUX) {
87
+ return cg::reduce(warp, compVal, cg::greater<TypeCmp>{});
88
+ } else if constexpr (sizeof(TypeCmp) == 8) {
89
+ uint32_t hi = static_cast<uint32_t>(compVal >> 32);
90
+ uint32_t lo = static_cast<uint32_t>(compVal & 0xffffffffu);
91
+ uint32_t maxHi;
92
+ asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n"
93
+ : "=r"(maxHi)
94
+ : "r"(hi));
95
+ uint32_t loContrib = hi == maxHi ? lo : 0u;
96
+ uint32_t maxLo;
97
+ asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n"
98
+ : "=r"(maxLo)
99
+ : "r"(loContrib));
100
+ return (static_cast<TypeCmp>(maxHi) << 32) | static_cast<TypeCmp>(maxLo);
101
+ } else {
102
+ TypeCmp result;
103
+ asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n"
104
+ : "=r"(result)
105
+ : "r"(compVal));
106
+ return result;
107
+ }
108
}
109
};
110
82
-////////////////////////////////////////////////////////////////////////////////////////////////////
83
-
84
-template <int K_, bool Enable_>
85
-struct TopKIdx {
86
- // by default, empty
111
+template <int N>
112
+struct IsPowerOf2 {
113
+ static constexpr bool value = N > 0 && (N & (N - 1)) == 0;
114
};
115
89
-template <int K_>
90
-struct TopKIdx<K_, true> {
91
- static constexpr int K = K_;
92
- int32_t val[K];
116
+template <int N>
117
+struct NextPow2 {
118
+ private:
119
+ static constexpr unsigned u = static_cast<unsigned>(N - 1);
120
+ static constexpr unsigned s1 = u | (u >> 1);
121
+ static constexpr unsigned s2 = s1 | (s1 >> 2);
122
+ static constexpr unsigned s3 = s2 | (s2 >> 4);
123
+ static constexpr unsigned s4 = s3 | (s3 >> 8);
124
+ static constexpr unsigned s5 = s4 | (s4 >> 16);
125
+
126
+ public:
127
+ static constexpr int value = N <= 1 ? 1 : static_cast<int>(s5 + 1);
128
};
129
95
-////////////////////////////////////////////////////////////////////////////////////////////////////
130
+template <int A, int B, int Size, typename T>
131
+__device__ __forceinline__ void topkCompareSwap(T* a) {
132
+ if constexpr (A < Size && B < Size) {
133
+ if (a[A] < a[B]) {
134
+ T tmp = a[A];
135
+ a[A] = a[B];
136
+ a[B] = tmp;
137
+ }
138
+ } else {
139
+ (void)a;
140
+ }
141
+}
142
97
-#define TOPK_SWAP(I, J) \
98
- { \
99
- auto pairMin = min(topK[I].compValIdx, topK[J].compValIdx); \
100
- auto pairMax = max(topK[I].compValIdx, topK[J].compValIdx); \
101
- topK[I].compValIdx = pairMax; \
102
- topK[J].compValIdx = pairMin; \
143
+template <int I, int End, int Step, int PairStride, int Size, typename T>
144
+__device__ __forceinline__ void topkMergePairs(T* a) {
145
+ if constexpr (I + Step < End) {
146
+ topkCompareSwap<I, I + Step, Size, T>(a);
147
+ topkMergePairs<I + PairStride, End, Step, PairStride, Size, T>(a);
148
+ } else {
149
+ (void)a;
150
+ }
151
+}
152
+
153
+template <int Lo, int N, int R, int Size, typename T>
154
+__device__ __forceinline__ void topkOEM(T* a) {
155
+ constexpr int M = R * 2;
156
+ if constexpr (M < N) {
157
+ topkOEM<Lo, N, M, Size, T>(a);
158
+ topkOEM<Lo + R, N - R, M, Size, T>(a);
159
+ topkMergePairs<Lo + R, Lo + N, R, M, Size, T>(a);
160
+ } else if constexpr (R < N) {
161
+ topkCompareSwap<Lo, Lo + R, Size, T>(a);
162
+ } else {
163
+ (void)a;
164
+ }
165
+}
166
+
167
+template <int Lo, int N, int Size, typename T>
168
+__device__ __forceinline__ void topkSortBatcher(T* a) {
169
+ if constexpr (N > 1) {
170
+ constexpr int Half = N / 2;
171
+ topkSortBatcher<Lo, Half, Size, T>(a);
172
+ topkSortBatcher<Lo + Half, N - Half, Size, T>(a);
173
+ topkOEM<Lo, N, 1, Size, T>(a);
174
+ } else {
175
+ (void)a;
176
}
177
+}
178
179
template <int N, typename RedType>
106
-struct Sort;
180
+struct Sort {
181
+ static_assert(N > 0 && N <= 64, "Sort only supports N in range [1, 64]");
182
+
183
+ static __device__ void run(RedType* topK) {
184
+ if constexpr (IsPowerOf2<N>::value) {
185
+#pragma unroll
186
+ for (int k = 2; k <= N; k *= 2) {
187
+#pragma unroll
188
+ for (int j = k / 2; j > 0; j /= 2) {
189
+#pragma unroll
190
+ for (int i = 0; i < N; ++i) {
191
+ int ixj = i ^ j;
192
+ if (ixj > i) {
193
+ if ((i & k) == 0) {
194
+ if (topK[i].compVal < topK[ixj].compVal) {
195
+ auto tmp = topK[i].compVal;
196
+ topK[i].compVal = topK[ixj].compVal;
197
+ topK[ixj].compVal = tmp;
198
+ }
199
+ } else {
200
+ if (topK[i].compVal > topK[ixj].compVal) {
201
+ auto tmp = topK[i].compVal;
202
+ topK[i].compVal = topK[ixj].compVal;
203
+ topK[ixj].compVal = tmp;
204
+ }
205
+ }
206
+ }
207
+ }
208
+ }
209
+ }
210
+ } else {
211
+ constexpr int P = NextPow2<N>::value;
212
+ topkSortBatcher<0, P, N, RedType>(topK);
213
+ }
214
+ }
215
+};
216
217
template <typename RedType>
218
struct Sort<1, RedType> {
110
- static __device__ void run(RedType* topK) {}
219
+ static __device__ void run(RedType*) {}
220
};
221
222
template <typename RedType>
223
struct Sort<2, RedType> {
115
- static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); }
224
+ static __device__ void run(RedType* topK) { topkCompareSwap<0, 1, 2>(topK); }
225
};
226
227
template <typename RedType>
228
struct Sort<3, RedType> {
229
static __device__ void run(RedType* topK) {
121
- TOPK_SWAP(0, 1);
122
- TOPK_SWAP(1, 2);
123
- TOPK_SWAP(0, 1);
230
+ topkCompareSwap<0, 1, 3>(topK);
231
+ topkCompareSwap<1, 2, 3>(topK);
232
+ topkCompareSwap<0, 1, 3>(topK);
233
}
234
};
235
236
template <typename RedType>
237
struct Sort<4, RedType> {
238
static __device__ void run(RedType* topK) {
130
- TOPK_SWAP(0, 2);
131
- TOPK_SWAP(1, 3);
132
- TOPK_SWAP(0, 1);
133
- TOPK_SWAP(2, 3);
134
- TOPK_SWAP(1, 2);
239
+ topkCompareSwap<0, 2, 4>(topK);
240
+ topkCompareSwap<1, 3, 4>(topK);
241
+ topkCompareSwap<0, 1, 4>(topK);
242
+ topkCompareSwap<2, 3, 4>(topK);
243
+ topkCompareSwap<1, 2, 4>(topK);
244
}
245
};
246
@@ -147,24 +256,23 @@ __forceinline__ __device__ void reduceTopK(
256
typename RedType::TypeCmp packedMax{};
257
#pragma unroll
258
for (int kk = 0; kk < actualK; ++kk) {
150
- topK =
151
- kk > 0 && packedMax == topK.compValIdx ? RedType{minValue, idx} : topK;
152
- // get the next largest value
259
+ topK = kk > 0 && packedMax == topK.compVal ? RedType{minValue, idx} : topK;
260
packedMax = topK.reduce(warp);
261
RedType::unpack(out[kk], outIdx[kk], packedMax);
262
}
263
};
264
158
-template <int K, typename Type, int N, bool IsSorted = false>
159
-__device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
160
- Type (&out)[K], int32_t (&outIdx)[K],
161
- Type (&value)[N], int32_t (&idx)[N],
162
- Type minValue, int actualK = K) {
265
+template <int K, typename Type, int N>
266
+__forceinline__ __device__ void reduceTopK(
267
+ cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K],
268
+ int32_t (&outIdx)[K], Type (&value)[N], int32_t (&idx)[N],
269
+ Type const minValue, int actualK = K) {
270
static_assert(K > 0, "Top K must have K > 0");
164
- static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
271
+ static_assert(K <= kWARP_SIZE, "Top K must have K <= kWARP_SIZE");
272
static_assert(N > 0, "Top K must have N > 0");
166
- static_assert(N < 5,
167
- "Only support candidates number less than or equal to 128");
273
+ static_assert(N <= 64,
274
+ "Only support candidates number less than or equal to "
275
+ "64*32=2048");
276
using RedType = TopKRedType<Type>;
277
RedType topK[N];
278
#pragma unroll
@@ -172,85 +280,88 @@ __device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
280
topK[nn] = RedType{value[nn], idx[nn]};
281
}
282
175
- if constexpr (!IsSorted) {
176
- Sort<N, RedType>::run(topK);
177
- }
283
+ Sort<N, RedType>::run(topK);
284
+
285
typename RedType::TypeCmp packedMax{};
179
-#pragma unroll
286
for (int kk = 0; kk < actualK; ++kk) {
181
- bool update = kk > 0 && packedMax == topK[0].compValIdx;
287
+ bool update = kk > 0 && packedMax == topK[0].compVal;
288
#pragma unroll
289
for (int nn = 0; nn < N; ++nn) {
290
topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]}
291
: update ? topK[nn + 1]
292
: topK[nn];
293
}
188
- // get the next largest value
294
packedMax = topK[0].reduce(warp);
295
RedType::unpack(out[kk], outIdx[kk], packedMax);
296
}
297
};
298
299
+template <int NumExperts, int NumTopExperts, int MinExperts, int MaxExperts,
300
+ int MinTopExperts, int MaxTopExperts>
301
+struct LaneOwnedTopKRange {
302
+ static_assert(MinExperts > 0 && MinExperts <= MaxExperts);
303
+ static_assert(MinTopExperts > 0 && MinTopExperts <= MaxTopExperts);
304
+ static constexpr bool kEnabled =
305
+ NumExperts >= MinExperts && NumExperts <= MaxExperts &&
306
+ NumTopExperts >= MinTopExperts && NumTopExperts <= MaxTopExperts;
307
+};
308
+
309
+static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_EXPERTS = 512;
310
+static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_EXPERTS = 1024;
311
+static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_TOP_EXPERTS = 9;
312
+static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_TOP_EXPERTS = 16;
313
+
314
+template <int NumExperts, int NumTopExperts>
315
+using HighExpertLaneOwnedTopKRange =
316
+ LaneOwnedTopKRange<NumExperts, NumTopExperts,
317
+ kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_EXPERTS,
318
+ kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_EXPERTS,
319
+ kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_TOP_EXPERTS,
320
+ kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_TOP_EXPERTS>;
321
+
322
template <int K, typename Type, int N>
195
-__forceinline__ __device__ void reduceTopK(
196
- cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K],
197
- int32_t (&outIdx)[K], Type (&value)[N], int32_t (&idx)[N],
198
- Type const minValue, int actualK = K) {
323
+__forceinline__ __device__ void reduceTopKForLane(
324
+ cg::thread_block_tile<kWARP_SIZE> const& warp, Type& out, int32_t& outIdx,
325
+ Type (&value)[N], int32_t (&idx)[N], Type const minValue, int32_t laneIdx) {
326
static_assert(K > 0, "Top K must have K > 0");
200
- static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
327
+ static_assert(K <= kWARP_SIZE, "Top K must have K <= kWARP_SIZE");
328
static_assert(N > 0, "Top K must have N > 0");
202
- static_assert(
203
- N <= 16,
204
- "Only support candidates number less than or equal to 16*32=512");
205
- static_assert(N <= 4 || N % 4 == 0,
206
- "Only support candidates number is a multiple of 4*32=128 or "
207
- "less than or equal to 4");
329
+ static_assert(N <= 64,
330
+ "Only support candidates number less than or equal to "
331
+ "64*32=2048");
332
using RedType = TopKRedType<Type>;
333
+ RedType topK[N];
334
+#pragma unroll
335
+ for (int nn = 0; nn < N; ++nn) {
336
+ topK[nn] = RedType{value[nn], idx[nn]};
337
+ }
338
210
- if constexpr (N <= 4) {
211
- reduceTopKFunc<K, Type, N>(warp, out, outIdx, value, idx, minValue,
212
- actualK);
213
- } else {
214
- constexpr int numLoops = N / 4;
215
- constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1;
216
-
217
- Type topKBufferValue[numResults];
218
- int32_t topKBufferIdx[numResults];
219
- int32_t laneIdx = threadIdx.x % kWARP_SIZE;
339
+ Sort<N, RedType>::run(topK);
340
221
- for (int ii = 0; ii < numResults; ++ii) {
222
- topKBufferValue[ii] = minValue;
223
- topKBufferIdx[ii] = ii * kWARP_SIZE - 1;
341
+ typename RedType::TypeCmp packedMax{};
342
+ typename RedType::TypeCmp lanePacked{};
343
+#pragma unroll
344
+ for (int kk = 0; kk < K; ++kk) {
345
+ bool update = kk > 0 && packedMax == topK[0].compVal;
346
+#pragma unroll
347
+ for (int nn = 0; nn < N; ++nn) {
348
+ topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]}
349
+ : update ? topK[nn + 1]
350
+ : topK[nn];
351
}
225
- for (int loop = 0; loop < numLoops; ++loop) {
226
- int start = loop * 4;
227
- Type topKValue[K];
228
- int32_t topKIdx[K];
229
- Type inValue[4];
230
- int32_t inIdx[4];
231
- for (int i = 0; i < 4; ++i) {
232
- inValue[i] = value[start + i];
233
- inIdx[i] = idx[start + i];
234
- }
235
- reduceTopKFunc<K, Type, 4>(warp, topKValue, topKIdx, inValue, inIdx,
236
- minValue, actualK);
237
- int inOffset = laneIdx % K;
238
- if (laneIdx >= loop * K && laneIdx < (loop + 1) * K) {
239
- topKBufferValue[0] = topKValue[inOffset];
240
- topKBufferIdx[0] = topKIdx[inOffset];
241
- }
242
- if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE))) {
243
- topKBufferValue[1] = topKValue[inOffset];
244
- topKBufferIdx[1] = topKIdx[inOffset];
245
- }
352
+ packedMax = topK[0].reduce(warp);
353
+ if (laneIdx == kk) {
354
+ lanePacked = packedMax;
355
}
247
-
248
- reduceTopKFunc<K, Type, numResults>(warp, out, outIdx, topKBufferValue,
249
- topKBufferIdx, minValue, actualK);
356
}
251
-};
357
253
-#undef TOPK_SWAP
358
+ if (laneIdx < K) {
359
+ RedType::unpack(out, outIdx, lanePacked);
360
+ } else {
361
+ out = minValue;
362
+ outIdx = -1;
363
+ }
364
+}
365
366
} // namespace reduce_topk
367
} // namespace moe
csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
+1
-1
@@ -1086,4 +1086,4 @@ void moe_lora_align_block_size(
1086
has_expert_map);
1087
}
1088
});
1089
-}
\ No newline at end of file
1089
+}
csrc/libtorch_stable/ops.h
+95
@@ -276,6 +276,61 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
276
torch::stable::Tensor const& cos_sin_cache, double eps,
277
int64_t cache_block_size);
278
279
+void fused_kimi_k3_mla_key_concat_kv_cache_insert(
280
+ torch::stable::Tensor& q, torch::stable::Tensor const& k_nope,
281
+ torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed,
282
+ torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache,
283
+ torch::stable::Tensor const& slot_mapping, int64_t cache_block_size,
284
+ std::optional<torch::stable::Tensor> position_ids,
285
+ std::optional<torch::stable::Tensor> cos_sin_cache);
286
+
287
+void fused_kimi_k3_mla_key_concat_ds_mla_insert(
288
+ torch::stable::Tensor& q, torch::stable::Tensor const& k_nope,
289
+ torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed,
290
+ torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache,
291
+ torch::stable::Tensor const& slot_mapping, int64_t cache_block_size,
292
+ std::optional<torch::stable::Tensor> position_ids,
293
+ std::optional<torch::stable::Tensor> cos_sin_cache);
294
+
295
+void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(
296
+ torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope,
297
+ torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed,
298
+ torch::stable::Tensor const& v, torch::stable::Tensor& q_fp8,
299
+ torch::stable::Tensor& k_fp8, torch::stable::Tensor& v_fp8,
300
+ torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
301
+ torch::stable::Tensor const& q_scale_inv,
302
+ torch::stable::Tensor const& k_scale_inv,
303
+ torch::stable::Tensor const& v_scale_inv,
304
+ torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size,
305
+ std::optional<torch::stable::Tensor> position_ids,
306
+ std::optional<torch::stable::Tensor> cos_sin_cache);
307
+
308
+void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(
309
+ torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe,
310
+ torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe,
311
+ torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache,
312
+ torch::stable::Tensor const& slot_mapping, int64_t cache_block_size,
313
+ std::optional<torch::stable::Tensor> position_ids,
314
+ std::optional<torch::stable::Tensor> cos_sin_cache);
315
+
316
+void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(
317
+ torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe,
318
+ torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe,
319
+ torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache,
320
+ torch::stable::Tensor const& slot_mapping,
321
+ torch::stable::Tensor const& q_scale_inv,
322
+ torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size,
323
+ std::optional<torch::stable::Tensor> position_ids,
324
+ std::optional<torch::stable::Tensor> cos_sin_cache);
325
+
326
+void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(
327
+ torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe,
328
+ torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe,
329
+ torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache,
330
+ torch::stable::Tensor const& slot_mapping, int64_t cache_block_size,
331
+ std::optional<torch::stable::Tensor> position_ids,
332
+ std::optional<torch::stable::Tensor> cos_sin_cache);
333
+
334
void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
335
torch::stable::Tensor const& q, torch::stable::Tensor const& kv,
336
torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache,
@@ -315,6 +370,19 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
370
std::optional<torch::stable::Tensor> index_q_out,
371
const std::string& kv_cache_dtype, bool skip_index_branch);
372
373
+#ifdef VLLM_ENABLE_FUSED_KDA_DECODE
374
+void fused_kda_decode(
375
+ torch::stable::Tensor const& x, torch::stable::Tensor const& weight,
376
+ std::optional<torch::stable::Tensor> bias,
377
+ torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g,
378
+ torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log,
379
+ torch::stable::Tensor const& dt_bias,
380
+ torch::stable::Tensor const& state_indices, torch::stable::Tensor& state,
381
+ torch::stable::Tensor& out, std::optional<double> lower_bound,
382
+ std::optional<torch::stable::Tensor> output_gate,
383
+ std::optional<torch::stable::Tensor> norm_weight, double norm_eps);
384
+#endif
385
+
386
#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
387
void kimi_k3_attn_res(torch::stable::Tensor& prefix,
388
torch::stable::Tensor const& delta,
@@ -383,6 +451,20 @@ fptr_t init_custom_ar(const std::vector<int64_t>& fake_ipc_ptrs,
451
void all_reduce(fptr_t _fa, torch::stable::Tensor& inp,
452
torch::stable::Tensor& out, fptr_t reg_buffer,
453
int64_t reg_buffer_sz_bytes);
454
+void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp,
455
+ torch::stable::Tensor& out, fptr_t reg_buffer,
456
+ int64_t reg_buffer_sz_bytes);
457
+void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp,
458
+ torch::stable::Tensor& out, fptr_t local_buffer,
459
+ fptr_t multicast_buffer, fptr_t epoch_buffer,
460
+ int64_t stage_sz_bytes);
461
+void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
462
+ torch::stable::Tensor& out, fptr_t reg_buffer,
463
+ int64_t reg_buffer_sz_bytes);
464
+void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
465
+ torch::stable::Tensor& out,
466
+ fptr_t local_buffer, fptr_t epoch_buffer,
467
+ int64_t stage_sz_bytes);
468
void dispose(fptr_t _fa);
469
int64_t meta_size();
470
void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);
@@ -421,6 +503,12 @@ void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
503
double threshold);
504
void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
505
double alpha = 1.702, double limit = 7.0);
506
+void situ_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
507
+ double beta = 1.0, double linear_beta = -1.0);
508
+void masked_situ_and_mul(torch::stable::Tensor& out,
509
+ torch::stable::Tensor& input,
510
+ const torch::stable::Tensor& expert_num_tokens,
511
+ double beta = 1.0, double linear_beta = -1.0);
512
void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input);
513
void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input);
514
void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input);
@@ -497,6 +585,13 @@ void concat_and_cache_mla(torch::stable::Tensor& kv_c,
585
const std::string& kv_cache_dtype,
586
torch::stable::Tensor& scale);
587
588
+void concat_and_cache_mla_grouped(torch::stable::Tensor& kv_c,
589
+ torch::stable::Tensor& k_pe,
590
+ torch::stable::Tensor& kv_cache_ptrs,
591
+ torch::stable::Tensor& slot_mapping,
592
+ int64_t block_size, int64_t block_stride,
593
+ int64_t entry_stride);
594
+
595
// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
596
void concat_and_cache_mla_rope_fused(
597
torch::stable::Tensor& positions, torch::stable::Tensor& q_pe,
csrc/libtorch_stable/torch_bindings.cpp
+91
-1
@@ -324,7 +324,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
324
// DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
325
// conditionally compiled so impl registration is in source file
326
ops.def(
327
- "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
327
+ "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b, "
328
+ "bool enable_pdl=False) -> ()");
329
330
// BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+).
331
// conditionally compiled so impl registration is in source file
@@ -447,6 +448,48 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
448
"Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, "
449
"int cache_block_size) -> ()");
450
451
+ // Kimi-K3 MLA epilogues: optional RoPE followed by concat/cache insertion.
452
+ ops.def(
453
+ "fused_kimi_k3_mla_key_concat_kv_cache_insert("
454
+ "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, "
455
+ "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, "
456
+ "int cache_block_size, Tensor? position_ids=None, "
457
+ "Tensor? cos_sin_cache=None) -> ()");
458
+ ops.def(
459
+ "fused_kimi_k3_mla_key_concat_ds_mla_insert("
460
+ "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, "
461
+ "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, "
462
+ "int cache_block_size, Tensor? position_ids=None, "
463
+ "Tensor? cos_sin_cache=None) -> ()");
464
+ ops.def(
465
+ "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert("
466
+ "Tensor q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, Tensor v, "
467
+ "Tensor! q_fp8, Tensor! k_fp8, Tensor! v_fp8, Tensor! k_cache, "
468
+ "Tensor slot_mapping, Tensor q_scale_inv, Tensor k_scale_inv, "
469
+ "Tensor v_scale_inv, Tensor cache_scale_inv, int cache_block_size, "
470
+ "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()");
471
+
472
+ // Kimi-K3 MLA decode epilogue: concat mqa_q = [ql_nope | q_pe] and insert the
473
+ // latent [kv_c_normed | k_pe] into the paged cache (bf16 / fp8 / fp8_ds_mla).
474
+ ops.def(
475
+ "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert("
476
+ "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, "
477
+ "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, "
478
+ "int cache_block_size, Tensor? position_ids=None, "
479
+ "Tensor? cos_sin_cache=None) -> ()");
480
+ ops.def(
481
+ "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert("
482
+ "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, "
483
+ "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, "
484
+ "Tensor q_scale_inv, Tensor cache_scale_inv, int cache_block_size, "
485
+ "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()");
486
+ ops.def(
487
+ "fused_kimi_k3_mla_decode_q_concat_ds_mla_insert("
488
+ "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, "
489
+ "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, "
490
+ "int cache_block_size, Tensor? position_ids=None, "
491
+ "Tensor? cos_sin_cache=None) -> ()");
492
+
493
#ifndef USE_ROCM
494
ops.def(
495
"minimax_allreduce_rms_qk("
@@ -468,6 +511,16 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
511
"int block_size, Tensor!? q_out, Tensor!? index_q_out, "
512
"str kv_cache_dtype, bool skip_index_branch=False) -> ()");
513
514
+#ifdef VLLM_ENABLE_FUSED_KDA_DECODE
515
+ ops.def(
516
+ "fused_kda_decode("
517
+ "Tensor x, Tensor weight, Tensor? bias, Tensor! conv_state, "
518
+ "Tensor raw_g, Tensor raw_beta, Tensor A_log, Tensor dt_bias, "
519
+ "Tensor state_indices, Tensor! state, Tensor! out, "
520
+ "float? lower_bound=None, Tensor? output_gate=None, "
521
+ "Tensor? norm_weight=None, float norm_eps=1e-5) -> ()");
522
+#endif
523
+
524
#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
525
ops.def(
526
"kimi_k3_attn_res("
@@ -538,6 +591,14 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
591
"limit=7.0) "
592
"-> ()");
593
594
+ // SituGLU implementation used in Kimi models.
595
+ ops.def(
596
+ "situ_and_mul(Tensor! out, Tensor input, float beta=1.0, float "
597
+ "linear_beta=-1.0) -> ()");
598
+ ops.def(
599
+ "masked_situ_and_mul(Tensor! out, Tensor input, Tensor "
600
+ "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()");
601
+
602
// GELU implementation used in GPT-2.
603
ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
604
@@ -696,11 +757,27 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
757
ops.impl(
758
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
759
TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert));
760
+ ops.impl("fused_kimi_k3_mla_key_concat_kv_cache_insert",
761
+ TORCH_BOX(&fused_kimi_k3_mla_key_concat_kv_cache_insert));
762
+ ops.impl("fused_kimi_k3_mla_key_concat_ds_mla_insert",
763
+ TORCH_BOX(&fused_kimi_k3_mla_key_concat_ds_mla_insert));
764
+ ops.impl("fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert",
765
+ TORCH_BOX(&fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert));
766
+ ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_insert",
767
+ TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_insert));
768
+ ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert",
769
+ TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert));
770
+ ops.impl("fused_kimi_k3_mla_decode_q_concat_ds_mla_insert",
771
+ TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_ds_mla_insert));
772
#ifndef USE_ROCM
773
ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk));
774
#endif
775
ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
776
TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert));
777
+#ifdef VLLM_ENABLE_FUSED_KDA_DECODE
778
+ ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode));
779
+#endif
780
+
781
#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
782
ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res));
783
#endif
@@ -726,6 +803,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
803
ops.impl("gelu_tanh_and_mul", TORCH_BOX(&gelu_tanh_and_mul));
804
ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul));
805
ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul));
806
+ ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul));
807
+ ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul));
808
ops.impl("gelu_new", TORCH_BOX(&gelu_new));
809
ops.impl("gelu_fast", TORCH_BOX(&gelu_fast));
810
ops.impl("gelu_quick", TORCH_BOX(&gelu_quick));
@@ -823,6 +902,15 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
902
" str kv_cache_dtype,"
903
" Tensor scale) -> ()");
904
905
+ // Grouped concat_and_cache_mla across all layers (bf16 only). Each
906
+ // layer's cache base pointer is read from kv_cache_ptrs.
907
+ ops.def(
908
+ "concat_and_cache_mla_grouped(Tensor kv_c, Tensor k_pe,"
909
+ " Tensor kv_cache_ptrs,"
910
+ " Tensor slot_mapping,"
911
+ " int block_size, int block_stride,"
912
+ " int entry_stride) -> ()");
913
+
914
// Rotate Q and K, then write to kv cache for MLA
915
ops.def(
916
"concat_and_cache_mla_rope_fused("
@@ -921,6 +1009,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) {
1009
ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache));
1010
ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash));
1011
ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla));
1012
+ ops.impl("concat_and_cache_mla_grouped",
1013
+ TORCH_BOX(&concat_and_cache_mla_grouped));
1014
ops.impl("concat_and_cache_mla_rope_fused",
1015
TORCH_BOX(&concat_and_cache_mla_rope_fused));
1016
ops.impl("convert_fp8", TORCH_BOX(&convert_fp8));
csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh
+31
-2
@@ -13,10 +13,18 @@ namespace vllm {
13
namespace fp8 {
14
#ifdef ENABLE_FP8
15
16
+// Unspecialized conversions are a compile error: the old passthrough
17
+// (`return x;`) silently skipped fp8 encoding for any (Tout, Tin) pair
18
+// without a specialization below (e.g. the torch stable-ABI scalar types),
19
+// corrupting quantized data with no runtime signal.
20
+template <typename>
21
+inline constexpr bool _no_conversion_specialization = false;
22
+
23
template <typename Tout, typename Tin>
24
__inline__ __device__ Tout vec_conversion(
25
const Tin& x, const __nv_fp8_interpretation_t fp8_type = __NV_E4M3) {
19
- return x;
26
+ static_assert(_no_conversion_specialization<Tin>,
27
+ "no vec_conversion specialization for this (Tout, Tin) pair");
28
}
29
30
// float -> c10::Float8_e4m3fn
@@ -301,7 +309,9 @@ __inline__ __device__ bf16_8_t vec_conversion<bf16_8_t, Float8_>(
309
template <typename Tout, typename Tin>
310
__inline__ __device__ Tout scaled_vec_conversion(
311
const Tin& x, const float scale, const __nv_fp8_interpretation_t fp8_type) {
304
- return x;
312
+ static_assert(
313
+ _no_conversion_specialization<Tin>,
314
+ "no scaled_vec_conversion specialization for this (Tout, Tin) pair");
315
}
316
317
// fp8 -> half
@@ -492,6 +502,25 @@ __inline__ __device__ uint8_t scaled_vec_conversion<uint8_t, __nv_bfloat16>(
502
__builtin_unreachable(); // Suppress missing return statement warning
503
}
504
505
+// torch stable-ABI (headeronly) scalar types delegate to the CUDA-native
506
+// conversions, so libtorch_stable kernels dispatched on c10::BFloat16 /
507
+// c10::Half quantize correctly without manual casts.
508
+template <>
509
+__inline__ __device__ uint8_t scaled_vec_conversion<uint8_t, c10::BFloat16>(
510
+ const c10::BFloat16& a, const float scale,
511
+ const __nv_fp8_interpretation_t fp8_type) {
512
+ return scaled_vec_conversion<uint8_t, __nv_bfloat16>(
513
+ reinterpret_cast<const __nv_bfloat16&>(a), scale, fp8_type);
514
+}
515
+
516
+template <>
517
+__inline__ __device__ uint8_t scaled_vec_conversion<uint8_t, c10::Half>(
518
+ const c10::Half& a, const float scale,
519
+ const __nv_fp8_interpretation_t fp8_type) {
520
+ return scaled_vec_conversion<uint8_t, uint16_t>(
521
+ reinterpret_cast<const uint16_t&>(a), scale, fp8_type);
522
+}
523
+
524
// float -> fp8
525
template <>
526
__inline__ __device__ uint8_t scaled_vec_conversion<uint8_t, float>(
setup.py
+5
@@ -783,6 +783,7 @@ class precompiled_wheel_utils:
783
"vllm/_qutlass_C.abi3.so",
784
"vllm/_flashmla_C.abi3.so",
785
"vllm/_flashmla_extension_C.abi3.so",
786
+ "vllm/_flashkda_C.abi3.so",
787
"vllm/_sparse_flashmla_C.abi3.so",
788
"vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
789
"vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
@@ -1150,6 +1151,10 @@ if _is_cuda():
1151
ext_modules.append(
1152
CMakeExtension(name="vllm._flashmla_extension_C", optional=True)
1153
)
1154
+ if USE_PRECOMPILED_EXTENSIONS or (
1155
+ CUDA_HOME and get_nvcc_cuda_version() >= Version("12.0")
1156
+ ):
1157
+ ext_modules.append(CMakeExtension(name="vllm._flashkda_C", optional=True))
1158
if envs.VLLM_USE_PRECOMPILED or (
1159
CUDA_HOME and get_nvcc_cuda_version() >= Version("12.3")
1160
):
tests/distributed/test_custom_all_reduce.py
+117
@@ -2,6 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import random
5
+from types import SimpleNamespace
6
+from unittest.mock import Mock
7
8
import pytest
9
import ray
@@ -9,6 +11,8 @@ import torch
11
import torch.distributed as dist
12
13
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa
14
+from vllm.distributed.device_communicators import custom_all_reduce
15
+from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
16
from vllm.distributed.parallel_state import get_tp_group, graph_capture
17
18
from ..utils import (
@@ -23,6 +27,59 @@ for i, v in enumerate(test_sizes):
27
test_sizes[i] -= v % 8
28
29
30
+def test_sp16_dispatches_only_to_mnnvl_lamport(
31
+ monkeypatch: pytest.MonkeyPatch,
32
+):
33
+ """SP16 uses the MNNVL Lamport kernels and rejects same-host dispatch."""
34
+ comm = CustomAllreduce.__new__(CustomAllreduce)
35
+ comm.disabled = False
36
+ comm.world_size = 16
37
+ comm.fully_connected = False
38
+ comm.mnnvl_only = True
39
+ comm._IS_CAPTURING = False
40
+ comm.max_mnnvl_all_gather_size = 2 * 1024 * 1024
41
+ comm.max_mnnvl_reduce_scatter_size = 16 * 1024 * 1024
42
+ comm.mnnvl_multicast_ptr = 1
43
+ comm.mnnvl_lamport_ag_local_ptr = 1
44
+ comm.mnnvl_lamport_ag_multicast_ptr = 1
45
+ comm.mnnvl_lamport_ag_epoch_ptr = 1
46
+ comm.mnnvl_lamport_rs_local_ptr = 1
47
+ comm.mnnvl_lamport_rs_epoch_ptr = 1
48
+ comm.mnnvl_buffer_size = 32 * 1024 * 1024
49
+ comm._ptr = 0
50
+
51
+ lamport_all_gather = Mock()
52
+ lamport_reduce_scatter = Mock()
53
+ monkeypatch.setattr(custom_all_reduce.current_platform, "is_cuda", lambda: True)
54
+ monkeypatch.setattr(
55
+ custom_all_reduce,
56
+ "ops",
57
+ SimpleNamespace(
58
+ mnnvl_lamport_all_gather=lamport_all_gather,
59
+ mnnvl_lamport_reduce_scatter=lamport_reduce_scatter,
60
+ ),
61
+ )
62
+
63
+ gathered = comm.custom_all_gather(torch.empty((8, 8), dtype=torch.bfloat16))
64
+ scattered = comm.custom_reduce_scatter(torch.empty((16, 8), dtype=torch.bfloat16))
65
+
66
+ assert gathered is not None
67
+ assert scattered is not None
68
+ lamport_all_gather.assert_called_once()
69
+ lamport_reduce_scatter.assert_called_once()
70
+ assert not comm.should_custom_ar(torch.empty(8, dtype=torch.bfloat16))
71
+ assert not comm.should_custom_all_gather(torch.empty((8, 8), dtype=torch.int32))
72
+ assert not comm.should_custom_all_gather(
73
+ torch.empty((131073, 8), dtype=torch.bfloat16)
74
+ )
75
+
76
+ comm.mnnvl_only = False
77
+ assert not comm.should_custom_all_gather(torch.empty((8, 8), dtype=torch.bfloat16))
78
+ assert not comm.should_custom_reduce_scatter(
79
+ torch.empty((16, 8), dtype=torch.bfloat16)
80
+ )
81
+
82
+
83
@ray.remote(num_gpus=1, max_calls=1)
84
def graph_allreduce(
85
monkeypatch: pytest.MonkeyPatch,
@@ -80,6 +137,32 @@ def graph_allreduce(
137
torch.testing.assert_close(out1, inp1)
138
torch.testing.assert_close(out2, inp2)
139
140
+ fa = get_tp_group().device_communicator.ca_comm
141
+ tp_rank = rank % tp_size
142
+ with graph_capture(device=device) as graph_capture_context:
143
+ local = torch.full(
144
+ (512, 4096), tp_rank + 1, dtype=torch.bfloat16, device=device
145
+ )
146
+ reduce_input = torch.full(
147
+ (512 * tp_size, 4096),
148
+ tp_rank + 1,
149
+ dtype=torch.bfloat16,
150
+ device=device,
151
+ )
152
+ graph = torch.cuda.CUDAGraph()
153
+ with torch.cuda.graph(graph, stream=graph_capture_context.stream):
154
+ gathered = fa.custom_all_gather(local)
155
+ scattered = fa.custom_reduce_scatter(reduce_input)
156
+ graph.replay()
157
+ assert gathered is not None
158
+ assert scattered is not None
159
+ expected_gather = torch.cat(
160
+ [torch.full_like(local, peer_rank + 1) for peer_rank in range(tp_size)]
161
+ )
162
+ expected_scatter = torch.full_like(local, tp_size * (tp_size + 1) // 2)
163
+ torch.testing.assert_close(gathered, expected_gather)
164
+ torch.testing.assert_close(scattered, expected_scatter)
165
+
166
167
@ray.remote(num_gpus=1, max_calls=1)
168
def eager_allreduce(
@@ -110,6 +193,29 @@ def eager_allreduce(
193
out = fa.all_reduce(out, registered=False)
194
torch.testing.assert_close(out, inp * (tp_size**num_communication))
195
196
+ group = get_tp_group().device_group
197
+ tp_rank = rank % tp_size
198
+ for dtype in [torch.float32, torch.float16, torch.bfloat16]:
199
+ local = torch.full((64, 4096), tp_rank + 1, dtype=dtype, device=device)
200
+ expected_gather = torch.empty(
201
+ (64 * tp_size, 4096), dtype=dtype, device=device
202
+ )
203
+ dist.all_gather_into_tensor(expected_gather, local, group=group)
204
+ gathered = fa.custom_all_gather(local)
205
+ assert gathered is not None
206
+ torch.testing.assert_close(gathered, expected_gather)
207
+
208
+ reduce_input = torch.full(
209
+ (64 * tp_size, 4096), tp_rank + 1, dtype=dtype, device=device
210
+ )
211
+ expected_scatter = torch.empty((64, 4096), dtype=dtype, device=device)
212
+ dist.reduce_scatter_tensor(
213
+ expected_scatter, reduce_input.clone(), group=group
214
+ )
215
+ scattered = fa.custom_reduce_scatter(reduce_input)
216
+ assert scattered is not None
217
+ torch.testing.assert_close(scattered, expected_scatter)
218
+
219
inp = torch.ones(sz * 4, dtype=torch.bfloat16, device=device)
220
out = inp
221
for _ in range(num_communication):
@@ -130,3 +236,14 @@ def test_custom_allreduce(
236
if world_size > torch.accelerator.device_count():
237
pytest.skip("Not enough GPUs to run the test.")
238
multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target)
239
+
240
+
241
+@pytest.mark.parametrize("test_target", [eager_allreduce, graph_allreduce])
242
+def test_custom_collectives_world_size_four(
243
+ monkeypatch: pytest.MonkeyPatch,
244
+ test_target,
245
+):
246
+ """Exercise the four-rank kernel specialization used by Kimi SP."""
247
+ if torch.accelerator.device_count() < 4:
248
+ pytest.skip("Not enough GPUs to run the test.")
249
+ multi_process_parallel(monkeypatch, 4, 1, test_target)
tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py
new
+221
@@ -0,0 +1,221 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""RoPE equivalence tests for the fused Kimi-K3 MLA epilogues."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import (
9
+ fused_mla_decode_q_concat_kv_cache_insert,
10
+ fused_mla_key_concat_ds_mla_insert,
11
+ fused_mla_key_concat_kv_cache_insert,
12
+ fused_mla_qkv_quant_kv_cache_fp8_insert,
13
+)
14
+from vllm.platforms import current_platform
15
+
16
+pytestmark = pytest.mark.skipif(
17
+ not current_platform.is_cuda(), reason="Kimi-K3 fused MLA requires CUDA"
18
+)
19
+
20
+_DTYPE = torch.bfloat16
21
+_NUM_TOKENS = 3
22
+_NUM_HEADS = 4
23
+_BLOCK_SIZE = 8
24
+_POSITIONS = (1, 7, 13)
25
+_SLOTS = (0, 3, 9)
26
+
27
+
28
+def _randn(*shape: int) -> torch.Tensor:
29
+ return torch.randn(*shape, device="cuda", dtype=_DTYPE) * 0.2
30
+
31
+
32
+def _rope_cache(max_position: int = 32) -> torch.Tensor:
33
+ inv_freq = 1.0 / (
34
+ 50000 ** (torch.arange(0, 64, 2, dtype=torch.float32, device="cuda") / 64)
35
+ )
36
+ positions = torch.arange(max_position, dtype=torch.float32, device="cuda")
37
+ freqs = torch.outer(positions, inv_freq)
38
+ # The fused epilogue reads the cos/sin table in fp32 (RoPE math runs in fp32).
39
+ return torch.cat((freqs.cos(), freqs.sin()), dim=-1)
40
+
41
+
42
+def _apply_gptj_rope(
43
+ x: torch.Tensor, positions: torch.Tensor, cos_sin_cache: torch.Tensor
44
+) -> torch.Tensor:
45
+ cos, sin = cos_sin_cache.index_select(0, positions).chunk(2, dim=-1)
46
+ for _ in range(x.ndim - 2):
47
+ cos = cos.unsqueeze(1)
48
+ sin = sin.unsqueeze(1)
49
+ x1 = x[..., ::2].float()
50
+ x2 = x[..., 1::2].float()
51
+ out1 = x1 * cos.float() - x2 * sin.float()
52
+ out2 = x2 * cos.float() + x1 * sin.float()
53
+ return torch.stack((out1, out2), dim=-1).flatten(-2).to(x.dtype)
54
+
55
+
56
+def _cache_rows(cache: torch.Tensor, slots: torch.Tensor) -> torch.Tensor:
57
+ return cache.reshape(-1, cache.shape[-1]).index_select(0, slots)
58
+
59
+
60
+def _assert_fp8_close(actual: torch.Tensor, expected: torch.Tensor) -> None:
61
+ torch.testing.assert_close(
62
+ actual.float(),
63
+ expected.to(torch.float8_e4m3fn).float(),
64
+ atol=0.03125,
65
+ rtol=0.15,
66
+ )
67
+
68
+
69
+@pytest.mark.parametrize("cache_kind", ["bf16", "fp8", "fp8_ds_mla"])
70
+@torch.inference_mode()
71
+def test_prefill_epilogue_fuses_gptj_rope(cache_kind: str) -> None:
72
+ torch.manual_seed(0)
73
+ positions = torch.tensor(_POSITIONS, device="cuda", dtype=torch.int64)
74
+ slots = torch.tensor(_SLOTS, device="cuda", dtype=torch.int64)
75
+ cos_sin_cache = _rope_cache()
76
+ q = _randn(_NUM_TOKENS, _NUM_HEADS, 192)
77
+ k_nope = _randn(_NUM_TOKENS, _NUM_HEADS, 128)
78
+ k_pe = _randn(_NUM_TOKENS, 64)
79
+ kv_c = _randn(_NUM_TOKENS, 512)
80
+ v = _randn(_NUM_TOKENS, _NUM_HEADS, 128)
81
+
82
+ q_expected = q.clone()
83
+ q_expected[..., 128:] = _apply_gptj_rope(
84
+ q_expected[..., 128:], positions, cos_sin_cache
85
+ )
86
+ k_pe_expected = _apply_gptj_rope(k_pe, positions, cos_sin_cache)
87
+ k_expected = torch.cat(
88
+ (k_nope, k_pe_expected[:, None, :].expand(-1, _NUM_HEADS, -1)), dim=-1
89
+ )
90
+ cache_expected = torch.cat((kv_c, k_pe_expected), dim=-1)
91
+
92
+ if cache_kind == "bf16":
93
+ cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE)
94
+ q_actual = q.clone()
95
+ k_actual = fused_mla_key_concat_kv_cache_insert(
96
+ q_actual,
97
+ k_nope,
98
+ k_pe,
99
+ kv_c,
100
+ cache,
101
+ slots,
102
+ positions,
103
+ cos_sin_cache,
104
+ )
105
+ torch.testing.assert_close(q_actual, q_expected)
106
+ torch.testing.assert_close(k_actual, k_expected)
107
+ torch.testing.assert_close(_cache_rows(cache, slots), cache_expected)
108
+ elif cache_kind == "fp8":
109
+ cache = torch.zeros(
110
+ 2, _BLOCK_SIZE, 576, device="cuda", dtype=torch.float8_e4m3fn
111
+ )
112
+ one = torch.ones(1, device="cuda", dtype=torch.float32)
113
+ q_actual, k_actual, v_actual = fused_mla_qkv_quant_kv_cache_fp8_insert(
114
+ q,
115
+ k_nope,
116
+ k_pe,
117
+ kv_c,
118
+ v,
119
+ cache,
120
+ slots,
121
+ one,
122
+ one,
123
+ one,
124
+ one,
125
+ positions,
126
+ cos_sin_cache,
127
+ )
128
+ _assert_fp8_close(q_actual, q_expected)
129
+ _assert_fp8_close(k_actual, k_expected)
130
+ _assert_fp8_close(v_actual, v)
131
+ _assert_fp8_close(_cache_rows(cache, slots), cache_expected)
132
+ else:
133
+ cache = torch.zeros(2, _BLOCK_SIZE, 656, device="cuda", dtype=torch.uint8)
134
+ q_actual = q.clone()
135
+ k_actual = fused_mla_key_concat_ds_mla_insert(
136
+ q_actual,
137
+ k_nope,
138
+ k_pe,
139
+ kv_c,
140
+ cache,
141
+ slots,
142
+ positions,
143
+ cos_sin_cache,
144
+ )
145
+ rope_cache = _cache_rows(cache, slots)[:, 528:656].view(_DTYPE)
146
+ torch.testing.assert_close(q_actual, q_expected)
147
+ torch.testing.assert_close(k_actual, k_expected)
148
+ torch.testing.assert_close(rope_cache, k_pe_expected)
149
+
150
+
151
+@pytest.mark.parametrize("cache_kind", ["bf16", "fp8", "fp8_ds_mla"])
152
+@torch.inference_mode()
153
+def test_decode_epilogue_fuses_gptj_rope(cache_kind: str) -> None:
154
+ torch.manual_seed(1)
155
+ positions = torch.tensor(_POSITIONS, device="cuda", dtype=torch.int64)
156
+ slots = torch.tensor(_SLOTS, device="cuda", dtype=torch.int64)
157
+ cos_sin_cache = _rope_cache()
158
+ ql_nope = _randn(_NUM_TOKENS, _NUM_HEADS, 512)
159
+ q_pe = _randn(_NUM_TOKENS, _NUM_HEADS, 64)
160
+ kv_c = _randn(_NUM_TOKENS, 512)
161
+ k_pe = _randn(_NUM_TOKENS, 64)
162
+
163
+ q_pe_expected = _apply_gptj_rope(q_pe, positions, cos_sin_cache)
164
+ k_pe_expected = _apply_gptj_rope(k_pe, positions, cos_sin_cache)
165
+ q_expected = torch.cat((ql_nope, q_pe_expected), dim=-1)
166
+ cache_expected = torch.cat((kv_c, k_pe_expected), dim=-1)
167
+
168
+ kwargs = {"positions": positions, "cos_sin_cache": cos_sin_cache}
169
+ if cache_kind == "bf16":
170
+ cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE)
171
+ q_actual = fused_mla_decode_q_concat_kv_cache_insert(
172
+ ql_nope, q_pe, kv_c, k_pe, cache, slots, **kwargs
173
+ )
174
+ torch.testing.assert_close(q_actual, q_expected)
175
+ torch.testing.assert_close(_cache_rows(cache, slots), cache_expected)
176
+ elif cache_kind == "fp8":
177
+ cache = torch.zeros(
178
+ 2, _BLOCK_SIZE, 576, device="cuda", dtype=torch.float8_e4m3fn
179
+ )
180
+ one = torch.ones(1, device="cuda", dtype=torch.float32)
181
+ q_actual = fused_mla_decode_q_concat_kv_cache_insert(
182
+ ql_nope,
183
+ q_pe,
184
+ kv_c,
185
+ k_pe,
186
+ cache,
187
+ slots,
188
+ q_scale_inv=one,
189
+ cache_scale_inv=one,
190
+ **kwargs,
191
+ )
192
+ _assert_fp8_close(q_actual, q_expected)
193
+ _assert_fp8_close(_cache_rows(cache, slots), cache_expected)
194
+ else:
195
+ cache = torch.zeros(2, _BLOCK_SIZE, 656, device="cuda", dtype=torch.uint8)
196
+ q_actual = fused_mla_decode_q_concat_kv_cache_insert(
197
+ ql_nope, q_pe, kv_c, k_pe, cache, slots, ds_mla=True, **kwargs
198
+ )
199
+ rope_cache = _cache_rows(cache, slots)[:, 528:656].view(_DTYPE)
200
+ torch.testing.assert_close(q_actual, q_expected)
201
+ torch.testing.assert_close(rope_cache, k_pe_expected)
202
+
203
+
204
+@torch.inference_mode()
205
+def test_decode_epilogue_preserves_nope_path() -> None:
206
+ torch.manual_seed(2)
207
+ slots = torch.tensor(_SLOTS, device="cuda", dtype=torch.int64)
208
+ ql_nope = _randn(_NUM_TOKENS, _NUM_HEADS, 512)
209
+ q_pe = _randn(_NUM_TOKENS, _NUM_HEADS, 64)
210
+ kv_c = _randn(_NUM_TOKENS, 512)
211
+ k_pe = _randn(_NUM_TOKENS, 64)
212
+ cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE)
213
+
214
+ q_actual = fused_mla_decode_q_concat_kv_cache_insert(
215
+ ql_nope, q_pe, kv_c, k_pe, cache, slots
216
+ )
217
+
218
+ torch.testing.assert_close(q_actual, torch.cat((ql_nope, q_pe), dim=-1))
219
+ torch.testing.assert_close(
220
+ _cache_rows(cache, slots), torch.cat((kv_c, k_pe), dim=-1)
221
+ )
tests/kernels/core/test_activation.py
+40
@@ -197,6 +197,46 @@ def test_silu_and_mul_with_clamp(
197
opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit))
198
199
200
+@pytest.mark.parametrize("linear_beta", [-1.0, 2.0])
201
+@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16])
202
+@torch.inference_mode()
203
+def test_masked_situ_and_mul(
204
+ default_vllm_config,
205
+ linear_beta: float,
206
+ dtype: torch.dtype,
207
+) -> None:
208
+ """Masked SITU computes valid expert rows and preserves padded zeros."""
209
+ device = CUDA_DEVICES[0]
210
+ num_experts, max_num_tokens, d = 4, 7, 512
211
+ beta = 1.5
212
+ input = torch.randn(num_experts, max_num_tokens, 2 * d, dtype=dtype, device=device)
213
+ expert_num_tokens = torch.tensor([0, 1, 4, 7], dtype=torch.int32, device=device)
214
+ output = torch.zeros(num_experts, max_num_tokens, d, dtype=dtype, device=device)
215
+
216
+ torch.ops._C.masked_situ_and_mul(
217
+ output, input, expert_num_tokens, beta, linear_beta
218
+ )
219
+
220
+ gate, up = input.float().chunk(2, dim=-1)
221
+ expected = beta * torch.tanh(gate / beta) * torch.sigmoid(gate)
222
+ if linear_beta > 0:
223
+ up = linear_beta * torch.tanh(up / linear_beta)
224
+ expected = (expected * up).to(dtype)
225
+ for expert, num_tokens in enumerate(expert_num_tokens.cpu().tolist()):
226
+ torch.testing.assert_close(
227
+ output[expert, :num_tokens],
228
+ expected[expert, :num_tokens],
229
+ atol=get_default_atol(output),
230
+ rtol=get_default_rtol(output),
231
+ )
232
+ assert torch.count_nonzero(output[expert, num_tokens:]) == 0
233
+
234
+ opcheck(
235
+ torch.ops._C.masked_situ_and_mul,
236
+ (output, input, expert_num_tokens, beta, linear_beta),
237
+ )
238
+
239
+
240
@pytest.mark.parametrize(
241
"activation",
242
[
tests/kernels/core/test_fused_q_kv_rmsnorm.py
+1
-1
@@ -13,7 +13,7 @@ from __future__ import annotations
13
import pytest
14
import torch
15
16
-from vllm.models.deepseek_v4.common.ops import fused_q_kv_rmsnorm
16
+from vllm.models.common.ops import fused_q_kv_rmsnorm
17
from vllm.platforms import current_platform
18
19
pytestmark = pytest.mark.skipif(
tests/kernels/core/test_fused_rms_norm_gated.py
+3
-1
@@ -7,7 +7,9 @@ matching the eager triton kernel output."""
7
import pytest
8
import torch
9
10
-from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
10
+from vllm.third_party.flash_linear_attention.ops.fused_norm_gate import (
11
+ FusedRMSNormGated,
12
+)
13
from vllm.utils.torch_utils import set_random_seed
14
15
DTYPES = [torch.bfloat16]
tests/kernels/moe/test_grouped_topk.py
+266
@@ -23,6 +23,53 @@ from vllm.platforms import current_platform
23
from vllm.utils.torch_utils import set_random_seed
24
25
26
+def _run_single_group_topk(
27
+ logits: torch.Tensor,
28
+ bias: torch.Tensor,
29
+ topk: int,
30
+ *,
31
+ scoring_func: str,
32
+ renormalize: bool,
33
+ routed_scaling_factor: float = 1.0,
34
+) -> tuple[torch.Tensor, torch.Tensor]:
35
+ return fused_grouped_topk(
36
+ hidden_states=torch.empty(
37
+ (logits.shape[0], 0), dtype=logits.dtype, device=logits.device
38
+ ),
39
+ gating_output=logits,
40
+ topk=topk,
41
+ renormalize=renormalize,
42
+ e_score_correction_bias=bias,
43
+ num_expert_group=1,
44
+ topk_group=1,
45
+ scoring_func=scoring_func,
46
+ routed_scaling_factor=routed_scaling_factor,
47
+ )
48
+
49
+
50
+def _single_group_reference(
51
+ logits: torch.Tensor,
52
+ bias: torch.Tensor,
53
+ topk: int,
54
+ *,
55
+ scoring_func: str,
56
+ renormalize: bool,
57
+ routed_scaling_factor: float = 1.0,
58
+) -> tuple[torch.Tensor, torch.Tensor]:
59
+ if scoring_func == "sigmoid":
60
+ scores = 0.5 * torch.tanh(0.5 * logits.float()) + 0.5
61
+ else:
62
+ scores = torch.softmax(logits, dim=-1).float()
63
+ indices = torch.argsort(
64
+ scores + bias.float(), dim=-1, descending=True, stable=True
65
+ )[:, :topk]
66
+ values = scores.gather(1, indices)
67
+ if renormalize:
68
+ values /= values.sum(dim=-1, keepdim=True) + 1e-20
69
+ values *= routed_scaling_factor
70
+ return values, indices.to(torch.int32)
71
+
72
+
73
@pytest.mark.skipif(
74
not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
75
)
@@ -101,3 +148,222 @@ def test_grouped_topk(
148
baseline_topk_weights, test_topk_weights, atol=2e-2, rtol=0
149
)
150
torch.testing.assert_close(baseline_topk_ids, test_topk_ids, atol=0, rtol=0)
151
+
152
+
153
+@pytest.mark.skipif(
154
+ not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
155
+)
156
+def test_grouped_topk_single_group_large_batch():
157
+ set_random_seed(0)
158
+ logits = torch.randn((1536, 896), dtype=torch.bfloat16, device="cuda")
159
+ bias = torch.randn((896,), dtype=torch.float32, device="cuda")
160
+
161
+ expected_values, expected_ids = _single_group_reference(
162
+ logits, bias, 16, scoring_func="sigmoid", renormalize=True
163
+ )
164
+ actual_values, actual_ids = _run_single_group_topk(
165
+ logits, bias, 16, scoring_func="sigmoid", renormalize=True
166
+ )
167
+
168
+ torch.testing.assert_close(actual_ids, expected_ids)
169
+ torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0)
170
+
171
+
172
+@pytest.mark.skipif(
173
+ not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
174
+)
175
+@pytest.mark.parametrize(
176
+ "num_experts,topk,input_dtype,bias_dtype",
177
+ [
178
+ (512, 9, torch.bfloat16, torch.float32),
179
+ (512, 16, torch.float16, torch.float16),
180
+ (513, 9, torch.float32, torch.bfloat16),
181
+ (513, 16, torch.bfloat16, torch.float32),
182
+ (895, 9, torch.float16, torch.bfloat16),
183
+ (896, 16, torch.float32, torch.float16),
184
+ (897, 9, torch.bfloat16, torch.bfloat16),
185
+ (897, 16, torch.float16, torch.float32),
186
+ (1024, 9, torch.float32, torch.bfloat16),
187
+ (1024, 16, torch.bfloat16, torch.float16),
188
+ ],
189
+)
190
+@pytest.mark.parametrize(
191
+ "scoring_func,renormalize,routed_scaling_factor",
192
+ [
193
+ ("sigmoid", True, 1.0),
194
+ ("sigmoid", False, 2.5),
195
+ ("softmax", True, 2.5),
196
+ ("softmax", False, 1.0),
197
+ ],
198
+)
199
+def test_grouped_topk_single_group_tiers(
200
+ num_experts: int,
201
+ topk: int,
202
+ input_dtype: torch.dtype,
203
+ bias_dtype: torch.dtype,
204
+ scoring_func: str,
205
+ renormalize: bool,
206
+ routed_scaling_factor: float,
207
+):
208
+ set_random_seed(7)
209
+ logits = torch.randn((17, num_experts), dtype=input_dtype, device="cuda")
210
+ bias = torch.randn((num_experts,), dtype=bias_dtype, device="cuda")
211
+
212
+ expected_values, expected_ids = _single_group_reference(
213
+ logits,
214
+ bias,
215
+ topk,
216
+ scoring_func=scoring_func,
217
+ renormalize=renormalize,
218
+ routed_scaling_factor=routed_scaling_factor,
219
+ )
220
+ actual_values, actual_ids = _run_single_group_topk(
221
+ logits,
222
+ bias,
223
+ topk,
224
+ scoring_func=scoring_func,
225
+ renormalize=renormalize,
226
+ routed_scaling_factor=routed_scaling_factor,
227
+ )
228
+
229
+ torch.testing.assert_close(actual_ids, expected_ids)
230
+ torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0)
231
+
232
+
233
+@pytest.mark.skipif(
234
+ not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
235
+)
236
+@pytest.mark.parametrize(
237
+ "num_experts,topk,scoring_func",
238
+ [
239
+ (128, 8, "sigmoid"),
240
+ (129, 8, "sigmoid"),
241
+ (257, 8, "sigmoid"),
242
+ (385, 8, "sigmoid"),
243
+ (512, 9, "sigmoid"),
244
+ (513, 9, "sigmoid"),
245
+ (769, 9, "sigmoid"),
246
+ (897, 16, "sigmoid"),
247
+ (1024, 16, "sigmoid"),
248
+ (128, 4, "softmax"),
249
+ (128, 5, "softmax"),
250
+ (129, 8, "softmax"),
251
+ (161, 8, "softmax"),
252
+ (256, 9, "softmax"),
253
+ (257, 8, "softmax"),
254
+ (512, 9, "softmax"),
255
+ (512, 17, "softmax"),
256
+ (512, 23, "softmax"),
257
+ (513, 8, "softmax"),
258
+ (577, 9, "softmax"),
259
+ (769, 9, "softmax"),
260
+ (897, 9, "softmax"),
261
+ (1024, 16, "softmax"),
262
+ ],
263
+)
264
+def test_grouped_topk_single_group_capacity_tiers(
265
+ num_experts: int,
266
+ topk: int,
267
+ scoring_func: str,
268
+):
269
+ set_random_seed(11)
270
+ logits = torch.randn((3, num_experts), dtype=torch.bfloat16, device="cuda")
271
+ bias = torch.randn((num_experts,), dtype=torch.float32, device="cuda")
272
+ expected_values, expected_ids = _single_group_reference(
273
+ logits,
274
+ bias,
275
+ topk,
276
+ scoring_func=scoring_func,
277
+ renormalize=True,
278
+ routed_scaling_factor=2.5,
279
+ )
280
+ actual_values, actual_ids = _run_single_group_topk(
281
+ logits,
282
+ bias,
283
+ topk,
284
+ scoring_func=scoring_func,
285
+ renormalize=True,
286
+ routed_scaling_factor=2.5,
287
+ )
288
+
289
+ torch.testing.assert_close(actual_ids, expected_ids)
290
+ torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0)
291
+
292
+
293
+@pytest.mark.skipif(
294
+ not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
295
+)
296
+@pytest.mark.parametrize("num_experts", [512, 896, 1024])
297
+def test_grouped_topk_single_group_stable_ties(num_experts: int):
298
+ logits = torch.zeros((1, num_experts), dtype=torch.bfloat16, device="cuda")
299
+ bias = torch.zeros((num_experts,), dtype=torch.float32, device="cuda")
300
+
301
+ actual_values, actual_ids = _run_single_group_topk(
302
+ logits,
303
+ bias,
304
+ 16,
305
+ scoring_func="sigmoid",
306
+ renormalize=True,
307
+ routed_scaling_factor=2.5,
308
+ )
309
+
310
+ expected_ids = torch.arange(16, dtype=torch.int32, device="cuda")[None]
311
+ expected_values = torch.full((1, 16), 2.5 / 16, dtype=torch.float32, device="cuda")
312
+ torch.testing.assert_close(actual_ids, expected_ids)
313
+ torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0)
314
+
315
+
316
+@pytest.mark.skipif(
317
+ not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
318
+)
319
+@pytest.mark.parametrize("num_experts", [512, 896, 1024])
320
+@pytest.mark.parametrize("num_finite", [0, 15])
321
+@pytest.mark.parametrize("renormalize", [False, True])
322
+def test_grouped_topk_single_group_nonfinite_scores(
323
+ num_experts: int, num_finite: int, renormalize: bool
324
+):
325
+ logits = torch.full(
326
+ (1, num_experts), float("nan"), dtype=torch.bfloat16, device="cuda"
327
+ )
328
+ if num_finite:
329
+ logits[0, :num_finite] = torch.arange(
330
+ num_finite, dtype=torch.bfloat16, device="cuda"
331
+ )
332
+ logits[0, num_finite] = torch.inf
333
+ logits[0, num_finite + 1] = -torch.inf
334
+ bias = torch.zeros((num_experts,), dtype=torch.float32, device="cuda")
335
+
336
+ actual_values, actual_ids = _run_single_group_topk(
337
+ logits,
338
+ bias,
339
+ 16,
340
+ scoring_func="sigmoid",
341
+ renormalize=renormalize,
342
+ routed_scaling_factor=2.5,
343
+ )
344
+
345
+ if num_finite == 0:
346
+ expected_ids = torch.arange(16, dtype=torch.int32, device="cuda")[None]
347
+ if renormalize:
348
+ expected_values = torch.full(
349
+ (1, 16), 1 / 16, dtype=torch.float32, device="cuda"
350
+ )
351
+ else:
352
+ expected_values = torch.zeros((1, 16), dtype=torch.float32, device="cuda")
353
+ else:
354
+ expected_ids = torch.cat(
355
+ (
356
+ torch.arange(num_finite - 1, -1, -1, dtype=torch.int32, device="cuda"),
357
+ torch.tensor([num_finite], dtype=torch.int32, device="cuda"),
358
+ )
359
+ )[None]
360
+ finite_values = logits[0, :num_finite].float().sigmoid().flip(0)
361
+ if renormalize:
362
+ finite_values /= finite_values.sum()
363
+ finite_values *= 2.5
364
+ expected_values = torch.cat(
365
+ (finite_values, torch.zeros(1, dtype=torch.float32, device="cuda"))
366
+ )[None]
367
+
368
+ torch.testing.assert_close(actual_ids, expected_ids)
369
+ torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0)
tests/kernels/moe/test_moe_align_block_size.py
+1
@@ -269,6 +269,7 @@ def test_moe_align_block_size_with_expert_map(
269
if (experts[k] in local_experts) or not mask_inactive_experts
270
else -1
271
)
272
+ topk_ids[0, 0] = -1
273
274
actual_sorted_ids, actual_expert_ids, actual_num_tokens = moe_align_block_size(
275
topk_ids=topk_ids,
tests/kernels/test_bf16_skinny_gemm.py
new
+649
@@ -0,0 +1,649 @@
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)."""
4
+
5
+from pathlib import Path
6
+from types import SimpleNamespace
7
+from typing import Any
8
+
9
+import pytest
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
15
+from vllm.models.kimi_k3.nvidia import low_latency_gemm as k3_gemm
16
+from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS
17
+
18
+# Keyed by local (N, K): (cute token counts, dsv3 token counts). 1536x7168 is
19
+# the unified shared_gate_up_proj/mla_g_proj entry (dsv3 M1..16).
20
+EXPECTED_SELECTIONS = {
21
+ (1536, 128): (set(), set(range(1, 17))),
22
+ (3072, 128): (set(), set(range(1, 17))),
23
+ (1536, 7168): (set(), set(range(1, 17))),
24
+ (3072, 7168): (set(range(1, 6)), set()),
25
+ (2112, 7168): (set(), set(range(1, 17))),
26
+ (2304, 1536): (set(), set(range(1, 17))),
27
+ (4608, 1536): (set(), set(range(1, 17))),
28
+ (3584, 7168): ({1}, set(range(2, 9))),
29
+ (6288, 7168): (set(range(1, 5)), set()),
30
+ (12448, 7168): (set(range(1, 4)), set()),
31
+ (7168, 768): (set(), set(range(1, 17))),
32
+ (7168, 1536): ({1}, set()),
33
+ (7168, 3072): ({1, 2}, set()),
34
+ (7168, 3584): ({1, 2}, set()),
35
+ (7168, 4224): ({1}, set()),
36
+ (7168, 8448): (set(range(1, 4)), set()),
37
+ (8448, 7168): ({1, 2}, set()),
38
+ (16896, 7168): ({1, 2}, set()),
39
+ (20480, 7168): (set(range(1, 5)), set()),
40
+ (40960, 7168): (set(range(1, 5)), set()),
41
+ # TP16.
42
+ (3216, 7168): (set(range(1, 6)), set(range(9, 16))),
43
+ (768, 7168): (set(range(1, 5)), set(range(5, 17))),
44
+ (1152, 1536): ({1}, set(range(2, 17))),
45
+ (768, 128): (set(), set(range(1, 17))),
46
+ (7168, 384): (set(), set(range(1, 9))),
47
+ (4224, 7168): (set(range(1, 4)), set(range(4, 9))),
48
+ (10240, 7168): (set(range(1, 5)), set()),
49
+}
50
+
51
+CUTE_CASES = [
52
+ (spec.n, spec.k, num_tokens)
53
+ for spec in k3_gemm.KIMI_K3_PROJECTIONS.values()
54
+ for num_tokens, _ in spec.cute_configs
55
+]
56
+
57
+RESIDUAL_CUTE_CASES = [
58
+ (spec.n, spec.k, num_tokens)
59
+ for spec in k3_gemm.KIMI_K3_PROJECTIONS.values()
60
+ for num_tokens, _ in spec.residual_configs
61
+]
62
+
63
+EXPECTED_CUTE_CONFIGS = {
64
+ (3072, 7168, 1): (224, 3, 4, 8),
65
+ (3072, 7168, 2): (128, 3, 2, 8),
66
+ (3072, 7168, 3): (128, 2, 1, 8),
67
+ (3072, 7168, 4): (64, 2, 2, 8),
68
+ (3072, 7168, 5): (128, 3, 1, 8),
69
+ (3584, 7168, 1): (224, 2, 4, 8),
70
+ (6288, 7168, 1): (224, 3, 4, 8),
71
+ (6288, 7168, 2): (64, 3, 2, 8),
72
+ (6288, 7168, 3): (32, 3, 4, 8),
73
+ (6288, 7168, 4): (128, 6, 1, 8),
74
+ (12448, 7168, 1): (224, 4, 2, 8),
75
+ (12448, 7168, 2): (64, 4, 2, 8),
76
+ (12448, 7168, 3): (64, 2, 2, 8),
77
+ (7168, 1536, 1): (96, 4, 2, 8),
78
+ (7168, 3072, 1): (96, 2, 4, 8),
79
+ (7168, 3072, 2): (32, 4, 4, 8),
80
+ (7168, 3584, 1): (224, 4, 2, 8),
81
+ (7168, 3584, 2): (64, 4, 2, 8),
82
+ (7168, 4224, 1): (96, 4, 2, 4),
83
+ (7168, 8448, 1): (32, 4, 4, 8),
84
+ (7168, 8448, 2): (96, 4, 1, 8),
85
+ (7168, 8448, 3): (96, 4, 1, 8),
86
+ (8448, 7168, 1): (224, 3, 4, 8),
87
+ (8448, 7168, 2): (32, 4, 4, 8),
88
+ (16896, 7168, 1): (224, 6, 4, 8),
89
+ (16896, 7168, 2): (32, 4, 4, 8),
90
+ (20480, 7168, 1): (224, 4, 2, 8),
91
+ (20480, 7168, 2): (64, 4, 2, 8),
92
+ (20480, 7168, 3): (64, 2, 2, 8),
93
+ (20480, 7168, 4): (64, 4, 1, 8),
94
+ (40960, 7168, 1): (128, 4, 2, 8),
95
+ (40960, 7168, 2): (64, 4, 2, 8),
96
+ (40960, 7168, 3): (64, 2, 2, 8),
97
+ (40960, 7168, 4): (64, 4, 1, 8),
98
+ # TP16.
99
+ (3216, 7168, 1): (224, 3, 4, 8),
100
+ (3216, 7168, 2): (128, 4, 2, 8),
101
+ (3216, 7168, 3): (128, 2, 1, 8),
102
+ (3216, 7168, 4): (64, 2, 2, 8),
103
+ (3216, 7168, 5): (128, 3, 1, 8),
104
+ (768, 7168, 1): (224, 2, 4, 8),
105
+ (768, 7168, 2): (224, 2, 2, 8),
106
+ (768, 7168, 3): (224, 2, 2, 8),
107
+ (768, 7168, 4): (224, 2, 2, 8),
108
+ (1152, 1536, 1): (192, 3, 4, 8),
109
+ (4224, 7168, 1): (224, 3, 4, 8),
110
+ (4224, 7168, 2): (128, 2, 1, 8),
111
+ (4224, 7168, 3): (64, 2, 2, 8),
112
+ (10240, 7168, 1): (224, 4, 2, 8),
113
+ (10240, 7168, 2): (32, 2, 4, 8),
114
+ (10240, 7168, 3): (64, 4, 1, 8),
115
+ (10240, 7168, 4): (64, 4, 1, 8),
116
+}
117
+
118
+EXPECTED_RESIDUAL_CUTE_CONFIGS = {
119
+ (7168, 3584, 1): (64, 4, 2, 8),
120
+ (7168, 3584, 2): (64, 7, 2, 8),
121
+ (7168, 3584, 3): (64, 2, 1, 8),
122
+ (7168, 3584, 4): (64, 2, 1, 8),
123
+}
124
+
125
+
126
+def _config_tuple(config) -> tuple[int, int, int, int]:
127
+ return (
128
+ config.block_size,
129
+ config.outputs_per_block,
130
+ config.k_unroll,
131
+ config.vector_width,
132
+ )
133
+
134
+
135
+def test_table_is_keyed_by_shape() -> None:
136
+ for (n, k), spec in k3_gemm.KIMI_K3_PROJECTIONS.items():
137
+ assert (spec.n, spec.k) == (n, k)
138
+
139
+
140
+def test_every_dsv3_routed_shape_is_instantiated() -> None:
141
+ """dsv3_fused_a_gemm specializes on (K, N); an unlisted shape raises.
142
+
143
+ The table routes by shape while the kernel is built per shape, so a missing
144
+ instantiation only shows up at the token counts that route to dsv3. Checking
145
+ it here needs no GPU, which is the point -- a GPU-only check is exactly what
146
+ let (3216, 7168) ship without its DISPATCH_DSV3_SHAPE(7168, 3216).
147
+ """
148
+ source = (
149
+ Path(__file__).resolve().parents[2]
150
+ / "csrc"
151
+ / "libtorch_stable"
152
+ / "dsv3_fused_a_gemm.cu"
153
+ ).read_text(encoding="utf-8")
154
+ # Benchmark-only shapes live behind VLLM_K3_BENCH_SHAPES and are not built
155
+ # by default, so they must not count as available.
156
+ production_macros = source.split("#ifdef VLLM_K3_BENCH_SHAPES")[0]
157
+ explicit = source.split("#undef DISPATCH_DSV3_SHAPE")[1].split(
158
+ "#ifdef VLLM_K3_BENCH_SHAPES"
159
+ )[0]
160
+ compiled = {
161
+ (int(hd_in), int(hd_out))
162
+ for hd_in, hd_out in re.findall(
163
+ r"DISPATCH_DSV3_SHAPE\((\d+),\s*(\d+)\)", production_macros
164
+ )
165
+ } | {
166
+ (int(hd_in), int(hd_out))
167
+ for hd_in, hd_out in re.findall(r"hd_in == (\d+) && hd_out == (\d+)", explicit)
168
+ }
169
+ assert compiled, "failed to parse the dispatch list"
170
+
171
+ missing = sorted(
172
+ (spec.n, spec.k)
173
+ for spec in KIMI_K3_PROJECTIONS.values()
174
+ if spec.dsv3_tokens and (spec.k, spec.n) not in compiled
175
+ )
176
+ assert not missing, (
177
+ f"routed to dsv3 with no instantiation: {missing}; add "
178
+ "DISPATCH_DSV3_SHAPE(K, N) for each"
179
+ )
180
+
181
+
182
+def test_packed_row_major_rejects_single_row_slice() -> None:
183
+ packed = torch.empty(1, 128)
184
+ sliced = torch.empty(1, 144)[:, :128]
185
+
186
+ assert packed.is_contiguous()
187
+ assert sliced.is_contiguous()
188
+ assert k3_gemm._is_packed_row_major(packed)
189
+ assert not k3_gemm._is_packed_row_major(sliced)
190
+
191
+
192
+def test_cute_configs_match_measured_table() -> None:
193
+ actual = {
194
+ (spec.n, spec.k, num_tokens): _config_tuple(config)
195
+ for spec in k3_gemm.KIMI_K3_PROJECTIONS.values()
196
+ for num_tokens, config in spec.cute_configs
197
+ }
198
+ assert actual == EXPECTED_CUTE_CONFIGS
199
+
200
+
201
+def test_residual_cute_configs_match_measured_table() -> None:
202
+ actual = {
203
+ (spec.n, spec.k, num_tokens): _config_tuple(config)
204
+ for spec in k3_gemm.KIMI_K3_PROJECTIONS.values()
205
+ for num_tokens, config in spec.residual_configs
206
+ }
207
+ assert actual == EXPECTED_RESIDUAL_CUTE_CONFIGS
208
+
209
+
210
+@pytest.mark.parametrize("key", EXPECTED_SELECTIONS)
211
+def test_sm103_selector_table(key: tuple[int, int]) -> None:
212
+ n, k = key
213
+ cute_tokens, dsv3_tokens = EXPECTED_SELECTIONS[key]
214
+ for num_tokens in range(1, 17):
215
+ backend = k3_gemm.select_kimi_k3_backend(num_tokens, n, k)
216
+ if num_tokens in cute_tokens:
217
+ assert backend == "cute"
218
+ elif num_tokens in dsv3_tokens:
219
+ assert backend == "dsv3_fused_a"
220
+ else:
221
+ assert backend is None
222
+
223
+
224
+@pytest.mark.parametrize("key", EXPECTED_SELECTIONS)
225
+def test_selector_requires_supported_shape_and_tokens(key: tuple[int, int]) -> None:
226
+ n, k = key
227
+ assert k3_gemm.select_kimi_k3_backend(0, n, k) is None
228
+ assert k3_gemm.select_kimi_k3_backend(17, n, k) is None
229
+ assert k3_gemm.select_kimi_k3_backend(1, n + 1, k) is None
230
+ assert k3_gemm.select_kimi_k3_backend(1, n, k + 1) is None
231
+
232
+
233
+def test_unlisted_shape_and_unselected_tokens_fall_back() -> None:
234
+ # Shape absent from the table.
235
+ assert k3_gemm.select_kimi_k3_backend(1, 1000, 1000) is None
236
+ # o_proj (7168,1536) is CuTe M1 only; M2+ falls back.
237
+ assert k3_gemm.select_kimi_k3_backend(2, 7168, 1536) is None
238
+
239
+
240
+@pytest.mark.parametrize("num_tokens", range(1, 17))
241
+def test_sm103_residual_selector_table(num_tokens: int) -> None:
242
+ backend = k3_gemm.select_kimi_k3_backend(num_tokens, 7168, 3584, has_residual=True)
243
+ assert backend == ("cute" if num_tokens <= 4 else None)
244
+
245
+
246
+def test_build_plan_matches_selector() -> None:
247
+ for spec in k3_gemm.KIMI_K3_PROJECTIONS.values():
248
+ plan = k3_gemm._build_plan(spec)
249
+ for num_tokens in range(1, 17):
250
+ backend = k3_gemm.select_kimi_k3_backend(num_tokens, spec.n, spec.k)
251
+ if backend is None:
252
+ assert num_tokens not in plan
253
+ else:
254
+ assert plan[num_tokens][0] == backend
255
+
256
+
257
+def test_installation_is_shape_specific_and_unquantized(
258
+ monkeypatch: pytest.MonkeyPatch,
259
+) -> None:
260
+ class FakeLinear(nn.Module):
261
+ def __init__(self, quant_method: object, n: int, k: int) -> None:
262
+ super().__init__()
263
+ self.quant_method = quant_method
264
+ self.weight = torch.empty(n, k)
265
+
266
+ class FakeHead(nn.Module):
267
+ def __init__(self, n: int, k: int) -> None:
268
+ super().__init__()
269
+ self.quant_method = k3_gemm.UnquantizedEmbeddingMethod()
270
+ self.weight = torch.empty(n, k)
271
+
272
+ root = nn.Module()
273
+ # dsv3-only shape (no cute warmup contribution).
274
+ root.dsv3_only = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 2304, 1536)
275
+ # quantized: must be left untouched.
276
+ quantized_method = object()
277
+ root.quantized = FakeLinear(quantized_method, 6288, 7168)
278
+ # cute shape.
279
+ root.cute = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 6288, 7168)
280
+ # cute + residual shape.
281
+ root.residual = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 7168, 3584)
282
+ # shape absent from the table: must be left untouched.
283
+ root.unlisted = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 1234, 5678)
284
+ root.lm_head = FakeHead(20480, 7168)
285
+
286
+ monkeypatch.setattr(k3_gemm, "LinearBase", FakeLinear)
287
+ monkeypatch.setattr(k3_gemm, "ParallelLMHead", FakeHead)
288
+ monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: True)
289
+ warmup_configs: set[SkinnyGemmConfig] = set()
290
+ residual_warmup_configs: set[SkinnyGemmConfig] = set()
291
+ monkeypatch.setattr(k3_gemm.shape_dynamic_skinny_gemm, "is_available", lambda: True)
292
+
293
+ def request_warmup_configs(dtype, configs, *, has_residual=False):
294
+ target = residual_warmup_configs if has_residual else warmup_configs
295
+ target.update(configs)
296
+
297
+ monkeypatch.setattr(
298
+ k3_gemm.shape_dynamic_skinny_gemm,
299
+ "request_warmup_configs",
300
+ request_warmup_configs,
301
+ )
302
+
303
+ k3_gemm.enable_kimi_k3_low_latency_gemm(root, torch.bfloat16)
304
+
305
+ assert isinstance(root.dsv3_only.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod)
306
+ assert isinstance(root.cute.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod)
307
+ assert isinstance(root.residual.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod)
308
+ assert root.quantized.quant_method is quantized_method
309
+ assert type(root.unlisted.quant_method) is k3_gemm.UnquantizedLinearMethod
310
+ assert isinstance(
311
+ root.lm_head.quant_method, k3_gemm.KimiK3LowLatencyEmbeddingMethod
312
+ )
313
+ # Warmup covers only the installed modules' local (N, K).
314
+ assert warmup_configs == {
315
+ config
316
+ for key in ((6288, 7168), (7168, 3584), (20480, 7168))
317
+ for _, config in k3_gemm.KIMI_K3_PROJECTIONS[key].cute_configs
318
+ }
319
+ assert residual_warmup_configs == {
320
+ config
321
+ for _, config in k3_gemm.KIMI_K3_PROJECTIONS[(7168, 3584)].residual_configs
322
+ }
323
+
324
+
325
+@pytest.mark.parametrize(
326
+ "dtype,platform_enabled",
327
+ [(torch.float16, True), (torch.bfloat16, False)],
328
+)
329
+def test_installation_requires_bf16_sm103(
330
+ monkeypatch: pytest.MonkeyPatch,
331
+ dtype: torch.dtype,
332
+ platform_enabled: bool,
333
+) -> None:
334
+ class FakeLinear(nn.Module):
335
+ def __init__(self) -> None:
336
+ super().__init__()
337
+ self.quant_method = k3_gemm.UnquantizedLinearMethod()
338
+ self.weight = torch.empty(2304, 1536)
339
+
340
+ root = nn.Module()
341
+ root.projection = FakeLinear()
342
+ monkeypatch.setattr(k3_gemm, "LinearBase", FakeLinear)
343
+ monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: platform_enabled)
344
+
345
+ k3_gemm.enable_kimi_k3_low_latency_gemm(root, dtype)
346
+
347
+ assert type(root.projection.quant_method) is k3_gemm.UnquantizedLinearMethod
348
+
349
+
350
+def _require_sm103_and_dsv3() -> None:
351
+ if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3):
352
+ pytest.skip("Kimi-K3 production selection requires SM103")
353
+ if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"):
354
+ pytest.skip("dsv3_fused_a_gemm was not built")
355
+
356
+
357
+def _require_sm103_and_cute() -> None:
358
+ if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3):
359
+ pytest.skip("Kimi-K3 production selection requires SM103")
360
+ if not k3_gemm.shape_dynamic_skinny_gemm.is_available():
361
+ pytest.skip("CuTe DSL is not available")
362
+
363
+
364
+@pytest.mark.parametrize("n,k,num_tokens", CUTE_CASES)
365
+def test_cute_selected_shapes(n: int, k: int, num_tokens: int) -> None:
366
+ _require_sm103_and_cute()
367
+ torch.manual_seed(42)
368
+ x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
369
+ weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
370
+
371
+ output = k3_gemm.try_low_latency_gemm(x, weight)
372
+
373
+ assert output is not None
374
+ reference = torch.nn.functional.linear(x, weight)
375
+ cosine = torch.nn.functional.cosine_similarity(
376
+ output.float().flatten(), reference.float().flatten(), dim=0
377
+ ).item()
378
+ assert cosine > 0.999
379
+
380
+
381
+def _dsv3_probe_tokens(tokens: frozenset[int]) -> set[int]:
382
+ """Extremes, plus both sides of the kernel's num_tokens<=8 tile_n branch."""
383
+ if not tokens:
384
+ return set()
385
+ return {min(tokens), max(tokens)} | ({8, 9} & set(tokens))
386
+
387
+
388
+# Derived from the table rather than hand-listed, so a shape routed to dsv3
389
+# cannot be added without being exercised here.
390
+DSV3_CASES = sorted(
391
+ (num_tokens, spec.n, spec.k)
392
+ for spec in KIMI_K3_PROJECTIONS.values()
393
+ for num_tokens in _dsv3_probe_tokens(spec.dsv3_tokens)
394
+)
395
+
396
+
397
+@pytest.mark.parametrize("num_tokens,n,k", DSV3_CASES)
398
+def test_dsv3_selected_shapes(num_tokens: int, n: int, k: int) -> None:
399
+ _require_sm103_and_dsv3()
400
+ spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)]
401
+ assert num_tokens in spec.dsv3_tokens
402
+ torch.manual_seed(42)
403
+ x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
404
+ weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
405
+
406
+ output = k3_gemm.try_low_latency_gemm(x, weight)
407
+
408
+ assert output is not None
409
+ reference = torch.nn.functional.linear(x, weight)
410
+ cosine = torch.nn.functional.cosine_similarity(
411
+ output.float().flatten(), reference.float().flatten(), dim=0
412
+ ).item()
413
+ assert cosine > 0.999
414
+
415
+
416
+def test_nonpacked_single_token_dsv3_falls_back() -> None:
417
+ _require_sm103_and_dsv3()
418
+ n, k = 1536, 128
419
+ storage = torch.randn(1, k + 16, dtype=torch.bfloat16, device="cuda")
420
+ x = storage[:, :k]
421
+ weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
422
+ spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)]
423
+ method = k3_gemm.KimiK3LowLatencyLinearMethod(
424
+ k3_gemm._build_plan(spec), k3_gemm._build_residual_plan(spec)
425
+ )
426
+
427
+ assert x.is_contiguous()
428
+ assert x.stride() == (k + 16, 1)
429
+ assert not k3_gemm._runtime_ok(x, weight) # strict guard rejects the view
430
+ output = method.apply(SimpleNamespace(weight=weight), x)
431
+
432
+ reference = torch.nn.functional.linear(x, weight)
433
+ torch.testing.assert_close(output, reference)
434
+
435
+
436
+def test_selected_kernels_cuda_graph_capture() -> None:
437
+ _require_sm103_and_cute()
438
+ _require_sm103_and_dsv3()
439
+ cute_spec = k3_gemm.KIMI_K3_PROJECTIONS[(6288, 7168)]
440
+ dsv3_spec = k3_gemm.KIMI_K3_PROJECTIONS[(1536, 128)]
441
+ cute_x = torch.randn(1, cute_spec.k, dtype=torch.bfloat16, device="cuda")
442
+ cute_weight = torch.randn(
443
+ cute_spec.n, cute_spec.k, dtype=torch.bfloat16, device="cuda"
444
+ )
445
+ dsv3_x = torch.randn(1, dsv3_spec.k, dtype=torch.bfloat16, device="cuda")
446
+ dsv3_weight = torch.randn(
447
+ dsv3_spec.n, dsv3_spec.k, dtype=torch.bfloat16, device="cuda"
448
+ )
449
+ k3_gemm.try_low_latency_gemm(cute_x, cute_weight)
450
+ k3_gemm.try_low_latency_gemm(dsv3_x, dsv3_weight)
451
+ torch.accelerator.synchronize()
452
+
453
+ graph = torch.cuda.CUDAGraph()
454
+ with torch.cuda.graph(graph):
455
+ cute_output = k3_gemm.try_low_latency_gemm(cute_x, cute_weight)
456
+ dsv3_output = k3_gemm.try_low_latency_gemm(dsv3_x, dsv3_weight)
457
+ graph.replay()
458
+ torch.accelerator.synchronize()
459
+
460
+ assert cute_output is not None
461
+ assert dsv3_output is not None
462
+ for output, activation, weight in (
463
+ (cute_output, cute_x, cute_weight),
464
+ (dsv3_output, dsv3_x, dsv3_weight),
465
+ ):
466
+ reference = torch.nn.functional.linear(activation, weight)
467
+ cosine = torch.nn.functional.cosine_similarity(
468
+ output.float().flatten(), reference.float().flatten(), dim=0
469
+ ).item()
470
+ assert cosine > 0.999
471
+
472
+
473
+@pytest.mark.parametrize("num_tokens", [1, 8, 9, 16])
474
+def test_dsv3_cuda_graph_capture_tile_branches(num_tokens: int) -> None:
475
+ """Capture DSV3 across the num_tokens<=8 vs >8 tile_n branch."""
476
+ _require_sm103_and_dsv3()
477
+ spec = k3_gemm.KIMI_K3_PROJECTIONS[(1536, 128)]
478
+ x = torch.randn(num_tokens, spec.k, dtype=torch.bfloat16, device="cuda")
479
+ weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda")
480
+ k3_gemm.try_low_latency_gemm(x, weight)
481
+ torch.accelerator.synchronize()
482
+
483
+ graph = torch.cuda.CUDAGraph()
484
+ with torch.cuda.graph(graph):
485
+ output = k3_gemm.try_low_latency_gemm(x, weight)
486
+ graph.replay()
487
+ torch.accelerator.synchronize()
488
+
489
+ assert output is not None
490
+ reference = torch.nn.functional.linear(x, weight)
491
+ cosine = torch.nn.functional.cosine_similarity(
492
+ output.float().flatten(), reference.float().flatten(), dim=0
493
+ ).item()
494
+ assert cosine > 0.999
495
+
496
+
497
+@pytest.mark.parametrize("n,k,num_tokens", RESIDUAL_CUTE_CASES)
498
+def test_cute_residual_epilogue(n: int, k: int, num_tokens: int) -> None:
499
+ _require_sm103_and_cute()
500
+ torch.manual_seed(42 + num_tokens)
501
+ x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
502
+ weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
503
+ residual = torch.randn(num_tokens, n, dtype=torch.bfloat16, device="cuda")
504
+ spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)]
505
+ config = spec.residual_config(num_tokens)
506
+ assert config is not None
507
+
508
+ output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual)
509
+
510
+ reference = x.float() @ weight.float().t() + residual.float()
511
+ cosine = torch.nn.functional.cosine_similarity(
512
+ output.float().flatten(), reference.flatten(), dim=0
513
+ ).item()
514
+ assert cosine > 0.999
515
+
516
+
517
+@pytest.mark.parametrize("num_tokens", range(1, 17))
518
+def test_cute_residual_epilogue_all_supported_token_counts(num_tokens: int) -> None:
519
+ _require_sm103_and_cute()
520
+ from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
521
+ ShapeDynamicSkinnyGemm,
522
+ )
523
+
524
+ n, k = 64, 512
525
+ x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
526
+ weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
527
+ residual = torch.randn(num_tokens, n, dtype=torch.bfloat16, device="cuda")
528
+ config = ShapeDynamicSkinnyGemm._config(num_tokens, n, k)
529
+
530
+ output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual)
531
+
532
+ reference = x.float() @ weight.float().t() + residual.float()
533
+ torch.testing.assert_close(output.float(), reference, rtol=2e-2, atol=2e-1)
534
+
535
+
536
+@pytest.mark.parametrize("num_tokens", range(1, 5))
537
+def test_cute_residual_epilogue_cuda_graph_capture(num_tokens: int) -> None:
538
+ _require_sm103_and_cute()
539
+ spec = k3_gemm.KIMI_K3_PROJECTIONS[(7168, 3584)]
540
+ config = spec.residual_config(num_tokens)
541
+ assert config is not None
542
+ x = torch.randn(num_tokens, spec.k, dtype=torch.bfloat16, device="cuda")
543
+ weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda")
544
+ residual = torch.randn(num_tokens, spec.n, dtype=torch.bfloat16, device="cuda")
545
+ k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual)
546
+ torch.accelerator.synchronize()
547
+
548
+ graph = torch.cuda.CUDAGraph()
549
+ with torch.cuda.graph(graph):
550
+ output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual)
551
+ graph.replay()
552
+ torch.accelerator.synchronize()
553
+
554
+ reference = x.float() @ weight.float().t() + residual.float()
555
+ cosine = torch.nn.functional.cosine_similarity(
556
+ output.float().flatten(), reference.flatten(), dim=0
557
+ ).item()
558
+ assert cosine > 0.999
559
+
560
+
561
+class _SkinnyGemmSpy:
562
+ """Wraps the skinny-GEMM singleton to record whether CuTe was invoked."""
563
+
564
+ def __init__(self, real: Any) -> None:
565
+ self._real = real
566
+ self.calls: list[int] = []
567
+
568
+ def __call__(self, a, b, config=None, residual=None):
569
+ self.calls.append(a.shape[0])
570
+ return self._real(a, b, config, residual)
571
+
572
+ def is_available(self) -> bool:
573
+ return self._real.is_available()
574
+
575
+
576
+@pytest.mark.parametrize("num_tokens", [1, 2, 3, 4])
577
+def test_latent_moe_production_layout_residual(
578
+ monkeypatch: pytest.MonkeyPatch,
579
+ num_tokens: int,
580
+) -> None:
581
+ """The real Latent-MoE residual is a non-packed slice of a cat buffer.
582
+
583
+ The strict packed-row-major guard rejects such a slice at every token count
584
+ (a size-1 leading dim reads as contiguous but its stride is not packed), so
585
+ the CuTe residual epilogue never fires for this production layout and the
586
+ method falls back to addmm. Output is correct regardless of the path.
587
+ """
588
+ _require_sm103_and_cute()
589
+ latent_dim, shared_dim = 3584, 7168 # routed_expert_up_proj K, N
590
+ torch.manual_seed(7 + num_tokens)
591
+ buf = torch.randn(
592
+ num_tokens, latent_dim + shared_dim, dtype=torch.bfloat16, device="cuda"
593
+ )
594
+ latent = buf[:, :latent_dim] # non-contiguous view (row stride = full width)
595
+ residual = buf[:, latent_dim:] # non-contiguous view
596
+ weight = torch.randn(shared_dim, latent_dim, dtype=torch.bfloat16, device="cuda")
597
+
598
+ spec = k3_gemm.KIMI_K3_PROJECTIONS[(shared_dim, latent_dim)]
599
+ method = k3_gemm.KimiK3LowLatencyLinearMethod(
600
+ k3_gemm._build_plan(spec), k3_gemm._build_residual_plan(spec)
601
+ )
602
+ spy = _SkinnyGemmSpy(k3_gemm.shape_dynamic_skinny_gemm)
603
+ monkeypatch.setattr(k3_gemm, "shape_dynamic_skinny_gemm", spy)
604
+
605
+ layer = SimpleNamespace(weight=weight)
606
+ output = method.apply_with_residual(layer, latent, residual)
607
+
608
+ reference = latent.float() @ weight.float().t() + residual.float()
609
+ cosine = torch.nn.functional.cosine_similarity(
610
+ output.float().flatten(), reference.flatten(), dim=0
611
+ ).item()
612
+ assert cosine > 0.999 # correct regardless of the path taken
613
+ assert not spy.calls, (
614
+ "non-packed buf-slice residual must fall back to addmm at every M"
615
+ )
616
+
617
+
618
+def test_residual_dispatch_falls_back_to_addmm(
619
+ monkeypatch: pytest.MonkeyPatch,
620
+) -> None:
621
+ fallback = torch.randn(2, 3)
622
+ residual = torch.randn(2, 3)
623
+ x = torch.randn(2, 4)
624
+ weight = torch.randn(3, 4)
625
+ monkeypatch.setattr(torch, "addmm", lambda *args: fallback)
626
+ # CPU tensors fail the runtime check, forcing the addmm fallback.
627
+ method = k3_gemm.KimiK3LowLatencyLinearMethod({}, {})
628
+
629
+ output = method.apply_with_residual(SimpleNamespace(weight=weight), x, residual)
630
+
631
+ assert output is fallback
632
+
633
+
634
+def test_fallback_preserves_default_method(monkeypatch: pytest.MonkeyPatch) -> None:
635
+ fallback = torch.empty(2, 8)
636
+ monkeypatch.setattr(
637
+ k3_gemm.UnquantizedLinearMethod,
638
+ "apply",
639
+ lambda *args: fallback,
640
+ )
641
+ # 1-D input fails the runtime check, forcing the base-method fallback.
642
+ method = k3_gemm.KimiK3LowLatencyLinearMethod({}, {})
643
+
644
+ output = method.apply(
645
+ SimpleNamespace(weight=torch.empty(0)),
646
+ torch.empty(0),
647
+ )
648
+
649
+ assert output is fallback
tests/kernels/test_kda.py
deleted
-226
@@ -1,226 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Precision tests for vllm's chunk_kda Triton operator.
4
-
5
-Compares chunk_kda against a naive recurrent reference (float32).
6
-Uses torch.rand for q/k/v to match FLA's test pattern.
7
-"""
8
-
9
-import pytest
10
-import torch
11
-import torch.nn.functional as F
12
-
13
-from vllm.third_party.flash_linear_attention.ops.kda import (
14
- chunk_kda,
15
- chunk_kda_with_fused_gate,
16
- fused_kda_gate,
17
-)
18
-from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
19
-
20
-DEVICE = "cuda"
21
-
22
-
23
-def naive_recurrent_kda(
24
- q: torch.Tensor,
25
- k: torch.Tensor,
26
- v: torch.Tensor,
27
- g: torch.Tensor,
28
- beta: torch.Tensor,
29
- scale: float | None = None,
30
- initial_state: torch.Tensor | None = None,
31
- output_final_state: bool = False,
32
-) -> tuple[torch.Tensor, torch.Tensor | None]:
33
- """Naive recurrent KDA reference, ported from FLA's naive.py."""
34
- dtype = v.dtype
35
- B, T, H, K = q.shape
36
- V = v.shape[-1]
37
- if scale is None:
38
- scale = K**-0.5
39
-
40
- q, k, v, g, beta = (x.to(torch.float) for x in [q, k, v, g, beta])
41
- q = q * scale
42
-
43
- S = k.new_zeros(B, H, K, V).to(q)
44
- if initial_state is not None:
45
- S += initial_state
46
- o = torch.zeros_like(v)
47
- for i in range(T):
48
- q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i]
49
- S = S * g_i[..., None].exp()
50
- S = S + torch.einsum(
51
- "bhk,bhv->bhkv",
52
- b_i[..., None] * k_i,
53
- v_i - (k_i[..., None] * S).sum(-2),
54
- )
55
- o[:, i] = torch.einsum("bhk,bhkv->bhv", q_i, S)
56
- if not output_final_state:
57
- S = None
58
- return o.to(dtype), S
59
-
60
-
61
-def assert_close(
62
- name: str,
63
- ref: torch.Tensor,
64
- tri: torch.Tensor,
65
- ratio: float,
66
- err_atol: float = 1e-6,
67
-):
68
- """RMSE-based relative error comparison."""
69
- abs_err = (ref.detach() - tri.detach()).flatten().abs().max().item()
70
- rmse_diff = (ref.detach() - tri.detach()).flatten().square().mean().sqrt().item()
71
- rmse_base = ref.detach().flatten().square().mean().sqrt().item()
72
- rel_err = rmse_diff / (rmse_base + 1e-8)
73
- print(f"{name:>4} | abs={abs_err:.6f} | rmse={rel_err:.6f} | thr={ratio}")
74
- if abs_err <= err_atol:
75
- return
76
- assert not torch.isnan(ref).any(), f"{name}: NaN detected in ref"
77
- assert not torch.isnan(tri).any(), f"{name}: NaN detected in tri"
78
- assert rel_err < ratio, (
79
- f"{name}: max abs err {abs_err:.6f}, rmse ratio {rel_err:.6f} >= {ratio}"
80
- )
81
-
82
-
83
-@pytest.mark.parametrize(
84
- ("H", "D", "cu_seqlens", "dtype"),
85
- [
86
- pytest.param(
87
- *test,
88
- id="H{}-D{}-cu{}-{}".format(*test),
89
- )
90
- for test in [
91
- (32, 128, [0, 64], torch.float16),
92
- (32, 128, [0, 1024], torch.float16),
93
- (32, 128, [0, 15], torch.float16),
94
- (32, 128, [0, 256, 512, 768, 1024], torch.float16),
95
- (32, 128, [0, 15, 100, 300, 1200], torch.float16),
96
- (64, 128, [0, 256, 500, 1000], torch.float16),
97
- (32, 128, [0, 8192], torch.float16),
98
- (32, 128, [0, 256, 500, 1000], torch.bfloat16),
99
- ]
100
- ],
101
-)
102
-@torch.inference_mode()
103
-def test_chunk_kda(
104
- H: int,
105
- D: int,
106
- cu_seqlens: list[int],
107
- dtype: torch.dtype,
108
-):
109
- T = cu_seqlens[-1]
110
- torch.manual_seed(42)
111
- B = 1
112
- cu_seqlens_t = torch.LongTensor(cu_seqlens).to(DEVICE)
113
- N = len(cu_seqlens) - 1
114
-
115
- q = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
116
- k = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
117
- v = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
118
- g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float32, device=DEVICE)).to(
119
- dtype
120
- )
121
- beta = torch.rand(B, T, H, dtype=dtype, device=DEVICE).sigmoid()
122
- h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE)
123
-
124
- # Naive reference with l2norm_fwd (same kernel as chunk_kda)
125
- ref_outputs = []
126
- ref_states = []
127
- for i in range(N):
128
- s, e = cu_seqlens[i], cu_seqlens[i + 1]
129
- q_i = l2norm_fwd(q[:, s:e].contiguous())
130
- k_i = l2norm_fwd(k[:, s:e].contiguous())
131
- o_i, ht_i = naive_recurrent_kda(
132
- q_i,
133
- k_i,
134
- v[:, s:e],
135
- g[:, s:e],
136
- beta[:, s:e],
137
- initial_state=h0[i],
138
- output_final_state=True,
139
- )
140
- ref_outputs.append(o_i)
141
- ref_states.append(ht_i)
142
- ref_o = torch.cat(ref_outputs, dim=1)
143
- ref_ht = torch.cat(ref_states, dim=0)
144
-
145
- # h0 transposed to (V, K) layout for the kernel; naive uses (K, V)
146
- tri_o, tri_ht = chunk_kda(
147
- q=q.clone(),
148
- k=k.clone(),
149
- v=v.clone(),
150
- g=g.clone(),
151
- beta=beta.clone(),
152
- initial_state=h0.transpose(-1, -2).contiguous().clone(),
153
- output_final_state=True,
154
- cu_seqlens=cu_seqlens_t,
155
- use_qk_l2norm_in_kernel=True,
156
- )
157
-
158
- assert not torch.isnan(tri_o).any(), "Triton output o contains NaN"
159
- assert not torch.isnan(tri_ht).any(), "Triton output ht contains NaN"
160
- assert_close("o", ref_o, tri_o, 0.005)
161
- assert_close("ht", ref_ht, tri_ht.transpose(-1, -2).contiguous(), 0.005)
162
-
163
-
164
-@pytest.mark.parametrize(
165
- ("cu_seqlens", "dtype"),
166
- [
167
- ([0, 64], torch.float16),
168
- ([0, 15, 100, 300], torch.bfloat16),
169
- ],
170
-)
171
-@torch.inference_mode()
172
-def test_chunk_kda_fused_gate_cumsum_matches_unfused(
173
- cu_seqlens: list[int],
174
- dtype: torch.dtype,
175
-):
176
- H, D = 8, 64
177
- T = cu_seqlens[-1]
178
- N = len(cu_seqlens) - 1
179
- torch.manual_seed(123)
180
-
181
- cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32, device=DEVICE)
182
- q = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
183
- k = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
184
- v = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
185
- raw_g = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
186
- beta = torch.rand(1, T, H, dtype=dtype, device=DEVICE).sigmoid()
187
- A_log = (torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5).contiguous()
188
- dt_bias = (
189
- torch.randn(H * D, dtype=torch.float32, device=DEVICE) * 0.1
190
- ).contiguous()
191
- h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE)
192
- initial_state = h0.transpose(-1, -2).contiguous()
193
-
194
- gate = fused_kda_gate(
195
- raw_g.reshape(T, H * D),
196
- A_log,
197
- D,
198
- g_bias=dt_bias,
199
- ).unsqueeze(0)
200
- old_o, old_ht = chunk_kda(
201
- q=q.clone(),
202
- k=k.clone(),
203
- v=v.clone(),
204
- g=gate,
205
- beta=beta.clone(),
206
- initial_state=initial_state.clone(),
207
- output_final_state=True,
208
- cu_seqlens=cu_seqlens_t,
209
- use_qk_l2norm_in_kernel=True,
210
- )
211
- new_o, new_ht = chunk_kda_with_fused_gate(
212
- q=q.clone(),
213
- k=k.clone(),
214
- v=v.clone(),
215
- raw_g=raw_g,
216
- beta=beta.clone(),
217
- A_log=A_log,
218
- g_bias=dt_bias,
219
- initial_state=initial_state.clone(),
220
- output_final_state=True,
221
- cu_seqlens=cu_seqlens_t,
222
- use_qk_l2norm_in_kernel=True,
223
- )
224
-
225
- assert_close("o", old_o, new_o, 1e-3, err_atol=1e-3)
226
- assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3)
tests/models/kimi_k3/test_attn_res.py
+33
@@ -5,6 +5,7 @@ import pytest
5
import torch
6
import torch.nn.functional as F
7
8
+from vllm.models.kimi_k3.common.mtp import fused_mtp_input
9
from vllm.models.kimi_k3.nvidia.ops import attn_res
10
from vllm.platforms import current_platform
11
@@ -191,3 +192,35 @@ def test_attn_res_without_output_norm():
192
)
193
194
torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
195
+
196
+
197
+@pytest.mark.parametrize("num_tokens", [0, 1, 17])
198
+def test_fused_mtp_input(num_tokens: int):
199
+ positions = torch.arange(num_tokens, device="cuda")
200
+ inputs_embeds = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=7)
201
+ previous_hidden_states = _randn_with_row_padding(
202
+ num_tokens, HIDDEN_SIZE, padding=11
203
+ )
204
+ enorm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
205
+ hnorm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
206
+
207
+ masked_inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
208
+ expected = torch.cat(
209
+ (
210
+ F.rms_norm(masked_inputs_embeds, (HIDDEN_SIZE,), enorm_weight, EPS),
211
+ F.rms_norm(previous_hidden_states, (HIDDEN_SIZE,), hnorm_weight, EPS),
212
+ ),
213
+ dim=-1,
214
+ )
215
+ actual = fused_mtp_input(
216
+ positions,
217
+ inputs_embeds,
218
+ previous_hidden_states,
219
+ enorm_weight,
220
+ hnorm_weight,
221
+ EPS,
222
+ )
223
+
224
+ torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
225
+ assert actual.shape == (num_tokens, 2 * HIDDEN_SIZE)
226
+ assert actual.is_contiguous()
tests/models/kimi_k3/test_eagle3.py
new
+130
@@ -0,0 +1,130 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+from unittest.mock import Mock
6
+
7
+import torch
8
+
9
+from vllm.model_executor.models.interfaces import supports_eagle3
10
+from vllm.models.kimi_k3.nvidia import model as kimi_model
11
+from vllm.models.kimi_k3.nvidia.model import (
12
+ KimiK3ForConditionalGeneration,
13
+ KimiLinearModel,
14
+)
15
+
16
+
17
+def _make_kimi_linear_model() -> KimiLinearModel:
18
+ model = object.__new__(KimiLinearModel)
19
+ object.__setattr__(model, "aux_hidden_state_layers", (2,))
20
+ object.__setattr__(model, "use_sequence_parallel", False)
21
+ return model
22
+
23
+
24
+def test_kimi_k3_advertises_eagle3_support():
25
+ assert supports_eagle3(KimiK3ForConditionalGeneration)
26
+
27
+
28
+def test_kimi_k3_uses_shared_eagle3_layer_configuration():
29
+ target = object.__new__(KimiK3ForConditionalGeneration)
30
+ torch.nn.Module.__init__(target)
31
+ model = _make_kimi_linear_model()
32
+ object.__setattr__(model, "layers", [None] * 93)
33
+ language_model = SimpleNamespace(
34
+ embed_input_ids=lambda _: None,
35
+ model=model,
36
+ )
37
+ object.__setattr__(target, "language_model", language_model)
38
+ object.__setattr__(target, "_language_model_names", ["language_model"])
39
+
40
+ target.set_aux_hidden_state_layers((2, 46, 90))
41
+
42
+ assert model.aux_hidden_state_layers == (2, 46, 90)
43
+ assert target.get_eagle3_default_aux_hidden_state_layers() == (
44
+ 2,
45
+ 46,
46
+ 90,
47
+ )
48
+
49
+
50
+def test_kimi_linear_forward_extracts_standard_aux_hidden_states(monkeypatch):
51
+ model = _make_kimi_linear_model()
52
+ initial_hidden_states = torch.tensor([[1.0, 2.0]])
53
+ layer_hidden_states = torch.tensor([[3.0, 4.0]])
54
+ layer_residual = torch.tensor([[5.0, 6.0]])
55
+
56
+ object.__setattr__(model, "start_layer", 0)
57
+ object.__setattr__(model, "end_layer", 1)
58
+ object.__setattr__(
59
+ model,
60
+ "layers",
61
+ [Mock(return_value=(layer_hidden_states, None, layer_residual))],
62
+ )
63
+ object.__setattr__(model, "aux_hidden_state_layers", (0, 1))
64
+ object.__setattr__(model, "use_attn_res", False)
65
+ monkeypatch.setattr(
66
+ kimi_model,
67
+ "get_pp_group",
68
+ lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
69
+ )
70
+
71
+ output, aux_hidden_states = model.forward(
72
+ input_ids=None,
73
+ positions=torch.tensor([0]),
74
+ intermediate_tensors=None,
75
+ inputs_embeds=initial_hidden_states,
76
+ )
77
+
78
+ expected_layer_output = layer_hidden_states + layer_residual
79
+ torch.testing.assert_close(output, expected_layer_output)
80
+ torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states)
81
+ torch.testing.assert_close(aux_hidden_states[1], expected_layer_output)
82
+
83
+
84
+def test_kimi_linear_forward_extracts_attn_res_aux_hidden_states(monkeypatch):
85
+ model = _make_kimi_linear_model()
86
+ initial_hidden_states = torch.tensor([[1.0, 2.0]])
87
+ layer_hidden_states = torch.tensor([[3.0, 4.0]])
88
+ prefix_sum = torch.tensor([[5.0, 6.0]])
89
+ block_residual = torch.tensor([[[7.0, 8.0]]])
90
+ final_hidden_states = torch.tensor([[9.0, 10.0]])
91
+
92
+ object.__setattr__(model, "start_layer", 0)
93
+ object.__setattr__(model, "end_layer", 1)
94
+ object.__setattr__(
95
+ model,
96
+ "layers",
97
+ [Mock(return_value=(layer_hidden_states, prefix_sum, block_residual))],
98
+ )
99
+ object.__setattr__(model, "aux_hidden_state_layers", (0, 1))
100
+ object.__setattr__(model, "use_attn_res", True)
101
+ object.__setattr__(model, "num_attn_res_blocks", 1)
102
+ object.__setattr__(
103
+ model,
104
+ "output_attn_res_norm",
105
+ SimpleNamespace(weight=torch.ones(2), variance_epsilon=1e-5),
106
+ )
107
+ object.__setattr__(
108
+ model,
109
+ "output_attn_res_proj",
110
+ SimpleNamespace(weight=torch.ones(1, 2)),
111
+ )
112
+ monkeypatch.setattr(
113
+ kimi_model,
114
+ "get_pp_group",
115
+ lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
116
+ )
117
+ final_attn_res = Mock(return_value=final_hidden_states)
118
+ monkeypatch.setattr(kimi_model, "attn_res", final_attn_res)
119
+
120
+ output, aux_hidden_states = model.forward(
121
+ input_ids=None,
122
+ positions=torch.tensor([0]),
123
+ intermediate_tensors=None,
124
+ inputs_embeds=initial_hidden_states,
125
+ )
126
+
127
+ torch.testing.assert_close(output, final_hidden_states)
128
+ torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states)
129
+ torch.testing.assert_close(aux_hidden_states[1], prefix_sum + layer_hidden_states)
130
+ assert final_attn_res.call_args.args[2] is block_residual
tests/models/kimi_k3/test_kda.py
new
+757
@@ -0,0 +1,757 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Precision tests for vllm's chunk_kda Triton operator.
4
+
5
+Compares chunk_kda against a naive recurrent reference (float32).
6
+Uses torch.rand for q/k/v to match FLA's test pattern.
7
+"""
8
+
9
+import pytest
10
+import torch
11
+import torch.nn.functional as F
12
+
13
+from vllm import _custom_ops as ops
14
+from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update
15
+from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
16
+ gather_initial_states,
17
+)
18
+from vllm.models.kimi_k3.nvidia.kda import (
19
+ is_flashkda_supported,
20
+ is_fused_kda_decode_supported,
21
+)
22
+from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (
23
+ chunk_kda,
24
+ chunk_kda_with_fused_gate,
25
+ fused_kda_gate,
26
+ fused_recurrent_kda,
27
+ fused_recurrent_kda_fwd,
28
+ fused_recurrent_kda_packed_decode,
29
+)
30
+from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
31
+
32
+DEVICE = "cuda"
33
+
34
+
35
+@torch.inference_mode()
36
+def test_gather_initial_states_correctness():
37
+ row_size = 8 * 128 * 128
38
+ storage = torch.randn(5, row_size + 256, dtype=torch.float32, device=DEVICE)
39
+ state = storage[:, :row_size].view(5, 8, 128, 128)
40
+ assert not state.is_contiguous()
41
+ assert state[0].is_contiguous()
42
+ indices = torch.tensor([4, 1, 3], dtype=torch.int32, device=DEVICE)
43
+ has_initial_state = torch.tensor([True, False, True], device=DEVICE)
44
+
45
+ expected = state[indices].clone()
46
+ expected[~has_initial_state] = 0
47
+
48
+ torch.testing.assert_close(
49
+ gather_initial_states(state, indices, has_initial_state),
50
+ expected,
51
+ )
52
+
53
+
54
+def naive_recurrent_kda(
55
+ q: torch.Tensor,
56
+ k: torch.Tensor,
57
+ v: torch.Tensor,
58
+ g: torch.Tensor,
59
+ beta: torch.Tensor,
60
+ scale: float | None = None,
61
+ initial_state: torch.Tensor | None = None,
62
+ output_final_state: bool = False,
63
+) -> tuple[torch.Tensor, torch.Tensor | None]:
64
+ """Naive recurrent KDA reference, ported from FLA's naive.py."""
65
+ dtype = v.dtype
66
+ B, T, H, K = q.shape
67
+ V = v.shape[-1]
68
+ if scale is None:
69
+ scale = K**-0.5
70
+
71
+ q, k, v, g, beta = (x.to(torch.float) for x in [q, k, v, g, beta])
72
+ q = q * scale
73
+
74
+ S = k.new_zeros(B, H, K, V).to(q)
75
+ if initial_state is not None:
76
+ S += initial_state
77
+ o = torch.zeros_like(v)
78
+ for i in range(T):
79
+ q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i]
80
+ S = S * g_i[..., None].exp()
81
+ S = S + torch.einsum(
82
+ "bhk,bhv->bhkv",
83
+ b_i[..., None] * k_i,
84
+ v_i - (k_i[..., None] * S).sum(-2),
85
+ )
86
+ o[:, i] = torch.einsum("bhk,bhkv->bhv", q_i, S)
87
+ if not output_final_state:
88
+ S = None
89
+ return o.to(dtype), S
90
+
91
+
92
+def assert_close(
93
+ name: str,
94
+ ref: torch.Tensor,
95
+ tri: torch.Tensor,
96
+ ratio: float,
97
+ err_atol: float = 1e-6,
98
+):
99
+ """RMSE-based relative error comparison."""
100
+ abs_err = (ref.detach() - tri.detach()).flatten().abs().max().item()
101
+ rmse_diff = (ref.detach() - tri.detach()).flatten().square().mean().sqrt().item()
102
+ rmse_base = ref.detach().flatten().square().mean().sqrt().item()
103
+ rel_err = rmse_diff / (rmse_base + 1e-8)
104
+ print(f"{name:>4} | abs={abs_err:.6f} | rmse={rel_err:.6f} | thr={ratio}")
105
+ if abs_err <= err_atol:
106
+ return
107
+ assert not torch.isnan(ref).any(), f"{name}: NaN detected in ref"
108
+ assert not torch.isnan(tri).any(), f"{name}: NaN detected in tri"
109
+ assert rel_err < ratio, (
110
+ f"{name}: max abs err {abs_err:.6f}, rmse ratio {rel_err:.6f} >= {ratio}"
111
+ )
112
+
113
+
114
+@pytest.mark.parametrize(
115
+ ("H", "D", "cu_seqlens", "dtype"),
116
+ [
117
+ pytest.param(
118
+ *test,
119
+ id="H{}-D{}-cu{}-{}".format(*test),
120
+ )
121
+ for test in [
122
+ (32, 128, [0, 64], torch.float16),
123
+ (32, 128, [0, 1024], torch.float16),
124
+ (32, 128, [0, 15], torch.float16),
125
+ (32, 128, [0, 256, 512, 768, 1024], torch.float16),
126
+ (32, 128, [0, 15, 100, 300, 1200], torch.float16),
127
+ (64, 128, [0, 256, 500, 1000], torch.float16),
128
+ (32, 128, [0, 8192], torch.float16),
129
+ (32, 128, [0, 256, 500, 1000], torch.bfloat16),
130
+ ]
131
+ ],
132
+)
133
+@torch.inference_mode()
134
+def test_chunk_kda(
135
+ H: int,
136
+ D: int,
137
+ cu_seqlens: list[int],
138
+ dtype: torch.dtype,
139
+):
140
+ T = cu_seqlens[-1]
141
+ torch.manual_seed(42)
142
+ B = 1
143
+ cu_seqlens_t = torch.LongTensor(cu_seqlens).to(DEVICE)
144
+ N = len(cu_seqlens) - 1
145
+
146
+ q = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
147
+ k = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
148
+ v = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
149
+ g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float32, device=DEVICE)).to(
150
+ dtype
151
+ )
152
+ beta = torch.rand(B, T, H, dtype=dtype, device=DEVICE).sigmoid()
153
+ h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE)
154
+
155
+ # Naive reference with l2norm_fwd (same kernel as chunk_kda)
156
+ ref_outputs = []
157
+ ref_states = []
158
+ for i in range(N):
159
+ s, e = cu_seqlens[i], cu_seqlens[i + 1]
160
+ q_i = l2norm_fwd(q[:, s:e].contiguous())
161
+ k_i = l2norm_fwd(k[:, s:e].contiguous())
162
+ o_i, ht_i = naive_recurrent_kda(
163
+ q_i,
164
+ k_i,
165
+ v[:, s:e],
166
+ g[:, s:e],
167
+ beta[:, s:e],
168
+ initial_state=h0[i],
169
+ output_final_state=True,
170
+ )
171
+ ref_outputs.append(o_i)
172
+ ref_states.append(ht_i)
173
+ ref_o = torch.cat(ref_outputs, dim=1)
174
+ ref_ht = torch.cat(ref_states, dim=0)
175
+
176
+ # h0 transposed to (V, K) layout for the kernel; naive uses (K, V)
177
+ tri_o, tri_ht = chunk_kda(
178
+ q=q.clone(),
179
+ k=k.clone(),
180
+ v=v.clone(),
181
+ g=g.clone(),
182
+ beta=beta.clone(),
183
+ initial_state=h0.transpose(-1, -2).contiguous().clone(),
184
+ output_final_state=True,
185
+ cu_seqlens=cu_seqlens_t,
186
+ use_qk_l2norm_in_kernel=True,
187
+ )
188
+
189
+ assert not torch.isnan(tri_o).any(), "Triton output o contains NaN"
190
+ assert not torch.isnan(tri_ht).any(), "Triton output ht contains NaN"
191
+ assert_close("o", ref_o, tri_o, 0.005)
192
+ assert_close("ht", ref_ht, tri_ht.transpose(-1, -2).contiguous(), 0.005)
193
+
194
+
195
+@pytest.mark.parametrize(
196
+ ("cu_seqlens", "dtype", "lower_bound"),
197
+ [
198
+ ([0, 64], torch.float16, None),
199
+ ([0, 15, 100, 300], torch.bfloat16, None),
200
+ ([0, 15, 100, 300], torch.bfloat16, -3.0),
201
+ ],
202
+)
203
+@torch.inference_mode()
204
+def test_chunk_kda_fused_gate_cumsum_matches_unfused(
205
+ cu_seqlens: list[int],
206
+ dtype: torch.dtype,
207
+ lower_bound: float | None,
208
+):
209
+ H, D = 8, 64
210
+ T = cu_seqlens[-1]
211
+ N = len(cu_seqlens) - 1
212
+ torch.manual_seed(123)
213
+
214
+ cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32, device=DEVICE)
215
+ q = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
216
+ k = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
217
+ v = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
218
+ raw_g = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
219
+ beta_storage = torch.randn(1, T, 2 * H + 3, dtype=dtype, device=DEVICE)
220
+ raw_beta = beta_storage[..., 1 : 2 * H + 1 : 2]
221
+ beta = raw_beta.float().sigmoid()
222
+ A_log = (torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5).contiguous()
223
+ dt_bias = (
224
+ torch.randn(H * D, dtype=torch.float32, device=DEVICE) * 0.1
225
+ ).contiguous()
226
+ h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE)
227
+ initial_state = h0.transpose(-1, -2).contiguous()
228
+
229
+ gate = fused_kda_gate(
230
+ raw_g.reshape(T, H * D),
231
+ A_log,
232
+ D,
233
+ g_bias=dt_bias,
234
+ lower_bound=lower_bound,
235
+ )
236
+ if lower_bound is not None:
237
+ expected_gate = lower_bound * torch.sigmoid(
238
+ A_log.exp()[None, :, None]
239
+ * (raw_g.float().view(T, H, D) + dt_bias.view(H, D))
240
+ )
241
+ torch.testing.assert_close(gate, expected_gate)
242
+ gate = gate.unsqueeze(0)
243
+ old_o, old_ht = chunk_kda(
244
+ q=q.clone(),
245
+ k=k.clone(),
246
+ v=v.clone(),
247
+ g=gate,
248
+ beta=beta,
249
+ initial_state=initial_state.clone(),
250
+ output_final_state=True,
251
+ cu_seqlens=cu_seqlens_t,
252
+ use_qk_l2norm_in_kernel=True,
253
+ )
254
+ new_o, new_ht = chunk_kda_with_fused_gate(
255
+ q=q.clone(),
256
+ k=k.clone(),
257
+ v=v.clone(),
258
+ raw_g=raw_g,
259
+ raw_beta=raw_beta,
260
+ A_log=A_log,
261
+ g_bias=dt_bias,
262
+ lower_bound=lower_bound,
263
+ initial_state=initial_state.clone(),
264
+ output_final_state=True,
265
+ cu_seqlens=cu_seqlens_t,
266
+ use_qk_l2norm_in_kernel=True,
267
+ )
268
+
269
+ assert_close("o", old_o, new_o, 1e-3, err_atol=1e-3)
270
+ assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3)
271
+
272
+
273
+@pytest.mark.parametrize("num_seqs", [1, 8, 32])
274
+@pytest.mark.parametrize("lower_bound", [-5.0, None])
275
+@pytest.mark.parametrize("state_indices_stride", [1, 8])
276
+@torch.inference_mode()
277
+def test_packed_kda_decode_correctness(
278
+ num_seqs: int,
279
+ lower_bound: float | None,
280
+ state_indices_stride: int,
281
+):
282
+ H, D = 8, 128
283
+ torch.manual_seed(321)
284
+
285
+ packed_storage = torch.randn(
286
+ num_seqs,
287
+ 3 * H * D + 1,
288
+ dtype=torch.bfloat16,
289
+ device=DEVICE,
290
+ )
291
+ mixed_qkv = packed_storage[:, : 3 * H * D]
292
+ assert mixed_qkv.stride(0) == 3 * H * D + 1
293
+ q, k, v = (
294
+ x.contiguous().view(1, num_seqs, H, D) for x in mixed_qkv.split(H * D, dim=-1)
295
+ )
296
+ raw_g = torch.randn(
297
+ 1,
298
+ num_seqs,
299
+ H,
300
+ D,
301
+ dtype=torch.bfloat16,
302
+ device=DEVICE,
303
+ )
304
+ raw_beta = torch.randn(
305
+ 1,
306
+ num_seqs,
307
+ H,
308
+ dtype=torch.bfloat16,
309
+ device=DEVICE,
310
+ )
311
+ beta = raw_beta.float().sigmoid()
312
+ A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5
313
+ dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1
314
+ state_storage = torch.randn(
315
+ num_seqs + 1,
316
+ H * D * D + 17,
317
+ dtype=torch.float32,
318
+ device=DEVICE,
319
+ )
320
+ state = state_storage[:, : H * D * D].view(num_seqs + 1, H, D, D)
321
+ assert not state.is_contiguous()
322
+ assert state.stride()[1:] == (D * D, D, 1)
323
+ state_indices_storage = torch.zeros(
324
+ num_seqs,
325
+ state_indices_stride,
326
+ dtype=torch.int32,
327
+ device=DEVICE,
328
+ )
329
+ state_indices = state_indices_storage[:, 0]
330
+ state_indices.copy_(
331
+ torch.arange(
332
+ 1,
333
+ num_seqs + 1,
334
+ dtype=torch.int32,
335
+ device=DEVICE,
336
+ )
337
+ )
338
+ gate = fused_kda_gate(
339
+ raw_g.reshape(num_seqs, H * D),
340
+ A_log,
341
+ D,
342
+ g_bias=dt_bias,
343
+ lower_bound=lower_bound,
344
+ ).unsqueeze(0)
345
+ dense_state = state.clone()
346
+ dense_out, _ = fused_recurrent_kda_fwd(
347
+ q=q,
348
+ k=k,
349
+ v=v,
350
+ g=gate,
351
+ beta=beta,
352
+ scale=D**-0.5,
353
+ initial_state=dense_state,
354
+ inplace_final_state=True,
355
+ cu_seqlens=torch.arange(
356
+ num_seqs + 1,
357
+ dtype=torch.int32,
358
+ device=DEVICE,
359
+ ),
360
+ ssm_state_indices=state_indices,
361
+ use_qk_l2norm_in_kernel=True,
362
+ )
363
+ packed_state = state
364
+ packed_out, _ = fused_recurrent_kda_packed_decode(
365
+ mixed_qkv=mixed_qkv,
366
+ raw_g=raw_g,
367
+ raw_beta=raw_beta,
368
+ A_log=A_log,
369
+ dt_bias=dt_bias,
370
+ lower_bound=lower_bound,
371
+ initial_state=packed_state,
372
+ state_indices=state_indices,
373
+ )
374
+
375
+ assert_close("o", dense_out, packed_out, 1e-3, err_atol=1e-3)
376
+ assert_close("ht", dense_state, packed_state, 1e-3, err_atol=1e-3)
377
+
378
+
379
+@pytest.mark.parametrize(
380
+ ("H", "fuse_gate"),
381
+ [(12, True), (12, False), (12, None), (96, None)],
382
+)
383
+@pytest.mark.parametrize("lower_bound", [-5.0, None])
384
+@torch.inference_mode()
385
+def test_kda_spec_decode_correctness(
386
+ H: int,
387
+ fuse_gate: bool | None,
388
+ lower_bound: float | None,
389
+):
390
+ num_seqs, query_len, D = 3, 3, 128
391
+ T = num_seqs * query_len
392
+ torch.manual_seed(1234)
393
+
394
+ qkv_storage = torch.randn(
395
+ 1,
396
+ T,
397
+ 3 * H * D + 7,
398
+ dtype=torch.bfloat16,
399
+ device=DEVICE,
400
+ )
401
+ packed_qkv = qkv_storage[..., : 3 * H * D]
402
+ q, k, v = (x.view(1, T, H, D) for x in packed_qkv.split(H * D, dim=-1))
403
+ gate_storage = torch.randn(
404
+ 1,
405
+ T,
406
+ H * D + 5,
407
+ dtype=torch.bfloat16,
408
+ device=DEVICE,
409
+ )
410
+ raw_g = gate_storage[..., : H * D].view(1, T, H, D)
411
+ beta_storage = torch.randn(
412
+ 1,
413
+ T,
414
+ H + 1,
415
+ dtype=torch.bfloat16,
416
+ device=DEVICE,
417
+ )
418
+ raw_beta = beta_storage[..., :H]
419
+ A_log = 0.5 * torch.randn(H, dtype=torch.float32, device=DEVICE)
420
+ dt_bias = 0.1 * torch.randn(H, D, dtype=torch.float32, device=DEVICE)
421
+ cu_seqlens = torch.arange(
422
+ 0,
423
+ T + 1,
424
+ query_len,
425
+ dtype=torch.int32,
426
+ device=DEVICE,
427
+ )
428
+ state_indices = torch.arange(
429
+ 1,
430
+ T + 1,
431
+ dtype=torch.int32,
432
+ device=DEVICE,
433
+ ).view(num_seqs, query_len)
434
+ num_accepted_tokens = torch.tensor(
435
+ [1, 2, 3],
436
+ dtype=torch.int32,
437
+ device=DEVICE,
438
+ )
439
+ state_storage = 0.01 * torch.randn(
440
+ T + 1,
441
+ H * D * D + 17,
442
+ dtype=torch.float32,
443
+ device=DEVICE,
444
+ )
445
+ state = state_storage[:, : H * D * D].view(T + 1, H, D, D)
446
+ output_storage = torch.full(
447
+ (1, T, H * D + 11),
448
+ torch.nan,
449
+ dtype=torch.bfloat16,
450
+ device=DEVICE,
451
+ )
452
+ output = output_storage[..., : H * D].view(1, T, H, D)
453
+
454
+ gate = fused_kda_gate(
455
+ raw_g.contiguous().view(T, H * D),
456
+ A_log,
457
+ D,
458
+ g_bias=dt_bias,
459
+ lower_bound=lower_bound,
460
+ ).unsqueeze(0)
461
+ beta = raw_beta.float().sigmoid()
462
+ q_norm = l2norm_fwd(q.contiguous())
463
+ k_norm = l2norm_fwd(k.contiguous())
464
+ expected_state = state.clone()
465
+ expected_outputs = []
466
+ for seq, accepted in enumerate(num_accepted_tokens.tolist()):
467
+ recurrent_state = expected_state[state_indices[seq, accepted - 1]].transpose(
468
+ -1, -2
469
+ )
470
+ start = seq * query_len
471
+ for token in range(query_len):
472
+ token_slice = slice(start + token, start + token + 1)
473
+ token_output, recurrent_state = naive_recurrent_kda(
474
+ q_norm[:, token_slice],
475
+ k_norm[:, token_slice],
476
+ v[:, token_slice],
477
+ gate[:, token_slice],
478
+ beta[:, token_slice],
479
+ initial_state=recurrent_state,
480
+ output_final_state=True,
481
+ )
482
+ assert recurrent_state is not None
483
+ expected_outputs.append(token_output)
484
+ expected_state[state_indices[seq, token]] = recurrent_state.transpose(
485
+ -1, -2
486
+ )
487
+ expected = torch.cat(expected_outputs, dim=1)
488
+
489
+ actual_state = state.clone()
490
+ actual, _ = fused_recurrent_kda(
491
+ q=q,
492
+ k=k,
493
+ v=v,
494
+ raw_g=raw_g,
495
+ raw_beta=raw_beta,
496
+ A_log=A_log,
497
+ dt_bias=dt_bias,
498
+ lower_bound=lower_bound,
499
+ initial_state=actual_state,
500
+ cu_seqlens=cu_seqlens,
501
+ ssm_state_indices=state_indices,
502
+ num_accepted_tokens=num_accepted_tokens,
503
+ out=output,
504
+ fuse_gate=fuse_gate,
505
+ )
506
+
507
+ assert actual.data_ptr() == output.data_ptr()
508
+ assert_close("o", expected, actual, 1e-3, err_atol=1e-3)
509
+ used_states = state_indices.flatten().long()
510
+ assert_close(
511
+ "ht",
512
+ expected_state[used_states],
513
+ actual_state[used_states],
514
+ 3e-3,
515
+ err_atol=3e-3,
516
+ )
517
+ assert torch.isnan(output_storage[..., H * D :]).all()
518
+
519
+
520
+@pytest.mark.parametrize(
521
+ ("num_heads", "num_seqs", "lower_bound", "fuse_output_norm"),
522
+ [
523
+ (12, 1, -5.0, True),
524
+ (12, 4, None, False),
525
+ (24, 4, None, False),
526
+ (48, 1, -5.0, True),
527
+ (96, 1, -5.0, True),
528
+ ],
529
+)
530
+@torch.inference_mode()
531
+def test_fused_kda_decode_correctness(
532
+ num_heads: int,
533
+ num_seqs: int,
534
+ lower_bound: float | None,
535
+ fuse_output_norm: bool,
536
+):
537
+ D, W = 128, 4
538
+ if not is_fused_kda_decode_supported(
539
+ num_heads,
540
+ D,
541
+ W,
542
+ num_spec=0,
543
+ input_dtype=torch.bfloat16,
544
+ conv_state_dtype=torch.bfloat16,
545
+ ):
546
+ pytest.skip("Fused KDA decode is not supported on this platform")
547
+ torch.manual_seed(967 + num_heads + num_seqs)
548
+ dim = num_heads * D
549
+ slots = num_seqs + 2
550
+ packed_x_storage = torch.randn(
551
+ num_seqs, 3 * dim + 17, dtype=torch.bfloat16, device=DEVICE
552
+ )
553
+ packed_x = packed_x_storage[:, : 3 * dim]
554
+ weight = 0.1 * torch.randn(3 * dim, W, dtype=torch.float32, device=DEVICE)
555
+ conv_seed = 0.1 * torch.randn(
556
+ slots,
557
+ W - 1,
558
+ 3 * dim,
559
+ dtype=torch.bfloat16,
560
+ device=DEVICE,
561
+ ).transpose(1, 2)
562
+ raw_g = torch.randn(
563
+ 1,
564
+ num_seqs,
565
+ num_heads,
566
+ D,
567
+ dtype=torch.bfloat16,
568
+ device=DEVICE,
569
+ )
570
+ raw_beta_storage = torch.randn(
571
+ 1,
572
+ num_seqs,
573
+ num_heads + 1,
574
+ dtype=torch.bfloat16,
575
+ device=DEVICE,
576
+ )
577
+ raw_beta = raw_beta_storage[:, :, :num_heads]
578
+ output_gate_storage = torch.randn(
579
+ num_seqs,
580
+ dim + 7,
581
+ dtype=torch.bfloat16,
582
+ device=DEVICE,
583
+ )
584
+ output_gate = output_gate_storage[:, :dim].view(num_seqs, num_heads, D)
585
+ norm_weight = torch.randn(D, dtype=torch.float32, device=DEVICE)
586
+ norm_eps = 1e-5
587
+ A_log = 0.5 * torch.randn(num_heads, dtype=torch.float32, device=DEVICE)
588
+ dt_bias = 0.1 * torch.randn(dim, dtype=torch.float32, device=DEVICE)
589
+ state_indices = torch.arange(
590
+ num_seqs,
591
+ 0,
592
+ -1,
593
+ dtype=torch.int32,
594
+ device=DEVICE,
595
+ )
596
+ state_seed = 0.01 * torch.randn(
597
+ slots,
598
+ num_heads,
599
+ D,
600
+ D,
601
+ dtype=torch.float32,
602
+ device=DEVICE,
603
+ )
604
+
605
+ conv_ref = conv_seed.clone()
606
+ state_ref = state_seed.clone()
607
+ mixed_qkv = causal_conv1d_update(
608
+ packed_x,
609
+ conv_ref,
610
+ weight,
611
+ activation="silu",
612
+ conv_state_indices=state_indices,
613
+ validate_data=True,
614
+ out=torch.empty_like(packed_x),
615
+ )
616
+ expected, _ = fused_recurrent_kda_packed_decode(
617
+ mixed_qkv=mixed_qkv,
618
+ raw_g=raw_g,
619
+ raw_beta=raw_beta,
620
+ A_log=A_log,
621
+ dt_bias=dt_bias,
622
+ lower_bound=lower_bound,
623
+ initial_state=state_ref,
624
+ state_indices=state_indices,
625
+ )
626
+ if fuse_output_norm:
627
+ expected_float = expected.float()
628
+ expected = (
629
+ expected_float
630
+ * torch.rsqrt(expected_float.square().mean(dim=-1, keepdim=True) + norm_eps)
631
+ * norm_weight
632
+ * output_gate.float().sigmoid().unsqueeze(0)
633
+ ).to(expected.dtype)
634
+
635
+ conv_slot_elements = 3 * dim * (W - 1)
636
+ state_slot_elements = num_heads * D * D
637
+ conv_slot_bytes = conv_slot_elements * torch.bfloat16.itemsize
638
+ page_bytes = conv_slot_bytes + state_slot_elements * torch.float32.itemsize
639
+ cache_storage = torch.empty(slots * page_bytes, dtype=torch.uint8, device=DEVICE)
640
+ conv_actual = torch.as_strided(
641
+ cache_storage.view(torch.bfloat16),
642
+ size=(slots, 3 * dim, W - 1),
643
+ stride=(page_bytes // torch.bfloat16.itemsize, 1, 3 * dim),
644
+ )
645
+ state_actual = torch.as_strided(
646
+ cache_storage.view(torch.float32),
647
+ size=(slots, num_heads, D, D),
648
+ stride=(page_bytes // torch.float32.itemsize, D * D, D, 1),
649
+ storage_offset=conv_slot_bytes // torch.float32.itemsize,
650
+ )
651
+ conv_actual.copy_(conv_seed)
652
+ state_actual.copy_(state_seed)
653
+ fused_weight = weight.reshape(3, dim, W).transpose(1, 2).contiguous()
654
+ actual = ops.fused_kda_decode(
655
+ x=packed_x,
656
+ weight=fused_weight,
657
+ bias=None,
658
+ conv_state=conv_actual,
659
+ raw_g=raw_g,
660
+ raw_beta=raw_beta,
661
+ A_log=A_log,
662
+ dt_bias=dt_bias,
663
+ state_indices=state_indices,
664
+ state=state_actual,
665
+ lower_bound=lower_bound,
666
+ output_gate=output_gate if fuse_output_norm else None,
667
+ norm_weight=norm_weight if fuse_output_norm else None,
668
+ norm_eps=norm_eps,
669
+ )
670
+
671
+ torch.testing.assert_close(actual, expected, atol=3e-2, rtol=3e-2)
672
+ torch.testing.assert_close(conv_actual, conv_ref, atol=0, rtol=0)
673
+ torch.testing.assert_close(state_actual, state_ref, atol=3e-2, rtol=3e-2)
674
+
675
+
676
+def test_fused_kda_decode_rejects_speculative_conv_state():
677
+ assert not is_fused_kda_decode_supported(
678
+ num_heads=12,
679
+ head_dim=128,
680
+ conv_width=4,
681
+ num_spec=2,
682
+ input_dtype=torch.bfloat16,
683
+ conv_state_dtype=torch.bfloat16,
684
+ )
685
+
686
+
687
+@torch.inference_mode()
688
+def test_flashkda_correctness():
689
+ if not is_flashkda_supported(128, torch.bfloat16, -3.0):
690
+ pytest.skip("FlashKDA is not supported on this platform")
691
+
692
+ import vllm._flashkda_C # noqa: F401
693
+
694
+ B, T, H, D = 1, 48, 2, 128
695
+ torch.manual_seed(11)
696
+ q, k, v, raw_g = [
697
+ torch.randn(B, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4)
698
+ ]
699
+ beta_logits = torch.randn(B, T, H, dtype=torch.bfloat16, device=DEVICE)
700
+ A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5
701
+ dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1
702
+ initial_state = torch.randn(2, H, D, D, dtype=torch.float32, device=DEVICE)
703
+ cu_seqlens = torch.tensor([0, 17, T], dtype=torch.int32, device=DEVICE)
704
+ lower_bound = -3.0
705
+
706
+ gate = lower_bound * torch.sigmoid(
707
+ A_log.exp()[None, None, :, None] * (raw_g.float() + dt_bias[None, None, :, :])
708
+ )
709
+ beta = beta_logits.float().sigmoid()
710
+ q_norm = l2norm_fwd(q.contiguous())
711
+ k_norm = l2norm_fwd(k.contiguous())
712
+
713
+ expected_outputs = []
714
+ expected_states = []
715
+ for i, (start, end) in enumerate(
716
+ zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist())
717
+ ):
718
+ output, final_state = naive_recurrent_kda(
719
+ q_norm[:, start:end],
720
+ k_norm[:, start:end],
721
+ v[:, start:end],
722
+ gate[:, start:end],
723
+ beta[:, start:end],
724
+ initial_state=initial_state[i].transpose(-1, -2),
725
+ output_final_state=True,
726
+ )
727
+ expected_outputs.append(output)
728
+ expected_states.append(final_state)
729
+ expected_out = torch.cat(expected_outputs, dim=1)
730
+ expected_state = torch.cat(expected_states).transpose(-1, -2).contiguous()
731
+
732
+ actual_out = torch.empty_like(v)
733
+ actual_state = torch.empty_like(initial_state)
734
+ workspace = torch.empty(
735
+ torch.ops._flashkda_C.get_workspace_size(T, H, cu_seqlens.numel() - 1),
736
+ dtype=torch.uint8,
737
+ device=DEVICE,
738
+ )
739
+ torch.ops._flashkda_C.fwd(
740
+ q,
741
+ k,
742
+ v,
743
+ raw_g,
744
+ beta_logits,
745
+ D**-0.5,
746
+ actual_out,
747
+ workspace,
748
+ A_log,
749
+ dt_bias,
750
+ lower_bound,
751
+ initial_state,
752
+ actual_state,
753
+ cu_seqlens,
754
+ )
755
+
756
+ assert_close("o", expected_out, actual_out, 0.01)
757
+ assert_close("ht", expected_state, actual_state, 0.01)
tests/models/kimi_k3/test_kda_metadata.py
new
+411
@@ -0,0 +1,411 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from dataclasses import fields
5
+
6
+import pytest
7
+import torch
8
+
9
+from tests.v1.attention.utils import (
10
+ BatchSpec,
11
+ create_common_attn_metadata,
12
+ create_vllm_config,
13
+)
14
+from vllm.config import SpeculativeConfig
15
+from vllm.config.compilation import CUDAGraphMode
16
+from vllm.models.kimi_k3.nvidia.kda_metadata import (
17
+ KimiK3KDAAttentionBackend,
18
+ KimiK3KDAMetadata,
19
+ KimiK3KDAMetadataBuilder,
20
+ _mamba_get_block_table_tensor,
21
+ stage_spec_decode_metadata,
22
+)
23
+from vllm.v1.attention.backend import AttentionMetadataBuilder
24
+from vllm.v1.attention.backends.gdn_attn import (
25
+ GDNAttentionBackend,
26
+ GDNAttentionMetadata,
27
+ GDNAttentionMetadataBuilder,
28
+)
29
+from vllm.v1.attention.backends.utils import (
30
+ NULL_BLOCK_ID,
31
+ mamba_get_block_table_tensor,
32
+)
33
+from vllm.v1.kv_cache_interface import MambaSpec
34
+
35
+BLOCK_SIZE = 16
36
+DEVICE = torch.device("cpu")
37
+PRUNED_METADATA_FIELDS = {
38
+ "chunk_indices",
39
+ "chunk_offsets",
40
+ "prefill_query_start_loc",
41
+ "prefill_state_indices",
42
+ "prefill_has_initial_state",
43
+ "spec_sequence_masks",
44
+}
45
+
46
+
47
+def _assert_matches_shared_gdn(reference, actual: KimiK3KDAMetadata):
48
+ for field in fields(KimiK3KDAMetadata):
49
+ actual_value = getattr(actual, field.name)
50
+ expected_value = getattr(reference, field.name)
51
+ if field.name in PRUNED_METADATA_FIELDS:
52
+ assert actual_value is None
53
+ continue
54
+ if (
55
+ field.name in {"spec_token_indx", "non_spec_token_indx"}
56
+ and actual.num_spec_decodes > 0
57
+ and actual.num_prefills == 0
58
+ and actual.num_decodes == 0
59
+ ):
60
+ assert actual_value is None
61
+ continue
62
+ if isinstance(actual_value, torch.Tensor):
63
+ torch.testing.assert_close(actual_value, expected_value)
64
+ elif field.name == "nums_dict":
65
+ assert (actual_value is None) == (expected_value is None)
66
+ if actual_value is not None:
67
+ assert actual_value[8]["tot"] == expected_value[8]["tot"]
68
+ torch.testing.assert_close(
69
+ actual_value[8]["nums"], expected_value[8]["nums"]
70
+ )
71
+ else:
72
+ assert actual_value == expected_value
73
+
74
+
75
+def _make_builder(
76
+ builder_cls: type[AttentionMetadataBuilder],
77
+ num_speculative_tokens: int,
78
+ full_cuda_graph: bool,
79
+ device: torch.device = DEVICE,
80
+ mamba_cache_mode: str = "none",
81
+) -> AttentionMetadataBuilder:
82
+ vllm_config = create_vllm_config(
83
+ model_name="Qwen/Qwen3.5-0.8B",
84
+ block_size=BLOCK_SIZE,
85
+ )
86
+ if num_speculative_tokens:
87
+ vllm_config.speculative_config = SpeculativeConfig(
88
+ method="ngram",
89
+ num_speculative_tokens=num_speculative_tokens,
90
+ )
91
+ vllm_config.compilation_config.cudagraph_mode = (
92
+ CUDAGraphMode.FULL_AND_PIECEWISE if full_cuda_graph else CUDAGraphMode.NONE
93
+ )
94
+ vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode
95
+ return builder_cls(
96
+ kv_cache_spec=MambaSpec(
97
+ block_size=BLOCK_SIZE,
98
+ shapes=((16, 64),),
99
+ dtypes=(torch.float16,),
100
+ num_speculative_blocks=num_speculative_tokens,
101
+ ),
102
+ layer_names=["layer.0"],
103
+ vllm_config=vllm_config,
104
+ device=device,
105
+ )
106
+
107
+
108
+@pytest.mark.parametrize(
109
+ (
110
+ "batch",
111
+ "num_decode_draft_tokens",
112
+ "num_speculative_tokens",
113
+ "full_cuda_graph",
114
+ "is_prefilling",
115
+ ),
116
+ [
117
+ pytest.param(
118
+ BatchSpec(seq_lens=[50, 30], query_lens=[3, 3]),
119
+ [2, 2],
120
+ 2,
121
+ False,
122
+ [False, False],
123
+ id="pure-spec-decode",
124
+ ),
125
+ pytest.param(
126
+ BatchSpec(seq_lens=[100, 65, 20], query_lens=[50, 1, 3]),
127
+ [-1, -1, 2],
128
+ 2,
129
+ False,
130
+ [True, False, False],
131
+ id="mixed-prefill-and-spec-decode",
132
+ ),
133
+ pytest.param(
134
+ BatchSpec(seq_lens=[40, 30], query_lens=[1, 1]),
135
+ None,
136
+ 0,
137
+ False,
138
+ [False, False],
139
+ id="regular-decode",
140
+ ),
141
+ pytest.param(
142
+ BatchSpec(seq_lens=[40, 30], query_lens=[1, 1]),
143
+ [0, 0],
144
+ 2,
145
+ False,
146
+ [False, False],
147
+ id="no-scheduled-draft-tokens",
148
+ ),
149
+ ],
150
+)
151
+def test_kimi_k3_kda_metadata_matches_shared_gdn(
152
+ batch: BatchSpec,
153
+ num_decode_draft_tokens: list[int] | None,
154
+ num_speculative_tokens: int,
155
+ full_cuda_graph: bool,
156
+ is_prefilling: list[bool],
157
+):
158
+ kwargs: dict[str, torch.Tensor] = {}
159
+ if num_decode_draft_tokens is not None:
160
+ kwargs = {
161
+ "num_decode_draft_tokens_cpu": torch.tensor(
162
+ num_decode_draft_tokens, dtype=torch.int32
163
+ ),
164
+ "num_accepted_tokens": torch.ones(
165
+ batch.batch_size, dtype=torch.int32, device=DEVICE
166
+ ),
167
+ }
168
+
169
+ common_attn_metadata = create_common_attn_metadata(
170
+ batch, BLOCK_SIZE, DEVICE
171
+ ).replace(is_prefilling=torch.tensor(is_prefilling, dtype=torch.bool))
172
+ reference = _make_builder(
173
+ GDNAttentionMetadataBuilder,
174
+ num_speculative_tokens,
175
+ full_cuda_graph,
176
+ ).build(
177
+ 0,
178
+ common_attn_metadata,
179
+ **kwargs,
180
+ )
181
+ actual = _make_builder(
182
+ KimiK3KDAMetadataBuilder,
183
+ num_speculative_tokens,
184
+ full_cuda_graph,
185
+ ).build(0, common_attn_metadata, **kwargs)
186
+
187
+ assert isinstance(actual, KimiK3KDAMetadata)
188
+ _assert_matches_shared_gdn(reference, actual)
189
+
190
+
191
+def test_mixed_regular_and_spec_decode_uses_packed_decode_metadata():
192
+ batch = BatchSpec(seq_lens=[100, 65, 20], query_lens=[1, 1, 3])
193
+ common_attn_metadata = create_common_attn_metadata(
194
+ batch, BLOCK_SIZE, DEVICE
195
+ ).replace(is_prefilling=torch.tensor([False, False, False]))
196
+ actual = _make_builder(
197
+ KimiK3KDAMetadataBuilder,
198
+ num_speculative_tokens=2,
199
+ full_cuda_graph=False,
200
+ ).build(
201
+ 0,
202
+ common_attn_metadata,
203
+ num_decode_draft_tokens_cpu=torch.tensor([-1, -1, 2], dtype=torch.int32),
204
+ num_accepted_tokens=torch.ones(3, dtype=torch.int32, device=DEVICE),
205
+ )
206
+
207
+ # The K3 layer dispatches the non-spec subgroup to packed decode whenever
208
+ # it contains no prefill request.
209
+ assert actual.num_decodes == 2
210
+ assert actual.num_decode_tokens == 2
211
+ assert actual.num_prefills == 0
212
+ assert actual.num_prefill_tokens == 0
213
+ assert actual.has_initial_state is None
214
+ assert actual.nums_dict is None
215
+ assert actual.non_spec_query_start_loc is None
216
+ torch.testing.assert_close(actual.non_spec_token_indx, torch.tensor([0, 1]))
217
+ torch.testing.assert_close(actual.spec_token_indx, torch.tensor([2, 3, 4]))
218
+ torch.testing.assert_close(
219
+ actual.spec_query_start_loc,
220
+ torch.tensor([0, 3], dtype=torch.int32),
221
+ )
222
+
223
+
224
+def test_mixed_regular_and_spec_decode_excludes_request_padding():
225
+ batch = BatchSpec(seq_lens=[16, 65, 20], query_lens=[0, 1, 3])
226
+ common_attn_metadata = create_common_attn_metadata(
227
+ batch, BLOCK_SIZE, DEVICE
228
+ ).replace(is_prefilling=torch.tensor([False, False, False]))
229
+ actual = _make_builder(
230
+ KimiK3KDAMetadataBuilder,
231
+ num_speculative_tokens=2,
232
+ full_cuda_graph=False,
233
+ ).build(
234
+ 0,
235
+ common_attn_metadata,
236
+ num_decode_draft_tokens_cpu=torch.tensor([-1, -1, 2], dtype=torch.int32),
237
+ num_accepted_tokens=torch.ones(3, dtype=torch.int32, device=DEVICE),
238
+ )
239
+
240
+ assert actual.num_decodes == 1
241
+ assert actual.non_spec_state_indices_tensor is not None
242
+ assert actual.non_spec_state_indices_tensor.shape == (1,)
243
+ torch.testing.assert_close(actual.non_spec_token_indx, torch.tensor([0]))
244
+ torch.testing.assert_close(actual.spec_token_indx, torch.tensor([1, 2, 3]))
245
+
246
+
247
+@pytest.mark.parametrize(
248
+ ("seq_len", "expected_has_initial_state"),
249
+ [
250
+ pytest.param(1, False, id="first-token-prefill"),
251
+ pytest.param(65, True, id="final-one-token-prefill-chunk"),
252
+ ],
253
+)
254
+def test_mixed_one_token_prefill_and_spec_decode_uses_prefill_metadata(
255
+ seq_len: int,
256
+ expected_has_initial_state: bool,
257
+):
258
+ batch = BatchSpec(seq_lens=[seq_len, 20], query_lens=[1, 3])
259
+ common_attn_metadata = create_common_attn_metadata(
260
+ batch, BLOCK_SIZE, DEVICE
261
+ ).replace(is_prefilling=torch.tensor([True, False]))
262
+ actual = _make_builder(
263
+ KimiK3KDAMetadataBuilder,
264
+ num_speculative_tokens=2,
265
+ full_cuda_graph=False,
266
+ ).build(
267
+ 0,
268
+ common_attn_metadata,
269
+ num_decode_draft_tokens_cpu=torch.tensor([-1, 2], dtype=torch.int32),
270
+ num_accepted_tokens=torch.ones(2, dtype=torch.int32, device=DEVICE),
271
+ )
272
+
273
+ assert actual.num_prefills == 1
274
+ assert actual.num_prefill_tokens == 1
275
+ assert actual.num_decodes == 0
276
+ assert actual.num_decode_tokens == 0
277
+ assert actual.has_initial_state is not None
278
+ assert actual.has_initial_state.tolist() == [expected_has_initial_state]
279
+ assert actual.non_spec_query_start_loc is not None
280
+ torch.testing.assert_close(
281
+ actual.non_spec_query_start_loc,
282
+ torch.tensor([0, 1], dtype=torch.int32),
283
+ )
284
+
285
+
286
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
287
+def test_kimi_k3_kda_cudagraph_capture_matches_shared_gdn():
288
+ device = torch.device("cuda")
289
+ batch = BatchSpec(seq_lens=[50, 30], query_lens=[3, 3])
290
+ common_attn_metadata = create_common_attn_metadata(
291
+ batch, BLOCK_SIZE, device
292
+ ).replace(is_prefilling=torch.tensor([False, False]))
293
+ reference = _make_builder(
294
+ GDNAttentionMetadataBuilder,
295
+ num_speculative_tokens=2,
296
+ full_cuda_graph=True,
297
+ device=device,
298
+ ).build_for_cudagraph_capture(common_attn_metadata)
299
+ actual = _make_builder(
300
+ KimiK3KDAMetadataBuilder,
301
+ num_speculative_tokens=2,
302
+ full_cuda_graph=True,
303
+ device=device,
304
+ ).build_for_cudagraph_capture(common_attn_metadata)
305
+
306
+ assert isinstance(actual, KimiK3KDAMetadata)
307
+ _assert_matches_shared_gdn(reference, actual)
308
+
309
+
310
+def test_kimi_k3_kda_backend_uses_private_metadata_builder():
311
+ assert KimiK3KDAAttentionBackend.get_builder_cls() is KimiK3KDAMetadataBuilder
312
+ assert KimiK3KDAAttentionBackend.is_ssm()
313
+ assert issubclass(KimiK3KDAAttentionBackend, GDNAttentionBackend)
314
+ assert issubclass(KimiK3KDAMetadata, GDNAttentionMetadata)
315
+ assert issubclass(KimiK3KDAMetadataBuilder, GDNAttentionMetadataBuilder)
316
+
317
+
318
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
319
+def test_stage_spec_decode_metadata_matches_pytorch():
320
+ device = torch.device("cuda")
321
+ num_spec_decodes = 33
322
+ batch_size = 65
323
+ num_state_slots = 3
324
+ state_indices = torch.arange(
325
+ num_spec_decodes * 32,
326
+ dtype=torch.int32,
327
+ device=device,
328
+ ).reshape(num_spec_decodes, 32)[:, :num_state_slots]
329
+ query_start_loc = (
330
+ torch.arange(num_spec_decodes + 1, dtype=torch.int32, device=device)
331
+ * num_state_slots
332
+ )
333
+ num_accepted_tokens = (
334
+ torch.arange(num_spec_decodes, dtype=torch.int32, device=device)
335
+ % num_state_slots
336
+ + 1
337
+ )
338
+
339
+ staged_state_indices = torch.empty(
340
+ (batch_size, num_state_slots), dtype=torch.int32, device=device
341
+ )
342
+ staged_query_start_loc = torch.empty(
343
+ batch_size + 1, dtype=torch.int32, device=device
344
+ )
345
+ staged_num_accepted_tokens = torch.empty(
346
+ batch_size, dtype=torch.int32, device=device
347
+ )
348
+ stage_spec_decode_metadata(
349
+ state_indices,
350
+ query_start_loc,
351
+ num_accepted_tokens,
352
+ staged_state_indices,
353
+ staged_query_start_loc,
354
+ staged_num_accepted_tokens,
355
+ num_spec_decodes=num_spec_decodes,
356
+ )
357
+
358
+ expected_state_indices = torch.full_like(staged_state_indices, NULL_BLOCK_ID)
359
+ expected_state_indices[:num_spec_decodes] = state_indices
360
+ expected_query_start_loc = torch.full(
361
+ (batch_size + 1,),
362
+ query_start_loc[-1],
363
+ dtype=torch.int32,
364
+ device=device,
365
+ )
366
+ expected_query_start_loc[: num_spec_decodes + 1] = query_start_loc
367
+ expected_num_accepted_tokens = torch.ones(
368
+ batch_size, dtype=torch.int32, device=device
369
+ )
370
+ expected_num_accepted_tokens[:num_spec_decodes] = num_accepted_tokens
371
+
372
+ torch.testing.assert_close(staged_state_indices, expected_state_indices)
373
+ torch.testing.assert_close(staged_query_start_loc, expected_query_start_loc)
374
+ torch.testing.assert_close(staged_num_accepted_tokens, expected_num_accepted_tokens)
375
+
376
+
377
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
378
+def test_aligned_block_table_matches_shared_gdn():
379
+ device = torch.device("cuda")
380
+ seq_lens = torch.tensor(
381
+ [0, 1, 15, 16, 17, 31, 32, 33, 511, 512, 513],
382
+ dtype=torch.int32,
383
+ device=device,
384
+ ).repeat(6)[:65]
385
+ block_table_storage = torch.arange(
386
+ seq_lens.numel() * 128,
387
+ dtype=torch.int32,
388
+ device=device,
389
+ ).reshape(seq_lens.numel(), 128)
390
+ block_table = block_table_storage[:, ::2]
391
+ kv_cache_spec = MambaSpec(
392
+ block_size=BLOCK_SIZE,
393
+ shapes=((16, 64),),
394
+ dtypes=(torch.float16,),
395
+ num_speculative_blocks=2,
396
+ )
397
+
398
+ expected = mamba_get_block_table_tensor(
399
+ block_table,
400
+ seq_lens,
401
+ kv_cache_spec,
402
+ "align",
403
+ )
404
+ actual = _mamba_get_block_table_tensor(
405
+ block_table,
406
+ seq_lens,
407
+ kv_cache_spec,
408
+ "align",
409
+ )
410
+
411
+ torch.testing.assert_close(actual, expected)
tests/models/kimi_k3/test_latent_moe_tail.py
new
+145
@@ -0,0 +1,145 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import ray
6
+import torch
7
+import torch.distributed as dist
8
+import torch.nn.functional as F
9
+
10
+from tests.utils import (
11
+ init_test_distributed_environment,
12
+ multi_gpu_test,
13
+ multi_process_parallel,
14
+)
15
+from vllm.distributed import get_tp_group
16
+from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
17
+from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import KimiK3LatentMoETailOp
18
+from vllm.platforms import current_platform
19
+
20
+HIDDEN_SIZE = 7168
21
+LATENT_SIZE = 3584
22
+EPS = 0.1
23
+
24
+
25
+@ray.remote(num_gpus=1, max_calls=1)
26
+def _test_latent_moe_tail_worker(
27
+ monkeypatch: pytest.MonkeyPatch,
28
+ tp_size: int,
29
+ pp_size: int,
30
+ rank: int,
31
+ distributed_init_port: str,
32
+) -> None:
33
+ monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
34
+ device = torch.device(f"cuda:{rank}")
35
+ torch.accelerator.set_device_index(device)
36
+ init_test_distributed_environment(
37
+ tp_size,
38
+ pp_size,
39
+ rank,
40
+ distributed_init_port,
41
+ )
42
+
43
+ torch.manual_seed(0)
44
+ rms_weight = 1 + 0.1 * torch.randn(
45
+ LATENT_SIZE,
46
+ device=device,
47
+ dtype=torch.bfloat16,
48
+ )
49
+ up_weight = (
50
+ torch.randn(
51
+ HIDDEN_SIZE,
52
+ LATENT_SIZE,
53
+ device=device,
54
+ dtype=torch.bfloat16,
55
+ )
56
+ / LATENT_SIZE**0.5
57
+ )
58
+
59
+ group = get_tp_group().device_group
60
+ op = KimiK3LatentMoETailOp.initialize(
61
+ hidden_size=HIDDEN_SIZE,
62
+ latent_size=LATENT_SIZE,
63
+ dtype=torch.bfloat16,
64
+ device=device,
65
+ rms_eps=EPS,
66
+ )
67
+ cutedsl_warmup()
68
+
69
+ for iteration, num_tokens in enumerate((1, 5, 8, 16, 5)):
70
+ torch.manual_seed(100 * iteration + rank + 1)
71
+ routed_output = torch.randn(
72
+ num_tokens,
73
+ LATENT_SIZE,
74
+ device=device,
75
+ dtype=torch.bfloat16,
76
+ ).mul_(0.01)
77
+ shared_output = torch.randn(
78
+ num_tokens,
79
+ HIDDEN_SIZE,
80
+ device=device,
81
+ dtype=torch.bfloat16,
82
+ )
83
+
84
+ routed_reference = routed_output.clone()
85
+ shared_reference = shared_output.clone()
86
+ dist.all_reduce(routed_reference, group=group)
87
+ dist.all_reduce(shared_reference, group=group)
88
+ expected = F.linear(
89
+ F.rms_norm(
90
+ routed_reference,
91
+ (LATENT_SIZE,),
92
+ rms_weight,
93
+ EPS,
94
+ ),
95
+ up_weight,
96
+ )
97
+ expected.add_(shared_reference)
98
+
99
+ actual = op(
100
+ routed_output,
101
+ shared_output,
102
+ rms_weight,
103
+ up_weight,
104
+ )
105
+ torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
106
+ assert actual.is_contiguous()
107
+
108
+ graph = torch.cuda.CUDAGraph()
109
+ with torch.cuda.graph(graph):
110
+ graph_output = op(
111
+ routed_output,
112
+ shared_output,
113
+ rms_weight,
114
+ up_weight,
115
+ )
116
+ graph.replay()
117
+ torch.testing.assert_close(graph_output, expected, atol=8e-2, rtol=3e-2)
118
+
119
+
120
+def _run_latent_moe_tail_test(
121
+ monkeypatch: pytest.MonkeyPatch,
122
+ tp_size: int,
123
+) -> None:
124
+ if not current_platform.is_device_capability_family(100):
125
+ pytest.skip("K3 latent-MoE tail fusion requires SM100")
126
+ multi_process_parallel(
127
+ monkeypatch,
128
+ tp_size,
129
+ 1,
130
+ _test_latent_moe_tail_worker,
131
+ )
132
+
133
+
134
+@multi_gpu_test(num_gpus=8)
135
+def test_latent_moe_tail_tp8_matches_native_path(
136
+ monkeypatch: pytest.MonkeyPatch,
137
+) -> None:
138
+ _run_latent_moe_tail_test(monkeypatch, 8)
139
+
140
+
141
+@multi_gpu_test(num_gpus=16)
142
+def test_latent_moe_tail_tp16_matches_native_path(
143
+ monkeypatch: pytest.MonkeyPatch,
144
+) -> None:
145
+ _run_latent_moe_tail_test(monkeypatch, 16)
tests/models/kimi_k3/test_sequence_parallel.py
new
+343
@@ -0,0 +1,343 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import MethodType, SimpleNamespace
5
+from unittest.mock import Mock
6
+
7
+import pytest
8
+import torch
9
+from torch import nn
10
+
11
+from vllm.config import ParallelConfig
12
+from vllm.models.kimi_k3.nvidia import model as kimi_model
13
+from vllm.models.kimi_k3.nvidia import mtp as kimi_mtp
14
+from vllm.models.kimi_k3.nvidia.ops import sequence_parallel as sp_ops
15
+from vllm.platforms import current_platform
16
+
17
+
18
+class _IdentityNorm(nn.Module):
19
+ def __init__(self, hidden_size: int = 2) -> None:
20
+ super().__init__()
21
+ self.weight = nn.Parameter(torch.ones(hidden_size), requires_grad=False)
22
+ self.variance_epsilon = 1e-5
23
+
24
+ def forward(
25
+ self,
26
+ hidden_states: torch.Tensor,
27
+ residual: torch.Tensor | None = None,
28
+ ):
29
+ if residual is None:
30
+ return hidden_states
31
+ return hidden_states, residual
32
+
33
+
34
+class _RecordingMoE(nn.Module):
35
+ def __init__(self) -> None:
36
+ super().__init__()
37
+ self.num_tokens = 0
38
+
39
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
40
+ self.num_tokens = hidden_states.shape[0]
41
+ return hidden_states
42
+
43
+
44
+class _Projection(nn.Module):
45
+ def __init__(self, hidden_size: int = 2) -> None:
46
+ super().__init__()
47
+ self.weight = nn.Parameter(
48
+ torch.ones(1, hidden_size),
49
+ requires_grad=False,
50
+ )
51
+
52
+
53
+class _SequenceParallelMTPBlock:
54
+ use_sequence_parallel = True
55
+
56
+ def __call__(
57
+ self,
58
+ *,
59
+ positions: torch.Tensor,
60
+ hidden_states: torch.Tensor,
61
+ residual: torch.Tensor | None,
62
+ ):
63
+ assert residual is None
64
+ return hidden_states * 2, None, hidden_states * 3
65
+
66
+
67
+def _mock_sequence_parallel_collectives(monkeypatch):
68
+ monkeypatch.setattr(
69
+ kimi_model,
70
+ "sp_reduce_scatter",
71
+ lambda tensor: tensor.chunk(2, dim=0)[0],
72
+ )
73
+ monkeypatch.setattr(
74
+ kimi_model,
75
+ "sp_shard",
76
+ lambda tensor: torch.nn.functional.pad(tensor, (0, 0, 0, 1))[:2],
77
+ )
78
+ monkeypatch.setattr(
79
+ kimi_model,
80
+ "sp_all_gather",
81
+ lambda tensor: torch.cat([tensor, tensor], dim=0),
82
+ )
83
+
84
+
85
+@pytest.mark.parametrize(
86
+ ("num_tokens", "is_padding", "tp_rank", "expected"),
87
+ [
88
+ (1, None, 0, [False]),
89
+ (1, None, 1, [True]),
90
+ (5, None, 2, [False, True]),
91
+ (5, None, 3, [True, True]),
92
+ (5, [False, True, False, False, False], 0, [False, True]),
93
+ ],
94
+)
95
+def test_sp_padding_mask_marks_added_rows(
96
+ monkeypatch,
97
+ num_tokens: int,
98
+ is_padding: list[bool] | None,
99
+ tp_rank: int,
100
+ expected: list[bool],
101
+):
102
+ monkeypatch.setattr(sp_ops, "get_tensor_model_parallel_world_size", lambda: 4)
103
+ monkeypatch.setattr(sp_ops, "get_tensor_model_parallel_rank", lambda: tp_rank)
104
+
105
+ hidden_states = torch.empty(num_tokens, 2)
106
+ padding = torch.tensor(is_padding) if is_padding is not None else None
107
+ actual = sp_ops.sp_padding_mask(padding, hidden_states)
108
+
109
+ torch.testing.assert_close(actual, torch.tensor(expected))
110
+
111
+
112
+@pytest.mark.parametrize(
113
+ ("data_parallel_size", "expected"),
114
+ [
115
+ (1, False),
116
+ (2, True),
117
+ ],
118
+)
119
+def test_moe_sequence_parallel_requires_data_parallel(
120
+ monkeypatch,
121
+ data_parallel_size: int,
122
+ expected: bool,
123
+):
124
+ monkeypatch.setattr(current_platform, "device_count", lambda: 2)
125
+ parallel_config = ParallelConfig(
126
+ tensor_parallel_size=2,
127
+ data_parallel_size=data_parallel_size,
128
+ enable_expert_parallel=True,
129
+ all2all_backend="allgather_reducescatter",
130
+ )
131
+
132
+ assert parallel_config.use_sequence_parallel_moe is expected
133
+
134
+
135
+def test_kimi_decoder_layer_keeps_moe_states_sequence_sharded(monkeypatch):
136
+ layer = object.__new__(kimi_model.KimiDecoderLayer)
137
+ nn.Module.__init__(layer)
138
+ layer.use_attn_res = False
139
+ layer.use_sequence_parallel = True
140
+ layer.input_layernorm = _IdentityNorm()
141
+ layer.post_attention_layernorm = _IdentityNorm()
142
+ layer.mlp = _RecordingMoE()
143
+ layer._run_self_attn = MethodType(
144
+ lambda self, positions, hidden_states: hidden_states,
145
+ layer,
146
+ )
147
+
148
+ _mock_sequence_parallel_collectives(monkeypatch)
149
+
150
+ positions = torch.arange(3)
151
+ full_hidden_states = torch.arange(6, dtype=torch.float32).view(3, 2)
152
+ hidden_states = kimi_model.sp_shard(full_hidden_states)
153
+ hidden_states, prefix_sum, residual = layer(
154
+ positions=positions,
155
+ hidden_states=hidden_states,
156
+ residual=None,
157
+ )
158
+
159
+ assert prefix_sum is None
160
+ assert hidden_states.shape == residual.shape == (2, 2)
161
+ assert layer.mlp.num_tokens == 2
162
+
163
+ hidden_states, prefix_sum, residual = layer(
164
+ positions=positions,
165
+ hidden_states=hidden_states,
166
+ residual=residual,
167
+ )
168
+
169
+ assert prefix_sum is None
170
+ assert hidden_states.shape == residual.shape == (2, 2)
171
+ assert layer.mlp.num_tokens == 2
172
+
173
+
174
+def test_kimi_attn_residual_states_stay_sequence_sharded(monkeypatch):
175
+ layer = object.__new__(kimi_model.KimiDecoderLayer)
176
+ nn.Module.__init__(layer)
177
+ layer.use_attn_res = True
178
+ layer.use_sequence_parallel = True
179
+ layer.prev_valid_blocks = 0
180
+ layer.block_write_idx = 0
181
+ layer.is_block_write_layer = False
182
+ layer.input_layernorm = _IdentityNorm()
183
+ layer.post_attention_layernorm = _IdentityNorm()
184
+ layer.self_attention_res_norm = _IdentityNorm()
185
+ layer.mlp_res_norm = _IdentityNorm()
186
+ layer.self_attention_res_proj = _Projection()
187
+ layer.mlp_res_proj = _Projection()
188
+ layer.mlp = _RecordingMoE()
189
+ layer._run_self_attn = MethodType(
190
+ lambda self, positions, hidden_states: hidden_states,
191
+ layer,
192
+ )
193
+
194
+ _mock_sequence_parallel_collectives(monkeypatch)
195
+ monkeypatch.setattr(
196
+ kimi_model,
197
+ "attn_res",
198
+ lambda prefix_sum, hidden_states, *args, **kwargs: (
199
+ prefix_sum if hidden_states is None else prefix_sum + hidden_states
200
+ ),
201
+ )
202
+
203
+ prefix_sum = kimi_model.sp_shard(torch.arange(6, dtype=torch.float32).view(3, 2))
204
+ block_residual = torch.zeros(2, 1, 2)
205
+ hidden_states, prefix_sum, block_residual = layer(
206
+ positions=torch.arange(3),
207
+ hidden_states=None,
208
+ prefix_sum=prefix_sum,
209
+ residual=block_residual,
210
+ )
211
+
212
+ assert hidden_states.shape == prefix_sum.shape == (2, 2)
213
+ assert block_residual.shape == (2, 1, 2)
214
+ assert layer.mlp.num_tokens == 2
215
+
216
+
217
+def test_kimi_mtp_restores_sequence_parallel_output(monkeypatch):
218
+ layer = object.__new__(kimi_mtp.KimiK3MultiTokenPredictorLayer)
219
+ nn.Module.__init__(layer)
220
+ layer.enorm = _IdentityNorm()
221
+ layer.hnorm = _IdentityNorm()
222
+ layer.eh_proj = nn.Identity()
223
+ object.__setattr__(layer, "mtp_block", _SequenceParallelMTPBlock())
224
+
225
+ final_norm = Mock(side_effect=lambda hidden_states: hidden_states + 1)
226
+ object.__setattr__(
227
+ layer,
228
+ "shared_head",
229
+ SimpleNamespace(norm=final_norm),
230
+ )
231
+
232
+ monkeypatch.setattr(
233
+ kimi_mtp,
234
+ "fused_mtp_input",
235
+ lambda positions, inputs_embeds, *args: inputs_embeds,
236
+ )
237
+ monkeypatch.setattr(
238
+ kimi_mtp,
239
+ "sp_shard",
240
+ lambda tensor: torch.nn.functional.pad(tensor, (0, 0, 0, 1))[:2],
241
+ )
242
+ monkeypatch.setattr(
243
+ kimi_mtp,
244
+ "sp_all_gather",
245
+ lambda tensor: torch.cat([tensor, tensor], dim=0),
246
+ )
247
+
248
+ inputs_embeds = torch.arange(6, dtype=torch.float32).view(3, 2)
249
+ logits_hidden_states, hidden_states = layer(
250
+ input_ids=torch.zeros(3, dtype=torch.long),
251
+ positions=torch.arange(3),
252
+ previous_hidden_states=torch.zeros_like(inputs_embeds),
253
+ inputs_embeds=inputs_embeds,
254
+ )
255
+
256
+ sharded_states = torch.nn.functional.pad(inputs_embeds, (0, 0, 0, 1))[:2]
257
+ expected_hidden_states = torch.cat(
258
+ [sharded_states * 5, sharded_states * 5],
259
+ dim=0,
260
+ )[:3]
261
+ torch.testing.assert_close(hidden_states, expected_hidden_states)
262
+ torch.testing.assert_close(logits_hidden_states, expected_hidden_states + 1)
263
+ final_norm.assert_called_once()
264
+ torch.testing.assert_close(final_norm.call_args.args[0], expected_hidden_states)
265
+
266
+
267
+def test_sp_all_gather_uses_custom_kernel(monkeypatch):
268
+ hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2)
269
+ expected = torch.cat([hidden_states, hidden_states])
270
+ custom_all_gather = Mock(return_value=expected)
271
+ device_communicator = SimpleNamespace(
272
+ custom_all_gather=custom_all_gather,
273
+ )
274
+ monkeypatch.setattr(
275
+ sp_ops,
276
+ "get_tp_group",
277
+ lambda: SimpleNamespace(device_communicator=device_communicator),
278
+ )
279
+ fallback = Mock(side_effect=AssertionError("unexpected fallback"))
280
+ monkeypatch.setattr(sp_ops, "tensor_model_parallel_all_gather", fallback)
281
+
282
+ output = sp_ops.sp_all_gather(hidden_states)
283
+
284
+ torch.testing.assert_close(output, expected)
285
+ custom_all_gather.assert_called_once_with(hidden_states)
286
+ fallback.assert_not_called()
287
+
288
+
289
+def test_sp_reduce_scatter_uses_custom_kernel_after_padding(monkeypatch):
290
+ hidden_states = torch.arange(6, dtype=torch.float32).view(3, 2)
291
+ expected = torch.arange(4, dtype=torch.float32).view(2, 2)
292
+ custom_reduce_scatter = Mock(return_value=expected)
293
+ device_communicator = SimpleNamespace(
294
+ custom_reduce_scatter=custom_reduce_scatter,
295
+ )
296
+ monkeypatch.setattr(
297
+ sp_ops,
298
+ "get_tp_group",
299
+ lambda: SimpleNamespace(device_communicator=device_communicator),
300
+ )
301
+ monkeypatch.setattr(
302
+ sp_ops,
303
+ "get_tensor_model_parallel_world_size",
304
+ lambda: 2,
305
+ )
306
+ fallback = Mock(side_effect=AssertionError("unexpected fallback"))
307
+ monkeypatch.setattr(sp_ops, "tensor_model_parallel_reduce_scatter", fallback)
308
+
309
+ output = sp_ops.sp_reduce_scatter(hidden_states)
310
+
311
+ torch.testing.assert_close(output, expected)
312
+ padded = custom_reduce_scatter.call_args.args[0]
313
+ assert padded.shape == (4, 2)
314
+ torch.testing.assert_close(padded[:3], hidden_states)
315
+ torch.testing.assert_close(padded[3], torch.zeros(2))
316
+ fallback.assert_not_called()
317
+
318
+
319
+def test_sp_collectives_fall_back_without_custom_kernel(monkeypatch):
320
+ hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2)
321
+ monkeypatch.setattr(
322
+ sp_ops,
323
+ "get_tp_group",
324
+ lambda: SimpleNamespace(device_communicator=None),
325
+ )
326
+ monkeypatch.setattr(
327
+ sp_ops,
328
+ "get_tensor_model_parallel_world_size",
329
+ lambda: 2,
330
+ )
331
+ all_gather = Mock(return_value=hidden_states)
332
+ reduce_scatter = Mock(return_value=hidden_states)
333
+ monkeypatch.setattr(sp_ops, "tensor_model_parallel_all_gather", all_gather)
334
+ monkeypatch.setattr(
335
+ sp_ops,
336
+ "tensor_model_parallel_reduce_scatter",
337
+ reduce_scatter,
338
+ )
339
+
340
+ torch.testing.assert_close(sp_ops.sp_all_gather(hidden_states), hidden_states)
341
+ torch.testing.assert_close(sp_ops.sp_reduce_scatter(hidden_states), hidden_states)
342
+ all_gather.assert_called_once_with(hidden_states, 0)
343
+ reduce_scatter.assert_called_once_with(hidden_states, 0)
tests/models/multimodal/generation/test_common.py
+1
-1
@@ -1242,7 +1242,7 @@ def test_custom_inputs_models(
1242
create_new_process_for_each_test=True,
1243
),
1244
)
1245
-@create_new_process_for_each_test()
1245
+@create_new_process_for_each_test("spawn")
1246
def test_single_image_models_heavy(
1247
tmp_path: PosixPath,
1248
model_type: str,
vllm/_custom_ops.py
+131
-8
@@ -2733,6 +2733,52 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
2733
)
2734
2735
2736
+def fused_kda_decode(
2737
+ x: torch.Tensor,
2738
+ weight: torch.Tensor,
2739
+ bias: torch.Tensor | None,
2740
+ conv_state: torch.Tensor,
2741
+ raw_g: torch.Tensor,
2742
+ raw_beta: torch.Tensor,
2743
+ A_log: torch.Tensor,
2744
+ dt_bias: torch.Tensor,
2745
+ state_indices: torch.Tensor,
2746
+ state: torch.Tensor,
2747
+ out: torch.Tensor | None = None,
2748
+ lower_bound: float | None = None,
2749
+ output_gate: torch.Tensor | None = None,
2750
+ norm_weight: torch.Tensor | None = None,
2751
+ norm_eps: float = 1e-5,
2752
+) -> torch.Tensor:
2753
+ if out is None:
2754
+ out = torch.empty(
2755
+ 1,
2756
+ x.shape[0],
2757
+ raw_g.shape[2],
2758
+ raw_g.shape[3],
2759
+ dtype=x.dtype,
2760
+ device=x.device,
2761
+ )
2762
+ torch.ops._C.fused_kda_decode(
2763
+ x,
2764
+ weight,
2765
+ bias,
2766
+ conv_state,
2767
+ raw_g,
2768
+ raw_beta,
2769
+ A_log,
2770
+ dt_bias,
2771
+ state_indices,
2772
+ state,
2773
+ out,
2774
+ lower_bound,
2775
+ output_gate,
2776
+ norm_weight,
2777
+ norm_eps,
2778
+ )
2779
+ return out
2780
+
2781
+
2782
def concat_and_cache_mla(
2783
kv_c: torch.Tensor,
2784
k_pe: torch.Tensor,
@@ -2746,6 +2792,26 @@ def concat_and_cache_mla(
2792
)
2793
2794
2795
+def concat_and_cache_mla_grouped(
2796
+ kv_c: torch.Tensor,
2797
+ k_pe: torch.Tensor,
2798
+ kv_cache_ptrs: torch.Tensor,
2799
+ slot_mapping: torch.Tensor,
2800
+ block_size: int,
2801
+ block_stride: int,
2802
+ entry_stride: int,
2803
+) -> None:
2804
+ torch.ops._C_cache_ops.concat_and_cache_mla_grouped(
2805
+ kv_c,
2806
+ k_pe,
2807
+ kv_cache_ptrs,
2808
+ slot_mapping,
2809
+ block_size,
2810
+ block_stride,
2811
+ entry_stride,
2812
+ )
2813
+
2814
+
2815
def kimi_k3_attn_res(
2816
prefix: torch.Tensor,
2817
delta: torch.Tensor,
@@ -3040,6 +3106,63 @@ def all_reduce(
3106
torch.ops._C_custom_ar.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes)
3107
3108
3109
+def custom_all_gather(
3110
+ fa: int,
3111
+ inp: torch.Tensor,
3112
+ out: torch.Tensor,
3113
+ reg_buffer: int,
3114
+ reg_buffer_sz_bytes: int,
3115
+) -> None:
3116
+ torch.ops._C_custom_ar.custom_all_gather(
3117
+ fa, inp, out, reg_buffer, reg_buffer_sz_bytes
3118
+ )
3119
+
3120
+
3121
+def mnnvl_lamport_all_gather(
3122
+ fa: int,
3123
+ inp: torch.Tensor,
3124
+ out: torch.Tensor,
3125
+ local_buffer: int,
3126
+ multicast_buffer: int,
3127
+ epoch_buffer: int,
3128
+ stage_sz_bytes: int,
3129
+) -> None:
3130
+ torch.ops._C_custom_ar.mnnvl_lamport_all_gather(
3131
+ fa,
3132
+ inp,
3133
+ out,
3134
+ local_buffer,
3135
+ multicast_buffer,
3136
+ epoch_buffer,
3137
+ stage_sz_bytes,
3138
+ )
3139
+
3140
+
3141
+def custom_reduce_scatter(
3142
+ fa: int,
3143
+ inp: torch.Tensor,
3144
+ out: torch.Tensor,
3145
+ reg_buffer: int,
3146
+ reg_buffer_sz_bytes: int,
3147
+) -> None:
3148
+ torch.ops._C_custom_ar.custom_reduce_scatter(
3149
+ fa, inp, out, reg_buffer, reg_buffer_sz_bytes
3150
+ )
3151
+
3152
+
3153
+def mnnvl_lamport_reduce_scatter(
3154
+ fa: int,
3155
+ inp: torch.Tensor,
3156
+ out: torch.Tensor,
3157
+ local_buffer: int,
3158
+ epoch_buffer: int,
3159
+ stage_sz_bytes: int,
3160
+) -> None:
3161
+ torch.ops._C_custom_ar.mnnvl_lamport_reduce_scatter(
3162
+ fa, inp, out, local_buffer, epoch_buffer, stage_sz_bytes
3163
+ )
3164
+
3165
+
3166
def dispose(fa: int) -> None:
3167
torch.ops._C_custom_ar.dispose(fa)
3168
@@ -3144,18 +3267,18 @@ def dsv3_fused_a_gemm(
3267
output: torch.Tensor,
3268
mat_a: torch.Tensor,
3269
mat_b: torch.Tensor,
3270
+ enable_pdl: bool = False,
3271
) -> None:
3148
- """DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
3272
+ """Low-latency fused-A-style GEMM (SM 9.0+, BF16, 1-16 tokens).
3273
3150
- Computes output = mat_a @ mat_b.T where:
3151
- mat_a: [num_tokens, 7168] row-major bf16 (hidden states)
3152
- mat_b: [7168, 2112] column-major bf16 (weight transposed)
3153
- output: [num_tokens, 2112] row-major bf16
3274
+ Computes ``output = mat_a @ mat_b`` for the compiled Kimi K3 and
3275
+ DeepSeek V3 projection shapes. ``mat_a`` and ``output`` are row-major;
3276
+ ``mat_b`` is the column-major transposed weight. ``enable_pdl`` permits
3277
+ programmatic dependent launch for callers that have validated it.
3278
3155
- Optimized for the DeepSeek V2/V3 QKV A-projection at small batch sizes.
3156
- Requires SM 9.0+ (Hopper).
3279
+ Requires SM 9.0+.
3280
"""
3158
- torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b)
3281
+ torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b, enable_pdl)
3282
3283
3284
if hasattr(torch.ops._C, "weight_packed_linear"):
vllm/envs.py
+6
@@ -269,6 +269,7 @@ if TYPE_CHECKING:
269
VLLM_NCCL_INCLUDE_PATH: str | None = None
270
VLLM_GC_DEBUG: str = ""
271
VLLM_DEBUG_WORKSPACE: bool = False
272
+ VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION: bool = False
273
VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
274
VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
275
VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024
@@ -1894,6 +1895,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
1895
# Debug workspace allocations.
1896
# logging of workspace resize operations.
1897
"VLLM_DEBUG_WORKSPACE": lambda: bool(int(os.getenv("VLLM_DEBUG_WORKSPACE", "0"))),
1898
+ # Enable the experimental Kimi K3 latent-MoE tail fusion.
1899
+ # Currently supported only on SM100 with TP=8/16 and BF16.
1900
+ "VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION": lambda: bool(
1901
+ int(os.getenv("VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION", "0"))
1902
+ ),
1903
# Disables parallel execution of shared_experts via separate cuda stream
1904
"VLLM_DISABLE_SHARED_EXPERTS_STREAM": lambda: bool(
1905
int(os.getenv("VLLM_DISABLE_SHARED_EXPERTS_STREAM", "0"))
vllm/model_executor/kernels/linear/cute_dsl/_skinny_gemm.py
new
+180
@@ -0,0 +1,180 @@
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 cutlass
7
+import cutlass.cute as cute
8
+from cuda.bindings.driver import CUstream
9
+from cutlass import const_expr
10
+
11
+
12
+class CuteSkinnyGemm:
13
+ """Shape-dynamic low-latency GEMM for small token counts.
14
+
15
+ Computes ``C[M, N] = A[M, K] @ B[N, K].T + residual`` with BF16 or FP16
16
+ inputs, FP32 accumulators, and an output matching the input dtype. The
17
+ 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
+ """
21
+
22
+ def __init__(
23
+ self,
24
+ *,
25
+ element_type,
26
+ num_rows: int,
27
+ block_size: int,
28
+ outputs_per_block: int,
29
+ vector_width: int = 8,
30
+ k_unroll: int = 1,
31
+ has_residual: bool = False,
32
+ use_pdl: bool = False,
33
+ ) -> None:
34
+ if block_size % cute.arch.WARP_SIZE != 0:
35
+ raise ValueError("block_size must be a multiple of the warp size")
36
+ self.element_type = element_type
37
+ self.num_rows = num_rows
38
+ self.block_size = block_size
39
+ self.outputs_per_block = outputs_per_block
40
+ self.vector_width = vector_width
41
+ self.k_unroll = k_unroll
42
+ self.has_residual = has_residual
43
+ self.use_pdl = use_pdl
44
+ self.num_warps = block_size // cute.arch.WARP_SIZE
45
+
46
+ @cute.jit
47
+ def __call__(
48
+ self,
49
+ gA: cute.Tensor,
50
+ gB: cute.Tensor,
51
+ gResidual: cute.Tensor,
52
+ gC: cute.Tensor,
53
+ stream: CUstream,
54
+ ) -> None:
55
+ n = cute.size(gB, mode=[0])
56
+ k = cute.size(gA, mode=[1])
57
+ copy_a = cute.make_copy_atom(
58
+ cute.nvgpu.CopyG2ROp(),
59
+ self.element_type,
60
+ num_bits_per_copy=self.vector_width * self.element_type.width,
61
+ load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
62
+ )
63
+ copy_b = cute.make_copy_atom(
64
+ cute.nvgpu.CopyG2ROp(),
65
+ self.element_type,
66
+ num_bits_per_copy=self.vector_width * self.element_type.width,
67
+ load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING,
68
+ )
69
+ self.kernel(gA, gB, gResidual, gC, k, copy_a, copy_b).launch(
70
+ grid=[cute.ceil_div(n, self.outputs_per_block), 1, 1],
71
+ block=[self.block_size, 1, 1],
72
+ smem=self.num_rows * self.outputs_per_block * self.num_warps * 4,
73
+ stream=stream,
74
+ use_pdl=self.use_pdl,
75
+ min_blocks_per_mp=1,
76
+ )
77
+
78
+ @cute.kernel
79
+ def kernel(
80
+ self,
81
+ gA: cute.Tensor,
82
+ gB: cute.Tensor,
83
+ gResidual: cute.Tensor,
84
+ gC: cute.Tensor,
85
+ k_extent: cutlass.Int32,
86
+ copy_a: cute.CopyAtom,
87
+ copy_b: cute.CopyAtom,
88
+ ) -> None:
89
+ tidx, _, _ = cute.arch.thread_idx()
90
+ block_idx, _, _ = cute.arch.block_idx()
91
+ warp_idx = cute.arch.warp_idx()
92
+
93
+ num_rows: cutlass.Constexpr = self.num_rows
94
+ outputs_per_block: cutlass.Constexpr = self.outputs_per_block
95
+ vector_width: cutlass.Constexpr = self.vector_width
96
+ block_size: cutlass.Constexpr = self.block_size
97
+ num_warps: cutlass.Constexpr = self.num_warps
98
+
99
+ acc_layout = cute.make_layout(
100
+ (num_rows, outputs_per_block), stride=(outputs_per_block, 1)
101
+ )
102
+ acc = cute.make_rmem_tensor(acc_layout, cutlass.Float32)
103
+ acc.fill(0.0)
104
+
105
+ if const_expr(self.use_pdl):
106
+ cute.arch.griddepcontrol_wait()
107
+
108
+ n_base = block_idx * outputs_per_block
109
+ k_tile_size: cutlass.Constexpr = block_size * vector_width
110
+ num_k_tiles = k_extent // k_tile_size
111
+
112
+ gA_vec = cute.logical_divide(gA, (None, vector_width))
113
+ gB_vec = cute.logical_divide(gB, (None, vector_width))
114
+ # Layout after both divides is (M/N, K_TILE, K_LANE, K_VEC).
115
+ tA_all = cute.logical_divide(gA_vec, (None, (None, block_size)))
116
+ tB_all = cute.logical_divide(gB_vec, (None, (None, block_size)))
117
+ tA = tA_all[None, (None, (tidx, None))]
118
+
119
+ a_regs = cute.make_rmem_tensor(
120
+ cute.make_layout((num_rows, vector_width), stride=(vector_width, 1)),
121
+ self.element_type,
122
+ )
123
+ b_regs = cute.make_rmem_tensor(
124
+ cute.make_layout(
125
+ (outputs_per_block, vector_width), stride=(vector_width, 1)
126
+ ),
127
+ self.element_type,
128
+ )
129
+
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])
133
+
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])
138
+
139
+ for vi in cutlass.range_constexpr(vector_width):
140
+ for mi in cutlass.range_constexpr(num_rows):
141
+ 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)
145
+
146
+ for mi in cutlass.range_constexpr(num_rows):
147
+ for ni in cutlass.range_constexpr(outputs_per_block):
148
+ acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni])
149
+
150
+ smem_layout = cute.make_layout(
151
+ (num_rows, outputs_per_block, num_warps),
152
+ stride=(outputs_per_block * num_warps, num_warps, 1),
153
+ )
154
+ smem = cutlass.utils.SmemAllocator()
155
+ partials = smem.allocate_tensor(cutlass.Float32, smem_layout, byte_alignment=16)
156
+ with cute.arch.elect_one():
157
+ for mi in cutlass.range_constexpr(num_rows):
158
+ for ni in cutlass.range_constexpr(outputs_per_block):
159
+ partials[mi, ni, warp_idx] = acc[mi, ni]
160
+
161
+ cute.arch.sync_threads()
162
+ if tidx == 0:
163
+ for mi in cutlass.range_constexpr(num_rows):
164
+ for ni in cutlass.range_constexpr(outputs_per_block):
165
+ n_idx = n_base + ni
166
+ total = (
167
+ partials[mi, ni, None]
168
+ .load()
169
+ .reduce(
170
+ cute.ReductionOp.ADD,
171
+ init_val=cutlass.Float32(0.0),
172
+ reduction_profile=0,
173
+ )
174
+ )
175
+ if const_expr(self.has_residual):
176
+ total += gResidual[mi, n_idx].to(cutlass.Float32)
177
+ gC[mi, n_idx] = cutlass.Float32(total).to(self.element_type)
178
+
179
+ if const_expr(self.use_pdl):
180
+ cute.arch.griddepcontrol_launch_dependents()
vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm.py
new
+252
@@ -0,0 +1,252 @@
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
+from collections.abc import Iterable
8
+from dataclasses import dataclass
9
+from functools import partial
10
+from typing import Any
11
+
12
+import torch
13
+
14
+logger = logging.getLogger(__name__)
15
+
16
+_cutedsl_available: bool | None = None
17
+
18
+
19
+@dataclass(frozen=True, slots=True)
20
+class SkinnyGemmConfig:
21
+ num_rows: int
22
+ block_size: int
23
+ outputs_per_block: int
24
+ k_unroll: int = 1
25
+ vector_width: int = 8
26
+
27
+
28
+class ShapeDynamicSkinnyGemm:
29
+ def __init__(self) -> None:
30
+ self._compiled: dict[tuple[torch.dtype, SkinnyGemmConfig, bool], Any] = {}
31
+ self._warmup_configs: set[tuple[torch.dtype, SkinnyGemmConfig, bool]] = set()
32
+ self._warmup_registered = False
33
+
34
+ @staticmethod
35
+ def is_available() -> bool:
36
+ global _cutedsl_available
37
+ if _cutedsl_available is not None:
38
+ return _cutedsl_available
39
+ try:
40
+ import cutlass # noqa: F401
41
+ import cutlass.cute # noqa: F401
42
+
43
+ _cutedsl_available = True
44
+ except ImportError:
45
+ _cutedsl_available = False
46
+ logger.info("cuteDSL is not available; skinny GEMM is disabled")
47
+ return _cutedsl_available
48
+
49
+ @staticmethod
50
+ def _config(m: int, n: int, k: int) -> SkinnyGemmConfig:
51
+ num_rows = m
52
+ wide_block = 224
53
+ if m == 1 and k >= 7168 and k % (wide_block * 8) == 0:
54
+ if n % 3 == 0:
55
+ k_unroll = 2 if n <= 2304 else 4
56
+ return SkinnyGemmConfig(num_rows, wide_block, 3, k_unroll)
57
+ if 2304 < n < 4096 and n % 2 == 0:
58
+ return SkinnyGemmConfig(num_rows, wide_block, 2, k_unroll=4)
59
+
60
+ if k <= 2048 or k % (128 * 8) != 0:
61
+ outputs_per_block = 2 if k <= 2048 else 4
62
+ if n % outputs_per_block:
63
+ outputs_per_block = 1
64
+ if k % (64 * 8) == 0:
65
+ return SkinnyGemmConfig(num_rows, 64, outputs_per_block, 2)
66
+ if k % (32 * 8) == 0:
67
+ return SkinnyGemmConfig(num_rows, 32, outputs_per_block, 2)
68
+ return SkinnyGemmConfig(num_rows, 32, outputs_per_block, 2, vector_width=4)
69
+
70
+ block_size = 64 if 4096 <= n < 8192 else 128
71
+ outputs_per_block = 1 if m == 1 and n <= 2304 else 2
72
+ if n % outputs_per_block:
73
+ outputs_per_block = 1
74
+ k_unroll = 2 if n <= 2304 or n >= 16384 else 1
75
+ return SkinnyGemmConfig(
76
+ num_rows,
77
+ block_size,
78
+ outputs_per_block,
79
+ k_unroll=k_unroll,
80
+ )
81
+
82
+ @staticmethod
83
+ def _cutlass_dtype(dtype: torch.dtype):
84
+ from cutlass import BFloat16, Float16
85
+
86
+ return BFloat16 if dtype == torch.bfloat16 else Float16
87
+
88
+ @staticmethod
89
+ def _stream():
90
+ from cuda.bindings.driver import CUstream
91
+
92
+ from vllm.utils.torch_utils import current_stream
93
+
94
+ return CUstream(current_stream().cuda_stream)
95
+
96
+ @staticmethod
97
+ def _use_pdl() -> bool:
98
+ from vllm.platforms import current_platform
99
+
100
+ return current_platform.is_arch_support_pdl()
101
+
102
+ def _compile(
103
+ self,
104
+ dtype: torch.dtype,
105
+ config: SkinnyGemmConfig,
106
+ has_residual: bool,
107
+ ) -> None:
108
+ import cutlass.cute as cute
109
+ from quack.compile_utils import make_fake_tensor
110
+
111
+ from ._skinny_gemm import CuteSkinnyGemm
112
+
113
+ element_type = self._cutlass_dtype(dtype)
114
+ n = cute.sym_int(divisibility=config.outputs_per_block)
115
+ k = cute.sym_int(divisibility=config.block_size * config.vector_width)
116
+ a = make_fake_tensor(
117
+ element_type,
118
+ (config.num_rows, k),
119
+ divisibility=config.vector_width,
120
+ )
121
+ b = make_fake_tensor(
122
+ element_type,
123
+ (n, k),
124
+ divisibility=config.vector_width,
125
+ )
126
+ c = make_fake_tensor(element_type, (config.num_rows, n), divisibility=1)
127
+ residual = make_fake_tensor(element_type, (config.num_rows, n), divisibility=1)
128
+ kernel = CuteSkinnyGemm(
129
+ element_type=element_type,
130
+ num_rows=config.num_rows,
131
+ block_size=config.block_size,
132
+ outputs_per_block=config.outputs_per_block,
133
+ vector_width=config.vector_width,
134
+ k_unroll=config.k_unroll,
135
+ has_residual=has_residual,
136
+ use_pdl=self._use_pdl(),
137
+ )
138
+ self._compiled[(dtype, config, has_residual)] = cute.compile(
139
+ kernel,
140
+ a,
141
+ b,
142
+ residual,
143
+ c,
144
+ self._stream(),
145
+ options="--enable-tvm-ffi --ptxas-options -maxrregcount=64",
146
+ )
147
+
148
+ def request_warmup(
149
+ self,
150
+ dtype: torch.dtype,
151
+ shapes: Iterable[tuple[int, int, int]],
152
+ ) -> None:
153
+ """Request compilation of ``(M, N, K)`` shapes before graph capture."""
154
+ self.request_warmup_configs(
155
+ dtype,
156
+ (self._config(m, n, k) for m, n, k in shapes),
157
+ )
158
+
159
+ def request_warmup_configs(
160
+ self,
161
+ dtype: torch.dtype,
162
+ configs: Iterable[SkinnyGemmConfig],
163
+ *,
164
+ has_residual: bool = False,
165
+ ) -> None:
166
+ """Request compilation of explicit measured configs before capture."""
167
+ self._warmup_configs.update((dtype, config, has_residual) for config in configs)
168
+ if self._warmup_registered:
169
+ return
170
+ from vllm.model_executor.warmup.cutedsl_warmup import (
171
+ register_cutedsl_warmup_provider,
172
+ )
173
+
174
+ register_cutedsl_warmup_provider(self)
175
+ self._warmup_registered = True
176
+
177
+ def get_cutedsl_warmup_compile_units(self):
178
+ from vllm.model_executor.warmup.cutedsl_warmup import CuTeDSLCompileUnit
179
+
180
+ return tuple(
181
+ CuTeDSLCompileUnit(
182
+ name=(
183
+ "shape-dynamic skinny GEMM with residual"
184
+ if has_residual
185
+ else "shape-dynamic skinny GEMM"
186
+ ),
187
+ key=("shape-dynamic-skinny-gemm", dtype, config, has_residual),
188
+ compile=partial(self._compile, dtype, config, has_residual),
189
+ )
190
+ for dtype, config, has_residual in sorted(
191
+ self._warmup_configs,
192
+ key=lambda item: (
193
+ str(item[0]),
194
+ item[1].num_rows,
195
+ item[1].block_size,
196
+ item[1].outputs_per_block,
197
+ item[1].k_unroll,
198
+ item[1].vector_width,
199
+ item[2],
200
+ ),
201
+ )
202
+ )
203
+
204
+ def __call__(
205
+ self,
206
+ a: torch.Tensor,
207
+ b: torch.Tensor,
208
+ config: SkinnyGemmConfig | None = None,
209
+ residual: torch.Tensor | None = None,
210
+ ) -> torch.Tensor:
211
+ if a.dim() != 2 or b.dim() != 2:
212
+ raise ValueError("a and b must be 2D tensors")
213
+ if a.dtype not in (torch.bfloat16, torch.float16) or b.dtype != a.dtype:
214
+ raise ValueError("a and b must have the same BF16 or FP16 dtype")
215
+ if not a.is_cuda or not b.is_cuda or a.device != b.device:
216
+ raise ValueError("a and b must be CUDA tensors on the same device")
217
+ if not a.is_contiguous() or not b.is_contiguous():
218
+ raise ValueError("a and b must be contiguous")
219
+ if a.shape[1] != b.shape[1]:
220
+ raise ValueError("a and b must have matching K dimensions")
221
+ if not 1 <= a.shape[0] <= 16:
222
+ raise ValueError("shape-dynamic skinny GEMM requires 1 <= M <= 16")
223
+ if residual is not None:
224
+ if residual.dim() != 2 or residual.shape != (a.shape[0], b.shape[0]):
225
+ raise ValueError("residual must have shape (M, N)")
226
+ if residual.dtype != a.dtype:
227
+ raise ValueError("residual must have the same dtype as a and b")
228
+ if residual.device != a.device or not residual.is_cuda:
229
+ raise ValueError("residual must be on the same CUDA device")
230
+ if not residual.is_contiguous():
231
+ raise ValueError("residual must be contiguous")
232
+
233
+ config = config or self._config(a.shape[0], b.shape[0], a.shape[1])
234
+ if config.num_rows != a.shape[0]:
235
+ raise ValueError("config num_rows must match M")
236
+ if b.shape[0] % config.outputs_per_block != 0:
237
+ raise ValueError("N must be divisible by outputs_per_block")
238
+ if a.shape[1] % (config.block_size * config.vector_width) != 0:
239
+ raise ValueError(
240
+ "K must be divisible by block_size * vector_width for this config"
241
+ )
242
+ has_residual = residual is not None
243
+ cache_key = (a.dtype, config, has_residual)
244
+ if cache_key not in self._compiled:
245
+ self._compile(a.dtype, config, has_residual)
246
+ output = torch.empty((a.shape[0], b.shape[0]), dtype=a.dtype, device=a.device)
247
+ residual_arg = output if residual is None else residual
248
+ self._compiled[cache_key](a, b, residual_arg, output, self._stream())
249
+ return output
250
+
251
+
252
+shape_dynamic_skinny_gemm = ShapeDynamicSkinnyGemm()
vllm/model_executor/layers/activation.py
+45
@@ -158,6 +158,51 @@ class SiluAndMul(CustomOp):
158
return self.forward_native(x)
159
160
161
+@CustomOp.register("situ_and_mul")
162
+class SituAndMul(CustomOp):
163
+ """SituGLU activation used by Kimi models.
164
+
165
+ Computes beta * tanh(gate / beta) * sigmoid(gate) * up. When
166
+ ``linear_beta`` is set, the up projection is also softly clipped with
167
+ linear_beta * tanh(up / linear_beta).
168
+ """
169
+
170
+ def __init__(
171
+ self,
172
+ beta: float = 1.0,
173
+ linear_beta: float | None = None,
174
+ *,
175
+ compile_native: bool = True,
176
+ ):
177
+ super().__init__(compile_native=compile_native)
178
+ self.beta = float(beta)
179
+ self.linear_beta = None if linear_beta is None else float(linear_beta)
180
+ if current_platform.is_cuda_alike():
181
+ self.op = torch.ops._C.situ_and_mul
182
+
183
+ def forward_native(self, x: torch.Tensor) -> torch.Tensor:
184
+ d = x.shape[-1] // 2
185
+ gate = x[..., :d].float()
186
+ up = x[..., d:].float()
187
+ gate = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate)
188
+ if self.linear_beta is not None:
189
+ up = self.linear_beta * torch.tanh(up / self.linear_beta)
190
+ return (gate * up).to(x.dtype)
191
+
192
+ def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
193
+ # Fused CUDA kernel: writes straight to `out`, no fp32 temporaries.
194
+ # linear_beta<=0 signals "unset" to the kernel (up passed through).
195
+ d = x.shape[-1] // 2
196
+ out = torch.empty(x.shape[:-1] + (d,), dtype=x.dtype, device=x.device)
197
+ self.op(
198
+ out, x, self.beta, -1.0 if self.linear_beta is None else self.linear_beta
199
+ )
200
+ return out
201
+
202
+ def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
203
+ return self.forward_native(x)
204
+
205
+
206
@CustomOp.register("silu_and_mul_with_clamp")
207
class SiluAndMulWithClamp(CustomOp):
208
"""SwiGLU activation with input clamping (used by some MoE shared experts).
vllm/model_executor/layers/fused_moe/activation.py
+23
@@ -22,6 +22,7 @@ class MoEActivation(Enum):
22
# expects the *packed* layout ([all gates; all ups]), as produced by a
23
# MergedColumnParallelLinear gate_up_proj (e.g. MiniMax-M3).
24
SWIGLUOAI = "swigluoai"
25
+ SITU = "situ"
26
SWIGLUOAI_UNINTERLEAVE = "swigluoai_uninterleave"
27
SWIGLUSTEP = "swiglustep"
28
@@ -77,6 +78,7 @@ _CUSTOM_OP_NAMES: dict[MoEActivation, str] = {
78
MoEActivation.SILU: "silu_and_mul",
79
MoEActivation.GELU: "gelu_and_mul",
80
MoEActivation.GELU_TANH: "gelu_tanh_and_mul",
81
+ MoEActivation.SITU: "situ_and_mul",
82
MoEActivation.SWIGLUOAI: "swigluoai_and_mul",
83
MoEActivation.SWIGLUOAI_UNINTERLEAVE: "silu_and_mul_with_clamp",
84
MoEActivation.SWIGLUSTEP: "swiglustep_and_mul",
@@ -133,6 +135,8 @@ def apply_moe_activation(
135
beta: float = 0.0,
136
topk_ids: torch.Tensor | None = None,
137
expert_map: torch.Tensor | None = None,
138
+ activation_situ_beta: float | None = None,
139
+ activation_situ_linear_beta: float | None = None,
140
) -> torch.Tensor:
141
"""Apply MoE activation function.
142
@@ -163,6 +167,25 @@ def apply_moe_activation(
167
torch.ops._C.gelu_and_mul(output, input)
168
elif activation == MoEActivation.GELU_TANH:
169
torch.ops._C.gelu_tanh_and_mul(output, input)
170
+ elif activation == MoEActivation.SITU:
171
+ # Fused CUDA kernel: writes straight to `output`, no fp32 temporaries.
172
+ # (The pure-torch fallback below upcast both halves to fp32 and
173
+ # allocated ~8 temporaries per call, blowing up MoE memory.)
174
+ # Both betas come from FusedMoEConfig; a missing beta means the caller
175
+ # bypassed the config plumbing, so fail rather than silently use 1.0.
176
+ # linear_beta is genuinely optional: <= 0 signals "unset" to the kernel
177
+ # (up passed through), matching SituAndMul(linear_beta=None).
178
+ assert activation_situ_beta is not None, (
179
+ "SITU requires activation_situ_beta from FusedMoEConfig"
180
+ )
181
+ torch.ops._C.situ_and_mul(
182
+ output,
183
+ input,
184
+ activation_situ_beta,
185
+ -1.0
186
+ if activation_situ_linear_beta is None
187
+ else activation_situ_linear_beta,
188
+ )
189
elif activation == MoEActivation.SWIGLUOAI:
190
torch.ops._C.swigluoai_and_mul(output, input)
191
elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
vllm/model_executor/layers/fused_moe/config.py
+4
@@ -1314,6 +1314,10 @@ class FusedMoEConfig:
1314
swiglu_alpha: float | None = None
1315
swiglu_beta: float | None = None
1316
1317
+ # SituGLU parameters used by Kimi sit(u/v2) activations.
1318
+ activation_situ_beta: float | None = None
1319
+ activation_situ_linear_beta: float | None = None
1320
+
1321
max_capture_size: int = 0
1322
1323
# Set by __post_init__
vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py
+29
-11
@@ -444,7 +444,14 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
444
445
@staticmethod
446
def _supports_activation(activation: MoEActivation) -> bool:
447
- return activation in [MoEActivation.SILU, MoEActivation.SWIGLUSTEP]
447
+ # SILU has fused gate+mul+quant kernels; SWIGLUSTEP/SITU take the
448
+ # general path (activation applied via self.activation, which forwards
449
+ # the situ betas, then FP8 requant).
450
+ return activation in [
451
+ MoEActivation.SILU,
452
+ MoEActivation.SWIGLUSTEP,
453
+ MoEActivation.SITU,
454
+ ]
455
456
@staticmethod
457
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -488,16 +495,15 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
495
M_sum, N = input.size()
496
activation_out_dim = self.adjust_N_for_activation(N, activation)
497
491
- if scale_fmt == DeepGemmQuantScaleFMT.UE8M0:
492
- assert activation == MoEActivation.SILU
493
- return fused_silu_mul_fp8_quant_packed(
494
- input=input,
495
- output_q=output,
496
- group_size=block_k,
497
- clamp_limit=self.gemm1_clamp_limit,
498
- )
499
-
498
if activation == MoEActivation.SILU:
499
+ # Fused gate+mul+quant kernels for the common SILU case.
500
+ if scale_fmt == DeepGemmQuantScaleFMT.UE8M0:
501
+ return fused_silu_mul_fp8_quant_packed(
502
+ input=input,
503
+ output_q=output,
504
+ group_size=block_k,
505
+ clamp_limit=self.gemm1_clamp_limit,
506
+ )
507
use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0
508
return silu_mul_per_token_group_quant_fp8_colmajor(
509
input=input,
@@ -506,12 +512,24 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
512
clamp_limit=self.gemm1_clamp_limit,
513
)
514
515
+ # General gated activations (SWIGLUSTEP, SITU): apply the activation
516
+ # (self.activation forwards the situ betas from moe_config) then
517
+ # FP8-requant into the layout DeepGEMM expects for this scale format.
518
act_out = torch.empty(
519
(M_sum, activation_out_dim), dtype=input.dtype, device=input.device
520
)
521
self.activation(activation, act_out, input)
522
+ if scale_fmt == DeepGemmQuantScaleFMT.UE8M0:
523
+ return per_token_group_quant_fp8_packed_for_deepgemm(
524
+ act_out, block_k, use_ue8m0=True, out_q=output
525
+ )
526
+ use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0
527
return per_token_group_quant_fp8(
514
- act_out, block_k, column_major_scales=True, out_q=output
528
+ act_out,
529
+ block_k,
530
+ column_major_scales=True,
531
+ out_q=output,
532
+ use_ue8m0=use_ue8m0,
533
)
534
535
def apply(
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
+63
@@ -93,6 +93,8 @@ def _fused_marlin_moe(
93
clamp_limit: float | None = None,
94
gemm1_alpha: float = 1.0,
95
gemm1_beta: float = 0.0,
96
+ activation_situ_beta: float | None = None,
97
+ activation_situ_linear_beta: float | None = None,
98
) -> torch.Tensor:
99
assert hidden_states.ndim == 2
100
M, K = hidden_states.size()
@@ -171,6 +173,8 @@ def _fused_marlin_moe(
173
beta=gemm1_beta,
174
topk_ids=topk_ids,
175
expert_map=expert_map,
176
+ activation_situ_beta=activation_situ_beta,
177
+ activation_situ_linear_beta=activation_situ_linear_beta,
178
)
179
180
if output is None:
@@ -256,6 +260,8 @@ def fused_marlin_moe(
260
clamp_limit: float | None = None,
261
gemm1_alpha: float = 1.0,
262
gemm1_beta: float = 0.0,
263
+ activation_situ_beta: float | None = None,
264
+ activation_situ_linear_beta: float | None = None,
265
) -> torch.Tensor:
266
"""
267
This function computes a Mixture of Experts (MoE) layer using two sets of
@@ -357,6 +363,8 @@ def fused_marlin_moe(
363
num_tokens_post_padded=num_tokens_post_padded,
364
activation=activation,
365
activation_func=activation_func,
366
+ activation_situ_beta=activation_situ_beta,
367
+ activation_situ_linear_beta=activation_situ_linear_beta,
368
input_global_scale1=input_global_scale1,
369
input_global_scale2=input_global_scale2,
370
global_scale1=global_scale1,
@@ -423,6 +431,9 @@ def batched_fused_marlin_moe(
431
clamp_limit: float | None = None,
432
gemm1_alpha: float = 1.0,
433
gemm1_beta: float = 0.0,
434
+ activation_func: Callable[..., None] = apply_moe_activation,
435
+ activation_situ_beta: float | None = None,
436
+ activation_situ_linear_beta: float | None = None,
437
) -> torch.Tensor:
438
"""
439
This function massages the inputs so the batched hidden_states can be
@@ -530,6 +541,9 @@ def batched_fused_marlin_moe(
541
quant_type=quant_type,
542
apply_router_weight_on_input=apply_router_weight_on_input,
543
activation=activation,
544
+ activation_func=activation_func,
545
+ activation_situ_beta=activation_situ_beta,
546
+ activation_situ_linear_beta=activation_situ_linear_beta,
547
expert_map=expert_map,
548
block_size_m=block_size_m,
549
sorted_token_ids=sorted_token_ids,
@@ -645,6 +659,7 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
659
MoEActivation.SILU,
660
MoEActivation.GELU,
661
MoEActivation.GELU_TANH,
662
+ MoEActivation.SITU,
663
MoEActivation.SWIGLUOAI,
664
MoEActivation.SWIGLUOAI_UNINTERLEAVE,
665
MoEActivation.SWIGLUSTEP,
@@ -793,6 +808,10 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
808
global_num_experts=global_num_experts,
809
activation=activation,
810
activation_func=self.activation,
811
+ activation_situ_beta=self.moe_config.activation_situ_beta,
812
+ activation_situ_linear_beta=(
813
+ self.moe_config.activation_situ_linear_beta
814
+ ),
815
moe_sum=self.moe_sum,
816
expert_map=expert_map,
817
output=output,
@@ -833,6 +852,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
852
beta: float = 0.0,
853
topk_ids: torch.Tensor | None = None,
854
expert_map: torch.Tensor | None = None,
855
+ activation_situ_beta: float | None = None,
856
+ activation_situ_linear_beta: float | None = None,
857
) -> None:
858
# act_input = intermediate_cache1 (M*topk, 2N for gated)
859
# act_output = intermediate_cache2 (M*topk, N)
@@ -871,6 +892,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
892
beta=beta,
893
topk_ids=topk_ids,
894
expert_map=expert_map,
895
+ activation_situ_beta=activation_situ_beta,
896
+ activation_situ_linear_beta=activation_situ_linear_beta,
897
)
898
lora_state["cache2"] = act_output
899
@@ -918,6 +941,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
941
global_num_experts=global_num_experts,
942
activation=activation,
943
activation_func=activation_with_lora,
944
+ activation_situ_beta=self.moe_config.activation_situ_beta,
945
+ activation_situ_linear_beta=self.moe_config.activation_situ_linear_beta,
946
moe_sum=moe_sum_with_lora,
947
expert_map=expert_map,
948
output=output,
@@ -1021,6 +1046,41 @@ class BatchedMarlinExperts(MarlinExpertsBase):
1046
apply_router_weight_on_input: bool,
1047
):
1048
assert expert_tokens_meta is not None, "Num valid tokens per batch is required"
1049
+
1050
+ def activation_func(
1051
+ act: MoEActivation,
1052
+ act_output: torch.Tensor,
1053
+ act_input: torch.Tensor,
1054
+ *,
1055
+ activation_situ_beta: float | None = None,
1056
+ activation_situ_linear_beta: float | None = None,
1057
+ **kwargs,
1058
+ ) -> None:
1059
+ if act != MoEActivation.SITU:
1060
+ self.activation(
1061
+ act,
1062
+ act_output,
1063
+ act_input,
1064
+ activation_situ_beta=activation_situ_beta,
1065
+ activation_situ_linear_beta=activation_situ_linear_beta,
1066
+ **kwargs,
1067
+ )
1068
+ return
1069
+
1070
+ num_experts, max_num_tokens = hidden_states.shape[:2]
1071
+ beta = activation_situ_beta
1072
+ linear_beta = activation_situ_linear_beta
1073
+ assert beta is not None, (
1074
+ "SITU requires activation_situ_beta from FusedMoEConfig"
1075
+ )
1076
+ torch.ops._C.masked_situ_and_mul(
1077
+ act_output.view(num_experts, max_num_tokens, -1),
1078
+ act_input.view(num_experts, max_num_tokens, -1),
1079
+ expert_tokens_meta.expert_num_tokens,
1080
+ beta,
1081
+ -1.0 if linear_beta is None else linear_beta,
1082
+ )
1083
+
1084
return batched_fused_marlin_moe(
1085
hidden_states=hidden_states,
1086
expert_num_tokens=expert_tokens_meta.expert_num_tokens,
@@ -1051,4 +1111,7 @@ class BatchedMarlinExperts(MarlinExpertsBase):
1111
clamp_limit=self.gemm1_clamp_limit,
1112
gemm1_alpha=self.gemm1_alpha,
1113
gemm1_beta=self.gemm1_beta,
1114
+ activation_func=activation_func,
1115
+ activation_situ_beta=self.moe_config.activation_situ_beta,
1116
+ activation_situ_linear_beta=self.moe_config.activation_situ_linear_beta,
1117
)
vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
+51
-6
@@ -79,6 +79,36 @@ class TrtLlmMxfp4ExpertsBase:
79
else:
80
self.gemm1_clamp_limit = None
81
82
+ # SITU (SituGLU) TRTLLM-Gen kernel computes
83
+ # left = alpha * tanh(x0 / alpha) * sigmoid(x0) # gate (x0)
84
+ # right = beta * tanh(x1 / beta) # up (x1)
85
+ # which matches vLLM's situ_and_mul with (beta, linear_beta), so map
86
+ # situ beta -> gatedActAlpha (gemm1_alpha) and situ linear_beta ->
87
+ # gatedActBeta (gemm1_beta). Both must be > 0.
88
+ if moe_config.activation == MoEActivation.SITU:
89
+ situ_beta = moe_config.activation_situ_beta
90
+ situ_linear_beta = moe_config.activation_situ_linear_beta
91
+ assert situ_beta is not None and situ_beta > 0, (
92
+ "SITU requires activation_situ_beta > 0"
93
+ )
94
+ assert situ_linear_beta is not None and situ_linear_beta > 0, (
95
+ "TRTLLM SiTuGlu requires activation_situ_linear_beta > 0 "
96
+ "(the private cubin has no up-passthrough path)"
97
+ )
98
+ self.gemm1_alpha = torch.full(
99
+ (self.local_num_experts,),
100
+ float(situ_beta),
101
+ dtype=torch.float32,
102
+ device=device,
103
+ )
104
+ self.gemm1_beta = torch.full(
105
+ (self.local_num_experts,),
106
+ float(situ_linear_beta),
107
+ dtype=torch.float32,
108
+ device=device,
109
+ )
110
+ self.gemm1_clamp_limit = None
111
+
112
self.max_capture_size = moe_config.max_capture_size
113
114
@staticmethod
@@ -103,7 +133,19 @@ class TrtLlmMxfp4ExpertsBase:
133
134
@staticmethod
135
def _supports_activation(activation: MoEActivation) -> bool:
106
- return activation in (MoEActivation.SWIGLUOAI, MoEActivation.SILU)
136
+ return activation in (
137
+ MoEActivation.SWIGLUOAI,
138
+ MoEActivation.SILU,
139
+ MoEActivation.SITU,
140
+ )
141
+
142
+ @staticmethod
143
+ def _flashinfer_activation_type(activation: MoEActivation) -> int:
144
+ from flashinfer.fused_moe.core import ActivationType
145
+
146
+ if activation == MoEActivation.SITU:
147
+ return ActivationType.Situ.value
148
+ return ActivationType.Swiglu.value
149
150
@staticmethod
151
def activation_format() -> mk.FusedMoEActivationFormat:
@@ -142,6 +184,7 @@ class TrtLlmMxfp4ExpertsMonolithic(
184
activation_key: QuantKey | None,
185
) -> bool:
186
return routing_method in [
187
+ RoutingMethodType.DeepSeekV3,
188
RoutingMethodType.Renormalize,
189
RoutingMethodType.RenormalizeNaive,
190
]
@@ -192,8 +235,8 @@ class TrtLlmMxfp4ExpertsMonolithic(
235
device=hidden_states.device,
236
)
237
trtllm_fp4_block_scale_moe(
195
- routing_logits=router_logits.to(torch.bfloat16),
196
- routing_bias=None,
238
+ routing_logits=router_logits,
239
+ routing_bias=e_score_correction_bias,
240
hidden_states=x_quant,
241
hidden_states_scale=x_scale,
242
gemm1_weights=w1,
@@ -210,14 +253,15 @@ class TrtLlmMxfp4ExpertsMonolithic(
253
output2_scale_scalar=None,
254
num_experts=global_num_experts,
255
top_k=self.topk,
213
- n_group=None,
214
- topk_group=None,
256
+ n_group=(num_expert_group or 0),
257
+ topk_group=(topk_group or 0),
258
intermediate_size=self.intermediate_size_per_partition,
259
local_expert_offset=self.ep_rank * self.local_num_experts,
260
local_num_experts=self.local_num_experts,
218
- routed_scaling_factor=None,
261
+ routed_scaling_factor=routed_scaling_factor,
262
routing_method_type=self.routing_method_type,
263
do_finalize=True,
264
+ activation_type=self._flashinfer_activation_type(activation),
265
tune_max_num_tokens=max(self.max_capture_size, 1),
266
output=output,
267
routing_replay_out=routing_replay_out,
@@ -338,6 +382,7 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
382
"routing_method_type": RoutingMethodType.Renormalize,
383
"do_finalize": True,
384
"enable_pdl": True,
385
+ "activation_type": self._flashinfer_activation_type(activation),
386
"output": output,
387
"tune_max_num_tokens": max(self.max_capture_size, 1),
388
}
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+33
-3
@@ -110,6 +110,27 @@ class TrtLlmNvFp4ExpertsBase:
110
self.gemm1_alpha = None
111
self.gemm1_beta = None
112
113
+ # SITU (Kimi SituGLU) TRTLLM-Gen kernel computes
114
+ # left=alpha*tanh(x0/alpha)*sigmoid(x0), right=beta*tanh(x1/beta),
115
+ # matching vLLM's situ_and_mul, so map situ beta -> gatedActAlpha
116
+ # (gemm1_alpha) and situ linear_beta -> gatedActBeta (gemm1_beta).
117
+ # These operate on the dequantized gate/up, so they are NOT folded by
118
+ # g1_alphas in process_weights_after_loading.
119
+ self.is_situ = moe_config.activation == MoEActivation.SITU
120
+ if self.is_situ:
121
+ situ_beta = moe_config.activation_situ_beta
122
+ situ_linear_beta = moe_config.activation_situ_linear_beta
123
+ assert situ_beta is not None and situ_beta > 0, (
124
+ "SITU requires activation_situ_beta > 0"
125
+ )
126
+ assert situ_linear_beta is not None and situ_linear_beta > 0, (
127
+ "TRTLLM SiTuGlu requires activation_situ_linear_beta > 0 "
128
+ "(the private cubin has no up-passthrough path)"
129
+ )
130
+ self.gemm1_alpha = _per_expert(situ_beta)
131
+ self.gemm1_beta = _per_expert(situ_linear_beta)
132
+ self.gemm1_clamp_limit = None
133
+
134
logger.debug_once(
135
"activation=%s, gemm1_alpha=%s, gemm1_beta=%s, gemm1_clamp_limit=%s",
136
moe_config.activation,
@@ -142,7 +163,10 @@ class TrtLlmNvFp4ExpertsBase:
163
# (via the in-place mul above) and never changes again, so this is a
164
# static, per-expert constant. Register on the layer so EPLB
165
# rearranges it alongside the other expert tensors.
145
- if self.gemm1_clamp_limit is not None:
166
+ # SITU alpha/beta act on the dequantized gate/up (tanh clamps), not the
167
+ # raw GEMM1 accumulator, so they are registered as-is without the
168
+ # g1_alphas fold used by the SwiGLU-OAI clamp/beta below.
169
+ if self.gemm1_clamp_limit is not None and not self.is_situ:
170
gemm1_clamp_limit = self.gemm1_clamp_limit / self.quant_config.g1_alphas
171
layer.register_parameter(
172
"gemm1_clamp_limit",
@@ -155,7 +179,11 @@ class TrtLlmNvFp4ExpertsBase:
179
# raw. Register both on the layer so EPLB rearranges them with the
180
# other per-expert tensors.
181
if self.gemm1_beta is not None:
158
- gemm1_beta = self.gemm1_beta / self.quant_config.g1_alphas
182
+ gemm1_beta = (
183
+ self.gemm1_beta
184
+ if self.is_situ
185
+ else self.gemm1_beta / self.quant_config.g1_alphas
186
+ )
187
layer.register_parameter(
188
"gemm1_beta",
189
torch.nn.Parameter(gemm1_beta, requires_grad=False),
@@ -197,13 +225,15 @@ class TrtLlmNvFp4ExpertsBase:
225
226
@staticmethod
227
def _supports_activation(activation: MoEActivation) -> bool:
200
- """Supports SiLU, RELU^2 non-gated, GELU, and clamped SwiGLU-OAI."""
228
+ """Supports SiLU, RELU^2 non-gated, GELU, clamped SwiGLU-OAI, and the
229
+ private SituGLU (SITU) cubin."""
230
return activation in [
231
MoEActivation.SILU,
232
MoEActivation.RELU2_NO_MUL,
233
MoEActivation.GELU,
234
MoEActivation.GELU_TANH,
235
MoEActivation.SWIGLUOAI_UNINTERLEAVE,
236
+ MoEActivation.SITU,
237
]
238
239
@staticmethod
vllm/model_executor/layers/fused_moe/layer.py
+6
@@ -120,6 +120,8 @@ def FusedMoE(
120
swiglu_limit: float | None = None,
121
swiglu_alpha: float | None = None,
122
swiglu_beta: float | None = None,
123
+ activation_situ_beta: float | None = None,
124
+ activation_situ_linear_beta: float | None = None,
125
e_score_correction_bias: torch.Tensor | None = None,
126
apply_router_weight_on_input: bool = False,
127
activation: str = "silu",
@@ -178,6 +180,8 @@ def FusedMoE(
180
scoring_func: Scoring function for routing ("softmax" or others)
181
routed_scaling_factor: Scaling factor applied to topk_weights or output
182
swiglu_limit: SwiGLU activation limit
183
+ activation_situ_beta: SituGLU activation beta
184
+ activation_situ_linear_beta: SituGLU linear beta
185
e_score_correction_bias: Expert score correction bias tensor
186
apply_router_weight_on_input: Whether to apply router weights on input
187
activation: Activation function name ("silu", "gelu", etc.)
@@ -352,6 +356,8 @@ def FusedMoE(
356
swiglu_limit=swiglu_limit,
357
swiglu_alpha=swiglu_alpha,
358
swiglu_beta=swiglu_beta,
359
+ activation_situ_beta=activation_situ_beta,
360
+ activation_situ_linear_beta=activation_situ_linear_beta,
361
max_capture_size=vllm_config.compilation_config.max_cudagraph_capture_size,
362
skip_final_all_reduce=skip_final_all_reduce,
363
)
vllm/model_executor/layers/fused_moe/modular_kernel.py
+23
-8
@@ -890,6 +890,8 @@ class FusedMoEExpertsModular(FusedMoEExperts):
890
beta: float = 0.0,
891
topk_ids: torch.Tensor | None = None,
892
expert_map: torch.Tensor | None = None,
893
+ activation_situ_beta: float | None = None,
894
+ activation_situ_linear_beta: float | None = None,
895
) -> None:
896
apply_moe_activation(
897
activation,
@@ -900,6 +902,16 @@ class FusedMoEExpertsModular(FusedMoEExperts):
902
beta=beta,
903
topk_ids=topk_ids,
904
expert_map=expert_map,
905
+ activation_situ_beta=(
906
+ self.moe_config.activation_situ_beta
907
+ if activation_situ_beta is None
908
+ else activation_situ_beta
909
+ ),
910
+ activation_situ_linear_beta=(
911
+ self.moe_config.activation_situ_linear_beta
912
+ if activation_situ_linear_beta is None
913
+ else activation_situ_linear_beta
914
+ ),
915
)
916
917
@abstractmethod
@@ -1307,6 +1319,14 @@ class FusedMoEKernelModularImpl:
1319
activation,
1320
)
1321
1322
+ use_output_alias = (
1323
+ output_alias is not None
1324
+ and output_alias.shape == fused_out.shape
1325
+ and output_alias.dtype == fused_out.dtype
1326
+ and output_alias.device == fused_out.device
1327
+ and output_alias.is_contiguous()
1328
+ )
1329
+
1330
# If caller's output buffer already matches fused_out shape/dtype, alias
1331
# to skip the redundant copy in TopKWeightAndReduceNoOP.apply downstream.
1332
# This eliminates ~94% of __amd_rocclr_copyBuffer events (Copy 2 of the
@@ -1314,15 +1334,10 @@ class FusedMoEKernelModularImpl:
1334
if current_platform.is_rocm():
1335
from vllm._aiter_ops import rocm_aiter_ops
1336
1317
- if (
1318
- rocm_aiter_ops.is_fused_moe_enabled()
1319
- and output_alias is not None
1320
- and output_alias.shape == fused_out.shape
1321
- and output_alias.dtype == fused_out.dtype
1322
- and output_alias.device == fused_out.device
1323
- and output_alias.is_contiguous()
1324
- ):
1337
+ if use_output_alias and rocm_aiter_ops.is_fused_moe_enabled():
1338
fused_out = output_alias
1339
+ elif use_output_alias:
1340
+ fused_out = output_alias
1341
1342
self.fused_experts.apply(
1343
output=fused_out,
vllm/model_executor/layers/fused_moe/runner/latent_moe_runner.py
new
+255
@@ -0,0 +1,255 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+import vllm.envs as envs
6
+from vllm.config import get_current_vllm_config
7
+from vllm.distributed import tensor_model_parallel_all_reduce
8
+from vllm.logger import init_logger
9
+from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
10
+ _AR_RESIDUAL_RMS_NORM,
11
+ _can_use_flashinfer,
12
+ flashinfer_trtllm_fused_allreduce_norm,
13
+)
14
+from vllm.model_executor.layers.layernorm import RMSNorm
15
+from vllm.utils.torch_utils import aux_stream, current_stream
16
+
17
+from .moe_runner import MoERunner, _unpack
18
+
19
+logger = init_logger(__name__)
20
+
21
+
22
+class LatentMoERunner(MoERunner):
23
+ """MoE runner for latent MoE with a replicated routed up-projection.
24
+
25
+ Fused path (tp>1, un-reduced combine output, shared expert, no SP):
26
+ concatenates the un-reduced latent partial (dim d) and the un-reduced
27
+ shared partial (dim D) into one contiguous buffer, all-reduces once, then
28
+ splits. The latent half is normed and up-projected locally (replicated
29
+ up-proj -> full hidden), and the shared add folds into the GEMM epilogue
30
+ (``torch.addmm``). One collective total, no post-reduction communication.
31
+
32
+ Native path: the replicated up-proj produces the full hidden dim on every
33
+ rank, so the base runner combines routed + shared correctly at any TP size
34
+ (using two collectives instead of the fused path's one).
35
+ """
36
+
37
+ def __init__(
38
+ self,
39
+ *args,
40
+ enable_k3_latent_moe_tail_fusion: bool = False,
41
+ **kwargs,
42
+ ) -> None:
43
+ super().__init__(*args, **kwargs)
44
+ self.enable_k3_latent_moe_tail_fusion = enable_k3_latent_moe_tail_fusion
45
+ use_fused_path = self._use_fused_path()
46
+ if (
47
+ self.enable_k3_latent_moe_tail_fusion
48
+ and use_fused_path
49
+ and self.moe_config.tp_size not in (8, 16)
50
+ ):
51
+ logger.warning_once(
52
+ "K3 latent-MoE tail fusion currently supports TP=8 and TP=16, "
53
+ "but TP=%d is configured. Falling back to the default path.",
54
+ self.moe_config.tp_size,
55
+ )
56
+ self.enable_k3_latent_moe_tail_fusion = False
57
+
58
+ if self.enable_k3_latent_moe_tail_fusion and use_fused_path:
59
+ vllm_config = get_current_vllm_config()
60
+ if vllm_config.parallel_config.use_ubatching:
61
+ raise ValueError(
62
+ "K3 latent-MoE tail fusion does not support DBO or ubatching."
63
+ )
64
+ if vllm_config.model_config.enable_sleep_mode:
65
+ raise ValueError(
66
+ "K3 latent-MoE tail fusion does not support sleep mode."
67
+ )
68
+ transform = self.routed_output_transform
69
+ assert transform is not None
70
+ norm = transform.norm
71
+ assert norm is not None
72
+ from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import (
73
+ KimiK3LatentMoETailOp,
74
+ )
75
+
76
+ op = KimiK3LatentMoETailOp.initialize(
77
+ hidden_size=transform.up_proj.weight.shape[0],
78
+ latent_size=norm.weight.shape[0],
79
+ dtype=norm.weight.dtype,
80
+ device=norm.weight.device,
81
+ rms_eps=norm.variance_epsilon,
82
+ )
83
+ self._k3_latent_moe_tail_op = op
84
+
85
+ def _get_zero_residual(self, hidden_states: torch.Tensor) -> torch.Tensor:
86
+ """Read-only zero ``residual_in`` for the fused AR+RMSNorm kernel.
87
+
88
+ flashinfer requires a residual buffer even when there is no residual to
89
+ add.
90
+ """
91
+ numel = hidden_states.numel()
92
+ buf = getattr(self, "_zero_residual", None)
93
+ if (
94
+ buf is None
95
+ or buf.dtype != hidden_states.dtype
96
+ or buf.device != hidden_states.device
97
+ or buf.numel() < numel
98
+ ):
99
+ buf = torch.zeros(
100
+ numel, dtype=hidden_states.dtype, device=hidden_states.device
101
+ )
102
+ self._zero_residual = buf
103
+ return buf[:numel].view_as(hidden_states)
104
+
105
+ def _use_fused_path(self) -> bool:
106
+ # The fused path merges the latent and shared reductions into one
107
+ # all-reduce, so it needs actual TP parallelism, a shared expert (to
108
+ # concat), an un-reduced combine output, and no sequence parallelism.
109
+ return (
110
+ self.moe_config.tp_size > 1
111
+ and self._shared_experts is not None
112
+ and not self._fused_output_is_reduced
113
+ and not self.moe_config.is_sequence_parallel
114
+ )
115
+
116
+ def forward(
117
+ self,
118
+ hidden_states: torch.Tensor,
119
+ router_logits: torch.Tensor,
120
+ input_ids: torch.Tensor | None = None,
121
+ shared_experts_input: torch.Tensor | None = None,
122
+ ) -> torch.Tensor:
123
+ if self._use_fused_path():
124
+ return self._fused_forward(
125
+ hidden_states, router_logits, input_ids, shared_experts_input
126
+ )
127
+ return super().forward(
128
+ hidden_states, router_logits, input_ids, shared_experts_input
129
+ )
130
+
131
+ def _fused_forward(
132
+ self,
133
+ hidden_states: torch.Tensor,
134
+ router_logits: torch.Tensor,
135
+ input_ids: torch.Tensor | None = None,
136
+ shared_experts_input: torch.Tensor | None = None,
137
+ ) -> torch.Tensor:
138
+ # When the caller pre-applies the routed input transform outside the
139
+ # runner (e.g. to overlap it on a separate stream), it passes the
140
+ # already-transformed routed input as ``hidden_states`` and the original
141
+ # hidden states as ``shared_experts_input``; skip the transform then.
142
+ if shared_experts_input is None:
143
+ hidden_states, shared_experts_input = self.apply_routed_input_transform(
144
+ hidden_states
145
+ )
146
+ hidden_states, og_hidden_dim_pre_xform, og_hidden_dim_post_xform = (
147
+ self._maybe_pad_hidden_states(
148
+ shared_experts_input,
149
+ hidden_states,
150
+ )
151
+ )
152
+
153
+ result = self._forward_entry(
154
+ hidden_states,
155
+ router_logits,
156
+ shared_experts_input,
157
+ input_ids,
158
+ self._encode_layer_name(),
159
+ self.moe_config.hidden_dim_unpadded
160
+ if self._quant_method.has_unpadded_output
161
+ else 0,
162
+ )
163
+
164
+ shared_output, fused_output = _unpack(result)
165
+ assert shared_output is not None
166
+
167
+ if og_hidden_dim_pre_xform is not None:
168
+ fused_output = fused_output[..., :og_hidden_dim_pre_xform]
169
+
170
+ transform = self.routed_output_transform
171
+ assert transform is not None
172
+
173
+ if self.enable_k3_latent_moe_tail_fusion:
174
+ op = self._k3_latent_moe_tail_op
175
+ if 0 < fused_output.shape[0] <= op.contract.max_num_tokens:
176
+ norm = transform.norm
177
+ assert norm is not None
178
+ result = op(
179
+ fused_output,
180
+ shared_output,
181
+ norm.weight,
182
+ transform.up_proj.weight,
183
+ )
184
+ result = self._maybe_reduce_final_output(
185
+ result, og_hidden_dim_post_xform, output_is_reduced=True
186
+ )
187
+ return self._maybe_add_zero_expert_output(result)
188
+
189
+ fused_latent = None
190
+ if transform.norm is not None:
191
+ fused_latent = self.allreduce_norm_latent_out(fused_output, transform.norm)
192
+ else:
193
+ fused_latent = tensor_model_parallel_all_reduce(fused_output)
194
+
195
+ shared_expert_stream = (
196
+ aux_stream()
197
+ if shared_output.size(0) <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
198
+ else None
199
+ )
200
+ if shared_expert_stream is not None:
201
+ # overlap shared expert allreduce with latent up_proj
202
+ main = current_stream()
203
+ shared_output.record_stream(shared_expert_stream)
204
+ shared_expert_stream.wait_stream(main)
205
+ with torch.cuda.stream(shared_expert_stream):
206
+ shared_output = tensor_model_parallel_all_reduce(shared_output)
207
+ result = torch.mm(fused_latent, transform.up_proj.weight.t())
208
+ main.wait_stream(shared_expert_stream)
209
+ else:
210
+ shared_output = tensor_model_parallel_all_reduce(shared_output)
211
+ result = torch.mm(fused_latent, transform.up_proj.weight.t())
212
+ result.add_(shared_output)
213
+
214
+ # Output is already fully reduced; this only strips padding.
215
+ result = self._maybe_reduce_final_output(
216
+ result, og_hidden_dim_post_xform, output_is_reduced=True
217
+ )
218
+
219
+ return self._maybe_add_zero_expert_output(result)
220
+
221
+ def allreduce_norm_latent_out(
222
+ self,
223
+ hidden_states: torch.Tensor,
224
+ norm: RMSNorm,
225
+ ) -> tuple[torch.Tensor, torch.Tensor]:
226
+ """All-reduce + add residual + (standard) RMSNorm, fused via flashinfer."""
227
+ if self.moe_config.tp_size == 1:
228
+ return norm(hidden_states)
229
+
230
+ if flashinfer_trtllm_fused_allreduce_norm is not None:
231
+ ok, max_token_num = _can_use_flashinfer(
232
+ hidden_states, self.moe_config.tp_size
233
+ )
234
+ if ok:
235
+ norm_out = torch.empty_like(hidden_states)
236
+ # With norm_out provided, the kernel writes the new residual
237
+ # (all_reduce(hidden_states) + residual) into the hidden_states
238
+ # buffer and the normalized result into norm_out.
239
+ flashinfer_trtllm_fused_allreduce_norm(
240
+ allreduce_in=hidden_states,
241
+ residual=self._get_zero_residual(hidden_states),
242
+ rms_gamma=norm.weight,
243
+ rms_eps=norm.variance_epsilon,
244
+ world_size=self.moe_config.tp_size,
245
+ weight_bias=0.0,
246
+ launch_with_pdl=True,
247
+ fp32_acc=True,
248
+ max_token_num=max_token_num,
249
+ pattern_code=_AR_RESIDUAL_RMS_NORM,
250
+ norm_out=norm_out,
251
+ )
252
+ return norm_out
253
+
254
+ reduced = tensor_model_parallel_all_reduce(hidden_states)
255
+ return norm(reduced)
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+62
-11
@@ -348,7 +348,8 @@ class MoERunner(MoERunnerInterface):
348
return self.routed_experts.quant_method
349
350
def apply_routed_input_transform(
351
- self, hidden_states: torch.Tensor
351
+ self,
352
+ hidden_states: torch.Tensor,
353
) -> tuple[torch.Tensor, torch.Tensor | None]:
354
"""Apply transform for routed experts (e.g., latent projection).
355
@@ -416,6 +417,7 @@ class MoERunner(MoERunnerInterface):
417
def _maybe_reduce_shared_expert_output(
418
self,
419
shared_output: torch.Tensor | None,
420
+ fused_output_is_reduced: bool | None = None,
421
) -> torch.Tensor | None:
422
"""All-reduce shared expert output when the combine kernel already
423
reduced fused output.
@@ -425,18 +427,43 @@ class MoERunner(MoERunnerInterface):
427
* If we have SP (TP=N, DP=M, EP), there is a separate AG step handled
428
in the model.
429
"""
430
+ if fused_output_is_reduced is None:
431
+ fused_output_is_reduced = self._fused_output_is_reduced
432
+
433
if (
434
shared_output is not None
435
and not self.moe_config.is_sequence_parallel
431
- and self._fused_output_is_reduced
436
+ and fused_output_is_reduced
437
):
438
shared_output = tensor_model_parallel_all_reduce(shared_output)
439
return shared_output
440
441
+ def _maybe_reduce_routed_output_before_transform(
442
+ self,
443
+ fused_output: torch.Tensor,
444
+ fused_output_is_reduced: bool,
445
+ ) -> tuple[torch.Tensor, bool]:
446
+ """All-reduce latent routed output before its output transform.
447
+
448
+ Latent MoE output transforms may contain non-linear ops, e.g. RMSNorm.
449
+ TP partial routed outputs must be summed in latent space before such
450
+ transforms are applied.
451
+ """
452
+ if (
453
+ self.routed_output_transform is not None
454
+ and not self.moe_config.is_sequence_parallel
455
+ and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
456
+ and not fused_output_is_reduced
457
+ ):
458
+ fused_output = tensor_model_parallel_all_reduce(fused_output)
459
+ fused_output_is_reduced = True
460
+ return fused_output, fused_output_is_reduced
461
+
462
def _maybe_reduce_final_output(
463
self,
464
states: torch.Tensor,
465
trunc_size: int | None,
466
+ output_is_reduced: bool | None = None,
467
) -> torch.Tensor:
468
"""All-reduce the combined output if needed.
469
@@ -455,11 +482,14 @@ class MoERunner(MoERunnerInterface):
482
# We don't need to reduce the final output if:
483
# - We are not running with TP or DP
484
# - The MK already reduced the fused output itself.
485
+ if output_is_reduced is None:
486
+ output_is_reduced = self._fused_output_is_reduced
487
+
488
if (
489
not self.moe_config.is_sequence_parallel
490
and not self.moe_config.skip_final_all_reduce
491
and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
462
- and not self._fused_output_is_reduced
492
+ and not output_is_reduced
493
):
494
states = tensor_model_parallel_all_reduce(states)
495
@@ -643,6 +673,7 @@ class MoERunner(MoERunnerInterface):
673
hidden_states: torch.Tensor,
674
router_logits: torch.Tensor,
675
input_ids: torch.Tensor | None = None,
676
+ shared_experts_input: torch.Tensor | None = None,
677
) -> torch.Tensor:
678
"""Invoke the fused moe layer.
679
@@ -664,11 +695,16 @@ class MoERunner(MoERunnerInterface):
695
_moe_forward and _moe_forward_shared must be split.
696
"""
697
667
- # Apply transform for routed experts (e.g., latent projection
668
- # for latent MoE)
669
- hidden_states, shared_experts_input = self.apply_routed_input_transform(
670
- hidden_states
671
- )
698
+ # Apply transform for routed experts (e.g., latent projection for
699
+ # latent MoE). When the caller pre-applies the routed input transform
700
+ # outside the runner (e.g. to overlap it on a separate stream), it
701
+ # passes the already-transformed routed input as ``hidden_states`` and
702
+ # the original hidden states as ``shared_experts_input``; skip the
703
+ # transform in that case so shared experts still see the original input.
704
+ if shared_experts_input is None:
705
+ hidden_states, shared_experts_input = self.apply_routed_input_transform(
706
+ hidden_states
707
+ )
708
709
# Record before `_maybe_pad_hidden_states` pads activations to match
710
# `moe_config.hidden_dim`, e.g. after `align_trtllm_fp4_moe_hidden_dim_for_fi`
@@ -708,9 +744,22 @@ class MoERunner(MoERunnerInterface):
744
if og_hidden_dim_pre_xform is not None:
745
fused_output = fused_output[..., :og_hidden_dim_pre_xform]
746
711
- # If combine kernel already reduced fused, reduce shared to match.
747
+ fused_output_is_reduced = self._fused_output_is_reduced
748
+
749
+ # Latent routed output has to be reduced before output transform,
750
+ # because the transform may include non-linear normalization.
751
+ fused_output, fused_output_is_reduced = (
752
+ self._maybe_reduce_routed_output_before_transform(
753
+ fused_output,
754
+ fused_output_is_reduced,
755
+ )
756
+ )
757
+
758
+ # If routed output is already reduced, reduce shared to match.
759
# See note above re: the two all-reduce points.
713
- shared_output = self._maybe_reduce_shared_expert_output(shared_output)
760
+ shared_output = self._maybe_reduce_shared_expert_output(
761
+ shared_output, fused_output_is_reduced
762
+ )
763
764
shared_output, fused_output = self._maybe_apply_routed_scale_to_output(
765
shared_output, fused_output
@@ -724,7 +773,9 @@ class MoERunner(MoERunnerInterface):
773
else:
774
result = fused_output
775
727
- result = self._maybe_reduce_final_output(result, og_hidden_dim_post_xform)
776
+ result = self._maybe_reduce_final_output(
777
+ result, og_hidden_dim_post_xform, fused_output_is_reduced
778
+ )
779
780
return self._maybe_add_zero_expert_output(result)
781
vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py
+1
@@ -36,6 +36,7 @@ class MoERunnerInterface(PluggableLayer, ABC):
36
hidden_states: torch.Tensor,
37
router_logits: torch.Tensor,
38
input_ids: torch.Tensor | None = None,
39
+ shared_experts_input: torch.Tensor | None = None,
40
) -> torch.Tensor:
41
raise NotImplementedError
42
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
+15
-3
@@ -1080,6 +1080,7 @@ def causal_conv1d_update(
1080
block_idx_last_scheduled_token: torch.Tensor | None = None,
1081
initial_state_idx: torch.Tensor | None = None,
1082
validate_data=False,
1083
+ out: torch.Tensor | None = None,
1084
):
1085
"""
1086
x: Input tensor which can take the following shapes:
@@ -1117,7 +1118,8 @@ def causal_conv1d_update(
1118
for example: conv_state_indices = [null_block_id, 1, 20, null_block_id]
1119
in this case, the kernel will not process entries at
1120
indices 0 and 3
1120
- out: (batch, dim) or (batch, dim, seqlen) or (num_tokens, dim), same shape as `x`
1121
+ out: optional output tensor with the same shape as `x`. When omitted,
1122
+ the input is overwritten.
1123
"""
1124
if validate_data:
1125
assert null_block_id is not None
@@ -1129,10 +1131,22 @@ def causal_conv1d_update(
1131
1132
original_x_dtype = x.dtype
1133
x = x.to(conv_state.dtype)
1134
+ if out is None:
1135
+ out = x
1136
+ else:
1137
+ if out.shape != x.shape:
1138
+ raise ValueError(
1139
+ f"`out` shape {tuple(out.shape)} must match `x` shape {tuple(x.shape)}."
1140
+ )
1141
+ if out.dtype != original_x_dtype or out.device != x.device:
1142
+ raise ValueError(
1143
+ "`out` must have the same dtype and device as the input `x`."
1144
+ )
1145
unsqueeze = query_start_loc is None and x.dim() == 2
1146
if unsqueeze:
1147
# make it (batch, dim, seqlen) with seqlen == 1
1148
x = x.unsqueeze(-1)
1149
+ out = out.unsqueeze(-1)
1150
if query_start_loc is None:
1151
batch, dim, seqlen = x.shape
1152
else:
@@ -1159,8 +1173,6 @@ def causal_conv1d_update(
1173
assert num_cache_lines >= batch
1174
assert weight.stride(1) == 1 # Need this
1175
1162
- # adopt the strategy in vLLM that overwrite on 'x' directly, rather than creating a new tensor 'o'
1163
- out = x
1176
stride_w_dim, stride_w_width = weight.stride()
1177
1178
if query_start_loc is None:
vllm/model_executor/layers/mamba/ops/gather_initial_states.py
new
+83
@@ -0,0 +1,83 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm.platforms import current_platform
7
+from vllm.triton_utils import tl, triton
8
+
9
+
10
+@triton.jit
11
+def _gather_initial_states_kernel(
12
+ state_ptr,
13
+ indices_ptr,
14
+ has_initial_state_ptr,
15
+ output_ptr,
16
+ stride_state_batch,
17
+ stride_indices,
18
+ stride_has_initial_state,
19
+ row_size: tl.constexpr,
20
+ BLOCK_SIZE: tl.constexpr,
21
+ launch_pdl: tl.constexpr,
22
+):
23
+ block_idx = tl.program_id(0)
24
+ batch_idx = tl.program_id(1)
25
+ offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
26
+ mask = offsets < row_size
27
+
28
+ if launch_pdl:
29
+ tl.extra.cuda.gdc_wait()
30
+ tl.extra.cuda.gdc_launch_dependents()
31
+
32
+ has_initial_state = tl.load(
33
+ has_initial_state_ptr + batch_idx * stride_has_initial_state
34
+ ).to(tl.int1)
35
+ state_idx = tl.load(indices_ptr + batch_idx * stride_indices).to(tl.int64)
36
+ state_idx = tl.where(has_initial_state, state_idx, 0)
37
+ values = tl.load(
38
+ state_ptr + state_idx * stride_state_batch + offsets,
39
+ mask=mask & has_initial_state,
40
+ other=0.0,
41
+ )
42
+ tl.store(output_ptr + batch_idx * row_size + offsets, values, mask=mask)
43
+
44
+
45
+def gather_initial_states(
46
+ state: torch.Tensor,
47
+ indices: torch.Tensor,
48
+ has_initial_state: torch.Tensor,
49
+) -> torch.Tensor:
50
+ """Gather dense state rows, replacing uninitialized rows with zeros."""
51
+ assert state.ndim >= 2
52
+ assert state.is_cuda
53
+ assert indices.ndim == 1 and has_initial_state.ndim == 1
54
+ assert indices.shape == has_initial_state.shape
55
+ assert indices.device == state.device
56
+ assert has_initial_state.device == state.device
57
+ assert indices.dtype in (torch.int32, torch.int64)
58
+ assert has_initial_state.dtype == torch.bool
59
+
60
+ row_size = state[0].numel()
61
+ # Mamba pages may pad stride(0), but each state row remains dense.
62
+ assert state[0].is_contiguous()
63
+ output = torch.empty(
64
+ (indices.numel(), *state.shape[1:]),
65
+ dtype=state.dtype,
66
+ device=state.device,
67
+ )
68
+ block_size = min(triton.next_power_of_2(row_size), 1024)
69
+ grid = (triton.cdiv(row_size, block_size), indices.numel())
70
+ _gather_initial_states_kernel[grid](
71
+ state,
72
+ indices,
73
+ has_initial_state,
74
+ output,
75
+ state.stride(0),
76
+ indices.stride(0),
77
+ has_initial_state.stride(0),
78
+ row_size=row_size,
79
+ BLOCK_SIZE=block_size,
80
+ num_warps=8,
81
+ launch_pdl=current_platform.is_arch_support_pdl(),
82
+ )
83
+ return output
vllm/models/common/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/common/ops/__init__.py
new
+9
@@ -0,0 +1,9 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Ops shared across model implementations."""
4
+
5
+from .fused_qk_rmsnorm import fused_q_kv_rmsnorm
6
+
7
+__all__ = [
8
+ "fused_q_kv_rmsnorm",
9
+]
vllm/models/common/ops/fused_qk_rmsnorm.py
new
+103
@@ -0,0 +1,103 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm.platforms import current_platform
6
+from vllm.triton_utils import tl, triton
7
+
8
+
9
+@triton.jit
10
+def _fused_q_kv_rmsnorm_kernel(
11
+ q_ptr,
12
+ q_out_ptr,
13
+ q_weight_ptr,
14
+ q_in_stride,
15
+ q_out_stride,
16
+ kv_ptr,
17
+ kv_out_ptr,
18
+ kv_weight_ptr,
19
+ kv_in_stride,
20
+ kv_out_stride,
21
+ eps,
22
+ Q_SIZE: tl.constexpr,
23
+ KV_SIZE: tl.constexpr,
24
+ BLOCK_SIZE: tl.constexpr,
25
+ launch_pdl: tl.constexpr,
26
+):
27
+ # num_tokens goes on grid-x (max 2**31 - 1); task goes on grid-y.
28
+ # CUDA's grid-y/z are capped at 65535, so putting num_tokens there crashes
29
+ # the launch at max-num-batched-tokens >= 65536 with "invalid argument".
30
+ # int64: q_in_stride can be ~24K (128 heads × 192) and overflows int32
31
+ # past num_tokens ~87K under large chunked prefill.
32
+ token_idx = tl.program_id(0).to(tl.int64)
33
+ pid_task = tl.program_id(1)
34
+
35
+ if pid_task == 0:
36
+ SIZE = Q_SIZE
37
+ row_in = q_ptr + token_idx * q_in_stride
38
+ weight_ptr = q_weight_ptr
39
+ row_out = q_out_ptr + token_idx * q_out_stride
40
+ else:
41
+ SIZE = KV_SIZE
42
+ row_in = kv_ptr + token_idx * kv_in_stride
43
+ weight_ptr = kv_weight_ptr
44
+ row_out = kv_out_ptr + token_idx * kv_out_stride
45
+
46
+ if launch_pdl:
47
+ tl.extra.cuda.gdc_wait()
48
+ tl.extra.cuda.gdc_launch_dependents()
49
+
50
+ # RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's
51
+ # `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which
52
+ # keep x, rrms, and w all in fp32 and perform a single cast at store.
53
+ block = tl.arange(0, BLOCK_SIZE)
54
+ mask = block < SIZE
55
+ x = tl.load(row_in + block, mask=mask, other=0.0).to(tl.float32)
56
+ variance = tl.sum(x * x, axis=0) / SIZE
57
+ rrms = tl.rsqrt(variance + eps)
58
+ w = tl.load(weight_ptr + block, mask=mask, other=0.0).to(tl.float32)
59
+ y = x * rrms * w
60
+ tl.store(row_out + block, y.to(row_out.dtype.element_ty), mask=mask)
61
+
62
+
63
+def fused_q_kv_rmsnorm(
64
+ qr: torch.Tensor,
65
+ kv: torch.Tensor,
66
+ q_weight: torch.Tensor,
67
+ kv_weight: torch.Tensor,
68
+ eps: float,
69
+) -> tuple[torch.Tensor, torch.Tensor]:
70
+ assert qr.ndim == 2 and kv.ndim == 2
71
+ assert qr.shape[0] == kv.shape[0], (
72
+ f"token dim mismatch: qr={qr.shape}, kv={kv.shape}"
73
+ )
74
+ assert qr.stride(-1) == 1 and kv.stride(-1) == 1
75
+ assert q_weight.is_contiguous() and kv_weight.is_contiguous()
76
+
77
+ q_size = qr.shape[1]
78
+ kv_size = kv.shape[1]
79
+ num_tokens = qr.shape[0]
80
+ qr_out = torch.empty_like(qr)
81
+ kv_out = torch.empty_like(kv)
82
+ if num_tokens == 0:
83
+ return qr_out, kv_out
84
+
85
+ block_size = triton.next_power_of_2(max(q_size, kv_size))
86
+ _fused_q_kv_rmsnorm_kernel[(num_tokens, 2)](
87
+ qr,
88
+ qr_out,
89
+ q_weight,
90
+ qr.stride(0),
91
+ qr_out.stride(0),
92
+ kv,
93
+ kv_out,
94
+ kv_weight,
95
+ kv.stride(0),
96
+ kv_out.stride(0),
97
+ eps,
98
+ Q_SIZE=q_size,
99
+ KV_SIZE=kv_size,
100
+ BLOCK_SIZE=block_size,
101
+ launch_pdl=current_platform.is_arch_support_pdl(),
102
+ )
103
+ return qr_out, kv_out
vllm/models/inkling/nvidia/ops/fa4_rel_attention.py
+9
-2
@@ -130,13 +130,20 @@ def inkling_fa4_rel_attention(
130
cute_window = (None, None) if window_size == (-1, -1) else window_size
131
132
rel_logits = rel_logits.contiguous()
133
+ flash_attn_varlen_func: Callable[..., Any]
134
if _use_sheared_bias():
134
- from vllm.third_party.tml_fa4 import flash_attn_varlen_func
135
+ from vllm.third_party.tml_fa4 import (
136
+ flash_attn_varlen_func as tml_flash_attn_varlen_func,
137
+ )
138
139
+ flash_attn_varlen_func = tml_flash_attn_varlen_func
140
bias_kwargs: dict[str, Any] = {"rel_bias": rel_logits}
141
else:
138
- from vllm.vllm_flash_attn.cute import flash_attn_varlen_func
142
+ from vllm.vllm_flash_attn.cute import (
143
+ flash_attn_varlen_func as cute_flash_attn_varlen_func,
144
+ )
145
146
+ flash_attn_varlen_func = cute_flash_attn_varlen_func
147
bias_kwargs = {
148
"score_mod": _get_score_mod(rel_extent),
149
"aux_tensors": [rel_logits],
vllm/models/kimi_k3/__init__.py
+26
@@ -1,2 +1,28 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Kimi K3 model — hardware-isolated entry point.
4
+
5
+The implementation lives under ``nvidia/`` and ``amd/``; this module picks the
6
+right one for the current platform and re-exports the public classes used by
7
+the model registry. (Mirrors ``vllm.models.minimax_m3``.)
8
+"""
9
+
10
+from typing import TYPE_CHECKING
11
+
12
+from vllm.platforms import current_platform
13
+
14
+# The NVIDIA branch is the static default that type-checkers see; the ROCm
15
+# branch overrides it at runtime (kept type-compatible via type: ignore).
16
+if TYPE_CHECKING or not current_platform.is_rocm():
17
+ from .nvidia.model import KimiK3ForConditionalGeneration, KimiLinearForCausalLM
18
+ from .nvidia.mtp import KimiK3MTP
19
+else:
20
+ from .amd.linear import KimiLinearForCausalLM # type: ignore[assignment]
21
+ from .amd.model import KimiK3ForConditionalGeneration # type: ignore[assignment]
22
+ from .amd.mtp import KimiK3MTP # type: ignore[assignment]
23
+
24
+__all__ = [
25
+ "KimiK3ForConditionalGeneration",
26
+ "KimiK3MTP",
27
+ "KimiLinearForCausalLM",
28
+]
vllm/models/kimi_k3/amd/linear.py
new
+1065
@@ -0,0 +1,1065 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections.abc import Iterable
5
+from typing import Any
6
+
7
+import torch
8
+from torch import nn
9
+
10
+from vllm.config import CacheConfig, VllmConfig
11
+from vllm.distributed import (
12
+ get_pp_group,
13
+ get_tensor_model_parallel_world_size,
14
+)
15
+from vllm.logger import init_logger
16
+from vllm.model_executor.layers.activation import SiluAndMul, SituAndMul
17
+from vllm.model_executor.layers.fused_moe import (
18
+ FusedMoE,
19
+ fused_moe_make_expert_params_mapping,
20
+)
21
+from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
22
+from vllm.model_executor.layers.layernorm import RMSNorm
23
+from vllm.model_executor.layers.linear import (
24
+ ColumnParallelLinear,
25
+ MergedColumnParallelLinear,
26
+ ReplicatedLinear,
27
+ RowParallelLinear,
28
+)
29
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
30
+from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import (
31
+ KimiGatedDeltaNetAttention,
32
+)
33
+from vllm.model_executor.layers.mamba.mamba_utils import (
34
+ MambaStateCopyFunc,
35
+ MambaStateCopyFuncCalculator,
36
+ MambaStateDtypeCalculator,
37
+ MambaStateShapeCalculator,
38
+)
39
+from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
40
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
41
+from vllm.model_executor.layers.vocab_parallel_embedding import (
42
+ ParallelLMHead,
43
+ VocabParallelEmbedding,
44
+)
45
+from vllm.model_executor.model_loader.weight_utils import (
46
+ default_weight_loader,
47
+ maybe_remap_kv_scale_name,
48
+)
49
+from vllm.model_executor.models.interfaces import (
50
+ EagleModelMixin,
51
+ HasInnerState,
52
+ IsHybrid,
53
+ MixtureOfExperts,
54
+ SupportsPP,
55
+)
56
+from vllm.model_executor.models.utils import (
57
+ AutoWeightsLoader,
58
+ PPMissingLayer,
59
+ get_spec_layer_idx_from_weight_name,
60
+ is_pp_missing_parameter,
61
+ make_layers,
62
+ maybe_prefix,
63
+)
64
+from vllm.models.kimi_k3.amd.ops.attn_res import attn_res
65
+from vllm.sequence import IntermediateTensors
66
+from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
67
+from vllm.utils.math_utils import cdiv
68
+
69
+logger = init_logger(__name__)
70
+
71
+
72
+class KimiMLP(nn.Module):
73
+ def __init__(
74
+ self,
75
+ hidden_size: int,
76
+ intermediate_size: int,
77
+ hidden_act: str,
78
+ quant_config: QuantizationConfig | None = None,
79
+ reduce_results: bool = True,
80
+ prefix: str = "",
81
+ activation_situ_beta: float | None = None,
82
+ activation_situ_linear_beta: float | None = None,
83
+ ) -> None:
84
+ super().__init__()
85
+
86
+ self.gate_up_proj = MergedColumnParallelLinear(
87
+ hidden_size,
88
+ [intermediate_size] * 2,
89
+ bias=False,
90
+ quant_config=quant_config,
91
+ prefix=f"{prefix}.gate_up_proj",
92
+ )
93
+ self.down_proj = RowParallelLinear(
94
+ intermediate_size,
95
+ hidden_size,
96
+ bias=False,
97
+ quant_config=quant_config,
98
+ reduce_results=reduce_results,
99
+ prefix=f"{prefix}.down_proj",
100
+ )
101
+ if hidden_act == "silu":
102
+ self.act_fn = SiluAndMul()
103
+ elif hidden_act == "situ":
104
+ self.act_fn = SituAndMul(
105
+ beta=activation_situ_beta or 1.0,
106
+ linear_beta=activation_situ_linear_beta,
107
+ )
108
+ else:
109
+ raise ValueError(
110
+ f"Unsupported activation: {hidden_act}. "
111
+ "Only silu and situ are supported."
112
+ )
113
+
114
+ def forward(self, x):
115
+ gate_up, _ = self.gate_up_proj(x)
116
+ x = self.act_fn(gate_up)
117
+ x, _ = self.down_proj(x)
118
+ return x
119
+
120
+
121
+class KimiRoutedOutputTransform(nn.Module):
122
+ def __init__(
123
+ self,
124
+ norm: RMSNorm | None,
125
+ up_proj: ReplicatedLinear,
126
+ ) -> None:
127
+ super().__init__()
128
+ self.norm = norm
129
+ self.up_proj = up_proj
130
+
131
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
132
+ if self.norm is not None:
133
+ hidden_states = self.norm(hidden_states)
134
+ hidden_states, _ = self.up_proj(hidden_states)
135
+ return hidden_states
136
+
137
+
138
+def _apply_attn_res(
139
+ prefix_sum: torch.Tensor,
140
+ block_residual: torch.Tensor,
141
+ proj: ReplicatedLinear,
142
+ norm: RMSNorm,
143
+ num_valid_blocks: int,
144
+) -> torch.Tensor:
145
+ if num_valid_blocks <= 0:
146
+ return prefix_sum
147
+
148
+ return attn_res(
149
+ prefix_sum,
150
+ block_residual,
151
+ norm.weight,
152
+ proj.weight.squeeze(0),
153
+ num_valid_blocks,
154
+ norm.variance_epsilon,
155
+ )
156
+
157
+
158
+class KimiMoE(nn.Module):
159
+ def __init__(
160
+ self,
161
+ config: KimiLinearConfig,
162
+ quant_config: QuantizationConfig | None = None,
163
+ prefix: str = "",
164
+ layer_idx: int = 0,
165
+ ):
166
+ super().__init__()
167
+ hidden_size = config.hidden_size
168
+ moe_intermediate_size = config.moe_intermediate_size
169
+ num_experts = config.num_experts
170
+ num_experts_per_token = config.num_experts_per_token
171
+ assert moe_intermediate_size is not None
172
+ assert num_experts is not None
173
+ assert num_experts_per_token is not None
174
+ moe_renormalize = config.moe_renormalize
175
+ routed_expert_hidden_size = config.routed_expert_hidden_size
176
+ self.use_latent_moe = routed_expert_hidden_size is not None
177
+ self.moe_hidden_size = (
178
+ routed_expert_hidden_size
179
+ if routed_expert_hidden_size is not None
180
+ else hidden_size
181
+ )
182
+ self.latent_moe_use_norm = config.latent_moe_use_norm
183
+ self.tp_size = get_tensor_model_parallel_world_size()
184
+ self.routed_scaling_factor = config.routed_scaling_factor
185
+ self.num_shared_experts = config.num_shared_experts
186
+ self.layer_idx = layer_idx
187
+ self.padded_moe_intermediate_size = moe_intermediate_size
188
+ min_moe_intermediate_per_partition = getattr(
189
+ config, "min_moe_intermediate_per_partition", 256
190
+ )
191
+ if self.tp_size > 1:
192
+ moe_intermediate_per_partition = moe_intermediate_size // self.tp_size
193
+ if moe_intermediate_per_partition < min_moe_intermediate_per_partition:
194
+ self.padded_moe_intermediate_size = (
195
+ min_moe_intermediate_per_partition * self.tp_size
196
+ )
197
+ activation_situ_beta = (
198
+ config.activation_situ_beta if config.hidden_act == "situ" else None
199
+ )
200
+ activation_situ_linear_beta = (
201
+ config.activation_situ_linear_beta if config.hidden_act == "situ" else None
202
+ )
203
+
204
+ # Route with fp32 logits for numerically stable expert selection.
205
+ self.gate = GateLinear(
206
+ input_size=hidden_size,
207
+ output_size=num_experts,
208
+ bias=False,
209
+ out_dtype=torch.float32,
210
+ prefix=f"{prefix}.gate",
211
+ )
212
+
213
+ self.gate.e_score_correction_bias = nn.Parameter(torch.empty(num_experts))
214
+
215
+ if self.num_shared_experts is not None:
216
+ shared_intermediate_size = moe_intermediate_size * self.num_shared_experts
217
+ self.shared_experts = KimiMLP(
218
+ hidden_size=config.hidden_size,
219
+ intermediate_size=shared_intermediate_size,
220
+ hidden_act=config.hidden_act,
221
+ quant_config=quant_config,
222
+ reduce_results=False,
223
+ prefix=f"{prefix}.shared_experts",
224
+ activation_situ_beta=activation_situ_beta,
225
+ activation_situ_linear_beta=activation_situ_linear_beta,
226
+ )
227
+ else:
228
+ self.shared_experts = None
229
+
230
+ self.routed_expert_down_proj: ReplicatedLinear | None
231
+ self.routed_expert_norm: RMSNorm | None
232
+ self.routed_expert_up_proj: ReplicatedLinear | None
233
+ self.routed_output_transform: KimiRoutedOutputTransform | None
234
+ if self.use_latent_moe:
235
+ self.routed_expert_down_proj = ReplicatedLinear(
236
+ hidden_size,
237
+ self.moe_hidden_size,
238
+ bias=False,
239
+ quant_config=None,
240
+ prefix=f"{prefix}.routed_expert_down_proj",
241
+ )
242
+ self.routed_expert_norm = (
243
+ RMSNorm(self.moe_hidden_size, eps=config.rms_norm_eps)
244
+ if self.latent_moe_use_norm
245
+ else None
246
+ )
247
+ self.routed_expert_up_proj = ReplicatedLinear(
248
+ self.moe_hidden_size,
249
+ hidden_size,
250
+ bias=False,
251
+ quant_config=None,
252
+ prefix=f"{prefix}.routed_expert_up_proj",
253
+ )
254
+ self.routed_output_transform = KimiRoutedOutputTransform(
255
+ self.routed_expert_norm, self.routed_expert_up_proj
256
+ )
257
+ else:
258
+ self.routed_expert_down_proj = None
259
+ self.routed_expert_norm = None
260
+ self.routed_expert_up_proj = None
261
+ self.routed_output_transform = None
262
+
263
+ self.experts = FusedMoE(
264
+ shared_experts=self.shared_experts,
265
+ num_experts=num_experts,
266
+ top_k=num_experts_per_token,
267
+ hidden_size=self.moe_hidden_size,
268
+ intermediate_size=self.padded_moe_intermediate_size,
269
+ activation=config.hidden_act,
270
+ activation_situ_beta=activation_situ_beta,
271
+ activation_situ_linear_beta=activation_situ_linear_beta,
272
+ renormalize=moe_renormalize,
273
+ quant_config=quant_config,
274
+ use_grouped_topk=config.use_grouped_topk,
275
+ num_expert_group=config.num_expert_group,
276
+ topk_group=config.topk_group,
277
+ prefix=f"{prefix}.experts",
278
+ scoring_func=config.moe_router_activation_func,
279
+ e_score_correction_bias=self.gate.e_score_correction_bias,
280
+ routed_scaling_factor=self.routed_scaling_factor,
281
+ routed_input_transform=self.routed_expert_down_proj,
282
+ routed_output_transform=self.routed_output_transform,
283
+ )
284
+ if self.padded_moe_intermediate_size != moe_intermediate_size:
285
+ w13_weight = getattr(self.experts, "w13_weight", None)
286
+ if w13_weight is None:
287
+ w13_weight = self.experts.w13_weight_packed
288
+ w2_weight = getattr(self.experts, "w2_weight", None)
289
+ if w2_weight is None:
290
+ w2_weight = self.experts.w2_weight_packed
291
+ w13_weight.data.zero_()
292
+ w2_weight.data.zero_()
293
+ self.experts.moe_config.intermediate_size_per_partition_unpadded = (
294
+ moe_intermediate_size // self.tp_size
295
+ )
296
+
297
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
298
+ num_tokens, hidden_size = hidden_states.shape
299
+ hidden_states = hidden_states.view(-1, hidden_size)
300
+ router_logits, _ = self.gate(hidden_states)
301
+ final_hidden_states = self.experts(
302
+ hidden_states=hidden_states, router_logits=router_logits
303
+ )
304
+ return final_hidden_states.view(num_tokens, hidden_size)
305
+
306
+
307
+class KimiMLAAttention(nn.Module):
308
+ """
309
+ Main reference: DeepseekV2 vllm Implementation
310
+ """
311
+
312
+ def __init__(
313
+ self,
314
+ config: KimiLinearConfig,
315
+ hidden_size: int,
316
+ num_heads: int,
317
+ qk_nope_head_dim: int,
318
+ qk_rope_head_dim: int,
319
+ v_head_dim: int,
320
+ q_lora_rank: int | None,
321
+ kv_lora_rank: int,
322
+ use_nope: bool = False,
323
+ cache_config: CacheConfig | None = None,
324
+ quant_config: QuantizationConfig | None = None,
325
+ prefix: str = "",
326
+ **kwargs,
327
+ ) -> None:
328
+ super().__init__()
329
+ self.hidden_size = hidden_size
330
+ self.qk_nope_head_dim = qk_nope_head_dim
331
+ self.qk_rope_head_dim = qk_rope_head_dim
332
+ self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
333
+ self.v_head_dim = v_head_dim
334
+ self.q_lora_rank = q_lora_rank
335
+ self.kv_lora_rank = kv_lora_rank
336
+ self.num_heads = num_heads
337
+ tp_size = get_tensor_model_parallel_world_size()
338
+ self.num_local_heads = num_heads // tp_size
339
+ self.scaling = self.qk_head_dim**-0.5
340
+ self.use_nope = use_nope
341
+ assert self.use_nope is True
342
+ assert num_heads % tp_size == 0
343
+ if self.q_lora_rank is not None:
344
+ self.fused_qkv_a_proj = MergedColumnParallelLinear(
345
+ self.hidden_size,
346
+ [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
347
+ bias=False,
348
+ quant_config=quant_config,
349
+ prefix=f"{prefix}.fused_qkv_a_proj",
350
+ disable_tp=True,
351
+ )
352
+ else:
353
+ self.kv_a_proj_with_mqa = ReplicatedLinear(
354
+ self.hidden_size,
355
+ self.kv_lora_rank + self.qk_rope_head_dim,
356
+ bias=False,
357
+ quant_config=quant_config,
358
+ prefix=f"{prefix}.kv_a_proj_with_mqa",
359
+ )
360
+ if self.q_lora_rank is not None:
361
+ self.q_a_layernorm = RMSNorm(
362
+ self.q_lora_rank,
363
+ eps=config.rms_norm_eps,
364
+ )
365
+ self.q_b_proj = ColumnParallelLinear(
366
+ self.q_lora_rank,
367
+ self.num_heads * self.qk_head_dim,
368
+ bias=False,
369
+ quant_config=quant_config,
370
+ prefix=f"{prefix}.q_b_proj",
371
+ )
372
+ else:
373
+ self.q_proj = ColumnParallelLinear(
374
+ self.hidden_size,
375
+ self.num_heads * self.qk_head_dim,
376
+ bias=False,
377
+ quant_config=quant_config,
378
+ prefix=f"{prefix}.q_proj",
379
+ )
380
+ self.kv_a_layernorm = RMSNorm(
381
+ self.kv_lora_rank,
382
+ eps=config.rms_norm_eps,
383
+ )
384
+ self.kv_b_proj = ColumnParallelLinear(
385
+ self.kv_lora_rank,
386
+ self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
387
+ bias=False,
388
+ quant_config=quant_config,
389
+ prefix=f"{prefix}.kv_b_proj",
390
+ )
391
+ self.o_proj = RowParallelLinear(
392
+ self.num_heads * self.v_head_dim,
393
+ self.hidden_size,
394
+ bias=False,
395
+ quant_config=quant_config,
396
+ prefix=f"{prefix}.o_proj",
397
+ )
398
+
399
+ self.use_output_gate = config.mla_use_output_gate
400
+ if self.use_output_gate:
401
+ projection_size = self.num_heads * self.v_head_dim
402
+ self.g_proj = ColumnParallelLinear(
403
+ self.hidden_size,
404
+ projection_size,
405
+ bias=False,
406
+ quant_config=quant_config,
407
+ prefix=f"{prefix}.g_proj",
408
+ )
409
+
410
+ # TODO: Remove this mypy workaround once the K3 PR is fully merged.
411
+ mla_modules = MLAModules( # type: ignore[call-arg]
412
+ kv_a_layernorm=self.kv_a_layernorm,
413
+ kv_b_proj=self.kv_b_proj,
414
+ rotary_emb=None,
415
+ o_proj=self.o_proj,
416
+ fused_qkv_a_proj=self.fused_qkv_a_proj
417
+ if self.q_lora_rank is not None
418
+ else None,
419
+ kv_a_proj_with_mqa=self.kv_a_proj_with_mqa
420
+ if self.q_lora_rank is None
421
+ else None,
422
+ q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None,
423
+ q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,
424
+ q_proj=self.q_proj if self.q_lora_rank is None else None,
425
+ indexer=None,
426
+ is_sparse=False,
427
+ topk_indices_buffer=None,
428
+ g_proj=getattr(self, "g_proj", None),
429
+ )
430
+ self.mla_attn = MultiHeadLatentAttentionWrapper(
431
+ self.hidden_size,
432
+ self.num_local_heads,
433
+ self.scaling,
434
+ self.qk_nope_head_dim,
435
+ self.qk_rope_head_dim,
436
+ self.v_head_dim,
437
+ self.q_lora_rank,
438
+ self.kv_lora_rank,
439
+ mla_modules,
440
+ cache_config,
441
+ quant_config,
442
+ prefix,
443
+ )
444
+
445
+ def forward(
446
+ self,
447
+ positions: torch.Tensor,
448
+ hidden_states: torch.Tensor,
449
+ output: torch.Tensor,
450
+ ) -> None:
451
+ output[:] = self.mla_attn(positions, hidden_states)
452
+
453
+
454
+class KimiDecoderLayer(nn.Module):
455
+ def __init__(
456
+ self,
457
+ config: KimiLinearConfig,
458
+ vllm_config: VllmConfig,
459
+ prefix: str = "",
460
+ ) -> None:
461
+ super().__init__()
462
+ self.hidden_size = config.hidden_size
463
+ self.layer_idx = int(prefix.rsplit(".", 1)[1])
464
+
465
+ self.is_moe = config.is_moe
466
+ layer_idx = self.layer_idx
467
+ model_config = vllm_config.model_config
468
+ cache_config = vllm_config.cache_config
469
+ quant_config = vllm_config.quant_config
470
+
471
+ if config.is_kda_layer(layer_idx):
472
+ self.self_attn = KimiGatedDeltaNetAttention(
473
+ config,
474
+ vllm_config,
475
+ prefix=f"{prefix}.self_attn",
476
+ )
477
+ else:
478
+ qk_nope_head_dim = config.qk_nope_head_dim
479
+ qk_rope_head_dim = config.qk_rope_head_dim
480
+ v_head_dim = config.v_head_dim
481
+ kv_lora_rank = config.kv_lora_rank
482
+ mla_use_nope = config.mla_use_nope
483
+ assert qk_nope_head_dim is not None
484
+ assert qk_rope_head_dim is not None
485
+ assert v_head_dim is not None
486
+ assert kv_lora_rank is not None
487
+ assert mla_use_nope is not None
488
+ self.self_attn = KimiMLAAttention(
489
+ layer_idx=layer_idx,
490
+ hidden_size=self.hidden_size,
491
+ num_heads=config.num_attention_heads,
492
+ quant_config=quant_config,
493
+ cache_config=cache_config,
494
+ model_config=model_config,
495
+ prefix=f"{prefix}.self_attn",
496
+ config=config,
497
+ qk_nope_head_dim=qk_nope_head_dim,
498
+ qk_rope_head_dim=qk_rope_head_dim,
499
+ v_head_dim=v_head_dim,
500
+ q_lora_rank=config.q_lora_rank,
501
+ kv_lora_rank=kv_lora_rank,
502
+ use_nope=mla_use_nope,
503
+ )
504
+
505
+ if (
506
+ self.is_moe
507
+ and config.num_experts is not None
508
+ and layer_idx >= config.first_k_dense_replace
509
+ and layer_idx % config.moe_layer_freq == 0
510
+ ):
511
+ self.block_sparse_moe = KimiMoE(
512
+ config=config,
513
+ quant_config=quant_config,
514
+ prefix=f"{prefix}.block_sparse_moe",
515
+ layer_idx=layer_idx,
516
+ )
517
+ self.mlp = self.block_sparse_moe
518
+ else:
519
+ self.mlp = KimiMLP(
520
+ hidden_size=self.hidden_size,
521
+ intermediate_size=config.intermediate_size,
522
+ hidden_act=config.hidden_act,
523
+ quant_config=quant_config,
524
+ prefix=f"{prefix}.mlp",
525
+ activation_situ_beta=config.activation_situ_beta,
526
+ activation_situ_linear_beta=config.activation_situ_linear_beta,
527
+ )
528
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
529
+ self.post_attention_layernorm = RMSNorm(
530
+ config.hidden_size, eps=config.rms_norm_eps
531
+ )
532
+
533
+ attn_res_block_size = config.attn_res_block_size
534
+ self.use_attn_residuals = attn_res_block_size is not None
535
+ if attn_res_block_size is not None:
536
+ self.attn_res_block_size = attn_res_block_size
537
+ self.is_block_write_layer = layer_idx % self.attn_res_block_size == 0
538
+ self.block_write_idx = layer_idx // self.attn_res_block_size
539
+ self.prev_valid_blocks = cdiv(layer_idx, self.attn_res_block_size)
540
+ self.self_attention_res_norm = RMSNorm(
541
+ config.hidden_size, eps=config.rms_norm_eps
542
+ )
543
+ self.mlp_res_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
544
+ self.self_attention_res_proj = ReplicatedLinear(
545
+ config.hidden_size,
546
+ 1,
547
+ bias=False,
548
+ quant_config=None,
549
+ prefix=f"{prefix}.self_attention_res_proj",
550
+ )
551
+ self.mlp_res_proj = ReplicatedLinear(
552
+ config.hidden_size,
553
+ 1,
554
+ bias=False,
555
+ quant_config=None,
556
+ prefix=f"{prefix}.mlp_res_proj",
557
+ )
558
+
559
+ def _run_self_attn(
560
+ self,
561
+ positions: torch.Tensor,
562
+ hidden_states: torch.Tensor,
563
+ ) -> torch.Tensor:
564
+ attn_output = torch.empty_like(hidden_states)
565
+ self.self_attn(
566
+ hidden_states=hidden_states,
567
+ positions=positions,
568
+ output=attn_output,
569
+ )
570
+ return attn_output
571
+
572
+ def forward(
573
+ self,
574
+ positions: torch.Tensor,
575
+ hidden_states: torch.Tensor,
576
+ residual: torch.Tensor | None,
577
+ **kwargs,
578
+ ) -> tuple[torch.Tensor, torch.Tensor]:
579
+ if self.use_attn_residuals:
580
+ assert residual is not None
581
+ return self.forward_attn_residual(positions, hidden_states, residual)
582
+
583
+ # Self Attention
584
+ if residual is None:
585
+ residual = hidden_states
586
+ hidden_states = self.input_layernorm(hidden_states)
587
+ else:
588
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
589
+
590
+ hidden_states = self._run_self_attn(positions, hidden_states)
591
+
592
+ # Fully Connected
593
+ hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
594
+ hidden_states = self.mlp(hidden_states)
595
+ return hidden_states, residual
596
+
597
+ def forward_attn_residual(
598
+ self,
599
+ positions: torch.Tensor,
600
+ hidden_states: torch.Tensor,
601
+ block_residual: torch.Tensor,
602
+ ) -> tuple[torch.Tensor, torch.Tensor]:
603
+ prefix_sum = hidden_states
604
+ hidden_states = _apply_attn_res(
605
+ prefix_sum,
606
+ block_residual,
607
+ self.self_attention_res_proj,
608
+ self.self_attention_res_norm,
609
+ self.prev_valid_blocks,
610
+ )
611
+
612
+ if self.is_block_write_layer:
613
+ block_residual[:, self.block_write_idx, :].copy_(prefix_sum)
614
+ prefix_sum = None
615
+
616
+ hidden_states = self.input_layernorm(hidden_states)
617
+ hidden_states = self._run_self_attn(positions, hidden_states)
618
+
619
+ if prefix_sum is not None:
620
+ prefix_sum = prefix_sum + hidden_states
621
+ else:
622
+ prefix_sum = hidden_states
623
+
624
+ mlp_valid_blocks = self.prev_valid_blocks + (
625
+ 1 if self.is_block_write_layer else 0
626
+ )
627
+ hidden_states = _apply_attn_res(
628
+ prefix_sum,
629
+ block_residual,
630
+ self.mlp_res_proj,
631
+ self.mlp_res_norm,
632
+ mlp_valid_blocks,
633
+ )
634
+
635
+ hidden_states = self.post_attention_layernorm(hidden_states)
636
+ hidden_states = self.mlp(hidden_states)
637
+ prefix_sum = prefix_sum + hidden_states
638
+ return prefix_sum, block_residual
639
+
640
+
641
+class KimiLinearModel(nn.Module, EagleModelMixin):
642
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
643
+ super().__init__()
644
+
645
+ config = vllm_config.model_config.hf_text_config
646
+ self.config = config
647
+
648
+ self.vocab_size = config.vocab_size
649
+
650
+ if get_pp_group().is_first_rank:
651
+ self.embed_tokens = VocabParallelEmbedding(
652
+ config.vocab_size,
653
+ config.hidden_size,
654
+ prefix=f"{prefix}.embed_tokens",
655
+ )
656
+ else:
657
+ self.embed_tokens = PPMissingLayer()
658
+
659
+ def get_layer(prefix: str):
660
+ return KimiDecoderLayer(
661
+ config,
662
+ vllm_config,
663
+ prefix,
664
+ )
665
+
666
+ self.start_layer, self.end_layer, self.layers = make_layers(
667
+ config.num_hidden_layers,
668
+ get_layer,
669
+ prefix=f"{prefix}.layers",
670
+ )
671
+
672
+ if get_pp_group().is_last_rank:
673
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
674
+ if config.attn_res_block_size is not None:
675
+ self.output_attn_res_norm = RMSNorm(
676
+ config.hidden_size, eps=config.rms_norm_eps
677
+ )
678
+ self.output_attn_res_proj = ReplicatedLinear(
679
+ config.hidden_size,
680
+ 1,
681
+ bias=False,
682
+ quant_config=None,
683
+ prefix=f"{prefix}.output_attn_res_proj",
684
+ )
685
+ else:
686
+ self.norm = PPMissingLayer()
687
+ if config.attn_res_block_size is not None:
688
+ self.output_attn_res_norm = PPMissingLayer()
689
+ self.output_attn_res_proj = PPMissingLayer()
690
+
691
+ world_size = get_tensor_model_parallel_world_size()
692
+ assert config.num_attention_heads % world_size == 0, (
693
+ "num_attention_heads must be divisible by world_size"
694
+ )
695
+
696
+ def make_empty_intermediate_tensors(
697
+ self,
698
+ batch_size: int,
699
+ dtype: torch.dtype,
700
+ device: torch.device,
701
+ ) -> IntermediateTensors:
702
+ residual_shape: tuple[int, ...] = (batch_size, self.config.hidden_size)
703
+ if self.config.attn_res_block_size is not None:
704
+ residual_shape = (
705
+ batch_size,
706
+ cdiv(self.start_layer, self.config.attn_res_block_size),
707
+ self.config.hidden_size,
708
+ )
709
+ return IntermediateTensors(
710
+ {
711
+ "hidden_states": torch.zeros(
712
+ (batch_size, self.config.hidden_size), dtype=dtype, device=device
713
+ ),
714
+ "residual": torch.zeros(residual_shape, dtype=dtype, device=device),
715
+ }
716
+ )
717
+
718
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
719
+ return self.embed_tokens(input_ids)
720
+
721
+ def _maybe_add_hidden_state(
722
+ self,
723
+ aux_hidden_states: list[torch.Tensor],
724
+ layer_idx: int,
725
+ hidden_states: torch.Tensor,
726
+ residual: torch.Tensor | None,
727
+ ) -> list[torch.Tensor]:
728
+ if self.config.attn_res_block_size is not None:
729
+ # attn-res `residual` is a block-state bank, not an additive
730
+ # residual; None makes the mixin capture the prefix sum directly.
731
+ residual = None
732
+ return super()._maybe_add_hidden_state(
733
+ aux_hidden_states, layer_idx, hidden_states, residual
734
+ )
735
+
736
+ def forward(
737
+ self,
738
+ input_ids: torch.Tensor | None,
739
+ positions: torch.Tensor,
740
+ intermediate_tensors: IntermediateTensors | None,
741
+ inputs_embeds: torch.Tensor | None = None,
742
+ **kwargs,
743
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
744
+ if get_pp_group().is_first_rank:
745
+ if inputs_embeds is not None:
746
+ hidden_states = inputs_embeds
747
+ else:
748
+ hidden_states = self.embed_input_ids(input_ids)
749
+ residual = None
750
+ else:
751
+ assert intermediate_tensors is not None
752
+ hidden_states = intermediate_tensors["hidden_states"]
753
+ residual = intermediate_tensors["residual"]
754
+
755
+ aux_hidden_states = self._maybe_add_hidden_state(
756
+ [], self.start_layer, hidden_states, residual
757
+ )
758
+
759
+ if self.config.attn_res_block_size is None:
760
+ for layer_idx, layer in enumerate(
761
+ self.layers[self.start_layer : self.end_layer],
762
+ start=self.start_layer,
763
+ ):
764
+ hidden_states, residual = layer(
765
+ positions=positions,
766
+ hidden_states=hidden_states,
767
+ residual=residual,
768
+ )
769
+ self._maybe_add_hidden_state(
770
+ aux_hidden_states, layer_idx + 1, hidden_states, residual
771
+ )
772
+
773
+ if not get_pp_group().is_last_rank:
774
+ return IntermediateTensors(
775
+ {"hidden_states": hidden_states, "residual": residual}
776
+ )
777
+
778
+ # NOTE: the final norm is applied in compute_logits instead of here,
779
+ # so the MTP draft model receives the pre-norm hidden states.
780
+ if residual is not None:
781
+ hidden_states = hidden_states + residual
782
+ if aux_hidden_states:
783
+ return hidden_states, aux_hidden_states
784
+ return hidden_states
785
+
786
+ attn_res_block_num = cdiv(self.end_layer, self.config.attn_res_block_size)
787
+ block_residual = hidden_states.new_empty(
788
+ hidden_states.size(0), attn_res_block_num, hidden_states.size(1)
789
+ )
790
+ if residual is not None:
791
+ block_residual[:, : residual.size(1), :].copy_(residual)
792
+ residual = block_residual
793
+
794
+ for layer_idx, layer in enumerate(
795
+ self.layers[self.start_layer : self.end_layer],
796
+ start=self.start_layer,
797
+ ):
798
+ hidden_states, residual = layer(
799
+ positions=positions,
800
+ hidden_states=hidden_states,
801
+ residual=residual,
802
+ )
803
+ if (layer_idx + 1) in self.aux_hidden_state_layers:
804
+ # AMD attn-res layer already returns prefix_sum + MLP delta as
805
+ # hidden_states; the override drops the block bank in residual.
806
+ self._maybe_add_hidden_state(
807
+ aux_hidden_states, layer_idx + 1, hidden_states, residual
808
+ )
809
+
810
+ if not get_pp_group().is_last_rank:
811
+ return IntermediateTensors(
812
+ {"hidden_states": hidden_states, "residual": residual}
813
+ )
814
+
815
+ hidden_states = _apply_attn_res(
816
+ hidden_states,
817
+ residual,
818
+ self.output_attn_res_proj,
819
+ self.output_attn_res_norm,
820
+ attn_res_block_num,
821
+ )
822
+ # NOTE: the final norm is applied in compute_logits instead of here, so
823
+ # the MTP draft model receives the pre-norm hidden states.
824
+ if aux_hidden_states:
825
+ return hidden_states, aux_hidden_states
826
+ return hidden_states
827
+
828
+ def load_weights(
829
+ self,
830
+ weights: Iterable[
831
+ tuple[str, torch.Tensor] | tuple[str, torch.Tensor, dict[str, Any]]
832
+ ],
833
+ ) -> set[str]:
834
+ kda_config = self.config.linear_attn_config
835
+ use_full_rank_gate = bool(
836
+ kda_config and kda_config.get("use_full_rank_gate", False)
837
+ )
838
+ beta_shard_id = 5 if use_full_rank_gate else 3
839
+ stacked_params_mapping = [
840
+ # (param_name, shard_name, shard_id)
841
+ (".in_proj_qkvgfab", ".q_proj", 0),
842
+ (".in_proj_qkvgfab", ".k_proj", 1),
843
+ (".in_proj_qkvgfab", ".v_proj", 2),
844
+ (".in_proj_qkvgfab", ".b_proj", beta_shard_id),
845
+ (".in_proj_qkvgfab", ".f_a_proj", 4),
846
+ (".conv1d", ".q_conv1d", 0),
847
+ (".conv1d", ".k_conv1d", 1),
848
+ (".conv1d", ".v_conv1d", 2),
849
+ (".gate_up_proj", ".gate_proj", 0),
850
+ (".gate_up_proj", ".up_proj", 1),
851
+ ]
852
+ if use_full_rank_gate:
853
+ stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3))
854
+ if getattr(self.config, "q_lora_rank", None) is not None:
855
+ stacked_params_mapping += [
856
+ (".fused_qkv_a_proj", ".q_a_proj", 0),
857
+ (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
858
+ ]
859
+ if self.config.is_moe:
860
+ # Params for weights, fp8 weight scales, fp8 activation scales
861
+ # (param_name, weight_name, expert_id, shard_id)
862
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
863
+ self,
864
+ ckpt_gate_proj_name="w1",
865
+ ckpt_down_proj_name="w2",
866
+ ckpt_up_proj_name="w3",
867
+ num_experts=self.config.num_experts,
868
+ )
869
+ else:
870
+ expert_params_mapping = []
871
+ params_dict = dict(self.named_parameters())
872
+ # Under the MXFP4 quant interface the routed experts register unpacked
873
+ # params (``w13_weight``), while the compressed-tensors checkpoint names
874
+ # them ``.weight_packed``. Rebind so the expert mapping resolves; scales
875
+ # already share the ``.weight_scale`` suffix.
876
+ experts_unpacked = not any(n.endswith("w13_weight_packed") for n in params_dict)
877
+ loaded_params: set[str] = set()
878
+ for args in weights:
879
+ name, loaded_weight = args[0], args[1]
880
+ kwargs: dict[str, Any] = args[2] if len(args) > 2 else {}
881
+ if "rotary_emb.inv_freq" in name:
882
+ continue
883
+ if experts_unpacked and name.endswith(".weight_packed"):
884
+ name = name.replace(".weight_packed", ".weight")
885
+
886
+ spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
887
+ if spec_layer is not None:
888
+ continue # skip spec decode layers for main model
889
+ if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
890
+ # Models trained using ColossalAI may include these tensors in
891
+ # the checkpoint. Skip them.
892
+ continue
893
+ for param_name, weight_name, shard_id in stacked_params_mapping:
894
+ if weight_name not in name:
895
+ continue
896
+ # We have mlp.experts[0].gate_proj in the checkpoint.
897
+ # Since we handle the experts below in expert_params_mapping,
898
+ # we need to skip here BEFORE we update the name, otherwise
899
+ # name will be updated to mlp.experts[0].gate_up_proj, which
900
+ # will then be updated below in expert_params_mapping
901
+ # for mlp.experts[0].gate_gate_up_proj, which breaks load.
902
+ if ("mlp.experts." in name) and name not in params_dict:
903
+ continue
904
+ name_mapped = name.replace(weight_name, param_name)
905
+ # Packed projections are only present on compatible layers.
906
+ if name_mapped not in params_dict:
907
+ continue
908
+ name = name_mapped
909
+ # Skip loading extra bias for GPTQ models.
910
+ if name.endswith(".bias") and name not in params_dict:
911
+ continue
912
+ if is_pp_missing_parameter(name, self):
913
+ continue
914
+ param = params_dict[name]
915
+ weight_loader = param.weight_loader
916
+ weight_loader(param, loaded_weight, shard_id)
917
+ break
918
+ else:
919
+ for (
920
+ expert_param_name,
921
+ expert_weight_name,
922
+ expert_id,
923
+ expert_shard_id,
924
+ ) in expert_params_mapping:
925
+ if expert_weight_name not in name:
926
+ continue
927
+ name = name.replace(expert_weight_name, expert_param_name)
928
+ if is_pp_missing_parameter(name, self):
929
+ continue
930
+ param = params_dict[name]
931
+ weight_loader = param.weight_loader
932
+ weight_loader(
933
+ param,
934
+ loaded_weight,
935
+ name,
936
+ expert_id=expert_id,
937
+ shard_id=expert_shard_id,
938
+ )
939
+ break
940
+ else:
941
+ # Skip loading extra bias for GPTQ models.
942
+ if (
943
+ name.endswith(".bias")
944
+ and name not in params_dict
945
+ and not self.config.is_linear_attn
946
+ ): # noqa: E501
947
+ continue
948
+ # Remapping the name of FP8 kv-scale.
949
+ remapped_name = maybe_remap_kv_scale_name(name, params_dict)
950
+ if remapped_name is None:
951
+ continue
952
+ name = remapped_name
953
+ if is_pp_missing_parameter(name, self):
954
+ continue
955
+
956
+ param = params_dict[name]
957
+ weight_loader = getattr(
958
+ param, "weight_loader", default_weight_loader
959
+ )
960
+ weight_loader(param, loaded_weight, **kwargs)
961
+ loaded_params.add(name)
962
+ return loaded_params
963
+
964
+
965
+class KimiLinearForCausalLM(
966
+ nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid
967
+):
968
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
969
+ super().__init__()
970
+ self.model_config = vllm_config.model_config
971
+ self.vllm_config = vllm_config
972
+ self.config = self.model_config.hf_config
973
+ quant_config = vllm_config.quant_config
974
+ self.quant_config = quant_config
975
+ self.model = KimiLinearModel(
976
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
977
+ )
978
+ if get_pp_group().is_last_rank:
979
+ self.lm_head = ParallelLMHead(
980
+ self.config.vocab_size,
981
+ self.config.hidden_size,
982
+ quant_config=quant_config,
983
+ prefix=maybe_prefix(prefix, "lm_head"),
984
+ )
985
+ else:
986
+ self.lm_head = PPMissingLayer()
987
+ logit_scale = getattr(self.config, "logit_scale", 1.0)
988
+ self.logits_processor = LogitsProcessor(
989
+ self.config.vocab_size, scale=logit_scale
990
+ )
991
+
992
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
993
+ return self.model.embed_input_ids(input_ids)
994
+
995
+ def make_empty_intermediate_tensors(
996
+ self,
997
+ batch_size: int,
998
+ dtype: torch.dtype,
999
+ device: torch.device,
1000
+ ) -> IntermediateTensors:
1001
+ return self.model.make_empty_intermediate_tensors(batch_size, dtype, device)
1002
+
1003
+ def forward(
1004
+ self,
1005
+ input_ids: torch.Tensor | None,
1006
+ positions: torch.Tensor,
1007
+ intermediate_tensors: IntermediateTensors | None = None,
1008
+ inputs_embeds: torch.Tensor | None = None,
1009
+ **kwargs,
1010
+ ) -> torch.Tensor | IntermediateTensors:
1011
+ hidden_states = self.model(
1012
+ input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
1013
+ )
1014
+ return hidden_states
1015
+
1016
+ @classmethod
1017
+ def get_mamba_state_dtype_from_config(
1018
+ cls,
1019
+ vllm_config: "VllmConfig",
1020
+ ) -> tuple[torch.dtype, torch.dtype]:
1021
+ return MambaStateDtypeCalculator.kda_state_dtype(
1022
+ vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype
1023
+ )
1024
+
1025
+ @classmethod
1026
+ def get_mamba_state_shape_from_config(
1027
+ cls, vllm_config: "VllmConfig"
1028
+ ) -> tuple[tuple[int, int], tuple[int, int, int]]:
1029
+ parallel_config = vllm_config.parallel_config
1030
+ hf_config = vllm_config.model_config.hf_config
1031
+ tp_size = parallel_config.tensor_parallel_size
1032
+ num_spec = (
1033
+ vllm_config.speculative_config.num_speculative_tokens
1034
+ if vllm_config.speculative_config
1035
+ else 0
1036
+ )
1037
+ return MambaStateShapeCalculator.kda_state_shape(
1038
+ tp_size,
1039
+ hf_config.linear_attn_config["num_heads"],
1040
+ hf_config.linear_attn_config["head_dim"],
1041
+ conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"],
1042
+ num_spec=num_spec,
1043
+ )
1044
+
1045
+ @classmethod
1046
+ def get_mamba_state_copy_func(
1047
+ cls,
1048
+ ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
1049
+ return MambaStateCopyFuncCalculator.kda_state_copy_func()
1050
+
1051
+ def compute_logits(
1052
+ self,
1053
+ hidden_states: torch.Tensor,
1054
+ ) -> torch.Tensor | None:
1055
+ # The model's final norm is applied here (not at the end of forward) so
1056
+ # that the pre-norm hidden states can be fed to the MTP draft model.
1057
+ hidden_states = self.model.norm(hidden_states, None)
1058
+ return self.logits_processor(self.lm_head, hidden_states)
1059
+
1060
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1061
+ loader = AutoWeightsLoader(
1062
+ self,
1063
+ skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
1064
+ )
1065
+ return loader.load_weights(weights)
vllm/models/kimi_k3/amd/model.py
new
+249
@@ -0,0 +1,249 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Kimi-K3 multimodal model implementation for vLLM."""
4
+
5
+from collections.abc import Iterable
6
+from typing import cast
7
+
8
+import torch
9
+from torch import nn
10
+
11
+from vllm.config import VllmConfig
12
+from vllm.model_executor.layers.quantization import QuantizationConfig
13
+from vllm.model_executor.layers.quantization.compressed_tensors import (
14
+ compressed_tensors,
15
+)
16
+from vllm.model_executor.models.interfaces import (
17
+ HasInnerState,
18
+ IsHybrid,
19
+ SupportsEagle3,
20
+ SupportsMultiModal,
21
+ SupportsPP,
22
+ SupportsQuant,
23
+)
24
+from vllm.model_executor.models.kimi_k25 import KimiK25MediaPixelInputs
25
+from vllm.model_executor.models.kimi_k25_vit import (
26
+ KimiK25MultiModalProjector,
27
+ MoonViT3dPretrainedModel,
28
+ vision_tower_forward,
29
+)
30
+from vllm.model_executor.models.utils import (
31
+ AutoWeightsLoader,
32
+ WeightsMapper,
33
+ init_vllm_registered_model,
34
+ maybe_prefix,
35
+)
36
+from vllm.model_executor.models.vision import is_vit_use_data_parallel
37
+from vllm.multimodal import MULTIMODAL_REGISTRY
38
+from vllm.multimodal.inputs import NestedTensors
39
+from vllm.platforms import current_platform
40
+from vllm.sequence import IntermediateTensors
41
+from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config
42
+
43
+from ..common.mm_preprocess import (
44
+ KimiK3DummyInputsBuilder,
45
+ KimiK3MultiModalProcessor,
46
+ KimiK3ProcessingInfo,
47
+)
48
+from .linear import KimiLinearForCausalLM
49
+
50
+
51
+@MULTIMODAL_REGISTRY.register_processor(
52
+ KimiK3MultiModalProcessor,
53
+ info=KimiK3ProcessingInfo,
54
+ dummy_inputs=KimiK3DummyInputsBuilder,
55
+)
56
+class KimiK3ForConditionalGeneration(
57
+ nn.Module,
58
+ SupportsMultiModal,
59
+ SupportsPP,
60
+ SupportsQuant,
61
+ SupportsEagle3,
62
+ HasInnerState,
63
+ IsHybrid,
64
+):
65
+ """Kimi-K3 model with Kimi-K2.5 vision and KimiLinear text."""
66
+
67
+ supports_encoder_tp_data = True
68
+
69
+ hf_to_vllm_mapper = WeightsMapper(
70
+ orig_to_new_prefix={
71
+ "language_model.layers.": "language_model.model.layers.",
72
+ "mm_projector.proj.0": "mm_projector.linear_1",
73
+ "mm_projector.proj.2": "mm_projector.linear_2",
74
+ }
75
+ )
76
+
77
+ @classmethod
78
+ def get_placeholder_str(cls, modality: str, i: int) -> str | None:
79
+ if modality == "image":
80
+ return "<|kimi_image_placeholder|>"
81
+ raise ValueError(f"Unsupported modality: {modality}")
82
+
83
+ def __init__(
84
+ self,
85
+ vllm_config: VllmConfig,
86
+ prefix: str = "",
87
+ ) -> None:
88
+ super().__init__()
89
+ model_config = vllm_config.model_config
90
+ config: KimiK3Config = model_config.hf_config
91
+ self.config = config
92
+ quant_config = vllm_config.quant_config
93
+
94
+ multimodal_config = model_config.multimodal_config
95
+ assert multimodal_config is not None
96
+ self.use_data_parallel = is_vit_use_data_parallel(
97
+ config.vision_config.num_attention_heads
98
+ )
99
+ self.hidden_size = config.text_config.hidden_size
100
+ self.device = current_platform.current_device()
101
+
102
+ with self._mark_tower_model(vllm_config, "image"):
103
+ self.vision_tower = MoonViT3dPretrainedModel(
104
+ config.vision_config,
105
+ quant_config=self._maybe_ignore_quant_config(quant_config),
106
+ prefix=maybe_prefix(prefix, "vision_tower"),
107
+ )
108
+ if self._maybe_ignore_quant_config(quant_config) is not None:
109
+ self.vision_tower = self.vision_tower.to(device=self.device)
110
+ else:
111
+ self.vision_tower = self.vision_tower.to(
112
+ device=self.device, dtype=model_config.dtype
113
+ )
114
+
115
+ self.mm_projector = KimiK25MultiModalProjector(
116
+ config=config.vision_config,
117
+ use_data_parallel=self.use_data_parallel,
118
+ quant_config=self._maybe_ignore_quant_config(quant_config),
119
+ prefix=maybe_prefix(prefix, "mm_projector"),
120
+ )
121
+ self.mm_projector = self.mm_projector.to(
122
+ device=self.device, dtype=model_config.dtype
123
+ )
124
+
125
+ self.quant_config = quant_config
126
+ with self._mark_language_model(vllm_config):
127
+ self.language_model = init_vllm_registered_model(
128
+ vllm_config=vllm_config,
129
+ hf_config=config.text_config,
130
+ prefix=maybe_prefix(prefix, "language_model"),
131
+ architectures=["KimiLinearForCausalLM"],
132
+ )
133
+ self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
134
+ self.language_model.make_empty_intermediate_tensors
135
+ )
136
+ self.media_placeholder: int = self.config.media_placeholder_token_id
137
+
138
+ def _maybe_ignore_quant_config(
139
+ self, quant_config: QuantizationConfig | None
140
+ ) -> QuantizationConfig | None:
141
+ if isinstance(quant_config, compressed_tensors.CompressedTensorsConfig):
142
+ return None
143
+ return quant_config
144
+
145
+ def _parse_and_validate_media_input(
146
+ self, **kwargs: object
147
+ ) -> KimiK25MediaPixelInputs | None:
148
+ pixel_values = kwargs.pop("pixel_values", None)
149
+ grid_thws = kwargs.pop("grid_thws", None)
150
+ if pixel_values is None:
151
+ return None
152
+
153
+ if isinstance(pixel_values, list):
154
+ pixel_values = torch.cat(cast(list[torch.Tensor], pixel_values), dim=0)
155
+ if not isinstance(pixel_values, torch.Tensor):
156
+ raise TypeError(
157
+ "pixel_values must be a tensor or a list of tensors, "
158
+ f"got {type(pixel_values)}"
159
+ )
160
+
161
+ if len(pixel_values.shape) == 5 or len(pixel_values.shape) == 3:
162
+ pixel_values = pixel_values.reshape(
163
+ pixel_values.shape[0] * pixel_values.shape[1], *pixel_values.shape[2:]
164
+ )
165
+
166
+ target_dtype = next(self.vision_tower.parameters()).dtype
167
+ pixel_values = pixel_values.to(target_dtype)
168
+ assert isinstance(grid_thws, torch.Tensor), (
169
+ f"expect grid_thws to be a tensor, got {type(grid_thws)}"
170
+ )
171
+ grid_thws = grid_thws.reshape(-1, grid_thws.shape[-1])
172
+ assert grid_thws.ndim == 2 and grid_thws.size(1) == 3, (
173
+ f"unexpected shape for grid_thws: {grid_thws.shape}"
174
+ )
175
+
176
+ return KimiK25MediaPixelInputs(
177
+ type="pixel_values",
178
+ pixel_values=pixel_values,
179
+ grid_thws=grid_thws,
180
+ )
181
+
182
+ def _process_media_input(
183
+ self, media_input: KimiK25MediaPixelInputs
184
+ ) -> list[torch.Tensor]:
185
+ media_features = vision_tower_forward(
186
+ self.vision_tower,
187
+ media_input["pixel_values"],
188
+ media_input["grid_thws"],
189
+ mm_projector=self.mm_projector,
190
+ use_data_parallel=self.use_data_parallel,
191
+ )
192
+ return media_features
193
+
194
+ def embed_multimodal(self, **kwargs: object) -> NestedTensors | None:
195
+ media_input = self._parse_and_validate_media_input(**kwargs)
196
+ if media_input is None:
197
+ return None
198
+ return self._process_media_input(media_input)
199
+
200
+ def forward(
201
+ self,
202
+ input_ids: torch.Tensor,
203
+ positions: torch.Tensor,
204
+ intermediate_tensors: IntermediateTensors | None = None,
205
+ inputs_embeds: torch.Tensor | None = None,
206
+ **kwargs: object,
207
+ ) -> IntermediateTensors:
208
+ if intermediate_tensors is not None:
209
+ inputs_embeds = None
210
+ hidden_states = self.language_model(
211
+ input_ids=input_ids,
212
+ positions=positions,
213
+ intermediate_tensors=intermediate_tensors,
214
+ inputs_embeds=inputs_embeds,
215
+ )
216
+ return hidden_states
217
+
218
+ def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
219
+ return self.language_model.compute_logits(hidden_states)
220
+
221
+ def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
222
+ return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs(
223
+ input_buffers, **kwargs
224
+ )
225
+
226
+ def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
227
+ return self.language_model.mamba_cache.get_seqlen_agnostic_capture_inputs(
228
+ batch_size
229
+ )
230
+
231
+ @classmethod
232
+ def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig):
233
+ text_config = vllm_config.model_config.hf_config.text_config
234
+ temp_vllm_config = vllm_config.with_hf_config(text_config)
235
+ return KimiLinearForCausalLM.get_mamba_state_dtype_from_config(temp_vllm_config)
236
+
237
+ @classmethod
238
+ def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig):
239
+ text_config = vllm_config.model_config.hf_config.text_config
240
+ temp_vllm_config = vllm_config.with_hf_config(text_config)
241
+ return KimiLinearForCausalLM.get_mamba_state_shape_from_config(temp_vllm_config)
242
+
243
+ @classmethod
244
+ def get_mamba_state_copy_func(cls):
245
+ return KimiLinearForCausalLM.get_mamba_state_copy_func()
246
+
247
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
248
+ loader = AutoWeightsLoader(self)
249
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/models/kimi_k3/amd/mtp.py
new
+403
@@ -0,0 +1,403 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inference-only Kimi-K3 Multi-Token-Prediction (MTP) draft model."""
4
+
5
+import copy
6
+from collections.abc import Iterable
7
+
8
+import torch
9
+import torch.nn as nn
10
+
11
+from vllm.config import VllmConfig
12
+from vllm.logger import init_logger
13
+from vllm.model_executor.layers.fused_moe import (
14
+ fused_moe_make_expert_params_mapping,
15
+)
16
+from vllm.model_executor.layers.layernorm import RMSNorm
17
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
18
+from vllm.model_executor.layers.quantization import QuantizationConfig
19
+from vllm.model_executor.layers.vocab_parallel_embedding import (
20
+ ParallelLMHead,
21
+ VocabParallelEmbedding,
22
+)
23
+from vllm.model_executor.model_loader.weight_utils import (
24
+ default_weight_loader,
25
+ maybe_remap_kv_scale_name,
26
+)
27
+from vllm.model_executor.models.utils import get_pp_missing_layer_names, maybe_prefix
28
+from vllm.sequence import IntermediateTensors
29
+from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
30
+
31
+from ..common.mtp import fused_mtp_input
32
+from .linear import KimiDecoderLayer, get_spec_layer_idx_from_weight_name
33
+
34
+logger = init_logger(__name__)
35
+
36
+
37
+class SharedHead(nn.Module):
38
+ def __init__(
39
+ self,
40
+ config: KimiLinearConfig,
41
+ prefix: str,
42
+ quant_config: QuantizationConfig | None = None,
43
+ ) -> None:
44
+ super().__init__()
45
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
46
+ self.head = ParallelLMHead(
47
+ config.vocab_size,
48
+ config.hidden_size,
49
+ quant_config=quant_config,
50
+ prefix=maybe_prefix(prefix, "head"),
51
+ )
52
+
53
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
54
+ return self.norm(hidden_states)
55
+
56
+
57
+class KimiK3MultiTokenPredictorLayer(nn.Module):
58
+ def __init__(
59
+ self,
60
+ config: KimiLinearConfig,
61
+ vllm_config: VllmConfig,
62
+ prefix: str,
63
+ ) -> None:
64
+ super().__init__()
65
+ self.config = config
66
+ quant_config = vllm_config.quant_config
67
+
68
+ self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
69
+ self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
70
+ self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
71
+
72
+ self.shared_head = SharedHead(
73
+ config=config, prefix=prefix, quant_config=quant_config
74
+ )
75
+ # The MTP block is a standard KimiDecoderLayer, but it must NOT use the
76
+ # attn-residual (block-residual) scheme even when the base model does:
77
+ # the draft starts from the target's hidden state, without the main
78
+ # model's accumulated per-block residual tensor. We disable it by
79
+ # shallow-copying the config with ``attn_res_block_size=None``.
80
+ block_config = copy.copy(config)
81
+ block_config.attn_res_block_size = None
82
+ # NOTE: the prefix must end in the numeric spec-layer index so that
83
+ # KimiDecoderLayer can parse ``layer_idx`` and pick MLA (full attn).
84
+ self.mtp_block = KimiDecoderLayer(block_config, vllm_config, prefix=prefix)
85
+
86
+ def forward(
87
+ self,
88
+ input_ids: torch.Tensor,
89
+ positions: torch.Tensor,
90
+ previous_hidden_states: torch.Tensor,
91
+ inputs_embeds: torch.Tensor | None = None,
92
+ spec_step_index: int = 0,
93
+ ) -> tuple[torch.Tensor, torch.Tensor]:
94
+ assert inputs_embeds is not None
95
+ hidden_states = self.eh_proj(
96
+ fused_mtp_input(
97
+ positions,
98
+ inputs_embeds,
99
+ previous_hidden_states,
100
+ self.enorm.weight,
101
+ self.hnorm.weight,
102
+ self.enorm.variance_epsilon,
103
+ )
104
+ )
105
+
106
+ hidden_states, residual = self.mtp_block(
107
+ positions=positions,
108
+ hidden_states=hidden_states,
109
+ residual=None,
110
+ )
111
+ # Produce the normalized logits input and the pre-norm recurrent state
112
+ # in one fused add-RMSNorm launch.
113
+ logits_hidden_states, hidden_states = self.shared_head.norm(
114
+ hidden_states, residual
115
+ )
116
+ return logits_hidden_states, hidden_states
117
+
118
+
119
+class KimiK3MultiTokenPredictor(nn.Module):
120
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
121
+ super().__init__()
122
+ config: KimiLinearConfig = vllm_config.model_config.hf_text_config
123
+ self.config = config
124
+ self.mtp_start_layer_idx = config.num_hidden_layers
125
+ self.num_mtp_layers = config.num_nextn_predict_layers
126
+
127
+ self.layers = torch.nn.ModuleDict(
128
+ {
129
+ str(idx): KimiK3MultiTokenPredictorLayer(
130
+ config, vllm_config, f"{prefix}.layers.{idx}"
131
+ )
132
+ for idx in range(
133
+ self.mtp_start_layer_idx,
134
+ self.mtp_start_layer_idx + self.num_mtp_layers,
135
+ )
136
+ }
137
+ )
138
+ self.embed_tokens = VocabParallelEmbedding(
139
+ config.vocab_size,
140
+ config.hidden_size,
141
+ prefix=maybe_prefix(prefix, "embed_tokens"),
142
+ )
143
+ self.logits_processor = LogitsProcessor(config.vocab_size)
144
+
145
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
146
+ return self.embed_tokens(input_ids)
147
+
148
+ def forward(
149
+ self,
150
+ input_ids: torch.Tensor,
151
+ positions: torch.Tensor,
152
+ previous_hidden_states: torch.Tensor,
153
+ inputs_embeds: torch.Tensor | None = None,
154
+ spec_step_idx: int = 0,
155
+ ) -> tuple[torch.Tensor, torch.Tensor]:
156
+ if inputs_embeds is None:
157
+ inputs_embeds = self.embed_tokens(input_ids)
158
+ current_step_idx = spec_step_idx % self.num_mtp_layers
159
+ return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
160
+ input_ids,
161
+ positions,
162
+ previous_hidden_states,
163
+ inputs_embeds,
164
+ current_step_idx,
165
+ )
166
+
167
+ def compute_logits(
168
+ self,
169
+ hidden_states: torch.Tensor,
170
+ spec_step_idx: int = 0,
171
+ ) -> torch.Tensor:
172
+ current_step_idx = spec_step_idx % self.num_mtp_layers
173
+ mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
174
+ logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states)
175
+ return logits
176
+
177
+
178
+class KimiK3MTP(nn.Module):
179
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
180
+ super().__init__()
181
+ self.config = vllm_config.model_config.hf_text_config
182
+ self.quant_config = vllm_config.quant_config
183
+ self.model = KimiK3MultiTokenPredictor(
184
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
185
+ )
186
+
187
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
188
+ return self.model.embed_input_ids(input_ids)
189
+
190
+ def forward(
191
+ self,
192
+ input_ids: torch.Tensor | None,
193
+ positions: torch.Tensor,
194
+ hidden_states: torch.Tensor,
195
+ intermediate_tensors: IntermediateTensors | None = None,
196
+ inputs_embeds: torch.Tensor | None = None,
197
+ spec_step_idx: int = 0,
198
+ ) -> tuple[torch.Tensor, torch.Tensor]:
199
+ return self.model(
200
+ input_ids,
201
+ positions,
202
+ hidden_states,
203
+ inputs_embeds,
204
+ spec_step_idx,
205
+ )
206
+
207
+ def compute_logits(
208
+ self,
209
+ hidden_states: torch.Tensor,
210
+ spec_step_idx: int = 0,
211
+ ) -> torch.Tensor | None:
212
+ return self.model.compute_logits(hidden_states, spec_step_idx)
213
+
214
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
215
+ # Mirror KimiLinearForCausalLM.load_weights naming: leading-dot shard
216
+ # names, q_lora-conditional fused QKV, and w1/w2/w3 expert weights.
217
+ kda_config = self.config.linear_attn_config
218
+ use_full_rank_gate = bool(
219
+ kda_config and kda_config.get("use_full_rank_gate", False)
220
+ )
221
+ beta_shard_id = 5 if use_full_rank_gate else 3
222
+ stacked_params_mapping = [
223
+ # (param_name, shard_name, shard_id)
224
+ (".in_proj_qkvgfab", ".q_proj", 0),
225
+ (".in_proj_qkvgfab", ".k_proj", 1),
226
+ (".in_proj_qkvgfab", ".v_proj", 2),
227
+ (".in_proj_qkvgfab", ".b_proj", beta_shard_id),
228
+ (".in_proj_qkvgfab", ".f_a_proj", 4),
229
+ (".conv1d", ".q_conv1d", 0),
230
+ (".conv1d", ".k_conv1d", 1),
231
+ (".conv1d", ".v_conv1d", 2),
232
+ (".gate_up_proj", ".gate_proj", 0),
233
+ (".gate_up_proj", ".up_proj", 1),
234
+ ]
235
+ if use_full_rank_gate:
236
+ stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3))
237
+ if getattr(self.config, "q_lora_rank", None) is not None:
238
+ stacked_params_mapping += [
239
+ (".fused_qkv_a_proj", ".q_a_proj", 0),
240
+ (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
241
+ ]
242
+
243
+ expert_params_mapping = (
244
+ fused_moe_make_expert_params_mapping(
245
+ self,
246
+ ckpt_gate_proj_name="w1",
247
+ ckpt_down_proj_name="w2",
248
+ ckpt_up_proj_name="w3",
249
+ num_experts=self.config.num_experts,
250
+ )
251
+ if self.config.is_moe
252
+ else []
253
+ )
254
+
255
+ pp_missing_layer_names = get_pp_missing_layer_names(self)
256
+ params_dict = dict(self.named_parameters())
257
+ # Under the MXFP4 quant interface the routed experts register unpacked
258
+ # params (``w13_weight``), while the compressed-tensors checkpoint names
259
+ # them ``.weight_packed``. Rebind so the expert mapping resolves; scales
260
+ # already share the ``.weight_scale`` suffix.
261
+ experts_unpacked = not any(n.endswith("w13_weight_packed") for n in params_dict)
262
+ loaded_params: set[str] = set()
263
+ for name, loaded_weight in weights:
264
+ if "rotary_emb.inv_freq" in name:
265
+ continue
266
+ # The multimodal checkpoint prefixes text weights with
267
+ # ``language_model.``; strip it so names match this draft model's
268
+ # parameter paths (``model.layers.{i}.``). Non-text weights
269
+ # (vision_tower, mm_projector, ...) never match a spec layer below.
270
+ if name.startswith("language_model."):
271
+ name = name[len("language_model.") :]
272
+ if experts_unpacked and name.endswith(".weight_packed"):
273
+ name = name.replace(".weight_packed", ".weight")
274
+ spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
275
+ if spec_layer is None:
276
+ continue
277
+ name = self._rewrite_spec_layer_name(spec_layer, name)
278
+
279
+ for param_name, weight_name, shard_id in stacked_params_mapping:
280
+ if weight_name not in name:
281
+ continue
282
+ # Routed experts (``.experts.{i}.w1/w2/w3``) are handled by the
283
+ # expert mapping below; skip them here. Shared experts
284
+ # (``.shared_experts.``) use gate/up_proj and fall through.
285
+ if ".experts." in name:
286
+ continue
287
+ name_mapped = name.replace(weight_name, param_name)
288
+ # Only take this mapping if the fused destination actually
289
+ # exists (e.g. QKV fusion is only present when q_lora is used).
290
+ if name_mapped not in params_dict:
291
+ continue
292
+ if name_mapped in pp_missing_layer_names:
293
+ continue
294
+ name = name_mapped
295
+ param = params_dict[name]
296
+ weight_loader = param.weight_loader
297
+ weight_loader(param, loaded_weight, shard_id)
298
+ break
299
+ else:
300
+ for (
301
+ expert_param_name,
302
+ expert_weight_name,
303
+ expert_id,
304
+ expert_shard_id,
305
+ ) in expert_params_mapping:
306
+ if expert_weight_name not in name:
307
+ continue
308
+ name_mapped = name.replace(expert_weight_name, expert_param_name)
309
+ if name_mapped in pp_missing_layer_names:
310
+ continue
311
+ param = params_dict[name_mapped]
312
+ weight_loader = param.weight_loader
313
+ weight_loader(
314
+ param,
315
+ loaded_weight,
316
+ name_mapped,
317
+ shard_id=expert_shard_id,
318
+ expert_id=expert_id,
319
+ )
320
+ name = name_mapped
321
+ break
322
+ else:
323
+ if name.endswith(".bias") and name not in params_dict:
324
+ continue
325
+ remapped_name = maybe_remap_kv_scale_name(name, params_dict)
326
+ if remapped_name is None:
327
+ continue
328
+ name = remapped_name
329
+
330
+ # The embedding is shared across MTP layers; only the first
331
+ # spec layer carries the hoisted (non-".layers") copy.
332
+ if spec_layer != self.model.mtp_start_layer_idx and (
333
+ ".layers" not in name
334
+ ):
335
+ continue
336
+ if name in pp_missing_layer_names:
337
+ continue
338
+ # The base model uses an attn-residual scheme whose per-layer
339
+ # weights (self_attention_res_*, mlp_res_*) are not used by
340
+ # the draft block; such names have no matching parameter and
341
+ # are safely skipped.
342
+ if name not in params_dict:
343
+ continue
344
+
345
+ param = params_dict[name]
346
+ weight_loader = getattr(
347
+ param, "weight_loader", default_weight_loader
348
+ )
349
+ weight_loader(param, loaded_weight)
350
+ loaded_params.add(name)
351
+
352
+ # Validate that weights were loaded for each expected MTP layer.
353
+ loaded_layers: set[int] = set()
354
+ for param_name in loaded_params:
355
+ spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name)
356
+ if spec_layer is not None:
357
+ loaded_layers.add(spec_layer)
358
+ for layer_idx in range(
359
+ self.model.mtp_start_layer_idx,
360
+ self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
361
+ ):
362
+ if layer_idx not in loaded_layers:
363
+ raise ValueError(
364
+ f"MTP speculative decoding layer {layer_idx} weights "
365
+ f"missing from checkpoint. The checkpoint may not include "
366
+ f"the MTP layer weights. Use a checkpoint that includes "
367
+ f"MTP layer weights, or disable speculative decoding."
368
+ )
369
+
370
+ return loaded_params
371
+
372
+ def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
373
+ """Rewrite a checkpoint weight name to this module's parameter path.
374
+
375
+ Top-level MTP submodules (enorm/hnorm/eh_proj/shared_head) stay under
376
+ ``model.layers.{spec_layer}.*``; the shared ``embed_tokens`` is hoisted
377
+ to ``model.*``; everything else is a transformer-block weight and gets
378
+ ``.mtp_block`` inserted.
379
+ """
380
+ spec_layer_weight_names = [
381
+ "embed_tokens",
382
+ "enorm",
383
+ "hnorm",
384
+ "eh_proj",
385
+ "shared_head",
386
+ ]
387
+ shared_weight_names = ["embed_tokens"]
388
+ spec_layer_weight = False
389
+ shared_weight = False
390
+ for weight_name in spec_layer_weight_names:
391
+ if weight_name in name:
392
+ spec_layer_weight = True
393
+ if weight_name in shared_weight_names:
394
+ shared_weight = True
395
+ break
396
+ if not spec_layer_weight:
397
+ name = name.replace(
398
+ f"model.layers.{spec_layer}.",
399
+ f"model.layers.{spec_layer}.mtp_block.",
400
+ )
401
+ elif shared_weight:
402
+ name = name.replace(f"model.layers.{spec_layer}.", "model.")
403
+ return name
vllm/models/kimi_k3/amd/ops/third_party/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/kimi_k3/amd/ops/third_party/kda/__init__.py
new
+47
@@ -0,0 +1,47 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+#
4
+# AMD/ROCm vendored copy of the Kimi-K3 KDA triton kernels.
5
+#
6
+# Provenance: mirror of vllm/models/kimi_k3/nvidia/ops/third_party/kda (tracker
7
+# mke-tracker @ 7adebfcf9; FLA vendored per PRs #39/#86). Split per-vendor so
8
+# AMD can carry gfx950-specific kernel changes without touching the NVIDIA copy.
9
+#
10
+# fla-org/flash-linear-attention#869 (the unmerged ROCm fixes our earlier
11
+# amd_fla shim carried against FLA 0.5.0) is covered here by the *newer* vendored
12
+# FLA rather than the literal patch:
13
+# - transpose-state-layout workaround: N/A (kernels rewritten; no
14
+# transpose_state_layout path remains),
15
+# - AMD autotune configs: present (is_amd num_warps/num_stages branches),
16
+# - OOB-mask correctness fix: present (all tl.load use mask=..., other=0).
17
+# Validated on gfx950: no core-dump, gsm8k 94.1%.
18
+#
19
+# AMD-specific deltas vs the NVIDIA copy: NONE yet (byte-identical). Keep in sync
20
+# with the NVIDIA copy on FLA updates; any divergence should be an intentional,
21
+# documented gfx950-specific change (a #869-style AMD-only fix).
22
+
23
+from .chunk import (
24
+ chunk_kda,
25
+ chunk_kda_fwd,
26
+ chunk_kda_with_fused_gate,
27
+ chunk_kda_with_fused_gate_fwd,
28
+ fused_kda_gate,
29
+ fused_kda_gate_chunk_cumsum,
30
+)
31
+from .fused_recurrent import (
32
+ fused_recurrent_kda,
33
+ fused_recurrent_kda_fwd,
34
+ fused_recurrent_kda_packed_decode,
35
+)
36
+
37
+__all__ = [
38
+ "chunk_kda",
39
+ "chunk_kda_fwd",
40
+ "chunk_kda_with_fused_gate",
41
+ "chunk_kda_with_fused_gate_fwd",
42
+ "fused_kda_gate",
43
+ "fused_kda_gate_chunk_cumsum",
44
+ "fused_recurrent_kda",
45
+ "fused_recurrent_kda_fwd",
46
+ "fused_recurrent_kda_packed_decode",
47
+]
vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py
new
+935
@@ -0,0 +1,935 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li
4
+#
5
+# This file contains code copied from the flash-linear-attention project.
6
+# The original source code was licensed under the MIT license and included
7
+# the following copyright notice:
8
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
9
+# ruff: noqa: E501
10
+
11
+
12
+import torch
13
+
14
+from vllm.third_party.flash_linear_attention.ops.chunk_delta_h import (
15
+ chunk_gated_delta_rule_fwd_h,
16
+)
17
+from vllm.third_party.flash_linear_attention.ops.cumsum import chunk_local_cumsum
18
+from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices
19
+from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
20
+from vllm.third_party.flash_linear_attention.ops.op import exp2, log
21
+from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE, is_amd
22
+from vllm.triton_utils import tl, triton
23
+from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2
24
+
25
+from .chunk_intra import chunk_kda_fwd_intra
26
+
27
+BT_LIST_AUTOTUNE = [32, 64, 128]
28
+NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32]
29
+
30
+
31
+@triton.heuristics(
32
+ {
33
+ "STORE_QG": lambda args: args["qg"] is not None,
34
+ "STORE_KG": lambda args: args["kg"] is not None,
35
+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
36
+ }
37
+)
38
+@triton.autotune(
39
+ configs=[
40
+ triton.Config({}, num_warps=num_warps, num_stages=num_stages)
41
+ for num_warps in [2, 4, 8]
42
+ for num_stages in [2, 3, 4]
43
+ ],
44
+ key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"],
45
+)
46
+@triton.jit(do_not_specialize=["T"])
47
+def recompute_w_u_fwd_kernel(
48
+ q,
49
+ k,
50
+ qg,
51
+ kg,
52
+ v,
53
+ beta,
54
+ w,
55
+ u,
56
+ A,
57
+ gk,
58
+ cu_seqlens,
59
+ chunk_indices,
60
+ T,
61
+ H: tl.constexpr,
62
+ K: tl.constexpr,
63
+ V: tl.constexpr,
64
+ BT: tl.constexpr,
65
+ BK: tl.constexpr,
66
+ BV: tl.constexpr,
67
+ STORE_QG: tl.constexpr,
68
+ STORE_KG: tl.constexpr,
69
+ IS_VARLEN: tl.constexpr,
70
+ DOT_PRECISION: tl.constexpr,
71
+):
72
+ i_t, i_bh = tl.program_id(0), tl.program_id(1)
73
+ i_b, i_h = i_bh // H, i_bh % H
74
+ if IS_VARLEN:
75
+ i_n, i_t = (
76
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
77
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
78
+ )
79
+ bos, eos = (
80
+ tl.load(cu_seqlens + i_n).to(tl.int32),
81
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
82
+ )
83
+ T = eos - bos
84
+ else:
85
+ bos, eos = i_b * T, i_b * T + T
86
+ p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
87
+ b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32)
88
+
89
+ p_A = tl.make_block_ptr(
90
+ A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)
91
+ )
92
+ b_A = tl.load(p_A, boundary_check=(0, 1))
93
+
94
+ for i_v in range(tl.cdiv(V, BV)):
95
+ p_v = tl.make_block_ptr(
96
+ v + (bos * H + i_h) * V,
97
+ (T, V),
98
+ (H * V, 1),
99
+ (i_t * BT, i_v * BV),
100
+ (BT, BV),
101
+ (1, 0),
102
+ )
103
+ p_u = tl.make_block_ptr(
104
+ u + (bos * H + i_h) * V,
105
+ (T, V),
106
+ (H * V, 1),
107
+ (i_t * BT, i_v * BV),
108
+ (BT, BV),
109
+ (1, 0),
110
+ )
111
+ b_v = tl.load(p_v, boundary_check=(0, 1))
112
+ b_vb = (b_v * b_b[:, None]).to(b_v.dtype)
113
+ b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION)
114
+ tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1))
115
+
116
+ for i_k in range(tl.cdiv(K, BK)):
117
+ p_w = tl.make_block_ptr(
118
+ w + (bos * H + i_h) * K,
119
+ (T, K),
120
+ (H * K, 1),
121
+ (i_t * BT, i_k * BK),
122
+ (BT, BK),
123
+ (1, 0),
124
+ )
125
+ p_k = tl.make_block_ptr(
126
+ k + (bos * H + i_h) * K,
127
+ (T, K),
128
+ (H * K, 1),
129
+ (i_t * BT, i_k * BK),
130
+ (BT, BK),
131
+ (1, 0),
132
+ )
133
+ b_k = tl.load(p_k, boundary_check=(0, 1))
134
+ b_kb = b_k * b_b[:, None]
135
+
136
+ p_gk = tl.make_block_ptr(
137
+ gk + (bos * H + i_h) * K,
138
+ (T, K),
139
+ (H * K, 1),
140
+ (i_t * BT, i_k * BK),
141
+ (BT, BK),
142
+ (1, 0),
143
+ )
144
+ b_gk = tl.load(p_gk, boundary_check=(0, 1))
145
+ b_kb *= exp2(b_gk)
146
+ if STORE_QG:
147
+ p_q = tl.make_block_ptr(
148
+ q + (bos * H + i_h) * K,
149
+ (T, K),
150
+ (H * K, 1),
151
+ (i_t * BT, i_k * BK),
152
+ (BT, BK),
153
+ (1, 0),
154
+ )
155
+ p_qg = tl.make_block_ptr(
156
+ qg + (bos * H + i_h) * K,
157
+ (T, K),
158
+ (H * K, 1),
159
+ (i_t * BT, i_k * BK),
160
+ (BT, BK),
161
+ (1, 0),
162
+ )
163
+ b_q = tl.load(p_q, boundary_check=(0, 1))
164
+ b_qg = b_q * exp2(b_gk)
165
+ tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1))
166
+ if STORE_KG:
167
+ last_idx = min(i_t * BT + BT, T) - 1
168
+
169
+ o_k = i_k * BK + tl.arange(0, BK)
170
+ m_k = o_k < K
171
+ b_gn = tl.load(
172
+ gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0
173
+ )
174
+ b_kg = b_k * exp2(b_gn - b_gk)
175
+
176
+ p_kg = tl.make_block_ptr(
177
+ kg + (bos * H + i_h) * K,
178
+ (T, K),
179
+ (H * K, 1),
180
+ (i_t * BT, i_k * BK),
181
+ (BT, BK),
182
+ (1, 0),
183
+ )
184
+ tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1))
185
+
186
+ b_w = tl.dot(b_A, b_kb.to(b_k.dtype))
187
+ tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))
188
+
189
+
190
+def recompute_w_u_fwd(
191
+ k: torch.Tensor,
192
+ v: torch.Tensor,
193
+ beta: torch.Tensor,
194
+ A: torch.Tensor,
195
+ q: torch.Tensor | None = None,
196
+ gk: torch.Tensor | None = None,
197
+ cu_seqlens: torch.Tensor | None = None,
198
+ chunk_indices: torch.Tensor | None = None,
199
+) -> tuple[torch.Tensor, torch.Tensor]:
200
+ B, T, H, K, V = *k.shape, v.shape[-1]
201
+ BT = A.shape[-1]
202
+ BK = 64
203
+ BV = 64
204
+
205
+ if chunk_indices is None and cu_seqlens is not None:
206
+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
207
+ NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
208
+
209
+ w = torch.empty_like(k)
210
+ u = torch.empty_like(v)
211
+ kg = torch.empty_like(k) if gk is not None else None
212
+ recompute_w_u_fwd_kernel[(NT, B * H)](
213
+ q=q,
214
+ k=k,
215
+ qg=None,
216
+ kg=kg,
217
+ v=v,
218
+ beta=beta,
219
+ w=w,
220
+ u=u,
221
+ A=A,
222
+ gk=gk,
223
+ cu_seqlens=cu_seqlens,
224
+ chunk_indices=chunk_indices,
225
+ T=T,
226
+ H=H,
227
+ K=K,
228
+ V=V,
229
+ BT=BT,
230
+ BK=BK,
231
+ BV=BV,
232
+ DOT_PRECISION="ieee",
233
+ )
234
+ return w, u, None, kg
235
+
236
+
237
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
238
+@triton.autotune(
239
+ configs=[
240
+ triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages)
241
+ for BK in [32, 64]
242
+ for BV in [64, 128]
243
+ for num_warps in [2, 4, 8]
244
+ for num_stages in [2, 3, 4]
245
+ ],
246
+ key=["BT"],
247
+)
248
+@triton.jit(do_not_specialize=["T"])
249
+def chunk_gla_fwd_kernel_o(
250
+ q,
251
+ v,
252
+ g,
253
+ h,
254
+ o,
255
+ A,
256
+ cu_seqlens,
257
+ chunk_indices,
258
+ scale,
259
+ T,
260
+ H: tl.constexpr,
261
+ K: tl.constexpr,
262
+ V: tl.constexpr,
263
+ BT: tl.constexpr,
264
+ BK: tl.constexpr,
265
+ BV: tl.constexpr,
266
+ IS_VARLEN: tl.constexpr,
267
+):
268
+ i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
269
+ i_b, i_h = i_bh // H, i_bh % H
270
+ if IS_VARLEN:
271
+ i_tg = i_t
272
+ i_n, i_t = (
273
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
274
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
275
+ )
276
+ bos, eos = (
277
+ tl.load(cu_seqlens + i_n).to(tl.int32),
278
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
279
+ )
280
+ T = eos - bos
281
+ NT = tl.cdiv(T, BT)
282
+ else:
283
+ NT = tl.cdiv(T, BT)
284
+ i_tg = i_b * NT + i_t
285
+ bos, eos = i_b * T, i_b * T + T
286
+
287
+ m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :]
288
+
289
+ b_o = tl.zeros([BT, BV], dtype=tl.float32)
290
+ for i_k in range(tl.cdiv(K, BK)):
291
+ p_q = tl.make_block_ptr(
292
+ q + (bos * H + i_h) * K,
293
+ (T, K),
294
+ (H * K, 1),
295
+ (i_t * BT, i_k * BK),
296
+ (BT, BK),
297
+ (1, 0),
298
+ )
299
+ p_g = tl.make_block_ptr(
300
+ g + (bos * H + i_h) * K,
301
+ (T, K),
302
+ (H * K, 1),
303
+ (i_t * BT, i_k * BK),
304
+ (BT, BK),
305
+ (1, 0),
306
+ )
307
+ p_h = tl.make_block_ptr(
308
+ h + (i_tg * H + i_h) * K * V,
309
+ (V, K),
310
+ (K, 1),
311
+ (i_v * BV, i_k * BK),
312
+ (BV, BK),
313
+ (1, 0),
314
+ )
315
+
316
+ # [BT, BK]
317
+ b_q = tl.load(p_q, boundary_check=(0, 1))
318
+ b_q = (b_q * scale).to(b_q.dtype)
319
+ # [BT, BK]
320
+ b_g = tl.load(p_g, boundary_check=(0, 1))
321
+ # [BT, BK]
322
+ b_qg = (b_q * exp2(b_g)).to(b_q.dtype)
323
+ # [BV, BK]
324
+ b_h = tl.load(p_h, boundary_check=(0, 1))
325
+ # [BT, BV]
326
+ if i_k >= 0:
327
+ b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype))
328
+ p_v = tl.make_block_ptr(
329
+ v + (bos * H + i_h) * V,
330
+ (T, V),
331
+ (H * V, 1),
332
+ (i_t * BT, i_v * BV),
333
+ (BT, BV),
334
+ (1, 0),
335
+ )
336
+ p_o = tl.make_block_ptr(
337
+ o + (bos * H + i_h) * V,
338
+ (T, V),
339
+ (H * V, 1),
340
+ (i_t * BT, i_v * BV),
341
+ (BT, BV),
342
+ (1, 0),
343
+ )
344
+ p_A = tl.make_block_ptr(
345
+ A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)
346
+ )
347
+ # [BT, BV]
348
+ b_v = tl.load(p_v, boundary_check=(0, 1))
349
+ # [BT, BT]
350
+ b_A = tl.load(p_A, boundary_check=(0, 1))
351
+ b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype)
352
+ b_o += tl.dot(b_A, b_v, allow_tf32=False)
353
+ tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
354
+
355
+
356
+def chunk_gla_fwd_o_gk(
357
+ q: torch.Tensor,
358
+ v: torch.Tensor,
359
+ g: torch.Tensor,
360
+ A: torch.Tensor,
361
+ h: torch.Tensor,
362
+ o: torch.Tensor,
363
+ scale: float,
364
+ cu_seqlens: torch.Tensor | None = None,
365
+ chunk_indices: torch.Tensor | None = None,
366
+ chunk_size: int = FLA_CHUNK_SIZE,
367
+):
368
+ B, T, H, K, V = *q.shape, v.shape[-1]
369
+ BT = chunk_size
370
+
371
+ if chunk_indices is None and cu_seqlens is not None:
372
+ chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
373
+ NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
374
+
375
+ def grid(meta):
376
+ return (cdiv(V, meta["BV"]), NT, B * H)
377
+
378
+ chunk_gla_fwd_kernel_o[grid](
379
+ q=q,
380
+ v=v,
381
+ g=g,
382
+ h=h,
383
+ o=o,
384
+ A=A,
385
+ cu_seqlens=cu_seqlens,
386
+ chunk_indices=chunk_indices,
387
+ scale=scale,
388
+ T=T,
389
+ H=H,
390
+ K=K,
391
+ V=V,
392
+ BT=BT,
393
+ )
394
+ return o
395
+
396
+
397
+@triton.heuristics(
398
+ {
399
+ "HAS_BIAS": lambda args: args["g_bias"] is not None,
400
+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
401
+ }
402
+)
403
+@triton.autotune(
404
+ configs=[
405
+ triton.Config({"BS": BS}, num_warps=num_warps)
406
+ for BS in [32, 64]
407
+ for num_warps in [2, 4, 8]
408
+ ],
409
+ key=["H", "S", "BT", "IS_VARLEN"],
410
+)
411
+@triton.jit(do_not_specialize=["T"])
412
+def kda_gate_chunk_cumsum_vector_kernel(
413
+ s,
414
+ raw_beta,
415
+ A_log,
416
+ g_bias,
417
+ o,
418
+ beta_out,
419
+ cu_seqlens,
420
+ chunk_indices,
421
+ cumsum_scale,
422
+ lower_bound,
423
+ beta,
424
+ threshold,
425
+ T,
426
+ stride_beta_batch,
427
+ stride_beta_token,
428
+ stride_beta_head,
429
+ H: tl.constexpr,
430
+ S: tl.constexpr,
431
+ BT: tl.constexpr,
432
+ BS: tl.constexpr,
433
+ HAS_BIAS: tl.constexpr,
434
+ IS_VARLEN: tl.constexpr,
435
+ USE_LOWER_BOUND: tl.constexpr,
436
+):
437
+ i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
438
+ i_b, i_h = i_bh // H, i_bh % H
439
+ if IS_VARLEN:
440
+ i_n, i_t = (
441
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
442
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
443
+ )
444
+ bos, eos = (
445
+ tl.load(cu_seqlens + i_n).to(tl.int32),
446
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
447
+ )
448
+ T = eos - bos
449
+ else:
450
+ bos = i_b * T
451
+
452
+ if i_s == 0:
453
+ o_beta_t = tl.arange(0, BT)
454
+ m_beta = i_t * BT + o_beta_t < T
455
+ if IS_VARLEN:
456
+ p_beta = (
457
+ raw_beta
458
+ + (bos + i_t * BT + o_beta_t) * stride_beta_token
459
+ + i_h * stride_beta_head
460
+ )
461
+ else:
462
+ p_beta = (
463
+ raw_beta
464
+ + i_b * stride_beta_batch
465
+ + (i_t * BT + o_beta_t) * stride_beta_token
466
+ + i_h * stride_beta_head
467
+ )
468
+ b_beta = tl.load(p_beta, mask=m_beta, other=0.0).to(tl.float32)
469
+ p_beta_out = beta_out + (bos + i_t * BT + o_beta_t) * H + i_h
470
+ tl.store(p_beta_out, tl.sigmoid(b_beta), mask=m_beta)
471
+ return
472
+
473
+ i_s -= 1
474
+
475
+ p_s = tl.make_block_ptr(
476
+ s + (bos * H + i_h) * S,
477
+ (T, S),
478
+ (H * S, 1),
479
+ (i_t * BT, i_s * BS),
480
+ (BT, BS),
481
+ (1, 0),
482
+ )
483
+ p_o = tl.make_block_ptr(
484
+ o + (bos * H + i_h) * S,
485
+ (T, S),
486
+ (H * S, 1),
487
+ (i_t * BT, i_s * BS),
488
+ (BT, BS),
489
+ (1, 0),
490
+ )
491
+
492
+ b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32)
493
+ if HAS_BIAS:
494
+ p_bias = tl.make_block_ptr(
495
+ g_bias + i_h * S,
496
+ (S,),
497
+ (1,),
498
+ (i_s * BS,),
499
+ (BS,),
500
+ (0,),
501
+ )
502
+ b_bias = tl.load(p_bias, boundary_check=(0,)).to(tl.float32)
503
+ b_s += b_bias[None, :]
504
+
505
+ b_a = tl.exp(tl.load(A_log + i_h).to(tl.float32))
506
+ if USE_LOWER_BOUND:
507
+ b_gate = lower_bound * tl.sigmoid(b_a * b_s)
508
+ else:
509
+ b_g_scaled = b_s * beta
510
+ b_softplus = tl.where(
511
+ b_g_scaled > threshold,
512
+ b_s,
513
+ (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)),
514
+ )
515
+ b_gate = -b_a * b_softplus
516
+
517
+ # Boundary loads return zero, but bias and gate activation can make padded
518
+ # rows nonzero. Padding trails valid rows, so it only affects masked stores.
519
+ b_o = tl.cumsum(b_gate, axis=0) * cumsum_scale
520
+ tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
521
+
522
+
523
+def fused_kda_gate_chunk_cumsum(
524
+ raw_g: torch.Tensor,
525
+ raw_beta: torch.Tensor,
526
+ A_log: torch.Tensor,
527
+ g_bias: torch.Tensor | None = None,
528
+ beta: float = 1.0,
529
+ threshold: float = 20.0,
530
+ lower_bound: float | None = None,
531
+ cu_seqlens: torch.Tensor | None = None,
532
+ chunk_indices: torch.Tensor | None = None,
533
+ chunk_size: int = FLA_CHUNK_SIZE,
534
+ output_dtype: torch.dtype | None = torch.float,
535
+) -> tuple[torch.Tensor, torch.Tensor]:
536
+ if cu_seqlens is not None:
537
+ assert raw_g.shape[0] == 1, (
538
+ "Only batch size 1 is supported when cu_seqlens are provided"
539
+ )
540
+ B, T, H, D = raw_g.shape
541
+ if raw_beta.shape != (B, T, H):
542
+ raise ValueError(f"Expected raw_beta shape {(B, T, H)}, got {raw_beta.shape}")
543
+ if chunk_indices is None and cu_seqlens is not None:
544
+ chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
545
+ NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices)
546
+
547
+ A_log = A_log.reshape(-1)
548
+ if g_bias is not None:
549
+ g_bias = g_bias.reshape(-1)
550
+ y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype)
551
+ beta_out = torch.empty(raw_beta.shape, device=raw_beta.device, dtype=torch.float32)
552
+
553
+ def grid(meta):
554
+ # For each (chunk, head), program 0 computes beta without extending a
555
+ # gate tile's critical path. The remaining programs cover the gate dim.
556
+ return (cdiv(meta["S"], meta["BS"]) + 1, NT, B * H)
557
+
558
+ kda_gate_chunk_cumsum_vector_kernel[grid](
559
+ s=raw_g,
560
+ raw_beta=raw_beta,
561
+ A_log=A_log,
562
+ g_bias=g_bias,
563
+ o=y,
564
+ beta_out=beta_out,
565
+ cu_seqlens=cu_seqlens,
566
+ chunk_indices=chunk_indices,
567
+ # RCP_LN2 folds in the natural-log -> log2 conversion so downstream
568
+ # exp2-based kernels reproduce exp(g). Keep this in sync with the
569
+ # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`.
570
+ cumsum_scale=RCP_LN2,
571
+ lower_bound=lower_bound or 0.0,
572
+ beta=beta,
573
+ threshold=threshold,
574
+ T=T,
575
+ stride_beta_batch=raw_beta.stride(0),
576
+ stride_beta_token=raw_beta.stride(1),
577
+ stride_beta_head=raw_beta.stride(2),
578
+ H=H,
579
+ S=D,
580
+ BT=chunk_size,
581
+ USE_LOWER_BOUND=lower_bound is not None,
582
+ )
583
+ return y, beta_out
584
+
585
+
586
+def _chunk_kda_fwd_with_cumulative_g(
587
+ q: torch.Tensor,
588
+ k: torch.Tensor,
589
+ v: torch.Tensor,
590
+ g: torch.Tensor,
591
+ beta: torch.Tensor,
592
+ scale: float,
593
+ initial_state: torch.Tensor,
594
+ output_final_state: bool,
595
+ cu_seqlens: torch.Tensor | None = None,
596
+ chunk_indices: torch.Tensor | None = None,
597
+ chunk_size: int = FLA_CHUNK_SIZE,
598
+ safe_gate: bool = False,
599
+):
600
+ # `g` must already be chunk-local cumulatively-summed AND scaled by
601
+ # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)).
602
+ # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of
603
+ # calling this helper directly unless that invariant is upheld.
604
+ Aqk, A = chunk_kda_fwd_intra(
605
+ q=q,
606
+ k=k,
607
+ gk=g,
608
+ beta=beta,
609
+ scale=scale,
610
+ cu_seqlens=cu_seqlens,
611
+ chunk_indices=chunk_indices,
612
+ chunk_size=chunk_size,
613
+ safe_gate=safe_gate,
614
+ )
615
+ w, u, _, kg = recompute_w_u_fwd(
616
+ k=k,
617
+ v=v,
618
+ beta=beta,
619
+ A=A,
620
+ gk=g,
621
+ cu_seqlens=cu_seqlens,
622
+ chunk_indices=chunk_indices,
623
+ )
624
+ del A
625
+ h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
626
+ k=kg,
627
+ w=w,
628
+ u=u,
629
+ gk=g,
630
+ initial_state=initial_state,
631
+ output_final_state=output_final_state,
632
+ cu_seqlens=cu_seqlens,
633
+ chunk_indices=chunk_indices,
634
+ use_exp2=True,
635
+ )
636
+ del w, u, kg
637
+ o = chunk_gla_fwd_o_gk(
638
+ q=q,
639
+ v=v_new,
640
+ g=g,
641
+ A=Aqk,
642
+ h=h,
643
+ o=v,
644
+ scale=scale,
645
+ cu_seqlens=cu_seqlens,
646
+ chunk_indices=chunk_indices,
647
+ chunk_size=chunk_size,
648
+ )
649
+ del Aqk, v_new, h
650
+ return o, final_state
651
+
652
+
653
+def chunk_kda_fwd(
654
+ q: torch.Tensor,
655
+ k: torch.Tensor,
656
+ v: torch.Tensor,
657
+ g: torch.Tensor,
658
+ beta: torch.Tensor,
659
+ scale: float,
660
+ initial_state: torch.Tensor,
661
+ output_final_state: bool,
662
+ cu_seqlens: torch.Tensor | None = None,
663
+):
664
+ chunk_size = FLA_CHUNK_SIZE
665
+ chunk_indices = (
666
+ prepare_chunk_indices(cu_seqlens, chunk_size)
667
+ if cu_seqlens is not None
668
+ else None
669
+ )
670
+ g = chunk_local_cumsum(
671
+ g,
672
+ chunk_size=chunk_size,
673
+ cu_seqlens=cu_seqlens,
674
+ chunk_indices=chunk_indices,
675
+ )
676
+ # KDA evaluates cumulative gate decays with exp2. Convert from natural-log
677
+ # space so exp(x) is preserved as exp2(x / ln(2)).
678
+ g = g * RCP_LN2
679
+ return _chunk_kda_fwd_with_cumulative_g(
680
+ q=q,
681
+ k=k,
682
+ v=v,
683
+ g=g,
684
+ beta=beta,
685
+ scale=scale,
686
+ initial_state=initial_state,
687
+ output_final_state=output_final_state,
688
+ cu_seqlens=cu_seqlens,
689
+ chunk_indices=chunk_indices,
690
+ chunk_size=chunk_size,
691
+ )
692
+
693
+
694
+def chunk_kda_with_fused_gate_fwd(
695
+ q: torch.Tensor,
696
+ k: torch.Tensor,
697
+ v: torch.Tensor,
698
+ raw_g: torch.Tensor,
699
+ raw_beta: torch.Tensor,
700
+ A_log: torch.Tensor,
701
+ g_bias: torch.Tensor | None,
702
+ scale: float,
703
+ initial_state: torch.Tensor,
704
+ output_final_state: bool,
705
+ lower_bound: float | None = None,
706
+ cu_seqlens: torch.Tensor | None = None,
707
+):
708
+ chunk_size = FLA_CHUNK_SIZE
709
+ chunk_indices = (
710
+ prepare_chunk_indices(cu_seqlens, chunk_size)
711
+ if cu_seqlens is not None
712
+ else None
713
+ )
714
+ g, beta = fused_kda_gate_chunk_cumsum(
715
+ raw_g,
716
+ raw_beta=raw_beta,
717
+ A_log=A_log,
718
+ g_bias=g_bias,
719
+ cu_seqlens=cu_seqlens,
720
+ chunk_indices=chunk_indices,
721
+ chunk_size=chunk_size,
722
+ lower_bound=lower_bound,
723
+ )
724
+ return _chunk_kda_fwd_with_cumulative_g(
725
+ q=q,
726
+ k=k,
727
+ v=v,
728
+ g=g,
729
+ beta=beta,
730
+ scale=scale,
731
+ initial_state=initial_state,
732
+ output_final_state=output_final_state,
733
+ cu_seqlens=cu_seqlens,
734
+ chunk_indices=chunk_indices,
735
+ chunk_size=chunk_size,
736
+ safe_gate=lower_bound is not None,
737
+ )
738
+
739
+
740
+def chunk_kda(
741
+ q: torch.Tensor,
742
+ k: torch.Tensor,
743
+ v: torch.Tensor,
744
+ g: torch.Tensor,
745
+ beta: torch.Tensor,
746
+ scale: float = None,
747
+ initial_state: torch.Tensor = None,
748
+ output_final_state: bool = False,
749
+ use_qk_l2norm_in_kernel: bool = False,
750
+ cu_seqlens: torch.Tensor | None = None,
751
+ **kwargs,
752
+):
753
+ if scale is None:
754
+ scale = k.shape[-1] ** -0.5
755
+
756
+ if use_qk_l2norm_in_kernel:
757
+ q = l2norm_fwd(q.contiguous())
758
+ k = l2norm_fwd(k.contiguous())
759
+
760
+ o, final_state = chunk_kda_fwd(
761
+ q=q,
762
+ k=k,
763
+ v=v.contiguous(),
764
+ g=g.contiguous(),
765
+ beta=beta.contiguous(),
766
+ scale=scale,
767
+ initial_state=initial_state.contiguous(),
768
+ output_final_state=output_final_state,
769
+ cu_seqlens=cu_seqlens,
770
+ )
771
+ return o, final_state
772
+
773
+
774
+def chunk_kda_with_fused_gate(
775
+ q: torch.Tensor,
776
+ k: torch.Tensor,
777
+ v: torch.Tensor,
778
+ raw_g: torch.Tensor,
779
+ raw_beta: torch.Tensor,
780
+ A_log: torch.Tensor,
781
+ g_bias: torch.Tensor | None,
782
+ scale: float | None = None,
783
+ initial_state: torch.Tensor | None = None,
784
+ output_final_state: bool = False,
785
+ lower_bound: float | None = None,
786
+ use_qk_l2norm_in_kernel: bool = False,
787
+ cu_seqlens: torch.Tensor | None = None,
788
+ **kwargs,
789
+):
790
+ """Run chunk KDA from raw gate and beta projections."""
791
+ if scale is None:
792
+ scale = k.shape[-1] ** -0.5
793
+
794
+ if use_qk_l2norm_in_kernel:
795
+ q = l2norm_fwd(q.contiguous())
796
+ k = l2norm_fwd(k.contiguous())
797
+
798
+ o, final_state = chunk_kda_with_fused_gate_fwd(
799
+ q=q,
800
+ k=k,
801
+ v=v.contiguous(),
802
+ raw_g=raw_g.contiguous(),
803
+ raw_beta=raw_beta,
804
+ A_log=A_log,
805
+ g_bias=g_bias,
806
+ scale=scale,
807
+ initial_state=initial_state.contiguous() if initial_state is not None else None,
808
+ output_final_state=output_final_state,
809
+ lower_bound=lower_bound,
810
+ cu_seqlens=cu_seqlens,
811
+ )
812
+ return o, final_state
813
+
814
+
815
+@triton.autotune(
816
+ configs=[
817
+ triton.Config({"BT": bt}, num_warps=nw, num_stages=ns)
818
+ for bt in BT_LIST_AUTOTUNE
819
+ for nw in NUM_WARPS_AUTOTUNE
820
+ for ns in [2, 3]
821
+ ],
822
+ key=["H", "D"],
823
+)
824
+@triton.jit
825
+def kda_gate_fwd_kernel(
826
+ g,
827
+ A,
828
+ y,
829
+ g_bias,
830
+ lower_bound,
831
+ beta: tl.constexpr,
832
+ threshold: tl.constexpr,
833
+ T,
834
+ H,
835
+ D: tl.constexpr,
836
+ BT: tl.constexpr,
837
+ BD: tl.constexpr,
838
+ HAS_BIAS: tl.constexpr,
839
+ USE_LOWER_BOUND: tl.constexpr,
840
+):
841
+ i_t, i_h = tl.program_id(0), tl.program_id(1)
842
+ n_t = i_t * BT
843
+
844
+ b_a = tl.exp(tl.load(A + i_h).to(tl.float32))
845
+
846
+ stride_row = H * D
847
+ stride_col = 1
848
+
849
+ g_ptr = tl.make_block_ptr(
850
+ base=g + i_h * D,
851
+ shape=(T, D),
852
+ strides=(stride_row, stride_col),
853
+ offsets=(n_t, 0),
854
+ block_shape=(BT, BD),
855
+ order=(1, 0),
856
+ )
857
+
858
+ y_ptr = tl.make_block_ptr(
859
+ base=y + i_h * D,
860
+ shape=(T, D),
861
+ strides=(stride_row, stride_col),
862
+ offsets=(n_t, 0),
863
+ block_shape=(BT, BD),
864
+ order=(1, 0),
865
+ )
866
+
867
+ b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32)
868
+
869
+ if HAS_BIAS:
870
+ n_d = tl.arange(0, BD)
871
+ bias_mask = n_d < D
872
+ b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to(
873
+ tl.float32
874
+ )
875
+ b_g = b_g + b_bias[None, :]
876
+
877
+ if USE_LOWER_BOUND:
878
+ b_y = lower_bound * tl.sigmoid(b_a * b_g)
879
+ else:
880
+ g_scaled = b_g * beta
881
+ use_linear = g_scaled > threshold
882
+ sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled)))
883
+ b_y = -b_a * sp
884
+
885
+ tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1))
886
+
887
+
888
+def fused_kda_gate(
889
+ g: torch.Tensor,
890
+ A: torch.Tensor,
891
+ head_k_dim: int,
892
+ g_bias: torch.Tensor | None = None,
893
+ beta: float = 1.0,
894
+ threshold: float = 20.0,
895
+ lower_bound: float | None = None,
896
+) -> torch.Tensor:
897
+ """
898
+ Forward pass for KDA gate:
899
+ input g: [..., H*D]
900
+ param A: [H] or [1, 1, H, 1]
901
+ beta: softplus beta parameter
902
+ threshold: softplus threshold parameter
903
+ return : [..., H, D]
904
+ """
905
+ orig_shape = g.shape[:-1]
906
+
907
+ g = g.view(-1, g.shape[-1])
908
+ T = g.shape[0]
909
+ HD = g.shape[1]
910
+ H = A.numel()
911
+ assert H * head_k_dim == HD
912
+
913
+ y = torch.empty_like(g, dtype=torch.float32)
914
+
915
+ def grid(meta):
916
+ return (cdiv(T, meta["BT"]), H)
917
+
918
+ kda_gate_fwd_kernel[grid](
919
+ g,
920
+ A,
921
+ y,
922
+ g_bias,
923
+ lower_bound or 0.0,
924
+ beta,
925
+ threshold,
926
+ T,
927
+ H,
928
+ head_k_dim,
929
+ BD=next_power_of_2(head_k_dim),
930
+ HAS_BIAS=g_bias is not None,
931
+ USE_LOWER_BOUND=lower_bound is not None,
932
+ )
933
+
934
+ y = y.view(*orig_shape, H, head_k_dim)
935
+ return y
vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py
new
+662
@@ -0,0 +1,662 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li
4
+#
5
+# This file contains code copied from the flash-linear-attention project.
6
+# The original source was licensed under the MIT license.
7
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
8
+# Forward-only adaptation of flash-linear-attention 0.5.0.
9
+# ruff: noqa: E501
10
+
11
+import torch
12
+
13
+from vllm.platforms import current_platform
14
+from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices
15
+from vllm.third_party.flash_linear_attention.ops.op import exp2, gather
16
+from vllm.third_party.flash_linear_attention.ops.utils import is_gather_supported
17
+from vllm.triton_utils import tl, triton
18
+
19
+from .chunk_intra_token_parallel import chunk_kda_fwd_intra_token_parallel
20
+
21
+################################################################################
22
+# Fused inter + solve_tril kernel: compute off-diagonal Akk and solve in one pass
23
+################################################################################
24
+
25
+
26
+@triton.heuristics(
27
+ {
28
+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
29
+ }
30
+)
31
+@triton.autotune(
32
+ configs=[
33
+ triton.Config({"BK": BK}, num_warps=num_warps)
34
+ for BK in [32, 64]
35
+ for num_warps in [1, 2, 4]
36
+ ],
37
+ key=["H", "HV", "K", "BC"],
38
+)
39
+@triton.jit(do_not_specialize=["T"])
40
+def chunk_kda_fwd_kernel_inter_solve_fused(
41
+ q,
42
+ k,
43
+ g,
44
+ beta,
45
+ Aqk,
46
+ Akkd,
47
+ Akk,
48
+ scale,
49
+ cu_seqlens,
50
+ chunk_indices,
51
+ T,
52
+ H: tl.constexpr,
53
+ HV: tl.constexpr,
54
+ K: tl.constexpr,
55
+ BT: tl.constexpr,
56
+ BC: tl.constexpr,
57
+ BK: tl.constexpr,
58
+ IS_VARLEN: tl.constexpr,
59
+ USE_SAFE_GATE: tl.constexpr,
60
+ SOLVE_TRIL_DOT_PRECISION: tl.constexpr,
61
+):
62
+ """
63
+ Fused kernel: compute inter-subchunk Akk + solve_tril in one pass.
64
+ Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd.
65
+
66
+ This kernel:
67
+ 1. Computes off-diagonal Aqk blocks -> writes to global
68
+ 2. Computes off-diagonal Akk blocks -> keeps in registers
69
+ 3. Loads diagonal Akk blocks from Akkd (fp32)
70
+ 4. Does forward substitution on diagonals
71
+ 5. Computes merged Akk_inv
72
+ 6. Writes Akk_inv to Akk
73
+ """
74
+ i_t, i_bh = tl.program_id(0), tl.program_id(1)
75
+ i_b, i_hv = i_bh // HV, i_bh % HV
76
+ i_h = i_hv // (HV // H)
77
+
78
+ if IS_VARLEN:
79
+ i_n, i_t = (
80
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
81
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
82
+ )
83
+ bos, eos = (
84
+ tl.load(cu_seqlens + i_n).to(tl.int32),
85
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
86
+ )
87
+ T = eos - bos
88
+ else:
89
+ bos, eos = i_b * T, i_b * T + T
90
+
91
+ if i_t * BT >= T:
92
+ return
93
+
94
+ i_tc0 = i_t * BT
95
+ i_tc1 = i_t * BT + BC
96
+ i_tc2 = i_t * BT + 2 * BC
97
+ i_tc3 = i_t * BT + 3 * BC
98
+
99
+ q += (bos * H + i_h) * K
100
+ k += (bos * H + i_h) * K
101
+ g += (bos * HV + i_hv) * K
102
+ Aqk += (bos * HV + i_hv) * BT
103
+ Akk += (bos * HV + i_hv) * BT
104
+ Akkd += (bos * HV + i_hv) * BC
105
+
106
+ o_i = tl.arange(0, BC)
107
+ m_tc1 = (i_tc1 + o_i) < T
108
+ m_tc2 = (i_tc2 + o_i) < T
109
+ m_tc3 = (i_tc3 + o_i) < T
110
+
111
+ b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32)
112
+ b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32)
113
+
114
+ b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32)
115
+ b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32)
116
+ b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32)
117
+ b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32)
118
+
119
+ b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32)
120
+ b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32)
121
+ b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32)
122
+ b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32)
123
+ b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32)
124
+ b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32)
125
+
126
+ ################################################################################
127
+ # off-diagonal blocks
128
+ ################################################################################
129
+ for i_k in range(tl.cdiv(K, BK)):
130
+ o_k = i_k * BK + tl.arange(0, BK)
131
+ m_k = o_k < K
132
+
133
+ p_k0 = tl.make_block_ptr(
134
+ k, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)
135
+ )
136
+ p_g0 = tl.make_block_ptr(
137
+ g, (T, K), (HV * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)
138
+ )
139
+ b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32)
140
+ b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32)
141
+
142
+ if i_tc1 < T:
143
+ p_q1 = tl.make_block_ptr(
144
+ q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)
145
+ )
146
+ p_k1 = tl.make_block_ptr(
147
+ k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)
148
+ )
149
+ p_g1 = tl.make_block_ptr(
150
+ g, (T, K), (HV * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)
151
+ )
152
+ # [BC, BK]
153
+ b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32)
154
+ b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32)
155
+ b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32)
156
+ # [BK]
157
+ b_gn1 = tl.load(g + i_tc1 * HV * K + o_k, mask=m_k, other=0).to(tl.float32)
158
+ # [BC, BK]
159
+ b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0)
160
+ # [BK, BC]
161
+ b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0))
162
+ # [BC, BC]
163
+ b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt)
164
+ b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt)
165
+
166
+ if i_tc2 < T:
167
+ p_q2 = tl.make_block_ptr(
168
+ q, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)
169
+ )
170
+ p_k2 = tl.make_block_ptr(
171
+ k, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)
172
+ )
173
+ p_g2 = tl.make_block_ptr(
174
+ g, (T, K), (HV * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)
175
+ )
176
+ # [BC, BK]
177
+ b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32)
178
+ b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32)
179
+ b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32)
180
+ # [BK]
181
+ b_gn2 = tl.load(g + i_tc2 * HV * K + o_k, mask=m_k, other=0).to(
182
+ tl.float32
183
+ )
184
+ # [BC, BK]
185
+ b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0)
186
+ b_qg2 = b_q2 * b_gqn2
187
+ b_kg2 = b_k2 * b_gqn2
188
+ # [BK, BC]
189
+ b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0))
190
+ b_Aqk20 += tl.dot(b_qg2, b_kgt)
191
+ b_Akk20 += tl.dot(b_kg2, b_kgt)
192
+ # [BC, BC]
193
+ b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1))
194
+ # [BC, BC]
195
+ b_Aqk21 += tl.dot(b_qg2, b_kgt)
196
+ b_Akk21 += tl.dot(b_kg2, b_kgt)
197
+
198
+ if i_tc3 < T:
199
+ p_q3 = tl.make_block_ptr(
200
+ q, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)
201
+ )
202
+ p_k3 = tl.make_block_ptr(
203
+ k, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)
204
+ )
205
+ p_g3 = tl.make_block_ptr(
206
+ g, (T, K), (HV * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)
207
+ )
208
+ # [BC, BK]
209
+ b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32)
210
+ b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32)
211
+ b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32)
212
+ # [BK]
213
+ b_gn3 = tl.load(g + i_tc3 * HV * K + o_k, mask=m_k, other=0).to(
214
+ tl.float32
215
+ )
216
+ # [BC, BK]
217
+ b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0)
218
+ b_qg3 = b_q3 * b_gqn3
219
+ b_kg3 = b_k3 * b_gqn3
220
+ # [BK, BC]
221
+ b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0))
222
+ # [BC, BC]
223
+ b_Aqk30 += tl.dot(b_qg3, b_kgt)
224
+ b_Akk30 += tl.dot(b_kg3, b_kgt)
225
+ # [BK, BC]
226
+ b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1))
227
+ # [BC, BC]
228
+ b_Aqk31 += tl.dot(b_qg3, b_kgt)
229
+ b_Akk31 += tl.dot(b_kg3, b_kgt)
230
+ # [BK, BC]
231
+ b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2))
232
+ # [BC, BC]
233
+ b_Aqk32 += tl.dot(b_qg3, b_kgt)
234
+ b_Akk32 += tl.dot(b_kg3, b_kgt)
235
+
236
+ ################################################################################
237
+ # save off-diagonal Aqk blocks and prepare Akk
238
+ ################################################################################
239
+ if i_tc1 < T:
240
+ p_Aqk10 = tl.make_block_ptr(
241
+ Aqk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)
242
+ )
243
+ tl.store(
244
+ p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
245
+ )
246
+
247
+ p_b1 = tl.make_block_ptr(
248
+ beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,)
249
+ )
250
+ b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32)
251
+ b_Akk10 = b_Akk10 * b_b1[:, None]
252
+ if i_tc2 < T:
253
+ p_Aqk20 = tl.make_block_ptr(
254
+ Aqk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)
255
+ )
256
+ p_Aqk21 = tl.make_block_ptr(
257
+ Aqk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)
258
+ )
259
+ tl.store(
260
+ p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
261
+ )
262
+ tl.store(
263
+ p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
264
+ )
265
+
266
+ p_b2 = tl.make_block_ptr(
267
+ beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,)
268
+ )
269
+ b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32)
270
+ b_Akk20 = b_Akk20 * b_b2[:, None]
271
+ b_Akk21 = b_Akk21 * b_b2[:, None]
272
+ if i_tc3 < T:
273
+ p_Aqk30 = tl.make_block_ptr(
274
+ Aqk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)
275
+ )
276
+ p_Aqk31 = tl.make_block_ptr(
277
+ Aqk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)
278
+ )
279
+ p_Aqk32 = tl.make_block_ptr(
280
+ Aqk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)
281
+ )
282
+ tl.store(
283
+ p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
284
+ )
285
+ tl.store(
286
+ p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
287
+ )
288
+ tl.store(
289
+ p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
290
+ )
291
+
292
+ p_b3 = tl.make_block_ptr(
293
+ beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,)
294
+ )
295
+ b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32)
296
+ b_Akk30 = b_Akk30 * b_b3[:, None]
297
+ b_Akk31 = b_Akk31 * b_b3[:, None]
298
+ b_Akk32 = b_Akk32 * b_b3[:, None]
299
+
300
+ p_Akk00 = tl.make_block_ptr(
301
+ Akkd, (T, BC), (HV * BC, 1), (i_tc0, 0), (BC, BC), (1, 0)
302
+ )
303
+ p_Akk11 = tl.make_block_ptr(
304
+ Akkd, (T, BC), (HV * BC, 1), (i_tc1, 0), (BC, BC), (1, 0)
305
+ )
306
+ p_Akk22 = tl.make_block_ptr(
307
+ Akkd, (T, BC), (HV * BC, 1), (i_tc2, 0), (BC, BC), (1, 0)
308
+ )
309
+ p_Akk33 = tl.make_block_ptr(
310
+ Akkd, (T, BC), (HV * BC, 1), (i_tc3, 0), (BC, BC), (1, 0)
311
+ )
312
+ b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32)
313
+ b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32)
314
+ b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32)
315
+ b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32)
316
+
317
+ ################################################################################
318
+ # forward substitution on diagonals
319
+ ################################################################################
320
+
321
+ if not USE_SAFE_GATE:
322
+ m_A = o_i[:, None] > o_i[None, :]
323
+ m_I = o_i[:, None] == o_i[None, :]
324
+
325
+ b_Ai00 = -tl.where(m_A, b_Ai00, 0)
326
+ b_Ai11 = -tl.where(m_A, b_Ai11, 0)
327
+ b_Ai22 = -tl.where(m_A, b_Ai22, 0)
328
+ b_Ai33 = -tl.where(m_A, b_Ai33, 0)
329
+
330
+ for i in range(2, min(BC, T - i_tc0)):
331
+ b_a00 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i)
332
+ b_a00 = tl.where(o_i < i, b_a00, 0.0)
333
+ b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0)
334
+ b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00)
335
+ for i in range(BC + 2, min(2 * BC, T - i_tc0)):
336
+ b_a11 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i)
337
+ b_a11 = tl.where(o_i < i - BC, b_a11, 0.0)
338
+ b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0)
339
+ b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11)
340
+ for i in range(2 * BC + 2, min(3 * BC, T - i_tc0)):
341
+ b_a22 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i)
342
+ b_a22 = tl.where(o_i < i - 2 * BC, b_a22, 0.0)
343
+ b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0)
344
+ b_Ai22 = tl.where((o_i == i - 2 * BC)[:, None], b_a22, b_Ai22)
345
+ for i in range(3 * BC + 2, min(4 * BC, T - i_tc0)):
346
+ b_a33 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i)
347
+ b_a33 = tl.where(o_i < i - 3 * BC, b_a33, 0.0)
348
+ b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0)
349
+ b_Ai33 = tl.where((o_i == i - 3 * BC)[:, None], b_a33, b_Ai33)
350
+
351
+ b_Ai00 += m_I
352
+ b_Ai11 += m_I
353
+ b_Ai22 += m_I
354
+ b_Ai33 += m_I
355
+
356
+ ################################################################################
357
+ # compute merged inverse using off-diagonals
358
+ ################################################################################
359
+
360
+ # we used tf32 to maintain matrix inverse's precision whenever possible.
361
+ b_Ai10 = -tl.dot(
362
+ tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION),
363
+ b_Ai00,
364
+ input_precision=SOLVE_TRIL_DOT_PRECISION,
365
+ )
366
+ b_Ai21 = -tl.dot(
367
+ tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION),
368
+ b_Ai11,
369
+ input_precision=SOLVE_TRIL_DOT_PRECISION,
370
+ )
371
+ b_Ai32 = -tl.dot(
372
+ tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION),
373
+ b_Ai22,
374
+ input_precision=SOLVE_TRIL_DOT_PRECISION,
375
+ )
376
+
377
+ b_Ai20 = -tl.dot(
378
+ b_Ai22,
379
+ tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION)
380
+ + tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION),
381
+ input_precision=SOLVE_TRIL_DOT_PRECISION,
382
+ )
383
+ b_Ai31 = -tl.dot(
384
+ b_Ai33,
385
+ tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION)
386
+ + tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION),
387
+ input_precision=SOLVE_TRIL_DOT_PRECISION,
388
+ )
389
+ b_Ai30 = -tl.dot(
390
+ b_Ai33,
391
+ tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION)
392
+ + tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION)
393
+ + tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION),
394
+ input_precision=SOLVE_TRIL_DOT_PRECISION,
395
+ )
396
+
397
+ ################################################################################
398
+ # store full Akk_inv to Akk
399
+ ################################################################################
400
+
401
+ p_Akk00 = tl.make_block_ptr(
402
+ Akk, (T, BT), (HV * BT, 1), (i_tc0, 0), (BC, BC), (1, 0)
403
+ )
404
+ p_Akk10 = tl.make_block_ptr(
405
+ Akk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)
406
+ )
407
+ p_Akk11 = tl.make_block_ptr(
408
+ Akk, (T, BT), (HV * BT, 1), (i_tc1, BC), (BC, BC), (1, 0)
409
+ )
410
+ p_Akk20 = tl.make_block_ptr(
411
+ Akk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)
412
+ )
413
+ p_Akk21 = tl.make_block_ptr(
414
+ Akk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)
415
+ )
416
+ p_Akk22 = tl.make_block_ptr(
417
+ Akk, (T, BT), (HV * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0)
418
+ )
419
+ p_Akk30 = tl.make_block_ptr(
420
+ Akk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)
421
+ )
422
+ p_Akk31 = tl.make_block_ptr(
423
+ Akk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)
424
+ )
425
+ p_Akk32 = tl.make_block_ptr(
426
+ Akk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)
427
+ )
428
+ p_Akk33 = tl.make_block_ptr(
429
+ Akk, (T, BT), (HV * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0)
430
+ )
431
+
432
+ tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1))
433
+ tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1))
434
+ tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1))
435
+ tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1))
436
+ tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1))
437
+ tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1))
438
+ tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1))
439
+ tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1))
440
+ tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1))
441
+ tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1))
442
+
443
+
444
+@triton.heuristics(
445
+ {
446
+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
447
+ }
448
+)
449
+@triton.autotune(
450
+ configs=[
451
+ triton.Config({}, num_warps=num_warps, num_stages=num_stages)
452
+ for num_warps in [1, 2, 4, 8]
453
+ for num_stages in [2, 3, 4]
454
+ ],
455
+ key=["BK", "NC", "BT", "HV"],
456
+)
457
+@triton.jit(do_not_specialize=["B", "T"])
458
+def chunk_kda_fwd_kernel_intra_sub_chunk(
459
+ q,
460
+ k,
461
+ g,
462
+ beta,
463
+ Aqk,
464
+ Akk,
465
+ scale,
466
+ cu_seqlens,
467
+ chunk_indices,
468
+ T,
469
+ H: tl.constexpr,
470
+ HV: tl.constexpr,
471
+ K: tl.constexpr,
472
+ BT: tl.constexpr,
473
+ BC: tl.constexpr,
474
+ BK: tl.constexpr,
475
+ IS_VARLEN: tl.constexpr,
476
+ USE_GATHER: tl.constexpr,
477
+):
478
+ i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
479
+ i_b, i_hv = i_bh // HV, i_bh % HV
480
+ i_h = i_hv // (HV // H)
481
+
482
+ if IS_VARLEN:
483
+ i_n, i_t = (
484
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
485
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
486
+ )
487
+ bos, eos = (
488
+ tl.load(cu_seqlens + i_n).to(tl.int32),
489
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
490
+ )
491
+ T = eos - bos
492
+ else:
493
+ bos, eos = i_b * T, i_b * T + T
494
+
495
+ i_ti = i_t * BT + i_i * BC
496
+ if i_ti >= T:
497
+ return
498
+
499
+ o_c = i_ti + tl.arange(0, BC)
500
+ m_c = o_c < T
501
+
502
+ q = q + (bos * H + i_h) * K
503
+ k = k + (bos * H + i_h) * K
504
+ g = g + (bos * HV + i_hv) * K
505
+ beta = beta + bos * HV + i_hv
506
+ Aqk = Aqk + (bos * HV + i_hv) * BT
507
+ Akk = Akk + (bos * HV + i_hv) * BC
508
+
509
+ p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0))
510
+ p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0))
511
+ p_g = tl.make_block_ptr(g, (T, K), (HV * K, 1), (i_ti, 0), (BC, BK), (1, 0))
512
+
513
+ p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,))
514
+
515
+ b_q = tl.load(p_q, boundary_check=(0, 1))
516
+ b_k = tl.load(p_k, boundary_check=(0, 1))
517
+ b_g = tl.load(p_g, boundary_check=(0, 1))
518
+ b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32)
519
+
520
+ if USE_GATHER:
521
+ b_gn = gather(
522
+ b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0
523
+ )
524
+ else:
525
+ # caculate offset
526
+ p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV * K + tl.arange(0, BK)
527
+ b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0)
528
+ b_gn = b_gn[None, :]
529
+
530
+ # current block, keep numerical stability by subtracting the left boundary
531
+ # less than 85 to avoid overflow in exp2
532
+ b_gm = (b_g - b_gn).to(tl.float32)
533
+
534
+ b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.0)
535
+ b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.0)
536
+
537
+ b_kgt = tl.trans(b_k * b_gk)
538
+
539
+ b_Aqk = tl.dot(b_q * b_gq, b_kgt) * scale
540
+ b_Akk = tl.dot(b_k * b_gq, b_kgt) * b_beta[:, None]
541
+
542
+ o_i = tl.arange(0, BC)
543
+ m_Aqk = o_i[:, None] >= o_i[None, :]
544
+ m_Akk = o_i[:, None] > o_i[None, :]
545
+ m_I = o_i[:, None] == o_i[None, :]
546
+
547
+ b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0)
548
+ b_Akk = tl.where(m_Akk, b_Akk, 0.0)
549
+
550
+ p_Aqk = tl.make_block_ptr(
551
+ Aqk, (T, BT), (HV * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)
552
+ )
553
+ p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV * BC, 1), (i_ti, 0), (BC, BC), (1, 0))
554
+ tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1))
555
+ tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1))
556
+
557
+ tl.debug_barrier()
558
+
559
+ ################################################################################
560
+ # forward substitution
561
+ ################################################################################
562
+
563
+ b_Ai = -b_Akk
564
+ for i in range(2, min(BC, T - i_ti)):
565
+ b_a = -tl.load(Akk + (i_ti + i) * HV * BC + o_i)
566
+ b_a = tl.where(o_i < i, b_a, 0.0)
567
+ b_a += tl.sum(b_a[:, None] * b_Ai, 0)
568
+ b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai)
569
+ b_Ai += m_I
570
+ tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1))
571
+
572
+
573
+def chunk_kda_fwd_intra(
574
+ q: torch.Tensor,
575
+ k: torch.Tensor,
576
+ gk: torch.Tensor | None = None,
577
+ beta: torch.Tensor | None = None,
578
+ scale: float | None = None,
579
+ cu_seqlens: torch.LongTensor | None = None,
580
+ chunk_size: int = 64,
581
+ chunk_indices: torch.LongTensor | None = None,
582
+ safe_gate: bool = False,
583
+):
584
+ B, T, H, K, HV = *k.shape, gk.shape[2]
585
+ BT = chunk_size
586
+ BC = 16
587
+ if chunk_indices is None and cu_seqlens is not None:
588
+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
589
+ NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
590
+ NC = triton.cdiv(BT, BC)
591
+
592
+ Aqk = torch.empty(B, T, HV, BT, device=k.device, dtype=k.dtype)
593
+ # Akk must be zero-initialized - kernel only writes lower triangular
594
+ Akk = torch.zeros(B, T, HV, BT, device=k.device, dtype=k.dtype)
595
+ # Separate fp32 buffer for diagonal 16x16 blocks (for precision in solve_tril)
596
+ Akkd = torch.empty(B, T, HV, BC, device=k.device, dtype=torch.float32)
597
+
598
+ # Compute diagonal blocks into Akkd in fp32.
599
+ if safe_gate:
600
+ grid = (NT, NC, B * HV)
601
+ BK = triton.next_power_of_2(K)
602
+ chunk_kda_fwd_kernel_intra_sub_chunk[grid](
603
+ q=q,
604
+ k=k,
605
+ g=gk,
606
+ beta=beta,
607
+ Aqk=Aqk,
608
+ Akk=Akkd,
609
+ scale=scale,
610
+ cu_seqlens=cu_seqlens,
611
+ chunk_indices=chunk_indices,
612
+ T=T,
613
+ H=H,
614
+ HV=HV,
615
+ K=K,
616
+ BT=BT,
617
+ BC=BC,
618
+ BK=BK,
619
+ USE_GATHER=is_gather_supported,
620
+ )
621
+ else:
622
+ Aqk, Akkd = chunk_kda_fwd_intra_token_parallel(
623
+ q=q,
624
+ k=k,
625
+ gk=gk,
626
+ beta=beta,
627
+ Aqk=Aqk,
628
+ Akk=Akkd,
629
+ scale=scale,
630
+ cu_seqlens=cu_seqlens,
631
+ chunk_size=BT,
632
+ sub_chunk_size=BC,
633
+ )
634
+
635
+ # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen)
636
+ solve_tril_dot_precision = (
637
+ "tf32"
638
+ if current_platform.is_cuda() and current_platform.has_device_capability(80)
639
+ else "ieee"
640
+ )
641
+ grid = (NT, B * HV)
642
+ chunk_kda_fwd_kernel_inter_solve_fused[grid](
643
+ q=q,
644
+ k=k,
645
+ g=gk,
646
+ beta=beta,
647
+ Aqk=Aqk,
648
+ Akkd=Akkd,
649
+ Akk=Akk,
650
+ scale=scale,
651
+ cu_seqlens=cu_seqlens,
652
+ chunk_indices=chunk_indices,
653
+ T=T,
654
+ H=H,
655
+ HV=HV,
656
+ K=K,
657
+ BT=BT,
658
+ BC=BC,
659
+ USE_SAFE_GATE=safe_gate,
660
+ SOLVE_TRIL_DOT_PRECISION=solve_tril_dot_precision,
661
+ )
662
+ return Aqk, Akk
vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py
new
+197
@@ -0,0 +1,197 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li
4
+#
5
+# This file contains code copied from the flash-linear-attention project.
6
+# The original source was licensed under the MIT license.
7
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
8
+# Forward-only adaptation of flash-linear-attention 0.5.0.
9
+# ruff: noqa: E501
10
+
11
+# Token-parallel implementation of KDA intra chunk kernel
12
+
13
+import torch
14
+
15
+from vllm.third_party.flash_linear_attention.ops.op import exp2
16
+from vllm.triton_utils import tl, triton
17
+
18
+
19
+@triton.heuristics(
20
+ {
21
+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
22
+ }
23
+)
24
+@triton.autotune(
25
+ configs=[
26
+ triton.Config({"BH": BH}, num_warps=num_warps)
27
+ for BH in [1, 2, 4, 8]
28
+ for num_warps in [1, 2, 4, 8]
29
+ ],
30
+ key=["K", "H", "HV"],
31
+)
32
+@triton.jit(do_not_specialize=["T", "N"])
33
+def chunk_kda_fwd_kernel_intra_token_parallel(
34
+ q,
35
+ k,
36
+ g,
37
+ beta,
38
+ Aqk,
39
+ Akk,
40
+ scale,
41
+ cu_seqlens,
42
+ N,
43
+ T,
44
+ H: tl.constexpr,
45
+ HV: tl.constexpr,
46
+ K: tl.constexpr,
47
+ BT: tl.constexpr,
48
+ BC: tl.constexpr,
49
+ BH: tl.constexpr,
50
+ IS_VARLEN: tl.constexpr,
51
+):
52
+ i_tg, i_hg = tl.program_id(0), tl.program_id(1)
53
+
54
+ if IS_VARLEN:
55
+ i_n = 0
56
+ left, right = 0, N
57
+
58
+ # Unrolled binary search (max B=2^32)
59
+ # We can limit iterations based on expected max batch size if needed
60
+ # 20 iterations covers B=1M, usually enough
61
+ for _ in range(20):
62
+ if left < right:
63
+ mid = (left + right) // 2
64
+ if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int32):
65
+ right = mid
66
+ else:
67
+ left = mid + 1
68
+ i_n = left
69
+
70
+ bos, eos = (
71
+ tl.load(cu_seqlens + i_n).to(tl.int32),
72
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
73
+ )
74
+ T = eos - bos
75
+ i_t = i_tg - bos
76
+ else:
77
+ bos = (i_tg // T) * T
78
+ i_t = i_tg % T
79
+
80
+ if i_t >= T:
81
+ return
82
+
83
+ i_c = i_t // BT
84
+ i_s = (i_t % BT) // BC
85
+ i_tc = i_c * BT
86
+ i_ts = i_tc + i_s * BC
87
+
88
+ G: tl.constexpr = HV // H
89
+
90
+ q += bos * H * K
91
+ k += bos * H * K
92
+ g += bos * HV * K
93
+ Aqk += bos * HV * BT
94
+ Akk += bos * HV * BC
95
+ beta += bos * HV
96
+
97
+ BK: tl.constexpr = triton.next_power_of_2(K)
98
+ o_hv = i_hg * BH + tl.arange(0, BH)
99
+ o_h = o_hv // G
100
+ o_k = tl.arange(0, BK)
101
+ m_hv = o_hv < HV
102
+ m_k = o_k < K
103
+ m_hk = m_hv[:, None] & m_k[None, :]
104
+
105
+ # q/k: [B, T, H, K], manual load via mapped qk head index
106
+ p_qk = o_h[:, None] * K + o_k[None, :]
107
+ b_q = tl.load(q + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
108
+ b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
109
+
110
+ # g: [B, T, HV, K], beta: [B, T, HV]
111
+ p_g = tl.make_block_ptr(
112
+ g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)
113
+ )
114
+ p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,))
115
+ b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)
116
+ b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32)
117
+ b_k *= b_beta[:, None]
118
+
119
+ for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))):
120
+ b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
121
+ p_gj = tl.make_block_ptr(
122
+ g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)
123
+ )
124
+ b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32)
125
+
126
+ b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0)
127
+ b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale
128
+ b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0)
129
+
130
+ tl.store(
131
+ Aqk + i_t * HV * BT + o_hv * BT + j % BT,
132
+ b_Aqk.to(Aqk.dtype.element_ty),
133
+ mask=m_hv,
134
+ )
135
+ tl.store(
136
+ Akk + i_t * HV * BC + o_hv * BC + j - i_ts,
137
+ b_Akk.to(Akk.dtype.element_ty),
138
+ mask=m_hv,
139
+ )
140
+
141
+
142
+def chunk_kda_fwd_intra_token_parallel(
143
+ q: torch.Tensor,
144
+ k: torch.Tensor,
145
+ gk: torch.Tensor,
146
+ beta: torch.Tensor,
147
+ Aqk: torch.Tensor,
148
+ Akk: torch.Tensor,
149
+ scale: float,
150
+ cu_seqlens: torch.LongTensor | None = None,
151
+ chunk_size: int = 64,
152
+ sub_chunk_size: int = 16,
153
+) -> None:
154
+ """
155
+ Token-parallel implementation: each token gets its own thread block.
156
+ Supports both fixed-length and variable-length sequences.
157
+ Reduces wasted computation on padding.
158
+
159
+ Writes directly to Aqk and Akk tensors (in-place).
160
+
161
+ Args:
162
+ q: [B, T, H, K]
163
+ k: [B, T, H, K]
164
+ gk: [B, T, HV, K] cumsum of gates (HV >= H for GVA)
165
+ beta: [B, T, HV]
166
+ Aqk: [B, T, HV, BT] output tensor to write to
167
+ Akk: [B, T, HV, BC] output tensor for diagonal blocks (fp32)
168
+ scale: attention scale
169
+ chunk_size: BT (default 64)
170
+ sub_chunk_size: BC (default 16)
171
+ """
172
+ B, T, H, K, HV = *q.shape, gk.shape[2]
173
+ N = len(cu_seqlens) - 1 if cu_seqlens is not None else B
174
+ BT = chunk_size
175
+ BC = sub_chunk_size
176
+
177
+ def grid(meta):
178
+ return (B * T, triton.cdiv(HV, meta["BH"]))
179
+
180
+ chunk_kda_fwd_kernel_intra_token_parallel[grid](
181
+ q=q,
182
+ k=k,
183
+ g=gk,
184
+ beta=beta,
185
+ Aqk=Aqk,
186
+ Akk=Akk,
187
+ scale=scale,
188
+ cu_seqlens=cu_seqlens,
189
+ N=N,
190
+ T=T,
191
+ H=H,
192
+ HV=HV,
193
+ K=K,
194
+ BT=BT,
195
+ BC=BC,
196
+ )
197
+ return Aqk, Akk
vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py
new
+621
@@ -0,0 +1,621 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li
4
+#
5
+# This file contains code adapted from the flash-linear-attention project.
6
+# The original source was licensed under the MIT license.
7
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
8
+# ruff: noqa: E501
9
+
10
+import torch
11
+
12
+from vllm.third_party.flash_linear_attention.ops.op import exp, log
13
+from vllm.triton_utils import tl, triton
14
+from vllm.utils.math_utils import cdiv, next_power_of_2
15
+
16
+
17
+@triton.heuristics(
18
+ {
19
+ "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None,
20
+ "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None,
21
+ }
22
+)
23
+@triton.jit
24
+def _kda_gate_beta_fwd_kernel(
25
+ raw_g,
26
+ raw_beta,
27
+ A_log,
28
+ dt_bias,
29
+ gate,
30
+ beta_out,
31
+ lower_bound,
32
+ softplus_beta: tl.constexpr,
33
+ softplus_threshold: tl.constexpr,
34
+ T,
35
+ stride_g_token: tl.constexpr,
36
+ stride_beta_token: tl.constexpr,
37
+ H: tl.constexpr,
38
+ D: tl.constexpr,
39
+ BT: tl.constexpr,
40
+ BD: tl.constexpr,
41
+ HAS_DT_BIAS: tl.constexpr,
42
+ USE_LOWER_BOUND: tl.constexpr,
43
+):
44
+ i_t, i_h = tl.program_id(0), tl.program_id(1)
45
+ o_t = i_t * BT + tl.arange(0, BT)
46
+ o_d = tl.arange(0, BD)
47
+ m_t = o_t < T
48
+ m_d = o_d < D
49
+
50
+ p_g = raw_g + o_t[:, None] * stride_g_token + i_h * D + o_d[None, :]
51
+ b_g = tl.load(p_g, mask=m_t[:, None] & m_d[None, :], other=0.0).to(tl.float32)
52
+ if HAS_DT_BIAS:
53
+ b_bias = tl.load(
54
+ dt_bias + i_h * D + o_d,
55
+ mask=m_d,
56
+ other=0.0,
57
+ ).to(tl.float32)
58
+ b_g += b_bias[None, :]
59
+
60
+ b_a = exp(tl.load(A_log + i_h).to(tl.float32))
61
+ if USE_LOWER_BOUND:
62
+ b_gate = lower_bound * tl.sigmoid(b_a * b_g)
63
+ else:
64
+ b_scaled = b_g * softplus_beta
65
+ b_softplus = tl.where(
66
+ b_scaled > softplus_threshold,
67
+ b_g,
68
+ log(1.0 + tl.exp(b_scaled)) / softplus_beta,
69
+ )
70
+ b_gate = -b_a * b_softplus
71
+
72
+ p_gate = gate + (o_t[:, None] * H + i_h) * D + o_d[None, :]
73
+ tl.store(
74
+ p_gate,
75
+ b_gate,
76
+ mask=m_t[:, None] & m_d[None, :],
77
+ )
78
+
79
+ b_beta = tl.load(
80
+ raw_beta + o_t * stride_beta_token + i_h,
81
+ mask=m_t,
82
+ other=0.0,
83
+ ).to(tl.float32)
84
+ tl.store(beta_out + o_t * H + i_h, tl.sigmoid(b_beta), mask=m_t)
85
+
86
+
87
+def _fused_kda_gate_beta(
88
+ raw_g: torch.Tensor,
89
+ raw_beta: torch.Tensor,
90
+ A_log: torch.Tensor,
91
+ dt_bias: torch.Tensor | None,
92
+ lower_bound: float | None,
93
+) -> tuple[torch.Tensor, torch.Tensor]:
94
+ B, T, H, D = raw_g.shape
95
+ assert B == 1
96
+ assert raw_beta.shape == (B, T, H)
97
+ assert raw_g.stride()[2:] == (D, 1)
98
+ assert raw_beta.stride(2) == 1
99
+ gate = torch.empty((B, T, H, D), dtype=torch.float32, device=raw_g.device)
100
+ beta = torch.empty((B, T, H), dtype=torch.float32, device=raw_beta.device)
101
+
102
+ BT = 16
103
+ _kda_gate_beta_fwd_kernel[(cdiv(T, BT), H)](
104
+ raw_g=raw_g,
105
+ raw_beta=raw_beta,
106
+ A_log=A_log,
107
+ dt_bias=dt_bias,
108
+ gate=gate,
109
+ beta_out=beta,
110
+ lower_bound=lower_bound,
111
+ softplus_beta=1.0,
112
+ softplus_threshold=20.0,
113
+ T=T,
114
+ stride_g_token=raw_g.stride(1),
115
+ stride_beta_token=raw_beta.stride(1),
116
+ H=H,
117
+ D=D,
118
+ BT=BT,
119
+ BD=next_power_of_2(D),
120
+ num_warps=4,
121
+ )
122
+ return gate, beta
123
+
124
+
125
+@triton.heuristics(
126
+ {
127
+ "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None,
128
+ "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None,
129
+ "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None,
130
+ }
131
+)
132
+@triton.jit(do_not_specialize=["N", "T"])
133
+def fused_recurrent_kda_fwd_kernel(
134
+ q,
135
+ k,
136
+ v,
137
+ g,
138
+ beta,
139
+ A_log,
140
+ dt_bias,
141
+ out,
142
+ state,
143
+ cu_seqlens,
144
+ state_indices,
145
+ num_accepted_tokens,
146
+ lower_bound,
147
+ scale: tl.constexpr,
148
+ N: tl.int64,
149
+ T: tl.int64,
150
+ H: tl.constexpr,
151
+ K: tl.constexpr,
152
+ V: tl.constexpr,
153
+ BK: tl.constexpr,
154
+ BV: tl.constexpr,
155
+ stride_qkv_token: tl.constexpr,
156
+ stride_g_token: tl.constexpr,
157
+ stride_beta_token: tl.constexpr,
158
+ stride_out_token: tl.constexpr,
159
+ stride_state_token: tl.constexpr,
160
+ stride_indices_seq: tl.constexpr,
161
+ IS_SPEC_DECODING: tl.constexpr,
162
+ USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
163
+ USE_GATE_IN_KERNEL: tl.constexpr,
164
+ APPLY_BETA_SIGMOID: tl.constexpr,
165
+ HAS_DT_BIAS: tl.constexpr,
166
+ USE_LOWER_BOUND: tl.constexpr,
167
+ num_stages: tl.constexpr,
168
+):
169
+ pid = tl.program_id(0)
170
+ i_v = pid % tl.cdiv(V, BV)
171
+ i_nh = pid // tl.cdiv(V, BV)
172
+ i_n, i_h = i_nh // H, i_nh % H
173
+ bos = tl.load(cu_seqlens + i_n).to(tl.int64)
174
+ eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64)
175
+ sequence_length = eos - bos
176
+ if sequence_length == 0:
177
+ return
178
+
179
+ o_k = tl.arange(0, BK)
180
+ o_v = i_v * BV + tl.arange(0, BV)
181
+ m_k = o_k < K
182
+ m_v = o_v < V
183
+ m_state = m_v[:, None] & m_k[None, :]
184
+
185
+ if IS_SPEC_DECODING:
186
+ initial_token = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1
187
+ else:
188
+ initial_token = 0
189
+ state_index = tl.load(state_indices + i_n * stride_indices_seq + initial_token).to(
190
+ tl.int64
191
+ )
192
+ p_out = out + bos * stride_out_token + i_h * V + o_v
193
+ if state_index <= 0:
194
+ tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=m_v)
195
+ return
196
+
197
+ p_state = (
198
+ state
199
+ + state_index * stride_state_token
200
+ + i_h * V * K
201
+ + o_v[:, None] * K
202
+ + o_k[None, :]
203
+ )
204
+ b_state = tl.load(p_state, mask=m_state, other=0.0).to(tl.float32)
205
+
206
+ p_q = q + bos * stride_qkv_token + i_h * K + o_k
207
+ p_k = k + bos * stride_qkv_token + i_h * K + o_k
208
+ p_v = v + bos * stride_qkv_token + i_h * V + o_v
209
+ p_g = g + bos * stride_g_token + i_h * K + o_k
210
+ p_beta = beta + bos * stride_beta_token + i_h
211
+ for i_t in tl.range(0, sequence_length, num_stages=num_stages):
212
+ b_q = tl.load(p_q, mask=m_k, other=0.0, eviction_policy="evict_last").to(
213
+ tl.float32
214
+ )
215
+ b_k = tl.load(p_k, mask=m_k, other=0.0, eviction_policy="evict_last").to(
216
+ tl.float32
217
+ )
218
+ b_v = tl.load(p_v, mask=m_v, other=0.0, eviction_policy="evict_first").to(
219
+ tl.float32
220
+ )
221
+ if USE_QK_L2NORM_IN_KERNEL:
222
+ b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
223
+ b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
224
+ b_q *= scale
225
+
226
+ b_gate = tl.load(
227
+ p_g,
228
+ mask=m_k,
229
+ other=0.0,
230
+ eviction_policy="evict_last",
231
+ ).to(tl.float32)
232
+ if USE_GATE_IN_KERNEL:
233
+ if HAS_DT_BIAS:
234
+ b_bias = tl.load(
235
+ dt_bias + i_h * K + o_k,
236
+ mask=m_k,
237
+ other=0.0,
238
+ ).to(tl.float32)
239
+ b_gate += b_bias
240
+ b_a = exp(tl.load(A_log + i_h).to(tl.float32))
241
+ if USE_LOWER_BOUND:
242
+ b_gate = lower_bound * tl.sigmoid(b_a * b_gate)
243
+ else:
244
+ b_softplus = tl.where(
245
+ b_gate > 20.0,
246
+ b_gate,
247
+ log(1.0 + tl.exp(b_gate)),
248
+ )
249
+ b_gate = -b_a * b_softplus
250
+
251
+ b_state *= exp(b_gate[None, :])
252
+ b_v -= tl.sum(b_state * b_k[None, :], axis=1)
253
+ b_beta = tl.load(p_beta, eviction_policy="evict_last").to(tl.float32)
254
+ if APPLY_BETA_SIGMOID:
255
+ b_beta = tl.sigmoid(b_beta)
256
+ b_v *= b_beta
257
+ b_state += b_v[:, None] * b_k[None, :]
258
+ b_out = tl.sum(b_state * b_q[None, :], axis=1)
259
+ tl.store(
260
+ p_out,
261
+ b_out.to(p_out.dtype.element_ty),
262
+ mask=m_v,
263
+ eviction_policy="evict_first",
264
+ )
265
+
266
+ final_state_index = tl.load(state_indices + i_n * stride_indices_seq + i_t).to(
267
+ tl.int64
268
+ )
269
+ if final_state_index > 0:
270
+ p_final_state = (
271
+ state
272
+ + final_state_index * stride_state_token
273
+ + i_h * V * K
274
+ + o_v[:, None] * K
275
+ + o_k[None, :]
276
+ )
277
+ tl.store(
278
+ p_final_state,
279
+ b_state.to(p_final_state.dtype.element_ty),
280
+ mask=m_state,
281
+ )
282
+
283
+ p_q += stride_qkv_token
284
+ p_k += stride_qkv_token
285
+ p_v += stride_qkv_token
286
+ p_g += stride_g_token
287
+ p_beta += stride_beta_token
288
+ p_out += stride_out_token
289
+
290
+
291
+def fused_recurrent_kda_fwd(
292
+ q: torch.Tensor,
293
+ k: torch.Tensor,
294
+ v: torch.Tensor,
295
+ g: torch.Tensor,
296
+ beta: torch.Tensor,
297
+ scale: float | None = None,
298
+ initial_state: torch.Tensor | None = None,
299
+ inplace_final_state: bool = True,
300
+ cu_seqlens: torch.Tensor | None = None,
301
+ ssm_state_indices: torch.Tensor | None = None,
302
+ num_accepted_tokens: torch.Tensor | None = None,
303
+ use_qk_l2norm_in_kernel: bool = True,
304
+ A_log: torch.Tensor | None = None,
305
+ dt_bias: torch.Tensor | None = None,
306
+ lower_bound: float | None = None,
307
+ use_gate_in_kernel: bool = False,
308
+ use_beta_sigmoid_in_kernel: bool = False,
309
+ out: torch.Tensor | None = None,
310
+) -> tuple[torch.Tensor, torch.Tensor]:
311
+ """Launch recurrent KDA with dense inner dimensions and row strides."""
312
+ B, T, H, K = q.shape
313
+ V = v.shape[-1]
314
+ assert B == 1 and k.shape == q.shape
315
+ assert v.shape == (B, T, H, V) and g.shape == (B, T, H, K)
316
+ assert beta.shape == (B, T, H)
317
+ assert initial_state is not None
318
+ assert cu_seqlens is not None
319
+ assert ssm_state_indices is not None
320
+ assert inplace_final_state
321
+ if out is None:
322
+ out = torch.empty_like(v)
323
+ assert out.shape == v.shape
324
+ assert initial_state.shape[1:] == (H, V, K)
325
+ assert ssm_state_indices.ndim in (1, 2)
326
+
327
+ assert q.stride()[2:] == k.stride()[2:] == (K, 1)
328
+ assert v.stride()[2:] == out.stride()[2:] == (V, 1)
329
+ assert g.stride()[2:] == (K, 1)
330
+ assert beta.stride(2) == 1
331
+ assert q.stride(1) == k.stride(1) == v.stride(1)
332
+ assert initial_state.stride()[1:] == (V * K, K, 1)
333
+ N = cu_seqlens.numel() - 1
334
+ if ssm_state_indices.ndim == 1:
335
+ assert T == N
336
+ assert num_accepted_tokens is None
337
+ else:
338
+ assert ssm_state_indices.stride(1) == 1
339
+ assert cu_seqlens.is_contiguous()
340
+ if use_gate_in_kernel:
341
+ assert A_log is not None and A_log.is_contiguous()
342
+ assert dt_bias is None or dt_bias.is_contiguous()
343
+
344
+ if scale is None:
345
+ scale = K**-0.5
346
+
347
+ BV = 32 if use_gate_in_kernel else 8
348
+ num_warps = 4 if use_gate_in_kernel else 1
349
+ grid = (cdiv(V, BV) * N * H,)
350
+ fused_recurrent_kda_fwd_kernel[grid](
351
+ q=q,
352
+ k=k,
353
+ v=v,
354
+ g=g,
355
+ beta=beta,
356
+ A_log=A_log,
357
+ dt_bias=dt_bias,
358
+ out=out,
359
+ state=initial_state,
360
+ cu_seqlens=cu_seqlens,
361
+ state_indices=ssm_state_indices,
362
+ num_accepted_tokens=num_accepted_tokens,
363
+ lower_bound=lower_bound,
364
+ scale=scale,
365
+ N=N,
366
+ T=T,
367
+ H=H,
368
+ K=K,
369
+ V=V,
370
+ BK=next_power_of_2(K),
371
+ BV=BV,
372
+ stride_qkv_token=q.stride(1),
373
+ stride_g_token=g.stride(1),
374
+ stride_beta_token=beta.stride(1),
375
+ stride_out_token=out.stride(1),
376
+ stride_state_token=initial_state.stride(0),
377
+ stride_indices_seq=ssm_state_indices.stride(0),
378
+ IS_SPEC_DECODING=num_accepted_tokens is not None,
379
+ USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
380
+ USE_GATE_IN_KERNEL=use_gate_in_kernel,
381
+ APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel,
382
+ num_warps=num_warps,
383
+ num_stages=2,
384
+ )
385
+ return out, initial_state
386
+
387
+
388
+def fused_recurrent_kda(
389
+ q: torch.Tensor,
390
+ k: torch.Tensor,
391
+ v: torch.Tensor,
392
+ raw_g: torch.Tensor,
393
+ raw_beta: torch.Tensor,
394
+ A_log: torch.Tensor,
395
+ dt_bias: torch.Tensor | None,
396
+ lower_bound: float | None,
397
+ initial_state: torch.Tensor,
398
+ cu_seqlens: torch.Tensor,
399
+ ssm_state_indices: torch.Tensor,
400
+ num_accepted_tokens: torch.Tensor | None = None,
401
+ out: torch.Tensor | None = None,
402
+ fuse_gate: bool | None = None,
403
+) -> tuple[torch.Tensor, torch.Tensor]:
404
+ """Run recurrent KDA from raw gate and beta inputs.
405
+
406
+ This vLLM wrapper applies the gate activation and beta sigmoid, selecting
407
+ whether to materialize them before launching the recurrent kernel.
408
+ """
409
+ if fuse_gate is None:
410
+ # gfx950: always fuse the gate.
411
+ fuse_gate = True
412
+
413
+ if fuse_gate:
414
+ gate = raw_g
415
+ beta = raw_beta
416
+ else:
417
+ gate, beta = _fused_kda_gate_beta(
418
+ raw_g,
419
+ raw_beta,
420
+ A_log,
421
+ dt_bias,
422
+ lower_bound,
423
+ )
424
+ return fused_recurrent_kda_fwd(
425
+ q=q,
426
+ k=k,
427
+ v=v,
428
+ g=gate,
429
+ beta=beta,
430
+ scale=q.shape[-1] ** -0.5,
431
+ initial_state=initial_state,
432
+ inplace_final_state=True,
433
+ cu_seqlens=cu_seqlens,
434
+ ssm_state_indices=ssm_state_indices,
435
+ num_accepted_tokens=num_accepted_tokens,
436
+ use_qk_l2norm_in_kernel=True,
437
+ A_log=A_log if fuse_gate else None,
438
+ dt_bias=dt_bias if fuse_gate else None,
439
+ lower_bound=lower_bound if fuse_gate else None,
440
+ use_gate_in_kernel=fuse_gate,
441
+ use_beta_sigmoid_in_kernel=fuse_gate,
442
+ out=out,
443
+ )
444
+
445
+
446
+@triton.jit
447
+def fused_recurrent_kda_packed_decode_kernel(
448
+ mixed_qkv,
449
+ raw_g,
450
+ raw_beta,
451
+ A_log,
452
+ dt_bias,
453
+ out,
454
+ state,
455
+ state_indices,
456
+ lower_bound,
457
+ scale: tl.constexpr,
458
+ stride_mixed_token: tl.constexpr,
459
+ stride_g_token: tl.constexpr,
460
+ stride_beta_token: tl.constexpr,
461
+ stride_state_token: tl.constexpr,
462
+ H: tl.constexpr,
463
+ K: tl.constexpr,
464
+ V: tl.constexpr,
465
+ BK: tl.constexpr,
466
+ BV: tl.constexpr,
467
+ SOFTPLUS_THRESHOLD: tl.constexpr,
468
+ USE_LOWER_BOUND: tl.constexpr,
469
+):
470
+ i_v, i_nh = tl.program_id(0), tl.program_id(1)
471
+ i_n, i_h = i_nh // H, i_nh % H
472
+
473
+ o_k = tl.arange(0, BK)
474
+ o_v = i_v * BV + tl.arange(0, BV)
475
+ mask_k = o_k < K
476
+ mask_v = o_v < V
477
+ mask_state = mask_v[:, None] & mask_k[None, :]
478
+
479
+ state_idx = tl.load(state_indices + i_n).to(tl.int64)
480
+ p_out = out + (i_n * H + i_h) * V + o_v
481
+ if state_idx <= 0:
482
+ tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=mask_v)
483
+ return
484
+
485
+ p_state = state + state_idx * stride_state_token
486
+ p_state += i_h * V * K + o_v[:, None] * K + o_k[None, :]
487
+ b_state = tl.load(p_state, mask=mask_state, other=0).to(tl.float32)
488
+
489
+ # Q, K, and V occupy consecutive channel ranges, while the token stride
490
+ # may also include the output-gate projection that follows packed QKV.
491
+ p_mixed = mixed_qkv + i_n * stride_mixed_token
492
+ b_q = tl.load(p_mixed + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32)
493
+ b_k = tl.load(
494
+ p_mixed + H * K + i_h * K + o_k,
495
+ mask=mask_k,
496
+ other=0,
497
+ ).to(tl.float32)
498
+ b_v = tl.load(
499
+ p_mixed + 2 * H * K + i_h * V + o_v,
500
+ mask=mask_v,
501
+ other=0,
502
+ ).to(tl.float32)
503
+
504
+ b_q /= tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
505
+ b_k /= tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
506
+ b_q *= scale
507
+
508
+ p_g = raw_g + i_n * stride_g_token + i_h * K + o_k
509
+ b_g = tl.load(p_g, mask=mask_k, other=0).to(tl.float32)
510
+ b_bias = tl.load(dt_bias + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32)
511
+ b_a = exp(tl.load(A_log + i_h).to(tl.float32))
512
+ b_g += b_bias
513
+ if USE_LOWER_BOUND:
514
+ b_gate = lower_bound * tl.sigmoid(b_a * b_g)
515
+ else:
516
+ b_softplus = tl.where(
517
+ b_g > SOFTPLUS_THRESHOLD,
518
+ b_g,
519
+ log(1.0 + tl.exp(b_g)),
520
+ )
521
+ b_gate = -b_a * b_softplus
522
+
523
+ b_state *= exp(b_gate[None, :])
524
+ b_v -= tl.sum(b_state * b_k[None, :], axis=1)
525
+ b_beta = tl.sigmoid(
526
+ tl.load(raw_beta + i_n * stride_beta_token + i_h).to(tl.float32)
527
+ )
528
+ b_v *= b_beta
529
+ b_state += b_v[:, None] * b_k[None, :]
530
+ b_out = tl.sum(b_state * b_q[None, :], axis=1)
531
+
532
+ tl.store(p_out, b_out.to(p_out.dtype.element_ty), mask=mask_v)
533
+ tl.store(p_state, b_state.to(p_state.dtype.element_ty), mask=mask_state)
534
+
535
+
536
+def fused_recurrent_kda_packed_decode(
537
+ mixed_qkv: torch.Tensor,
538
+ raw_g: torch.Tensor,
539
+ raw_beta: torch.Tensor,
540
+ A_log: torch.Tensor,
541
+ dt_bias: torch.Tensor,
542
+ lower_bound: float | None,
543
+ initial_state: torch.Tensor,
544
+ state_indices: torch.Tensor,
545
+ scale: float | None = None,
546
+) -> tuple[torch.Tensor, torch.Tensor]:
547
+ """Run one-token KDA decode directly from packed post-conv QKV."""
548
+ if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1:
549
+ raise ValueError("`mixed_qkv` must be 2D and contiguous in its last dim.")
550
+ if raw_g.ndim != 4 or raw_g.shape[0] != 1:
551
+ raise ValueError("`raw_g` must have shape [1, B, H, K].")
552
+ if raw_beta.ndim != 3 or raw_beta.shape[0] != 1:
553
+ raise ValueError("`raw_beta` must have shape [1, B, H].")
554
+ if initial_state.ndim != 4:
555
+ raise ValueError("`initial_state` must have shape [cache, H, V, K].")
556
+ _, H, V, K = initial_state.shape
557
+ if raw_g.stride()[2:] != (K, 1):
558
+ raise ValueError("`raw_g` must be contiguous within each token.")
559
+ if raw_beta.stride(2) != 1:
560
+ raise ValueError("`raw_beta` heads must be contiguous.")
561
+ if initial_state.stride()[1:] != (V * K, K, 1):
562
+ raise ValueError("`initial_state` must be contiguous within each cache slot.")
563
+ if state_indices.ndim != 1 or state_indices.stride(0) != 1:
564
+ raise ValueError("`state_indices` must be contiguous and one-dimensional.")
565
+ if A_log.ndim != 1 or not A_log.is_contiguous():
566
+ raise ValueError("`A_log` must be contiguous and one-dimensional.")
567
+ if not dt_bias.is_contiguous():
568
+ raise ValueError("`dt_bias` must be contiguous.")
569
+
570
+ device = mixed_qkv.device
571
+ if any(
572
+ x.device != device
573
+ for x in (raw_g, raw_beta, A_log, dt_bias, initial_state, state_indices)
574
+ ):
575
+ raise ValueError("All packed KDA inputs must be on the same device.")
576
+
577
+ B = mixed_qkv.shape[0]
578
+ if raw_g.shape != (1, B, H, K):
579
+ raise ValueError(f"Unexpected raw gate shape {tuple(raw_g.shape)}.")
580
+ if raw_beta.shape != (1, B, H):
581
+ raise ValueError(f"Unexpected raw beta shape {tuple(raw_beta.shape)}.")
582
+ if mixed_qkv.shape[1] != 2 * H * K + H * V:
583
+ raise ValueError(f"Unexpected packed QKV shape {tuple(mixed_qkv.shape)}.")
584
+ if A_log.numel() != H or dt_bias.numel() != H * K:
585
+ raise ValueError("`A_log` or `dt_bias` has an incompatible shape.")
586
+ if state_indices.shape[0] != B:
587
+ raise ValueError("`state_indices` must contain one entry per token.")
588
+
589
+ BK = next_power_of_2(K)
590
+ BV = min(next_power_of_2(V), 32)
591
+ if scale is None:
592
+ scale = K**-0.5
593
+
594
+ out = torch.empty((1, B, H, V), dtype=mixed_qkv.dtype, device=device)
595
+ grid = (cdiv(V, BV), B * H)
596
+ fused_recurrent_kda_packed_decode_kernel[grid](
597
+ mixed_qkv=mixed_qkv,
598
+ raw_g=raw_g,
599
+ raw_beta=raw_beta,
600
+ A_log=A_log,
601
+ dt_bias=dt_bias,
602
+ out=out,
603
+ state=initial_state,
604
+ state_indices=state_indices,
605
+ lower_bound=lower_bound or 0.0,
606
+ scale=scale,
607
+ stride_mixed_token=mixed_qkv.stride(0),
608
+ stride_g_token=raw_g.stride(1),
609
+ stride_beta_token=raw_beta.stride(1),
610
+ stride_state_token=initial_state.stride(0),
611
+ H=H,
612
+ K=K,
613
+ V=V,
614
+ BK=BK,
615
+ BV=BV,
616
+ SOFTPLUS_THRESHOLD=20.0,
617
+ USE_LOWER_BOUND=lower_bound is not None,
618
+ num_warps=4,
619
+ num_stages=2,
620
+ )
621
+ return out, initial_state
vllm/models/kimi_k3/common/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/kimi_k3/common/mm_preprocess.py
new
+330
@@ -0,0 +1,330 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Shared Kimi-K3 multimodal preprocessing."""
4
+
5
+import math
6
+from collections.abc import Mapping, Sequence
7
+from typing import Any, cast
8
+
9
+import torch
10
+from transformers import BatchFeature
11
+
12
+from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
13
+from vllm.inputs import MultiModalDataDict
14
+from vllm.logger import init_logger
15
+from vllm.multimodal.inputs import (
16
+ MultiModalFieldConfig,
17
+ MultiModalKwargsItems,
18
+)
19
+from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
20
+from vllm.multimodal.processing import (
21
+ BaseDummyInputsBuilder,
22
+ BaseMultiModalProcessor,
23
+ BaseProcessingInfo,
24
+ InputProcessingContext,
25
+ PromptReplacement,
26
+ PromptUpdate,
27
+ PromptUpdateDetails,
28
+)
29
+from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config
30
+from vllm.transformers_utils.processor import cached_get_image_processor
31
+from vllm.transformers_utils.processors.kimi_k3 import KimiK3Processor
32
+
33
+logger = init_logger(__name__)
34
+
35
+
36
+def navit_resize_image(
37
+ width: int,
38
+ height: int,
39
+ patch_size: int,
40
+ merge_kernel_size: int,
41
+ in_patch_limit: int,
42
+ patch_limit_on_one_side: int,
43
+ fixed_output_tokens: int | None,
44
+):
45
+ # Apply the patch limits.
46
+ s1 = math.sqrt(
47
+ in_patch_limit
48
+ / (max(1.0, width // patch_size) * max(1.0, height // patch_size))
49
+ )
50
+ s2 = patch_limit_on_one_side * patch_size / width
51
+ s3 = patch_limit_on_one_side * patch_size / height
52
+ scale = min(1.0, s1, s2, s3)
53
+ new_w, new_h = max(1, int(width * scale)), max(1, int(height * scale))
54
+ new_w = min(new_w, patch_limit_on_one_side * patch_size)
55
+ new_h = min(new_h, patch_limit_on_one_side * patch_size)
56
+
57
+ factor = merge_kernel_size * patch_size
58
+
59
+ pad_height = (factor - new_h % factor) % factor
60
+ pad_width = (factor - new_w % factor) % factor
61
+
62
+ if fixed_output_tokens is not None:
63
+ num_tokens = fixed_output_tokens
64
+ else:
65
+ # Calculate new dimensions after padding and patching
66
+ token_height = (new_h + pad_height) // factor
67
+ token_width = (new_w + pad_width) // factor
68
+
69
+ assert token_height * merge_kernel_size <= patch_limit_on_one_side, (
70
+ f"token_height {token_height} * merge_kernel_size {merge_kernel_size} > "
71
+ f"patch_limit_on_one_side {patch_limit_on_one_side}"
72
+ )
73
+ assert token_width * merge_kernel_size <= patch_limit_on_one_side, (
74
+ f"token_width {token_width} * merge_kernel_size {merge_kernel_size} > "
75
+ f"patch_limit_on_one_side {patch_limit_on_one_side}"
76
+ )
77
+
78
+ num_tokens = token_height * token_width
79
+ return {
80
+ "num_tokens": num_tokens,
81
+ "new_width": new_w,
82
+ "new_height": new_h,
83
+ "pad_width": pad_width,
84
+ "pad_height": pad_height,
85
+ "sampled_nframes": 1,
86
+ }
87
+
88
+
89
+class KimiK3ProcessingInfo(BaseProcessingInfo):
90
+ """Processing information for the image-only Kimi-K3 model.
91
+
92
+ K3 uses the standard ``image`` modality (unlike K2.5's unified
93
+ ``vision_chunk``), so it builds its own ``KimiK3Processor`` wrapper around
94
+ the checkpoint's image processor and resolves the ``<|media_pad|>`` token
95
+ id the same way K2.5 does.
96
+ """
97
+
98
+ def __init__(self, ctx: InputProcessingContext) -> None:
99
+ super().__init__(ctx)
100
+
101
+ self.hf_config = hf_config = self.get_hf_config()
102
+
103
+ tokenizer = self.get_tokenizer()
104
+ image_processor = cached_get_image_processor(
105
+ self.ctx.model_config.model,
106
+ revision=self.ctx.model_config.revision,
107
+ trust_remote_code=self.ctx.model_config.trust_remote_code,
108
+ )
109
+
110
+ # Resolve token ID from the tokenizer because transformers v5
111
+ # may remap token IDs vs config.json.
112
+ config_token_id = hf_config.media_placeholder_token_id
113
+ resolved_token_id = tokenizer.convert_tokens_to_ids("<|media_pad|>")
114
+ unk_token_id = getattr(tokenizer, "unk_token_id", None)
115
+ is_valid_resolved = isinstance(resolved_token_id, int) and (
116
+ unk_token_id is None or resolved_token_id != unk_token_id
117
+ )
118
+ if is_valid_resolved and resolved_token_id != config_token_id:
119
+ logger.warning_once(
120
+ "Kimi-K3 config.media_placeholder_token_id (%d) disagrees "
121
+ "with tokenizer mapping for <|media_pad|> (%d). "
122
+ "Using tokenizer value.",
123
+ config_token_id,
124
+ resolved_token_id,
125
+ )
126
+ media_token_id = resolved_token_id
127
+ # Patch config so downstream code also sees the correct ID.
128
+ hf_config.media_placeholder_token_id = resolved_token_id
129
+ else:
130
+ media_token_id = config_token_id
131
+
132
+ self.media_token_id = media_token_id
133
+ self.media_token = tokenizer.decode(media_token_id)
134
+
135
+ self.image_processor = image_processor
136
+ self.hf_processor = KimiK3Processor(
137
+ tokenizer=tokenizer,
138
+ image_processor=image_processor,
139
+ )
140
+ self.media_tokens_calculator = image_processor.media_tokens_calculator
141
+
142
+ def get_hf_processor(self, **kwargs: object) -> KimiK3Processor:
143
+ return self.hf_processor
144
+
145
+ def get_hf_config(self) -> KimiK3Config:
146
+ return self.ctx.get_hf_config(KimiK3Config)
147
+
148
+ def get_supported_mm_limits(self) -> Mapping[str, int | None]:
149
+ # None means unlimited
150
+ return {"image": None}
151
+
152
+ @classmethod
153
+ def get_max_image_size(
154
+ cls,
155
+ patch_size: int,
156
+ merge_kernel_size: int,
157
+ in_patch_limit: int,
158
+ patch_limit_on_one_side: int,
159
+ fixed_output_tokens: int | None,
160
+ ) -> ImageSize:
161
+ max_side = patch_limit_on_one_side * patch_size
162
+ best_score = (-1, -1)
163
+ best_size = (max_side, max_side)
164
+
165
+ for width_patches in range(patch_limit_on_one_side + 1):
166
+ width = min((width_patches + 1) * patch_size - 1, max_side)
167
+ for height_patches in range(width_patches, patch_limit_on_one_side + 1):
168
+ height = min((height_patches + 1) * patch_size - 1, max_side)
169
+ resize_config = navit_resize_image(
170
+ width,
171
+ height,
172
+ patch_size,
173
+ merge_kernel_size,
174
+ in_patch_limit,
175
+ patch_limit_on_one_side,
176
+ fixed_output_tokens,
177
+ )
178
+ padded_width = resize_config["new_width"] + resize_config["pad_width"]
179
+ padded_height = (
180
+ resize_config["new_height"] + resize_config["pad_height"]
181
+ )
182
+ num_patches = padded_width // patch_size * (padded_height // patch_size)
183
+ score = (resize_config["num_tokens"], num_patches)
184
+ if score > best_score:
185
+ best_score = score
186
+ best_size = (width, height)
187
+ return ImageSize(width=best_size[0], height=best_size[1])
188
+
189
+
190
+class KimiK3DummyInputsBuilder(BaseDummyInputsBuilder[KimiK3ProcessingInfo]):
191
+ """Builds image-based dummy inputs for K3 profiling.
192
+
193
+ The dummy text is made of ``<|kimi_image_placeholder|>`` tokens — exactly
194
+ the placeholder that K3's ``_get_prompt_updates`` expands — and the dummy
195
+ mm data is a plain list of PIL images under the ``image`` key.
196
+ """
197
+
198
+ def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
199
+ num_images = mm_counts.get("image", 0)
200
+ return self.info.get_hf_config().image_placeholder * num_images
201
+
202
+ def get_dummy_mm_data(
203
+ self,
204
+ seq_len: int,
205
+ mm_counts: Mapping[str, int],
206
+ mm_options: Mapping[str, BaseDummyOptions] | None = None,
207
+ ) -> MultiModalDataDict:
208
+ media_proc_cfg = self.info.image_processor.media_proc_cfg
209
+ max_size = self.info.get_max_image_size(
210
+ media_proc_cfg["patch_size"],
211
+ media_proc_cfg["merge_kernel_size"],
212
+ media_proc_cfg["in_patch_limit"],
213
+ media_proc_cfg["patch_limit_on_one_side"],
214
+ media_proc_cfg["fixed_output_tokens"],
215
+ )
216
+ num_images = mm_counts.get("image", 0)
217
+ image_overrides = cast(
218
+ ImageDummyOptions | None,
219
+ mm_options.get("image") if mm_options else None,
220
+ )
221
+ return {
222
+ "image": self._get_dummy_images(
223
+ width=max_size.width,
224
+ height=max_size.height,
225
+ num_images=num_images,
226
+ overrides=image_overrides,
227
+ )
228
+ }
229
+
230
+
231
+class KimiK3MultiModalProcessor(BaseMultiModalProcessor[KimiK3ProcessingInfo]):
232
+ """Image-only multi-modal processor for Kimi-K3."""
233
+
234
+ def _call_hf_processor(
235
+ self,
236
+ prompt: str,
237
+ mm_data: Mapping[str, object],
238
+ mm_kwargs: Mapping[str, object],
239
+ tok_kwargs: Mapping[str, object],
240
+ ) -> BatchFeature:
241
+ # Override so the base always routes through the text+mm path
242
+ # (`KimiK3Processor.__call__`). Otherwise the mm-only fast path calls
243
+ # the checkpoint image processor directly with bare PIL images, but it
244
+ # requires `{"type": "image", "image": PIL}` media dicts that only our
245
+ # wrapper builds.
246
+ return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs)
247
+
248
+ def _hf_processor_applies_updates(
249
+ self,
250
+ prompt_text: str,
251
+ mm_items: MultiModalDataItems,
252
+ hf_processor_mm_kwargs: Mapping[str, object],
253
+ tokenization_kwargs: Mapping[str, object],
254
+ ) -> bool:
255
+ return False
256
+
257
+ def _get_mm_fields_config(
258
+ self,
259
+ hf_inputs: BatchFeature,
260
+ hf_processor_mm_kwargs: Mapping[str, object],
261
+ ) -> Mapping[str, MultiModalFieldConfig]:
262
+ """Slice the flattened patch tensor back into per-image items.
263
+
264
+ ``pixel_values`` holds all patches from every image concatenated; each
265
+ image's patch count is ``prod(grid_thws[i])``. ``grid_thws`` is one
266
+ ``[N_t, N_h, N_w]`` row per image.
267
+ """
268
+ grid_thws = hf_inputs.get("grid_thws", torch.empty((0, 3)))
269
+ grid_sizes = grid_thws.prod(-1)
270
+
271
+ return dict(
272
+ pixel_values=MultiModalFieldConfig.flat_from_sizes("image", grid_sizes),
273
+ grid_thws=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
274
+ )
275
+
276
+ def _get_prompt_updates(
277
+ self,
278
+ mm_items: MultiModalDataItems,
279
+ hf_processor_mm_kwargs: Mapping[str, Any],
280
+ out_mm_kwargs: MultiModalKwargsItems,
281
+ ) -> Sequence[PromptUpdate]:
282
+ """Expand each K3 image placeholder into a resolution-aware update.
283
+
284
+ K3's prompt carries a single ``<|kimi_image_placeholder|>`` token per
285
+ image. This replaces that token with
286
+ ``<|media_begin|>image {w}x{h}<|media_content|>{pads}<|media_end|>``,
287
+ embedding the per-image resolution in the prompt and marking only the
288
+ ``<|media_pad|>`` positions as embedding slots (the number of pads is
289
+ the feature size returned by ``media_tokens_calculator``).
290
+ """
291
+ media_token_id = self.info.media_token_id
292
+ media_token = self.info.media_token
293
+ image_placeholder = self.info.get_hf_config().image_placeholder
294
+
295
+ def get_replacement(item_idx: int) -> PromptUpdateDetails[str]:
296
+ images = mm_items.get_items("image", ImageProcessorItems)
297
+ image = images.get(item_idx)
298
+ if image is None:
299
+ raise ValueError(f"Missing image data at index {item_idx}")
300
+
301
+ # The checkpoint image processor works on media dicts, so wrap the
302
+ # PIL image before asking it for the token count.
303
+ num_media_token = self.info.media_tokens_calculator(
304
+ {"type": "image", "image": image}
305
+ )
306
+ pads = media_token * num_media_token
307
+
308
+ # NOTE: `width`/`height` are the ORIGINAL upload dimensions, not the
309
+ # post-preprocess (smart-resized) ones. `image` comes from the
310
+ # untouched parsed `mm_items`; the checkpoint image processor
311
+ # (`KimiK3VisionProcessor.preprocess`) only produces new tensors via
312
+ # `image.resize(...)` and never mutates the stored PIL. This matches
313
+ # the reference HF processor (`KimiK3Processor.preprocess_medias`),
314
+ # which also builds the prompt from the original `img.size`. The
315
+ # resize is reflected only in the pad count above.
316
+ width, height = images.get_image_size(item_idx)
317
+ full = (
318
+ f"<|media_begin|>image {width}x{height}<|media_content|>"
319
+ f"{pads}<|media_end|>"
320
+ )
321
+
322
+ return PromptUpdateDetails.select_token_id(full, media_token_id)
323
+
324
+ return [
325
+ PromptReplacement(
326
+ modality="image",
327
+ target=image_placeholder,
328
+ replacement=get_replacement,
329
+ ),
330
+ ]
vllm/models/kimi_k3/common/mtp.py
new
+96
@@ -0,0 +1,96 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Fused Kimi-K3 MTP input preparation."""
4
+
5
+import torch
6
+
7
+from vllm.triton_utils import tl, triton
8
+
9
+
10
+@triton.jit
11
+def _rms_norm(x, weight, eps, hidden_size: tl.constexpr):
12
+ x = x.to(tl.float32)
13
+ variance = tl.sum(x * x, axis=0) / hidden_size
14
+ return x * tl.rsqrt(variance + eps) * weight.to(tl.float32)
15
+
16
+
17
+@triton.jit
18
+def _fused_mtp_input_kernel(
19
+ positions_ptr,
20
+ inputs_embeds_ptr,
21
+ previous_hidden_states_ptr,
22
+ enorm_weight_ptr,
23
+ hnorm_weight_ptr,
24
+ output_ptr,
25
+ eps,
26
+ inputs_embeds_stride,
27
+ previous_hidden_states_stride,
28
+ output_stride,
29
+ hidden_size: tl.constexpr,
30
+ block_size: tl.constexpr,
31
+):
32
+ token_idx = tl.program_id(0).to(tl.int64)
33
+ input_idx = tl.program_id(1)
34
+ offsets = tl.arange(0, block_size)
35
+ mask = offsets < hidden_size
36
+
37
+ if input_idx == 0:
38
+ position = tl.load(positions_ptr + token_idx)
39
+ values = tl.load(
40
+ inputs_embeds_ptr + token_idx * inputs_embeds_stride + offsets,
41
+ mask=mask & (position != 0),
42
+ other=0.0,
43
+ )
44
+ weight = tl.load(enorm_weight_ptr + offsets, mask=mask, other=0.0)
45
+ else:
46
+ values = tl.load(
47
+ previous_hidden_states_ptr
48
+ + token_idx * previous_hidden_states_stride
49
+ + offsets,
50
+ mask=mask,
51
+ other=0.0,
52
+ )
53
+ weight = tl.load(hnorm_weight_ptr + offsets, mask=mask, other=0.0)
54
+
55
+ output = _rms_norm(values, weight, eps, hidden_size)
56
+ tl.store(
57
+ output_ptr + token_idx * output_stride + input_idx * hidden_size + offsets,
58
+ output,
59
+ mask=mask,
60
+ )
61
+
62
+
63
+def fused_mtp_input(
64
+ positions: torch.Tensor,
65
+ inputs_embeds: torch.Tensor,
66
+ previous_hidden_states: torch.Tensor,
67
+ enorm_weight: torch.Tensor,
68
+ hnorm_weight: torch.Tensor,
69
+ eps: float,
70
+) -> torch.Tensor:
71
+ """Mask and normalize both MTP inputs into the projection layout."""
72
+ num_tokens, hidden_size = inputs_embeds.shape
73
+ output = torch.empty(
74
+ num_tokens,
75
+ 2 * hidden_size,
76
+ dtype=inputs_embeds.dtype,
77
+ device=inputs_embeds.device,
78
+ )
79
+ if num_tokens == 0:
80
+ return output
81
+
82
+ _fused_mtp_input_kernel[(num_tokens, 2)](
83
+ positions,
84
+ inputs_embeds,
85
+ previous_hidden_states,
86
+ enorm_weight,
87
+ hnorm_weight,
88
+ output,
89
+ eps,
90
+ inputs_embeds.stride(0),
91
+ previous_hidden_states.stride(0),
92
+ output.stride(0),
93
+ hidden_size,
94
+ triton.next_power_of_2(hidden_size),
95
+ )
96
+ return output
vllm/models/kimi_k3/nvidia/dspark_mla.py
new
+527
@@ -0,0 +1,527 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""K3 dense MLA draft model for DSpark speculative decoding."""
4
+
5
+from collections.abc import Iterable
6
+
7
+import torch
8
+import torch.nn as nn
9
+import torch.nn.functional as F
10
+
11
+import vllm._custom_ops as ops
12
+from vllm.config import VllmConfig
13
+from vllm.model_executor.layers.layernorm import RMSNorm
14
+from vllm.model_executor.layers.linear import ReplicatedLinear
15
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
16
+from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead
17
+from vllm.model_executor.models.utils import (
18
+ AutoWeightsLoader,
19
+ WeightsMapper,
20
+ get_draft_quant_config,
21
+ maybe_prefix,
22
+)
23
+from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention
24
+from vllm.models.kimi_k3.nvidia.model import KimiMLP
25
+from vllm.utils.torch_utils import is_quantized_kv_cache
26
+
27
+
28
+class ReplicatedDSparkMarkovHead(DSparkMarkovHead):
29
+ """DSpark Markov head with full weights on every TP rank."""
30
+
31
+ def __init__(
32
+ self,
33
+ vocab_size: int,
34
+ draft_vocab_size: int,
35
+ markov_rank: int,
36
+ prefix: str,
37
+ ) -> None:
38
+ # TODO: Remove this mypy workaround once the K3 PR is fully merged.
39
+ super().__init__( # type: ignore[call-arg]
40
+ vocab_size,
41
+ draft_vocab_size,
42
+ markov_rank,
43
+ prefix,
44
+ replicated=True,
45
+ )
46
+
47
+
48
+class K3DSparkDecoderLayer(nn.Module):
49
+ def __init__(
50
+ self,
51
+ *,
52
+ vllm_config: VllmConfig,
53
+ config,
54
+ layer_idx: int,
55
+ start_layer_id: int,
56
+ prefix: str,
57
+ ) -> None:
58
+ super().__init__()
59
+ quant_config = get_draft_quant_config(vllm_config)
60
+ self.self_attn = MultiHeadLatentAttention(
61
+ config=config,
62
+ hidden_size=config.hidden_size,
63
+ num_heads=config.num_attention_heads,
64
+ qk_nope_head_dim=config.qk_nope_head_dim,
65
+ qk_rope_head_dim=config.qk_rope_head_dim,
66
+ v_head_dim=config.v_head_dim,
67
+ q_lora_rank=config.q_lora_rank,
68
+ kv_lora_rank=config.kv_lora_rank,
69
+ cache_config=vllm_config.cache_config,
70
+ quant_config=quant_config,
71
+ prefix=maybe_prefix(
72
+ prefix, f"layers.{start_layer_id + layer_idx}.self_attn"
73
+ ),
74
+ use_rope=True,
75
+ non_causal_multi_token_decode=True,
76
+ )
77
+ self.mlp = KimiMLP(
78
+ hidden_size=config.hidden_size,
79
+ intermediate_size=config.intermediate_size,
80
+ hidden_act=config.hidden_act,
81
+ quant_config=quant_config,
82
+ prefix=maybe_prefix(prefix, f"layers.{start_layer_id + layer_idx}.mlp"),
83
+ )
84
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
85
+ self.post_attention_layernorm = RMSNorm(
86
+ config.hidden_size, eps=config.rms_norm_eps
87
+ )
88
+
89
+ def forward(
90
+ self,
91
+ positions: torch.Tensor,
92
+ hidden_states: torch.Tensor,
93
+ residual: torch.Tensor | None,
94
+ ) -> tuple[torch.Tensor, torch.Tensor]:
95
+ if residual is None:
96
+ residual = hidden_states
97
+ hidden_states = self.input_layernorm(hidden_states)
98
+ else:
99
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
100
+
101
+ hidden_states = self.self_attn(
102
+ positions=positions,
103
+ hidden_states=hidden_states,
104
+ )
105
+ hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
106
+ hidden_states = self.mlp(hidden_states)
107
+ return hidden_states, residual
108
+
109
+
110
+class K3DSparkModel(nn.Module):
111
+ def __init__(
112
+ self,
113
+ *,
114
+ vllm_config: VllmConfig,
115
+ start_layer_id: int,
116
+ prefix: str,
117
+ ) -> None:
118
+ super().__init__()
119
+ assert vllm_config.speculative_config is not None
120
+ self.config = vllm_config.speculative_config.draft_model_config.hf_config
121
+ self.quant_config = get_draft_quant_config(vllm_config)
122
+
123
+ # The frozen target embedding is aliased after the draft checkpoint loads.
124
+ self.embed_tokens: nn.Module | None = None
125
+
126
+ self.context_proj = ReplicatedLinear(
127
+ self.config.target_hidden_size * self.config.num_target_layers,
128
+ self.config.hidden_size,
129
+ bias=False,
130
+ return_bias=False,
131
+ quant_config=self.quant_config,
132
+ prefix=maybe_prefix(prefix, "context_proj"),
133
+ )
134
+ self.context_norm = RMSNorm(
135
+ self.config.hidden_size, eps=self.config.rms_norm_eps
136
+ )
137
+
138
+ self.layers = nn.ModuleList(
139
+ [
140
+ K3DSparkDecoderLayer(
141
+ vllm_config=vllm_config,
142
+ config=self.config,
143
+ layer_idx=layer_idx,
144
+ start_layer_id=start_layer_id,
145
+ prefix=prefix,
146
+ )
147
+ for layer_idx in range(self.config.num_hidden_layers)
148
+ ]
149
+ )
150
+ self.final_norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)
151
+ self.markov_head = ReplicatedDSparkMarkovHead(
152
+ self.config.vocab_size,
153
+ self.config.draft_vocab_size,
154
+ self.config.markov_rank,
155
+ prefix=maybe_prefix(prefix, "markov_head"),
156
+ )
157
+ self._context_kv_fusion_available: bool | None = None
158
+ self._max_num_context_tokens = (
159
+ vllm_config.scheduler_config.max_num_batched_tokens
160
+ )
161
+
162
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
163
+ assert self.embed_tokens is not None
164
+ return self.embed_tokens(input_ids)
165
+
166
+ def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor:
167
+ return self.context_norm(self.context_proj(hidden_states))
168
+
169
+ @torch.inference_mode()
170
+ def precompute_and_store_context_kv(
171
+ self,
172
+ context_states: torch.Tensor,
173
+ context_positions: torch.Tensor,
174
+ context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
175
+ ) -> None:
176
+ """Project target-derived context into each draft layer's latent cache."""
177
+ if self._context_kv_fusion_available is None:
178
+ self._build_fused_context_kv_buffers()
179
+ if self._context_kv_fusion_available:
180
+ self._precompute_fused_context_kv(
181
+ context_states, context_positions, context_slot_mapping
182
+ )
183
+ return
184
+
185
+ # Quantized fallback. Directly invoking the projection modules preserves
186
+ # their quantization methods, at the cost of also computing unused Q rows.
187
+ for layer_idx, layer in enumerate(self.layers):
188
+ attn = layer.self_attn
189
+ assert attn.fused_qkv_a_proj is not None
190
+ assert attn.q_lora_rank is not None
191
+ assert attn.rotary_emb is not None
192
+ qkv_lora = attn.fused_qkv_a_proj(context_states)[0]
193
+ kv_lora = qkv_lora[..., attn.q_lora_rank :]
194
+ kv_c, k_pe = kv_lora.split(
195
+ [attn.kv_lora_rank, attn.qk_rope_head_dim], dim=-1
196
+ )
197
+ kv_c = attn.kv_a_layernorm(kv_c)
198
+ k_pe = k_pe.unsqueeze(1)
199
+ # DeepSeek YaRN's FlashInfer path requires paired Q/K tensors.
200
+ # The vLLM CUDA op supports rotating one tensor in place and
201
+ # consumes the same (possibly scaled fp32) cos/sin cache.
202
+ rotary_emb = attn.rotary_emb
203
+ ops.rotary_embedding(
204
+ context_positions,
205
+ k_pe,
206
+ None,
207
+ rotary_emb.head_size,
208
+ rotary_emb.cos_sin_cache,
209
+ rotary_emb.is_neox_style,
210
+ )
211
+
212
+ slot_mapping = (
213
+ context_slot_mapping[layer_idx]
214
+ if isinstance(context_slot_mapping, (list, tuple))
215
+ else context_slot_mapping
216
+ )
217
+ if slot_mapping is None:
218
+ continue
219
+ attn.impl.do_kv_cache_update(
220
+ kv_c,
221
+ k_pe,
222
+ attn.kv_cache,
223
+ slot_mapping,
224
+ attn.kv_cache_dtype,
225
+ attn._k_scale,
226
+ )
227
+
228
+ def _build_fused_context_kv_buffers(self) -> None:
229
+ """Build a cross-layer KV-only A projection after checkpoint loading."""
230
+ if self.quant_config is not None:
231
+ self._context_kv_fusion_available = False
232
+ return
233
+
234
+ attentions = [layer.self_attn for layer in self.layers]
235
+ if not attentions or any(
236
+ attn.fused_qkv_a_proj is None
237
+ or not hasattr(attn.fused_qkv_a_proj, "weight")
238
+ for attn in attentions
239
+ ):
240
+ self._context_kv_fusion_available = False
241
+ return
242
+
243
+ attn0 = attentions[0]
244
+ assert attn0.q_lora_rank is not None
245
+ kv_width = attn0.kv_lora_rank + attn0.qk_rope_head_dim
246
+ kv_weights = []
247
+ for attn in attentions:
248
+ assert attn.q_lora_rank is not None
249
+ assert (
250
+ attn.q_lora_rank == attn0.q_lora_rank
251
+ and attn.kv_lora_rank == attn0.kv_lora_rank
252
+ and attn.qk_rope_head_dim == attn0.qk_rope_head_dim
253
+ and attn.kv_a_layernorm.variance_epsilon
254
+ == attn0.kv_a_layernorm.variance_epsilon
255
+ ), "All MLA DSpark layers must share their latent KV geometry."
256
+ kv_weights.append(
257
+ attn.fused_qkv_a_proj.weight.detach().narrow(
258
+ 0, attn.q_lora_rank, kv_width
259
+ )
260
+ )
261
+
262
+ # [L * (kv_lora_rank + rope_dim), hidden_size]. The underlying fused
263
+ # A weights are replicated (`disable_tp=True`), so this is valid on
264
+ # every TP rank without communication.
265
+ self._fused_context_kv_weight = torch.cat(kv_weights, dim=0)
266
+ self._context_kv_norm_weights = torch.stack(
267
+ [attn.kv_a_layernorm.weight.detach() for attn in attentions], dim=0
268
+ ).contiguous()
269
+ self._num_context_layers = len(attentions)
270
+ self._context_kv_width = kv_width
271
+ self._context_kv_lora_rank = attn0.kv_lora_rank
272
+ self._context_rope_dim = attn0.qk_rope_head_dim
273
+ self._context_rms_norm_eps = attn0.kv_a_layernorm.variance_epsilon
274
+ self._context_positions_repeated = torch.empty(
275
+ self._num_context_layers * self._max_num_context_tokens,
276
+ dtype=torch.int64,
277
+ device=self._fused_context_kv_weight.device,
278
+ )
279
+ self._context_kv_fusion_available = True
280
+
281
+ def _precompute_fused_context_kv(
282
+ self,
283
+ context_states: torch.Tensor,
284
+ context_positions: torch.Tensor,
285
+ context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None,
286
+ ) -> None:
287
+ num_ctx = context_states.shape[0]
288
+ num_layers = self._num_context_layers
289
+
290
+ # One KV-only GEMM replaces five full Q+KV GEMMs. For K3 this projects
291
+ # 5*576 rows rather than 5*2112 rows (72.7% fewer A-projection FLOPs).
292
+ all_kv = F.linear(context_states, self._fused_context_kv_weight)
293
+ all_kv = all_kv.view(num_ctx, num_layers, self._context_kv_width)
294
+ all_kv_c = all_kv[..., : self._context_kv_lora_rank]
295
+ all_k_pe = all_kv[..., self._context_kv_lora_rank :]
296
+
297
+ # Layer-major layout lets the 2-D RMSNorm weights select a distinct row
298
+ # for each draft layer in one grouped kernel.
299
+ all_kv_c = all_kv_c.permute(1, 0, 2).contiguous()
300
+ all_kv_c_normed = torch.empty_like(all_kv_c)
301
+ ops.rms_norm(
302
+ all_kv_c_normed,
303
+ all_kv_c,
304
+ self._context_kv_norm_weights,
305
+ self._context_rms_norm_eps,
306
+ )
307
+
308
+ all_k_pe = all_k_pe.permute(1, 0, 2).contiguous()
309
+ all_k_pe_flat = all_k_pe.view(num_layers * num_ctx, 1, self._context_rope_dim)
310
+ repeated_positions = self._context_positions_repeated[: num_layers * num_ctx]
311
+ repeated_positions.view(num_layers, num_ctx).copy_(context_positions)
312
+ # Keep the single-tensor context RoPE on vLLM's optimized CUDA op;
313
+ # DeepSeek YaRN's FlashInfer wrapper assumes a non-null key tensor.
314
+ rotary_emb = self.layers[0].self_attn.rotary_emb
315
+ assert rotary_emb is not None
316
+ ops.rotary_embedding(
317
+ repeated_positions,
318
+ all_k_pe_flat,
319
+ None,
320
+ rotary_emb.head_size,
321
+ rotary_emb.cos_sin_cache,
322
+ rotary_emb.is_neox_style,
323
+ )
324
+ all_k_pe = all_k_pe_flat.view(num_layers, num_ctx, 1, self._context_rope_dim)
325
+
326
+ if context_slot_mapping is None:
327
+ return
328
+
329
+ cache_layers = [layer.self_attn for layer in self.layers]
330
+ if (
331
+ not is_quantized_kv_cache(cache_layers[0].kv_cache_dtype)
332
+ and self._has_uniform_block_layout(cache_layers)
333
+ and (
334
+ isinstance(context_slot_mapping, torch.Tensor)
335
+ or all(s is not None for s in context_slot_mapping)
336
+ )
337
+ ):
338
+ # Grouped context KV insert only supports unquantized (bf16) KV cache
339
+ # and assumes that all layers share the same block layout.
340
+
341
+ if isinstance(context_slot_mapping, (list, tuple)):
342
+ per_layer_slot_mappings = [
343
+ s for s in context_slot_mapping if s is not None
344
+ ]
345
+ if len({s.data_ptr() for s in per_layer_slot_mappings}) == 1:
346
+ # All rows alias to the same slot mapping.
347
+ slot_mapping = (
348
+ per_layer_slot_mappings[0].unsqueeze(0).expand(num_layers, -1)
349
+ )
350
+ else:
351
+ slot_mapping = torch.stack(per_layer_slot_mappings, dim=0)
352
+ else:
353
+ # Broadcast the single shared context_slot_mapping tensor.
354
+ slot_mapping = context_slot_mapping.unsqueeze(0).expand(num_layers, -1)
355
+
356
+ ref_cache = cache_layers[0].kv_cache
357
+ ops.concat_and_cache_mla_grouped(
358
+ all_kv_c_normed,
359
+ all_k_pe.squeeze(2),
360
+ self._get_context_kv_cache_ptrs(cache_layers),
361
+ slot_mapping,
362
+ ref_cache.size(1),
363
+ ref_cache.stride(0),
364
+ ref_cache.stride(1),
365
+ )
366
+ return
367
+
368
+ for layer_idx, layer in enumerate(self.layers):
369
+ slot_mapping = (
370
+ context_slot_mapping[layer_idx]
371
+ if isinstance(context_slot_mapping, (list, tuple))
372
+ else context_slot_mapping
373
+ )
374
+ if slot_mapping is None:
375
+ continue
376
+ attn = layer.self_attn
377
+ attn.impl.do_kv_cache_update(
378
+ all_kv_c_normed[layer_idx],
379
+ all_k_pe[layer_idx],
380
+ attn.kv_cache,
381
+ slot_mapping,
382
+ attn.kv_cache_dtype,
383
+ attn._k_scale,
384
+ )
385
+
386
+ def _has_uniform_block_layout(
387
+ self,
388
+ cache_layers: list[MultiHeadLatentAttention],
389
+ ) -> bool:
390
+ if not hasattr(self, "_layers_share_kv_block_layout"):
391
+ ref_cache = cache_layers[0].kv_cache
392
+ self._layers_share_kv_block_layout = all(
393
+ cl.kv_cache.size(1) == ref_cache.size(1)
394
+ and cl.kv_cache.stride(0) == ref_cache.stride(0)
395
+ and cl.kv_cache.stride(1) == ref_cache.stride(1)
396
+ for cl in cache_layers
397
+ )
398
+ return self._layers_share_kv_block_layout
399
+
400
+ def _get_context_kv_cache_ptrs(
401
+ self,
402
+ cache_layers: list[MultiHeadLatentAttention],
403
+ ) -> torch.Tensor:
404
+ # The per-layer KV cache base pointers are stable after allocation, so
405
+ # build the pointer array once and return it on every call.
406
+ if not hasattr(self, "_context_cache_ptrs"):
407
+ ref_cache = cache_layers[0].kv_cache
408
+ cache_ptrs = torch.tensor(
409
+ [cl.kv_cache.data_ptr() for cl in cache_layers],
410
+ dtype=torch.int64,
411
+ device=ref_cache.device,
412
+ )
413
+ self._context_cache_ptrs = cache_ptrs
414
+ return self._context_cache_ptrs
415
+
416
+ def forward(
417
+ self,
418
+ input_ids: torch.Tensor,
419
+ positions: torch.Tensor,
420
+ inputs_embeds: torch.Tensor | None = None,
421
+ ) -> torch.Tensor:
422
+ if inputs_embeds is None:
423
+ inputs_embeds = self.embed_input_ids(input_ids)
424
+
425
+ hidden_states = inputs_embeds
426
+ residual = None
427
+ for layer in self.layers:
428
+ hidden_states, residual = layer(
429
+ positions=positions,
430
+ hidden_states=hidden_states,
431
+ residual=residual,
432
+ )
433
+ hidden_states, _ = self.final_norm(hidden_states, residual)
434
+ return hidden_states
435
+
436
+
437
+class K3DSparkForCausalLM(nn.Module):
438
+ has_own_embed_tokens = False
439
+ has_own_lm_head = False
440
+ draft_id_to_target_id = None
441
+ checkpoint_skip_substrs = ("confidence_head", "embed_tokens", "lm_head")
442
+
443
+ hf_to_vllm_mapper = WeightsMapper(
444
+ orig_to_new_prefix={"": "model."},
445
+ orig_to_new_stacked={
446
+ ".gate_proj": (".gate_up_proj", 0),
447
+ ".up_proj": (".gate_up_proj", 1),
448
+ ".q_a_proj": (".fused_qkv_a_proj", 0),
449
+ ".kv_a_proj_with_mqa": (".fused_qkv_a_proj", 1),
450
+ },
451
+ )
452
+
453
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
454
+ super().__init__()
455
+ assert vllm_config.speculative_config is not None
456
+ self.draft_model_config = vllm_config.speculative_config.draft_model_config
457
+ self.config = self.draft_model_config.hf_config
458
+ target_layer_num = vllm_config.model_config.get_num_layers(
459
+ vllm_config.parallel_config
460
+ )
461
+ self.model = K3DSparkModel(
462
+ vllm_config=vllm_config,
463
+ start_layer_id=target_layer_num,
464
+ prefix=maybe_prefix(prefix, "model"),
465
+ )
466
+
467
+ # Assigned by load_dspark_model from the target. Keeping no placeholder
468
+ # avoids a transient full-vocabulary allocation for this 163k-vocab model.
469
+ self.lm_head: nn.Module | None = None
470
+ logit_scale = getattr(self.config, "logit_scale", 1.0)
471
+ self.logits_processor = LogitsProcessor(
472
+ self.config.draft_vocab_size, scale=logit_scale
473
+ )
474
+
475
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
476
+ return self.model.embed_input_ids(input_ids)
477
+
478
+ def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor:
479
+ return self.model.combine_hidden_states(hidden_states)
480
+
481
+ def get_draft_kv_cache_layer_names(self) -> list[str]:
482
+ return [layer.self_attn.layer_name for layer in self.model.layers]
483
+
484
+ def precompute_and_store_context_kv(
485
+ self,
486
+ context_states: torch.Tensor,
487
+ context_positions: torch.Tensor,
488
+ context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
489
+ ) -> None:
490
+ self.model.precompute_and_store_context_kv(
491
+ context_states, context_positions, context_slot_mapping
492
+ )
493
+
494
+ def forward(
495
+ self,
496
+ input_ids: torch.Tensor,
497
+ positions: torch.Tensor,
498
+ inputs_embeds: torch.Tensor | None = None,
499
+ ) -> torch.Tensor:
500
+ return self.model(input_ids, positions, inputs_embeds)
501
+
502
+ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
503
+ assert self.lm_head is not None
504
+ return self.logits_processor(self.lm_head, hidden_states)
505
+
506
+ def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
507
+ return self.compute_logits(hidden_states)
508
+
509
+ def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor:
510
+ return draft_ids
511
+
512
+ def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
513
+ return self.model.markov_head.embed(token_ids)
514
+
515
+ def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
516
+ return self.model.markov_head.bias(markov_embed, self.logits_processor)
517
+
518
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
519
+ # confidence_head is training-only. The frozen target embedding and LM
520
+ # head are shared after this draft-specific checkpoint is loaded.
521
+ loader = AutoWeightsLoader(
522
+ self,
523
+ skip_substrs=list(self.checkpoint_skip_substrs),
524
+ )
525
+ loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
526
+ self.model._build_fused_context_kv_buffers()
527
+ return loaded_weights
vllm/models/kimi_k3/nvidia/kda.py
new
+775
@@ -0,0 +1,775 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections.abc import Callable
5
+
6
+import torch
7
+from einops import rearrange
8
+from torch import nn
9
+from torch.nn.parameter import Parameter
10
+
11
+from vllm import _custom_ops as ops
12
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
13
+from vllm.config import VllmConfig
14
+from vllm.distributed import divide, get_tensor_model_parallel_rank
15
+from vllm.forward_context import get_forward_context
16
+from vllm.logger import init_logger
17
+from vllm.model_executor.layers.linear import (
18
+ ColumnParallelLinear,
19
+ MergedColumnParallelLinear,
20
+ RowParallelLinear,
21
+)
22
+from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
23
+from vllm.model_executor.layers.mamba.mamba_utils import (
24
+ MambaStateDtypeCalculator,
25
+ MambaStateShapeCalculator,
26
+ is_conv_state_dim_first,
27
+)
28
+from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
29
+ causal_conv1d_fn,
30
+ causal_conv1d_update,
31
+)
32
+from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
33
+ gather_initial_states,
34
+)
35
+from vllm.model_executor.model_loader.weight_utils import (
36
+ default_weight_loader,
37
+ sharded_weight_loader,
38
+)
39
+from vllm.model_executor.parameter import BasevLLMParameter
40
+from vllm.model_executor.utils import set_weight_attrs
41
+from vllm.models.kimi_k3.nvidia.kda_metadata import (
42
+ KimiK3KDAAttentionBackend,
43
+ KimiK3KDAMetadata,
44
+)
45
+from vllm.platforms import current_platform
46
+from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
47
+from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
48
+from vllm.v1.attention.backend import AttentionBackend
49
+
50
+logger = init_logger(__name__)
51
+
52
+_KDA_GATE_LOGBOUND_MIN = -5.0
53
+
54
+
55
+def a_log_weight_loader(
56
+ shard_axis: int,
57
+) -> Callable[[torch.Tensor, torch.Tensor], None]:
58
+ """Load KDA A_log stored as either old 4D or current 1D weights."""
59
+
60
+ def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
61
+ tp_rank = get_tensor_model_parallel_rank()
62
+ shard_size = param.data.shape[shard_axis]
63
+ start_idx = tp_rank * shard_size
64
+
65
+ if loaded_weight.dim() == 4:
66
+ assert loaded_weight.shape[:2] == (1, 1), (
67
+ f"Expected old A_log shape (1, 1, H, 1), got {loaded_weight.shape}"
68
+ )
69
+ assert loaded_weight.shape[-1] == 1, (
70
+ f"Expected old A_log last dim to be 1, got {loaded_weight.shape}"
71
+ )
72
+ loaded_weight = loaded_weight.view(loaded_weight.shape[2])
73
+
74
+ loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size)
75
+ return default_weight_loader(param, loaded_weight)
76
+
77
+ return loader
78
+
79
+
80
+class _KimiGDNMergedColumnParallelLinear(MergedColumnParallelLinear):
81
+ """Merged projection with one output replicated across TP ranks."""
82
+
83
+ def __init__(
84
+ self,
85
+ input_size: int,
86
+ output_sizes: list[int],
87
+ replicated_shard_id: int,
88
+ tp_size: int,
89
+ **kwargs,
90
+ ) -> None:
91
+ self.replicated_shard_id = replicated_shard_id
92
+ output_sizes = output_sizes.copy()
93
+ output_sizes[replicated_shard_id] *= tp_size
94
+ super().__init__(input_size, output_sizes, **kwargs)
95
+
96
+ def weight_loader(
97
+ self,
98
+ param: Parameter,
99
+ loaded_weight: torch.Tensor,
100
+ loaded_shard_id: tuple[int, ...] | int | None = None,
101
+ ) -> None:
102
+ tp_rank = self.tp_rank
103
+ param_tp_rank = getattr(param, "tp_rank", None)
104
+ if loaded_shard_id == self.replicated_shard_id:
105
+ self.tp_rank = 0
106
+ if param_tp_rank is not None:
107
+ param.tp_rank = 0
108
+ try:
109
+ super().weight_loader(param, loaded_weight, loaded_shard_id)
110
+ finally:
111
+ self.tp_rank = tp_rank
112
+ if param_tp_rank is not None:
113
+ param.tp_rank = param_tp_rank
114
+
115
+ def weight_loader_v2(
116
+ self,
117
+ param: BasevLLMParameter,
118
+ loaded_weight: torch.Tensor,
119
+ loaded_shard_id: tuple[int, ...] | int | None = None,
120
+ ) -> None:
121
+ tp_rank = self.tp_rank
122
+ param_tp_rank = getattr(param, "tp_rank", None)
123
+ if loaded_shard_id == self.replicated_shard_id:
124
+ self.tp_rank = 0
125
+ if param_tp_rank is not None:
126
+ param.tp_rank = 0
127
+ try:
128
+ super().weight_loader_v2(param, loaded_weight, loaded_shard_id)
129
+ finally:
130
+ self.tp_rank = tp_rank
131
+ if param_tp_rank is not None:
132
+ param.tp_rank = param_tp_rank
133
+
134
+
135
+def is_fused_kda_decode_supported(
136
+ num_heads: int,
137
+ head_dim: int,
138
+ conv_width: int,
139
+ num_spec: int,
140
+ input_dtype: torch.dtype,
141
+ conv_state_dtype: torch.dtype,
142
+) -> bool:
143
+ if (
144
+ num_heads not in (12, 24, 48, 96)
145
+ or head_dim != 128
146
+ or conv_width != 4
147
+ or num_spec != 0
148
+ or input_dtype != torch.bfloat16
149
+ or conv_state_dtype != torch.bfloat16
150
+ or is_conv_state_dim_first()
151
+ or not hasattr(torch.ops._C, "fused_kda_decode")
152
+ ):
153
+ return False
154
+ # SM90 is architecture-specific; SM10x and SM12x use family binaries.
155
+ return (
156
+ current_platform.is_device_capability(90)
157
+ or current_platform.is_device_capability_family(100)
158
+ or current_platform.is_device_capability_family(120)
159
+ )
160
+
161
+
162
+def is_flashkda_supported(
163
+ head_dim: int,
164
+ dtype: torch.dtype,
165
+ lower_bound: float | None,
166
+) -> bool:
167
+ capability = current_platform.get_device_capability()
168
+ return (
169
+ capability is not None
170
+ and capability.major in (9, 10, 12)
171
+ and head_dim == 128
172
+ and dtype == torch.bfloat16
173
+ and lower_bound is not None
174
+ )
175
+
176
+
177
+def _flashkda_prefill(
178
+ q: torch.Tensor,
179
+ k: torch.Tensor,
180
+ v: torch.Tensor,
181
+ g: torch.Tensor,
182
+ beta: torch.Tensor,
183
+ A_log: torch.Tensor,
184
+ dt_bias: torch.Tensor,
185
+ lower_bound: float,
186
+ initial_state: torch.Tensor,
187
+ cu_seqlens: torch.Tensor,
188
+) -> tuple[torch.Tensor, torch.Tensor]:
189
+ import vllm._flashkda_C # noqa: F401
190
+
191
+ out = torch.empty(v.shape, dtype=v.dtype, device=v.device)
192
+ final_state = torch.empty_like(initial_state)
193
+ workspace = torch.empty(
194
+ torch.ops._flashkda_C.get_workspace_size(
195
+ q.shape[0] * q.shape[1],
196
+ q.shape[2],
197
+ cu_seqlens.numel() - 1,
198
+ ),
199
+ dtype=torch.uint8,
200
+ device=q.device,
201
+ )
202
+ # FlashKDA hardcodes dense Q/K/V/G strides. Beta may be row-strided because
203
+ # FlashKDA materializes its transposed [H, T] layout internally.
204
+ # TODO: Teach FlashKDA to consume beta in [T, H] layout directly instead
205
+ # of transposing it to contiguous [H, T] storage internally.
206
+ torch.ops._flashkda_C.fwd(
207
+ q.contiguous(),
208
+ k.contiguous(),
209
+ v.contiguous(),
210
+ g.contiguous(),
211
+ beta,
212
+ q.shape[-1] ** -0.5,
213
+ out,
214
+ workspace,
215
+ A_log.contiguous(),
216
+ dt_bias.view(-1, q.shape[-1]).contiguous(),
217
+ lower_bound,
218
+ initial_state.contiguous(),
219
+ final_state,
220
+ cu_seqlens.contiguous(),
221
+ )
222
+ return out, final_state
223
+
224
+
225
+def resolve_kda_prefill_backend(
226
+ backend: str,
227
+ head_dim: int,
228
+ dtype: torch.dtype,
229
+ lower_bound: float | None,
230
+) -> str:
231
+ if backend not in ("auto", "triton", "flashkda"):
232
+ raise ValueError(f"Unsupported KDA prefill backend: {backend}")
233
+ supported = is_flashkda_supported(head_dim, dtype, lower_bound)
234
+ if backend == "flashkda" and not supported:
235
+ raise RuntimeError(
236
+ "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16, "
237
+ "head_dim=128, and a bounded KDA gate."
238
+ )
239
+ if supported and backend != "triton":
240
+ logger.info_once("Using FlashKDA KDA prefill backend.")
241
+ return "flashkda"
242
+ return "triton"
243
+
244
+
245
+def _make_decode_conv1d_weight_loader(
246
+ dims: list[int],
247
+ tp_size: int,
248
+ tp_rank: int,
249
+ decode_conv1d_weight: torch.Tensor | None,
250
+) -> Callable[..., None]:
251
+ sharded_dims = [dim // tp_size for dim in dims]
252
+
253
+ def weight_loader(
254
+ param: torch.Tensor,
255
+ loaded_weight: torch.Tensor,
256
+ loaded_shard_id: int,
257
+ ) -> None:
258
+ if loaded_weight.dim() == 2:
259
+ loaded_weight = loaded_weight.unsqueeze(1)
260
+ shard_size = sharded_dims[loaded_shard_id]
261
+ source_start = tp_rank * shard_size
262
+ target_start = sum(sharded_dims[:loaded_shard_id])
263
+ loaded_shard = loaded_weight[source_start : source_start + shard_size]
264
+ param.data[target_start : target_start + shard_size].copy_(loaded_shard)
265
+ if decode_conv1d_weight is not None and not param.is_meta:
266
+ decode_conv1d_weight[loaded_shard_id].copy_(
267
+ loaded_shard.squeeze(1).transpose(0, 1)
268
+ )
269
+
270
+ return weight_loader
271
+
272
+
273
+def _make_decode_norm_weight_loader(
274
+ decode_norm_weight: torch.Tensor,
275
+) -> Callable[..., None]:
276
+ def weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
277
+ default_weight_loader(param, loaded_weight)
278
+ if not param.is_meta:
279
+ decode_norm_weight.copy_(param.data)
280
+
281
+ return weight_loader
282
+
283
+
284
+class KimiK3DeltaAttention(GatedDeltaNetAttention):
285
+ def get_attn_backend(self) -> type[AttentionBackend]:
286
+ return KimiK3KDAAttentionBackend
287
+
288
+ def get_state_dtype(
289
+ self,
290
+ ) -> tuple[torch.dtype, torch.dtype]:
291
+ if self.model_config is None or self.cache_config is None:
292
+ raise ValueError("model_config and cache_config must be set")
293
+ return MambaStateDtypeCalculator.kda_state_dtype(
294
+ self.model_config.dtype, self.cache_config.mamba_cache_dtype
295
+ )
296
+
297
+ def get_state_shape(
298
+ self,
299
+ ) -> tuple[tuple[int, ...], tuple[int, ...]]:
300
+ return MambaStateShapeCalculator.kda_state_shape(
301
+ self.tp_size,
302
+ self.num_heads,
303
+ self.head_dim,
304
+ conv_kernel_size=self.conv_size,
305
+ num_spec=self.num_spec,
306
+ )
307
+
308
+ def __init__(
309
+ self,
310
+ config: KimiLinearConfig,
311
+ vllm_config: VllmConfig,
312
+ prefix: str = "",
313
+ ) -> None:
314
+ super().__init__(config, vllm_config, prefix)
315
+
316
+ kda_config = config.linear_attn_config # type: ignore[attr-defined]
317
+ assert kda_config is not None, "linear_attn_config must be set"
318
+ self.head_dim = kda_config["head_dim"]
319
+ self.num_heads = kda_config["num_heads"]
320
+ assert self.num_heads % self.tp_size == 0
321
+ self.local_num_heads = divide(self.num_heads, self.tp_size)
322
+ self.projection_size = self.head_dim * self.num_heads
323
+ self.local_projection_size = divide(self.projection_size, self.tp_size)
324
+ self.conv_size = kda_config["short_conv_kernel_size"]
325
+ assert kda_config.get("use_full_rank_gate", False), (
326
+ "KimiK3DeltaAttention requires a full-rank gate"
327
+ )
328
+
329
+ # Keep f_a before the narrow beta shard, then pad each TP-local row
330
+ # to select the aligned BF16 GEMM path.
331
+ qkvg_output_sizes = [self.projection_size] * 4
332
+ in_proj_output_sizes = qkvg_output_sizes + [
333
+ self.head_dim,
334
+ self.num_heads,
335
+ ]
336
+ local_output_size = (
337
+ 4 * self.local_projection_size + self.head_dim + self.local_num_heads
338
+ )
339
+ self.in_proj_padding = -local_output_size % 16
340
+ if self.in_proj_padding:
341
+ in_proj_output_sizes.append(self.in_proj_padding * self.tp_size)
342
+ self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear(
343
+ self.hidden_size,
344
+ in_proj_output_sizes,
345
+ replicated_shard_id=4,
346
+ tp_size=self.tp_size,
347
+ bias=False,
348
+ quant_config=self.quant_config,
349
+ prefix=f"{prefix}.in_proj_qkvgfab",
350
+ )
351
+ if self.in_proj_padding:
352
+ self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_()
353
+
354
+ self.f_b_proj = ColumnParallelLinear(
355
+ self.head_dim,
356
+ self.projection_size,
357
+ bias=False,
358
+ quant_config=self.quant_config,
359
+ prefix=f"{prefix}.f_b_proj",
360
+ )
361
+ self.dt_bias = nn.Parameter(
362
+ torch.empty(self.local_projection_size, dtype=torch.float32)
363
+ )
364
+ set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
365
+
366
+ # One packed parameter and cache let decode run a single conv update.
367
+ # Prefill slices them back into Q/K/V to obtain dense outputs cheaply.
368
+ self.conv1d = ColumnParallelLinear(
369
+ input_size=self.conv_size,
370
+ output_size=3 * self.projection_size,
371
+ bias=False,
372
+ params_dtype=torch.float32,
373
+ prefix=f"{prefix}.conv1d",
374
+ )
375
+ self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
376
+ # Keep a width-major copy for fused decode without changing the layout
377
+ # consumed by the prefill and fallback decode kernels.
378
+ conv_state_dtype, _ = self.get_state_dtype()
379
+ decode_conv1d_weight = None
380
+ if is_fused_kda_decode_supported(
381
+ self.local_num_heads,
382
+ self.head_dim,
383
+ self.conv_size,
384
+ self.num_spec,
385
+ vllm_config.model_config.dtype,
386
+ conv_state_dtype,
387
+ ):
388
+ logger.info_once("Fused KDA decode kernel (conv+KDA+norm) is enabled.")
389
+ decode_conv1d_weight = torch.empty(
390
+ 3,
391
+ self.conv_size,
392
+ self.local_projection_size,
393
+ dtype=self.conv1d.weight.dtype,
394
+ device=self.conv1d.weight.device,
395
+ )
396
+ self.register_buffer(
397
+ "decode_conv1d_weight", decode_conv1d_weight, persistent=False
398
+ )
399
+ delattr(self.conv1d.weight, "weight_loader")
400
+ set_weight_attrs(
401
+ self.conv1d.weight,
402
+ {
403
+ "weight_loader": _make_decode_conv1d_weight_loader(
404
+ [self.projection_size] * 3,
405
+ self.tp_size,
406
+ self.tp_rank,
407
+ decode_conv1d_weight,
408
+ )
409
+ },
410
+ )
411
+
412
+ self.A_log = nn.Parameter(
413
+ torch.empty(self.local_num_heads, dtype=torch.float32)
414
+ )
415
+ set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)})
416
+
417
+ self.gate_lower_bound: float | None = kda_config.get("gate_lower_bound", None)
418
+ if self.gate_lower_bound is not None:
419
+ assert _KDA_GATE_LOGBOUND_MIN <= self.gate_lower_bound < 0, (
420
+ "KDA gate lower bound must be in "
421
+ f"[{_KDA_GATE_LOGBOUND_MIN}, 0). "
422
+ f"Got {self.gate_lower_bound}."
423
+ )
424
+
425
+ additional_config = vllm_config.additional_config
426
+ backend = (
427
+ additional_config.get("kda_prefill_backend", "auto")
428
+ if isinstance(additional_config, dict)
429
+ else "auto"
430
+ )
431
+ self.kda_prefill_backend = resolve_kda_prefill_backend(
432
+ backend,
433
+ self.head_dim,
434
+ vllm_config.model_config.dtype,
435
+ self.gate_lower_bound,
436
+ )
437
+
438
+ self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
439
+ decode_norm_weight = None
440
+ if decode_conv1d_weight is not None:
441
+ decode_norm_weight = torch.empty(
442
+ self.head_dim,
443
+ dtype=torch.float32,
444
+ device=self.o_norm.weight.device,
445
+ )
446
+ self.register_buffer("decode_norm_weight", decode_norm_weight, persistent=False)
447
+ if decode_norm_weight is not None:
448
+ # Upcast once while loading; direct BF16 norm weights slow the
449
+ # fully fused decode kernel.
450
+ if hasattr(self.o_norm.weight, "weight_loader"):
451
+ delattr(self.o_norm.weight, "weight_loader")
452
+ set_weight_attrs(
453
+ self.o_norm.weight,
454
+ {"weight_loader": _make_decode_norm_weight_loader(decode_norm_weight)},
455
+ )
456
+ self.o_proj = RowParallelLinear(
457
+ self.projection_size,
458
+ self.hidden_size,
459
+ bias=False,
460
+ quant_config=self.quant_config,
461
+ prefix=f"{prefix}.o_proj",
462
+ )
463
+
464
+ compilation_config = vllm_config.compilation_config
465
+ if prefix in compilation_config.static_forward_context:
466
+ raise ValueError(f"Duplicate layer name: {prefix}")
467
+ compilation_config.static_forward_context[prefix] = self
468
+
469
+ def forward(
470
+ self,
471
+ hidden_states: torch.Tensor,
472
+ positions: torch.Tensor,
473
+ ) -> torch.Tensor:
474
+ num_tokens = hidden_states.size(0)
475
+ projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]
476
+ split_sizes = [
477
+ 3 * self.local_projection_size,
478
+ self.local_projection_size,
479
+ self.head_dim,
480
+ self.local_num_heads,
481
+ ]
482
+ if self.in_proj_padding:
483
+ split_sizes.append(self.in_proj_padding)
484
+ projected = projected_qkvgfab.split(split_sizes, dim=-1)
485
+ mixed_qkv, g_proj_states, f_a, beta = projected[:4]
486
+
487
+ g1 = self.f_b_proj(f_a)[0]
488
+ beta = beta.unsqueeze(0)
489
+ g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim)
490
+ g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)
491
+ core_attn_out = torch.empty(
492
+ (1, num_tokens, self.local_num_heads, self.head_dim),
493
+ dtype=hidden_states.dtype,
494
+ device=hidden_states.device,
495
+ )
496
+ self._forward(
497
+ mixed_qkv=mixed_qkv,
498
+ g1=g1,
499
+ g2=g2,
500
+ beta=beta,
501
+ core_attn_out=core_attn_out,
502
+ )
503
+ core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
504
+ return self.o_proj(core_attn_out)[0]
505
+
506
+ @eager_break_during_capture
507
+ def _forward(
508
+ self,
509
+ mixed_qkv: torch.Tensor,
510
+ g1: torch.Tensor,
511
+ g2: torch.Tensor,
512
+ beta: torch.Tensor,
513
+ core_attn_out: torch.Tensor,
514
+ ) -> None:
515
+ forward_context = get_forward_context()
516
+ attn_metadata_raw = forward_context.attn_metadata
517
+ if attn_metadata_raw is None:
518
+ return
519
+
520
+ from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (
521
+ chunk_kda_with_fused_gate,
522
+ fused_recurrent_kda,
523
+ fused_recurrent_kda_packed_decode,
524
+ )
525
+
526
+ assert isinstance(attn_metadata_raw, dict)
527
+ attn_metadata_narrowed = attn_metadata_raw[self.prefix]
528
+ assert isinstance(attn_metadata_narrowed, KimiK3KDAMetadata)
529
+ m = attn_metadata_narrowed
530
+ has_initial_state = m.has_initial_state
531
+ non_spec_query_start_loc = m.non_spec_query_start_loc
532
+ non_spec_state_indices_tensor = m.non_spec_state_indices_tensor
533
+ spec_token_indx = m.spec_token_indx
534
+ non_spec_token_indx = m.non_spec_token_indx
535
+ spec_state_indices_tensor = m.spec_state_indices_tensor
536
+ spec_query_start_loc = m.spec_query_start_loc
537
+ num_accepted_tokens = m.num_accepted_tokens
538
+ num_actual_tokens = m.num_actual_tokens
539
+ has_spec_decode = m.num_spec_decodes > 0
540
+ mixed_qkv = mixed_qkv[:num_actual_tokens]
541
+ g1 = g1[:, :num_actual_tokens]
542
+ beta = beta[:, :num_actual_tokens]
543
+
544
+ conv_state, recurrent_state = self.kv_cache
545
+ # The convolution kernels consume (..., dim, width - 1).
546
+ if not is_conv_state_dim_first():
547
+ conv_state = conv_state.transpose(-1, -2)
548
+
549
+ if (
550
+ self.decode_conv1d_weight is not None
551
+ and self.decode_norm_weight is not None
552
+ and not has_spec_decode
553
+ and m.num_prefills == 0
554
+ and m.num_decodes > 0
555
+ ):
556
+ assert non_spec_state_indices_tensor is not None
557
+ ops.fused_kda_decode(
558
+ x=mixed_qkv,
559
+ weight=self.decode_conv1d_weight,
560
+ bias=self.conv1d.bias,
561
+ conv_state=conv_state,
562
+ raw_g=g1,
563
+ raw_beta=beta,
564
+ A_log=self.A_log,
565
+ dt_bias=self.dt_bias,
566
+ state_indices=non_spec_state_indices_tensor[:num_actual_tokens],
567
+ state=recurrent_state,
568
+ out=core_attn_out[:, :num_actual_tokens],
569
+ lower_bound=self.gate_lower_bound,
570
+ output_gate=g2[:num_actual_tokens],
571
+ norm_weight=self.decode_norm_weight,
572
+ norm_eps=self.o_norm.eps,
573
+ )
574
+ return
575
+
576
+ conv_weights = self.conv1d.weight.view(
577
+ self.conv1d.weight.size(0), self.conv1d.weight.size(2)
578
+ )
579
+ q_conv_weight, k_conv_weight, v_conv_weight = conv_weights.split(
580
+ self.local_projection_size, dim=0
581
+ )
582
+ q_conv_state, k_conv_state, v_conv_state = conv_state.split(
583
+ self.local_projection_size, dim=-2
584
+ )
585
+
586
+ # Separate multi-query speculative tokens from prefill/plain decode.
587
+ if has_spec_decode:
588
+ if m.num_prefills == 0 and m.num_decodes == 0:
589
+ mixed_qkv_spec = mixed_qkv
590
+ g1_spec, beta_spec = g1, beta
591
+ mixed_qkv_ns = g1_ns = beta_ns = None
592
+ else:
593
+ assert spec_token_indx is not None
594
+ assert non_spec_token_indx is not None
595
+ mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)
596
+ g1_spec = g1.index_select(1, spec_token_indx)
597
+ beta_spec = beta.index_select(1, spec_token_indx)
598
+ mixed_qkv_ns = mixed_qkv.index_select(0, non_spec_token_indx)
599
+ g1_ns = g1.index_select(1, non_spec_token_indx)
600
+ beta_ns = beta.index_select(1, non_spec_token_indx)
601
+ else:
602
+ mixed_qkv_spec = g1_spec = beta_spec = None
603
+ mixed_qkv_ns, g1_ns, beta_ns = mixed_qkv, g1, beta
604
+
605
+ # Spec-decode multi-query path.
606
+ core_attn_out_spec = None
607
+ if has_spec_decode:
608
+ assert spec_state_indices_tensor is not None
609
+ assert spec_query_start_loc is not None
610
+ spec_conv_indices = spec_state_indices_tensor[:, 0][: m.num_spec_decodes]
611
+ spec_max_query_len = spec_state_indices_tensor.size(-1)
612
+ spec_conv_out = torch.empty_like(mixed_qkv_spec)
613
+ mixed_qkv_spec = causal_conv1d_update(
614
+ mixed_qkv_spec,
615
+ conv_state,
616
+ conv_weights,
617
+ self.conv1d.bias,
618
+ activation="silu",
619
+ conv_state_indices=spec_conv_indices,
620
+ num_accepted_tokens=num_accepted_tokens,
621
+ query_start_loc=spec_query_start_loc,
622
+ max_query_len=spec_max_query_len,
623
+ validate_data=False,
624
+ out=spec_conv_out,
625
+ )
626
+ q_spec, k_spec, v_spec = (
627
+ rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim)
628
+ for x in mixed_qkv_spec.split(self.local_projection_size, dim=-1)
629
+ )
630
+ spec_cu_seqlens = spec_query_start_loc[: m.num_spec_decodes + 1]
631
+ spec_out = (
632
+ core_attn_out[:, : q_spec.shape[1]]
633
+ if m.num_prefills == 0 and m.num_decodes == 0
634
+ else None
635
+ )
636
+ core_attn_out_spec, _ = fused_recurrent_kda(
637
+ q=q_spec,
638
+ k=k_spec,
639
+ v=v_spec,
640
+ raw_g=g1_spec,
641
+ raw_beta=beta_spec,
642
+ A_log=self.A_log,
643
+ dt_bias=self.dt_bias,
644
+ lower_bound=self.gate_lower_bound,
645
+ initial_state=recurrent_state,
646
+ cu_seqlens=spec_cu_seqlens,
647
+ ssm_state_indices=spec_state_indices_tensor,
648
+ num_accepted_tokens=num_accepted_tokens,
649
+ out=spec_out,
650
+ )
651
+
652
+ # Prefill or plain-decode path.
653
+ core_attn_out_non_spec = None
654
+ if mixed_qkv_ns is not None:
655
+ assert g1_ns is not None and beta_ns is not None
656
+ if m.num_prefills > 0:
657
+ q_ns, k_ns, v_ns = mixed_qkv_ns.split(
658
+ self.local_projection_size, dim=-1
659
+ )
660
+
661
+ # Separate convolution calls accept row-strided packed inputs
662
+ # and produce dense Q/K/V without an additional V copy.
663
+ def _prefill_conv(
664
+ x: torch.Tensor,
665
+ state: torch.Tensor,
666
+ weight: torch.Tensor,
667
+ ) -> torch.Tensor:
668
+ return causal_conv1d_fn(
669
+ x.transpose(0, 1),
670
+ weight,
671
+ None,
672
+ activation="silu",
673
+ conv_states=state,
674
+ has_initial_state=has_initial_state,
675
+ cache_indices=non_spec_state_indices_tensor,
676
+ query_start_loc=non_spec_query_start_loc,
677
+ metadata=m,
678
+ ).transpose(0, 1)
679
+
680
+ q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight)
681
+ k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight)
682
+ v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight)
683
+ q_ns, k_ns, v_ns = (
684
+ rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim)
685
+ for x in (q_ns, k_ns, v_ns)
686
+ )
687
+
688
+ assert non_spec_state_indices_tensor is not None
689
+ assert has_initial_state is not None
690
+ initial_state = gather_initial_states(
691
+ recurrent_state,
692
+ non_spec_state_indices_tensor,
693
+ has_initial_state,
694
+ )
695
+ if self.kda_prefill_backend == "flashkda":
696
+ assert self.gate_lower_bound is not None
697
+ (
698
+ core_attn_out_non_spec,
699
+ last_recurrent_state,
700
+ ) = _flashkda_prefill(
701
+ q=q_ns,
702
+ k=k_ns,
703
+ v=v_ns,
704
+ g=g1_ns,
705
+ beta=beta_ns,
706
+ A_log=self.A_log,
707
+ dt_bias=self.dt_bias,
708
+ lower_bound=self.gate_lower_bound,
709
+ initial_state=initial_state,
710
+ cu_seqlens=non_spec_query_start_loc,
711
+ )
712
+ else:
713
+ (
714
+ core_attn_out_non_spec,
715
+ last_recurrent_state,
716
+ ) = chunk_kda_with_fused_gate(
717
+ q=q_ns,
718
+ k=k_ns,
719
+ v=v_ns,
720
+ raw_g=g1_ns,
721
+ raw_beta=beta_ns,
722
+ A_log=self.A_log,
723
+ g_bias=self.dt_bias,
724
+ lower_bound=self.gate_lower_bound,
725
+ initial_state=initial_state,
726
+ output_final_state=True,
727
+ use_qk_l2norm_in_kernel=True,
728
+ cu_seqlens=non_spec_query_start_loc,
729
+ )
730
+ recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state
731
+ else:
732
+ # Pure non-speculative decode.
733
+ assert non_spec_state_indices_tensor is not None
734
+ decode_conv_indices = non_spec_state_indices_tensor[
735
+ : mixed_qkv_ns.size(0)
736
+ ]
737
+ packed_conv_out = torch.empty_like(mixed_qkv_ns)
738
+ mixed_qkv_ns = causal_conv1d_update(
739
+ mixed_qkv_ns,
740
+ conv_state,
741
+ conv_weights,
742
+ self.conv1d.bias,
743
+ activation="silu",
744
+ conv_state_indices=decode_conv_indices,
745
+ validate_data=True,
746
+ out=packed_conv_out,
747
+ )
748
+ (
749
+ core_attn_out_non_spec,
750
+ _,
751
+ ) = fused_recurrent_kda_packed_decode(
752
+ mixed_qkv=mixed_qkv_ns,
753
+ raw_g=g1_ns,
754
+ raw_beta=beta_ns,
755
+ A_log=self.A_log,
756
+ dt_bias=self.dt_bias,
757
+ lower_bound=self.gate_lower_bound,
758
+ initial_state=recurrent_state,
759
+ state_indices=decode_conv_indices,
760
+ )
761
+
762
+ # Restore the scheduler's original token order for mixed batches.
763
+ if core_attn_out_spec is not None and core_attn_out_non_spec is not None:
764
+ core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec)
765
+ core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
766
+ elif core_attn_out_non_spec is not None:
767
+ # TODO: prefill and decode kernels write directly to core_attn_out
768
+ core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
769
+ 0, :num_actual_tokens
770
+ ]
771
+ else:
772
+ assert core_attn_out_spec is not None
773
+ # Triton normalizes in place, so this is a self-copy with no device
774
+ # work. Keep it for the out-of-place native implementation.
775
+ core_attn_out.copy_(self.o_norm(core_attn_out, g2))
vllm/models/kimi_k3/nvidia/kda_metadata.py
new
+496
@@ -0,0 +1,496 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Kimi-K3 specialization of GDN attention metadata.
4
+
5
+The request classification and cudagraph staging intentionally mirror
6
+``GDNAttentionMetadataBuilder``. Kimi-K3 builds the metadata required by its
7
+prefill KDA kernel internally, so this builder omits the shared FLA chunk
8
+metadata construction.
9
+"""
10
+
11
+from dataclasses import dataclass
12
+from functools import cache
13
+
14
+import torch
15
+
16
+from vllm.platforms import current_platform
17
+from vllm.triton_utils import tl, triton
18
+from vllm.utils.torch_utils import async_tensor_h2d
19
+from vllm.v1.attention.backend import CommonAttentionMetadata
20
+from vllm.v1.attention.backends.gdn_attn import (
21
+ GDNAttentionBackend,
22
+ GDNAttentionMetadata,
23
+ GDNAttentionMetadataBuilder,
24
+)
25
+from vllm.v1.attention.backends.utils import (
26
+ NULL_BLOCK_ID,
27
+ compute_causal_conv1d_metadata,
28
+ split_decodes_and_prefills,
29
+)
30
+from vllm.v1.kv_cache_interface import MambaSpec
31
+
32
+
33
+@cache
34
+def _metadata_launch_pdl() -> bool:
35
+ return current_platform.is_arch_support_pdl()
36
+
37
+
38
+@triton.jit(do_not_specialize=["num_requests"])
39
+def _get_aligned_state_indices_kernel(
40
+ block_table_ptr,
41
+ seq_lens_ptr,
42
+ state_indices_ptr,
43
+ block_table_stride_0: tl.constexpr,
44
+ block_table_stride_1: tl.constexpr,
45
+ seq_lens_stride: tl.constexpr,
46
+ state_indices_stride_0: tl.constexpr,
47
+ state_indices_stride_1: tl.constexpr,
48
+ num_requests,
49
+ CACHE_BLOCK_SIZE: tl.constexpr,
50
+ NUM_STATE_SLOTS: tl.constexpr,
51
+ BLOCK_STATE_SLOTS: tl.constexpr,
52
+ BLOCK_ROWS: tl.constexpr,
53
+ launch_pdl: tl.constexpr,
54
+):
55
+ if launch_pdl:
56
+ tl.extra.cuda.gdc_wait()
57
+ tl.extra.cuda.gdc_launch_dependents()
58
+
59
+ rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
60
+ valid_row = rows < num_requests
61
+ seq_lens = tl.load(
62
+ seq_lens_ptr + rows * seq_lens_stride,
63
+ mask=valid_row,
64
+ other=1,
65
+ )
66
+ # Triton truncates signed division toward zero, unlike PyTorch floor
67
+ # division. Clamping makes both semantics equivalent for seq_lens <= 0.
68
+ first_state_slot = tl.maximum((seq_lens - 1) // CACHE_BLOCK_SIZE, 0)
69
+
70
+ state_slots = tl.arange(0, BLOCK_STATE_SLOTS)
71
+ valid_state_slot = state_slots < NUM_STATE_SLOTS
72
+ state_indices = tl.load(
73
+ block_table_ptr
74
+ + rows[:, None] * block_table_stride_0
75
+ + (first_state_slot[:, None] + state_slots[None, :]) * block_table_stride_1,
76
+ mask=valid_row[:, None] & valid_state_slot[None, :],
77
+ )
78
+ tl.store(
79
+ state_indices_ptr
80
+ + rows[:, None] * state_indices_stride_0
81
+ + state_slots[None, :] * state_indices_stride_1,
82
+ state_indices,
83
+ mask=valid_row[:, None] & valid_state_slot[None, :],
84
+ )
85
+
86
+
87
+def _mamba_get_block_table_tensor(
88
+ block_table: torch.Tensor,
89
+ seq_lens: torch.Tensor,
90
+ kv_cache_spec: MambaSpec,
91
+ mamba_cache_mode: str,
92
+) -> torch.Tensor:
93
+ if mamba_cache_mode in ("all", "none"):
94
+ return block_table
95
+
96
+ assert block_table.is_cuda and seq_lens.is_cuda
97
+ num_requests = block_table.shape[0]
98
+ num_state_slots = 1 + kv_cache_spec.num_speculative_blocks
99
+ state_indices = torch.empty(
100
+ (num_requests, num_state_slots),
101
+ dtype=block_table.dtype,
102
+ device=block_table.device,
103
+ )
104
+ BLOCK_ROWS = 32
105
+ grid = (triton.cdiv(num_requests, BLOCK_ROWS),)
106
+ _get_aligned_state_indices_kernel[grid](
107
+ block_table,
108
+ seq_lens,
109
+ state_indices,
110
+ block_table.stride(0),
111
+ block_table.stride(1),
112
+ seq_lens.stride(0),
113
+ state_indices.stride(0),
114
+ state_indices.stride(1),
115
+ num_requests,
116
+ CACHE_BLOCK_SIZE=kv_cache_spec.block_size,
117
+ NUM_STATE_SLOTS=num_state_slots,
118
+ BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots),
119
+ BLOCK_ROWS=BLOCK_ROWS,
120
+ num_warps=1,
121
+ launch_pdl=_metadata_launch_pdl(),
122
+ )
123
+ return state_indices
124
+
125
+
126
+@triton.jit(do_not_specialize=["num_spec_decodes", "batch_size"])
127
+def _stage_spec_decode_metadata_kernel(
128
+ state_indices_ptr,
129
+ query_start_loc_ptr,
130
+ num_accepted_tokens_ptr,
131
+ staged_state_indices_ptr,
132
+ staged_query_start_loc_ptr,
133
+ staged_num_accepted_tokens_ptr,
134
+ state_indices_stride_0: tl.constexpr,
135
+ state_indices_stride_1: tl.constexpr,
136
+ staged_state_indices_stride_0: tl.constexpr,
137
+ staged_state_indices_stride_1: tl.constexpr,
138
+ num_spec_decodes,
139
+ batch_size,
140
+ NUM_STATE_SLOTS: tl.constexpr,
141
+ BLOCK_STATE_SLOTS: tl.constexpr,
142
+ NULL_STATE_ID: tl.constexpr,
143
+ BLOCK_ROWS: tl.constexpr,
144
+ launch_pdl: tl.constexpr,
145
+):
146
+ if launch_pdl:
147
+ tl.extra.cuda.gdc_wait()
148
+ tl.extra.cuda.gdc_launch_dependents()
149
+
150
+ rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
151
+ real_request = rows < num_spec_decodes
152
+
153
+ state_slots = tl.arange(0, BLOCK_STATE_SLOTS)
154
+ valid_state_slot = state_slots < NUM_STATE_SLOTS
155
+ state_indices = tl.load(
156
+ state_indices_ptr
157
+ + rows[:, None] * state_indices_stride_0
158
+ + state_slots[None, :] * state_indices_stride_1,
159
+ mask=real_request[:, None] & valid_state_slot[None, :],
160
+ other=NULL_STATE_ID,
161
+ )
162
+ tl.store(
163
+ staged_state_indices_ptr
164
+ + rows[:, None] * staged_state_indices_stride_0
165
+ + state_slots[None, :] * staged_state_indices_stride_1,
166
+ state_indices,
167
+ mask=(rows < batch_size)[:, None] & valid_state_slot[None, :],
168
+ )
169
+
170
+ query_row = tl.minimum(rows, num_spec_decodes)
171
+ query_start_loc = tl.load(
172
+ query_start_loc_ptr + query_row,
173
+ mask=rows <= batch_size,
174
+ )
175
+ tl.store(
176
+ staged_query_start_loc_ptr + rows,
177
+ query_start_loc,
178
+ mask=rows <= batch_size,
179
+ )
180
+
181
+ num_accepted_tokens = tl.load(
182
+ num_accepted_tokens_ptr + rows,
183
+ mask=real_request,
184
+ other=1,
185
+ )
186
+ tl.store(
187
+ staged_num_accepted_tokens_ptr + rows,
188
+ num_accepted_tokens,
189
+ mask=rows < batch_size,
190
+ )
191
+
192
+
193
+def stage_spec_decode_metadata(
194
+ state_indices: torch.Tensor,
195
+ query_start_loc: torch.Tensor,
196
+ num_accepted_tokens: torch.Tensor,
197
+ staged_state_indices: torch.Tensor,
198
+ staged_query_start_loc: torch.Tensor,
199
+ staged_num_accepted_tokens: torch.Tensor,
200
+ *,
201
+ num_spec_decodes: int,
202
+) -> None:
203
+ """Stage speculative-decode metadata into CUDA-graph buffers."""
204
+ assert state_indices.is_cuda
205
+ assert state_indices.ndim == 2
206
+ batch_size, num_state_slots = staged_state_indices.shape
207
+ BLOCK_ROWS = 32
208
+ grid = (triton.cdiv(batch_size + 1, BLOCK_ROWS),)
209
+ _stage_spec_decode_metadata_kernel[grid](
210
+ state_indices,
211
+ query_start_loc,
212
+ num_accepted_tokens,
213
+ staged_state_indices,
214
+ staged_query_start_loc,
215
+ staged_num_accepted_tokens,
216
+ state_indices.stride(0),
217
+ state_indices.stride(1),
218
+ staged_state_indices.stride(0),
219
+ staged_state_indices.stride(1),
220
+ num_spec_decodes,
221
+ batch_size,
222
+ NUM_STATE_SLOTS=num_state_slots,
223
+ BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots),
224
+ NULL_STATE_ID=NULL_BLOCK_ID,
225
+ BLOCK_ROWS=BLOCK_ROWS,
226
+ num_warps=1,
227
+ launch_pdl=_metadata_launch_pdl(),
228
+ )
229
+
230
+
231
+@dataclass
232
+class KimiK3KDAMetadata(GDNAttentionMetadata):
233
+ pass
234
+
235
+
236
+class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
237
+ def build( # type: ignore[override]
238
+ self,
239
+ common_prefix_len: int,
240
+ common_attn_metadata: CommonAttentionMetadata,
241
+ num_accepted_tokens: torch.Tensor | None = None,
242
+ num_decode_draft_tokens_cpu: torch.Tensor | None = None,
243
+ fast_build: bool = False,
244
+ ) -> KimiK3KDAMetadata:
245
+ m = common_attn_metadata
246
+ query_start_loc = m.query_start_loc
247
+ query_start_loc_cpu = m.query_start_loc_cpu
248
+ assert isinstance(self.kv_cache_spec, MambaSpec)
249
+ # Equivalent PyTorch "align" path:
250
+ # start = ((seq_lens - 1) // block_size).clamp_(min=0)
251
+ # offsets = torch.arange(1 + num_speculative_blocks, dtype=torch.int32)
252
+ # indices = (start[:, None] + offsets).to(torch.int64)
253
+ # block_table_tensor = torch.gather(block_table, 1, indices)
254
+ block_table_tensor = _mamba_get_block_table_tensor(
255
+ m.block_table_tensor,
256
+ m.seq_lens,
257
+ self.kv_cache_spec,
258
+ self.vllm_config.cache_config.mamba_cache_mode,
259
+ )
260
+
261
+ if not self.use_spec_decode or num_decode_draft_tokens_cpu is None:
262
+ spec_sequence_masks_cpu = None
263
+ num_spec_decodes = 0
264
+ else:
265
+ spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0
266
+ # A nonnegative entry identifies a spec request. If no draft token
267
+ # was scheduled, process the whole batch as non-spec instead.
268
+ if num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() == 0:
269
+ spec_sequence_masks_cpu = None
270
+ num_spec_decodes = 0
271
+ else:
272
+ num_spec_decodes = spec_sequence_masks_cpu.sum().item()
273
+
274
+ if num_spec_decodes == 0:
275
+ # The runner orders ordinary decodes before prefills.
276
+ num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
277
+ split_decodes_and_prefills(m, decode_threshold=1)
278
+ )
279
+ num_spec_decode_tokens = 0
280
+ spec_token_indx = None
281
+ non_spec_token_indx = None
282
+ spec_state_indices_tensor = None
283
+ non_spec_state_indices_tensor = block_table_tensor[:, 0]
284
+ spec_query_start_loc = None
285
+ non_spec_query_start_loc = query_start_loc
286
+ non_spec_query_start_loc_cpu = query_start_loc_cpu
287
+ num_accepted_tokens = None
288
+ else:
289
+ assert spec_sequence_masks_cpu is not None
290
+ assert num_accepted_tokens is not None
291
+ query_lens_cpu = query_start_loc_cpu.diff()
292
+ num_query_tokens = query_start_loc_cpu[-1].item()
293
+
294
+ # Exclude zero-length cudagraph padding from request-indexed
295
+ # non-spec metadata.
296
+ active_non_spec_mask_cpu = (~spec_sequence_masks_cpu) & (query_lens_cpu > 0)
297
+ non_spec_query_lens_cpu = query_lens_cpu[active_non_spec_mask_cpu]
298
+ num_non_spec_requests = non_spec_query_lens_cpu.numel()
299
+ num_non_spec_tokens = non_spec_query_lens_cpu.sum().item()
300
+
301
+ # Query length alone cannot distinguish a true decode from a
302
+ # one-token prefill chunk. Packed decode is safe only when every
303
+ # active non-spec request is a true, one-token decode.
304
+ assert m.is_prefilling is not None
305
+ assert m.is_prefilling.device.type == "cpu"
306
+ non_spec_is_prefilling = m.is_prefilling[active_non_spec_mask_cpu]
307
+ use_prefill = torch.any(
308
+ non_spec_is_prefilling | (non_spec_query_lens_cpu != 1)
309
+ ).item()
310
+ if use_prefill:
311
+ num_prefills = num_non_spec_requests
312
+ num_prefill_tokens = num_non_spec_tokens
313
+ num_decodes = 0
314
+ num_decode_tokens = 0
315
+ else:
316
+ num_prefills = 0
317
+ num_prefill_tokens = 0
318
+ num_decodes = num_non_spec_requests
319
+ num_decode_tokens = num_non_spec_tokens
320
+
321
+ num_spec_decode_tokens = num_query_tokens - num_non_spec_tokens
322
+
323
+ if num_prefills == 0 and num_decodes == 0:
324
+ spec_token_indx = None
325
+ non_spec_token_indx = None
326
+ # Real requests precede trailing cudagraph padding.
327
+ spec_state_indices_tensor = block_table_tensor[
328
+ :num_spec_decodes, : self.num_spec + 1
329
+ ]
330
+ non_spec_state_indices_tensor = None
331
+ # Padding trails real requests, so this prefix already contains
332
+ # the correct cumulative token counts.
333
+ spec_query_start_loc = query_start_loc[: num_spec_decodes + 1]
334
+ non_spec_query_start_loc = None
335
+ non_spec_query_start_loc_cpu = None
336
+ num_accepted_tokens = num_accepted_tokens[:num_spec_decodes]
337
+ else:
338
+ query_lens = query_start_loc.diff()
339
+ spec_sequence_masks_gpu = async_tensor_h2d(
340
+ spec_sequence_masks_cpu, device=query_start_loc.device
341
+ )
342
+ spec_token_masks = torch.repeat_interleave(
343
+ spec_sequence_masks_gpu,
344
+ query_lens,
345
+ output_size=num_query_tokens,
346
+ )
347
+ # Stable partitioning preserves request-local token order in
348
+ # both subgroup tensors.
349
+ index = torch.argsort(spec_token_masks, stable=True)
350
+ num_non_spec_tokens = num_prefill_tokens + num_decode_tokens
351
+ non_spec_token_indx = index[:num_non_spec_tokens]
352
+ spec_token_indx = index[num_non_spec_tokens:]
353
+
354
+ # Spec requests carry one state slot per speculative step;
355
+ # non-spec requests use only their current state slot.
356
+ spec_state_indices_tensor = block_table_tensor[
357
+ spec_sequence_masks_cpu, : self.num_spec + 1
358
+ ]
359
+ non_spec_state_indices_tensor = block_table_tensor[
360
+ active_non_spec_mask_cpu, 0
361
+ ]
362
+
363
+ spec_query_lens = query_lens[spec_sequence_masks_cpu]
364
+ spec_query_start_loc = torch.zeros(
365
+ num_spec_decodes + 1,
366
+ dtype=torch.int32,
367
+ device=query_start_loc.device,
368
+ )
369
+ torch.cumsum(
370
+ spec_query_lens,
371
+ dim=0,
372
+ out=spec_query_start_loc[1:],
373
+ )
374
+ if num_prefills > 0:
375
+ non_spec_query_lens = query_lens[active_non_spec_mask_cpu]
376
+ non_spec_query_start_loc = torch.zeros(
377
+ non_spec_query_lens.size(0) + 1,
378
+ dtype=torch.int32,
379
+ device=query_start_loc.device,
380
+ )
381
+ torch.cumsum(
382
+ non_spec_query_lens,
383
+ dim=0,
384
+ out=non_spec_query_start_loc[1:],
385
+ )
386
+ non_spec_query_start_loc_cpu = torch.zeros(
387
+ non_spec_query_lens_cpu.size(0) + 1,
388
+ dtype=torch.int32,
389
+ )
390
+ torch.cumsum(
391
+ non_spec_query_lens_cpu,
392
+ dim=0,
393
+ out=non_spec_query_start_loc_cpu[1:],
394
+ )
395
+ else:
396
+ # Packed decode consumes one row per request and does not
397
+ # use cumulative sequence lengths.
398
+ non_spec_query_start_loc = None
399
+ non_spec_query_start_loc_cpu = None
400
+
401
+ num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu]
402
+
403
+ # Unlike the shared GDN layer, Kimi-K3's prefill KDA wrapper prepares
404
+ # its own chunk indices. Only causal-convolution metadata is needed here.
405
+ nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
406
+ if num_prefills > 0:
407
+ has_initial_state = m.compute_num_computed_tokens() > 0
408
+ if spec_sequence_masks_cpu is not None:
409
+ has_initial_state = has_initial_state[active_non_spec_mask_cpu]
410
+ assert non_spec_query_start_loc_cpu is not None
411
+ nums_dict, batch_ptr, token_chunk_offset_ptr = (
412
+ compute_causal_conv1d_metadata(
413
+ non_spec_query_start_loc_cpu,
414
+ device=query_start_loc.device,
415
+ )
416
+ )
417
+ else:
418
+ has_initial_state = None
419
+
420
+ # Prepare per-request tensors for cudagraph replay. num_actual_tokens
421
+ # may be token-padded, while state/query/acceptance metadata is indexed
422
+ # by request.
423
+ batch_size = m.num_reqs
424
+ if (
425
+ self.use_full_cuda_graph
426
+ and num_spec_decodes > 0
427
+ and num_prefills == 0
428
+ and num_decodes == 0
429
+ and num_spec_decodes <= self.decode_cudagraph_max_bs
430
+ and num_spec_decode_tokens <= self.decode_cudagraph_max_bs
431
+ ):
432
+ # Equivalent PyTorch staging:
433
+ # state[:N].copy_(state_src); state[N:].fill_(NULL_BLOCK_ID)
434
+ # qsl[:N + 1].copy_(qsl_src); qsl[N + 1:].fill_(qsl_src[-1])
435
+ # accepted[:N].copy_(accepted_src); accepted[N:].fill_(1)
436
+ stage_spec_decode_metadata(
437
+ state_indices=spec_state_indices_tensor,
438
+ query_start_loc=spec_query_start_loc,
439
+ num_accepted_tokens=num_accepted_tokens,
440
+ staged_state_indices=self.spec_state_indices_tensor[:batch_size],
441
+ staged_query_start_loc=self.spec_query_start_loc[: batch_size + 1],
442
+ staged_num_accepted_tokens=self.num_accepted_tokens[:batch_size],
443
+ num_spec_decodes=num_spec_decodes,
444
+ )
445
+
446
+ spec_state_indices_tensor = self.spec_state_indices_tensor[:batch_size]
447
+ spec_query_start_loc = self.spec_query_start_loc[: batch_size + 1]
448
+ num_accepted_tokens = self.num_accepted_tokens[:batch_size]
449
+
450
+ if (
451
+ self.use_full_cuda_graph
452
+ and num_prefills == 0
453
+ and num_spec_decodes == 0
454
+ and num_decodes <= self.decode_cudagraph_max_bs
455
+ ):
456
+ self.non_spec_state_indices_tensor[:num_decodes].copy_(
457
+ non_spec_state_indices_tensor, non_blocking=True
458
+ )
459
+ self.non_spec_state_indices_tensor[num_decodes:batch_size].fill_(
460
+ NULL_BLOCK_ID
461
+ )
462
+ non_spec_state_indices_tensor = self.non_spec_state_indices_tensor[
463
+ :batch_size
464
+ ]
465
+
466
+ return KimiK3KDAMetadata(
467
+ num_prefills=num_prefills,
468
+ num_prefill_tokens=num_prefill_tokens,
469
+ num_decodes=num_decodes,
470
+ num_decode_tokens=num_decode_tokens,
471
+ num_spec_decodes=num_spec_decodes,
472
+ num_spec_decode_tokens=num_spec_decode_tokens,
473
+ num_actual_tokens=m.num_actual_tokens,
474
+ has_initial_state=has_initial_state,
475
+ spec_query_start_loc=spec_query_start_loc,
476
+ non_spec_query_start_loc=non_spec_query_start_loc,
477
+ spec_state_indices_tensor=spec_state_indices_tensor,
478
+ non_spec_state_indices_tensor=non_spec_state_indices_tensor,
479
+ spec_sequence_masks=None,
480
+ spec_token_indx=spec_token_indx,
481
+ non_spec_token_indx=non_spec_token_indx,
482
+ num_accepted_tokens=num_accepted_tokens,
483
+ nums_dict=nums_dict,
484
+ batch_ptr=batch_ptr,
485
+ token_chunk_offset_ptr=token_chunk_offset_ptr,
486
+ )
487
+
488
+
489
+class KimiK3KDAAttentionBackend(GDNAttentionBackend):
490
+ @staticmethod
491
+ def get_name() -> str:
492
+ return "KIMI_K3_KDA"
493
+
494
+ @staticmethod
495
+ def get_builder_cls() -> type[KimiK3KDAMetadataBuilder]:
496
+ return KimiK3KDAMetadataBuilder
vllm/models/kimi_k3/nvidia/low_latency_gemm.py
new
+513
@@ -0,0 +1,513 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Kimi-K3 decode GEMM selection for unquantized BF16 on SM103.
4
+
5
+Dispatch is purely by local ``(N, K)`` shape and token count ``M`` — the module
6
+name plays no role. Each measured shape maps to a :class:`ProjectionSpec`
7
+holding the winning backend per token count. The static part of the decision is
8
+resolved once per module at install time into a small ``{M: call}`` plan, so the
9
+per-forward path is a single dict lookup.
10
+"""
11
+
12
+from __future__ import annotations
13
+
14
+from dataclasses import dataclass
15
+from typing import Literal
16
+
17
+import torch
18
+from torch import nn
19
+
20
+import vllm.envs as envs
21
+from vllm import _custom_ops as ops
22
+from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
23
+ SkinnyGemmConfig,
24
+ shape_dynamic_skinny_gemm,
25
+)
26
+from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
27
+from vllm.model_executor.layers.vocab_parallel_embedding import (
28
+ ParallelLMHead,
29
+ UnquantizedEmbeddingMethod,
30
+)
31
+from vllm.platforms import current_platform
32
+
33
+Backend = Literal["cute", "dsv3_fused_a"]
34
+# A resolved per-token-count call: the backend plus its CuTe config (None for
35
+# dsv3, which needs no config).
36
+ResolvedCall = tuple[Backend, SkinnyGemmConfig | None]
37
+
38
+
39
+@dataclass(frozen=True, slots=True)
40
+class ProjectionSpec:
41
+ n: int
42
+ k: int
43
+ dsv3_tokens: frozenset[int] = frozenset()
44
+ cute_configs: tuple[tuple[int, SkinnyGemmConfig], ...] = ()
45
+ residual_configs: tuple[tuple[int, SkinnyGemmConfig], ...] = ()
46
+ name: str = "" # optional debug label; never used for dispatch
47
+
48
+ def cute_config(self, num_tokens: int) -> SkinnyGemmConfig | None:
49
+ return dict(self.cute_configs).get(num_tokens)
50
+
51
+ def residual_config(self, num_tokens: int) -> SkinnyGemmConfig | None:
52
+ return dict(self.residual_configs).get(num_tokens)
53
+
54
+
55
+def _cute(
56
+ num_tokens: int,
57
+ block_size: int,
58
+ outputs_per_block: int,
59
+ k_unroll: int,
60
+ vector_width: int = 8,
61
+) -> SkinnyGemmConfig:
62
+ return SkinnyGemmConfig(
63
+ num_tokens,
64
+ block_size,
65
+ outputs_per_block,
66
+ k_unroll,
67
+ vector_width,
68
+ )
69
+
70
+
71
+_M1_TO_16 = frozenset(range(1, 17))
72
+_M1 = frozenset({1})
73
+
74
+# Keyed by local (N, K). Where two projections share a shape (only 1536x7168:
75
+# shared_gate_up_proj and mla_g_proj) the entry is unified.
76
+KIMI_K3_PROJECTIONS: dict[tuple[int, int], ProjectionSpec] = {
77
+ (1536, 128): ProjectionSpec(1536, 128, _M1_TO_16, name="f_b_proj"),
78
+ (3072, 128): ProjectionSpec(3072, 128, _M1_TO_16, name="f_b_proj"),
79
+ # 1536x7168 is shared by shared_gate_up_proj and mla_g_proj. dsv3 M1..16 is
80
+ # only crash-safe once the mla_g aux-stream/PDL capture fix lands (subtask
81
+ # task_7388aba1); the fallback if it cannot be fixed is dsv3_tokens=_M1.
82
+ (1536, 7168): ProjectionSpec(
83
+ 1536, 7168, _M1_TO_16, name="shared_gate_up_proj/mla_g_proj"
84
+ ),
85
+ (3072, 7168): ProjectionSpec(
86
+ 3072,
87
+ 7168,
88
+ cute_configs=(
89
+ (1, _cute(1, 224, 3, 4)),
90
+ (2, _cute(2, 128, 3, 2)),
91
+ (3, _cute(3, 128, 2, 1)),
92
+ (4, _cute(4, 64, 2, 2)),
93
+ (5, _cute(5, 128, 3, 1)),
94
+ ),
95
+ name="shared_gate_up_proj",
96
+ ),
97
+ (2112, 7168): ProjectionSpec(2112, 7168, _M1_TO_16, name="fused_qkv_a_proj"),
98
+ (2304, 1536): ProjectionSpec(2304, 1536, _M1_TO_16, name="q_b_proj"),
99
+ (4608, 1536): ProjectionSpec(4608, 1536, _M1_TO_16, name="q_b_proj"),
100
+ (3584, 7168): ProjectionSpec(
101
+ 3584,
102
+ 7168,
103
+ frozenset(range(2, 9)),
104
+ ((1, _cute(1, 224, 2, 4)),),
105
+ name="routed_expert_down_proj",
106
+ ),
107
+ (6288, 7168): ProjectionSpec(
108
+ 6288,
109
+ 7168,
110
+ cute_configs=(
111
+ (1, _cute(1, 224, 3, 4)),
112
+ (2, _cute(2, 64, 3, 2)),
113
+ (3, _cute(3, 32, 3, 4)),
114
+ (4, _cute(4, 128, 6, 1)),
115
+ ),
116
+ name="in_proj_qkvgfab",
117
+ ),
118
+ (12448, 7168): ProjectionSpec(
119
+ 12448,
120
+ 7168,
121
+ cute_configs=(
122
+ (1, _cute(1, 224, 4, 2)),
123
+ (2, _cute(2, 64, 4, 2)),
124
+ (3, _cute(3, 64, 2, 2)),
125
+ ),
126
+ name="in_proj_qkvgfab",
127
+ ),
128
+ (7168, 768): ProjectionSpec(7168, 768, _M1_TO_16, name="shared_down_proj"),
129
+ (7168, 1536): ProjectionSpec(
130
+ 7168, 1536, cute_configs=((1, _cute(1, 96, 4, 2)),), name="o_proj"
131
+ ),
132
+ (7168, 3072): ProjectionSpec(
133
+ 7168,
134
+ 3072,
135
+ cute_configs=(
136
+ (1, _cute(1, 96, 2, 4)),
137
+ (2, _cute(2, 32, 4, 4)),
138
+ ),
139
+ name="o_proj",
140
+ ),
141
+ (7168, 3584): ProjectionSpec(
142
+ 7168,
143
+ 3584,
144
+ cute_configs=(
145
+ (1, _cute(1, 224, 4, 2)),
146
+ (2, _cute(2, 64, 4, 2)),
147
+ ),
148
+ residual_configs=(
149
+ (1, _cute(1, 64, 4, 2)),
150
+ (2, _cute(2, 64, 7, 2)),
151
+ (3, _cute(3, 64, 2, 1)),
152
+ (4, _cute(4, 64, 2, 1)),
153
+ ),
154
+ name="routed_expert_up_proj",
155
+ ),
156
+ (7168, 4224): ProjectionSpec(
157
+ 7168,
158
+ 4224,
159
+ cute_configs=((1, _cute(1, 96, 4, 2, 4)),),
160
+ name="dense_down_proj",
161
+ ),
162
+ (7168, 8448): ProjectionSpec(
163
+ 7168,
164
+ 8448,
165
+ cute_configs=(
166
+ (1, _cute(1, 32, 4, 4)),
167
+ (2, _cute(2, 96, 4, 1)),
168
+ (3, _cute(3, 96, 4, 1)),
169
+ ),
170
+ name="dense_down_proj",
171
+ ),
172
+ (8448, 7168): ProjectionSpec(
173
+ 8448,
174
+ 7168,
175
+ cute_configs=(
176
+ (1, _cute(1, 224, 3, 4)),
177
+ (2, _cute(2, 32, 4, 4)),
178
+ ),
179
+ name="dense_gate_up_proj",
180
+ ),
181
+ (16896, 7168): ProjectionSpec(
182
+ 16896,
183
+ 7168,
184
+ cute_configs=(
185
+ (1, _cute(1, 224, 6, 4)),
186
+ (2, _cute(2, 32, 4, 4)),
187
+ ),
188
+ name="dense_gate_up_proj",
189
+ ),
190
+ (20480, 7168): ProjectionSpec(
191
+ 20480,
192
+ 7168,
193
+ cute_configs=(
194
+ (1, _cute(1, 224, 4, 2)),
195
+ (2, _cute(2, 64, 4, 2)),
196
+ (3, _cute(3, 64, 2, 2)),
197
+ (4, _cute(4, 64, 4, 1)),
198
+ ),
199
+ name="lm_head",
200
+ ),
201
+ (40960, 7168): ProjectionSpec(
202
+ 40960,
203
+ 7168,
204
+ cute_configs=(
205
+ (1, _cute(1, 128, 4, 2)),
206
+ (2, _cute(2, 64, 4, 2)),
207
+ (3, _cute(3, 64, 2, 2)),
208
+ (4, _cute(4, 64, 4, 1)),
209
+ ),
210
+ name="lm_head",
211
+ ),
212
+ # TP16. Measured on B300 over M=1..16 with the same >=5% threshold as the
213
+ # entries above. The replicated projections (2112x7168, 3584x7168,
214
+ # 7168x3584) keep their shapes at TP16 and reuse the entries above, and
215
+ # o_proj lands on 7168x768, which shared_down_proj already covers.
216
+ (3216, 7168): ProjectionSpec(
217
+ 3216,
218
+ 7168,
219
+ # Both gaps in this range are measured, not oversights: dsv3 is only
220
+ # 4% ahead at M6..M8, and at M16 cuBLAS switches to a faster kernel
221
+ # (11.42us vs dsv3's 11.83us) after trailing it by 6-8% at M9..M15.
222
+ frozenset(range(9, 16)),
223
+ cute_configs=(
224
+ (1, _cute(1, 224, 3, 4)),
225
+ (2, _cute(2, 128, 4, 2)),
226
+ (3, _cute(3, 128, 2, 1)),
227
+ (4, _cute(4, 64, 2, 2)),
228
+ (5, _cute(5, 128, 3, 1)),
229
+ ),
230
+ name="in_proj_qkvgfab",
231
+ ),
232
+ (768, 7168): ProjectionSpec(
233
+ 768,
234
+ 7168,
235
+ frozenset(range(5, 17)),
236
+ cute_configs=(
237
+ (1, _cute(1, 224, 2, 4)),
238
+ (2, _cute(2, 224, 2, 2)),
239
+ (3, _cute(3, 224, 2, 2)),
240
+ (4, _cute(4, 224, 2, 2)),
241
+ ),
242
+ name="mla_g_proj/shared_gate_up_proj",
243
+ ),
244
+ (1152, 1536): ProjectionSpec(
245
+ 1152,
246
+ 1536,
247
+ frozenset(range(2, 17)),
248
+ ((1, _cute(1, 192, 3, 4)),),
249
+ name="q_b_proj",
250
+ ),
251
+ (768, 128): ProjectionSpec(768, 128, _M1_TO_16, name="f_b_proj"),
252
+ # dsv3 drops under 5% from M9 on for this shape.
253
+ (7168, 384): ProjectionSpec(
254
+ 7168, 384, frozenset(range(1, 9)), name="shared_down_proj"
255
+ ),
256
+ (4224, 7168): ProjectionSpec(
257
+ 4224,
258
+ 7168,
259
+ frozenset(range(4, 9)),
260
+ cute_configs=(
261
+ (1, _cute(1, 224, 3, 4)),
262
+ (2, _cute(2, 128, 2, 1)),
263
+ (3, _cute(3, 64, 2, 2)),
264
+ ),
265
+ name="dense_gate_up_proj",
266
+ ),
267
+ (10240, 7168): ProjectionSpec(
268
+ 10240,
269
+ 7168,
270
+ cute_configs=(
271
+ (1, _cute(1, 224, 4, 2)),
272
+ (2, _cute(2, 32, 2, 4)),
273
+ (3, _cute(3, 64, 4, 1)),
274
+ (4, _cute(4, 64, 4, 1)),
275
+ ),
276
+ name="lm_head",
277
+ ),
278
+ # 7168x2112 (TP16 dense down_proj) has no entry on purpose: K=2112 divides
279
+ # none of the fused-A tile_k values, and the CuTe kernel is left with
280
+ # vector_width=2, which measured slower than cuBLAS.
281
+}
282
+
283
+
284
+def _backend_for(
285
+ spec: ProjectionSpec, num_tokens: int, has_residual: bool
286
+) -> Backend | None:
287
+ if has_residual:
288
+ return "cute" if spec.residual_config(num_tokens) is not None else None
289
+ if spec.cute_config(num_tokens) is not None:
290
+ return "cute"
291
+ if num_tokens in spec.dsv3_tokens:
292
+ return "dsv3_fused_a"
293
+ return None
294
+
295
+
296
+def select_kimi_k3_backend(
297
+ num_tokens: int,
298
+ n: int,
299
+ k: int,
300
+ *,
301
+ has_residual: bool = False,
302
+) -> Backend | None:
303
+ """Backend for a local ``(N, K)`` at ``num_tokens``, or None to fall back."""
304
+ spec = KIMI_K3_PROJECTIONS.get((n, k))
305
+ return _backend_for(spec, num_tokens, has_residual) if spec is not None else None
306
+
307
+
308
+def _build_plan(spec: ProjectionSpec) -> dict[int, ResolvedCall]:
309
+ plan: dict[int, ResolvedCall] = {}
310
+ for num_tokens in range(1, 17):
311
+ backend = _backend_for(spec, num_tokens, has_residual=False)
312
+ if backend == "cute":
313
+ plan[num_tokens] = ("cute", spec.cute_config(num_tokens))
314
+ elif backend == "dsv3_fused_a":
315
+ plan[num_tokens] = ("dsv3_fused_a", None)
316
+ return plan
317
+
318
+
319
+def _build_residual_plan(spec: ProjectionSpec) -> dict[int, SkinnyGemmConfig]:
320
+ return {num_tokens: config for num_tokens, config in spec.residual_configs}
321
+
322
+
323
+def _is_sm103() -> bool:
324
+ return current_platform.is_device_capability((10, 3))
325
+
326
+
327
+def _is_packed_row_major(tensor: torch.Tensor) -> bool:
328
+ return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1)
329
+
330
+
331
+def _runtime_ok(x: torch.Tensor, weight: torch.Tensor) -> bool:
332
+ return (
333
+ _is_packed_row_major(x)
334
+ and _is_packed_row_major(weight)
335
+ and x.dtype == torch.bfloat16
336
+ and weight.dtype == torch.bfloat16
337
+ and x.is_cuda
338
+ and weight.is_cuda
339
+ and x.device == weight.device
340
+ and x.shape[1] == weight.shape[1]
341
+ )
342
+
343
+
344
+def _residual_ok(x: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor) -> bool:
345
+ return (
346
+ residual.dim() == 2
347
+ and residual.dtype == torch.bfloat16
348
+ and residual.is_cuda
349
+ and residual.device == x.device
350
+ and residual.is_contiguous()
351
+ and residual.shape == (x.shape[0], weight.shape[0])
352
+ )
353
+
354
+
355
+def _run_plan(
356
+ plan: dict[int, ResolvedCall], x: torch.Tensor, weight: torch.Tensor
357
+) -> torch.Tensor | None:
358
+ entry = plan.get(x.shape[0])
359
+ if entry is None:
360
+ return None
361
+ backend, config = entry
362
+ if backend == "cute":
363
+ if not shape_dynamic_skinny_gemm.is_available():
364
+ return None
365
+ return shape_dynamic_skinny_gemm(x, weight, config, None)
366
+ if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"):
367
+ return None
368
+ output = torch.empty((x.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device)
369
+ ops.dsv3_fused_a_gemm(output, x, weight.t(), enable_pdl=True)
370
+ return output
371
+
372
+
373
+def _run_residual_plan(
374
+ residual_plan: dict[int, SkinnyGemmConfig],
375
+ x: torch.Tensor,
376
+ weight: torch.Tensor,
377
+ residual: torch.Tensor,
378
+) -> torch.Tensor | None:
379
+ config = residual_plan.get(x.shape[0])
380
+ if config is None or not shape_dynamic_skinny_gemm.is_available():
381
+ return None
382
+ return shape_dynamic_skinny_gemm(x, weight, config, residual)
383
+
384
+
385
+def try_low_latency_gemm(
386
+ x: torch.Tensor,
387
+ weight: torch.Tensor,
388
+ residual: torch.Tensor | None = None,
389
+) -> torch.Tensor | None:
390
+ """Run the shape-selected low-latency kernel, or None to fall back.
391
+
392
+ Resolves the plan from the shape table on each call; production installs a
393
+ precomputed plan (see :func:`enable_kimi_k3_low_latency_gemm`) and does not
394
+ use this path.
395
+ """
396
+ if envs.VLLM_BATCH_INVARIANT or not _is_sm103() or not _runtime_ok(x, weight):
397
+ return None
398
+ spec = KIMI_K3_PROJECTIONS.get((weight.shape[0], weight.shape[1]))
399
+ if spec is None:
400
+ return None
401
+ if residual is None:
402
+ return _run_plan(_build_plan(spec), x, weight)
403
+ if not _residual_ok(x, weight, residual):
404
+ return None
405
+ return _run_residual_plan(_build_residual_plan(spec), x, weight, residual)
406
+
407
+
408
+class _KimiK3LowLatencyApply:
409
+ """Mixin: try the precomputed plan, else defer to the base method."""
410
+
411
+ def __init__(self, plan: dict[int, ResolvedCall]) -> None:
412
+ self._plan = plan
413
+
414
+ def apply(
415
+ self,
416
+ layer: nn.Module,
417
+ x: torch.Tensor,
418
+ bias: torch.Tensor | None = None,
419
+ ) -> torch.Tensor:
420
+ if (
421
+ bias is None
422
+ and not envs.VLLM_BATCH_INVARIANT
423
+ and _runtime_ok(x, layer.weight)
424
+ ):
425
+ output = _run_plan(self._plan, x, layer.weight)
426
+ if output is not None:
427
+ return output
428
+ return super().apply(layer, x, bias) # type: ignore[misc]
429
+
430
+
431
+class KimiK3LowLatencyLinearMethod(_KimiK3LowLatencyApply, UnquantizedLinearMethod):
432
+ def __init__(
433
+ self,
434
+ plan: dict[int, ResolvedCall],
435
+ residual_plan: dict[int, SkinnyGemmConfig],
436
+ ) -> None:
437
+ super().__init__(plan)
438
+ self._residual_plan = residual_plan
439
+
440
+ def apply_with_residual(
441
+ self,
442
+ layer: nn.Module,
443
+ x: torch.Tensor,
444
+ residual: torch.Tensor,
445
+ ) -> torch.Tensor:
446
+ if (
447
+ not envs.VLLM_BATCH_INVARIANT
448
+ and _runtime_ok(x, layer.weight)
449
+ and _residual_ok(x, layer.weight, residual)
450
+ ):
451
+ output = _run_residual_plan(self._residual_plan, x, layer.weight, residual)
452
+ if output is not None:
453
+ return output
454
+ return torch.addmm(residual, x, layer.weight.t())
455
+
456
+
457
+class KimiK3LowLatencyEmbeddingMethod(
458
+ _KimiK3LowLatencyApply, UnquantizedEmbeddingMethod
459
+):
460
+ pass
461
+
462
+
463
+def enable_kimi_k3_low_latency_gemm(
464
+ module: nn.Module,
465
+ dtype: torch.dtype,
466
+) -> None:
467
+ """Install shape-selected low-latency GEMMs and register CuTe warmups.
468
+
469
+ Modules are matched purely by type, an exactly-unquantized method, and a
470
+ local ``(N, K)`` present in :data:`KIMI_K3_PROJECTIONS`.
471
+ """
472
+ if dtype != torch.bfloat16 or not _is_sm103():
473
+ return
474
+
475
+ warmup_configs: set[SkinnyGemmConfig] = set()
476
+ residual_warmup_configs: set[SkinnyGemmConfig] = set()
477
+ for child in module.modules():
478
+ is_linear = (
479
+ isinstance(child, LinearBase)
480
+ and type(child.quant_method) is UnquantizedLinearMethod
481
+ )
482
+ # ParallelLMHead is a VocabParallelEmbedding subclass; embed_tokens is
483
+ # the parent type, so isinstance already excludes it.
484
+ is_head = (
485
+ isinstance(child, ParallelLMHead)
486
+ and type(child.quant_method) is UnquantizedEmbeddingMethod
487
+ )
488
+ if not (is_linear or is_head):
489
+ continue
490
+ weight = getattr(child, "weight", None)
491
+ if weight is None or weight.dim() != 2:
492
+ continue
493
+ spec = KIMI_K3_PROJECTIONS.get((weight.shape[0], weight.shape[1]))
494
+ if spec is None:
495
+ continue
496
+ if is_linear:
497
+ child.quant_method = KimiK3LowLatencyLinearMethod(
498
+ _build_plan(spec), _build_residual_plan(spec)
499
+ )
500
+ else:
501
+ child.quant_method = KimiK3LowLatencyEmbeddingMethod(_build_plan(spec))
502
+ # Warm up only the configs measured for this module's local (N, K) so a
503
+ # TP8 deployment does not compile TP4 configs and vice versa.
504
+ warmup_configs.update(config for _, config in spec.cute_configs)
505
+ residual_warmup_configs.update(config for _, config in spec.residual_configs)
506
+
507
+ if shape_dynamic_skinny_gemm.is_available():
508
+ if warmup_configs:
509
+ shape_dynamic_skinny_gemm.request_warmup_configs(dtype, warmup_configs)
510
+ if residual_warmup_configs:
511
+ shape_dynamic_skinny_gemm.request_warmup_configs(
512
+ dtype, residual_warmup_configs, has_residual=True
513
+ )
vllm/models/kimi_k3/nvidia/mla.py
new
+764
@@ -0,0 +1,764 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Clean Multi-head Latent Attention for Kimi-K3 (NVIDIA).
4
+
5
+This is a self-contained MLA layer that owns the full attention path:
6
+
7
+ hidden_states
8
+ -> fused pre-attention ops (fused_qkv_a_proj / norms / q_b_proj)
9
+ -> explicit prefill / decode split
10
+ prefill: fused key-concat + cache-insert kernel -> run_prefill_new_tokens
11
+ (+ chunked-context merge); dispatched by cache dtype
12
+ (bf16 / plain fp8 / fp8_ds_mla)
13
+ decode : W_UK absorb (BMM1) -> fused q-concat + cache-insert kernel
14
+ -> impl.forward_mqa -> W_UV up-proj (MQA)
15
+ -> optional output gate
16
+ -> o_proj
17
+
18
+Unlike ``MultiHeadLatentAttentionWrapper`` (which delegates orchestration to
19
+``MLAAttention.forward``), this class *is* the ``AttentionLayerBase``: it selects
20
+the backend, builds the impl, registers itself in the forward context, owns the
21
+KV cache, and absorbs ``kv_b_proj`` into ``W_UK_T`` / ``W_UV`` -- mirroring the
22
+``DeepseekV4Attention`` structure.
23
+
24
+K3 specifics: optional rotary embedding (disabled for the target model's NoPE
25
+layers, enabled for DSpark) and an optional sigmoid output gate (``g_proj``).
26
+
27
+Out of scope (extension points, not wired here): context parallelism (DCP/PCP),
28
+sparse/indexer MLA, and the ROCm/aiter fp8/fp4 BMM fast paths.
29
+"""
30
+
31
+import math
32
+from typing import TYPE_CHECKING, cast
33
+
34
+import torch
35
+from torch import nn
36
+
37
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
38
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
39
+from vllm.distributed import get_tensor_model_parallel_world_size
40
+from vllm.forward_context import get_forward_context
41
+from vllm.logger import init_logger
42
+from vllm.model_executor.layers.attention.attention import (
43
+ _init_kv_cache_quant,
44
+ set_default_quant_scales,
45
+ should_load_quant_weights,
46
+)
47
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
48
+from vllm.model_executor.layers.layernorm import RMSNorm
49
+from vllm.model_executor.layers.linear import (
50
+ ColumnParallelLinear,
51
+ MergedColumnParallelLinear,
52
+ ReplicatedLinear,
53
+ RowParallelLinear,
54
+)
55
+from vllm.model_executor.layers.quantization import QuantizationConfig
56
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
57
+ get_and_maybe_dequant_weights,
58
+)
59
+from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope
60
+from vllm.model_executor.utils import replace_parameter
61
+from vllm.models.common.ops import fused_q_kv_rmsnorm
62
+from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import (
63
+ fused_mla_decode_q_concat_kv_cache_insert,
64
+ fused_mla_key_concat_ds_mla_insert,
65
+ fused_mla_key_concat_kv_cache_insert,
66
+ fused_mla_qkv_quant_kv_cache_fp8_insert,
67
+)
68
+from vllm.platforms import current_platform
69
+from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
70
+from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
71
+from vllm.utils.torch_utils import (
72
+ is_quantized_kv_cache,
73
+ kv_cache_dtype_str_to_dtype,
74
+)
75
+from vllm.v1.attention.backend import (
76
+ AttentionBackend,
77
+ AttentionType,
78
+ MLAAttentionImpl,
79
+)
80
+from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
81
+from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
82
+from vllm.v1.attention.selector import get_attn_backend
83
+from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec, get_kv_quant_mode
84
+
85
+if TYPE_CHECKING:
86
+ from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata
87
+
88
+logger = init_logger(__name__)
89
+
90
+# Below this many tokens, overlap the g_proj GEMM on the aux stream with the
91
+# attention front-end (the GEMM is small and launch-bound, so the overlap
92
+# hides it); at or above it, run the gate on the main stream.
93
+_GATE_MULTI_STREAM_TOKEN_THRESHOLD = 512
94
+
95
+
96
+@torch.compile(backend=current_platform.simple_compile_backend)
97
+def _gate_sigmoid_mul(attn_out: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
98
+ """Apply the sigmoid output gate to a precomputed ``g_proj`` projection."""
99
+ return attn_out * gate.sigmoid()
100
+
101
+
102
+class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
103
+ """Kimi-K3 Multi-head Latent Attention with optional RoPE and output gate."""
104
+
105
+ def __init__(
106
+ self,
107
+ config: KimiLinearConfig,
108
+ hidden_size: int,
109
+ num_heads: int,
110
+ qk_nope_head_dim: int,
111
+ qk_rope_head_dim: int,
112
+ v_head_dim: int,
113
+ q_lora_rank: int | None,
114
+ kv_lora_rank: int,
115
+ use_output_gate: bool = False,
116
+ cache_config: CacheConfig | None = None,
117
+ quant_config: QuantizationConfig | None = None,
118
+ prefix: str = "",
119
+ aux_stream: torch.cuda.Stream | None = None,
120
+ use_rope: bool = False,
121
+ non_causal_multi_token_decode: bool = False,
122
+ ) -> None:
123
+ super().__init__()
124
+ self.hidden_size = hidden_size
125
+ self.qk_nope_head_dim = qk_nope_head_dim
126
+ self.qk_rope_head_dim = qk_rope_head_dim
127
+ self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
128
+ self.v_head_dim = v_head_dim
129
+ self.q_lora_rank = q_lora_rank
130
+ self.kv_lora_rank = kv_lora_rank
131
+ self.non_causal_multi_token_decode = non_causal_multi_token_decode
132
+ # Latent "head" seen by the attention kernel / KV cache.
133
+ self.head_size = kv_lora_rank + qk_rope_head_dim
134
+ self.scale = self.qk_head_dim**-0.5
135
+ self.rms_norm_eps = config.rms_norm_eps
136
+ self.layer_name = prefix
137
+
138
+ self.rotary_emb: RotaryEmbedding | None = None
139
+ if use_rope:
140
+ rope_parameters = dict(config.rope_parameters)
141
+ if rope_parameters["rope_type"] != "default":
142
+ rope_parameters["rope_type"] = (
143
+ "deepseek_yarn"
144
+ if rope_parameters.get("apply_yarn_scaling", True)
145
+ else "deepseek_llama_scaling"
146
+ )
147
+ self.rotary_emb = get_rope(
148
+ qk_rope_head_dim,
149
+ max_position=config.max_position_embeddings,
150
+ rope_parameters=rope_parameters,
151
+ is_neox_style=False,
152
+ dtype=torch.float32,
153
+ )
154
+ if rope_parameters["rope_type"] == "deepseek_yarn":
155
+ mscale_all_dim = rope_parameters.get("mscale_all_dim", False)
156
+ scaling_factor = rope_parameters["factor"]
157
+ mscale = (
158
+ 1.0
159
+ if scaling_factor <= 1
160
+ else 0.1 * float(mscale_all_dim) * math.log(scaling_factor) + 1.0
161
+ )
162
+ self.scale *= mscale * mscale
163
+ # The fused epilogues read the cos/sin table directly in fp32 and run
164
+ # the RoPE math in fp32, so there is no per-forward dtype cast (and no
165
+ # precision loss). deepseek_yarn builds cos_sin_cache in fp32 already;
166
+ # dtype=torch.float32 above forces it for the default rope too (the
167
+ # DSpark draft, which has no yarn scaling).
168
+ assert self.rotary_emb.cos_sin_cache.dtype == torch.float32, (
169
+ "K3 fused MLA RoPE requires an fp32 cos/sin cache; got "
170
+ f"{self.rotary_emb.cos_sin_cache.dtype}."
171
+ )
172
+
173
+ tp_size = get_tensor_model_parallel_world_size()
174
+ assert num_heads % tp_size == 0
175
+ self.num_heads = num_heads
176
+ self.num_local_heads = num_heads // tp_size
177
+
178
+ # ---- Pre-attention projections (fusable front-end) ----
179
+ # Two query variants: a low-rank q-LoRA path (Kimi-K3) fused with the
180
+ # kv-down proj, or an uncompressed q path (Kimi-Linear, ``q_lora_rank``
181
+ # None) with a standalone ``q_proj`` and separate ``kv_a_proj_with_mqa``.
182
+ if self.q_lora_rank is not None:
183
+ # Fused q-down + kv-down projection. Replicated (disable_tp) because
184
+ # the low-rank latents are shared across TP ranks; TP splitting
185
+ # happens at q_b_proj / kv_b_proj. Checkpoint weights ``q_a_proj``
186
+ # and ``kv_a_proj_with_mqa`` map onto shards 0 and 1 respectively.
187
+ self.fused_qkv_a_proj = MergedColumnParallelLinear(
188
+ self.hidden_size,
189
+ [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
190
+ bias=False,
191
+ quant_config=quant_config,
192
+ prefix=f"{prefix}.fused_qkv_a_proj",
193
+ disable_tp=True,
194
+ )
195
+ self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
196
+ self.q_b_proj = ColumnParallelLinear(
197
+ self.q_lora_rank,
198
+ self.num_heads * self.qk_head_dim,
199
+ bias=False,
200
+ quant_config=quant_config,
201
+ prefix=f"{prefix}.q_b_proj",
202
+ )
203
+ else:
204
+ # Uncompressed query: full-rank q_proj (TP-split over heads) plus a
205
+ # replicated kv-down projection (shared latent across TP ranks).
206
+ self.q_proj = ColumnParallelLinear(
207
+ self.hidden_size,
208
+ self.num_heads * self.qk_head_dim,
209
+ bias=False,
210
+ quant_config=quant_config,
211
+ prefix=f"{prefix}.q_proj",
212
+ )
213
+ self.kv_a_proj_with_mqa = ReplicatedLinear(
214
+ self.hidden_size,
215
+ self.kv_lora_rank + self.qk_rope_head_dim,
216
+ bias=False,
217
+ quant_config=quant_config,
218
+ prefix=f"{prefix}.kv_a_proj_with_mqa",
219
+ )
220
+ self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
221
+ self.kv_b_proj = ColumnParallelLinear(
222
+ self.kv_lora_rank,
223
+ self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
224
+ bias=False,
225
+ quant_config=quant_config,
226
+ prefix=f"{prefix}.kv_b_proj",
227
+ )
228
+
229
+ # ---- Post-attention projections ----
230
+ self.use_output_gate = use_output_gate
231
+ self.g_proj = (
232
+ ColumnParallelLinear(
233
+ self.hidden_size,
234
+ self.num_heads * self.v_head_dim,
235
+ bias=False,
236
+ quant_config=quant_config,
237
+ prefix=f"{prefix}.g_proj",
238
+ )
239
+ if use_output_gate
240
+ else None
241
+ )
242
+ # Aux stream (created at the model level, DeepseekV4 convention) for
243
+ # overlapping the g_proj GEMM with the attention front-end. None on
244
+ # ROCm/non-cuda -> maybe_execute_in_parallel falls back to sequential.
245
+ self.aux_stream = aux_stream
246
+ self._gate_events = (
247
+ [torch.cuda.Event(), torch.cuda.Event()]
248
+ if self.g_proj is not None and current_platform.is_cuda_alike()
249
+ else None
250
+ )
251
+ self.o_proj = RowParallelLinear(
252
+ self.num_heads * self.v_head_dim,
253
+ self.hidden_size,
254
+ bias=False,
255
+ quant_config=quant_config,
256
+ prefix=f"{prefix}.o_proj",
257
+ )
258
+
259
+ # ---- Attention backend / impl / KV cache ----
260
+ self.quant_config = quant_config
261
+ if cache_config is not None:
262
+ self.kv_cache_dtype = cache_config.cache_dtype
263
+ else:
264
+ self.kv_cache_dtype = "auto"
265
+
266
+ dtype = torch.get_default_dtype()
267
+ self.attn_backend = get_attn_backend(
268
+ self.head_size,
269
+ dtype,
270
+ self.kv_cache_dtype,
271
+ use_mla=True,
272
+ use_sparse=False,
273
+ num_heads=self.num_local_heads,
274
+ )
275
+ _init_kv_cache_quant(self, quant_config, prefix)
276
+ # Unit (1.0) scale for the fused fp8 prefill path: q/k/v are cast
277
+ # unscaled to match forward_mha (the prefill flash path does not
278
+ # dequantize); only the cache uses _k_scale.
279
+ self.register_buffer(
280
+ "_one_scale", torch.ones(1, dtype=torch.float32), persistent=False
281
+ )
282
+
283
+ impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls())
284
+ self.impl = impl_cls( # type: ignore[assignment]
285
+ num_heads=self.num_local_heads,
286
+ head_size=self.head_size,
287
+ scale=self.scale,
288
+ num_kv_heads=1,
289
+ alibi_slopes=None,
290
+ sliding_window=None,
291
+ kv_cache_dtype=self.kv_cache_dtype,
292
+ logits_soft_cap=None,
293
+ attn_type=AttentionType.DECODER,
294
+ kv_sharing_target_layer_name=None,
295
+ q_lora_rank=self.q_lora_rank,
296
+ kv_lora_rank=self.kv_lora_rank,
297
+ qk_nope_head_dim=self.qk_nope_head_dim,
298
+ qk_rope_head_dim=self.qk_rope_head_dim,
299
+ qk_head_dim=self.qk_head_dim,
300
+ v_head_dim=self.v_head_dim,
301
+ kv_b_proj=self.kv_b_proj,
302
+ indexer=None,
303
+ )
304
+ self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None)
305
+
306
+ vllm_config = get_current_vllm_config()
307
+ parallel_config = vllm_config.parallel_config
308
+ assert (
309
+ parallel_config.decode_context_parallel_size <= 1
310
+ and parallel_config.prefill_context_parallel_size <= 1
311
+ ), "Kimi-K3 MultiHeadLatentAttention does not support context parallelism."
312
+ self.prefill_backend = get_mla_prefill_backend(vllm_config)(
313
+ num_heads=self.num_local_heads,
314
+ scale=self.scale,
315
+ kv_lora_rank=self.kv_lora_rank,
316
+ qk_nope_head_dim=self.qk_nope_head_dim,
317
+ qk_rope_head_dim=self.qk_rope_head_dim,
318
+ v_head_dim=self.v_head_dim,
319
+ vllm_config=vllm_config,
320
+ )
321
+
322
+ compilation_config = vllm_config.compilation_config
323
+ if prefix in compilation_config.static_forward_context:
324
+ raise ValueError(f"Duplicate layer name: {prefix}")
325
+ compilation_config.static_forward_context[prefix] = self
326
+ self.kv_cache = torch.tensor([])
327
+
328
+ # ------------------------------------------------------------------
329
+ # AttentionLayerBase interface
330
+ # ------------------------------------------------------------------
331
+ def get_attn_backend(self) -> type[AttentionBackend]:
332
+ return self.attn_backend
333
+
334
+ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
335
+ kv_cache_dtype = kv_cache_dtype_str_to_dtype(
336
+ self.kv_cache_dtype, vllm_config.model_config
337
+ )
338
+ # TODO: Remove this mypy workaround once the K3 PR is fully merged.
339
+ return MLAAttentionSpec( # type: ignore[call-arg]
340
+ block_size=vllm_config.cache_config.block_size,
341
+ num_kv_heads=1,
342
+ head_size=self.head_size,
343
+ dtype=kv_cache_dtype,
344
+ cache_dtype_str=self.kv_cache_dtype,
345
+ kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
346
+ non_causal_multi_token_decode=self.non_causal_multi_token_decode,
347
+ )
348
+
349
+ def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
350
+ """Absorb ``kv_b_proj`` into decode-time ``W_UK_T`` / ``W_UV`` bmm weights.
351
+
352
+ ``kv_b_proj`` produces ``[k_nope; v]`` per head from the ``kv_lora_rank``
353
+ latent. For the MQA decode path we pre-split it so that queries are
354
+ projected into latent space by ``W_UK_T`` and the attention output is
355
+ projected back to ``v`` by ``W_UV`` -- avoiding materializing full K/V.
356
+ """
357
+ kv_b_proj_weight = get_and_maybe_dequant_weights(
358
+ self.kv_b_proj, out_dtype=act_dtype
359
+ ).T
360
+ assert kv_b_proj_weight.shape == (
361
+ self.kv_lora_rank,
362
+ self.num_local_heads * (self.qk_nope_head_dim + self.v_head_dim),
363
+ ), f"{kv_b_proj_weight.shape=}"
364
+ kv_b_proj_weight = kv_b_proj_weight.view(
365
+ self.kv_lora_rank,
366
+ self.num_local_heads,
367
+ self.qk_nope_head_dim + self.v_head_dim,
368
+ )
369
+ W_UK, W_UV = kv_b_proj_weight.split(
370
+ [self.qk_nope_head_dim, self.v_head_dim], dim=-1
371
+ )
372
+ # (L, N, V) -> (N, L, V)
373
+ replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True)
374
+ # (L, N, P) -> (N, P, L)
375
+ replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True)
376
+
377
+ quant_method = (
378
+ self.quant_config.get_quant_method(self, prefix=self.layer_name)
379
+ if self.quant_config
380
+ else None
381
+ )
382
+ if not should_load_quant_weights(quant_method):
383
+ set_default_quant_scales(self, register_buffer=False)
384
+
385
+ # Precompute reciprocal scales once here (scales are final after load;
386
+ # K3 has no runtime calculate_kv_scales path) so the fp8 fused kernels
387
+ # in the decode/prefill hot path take a ready inverse instead of
388
+ # launching a per-step reciprocal kernel.
389
+ self.register_buffer(
390
+ "_q_scale_inv", self._q_scale.reciprocal().reshape(1), persistent=False
391
+ )
392
+ self.register_buffer(
393
+ "_k_scale_inv", self._k_scale.reciprocal().reshape(1), persistent=False
394
+ )
395
+
396
+ def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor) -> None:
397
+ """Project latent attention output back to ``v`` via ``W_UV`` (bmm)."""
398
+ # (B, N, L) -> (N, B, L)
399
+ x = x.view(-1, self.num_local_heads, self.kv_lora_rank).transpose(0, 1)
400
+ out = out.view(-1, self.num_local_heads, self.v_head_dim)
401
+ # (N, B, L) x (N, L, V) -> (N, B, V) written transposed into (B, N, V)
402
+ torch.bmm(x, self.W_UV, out=out.transpose(0, 1))
403
+
404
+ def _attn_read_kv_cache(self) -> torch.Tensor:
405
+ """Latent cache as seen by the attention read kernels (decode / context).
406
+
407
+ A plain per-tensor fp8 cache is stored as ``uint8``; view it as fp8 so
408
+ the backend reads it as E4M3 rather than fp4/E2M1 -- the latter doubles
409
+ the perceived head dim (``head_size * 2``) and fails the kernel's
410
+ ``head_dim_k == head_dim_q`` check. Mirrors ``MLAAttention.forward``;
411
+ the fp8_ds_mla layout keeps its native uint8 view.
412
+ """
413
+ cache = self.kv_cache
414
+ if (
415
+ is_quantized_kv_cache(self.kv_cache_dtype)
416
+ and self.kv_cache_dtype != "fp8_ds_mla"
417
+ ):
418
+ return cache.view(current_platform.fp8_dtype())
419
+ return cache
420
+
421
+ # ------------------------------------------------------------------
422
+ # Forward
423
+ # ------------------------------------------------------------------
424
+ def _forward_attn(
425
+ self,
426
+ positions: torch.Tensor,
427
+ hidden_states: torch.Tensor,
428
+ ) -> torch.Tensor:
429
+ """Attention front-end: fused qkv-a proj -> norms -> q_b -> attention.
430
+
431
+ Returns the pre-gate attention output ``[num_tokens,
432
+ num_local_heads * v_head_dim]``. On a profile/dummy run
433
+ it returns a zeroed buffer.
434
+ """
435
+ if self.q_lora_rank is not None:
436
+ qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
437
+ q_c, kv_c, k_pe = qkv_lora.split(
438
+ [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
439
+ )
440
+ q_c, kv_c_normed = fused_q_kv_rmsnorm(
441
+ q_c,
442
+ kv_c,
443
+ self.q_a_layernorm.weight.data,
444
+ self.kv_a_layernorm.weight.data,
445
+ self.rms_norm_eps,
446
+ )
447
+ q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim)
448
+ else:
449
+ # Uncompressed query: project directly (no q-LoRA, no q norm) and
450
+ # normalize only the kv latent.
451
+ q = self.q_proj(hidden_states)[0].view(
452
+ -1, self.num_local_heads, self.qk_head_dim
453
+ )
454
+ kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0]
455
+ kv_c, k_pe = kv_lora.split(
456
+ [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
457
+ )
458
+ kv_c_normed = self.kv_a_layernorm(kv_c)
459
+ k_pe = k_pe.unsqueeze(1)
460
+
461
+ attn_out = torch.empty(
462
+ (hidden_states.shape[0], self.num_local_heads * self.v_head_dim),
463
+ dtype=hidden_states.dtype,
464
+ device=hidden_states.device,
465
+ )
466
+ self._attention(positions, q, kv_c_normed, k_pe, attn_out)
467
+ return attn_out
468
+
469
+ def forward(
470
+ self,
471
+ positions: torch.Tensor,
472
+ hidden_states: torch.Tensor,
473
+ ) -> torch.Tensor:
474
+ # Both branches produce (attn_out, gate); they differ only in whether
475
+ # the g_proj GEMM is overlapped on the aux stream.
476
+ g_proj = self.g_proj
477
+ events = self._gate_events
478
+ if (
479
+ g_proj is not None
480
+ and events is not None
481
+ and self.aux_stream is not None
482
+ and hidden_states.shape[0] < _GATE_MULTI_STREAM_TOKEN_THRESHOLD
483
+ ):
484
+ attn_out, gate = maybe_execute_in_parallel(
485
+ lambda: self._forward_attn(positions, hidden_states),
486
+ lambda: g_proj(hidden_states)[0],
487
+ events[0],
488
+ events[1],
489
+ self.aux_stream,
490
+ )
491
+ else:
492
+ attn_out = self._forward_attn(positions, hidden_states)
493
+ gate = g_proj(hidden_states)[0] if g_proj is not None else None
494
+
495
+ if gate is not None:
496
+ attn_out = _gate_sigmoid_mul(attn_out, gate)
497
+
498
+ # ``o_proj`` (RowParallelLinear + out-of-place all-reduce) returns a
499
+ # fresh private tensor, so return it directly rather than copying into a
500
+ # caller buffer -- the previous ``output[:] = ...`` convention forced an
501
+ # extra [num_tokens, hidden] copy per layer.
502
+ return self.o_proj(attn_out)[0]
503
+
504
+ @eager_break_during_capture
505
+ def _attention(
506
+ self,
507
+ positions: torch.Tensor,
508
+ q: torch.Tensor,
509
+ kv_c_normed: torch.Tensor,
510
+ k_pe: torch.Tensor,
511
+ attn_out: torch.Tensor,
512
+ ) -> None:
513
+ forward_context = get_forward_context()
514
+ attn_metadata_by_layer = forward_context.attn_metadata
515
+ if attn_metadata_by_layer is None:
516
+ attn_out.zero_()
517
+ return
518
+ assert isinstance(attn_metadata_by_layer, dict)
519
+ attn_metadata = cast(
520
+ "MLACommonMetadata", attn_metadata_by_layer[self.layer_name]
521
+ )
522
+
523
+ num_actual_toks = attn_metadata.num_actual_tokens
524
+ slot_mapping_by_layer = forward_context.slot_mapping
525
+ assert isinstance(slot_mapping_by_layer, dict)
526
+ slot_mapping = slot_mapping_by_layer[self.layer_name]
527
+
528
+ q = q[:num_actual_toks]
529
+ kv_c_normed = kv_c_normed[:num_actual_toks]
530
+ k_pe = k_pe[:num_actual_toks]
531
+ positions = positions[:num_actual_toks]
532
+ attn_out = attn_out[:num_actual_toks]
533
+
534
+ cos_sin_cache = None
535
+ rope_positions = None
536
+ if self.rotary_emb is not None:
537
+ # Pass the fp32 cos/sin table straight to the fused epilogue (it reads
538
+ # fp32 and does the RoPE math in fp32) -- no per-forward dtype cast.
539
+ cos_sin_cache = self.rotary_emb.cos_sin_cache
540
+ rope_positions = positions
541
+
542
+ # Decode tokens are laid out first, prefill tokens after. The fused
543
+ # prefill covers every supported config (bf16 / plain-fp8 /
544
+ # fp8_ds_mla), so there is no dense-MHA (forward_mha) fallback.
545
+ num_mqa_tokens = attn_metadata.num_decode_tokens
546
+ num_mha_tokens = q.size(0) - num_mqa_tokens
547
+
548
+ # Both the prefill and decode fused epilogues write their own cache
549
+ # slice, so there is no separate do_kv_cache_update.
550
+
551
+ # ---- Prefill: fused key-concat + cache-insert + attention ----
552
+ if num_mha_tokens > 0:
553
+ self._forward_prefill_fused(
554
+ q[num_mqa_tokens:],
555
+ kv_c_normed[num_mqa_tokens:],
556
+ k_pe[num_mqa_tokens:],
557
+ rope_positions[num_mqa_tokens:] if rope_positions is not None else None,
558
+ cos_sin_cache,
559
+ slot_mapping[num_mqa_tokens:num_actual_toks],
560
+ attn_metadata,
561
+ attn_out[num_mqa_tokens:],
562
+ )
563
+
564
+ # ---- Decode: latent multi-query attention ----
565
+ if num_mqa_tokens > 0:
566
+ mqa_q_nope, mqa_q_pe = q[:num_mqa_tokens].split(
567
+ [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
568
+ )
569
+ # BMM1: absorb q_nope into latent space. (N,B,P) x (N,P,L) -> (B,N,L)
570
+ ql_nope = torch.bmm(mqa_q_nope.transpose(0, 1), self.W_UK_T).transpose(0, 1)
571
+ # Fused: concat mqa_q = [ql_nope | q_pe] and insert the decode-token
572
+ # latent into the paged cache (one launch, right before forward_mqa).
573
+ mqa_q = self._decode_concat_cache(
574
+ ql_nope,
575
+ mqa_q_pe,
576
+ kv_c_normed[:num_mqa_tokens],
577
+ k_pe[:num_mqa_tokens],
578
+ rope_positions[:num_mqa_tokens] if rope_positions is not None else None,
579
+ cos_sin_cache,
580
+ slot_mapping[:num_mqa_tokens],
581
+ )
582
+ latent_out, _lse = self.impl.forward_mqa( # type: ignore[attr-defined]
583
+ mqa_q, self._attn_read_kv_cache(), attn_metadata, self
584
+ )
585
+ self._v_up_proj(latent_out, out=attn_out[:num_mqa_tokens])
586
+
587
+ def _decode_concat_cache(
588
+ self,
589
+ ql_nope: torch.Tensor,
590
+ q_pe: torch.Tensor,
591
+ kv_c_normed: torch.Tensor,
592
+ k_pe: torch.Tensor,
593
+ positions: torch.Tensor | None,
594
+ cos_sin_cache: torch.Tensor | None,
595
+ slot_mapping: torch.Tensor,
596
+ ) -> torch.Tensor:
597
+ """Fused decode query-concat + latent cache insert, dispatched by cache
598
+ dtype (same policy as prefill: fp8 cache -> fp8 query)."""
599
+ if self.kv_cache_dtype == "fp8_ds_mla":
600
+ cache = self.kv_cache
601
+ if cache.dtype != torch.uint8:
602
+ cache = cache.view(torch.uint8)
603
+ return fused_mla_decode_q_concat_kv_cache_insert(
604
+ ql_nope,
605
+ q_pe,
606
+ kv_c_normed,
607
+ k_pe,
608
+ cache,
609
+ slot_mapping,
610
+ ds_mla=True,
611
+ positions=positions,
612
+ cos_sin_cache=cos_sin_cache,
613
+ )
614
+ if is_quantized_kv_cache(self.kv_cache_dtype):
615
+ assert self.impl.supports_quant_query_input, ( # type: ignore[attr-defined]
616
+ "Kimi-K3 fp8 KV cache decode requires a backend that accepts an "
617
+ "fp8 (quantized) query input."
618
+ )
619
+ cache = self.kv_cache
620
+ if cache.dtype != torch.float8_e4m3fn:
621
+ cache = cache.view(torch.float8_e4m3fn)
622
+ return fused_mla_decode_q_concat_kv_cache_insert(
623
+ ql_nope,
624
+ q_pe,
625
+ kv_c_normed,
626
+ k_pe,
627
+ cache,
628
+ slot_mapping,
629
+ q_scale_inv=self._q_scale_inv,
630
+ cache_scale_inv=self._k_scale_inv,
631
+ positions=positions,
632
+ cos_sin_cache=cos_sin_cache,
633
+ )
634
+ return fused_mla_decode_q_concat_kv_cache_insert(
635
+ ql_nope,
636
+ q_pe,
637
+ kv_c_normed,
638
+ k_pe,
639
+ self.kv_cache,
640
+ slot_mapping,
641
+ positions=positions,
642
+ cos_sin_cache=cos_sin_cache,
643
+ )
644
+
645
+ def _forward_prefill_fused(
646
+ self,
647
+ q: torch.Tensor,
648
+ kv_c_normed: torch.Tensor,
649
+ k_pe: torch.Tensor,
650
+ positions: torch.Tensor | None,
651
+ cos_sin_cache: torch.Tensor | None,
652
+ slot_mapping: torch.Tensor,
653
+ attn_metadata,
654
+ out: torch.Tensor,
655
+ ) -> None:
656
+ """Prefill using the fused key-concat + cache-insert kernel.
657
+
658
+ Replaces ``_concat_k_nope_k_pe`` and the prefill cache write with one
659
+ fused kernel launch, dispatched by cache dtype. The chunked context
660
+ gather + online-softmax merge are delegated to the impl.
661
+
662
+ Supported configs (K3 fp8 policy):
663
+ - bf16 cache -> bf16 prefill query
664
+ - plain fp8 cache -> fp8 prefill query (unscaled q/k/v; cache _k_scale)
665
+ - fp8_ds_mla cache -> bf16 prefill query (656B per-tile self-scaled)
666
+ """
667
+ prefill = attn_metadata.prefill
668
+ has_context = prefill.chunked_context is not None
669
+ fp8_prefill = prefill.q_data_type == current_platform.fp8_dtype()
670
+
671
+ kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
672
+ -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim
673
+ )
674
+ k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
675
+
676
+ if self.kv_cache_dtype == "fp8_ds_mla":
677
+ # fp8_ds_mla cache (656B, per-tile self-scaled); bf16 attention.
678
+ assert not fp8_prefill, (
679
+ "Kimi-K3 fp8_ds_mla uses a bf16 prefill query; fp8 prefill "
680
+ "query is not supported with fp8_ds_mla."
681
+ )
682
+ kv_cache = self.kv_cache
683
+ if kv_cache.dtype != torch.uint8:
684
+ kv_cache = kv_cache.view(torch.uint8)
685
+ k = fused_mla_key_concat_ds_mla_insert(
686
+ q,
687
+ k_nope,
688
+ k_pe,
689
+ kv_c_normed,
690
+ kv_cache,
691
+ slot_mapping,
692
+ positions,
693
+ cos_sin_cache,
694
+ )
695
+ elif is_quantized_kv_cache(self.kv_cache_dtype):
696
+ assert fp8_prefill, (
697
+ "Kimi-K3 fp8 KV cache requires an fp8 prefill query; enable "
698
+ "--attention-config '{\"use_prefill_query_quantization\": true}'."
699
+ )
700
+ # Plain per-tensor fp8: quant q/k/v (unscaled, matching forward_mha's
701
+ # unscaled `.to(fp8)`) and insert the fp8 latent (scaled by _k_scale).
702
+ kv_cache = self.kv_cache
703
+ if kv_cache.dtype != torch.float8_e4m3fn:
704
+ kv_cache = kv_cache.view(torch.float8_e4m3fn)
705
+ q, k, v = fused_mla_qkv_quant_kv_cache_fp8_insert(
706
+ q,
707
+ k_nope,
708
+ k_pe,
709
+ kv_c_normed,
710
+ v,
711
+ kv_cache,
712
+ slot_mapping,
713
+ self._one_scale,
714
+ self._one_scale,
715
+ self._one_scale,
716
+ self._k_scale_inv,
717
+ positions,
718
+ cos_sin_cache,
719
+ )
720
+ else:
721
+ # Concat full K = [k_nope | k_pe] and insert [kv_c_normed | k_pe]
722
+ # into the paged cache for these prefill tokens, in one launch.
723
+ k = fused_mla_key_concat_kv_cache_insert(
724
+ q,
725
+ k_nope,
726
+ k_pe,
727
+ kv_c_normed,
728
+ self.kv_cache,
729
+ slot_mapping,
730
+ positions,
731
+ cos_sin_cache,
732
+ )
733
+
734
+ # When there is no chunked context, backends that honor `out` write the
735
+ # attention result straight into it, avoiding a slice+flatten+copy.
736
+ writes_out = not has_context and prefill.prefill_backend.supports_out()
737
+ output_prefill = prefill.prefill_backend.run_prefill_new_tokens(
738
+ q=q,
739
+ k=k,
740
+ v=v,
741
+ return_softmax_lse=has_context,
742
+ out=(
743
+ out.view(-1, self.num_local_heads, self.v_head_dim)
744
+ if writes_out
745
+ else None
746
+ ),
747
+ )
748
+
749
+ if has_context:
750
+ context_output, context_lse = self.impl._compute_prefill_context( # type: ignore[attr-defined]
751
+ q, self._attn_read_kv_cache(), attn_metadata, self._k_scale
752
+ )
753
+ suffix_output, suffix_lse = output_prefill
754
+ out = out.view(-1, self.num_local_heads, self.v_head_dim)
755
+ merge_attn_states(
756
+ output=out,
757
+ prefix_output=context_output[..., : self.v_head_dim],
758
+ prefix_lse=context_lse,
759
+ suffix_output=suffix_output[..., : self.v_head_dim],
760
+ suffix_lse=suffix_lse,
761
+ prefill_tokens_with_context=prefill.chunked_context.prefill_tokens_with_context,
762
+ )
763
+ elif not writes_out:
764
+ out.copy_(output_prefill[..., : self.v_head_dim].flatten(start_dim=-2))
vllm/models/kimi_k3/nvidia/model.py
new
+1859
@@ -0,0 +1,1859 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Kimi-K3 multimodal model implementation for vLLM."""
4
+
5
+import math
6
+from collections.abc import Iterable
7
+from typing import Any, cast
8
+
9
+import torch
10
+from torch import nn
11
+
12
+import vllm.envs as envs
13
+from vllm.config import VllmConfig
14
+from vllm.distributed import (
15
+ get_ep_group,
16
+ get_pp_group,
17
+ get_tensor_model_parallel_world_size,
18
+)
19
+from vllm.forward_context import get_forward_context, is_forward_context_available
20
+from vllm.logger import init_logger
21
+from vllm.model_executor.layers.activation import SiluAndMul, SituAndMul
22
+from vllm.model_executor.layers.fused_moe import (
23
+ FusedMoE,
24
+ fused_moe_make_expert_params_mapping,
25
+)
26
+from vllm.model_executor.layers.fused_moe.router.base_router import (
27
+ eplb_map_to_physical_and_record,
28
+)
29
+from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
30
+from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import (
31
+ fused_grouped_topk,
32
+)
33
+from vllm.model_executor.layers.fused_moe.runner.latent_moe_runner import (
34
+ LatentMoERunner,
35
+)
36
+from vllm.model_executor.layers.layernorm import RMSNorm
37
+from vllm.model_executor.layers.linear import (
38
+ MergedColumnParallelLinear,
39
+ ReplicatedLinear,
40
+ RowParallelLinear,
41
+)
42
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
43
+from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import (
44
+ KimiGatedDeltaNetAttention as KimiLinearGatedDeltaNetAttention,
45
+)
46
+from vllm.model_executor.layers.mamba.mamba_utils import (
47
+ MambaStateCopyFunc,
48
+ MambaStateCopyFuncCalculator,
49
+ MambaStateDtypeCalculator,
50
+ MambaStateShapeCalculator,
51
+)
52
+from vllm.model_executor.layers.quantization import QuantizationConfig
53
+from vllm.model_executor.layers.quantization.compressed_tensors import (
54
+ compressed_tensors,
55
+)
56
+from vllm.model_executor.layers.vocab_parallel_embedding import (
57
+ ParallelLMHead,
58
+ VocabParallelEmbedding,
59
+)
60
+from vllm.model_executor.model_loader.weight_utils import (
61
+ default_weight_loader,
62
+ maybe_remap_kv_scale_name,
63
+)
64
+from vllm.model_executor.models.interfaces import (
65
+ EagleModelMixin,
66
+ HasInnerState,
67
+ IsHybrid,
68
+ MixtureOfExperts,
69
+ SupportsEagle3,
70
+ SupportsEncoderCudaGraph,
71
+ SupportsMultiModal,
72
+ SupportsPP,
73
+ SupportsQuant,
74
+)
75
+from vllm.model_executor.models.kimi_k25 import KimiK25MediaPixelInputs
76
+from vllm.model_executor.models.kimi_k25_vit import (
77
+ KimiK25MultiModalProjector,
78
+ MoonViT3dPretrainedModel,
79
+ vision_tower_forward,
80
+)
81
+from vllm.model_executor.models.utils import (
82
+ AutoWeightsLoader,
83
+ PPMissingLayer,
84
+ WeightsMapper,
85
+ init_vllm_registered_model,
86
+ is_pp_missing_parameter,
87
+ make_layers,
88
+ maybe_prefix,
89
+)
90
+from vllm.model_executor.models.vision import is_vit_use_data_parallel
91
+from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MegaMoEExperts
92
+from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
93
+from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
94
+from vllm.models.kimi_k3.nvidia.low_latency_gemm import (
95
+ enable_kimi_k3_low_latency_gemm,
96
+)
97
+from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention
98
+from vllm.models.kimi_k3.nvidia.ops import attn_res
99
+from vllm.models.kimi_k3.nvidia.ops.sequence_parallel import (
100
+ sp_all_gather,
101
+ sp_padding_mask,
102
+ sp_reduce_scatter,
103
+ sp_shard,
104
+)
105
+from vllm.multimodal import MULTIMODAL_REGISTRY
106
+from vllm.multimodal.inputs import NestedTensors
107
+from vllm.platforms import current_platform
108
+from vllm.sequence import IntermediateTensors
109
+from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config
110
+from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
111
+from vllm.utils.math_utils import cdiv
112
+from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
113
+from vllm.utils.torch_utils import aux_stream
114
+from vllm.v1.worker.ubatching import dbo_current_ubatch_id
115
+
116
+from ..common.mm_preprocess import (
117
+ KimiK3DummyInputsBuilder,
118
+ KimiK3MultiModalProcessor,
119
+ KimiK3ProcessingInfo,
120
+)
121
+
122
+logger = init_logger(__name__)
123
+
124
+
125
+class KimiMLP(nn.Module):
126
+ def __init__(
127
+ self,
128
+ hidden_size: int,
129
+ intermediate_size: int,
130
+ hidden_act: str,
131
+ quant_config: QuantizationConfig | None = None,
132
+ reduce_results: bool = True,
133
+ use_sequence_parallel: bool = False,
134
+ prefix: str = "",
135
+ activation_situ_beta: float | None = None,
136
+ activation_situ_linear_beta: float | None = None,
137
+ ) -> None:
138
+ super().__init__()
139
+
140
+ self.gate_up_proj = MergedColumnParallelLinear(
141
+ hidden_size,
142
+ [intermediate_size] * 2,
143
+ bias=False,
144
+ quant_config=quant_config,
145
+ disable_tp=use_sequence_parallel,
146
+ prefix=f"{prefix}.gate_up_proj",
147
+ )
148
+ self.down_proj = RowParallelLinear(
149
+ intermediate_size,
150
+ hidden_size,
151
+ bias=False,
152
+ quant_config=quant_config,
153
+ reduce_results=reduce_results,
154
+ disable_tp=use_sequence_parallel,
155
+ prefix=f"{prefix}.down_proj",
156
+ )
157
+ if hidden_act == "silu":
158
+ self.act_fn = SiluAndMul()
159
+ elif hidden_act == "situ":
160
+ self.act_fn = SituAndMul(
161
+ beta=activation_situ_beta or 1.0,
162
+ linear_beta=activation_situ_linear_beta,
163
+ )
164
+ else:
165
+ raise ValueError(
166
+ f"Unsupported activation: {hidden_act}. "
167
+ "Only silu and situ are supported."
168
+ )
169
+
170
+ def forward(self, x):
171
+ gate_up, _ = self.gate_up_proj(x)
172
+ x = self.act_fn(gate_up)
173
+ x, _ = self.down_proj(x)
174
+ return x
175
+
176
+
177
+class KimiRoutedOutputTransform(nn.Module):
178
+ def __init__(
179
+ self,
180
+ norm: RMSNorm | None,
181
+ up_proj: ReplicatedLinear,
182
+ ) -> None:
183
+ super().__init__()
184
+ self.norm = norm
185
+ self.up_proj = up_proj
186
+
187
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
188
+ if self.norm is not None:
189
+ hidden_states = self.norm(hidden_states)
190
+ hidden_states, _ = self.up_proj(hidden_states)
191
+ return hidden_states
192
+
193
+
194
+class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts):
195
+ """Kimi K3 adapter for the DeepGEMM MegaMoE kernel."""
196
+
197
+ _kimi_symm_buffer_cache: dict[tuple[object, ...], object] = {}
198
+ _synchronized_ep_groups: set[tuple[int, int]] = set()
199
+
200
+ def __init__(
201
+ self,
202
+ *args,
203
+ activation: str,
204
+ activation_beta: float | None,
205
+ activation_linear_beta: float | None,
206
+ **kwargs,
207
+ ):
208
+ super().__init__(*args, **kwargs)
209
+ self.activation = activation
210
+ self.activation_beta = activation_beta
211
+ self.activation_linear_beta = activation_linear_beta
212
+
213
+ def synchronize_first_launch(self) -> None:
214
+ ep_group = get_ep_group()
215
+ device = torch.accelerator.current_device_index()
216
+ key = (id(ep_group.cpu_group), device)
217
+ if key in self._synchronized_ep_groups:
218
+ return
219
+ torch.accelerator.synchronize()
220
+ torch.distributed.barrier(group=ep_group.cpu_group)
221
+ self._synchronized_ep_groups.add(key)
222
+
223
+ def finalize_weights(self) -> None:
224
+ if self._transformed_l1_weights is not None:
225
+ return
226
+
227
+ self._check_runtime_supported()
228
+ from vllm.utils.deep_gemm import _import_deep_gemm
229
+
230
+ deep_gemm = _import_deep_gemm()
231
+ w13_scale = deep_gemm.transform_sf_into_required_layout(
232
+ self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),
233
+ 2 * self.intermediate_size,
234
+ self.hidden_size,
235
+ (1, 32),
236
+ self.num_local_experts,
237
+ )
238
+ w2_scale = deep_gemm.transform_sf_into_required_layout(
239
+ self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(),
240
+ self.hidden_size,
241
+ self.intermediate_size,
242
+ (1, 32),
243
+ self.num_local_experts,
244
+ )
245
+ self._transformed_l1_weights, self._transformed_l2_weights = (
246
+ deep_gemm.transform_weights_for_mega_moe(
247
+ (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale),
248
+ (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale),
249
+ activation=self.activation,
250
+ )
251
+ )
252
+ self.w13_weight = None
253
+ self.w13_weight_scale = None
254
+ self.w2_weight = None
255
+ self.w2_weight_scale = None
256
+
257
+ def get_symm_buffer(self):
258
+ from vllm.utils.deep_gemm import _import_deep_gemm
259
+
260
+ deep_gemm = _import_deep_gemm()
261
+ group = get_ep_group().device_group
262
+ device = torch.accelerator.current_device_index()
263
+ key = (
264
+ id(group),
265
+ device,
266
+ self.num_experts,
267
+ self.max_num_tokens,
268
+ self.top_k,
269
+ self.hidden_size,
270
+ self.intermediate_size,
271
+ self.activation,
272
+ )
273
+ symm_buffer = self._kimi_symm_buffer_cache.get(key)
274
+ if symm_buffer is None:
275
+ symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe(
276
+ group,
277
+ self.num_experts,
278
+ self.max_num_tokens,
279
+ self.top_k,
280
+ self.hidden_size,
281
+ self.intermediate_size,
282
+ activation=self.activation,
283
+ )
284
+ self._kimi_symm_buffer_cache[key] = symm_buffer
285
+ return symm_buffer
286
+
287
+ def forward(
288
+ self,
289
+ hidden_states: torch.Tensor,
290
+ topk_weights: torch.Tensor,
291
+ topk_ids: torch.Tensor,
292
+ *,
293
+ activation_clamp: float | None,
294
+ fast_math: bool = True,
295
+ ) -> torch.Tensor:
296
+ self.synchronize_first_launch()
297
+ if hidden_states.shape[0] > self.max_num_tokens:
298
+ raise ValueError(
299
+ f"Kimi K3 MegaMoE got {hidden_states.shape[0]} tokens, "
300
+ f"but its symmetric buffer supports {self.max_num_tokens}."
301
+ )
302
+ y = torch.empty_like(hidden_states, dtype=torch.bfloat16)
303
+ from vllm.utils.deep_gemm import _import_deep_gemm
304
+
305
+ deep_gemm = _import_deep_gemm()
306
+ symm_buffer = self.get_symm_buffer()
307
+ num_tokens = hidden_states.shape[0]
308
+ is_padding = None
309
+ if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
310
+ is_padding = get_forward_context().is_padding
311
+ if is_padding is not None:
312
+ is_padding = is_padding[:num_tokens]
313
+
314
+ eplb_state = self.eplb_state
315
+ if eplb_state.logical_to_physical_map is not None:
316
+ assert eplb_state.expert_load_view is not None
317
+ assert eplb_state.logical_replica_count is not None
318
+ assert eplb_state.should_record_tensor is not None
319
+ if is_padding is not None:
320
+ topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids)
321
+ topk_ids = eplb_map_to_physical_and_record(
322
+ topk_ids=topk_ids,
323
+ expert_load_view=eplb_state.expert_load_view,
324
+ logical_to_physical_map=eplb_state.logical_to_physical_map,
325
+ logical_replica_count=eplb_state.logical_replica_count,
326
+ record_enabled=eplb_state.should_record_tensor,
327
+ num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[
328
+ dbo_current_ubatch_id()
329
+ ]
330
+ if eplb_state.num_unpadded_tokens_tensors is not None
331
+ else None,
332
+ )
333
+
334
+ prepare_megamoe_inputs(
335
+ hidden_states,
336
+ topk_weights,
337
+ topk_ids,
338
+ symm_buffer.x[:num_tokens],
339
+ symm_buffer.x_sf[:num_tokens],
340
+ symm_buffer.topk_idx[:num_tokens],
341
+ symm_buffer.topk_weights[:num_tokens],
342
+ is_padding=is_padding,
343
+ )
344
+ self.finalize_weights()
345
+ assert self._transformed_l1_weights is not None
346
+ assert self._transformed_l2_weights is not None
347
+ deep_gemm.fp8_fp4_mega_moe(
348
+ y,
349
+ self._transformed_l1_weights,
350
+ self._transformed_l2_weights,
351
+ symm_buffer,
352
+ activation_clamp=activation_clamp,
353
+ activation=self.activation,
354
+ activation_beta=self.activation_beta,
355
+ activation_linear_beta=self.activation_linear_beta,
356
+ fast_math=fast_math,
357
+ )
358
+ return y
359
+
360
+
361
+def make_kimi_k3_mega_moe_expert_params_mapping(
362
+ num_experts: int,
363
+) -> list[tuple[str, str, int, str]]:
364
+ mapping = []
365
+ for expert_id in range(num_experts):
366
+ for shard_id in ("w1", "w2", "w3"):
367
+ param_prefix = "w13" if shard_id in ("w1", "w3") else "w2"
368
+ for suffix in ("weight_packed", "weight_scale"):
369
+ param_suffix = "weight" if suffix == "weight_packed" else suffix
370
+ mapping.append(
371
+ (
372
+ f"experts.{param_prefix}_{param_suffix}",
373
+ f"experts.{expert_id}.{shard_id}.{suffix}",
374
+ expert_id,
375
+ shard_id,
376
+ )
377
+ )
378
+ return mapping
379
+
380
+
381
+class KimiMoE(nn.Module):
382
+ def __init__(
383
+ self,
384
+ config: KimiLinearConfig,
385
+ vllm_config: VllmConfig,
386
+ quant_config: QuantizationConfig | None = None,
387
+ prefix: str = "",
388
+ layer_idx: int = 0,
389
+ use_sequence_parallel: bool = False,
390
+ ):
391
+ super().__init__()
392
+ hidden_size = config.hidden_size
393
+ moe_intermediate_size = config.moe_intermediate_size
394
+ num_experts = config.num_experts
395
+ num_experts_per_token = config.num_experts_per_token
396
+ assert moe_intermediate_size is not None
397
+ assert num_experts is not None
398
+ assert num_experts_per_token is not None
399
+ moe_renormalize = config.moe_renormalize
400
+ routed_expert_hidden_size = config.routed_expert_hidden_size
401
+ self.use_latent_moe = routed_expert_hidden_size is not None
402
+ self.moe_hidden_size = (
403
+ routed_expert_hidden_size
404
+ if routed_expert_hidden_size is not None
405
+ else hidden_size
406
+ )
407
+ self.latent_moe_use_norm = config.latent_moe_use_norm
408
+ self.tp_size = get_tensor_model_parallel_world_size()
409
+ self.routed_scaling_factor = config.routed_scaling_factor
410
+ self.moe_renormalize = moe_renormalize
411
+ self.use_grouped_topk = config.use_grouped_topk
412
+ self.num_expert_group = config.num_expert_group
413
+ self.topk_group = config.topk_group
414
+ self.moe_router_activation_func = config.moe_router_activation_func
415
+ self.num_shared_experts = config.num_shared_experts
416
+ self.layer_idx = layer_idx
417
+ self.use_mega_moe = (
418
+ vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
419
+ )
420
+ if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
421
+ raise NotImplementedError(
422
+ "Kimi K3 MegaMoE requires expert parallel. Enable it with "
423
+ "--enable-expert-parallel."
424
+ )
425
+ if self.use_mega_moe and config.hidden_act != "situ":
426
+ raise ValueError("Kimi K3 MegaMoE requires SITU activation.")
427
+ if self.use_mega_moe and not self.use_latent_moe:
428
+ raise ValueError("Kimi K3 MegaMoE requires latent MoE projections.")
429
+ if self.use_mega_moe and not self.use_grouped_topk:
430
+ raise ValueError("Kimi K3 MegaMoE requires grouped top-k routing.")
431
+ if self.use_mega_moe and (self.num_expert_group != 1 or self.topk_group != 1):
432
+ raise NotImplementedError(
433
+ "Kimi K3 MegaMoE currently requires one expert group."
434
+ )
435
+ self.padded_moe_intermediate_size = moe_intermediate_size
436
+ min_moe_intermediate_per_partition = getattr(
437
+ config, "min_moe_intermediate_per_partition", 256
438
+ )
439
+ if self.tp_size > 1:
440
+ moe_intermediate_per_partition = moe_intermediate_size // self.tp_size
441
+ if moe_intermediate_per_partition < min_moe_intermediate_per_partition:
442
+ self.padded_moe_intermediate_size = (
443
+ min_moe_intermediate_per_partition * self.tp_size
444
+ )
445
+ activation_situ_beta = (
446
+ config.activation_situ_beta if config.hidden_act == "situ" else None
447
+ )
448
+ activation_situ_linear_beta = (
449
+ config.activation_situ_linear_beta if config.hidden_act == "situ" else None
450
+ )
451
+
452
+ # Route with fp32 logits for numerically stable expert selection.
453
+ self.gate = GateLinear(
454
+ input_size=hidden_size,
455
+ output_size=num_experts,
456
+ bias=False,
457
+ out_dtype=torch.float32,
458
+ prefix=f"{prefix}.gate",
459
+ )
460
+
461
+ self.gate.e_score_correction_bias = nn.Parameter(
462
+ torch.empty(num_experts, dtype=torch.float32)
463
+ )
464
+
465
+ if self.num_shared_experts is not None:
466
+ shared_intermediate_size = moe_intermediate_size * self.num_shared_experts
467
+ self.shared_experts = KimiMLP(
468
+ hidden_size=config.hidden_size,
469
+ intermediate_size=shared_intermediate_size,
470
+ hidden_act=config.hidden_act,
471
+ quant_config=quant_config,
472
+ reduce_results=False,
473
+ use_sequence_parallel=use_sequence_parallel,
474
+ prefix=f"{prefix}.shared_experts",
475
+ activation_situ_beta=activation_situ_beta,
476
+ activation_situ_linear_beta=activation_situ_linear_beta,
477
+ )
478
+ else:
479
+ self.shared_experts = None
480
+
481
+ self.routed_expert_down_proj: ReplicatedLinear | None
482
+ self.routed_expert_norm: RMSNorm | None
483
+ self.routed_expert_up_proj: ReplicatedLinear | None
484
+ self.routed_output_transform: KimiRoutedOutputTransform | None
485
+ if self.use_latent_moe:
486
+ self.routed_expert_down_proj = ReplicatedLinear(
487
+ hidden_size,
488
+ self.moe_hidden_size,
489
+ bias=False,
490
+ quant_config=None,
491
+ prefix=f"{prefix}.routed_expert_down_proj",
492
+ )
493
+ self.routed_expert_norm = (
494
+ RMSNorm(self.moe_hidden_size, eps=config.rms_norm_eps)
495
+ if self.latent_moe_use_norm
496
+ else None
497
+ )
498
+ # Replicated up-proj: the full weight lives on every rank and
499
+ # produces the full hidden dim locally. This lets LatentMoERunner
500
+ # fuse the latent and shared reductions into a single all-reduce
501
+ # (concat the two partials, reduce once), then run the up-proj and
502
+ # shared add locally with no further collective.
503
+ self.routed_expert_up_proj = ReplicatedLinear(
504
+ self.moe_hidden_size,
505
+ hidden_size,
506
+ bias=False,
507
+ quant_config=None,
508
+ prefix=f"{prefix}.routed_expert_up_proj",
509
+ )
510
+
511
+ self.routed_output_transform = KimiRoutedOutputTransform(
512
+ self.routed_expert_norm, self.routed_expert_up_proj
513
+ )
514
+ # Auxiliary CUDA stream to overlap the router gate with the routed
515
+ # down projection on decode-sized batches (gated by
516
+ # VLLM_ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD).
517
+ self._down_proj_stream: torch.cuda.Stream | None = aux_stream()
518
+ self._down_proj_events = (torch.cuda.Event(), torch.cuda.Event())
519
+ else:
520
+ self.routed_expert_down_proj = None
521
+ self.routed_expert_norm = None
522
+ self.routed_expert_up_proj = None
523
+ self.routed_output_transform = None
524
+
525
+ if self.use_mega_moe:
526
+ ep_group = get_ep_group()
527
+ ep_size = ep_group.world_size
528
+ ep_rank = ep_group.rank_in_group
529
+ if num_experts % ep_size != 0:
530
+ raise ValueError(
531
+ f"Kimi K3 num_experts={num_experts} must be divisible by "
532
+ f"EP size {ep_size}."
533
+ )
534
+ num_local_experts = num_experts // ep_size
535
+ self.experts = KimiK3MegaMoEExperts(
536
+ vllm_config,
537
+ num_experts=num_experts,
538
+ num_local_experts=num_local_experts,
539
+ experts_start_idx=ep_rank * num_local_experts,
540
+ top_k=num_experts_per_token,
541
+ hidden_size=self.moe_hidden_size,
542
+ intermediate_size=self.padded_moe_intermediate_size,
543
+ prefix=f"{prefix}.experts",
544
+ activation="situ",
545
+ activation_beta=activation_situ_beta,
546
+ activation_linear_beta=activation_situ_linear_beta,
547
+ )
548
+ else:
549
+ enable_tail_fusion = envs.VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION
550
+ self.experts = FusedMoE(
551
+ shared_experts=self.shared_experts,
552
+ num_experts=num_experts,
553
+ top_k=num_experts_per_token,
554
+ hidden_size=self.moe_hidden_size,
555
+ intermediate_size=self.padded_moe_intermediate_size,
556
+ activation=config.hidden_act,
557
+ activation_situ_beta=activation_situ_beta,
558
+ activation_situ_linear_beta=activation_situ_linear_beta,
559
+ renormalize=moe_renormalize,
560
+ quant_config=quant_config,
561
+ use_grouped_topk=config.use_grouped_topk,
562
+ num_expert_group=config.num_expert_group,
563
+ topk_group=config.topk_group,
564
+ prefix=f"{prefix}.experts",
565
+ scoring_func=config.moe_router_activation_func,
566
+ e_score_correction_bias=self.gate.e_score_correction_bias,
567
+ routed_scaling_factor=self.routed_scaling_factor,
568
+ # Down projection runs outside FusedMoE so it can overlap the
569
+ # router gate on the aux stream (see forward()); the original
570
+ # hidden states are passed to forward() as shared_experts_input
571
+ # so shared experts still see the untransformed input.
572
+ routed_input_transform=None,
573
+ routed_output_transform=self.routed_output_transform,
574
+ is_sequence_parallel=use_sequence_parallel,
575
+ runner_cls=LatentMoERunner if self.use_latent_moe else None,
576
+ runner_args=(
577
+ {"enable_k3_latent_moe_tail_fusion": enable_tail_fusion}
578
+ if self.use_latent_moe
579
+ else None
580
+ ),
581
+ )
582
+ if self.padded_moe_intermediate_size != moe_intermediate_size:
583
+ w13_weight = getattr(self.experts, "w13_weight", None)
584
+ if w13_weight is None:
585
+ w13_weight = getattr(self.experts, "w13_weight_packed", None)
586
+ w2_weight = getattr(self.experts, "w2_weight", None)
587
+ if w2_weight is None:
588
+ w2_weight = getattr(self.experts, "w2_weight_packed", None)
589
+ if w13_weight is not None:
590
+ w13_weight.data.zero_()
591
+ if w2_weight is not None:
592
+ w2_weight.data.zero_()
593
+ self.experts.moe_config.intermediate_size_per_partition_unpadded = (
594
+ moe_intermediate_size // self.tp_size
595
+ )
596
+
597
+ def _maybe_overlap_router_and_down_proj(
598
+ self, hidden_states: torch.Tensor
599
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
600
+ """Compute the routed-expert down projection alongside the router,
601
+ overlapping them on separate CUDA streams when latent MoE is enabled.
602
+
603
+ The router gate and the down projection both read ``hidden_states``, so
604
+ the gate runs on the default stream and the down projection on the aux
605
+ stream, joined via ``maybe_execute_in_parallel``. For MegaMoE the
606
+ grouped top-k selection consumes only the gate logits, so it also runs
607
+ on the default stream and overlaps the down projection.
608
+
609
+ Returns:
610
+ ``(routed_hidden_states, router_output, topk_ids)``.
611
+ ``routed_hidden_states`` is the down-projected latent (or the
612
+ original ``hidden_states`` when latent MoE is disabled). For MegaMoE
613
+ ``router_output`` holds the grouped top-k weights and ``topk_ids``
614
+ the selected experts; otherwise ``router_output`` holds the raw gate
615
+ logits and ``topk_ids`` is ``None``.
616
+ """
617
+
618
+ def _router(
619
+ hidden_states: torch.Tensor,
620
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
621
+ router_logits, _ = self.gate(hidden_states)
622
+ if not self.use_mega_moe:
623
+ return router_logits, None
624
+ return fused_grouped_topk(
625
+ hidden_states=hidden_states,
626
+ gating_output=router_logits,
627
+ topk=self.experts.top_k,
628
+ renormalize=self.moe_renormalize,
629
+ e_score_correction_bias=self.gate.e_score_correction_bias.data,
630
+ num_expert_group=self.num_expert_group,
631
+ topk_group=self.topk_group,
632
+ scoring_func=self.moe_router_activation_func,
633
+ routed_scaling_factor=self.routed_scaling_factor,
634
+ )
635
+
636
+ down_proj = self.routed_expert_down_proj
637
+ if down_proj is None:
638
+ router_output, topk_ids = _router(hidden_states)
639
+ return hidden_states, router_output, topk_ids
640
+ num_tokens = hidden_states.shape[0]
641
+ (router_output, topk_ids), (routed_hidden_states, _) = (
642
+ maybe_execute_in_parallel(
643
+ lambda: _router(hidden_states),
644
+ lambda: down_proj(hidden_states),
645
+ self._down_proj_events[0],
646
+ self._down_proj_events[1],
647
+ self._down_proj_stream
648
+ if num_tokens <= envs.VLLM_ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD
649
+ else None,
650
+ )
651
+ )
652
+ return routed_hidden_states, router_output, topk_ids
653
+
654
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
655
+ num_tokens, hidden_size = hidden_states.shape
656
+ hidden_states = hidden_states.view(-1, hidden_size)
657
+ # Overlap the gate with the routed down projection; the returned hidden
658
+ # states are already down-projected. Keep the original ``hidden_states``
659
+ # for the shared experts.
660
+ routed_hidden_states, router_output, topk_ids = (
661
+ self._maybe_overlap_router_and_down_proj(hidden_states)
662
+ )
663
+ if self.use_mega_moe:
664
+ assert self.routed_output_transform is not None
665
+ assert topk_ids is not None
666
+ final_hidden_states = self.experts(
667
+ routed_hidden_states,
668
+ router_output,
669
+ topk_ids,
670
+ activation_clamp=None,
671
+ )
672
+ final_hidden_states = self.routed_output_transform(final_hidden_states)
673
+ if self.shared_experts is not None:
674
+ final_hidden_states = final_hidden_states + self.shared_experts(
675
+ hidden_states
676
+ )
677
+ else:
678
+ # Routed experts consume the down-projected latent; shared experts
679
+ # (inside FusedMoE) get the original hidden states via
680
+ # shared_experts_input.
681
+ final_hidden_states = self.experts(
682
+ hidden_states=routed_hidden_states,
683
+ router_logits=router_output,
684
+ shared_experts_input=hidden_states,
685
+ )
686
+ return final_hidden_states.view(num_tokens, hidden_size)
687
+
688
+
689
+class KimiDecoderLayer(nn.Module):
690
+ def __init__(
691
+ self,
692
+ config: KimiLinearConfig,
693
+ vllm_config: VllmConfig,
694
+ prefix: str = "",
695
+ aux_stream: torch.cuda.Stream | None = None,
696
+ ) -> None:
697
+ super().__init__()
698
+ self.hidden_size = config.hidden_size
699
+ self.layer_idx = int(prefix.rsplit(".", 1)[1])
700
+
701
+ self.is_moe = config.is_moe
702
+ layer_idx = self.layer_idx
703
+ cache_config = vllm_config.cache_config
704
+ quant_config = vllm_config.quant_config
705
+ parallel_config = vllm_config.parallel_config
706
+ self.is_moe_layer = (
707
+ self.is_moe
708
+ and config.num_experts is not None
709
+ and layer_idx >= config.first_k_dense_replace
710
+ and layer_idx % config.moe_layer_freq == 0
711
+ )
712
+
713
+ use_mega_moe = vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
714
+ self.use_sequence_parallel = (
715
+ parallel_config.pipeline_parallel_size == 1
716
+ and parallel_config.enable_expert_parallel
717
+ and parallel_config.tensor_parallel_size > 1
718
+ and (use_mega_moe or parallel_config.data_parallel_size > 1)
719
+ )
720
+ if config.is_kda_layer(layer_idx):
721
+ kda_config = config.linear_attn_config
722
+ assert kda_config is not None
723
+ # This class also serves standalone Kimi-Linear through the model
724
+ # registry. Only Kimi-K3's full-rank gate uses the private KDA path.
725
+ if kda_config.get("use_full_rank_gate", False):
726
+ self.self_attn = KimiK3DeltaAttention(
727
+ config,
728
+ vllm_config,
729
+ prefix=f"{prefix}.self_attn",
730
+ )
731
+ self._self_attn_writes_output = False
732
+ else:
733
+ self.self_attn = KimiLinearGatedDeltaNetAttention(
734
+ config,
735
+ vllm_config,
736
+ prefix=f"{prefix}.self_attn",
737
+ )
738
+ self._self_attn_writes_output = True
739
+ else:
740
+ qk_nope_head_dim = config.qk_nope_head_dim
741
+ qk_rope_head_dim = config.qk_rope_head_dim
742
+ v_head_dim = config.v_head_dim
743
+ kv_lora_rank = config.kv_lora_rank
744
+ mla_use_nope = config.mla_use_nope
745
+ assert qk_nope_head_dim is not None
746
+ assert qk_rope_head_dim is not None
747
+ assert v_head_dim is not None
748
+ assert kv_lora_rank is not None
749
+ assert mla_use_nope, "Kimi-K3 MLA (MultiHeadLatentAttention) is NoPE-only"
750
+ # q_lora_rank may be None (Kimi-Linear): the MLA layer then uses an
751
+ # uncompressed q_proj instead of the fused q-LoRA front-end.
752
+ self.self_attn = MultiHeadLatentAttention(
753
+ config=config,
754
+ hidden_size=self.hidden_size,
755
+ num_heads=config.num_attention_heads,
756
+ qk_nope_head_dim=qk_nope_head_dim,
757
+ qk_rope_head_dim=qk_rope_head_dim,
758
+ v_head_dim=v_head_dim,
759
+ q_lora_rank=config.q_lora_rank,
760
+ kv_lora_rank=kv_lora_rank,
761
+ use_output_gate=bool(config.mla_use_output_gate),
762
+ cache_config=cache_config,
763
+ quant_config=quant_config,
764
+ prefix=f"{prefix}.self_attn",
765
+ aux_stream=aux_stream,
766
+ )
767
+ self._self_attn_writes_output = False
768
+
769
+ if self.use_sequence_parallel:
770
+ self.self_attn.o_proj.reduce_results = False
771
+
772
+ if self.is_moe_layer:
773
+ self.block_sparse_moe = KimiMoE(
774
+ config=config,
775
+ vllm_config=vllm_config,
776
+ quant_config=quant_config,
777
+ prefix=f"{prefix}.block_sparse_moe",
778
+ layer_idx=layer_idx,
779
+ use_sequence_parallel=self.use_sequence_parallel,
780
+ )
781
+ self.mlp = self.block_sparse_moe
782
+ else:
783
+ self.mlp = KimiMLP(
784
+ hidden_size=self.hidden_size,
785
+ intermediate_size=config.intermediate_size,
786
+ hidden_act=config.hidden_act,
787
+ quant_config=quant_config,
788
+ prefix=f"{prefix}.mlp",
789
+ use_sequence_parallel=self.use_sequence_parallel,
790
+ activation_situ_beta=config.activation_situ_beta,
791
+ activation_situ_linear_beta=config.activation_situ_linear_beta,
792
+ )
793
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
794
+ self.post_attention_layernorm = RMSNorm(
795
+ config.hidden_size, eps=config.rms_norm_eps
796
+ )
797
+
798
+ attn_res_block_size = config.attn_res_block_size
799
+ self.use_attn_res = attn_res_block_size is not None
800
+ if self.use_attn_res:
801
+ assert attn_res_block_size is not None
802
+ self.attn_res_block_size = attn_res_block_size
803
+ self.is_block_write_layer = layer_idx % self.attn_res_block_size == 0
804
+ self.block_write_idx = layer_idx // self.attn_res_block_size
805
+ self.prev_valid_blocks = cdiv(layer_idx, self.attn_res_block_size)
806
+ self.self_attention_res_norm = RMSNorm(
807
+ config.hidden_size, eps=config.rms_norm_eps
808
+ )
809
+ self.mlp_res_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
810
+ self.self_attention_res_proj = ReplicatedLinear(
811
+ config.hidden_size,
812
+ 1,
813
+ bias=False,
814
+ quant_config=None,
815
+ prefix=f"{prefix}.self_attention_res_proj",
816
+ )
817
+ self.mlp_res_proj = ReplicatedLinear(
818
+ config.hidden_size,
819
+ 1,
820
+ bias=False,
821
+ quant_config=None,
822
+ prefix=f"{prefix}.mlp_res_proj",
823
+ )
824
+
825
+ def _run_self_attn(
826
+ self,
827
+ positions: torch.Tensor,
828
+ hidden_states: torch.Tensor,
829
+ ) -> torch.Tensor:
830
+ if self._self_attn_writes_output:
831
+ output = torch.empty_like(hidden_states)
832
+ self.self_attn(
833
+ hidden_states=hidden_states,
834
+ positions=positions,
835
+ output=output,
836
+ )
837
+ return output
838
+ return self.self_attn(
839
+ hidden_states=hidden_states,
840
+ positions=positions,
841
+ )
842
+
843
+ def _pre_attn_norm(
844
+ self,
845
+ hidden_states: torch.Tensor | None,
846
+ residual: torch.Tensor | None,
847
+ prefix_sum: torch.Tensor | None,
848
+ ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
849
+ if not self.use_attn_res:
850
+ assert hidden_states is not None
851
+ if residual is None:
852
+ residual = hidden_states
853
+ hidden_states = self.input_layernorm(hidden_states)
854
+ else:
855
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
856
+ return hidden_states, prefix_sum, residual
857
+
858
+ assert prefix_sum is not None
859
+ assert residual is not None
860
+ hidden_states = attn_res(
861
+ prefix_sum,
862
+ hidden_states,
863
+ residual,
864
+ self.self_attention_res_norm.weight,
865
+ self.self_attention_res_proj.weight.squeeze(0),
866
+ self.input_layernorm.weight,
867
+ num_blocks=self.prev_valid_blocks,
868
+ block_write_idx=(self.block_write_idx if self.is_block_write_layer else -1),
869
+ eps=self.self_attention_res_norm.variance_epsilon,
870
+ output_norm_eps=self.input_layernorm.variance_epsilon,
871
+ )
872
+ return hidden_states, prefix_sum, residual
873
+
874
+ def _post_attn_norm(
875
+ self,
876
+ hidden_states: torch.Tensor,
877
+ residual: torch.Tensor,
878
+ prefix_sum: torch.Tensor | None,
879
+ ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
880
+ if not self.use_attn_res:
881
+ hidden_states, residual = self.post_attention_layernorm(
882
+ hidden_states, residual
883
+ )
884
+ return hidden_states, prefix_sum, residual
885
+
886
+ assert prefix_sum is not None
887
+ if self.is_block_write_layer:
888
+ prefix_sum = hidden_states
889
+ prefix_delta = None
890
+ else:
891
+ prefix_delta = hidden_states
892
+ mlp_valid_blocks = self.prev_valid_blocks + self.is_block_write_layer
893
+ hidden_states = attn_res(
894
+ prefix_sum,
895
+ prefix_delta,
896
+ residual,
897
+ self.mlp_res_norm.weight,
898
+ self.mlp_res_proj.weight.squeeze(0),
899
+ self.post_attention_layernorm.weight,
900
+ num_blocks=mlp_valid_blocks,
901
+ block_write_idx=-1,
902
+ eps=self.mlp_res_norm.variance_epsilon,
903
+ output_norm_eps=self.post_attention_layernorm.variance_epsilon,
904
+ )
905
+ return hidden_states, prefix_sum, residual
906
+
907
+ def forward(
908
+ self,
909
+ positions: torch.Tensor,
910
+ hidden_states: torch.Tensor | None,
911
+ residual: torch.Tensor | None,
912
+ prefix_sum: torch.Tensor | None = None,
913
+ **kwargs,
914
+ ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
915
+ hidden_states, prefix_sum, residual = self._pre_attn_norm(
916
+ hidden_states, residual, prefix_sum
917
+ )
918
+ assert hidden_states is not None
919
+
920
+ if self.use_sequence_parallel:
921
+ hidden_states = sp_all_gather(hidden_states)
922
+ # Remove SP padding before attention.
923
+ hidden_states = hidden_states[: positions.shape[0]]
924
+
925
+ # Attention.
926
+ hidden_states = self._run_self_attn(positions, hidden_states)
927
+
928
+ if self.use_sequence_parallel:
929
+ # Add SP padding if needed, and then perform reduce scatter.
930
+ hidden_states = sp_reduce_scatter(hidden_states)
931
+
932
+ hidden_states, prefix_sum, residual = self._post_attn_norm(
933
+ hidden_states, residual, prefix_sum
934
+ )
935
+
936
+ # MoE/MLP.
937
+ hidden_states = self.mlp(hidden_states)
938
+ return hidden_states, prefix_sum, residual
939
+
940
+
941
+class KimiLinearModel(nn.Module, EagleModelMixin):
942
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
943
+ super().__init__()
944
+
945
+ config = vllm_config.model_config.hf_text_config
946
+ self.config = config
947
+ self.attn_res_block_size: int | None = config.attn_res_block_size
948
+ self.use_attn_res = self.attn_res_block_size is not None
949
+ parallel_config = vllm_config.parallel_config
950
+ use_mega_moe = vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
951
+ self.use_sequence_parallel = (
952
+ parallel_config.pipeline_parallel_size == 1
953
+ and parallel_config.enable_expert_parallel
954
+ and parallel_config.tensor_parallel_size > 1
955
+ and (use_mega_moe or parallel_config.data_parallel_size > 1)
956
+ )
957
+
958
+ self.vocab_size = config.vocab_size
959
+
960
+ if get_pp_group().is_first_rank:
961
+ self.embed_tokens = VocabParallelEmbedding(
962
+ config.vocab_size,
963
+ config.hidden_size,
964
+ prefix=f"{prefix}.embed_tokens",
965
+ )
966
+ else:
967
+ self.embed_tokens = PPMissingLayer()
968
+
969
+ # Aux stream for overlapping the MLA g_proj output-gate GEMM with the
970
+ # attention front-end (DeepseekV4 convention: created at the model
971
+ # level and threaded into each attention layer).
972
+ aux_stream = torch.cuda.Stream()
973
+
974
+ def get_layer(prefix: str):
975
+ return KimiDecoderLayer(
976
+ config,
977
+ vllm_config,
978
+ prefix,
979
+ aux_stream=aux_stream,
980
+ )
981
+
982
+ self.start_layer, self.end_layer, self.layers = make_layers(
983
+ config.num_hidden_layers,
984
+ get_layer,
985
+ prefix=f"{prefix}.layers",
986
+ )
987
+ self.num_attn_res_blocks = (
988
+ cdiv(self.end_layer, self.attn_res_block_size)
989
+ if self.attn_res_block_size is not None
990
+ else 0
991
+ )
992
+
993
+ if get_pp_group().is_last_rank:
994
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
995
+ if self.use_attn_res:
996
+ self.output_attn_res_norm = RMSNorm(
997
+ config.hidden_size, eps=config.rms_norm_eps
998
+ )
999
+ self.output_attn_res_proj = ReplicatedLinear(
1000
+ config.hidden_size,
1001
+ 1,
1002
+ bias=False,
1003
+ quant_config=None,
1004
+ prefix=f"{prefix}.output_attn_res_proj",
1005
+ )
1006
+ else:
1007
+ self.norm = PPMissingLayer()
1008
+ if self.use_attn_res:
1009
+ self.output_attn_res_norm = PPMissingLayer()
1010
+ self.output_attn_res_proj = PPMissingLayer()
1011
+
1012
+ world_size = get_tensor_model_parallel_world_size()
1013
+ assert config.num_attention_heads % world_size == 0, (
1014
+ "num_attention_heads must be divisible by world_size"
1015
+ )
1016
+
1017
+ def make_empty_intermediate_tensors(
1018
+ self,
1019
+ batch_size: int,
1020
+ dtype: torch.dtype,
1021
+ device: torch.device,
1022
+ ) -> IntermediateTensors:
1023
+ residual_shape: tuple[int, ...] = (batch_size, self.config.hidden_size)
1024
+ if self.use_attn_res:
1025
+ assert self.attn_res_block_size is not None
1026
+ residual_shape = (
1027
+ batch_size,
1028
+ cdiv(self.start_layer, self.attn_res_block_size),
1029
+ self.config.hidden_size,
1030
+ )
1031
+ return IntermediateTensors(
1032
+ {
1033
+ "hidden_states": torch.zeros(
1034
+ (batch_size, self.config.hidden_size), dtype=dtype, device=device
1035
+ ),
1036
+ "residual": torch.zeros(residual_shape, dtype=dtype, device=device),
1037
+ }
1038
+ )
1039
+
1040
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1041
+ return self.embed_tokens(input_ids)
1042
+
1043
+ def forward(
1044
+ self,
1045
+ input_ids: torch.Tensor | None,
1046
+ positions: torch.Tensor,
1047
+ intermediate_tensors: IntermediateTensors | None,
1048
+ inputs_embeds: torch.Tensor | None = None,
1049
+ **kwargs,
1050
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
1051
+ if get_pp_group().is_first_rank:
1052
+ if inputs_embeds is not None:
1053
+ hidden_states = inputs_embeds
1054
+ else:
1055
+ hidden_states = self.embed_input_ids(input_ids)
1056
+ residual = None
1057
+ else:
1058
+ assert intermediate_tensors is not None
1059
+ hidden_states = intermediate_tensors["hidden_states"]
1060
+ residual = intermediate_tensors["residual"]
1061
+ assert hidden_states is not None
1062
+
1063
+ aux_hidden_states: list[torch.Tensor] = []
1064
+ if self.start_layer in self.aux_hidden_state_layers:
1065
+ if self.use_attn_res or residual is None:
1066
+ aux_hidden_states.append(hidden_states)
1067
+ else:
1068
+ aux_hidden_states.append(hidden_states + residual)
1069
+
1070
+ full_num_tokens = positions.shape[0]
1071
+ if self.use_sequence_parallel:
1072
+ if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
1073
+ forward_context = get_forward_context()
1074
+ forward_context.is_padding = sp_padding_mask(
1075
+ forward_context.is_padding, hidden_states
1076
+ )
1077
+ hidden_states = sp_shard(hidden_states)
1078
+ assert residual is None, "Currently, SP is not supported with PP"
1079
+
1080
+ prefix_sum = None
1081
+ if self.use_attn_res:
1082
+ block_residual = hidden_states.new_empty(
1083
+ hidden_states.size(0),
1084
+ self.num_attn_res_blocks,
1085
+ hidden_states.size(1),
1086
+ )
1087
+ if residual is not None:
1088
+ block_residual[:, : residual.size(1), :].copy_(residual)
1089
+ prefix_sum = hidden_states
1090
+ hidden_states = None
1091
+ residual = block_residual
1092
+
1093
+ for layer_idx, layer in enumerate(
1094
+ self.layers[self.start_layer : self.end_layer],
1095
+ start=self.start_layer,
1096
+ ):
1097
+ hidden_states, prefix_sum, residual = layer(
1098
+ positions=positions,
1099
+ hidden_states=hidden_states,
1100
+ prefix_sum=prefix_sum,
1101
+ residual=residual,
1102
+ )
1103
+ if (layer_idx + 1) in self.aux_hidden_state_layers:
1104
+ if self.use_attn_res:
1105
+ assert prefix_sum is not None
1106
+ aux_hidden_state = prefix_sum + hidden_states
1107
+ else:
1108
+ assert residual is not None
1109
+ aux_hidden_state = hidden_states + residual
1110
+
1111
+ if self.use_sequence_parallel:
1112
+ # Gather SP-sharded aux hidden states.
1113
+ # TODO: Optimize this.
1114
+ aux_hidden_state = sp_all_gather(aux_hidden_state)
1115
+ aux_hidden_state = aux_hidden_state[:full_num_tokens]
1116
+ aux_hidden_states.append(aux_hidden_state)
1117
+
1118
+ assert hidden_states is not None
1119
+ assert residual is not None
1120
+ if not get_pp_group().is_last_rank:
1121
+ assert not self.use_sequence_parallel, (
1122
+ "Currently, SP is not supported with PP"
1123
+ )
1124
+ if prefix_sum is not None:
1125
+ hidden_states = hidden_states + prefix_sum
1126
+ return IntermediateTensors(
1127
+ {"hidden_states": hidden_states, "residual": residual}
1128
+ )
1129
+
1130
+ if self.use_attn_res:
1131
+ assert prefix_sum is not None
1132
+ hidden_states = attn_res(
1133
+ prefix_sum,
1134
+ hidden_states,
1135
+ residual,
1136
+ self.output_attn_res_norm.weight,
1137
+ self.output_attn_res_proj.weight.squeeze(0),
1138
+ None,
1139
+ num_blocks=self.num_attn_res_blocks,
1140
+ block_write_idx=-1,
1141
+ eps=self.output_attn_res_norm.variance_epsilon,
1142
+ output_norm_eps=0.0,
1143
+ )
1144
+ else:
1145
+ hidden_states = hidden_states + residual
1146
+
1147
+ if self.use_sequence_parallel:
1148
+ # Gather SP-sharded hidden states.
1149
+ hidden_states = sp_all_gather(hidden_states)
1150
+ hidden_states = hidden_states[:full_num_tokens]
1151
+
1152
+ # NOTE: the final norm is applied in compute_logits instead of here, so
1153
+ # the MTP draft model receives the pre-norm hidden states.
1154
+ if aux_hidden_states:
1155
+ return hidden_states, aux_hidden_states
1156
+ return hidden_states
1157
+
1158
+ def load_weights(
1159
+ self,
1160
+ weights: Iterable[
1161
+ tuple[str, torch.Tensor] | tuple[str, torch.Tensor, dict[str, Any]]
1162
+ ],
1163
+ ) -> set[str]:
1164
+ kda_config = self.config.linear_attn_config
1165
+ use_full_rank_gate = bool(
1166
+ kda_config and kda_config.get("use_full_rank_gate", False)
1167
+ )
1168
+ beta_shard_id = 5 if use_full_rank_gate else 3
1169
+ stacked_params_mapping = [
1170
+ # (param_name, shard_name, shard_id)
1171
+ (".in_proj_qkvgfab", ".q_proj", 0),
1172
+ (".in_proj_qkvgfab", ".k_proj", 1),
1173
+ (".in_proj_qkvgfab", ".v_proj", 2),
1174
+ (".in_proj_qkvgfab", ".b_proj", beta_shard_id),
1175
+ (".in_proj_qkvgfab", ".f_a_proj", 4),
1176
+ (".conv1d", ".q_conv1d", 0),
1177
+ (".conv1d", ".k_conv1d", 1),
1178
+ (".conv1d", ".v_conv1d", 2),
1179
+ (".gate_up_proj", ".gate_proj", 0),
1180
+ (".gate_up_proj", ".up_proj", 1),
1181
+ ]
1182
+ if use_full_rank_gate:
1183
+ stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3))
1184
+ if getattr(self.config, "q_lora_rank", None) is not None:
1185
+ stacked_params_mapping += [
1186
+ (".fused_qkv_a_proj", ".q_a_proj", 0),
1187
+ (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
1188
+ ]
1189
+ use_mega_moe = any(
1190
+ module.use_mega_moe
1191
+ for module in self.modules()
1192
+ if isinstance(module, KimiMoE)
1193
+ )
1194
+ if self.config.is_moe and use_mega_moe:
1195
+ expert_params_mapping = make_kimi_k3_mega_moe_expert_params_mapping(
1196
+ self.config.num_experts
1197
+ )
1198
+ elif self.config.is_moe:
1199
+ # Params for weights, fp8 weight scales, fp8 activation scales
1200
+ # (param_name, weight_name, expert_id, shard_id)
1201
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
1202
+ self,
1203
+ ckpt_gate_proj_name="w1",
1204
+ ckpt_down_proj_name="w2",
1205
+ ckpt_up_proj_name="w3",
1206
+ num_experts=self.config.num_experts,
1207
+ )
1208
+ else:
1209
+ expert_params_mapping = []
1210
+ params_dict = dict(self.named_parameters())
1211
+ # Under the MXFP4 quant interface the routed experts register unpacked
1212
+ # params (``w13_weight``), while the compressed-tensors checkpoint names
1213
+ # them ``.weight_packed``. Rebind so the expert mapping resolves; scales
1214
+ # already share the ``.weight_scale`` suffix.
1215
+ experts_unpacked = not use_mega_moe and not any(
1216
+ n.endswith("w13_weight_packed") for n in params_dict
1217
+ )
1218
+ loaded_params: set[str] = set()
1219
+ for args in weights:
1220
+ name, loaded_weight = args[0], args[1]
1221
+ kwargs: dict[str, Any] = args[2] if len(args) > 2 else {}
1222
+ if "rotary_emb.inv_freq" in name:
1223
+ continue
1224
+ if experts_unpacked and name.endswith(".weight_packed"):
1225
+ name = name.replace(".weight_packed", ".weight")
1226
+
1227
+ spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
1228
+ if spec_layer is not None:
1229
+ continue # skip spec decode layers for main model
1230
+ if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
1231
+ # Models trained using ColossalAI may include these tensors in
1232
+ # the checkpoint. Skip them.
1233
+ continue
1234
+ for param_name, weight_name, shard_id in stacked_params_mapping:
1235
+ if weight_name not in name:
1236
+ continue
1237
+ # We have mlp.experts[0].gate_proj in the checkpoint.
1238
+ # Since we handle the experts below in expert_params_mapping,
1239
+ # we need to skip here BEFORE we update the name, otherwise
1240
+ # name will be updated to mlp.experts[0].gate_up_proj, which
1241
+ # will then be updated below in expert_params_mapping
1242
+ # for mlp.experts[0].gate_gate_up_proj, which breaks load.
1243
+ if ("mlp.experts." in name) and name not in params_dict:
1244
+ continue
1245
+ name_mapped = name.replace(weight_name, param_name)
1246
+ # Packed projections are only present on compatible layers.
1247
+ if name_mapped not in params_dict:
1248
+ continue
1249
+ name = name_mapped
1250
+ # Skip loading extra bias for GPTQ models.
1251
+ if name.endswith(".bias") and name not in params_dict:
1252
+ continue
1253
+ if is_pp_missing_parameter(name, self):
1254
+ continue
1255
+ param = params_dict[name]
1256
+ weight_loader = param.weight_loader
1257
+ weight_loader(param, loaded_weight, shard_id)
1258
+ break
1259
+ else:
1260
+ for (
1261
+ expert_param_name,
1262
+ expert_weight_name,
1263
+ expert_id,
1264
+ expert_shard_id,
1265
+ ) in expert_params_mapping:
1266
+ if expert_weight_name not in name:
1267
+ continue
1268
+ name = name.replace(expert_weight_name, expert_param_name)
1269
+ if is_pp_missing_parameter(name, self):
1270
+ continue
1271
+ param = params_dict[name]
1272
+ weight_loader = param.weight_loader
1273
+ weight_loader(
1274
+ param,
1275
+ loaded_weight,
1276
+ name,
1277
+ expert_id=expert_id,
1278
+ shard_id=expert_shard_id,
1279
+ )
1280
+ break
1281
+ else:
1282
+ # Skip loading extra bias for GPTQ models.
1283
+ if (
1284
+ name.endswith(".bias")
1285
+ and name not in params_dict
1286
+ and not self.config.is_linear_attn
1287
+ ): # noqa: E501
1288
+ continue
1289
+ # Remapping the name of FP8 kv-scale.
1290
+ remapped_name = maybe_remap_kv_scale_name(name, params_dict)
1291
+ if remapped_name is None:
1292
+ continue
1293
+ name = remapped_name
1294
+ if is_pp_missing_parameter(name, self):
1295
+ continue
1296
+
1297
+ param = params_dict[name]
1298
+ weight_loader = getattr(
1299
+ param, "weight_loader", default_weight_loader
1300
+ )
1301
+ weight_loader(param, loaded_weight, **kwargs)
1302
+ loaded_params.add(name)
1303
+ return loaded_params
1304
+
1305
+ def finalize_mega_moe_weights(self) -> None:
1306
+ for module in self.modules():
1307
+ if isinstance(module, KimiMoE) and module.use_mega_moe:
1308
+ module.experts.finalize_weights()
1309
+
1310
+
1311
+class KimiLinearForCausalLM(
1312
+ nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid
1313
+):
1314
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1315
+ super().__init__()
1316
+ self.model_config = vllm_config.model_config
1317
+ self.vllm_config = vllm_config
1318
+ self.config = self.model_config.hf_config
1319
+ quant_config = vllm_config.quant_config
1320
+ self.quant_config = quant_config
1321
+ self.model = KimiLinearModel(
1322
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
1323
+ )
1324
+ if get_pp_group().is_last_rank:
1325
+ self.lm_head = ParallelLMHead(
1326
+ self.config.vocab_size,
1327
+ self.config.hidden_size,
1328
+ quant_config=quant_config,
1329
+ prefix=maybe_prefix(prefix, "lm_head"),
1330
+ )
1331
+ else:
1332
+ self.lm_head = PPMissingLayer()
1333
+ enable_kimi_k3_low_latency_gemm(self, self.model_config.dtype)
1334
+ logit_scale = getattr(self.config, "logit_scale", 1.0)
1335
+ self.logits_processor = LogitsProcessor(
1336
+ self.config.vocab_size, scale=logit_scale
1337
+ )
1338
+
1339
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1340
+ return self.model.embed_input_ids(input_ids)
1341
+
1342
+ def make_empty_intermediate_tensors(
1343
+ self,
1344
+ batch_size: int,
1345
+ dtype: torch.dtype,
1346
+ device: torch.device,
1347
+ ) -> IntermediateTensors:
1348
+ return self.model.make_empty_intermediate_tensors(batch_size, dtype, device)
1349
+
1350
+ def forward( # type: ignore[override]
1351
+ self,
1352
+ input_ids: torch.Tensor | None,
1353
+ positions: torch.Tensor,
1354
+ intermediate_tensors: IntermediateTensors | None = None,
1355
+ inputs_embeds: torch.Tensor | None = None,
1356
+ **kwargs,
1357
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
1358
+ return self.model(
1359
+ input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
1360
+ )
1361
+
1362
+ @classmethod
1363
+ def get_mamba_state_dtype_from_config(
1364
+ cls,
1365
+ vllm_config: "VllmConfig",
1366
+ ) -> tuple[torch.dtype, torch.dtype]:
1367
+ return MambaStateDtypeCalculator.kda_state_dtype(
1368
+ vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype
1369
+ )
1370
+
1371
+ @classmethod
1372
+ def get_mamba_state_shape_from_config(
1373
+ cls, vllm_config: "VllmConfig"
1374
+ ) -> tuple[tuple[int, int], tuple[int, int, int]]:
1375
+ parallel_config = vllm_config.parallel_config
1376
+ hf_config = vllm_config.model_config.hf_config
1377
+ tp_size = parallel_config.tensor_parallel_size
1378
+ num_spec = (
1379
+ vllm_config.speculative_config.num_speculative_tokens
1380
+ if vllm_config.speculative_config
1381
+ else 0
1382
+ )
1383
+ return MambaStateShapeCalculator.kda_state_shape(
1384
+ tp_size,
1385
+ hf_config.linear_attn_config["num_heads"],
1386
+ hf_config.linear_attn_config["head_dim"],
1387
+ conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"],
1388
+ num_spec=num_spec,
1389
+ )
1390
+
1391
+ @classmethod
1392
+ def get_mamba_state_copy_func(
1393
+ cls,
1394
+ ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
1395
+ return MambaStateCopyFuncCalculator.kda_state_copy_func()
1396
+
1397
+ def compute_logits(
1398
+ self,
1399
+ hidden_states: torch.Tensor,
1400
+ ) -> torch.Tensor | None:
1401
+ # The model's final norm is applied here (not at the end of forward) so
1402
+ # that the pre-norm hidden states can be fed to the MTP draft model.
1403
+ hidden_states = self.model.norm(hidden_states, None)
1404
+ return self.logits_processor(self.lm_head, hidden_states)
1405
+
1406
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1407
+ loader = AutoWeightsLoader(
1408
+ self,
1409
+ skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
1410
+ )
1411
+ loaded = loader.load_weights(weights)
1412
+ self.model.finalize_mega_moe_weights()
1413
+ # The fused MultiHeadLatentAttention's process_weights_after_loading
1414
+ # (W_UK_T / W_UV absorption) is driven by the loader's generic post-load
1415
+ # hook for any AttentionLayerBase, so no manual trigger is needed here.
1416
+ return loaded
1417
+
1418
+
1419
+def get_spec_layer_idx_from_weight_name(
1420
+ config: KimiLinearConfig, weight_name: str
1421
+) -> int | None:
1422
+ if hasattr(config, "num_nextn_predict_layers") and (
1423
+ config.num_nextn_predict_layers > 0
1424
+ ):
1425
+ layer_idx = config.num_hidden_layers
1426
+ for i in range(config.num_nextn_predict_layers):
1427
+ # Match regardless of the surrounding prefix. The name may arrive as
1428
+ # ``model.layers.{i}.``, a bare ``layers.{i}.`` (after AutoWeightsLoader
1429
+ # has stripped the ``model.`` prefix in the main model), or with the
1430
+ # multimodal ``language_model.model.layers.{i}.`` prefix.
1431
+ if f"layers.{layer_idx + i}." in weight_name:
1432
+ return layer_idx + i
1433
+ return None
1434
+
1435
+
1436
+@MULTIMODAL_REGISTRY.register_processor(
1437
+ KimiK3MultiModalProcessor,
1438
+ info=KimiK3ProcessingInfo,
1439
+ dummy_inputs=KimiK3DummyInputsBuilder,
1440
+)
1441
+class KimiK3ForConditionalGeneration(
1442
+ nn.Module,
1443
+ SupportsMultiModal,
1444
+ SupportsEncoderCudaGraph,
1445
+ SupportsPP,
1446
+ SupportsQuant,
1447
+ SupportsEagle3,
1448
+ HasInnerState,
1449
+ IsHybrid,
1450
+):
1451
+ """Kimi-K3 model with Kimi-K2.5 vision and KimiLinear text."""
1452
+
1453
+ supports_encoder_tp_data = True
1454
+
1455
+ hf_to_vllm_mapper = WeightsMapper(
1456
+ orig_to_new_prefix={
1457
+ "language_model.layers.": "language_model.model.layers.",
1458
+ "mm_projector.proj.0": "mm_projector.linear_1",
1459
+ "mm_projector.proj.2": "mm_projector.linear_2",
1460
+ }
1461
+ )
1462
+
1463
+ @classmethod
1464
+ def get_placeholder_str(cls, modality: str, i: int) -> str | None:
1465
+ if modality == "image":
1466
+ return "<|kimi_image_placeholder|>"
1467
+ raise ValueError(f"Unsupported modality: {modality}")
1468
+
1469
+ def __init__(
1470
+ self,
1471
+ vllm_config: VllmConfig,
1472
+ prefix: str = "",
1473
+ ) -> None:
1474
+ super().__init__()
1475
+ model_config = vllm_config.model_config
1476
+ config: KimiK3Config = model_config.hf_config
1477
+ self.config = config
1478
+ self.model_config = model_config
1479
+ quant_config = vllm_config.quant_config
1480
+
1481
+ multimodal_config = model_config.multimodal_config
1482
+ assert multimodal_config is not None
1483
+ self.use_data_parallel = is_vit_use_data_parallel(
1484
+ config.vision_config.num_attention_heads
1485
+ )
1486
+ self.hidden_size = config.text_config.hidden_size
1487
+ self.device = current_platform.current_device()
1488
+
1489
+ with self._mark_tower_model(vllm_config, "image"):
1490
+ self.vision_tower = MoonViT3dPretrainedModel(
1491
+ config.vision_config,
1492
+ quant_config=self._maybe_ignore_quant_config(quant_config),
1493
+ prefix=maybe_prefix(prefix, "vision_tower"),
1494
+ )
1495
+ if self._maybe_ignore_quant_config(quant_config) is not None:
1496
+ self.vision_tower = self.vision_tower.to(device=self.device)
1497
+ else:
1498
+ self.vision_tower = self.vision_tower.to(
1499
+ device=self.device, dtype=model_config.dtype
1500
+ )
1501
+
1502
+ vision_attn = self.vision_tower.encoder.blocks[0].attn
1503
+ if vision_attn.is_flash_attn_backend and vision_attn._fa_version == 4:
1504
+ from vllm.models.kimi_k3.nvidia.ops.vision_fa4_warmup import (
1505
+ KimiK3VisionFA4WarmupConfig,
1506
+ register_kimi_k3_vision_fa4_warmup,
1507
+ )
1508
+
1509
+ merge_height, merge_width = config.vision_config.merge_kernel_size
1510
+ mm_config = model_config.get_multimodal_config()
1511
+ assert mm_config is not None
1512
+ register_kimi_k3_vision_fa4_warmup(
1513
+ KimiK3VisionFA4WarmupConfig(
1514
+ num_heads=vision_attn.num_heads,
1515
+ head_dim=vision_attn.head_size,
1516
+ dtype=vision_attn.dtype,
1517
+ max_batch_size=(
1518
+ vllm_config.scheduler_config.max_num_seqs
1519
+ * mm_config.get_limit_per_prompt("image")
1520
+ ),
1521
+ max_seqlen=(
1522
+ vllm_config.scheduler_config.max_num_encoder_input_tokens
1523
+ * merge_height
1524
+ * merge_width
1525
+ ),
1526
+ )
1527
+ )
1528
+
1529
+ self.mm_projector = KimiK25MultiModalProjector(
1530
+ config=config.vision_config,
1531
+ use_data_parallel=self.use_data_parallel,
1532
+ quant_config=self._maybe_ignore_quant_config(quant_config),
1533
+ prefix=maybe_prefix(prefix, "mm_projector"),
1534
+ )
1535
+ self.mm_projector = self.mm_projector.to(
1536
+ device=self.device, dtype=model_config.dtype
1537
+ )
1538
+
1539
+ self.quant_config = quant_config
1540
+ with self._mark_language_model(vllm_config):
1541
+ self.language_model = init_vllm_registered_model(
1542
+ vllm_config=vllm_config,
1543
+ hf_config=config.text_config,
1544
+ prefix=maybe_prefix(prefix, "language_model"),
1545
+ architectures=["KimiLinearForCausalLM"],
1546
+ )
1547
+ self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
1548
+ self.language_model.make_empty_intermediate_tensors
1549
+ )
1550
+ self.media_placeholder: int = self.config.media_placeholder_token_id
1551
+
1552
+ # -- SupportsEncoderCudaGraph protocol methods --
1553
+
1554
+ def get_encoder_cudagraph_config(self):
1555
+ from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig
1556
+
1557
+ return EncoderCudaGraphConfig(
1558
+ modalities=["image"],
1559
+ buffer_keys=[
1560
+ "pixel_values",
1561
+ "pos_embeds",
1562
+ "rope_freqs_cis",
1563
+ "cu_seqlens",
1564
+ "max_seqlen",
1565
+ "sequence_lengths",
1566
+ "merge_gather_idx",
1567
+ ],
1568
+ out_hidden_size=self.hidden_size,
1569
+ )
1570
+
1571
+ def get_encoder_cudagraph_budget_range(
1572
+ self, vllm_config: VllmConfig
1573
+ ) -> tuple[int, int]:
1574
+ min_budget = 64
1575
+ max_budget = min(
1576
+ vllm_config.scheduler_config.max_num_batched_tokens,
1577
+ self.model_config.max_model_len,
1578
+ )
1579
+ return min_budget, max_budget
1580
+
1581
+ @staticmethod
1582
+ def _get_grid_thws(mm_kwargs: dict[str, Any]) -> list[list[int]]:
1583
+ grid_thws = mm_kwargs["grid_thws"]
1584
+ if not isinstance(grid_thws, list):
1585
+ grid_thws = grid_thws.tolist()
1586
+ return grid_thws
1587
+
1588
+ @staticmethod
1589
+ def _get_pixel_values(mm_kwargs: dict[str, Any]) -> torch.Tensor:
1590
+ pixel_values = mm_kwargs["pixel_values"]
1591
+ if isinstance(pixel_values, list):
1592
+ pixel_values = torch.cat(pixel_values)
1593
+ if pixel_values.ndim in (3, 5):
1594
+ pixel_values = pixel_values.reshape(
1595
+ pixel_values.shape[0] * pixel_values.shape[1],
1596
+ *pixel_values.shape[2:],
1597
+ )
1598
+ return pixel_values
1599
+
1600
+ def get_encoder_cudagraph_item_specs(self, mm_kwargs: dict[str, Any]):
1601
+ from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
1602
+
1603
+ kh, kw = self.config.vision_config.merge_kernel_size
1604
+ return [
1605
+ EncoderItemSpec(
1606
+ input_size=t * h * w,
1607
+ output_tokens=(h // kh) * (w // kw),
1608
+ )
1609
+ for t, h, w in self._get_grid_thws(mm_kwargs)
1610
+ ]
1611
+
1612
+ def select_encoder_cudagraph_items(
1613
+ self, mm_kwargs: dict[str, Any], indices: list[int]
1614
+ ) -> dict[str, Any]:
1615
+ grid_thws = self._get_grid_thws(mm_kwargs)
1616
+ pixel_values = self._get_pixel_values(mm_kwargs)
1617
+ source_grid = mm_kwargs["grid_thws"]
1618
+
1619
+ if not indices:
1620
+ empty_grid = (
1621
+ source_grid[:0] if isinstance(source_grid, torch.Tensor) else []
1622
+ )
1623
+ return {"pixel_values": pixel_values[:0], "grid_thws": empty_grid}
1624
+
1625
+ patch_counts = [t * h * w for t, h, w in grid_thws]
1626
+ offsets = [0]
1627
+ for count in patch_counts:
1628
+ offsets.append(offsets[-1] + count)
1629
+ selected_pixel_values = torch.cat(
1630
+ [pixel_values[offsets[i] : offsets[i + 1]] for i in indices]
1631
+ )
1632
+ grid_device = (
1633
+ source_grid.device if isinstance(source_grid, torch.Tensor) else None
1634
+ )
1635
+ selected_grid = torch.tensor(
1636
+ [grid_thws[i] for i in indices],
1637
+ dtype=torch.long,
1638
+ device=grid_device,
1639
+ )
1640
+ return {"pixel_values": selected_pixel_values, "grid_thws": selected_grid}
1641
+
1642
+ def prepare_encoder_cudagraph_capture_inputs(
1643
+ self,
1644
+ token_budget: int,
1645
+ max_batch_size: int,
1646
+ max_frames_per_batch: int,
1647
+ device: torch.device,
1648
+ dtype: torch.dtype,
1649
+ path: str = "default",
1650
+ ):
1651
+ from vllm.v1.worker.encoder_cudagraph_defs import (
1652
+ EncoderCudaGraphCaptureInputs,
1653
+ )
1654
+
1655
+ kh, kw = self.config.vision_config.merge_kernel_size
1656
+ per_item_output = (token_budget + max_batch_size - 1) // max_batch_size
1657
+ rope = self.vision_tower.encoder.rope_2d
1658
+ max_output_width = rope.max_width // kw
1659
+ max_output_height = rope.max_height // kh
1660
+ output_width = min(math.ceil(math.sqrt(per_item_output)), max_output_width)
1661
+ output_height = (per_item_output + output_width - 1) // output_width
1662
+ if output_height > max_output_height:
1663
+ output_height = max_output_height
1664
+ output_width = (per_item_output + output_height - 1) // output_height
1665
+ if output_width > max_output_width:
1666
+ raise ValueError(
1667
+ f"Encoder CUDA graph budget {token_budget} exceeds K3 RoPE "
1668
+ f"capacity for max_batch_size={max_batch_size}"
1669
+ )
1670
+ grid_thws = [
1671
+ [1, output_height * kh, output_width * kw] for _ in range(max_batch_size)
1672
+ ]
1673
+
1674
+ patch_size: int | tuple[int, int] = self.config.vision_config.patch_size
1675
+ if isinstance(patch_size, int):
1676
+ patch_size = (patch_size, patch_size)
1677
+ total_patches = sum(t * h * w for t, h, w in grid_thws)
1678
+ pixel_values = torch.randn(
1679
+ total_patches,
1680
+ 3,
1681
+ patch_size[0],
1682
+ patch_size[1],
1683
+ device=device,
1684
+ dtype=dtype,
1685
+ )
1686
+ metadata = self.vision_tower.prepare_encoder_cudagraph_metadata(
1687
+ grid_thws,
1688
+ max_batch_size=max_batch_size,
1689
+ max_seqlen_override=max(
1690
+ token_budget * kh * kw,
1691
+ max(t * h * w for t, h, w in grid_thws),
1692
+ ),
1693
+ device=device,
1694
+ )
1695
+ return EncoderCudaGraphCaptureInputs(
1696
+ values=metadata | {"pixel_values": pixel_values}
1697
+ )
1698
+
1699
+ def prepare_encoder_cudagraph_replay_buffers(
1700
+ self,
1701
+ mm_kwargs: dict[str, Any],
1702
+ max_batch_size: int,
1703
+ max_frames_per_batch: int,
1704
+ path: str = "default",
1705
+ ):
1706
+ from vllm.v1.worker.encoder_cudagraph_defs import (
1707
+ EncoderCudaGraphReplayBuffers,
1708
+ )
1709
+
1710
+ pixel_values = self._get_pixel_values(mm_kwargs)
1711
+ metadata = self.vision_tower.prepare_encoder_cudagraph_metadata(
1712
+ self._get_grid_thws(mm_kwargs),
1713
+ max_batch_size=max_batch_size,
1714
+ device=pixel_values.device,
1715
+ )
1716
+ return EncoderCudaGraphReplayBuffers(
1717
+ values=metadata | {"pixel_values": pixel_values}
1718
+ )
1719
+
1720
+ def _project_encoder_features(self, image_features: torch.Tensor) -> torch.Tensor:
1721
+ projector_dtype = next(self.mm_projector.parameters()).dtype
1722
+ if image_features.dtype != projector_dtype:
1723
+ image_features = image_features.to(projector_dtype)
1724
+ output = self.mm_projector(image_features)
1725
+ return output.reshape(-1, output.shape[-1])
1726
+
1727
+ def encoder_cudagraph_forward(
1728
+ self,
1729
+ values: dict[str, torch.Tensor],
1730
+ path: str = "default",
1731
+ ) -> torch.Tensor:
1732
+ pixel_values = values.pop("pixel_values")
1733
+ image_features = self.vision_tower(pixel_values, None, encoder_metadata=values)
1734
+ return self._project_encoder_features(image_features)
1735
+
1736
+ def encoder_eager_forward(
1737
+ self,
1738
+ mm_kwargs: dict[str, Any],
1739
+ path: str = "default",
1740
+ ) -> torch.Tensor:
1741
+ image_features = self.vision_tower(
1742
+ self._get_pixel_values(mm_kwargs).to(
1743
+ next(self.vision_tower.parameters()).dtype
1744
+ ),
1745
+ self._get_grid_thws(mm_kwargs),
1746
+ )
1747
+ return self._project_encoder_features(torch.cat(image_features))
1748
+
1749
+ def _maybe_ignore_quant_config(
1750
+ self, quant_config: QuantizationConfig | None
1751
+ ) -> QuantizationConfig | None:
1752
+ if isinstance(quant_config, compressed_tensors.CompressedTensorsConfig):
1753
+ return None
1754
+ return quant_config
1755
+
1756
+ def _parse_and_validate_media_input(
1757
+ self, **kwargs: object
1758
+ ) -> KimiK25MediaPixelInputs | None:
1759
+ pixel_values = kwargs.pop("pixel_values", None)
1760
+ grid_thws = kwargs.pop("grid_thws", None)
1761
+ if pixel_values is None:
1762
+ return None
1763
+
1764
+ if isinstance(pixel_values, list):
1765
+ pixel_values = torch.cat(cast(list[torch.Tensor], pixel_values), dim=0)
1766
+ if not isinstance(pixel_values, torch.Tensor):
1767
+ raise TypeError(
1768
+ "pixel_values must be a tensor or a list of tensors, "
1769
+ f"got {type(pixel_values)}"
1770
+ )
1771
+
1772
+ if len(pixel_values.shape) == 5 or len(pixel_values.shape) == 3:
1773
+ pixel_values = pixel_values.reshape(
1774
+ pixel_values.shape[0] * pixel_values.shape[1], *pixel_values.shape[2:]
1775
+ )
1776
+
1777
+ target_dtype = next(self.vision_tower.parameters()).dtype
1778
+ pixel_values = pixel_values.to(target_dtype)
1779
+ assert isinstance(grid_thws, torch.Tensor), (
1780
+ f"expect grid_thws to be a tensor, got {type(grid_thws)}"
1781
+ )
1782
+ grid_thws = grid_thws.reshape(-1, grid_thws.shape[-1])
1783
+ assert grid_thws.ndim == 2 and grid_thws.size(1) == 3, (
1784
+ f"unexpected shape for grid_thws: {grid_thws.shape}"
1785
+ )
1786
+
1787
+ return KimiK25MediaPixelInputs(
1788
+ type="pixel_values",
1789
+ pixel_values=pixel_values,
1790
+ grid_thws=grid_thws,
1791
+ )
1792
+
1793
+ def _process_media_input(
1794
+ self, media_input: KimiK25MediaPixelInputs
1795
+ ) -> list[torch.Tensor]:
1796
+ media_features = vision_tower_forward(
1797
+ self.vision_tower,
1798
+ media_input["pixel_values"],
1799
+ media_input["grid_thws"],
1800
+ mm_projector=self.mm_projector,
1801
+ use_data_parallel=self.use_data_parallel,
1802
+ )
1803
+ return media_features
1804
+
1805
+ def embed_multimodal(self, **kwargs: object) -> NestedTensors | None:
1806
+ media_input = self._parse_and_validate_media_input(**kwargs)
1807
+ if media_input is None:
1808
+ return None
1809
+ return self._process_media_input(media_input)
1810
+
1811
+ def forward( # type: ignore[override]
1812
+ self,
1813
+ input_ids: torch.Tensor,
1814
+ positions: torch.Tensor,
1815
+ intermediate_tensors: IntermediateTensors | None = None,
1816
+ inputs_embeds: torch.Tensor | None = None,
1817
+ **kwargs: object,
1818
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
1819
+ if intermediate_tensors is not None:
1820
+ inputs_embeds = None
1821
+ return self.language_model(
1822
+ input_ids=input_ids,
1823
+ positions=positions,
1824
+ intermediate_tensors=intermediate_tensors,
1825
+ inputs_embeds=inputs_embeds,
1826
+ )
1827
+
1828
+ def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
1829
+ return self.language_model.compute_logits(hidden_states)
1830
+
1831
+ def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
1832
+ return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs(
1833
+ input_buffers, **kwargs
1834
+ )
1835
+
1836
+ def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
1837
+ return self.language_model.mamba_cache.get_seqlen_agnostic_capture_inputs(
1838
+ batch_size
1839
+ )
1840
+
1841
+ @classmethod
1842
+ def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig):
1843
+ text_config = vllm_config.model_config.hf_config.text_config
1844
+ temp_vllm_config = vllm_config.with_hf_config(text_config)
1845
+ return KimiLinearForCausalLM.get_mamba_state_dtype_from_config(temp_vllm_config)
1846
+
1847
+ @classmethod
1848
+ def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig):
1849
+ text_config = vllm_config.model_config.hf_config.text_config
1850
+ temp_vllm_config = vllm_config.with_hf_config(text_config)
1851
+ return KimiLinearForCausalLM.get_mamba_state_shape_from_config(temp_vllm_config)
1852
+
1853
+ @classmethod
1854
+ def get_mamba_state_copy_func(cls):
1855
+ return KimiLinearForCausalLM.get_mamba_state_copy_func()
1856
+
1857
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
1858
+ loader = AutoWeightsLoader(self)
1859
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/models/kimi_k3/nvidia/mtp.py
new
+443
@@ -0,0 +1,443 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inference-only Kimi-K3 Multi-Token-Prediction (MTP) draft model."""
4
+
5
+import copy
6
+from collections.abc import Iterable
7
+
8
+import torch
9
+import torch.nn as nn
10
+
11
+import vllm.envs as envs
12
+from vllm.config import VllmConfig
13
+from vllm.forward_context import get_forward_context, is_forward_context_available
14
+from vllm.logger import init_logger
15
+from vllm.model_executor.layers.fused_moe import (
16
+ fused_moe_make_expert_params_mapping,
17
+)
18
+from vllm.model_executor.layers.layernorm import RMSNorm
19
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
20
+from vllm.model_executor.layers.quantization import QuantizationConfig
21
+from vllm.model_executor.layers.vocab_parallel_embedding import (
22
+ ParallelLMHead,
23
+ VocabParallelEmbedding,
24
+)
25
+from vllm.model_executor.model_loader.weight_utils import (
26
+ default_weight_loader,
27
+ maybe_remap_kv_scale_name,
28
+)
29
+from vllm.model_executor.models.utils import get_pp_missing_layer_names, maybe_prefix
30
+from vllm.sequence import IntermediateTensors
31
+from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
32
+
33
+from ..common.mtp import fused_mtp_input
34
+from .low_latency_gemm import enable_kimi_k3_low_latency_gemm
35
+from .model import (
36
+ KimiDecoderLayer,
37
+ KimiMoE,
38
+ get_spec_layer_idx_from_weight_name,
39
+ make_kimi_k3_mega_moe_expert_params_mapping,
40
+)
41
+from .ops.sequence_parallel import sp_all_gather, sp_padding_mask, sp_shard
42
+
43
+logger = init_logger(__name__)
44
+
45
+
46
+class SharedHead(nn.Module):
47
+ def __init__(
48
+ self,
49
+ config: KimiLinearConfig,
50
+ prefix: str,
51
+ quant_config: QuantizationConfig | None = None,
52
+ ) -> None:
53
+ super().__init__()
54
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
55
+ self.head = ParallelLMHead(
56
+ config.vocab_size,
57
+ config.hidden_size,
58
+ quant_config=quant_config,
59
+ prefix=maybe_prefix(prefix, "head"),
60
+ )
61
+
62
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
63
+ return self.norm(hidden_states)
64
+
65
+
66
+class KimiK3MultiTokenPredictorLayer(nn.Module):
67
+ def __init__(
68
+ self,
69
+ config: KimiLinearConfig,
70
+ vllm_config: VllmConfig,
71
+ prefix: str,
72
+ ) -> None:
73
+ super().__init__()
74
+ self.config = config
75
+ quant_config = vllm_config.quant_config
76
+
77
+ self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
78
+ self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
79
+ self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
80
+
81
+ self.shared_head = SharedHead(
82
+ config=config, prefix=prefix, quant_config=quant_config
83
+ )
84
+ # The MTP block starts without the base model's AttnRes state.
85
+ block_config = copy.copy(config)
86
+ block_config.attn_res_block_size = None
87
+ # NOTE: the prefix must end in the numeric spec-layer index so that
88
+ # KimiDecoderLayer can parse ``layer_idx`` and pick MLA (full attn).
89
+ # Own aux stream for the MLA g_proj output-gate overlap (DeepseekV4
90
+ # convention: the MTP block creates its own stream).
91
+ aux_stream = torch.cuda.Stream()
92
+ self.mtp_block = KimiDecoderLayer(
93
+ block_config, vllm_config, prefix=prefix, aux_stream=aux_stream
94
+ )
95
+
96
+ def forward(
97
+ self,
98
+ input_ids: torch.Tensor,
99
+ positions: torch.Tensor,
100
+ previous_hidden_states: torch.Tensor,
101
+ inputs_embeds: torch.Tensor | None = None,
102
+ spec_step_index: int = 0,
103
+ ) -> tuple[torch.Tensor, torch.Tensor]:
104
+ assert inputs_embeds is not None
105
+ hidden_states = self.eh_proj(
106
+ fused_mtp_input(
107
+ positions,
108
+ inputs_embeds,
109
+ previous_hidden_states,
110
+ self.enorm.weight,
111
+ self.hnorm.weight,
112
+ self.enorm.variance_epsilon,
113
+ )
114
+ )
115
+ if self.mtp_block.use_sequence_parallel:
116
+ if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
117
+ forward_context = get_forward_context()
118
+ forward_context.is_padding = sp_padding_mask(
119
+ forward_context.is_padding, hidden_states
120
+ )
121
+ hidden_states = sp_shard(hidden_states)
122
+
123
+ hidden_states, _, residual = self.mtp_block(
124
+ positions=positions,
125
+ hidden_states=hidden_states,
126
+ residual=None,
127
+ )
128
+ if self.mtp_block.use_sequence_parallel:
129
+ assert residual is not None
130
+ hidden_states = hidden_states + residual
131
+ hidden_states = sp_all_gather(hidden_states)[: positions.shape[0]]
132
+ logits_hidden_states = self.shared_head.norm(hidden_states)
133
+ return logits_hidden_states, hidden_states
134
+
135
+ # Produce the normalized logits input and the pre-norm recurrent state
136
+ # in one fused add-RMSNorm launch.
137
+ logits_hidden_states, hidden_states = self.shared_head.norm(
138
+ hidden_states, residual
139
+ )
140
+ return logits_hidden_states, hidden_states
141
+
142
+
143
+class KimiK3MultiTokenPredictor(nn.Module):
144
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
145
+ super().__init__()
146
+ config: KimiLinearConfig = vllm_config.model_config.hf_text_config
147
+ self.config = config
148
+ self.mtp_start_layer_idx = config.num_hidden_layers
149
+ self.num_mtp_layers = config.num_nextn_predict_layers
150
+
151
+ self.layers = torch.nn.ModuleDict(
152
+ {
153
+ str(idx): KimiK3MultiTokenPredictorLayer(
154
+ config, vllm_config, f"{prefix}.layers.{idx}"
155
+ )
156
+ for idx in range(
157
+ self.mtp_start_layer_idx,
158
+ self.mtp_start_layer_idx + self.num_mtp_layers,
159
+ )
160
+ }
161
+ )
162
+ self.embed_tokens = VocabParallelEmbedding(
163
+ config.vocab_size,
164
+ config.hidden_size,
165
+ prefix=maybe_prefix(prefix, "embed_tokens"),
166
+ )
167
+ self.logits_processor = LogitsProcessor(config.vocab_size)
168
+
169
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
170
+ return self.embed_tokens(input_ids)
171
+
172
+ def forward(
173
+ self,
174
+ input_ids: torch.Tensor,
175
+ positions: torch.Tensor,
176
+ previous_hidden_states: torch.Tensor,
177
+ inputs_embeds: torch.Tensor | None = None,
178
+ spec_step_idx: int = 0,
179
+ ) -> tuple[torch.Tensor, torch.Tensor]:
180
+ if inputs_embeds is None:
181
+ inputs_embeds = self.embed_tokens(input_ids)
182
+ current_step_idx = spec_step_idx % self.num_mtp_layers
183
+ return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
184
+ input_ids,
185
+ positions,
186
+ previous_hidden_states,
187
+ inputs_embeds,
188
+ current_step_idx,
189
+ )
190
+
191
+ def compute_logits(
192
+ self,
193
+ hidden_states: torch.Tensor,
194
+ spec_step_idx: int = 0,
195
+ ) -> torch.Tensor:
196
+ current_step_idx = spec_step_idx % self.num_mtp_layers
197
+ mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
198
+ logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states)
199
+ return logits
200
+
201
+
202
+class KimiK3MTP(nn.Module):
203
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
204
+ super().__init__()
205
+ self.config = vllm_config.model_config.hf_text_config
206
+ self.quant_config = vllm_config.quant_config
207
+ self.model = KimiK3MultiTokenPredictor(
208
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
209
+ )
210
+ enable_kimi_k3_low_latency_gemm(self, vllm_config.model_config.dtype)
211
+
212
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
213
+ return self.model.embed_input_ids(input_ids)
214
+
215
+ def forward(
216
+ self,
217
+ input_ids: torch.Tensor | None,
218
+ positions: torch.Tensor,
219
+ hidden_states: torch.Tensor,
220
+ intermediate_tensors: IntermediateTensors | None = None,
221
+ inputs_embeds: torch.Tensor | None = None,
222
+ spec_step_idx: int = 0,
223
+ ) -> tuple[torch.Tensor, torch.Tensor]:
224
+ return self.model(
225
+ input_ids,
226
+ positions,
227
+ hidden_states,
228
+ inputs_embeds,
229
+ spec_step_idx,
230
+ )
231
+
232
+ def compute_logits(
233
+ self,
234
+ hidden_states: torch.Tensor,
235
+ spec_step_idx: int = 0,
236
+ ) -> torch.Tensor | None:
237
+ return self.model.compute_logits(hidden_states, spec_step_idx)
238
+
239
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
240
+ # Mirror KimiLinearForCausalLM.load_weights naming: leading-dot shard
241
+ # names, q_lora-conditional fused QKV, and w1/w2/w3 expert weights.
242
+ kda_config = self.config.linear_attn_config
243
+ use_full_rank_gate = bool(
244
+ kda_config and kda_config.get("use_full_rank_gate", False)
245
+ )
246
+ beta_shard_id = 5 if use_full_rank_gate else 3
247
+ stacked_params_mapping = [
248
+ # (param_name, shard_name, shard_id)
249
+ (".in_proj_qkvgfab", ".q_proj", 0),
250
+ (".in_proj_qkvgfab", ".k_proj", 1),
251
+ (".in_proj_qkvgfab", ".v_proj", 2),
252
+ (".in_proj_qkvgfab", ".b_proj", beta_shard_id),
253
+ (".in_proj_qkvgfab", ".f_a_proj", 4),
254
+ (".conv1d", ".q_conv1d", 0),
255
+ (".conv1d", ".k_conv1d", 1),
256
+ (".conv1d", ".v_conv1d", 2),
257
+ (".gate_up_proj", ".gate_proj", 0),
258
+ (".gate_up_proj", ".up_proj", 1),
259
+ ]
260
+ if use_full_rank_gate:
261
+ stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3))
262
+ if getattr(self.config, "q_lora_rank", None) is not None:
263
+ stacked_params_mapping += [
264
+ (".fused_qkv_a_proj", ".q_a_proj", 0),
265
+ (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
266
+ ]
267
+
268
+ use_mega_moe = any(
269
+ module.use_mega_moe
270
+ for module in self.modules()
271
+ if isinstance(module, KimiMoE)
272
+ )
273
+ if self.config.is_moe and use_mega_moe:
274
+ expert_params_mapping = make_kimi_k3_mega_moe_expert_params_mapping(
275
+ self.config.num_experts
276
+ )
277
+ elif self.config.is_moe:
278
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
279
+ self,
280
+ ckpt_gate_proj_name="w1",
281
+ ckpt_down_proj_name="w2",
282
+ ckpt_up_proj_name="w3",
283
+ num_experts=self.config.num_experts,
284
+ )
285
+ else:
286
+ expert_params_mapping = []
287
+
288
+ pp_missing_layer_names = get_pp_missing_layer_names(self)
289
+ params_dict = dict(self.named_parameters())
290
+ # Under the MXFP4 quant interface the routed experts register unpacked
291
+ # params (``w13_weight``), while the compressed-tensors checkpoint names
292
+ # them ``.weight_packed``. Rebind so the expert mapping resolves; scales
293
+ # already share the ``.weight_scale`` suffix.
294
+ experts_unpacked = not use_mega_moe and not any(
295
+ n.endswith("w13_weight_packed") for n in params_dict
296
+ )
297
+ loaded_params: set[str] = set()
298
+ for name, loaded_weight in weights:
299
+ if "rotary_emb.inv_freq" in name:
300
+ continue
301
+ # The multimodal checkpoint prefixes text weights with
302
+ # ``language_model.``; strip it so names match this draft model's
303
+ # parameter paths (``model.layers.{i}.``). Non-text weights
304
+ # (vision_tower, mm_projector, ...) never match a spec layer below.
305
+ if name.startswith("language_model."):
306
+ name = name[len("language_model.") :]
307
+ if experts_unpacked and name.endswith(".weight_packed"):
308
+ name = name.replace(".weight_packed", ".weight")
309
+ spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
310
+ if spec_layer is None:
311
+ continue
312
+ name = self._rewrite_spec_layer_name(spec_layer, name)
313
+
314
+ for param_name, weight_name, shard_id in stacked_params_mapping:
315
+ if weight_name not in name:
316
+ continue
317
+ # Routed experts (``.experts.{i}.w1/w2/w3``) are handled by the
318
+ # expert mapping below; skip them here. Shared experts
319
+ # (``.shared_experts.``) use gate/up_proj and fall through.
320
+ if ".experts." in name:
321
+ continue
322
+ name_mapped = name.replace(weight_name, param_name)
323
+ # Only take this mapping if the fused destination actually
324
+ # exists (e.g. QKV fusion is only present when q_lora is used).
325
+ if name_mapped not in params_dict:
326
+ continue
327
+ if name_mapped in pp_missing_layer_names:
328
+ continue
329
+ name = name_mapped
330
+ param = params_dict[name]
331
+ weight_loader = param.weight_loader
332
+ weight_loader(param, loaded_weight, shard_id)
333
+ break
334
+ else:
335
+ for (
336
+ expert_param_name,
337
+ expert_weight_name,
338
+ expert_id,
339
+ expert_shard_id,
340
+ ) in expert_params_mapping:
341
+ if expert_weight_name not in name:
342
+ continue
343
+ name_mapped = name.replace(expert_weight_name, expert_param_name)
344
+ if name_mapped in pp_missing_layer_names:
345
+ continue
346
+ param = params_dict[name_mapped]
347
+ weight_loader = param.weight_loader
348
+ weight_loader(
349
+ param,
350
+ loaded_weight,
351
+ name_mapped,
352
+ shard_id=expert_shard_id,
353
+ expert_id=expert_id,
354
+ )
355
+ name = name_mapped
356
+ break
357
+ else:
358
+ if name.endswith(".bias") and name not in params_dict:
359
+ continue
360
+ remapped_name = maybe_remap_kv_scale_name(name, params_dict)
361
+ if remapped_name is None:
362
+ continue
363
+ name = remapped_name
364
+
365
+ # The embedding is shared across MTP layers; only the first
366
+ # spec layer carries the hoisted (non-".layers") copy.
367
+ if spec_layer != self.model.mtp_start_layer_idx and (
368
+ ".layers" not in name
369
+ ):
370
+ continue
371
+ if name in pp_missing_layer_names:
372
+ continue
373
+ # The base model uses an attn-residual scheme whose per-layer
374
+ # weights (self_attention_res_*, mlp_res_*) are not used by
375
+ # the draft block; such names have no matching parameter and
376
+ # are safely skipped.
377
+ if name not in params_dict:
378
+ continue
379
+
380
+ param = params_dict[name]
381
+ weight_loader = getattr(
382
+ param, "weight_loader", default_weight_loader
383
+ )
384
+ weight_loader(param, loaded_weight)
385
+ loaded_params.add(name)
386
+
387
+ # Validate that weights were loaded for each expected MTP layer.
388
+ loaded_layers: set[int] = set()
389
+ for param_name in loaded_params:
390
+ spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name)
391
+ if spec_layer is not None:
392
+ loaded_layers.add(spec_layer)
393
+ for layer_idx in range(
394
+ self.model.mtp_start_layer_idx,
395
+ self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
396
+ ):
397
+ if layer_idx not in loaded_layers:
398
+ raise ValueError(
399
+ f"MTP speculative decoding layer {layer_idx} weights "
400
+ f"missing from checkpoint. The checkpoint may not include "
401
+ f"the MTP layer weights. Use a checkpoint that includes "
402
+ f"MTP layer weights, or disable speculative decoding."
403
+ )
404
+
405
+ if use_mega_moe:
406
+ for module in self.modules():
407
+ if isinstance(module, KimiMoE) and module.use_mega_moe:
408
+ module.experts.finalize_weights()
409
+
410
+ return loaded_params
411
+
412
+ def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
413
+ """Rewrite a checkpoint weight name to this module's parameter path.
414
+
415
+ Top-level MTP submodules (enorm/hnorm/eh_proj/shared_head) stay under
416
+ ``model.layers.{spec_layer}.*``; the shared ``embed_tokens`` is hoisted
417
+ to ``model.*``; everything else is a transformer-block weight and gets
418
+ ``.mtp_block`` inserted.
419
+ """
420
+ spec_layer_weight_names = [
421
+ "embed_tokens",
422
+ "enorm",
423
+ "hnorm",
424
+ "eh_proj",
425
+ "shared_head",
426
+ ]
427
+ shared_weight_names = ["embed_tokens"]
428
+ spec_layer_weight = False
429
+ shared_weight = False
430
+ for weight_name in spec_layer_weight_names:
431
+ if weight_name in name:
432
+ spec_layer_weight = True
433
+ if weight_name in shared_weight_names:
434
+ shared_weight = True
435
+ break
436
+ if not spec_layer_weight:
437
+ name = name.replace(
438
+ f"model.layers.{spec_layer}.",
439
+ f"model.layers.{spec_layer}.mtp_block.",
440
+ )
441
+ elif shared_weight:
442
+ name = name.replace(f"model.layers.{spec_layer}.", "model.")
443
+ return name
vllm/models/kimi_k3/nvidia/ops/cute_dsl/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/__init__.py
new
+14
@@ -0,0 +1,14 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""CuTe DSL kernels for KimiK3LatentMoETailOp."""
5
+
6
+from .allreduce_rmsnorm_reduce_scatter_early_exit import CollectiveKernel
7
+from .fused_add_multicast_gemm import AdaptiveUpProjectionKernel
8
+from .lamport_copy import LamportCopyKernel
9
+
10
+__all__ = [
11
+ "AdaptiveUpProjectionKernel",
12
+ "CollectiveKernel",
13
+ "LamportCopyKernel",
14
+]
vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py
new
+1012
@@ -0,0 +1,1012 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# Specialized from FlashInfer's oneshotAllreduceFusionKernel.
4
+
5
+"""Routed AllReduce/RMSNorm with CTA-specialized ReduceScatter early exit."""
6
+
7
+from __future__ import annotations
8
+
9
+from typing import Any
10
+
11
+import cuda.bindings.driver as cuda
12
+import cutlass
13
+import cutlass.cute as cute
14
+import torch
15
+import torch.distributed as dist
16
+import torch.distributed._symmetric_memory as symm_mem
17
+from cutlass import BFloat16, Float32, Int32, Int64, Uint32
18
+
19
+from .primitives import (
20
+ NUM_LAMPORT_BUFFERS,
21
+ PACKED_BYTES,
22
+ VEC_BF16,
23
+ bf16x8_to_packed_u32x4,
24
+ block_sum_specialized,
25
+ fragment_is_dirty,
26
+ load_global_u32x4,
27
+ load_volatile_u32,
28
+ map_shared_to_peer,
29
+ packed_u32x4_to_bf16x8,
30
+ red_async_release_gpu_add_u32,
31
+ sanitize_negative_zero,
32
+ store_global_u32x4,
33
+ store_lamport_sentinel_128,
34
+ store_shared_cluster_f32,
35
+ to_cute,
36
+ to_cute_dynamic_m,
37
+)
38
+
39
+
40
+def _mapping(
41
+ tp_size: int,
42
+ latent_dim: int,
43
+ hidden_dim: int,
44
+) -> tuple[int, int, int, int]:
45
+ shard_dim = hidden_dim // tp_size
46
+ cluster_ctas = latent_dim // shard_dim
47
+ threads = shard_dim // VEC_BF16
48
+ shared_roles = tp_size // cluster_ctas
49
+ return shard_dim, cluster_ctas, threads, shared_roles
50
+
51
+
52
+def validate_shape(*, tp_size: int, latent_dim: int, hidden_dim: int) -> None:
53
+ """Validate constraints imposed by the fused collective mapping."""
54
+
55
+ if tp_size <= 0 or latent_dim <= 0 or hidden_dim <= 0:
56
+ raise ValueError("collective dimensions must be positive")
57
+ if hidden_dim % tp_size:
58
+ raise ValueError("hidden_dim must be divisible by tp_size")
59
+ shard, cluster, threads, _ = _mapping(tp_size, latent_dim, hidden_dim)
60
+ if shard % VEC_BF16:
61
+ raise ValueError("hidden_dim / tp_size must be divisible by 8")
62
+ if latent_dim % shard:
63
+ raise ValueError("latent_dim must be an integer multiple of shard_dim")
64
+ if cluster > 16 or cluster & (cluster - 1):
65
+ raise ValueError("collective cluster width must be a power of two <= 16")
66
+ if tp_size % cluster:
67
+ raise ValueError("tp_size must be divisible by collective cluster width")
68
+ if not 32 <= threads <= 1024:
69
+ raise ValueError("collective threads per CTA must be in [32, 1024]")
70
+ if threads < cluster:
71
+ raise ValueError("collective threads must cover every cluster CTA")
72
+
73
+
74
+def _select_routed_schedule(
75
+ tp_size: int,
76
+ latent_dim: int,
77
+ hidden_dim: int,
78
+ max_m: int,
79
+) -> tuple[int, int]:
80
+ """Match upstream MNNVL's one-cluster-per-token occupancy policy."""
81
+
82
+ _, cluster_ctas, threads, _ = _mapping(tp_size, latent_dim, hidden_dim)
83
+ sm_count = torch.cuda.get_device_properties(
84
+ torch.accelerator.current_device_index()
85
+ ).multi_processor_count
86
+ while max_m * cluster_ctas > sm_count and cluster_ctas > 1 and threads <= 512:
87
+ cluster_ctas //= 2
88
+ threads *= 2
89
+ return threads, cluster_ctas
90
+
91
+
92
+class AllReduceRMSNormWithReduceScatterEarlyExit:
93
+ """One routed role plus one ReduceScatter role per destination group."""
94
+
95
+ def __init__(
96
+ self,
97
+ *,
98
+ rank: int,
99
+ tp_size: int,
100
+ latent_dim: int,
101
+ hidden_dim: int,
102
+ max_m: int,
103
+ max_token_ctas: int,
104
+ fp32_internal: bool = False,
105
+ include_reduce_scatter: bool = True,
106
+ include_routed: bool = True,
107
+ ):
108
+ validate_shape(
109
+ tp_size=tp_size,
110
+ latent_dim=latent_dim,
111
+ hidden_dim=hidden_dim,
112
+ )
113
+ if not 0 <= rank < tp_size:
114
+ raise ValueError(f"rank must be in [0,{tp_size}), got {rank}")
115
+ if not include_routed and not include_reduce_scatter:
116
+ raise ValueError("at least one collective role must be enabled")
117
+ self.rank = rank
118
+ self.tp_size = tp_size
119
+ self.latent_dim = latent_dim
120
+ self.hidden_dim = hidden_dim
121
+ (
122
+ self.shard_dim,
123
+ mapped_cluster,
124
+ mapped_threads,
125
+ shared_roles,
126
+ ) = _mapping(tp_size, latent_dim, hidden_dim)
127
+ if include_reduce_scatter:
128
+ self.cluster_ctas = mapped_cluster
129
+ self.threads = mapped_threads
130
+ else:
131
+ self.threads, self.cluster_ctas = _select_routed_schedule(
132
+ tp_size, latent_dim, hidden_dim, max_m
133
+ )
134
+ # Diagnostic specialization: a one-role grid executes only the routed
135
+ # AllReduce/RMSNorm path. Keep this compile-time so the production
136
+ # fused path is unchanged when include_reduce_scatter=True.
137
+ if include_routed and include_reduce_scatter:
138
+ self.roles = 1 + shared_roles
139
+ elif include_routed:
140
+ self.roles = 1
141
+ else:
142
+ self.roles = shared_roles
143
+ self.warps = (self.threads + 31) // 32
144
+ self.last_warp_lanes = self.threads - (self.warps - 1) * 32
145
+ self.last_warp_mask = (1 << self.last_warp_lanes) - 1
146
+ # The upstream MNNVL oneshot protocol assigns one cluster to each
147
+ # token. Keep that exact ownership in the routed-only diagnostic;
148
+ # reusing a cluster for multiple token waves allows the Lamport
149
+ # generation metadata to change between waves.
150
+ self.token_ctas = (
151
+ min(max_m, max_token_ctas) if include_reduce_scatter else max_m
152
+ )
153
+ self.fp32_internal = fp32_internal
154
+ self.include_reduce_scatter = include_reduce_scatter
155
+ self.include_routed = include_routed
156
+
157
+ @cute.jit
158
+ def __call__(
159
+ self,
160
+ latent_source: cute.Tensor,
161
+ gamma: cute.Tensor,
162
+ latent_output: cute.Tensor,
163
+ routed_workspace: cute.Tensor,
164
+ latent_flags: cute.Tensor,
165
+ latent_multicast_ptr: Int64,
166
+ shared_source: cute.Tensor,
167
+ shared_output: cute.Tensor,
168
+ shared_workspace: cute.Tensor,
169
+ shared_flags: cute.Tensor,
170
+ shared_peer_ptrs: cute.Tensor,
171
+ m: Int32,
172
+ epsilon: Float32,
173
+ stream: cuda.CUstream,
174
+ ):
175
+ self.kernel(
176
+ latent_source,
177
+ gamma,
178
+ latent_output,
179
+ routed_workspace,
180
+ latent_flags,
181
+ latent_multicast_ptr,
182
+ shared_source,
183
+ shared_output,
184
+ shared_workspace,
185
+ shared_flags,
186
+ shared_peer_ptrs,
187
+ m,
188
+ epsilon,
189
+ ).launch(
190
+ grid=(self.token_ctas, self.cluster_ctas, self.roles),
191
+ block=(self.threads, 1, 1),
192
+ cluster=(1, self.cluster_ctas, 1),
193
+ smem=(self.warps + self.cluster_ctas) * 4,
194
+ stream=stream,
195
+ use_pdl=True,
196
+ )
197
+
198
+ @cute.kernel
199
+ def kernel(
200
+ self,
201
+ latent_source: cute.Tensor,
202
+ gamma: cute.Tensor,
203
+ latent_output: cute.Tensor,
204
+ routed_workspace: cute.Tensor,
205
+ latent_flags: cute.Tensor,
206
+ latent_multicast_ptr: Int64,
207
+ shared_source: cute.Tensor,
208
+ shared_output: cute.Tensor,
209
+ shared_workspace: cute.Tensor,
210
+ shared_flags: cute.Tensor,
211
+ shared_peer_ptrs: cute.Tensor,
212
+ m: Int32,
213
+ epsilon: Float32,
214
+ ):
215
+ tidx, _, _ = cute.arch.thread_idx()
216
+ token_cta, cta_y, role = cute.arch.block_idx()
217
+ logical_role = role
218
+ if cutlass.const_expr(not self.include_routed):
219
+ # A shared-only grid starts at z=0, while _token_device reserves
220
+ # logical role 0 for the routed collective.
221
+ logical_role = role + Int32(1)
222
+ cluster_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
223
+
224
+ cute.arch.griddepcontrol_wait()
225
+ token = token_cta
226
+ while token < m:
227
+ self._token_device(
228
+ latent_source,
229
+ gamma,
230
+ latent_output,
231
+ routed_workspace,
232
+ latent_flags,
233
+ latent_multicast_ptr,
234
+ shared_source,
235
+ shared_output,
236
+ shared_workspace,
237
+ shared_flags,
238
+ shared_peer_ptrs,
239
+ m,
240
+ epsilon,
241
+ token,
242
+ token_cta,
243
+ cta_y,
244
+ logical_role,
245
+ cluster_rank,
246
+ tidx,
247
+ )
248
+ token = token + self.token_ctas
249
+
250
+ @cute.jit
251
+ def _token_device(
252
+ self,
253
+ latent_source: cute.Tensor,
254
+ gamma: cute.Tensor,
255
+ latent_output: cute.Tensor,
256
+ routed_workspace: cute.Tensor,
257
+ latent_flags: cute.Tensor,
258
+ latent_multicast_ptr: Int64,
259
+ shared_source: cute.Tensor,
260
+ shared_output: cute.Tensor,
261
+ shared_workspace: cute.Tensor,
262
+ shared_flags: cute.Tensor,
263
+ shared_peer_ptrs: cute.Tensor,
264
+ m: Int32,
265
+ epsilon: Float32,
266
+ token: Int32,
267
+ token_cta: Int32,
268
+ cta_y: Int32,
269
+ role: Int32,
270
+ cluster_rank: Int32,
271
+ tidx: Int32,
272
+ ):
273
+ if role == 0:
274
+ # ---------------- routed AllReduce + RMSNorm ----------------
275
+ packed_idx = cluster_rank * self.threads + tidx
276
+ element_offset = (
277
+ Int64(token) * self.latent_dim + Int64(packed_idx) * VEC_BF16
278
+ )
279
+ current_index = cute.arch.load((latent_flags.iterator + 0).llvm_ptr, Uint32)
280
+ dirty_index = cute.arch.load((latent_flags.iterator + 1).llvm_ptr, Uint32)
281
+ bytes_per_buffer = cute.arch.load(
282
+ (latent_flags.iterator + 2).llvm_ptr, Uint32
283
+ )
284
+ dirty_num_stages = cute.arch.load(
285
+ (latent_flags.iterator + 3).llvm_ptr, Uint32
286
+ )
287
+ bytes_to_clear = cute.arch.load(
288
+ (latent_flags.iterator + 4).llvm_ptr, Uint32
289
+ )
290
+ current_elements = Int64(current_index) * (
291
+ Int64(bytes_per_buffer) // Int64(2)
292
+ )
293
+ dirty_elements = Int64(dirty_index) * (Int64(bytes_per_buffer) // Int64(2))
294
+
295
+ local_ptr = cute.make_ptr(
296
+ BFloat16,
297
+ (latent_source.iterator + element_offset).llvm_ptr,
298
+ cute.AddressSpace.gmem,
299
+ assumed_align=16,
300
+ )
301
+ local_packed = sanitize_negative_zero(
302
+ load_global_u32x4(local_ptr, volatile=False)
303
+ )
304
+ multicast_offset = (
305
+ Int64(current_index) * Int64(bytes_per_buffer)
306
+ + (
307
+ (Int64(token) * self.tp_size + self.rank) * self.latent_dim
308
+ + Int64(packed_idx) * VEC_BF16
309
+ )
310
+ * 2
311
+ )
312
+ store_global_u32x4(
313
+ latent_multicast_ptr + multicast_offset,
314
+ local_packed,
315
+ volatile=False,
316
+ )
317
+
318
+ cute.arch.cluster_arrive()
319
+ if cluster_rank == 0 and tidx < 32:
320
+ cute.arch.cluster_wait()
321
+ if tidx == 0:
322
+ red_async_release_gpu_add_u32(latent_flags.iterator + 8, Uint32(1))
323
+
324
+ global_tid = (
325
+ Int64(token) * self.cluster_ctas + Int64(cta_y)
326
+ ) * self.threads + Int64(tidx)
327
+ total_threads = Int64(m) * self.cluster_ctas * self.threads
328
+ clear_fragments = (Int64(bytes_to_clear) + PACKED_BYTES - 1) // PACKED_BYTES
329
+ clear_idx = global_tid
330
+ if dirty_num_stages > Uint32(0):
331
+ while clear_idx < clear_fragments:
332
+ clear_ptr = cute.make_ptr(
333
+ BFloat16,
334
+ (
335
+ routed_workspace.iterator
336
+ + dirty_elements
337
+ + clear_idx * VEC_BF16
338
+ ).llvm_ptr,
339
+ cute.AddressSpace.gmem,
340
+ assumed_align=16,
341
+ )
342
+ store_lamport_sentinel_128(clear_ptr)
343
+ clear_idx = clear_idx + total_threads
344
+
345
+ rank_words = cute.make_rmem_tensor(
346
+ cute.make_layout((self.tp_size, 4), stride=(4, 1)), Uint32
347
+ )
348
+ for word in cutlass.range_constexpr(4):
349
+ rank_words[self.rank, word] = local_packed[word]
350
+ valid = False
351
+ while not valid:
352
+ valid = True
353
+ for source_rank in cutlass.range_constexpr(self.tp_size):
354
+ if cutlass.const_expr(source_rank != self.rank):
355
+ remote_element = current_elements + (
356
+ (Int64(token) * self.tp_size + source_rank)
357
+ * self.latent_dim
358
+ + Int64(packed_idx) * VEC_BF16
359
+ )
360
+ remote_ptr = cute.make_ptr(
361
+ BFloat16,
362
+ (routed_workspace.iterator + remote_element).llvm_ptr,
363
+ cute.AddressSpace.gmem,
364
+ assumed_align=16,
365
+ )
366
+ remote = load_global_u32x4(remote_ptr, volatile=True)
367
+ for word in cutlass.range_constexpr(4):
368
+ rank_words[source_rank, word] = remote[word]
369
+ valid = valid & (not fragment_is_dirty(remote))
370
+
371
+ accum = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32)
372
+ for element in cutlass.range_constexpr(VEC_BF16):
373
+ accum[element] = Float32(0.0)
374
+ for source_rank in cutlass.range_constexpr(self.tp_size):
375
+ values = packed_u32x4_to_bf16x8(
376
+ rank_words[source_rank, None].load()
377
+ ).to(Float32)
378
+ for element in cutlass.range_constexpr(VEC_BF16):
379
+ accum[element] = accum[element] + values[element]
380
+
381
+ # Preserve the original early PDL point before RMSNorm.
382
+ cute.arch.griddepcontrol_launch_dependents()
383
+
384
+ if cutlass.const_expr(self.fp32_internal):
385
+ # High-precision fused mode: retain the rank reduction in
386
+ # FP32 through the RMS square and row reduction.
387
+ norm_input = accum.load()
388
+ norm_square = norm_input * norm_input
389
+ else:
390
+ norm_input_bf16 = accum.load().to(BFloat16)
391
+ norm_input = norm_input_bf16.to(Float32)
392
+ # Upstream-compatible mode: FlashInfer evaluates BF16 *
393
+ # BF16 first, then promotes the rounded square to FP32.
394
+ norm_square = (norm_input_bf16 * norm_input_bf16).to(Float32)
395
+ thread_sum = norm_square.reduce(
396
+ cute.ReductionOp.ADD,
397
+ init_val=Float32(0.0),
398
+ reduction_profile=0,
399
+ )
400
+ smem = cutlass.utils.SmemAllocator()
401
+ warp_sums = smem.allocate_tensor(
402
+ Float32, cute.make_layout((self.warps,)), byte_alignment=4
403
+ )
404
+ cluster_sums = smem.allocate_tensor(
405
+ Float32,
406
+ cute.make_layout((self.cluster_ctas,)),
407
+ byte_alignment=4,
408
+ )
409
+ block_sum = block_sum_specialized(
410
+ thread_sum,
411
+ warp_sums,
412
+ tidx,
413
+ self.warps,
414
+ self.last_warp_lanes,
415
+ self.last_warp_mask,
416
+ )
417
+ if tidx < self.cluster_ctas:
418
+ local_slot = cluster_sums.iterator + cluster_rank
419
+ remote_slot = map_shared_to_peer(local_slot, Int32(tidx))
420
+ store_shared_cluster_f32(remote_slot, block_sum)
421
+ cute.arch.cluster_arrive()
422
+ cute.arch.cluster_wait()
423
+
424
+ full_sum = Float32(0.0)
425
+ for peer in cutlass.range_constexpr(self.cluster_ctas):
426
+ full_sum = full_sum + cluster_sums[peer]
427
+ inv_rms = cute.math.rsqrt(
428
+ full_sum / Float32(self.latent_dim) + epsilon, fastmath=True
429
+ )
430
+ gamma_ptr = cute.make_ptr(
431
+ BFloat16,
432
+ (gamma.iterator + Int64(packed_idx) * VEC_BF16).llvm_ptr,
433
+ cute.AddressSpace.gmem,
434
+ assumed_align=16,
435
+ )
436
+ gamma_values = packed_u32x4_to_bf16x8(
437
+ load_global_u32x4(gamma_ptr, volatile=False)
438
+ )
439
+ result = (norm_input * inv_rms * gamma_values.to(Float32)).to(BFloat16)
440
+ store_global_u32x4(
441
+ Int64((latent_output.iterator + element_offset).toint()),
442
+ bf16x8_to_packed_u32x4(result),
443
+ volatile=False,
444
+ )
445
+
446
+ # The x=0 CTA rotates only after reaching its final token wave.
447
+ # Waiting for all M arrivals then guarantees every token-wave CTA
448
+ # loaded the current generation before the metadata is advanced.
449
+ if (
450
+ token_cta == 0
451
+ and token + self.token_ctas >= m
452
+ and cta_y == 0
453
+ and tidx == 0
454
+ ):
455
+ access_counter = latent_flags.iterator + 8
456
+ arrived = load_volatile_u32(access_counter)
457
+ while arrived < Uint32(m):
458
+ arrived = load_volatile_u32(access_counter)
459
+ next_index = (current_index + Uint32(1)) % Uint32(NUM_LAMPORT_BUFFERS)
460
+ actual_bytes = Uint32(m) * Uint32(self.tp_size * self.latent_dim * 2)
461
+ cute.arch.store((latent_flags.iterator + 0).llvm_ptr, next_index)
462
+ cute.arch.store((latent_flags.iterator + 1).llvm_ptr, current_index)
463
+ cute.arch.store(
464
+ (latent_flags.iterator + 2).llvm_ptr,
465
+ bytes_per_buffer,
466
+ )
467
+ cute.arch.store((latent_flags.iterator + 3).llvm_ptr, Uint32(1))
468
+ cute.arch.store((latent_flags.iterator + 4).llvm_ptr, actual_bytes)
469
+ for index in cutlass.range_constexpr(5, 8):
470
+ cute.arch.store(
471
+ (latent_flags.iterator + index).llvm_ptr,
472
+ Uint32(0),
473
+ )
474
+ cute.arch.store(access_counter.llvm_ptr, Uint32(0))
475
+
476
+ else:
477
+ # ---------------- shared ReduceScatter ----------------
478
+ shared_group = role - 1
479
+ destination = shared_group * self.cluster_ctas + cluster_rank
480
+ current_index = cute.arch.load((shared_flags.iterator + 0).llvm_ptr, Uint32)
481
+ dirty_index = cute.arch.load((shared_flags.iterator + 1).llvm_ptr, Uint32)
482
+ bytes_per_buffer = cute.arch.load(
483
+ (shared_flags.iterator + 2).llvm_ptr, Uint32
484
+ )
485
+ dirty_num_stages = cute.arch.load(
486
+ (shared_flags.iterator + 3).llvm_ptr, Uint32
487
+ )
488
+ bytes_to_clear = cute.arch.load(
489
+ (shared_flags.iterator + 4).llvm_ptr, Uint32
490
+ )
491
+ current_elements = Int64(current_index) * (
492
+ Int64(bytes_per_buffer) // Int64(2)
493
+ )
494
+ dirty_elements = Int64(dirty_index) * (Int64(bytes_per_buffer) // Int64(2))
495
+
496
+ source_element = (
497
+ Int64(token) * self.hidden_dim
498
+ + Int64(destination) * self.shard_dim
499
+ + Int64(tidx) * VEC_BF16
500
+ )
501
+ source_ptr = cute.make_ptr(
502
+ BFloat16,
503
+ (shared_source.iterator + source_element).llvm_ptr,
504
+ cute.AddressSpace.gmem,
505
+ assumed_align=16,
506
+ )
507
+ local_packed = sanitize_negative_zero(
508
+ load_global_u32x4(source_ptr, volatile=False)
509
+ )
510
+ peer_base = cute.arch.load(
511
+ (shared_peer_ptrs.iterator + destination).llvm_ptr,
512
+ Int64,
513
+ )
514
+ destination_element = current_elements + (
515
+ (Int64(token) * self.tp_size + self.rank) * self.shard_dim
516
+ + Int64(tidx) * VEC_BF16
517
+ )
518
+ store_global_u32x4(
519
+ peer_base + destination_element * 2,
520
+ local_packed,
521
+ volatile=False,
522
+ )
523
+
524
+ # One arrival per shared destination group and token.
525
+ cute.arch.cluster_arrive()
526
+ if cluster_rank == 0 and tidx < 32:
527
+ cute.arch.cluster_wait()
528
+ if tidx == 0:
529
+ red_async_release_gpu_add_u32(shared_flags.iterator + 8, Uint32(1))
530
+
531
+ global_tid = (
532
+ Int64(token) * self.tp_size + Int64(destination)
533
+ ) * self.threads + Int64(tidx)
534
+ total_threads = Int64(m) * self.tp_size * self.threads
535
+ clear_fragments = (Int64(bytes_to_clear) + PACKED_BYTES - 1) // PACKED_BYTES
536
+ clear_idx = global_tid
537
+ if dirty_num_stages > Uint32(0):
538
+ while clear_idx < clear_fragments:
539
+ clear_ptr = cute.make_ptr(
540
+ BFloat16,
541
+ (
542
+ shared_workspace.iterator
543
+ + dirty_elements
544
+ + clear_idx * VEC_BF16
545
+ ).llvm_ptr,
546
+ cute.AddressSpace.gmem,
547
+ assumed_align=16,
548
+ )
549
+ store_lamport_sentinel_128(clear_ptr)
550
+ clear_idx = clear_idx + total_threads
551
+
552
+ if destination == self.rank:
553
+ rank_words = cute.make_rmem_tensor(
554
+ cute.make_layout((self.tp_size, 4), stride=(4, 1)), Uint32
555
+ )
556
+ valid = False
557
+ while not valid:
558
+ valid = True
559
+ for source_rank in cutlass.range_constexpr(self.tp_size):
560
+ remote_element = current_elements + (
561
+ (Int64(token) * self.tp_size + source_rank) * self.shard_dim
562
+ + Int64(tidx) * VEC_BF16
563
+ )
564
+ remote_ptr = cute.make_ptr(
565
+ BFloat16,
566
+ (shared_workspace.iterator + remote_element).llvm_ptr,
567
+ cute.AddressSpace.gmem,
568
+ assumed_align=16,
569
+ )
570
+ remote = load_global_u32x4(remote_ptr, volatile=True)
571
+ for word in cutlass.range_constexpr(4):
572
+ rank_words[source_rank, word] = remote[word]
573
+ valid = valid & (not fragment_is_dirty(remote))
574
+
575
+ accum = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32)
576
+ for element in cutlass.range_constexpr(VEC_BF16):
577
+ accum[element] = Float32(0.0)
578
+ for source_rank in cutlass.range_constexpr(self.tp_size):
579
+ values = packed_u32x4_to_bf16x8(
580
+ rank_words[source_rank, None].load()
581
+ ).to(Float32)
582
+ for element in cutlass.range_constexpr(VEC_BF16):
583
+ accum[element] = accum[element] + values[element]
584
+ result = accum.load().to(BFloat16)
585
+ output_element = (
586
+ Int64(token) * self.hidden_dim
587
+ + self.rank * self.shard_dim
588
+ + Int64(tidx) * VEC_BF16
589
+ )
590
+ store_global_u32x4(
591
+ Int64((shared_output.iterator + output_element).toint()),
592
+ bf16x8_to_packed_u32x4(result),
593
+ volatile=False,
594
+ )
595
+
596
+ if destination == self.rank:
597
+ cute.arch.barrier()
598
+
599
+ cute.arch.griddepcontrol_launch_dependents()
600
+
601
+ if (
602
+ token_cta == 0
603
+ and token + self.token_ctas >= m
604
+ and shared_group == 0
605
+ and cluster_rank == 0
606
+ and tidx == 0
607
+ ):
608
+ access_counter = shared_flags.iterator + 8
609
+ arrived = load_volatile_u32(access_counter)
610
+ target = Uint32(m) * Uint32(self.tp_size // self.cluster_ctas)
611
+ while arrived < target:
612
+ arrived = load_volatile_u32(access_counter)
613
+ next_index = (current_index + Uint32(1)) % Uint32(NUM_LAMPORT_BUFFERS)
614
+ actual_bytes = Uint32(m) * Uint32(self.tp_size * self.shard_dim * 2)
615
+ cute.arch.store((shared_flags.iterator + 0).llvm_ptr, next_index)
616
+ cute.arch.store((shared_flags.iterator + 1).llvm_ptr, current_index)
617
+ cute.arch.store(
618
+ (shared_flags.iterator + 2).llvm_ptr,
619
+ bytes_per_buffer,
620
+ )
621
+ cute.arch.store((shared_flags.iterator + 3).llvm_ptr, Uint32(1))
622
+ cute.arch.store((shared_flags.iterator + 4).llvm_ptr, actual_bytes)
623
+ for index in cutlass.range_constexpr(5, 8):
624
+ cute.arch.store(
625
+ (shared_flags.iterator + index).llvm_ptr,
626
+ Uint32(0),
627
+ )
628
+ cute.arch.store(access_counter.llvm_ptr, Uint32(0))
629
+
630
+
631
+_COMPILED: dict[tuple[object, ...], Any] = {}
632
+
633
+
634
+def _routed_workspace_cute(workspace: torch.Tensor):
635
+ return to_cute(workspace.view(torch.bfloat16), 16)
636
+
637
+
638
+def _compile_key(
639
+ rank: int,
640
+ tp_size: int,
641
+ latent_dim: int,
642
+ hidden_dim: int,
643
+ max_m: int,
644
+ max_token_ctas: int,
645
+ fp32_internal: bool,
646
+ include_reduce_scatter: bool,
647
+ include_routed: bool,
648
+):
649
+ return (
650
+ torch.accelerator.current_device_index(),
651
+ rank,
652
+ tp_size,
653
+ latent_dim,
654
+ hidden_dim,
655
+ max_m,
656
+ max_token_ctas,
657
+ fp32_internal,
658
+ include_reduce_scatter,
659
+ include_routed,
660
+ )
661
+
662
+
663
+def _runtime_args(
664
+ latent_source: torch.Tensor,
665
+ gamma: torch.Tensor,
666
+ latent_output: torch.Tensor,
667
+ routed_workspace: torch.Tensor,
668
+ routed_flags: torch.Tensor,
669
+ routed_multicast_ptr: int,
670
+ shared_source: torch.Tensor,
671
+ shared_output: torch.Tensor,
672
+ shared_workspace: torch.Tensor,
673
+ shared_flags: torch.Tensor,
674
+ shared_peer_ptrs: torch.Tensor,
675
+ rms_eps: float,
676
+):
677
+ return (
678
+ to_cute_dynamic_m(latent_source, mode=0, assumed_align=16),
679
+ to_cute(gamma, 16),
680
+ to_cute(latent_output, 16),
681
+ _routed_workspace_cute(routed_workspace),
682
+ to_cute(routed_flags, 16),
683
+ Int64(routed_multicast_ptr),
684
+ to_cute_dynamic_m(shared_source, mode=0, assumed_align=16),
685
+ to_cute(shared_output, 16),
686
+ to_cute(shared_workspace, 16),
687
+ to_cute(shared_flags, 16),
688
+ to_cute(shared_peer_ptrs, 16),
689
+ Int32(latent_source.shape[0]),
690
+ Float32(rms_eps),
691
+ cuda.CUstream(torch.cuda.current_stream(latent_source.device).cuda_stream),
692
+ )
693
+
694
+
695
+def compile_kernel(
696
+ *,
697
+ rank: int,
698
+ tp_size: int,
699
+ latent_dim: int,
700
+ hidden_dim: int,
701
+ max_m: int,
702
+ max_token_ctas: int,
703
+ latent_output: torch.Tensor,
704
+ routed_workspace: torch.Tensor,
705
+ routed_flags: torch.Tensor,
706
+ routed_multicast_ptr: int,
707
+ shared_output: torch.Tensor,
708
+ shared_workspace: torch.Tensor,
709
+ shared_flags: torch.Tensor,
710
+ shared_peer_ptrs: torch.Tensor,
711
+ rms_eps: float,
712
+ fp32_internal: bool,
713
+ include_reduce_scatter: bool = True,
714
+ include_routed: bool = True,
715
+) -> None:
716
+ """Compile the rank/M specialization without retaining caller tensors."""
717
+
718
+ key = _compile_key(
719
+ rank,
720
+ tp_size,
721
+ latent_dim,
722
+ hidden_dim,
723
+ max_m,
724
+ max_token_ctas,
725
+ fp32_internal,
726
+ include_reduce_scatter,
727
+ include_routed,
728
+ )
729
+ if key in _COMPILED:
730
+ return
731
+ device = latent_output.device
732
+ latent = torch.empty((max_m, latent_dim), dtype=torch.bfloat16, device=device)
733
+ gamma = torch.empty((latent_dim,), dtype=torch.bfloat16, device=device)
734
+ shared = torch.empty((max_m, hidden_dim), dtype=torch.bfloat16, device=device)
735
+ kernel = AllReduceRMSNormWithReduceScatterEarlyExit(
736
+ rank=rank,
737
+ tp_size=tp_size,
738
+ latent_dim=latent_dim,
739
+ hidden_dim=hidden_dim,
740
+ max_m=max_m,
741
+ max_token_ctas=max_token_ctas,
742
+ fp32_internal=fp32_internal,
743
+ include_reduce_scatter=include_reduce_scatter,
744
+ include_routed=include_routed,
745
+ )
746
+ _COMPILED[key] = cute.compile(
747
+ kernel,
748
+ *_runtime_args(
749
+ latent,
750
+ gamma,
751
+ latent_output,
752
+ routed_workspace,
753
+ routed_flags,
754
+ routed_multicast_ptr,
755
+ shared,
756
+ shared_output,
757
+ shared_workspace,
758
+ shared_flags,
759
+ shared_peer_ptrs,
760
+ rms_eps,
761
+ ),
762
+ )
763
+
764
+
765
+def launch(
766
+ latent_source: torch.Tensor,
767
+ gamma: torch.Tensor,
768
+ latent_output: torch.Tensor,
769
+ routed_workspace: torch.Tensor,
770
+ routed_flags: torch.Tensor,
771
+ routed_multicast_ptr: int,
772
+ shared_source: torch.Tensor,
773
+ shared_output: torch.Tensor,
774
+ shared_workspace: torch.Tensor,
775
+ shared_flags: torch.Tensor,
776
+ shared_peer_ptrs: torch.Tensor,
777
+ rms_eps: float,
778
+ *,
779
+ rank: int,
780
+ tp_size: int,
781
+ latent_dim: int,
782
+ hidden_dim: int,
783
+ max_m: int,
784
+ max_token_ctas: int,
785
+ fp32_internal: bool,
786
+ include_reduce_scatter: bool = True,
787
+ include_routed: bool = True,
788
+) -> None:
789
+ compile_kernel(
790
+ rank=rank,
791
+ tp_size=tp_size,
792
+ latent_dim=latent_dim,
793
+ hidden_dim=hidden_dim,
794
+ max_m=max_m,
795
+ max_token_ctas=max_token_ctas,
796
+ latent_output=latent_output,
797
+ routed_workspace=routed_workspace,
798
+ routed_flags=routed_flags,
799
+ routed_multicast_ptr=routed_multicast_ptr,
800
+ shared_output=shared_output,
801
+ shared_workspace=shared_workspace,
802
+ shared_flags=shared_flags,
803
+ shared_peer_ptrs=shared_peer_ptrs,
804
+ rms_eps=rms_eps,
805
+ fp32_internal=fp32_internal,
806
+ include_reduce_scatter=include_reduce_scatter,
807
+ include_routed=include_routed,
808
+ )
809
+ _COMPILED[
810
+ _compile_key(
811
+ rank,
812
+ tp_size,
813
+ latent_dim,
814
+ hidden_dim,
815
+ max_m,
816
+ max_token_ctas,
817
+ fp32_internal,
818
+ include_reduce_scatter,
819
+ include_routed,
820
+ )
821
+ ](
822
+ *_runtime_args(
823
+ latent_source,
824
+ gamma,
825
+ latent_output,
826
+ routed_workspace,
827
+ routed_flags,
828
+ routed_multicast_ptr,
829
+ shared_source,
830
+ shared_output,
831
+ shared_workspace,
832
+ shared_flags,
833
+ shared_peer_ptrs,
834
+ rms_eps,
835
+ )
836
+ )
837
+
838
+
839
+class CollectiveKernel:
840
+ """Own and launch the routed AllReduce/RMSNorm plus shared ReduceScatter."""
841
+
842
+ def __init__(
843
+ self,
844
+ *,
845
+ group: dist.ProcessGroup,
846
+ rank: int,
847
+ tp_size: int,
848
+ latent_dim: int,
849
+ hidden_dim: int,
850
+ max_m: int,
851
+ max_token_ctas: int,
852
+ rms_eps: float,
853
+ fp32_internal: bool,
854
+ ) -> None:
855
+ validate_shape(
856
+ tp_size=tp_size,
857
+ latent_dim=latent_dim,
858
+ hidden_dim=hidden_dim,
859
+ )
860
+ self.rank = rank
861
+ self.tp_size = tp_size
862
+ self.latent_dim = latent_dim
863
+ self.hidden_dim = hidden_dim
864
+ self.shard_dim = hidden_dim // tp_size
865
+ self.max_m = max_m
866
+ self.max_token_ctas = max_token_ctas
867
+ self.rms_eps = float(rms_eps)
868
+ self.fp32_internal = fp32_internal
869
+ device = torch.device("cuda", torch.accelerator.current_device_index())
870
+
871
+ bytes_per_routed_buffer = max_m * tp_size * latent_dim * 2
872
+ routed_bytes = NUM_LAMPORT_BUFFERS * bytes_per_routed_buffer
873
+ self._routed_workspace = symm_mem.empty(
874
+ routed_bytes // 4,
875
+ dtype=torch.float32,
876
+ device=device,
877
+ )
878
+ self._routed_symm_mem = symm_mem.rendezvous(self._routed_workspace, group)
879
+ self._routed_workspace.fill_(-0.0)
880
+ actual_bytes_per_buffer = (
881
+ self._routed_symm_mem.buffer_size // NUM_LAMPORT_BUFFERS // 16 * 16
882
+ )
883
+ if actual_bytes_per_buffer < bytes_per_routed_buffer:
884
+ raise RuntimeError("routed symmetric workspace is too small")
885
+ self._routed_flags = torch.tensor(
886
+ [0, 2, actual_bytes_per_buffer, 0, 0, 0, 0, 0, 0],
887
+ dtype=torch.uint32,
888
+ device=device,
889
+ )
890
+ routed_multicast_ptr = self._routed_symm_mem.multicast_ptr
891
+ if routed_multicast_ptr is None or routed_multicast_ptr == 0:
892
+ raise RuntimeError("routed NVLS multicast mapping is unavailable")
893
+ self._routed_multicast_ptr = int(routed_multicast_ptr)
894
+
895
+ self._latent_output = torch.empty(
896
+ (max_m, latent_dim), dtype=torch.bfloat16, device=device
897
+ )
898
+ self._shared_output = torch.empty(
899
+ (max_m, hidden_dim), dtype=torch.bfloat16, device=device
900
+ )
901
+ shard_start = rank * self.shard_dim
902
+ shard_end = shard_start + self.shard_dim
903
+ self._shared_shard = self._shared_output[:, shard_start:shard_end]
904
+
905
+ self._shared_workspace = symm_mem.empty(
906
+ (NUM_LAMPORT_BUFFERS, max_m, tp_size, self.shard_dim),
907
+ dtype=torch.bfloat16,
908
+ device=device,
909
+ )
910
+ self._shared_symm_mem = symm_mem.rendezvous(self._shared_workspace, group)
911
+ self._shared_workspace.view(torch.int32).fill_(-0x80000000)
912
+ self._shared_flags = torch.zeros(12, dtype=torch.int32, device=device)
913
+ self._shared_flags[1] = 1
914
+ self._shared_flags[2] = max_m * tp_size * self.shard_dim * 2
915
+ peer_ptrs = [
916
+ self._shared_symm_mem.get_buffer(
917
+ peer,
918
+ self._shared_workspace.shape,
919
+ torch.bfloat16,
920
+ ).data_ptr()
921
+ for peer in range(tp_size)
922
+ ]
923
+ if any(pointer == 0 for pointer in peer_ptrs):
924
+ raise RuntimeError("shared LSA peer mapping is unavailable")
925
+ self._shared_peer_ptrs = torch.tensor(
926
+ peer_ptrs, dtype=torch.int64, device=device
927
+ )
928
+
929
+ torch.accelerator.synchronize(device)
930
+ dist.barrier(group=group, device_ids=[device.index])
931
+ for owner in range(tp_size):
932
+ if rank == owner:
933
+ compile_kernel(
934
+ rank=rank,
935
+ tp_size=tp_size,
936
+ latent_dim=latent_dim,
937
+ hidden_dim=hidden_dim,
938
+ max_m=max_m,
939
+ max_token_ctas=max_token_ctas,
940
+ latent_output=self._latent_output,
941
+ routed_workspace=self._routed_workspace,
942
+ routed_flags=self._routed_flags,
943
+ routed_multicast_ptr=self._routed_multicast_ptr,
944
+ shared_output=self._shared_output,
945
+ shared_workspace=self._shared_workspace,
946
+ shared_flags=self._shared_flags,
947
+ shared_peer_ptrs=self._shared_peer_ptrs,
948
+ rms_eps=self.rms_eps,
949
+ fp32_internal=fp32_internal,
950
+ )
951
+ dist.barrier(group=group, device_ids=[device.index])
952
+
953
+ def __call__(
954
+ self,
955
+ latent_source: torch.Tensor,
956
+ shared_source: torch.Tensor,
957
+ gamma: torch.Tensor,
958
+ ) -> tuple[torch.Tensor, torch.Tensor]:
959
+ if latent_source.ndim != 2 or shared_source.ndim != 2:
960
+ raise ValueError("latent_source and shared_source must be rank-2")
961
+ m = latent_source.shape[0]
962
+ device = self._routed_workspace.device
963
+ expected = (
964
+ (latent_source, (m, self.latent_dim), "latent_source"),
965
+ (shared_source, (m, self.hidden_dim), "shared_source"),
966
+ (gamma, (self.latent_dim,), "gamma"),
967
+ )
968
+ for tensor, shape, name in expected:
969
+ if (
970
+ tensor.shape != shape
971
+ or tensor.dtype != torch.bfloat16
972
+ or tensor.device != device
973
+ or not tensor.is_contiguous()
974
+ ):
975
+ raise ValueError(f"{name} must be contiguous CUDA BF16 {list(shape)}")
976
+ if not 1 <= m <= self.max_m:
977
+ raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]")
978
+
979
+ with torch.accelerator.device_index(device.index):
980
+ launch(
981
+ latent_source,
982
+ gamma,
983
+ self._latent_output,
984
+ self._routed_workspace,
985
+ self._routed_flags,
986
+ self._routed_multicast_ptr,
987
+ shared_source,
988
+ self._shared_output,
989
+ self._shared_workspace,
990
+ self._shared_flags,
991
+ self._shared_peer_ptrs,
992
+ self.rms_eps,
993
+ rank=self.rank,
994
+ tp_size=self.tp_size,
995
+ latent_dim=self.latent_dim,
996
+ hidden_dim=self.hidden_dim,
997
+ max_m=self.max_m,
998
+ max_token_ctas=self.max_token_ctas,
999
+ fp32_internal=self.fp32_internal,
1000
+ )
1001
+ return (
1002
+ self._latent_output[:m],
1003
+ self._shared_shard,
1004
+ )
1005
+
1006
+ @property
1007
+ def latent_output(self) -> torch.Tensor:
1008
+ return self._latent_output
1009
+
1010
+ @property
1011
+ def shared_output(self) -> torch.Tensor:
1012
+ return self._shared_output
vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py
new
+1291
@@ -0,0 +1,1291 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
5
+# SPDX-License-Identifier: BSD-3-Clause
6
+
7
+# Redistribution and use in source and binary forms, with or without
8
+# modification, are permitted provided that the following conditions are met:
9
+
10
+# 1. Redistributions of source code must retain the above copyright notice, this
11
+# list of conditions and the following disclaimer.
12
+
13
+# 2. Redistributions in binary form must reproduce the above copyright notice,
14
+# this list of conditions and the following disclaimer in the documentation
15
+# and/or other materials provided with the distribution.
16
+
17
+# 3. Neither the name of the copyright holder nor the names of its
18
+# contributors may be used to endorse or promote products derived from
19
+# this software without specific prior written permission.
20
+
21
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
22
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
23
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
24
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
25
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
26
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
27
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
28
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
29
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
30
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
31
+
32
+"""Blackwell GEMM with a fused shared-shard add and multicast epilogue.
33
+
34
+Modified from the original CUTLASS CuTe DSL SM100 persistent GEMM tutorial.
35
+The PDL wait before A loading orders the up-projection and shared-expert
36
+inputs after the producer collective.
37
+"""
38
+
39
+import math
40
+from typing import Any
41
+
42
+import cuda.bindings.driver as cuda
43
+import cutlass
44
+import cutlass.cute as cute
45
+import cutlass.pipeline as pipeline
46
+import cutlass.utils as utils
47
+import torch
48
+import torch.distributed as dist
49
+import torch.distributed._symmetric_memory as symm_mem
50
+from cutlass.cute.nvgpu import cpasync, tcgen05
51
+from cutlass.cute.nvgpu.common import CacheEvictionPriority
52
+from cutlass.cute.runtime import from_dlpack
53
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
54
+
55
+from .fused_add_multicast_skinny_gemm import (
56
+ FusedAddMulticastSkinnyGemmKernel,
57
+)
58
+from .primitives import CUDAGraphCompatibleWrapper
59
+
60
+
61
+def _as_cute(tensor: torch.Tensor, *, dynamic_m: bool = False):
62
+ converted = from_dlpack(
63
+ CUDAGraphCompatibleWrapper(tensor.detach()), assumed_align=16
64
+ )
65
+ if dynamic_m:
66
+ converted = converted.mark_compact_shape_dynamic(
67
+ mode=1,
68
+ stride_order=tensor.dim_order(),
69
+ )
70
+ return converted
71
+
72
+
73
+def validate_configuration(
74
+ *,
75
+ latent_dim: int,
76
+ shard_dim: int,
77
+ mma_tiler_mn: tuple[int, int],
78
+ cluster_shape_mn: tuple[int, int],
79
+ b_prime_stages: int,
80
+) -> None:
81
+ """Validate constraints imposed by this BF16 SM100 GEMM."""
82
+
83
+ if latent_dim <= 0 or shard_dim <= 0:
84
+ raise ValueError("GEMM K and N must be positive")
85
+ if latent_dim % 8 or shard_dim % 8:
86
+ raise ValueError("GEMM K and N must be divisible by 8 BF16 values")
87
+ if mma_tiler_mn[0] not in (64, 128):
88
+ raise ValueError("MMA M tile must be 64 or 128")
89
+ if mma_tiler_mn[1] not in range(32, 257, 32):
90
+ raise ValueError("MMA N tile must be a multiple of 32 in [32, 256]")
91
+ if (
92
+ len(cluster_shape_mn) != 2
93
+ or any(value <= 0 or value & (value - 1) for value in cluster_shape_mn)
94
+ or math.prod(cluster_shape_mn) > 16
95
+ ):
96
+ raise ValueError("cluster dimensions must be powers of two with product <= 16")
97
+ if not 0 <= b_prime_stages <= math.ceil(latent_dim / 128):
98
+ raise ValueError("b_prime_stages exceeds the GEMM K-tile count")
99
+
100
+
101
+@cute.jit
102
+def _epilogue_tma_store_add_shared(
103
+ gemm_kernel,
104
+ epi_tidx: cutlass.Int32,
105
+ warp_idx: cutlass.Int32,
106
+ tma_atom_c: cute.CopyAtom,
107
+ tCtAcc_base: cute.Tensor,
108
+ sC: cute.Tensor,
109
+ tCgC_base: cute.Tensor,
110
+ tCgShared_base: cute.Tensor,
111
+ tCcC_base: cute.Tensor,
112
+ mC_mnl: cute.Tensor,
113
+ mShared_mnl: cute.Tensor,
114
+ epi_tile: cute.Tile,
115
+ num_tiles_executed: cutlass.Int32,
116
+ mma_tile_coord_mnl,
117
+ acc_consumer_state: pipeline.PipelineState,
118
+ acc_pipeline: pipeline.PipelineAsync,
119
+ c_pipeline: pipeline.PipelineTmaStore,
120
+) -> pipeline.PipelineState:
121
+ """BF16 GEMM rounding + BF16 shared shard addition, then swizzled S2G TMA."""
122
+ sm100 = utils.gemm.sm100
123
+ tCgC = sm100.transform_partitioned_tensor_layout(tCgC_base)
124
+ tCgShared = sm100.transform_partitioned_tensor_layout(tCgShared_base)
125
+ tCcC = sm100.transform_partitioned_tensor_layout(tCcC_base)
126
+ tCtAcc = sm100.transform_partitioned_tensor_layout(tCtAcc_base)
127
+
128
+ tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = sm100.epilogue_tmem_copy_and_partition(
129
+ gemm_kernel,
130
+ epi_tidx,
131
+ tCtAcc,
132
+ tCgC,
133
+ epi_tile,
134
+ False,
135
+ )
136
+ tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, gemm_kernel.c_dtype)
137
+ tTR_rShared = cute.make_rmem_tensor(tTR_rAcc.shape, gemm_kernel.c_dtype)
138
+ tiled_copy_r2s, tRS_rC, tRS_sC = sm100.epilogue_smem_copy_and_partition(
139
+ gemm_kernel, tiled_copy_t2r, tTR_rC, epi_tidx, sC
140
+ )
141
+
142
+ tCgC_epi = cute.flat_divide(tCgC, epi_tile)
143
+ bSG_sC, bSG_gC_partitioned = cpasync.tma_partition(
144
+ tma_atom_c,
145
+ 0,
146
+ cute.make_layout(1),
147
+ cute.group_modes(sC, 0, 2),
148
+ cute.group_modes(tCgC_epi, 0, 2),
149
+ )
150
+ epilog_sync_barrier = pipeline.NamedBarrier(
151
+ barrier_id=gemm_kernel.epilog_sync_bar_id,
152
+ num_threads=32 * len(gemm_kernel.epilogue_warp_id),
153
+ )
154
+
155
+ bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)]
156
+ tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_consumer_state.index)]
157
+ thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx)
158
+ tTR_gShared_partitioned = thr_copy_t2r.partition_D(
159
+ cute.flat_divide(tCgShared, epi_tile)
160
+ )
161
+ tTR_cC_partitioned = thr_copy_t2r.partition_D(cute.flat_divide(tCcC, epi_tile))
162
+ tTR_gShared = tTR_gShared_partitioned[
163
+ (None, None, None, None, None, *mma_tile_coord_mnl)
164
+ ]
165
+ tTR_cC = tTR_cC_partitioned[(None, None, None, None, None, *mma_tile_coord_mnl)]
166
+
167
+ exemplar = tTR_gShared_partitioned[(None, None, None, 0, 0, 0, 0, 0)]
168
+ mcl_r = cute.max_common_layout(tTR_rShared.layout, exemplar.layout)
169
+ shared_copy_bits = min(
170
+ exemplar.iterator.alignment * 8,
171
+ cute.size(mcl_r) * gemm_kernel.c_dtype.width,
172
+ 128,
173
+ )
174
+ shared_g2r_atom = cute.make_copy_atom(
175
+ cute.nvgpu.CopyG2ROp(),
176
+ gemm_kernel.c_dtype,
177
+ num_bits_per_copy=shared_copy_bits,
178
+ l1c_evict_priority=CacheEvictionPriority.NO_ALLOCATE,
179
+ )
180
+
181
+ tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
182
+ tTR_gShared = cute.group_modes(tTR_gShared, 3, cute.rank(tTR_gShared))
183
+ tTR_cC = cute.group_modes(tTR_cC, 3, cute.rank(tTR_cC))
184
+ bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))
185
+
186
+ subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
187
+ previous_subtile_count = num_tiles_executed * subtile_cnt
188
+ cute.arch.griddepcontrol_wait()
189
+ for subtile_idx in range(subtile_cnt):
190
+ tTR_gShared_subtile = tTR_gShared[(None, None, None, subtile_idx)]
191
+ tTR_cC_subtile = tTR_cC[(None, None, None, subtile_idx)]
192
+ pred_shape = (1, *tTR_cC_subtile.shape[1:])
193
+ pred = cute.make_rmem_tensor(pred_shape, cutlass.Boolean)
194
+ for m_idx in range(tTR_cC_subtile.shape[1]):
195
+ for n_idx in range(tTR_cC_subtile.shape[2]):
196
+ pred[(0, m_idx, n_idx)] = cute.elem_less(
197
+ tTR_cC_subtile[(0, m_idx, n_idx)], mC_mnl.shape
198
+ )
199
+ tTR_rShared.store(cute.zeros_like(tTR_rShared, dtype=gemm_kernel.c_dtype))
200
+ cute.copy(
201
+ shared_g2r_atom,
202
+ tTR_gShared_subtile,
203
+ tTR_rShared,
204
+ pred=pred,
205
+ )
206
+
207
+ # Load the shared addend before waiting for the accumulator.
208
+ if subtile_idx == 0:
209
+ acc_pipeline.consumer_wait(acc_consumer_state)
210
+ tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)]
211
+ cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)
212
+
213
+ gemm_vec = tiled_copy_r2s.retile(tTR_rAcc).load().to(gemm_kernel.c_dtype)
214
+ shared_vec = tiled_copy_r2s.retile(tTR_rShared).load()
215
+ fused_vec = (gemm_vec.to(cutlass.Float32) + shared_vec.to(cutlass.Float32)).to(
216
+ gemm_kernel.c_dtype
217
+ )
218
+ # The symmetric output is an in-band Lamport mailbox whose empty
219
+ # marker contains BF16 -0. Normalize either signed zero to +0 so a
220
+ # legitimate result can never be mistaken for an unwritten fragment.
221
+ fused_vec = cute.where(
222
+ fused_vec == cute.zeros_like(fused_vec),
223
+ cute.zeros_like(fused_vec),
224
+ fused_vec,
225
+ )
226
+ tRS_rC.store(fused_vec)
227
+
228
+ c_buffer = (previous_subtile_count + subtile_idx) % gemm_kernel.num_c_stage
229
+ cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)])
230
+ cute.arch.fence_proxy("async.shared", space="cta")
231
+ epilog_sync_barrier.arrive_and_wait()
232
+ if warp_idx == gemm_kernel.epilogue_warp_id[0]:
233
+ cute.copy(
234
+ tma_atom_c,
235
+ bSG_sC[(None, c_buffer)],
236
+ bSG_gC[(None, subtile_idx)],
237
+ )
238
+ c_pipeline.producer_commit()
239
+ c_pipeline.producer_acquire()
240
+ epilog_sync_barrier.arrive_and_wait()
241
+
242
+ epilog_sync_barrier.arrive_and_wait()
243
+ with cute.arch.elect_one():
244
+ acc_pipeline.consumer_release(acc_consumer_state)
245
+ acc_consumer_state.advance()
246
+ return acc_consumer_state
247
+
248
+
249
+def _compute_stages(
250
+ tiled_mma: cute.TiledMma,
251
+ mma_tiler_mnk: tuple[int, int, int],
252
+ a_dtype,
253
+ b_dtype,
254
+ c_dtype,
255
+ smem_capacity: int,
256
+ c_smem_layout,
257
+) -> tuple[int, int, int]:
258
+ """Choose accumulator, mainloop, and epilogue stage counts."""
259
+ num_acc_stage = 2
260
+ num_c_stage = 2
261
+ a_smem_layout_stage_one = utils.sm100.make_smem_layout_a(
262
+ tiled_mma, mma_tiler_mnk, a_dtype, 1
263
+ )
264
+ b_smem_layout_staged_one = utils.sm100.make_smem_layout_b(
265
+ tiled_mma, mma_tiler_mnk, b_dtype, 1
266
+ )
267
+
268
+ ab_bytes_per_stage = cute.size_in_bytes(
269
+ a_dtype, a_smem_layout_stage_one
270
+ ) + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one)
271
+ mbar_helpers_bytes = 1024
272
+
273
+ c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout)
274
+ c_bytes = c_bytes_per_stage * num_c_stage
275
+ num_ab_stage = (
276
+ smem_capacity - (mbar_helpers_bytes + c_bytes)
277
+ ) // ab_bytes_per_stage
278
+ num_c_stage += (
279
+ smem_capacity
280
+ - ab_bytes_per_stage * num_ab_stage
281
+ - (mbar_helpers_bytes + c_bytes)
282
+ ) // c_bytes_per_stage
283
+ return num_acc_stage, num_ab_stage, num_c_stage
284
+
285
+
286
+class FusedAddMulticastGemm:
287
+ """Persistent Blackwell GEMM with a shared-add epilogue.
288
+
289
+ B priming may overlap the producer collective. The PDL wait before A
290
+ loading orders both inputs before the shared-add epilogue.
291
+ """
292
+
293
+ def __init__(
294
+ self,
295
+ mma_tiler_mn: tuple[int, int],
296
+ cluster_shape_mn: tuple[int, int],
297
+ b_prime_stages: int = 2,
298
+ ):
299
+ self.acc_dtype = cutlass.Float32
300
+ self.cluster_shape_mn = cluster_shape_mn
301
+ self.mma_tiler = (*mma_tiler_mn, 1)
302
+ # B primes the combined A+B pipeline before the PDL wait.
303
+ self.b_prime_stages = b_prime_stages
304
+ self.cta_group = tcgen05.CtaGroup.ONE
305
+ self.epilogue_warp_id = (0, 1, 2, 3)
306
+ self.mma_warp_id = 4
307
+ self.tma_warp_id = 5
308
+ self.threads_per_cta = 32 * len(
309
+ (self.mma_warp_id, self.tma_warp_id, *self.epilogue_warp_id)
310
+ )
311
+ self.epilog_sync_bar_id = 1
312
+ self.tmem_alloc_sync_bar_id = 2
313
+
314
+ def _create_tiled_mma(self):
315
+ return utils.sm100.make_trivial_tiled_mma(
316
+ self.a_dtype,
317
+ self.a_major_mode,
318
+ self.b_major_mode,
319
+ self.acc_dtype,
320
+ self.cta_group,
321
+ self.mma_tiler[:2],
322
+ )
323
+
324
+ def _setup_attributes(self):
325
+ """Derive layouts and stage counts from the compiled tensor shapes."""
326
+ tiled_mma = self._create_tiled_mma()
327
+
328
+ mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
329
+ mma_inst_tile_k = 4
330
+ self.mma_tiler = (
331
+ self.mma_tiler[0],
332
+ self.mma_tiler[1],
333
+ mma_inst_shape_k * mma_inst_tile_k,
334
+ )
335
+ self.cta_tile_shape_mnk = (
336
+ self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
337
+ self.mma_tiler[1],
338
+ self.mma_tiler[2],
339
+ )
340
+
341
+ self.cluster_layout_vmnk = cute.tiled_divide(
342
+ cute.make_layout((*self.cluster_shape_mn, 1)),
343
+ (tiled_mma.thr_id.shape,),
344
+ )
345
+
346
+ self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2])
347
+ self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1])
348
+ self.is_a_mcast = self.num_mcast_ctas_a > 1
349
+ self.is_b_mcast = self.num_mcast_ctas_b > 1
350
+
351
+ self.epi_tile = utils.sm100.compute_epilogue_tile_shape(
352
+ self.cta_tile_shape_mnk,
353
+ False,
354
+ self.c_layout,
355
+ self.c_dtype,
356
+ )
357
+ c_smem_layout = utils.sm100.make_smem_layout_epi(
358
+ self.c_dtype, self.c_layout, self.epi_tile, 1
359
+ )
360
+
361
+ self.num_acc_stage, self.num_ab_stage, self.num_c_stage = _compute_stages(
362
+ tiled_mma,
363
+ self.mma_tiler,
364
+ self.a_dtype,
365
+ self.b_dtype,
366
+ self.c_dtype,
367
+ utils.get_smem_capacity_in_bytes(),
368
+ c_smem_layout,
369
+ )
370
+
371
+ self.a_smem_layout_staged = utils.sm100.make_smem_layout_a(
372
+ tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage
373
+ )
374
+ self.b_smem_layout_staged = utils.sm100.make_smem_layout_b(
375
+ tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage
376
+ )
377
+
378
+ self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi(
379
+ self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage
380
+ )
381
+
382
+ self.num_tmem_alloc_cols = self._compute_num_tmem_alloc_cols(
383
+ tiled_mma, self.mma_tiler, self.num_acc_stage, "sm_100"
384
+ )
385
+
386
+ @cute.jit
387
+ def __call__(
388
+ self,
389
+ a: cute.Tensor,
390
+ b: cute.Tensor,
391
+ c: cute.Tensor,
392
+ shared_shard: cute.Tensor,
393
+ c_multicast_i64: cutlass.Int64,
394
+ max_active_clusters: cutlass.Constexpr,
395
+ stream: cuda.CUstream,
396
+ ):
397
+ """Launch the persistent GEMM."""
398
+ # Preserve C's logical strided layout but point the TMA descriptor at
399
+ # this rank's shard inside the LSA multicast mapping. One TMA store is
400
+ # therefore replicated into the same shard on all eight ranks.
401
+ c = cute.make_tensor(
402
+ cute.make_ptr(
403
+ c.element_type,
404
+ c_multicast_i64,
405
+ cute.AddressSpace.gmem,
406
+ assumed_align=16,
407
+ ),
408
+ c.layout,
409
+ )
410
+
411
+ self.a_dtype: type[cutlass.Numeric] = a.element_type
412
+ self.b_dtype: type[cutlass.Numeric] = b.element_type
413
+ self.c_dtype: type[cutlass.Numeric] = c.element_type
414
+ self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode()
415
+ self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode()
416
+ self.c_layout = utils.LayoutEnum.from_tensor(c)
417
+
418
+ if cutlass.const_expr(self.a_dtype != self.b_dtype):
419
+ raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}")
420
+
421
+ tiled_mma = self._create_tiled_mma()
422
+
423
+ self._setup_attributes()
424
+ if cutlass.const_expr(self.b_prime_stages > self.num_ab_stage):
425
+ raise ValueError(
426
+ "b_prime_stages exceeds the compiled A/B pipeline stage count"
427
+ )
428
+
429
+ atom_thr_size = cute.size(tiled_mma.thr_id.shape)
430
+
431
+ a_op = utils.sm100.cluster_shape_to_tma_atom_A(
432
+ self.cluster_shape_mn, tiled_mma.thr_id
433
+ )
434
+ a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
435
+ tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
436
+ a_op,
437
+ a,
438
+ a_smem_layout,
439
+ self.mma_tiler,
440
+ tiled_mma,
441
+ self.cluster_layout_vmnk.shape,
442
+ internal_type=(
443
+ cutlass.TFloat32 if a.element_type is cutlass.Float32 else None
444
+ ),
445
+ )
446
+
447
+ b_op = utils.sm100.cluster_shape_to_tma_atom_B(
448
+ self.cluster_shape_mn, tiled_mma.thr_id
449
+ )
450
+ b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
451
+ tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
452
+ b_op,
453
+ b,
454
+ b_smem_layout,
455
+ self.mma_tiler,
456
+ tiled_mma,
457
+ self.cluster_layout_vmnk.shape,
458
+ internal_type=(
459
+ cutlass.TFloat32 if b.element_type is cutlass.Float32 else None
460
+ ),
461
+ )
462
+
463
+ a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout)
464
+ b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout)
465
+ self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size
466
+
467
+ epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1])
468
+ tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
469
+ cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile
470
+ )
471
+
472
+ tile_sched_params, grid = self._compute_grid(
473
+ c, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters
474
+ )
475
+ self.kernel(
476
+ tiled_mma,
477
+ tma_atom_a,
478
+ tma_tensor_a,
479
+ tma_atom_b,
480
+ tma_tensor_b,
481
+ tma_atom_c,
482
+ tma_tensor_c,
483
+ self.cluster_layout_vmnk,
484
+ self.a_smem_layout_staged,
485
+ self.b_smem_layout_staged,
486
+ self.c_smem_layout_staged,
487
+ self.epi_tile,
488
+ tile_sched_params,
489
+ shared_shard,
490
+ ).launch(
491
+ grid=grid,
492
+ block=[self.threads_per_cta, 1, 1],
493
+ cluster=(*self.cluster_shape_mn, 1),
494
+ stream=stream,
495
+ use_pdl=True,
496
+ )
497
+ return
498
+
499
+ @cute.kernel
500
+ def kernel(
501
+ self,
502
+ tiled_mma: cute.TiledMma,
503
+ tma_atom_a: cute.CopyAtom,
504
+ mA_mkl: cute.Tensor,
505
+ tma_atom_b: cute.CopyAtom,
506
+ mB_nkl: cute.Tensor,
507
+ tma_atom_c: cute.CopyAtom,
508
+ mC_mnl: cute.Tensor,
509
+ cluster_layout_vmnk: cute.Layout,
510
+ a_smem_layout_staged: cute.ComposedLayout,
511
+ b_smem_layout_staged: cute.ComposedLayout,
512
+ c_smem_layout_staged: cute.Layout | cute.ComposedLayout,
513
+ epi_tile: cute.Tile,
514
+ tile_sched_params: utils.PersistentTileSchedulerParams,
515
+ shared_shard: cute.Tensor,
516
+ ):
517
+ self._gemm_device(
518
+ tiled_mma,
519
+ tma_atom_a,
520
+ mA_mkl,
521
+ tma_atom_b,
522
+ mB_nkl,
523
+ tma_atom_c,
524
+ mC_mnl,
525
+ cluster_layout_vmnk,
526
+ a_smem_layout_staged,
527
+ b_smem_layout_staged,
528
+ c_smem_layout_staged,
529
+ epi_tile,
530
+ tile_sched_params,
531
+ shared_shard,
532
+ )
533
+
534
+ @cute.jit
535
+ def _gemm_device(
536
+ self,
537
+ tiled_mma: cute.TiledMma,
538
+ tma_atom_a: cute.CopyAtom,
539
+ mA_mkl: cute.Tensor,
540
+ tma_atom_b: cute.CopyAtom,
541
+ mB_nkl: cute.Tensor,
542
+ tma_atom_c: cute.CopyAtom,
543
+ mC_mnl: cute.Tensor,
544
+ cluster_layout_vmnk: cute.Layout,
545
+ a_smem_layout_staged: cute.ComposedLayout,
546
+ b_smem_layout_staged: cute.ComposedLayout,
547
+ c_smem_layout_staged: cute.Layout | cute.ComposedLayout,
548
+ epi_tile: cute.Tile,
549
+ tile_sched_params: utils.PersistentTileSchedulerParams,
550
+ shared_shard: cute.Tensor,
551
+ ):
552
+ warp_idx = cute.arch.warp_idx()
553
+ warp_idx = cute.arch.make_warp_uniform(warp_idx)
554
+
555
+ if warp_idx == self.tma_warp_id:
556
+ cpasync.prefetch_descriptor(tma_atom_a)
557
+ cpasync.prefetch_descriptor(tma_atom_b)
558
+ cpasync.prefetch_descriptor(tma_atom_c)
559
+
560
+ bidx, bidy, bidz = cute.arch.block_idx()
561
+ mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
562
+ is_leader_cta = mma_tile_coord_v == 0
563
+ cta_rank_in_cluster = cute.arch.make_warp_uniform(
564
+ cute.arch.block_idx_in_cluster()
565
+ )
566
+ block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(
567
+ cta_rank_in_cluster
568
+ )
569
+ tidx, _, _ = cute.arch.thread_idx()
570
+
571
+ @cute.struct
572
+ class SharedStorage:
573
+ ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]
574
+ acc_full_mbar_ptr: cute.struct.MemRange[
575
+ cutlass.Int64, self.num_acc_stage * 2
576
+ ]
577
+ tmem_dealloc_mbar_ptr: cutlass.Int64
578
+ tmem_holding_buf: cutlass.Int32
579
+
580
+ smem = utils.SmemAllocator()
581
+ storage = smem.allocate(SharedStorage)
582
+
583
+ ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
584
+ num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
585
+ ab_pipeline_consumer_group = pipeline.CooperativeGroup(
586
+ pipeline.Agent.Thread, num_tma_producer
587
+ )
588
+ ab_pipeline = pipeline.PipelineTmaUmma.create(
589
+ barrier_storage=storage.ab_full_mbar_ptr.data_ptr(),
590
+ num_stages=self.num_ab_stage,
591
+ producer_group=ab_pipeline_producer_group,
592
+ consumer_group=ab_pipeline_consumer_group,
593
+ tx_count=self.num_tma_load_bytes,
594
+ cta_layout_vmnk=cluster_layout_vmnk,
595
+ defer_sync=True,
596
+ )
597
+ ab_producer, ab_consumer = ab_pipeline.make_participants()
598
+
599
+ acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
600
+ num_acc_consumer_threads = len(self.epilogue_warp_id)
601
+ acc_pipeline_consumer_group = pipeline.CooperativeGroup(
602
+ pipeline.Agent.Thread, num_acc_consumer_threads
603
+ )
604
+ acc_pipeline = pipeline.PipelineUmmaAsync.create(
605
+ barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
606
+ num_stages=self.num_acc_stage,
607
+ producer_group=acc_pipeline_producer_group,
608
+ consumer_group=acc_pipeline_consumer_group,
609
+ cta_layout_vmnk=cluster_layout_vmnk,
610
+ defer_sync=True,
611
+ )
612
+
613
+ tmem_alloc_barrier = pipeline.NamedBarrier(
614
+ barrier_id=self.tmem_alloc_sync_bar_id,
615
+ num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)),
616
+ )
617
+ tmem = utils.TmemAllocator(
618
+ storage.tmem_holding_buf,
619
+ barrier_for_retrieve=tmem_alloc_barrier,
620
+ allocator_warp_id=self.epilogue_warp_id[0],
621
+ is_two_cta=False,
622
+ two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
623
+ )
624
+
625
+ pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True)
626
+
627
+ sA = smem.allocate_tensor(
628
+ element_type=self.a_dtype,
629
+ layout=a_smem_layout_staged.outer,
630
+ byte_alignment=128,
631
+ swizzle=a_smem_layout_staged.inner,
632
+ )
633
+ sB = smem.allocate_tensor(
634
+ element_type=self.b_dtype,
635
+ layout=b_smem_layout_staged.outer,
636
+ byte_alignment=128,
637
+ swizzle=b_smem_layout_staged.inner,
638
+ )
639
+
640
+ a_full_mcast_mask = None
641
+ b_full_mcast_mask = None
642
+ if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast):
643
+ a_full_mcast_mask = cpasync.create_tma_multicast_mask(
644
+ cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
645
+ )
646
+ b_full_mcast_mask = cpasync.create_tma_multicast_mask(
647
+ cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
648
+ )
649
+
650
+ gA_mkl = cute.local_tile(
651
+ mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
652
+ )
653
+ gB_nkl = cute.local_tile(
654
+ mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
655
+ )
656
+ gC_mnl = cute.local_tile(
657
+ mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
658
+ )
659
+ # Shared shard is physically [M, shard_dim]. Give it the same logical MNL
660
+ # view as C so its epilogue partition is coordinate-identical.
661
+ mShared_mnl = cute.make_tensor(
662
+ shared_shard.iterator,
663
+ cute.append(shared_shard.layout, cute.make_layout((1,), stride=(0,))),
664
+ )
665
+ gShared_mnl = cute.local_tile(
666
+ mShared_mnl,
667
+ cute.slice_(self.mma_tiler, (None, None, 0)),
668
+ (None, None, None),
669
+ )
670
+ k_tile_cnt = cute.size(gA_mkl, mode=[3])
671
+
672
+ thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
673
+ tCgA = thr_mma.partition_A(gA_mkl)
674
+ tCgB = thr_mma.partition_B(gB_nkl)
675
+ tCgC = thr_mma.partition_C(gC_mnl)
676
+ tCgShared = thr_mma.partition_C(gShared_mnl)
677
+
678
+ # Predicate the partial M tile when M is smaller than the MMA tile.
679
+ idC = cute.make_identity_tensor(mC_mnl.shape)
680
+ cC_mnl = cute.local_tile(
681
+ idC, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
682
+ )
683
+ tCcC = thr_mma.partition_C(cC_mnl)
684
+
685
+ a_cta_layout = cute.make_layout(
686
+ cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
687
+ )
688
+ tAsA, tAgA = cpasync.tma_partition(
689
+ tma_atom_a,
690
+ block_in_cluster_coord_vmnk[2],
691
+ a_cta_layout,
692
+ cute.group_modes(sA, 0, 3),
693
+ cute.group_modes(tCgA, 0, 3),
694
+ )
695
+ b_cta_layout = cute.make_layout(
696
+ cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
697
+ )
698
+ tBsB, tBgB = cpasync.tma_partition(
699
+ tma_atom_b,
700
+ block_in_cluster_coord_vmnk[1],
701
+ b_cta_layout,
702
+ cute.group_modes(sB, 0, 3),
703
+ cute.group_modes(tCgB, 0, 3),
704
+ )
705
+
706
+ tCrA = tiled_mma.make_fragment_A(sA)
707
+ tCrB = tiled_mma.make_fragment_B(sB)
708
+ acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
709
+ tCtAcc_fake = tiled_mma.make_fragment_C(
710
+ cute.append(acc_shape, self.num_acc_stage)
711
+ )
712
+
713
+ pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk)
714
+
715
+ gemm_grid_z = cute.arch.grid_dim()[2]
716
+ tile_sched = utils.StaticPersistentTileScheduler.create(
717
+ tile_sched_params,
718
+ cute.arch.block_idx(),
719
+ (
720
+ cute.arch.grid_dim()[0],
721
+ cute.arch.grid_dim()[1],
722
+ gemm_grid_z,
723
+ ),
724
+ )
725
+ work_tile = tile_sched.initial_work_tile_info()
726
+
727
+ if warp_idx == self.tma_warp_id:
728
+ while work_tile.is_valid_tile:
729
+ cur_tile_coord = work_tile.tile_idx
730
+ mma_tile_coord_mnl = (
731
+ cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
732
+ cur_tile_coord[1],
733
+ cur_tile_coord[2],
734
+ )
735
+
736
+ tAgA_slice = tAgA[
737
+ (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])
738
+ ]
739
+ tBgB_slice = tBgB[
740
+ (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])
741
+ ]
742
+
743
+ # Prime a short prefix of the existing combined A+B ring with
744
+ # B only. Its barrier still expects A+B bytes, so the MMA
745
+ # consumer cannot observe a half-filled stage.
746
+ ab_producer.reset()
747
+ peek_ab_empty_status = ab_producer.try_acquire()
748
+
749
+ for k_tile in cutlass.range(0, self.b_prime_stages, 1, unroll=1):
750
+ handle = ab_producer.acquire_and_advance(peek_ab_empty_status)
751
+ cute.copy(
752
+ tma_atom_b,
753
+ tBgB_slice[(None, handle.count)],
754
+ tBsB[(None, handle.index)],
755
+ tma_bar_ptr=handle.barrier,
756
+ mcast_mask=b_full_mcast_mask,
757
+ )
758
+ peek_ab_empty_status = cutlass.Boolean(1)
759
+ if handle.count + 1 < self.b_prime_stages:
760
+ peek_ab_empty_status = ab_producer.try_acquire()
761
+
762
+ cute.arch.griddepcontrol_wait()
763
+
764
+ # Supply A to the very same stages after the producer AR has
765
+ # programmatically released this dependent kernel.
766
+ a_fill_state = pipeline.make_pipeline_state(
767
+ pipeline.PipelineUserType.Producer, self.num_ab_stage
768
+ )
769
+ for k_tile in cutlass.range(0, self.b_prime_stages, 1, unroll=1):
770
+ a_barrier = ab_pipeline.producer_get_barrier(a_fill_state)
771
+ cute.copy(
772
+ tma_atom_a,
773
+ tAgA_slice[(None, k_tile)],
774
+ tAsA[(None, a_fill_state.index)],
775
+ tma_bar_ptr=a_barrier,
776
+ mcast_mask=a_full_mcast_mask,
777
+ )
778
+ a_fill_state.advance()
779
+
780
+ peek_ab_empty_status = ab_producer.try_acquire()
781
+
782
+ for k_tile in cutlass.range(
783
+ self.b_prime_stages, k_tile_cnt, 1, unroll=1
784
+ ):
785
+ handle = ab_producer.acquire_and_advance(peek_ab_empty_status)
786
+
787
+ cute.copy(
788
+ tma_atom_a,
789
+ tAgA_slice[(None, handle.count)],
790
+ tAsA[(None, handle.index)],
791
+ tma_bar_ptr=handle.barrier,
792
+ mcast_mask=a_full_mcast_mask,
793
+ )
794
+ cute.copy(
795
+ tma_atom_b,
796
+ tBgB_slice[(None, handle.count)],
797
+ tBsB[(None, handle.index)],
798
+ tma_bar_ptr=handle.barrier,
799
+ mcast_mask=b_full_mcast_mask,
800
+ )
801
+
802
+ peek_ab_empty_status = cutlass.Boolean(1)
803
+ if handle.count + 1 < k_tile_cnt:
804
+ peek_ab_empty_status = ab_producer.try_acquire()
805
+
806
+ tile_sched.advance_to_next_work()
807
+ work_tile = tile_sched.get_current_work()
808
+
809
+ ab_producer.tail()
810
+
811
+ if warp_idx == self.mma_warp_id:
812
+ tmem.wait_for_alloc()
813
+ tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
814
+ tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
815
+
816
+ acc_producer_state = pipeline.make_pipeline_state(
817
+ pipeline.PipelineUserType.Producer, self.num_acc_stage
818
+ )
819
+
820
+ while work_tile.is_valid_tile:
821
+ cur_tile_coord = work_tile.tile_idx
822
+ mma_tile_coord_mnl = (
823
+ cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
824
+ cur_tile_coord[1],
825
+ cur_tile_coord[2],
826
+ )
827
+
828
+ tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)]
829
+
830
+ ab_consumer.reset()
831
+ peek_ab_full_status = cutlass.Boolean(1)
832
+ if is_leader_cta:
833
+ peek_ab_full_status = ab_consumer.try_wait()
834
+
835
+ if is_leader_cta:
836
+ acc_pipeline.producer_acquire(acc_producer_state)
837
+
838
+ tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
839
+
840
+ for k_tile in range(k_tile_cnt):
841
+ if is_leader_cta:
842
+ handle = ab_consumer.wait_and_advance(peek_ab_full_status)
843
+
844
+ num_kblocks = cute.size(tCrA, mode=[2])
845
+ for kblk_idx in cutlass.range(num_kblocks, unroll_full=True):
846
+ kblk_crd = (None, None, kblk_idx, handle.index)
847
+
848
+ cute.gemm(
849
+ tiled_mma,
850
+ tCtAcc,
851
+ tCrA[kblk_crd],
852
+ tCrB[kblk_crd],
853
+ tCtAcc,
854
+ )
855
+ tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
856
+
857
+ handle.release()
858
+
859
+ peek_ab_full_status = cutlass.Boolean(1)
860
+ if handle.count + 1 < k_tile_cnt:
861
+ peek_ab_full_status = ab_consumer.try_wait()
862
+
863
+ if is_leader_cta:
864
+ acc_pipeline.producer_commit(acc_producer_state)
865
+ acc_producer_state.advance()
866
+
867
+ tile_sched.advance_to_next_work()
868
+ work_tile = tile_sched.get_current_work()
869
+
870
+ acc_pipeline.producer_tail(acc_producer_state)
871
+
872
+ sC = smem.allocate_tensor(
873
+ element_type=self.c_dtype,
874
+ layout=c_smem_layout_staged.outer,
875
+ byte_alignment=128,
876
+ swizzle=c_smem_layout_staged.inner,
877
+ )
878
+
879
+ if warp_idx < self.mma_warp_id:
880
+ tmem.allocate(self.num_tmem_alloc_cols)
881
+
882
+ tmem.wait_for_alloc()
883
+ tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
884
+ tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
885
+
886
+ acc_consumer_state = pipeline.make_pipeline_state(
887
+ pipeline.PipelineUserType.Consumer, self.num_acc_stage
888
+ )
889
+ c_producer_group = pipeline.CooperativeGroup(
890
+ pipeline.Agent.Thread,
891
+ 32 * len(self.epilogue_warp_id),
892
+ )
893
+ c_pipeline = pipeline.PipelineTmaStore.create(
894
+ num_stages=self.num_c_stage, producer_group=c_producer_group
895
+ )
896
+ while work_tile.is_valid_tile:
897
+ cur_tile_coord = work_tile.tile_idx
898
+ mma_tile_coord_mnl = (
899
+ cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
900
+ cur_tile_coord[1],
901
+ cur_tile_coord[2],
902
+ )
903
+ tile_sched.advance_to_next_work()
904
+ work_tile = tile_sched.get_current_work()
905
+
906
+ num_tiles_executed = tile_sched.num_tiles_executed
907
+ acc_consumer_state = _epilogue_tma_store_add_shared(
908
+ self,
909
+ tidx,
910
+ warp_idx,
911
+ tma_atom_c,
912
+ tCtAcc_base,
913
+ sC,
914
+ tCgC,
915
+ tCgShared,
916
+ tCcC,
917
+ mC_mnl,
918
+ mShared_mnl,
919
+ epi_tile,
920
+ num_tiles_executed,
921
+ mma_tile_coord_mnl,
922
+ acc_consumer_state,
923
+ acc_pipeline,
924
+ c_pipeline,
925
+ )
926
+
927
+ c_pipeline.producer_tail()
928
+
929
+ tmem.relinquish_alloc_permit()
930
+ tmem.free(tmem_ptr)
931
+
932
+ # Allow the Lamport copy to become resident before this grid
933
+ # fully retires. Its griddepcontrol.wait still enforces complete
934
+ # producer-grid ordering before mailbox inspection.
935
+ cute.arch.griddepcontrol_launch_dependents()
936
+
937
+ @staticmethod
938
+ def _compute_grid(
939
+ c: cute.Tensor,
940
+ cta_tile_shape_mnk: tuple[int, int, int],
941
+ cluster_shape_mn: tuple[int, int],
942
+ max_active_clusters: cutlass.Constexpr,
943
+ ) -> tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]:
944
+ """Build the static persistent schedule."""
945
+ c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0))
946
+ gc = cute.zipped_divide(c, tiler=c_shape)
947
+ num_ctas_mnl = gc[(0, (None, None, None))].shape
948
+ cluster_shape_mnl = (*cluster_shape_mn, 1)
949
+
950
+ tile_sched_params = utils.PersistentTileSchedulerParams(
951
+ num_ctas_mnl, cluster_shape_mnl
952
+ )
953
+ grid = utils.StaticPersistentTileScheduler.get_grid_shape(
954
+ tile_sched_params, max_active_clusters
955
+ )
956
+
957
+ return tile_sched_params, grid
958
+
959
+ @staticmethod
960
+ def _compute_num_tmem_alloc_cols(
961
+ tiled_mma: cute.TiledMma,
962
+ mma_tiler: tuple[int, int, int],
963
+ num_acc_stage: int,
964
+ arch: str,
965
+ ) -> int:
966
+ """Return the required tensor-memory column count."""
967
+ acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2])
968
+ tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage))
969
+ num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch=arch)
970
+
971
+ return num_tmem_alloc_cols
972
+
973
+
974
+@cute.jit
975
+def launch_kernel(
976
+ gemm_op: cutlass.Constexpr,
977
+ a: cute.Tensor, # (l, m, k)
978
+ b: cute.Tensor, # (l, n, k)
979
+ c: cute.Tensor, # (l, m, n)
980
+ shared_shard: cute.Tensor, # (m, shard_dim), private TP shard
981
+ rows: cutlass.Int64,
982
+ c_multicast_i64: cutlass.Int64,
983
+ full_hidden_dim: cutlass.Constexpr,
984
+ shard_dim: cutlass.Constexpr,
985
+ max_active_clusters: cutlass.Constexpr,
986
+ stream: cuda.CUstream,
987
+):
988
+ """Launch the fused-add multicast GEMM using PyTorch BMM tensor order."""
989
+ # C is passed as the fixed-capacity [1,max_m,H] symmetric mailbox. Only M
990
+ # is runtime-variable; construct this rank's logical [1,M,S] view here so
991
+ # H and S remain compile-time constants and no dynamic strided host view is
992
+ # needed on every call. __call__ later replaces the local base pointer with
993
+ # the already rank-offset multicast address.
994
+ c = cute.make_tensor(
995
+ c.iterator,
996
+ cute.make_layout(
997
+ (1, rows, shard_dim),
998
+ stride=(0, full_hidden_dim, 1),
999
+ ),
1000
+ )
1001
+ # (l,m,k) -> (m,k,l)
1002
+ a = cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0]))
1003
+ # (l,n,k) -> (n,k,l)
1004
+ b = cute.make_tensor(b.iterator, cute.select(b.layout, mode=[1, 2, 0]))
1005
+ # (l,m,n) -> (m,n,l)
1006
+ c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0]))
1007
+
1008
+ gemm_op(
1009
+ a,
1010
+ b,
1011
+ c,
1012
+ shared_shard,
1013
+ c_multicast_i64,
1014
+ max_active_clusters,
1015
+ stream,
1016
+ )
1017
+
1018
+
1019
+_COMPILED: dict[tuple[object, ...], object] = {}
1020
+
1021
+
1022
+def compile_kernel(
1023
+ mnkl: tuple[int, int, int, int],
1024
+ a: cute.Tensor,
1025
+ b: cute.Tensor,
1026
+ c: cute.Tensor,
1027
+ shared_shard: cute.Tensor,
1028
+ full_hidden_dim: int,
1029
+ shard_dim: int,
1030
+ mma_tiler_mn: tuple[int, int] = (64, 32),
1031
+ cluster_shape_mn: tuple[int, int] = (1, 8),
1032
+ max_active_clusters: cutlass.Constexpr = None,
1033
+ b_prime_stages: int = 2,
1034
+):
1035
+ key = (
1036
+ torch.accelerator.current_device_index(),
1037
+ mnkl,
1038
+ full_hidden_dim,
1039
+ shard_dim,
1040
+ mma_tiler_mn,
1041
+ cluster_shape_mn,
1042
+ max_active_clusters,
1043
+ b_prime_stages,
1044
+ )
1045
+ if key in _COMPILED:
1046
+ return _COMPILED[key]
1047
+
1048
+ from cutlass.cute.runtime import make_fake_stream, make_fake_tensor
1049
+
1050
+ gemm = FusedAddMulticastGemm(
1051
+ mma_tiler_mn,
1052
+ cluster_shape_mn,
1053
+ b_prime_stages,
1054
+ )
1055
+ validate_configuration(
1056
+ latent_dim=mnkl[2],
1057
+ shard_dim=mnkl[1],
1058
+ mma_tiler_mn=mma_tiler_mn,
1059
+ cluster_shape_mn=cluster_shape_mn,
1060
+ b_prime_stages=b_prime_stages,
1061
+ )
1062
+ if any(tensor.element_type is not cutlass.BFloat16 for tensor in (a, b, c)):
1063
+ raise ValueError("up-projection tensors must be BF16")
1064
+
1065
+ # The producer writes [M, full_hidden_dim]. GEMM receives a rank-local
1066
+ # [M, shard_dim] view: contiguous within a row, with full_hidden_dim as
1067
+ # its leading dimension.
1068
+ shared_shard_compile = make_fake_tensor(
1069
+ shared_shard.element_type,
1070
+ (mnkl[0], shard_dim),
1071
+ stride=(full_hidden_dim, 1),
1072
+ assumed_align=16,
1073
+ )
1074
+ stream = make_fake_stream()
1075
+ compiled = cute.compile(
1076
+ launch_kernel,
1077
+ gemm,
1078
+ a,
1079
+ b,
1080
+ c,
1081
+ shared_shard_compile,
1082
+ cutlass.Int64(mnkl[0]),
1083
+ cutlass.Int64(0),
1084
+ full_hidden_dim,
1085
+ shard_dim,
1086
+ max_active_clusters,
1087
+ stream,
1088
+ )
1089
+ _COMPILED[key] = compiled
1090
+ return compiled
1091
+
1092
+
1093
+class AdaptiveUpProjectionKernel:
1094
+ """Dispatch static-M Skinny or dynamic-M WGMMA into one mailbox."""
1095
+
1096
+ def __init__(
1097
+ self,
1098
+ *,
1099
+ group: dist.ProcessGroup,
1100
+ rank: int,
1101
+ tp_size: int,
1102
+ latent_dim: int,
1103
+ hidden_dim: int,
1104
+ max_m: int,
1105
+ skinny_max_m: int,
1106
+ mma_tiler_mn: tuple[int, int],
1107
+ cluster_shape_mn: tuple[int, int],
1108
+ b_prime_stages: int,
1109
+ ) -> None:
1110
+ if hidden_dim % tp_size:
1111
+ raise ValueError("hidden_dim must be divisible by TP size")
1112
+ if not 0 <= skinny_max_m <= min(8, max_m):
1113
+ raise ValueError("skinny_max_m must be in [0, min(8, max_m)]")
1114
+ self.rank = rank
1115
+ self.tp_size = tp_size
1116
+ self.latent_dim = latent_dim
1117
+ self.hidden_dim = hidden_dim
1118
+ self.shard_dim = hidden_dim // tp_size
1119
+ self.max_m = max_m
1120
+ self.skinny_max_m = skinny_max_m
1121
+ self.mma_tiler_mn = mma_tiler_mn
1122
+ self.cluster_shape_mn = cluster_shape_mn
1123
+ self.b_prime_stages = b_prime_stages
1124
+ device = torch.device("cuda", torch.accelerator.current_device_index())
1125
+ self._device = device
1126
+ self._dynamic: Any | None = None
1127
+ self._skinny_by_m: dict[int, FusedAddMulticastSkinnyGemmKernel] = {}
1128
+ validate_configuration(
1129
+ latent_dim=latent_dim,
1130
+ shard_dim=self.shard_dim,
1131
+ mma_tiler_mn=mma_tiler_mn,
1132
+ cluster_shape_mn=cluster_shape_mn,
1133
+ b_prime_stages=b_prime_stages,
1134
+ )
1135
+ if skinny_max_m and self.latent_dim % (224 * 8):
1136
+ raise ValueError(
1137
+ "Skinny up-projection requires latent_dim divisible by 1792."
1138
+ )
1139
+
1140
+ self._mailbox = symm_mem.empty(
1141
+ (1, max_m, hidden_dim),
1142
+ dtype=torch.bfloat16,
1143
+ device=device,
1144
+ )
1145
+ self._mailbox_symm_mem = symm_mem.rendezvous(self._mailbox, group)
1146
+ self._mailbox.view(torch.int32).fill_(-0x80000000)
1147
+ multicast_ptr = self._mailbox_symm_mem.multicast_ptr
1148
+ if multicast_ptr is None or multicast_ptr == 0:
1149
+ raise RuntimeError("mailbox NVLS multicast mapping is unavailable")
1150
+ self._mailbox_multicast_ptr = int(multicast_ptr)
1151
+
1152
+ cluster_size = math.prod(cluster_shape_mn)
1153
+ self._max_active_clusters = utils.HardwareInfo().get_max_active_clusters(
1154
+ cluster_size
1155
+ )
1156
+ self._mailbox_c = _as_cute(self._mailbox)
1157
+
1158
+ def compile_dynamic(self) -> None:
1159
+ if self._dynamic is not None:
1160
+ return
1161
+ device = self._device
1162
+ with torch.accelerator.device_index(device.index):
1163
+ compile_latent = torch.empty(
1164
+ (1, self.max_m, self.latent_dim),
1165
+ dtype=torch.bfloat16,
1166
+ device=device,
1167
+ )
1168
+ compile_weight = torch.empty(
1169
+ (self.shard_dim, self.latent_dim),
1170
+ dtype=torch.bfloat16,
1171
+ device=device,
1172
+ )
1173
+ compile_shared = torch.empty(
1174
+ (self.max_m, self.hidden_dim),
1175
+ dtype=torch.bfloat16,
1176
+ device=device,
1177
+ )[
1178
+ :,
1179
+ self.rank * self.shard_dim : (self.rank + 1) * self.shard_dim,
1180
+ ]
1181
+ compile_latent_c = _as_cute(compile_latent, dynamic_m=True)
1182
+ compile_weight_c = _as_cute(compile_weight.unsqueeze(0))
1183
+ compile_shared_c = _as_cute(compile_shared)
1184
+
1185
+ self._dynamic = compile_kernel(
1186
+ (self.max_m, self.shard_dim, self.latent_dim, 1),
1187
+ compile_latent_c,
1188
+ compile_weight_c,
1189
+ self._mailbox_c,
1190
+ compile_shared_c,
1191
+ self.hidden_dim,
1192
+ self.shard_dim,
1193
+ self.mma_tiler_mn,
1194
+ self.cluster_shape_mn,
1195
+ self._max_active_clusters,
1196
+ self.b_prime_stages,
1197
+ )
1198
+
1199
+ def compile_skinny(self, m: int) -> None:
1200
+ if not 1 <= m <= self.skinny_max_m:
1201
+ raise ValueError(
1202
+ f"Skinny up-projection requires M in [1, {self.skinny_max_m}]."
1203
+ )
1204
+ if m in self._skinny_by_m:
1205
+ return
1206
+ with torch.accelerator.device_index(self._device.index):
1207
+ self._skinny_by_m[m] = FusedAddMulticastSkinnyGemmKernel(
1208
+ rank=self.rank,
1209
+ tp_size=self.tp_size,
1210
+ latent_dim=self.latent_dim,
1211
+ hidden_dim=self.hidden_dim,
1212
+ num_rows=m,
1213
+ )
1214
+
1215
+ def ensure_compiled(self, m: int) -> None:
1216
+ if not 1 <= m <= self.max_m:
1217
+ raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]")
1218
+ if m <= self.skinny_max_m:
1219
+ self.compile_skinny(m)
1220
+ else:
1221
+ self.compile_dynamic()
1222
+
1223
+ def __call__(
1224
+ self,
1225
+ latent: torch.Tensor,
1226
+ weight: torch.Tensor,
1227
+ shared_shard: torch.Tensor,
1228
+ ) -> torch.Tensor:
1229
+ if latent.ndim != 2:
1230
+ raise ValueError("latent must be rank-2")
1231
+ m = latent.shape[0]
1232
+ device = self._mailbox.device
1233
+ expected = (
1234
+ (latent, (m, self.latent_dim), "latent"),
1235
+ (
1236
+ weight,
1237
+ (self.shard_dim, self.latent_dim),
1238
+ "weight",
1239
+ ),
1240
+ (
1241
+ shared_shard,
1242
+ (self.max_m, self.shard_dim),
1243
+ "shared_shard",
1244
+ ),
1245
+ )
1246
+ for tensor, shape, name in expected:
1247
+ if (
1248
+ tensor.shape != shape
1249
+ or tensor.dtype != torch.bfloat16
1250
+ or tensor.device != device
1251
+ ):
1252
+ raise ValueError(f"{name} must be CUDA torch.bfloat16 {list(shape)}")
1253
+ if (
1254
+ not latent.is_contiguous()
1255
+ or not weight.is_contiguous()
1256
+ or shared_shard.stride() != (self.hidden_dim, 1)
1257
+ ):
1258
+ raise ValueError("up-projection inputs have unsupported strides")
1259
+ if not 1 <= m <= self.max_m:
1260
+ raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]")
1261
+
1262
+ if m <= self.skinny_max_m:
1263
+ skinny = self._skinny_by_m.get(m)
1264
+ if skinny is None:
1265
+ raise RuntimeError(
1266
+ f"Skinny up-projection M={m} was not compiled before launch."
1267
+ )
1268
+ return skinny(
1269
+ latent,
1270
+ weight,
1271
+ shared_shard,
1272
+ self._mailbox,
1273
+ self._mailbox_multicast_ptr,
1274
+ )
1275
+
1276
+ if self._dynamic is None:
1277
+ raise RuntimeError("Dynamic up-projection was not compiled before launch.")
1278
+ with torch.accelerator.device_index(device.index):
1279
+ stream = cuda.CUstream(torch.cuda.current_stream(device).cuda_stream)
1280
+ self._dynamic(
1281
+ _as_cute(latent.unsqueeze(0), dynamic_m=True),
1282
+ _as_cute(weight.unsqueeze(0)),
1283
+ self._mailbox_c,
1284
+ _as_cute(shared_shard),
1285
+ cutlass.Int64(m),
1286
+ cutlass.Int64(
1287
+ self._mailbox_multicast_ptr + self.rank * self.shard_dim * 2
1288
+ ),
1289
+ stream,
1290
+ )
1291
+ return self._mailbox
vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py
new
+438
@@ -0,0 +1,438 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""Static-M SIMT skinny GEMM with a shared-add multicast epilogue."""
5
+
6
+from __future__ import annotations
7
+
8
+from dataclasses import dataclass
9
+
10
+import cuda.bindings.driver as cuda
11
+import cutlass
12
+import cutlass.cute as cute
13
+import torch
14
+from cutlass import BFloat16, Float32, Int64, const_expr
15
+from cutlass.cute.runtime import from_dlpack
16
+
17
+from .primitives import (
18
+ CUDAGraphCompatibleWrapper,
19
+ bf16x2_to_u32,
20
+ bf16x4_to_packed_u32x2,
21
+ bf16x8_to_packed_u32x4,
22
+ sanitize_negative_zero,
23
+ sanitize_negative_zero_u32,
24
+ sanitize_negative_zero_u32x2,
25
+ store_global_u32,
26
+ store_global_u32x2,
27
+ store_global_u32x4,
28
+)
29
+
30
+
31
+@dataclass(frozen=True)
32
+class SkinnyConfig:
33
+ block_size: int = 224
34
+ outputs_per_block: int = 2
35
+ k_unroll: int = 2
36
+ vector_width: int = 8
37
+ prefetch_b_before_pdl: bool = True
38
+
39
+
40
+def config_for_m(num_rows: int, shard_dim: int = 896) -> SkinnyConfig:
41
+ if shard_dim == 448:
42
+ if num_rows >= 6:
43
+ return SkinnyConfig(
44
+ block_size=224,
45
+ outputs_per_block=2,
46
+ k_unroll=1,
47
+ )
48
+ outputs_per_block = 2 if num_rows <= 3 else 4
49
+ return SkinnyConfig(
50
+ block_size=448,
51
+ outputs_per_block=outputs_per_block,
52
+ k_unroll=1,
53
+ )
54
+ if num_rows == 1:
55
+ return SkinnyConfig(outputs_per_block=8, k_unroll=1)
56
+ return SkinnyConfig(outputs_per_block=4)
57
+
58
+
59
+def _as_cute(tensor: torch.Tensor):
60
+ return from_dlpack(
61
+ CUDAGraphCompatibleWrapper(tensor.detach()),
62
+ assumed_align=16,
63
+ )
64
+
65
+
66
+class FusedAddMulticastSkinnyGemm:
67
+ """SIMT GEMM adapted from the existing Skinny GEMM."""
68
+
69
+ def __init__(
70
+ self,
71
+ *,
72
+ num_rows: int,
73
+ hidden_dim: int,
74
+ config: SkinnyConfig,
75
+ ) -> None:
76
+ if config.block_size % 32:
77
+ raise ValueError("skinny block_size must be a multiple of 32")
78
+ self.num_rows = num_rows
79
+ self.hidden_dim = hidden_dim
80
+ self.block_size = config.block_size
81
+ self.outputs_per_block = config.outputs_per_block
82
+ self.k_unroll = config.k_unroll
83
+ self.vector_width = config.vector_width
84
+ self.prefetch_b_before_pdl = config.prefetch_b_before_pdl
85
+ self.num_warps = config.block_size // 32
86
+
87
+ @cute.jit
88
+ def __call__(
89
+ self,
90
+ gA: cute.Tensor,
91
+ gB: cute.Tensor,
92
+ gShared: cute.Tensor,
93
+ output_multicast_ptr: Int64,
94
+ stream: cuda.CUstream,
95
+ ) -> None:
96
+ n = cute.size(gB, mode=[0])
97
+ k = cute.size(gA, mode=[1])
98
+ copy_a = cute.make_copy_atom(
99
+ cute.nvgpu.CopyG2ROp(),
100
+ BFloat16,
101
+ num_bits_per_copy=self.vector_width * BFloat16.width,
102
+ load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
103
+ )
104
+ copy_b = cute.make_copy_atom(
105
+ cute.nvgpu.CopyG2ROp(),
106
+ BFloat16,
107
+ num_bits_per_copy=self.vector_width * BFloat16.width,
108
+ load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING,
109
+ )
110
+ self.kernel(
111
+ gA,
112
+ gB,
113
+ gShared,
114
+ output_multicast_ptr,
115
+ k,
116
+ copy_a,
117
+ copy_b,
118
+ ).launch(
119
+ grid=[cute.ceil_div(n, self.outputs_per_block), 1, 1],
120
+ block=[self.block_size, 1, 1],
121
+ smem=(self.num_rows * self.outputs_per_block * self.num_warps * 4),
122
+ stream=stream,
123
+ use_pdl=True,
124
+ min_blocks_per_mp=1,
125
+ )
126
+
127
+ @cute.kernel
128
+ def kernel(
129
+ self,
130
+ gA: cute.Tensor,
131
+ gB: cute.Tensor,
132
+ gShared: cute.Tensor,
133
+ output_multicast_ptr: Int64,
134
+ k_extent: cutlass.Int32,
135
+ copy_a: cute.CopyAtom,
136
+ copy_b: cute.CopyAtom,
137
+ ) -> None:
138
+ tidx, _, _ = cute.arch.thread_idx()
139
+ block_idx, _, _ = cute.arch.block_idx()
140
+ warp_idx = cute.arch.warp_idx()
141
+
142
+ outputs_per_block: cutlass.Constexpr = self.outputs_per_block
143
+ vector_width: cutlass.Constexpr = self.vector_width
144
+ block_size: cutlass.Constexpr = self.block_size
145
+ num_warps: cutlass.Constexpr = self.num_warps
146
+ num_rows: cutlass.Constexpr = self.num_rows
147
+
148
+ acc = cute.make_rmem_tensor(
149
+ cute.make_layout(
150
+ (num_rows, outputs_per_block),
151
+ stride=(outputs_per_block, 1),
152
+ ),
153
+ Float32,
154
+ )
155
+ acc.fill(0.0)
156
+
157
+ n_base = block_idx * outputs_per_block
158
+ k_tile_size: cutlass.Constexpr = block_size * vector_width
159
+ num_k_tiles = k_extent // k_tile_size
160
+ gA_vec = cute.logical_divide(gA, (None, vector_width))
161
+ gB_vec = cute.logical_divide(gB, (None, vector_width))
162
+ tA_all = cute.logical_divide(gA_vec, (None, (None, block_size)))
163
+ tB_all = cute.logical_divide(gB_vec, (None, (None, block_size)))
164
+ tA = tA_all[None, (None, (tidx, None))]
165
+
166
+ a_regs = cute.make_rmem_tensor(
167
+ cute.make_layout(
168
+ (num_rows, vector_width),
169
+ stride=(vector_width, 1),
170
+ ),
171
+ BFloat16,
172
+ )
173
+ b_regs = cute.make_rmem_tensor(
174
+ cute.make_layout(
175
+ (outputs_per_block, vector_width),
176
+ stride=(vector_width, 1),
177
+ ),
178
+ BFloat16,
179
+ )
180
+
181
+ if const_expr(self.prefetch_b_before_pdl):
182
+ for ni in cutlass.range_constexpr(outputs_per_block):
183
+ tB = tB_all[n_base + ni, (None, (tidx, None))]
184
+ cute.copy(copy_b, tB[None, 0], b_regs[ni, None])
185
+
186
+ cute.arch.griddepcontrol_wait()
187
+
188
+ for mi in cutlass.range_constexpr(num_rows):
189
+ cute.copy(copy_a, tA[mi, None, 0], a_regs[mi, None])
190
+ if const_expr(not self.prefetch_b_before_pdl):
191
+ for ni in cutlass.range_constexpr(outputs_per_block):
192
+ tB = tB_all[n_base + ni, (None, (tidx, None))]
193
+ cute.copy(copy_b, tB[None, 0], b_regs[ni, None])
194
+ for vi in cutlass.range_constexpr(vector_width):
195
+ for mi in cutlass.range_constexpr(num_rows):
196
+ for ni in cutlass.range_constexpr(outputs_per_block):
197
+ acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to(Float32) * b_regs[
198
+ ni, vi
199
+ ].to(Float32)
200
+
201
+ for k_tile in cutlass.range(1, num_k_tiles, unroll=self.k_unroll):
202
+ for mi in cutlass.range_constexpr(num_rows):
203
+ cute.copy(
204
+ copy_a,
205
+ tA[mi, None, k_tile],
206
+ a_regs[mi, None],
207
+ )
208
+ for ni in cutlass.range_constexpr(outputs_per_block):
209
+ tB = tB_all[n_base + ni, (None, (tidx, None))]
210
+ cute.copy(copy_b, tB[None, k_tile], b_regs[ni, None])
211
+ for vi in cutlass.range_constexpr(vector_width):
212
+ for mi in cutlass.range_constexpr(num_rows):
213
+ for ni in cutlass.range_constexpr(outputs_per_block):
214
+ acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to(Float32) * b_regs[
215
+ ni, vi
216
+ ].to(Float32)
217
+
218
+ for mi in cutlass.range_constexpr(num_rows):
219
+ for ni in cutlass.range_constexpr(outputs_per_block):
220
+ acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni])
221
+
222
+ smem_layout = cute.make_layout(
223
+ (num_rows, outputs_per_block, num_warps),
224
+ stride=(outputs_per_block * num_warps, num_warps, 1),
225
+ )
226
+ smem = cutlass.utils.SmemAllocator()
227
+ partials = smem.allocate_tensor(
228
+ Float32,
229
+ smem_layout,
230
+ byte_alignment=16,
231
+ )
232
+ with cute.arch.elect_one():
233
+ for mi in cutlass.range_constexpr(num_rows):
234
+ for ni in cutlass.range_constexpr(outputs_per_block):
235
+ partials[mi, ni, warp_idx] = acc[mi, ni]
236
+
237
+ cute.arch.sync_threads()
238
+ if tidx == 0:
239
+ fused = cute.make_rmem_tensor(
240
+ cute.make_layout((outputs_per_block,)),
241
+ BFloat16,
242
+ )
243
+ for mi in cutlass.range_constexpr(num_rows):
244
+ for ni in cutlass.range_constexpr(outputs_per_block):
245
+ total = (
246
+ partials[mi, ni, None]
247
+ .load()
248
+ .reduce(
249
+ cute.ReductionOp.ADD,
250
+ init_val=Float32(0.0),
251
+ reduction_profile=0,
252
+ )
253
+ )
254
+ gemm_value = Float32(total).to(BFloat16)
255
+ fused[ni] = (
256
+ gemm_value.to(Float32) + gShared[mi, n_base + ni].to(Float32)
257
+ ).to(BFloat16)
258
+ output_offset = Int64((mi * self.hidden_dim + n_base) * 2)
259
+ if const_expr(outputs_per_block == 2):
260
+ packed = sanitize_negative_zero_u32(bf16x2_to_u32(fused.load()))
261
+ store_global_u32(
262
+ output_multicast_ptr + output_offset,
263
+ packed,
264
+ )
265
+ elif const_expr(outputs_per_block == 4):
266
+ packed = sanitize_negative_zero_u32x2(
267
+ bf16x4_to_packed_u32x2(fused.load())
268
+ )
269
+ store_global_u32x2(
270
+ output_multicast_ptr + output_offset,
271
+ packed,
272
+ )
273
+ else:
274
+ packed = sanitize_negative_zero(
275
+ bf16x8_to_packed_u32x4(fused.load())
276
+ )
277
+ store_global_u32x4(
278
+ output_multicast_ptr + output_offset,
279
+ packed,
280
+ )
281
+
282
+ cute.arch.griddepcontrol_launch_dependents()
283
+
284
+
285
+_COMPILED: dict[tuple[object, ...], object] = {}
286
+
287
+
288
+def compile_kernel(
289
+ *,
290
+ num_rows: int,
291
+ latent_dim: int,
292
+ hidden_dim: int,
293
+ shard_dim: int,
294
+ config: SkinnyConfig,
295
+):
296
+ key = (
297
+ torch.accelerator.current_device_index(),
298
+ num_rows,
299
+ latent_dim,
300
+ hidden_dim,
301
+ shard_dim,
302
+ config,
303
+ )
304
+ if key in _COMPILED:
305
+ return _COMPILED[key]
306
+
307
+ from cutlass.cute.runtime import make_fake_stream, make_fake_tensor
308
+
309
+ if shard_dim % config.outputs_per_block:
310
+ raise ValueError("shard_dim must be divisible by outputs_per_block")
311
+ if latent_dim % (config.block_size * config.vector_width):
312
+ raise ValueError("latent_dim must be divisible by block_size * vector_width")
313
+
314
+ a = make_fake_tensor(
315
+ BFloat16,
316
+ (num_rows, latent_dim),
317
+ stride=(latent_dim, 1),
318
+ assumed_align=16,
319
+ )
320
+ b = make_fake_tensor(
321
+ BFloat16,
322
+ (shard_dim, latent_dim),
323
+ stride=(latent_dim, 1),
324
+ assumed_align=16,
325
+ )
326
+ shared = make_fake_tensor(
327
+ BFloat16,
328
+ (num_rows, shard_dim),
329
+ stride=(hidden_dim, 1),
330
+ assumed_align=16,
331
+ )
332
+ compiled = cute.compile(
333
+ FusedAddMulticastSkinnyGemm(
334
+ num_rows=num_rows,
335
+ hidden_dim=hidden_dim,
336
+ config=config,
337
+ ),
338
+ a,
339
+ b,
340
+ shared,
341
+ Int64(0),
342
+ make_fake_stream(),
343
+ options="--ptxas-options -maxrregcount=128",
344
+ )
345
+ _COMPILED[key] = compiled
346
+ return compiled
347
+
348
+
349
+class FusedAddMulticastSkinnyGemmKernel:
350
+ """Buffer-free compiled launcher for one static M."""
351
+
352
+ def __init__(
353
+ self,
354
+ *,
355
+ rank: int,
356
+ tp_size: int,
357
+ latent_dim: int,
358
+ hidden_dim: int,
359
+ num_rows: int,
360
+ ) -> None:
361
+ if not 1 <= num_rows <= 8:
362
+ raise ValueError("skinny backend requires static M in [1, 8]")
363
+ if hidden_dim % tp_size:
364
+ raise ValueError("hidden_dim must be divisible by TP size")
365
+ self.rank = rank
366
+ self.latent_dim = latent_dim
367
+ self.hidden_dim = hidden_dim
368
+ self.shard_dim = hidden_dim // tp_size
369
+ self.num_rows = num_rows
370
+ self._skinny = compile_kernel(
371
+ num_rows=num_rows,
372
+ latent_dim=latent_dim,
373
+ hidden_dim=hidden_dim,
374
+ shard_dim=self.shard_dim,
375
+ config=config_for_m(num_rows, self.shard_dim),
376
+ )
377
+
378
+ def __call__(
379
+ self,
380
+ latent: torch.Tensor,
381
+ weight: torch.Tensor,
382
+ shared_shard: torch.Tensor,
383
+ mailbox: torch.Tensor,
384
+ mailbox_multicast_ptr: int,
385
+ ) -> torch.Tensor:
386
+ device = mailbox.device
387
+ expected = (
388
+ (
389
+ latent,
390
+ (self.num_rows, self.latent_dim),
391
+ torch.bfloat16,
392
+ "latent",
393
+ ),
394
+ (
395
+ weight,
396
+ (self.shard_dim, self.latent_dim),
397
+ torch.bfloat16,
398
+ "weight",
399
+ ),
400
+ (
401
+ shared_shard,
402
+ (mailbox.shape[1], self.shard_dim),
403
+ torch.bfloat16,
404
+ "shared_shard",
405
+ ),
406
+ (
407
+ mailbox,
408
+ (1, mailbox.shape[1], self.hidden_dim),
409
+ torch.bfloat16,
410
+ "mailbox",
411
+ ),
412
+ )
413
+ for tensor, shape, dtype, name in expected:
414
+ if (
415
+ tensor.shape != shape
416
+ or tensor.dtype != dtype
417
+ or tensor.device != device
418
+ ):
419
+ raise ValueError(f"{name} must be CUDA {dtype} {list(shape)}")
420
+ if (
421
+ not latent.is_contiguous()
422
+ or not weight.is_contiguous()
423
+ or shared_shard.stride() != (self.hidden_dim, 1)
424
+ or not mailbox.is_contiguous()
425
+ ):
426
+ raise ValueError("skinny up-projection inputs have unsupported strides")
427
+ if mailbox.shape[1] < self.num_rows:
428
+ raise ValueError("mailbox capacity is smaller than runtime M")
429
+
430
+ with torch.accelerator.device_index(device.index):
431
+ self._skinny(
432
+ _as_cute(latent),
433
+ _as_cute(weight),
434
+ _as_cute(shared_shard[: self.num_rows]),
435
+ Int64(mailbox_multicast_ptr + self.rank * self.shard_dim * 2),
436
+ cuda.CUstream(torch.cuda.current_stream(device).cuda_stream),
437
+ )
438
+ return mailbox
vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py
new
+226
@@ -0,0 +1,226 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""Lamport mailbox-to-local copy and reset."""
5
+
6
+from __future__ import annotations
7
+
8
+import functools
9
+
10
+import cuda.bindings.driver as cuda
11
+import cutlass
12
+import cutlass.cute as cute
13
+import torch
14
+from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream
15
+
16
+from .primitives import (
17
+ VEC_BF16,
18
+ fragment_is_dirty,
19
+ load_global_u32x4,
20
+ store_global_u32x4,
21
+ store_lamport_sentinel_128,
22
+ to_cute,
23
+ to_cute_dynamic_m,
24
+)
25
+
26
+
27
+class LamportCopy:
28
+ """Consume the local physical copy of an NVLS-multicast mailbox."""
29
+
30
+ def __init__(self, hidden_dim: int, ctas: int, threads: int):
31
+ self.hidden_dim = hidden_dim
32
+ self.ctas = ctas
33
+ self.threads = threads
34
+
35
+ @cute.jit
36
+ def __call__(
37
+ self,
38
+ symmetric_mailbox: cute.Tensor,
39
+ local_output: cute.Tensor,
40
+ m: cutlass.Int32,
41
+ stream: cuda.CUstream,
42
+ ):
43
+ self.kernel(symmetric_mailbox, local_output, m).launch(
44
+ grid=(self.ctas, 1, 1),
45
+ block=(self.threads, 1, 1),
46
+ stream=stream,
47
+ use_pdl=True,
48
+ )
49
+
50
+ @cute.kernel
51
+ def kernel(
52
+ self,
53
+ symmetric_mailbox: cute.Tensor,
54
+ local_output: cute.Tensor,
55
+ m: cutlass.Int32,
56
+ ):
57
+ # The CTA may be scheduled early, but mailbox inspection must not pass
58
+ # the producer GEMM's programmatic completion point.
59
+ cute.arch.griddepcontrol_wait()
60
+
61
+ tidx, _, _ = cute.arch.thread_idx()
62
+ block, _, _ = cute.arch.block_idx()
63
+ thread = cutlass.Int64(block * self.threads + tidx)
64
+ stride = cutlass.Int64(self.ctas * self.threads)
65
+ fragments = cutlass.Int64(m) * cutlass.Int64(self.hidden_dim // VEC_BF16)
66
+
67
+ fragment = thread
68
+ while fragment < fragments:
69
+ element = fragment * VEC_BF16
70
+ source = cute.make_ptr(
71
+ cutlass.BFloat16,
72
+ (symmetric_mailbox.iterator + element).llvm_ptr,
73
+ cute.AddressSpace.gmem,
74
+ assumed_align=16,
75
+ )
76
+ packed = load_global_u32x4(source, volatile=True)
77
+ while fragment_is_dirty(packed):
78
+ packed = load_global_u32x4(source, volatile=True)
79
+
80
+ destination = cutlass.Int64((local_output.iterator + element).toint())
81
+ store_global_u32x4(destination, packed, volatile=False)
82
+ fragment = fragment + stride
83
+
84
+ # The returned ordinary tensor is complete. A same-stream successor
85
+ # may overlap the mailbox cleanup below.
86
+ cute.arch.griddepcontrol_launch_dependents()
87
+
88
+ fragment = thread
89
+ while fragment < fragments:
90
+ element = fragment * VEC_BF16
91
+ source = cute.make_ptr(
92
+ cutlass.BFloat16,
93
+ (symmetric_mailbox.iterator + element).llvm_ptr,
94
+ cute.AddressSpace.gmem,
95
+ assumed_align=16,
96
+ )
97
+ store_lamport_sentinel_128(source)
98
+ fragment = fragment + stride
99
+
100
+
101
+@functools.cache
102
+def compile_kernel(
103
+ hidden_dim: int,
104
+ max_m: int,
105
+ ctas: int,
106
+ threads: int,
107
+ device_index: int,
108
+):
109
+ if hidden_dim <= 0 or hidden_dim % VEC_BF16:
110
+ raise ValueError("hidden_dim must be a positive multiple of 8")
111
+ if max_m <= 0:
112
+ raise ValueError("max_m must be positive")
113
+ if ctas <= 0:
114
+ raise ValueError("Lamport copy ctas must be positive")
115
+ if not 1 <= threads <= 1024:
116
+ raise ValueError("Lamport copy threads must be in [1, 1024]")
117
+ with torch.accelerator.device_index(device_index):
118
+ mailbox = make_fake_compact_tensor(
119
+ cutlass.BFloat16, (max_m * hidden_dim,), assumed_align=16
120
+ )
121
+ output = make_fake_compact_tensor(
122
+ cutlass.BFloat16,
123
+ (cute.sym_int32(divisibility=VEC_BF16),),
124
+ assumed_align=16,
125
+ )
126
+ return cute.compile(
127
+ LamportCopy(hidden_dim, ctas, threads),
128
+ mailbox,
129
+ output,
130
+ cutlass.Int32(max_m),
131
+ make_fake_stream(),
132
+ )
133
+
134
+
135
+def launch(
136
+ symmetric_mailbox: torch.Tensor,
137
+ local_output: torch.Tensor,
138
+ *,
139
+ m: int,
140
+ hidden_dim: int,
141
+ max_m: int,
142
+ ctas: int,
143
+ threads: int,
144
+) -> None:
145
+ if not 1 <= m <= max_m:
146
+ raise ValueError(f"runtime M={m} must be in [1, {max_m}]")
147
+ if (
148
+ symmetric_mailbox.ndim != 3
149
+ or symmetric_mailbox.dtype != torch.bfloat16
150
+ or not symmetric_mailbox.is_cuda
151
+ or not symmetric_mailbox.is_contiguous()
152
+ or symmetric_mailbox.shape[0] != 1
153
+ or symmetric_mailbox.shape[2] != hidden_dim
154
+ ):
155
+ raise ValueError(
156
+ f"symmetric_mailbox must be contiguous CUDA BF16 [1,M,{hidden_dim}]"
157
+ )
158
+ if symmetric_mailbox.shape[1] != max_m:
159
+ raise ValueError("symmetric mailbox must retain its full max_m capacity")
160
+ if (
161
+ local_output.shape != (1, m, hidden_dim)
162
+ or local_output.dtype != torch.bfloat16
163
+ or local_output.device != symmetric_mailbox.device
164
+ or not local_output.is_contiguous()
165
+ ):
166
+ raise ValueError("local_output must be contiguous CUDA BF16 [1,M,H]")
167
+
168
+ device_index = symmetric_mailbox.device.index
169
+ stream = cuda.CUstream(
170
+ torch.cuda.current_stream(symmetric_mailbox.device).cuda_stream
171
+ )
172
+ compile_kernel(hidden_dim, max_m, ctas, threads, device_index)(
173
+ to_cute(symmetric_mailbox.flatten(), 16),
174
+ to_cute_dynamic_m(
175
+ local_output.flatten(),
176
+ mode=0,
177
+ assumed_align=16,
178
+ ),
179
+ cutlass.Int32(m),
180
+ stream,
181
+ )
182
+
183
+
184
+class LamportCopyKernel:
185
+ """Copy a borrowed symmetric mailbox into a fresh local tensor."""
186
+
187
+ def __init__(
188
+ self,
189
+ *,
190
+ hidden_dim: int,
191
+ max_m: int,
192
+ ctas: int,
193
+ threads: int,
194
+ ) -> None:
195
+ self.hidden_dim = hidden_dim
196
+ self.max_m = max_m
197
+ self.ctas = ctas
198
+ self.threads = threads
199
+ compile_kernel(
200
+ hidden_dim,
201
+ max_m,
202
+ ctas,
203
+ threads,
204
+ torch.accelerator.current_device_index(),
205
+ )
206
+
207
+ def __call__(self, symmetric_mailbox: torch.Tensor, *, m: int) -> torch.Tensor:
208
+ if not symmetric_mailbox.is_cuda:
209
+ raise ValueError("symmetric_mailbox must be a CUDA tensor")
210
+ device = symmetric_mailbox.device
211
+ with torch.accelerator.device_index(device.index):
212
+ output = torch.empty(
213
+ (1, m, self.hidden_dim),
214
+ dtype=torch.bfloat16,
215
+ device=device,
216
+ )
217
+ launch(
218
+ symmetric_mailbox,
219
+ output,
220
+ m=m,
221
+ hidden_dim=self.hidden_dim,
222
+ max_m=self.max_m,
223
+ ctas=self.ctas,
224
+ threads=self.threads,
225
+ )
226
+ return output
vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py
new
+437
@@ -0,0 +1,437 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""Shared CuTe DSL primitives; this module does not define a CUDA kernel."""
5
+
6
+from __future__ import annotations
7
+
8
+import cutlass
9
+import cutlass.cute as cute
10
+import torch
11
+from cutlass import BFloat16, Float32, Int32, Int64, Uint16, Uint32
12
+from cutlass._mlir import ir
13
+from cutlass._mlir.dialects import llvm, vector
14
+from cutlass.cute.runtime import from_dlpack
15
+from cutlass.cutlass_dsl import T, dsl_user_op
16
+
17
+VEC_BF16 = 8
18
+PACKED_BYTES = 16
19
+NUM_LAMPORT_BUFFERS = 3
20
+NEG_ZERO_F32_BITS = 0x80000000
21
+NEG_ZERO_BF16_BITS = 0x8000
22
+
23
+
24
+class CUDAGraphCompatibleWrapper:
25
+ """DLPack view that does not synchronize with the producer stream."""
26
+
27
+ def __init__(self, tensor: torch.Tensor):
28
+ self.tensor = tensor
29
+
30
+ def __dlpack__(self, stream=None):
31
+ return self.tensor.__dlpack__(stream=-1)
32
+
33
+ def __dlpack_device__(self):
34
+ return self.tensor.__dlpack_device__()
35
+
36
+
37
+def to_cute(tensor: torch.Tensor, assumed_align: int = 16) -> cute.Tensor:
38
+ return from_dlpack(
39
+ CUDAGraphCompatibleWrapper(tensor.detach()), assumed_align=assumed_align
40
+ )
41
+
42
+
43
+def to_cute_dynamic_m(
44
+ tensor: torch.Tensor,
45
+ *,
46
+ mode: int,
47
+ assumed_align: int = 16,
48
+) -> cute.Tensor:
49
+ """Expose exactly one compact tensor mode as a runtime shape.
50
+
51
+ Model dimensions remain part of the compiled tensor type. Only the token
52
+ mode is symbolic, so changing M within an Op's capacity reuses the same
53
+ compiled kernel.
54
+ """
55
+
56
+ return to_cute(tensor, assumed_align).mark_compact_shape_dynamic(
57
+ mode=mode,
58
+ stride_order=tensor.dim_order(),
59
+ )
60
+
61
+
62
+@dsl_user_op
63
+def load_global_u32x4(
64
+ pointer: cute.Pointer,
65
+ *,
66
+ volatile: cutlass.Constexpr[bool] = False,
67
+ loc=None,
68
+ ip=None,
69
+):
70
+ """Load one 128-bit fragment as four u32 registers.
71
+
72
+ The volatile form is the Lamport polling load. Marking the asm
73
+ side-effecting prevents loop-invariant motion and common-subexpression
74
+ elimination across polling iterations.
75
+ """
76
+
77
+ address = pointer.toint(loc=loc, ip=ip)
78
+ opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32"
79
+ out = llvm.inline_asm(
80
+ llvm.StructType.get_literal([T.i32()] * 4),
81
+ [address.ir_value(loc=loc, ip=ip)],
82
+ f"{opcode} {{$0, $1, $2, $3}}, [$4];",
83
+ "=r,=r,=r,=r,l",
84
+ has_side_effects=volatile,
85
+ is_align_stack=False,
86
+ asm_dialect=llvm.AsmDialect.AD_ATT,
87
+ loc=loc,
88
+ ip=ip,
89
+ )
90
+ packed = vector.from_elements(
91
+ ir.VectorType.get([4], T.i32(), loc=loc),
92
+ [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(4)],
93
+ loc=loc,
94
+ ip=ip,
95
+ )
96
+ return cute.TensorSSA(packed, 4, Uint32)
97
+
98
+
99
+@dsl_user_op
100
+def store_global_u32x4(
101
+ address: Int64,
102
+ packed,
103
+ *,
104
+ volatile: cutlass.Constexpr[bool] = False,
105
+ loc=None,
106
+ ip=None,
107
+) -> None:
108
+ """Store four packed words to an ordinary or NVLS multicast global VA."""
109
+
110
+ words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(4)]
111
+ opcode = "st.volatile.global.v4.u32" if volatile else "st.global.v4.u32"
112
+ llvm.inline_asm(
113
+ None,
114
+ [address.ir_value(loc=loc, ip=ip), *words],
115
+ f"{opcode} [$0], {{$1, $2, $3, $4}};",
116
+ "l,r,r,r,r",
117
+ has_side_effects=True,
118
+ is_align_stack=False,
119
+ asm_dialect=llvm.AsmDialect.AD_ATT,
120
+ loc=loc,
121
+ ip=ip,
122
+ )
123
+
124
+
125
+@dsl_user_op
126
+def store_global_u32x2(
127
+ address: Int64,
128
+ packed,
129
+ *,
130
+ loc=None,
131
+ ip=None,
132
+) -> None:
133
+ words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(2)]
134
+ llvm.inline_asm(
135
+ None,
136
+ [address.ir_value(loc=loc, ip=ip), *words],
137
+ "st.global.v2.u32 [$0], {$1, $2};",
138
+ "l,r,r",
139
+ has_side_effects=True,
140
+ is_align_stack=False,
141
+ asm_dialect=llvm.AsmDialect.AD_ATT,
142
+ loc=loc,
143
+ ip=ip,
144
+ )
145
+
146
+
147
+@dsl_user_op
148
+def store_global_u32(
149
+ address: Int64,
150
+ word: Uint32,
151
+ *,
152
+ loc=None,
153
+ ip=None,
154
+) -> None:
155
+ llvm.inline_asm(
156
+ None,
157
+ [
158
+ address.ir_value(loc=loc, ip=ip),
159
+ word.ir_value(loc=loc, ip=ip),
160
+ ],
161
+ "st.global.u32 [$0], $1;",
162
+ "l,r",
163
+ has_side_effects=True,
164
+ is_align_stack=False,
165
+ asm_dialect=llvm.AsmDialect.AD_ATT,
166
+ loc=loc,
167
+ ip=ip,
168
+ )
169
+
170
+
171
+@dsl_user_op
172
+def store_lamport_sentinel_128(pointer: cute.Pointer, *, loc=None, ip=None) -> None:
173
+ """Reset one Lamport fragment to four FP32 negative-zero bit patterns."""
174
+
175
+ address = pointer.toint(loc=loc, ip=ip)
176
+ value = Uint32(NEG_ZERO_F32_BITS).ir_value(loc=loc, ip=ip)
177
+ llvm.inline_asm(
178
+ None,
179
+ [address.ir_value(loc=loc, ip=ip), value, value, value, value],
180
+ "st.global.v4.u32 [$0], {$1, $2, $3, $4};",
181
+ "l,r,r,r,r",
182
+ has_side_effects=True,
183
+ is_align_stack=False,
184
+ asm_dialect=llvm.AsmDialect.AD_ATT,
185
+ loc=loc,
186
+ ip=ip,
187
+ )
188
+
189
+
190
+@dsl_user_op
191
+def red_async_release_gpu_add_u32(
192
+ pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None
193
+) -> None:
194
+ """The exact SM100 arrival primitive used by upstream LamportFlags."""
195
+
196
+ address = pointer.toint(loc=loc, ip=ip)
197
+ llvm.inline_asm(
198
+ None,
199
+ [
200
+ address.ir_value(loc=loc, ip=ip),
201
+ value.ir_value(loc=loc, ip=ip),
202
+ ],
203
+ "red.async.release.global.gpu.add.u32 [$0], $1;",
204
+ "l,r",
205
+ has_side_effects=True,
206
+ is_align_stack=False,
207
+ asm_dialect=llvm.AsmDialect.AD_ATT,
208
+ loc=loc,
209
+ ip=ip,
210
+ )
211
+
212
+
213
+@dsl_user_op
214
+def load_volatile_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32:
215
+ address = pointer.toint(loc=loc, ip=ip)
216
+ return Uint32(
217
+ llvm.inline_asm(
218
+ T.i32(),
219
+ [address.ir_value(loc=loc, ip=ip)],
220
+ "ld.volatile.global.u32 $0, [$1];",
221
+ "=r,l",
222
+ has_side_effects=True,
223
+ is_align_stack=False,
224
+ asm_dialect=llvm.AsmDialect.AD_ATT,
225
+ loc=loc,
226
+ ip=ip,
227
+ )
228
+ )
229
+
230
+
231
+@dsl_user_op
232
+def map_shared_to_peer(
233
+ smem_ptr: cute.Pointer,
234
+ peer_rank: Int32,
235
+ *,
236
+ loc=None,
237
+ ip=None,
238
+) -> Int32:
239
+ """Map a local shared-memory slot to the same slot in a peer CTA."""
240
+
241
+ smem_address = smem_ptr.toint(loc=loc, ip=ip).ir_value()
242
+ return Int32(
243
+ llvm.inline_asm(
244
+ T.i32(),
245
+ [smem_address, peer_rank.ir_value(loc=loc, ip=ip)],
246
+ "mapa.shared::cluster.u32 $0, $1, $2;",
247
+ "=r,r,r",
248
+ has_side_effects=False,
249
+ is_align_stack=False,
250
+ asm_dialect=llvm.AsmDialect.AD_ATT,
251
+ loc=loc,
252
+ ip=ip,
253
+ )
254
+ )
255
+
256
+
257
+@dsl_user_op
258
+def store_shared_cluster_f32(
259
+ remote_address: Int32,
260
+ value: Float32,
261
+ *,
262
+ loc=None,
263
+ ip=None,
264
+) -> None:
265
+ llvm.inline_asm(
266
+ None,
267
+ [
268
+ remote_address.ir_value(loc=loc, ip=ip),
269
+ value.ir_value(loc=loc, ip=ip),
270
+ ],
271
+ "st.shared::cluster.f32 [$0], $1;",
272
+ "r,f",
273
+ has_side_effects=True,
274
+ is_align_stack=False,
275
+ asm_dialect=llvm.AsmDialect.AD_ATT,
276
+ loc=loc,
277
+ ip=ip,
278
+ )
279
+
280
+
281
+@dsl_user_op
282
+def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None):
283
+ target_type = ir.VectorType.get([VEC_BF16], BFloat16.mlir_type, loc=loc)
284
+ values = llvm.bitcast(
285
+ target_type,
286
+ packed.ir_value(loc=loc, ip=ip),
287
+ loc=loc,
288
+ ip=ip,
289
+ )
290
+ return cute.TensorSSA(values, VEC_BF16, BFloat16)
291
+
292
+
293
+@dsl_user_op
294
+def bf16x8_to_packed_u32x4(values, *, loc=None, ip=None):
295
+ target_type = ir.VectorType.get([4], T.i32(), loc=loc)
296
+ packed = llvm.bitcast(
297
+ target_type,
298
+ values.ir_value(loc=loc, ip=ip),
299
+ loc=loc,
300
+ ip=ip,
301
+ )
302
+ return cute.TensorSSA(packed, 4, Uint32)
303
+
304
+
305
+@dsl_user_op
306
+def bf16x4_to_packed_u32x2(values, *, loc=None, ip=None):
307
+ target_type = ir.VectorType.get([2], T.i32(), loc=loc)
308
+ packed = llvm.bitcast(
309
+ target_type,
310
+ values.ir_value(loc=loc, ip=ip),
311
+ loc=loc,
312
+ ip=ip,
313
+ )
314
+ return cute.TensorSSA(packed, 2, Uint32)
315
+
316
+
317
+@dsl_user_op
318
+def bf16x2_to_u32(values, *, loc=None, ip=None):
319
+ packed = llvm.bitcast(
320
+ T.i32(),
321
+ values.ir_value(loc=loc, ip=ip),
322
+ loc=loc,
323
+ ip=ip,
324
+ )
325
+ return Uint32(packed)
326
+
327
+
328
+@cute.jit
329
+def sanitize_negative_zero_u32(word):
330
+ low = Uint16(word & Uint32(0xFFFF))
331
+ high = Uint16(word >> Uint32(16))
332
+ if low == Uint16(NEG_ZERO_BF16_BITS):
333
+ word = word & Uint32(0xFFFF0000)
334
+ if high == Uint16(NEG_ZERO_BF16_BITS):
335
+ word = word & Uint32(0x0000FFFF)
336
+ return word
337
+
338
+
339
+@cute.jit
340
+def sanitize_negative_zero_u32x2(packed):
341
+ result = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32)
342
+ for i in cutlass.range_constexpr(2):
343
+ result[i] = sanitize_negative_zero_u32(packed[i])
344
+ return result.load()
345
+
346
+
347
+@cute.jit
348
+def sanitize_negative_zero(packed):
349
+ """Turn real BF16 -0 into +0 so it cannot equal the empty sentinel."""
350
+
351
+ result = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32)
352
+ for i in cutlass.range_constexpr(4):
353
+ word = packed[i]
354
+ low = Uint16(word & Uint32(0xFFFF))
355
+ high = Uint16(word >> Uint32(16))
356
+ if low == Uint16(NEG_ZERO_BF16_BITS):
357
+ word = word & Uint32(0xFFFF0000)
358
+ if high == Uint16(NEG_ZERO_BF16_BITS):
359
+ word = word & Uint32(0x0000FFFF)
360
+ result[i] = word
361
+ return result.load()
362
+
363
+
364
+@cute.jit
365
+def fragment_is_dirty(packed):
366
+ """Bit-exact upstream sentinel check: one comparison per 32-bit word."""
367
+
368
+ dirty = packed[0] == Uint32(NEG_ZERO_F32_BITS)
369
+ for i in cutlass.range_constexpr(1, 4):
370
+ dirty = dirty | (packed[i] == Uint32(NEG_ZERO_F32_BITS))
371
+ return dirty
372
+
373
+
374
+@cute.jit
375
+def warp_sum_specialized(
376
+ value: Float32,
377
+ warp_idx: Int32,
378
+ lane: Int32,
379
+ warps: cutlass.Constexpr[int],
380
+ last_warp_lanes: cutlass.Constexpr[int],
381
+ last_warp_mask: cutlass.Constexpr[int],
382
+) -> Float32:
383
+ """Warp sum supporting a compile-time partial final warp."""
384
+
385
+ if warp_idx == Int32(warps - 1) and cutlass.const_expr(last_warp_lanes < 32):
386
+ for offset in cutlass.range_constexpr(16, 0, -1):
387
+ # range_constexpr does not provide powers-of-two stepping.
388
+ if cutlass.const_expr(offset in (16, 8, 4, 2, 1)):
389
+ other = cute.arch.shuffle_sync_bfly(
390
+ value,
391
+ offset=offset,
392
+ mask=last_warp_mask,
393
+ mask_and_clamp=31,
394
+ )
395
+ if (lane ^ Int32(offset)) < Int32(last_warp_lanes):
396
+ value = value + other
397
+ else:
398
+ for offset in cutlass.range_constexpr(16, 0, -1):
399
+ if cutlass.const_expr(offset in (16, 8, 4, 2, 1)):
400
+ value = value + cute.arch.shuffle_sync_bfly(
401
+ value,
402
+ offset=offset,
403
+ mask=-1,
404
+ mask_and_clamp=31,
405
+ )
406
+ return value
407
+
408
+
409
+@cute.jit
410
+def block_sum_specialized(
411
+ value: Float32,
412
+ warp_sums: cute.Tensor,
413
+ tidx: Int32,
414
+ warps: cutlass.Constexpr[int],
415
+ last_warp_lanes: cutlass.Constexpr[int],
416
+ last_warp_mask: cutlass.Constexpr[int],
417
+) -> Float32:
418
+ """Upstream-equivalent FP32 block reduction."""
419
+
420
+ lane = cute.arch.lane_idx()
421
+ warp_idx = cute.arch.warp_idx()
422
+ value = warp_sum_specialized(
423
+ value, warp_idx, lane, warps, last_warp_lanes, last_warp_mask
424
+ )
425
+ if lane == 0:
426
+ warp_sums[warp_idx] = value
427
+ cute.arch.barrier()
428
+
429
+ block_sum = Float32(0.0)
430
+ if warp_idx == 0:
431
+ if lane < Int32(warps):
432
+ block_sum = warp_sums[lane]
433
+ block_sum = cute.arch.warp_reduction_sum(block_sum)
434
+ if lane == 0:
435
+ warp_sums[0] = block_sum
436
+ cute.arch.barrier()
437
+ return warp_sums[0]
vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py
new
+244
@@ -0,0 +1,244 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Fused MLA prefill and decode epilogues for Kimi-K3.
4
+
5
+Thin wrappers over the CUDA ops in
6
+``csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu``, which
7
+mirror ``fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_{bf16,fp8}_insert``.
8
+
9
+- ``fused_mla_key_concat_kv_cache_insert`` (bf16): optionally apply RoPE,
10
+ concat the full per-head key ``[k_nope | k_pe]`` into ``k_out``, and insert
11
+ the latent ``[kv_c_normed | k_pe]`` into the paged cache.
12
+- ``fused_mla_qkv_quant_kv_cache_fp8_insert`` (fp8): additionally quantize
13
+ ``q``/``k``/``v`` to E4M3 with ``q_scale`` / ``k_scale`` / ``v_scale`` (the
14
+ cache shares ``k_scale``, as in ``concat_and_cache_mla``).
15
+
16
+The optional ``positions`` / ``cos_sin_cache`` pair enables GPT-J-style RoPE
17
+inside the epilogue. Omitting both keeps the K3 NoPE fast path. The kernels use
18
+Programmatic Dependent Launch to overlap the tail of the producing GEMMs on
19
+sm_90+.
20
+"""
21
+
22
+import torch
23
+
24
+
25
+def fused_mla_key_concat_kv_cache_insert(
26
+ q: torch.Tensor, # [Tp, H, qk_head_dim], RoPE is applied in place
27
+ k_nope: torch.Tensor, # [Tp, H, qk_nope_head_dim]
28
+ k_pe: torch.Tensor, # [Tp, qk_rope_head_dim] or [Tp, 1, qk_rope_head_dim]
29
+ kv_c_normed: torch.Tensor, # [Tp, kv_lora_rank]
30
+ kv_cache: torch.Tensor, # [num_blocks, block_size, kv_lora_rank + rope]
31
+ slot_mapping: torch.Tensor, # [Tp] int64
32
+ positions: torch.Tensor | None = None, # [Tp] int64
33
+ cos_sin_cache: torch.Tensor | None = None, # [max_position, rope]
34
+) -> torch.Tensor:
35
+ """Apply optional RoPE, concat K, and insert the paged latent (bf16).
36
+
37
+ Returns the full key ``[Tp, H, qk_nope_head_dim + qk_rope_head_dim]``;
38
+ optionally rotates ``q`` and writes ``kv_cache`` in place.
39
+ """
40
+ k_pe = k_pe.reshape(k_pe.shape[0], -1)
41
+ tp, num_heads, qk_nope_head_dim = k_nope.shape
42
+ qk_head_dim = qk_nope_head_dim + k_pe.shape[1]
43
+ k_out = torch.empty(
44
+ (tp, num_heads, qk_head_dim), dtype=k_nope.dtype, device=k_nope.device
45
+ )
46
+ if tp == 0:
47
+ return k_out
48
+ torch.ops._C.fused_kimi_k3_mla_key_concat_kv_cache_insert(
49
+ q,
50
+ k_nope,
51
+ k_pe,
52
+ kv_c_normed,
53
+ k_out,
54
+ kv_cache,
55
+ slot_mapping,
56
+ kv_cache.shape[1],
57
+ positions,
58
+ cos_sin_cache,
59
+ )
60
+ return k_out
61
+
62
+
63
+def fused_mla_key_concat_ds_mla_insert(
64
+ q: torch.Tensor, # [Tp, H, qk_head_dim], RoPE is applied in place
65
+ k_nope: torch.Tensor, # [Tp, H, qk_nope_head_dim]
66
+ k_pe: torch.Tensor, # [Tp, qk_rope_head_dim] or [Tp, 1, qk_rope_head_dim]
67
+ kv_c_normed: torch.Tensor, # [Tp, kv_lora_rank]
68
+ kv_cache: torch.Tensor, # [num_blocks, block_size, 656] uint8 (fp8_ds_mla)
69
+ slot_mapping: torch.Tensor, # [Tp] int64
70
+ positions: torch.Tensor | None = None, # [Tp] int64
71
+ cos_sin_cache: torch.Tensor | None = None, # [max_position, rope]
72
+) -> torch.Tensor:
73
+ """Concat full K (bf16) and insert the latent in the fp8_ds_mla layout.
74
+
75
+ The cache uses DeepSeek's 656-byte block-scaled layout (NoPE in 4 tiles of
76
+ 128 with per-tile dynamic fp8 scales, RoPE as bf16) -- self-scaling, so no
77
+ scale argument. Returns the bf16 full key; optionally rotates ``q`` and
78
+ writes ``kv_cache`` in place.
79
+ """
80
+ k_pe = k_pe.reshape(k_pe.shape[0], -1)
81
+ tp, num_heads, qk_nope_head_dim = k_nope.shape
82
+ qk_head_dim = qk_nope_head_dim + k_pe.shape[1]
83
+ k_out = torch.empty(
84
+ (tp, num_heads, qk_head_dim), dtype=k_nope.dtype, device=k_nope.device
85
+ )
86
+ if tp == 0:
87
+ return k_out
88
+ torch.ops._C.fused_kimi_k3_mla_key_concat_ds_mla_insert(
89
+ q,
90
+ k_nope,
91
+ k_pe,
92
+ kv_c_normed,
93
+ k_out,
94
+ kv_cache,
95
+ slot_mapping,
96
+ kv_cache.shape[1],
97
+ positions,
98
+ cos_sin_cache,
99
+ )
100
+ return k_out
101
+
102
+
103
+def fused_mla_qkv_quant_kv_cache_fp8_insert(
104
+ q: torch.Tensor, # [Tp, H, qk_head_dim]
105
+ k_nope: torch.Tensor, # [Tp, H, qk_nope_head_dim]
106
+ k_pe: torch.Tensor, # [Tp, qk_rope_head_dim] or [Tp, 1, qk_rope_head_dim]
107
+ kv_c_normed: torch.Tensor, # [Tp, kv_lora_rank]
108
+ v: torch.Tensor, # [Tp, H, v_head_dim]
109
+ kv_cache: torch.Tensor, # [num_blocks, block_size, kv_lora_rank + rope] fp8
110
+ slot_mapping: torch.Tensor, # [Tp] int64
111
+ q_scale_inv: torch.Tensor, # scalar fp32, 1 / q scale (attention query)
112
+ k_scale_inv: torch.Tensor, # scalar fp32, 1 / k scale (attention key)
113
+ v_scale_inv: torch.Tensor, # scalar fp32, 1 / v scale (attention value)
114
+ cache_scale_inv: torch.Tensor, # scalar fp32, 1 / kv scale (cache latent)
115
+ positions: torch.Tensor | None = None, # [Tp] int64
116
+ cos_sin_cache: torch.Tensor | None = None, # [max_position, rope]
117
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
118
+ """Quantize q/k/v to fp8 and insert the fp8 latent into the paged cache.
119
+
120
+ The attention key ``k_fp8`` and the cache latent use *separate* scales
121
+ (``k_scale_inv`` vs ``cache_scale_inv``): the cache must be quantized with
122
+ ``_k_scale`` (read back by decode / context), while the prefill attention
123
+ q/k/v currently stay unscaled (the prefill flash path does not dequantize).
124
+
125
+ Returns ``(q_fp8, k_fp8, v_fp8)``; writes the fp8 ``kv_cache`` in place.
126
+ """
127
+ k_pe = k_pe.reshape(k_pe.shape[0], -1)
128
+ tp, num_heads, _ = q.shape
129
+ qk_head_dim = q.shape[2]
130
+ v_head_dim = v.shape[2]
131
+ fp8 = torch.float8_e4m3fn
132
+ q_fp8 = torch.empty((tp, num_heads, qk_head_dim), dtype=fp8, device=q.device)
133
+ k_fp8 = torch.empty((tp, num_heads, qk_head_dim), dtype=fp8, device=q.device)
134
+ v_fp8 = torch.empty((tp, num_heads, v_head_dim), dtype=fp8, device=q.device)
135
+ if tp == 0:
136
+ return q_fp8, k_fp8, v_fp8
137
+ torch.ops._C.fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(
138
+ q,
139
+ k_nope,
140
+ k_pe,
141
+ kv_c_normed,
142
+ v,
143
+ q_fp8,
144
+ k_fp8,
145
+ v_fp8,
146
+ kv_cache,
147
+ slot_mapping,
148
+ q_scale_inv,
149
+ k_scale_inv,
150
+ v_scale_inv,
151
+ cache_scale_inv,
152
+ kv_cache.shape[1],
153
+ positions,
154
+ cos_sin_cache,
155
+ )
156
+ return q_fp8, k_fp8, v_fp8
157
+
158
+
159
+def fused_mla_decode_q_concat_kv_cache_insert(
160
+ ql_nope: torch.Tensor, # [B, H, kv_lora_rank] (BMM1 output, absorbed q)
161
+ q_pe: torch.Tensor, # [B, H, qk_rope_head_dim]
162
+ kv_c_normed: torch.Tensor, # [B, kv_lora_rank]
163
+ k_pe: torch.Tensor, # [B, qk_rope_head_dim] or [B, 1, qk_rope_head_dim]
164
+ kv_cache: torch.Tensor, # [num_blocks, block_size, entry]
165
+ slot_mapping: torch.Tensor, # [B] int64
166
+ *,
167
+ ds_mla: bool = False,
168
+ q_scale_inv: torch.Tensor | None = None, # scalar fp32, 1 / q scale
169
+ cache_scale_inv: torch.Tensor | None = None, # scalar fp32, 1 / kv scale
170
+ positions: torch.Tensor | None = None, # [B] int64
171
+ cos_sin_cache: torch.Tensor | None = None, # [max_position, rope]
172
+) -> torch.Tensor:
173
+ """Concat the latent decode query ``mqa_q = [ql_nope | q_pe]`` and insert the
174
+ latent ``[kv_c_normed | k_pe]`` into the paged cache, in one launch (runs
175
+ right before ``forward_mqa``).
176
+
177
+ Dispatched by cache format:
178
+ - bf16 -> bf16 mqa_q, bf16 cache
179
+ - plain fp8 -> fp8 mqa_q (q_scale_inv), fp8 cache (cache_scale_inv)
180
+ - fp8_ds_mla -> bf16 mqa_q, 656B block-scaled cache
181
+
182
+ Returns ``mqa_q`` of shape ``[B, H, kv_lora_rank + qk_rope_head_dim]``;
183
+ writes ``kv_cache`` in place.
184
+ """
185
+ k_pe = k_pe.reshape(k_pe.shape[0], -1)
186
+ b, num_heads, kv_lora_rank = ql_nope.shape
187
+ entry = kv_lora_rank + q_pe.shape[-1]
188
+ fp8_q = q_scale_inv is not None
189
+ out_dtype = torch.float8_e4m3fn if fp8_q else ql_nope.dtype
190
+ mqa_q = torch.empty((b, num_heads, entry), dtype=out_dtype, device=ql_nope.device)
191
+ if b == 0:
192
+ return mqa_q
193
+
194
+ if ds_mla:
195
+ cache = (
196
+ kv_cache if kv_cache.dtype == torch.uint8 else kv_cache.view(torch.uint8)
197
+ )
198
+ torch.ops._C.fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(
199
+ ql_nope,
200
+ q_pe,
201
+ kv_c_normed,
202
+ k_pe,
203
+ mqa_q,
204
+ cache,
205
+ slot_mapping,
206
+ cache.shape[1],
207
+ positions,
208
+ cos_sin_cache,
209
+ )
210
+ elif fp8_q:
211
+ assert cache_scale_inv is not None, "fp8 decode requires cache_scale_inv"
212
+ cache = (
213
+ kv_cache
214
+ if kv_cache.dtype == torch.float8_e4m3fn
215
+ else kv_cache.view(torch.float8_e4m3fn)
216
+ )
217
+ torch.ops._C.fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(
218
+ ql_nope,
219
+ q_pe,
220
+ kv_c_normed,
221
+ k_pe,
222
+ mqa_q,
223
+ cache,
224
+ slot_mapping,
225
+ q_scale_inv,
226
+ cache_scale_inv,
227
+ cache.shape[1],
228
+ positions,
229
+ cos_sin_cache,
230
+ )
231
+ else:
232
+ torch.ops._C.fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(
233
+ ql_nope,
234
+ q_pe,
235
+ kv_c_normed,
236
+ k_pe,
237
+ mqa_q,
238
+ kv_cache,
239
+ slot_mapping,
240
+ kv_cache.shape[1],
241
+ positions,
242
+ cos_sin_cache,
243
+ )
244
+ return mqa_q
vllm/models/kimi_k3/nvidia/ops/latent_moe_tail.py
new
+265
@@ -0,0 +1,265 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from dataclasses import dataclass
5
+from functools import partial
6
+from typing import ClassVar
7
+
8
+import torch
9
+import torch.distributed as dist
10
+
11
+from vllm.distributed import get_tp_group
12
+from vllm.model_executor.warmup.cutedsl_warmup import (
13
+ CuTeDSLCompileUnit,
14
+ register_cutedsl_warmup_provider,
15
+)
16
+
17
+_MAX_NUM_TOKENS = 16
18
+_SKINNY_MAX_NUM_TOKENS = 5
19
+_MMA_TILER_MN = (64, 32)
20
+_GEMM_CLUSTER_MN = (1, 8)
21
+_B_PRIME_STAGES = 2
22
+_COLLECTIVE_TOKEN_CTAS = 8
23
+_LAMPORT_COPY_CTAS = 32
24
+_LAMPORT_COPY_THREADS = 224
25
+_SUPPORTED_TP_SIZES = (8, 16)
26
+
27
+
28
+@dataclass(frozen=True)
29
+class KimiK3LatentMoETailContract:
30
+ tp_group_id: int
31
+ tp_size: int
32
+ device: torch.device
33
+ dtype: torch.dtype
34
+ hidden_size: int
35
+ latent_size: int
36
+ max_num_tokens: int
37
+ rms_eps: float
38
+
39
+
40
+class KimiK3LatentMoETailOp:
41
+ """Process-wide cached K3 latent-MoE tail implementation."""
42
+
43
+ _instances: ClassVar[
44
+ dict[KimiK3LatentMoETailContract, "KimiK3LatentMoETailOp"]
45
+ ] = {}
46
+
47
+ @classmethod
48
+ def _contract_and_group(
49
+ cls,
50
+ *,
51
+ hidden_size: int,
52
+ latent_size: int,
53
+ dtype: torch.dtype,
54
+ device: torch.device,
55
+ rms_eps: float,
56
+ ) -> tuple[KimiK3LatentMoETailContract, dist.ProcessGroup]:
57
+ tp = get_tp_group()
58
+ group = tp.device_group
59
+ device = torch.device(device)
60
+ tp_device = torch.device(tp.device)
61
+ if device != tp_device:
62
+ raise ValueError(
63
+ f"Input device {device} does not match TP device {tp_device}."
64
+ )
65
+ return (
66
+ KimiK3LatentMoETailContract(
67
+ tp_group_id=id(group),
68
+ tp_size=dist.get_world_size(group),
69
+ device=device,
70
+ dtype=dtype,
71
+ hidden_size=hidden_size,
72
+ latent_size=latent_size,
73
+ max_num_tokens=_MAX_NUM_TOKENS,
74
+ rms_eps=float(rms_eps),
75
+ ),
76
+ group,
77
+ )
78
+
79
+ @classmethod
80
+ def initialize(
81
+ cls,
82
+ *,
83
+ hidden_size: int,
84
+ latent_size: int,
85
+ dtype: torch.dtype,
86
+ device: torch.device,
87
+ rms_eps: float,
88
+ ) -> "KimiK3LatentMoETailOp":
89
+ contract, group = cls._contract_and_group(
90
+ hidden_size=hidden_size,
91
+ latent_size=latent_size,
92
+ dtype=dtype,
93
+ device=device,
94
+ rms_eps=rms_eps,
95
+ )
96
+ op = cls._instances.get(contract)
97
+ if op is None:
98
+ op = cls(contract, group)
99
+ cls._instances[contract] = op
100
+ return op
101
+
102
+ def __init__(
103
+ self,
104
+ contract: KimiK3LatentMoETailContract,
105
+ group: dist.ProcessGroup,
106
+ ) -> None:
107
+ if contract.tp_size not in _SUPPORTED_TP_SIZES:
108
+ raise ValueError(
109
+ "K3 latent-MoE tail fusion requires TP in "
110
+ f"{_SUPPORTED_TP_SIZES}, got {contract.tp_size}."
111
+ )
112
+ if contract.device.type != "cuda":
113
+ raise ValueError("K3 latent-MoE tail fusion requires a CUDA device.")
114
+ if torch.cuda.get_device_capability(contract.device)[0] != 10:
115
+ raise ValueError("K3 latent-MoE tail fusion requires SM100.")
116
+ if contract.dtype != torch.bfloat16:
117
+ raise ValueError("K3 latent-MoE tail fusion requires bfloat16.")
118
+ if (contract.hidden_size, contract.latent_size) != (7168, 3584):
119
+ raise ValueError(
120
+ "K3 latent-MoE tail fusion requires hidden_size=7168 and "
121
+ "latent_size=3584."
122
+ )
123
+
124
+ self.contract = contract
125
+ self.rank = dist.get_rank(group)
126
+ from .cute_dsl.latent_moe_tail import (
127
+ AdaptiveUpProjectionKernel,
128
+ CollectiveKernel,
129
+ LamportCopyKernel,
130
+ )
131
+
132
+ with torch.accelerator.device_index(contract.device.index):
133
+ self._collective = CollectiveKernel(
134
+ group=group,
135
+ rank=self.rank,
136
+ tp_size=contract.tp_size,
137
+ latent_dim=contract.latent_size,
138
+ hidden_dim=contract.hidden_size,
139
+ max_m=contract.max_num_tokens,
140
+ max_token_ctas=_COLLECTIVE_TOKEN_CTAS,
141
+ rms_eps=contract.rms_eps,
142
+ fp32_internal=False,
143
+ )
144
+ self._up_projection = AdaptiveUpProjectionKernel(
145
+ group=group,
146
+ rank=self.rank,
147
+ tp_size=contract.tp_size,
148
+ latent_dim=contract.latent_size,
149
+ hidden_dim=contract.hidden_size,
150
+ max_m=contract.max_num_tokens,
151
+ skinny_max_m=_SKINNY_MAX_NUM_TOKENS,
152
+ mma_tiler_mn=_MMA_TILER_MN,
153
+ cluster_shape_mn=_GEMM_CLUSTER_MN,
154
+ b_prime_stages=_B_PRIME_STAGES,
155
+ )
156
+ self._lamport_copy = LamportCopyKernel(
157
+ hidden_dim=contract.hidden_size,
158
+ max_m=contract.max_num_tokens,
159
+ ctas=_LAMPORT_COPY_CTAS,
160
+ threads=_LAMPORT_COPY_THREADS,
161
+ )
162
+ register_cutedsl_warmup_provider(self)
163
+
164
+ def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]:
165
+ contract = self.contract
166
+ skinny_units = tuple(
167
+ CuTeDSLCompileUnit(
168
+ name="K3 latent MoE tail Skinny up-projection",
169
+ key=(
170
+ "k3-latent-moe-tail-skinny-up-projection",
171
+ contract,
172
+ m,
173
+ ),
174
+ compile=partial(self._up_projection.compile_skinny, m),
175
+ )
176
+ for m in range(1, self._up_projection.skinny_max_m + 1)
177
+ )
178
+ return skinny_units + (
179
+ CuTeDSLCompileUnit(
180
+ name="K3 latent MoE tail dynamic up-projection",
181
+ key=(
182
+ "k3-latent-moe-tail-dynamic-up-projection",
183
+ contract,
184
+ _MMA_TILER_MN,
185
+ _GEMM_CLUSTER_MN,
186
+ _B_PRIME_STAGES,
187
+ ),
188
+ compile=self._up_projection.compile_dynamic,
189
+ ),
190
+ )
191
+
192
+ def __call__(
193
+ self,
194
+ routed_output: torch.Tensor,
195
+ shared_output: torch.Tensor,
196
+ rms_weight: torch.Tensor,
197
+ up_weight: torch.Tensor,
198
+ ) -> torch.Tensor:
199
+ self._validate_inputs(
200
+ routed_output,
201
+ shared_output,
202
+ rms_weight,
203
+ up_weight,
204
+ )
205
+ self._up_projection.ensure_compiled(routed_output.shape[0])
206
+ latent, shared_shard = self._collective(
207
+ routed_output,
208
+ shared_output,
209
+ rms_weight,
210
+ )
211
+ local_hidden_size = self.contract.hidden_size // self.contract.tp_size
212
+ local_up_weight = up_weight.narrow(
213
+ 0,
214
+ self.rank * local_hidden_size,
215
+ local_hidden_size,
216
+ )
217
+ mailbox = self._up_projection(
218
+ latent,
219
+ local_up_weight,
220
+ shared_shard,
221
+ )
222
+ return self._lamport_copy(
223
+ mailbox,
224
+ m=routed_output.shape[0],
225
+ ).squeeze(0)
226
+
227
+ def _validate_inputs(
228
+ self,
229
+ routed_output: torch.Tensor,
230
+ shared_output: torch.Tensor,
231
+ rms_weight: torch.Tensor,
232
+ up_weight: torch.Tensor,
233
+ ) -> None:
234
+ contract = self.contract
235
+ if routed_output.ndim != 2:
236
+ raise ValueError("routed_output must be a 2D tensor.")
237
+ num_tokens = routed_output.shape[0]
238
+ if routed_output.shape != (num_tokens, contract.latent_size):
239
+ raise ValueError(
240
+ f"routed_output must have shape [M, {contract.latent_size}]."
241
+ )
242
+ if shared_output.shape != (num_tokens, contract.hidden_size):
243
+ raise ValueError(
244
+ f"shared_output must have shape [M, {contract.hidden_size}]."
245
+ )
246
+ if rms_weight.shape != (contract.latent_size,):
247
+ raise ValueError(f"rms_weight must have shape [{contract.latent_size}].")
248
+ if up_weight.shape != (contract.hidden_size, contract.latent_size):
249
+ raise ValueError(
250
+ "up_weight must have shape "
251
+ f"[{contract.hidden_size}, {contract.latent_size}]."
252
+ )
253
+ if not 1 <= num_tokens <= contract.max_num_tokens:
254
+ raise ValueError(
255
+ "K3 latent-MoE tail fusion requires between 1 and "
256
+ f"{contract.max_num_tokens} tokens."
257
+ )
258
+
259
+ tensors = (routed_output, shared_output, rms_weight, up_weight)
260
+ if any(tensor.device != contract.device for tensor in tensors):
261
+ raise ValueError("All inputs must be on the contract device.")
262
+ if any(tensor.dtype != contract.dtype for tensor in tensors):
263
+ raise ValueError("All inputs must use the contract dtype.")
264
+ if any(not tensor.is_contiguous() for tensor in tensors):
265
+ raise ValueError("All inputs must be contiguous.")
vllm/models/kimi_k3/nvidia/ops/sequence_parallel.py
new
+73
@@ -0,0 +1,73 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm.distributed import (
7
+ get_tensor_model_parallel_rank,
8
+ get_tensor_model_parallel_world_size,
9
+ get_tp_group,
10
+ tensor_model_parallel_all_gather,
11
+ tensor_model_parallel_reduce_scatter,
12
+)
13
+
14
+
15
+def _custom_collective(
16
+ name: str,
17
+ x: torch.Tensor,
18
+) -> torch.Tensor | None:
19
+ device_communicator = get_tp_group().device_communicator
20
+ if device_communicator is None:
21
+ return None
22
+ collective = getattr(device_communicator, name, None)
23
+ if collective is None:
24
+ return None
25
+ return collective(x)
26
+
27
+
28
+def sp_all_gather(x: torch.Tensor) -> torch.Tensor:
29
+ output = _custom_collective("custom_all_gather", x)
30
+ if output is not None:
31
+ return output
32
+ return tensor_model_parallel_all_gather(x, 0)
33
+
34
+
35
+def sp_reduce_scatter(x: torch.Tensor) -> torch.Tensor:
36
+ assert x.ndim == 2
37
+ tp_size = get_tensor_model_parallel_world_size()
38
+ sp_pad = (-x.shape[0]) % tp_size
39
+ if sp_pad > 0:
40
+ x = torch.nn.functional.pad(x, (0, 0, 0, sp_pad))
41
+ output = _custom_collective("custom_reduce_scatter", x)
42
+ if output is not None:
43
+ return output
44
+ return tensor_model_parallel_reduce_scatter(x, 0)
45
+
46
+
47
+def sp_shard(x: torch.Tensor) -> torch.Tensor:
48
+ assert x.ndim == 2
49
+ tp_size = get_tensor_model_parallel_world_size()
50
+ tp_rank = get_tensor_model_parallel_rank()
51
+ sp_pad = (-x.shape[0]) % tp_size
52
+ if sp_pad > 0:
53
+ x = torch.nn.functional.pad(x, (0, 0, 0, sp_pad))
54
+ chunk = x.shape[0] // tp_size
55
+ return x[tp_rank * chunk : (tp_rank + 1) * chunk]
56
+
57
+
58
+def sp_padding_mask(
59
+ is_padding: torch.Tensor | None,
60
+ hidden_states: torch.Tensor,
61
+) -> torch.Tensor:
62
+ num_tokens = hidden_states.shape[0]
63
+ if is_padding is None:
64
+ is_padding = hidden_states.new_zeros(num_tokens, dtype=torch.bool)
65
+ assert is_padding.shape[0] == num_tokens
66
+
67
+ tp_size = get_tensor_model_parallel_world_size()
68
+ sp_pad = (-num_tokens) % tp_size
69
+ if sp_pad > 0:
70
+ is_padding = torch.nn.functional.pad(is_padding, (0, sp_pad), value=True)
71
+ chunk = is_padding.shape[0] // tp_size
72
+ tp_rank = get_tensor_model_parallel_rank()
73
+ return is_padding[tp_rank * chunk : (tp_rank + 1) * chunk]
vllm/models/kimi_k3/nvidia/ops/third_party/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/kimi_k3/nvidia/ops/third_party/kda/__init__.py
new
+28
@@ -0,0 +1,28 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from .chunk import (
5
+ chunk_kda,
6
+ chunk_kda_fwd,
7
+ chunk_kda_with_fused_gate,
8
+ chunk_kda_with_fused_gate_fwd,
9
+ fused_kda_gate,
10
+ fused_kda_gate_chunk_cumsum,
11
+)
12
+from .fused_recurrent import (
13
+ fused_recurrent_kda,
14
+ fused_recurrent_kda_fwd,
15
+ fused_recurrent_kda_packed_decode,
16
+)
17
+
18
+__all__ = [
19
+ "chunk_kda",
20
+ "chunk_kda_fwd",
21
+ "chunk_kda_with_fused_gate",
22
+ "chunk_kda_with_fused_gate_fwd",
23
+ "fused_kda_gate",
24
+ "fused_kda_gate_chunk_cumsum",
25
+ "fused_recurrent_kda",
26
+ "fused_recurrent_kda_fwd",
27
+ "fused_recurrent_kda_packed_decode",
28
+]
vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py
new
+938
@@ -0,0 +1,938 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li
4
+#
5
+# This file contains code copied from the flash-linear-attention project.
6
+# The original source code was licensed under the MIT license and included
7
+# the following copyright notice:
8
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
9
+# ruff: noqa: E501
10
+
11
+
12
+import torch
13
+
14
+from vllm.third_party.flash_linear_attention.ops.chunk_delta_h import (
15
+ chunk_gated_delta_rule_fwd_h,
16
+)
17
+from vllm.third_party.flash_linear_attention.ops.cumsum import chunk_local_cumsum
18
+from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices
19
+from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
20
+from vllm.third_party.flash_linear_attention.ops.op import exp2, log
21
+from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE, is_amd
22
+from vllm.triton_utils import tl, triton
23
+from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2
24
+
25
+from .chunk_intra import chunk_kda_fwd_intra
26
+
27
+BT_LIST_AUTOTUNE = [32, 64, 128]
28
+NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32]
29
+
30
+
31
+@triton.heuristics(
32
+ {
33
+ "STORE_QG": lambda args: args["qg"] is not None,
34
+ "STORE_KG": lambda args: args["kg"] is not None,
35
+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
36
+ }
37
+)
38
+@triton.autotune(
39
+ configs=[
40
+ triton.Config({}, num_warps=num_warps, num_stages=num_stages)
41
+ for num_warps in [2, 4, 8]
42
+ for num_stages in [2, 3, 4]
43
+ ],
44
+ key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"],
45
+)
46
+@triton.jit(do_not_specialize=["T"])
47
+def recompute_w_u_fwd_kernel(
48
+ q,
49
+ k,
50
+ qg,
51
+ kg,
52
+ v,
53
+ beta,
54
+ w,
55
+ u,
56
+ A,
57
+ gk,
58
+ cu_seqlens,
59
+ chunk_indices,
60
+ T,
61
+ H: tl.constexpr,
62
+ K: tl.constexpr,
63
+ V: tl.constexpr,
64
+ BT: tl.constexpr,
65
+ BK: tl.constexpr,
66
+ BV: tl.constexpr,
67
+ STORE_QG: tl.constexpr,
68
+ STORE_KG: tl.constexpr,
69
+ IS_VARLEN: tl.constexpr,
70
+ DOT_PRECISION: tl.constexpr,
71
+):
72
+ i_t, i_bh = tl.program_id(0), tl.program_id(1)
73
+ i_b, i_h = i_bh // H, i_bh % H
74
+ if IS_VARLEN:
75
+ i_n, i_t = (
76
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
77
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
78
+ )
79
+ bos, eos = (
80
+ tl.load(cu_seqlens + i_n).to(tl.int32),
81
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
82
+ )
83
+ T = eos - bos
84
+ else:
85
+ bos, eos = i_b * T, i_b * T + T
86
+ p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
87
+ b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32)
88
+
89
+ p_A = tl.make_block_ptr(
90
+ A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)
91
+ )
92
+ b_A = tl.load(p_A, boundary_check=(0, 1))
93
+
94
+ for i_v in range(tl.cdiv(V, BV)):
95
+ p_v = tl.make_block_ptr(
96
+ v + (bos * H + i_h) * V,
97
+ (T, V),
98
+ (H * V, 1),
99
+ (i_t * BT, i_v * BV),
100
+ (BT, BV),
101
+ (1, 0),
102
+ )
103
+ p_u = tl.make_block_ptr(
104
+ u + (bos * H + i_h) * V,
105
+ (T, V),
106
+ (H * V, 1),
107
+ (i_t * BT, i_v * BV),
108
+ (BT, BV),
109
+ (1, 0),
110
+ )
111
+ b_v = tl.load(p_v, boundary_check=(0, 1))
112
+ b_vb = (b_v * b_b[:, None]).to(b_v.dtype)
113
+ b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION)
114
+ tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1))
115
+
116
+ for i_k in range(tl.cdiv(K, BK)):
117
+ p_w = tl.make_block_ptr(
118
+ w + (bos * H + i_h) * K,
119
+ (T, K),
120
+ (H * K, 1),
121
+ (i_t * BT, i_k * BK),
122
+ (BT, BK),
123
+ (1, 0),
124
+ )
125
+ p_k = tl.make_block_ptr(
126
+ k + (bos * H + i_h) * K,
127
+ (T, K),
128
+ (H * K, 1),
129
+ (i_t * BT, i_k * BK),
130
+ (BT, BK),
131
+ (1, 0),
132
+ )
133
+ b_k = tl.load(p_k, boundary_check=(0, 1))
134
+ b_kb = b_k * b_b[:, None]
135
+
136
+ p_gk = tl.make_block_ptr(
137
+ gk + (bos * H + i_h) * K,
138
+ (T, K),
139
+ (H * K, 1),
140
+ (i_t * BT, i_k * BK),
141
+ (BT, BK),
142
+ (1, 0),
143
+ )
144
+ b_gk = tl.load(p_gk, boundary_check=(0, 1))
145
+ b_kb *= exp2(b_gk)
146
+ if STORE_QG:
147
+ p_q = tl.make_block_ptr(
148
+ q + (bos * H + i_h) * K,
149
+ (T, K),
150
+ (H * K, 1),
151
+ (i_t * BT, i_k * BK),
152
+ (BT, BK),
153
+ (1, 0),
154
+ )
155
+ p_qg = tl.make_block_ptr(
156
+ qg + (bos * H + i_h) * K,
157
+ (T, K),
158
+ (H * K, 1),
159
+ (i_t * BT, i_k * BK),
160
+ (BT, BK),
161
+ (1, 0),
162
+ )
163
+ b_q = tl.load(p_q, boundary_check=(0, 1))
164
+ b_qg = b_q * exp2(b_gk)
165
+ tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1))
166
+ if STORE_KG:
167
+ last_idx = min(i_t * BT + BT, T) - 1
168
+
169
+ o_k = i_k * BK + tl.arange(0, BK)
170
+ m_k = o_k < K
171
+ b_gn = tl.load(
172
+ gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0
173
+ )
174
+ b_kg = b_k * exp2(b_gn - b_gk)
175
+
176
+ p_kg = tl.make_block_ptr(
177
+ kg + (bos * H + i_h) * K,
178
+ (T, K),
179
+ (H * K, 1),
180
+ (i_t * BT, i_k * BK),
181
+ (BT, BK),
182
+ (1, 0),
183
+ )
184
+ tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1))
185
+
186
+ b_w = tl.dot(b_A, b_kb.to(b_k.dtype))
187
+ tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))
188
+
189
+
190
+def recompute_w_u_fwd(
191
+ k: torch.Tensor,
192
+ v: torch.Tensor,
193
+ beta: torch.Tensor,
194
+ A: torch.Tensor,
195
+ q: torch.Tensor | None = None,
196
+ gk: torch.Tensor | None = None,
197
+ cu_seqlens: torch.Tensor | None = None,
198
+ chunk_indices: torch.Tensor | None = None,
199
+) -> tuple[torch.Tensor, torch.Tensor]:
200
+ B, T, H, K, V = *k.shape, v.shape[-1]
201
+ BT = A.shape[-1]
202
+ BK = 64
203
+ BV = 64
204
+
205
+ if chunk_indices is None and cu_seqlens is not None:
206
+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
207
+ NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
208
+
209
+ w = torch.empty_like(k)
210
+ u = torch.empty_like(v)
211
+ kg = torch.empty_like(k) if gk is not None else None
212
+ recompute_w_u_fwd_kernel[(NT, B * H)](
213
+ q=q,
214
+ k=k,
215
+ qg=None,
216
+ kg=kg,
217
+ v=v,
218
+ beta=beta,
219
+ w=w,
220
+ u=u,
221
+ A=A,
222
+ gk=gk,
223
+ cu_seqlens=cu_seqlens,
224
+ chunk_indices=chunk_indices,
225
+ T=T,
226
+ H=H,
227
+ K=K,
228
+ V=V,
229
+ BT=BT,
230
+ BK=BK,
231
+ BV=BV,
232
+ DOT_PRECISION="ieee",
233
+ )
234
+ return w, u, None, kg
235
+
236
+
237
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
238
+@triton.autotune(
239
+ configs=[
240
+ triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages)
241
+ for BK in [32, 64]
242
+ for BV in [64, 128]
243
+ for num_warps in [2, 4, 8]
244
+ for num_stages in [2, 3, 4]
245
+ ],
246
+ key=["BT"],
247
+)
248
+@triton.jit(do_not_specialize=["T"])
249
+def chunk_gla_fwd_kernel_o(
250
+ q,
251
+ v,
252
+ g,
253
+ h,
254
+ o,
255
+ A,
256
+ cu_seqlens,
257
+ chunk_indices,
258
+ scale,
259
+ T,
260
+ H: tl.constexpr,
261
+ K: tl.constexpr,
262
+ V: tl.constexpr,
263
+ BT: tl.constexpr,
264
+ BK: tl.constexpr,
265
+ BV: tl.constexpr,
266
+ IS_VARLEN: tl.constexpr,
267
+):
268
+ i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
269
+ i_b, i_h = i_bh // H, i_bh % H
270
+ if IS_VARLEN:
271
+ i_tg = i_t
272
+ i_n, i_t = (
273
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
274
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
275
+ )
276
+ bos, eos = (
277
+ tl.load(cu_seqlens + i_n).to(tl.int32),
278
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
279
+ )
280
+ T = eos - bos
281
+ NT = tl.cdiv(T, BT)
282
+ else:
283
+ NT = tl.cdiv(T, BT)
284
+ i_tg = i_b * NT + i_t
285
+ bos, eos = i_b * T, i_b * T + T
286
+
287
+ m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :]
288
+
289
+ b_o = tl.zeros([BT, BV], dtype=tl.float32)
290
+ for i_k in range(tl.cdiv(K, BK)):
291
+ p_q = tl.make_block_ptr(
292
+ q + (bos * H + i_h) * K,
293
+ (T, K),
294
+ (H * K, 1),
295
+ (i_t * BT, i_k * BK),
296
+ (BT, BK),
297
+ (1, 0),
298
+ )
299
+ p_g = tl.make_block_ptr(
300
+ g + (bos * H + i_h) * K,
301
+ (T, K),
302
+ (H * K, 1),
303
+ (i_t * BT, i_k * BK),
304
+ (BT, BK),
305
+ (1, 0),
306
+ )
307
+ p_h = tl.make_block_ptr(
308
+ h + (i_tg * H + i_h) * K * V,
309
+ (V, K),
310
+ (K, 1),
311
+ (i_v * BV, i_k * BK),
312
+ (BV, BK),
313
+ (1, 0),
314
+ )
315
+
316
+ # [BT, BK]
317
+ b_q = tl.load(p_q, boundary_check=(0, 1))
318
+ b_q = (b_q * scale).to(b_q.dtype)
319
+ # [BT, BK]
320
+ b_g = tl.load(p_g, boundary_check=(0, 1))
321
+ # [BT, BK]
322
+ b_qg = (b_q * exp2(b_g)).to(b_q.dtype)
323
+ # [BV, BK]
324
+ b_h = tl.load(p_h, boundary_check=(0, 1))
325
+ # [BT, BV]
326
+ if i_k >= 0:
327
+ b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype))
328
+ p_v = tl.make_block_ptr(
329
+ v + (bos * H + i_h) * V,
330
+ (T, V),
331
+ (H * V, 1),
332
+ (i_t * BT, i_v * BV),
333
+ (BT, BV),
334
+ (1, 0),
335
+ )
336
+ p_o = tl.make_block_ptr(
337
+ o + (bos * H + i_h) * V,
338
+ (T, V),
339
+ (H * V, 1),
340
+ (i_t * BT, i_v * BV),
341
+ (BT, BV),
342
+ (1, 0),
343
+ )
344
+ p_A = tl.make_block_ptr(
345
+ A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)
346
+ )
347
+ # [BT, BV]
348
+ b_v = tl.load(p_v, boundary_check=(0, 1))
349
+ # [BT, BT]
350
+ b_A = tl.load(p_A, boundary_check=(0, 1))
351
+ b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype)
352
+ b_o += tl.dot(b_A, b_v, allow_tf32=False)
353
+ tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
354
+
355
+
356
+def chunk_gla_fwd_o_gk(
357
+ q: torch.Tensor,
358
+ v: torch.Tensor,
359
+ g: torch.Tensor,
360
+ A: torch.Tensor,
361
+ h: torch.Tensor,
362
+ o: torch.Tensor,
363
+ scale: float,
364
+ cu_seqlens: torch.Tensor | None = None,
365
+ chunk_indices: torch.Tensor | None = None,
366
+ chunk_size: int = FLA_CHUNK_SIZE,
367
+):
368
+ B, T, H, K, V = *q.shape, v.shape[-1]
369
+ BT = chunk_size
370
+
371
+ if chunk_indices is None and cu_seqlens is not None:
372
+ chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
373
+ NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
374
+
375
+ def grid(meta):
376
+ return (cdiv(V, meta["BV"]), NT, B * H)
377
+
378
+ chunk_gla_fwd_kernel_o[grid](
379
+ q=q,
380
+ v=v,
381
+ g=g,
382
+ h=h,
383
+ o=o,
384
+ A=A,
385
+ cu_seqlens=cu_seqlens,
386
+ chunk_indices=chunk_indices,
387
+ scale=scale,
388
+ T=T,
389
+ H=H,
390
+ K=K,
391
+ V=V,
392
+ BT=BT,
393
+ )
394
+ return o
395
+
396
+
397
+@triton.heuristics(
398
+ {
399
+ "HAS_BIAS": lambda args: args["g_bias"] is not None,
400
+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
401
+ }
402
+)
403
+@triton.autotune(
404
+ configs=[
405
+ triton.Config({"BS": BS}, num_warps=num_warps)
406
+ for BS in [32, 64]
407
+ for num_warps in [2, 4, 8]
408
+ ],
409
+ key=["H", "S", "BT", "IS_VARLEN"],
410
+)
411
+@triton.jit(do_not_specialize=["T"])
412
+def kda_gate_chunk_cumsum_vector_kernel(
413
+ s,
414
+ raw_beta,
415
+ A_log,
416
+ g_bias,
417
+ o,
418
+ beta_out,
419
+ cu_seqlens,
420
+ chunk_indices,
421
+ cumsum_scale,
422
+ lower_bound,
423
+ beta,
424
+ threshold,
425
+ T,
426
+ stride_beta_batch,
427
+ stride_beta_token,
428
+ stride_beta_head,
429
+ H: tl.constexpr,
430
+ S: tl.constexpr,
431
+ BT: tl.constexpr,
432
+ BS: tl.constexpr,
433
+ HAS_BIAS: tl.constexpr,
434
+ IS_VARLEN: tl.constexpr,
435
+ USE_LOWER_BOUND: tl.constexpr,
436
+):
437
+ i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
438
+ i_b, i_h = i_bh // H, i_bh % H
439
+ if IS_VARLEN:
440
+ i_n, i_t = (
441
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
442
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
443
+ )
444
+ bos, eos = (
445
+ tl.load(cu_seqlens + i_n).to(tl.int32),
446
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
447
+ )
448
+ T = eos - bos
449
+ else:
450
+ bos = i_b * T
451
+
452
+ if i_s == 0:
453
+ o_beta_t = tl.arange(0, BT)
454
+ m_beta = i_t * BT + o_beta_t < T
455
+ if IS_VARLEN:
456
+ p_beta = (
457
+ raw_beta
458
+ + (bos + i_t * BT + o_beta_t) * stride_beta_token
459
+ + i_h * stride_beta_head
460
+ )
461
+ else:
462
+ p_beta = (
463
+ raw_beta
464
+ + i_b * stride_beta_batch
465
+ + (i_t * BT + o_beta_t) * stride_beta_token
466
+ + i_h * stride_beta_head
467
+ )
468
+ b_beta = tl.load(p_beta, mask=m_beta, other=0.0).to(tl.float32)
469
+ p_beta_out = beta_out + (bos + i_t * BT + o_beta_t) * H + i_h
470
+ tl.store(p_beta_out, tl.sigmoid(b_beta), mask=m_beta)
471
+ return
472
+
473
+ i_s -= 1
474
+
475
+ p_s = tl.make_block_ptr(
476
+ s + (bos * H + i_h) * S,
477
+ (T, S),
478
+ (H * S, 1),
479
+ (i_t * BT, i_s * BS),
480
+ (BT, BS),
481
+ (1, 0),
482
+ )
483
+ p_o = tl.make_block_ptr(
484
+ o + (bos * H + i_h) * S,
485
+ (T, S),
486
+ (H * S, 1),
487
+ (i_t * BT, i_s * BS),
488
+ (BT, BS),
489
+ (1, 0),
490
+ )
491
+
492
+ b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32)
493
+ if HAS_BIAS:
494
+ p_bias = tl.make_block_ptr(
495
+ g_bias + i_h * S,
496
+ (S,),
497
+ (1,),
498
+ (i_s * BS,),
499
+ (BS,),
500
+ (0,),
501
+ )
502
+ b_bias = tl.load(p_bias, boundary_check=(0,)).to(tl.float32)
503
+ b_s += b_bias[None, :]
504
+
505
+ b_a = tl.exp(tl.load(A_log + i_h).to(tl.float32))
506
+ if USE_LOWER_BOUND:
507
+ b_gate = lower_bound * tl.sigmoid(b_a * b_s)
508
+ else:
509
+ b_g_scaled = b_s * beta
510
+ b_softplus = tl.where(
511
+ b_g_scaled > threshold,
512
+ b_s,
513
+ (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)),
514
+ )
515
+ b_gate = -b_a * b_softplus
516
+
517
+ # Boundary loads return zero, but bias and gate activation can make padded
518
+ # rows nonzero. Padding trails valid rows, so it only affects masked stores.
519
+ b_o = tl.cumsum(b_gate, axis=0) * cumsum_scale
520
+ tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
521
+
522
+
523
+def fused_kda_gate_chunk_cumsum(
524
+ raw_g: torch.Tensor,
525
+ raw_beta: torch.Tensor,
526
+ A_log: torch.Tensor,
527
+ g_bias: torch.Tensor | None = None,
528
+ beta: float = 1.0,
529
+ threshold: float = 20.0,
530
+ lower_bound: float | None = None,
531
+ cu_seqlens: torch.Tensor | None = None,
532
+ chunk_indices: torch.Tensor | None = None,
533
+ chunk_size: int = FLA_CHUNK_SIZE,
534
+ output_dtype: torch.dtype | None = torch.float,
535
+) -> tuple[torch.Tensor, torch.Tensor]:
536
+ if cu_seqlens is not None:
537
+ assert raw_g.shape[0] == 1, (
538
+ "Only batch size 1 is supported when cu_seqlens are provided"
539
+ )
540
+ B, T, H, D = raw_g.shape
541
+ if raw_beta.shape != (B, T, H):
542
+ raise ValueError(
543
+ f"Expected raw_beta shape {(B, T, H)}, got {raw_beta.shape}"
544
+ )
545
+ if chunk_indices is None and cu_seqlens is not None:
546
+ chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
547
+ NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices)
548
+
549
+ A_log = A_log.reshape(-1)
550
+ if g_bias is not None:
551
+ g_bias = g_bias.reshape(-1)
552
+ y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype)
553
+ beta_out = torch.empty(raw_beta.shape, device=raw_beta.device, dtype=torch.float32)
554
+
555
+ def grid(meta):
556
+ # For each (chunk, head), program 0 computes beta without extending a
557
+ # gate tile's critical path. The remaining programs cover the gate dim.
558
+ return (cdiv(meta["S"], meta["BS"]) + 1, NT, B * H)
559
+
560
+ kda_gate_chunk_cumsum_vector_kernel[grid](
561
+ s=raw_g,
562
+ raw_beta=raw_beta,
563
+ A_log=A_log,
564
+ g_bias=g_bias,
565
+ o=y,
566
+ beta_out=beta_out,
567
+ cu_seqlens=cu_seqlens,
568
+ chunk_indices=chunk_indices,
569
+ # RCP_LN2 folds in the natural-log -> log2 conversion so downstream
570
+ # exp2-based kernels reproduce exp(g). Keep this in sync with the
571
+ # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`.
572
+ cumsum_scale=RCP_LN2,
573
+ lower_bound=lower_bound or 0.0,
574
+ beta=beta,
575
+ threshold=threshold,
576
+ T=T,
577
+ stride_beta_batch=raw_beta.stride(0),
578
+ stride_beta_token=raw_beta.stride(1),
579
+ stride_beta_head=raw_beta.stride(2),
580
+ H=H,
581
+ S=D,
582
+ BT=chunk_size,
583
+ USE_LOWER_BOUND=lower_bound is not None,
584
+ )
585
+ return y, beta_out
586
+
587
+
588
+def _chunk_kda_fwd_with_cumulative_g(
589
+ q: torch.Tensor,
590
+ k: torch.Tensor,
591
+ v: torch.Tensor,
592
+ g: torch.Tensor,
593
+ beta: torch.Tensor,
594
+ scale: float,
595
+ initial_state: torch.Tensor,
596
+ output_final_state: bool,
597
+ cu_seqlens: torch.Tensor | None = None,
598
+ chunk_indices: torch.Tensor | None = None,
599
+ chunk_size: int = FLA_CHUNK_SIZE,
600
+ safe_gate: bool = False,
601
+):
602
+ # `g` must already be chunk-local cumulatively-summed AND scaled by
603
+ # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)).
604
+ # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of
605
+ # calling this helper directly unless that invariant is upheld.
606
+ Aqk, A = chunk_kda_fwd_intra(
607
+ q=q,
608
+ k=k,
609
+ gk=g,
610
+ beta=beta,
611
+ scale=scale,
612
+ cu_seqlens=cu_seqlens,
613
+ chunk_indices=chunk_indices,
614
+ chunk_size=chunk_size,
615
+ safe_gate=safe_gate,
616
+ )
617
+ w, u, _, kg = recompute_w_u_fwd(
618
+ k=k,
619
+ v=v,
620
+ beta=beta,
621
+ A=A,
622
+ gk=g,
623
+ cu_seqlens=cu_seqlens,
624
+ chunk_indices=chunk_indices,
625
+ )
626
+ del A
627
+ h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
628
+ k=kg,
629
+ w=w,
630
+ u=u,
631
+ gk=g,
632
+ initial_state=initial_state,
633
+ output_final_state=output_final_state,
634
+ cu_seqlens=cu_seqlens,
635
+ chunk_indices=chunk_indices,
636
+ use_exp2=True,
637
+ )
638
+ del w, u, kg
639
+ o = chunk_gla_fwd_o_gk(
640
+ q=q,
641
+ v=v_new,
642
+ g=g,
643
+ A=Aqk,
644
+ h=h,
645
+ o=v,
646
+ scale=scale,
647
+ cu_seqlens=cu_seqlens,
648
+ chunk_indices=chunk_indices,
649
+ chunk_size=chunk_size,
650
+ )
651
+ del Aqk, v_new, h
652
+ return o, final_state
653
+
654
+
655
+def chunk_kda_fwd(
656
+ q: torch.Tensor,
657
+ k: torch.Tensor,
658
+ v: torch.Tensor,
659
+ g: torch.Tensor,
660
+ beta: torch.Tensor,
661
+ scale: float,
662
+ initial_state: torch.Tensor,
663
+ output_final_state: bool,
664
+ cu_seqlens: torch.Tensor | None = None,
665
+):
666
+ chunk_size = FLA_CHUNK_SIZE
667
+ chunk_indices = (
668
+ prepare_chunk_indices(cu_seqlens, chunk_size)
669
+ if cu_seqlens is not None
670
+ else None
671
+ )
672
+ g = chunk_local_cumsum(
673
+ g,
674
+ chunk_size=chunk_size,
675
+ cu_seqlens=cu_seqlens,
676
+ chunk_indices=chunk_indices,
677
+ )
678
+ # KDA evaluates cumulative gate decays with exp2. Convert from natural-log
679
+ # space so exp(x) is preserved as exp2(x / ln(2)).
680
+ g = g * RCP_LN2
681
+ return _chunk_kda_fwd_with_cumulative_g(
682
+ q=q,
683
+ k=k,
684
+ v=v,
685
+ g=g,
686
+ beta=beta,
687
+ scale=scale,
688
+ initial_state=initial_state,
689
+ output_final_state=output_final_state,
690
+ cu_seqlens=cu_seqlens,
691
+ chunk_indices=chunk_indices,
692
+ chunk_size=chunk_size,
693
+ )
694
+
695
+
696
+def chunk_kda_with_fused_gate_fwd(
697
+ q: torch.Tensor,
698
+ k: torch.Tensor,
699
+ v: torch.Tensor,
700
+ raw_g: torch.Tensor,
701
+ raw_beta: torch.Tensor,
702
+ A_log: torch.Tensor,
703
+ g_bias: torch.Tensor | None,
704
+ scale: float,
705
+ initial_state: torch.Tensor,
706
+ output_final_state: bool,
707
+ lower_bound: float | None = None,
708
+ cu_seqlens: torch.Tensor | None = None,
709
+):
710
+ chunk_size = FLA_CHUNK_SIZE
711
+ chunk_indices = (
712
+ prepare_chunk_indices(cu_seqlens, chunk_size)
713
+ if cu_seqlens is not None
714
+ else None
715
+ )
716
+ g, beta = fused_kda_gate_chunk_cumsum(
717
+ raw_g,
718
+ raw_beta=raw_beta,
719
+ A_log=A_log,
720
+ g_bias=g_bias,
721
+ cu_seqlens=cu_seqlens,
722
+ chunk_indices=chunk_indices,
723
+ chunk_size=chunk_size,
724
+ lower_bound=lower_bound,
725
+ )
726
+ return _chunk_kda_fwd_with_cumulative_g(
727
+ q=q,
728
+ k=k,
729
+ v=v,
730
+ g=g,
731
+ beta=beta,
732
+ scale=scale,
733
+ initial_state=initial_state,
734
+ output_final_state=output_final_state,
735
+ cu_seqlens=cu_seqlens,
736
+ chunk_indices=chunk_indices,
737
+ chunk_size=chunk_size,
738
+ safe_gate=lower_bound is not None,
739
+ )
740
+
741
+
742
+def chunk_kda(
743
+ q: torch.Tensor,
744
+ k: torch.Tensor,
745
+ v: torch.Tensor,
746
+ g: torch.Tensor,
747
+ beta: torch.Tensor,
748
+ scale: float = None,
749
+ initial_state: torch.Tensor = None,
750
+ output_final_state: bool = False,
751
+ use_qk_l2norm_in_kernel: bool = False,
752
+ cu_seqlens: torch.Tensor | None = None,
753
+ **kwargs,
754
+):
755
+ if scale is None:
756
+ scale = k.shape[-1] ** -0.5
757
+
758
+ if use_qk_l2norm_in_kernel:
759
+ q = l2norm_fwd(q.contiguous())
760
+ k = l2norm_fwd(k.contiguous())
761
+
762
+ o, final_state = chunk_kda_fwd(
763
+ q=q,
764
+ k=k,
765
+ v=v.contiguous(),
766
+ g=g.contiguous(),
767
+ beta=beta.contiguous(),
768
+ scale=scale,
769
+ initial_state=initial_state.contiguous(),
770
+ output_final_state=output_final_state,
771
+ cu_seqlens=cu_seqlens,
772
+ )
773
+ return o, final_state
774
+
775
+
776
+def chunk_kda_with_fused_gate(
777
+ q: torch.Tensor,
778
+ k: torch.Tensor,
779
+ v: torch.Tensor,
780
+ raw_g: torch.Tensor,
781
+ raw_beta: torch.Tensor,
782
+ A_log: torch.Tensor,
783
+ g_bias: torch.Tensor | None,
784
+ scale: float | None = None,
785
+ initial_state: torch.Tensor | None = None,
786
+ output_final_state: bool = False,
787
+ lower_bound: float | None = None,
788
+ use_qk_l2norm_in_kernel: bool = False,
789
+ cu_seqlens: torch.Tensor | None = None,
790
+ **kwargs,
791
+):
792
+ """Run chunk KDA from raw gate and beta projections."""
793
+ if scale is None:
794
+ scale = k.shape[-1] ** -0.5
795
+
796
+ if use_qk_l2norm_in_kernel:
797
+ q = l2norm_fwd(q.contiguous())
798
+ k = l2norm_fwd(k.contiguous())
799
+
800
+ o, final_state = chunk_kda_with_fused_gate_fwd(
801
+ q=q,
802
+ k=k,
803
+ v=v.contiguous(),
804
+ raw_g=raw_g.contiguous(),
805
+ raw_beta=raw_beta,
806
+ A_log=A_log,
807
+ g_bias=g_bias,
808
+ scale=scale,
809
+ initial_state=initial_state.contiguous() if initial_state is not None else None,
810
+ output_final_state=output_final_state,
811
+ lower_bound=lower_bound,
812
+ cu_seqlens=cu_seqlens,
813
+ )
814
+ return o, final_state
815
+
816
+
817
+@triton.autotune(
818
+ configs=[
819
+ triton.Config({"BT": bt}, num_warps=nw, num_stages=ns)
820
+ for bt in BT_LIST_AUTOTUNE
821
+ for nw in NUM_WARPS_AUTOTUNE
822
+ for ns in [2, 3]
823
+ ],
824
+ key=["H", "D"],
825
+)
826
+@triton.jit
827
+def kda_gate_fwd_kernel(
828
+ g,
829
+ A,
830
+ y,
831
+ g_bias,
832
+ lower_bound,
833
+ beta: tl.constexpr,
834
+ threshold: tl.constexpr,
835
+ T,
836
+ H,
837
+ D: tl.constexpr,
838
+ BT: tl.constexpr,
839
+ BD: tl.constexpr,
840
+ HAS_BIAS: tl.constexpr,
841
+ USE_LOWER_BOUND: tl.constexpr,
842
+):
843
+ i_t, i_h = tl.program_id(0), tl.program_id(1)
844
+ n_t = i_t * BT
845
+
846
+ b_a = tl.exp(tl.load(A + i_h).to(tl.float32))
847
+
848
+ stride_row = H * D
849
+ stride_col = 1
850
+
851
+ g_ptr = tl.make_block_ptr(
852
+ base=g + i_h * D,
853
+ shape=(T, D),
854
+ strides=(stride_row, stride_col),
855
+ offsets=(n_t, 0),
856
+ block_shape=(BT, BD),
857
+ order=(1, 0),
858
+ )
859
+
860
+ y_ptr = tl.make_block_ptr(
861
+ base=y + i_h * D,
862
+ shape=(T, D),
863
+ strides=(stride_row, stride_col),
864
+ offsets=(n_t, 0),
865
+ block_shape=(BT, BD),
866
+ order=(1, 0),
867
+ )
868
+
869
+ b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32)
870
+
871
+ if HAS_BIAS:
872
+ n_d = tl.arange(0, BD)
873
+ bias_mask = n_d < D
874
+ b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to(
875
+ tl.float32
876
+ )
877
+ b_g = b_g + b_bias[None, :]
878
+
879
+ if USE_LOWER_BOUND:
880
+ b_y = lower_bound * tl.sigmoid(b_a * b_g)
881
+ else:
882
+ g_scaled = b_g * beta
883
+ use_linear = g_scaled > threshold
884
+ sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled)))
885
+ b_y = -b_a * sp
886
+
887
+ tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1))
888
+
889
+
890
+def fused_kda_gate(
891
+ g: torch.Tensor,
892
+ A: torch.Tensor,
893
+ head_k_dim: int,
894
+ g_bias: torch.Tensor | None = None,
895
+ beta: float = 1.0,
896
+ threshold: float = 20.0,
897
+ lower_bound: float | None = None,
898
+) -> torch.Tensor:
899
+ """
900
+ Forward pass for KDA gate:
901
+ input g: [..., H*D]
902
+ param A: [H] or [1, 1, H, 1]
903
+ beta: softplus beta parameter
904
+ threshold: softplus threshold parameter
905
+ return : [..., H, D]
906
+ """
907
+ orig_shape = g.shape[:-1]
908
+
909
+ g = g.view(-1, g.shape[-1])
910
+ T = g.shape[0]
911
+ HD = g.shape[1]
912
+ H = A.numel()
913
+ assert H * head_k_dim == HD
914
+ assert g.stride() == (HD, 1)
915
+
916
+ y = torch.empty_like(g, dtype=torch.float32)
917
+
918
+ def grid(meta):
919
+ return (cdiv(T, meta["BT"]), H)
920
+
921
+ kda_gate_fwd_kernel[grid](
922
+ g,
923
+ A,
924
+ y,
925
+ g_bias,
926
+ lower_bound or 0.0,
927
+ beta,
928
+ threshold,
929
+ T,
930
+ H,
931
+ head_k_dim,
932
+ BD=next_power_of_2(head_k_dim),
933
+ HAS_BIAS=g_bias is not None,
934
+ USE_LOWER_BOUND=lower_bound is not None,
935
+ )
936
+
937
+ y = y.view(*orig_shape, H, head_k_dim)
938
+ return y
vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py
new
+559
@@ -0,0 +1,559 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li
4
+#
5
+# This file contains code copied from the flash-linear-attention project.
6
+# The original source was licensed under the MIT license.
7
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
8
+# Forward-only adaptation of flash-linear-attention 0.5.0.
9
+# ruff: noqa: E501
10
+
11
+import torch
12
+
13
+from vllm.platforms import current_platform
14
+from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices
15
+from vllm.third_party.flash_linear_attention.ops.op import exp2, gather
16
+from vllm.third_party.flash_linear_attention.ops.utils import is_gather_supported
17
+from vllm.triton_utils import tl, triton
18
+
19
+from .chunk_intra_token_parallel import chunk_kda_fwd_intra_token_parallel
20
+
21
+################################################################################
22
+# Fused inter + solve_tril kernel: compute off-diagonal Akk and solve in one pass
23
+################################################################################
24
+
25
+
26
+@triton.heuristics({
27
+ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
28
+})
29
+@triton.autotune(
30
+ configs=[
31
+ triton.Config({'BK': BK}, num_warps=num_warps)
32
+ for BK in [32, 64]
33
+ for num_warps in [1, 2, 4]
34
+ ],
35
+ key=["H", "HV", "K", "BC"],
36
+)
37
+@triton.jit(do_not_specialize=['T'])
38
+def chunk_kda_fwd_kernel_inter_solve_fused(
39
+ q,
40
+ k,
41
+ g,
42
+ beta,
43
+ Aqk,
44
+ Akkd,
45
+ Akk,
46
+ scale,
47
+ cu_seqlens,
48
+ chunk_indices,
49
+ T,
50
+ H: tl.constexpr,
51
+ HV: tl.constexpr,
52
+ K: tl.constexpr,
53
+ BT: tl.constexpr,
54
+ BC: tl.constexpr,
55
+ BK: tl.constexpr,
56
+ IS_VARLEN: tl.constexpr,
57
+ USE_SAFE_GATE: tl.constexpr,
58
+ SOLVE_TRIL_DOT_PRECISION: tl.constexpr,
59
+):
60
+ """
61
+ Fused kernel: compute inter-subchunk Akk + solve_tril in one pass.
62
+ Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd.
63
+
64
+ This kernel:
65
+ 1. Computes off-diagonal Aqk blocks -> writes to global
66
+ 2. Computes off-diagonal Akk blocks -> keeps in registers
67
+ 3. Loads diagonal Akk blocks from Akkd (fp32)
68
+ 4. Does forward substitution on diagonals
69
+ 5. Computes merged Akk_inv
70
+ 6. Writes Akk_inv to Akk
71
+ """
72
+ i_t, i_bh = tl.program_id(0), tl.program_id(1)
73
+ i_b, i_hv = i_bh // HV, i_bh % HV
74
+ i_h = i_hv // (HV // H)
75
+
76
+ if IS_VARLEN:
77
+ i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
78
+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
79
+ T = eos - bos
80
+ else:
81
+ bos, eos = i_b * T, i_b * T + T
82
+
83
+ if i_t * BT >= T:
84
+ return
85
+
86
+ i_tc0 = i_t * BT
87
+ i_tc1 = i_t * BT + BC
88
+ i_tc2 = i_t * BT + 2 * BC
89
+ i_tc3 = i_t * BT + 3 * BC
90
+
91
+ q += (bos * H + i_h) * K
92
+ k += (bos * H + i_h) * K
93
+ g += (bos * HV + i_hv) * K
94
+ Aqk += (bos * HV + i_hv) * BT
95
+ Akk += (bos * HV + i_hv) * BT
96
+ Akkd += (bos * HV + i_hv) * BC
97
+
98
+ o_i = tl.arange(0, BC)
99
+ m_tc1 = (i_tc1 + o_i) < T
100
+ m_tc2 = (i_tc2 + o_i) < T
101
+ m_tc3 = (i_tc3 + o_i) < T
102
+
103
+ b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32)
104
+ b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32)
105
+
106
+ b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32)
107
+ b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32)
108
+ b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32)
109
+ b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32)
110
+
111
+ b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32)
112
+ b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32)
113
+ b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32)
114
+ b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32)
115
+ b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32)
116
+ b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32)
117
+
118
+ ################################################################################
119
+ # off-diagonal blocks
120
+ ################################################################################
121
+ for i_k in range(tl.cdiv(K, BK)):
122
+ o_k = i_k * BK + tl.arange(0, BK)
123
+ m_k = o_k < K
124
+
125
+ p_k0 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0))
126
+ p_g0 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0))
127
+ b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32)
128
+ b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32)
129
+
130
+ if i_tc1 < T:
131
+ p_q1 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0))
132
+ p_k1 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0))
133
+ p_g1 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0))
134
+ # [BC, BK]
135
+ b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32)
136
+ b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32)
137
+ b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32)
138
+ # [BK]
139
+ b_gn1 = tl.load(g + i_tc1 * HV*K + o_k, mask=m_k, other=0).to(tl.float32)
140
+ # [BC, BK]
141
+ b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0)
142
+ # [BK, BC]
143
+ b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0))
144
+ # [BC, BC]
145
+ b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt)
146
+ b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt)
147
+
148
+ if i_tc2 < T:
149
+ p_q2 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0))
150
+ p_k2 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0))
151
+ p_g2 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0))
152
+ # [BC, BK]
153
+ b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32)
154
+ b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32)
155
+ b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32)
156
+ # [BK]
157
+ b_gn2 = tl.load(g + i_tc2 * HV*K + o_k, mask=m_k, other=0).to(tl.float32)
158
+ # [BC, BK]
159
+ b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0)
160
+ b_qg2 = b_q2 * b_gqn2
161
+ b_kg2 = b_k2 * b_gqn2
162
+ # [BK, BC]
163
+ b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0))
164
+ b_Aqk20 += tl.dot(b_qg2, b_kgt)
165
+ b_Akk20 += tl.dot(b_kg2, b_kgt)
166
+ # [BC, BC]
167
+ b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1))
168
+ # [BC, BC]
169
+ b_Aqk21 += tl.dot(b_qg2, b_kgt)
170
+ b_Akk21 += tl.dot(b_kg2, b_kgt)
171
+
172
+ if i_tc3 < T:
173
+ p_q3 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0))
174
+ p_k3 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0))
175
+ p_g3 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0))
176
+ # [BC, BK]
177
+ b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32)
178
+ b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32)
179
+ b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32)
180
+ # [BK]
181
+ b_gn3 = tl.load(g + i_tc3 * HV*K + o_k, mask=m_k, other=0).to(tl.float32)
182
+ # [BC, BK]
183
+ b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0)
184
+ b_qg3 = b_q3 * b_gqn3
185
+ b_kg3 = b_k3 * b_gqn3
186
+ # [BK, BC]
187
+ b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0))
188
+ # [BC, BC]
189
+ b_Aqk30 += tl.dot(b_qg3, b_kgt)
190
+ b_Akk30 += tl.dot(b_kg3, b_kgt)
191
+ # [BK, BC]
192
+ b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1))
193
+ # [BC, BC]
194
+ b_Aqk31 += tl.dot(b_qg3, b_kgt)
195
+ b_Akk31 += tl.dot(b_kg3, b_kgt)
196
+ # [BK, BC]
197
+ b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2))
198
+ # [BC, BC]
199
+ b_Aqk32 += tl.dot(b_qg3, b_kgt)
200
+ b_Akk32 += tl.dot(b_kg3, b_kgt)
201
+
202
+ ################################################################################
203
+ # save off-diagonal Aqk blocks and prepare Akk
204
+ ################################################################################
205
+ if i_tc1 < T:
206
+ p_Aqk10 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0))
207
+ tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
208
+
209
+ p_b1 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,))
210
+ b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32)
211
+ b_Akk10 = b_Akk10 * b_b1[:, None]
212
+ if i_tc2 < T:
213
+ p_Aqk20 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0))
214
+ p_Aqk21 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0))
215
+ tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
216
+ tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
217
+
218
+ p_b2 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,))
219
+ b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32)
220
+ b_Akk20 = b_Akk20 * b_b2[:, None]
221
+ b_Akk21 = b_Akk21 * b_b2[:, None]
222
+ if i_tc3 < T:
223
+ p_Aqk30 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0))
224
+ p_Aqk31 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0))
225
+ p_Aqk32 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0))
226
+ tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
227
+ tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
228
+ tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
229
+
230
+ p_b3 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,))
231
+ b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32)
232
+ b_Akk30 = b_Akk30 * b_b3[:, None]
233
+ b_Akk31 = b_Akk31 * b_b3[:, None]
234
+ b_Akk32 = b_Akk32 * b_b3[:, None]
235
+
236
+ p_Akk00 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc0, 0), (BC, BC), (1, 0))
237
+ p_Akk11 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc1, 0), (BC, BC), (1, 0))
238
+ p_Akk22 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc2, 0), (BC, BC), (1, 0))
239
+ p_Akk33 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc3, 0), (BC, BC), (1, 0))
240
+ b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32)
241
+ b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32)
242
+ b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32)
243
+ b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32)
244
+
245
+ ################################################################################
246
+ # forward substitution on diagonals
247
+ ################################################################################
248
+
249
+ if not USE_SAFE_GATE:
250
+ m_A = o_i[:, None] > o_i[None, :]
251
+ m_I = o_i[:, None] == o_i[None, :]
252
+
253
+ b_Ai00 = -tl.where(m_A, b_Ai00, 0)
254
+ b_Ai11 = -tl.where(m_A, b_Ai11, 0)
255
+ b_Ai22 = -tl.where(m_A, b_Ai22, 0)
256
+ b_Ai33 = -tl.where(m_A, b_Ai33, 0)
257
+
258
+ for i in range(2, min(BC, T - i_tc0)):
259
+ b_a00 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i)
260
+ b_a00 = tl.where(o_i < i, b_a00, 0.)
261
+ b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0)
262
+ b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00)
263
+ for i in range(BC + 2, min(2*BC, T - i_tc0)):
264
+ b_a11 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i)
265
+ b_a11 = tl.where(o_i < i - BC, b_a11, 0.)
266
+ b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0)
267
+ b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11)
268
+ for i in range(2*BC + 2, min(3*BC, T - i_tc0)):
269
+ b_a22 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i)
270
+ b_a22 = tl.where(o_i < i - 2*BC, b_a22, 0.)
271
+ b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0)
272
+ b_Ai22 = tl.where((o_i == i - 2*BC)[:, None], b_a22, b_Ai22)
273
+ for i in range(3*BC + 2, min(4*BC, T - i_tc0)):
274
+ b_a33 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i)
275
+ b_a33 = tl.where(o_i < i - 3*BC, b_a33, 0.)
276
+ b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0)
277
+ b_Ai33 = tl.where((o_i == i - 3*BC)[:, None], b_a33, b_Ai33)
278
+
279
+ b_Ai00 += m_I
280
+ b_Ai11 += m_I
281
+ b_Ai22 += m_I
282
+ b_Ai33 += m_I
283
+
284
+ ################################################################################
285
+ # compute merged inverse using off-diagonals
286
+ ################################################################################
287
+
288
+ # we used tf32 to maintain matrix inverse's precision whenever possible.
289
+ b_Ai10 = -tl.dot(
290
+ tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION),
291
+ b_Ai00,
292
+ input_precision=SOLVE_TRIL_DOT_PRECISION
293
+ )
294
+ b_Ai21 = -tl.dot(
295
+ tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION),
296
+ b_Ai11,
297
+ input_precision=SOLVE_TRIL_DOT_PRECISION
298
+ )
299
+ b_Ai32 = -tl.dot(
300
+ tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION),
301
+ b_Ai22,
302
+ input_precision=SOLVE_TRIL_DOT_PRECISION
303
+ )
304
+
305
+ b_Ai20 = -tl.dot(
306
+ b_Ai22,
307
+ tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) +
308
+ tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION),
309
+ input_precision=SOLVE_TRIL_DOT_PRECISION
310
+ )
311
+ b_Ai31 = -tl.dot(
312
+ b_Ai33,
313
+ tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) +
314
+ tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION),
315
+ input_precision=SOLVE_TRIL_DOT_PRECISION
316
+ )
317
+ b_Ai30 = -tl.dot(
318
+ b_Ai33,
319
+ tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) +
320
+ tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) +
321
+ tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION),
322
+ input_precision=SOLVE_TRIL_DOT_PRECISION
323
+ )
324
+
325
+ ################################################################################
326
+ # store full Akk_inv to Akk
327
+ ################################################################################
328
+
329
+ p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc0, 0), (BC, BC), (1, 0))
330
+ p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0))
331
+ p_Akk11 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, BC), (BC, BC), (1, 0))
332
+ p_Akk20 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0))
333
+ p_Akk21 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0))
334
+ p_Akk22 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 2*BC), (BC, BC), (1, 0))
335
+ p_Akk30 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0))
336
+ p_Akk31 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0))
337
+ p_Akk32 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0))
338
+ p_Akk33 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 3*BC), (BC, BC), (1, 0))
339
+
340
+ tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1))
341
+ tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1))
342
+ tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1))
343
+ tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1))
344
+ tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1))
345
+ tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1))
346
+ tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1))
347
+ tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1))
348
+ tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1))
349
+ tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1))
350
+
351
+
352
+@triton.heuristics({
353
+ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
354
+})
355
+@triton.autotune(
356
+ configs=[
357
+ triton.Config({}, num_warps=num_warps, num_stages=num_stages)
358
+ for num_warps in [1, 2, 4, 8]
359
+ for num_stages in [2, 3, 4]
360
+ ],
361
+ key=['BK', 'NC', 'BT', 'HV'],
362
+)
363
+@triton.jit(do_not_specialize=['B', 'T'])
364
+def chunk_kda_fwd_kernel_intra_sub_chunk(
365
+ q,
366
+ k,
367
+ g,
368
+ beta,
369
+ Aqk,
370
+ Akk,
371
+ scale,
372
+ cu_seqlens,
373
+ chunk_indices,
374
+ T,
375
+ H: tl.constexpr,
376
+ HV: tl.constexpr,
377
+ K: tl.constexpr,
378
+ BT: tl.constexpr,
379
+ BC: tl.constexpr,
380
+ BK: tl.constexpr,
381
+ IS_VARLEN: tl.constexpr,
382
+ USE_GATHER: tl.constexpr,
383
+):
384
+ i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
385
+ i_b, i_hv = i_bh // HV, i_bh % HV
386
+ i_h = i_hv // (HV // H)
387
+
388
+ if IS_VARLEN:
389
+ i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
390
+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
391
+ T = eos - bos
392
+ else:
393
+ bos, eos = i_b * T, i_b * T + T
394
+
395
+ i_ti = i_t * BT + i_i * BC
396
+ if i_ti >= T:
397
+ return
398
+
399
+ o_c = i_ti + tl.arange(0, BC)
400
+ m_c = o_c < T
401
+
402
+ q = q + (bos * H + i_h) * K
403
+ k = k + (bos * H + i_h) * K
404
+ g = g + (bos * HV + i_hv) * K
405
+ beta = beta + bos * HV + i_hv
406
+ Aqk = Aqk + (bos * HV + i_hv) * BT
407
+ Akk = Akk + (bos * HV + i_hv) * BC
408
+
409
+ p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0))
410
+ p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0))
411
+ p_g = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_ti, 0), (BC, BK), (1, 0))
412
+
413
+ p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,))
414
+
415
+ b_q = tl.load(p_q, boundary_check=(0, 1))
416
+ b_k = tl.load(p_k, boundary_check=(0, 1))
417
+ b_g = tl.load(p_g, boundary_check=(0, 1))
418
+ b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32)
419
+
420
+ if USE_GATHER:
421
+ b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0)
422
+ else:
423
+ # caculate offset
424
+ p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV*K + tl.arange(0, BK)
425
+ b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0)
426
+ b_gn = b_gn[None, :]
427
+
428
+ # current block, keep numerical stability by subtracting the left boundary
429
+ # less than 85 to avoid overflow in exp2
430
+ b_gm = (b_g - b_gn).to(tl.float32)
431
+
432
+ b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.)
433
+ b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.)
434
+
435
+ b_kgt = tl.trans(b_k * b_gk)
436
+
437
+ b_Aqk = tl.dot(b_q * b_gq, b_kgt) * scale
438
+ b_Akk = tl.dot(b_k * b_gq, b_kgt) * b_beta[:, None]
439
+
440
+ o_i = tl.arange(0, BC)
441
+ m_Aqk = o_i[:, None] >= o_i[None, :]
442
+ m_Akk = o_i[:, None] > o_i[None, :]
443
+ m_I = o_i[:, None] == o_i[None, :]
444
+
445
+ b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0)
446
+ b_Akk = tl.where(m_Akk, b_Akk, 0.0)
447
+
448
+ p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0))
449
+ p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV*BC, 1), (i_ti, 0), (BC, BC), (1, 0))
450
+ tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1))
451
+ tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1))
452
+
453
+ tl.debug_barrier()
454
+
455
+ ################################################################################
456
+ # forward substitution
457
+ ################################################################################
458
+
459
+ b_Ai = -b_Akk
460
+ for i in range(2, min(BC, T - i_ti)):
461
+ b_a = -tl.load(Akk + (i_ti + i) * HV*BC + o_i)
462
+ b_a = tl.where(o_i < i, b_a, 0.)
463
+ b_a += tl.sum(b_a[:, None] * b_Ai, 0)
464
+ b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai)
465
+ b_Ai += m_I
466
+ tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1))
467
+
468
+
469
+def chunk_kda_fwd_intra(
470
+ q: torch.Tensor,
471
+ k: torch.Tensor,
472
+ gk: torch.Tensor | None = None,
473
+ beta: torch.Tensor | None = None,
474
+ scale: float | None = None,
475
+ cu_seqlens: torch.LongTensor | None = None,
476
+ chunk_size: int = 64,
477
+ chunk_indices: torch.LongTensor | None = None,
478
+ safe_gate: bool = False,
479
+):
480
+ B, T, H, K, HV = *k.shape, gk.shape[2]
481
+ BT = chunk_size
482
+ BC = 16
483
+ if chunk_indices is None and cu_seqlens is not None:
484
+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
485
+ NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
486
+ NC = triton.cdiv(BT, BC)
487
+
488
+ Aqk = torch.empty(B, T, HV, BT, device=k.device, dtype=k.dtype)
489
+ # Akk must be zero-initialized - kernel only writes lower triangular
490
+ Akk = torch.zeros(B, T, HV, BT, device=k.device, dtype=k.dtype)
491
+ # Separate fp32 buffer for diagonal 16x16 blocks (for precision in solve_tril)
492
+ Akkd = torch.empty(B, T, HV, BC, device=k.device, dtype=torch.float32)
493
+
494
+ # Compute diagonal blocks into Akkd in fp32.
495
+ if safe_gate:
496
+ grid = (NT, NC, B * HV)
497
+ BK = triton.next_power_of_2(K)
498
+ chunk_kda_fwd_kernel_intra_sub_chunk[grid](
499
+ q=q,
500
+ k=k,
501
+ g=gk,
502
+ beta=beta,
503
+ Aqk=Aqk,
504
+ Akk=Akkd,
505
+ scale=scale,
506
+ cu_seqlens=cu_seqlens,
507
+ chunk_indices=chunk_indices,
508
+ T=T,
509
+ H=H,
510
+ HV=HV,
511
+ K=K,
512
+ BT=BT,
513
+ BC=BC,
514
+ BK=BK,
515
+ USE_GATHER=is_gather_supported,
516
+ )
517
+ else:
518
+ Aqk, Akkd = chunk_kda_fwd_intra_token_parallel(
519
+ q=q,
520
+ k=k,
521
+ gk=gk,
522
+ beta=beta,
523
+ Aqk=Aqk,
524
+ Akk=Akkd,
525
+ scale=scale,
526
+ cu_seqlens=cu_seqlens,
527
+ chunk_size=BT,
528
+ sub_chunk_size=BC,
529
+ )
530
+
531
+ # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen)
532
+ solve_tril_dot_precision = (
533
+ "tf32"
534
+ if current_platform.is_cuda()
535
+ and current_platform.has_device_capability(80)
536
+ else "ieee"
537
+ )
538
+ grid = (NT, B * HV)
539
+ chunk_kda_fwd_kernel_inter_solve_fused[grid](
540
+ q=q,
541
+ k=k,
542
+ g=gk,
543
+ beta=beta,
544
+ Aqk=Aqk,
545
+ Akkd=Akkd,
546
+ Akk=Akk,
547
+ scale=scale,
548
+ cu_seqlens=cu_seqlens,
549
+ chunk_indices=chunk_indices,
550
+ T=T,
551
+ H=H,
552
+ HV=HV,
553
+ K=K,
554
+ BT=BT,
555
+ BC=BC,
556
+ USE_SAFE_GATE=safe_gate,
557
+ SOLVE_TRIL_DOT_PRECISION=solve_tril_dot_precision,
558
+ )
559
+ return Aqk, Akk
vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py
new
+178
@@ -0,0 +1,178 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li
4
+#
5
+# This file contains code copied from the flash-linear-attention project.
6
+# The original source was licensed under the MIT license.
7
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
8
+# Forward-only adaptation of flash-linear-attention 0.5.0.
9
+# ruff: noqa: E501
10
+
11
+# Token-parallel implementation of KDA intra chunk kernel
12
+
13
+import torch
14
+
15
+from vllm.third_party.flash_linear_attention.ops.op import exp2
16
+from vllm.triton_utils import tl, triton
17
+
18
+
19
+@triton.heuristics({
20
+ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
21
+})
22
+@triton.autotune(
23
+ configs=[
24
+ triton.Config({'BH': BH}, num_warps=num_warps)
25
+ for BH in [1, 2, 4, 8]
26
+ for num_warps in [1, 2, 4, 8]
27
+ ],
28
+ key=["K", "H", "HV"],
29
+)
30
+@triton.jit(do_not_specialize=['T', 'N'])
31
+def chunk_kda_fwd_kernel_intra_token_parallel(
32
+ q,
33
+ k,
34
+ g,
35
+ beta,
36
+ Aqk,
37
+ Akk,
38
+ scale,
39
+ cu_seqlens,
40
+ N,
41
+ T,
42
+ H: tl.constexpr,
43
+ HV: tl.constexpr,
44
+ K: tl.constexpr,
45
+ BT: tl.constexpr,
46
+ BC: tl.constexpr,
47
+ BH: tl.constexpr,
48
+ IS_VARLEN: tl.constexpr,
49
+):
50
+ i_tg, i_hg = tl.program_id(0), tl.program_id(1)
51
+
52
+ if IS_VARLEN:
53
+ i_n = 0
54
+ left, right = 0, N
55
+
56
+ # Unrolled binary search (max B=2^32)
57
+ # We can limit iterations based on expected max batch size if needed
58
+ # 20 iterations covers B=1M, usually enough
59
+ for _ in range(20):
60
+ if left < right:
61
+ mid = (left + right) // 2
62
+ if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int32):
63
+ right = mid
64
+ else:
65
+ left = mid + 1
66
+ i_n = left
67
+
68
+ bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
69
+ T = eos - bos
70
+ i_t = i_tg - bos
71
+ else:
72
+ bos = (i_tg // T) * T
73
+ i_t = i_tg % T
74
+
75
+ if i_t >= T:
76
+ return
77
+
78
+ i_c = i_t // BT
79
+ i_s = (i_t % BT) // BC
80
+ i_tc = i_c * BT
81
+ i_ts = i_tc + i_s * BC
82
+
83
+ G: tl.constexpr = HV // H
84
+
85
+ q += bos * H*K
86
+ k += bos * H*K
87
+ g += bos * HV*K
88
+ Aqk += bos * HV*BT
89
+ Akk += bos * HV*BC
90
+ beta += bos * HV
91
+
92
+ BK: tl.constexpr = triton.next_power_of_2(K)
93
+ o_hv = i_hg * BH + tl.arange(0, BH)
94
+ o_h = o_hv // G
95
+ o_k = tl.arange(0, BK)
96
+ m_hv = o_hv < HV
97
+ m_k = o_k < K
98
+ m_hk = m_hv[:, None] & m_k[None, :]
99
+
100
+ # q/k: [B, T, H, K], manual load via mapped qk head index
101
+ p_qk = o_h[:, None] * K + o_k[None, :]
102
+ b_q = tl.load(q + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
103
+ b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
104
+
105
+ # g: [B, T, HV, K], beta: [B, T, HV]
106
+ p_g = tl.make_block_ptr(g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0))
107
+ p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,))
108
+ b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)
109
+ b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32)
110
+ b_k *= b_beta[:, None]
111
+
112
+ for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))):
113
+ b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
114
+ p_gj = tl.make_block_ptr(g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0))
115
+ b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32)
116
+
117
+ b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0)
118
+ b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale
119
+ b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0)
120
+
121
+ tl.store(Aqk + i_t * HV * BT + o_hv * BT + j % BT, b_Aqk.to(Aqk.dtype.element_ty), mask=m_hv)
122
+ tl.store(Akk + i_t * HV * BC + o_hv * BC + j - i_ts, b_Akk.to(Akk.dtype.element_ty), mask=m_hv)
123
+
124
+
125
+def chunk_kda_fwd_intra_token_parallel(
126
+ q: torch.Tensor,
127
+ k: torch.Tensor,
128
+ gk: torch.Tensor,
129
+ beta: torch.Tensor,
130
+ Aqk: torch.Tensor,
131
+ Akk: torch.Tensor,
132
+ scale: float,
133
+ cu_seqlens: torch.LongTensor | None = None,
134
+ chunk_size: int = 64,
135
+ sub_chunk_size: int = 16,
136
+) -> None:
137
+ """
138
+ Token-parallel implementation: each token gets its own thread block.
139
+ Supports both fixed-length and variable-length sequences.
140
+ Reduces wasted computation on padding.
141
+
142
+ Writes directly to Aqk and Akk tensors (in-place).
143
+
144
+ Args:
145
+ q: [B, T, H, K]
146
+ k: [B, T, H, K]
147
+ gk: [B, T, HV, K] cumsum of gates (HV >= H for GVA)
148
+ beta: [B, T, HV]
149
+ Aqk: [B, T, HV, BT] output tensor to write to
150
+ Akk: [B, T, HV, BC] output tensor for diagonal blocks (fp32)
151
+ scale: attention scale
152
+ chunk_size: BT (default 64)
153
+ sub_chunk_size: BC (default 16)
154
+ """
155
+ B, T, H, K, HV = *q.shape, gk.shape[2]
156
+ N = len(cu_seqlens) - 1 if cu_seqlens is not None else B
157
+ BT = chunk_size
158
+ BC = sub_chunk_size
159
+
160
+ def grid(meta): return (B * T, triton.cdiv(HV, meta['BH']))
161
+ chunk_kda_fwd_kernel_intra_token_parallel[grid](
162
+ q=q,
163
+ k=k,
164
+ g=gk,
165
+ beta=beta,
166
+ Aqk=Aqk,
167
+ Akk=Akk,
168
+ scale=scale,
169
+ cu_seqlens=cu_seqlens,
170
+ N=N,
171
+ T=T,
172
+ H=H,
173
+ HV=HV,
174
+ K=K,
175
+ BT=BT,
176
+ BC=BC,
177
+ )
178
+ return Aqk, Akk
vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py
new
+671
@@ -0,0 +1,671 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li
4
+#
5
+# This file contains code adapted from the flash-linear-attention project.
6
+# The original source was licensed under the MIT license.
7
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
8
+# ruff: noqa: E501
9
+
10
+import torch
11
+
12
+from vllm.platforms import current_platform
13
+from vllm.third_party.flash_linear_attention.ops.op import exp, log
14
+from vllm.triton_utils import tl, triton
15
+from vllm.utils.math_utils import cdiv, next_power_of_2
16
+
17
+
18
+@triton.heuristics(
19
+ {
20
+ "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None,
21
+ "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None,
22
+ }
23
+)
24
+@triton.jit
25
+def _kda_gate_beta_fwd_kernel(
26
+ raw_g,
27
+ raw_beta,
28
+ A_log,
29
+ dt_bias,
30
+ gate,
31
+ beta_out,
32
+ lower_bound,
33
+ softplus_beta: tl.constexpr,
34
+ softplus_threshold: tl.constexpr,
35
+ T,
36
+ stride_g_token: tl.constexpr,
37
+ stride_beta_token: tl.constexpr,
38
+ H: tl.constexpr,
39
+ D: tl.constexpr,
40
+ BT: tl.constexpr,
41
+ BD: tl.constexpr,
42
+ HAS_DT_BIAS: tl.constexpr,
43
+ USE_LOWER_BOUND: tl.constexpr,
44
+ launch_pdl: tl.constexpr,
45
+):
46
+ if launch_pdl:
47
+ tl.extra.cuda.gdc_wait()
48
+ tl.extra.cuda.gdc_launch_dependents()
49
+
50
+ i_t, i_h = tl.program_id(0), tl.program_id(1)
51
+ o_t = i_t * BT + tl.arange(0, BT)
52
+ o_d = tl.arange(0, BD)
53
+ m_t = o_t < T
54
+ m_d = o_d < D
55
+
56
+ p_g = raw_g + o_t[:, None] * stride_g_token + i_h * D + o_d[None, :]
57
+ b_g = tl.load(p_g, mask=m_t[:, None] & m_d[None, :], other=0.0).to(tl.float32)
58
+ if HAS_DT_BIAS:
59
+ b_bias = tl.load(
60
+ dt_bias + i_h * D + o_d,
61
+ mask=m_d,
62
+ other=0.0,
63
+ ).to(tl.float32)
64
+ b_g += b_bias[None, :]
65
+
66
+ b_a = exp(tl.load(A_log + i_h).to(tl.float32))
67
+ if USE_LOWER_BOUND:
68
+ b_gate = lower_bound * tl.sigmoid(b_a * b_g)
69
+ else:
70
+ b_scaled = b_g * softplus_beta
71
+ b_softplus = tl.where(
72
+ b_scaled > softplus_threshold,
73
+ b_g,
74
+ log(1.0 + tl.exp(b_scaled)) / softplus_beta,
75
+ )
76
+ b_gate = -b_a * b_softplus
77
+
78
+ p_gate = gate + (o_t[:, None] * H + i_h) * D + o_d[None, :]
79
+ tl.store(
80
+ p_gate,
81
+ b_gate,
82
+ mask=m_t[:, None] & m_d[None, :],
83
+ )
84
+
85
+ b_beta = tl.load(
86
+ raw_beta + o_t * stride_beta_token + i_h,
87
+ mask=m_t,
88
+ other=0.0,
89
+ ).to(tl.float32)
90
+ tl.store(beta_out + o_t * H + i_h, tl.sigmoid(b_beta), mask=m_t)
91
+
92
+
93
+def _fused_kda_gate_beta(
94
+ raw_g: torch.Tensor,
95
+ raw_beta: torch.Tensor,
96
+ A_log: torch.Tensor,
97
+ dt_bias: torch.Tensor | None,
98
+ lower_bound: float | None,
99
+) -> tuple[torch.Tensor, torch.Tensor]:
100
+ B, T, H, D = raw_g.shape
101
+ assert B == 1
102
+ assert raw_beta.shape == (B, T, H)
103
+ assert raw_g.stride()[2:] == (D, 1)
104
+ assert raw_beta.stride(2) == 1
105
+ gate = torch.empty((B, T, H, D), dtype=torch.float32, device=raw_g.device)
106
+ beta = torch.empty((B, T, H), dtype=torch.float32, device=raw_beta.device)
107
+
108
+ BT = 16
109
+ _kda_gate_beta_fwd_kernel[(cdiv(T, BT), H)](
110
+ raw_g=raw_g,
111
+ raw_beta=raw_beta,
112
+ A_log=A_log,
113
+ dt_bias=dt_bias,
114
+ gate=gate,
115
+ beta_out=beta,
116
+ lower_bound=lower_bound,
117
+ softplus_beta=1.0,
118
+ softplus_threshold=20.0,
119
+ T=T,
120
+ stride_g_token=raw_g.stride(1),
121
+ stride_beta_token=raw_beta.stride(1),
122
+ H=H,
123
+ D=D,
124
+ BT=BT,
125
+ BD=next_power_of_2(D),
126
+ num_warps=4,
127
+ launch_pdl=current_platform.is_arch_support_pdl(),
128
+ )
129
+ return gate, beta
130
+
131
+
132
+@triton.heuristics(
133
+ {
134
+ "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None,
135
+ "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None,
136
+ "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None,
137
+ }
138
+)
139
+@triton.jit(do_not_specialize=["N", "T", "stride_beta_token"])
140
+def fused_recurrent_kda_fwd_kernel(
141
+ q,
142
+ k,
143
+ v,
144
+ g,
145
+ beta,
146
+ A_log,
147
+ dt_bias,
148
+ out,
149
+ state,
150
+ cu_seqlens,
151
+ state_indices,
152
+ num_accepted_tokens,
153
+ lower_bound,
154
+ scale: tl.constexpr,
155
+ N: tl.int64,
156
+ T: tl.int64,
157
+ H: tl.constexpr,
158
+ K: tl.constexpr,
159
+ V: tl.constexpr,
160
+ BK: tl.constexpr,
161
+ BV: tl.constexpr,
162
+ stride_qkv_token: tl.constexpr,
163
+ stride_g_token: tl.constexpr,
164
+ stride_beta_token,
165
+ stride_out_token: tl.constexpr,
166
+ stride_state_token: tl.constexpr,
167
+ stride_indices_seq: tl.constexpr,
168
+ IS_SPEC_DECODING: tl.constexpr,
169
+ USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
170
+ USE_GATE_IN_KERNEL: tl.constexpr,
171
+ APPLY_BETA_SIGMOID: tl.constexpr,
172
+ HAS_DT_BIAS: tl.constexpr,
173
+ USE_LOWER_BOUND: tl.constexpr,
174
+ num_stages: tl.constexpr,
175
+ launch_pdl: tl.constexpr,
176
+):
177
+ if launch_pdl:
178
+ tl.extra.cuda.gdc_wait()
179
+
180
+ pid = tl.program_id(0)
181
+ i_v = pid % tl.cdiv(V, BV)
182
+ i_nh = pid // tl.cdiv(V, BV)
183
+ i_n, i_h = i_nh // H, i_nh % H
184
+ bos = tl.load(cu_seqlens + i_n).to(tl.int64)
185
+ eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64)
186
+ sequence_length = eos - bos
187
+ if sequence_length == 0:
188
+ return
189
+
190
+ o_k = tl.arange(0, BK)
191
+ o_v = i_v * BV + tl.arange(0, BV)
192
+ m_k = o_k < K
193
+ m_v = o_v < V
194
+ m_state = m_v[:, None] & m_k[None, :]
195
+
196
+ if IS_SPEC_DECODING:
197
+ initial_token = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1
198
+ else:
199
+ initial_token = 0
200
+ state_index = tl.load(state_indices + i_n * stride_indices_seq + initial_token).to(
201
+ tl.int64
202
+ )
203
+ p_out = out + bos * stride_out_token + i_h * V + o_v
204
+ if state_index <= 0:
205
+ tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=m_v)
206
+ return
207
+
208
+ p_state = (
209
+ state
210
+ + state_index * stride_state_token
211
+ + i_h * V * K
212
+ + o_v[:, None] * K
213
+ + o_k[None, :]
214
+ )
215
+ b_state = tl.load(p_state, mask=m_state, other=0.0).to(tl.float32)
216
+
217
+ p_q = q + bos * stride_qkv_token + i_h * K + o_k
218
+ p_k = k + bos * stride_qkv_token + i_h * K + o_k
219
+ p_v = v + bos * stride_qkv_token + i_h * V + o_v
220
+ p_g = g + bos * stride_g_token + i_h * K + o_k
221
+ p_beta = beta + bos * stride_beta_token + i_h
222
+ for i_t in tl.range(0, sequence_length, num_stages=num_stages):
223
+ b_q = tl.load(p_q, mask=m_k, other=0.0, eviction_policy="evict_last").to(
224
+ tl.float32
225
+ )
226
+ b_k = tl.load(p_k, mask=m_k, other=0.0, eviction_policy="evict_last").to(
227
+ tl.float32
228
+ )
229
+ b_v = tl.load(p_v, mask=m_v, other=0.0, eviction_policy="evict_first").to(
230
+ tl.float32
231
+ )
232
+ if USE_QK_L2NORM_IN_KERNEL:
233
+ b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
234
+ b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
235
+ b_q *= scale
236
+
237
+ b_gate = tl.load(
238
+ p_g,
239
+ mask=m_k,
240
+ other=0.0,
241
+ eviction_policy="evict_last",
242
+ ).to(tl.float32)
243
+ if USE_GATE_IN_KERNEL:
244
+ if HAS_DT_BIAS:
245
+ b_bias = tl.load(
246
+ dt_bias + i_h * K + o_k,
247
+ mask=m_k,
248
+ other=0.0,
249
+ ).to(tl.float32)
250
+ b_gate += b_bias
251
+ b_a = exp(tl.load(A_log + i_h).to(tl.float32))
252
+ if USE_LOWER_BOUND:
253
+ b_gate = lower_bound * tl.sigmoid(b_a * b_gate)
254
+ else:
255
+ b_softplus = tl.where(
256
+ b_gate > 20.0,
257
+ b_gate,
258
+ log(1.0 + tl.exp(b_gate)),
259
+ )
260
+ b_gate = -b_a * b_softplus
261
+
262
+ b_state *= exp(b_gate[None, :])
263
+ b_v -= tl.sum(b_state * b_k[None, :], axis=1)
264
+ b_beta = tl.load(p_beta, eviction_policy="evict_last").to(tl.float32)
265
+ if APPLY_BETA_SIGMOID:
266
+ b_beta = tl.sigmoid(b_beta)
267
+ b_v *= b_beta
268
+ b_state += b_v[:, None] * b_k[None, :]
269
+ b_out = tl.sum(b_state * b_q[None, :], axis=1)
270
+ tl.store(
271
+ p_out,
272
+ b_out.to(p_out.dtype.element_ty),
273
+ mask=m_v,
274
+ eviction_policy="evict_first",
275
+ )
276
+
277
+ final_state_index = tl.load(state_indices + i_n * stride_indices_seq + i_t).to(
278
+ tl.int64
279
+ )
280
+ if final_state_index > 0:
281
+ p_final_state = (
282
+ state
283
+ + final_state_index * stride_state_token
284
+ + i_h * V * K
285
+ + o_v[:, None] * K
286
+ + o_k[None, :]
287
+ )
288
+ tl.store(
289
+ p_final_state,
290
+ b_state.to(p_final_state.dtype.element_ty),
291
+ mask=m_state,
292
+ )
293
+
294
+ p_q += stride_qkv_token
295
+ p_k += stride_qkv_token
296
+ p_v += stride_qkv_token
297
+ p_g += stride_g_token
298
+ p_beta += stride_beta_token
299
+ p_out += stride_out_token
300
+
301
+ if launch_pdl:
302
+ tl.extra.cuda.gdc_launch_dependents()
303
+
304
+
305
+# Consumed by kimi_k3_triton_warmup.py during kernel_warmup().
306
+def get_fused_recurrent_kda_fwd_warmup_profiles(
307
+ num_heads: int,
308
+) -> tuple[int, ...]:
309
+ """Return representative sequence counts for gated launch variants."""
310
+ # The region above 192 head-sequences reuses the second launch variant.
311
+ return (
312
+ 1,
313
+ 48 // num_heads + 1,
314
+ 96 // num_heads + 1,
315
+ )
316
+
317
+
318
+def fused_recurrent_kda_fwd(
319
+ q: torch.Tensor,
320
+ k: torch.Tensor,
321
+ v: torch.Tensor,
322
+ g: torch.Tensor,
323
+ beta: torch.Tensor,
324
+ scale: float | None = None,
325
+ initial_state: torch.Tensor | None = None,
326
+ inplace_final_state: bool = True,
327
+ cu_seqlens: torch.Tensor | None = None,
328
+ ssm_state_indices: torch.Tensor | None = None,
329
+ num_accepted_tokens: torch.Tensor | None = None,
330
+ use_qk_l2norm_in_kernel: bool = True,
331
+ A_log: torch.Tensor | None = None,
332
+ dt_bias: torch.Tensor | None = None,
333
+ lower_bound: float | None = None,
334
+ use_gate_in_kernel: bool = False,
335
+ use_beta_sigmoid_in_kernel: bool = False,
336
+ out: torch.Tensor | None = None,
337
+) -> tuple[torch.Tensor, torch.Tensor]:
338
+ """Launch recurrent KDA with dense inner dimensions and row strides."""
339
+ B, T, H, K = q.shape
340
+ V = v.shape[-1]
341
+ assert B == 1 and k.shape == q.shape
342
+ assert v.shape == (B, T, H, V) and g.shape == (B, T, H, K)
343
+ assert beta.shape == (B, T, H)
344
+ assert initial_state is not None
345
+ assert cu_seqlens is not None
346
+ assert ssm_state_indices is not None
347
+ assert inplace_final_state
348
+ if out is None:
349
+ out = torch.empty_like(v)
350
+ assert out.shape == v.shape
351
+ assert initial_state.shape[1:] == (H, V, K)
352
+ assert ssm_state_indices.ndim in (1, 2)
353
+
354
+ assert q.stride()[2:] == k.stride()[2:] == (K, 1)
355
+ assert v.stride()[2:] == out.stride()[2:] == (V, 1)
356
+ assert g.stride()[2:] == (K, 1)
357
+ assert beta.stride(2) == 1
358
+ assert q.stride(1) == k.stride(1) == v.stride(1)
359
+ assert initial_state.stride()[1:] == (V * K, K, 1)
360
+ N = cu_seqlens.numel() - 1
361
+ if ssm_state_indices.ndim == 1:
362
+ assert T == N
363
+ assert num_accepted_tokens is None
364
+ else:
365
+ assert ssm_state_indices.stride(1) == 1
366
+ assert cu_seqlens.is_contiguous()
367
+ if use_gate_in_kernel:
368
+ assert A_log is not None and A_log.is_contiguous()
369
+ assert dt_bias is None or dt_bias.is_contiguous()
370
+
371
+ if scale is None:
372
+ scale = K**-0.5
373
+
374
+ if use_gate_in_kernel:
375
+ # Tuned on GB300 for Kimi-K3 shapes. Keep the warmup profiles above in
376
+ # sync with these boundaries.
377
+ head_sequences = H * N
378
+ if head_sequences <= 48:
379
+ BV, num_stages = 4, 4
380
+ elif head_sequences <= 96:
381
+ BV, num_stages = 8, 3
382
+ elif head_sequences <= 192:
383
+ BV, num_stages = 16, 3
384
+ else:
385
+ BV, num_stages = 8, 3
386
+ num_warps = 1
387
+ else:
388
+ BV, num_warps, num_stages = 8, 1, 2
389
+ grid = (cdiv(V, BV) * N * H,)
390
+ fused_recurrent_kda_fwd_kernel[grid](
391
+ q=q,
392
+ k=k,
393
+ v=v,
394
+ g=g,
395
+ beta=beta,
396
+ A_log=A_log,
397
+ dt_bias=dt_bias,
398
+ out=out,
399
+ state=initial_state,
400
+ cu_seqlens=cu_seqlens,
401
+ state_indices=ssm_state_indices,
402
+ num_accepted_tokens=num_accepted_tokens,
403
+ lower_bound=lower_bound,
404
+ scale=scale,
405
+ N=N,
406
+ T=T,
407
+ H=H,
408
+ K=K,
409
+ V=V,
410
+ BK=next_power_of_2(K),
411
+ BV=BV,
412
+ stride_qkv_token=q.stride(1),
413
+ stride_g_token=g.stride(1),
414
+ stride_beta_token=beta.stride(1),
415
+ stride_out_token=out.stride(1),
416
+ stride_state_token=initial_state.stride(0),
417
+ stride_indices_seq=ssm_state_indices.stride(0),
418
+ IS_SPEC_DECODING=num_accepted_tokens is not None,
419
+ USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
420
+ USE_GATE_IN_KERNEL=use_gate_in_kernel,
421
+ APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel,
422
+ num_warps=num_warps,
423
+ num_stages=num_stages,
424
+ launch_pdl=current_platform.is_arch_support_pdl(),
425
+ )
426
+ return out, initial_state
427
+
428
+
429
+def fused_recurrent_kda(
430
+ q: torch.Tensor,
431
+ k: torch.Tensor,
432
+ v: torch.Tensor,
433
+ raw_g: torch.Tensor,
434
+ raw_beta: torch.Tensor,
435
+ A_log: torch.Tensor,
436
+ dt_bias: torch.Tensor | None,
437
+ lower_bound: float | None,
438
+ initial_state: torch.Tensor,
439
+ cu_seqlens: torch.Tensor,
440
+ ssm_state_indices: torch.Tensor,
441
+ num_accepted_tokens: torch.Tensor | None = None,
442
+ out: torch.Tensor | None = None,
443
+ fuse_gate: bool | None = None,
444
+) -> tuple[torch.Tensor, torch.Tensor]:
445
+ """Run recurrent KDA from raw gate and beta inputs.
446
+
447
+ This vLLM wrapper applies the gate activation and beta sigmoid, selecting
448
+ whether to materialize them before launching the recurrent kernel.
449
+ """
450
+ if fuse_gate is None:
451
+ fuse_gate = True
452
+
453
+ if fuse_gate:
454
+ gate = raw_g
455
+ beta = raw_beta
456
+ else:
457
+ gate, beta = _fused_kda_gate_beta(
458
+ raw_g,
459
+ raw_beta,
460
+ A_log,
461
+ dt_bias,
462
+ lower_bound,
463
+ )
464
+ return fused_recurrent_kda_fwd(
465
+ q=q,
466
+ k=k,
467
+ v=v,
468
+ g=gate,
469
+ beta=beta,
470
+ scale=q.shape[-1] ** -0.5,
471
+ initial_state=initial_state,
472
+ inplace_final_state=True,
473
+ cu_seqlens=cu_seqlens,
474
+ ssm_state_indices=ssm_state_indices,
475
+ num_accepted_tokens=num_accepted_tokens,
476
+ use_qk_l2norm_in_kernel=True,
477
+ A_log=A_log if fuse_gate else None,
478
+ dt_bias=dt_bias if fuse_gate else None,
479
+ lower_bound=lower_bound if fuse_gate else None,
480
+ use_gate_in_kernel=fuse_gate,
481
+ use_beta_sigmoid_in_kernel=fuse_gate,
482
+ out=out,
483
+ )
484
+
485
+
486
+@triton.jit(
487
+ do_not_specialize=["stride_beta_token", "stride_state_indices"]
488
+)
489
+def fused_recurrent_kda_packed_decode_kernel(
490
+ mixed_qkv,
491
+ raw_g,
492
+ raw_beta,
493
+ A_log,
494
+ dt_bias,
495
+ out,
496
+ state,
497
+ state_indices,
498
+ lower_bound,
499
+ scale: tl.constexpr,
500
+ stride_mixed_token: tl.constexpr,
501
+ stride_g_token: tl.constexpr,
502
+ stride_beta_token,
503
+ stride_state_token: tl.constexpr,
504
+ stride_state_indices,
505
+ H: tl.constexpr,
506
+ K: tl.constexpr,
507
+ V: tl.constexpr,
508
+ BK: tl.constexpr,
509
+ BV: tl.constexpr,
510
+ SOFTPLUS_THRESHOLD: tl.constexpr,
511
+ USE_LOWER_BOUND: tl.constexpr,
512
+ launch_pdl: tl.constexpr,
513
+):
514
+ i_v, i_nh = tl.program_id(0), tl.program_id(1)
515
+ i_n, i_h = i_nh // H, i_nh % H
516
+
517
+ o_k = tl.arange(0, BK)
518
+ o_v = i_v * BV + tl.arange(0, BV)
519
+ mask_k = o_k < K
520
+ mask_v = o_v < V
521
+ mask_state = mask_v[:, None] & mask_k[None, :]
522
+
523
+ if launch_pdl:
524
+ tl.extra.cuda.gdc_wait()
525
+ tl.extra.cuda.gdc_launch_dependents()
526
+
527
+ state_idx = tl.load(state_indices + i_n * stride_state_indices).to(tl.int64)
528
+ p_out = out + (i_n * H + i_h) * V + o_v
529
+ if state_idx <= 0:
530
+ tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=mask_v)
531
+ return
532
+
533
+ p_state = state + state_idx * stride_state_token
534
+ p_state += i_h * V * K + o_v[:, None] * K + o_k[None, :]
535
+ b_state = tl.load(p_state, mask=mask_state, other=0).to(tl.float32)
536
+
537
+ # Q, K, and V occupy consecutive channel ranges, while the token stride
538
+ # may also include the output-gate projection that follows packed QKV.
539
+ p_mixed = mixed_qkv + i_n * stride_mixed_token
540
+ b_q = tl.load(p_mixed + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32)
541
+ b_k = tl.load(
542
+ p_mixed + H * K + i_h * K + o_k,
543
+ mask=mask_k,
544
+ other=0,
545
+ ).to(tl.float32)
546
+ b_v = tl.load(
547
+ p_mixed + 2 * H * K + i_h * V + o_v,
548
+ mask=mask_v,
549
+ other=0,
550
+ ).to(tl.float32)
551
+
552
+ b_q /= tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
553
+ b_k /= tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
554
+ b_q *= scale
555
+
556
+ p_g = raw_g + i_n * stride_g_token + i_h * K + o_k
557
+ b_g = tl.load(p_g, mask=mask_k, other=0).to(tl.float32)
558
+ b_bias = tl.load(dt_bias + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32)
559
+ b_a = exp(tl.load(A_log + i_h).to(tl.float32))
560
+ b_g += b_bias
561
+ if USE_LOWER_BOUND:
562
+ b_gate = lower_bound * tl.sigmoid(b_a * b_g)
563
+ else:
564
+ b_softplus = tl.where(
565
+ b_g > SOFTPLUS_THRESHOLD,
566
+ b_g,
567
+ log(1.0 + tl.exp(b_g)),
568
+ )
569
+ b_gate = -b_a * b_softplus
570
+
571
+ b_state *= exp(b_gate[None, :])
572
+ b_v -= tl.sum(b_state * b_k[None, :], axis=1)
573
+ b_beta = tl.sigmoid(
574
+ tl.load(raw_beta + i_n * stride_beta_token + i_h).to(tl.float32)
575
+ )
576
+ b_v *= b_beta
577
+ b_state += b_v[:, None] * b_k[None, :]
578
+ b_out = tl.sum(b_state * b_q[None, :], axis=1)
579
+
580
+ tl.store(p_out, b_out.to(p_out.dtype.element_ty), mask=mask_v)
581
+ tl.store(p_state, b_state.to(p_state.dtype.element_ty), mask=mask_state)
582
+
583
+
584
+def fused_recurrent_kda_packed_decode(
585
+ mixed_qkv: torch.Tensor,
586
+ raw_g: torch.Tensor,
587
+ raw_beta: torch.Tensor,
588
+ A_log: torch.Tensor,
589
+ dt_bias: torch.Tensor,
590
+ lower_bound: float | None,
591
+ initial_state: torch.Tensor,
592
+ state_indices: torch.Tensor,
593
+ scale: float | None = None,
594
+) -> tuple[torch.Tensor, torch.Tensor]:
595
+ """Run one-token KDA decode directly from packed post-conv QKV."""
596
+ if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1:
597
+ raise ValueError("`mixed_qkv` must be 2D and contiguous in its last dim.")
598
+ if raw_g.ndim != 4 or raw_g.shape[0] != 1:
599
+ raise ValueError("`raw_g` must have shape [1, B, H, K].")
600
+ if raw_beta.ndim != 3 or raw_beta.shape[0] != 1:
601
+ raise ValueError("`raw_beta` must have shape [1, B, H].")
602
+ if initial_state.ndim != 4:
603
+ raise ValueError("`initial_state` must have shape [cache, H, V, K].")
604
+ _, H, V, K = initial_state.shape
605
+ if raw_g.stride()[2:] != (K, 1):
606
+ raise ValueError("`raw_g` must be contiguous within each token.")
607
+ if raw_beta.stride(2) != 1:
608
+ raise ValueError("`raw_beta` heads must be contiguous.")
609
+ if initial_state.stride()[1:] != (V * K, K, 1):
610
+ raise ValueError("`initial_state` must be contiguous within each cache slot.")
611
+ if state_indices.ndim != 1:
612
+ raise ValueError("`state_indices` must be one-dimensional.")
613
+ if A_log.ndim != 1 or not A_log.is_contiguous():
614
+ raise ValueError("`A_log` must be contiguous and one-dimensional.")
615
+ if not dt_bias.is_contiguous():
616
+ raise ValueError("`dt_bias` must be contiguous.")
617
+
618
+ device = mixed_qkv.device
619
+ if any(
620
+ x.device != device
621
+ for x in (raw_g, raw_beta, A_log, dt_bias, initial_state, state_indices)
622
+ ):
623
+ raise ValueError("All packed KDA inputs must be on the same device.")
624
+
625
+ B = mixed_qkv.shape[0]
626
+ if raw_g.shape != (1, B, H, K):
627
+ raise ValueError(f"Unexpected raw gate shape {tuple(raw_g.shape)}.")
628
+ if raw_beta.shape != (1, B, H):
629
+ raise ValueError(f"Unexpected raw beta shape {tuple(raw_beta.shape)}.")
630
+ if mixed_qkv.shape[1] != 2 * H * K + H * V:
631
+ raise ValueError(f"Unexpected packed QKV shape {tuple(mixed_qkv.shape)}.")
632
+ if A_log.numel() != H or dt_bias.numel() != H * K:
633
+ raise ValueError("`A_log` or `dt_bias` has an incompatible shape.")
634
+ if state_indices.shape[0] != B:
635
+ raise ValueError("`state_indices` must contain one entry per token.")
636
+
637
+ BK = next_power_of_2(K)
638
+ BV = min(next_power_of_2(V), 32)
639
+ if scale is None:
640
+ scale = K**-0.5
641
+
642
+ out = torch.empty((1, B, H, V), dtype=mixed_qkv.dtype, device=device)
643
+ grid = (cdiv(V, BV), B * H)
644
+ fused_recurrent_kda_packed_decode_kernel[grid](
645
+ mixed_qkv=mixed_qkv,
646
+ raw_g=raw_g,
647
+ raw_beta=raw_beta,
648
+ A_log=A_log,
649
+ dt_bias=dt_bias,
650
+ out=out,
651
+ state=initial_state,
652
+ state_indices=state_indices,
653
+ lower_bound=lower_bound or 0.0,
654
+ scale=scale,
655
+ stride_mixed_token=mixed_qkv.stride(0),
656
+ stride_g_token=raw_g.stride(1),
657
+ stride_beta_token=raw_beta.stride(1),
658
+ stride_state_token=initial_state.stride(0),
659
+ stride_state_indices=state_indices.stride(0),
660
+ H=H,
661
+ K=K,
662
+ V=V,
663
+ BK=BK,
664
+ BV=BV,
665
+ SOFTPLUS_THRESHOLD=20.0,
666
+ USE_LOWER_BOUND=lower_bound is not None,
667
+ num_warps=4,
668
+ num_stages=2,
669
+ launch_pdl=current_platform.is_arch_support_pdl(),
670
+ )
671
+ return out, initial_state
vllm/models/kimi_k3/nvidia/ops/vision_fa4_warmup.py
new
+195
@@ -0,0 +1,195 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Startup compilation of Kimi-K3 vision FA4 kernels."""
4
+
5
+from __future__ import annotations
6
+
7
+import math
8
+from collections.abc import Iterator
9
+from dataclasses import dataclass
10
+from functools import partial
11
+
12
+import torch
13
+
14
+from vllm.model_executor.warmup.cutedsl_warmup import (
15
+ CuTeDSLCompileUnit,
16
+ register_cutedsl_warmup_provider,
17
+)
18
+from vllm.platforms import current_platform
19
+
20
+_FA4_TILE_SIZE = 128
21
+_FA4_MAX_SPLITS = 128
22
+
23
+
24
+@dataclass(frozen=True)
25
+class KimiK3VisionFA4WarmupConfig:
26
+ num_heads: int
27
+ head_dim: int
28
+ dtype: torch.dtype
29
+ max_batch_size: int
30
+ max_seqlen: int
31
+
32
+
33
+@dataclass(frozen=True)
34
+class _FA4WarmupProbe:
35
+ batch_size: int
36
+ max_seqlen: int
37
+ dispatch_key: tuple[int, bool, int | None]
38
+
39
+
40
+def _combine_log_max_splits(head_dim: int, num_splits: int) -> int:
41
+ k_block_size = 64 if head_dim <= 64 else 128
42
+ tile_m = 8 if k_block_size == 128 else 16
43
+ min_log_max_splits = 5 if tile_m == 8 else 4
44
+ return max(math.ceil(math.log2(num_splits)), min_log_max_splits)
45
+
46
+
47
+def _get_dispatch(
48
+ config: KimiK3VisionFA4WarmupConfig,
49
+ *,
50
+ batch_size: int,
51
+ max_seqlen: int,
52
+ num_sms: int,
53
+) -> tuple[int, bool, int | None]:
54
+ # These are the shape-dependent fields in FA4's forward and combine
55
+ # compile keys for noncausal varlen MHA. Compile units still pass
56
+ # num_splits=0 so the production selector makes the final decision.
57
+ q_stage = 2 if max_seqlen > _FA4_TILE_SIZE else 1
58
+ effective_q_tile = q_stage * _FA4_TILE_SIZE
59
+ num_m_blocks = math.ceil(max_seqlen / effective_q_tile)
60
+ num_n_blocks = math.ceil(max_seqlen / _FA4_TILE_SIZE)
61
+
62
+ if num_n_blocks <= 4:
63
+ num_splits = 1
64
+ else:
65
+ total_m_blocks = batch_size * config.num_heads * num_m_blocks
66
+ num_splits = min(
67
+ num_sms // total_m_blocks,
68
+ _FA4_MAX_SPLITS,
69
+ num_n_blocks,
70
+ )
71
+
72
+ is_split_kv = num_splits > 1
73
+ combine_key = (
74
+ _combine_log_max_splits(config.head_dim, num_splits) if is_split_kv else None
75
+ )
76
+ return q_stage, is_split_kv, combine_key
77
+
78
+
79
+def _get_warmup_probes(
80
+ config: KimiK3VisionFA4WarmupConfig,
81
+ *,
82
+ num_sms: int,
83
+) -> tuple[_FA4WarmupProbe, ...]:
84
+ if config.max_batch_size <= 0 or config.max_seqlen <= 0:
85
+ return ()
86
+
87
+ probes: dict[tuple[int, bool, int | None], _FA4WarmupProbe] = {}
88
+
89
+ def add_probe(batch_size: int, max_seqlen: int) -> None:
90
+ dispatch_key = _get_dispatch(
91
+ config,
92
+ batch_size=batch_size,
93
+ max_seqlen=max_seqlen,
94
+ num_sms=num_sms,
95
+ )
96
+ if dispatch_key not in probes:
97
+ probes[dispatch_key] = _FA4WarmupProbe(
98
+ batch_size=batch_size,
99
+ max_seqlen=max_seqlen,
100
+ dispatch_key=dispatch_key,
101
+ )
102
+
103
+ add_probe(batch_size=1, max_seqlen=1)
104
+ if config.max_seqlen > _FA4_TILE_SIZE:
105
+ add_probe(batch_size=1, max_seqlen=_FA4_TILE_SIZE + 1)
106
+
107
+ max_n_blocks = math.ceil(config.max_seqlen / _FA4_TILE_SIZE)
108
+ for num_n_blocks in range(5, max_n_blocks + 1):
109
+ max_seqlen = min(
110
+ (num_n_blocks - 1) * _FA4_TILE_SIZE + 1,
111
+ config.max_seqlen,
112
+ )
113
+ num_m_blocks = math.ceil(max_seqlen / (2 * _FA4_TILE_SIZE))
114
+ max_split_batch_size = min(
115
+ config.max_batch_size,
116
+ num_sms // (2 * config.num_heads * num_m_blocks),
117
+ )
118
+ if max_split_batch_size == 0:
119
+ break
120
+ for batch_size in range(1, max_split_batch_size + 1):
121
+ add_probe(batch_size, max_seqlen)
122
+
123
+ return tuple(probes.values())
124
+
125
+
126
+def _iter_compile_units(
127
+ config: KimiK3VisionFA4WarmupConfig,
128
+) -> Iterator[CuTeDSLCompileUnit]:
129
+ device_id = current_platform.current_device()
130
+ num_sms = current_platform.num_compute_units(device_id)
131
+ for probe in _get_warmup_probes(config, num_sms=num_sms):
132
+ yield CuTeDSLCompileUnit(
133
+ name="kimi_k3_vision_fa4",
134
+ key=("kimi_k3_vision_fa4", config, probe.dispatch_key),
135
+ compile=partial(_compile, config, probe),
136
+ )
137
+
138
+
139
+def _compile(
140
+ config: KimiK3VisionFA4WarmupConfig,
141
+ probe: _FA4WarmupProbe,
142
+) -> None:
143
+ from vllm.vllm_flash_attn import flash_attn_varlen_func
144
+
145
+ device = current_platform.current_device()
146
+ total_tokens = probe.batch_size * probe.max_seqlen
147
+ qkv = torch.empty(
148
+ 3,
149
+ total_tokens,
150
+ config.num_heads,
151
+ config.head_dim,
152
+ device=device,
153
+ dtype=config.dtype,
154
+ )
155
+ cu_seqlens = torch.arange(
156
+ 0,
157
+ total_tokens + 1,
158
+ probe.max_seqlen,
159
+ device=qkv.device,
160
+ dtype=torch.int32,
161
+ )
162
+ flash_attn_varlen_func(
163
+ qkv[0],
164
+ qkv[1],
165
+ qkv[2],
166
+ cu_seqlens_q=cu_seqlens,
167
+ cu_seqlens_k=cu_seqlens,
168
+ max_seqlen_q=probe.max_seqlen,
169
+ max_seqlen_k=probe.max_seqlen,
170
+ dropout_p=0.0,
171
+ causal=False,
172
+ softmax_scale=config.head_dim**-0.5,
173
+ fa_version=4,
174
+ num_splits=0,
175
+ )
176
+
177
+
178
+class _WarmupProvider:
179
+ def __init__(self) -> None:
180
+ self.configs: set[KimiK3VisionFA4WarmupConfig] = set()
181
+
182
+ def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]:
183
+ return tuple(
184
+ unit for config in self.configs for unit in _iter_compile_units(config)
185
+ )
186
+
187
+
188
+_PROVIDER = _WarmupProvider()
189
+
190
+
191
+def register_kimi_k3_vision_fa4_warmup(
192
+ config: KimiK3VisionFA4WarmupConfig,
193
+) -> None:
194
+ _PROVIDER.configs.add(config)
195
+ register_cutedsl_warmup_provider(_PROVIDER)
vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py
new
+412
@@ -0,0 +1,412 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
4
+#
5
+# This file contains code copied from the flash-linear-attention project.
6
+# The original source was licensed under the MIT license.
7
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
8
+# ruff: noqa: E501
9
+
10
+import torch
11
+import torch.nn as nn
12
+
13
+from vllm.model_executor.custom_op import CustomOp
14
+from vllm.platforms import current_platform
15
+from vllm.triton_utils import tl, triton
16
+from vllm.utils.math_utils import cdiv, next_power_of_2
17
+
18
+@triton.heuristics(
19
+ {
20
+ "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None,
21
+ "HAS_RESIDUAL": lambda args: args["residual"] is not None,
22
+ "HAS_WEIGHT": lambda args: args["w"] is not None,
23
+ "HAS_BIAS": lambda args: args["b"] is not None,
24
+ }
25
+)
26
+@triton.jit
27
+def layer_norm_gated_fwd_kernel(
28
+ x, # pointer to the input
29
+ g, # pointer to the gate
30
+ y, # pointer to the output
31
+ w, # pointer to the weights
32
+ b, # pointer to the biases
33
+ residual, # pointer to the residual
34
+ residual_out, # pointer to the residual
35
+ mean, # pointer to the mean
36
+ rstd, # pointer to the 1/std
37
+ eps, # epsilon to avoid division by zero
38
+ T, # number of rows in x
39
+ H: tl.constexpr, # number of heads
40
+ g_stride_n,
41
+ D: tl.constexpr, # number of columns in x
42
+ BT: tl.constexpr,
43
+ BD: tl.constexpr,
44
+ ACTIVATION: tl.constexpr,
45
+ IS_RMS_NORM: tl.constexpr,
46
+ STORE_RESIDUAL_OUT: tl.constexpr,
47
+ HAS_RESIDUAL: tl.constexpr,
48
+ HAS_WEIGHT: tl.constexpr,
49
+ HAS_BIAS: tl.constexpr,
50
+ launch_pdl: tl.constexpr,
51
+):
52
+ i_t = tl.program_id(0)
53
+
54
+ o_d = tl.arange(0, BD)
55
+ m_d = o_d < D
56
+
57
+ if launch_pdl:
58
+ tl.extra.cuda.gdc_wait()
59
+ tl.extra.cuda.gdc_launch_dependents()
60
+
61
+ p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
62
+ b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32)
63
+ if HAS_RESIDUAL:
64
+ p_res = tl.make_block_ptr(
65
+ residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)
66
+ )
67
+ b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32)
68
+ if STORE_RESIDUAL_OUT:
69
+ p_res_out = tl.make_block_ptr(
70
+ residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)
71
+ )
72
+ tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1))
73
+ if not IS_RMS_NORM:
74
+ b_mean = tl.sum(b_x, axis=1) / D
75
+ p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,))
76
+ tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,))
77
+ b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0)
78
+ b_var = tl.sum(b_xbar * b_xbar, axis=1) / D
79
+ else:
80
+ b_xbar = tl.where(m_d[None, :], b_x, 0.0)
81
+ b_var = tl.sum(b_xbar * b_xbar, axis=1) / D
82
+ b_rstd = 1 / tl.sqrt(b_var + eps)
83
+
84
+ p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,))
85
+ tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,))
86
+
87
+ if HAS_WEIGHT:
88
+ b_w = tl.load(w + o_d, mask=m_d).to(tl.float32)
89
+ if HAS_BIAS:
90
+ b_b = tl.load(b + o_d, mask=m_d).to(tl.float32)
91
+ b_x_hat = (
92
+ (b_x - b_mean[:, None]) * b_rstd[:, None]
93
+ if not IS_RMS_NORM
94
+ else b_x * b_rstd[:, None]
95
+ )
96
+ b_y = b_x_hat * b_w[None, :] if HAS_WEIGHT else b_x_hat
97
+ if HAS_BIAS:
98
+ b_y = b_y + b_b[None, :]
99
+
100
+ # swish/sigmoid output gate
101
+ o_t = i_t * BT + tl.arange(0, BT)
102
+ o_g = (o_t // H) * g_stride_n + (o_t % H) * D
103
+ b_g = tl.load(
104
+ g + o_g[:, None] + o_d[None, :],
105
+ mask=(o_t[:, None] < T) & m_d[None, :],
106
+ other=0.0,
107
+ ).to(tl.float32)
108
+ if ACTIVATION == "swish" or ACTIVATION == "silu":
109
+ b_y = b_y * b_g * tl.sigmoid(b_g)
110
+ elif ACTIVATION == "sigmoid":
111
+ b_y = b_y * tl.sigmoid(b_g)
112
+
113
+ # Write output
114
+ p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
115
+ tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1))
116
+
117
+
118
+@triton.heuristics(
119
+ {
120
+ "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None,
121
+ "HAS_RESIDUAL": lambda args: args["residual"] is not None,
122
+ "HAS_WEIGHT": lambda args: args["w"] is not None,
123
+ "HAS_BIAS": lambda args: args["b"] is not None,
124
+ }
125
+)
126
+@triton.jit
127
+def layer_norm_gated_fwd_kernel1(
128
+ x, # pointer to the input
129
+ g, # pointer to the gate
130
+ y, # pointer to the output
131
+ w, # pointer to the weights
132
+ b, # pointer to the biases
133
+ residual, # pointer to the residual
134
+ residual_out, # pointer to the residual
135
+ mean, # pointer to the mean
136
+ rstd, # pointer to the 1/std
137
+ eps, # epsilon to avoid division by zero
138
+ D: tl.constexpr, # number of columns in x
139
+ BD: tl.constexpr,
140
+ ACTIVATION: tl.constexpr,
141
+ IS_RMS_NORM: tl.constexpr,
142
+ STORE_RESIDUAL_OUT: tl.constexpr,
143
+ HAS_RESIDUAL: tl.constexpr,
144
+ HAS_WEIGHT: tl.constexpr,
145
+ HAS_BIAS: tl.constexpr,
146
+ launch_pdl: tl.constexpr,
147
+):
148
+ i_t = tl.program_id(0)
149
+ x += i_t * D
150
+ y += i_t * D
151
+ g += i_t * D
152
+ if HAS_RESIDUAL:
153
+ residual += i_t * D
154
+ if STORE_RESIDUAL_OUT:
155
+ residual_out += i_t * D
156
+
157
+ if launch_pdl:
158
+ tl.extra.cuda.gdc_wait()
159
+ tl.extra.cuda.gdc_launch_dependents()
160
+
161
+ o_d = tl.arange(0, BD)
162
+ m_d = o_d < D
163
+ b_x = tl.load(x + o_d, mask=m_d, other=0.0).to(tl.float32)
164
+ if HAS_RESIDUAL:
165
+ b_x += tl.load(residual + o_d, mask=m_d, other=0.0).to(tl.float32)
166
+ if STORE_RESIDUAL_OUT:
167
+ tl.store(residual_out + o_d, b_x, mask=m_d)
168
+ if not IS_RMS_NORM:
169
+ b_mean = tl.sum(b_x, axis=0) / D
170
+ tl.store(mean + i_t, b_mean)
171
+ b_xbar = tl.where(m_d, b_x - b_mean, 0.0)
172
+ b_var = tl.sum(b_xbar * b_xbar, axis=0) / D
173
+ else:
174
+ b_xbar = tl.where(m_d, b_x, 0.0)
175
+ b_var = tl.sum(b_xbar * b_xbar, axis=0) / D
176
+ b_rstd = 1 / tl.sqrt(b_var + eps)
177
+ tl.store(rstd + i_t, b_rstd)
178
+
179
+ if HAS_WEIGHT:
180
+ b_w = tl.load(w + o_d, mask=m_d).to(tl.float32)
181
+ if HAS_BIAS:
182
+ b_b = tl.load(b + o_d, mask=m_d).to(tl.float32)
183
+ b_x_hat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd
184
+ b_y = b_x_hat * b_w if HAS_WEIGHT else b_x_hat
185
+ if HAS_BIAS:
186
+ b_y = b_y + b_b
187
+
188
+ # swish/sigmoid output gate
189
+ b_g = tl.load(g + o_d, mask=m_d, other=0.0).to(tl.float32)
190
+ if ACTIVATION == "swish" or ACTIVATION == "silu":
191
+ b_y = b_y * b_g * tl.sigmoid(b_g)
192
+ elif ACTIVATION == "sigmoid":
193
+ b_y = b_y * tl.sigmoid(b_g)
194
+
195
+ # Write output
196
+ tl.store(y + o_d, b_y, mask=m_d)
197
+
198
+
199
+def layer_norm_gated_fwd(
200
+ x: torch.Tensor,
201
+ g: torch.Tensor,
202
+ weight: torch.Tensor,
203
+ bias: torch.Tensor,
204
+ activation: str = "swish",
205
+ eps: float = 1e-5,
206
+ residual: torch.Tensor = None,
207
+ out_dtype: torch.dtype = None,
208
+ residual_dtype: torch.dtype = None,
209
+ is_rms_norm: bool = False,
210
+ H: int = 1,
211
+ g_stride_n: int | None = None,
212
+):
213
+ if residual is not None:
214
+ residual_dtype = residual.dtype
215
+ T, D = x.shape
216
+ if g_stride_n is None:
217
+ g_stride_n = D
218
+ assert T % H == 0
219
+ if residual is not None:
220
+ assert residual.shape == (T, D)
221
+ if weight is not None:
222
+ assert weight.shape == (D,)
223
+ if bias is not None:
224
+ assert bias.shape == (D,)
225
+ # allocate output
226
+ y = x if out_dtype is None else torch.empty_like(x, dtype=out_dtype)
227
+ if residual is not None or (
228
+ residual_dtype is not None and residual_dtype != x.dtype
229
+ ):
230
+ residual_out = torch.empty(T, D, device=x.device, dtype=residual_dtype)
231
+ else:
232
+ residual_out = None
233
+ mean = (
234
+ torch.empty((T,), dtype=torch.float, device=x.device)
235
+ if not is_rms_norm
236
+ else None
237
+ )
238
+ rstd = torch.empty((T,), dtype=torch.float, device=x.device)
239
+ # Less than 64KB per feature: enqueue fused kernel
240
+ MAX_FUSED_SIZE = 65536 // x.element_size()
241
+ BD = min(MAX_FUSED_SIZE, next_power_of_2(D))
242
+ if D > BD:
243
+ raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
244
+ if D <= 512:
245
+ BT = 16
246
+ layer_norm_gated_fwd_kernel[(cdiv(T, BT),)](
247
+ x=x,
248
+ g=g,
249
+ y=y,
250
+ w=weight,
251
+ b=bias,
252
+ residual=residual,
253
+ residual_out=residual_out,
254
+ mean=mean,
255
+ rstd=rstd,
256
+ eps=eps,
257
+ T=T,
258
+ H=H,
259
+ g_stride_n=g_stride_n,
260
+ D=D,
261
+ BD=BD,
262
+ BT=BT,
263
+ ACTIVATION=activation,
264
+ IS_RMS_NORM=is_rms_norm,
265
+ num_warps=8,
266
+ launch_pdl=current_platform.is_arch_support_pdl(),
267
+ )
268
+ else:
269
+ layer_norm_gated_fwd_kernel1[(T,)](
270
+ x=x,
271
+ g=g,
272
+ y=y,
273
+ w=weight,
274
+ b=bias,
275
+ residual=residual,
276
+ residual_out=residual_out,
277
+ mean=mean,
278
+ rstd=rstd,
279
+ eps=eps,
280
+ D=D,
281
+ BD=BD,
282
+ ACTIVATION=activation,
283
+ IS_RMS_NORM=is_rms_norm,
284
+ num_warps=4,
285
+ launch_pdl=current_platform.is_arch_support_pdl(),
286
+ )
287
+ # residual_out is None if residual is None and residual_dtype == input_dtype
288
+ return y, mean, rstd, residual_out if residual_out is not None else x
289
+
290
+
291
+def rms_norm_gated(
292
+ x: torch.Tensor,
293
+ g: torch.Tensor,
294
+ weight: torch.Tensor,
295
+ bias: torch.Tensor,
296
+ activation: str = "swish",
297
+ residual: torch.Tensor | None = None,
298
+ prenorm: bool = False,
299
+ residual_in_fp32: bool = False,
300
+ eps: float = 1e-6,
301
+):
302
+ x_shape_og = x.shape
303
+ # reshape input data into 2D tensor
304
+ x = x.contiguous().reshape(-1, x.shape[-1])
305
+ D = x.shape[-1]
306
+ # The tiled kernel supports row-strided gates; kernel1 does not.
307
+ if D <= 512:
308
+ H = 1 if g.ndim == 2 else g.shape[-2]
309
+ g = g.view(-1, H, D)
310
+ g_stride_n = g.stride(0)
311
+ else:
312
+ g = g.contiguous()
313
+ H = 1
314
+ g_stride_n = D
315
+ if residual is not None:
316
+ assert residual.shape == x_shape_og
317
+ residual = residual.contiguous().reshape(-1, residual.shape[-1])
318
+ residual_dtype = (
319
+ residual.dtype
320
+ if residual is not None
321
+ else (torch.float if residual_in_fp32 else None)
322
+ )
323
+ y, _, _, residual_out = layer_norm_gated_fwd(
324
+ x=x,
325
+ g=g,
326
+ weight=weight,
327
+ bias=bias,
328
+ activation=activation,
329
+ eps=eps,
330
+ residual=residual,
331
+ residual_dtype=residual_dtype,
332
+ is_rms_norm=True,
333
+ H=H,
334
+ g_stride_n=g_stride_n,
335
+ )
336
+ y = y.reshape(x_shape_og)
337
+ return y if not prenorm else (y, residual_out.reshape(x_shape_og))
338
+
339
+
340
+@CustomOp.register("fused_rms_norm_gated")
341
+class FusedRMSNormGated(CustomOp):
342
+ def __init__(
343
+ self,
344
+ hidden_size: int,
345
+ elementwise_affine: bool = True,
346
+ eps: float = 1e-5,
347
+ activation: str = "swish",
348
+ device: torch.device | None = None,
349
+ dtype: torch.dtype | None = None,
350
+ ) -> None:
351
+ factory_kwargs = {"device": device, "dtype": dtype}
352
+ super().__init__()
353
+
354
+ self.hidden_size = hidden_size
355
+ self.elementwise_affine = elementwise_affine
356
+ self.eps = eps
357
+ self.activation = activation
358
+
359
+ if self.activation not in ["swish", "silu", "sigmoid"]:
360
+ raise ValueError(f"Unsupported activation: {self.activation}")
361
+
362
+ if elementwise_affine:
363
+ self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs))
364
+ else:
365
+ self.register_parameter("weight", None)
366
+ self.register_parameter("bias", None)
367
+
368
+ def forward_native(
369
+ self,
370
+ x: torch.Tensor,
371
+ g: torch.Tensor,
372
+ residual: torch.Tensor | None = None,
373
+ prenorm: bool = False,
374
+ residual_in_fp32: bool = False,
375
+ ) -> torch.Tensor:
376
+ """Decomposed PyTorch ops for torch.compile/inductor fusion."""
377
+ # TODO(https://github.com/vllm-project/vllm/issues/36175): implement
378
+ # native residual/prenorm path and unify with RMSNormGated.
379
+ # For now, fall back to the triton kernel.
380
+ if residual is not None or prenorm:
381
+ return self.forward_cuda(x, g, residual, prenorm, residual_in_fp32)
382
+ x_float = x.float()
383
+ variance = x_float.pow(2).mean(dim=-1, keepdim=True)
384
+ x_normed = x_float * torch.rsqrt(variance + self.eps)
385
+ if self.weight is not None:
386
+ x_normed = x_normed * self.weight.float()
387
+ g_float = g.float()
388
+ if self.activation in ("swish", "silu"):
389
+ out = x_normed * g_float * torch.sigmoid(g_float)
390
+ else: # sigmoid
391
+ out = x_normed * torch.sigmoid(g_float)
392
+ return out.to(x.dtype)
393
+
394
+ def forward_cuda(
395
+ self,
396
+ x: torch.Tensor,
397
+ g: torch.Tensor,
398
+ residual: torch.Tensor | None = None,
399
+ prenorm: bool = False,
400
+ residual_in_fp32: bool = False,
401
+ ) -> torch.Tensor:
402
+ return rms_norm_gated(
403
+ x,
404
+ g,
405
+ self.weight,
406
+ self.bias,
407
+ self.activation,
408
+ residual=residual,
409
+ eps=self.eps,
410
+ prenorm=prenorm,
411
+ residual_in_fp32=residual_in_fp32,
412
+ )
vllm/third_party/flash_linear_attention/ops/kda.py
+63
-25
@@ -166,6 +166,8 @@ def layer_norm_gated_fwd_kernel(
166
rstd, # pointer to the 1/std
167
eps, # epsilon to avoid division by zero
168
T, # number of rows in x
169
+ H: tl.constexpr, # number of heads
170
+ g_stride_n: tl.constexpr,
171
D: tl.constexpr, # number of columns in x
172
BT: tl.constexpr,
173
BD: tl.constexpr,
@@ -221,8 +223,13 @@ def layer_norm_gated_fwd_kernel(
223
b_y = b_y + b_b[None, :]
224
225
# swish/sigmoid output gate
224
- p_g = tl.make_block_ptr(g, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
225
- b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)
226
+ o_t = i_t * BT + tl.arange(0, BT)
227
+ o_g = (o_t // H) * g_stride_n + (o_t % H) * D
228
+ b_g = tl.load(
229
+ g + o_g[:, None] + o_d[None, :],
230
+ mask=(o_t[:, None] < T) & m_d[None, :],
231
+ other=0.0,
232
+ ).to(tl.float32)
233
if ACTIVATION == "swish" or ACTIVATION == "silu":
234
b_y = b_y * b_g * tl.sigmoid(b_g)
235
elif ACTIVATION == "sigmoid":
@@ -320,10 +327,15 @@ def layer_norm_gated_fwd(
327
out_dtype: torch.dtype = None,
328
residual_dtype: torch.dtype = None,
329
is_rms_norm: bool = False,
330
+ H: int = 1,
331
+ g_stride_n: int | None = None,
332
):
333
if residual is not None:
334
residual_dtype = residual.dtype
335
T, D = x.shape
336
+ if g_stride_n is None:
337
+ g_stride_n = D
338
+ assert T % H == 0
339
if residual is not None:
340
assert residual.shape == (T, D)
341
if weight is not None:
@@ -349,10 +361,8 @@ def layer_norm_gated_fwd(
361
BD = min(MAX_FUSED_SIZE, next_power_of_2(D))
362
if D > BD:
363
raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
352
- # heuristics for number of warps
353
-
364
if D <= 512:
355
- BT = 32
365
+ BT = 16
366
layer_norm_gated_fwd_kernel[(cdiv(T, BT),)](
367
x=x,
368
g=g,
@@ -365,12 +375,14 @@ def layer_norm_gated_fwd(
375
rstd=rstd,
376
eps=eps,
377
T=T,
378
+ H=H,
379
+ g_stride_n=g_stride_n,
380
D=D,
381
BD=BD,
382
BT=BT,
383
ACTIVATION=activation,
384
IS_RMS_NORM=is_rms_norm,
373
- num_warps=4,
385
+ num_warps=8,
386
)
387
else:
388
layer_norm_gated_fwd_kernel1[(T,)](
@@ -408,7 +420,16 @@ def rms_norm_gated(
420
x_shape_og = x.shape
421
# reshape input data into 2D tensor
422
x = x.contiguous().reshape(-1, x.shape[-1])
411
- g = g.contiguous().reshape(-1, g.shape[-1])
423
+ D = x.shape[-1]
424
+ # The tiled kernel supports row-strided gates; kernel1 does not.
425
+ if D <= 512:
426
+ H = 1 if g.ndim == 2 else g.shape[-2]
427
+ g = g.view(-1, H, D)
428
+ g_stride_n = g.stride(0)
429
+ else:
430
+ g = g.contiguous()
431
+ H = 1
432
+ g_stride_n = D
433
if residual is not None:
434
assert residual.shape == x_shape_og
435
residual = residual.contiguous().reshape(-1, residual.shape[-1])
@@ -427,6 +448,8 @@ def rms_norm_gated(
448
residual=residual,
449
residual_dtype=residual_dtype,
450
is_rms_norm=True,
451
+ H=H,
452
+ g_stride_n=g_stride_n,
453
)
454
y = y.reshape(x_shape_og)
455
return y if not prenorm else (y, residual_out.reshape(x_shape_og))
@@ -1187,6 +1210,7 @@ def kda_gate_cumsum_fwd_kernel(
1210
cu_seqlens,
1211
chunk_indices,
1212
cumsum_scale,
1213
+ lower_bound,
1214
beta,
1215
threshold,
1216
T,
@@ -1196,6 +1220,7 @@ def kda_gate_cumsum_fwd_kernel(
1220
BD: tl.constexpr,
1221
HAS_BIAS: tl.constexpr,
1222
IS_VARLEN: tl.constexpr,
1223
+ USE_LOWER_BOUND: tl.constexpr,
1224
):
1225
i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
1226
i_b, i_h = i_bh // H, i_bh % H
@@ -1235,14 +1260,17 @@ def kda_gate_cumsum_fwd_kernel(
1260
b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32)
1261
b_g = b_g + b_bias[None, :]
1262
1238
- b_a = -tl.exp(tl.load(A + i_h).to(tl.float32))
1239
- b_g_scaled = b_g * beta
1240
- b_softplus = tl.where(
1241
- b_g_scaled > threshold,
1242
- b_g,
1243
- (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)),
1244
- )
1245
- b_gate = b_a * b_softplus
1263
+ b_a = tl.exp(tl.load(A + i_h).to(tl.float32))
1264
+ if USE_LOWER_BOUND:
1265
+ b_gate = lower_bound * tl.sigmoid(b_a * b_g)
1266
+ else:
1267
+ b_g_scaled = b_g * beta
1268
+ b_softplus = tl.where(
1269
+ b_g_scaled > threshold,
1270
+ b_g,
1271
+ (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)),
1272
+ )
1273
+ b_gate = -b_a * b_softplus
1274
1275
# Out-of-bounds rows (load returns 0, but softplus/bias can still make
1276
# b_gate non-zero) participate in the dot product. They only contribute to
@@ -1260,6 +1288,7 @@ def fused_kda_gate_chunk_cumsum(
1288
g_bias: torch.Tensor | None = None,
1289
beta: float = 1.0,
1290
threshold: float = 20.0,
1291
+ lower_bound: float | None = None,
1292
cu_seqlens: torch.Tensor | None = None,
1293
chunk_indices: torch.Tensor | None = None,
1294
chunk_size: int = FLA_CHUNK_SIZE,
@@ -1293,16 +1322,17 @@ def fused_kda_gate_chunk_cumsum(
1322
# exp2-based kernels reproduce exp(g). Keep this in sync with the
1323
# `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`.
1324
cumsum_scale=RCP_LN2,
1325
+ lower_bound=lower_bound or 0.0,
1326
beta=beta,
1327
threshold=threshold,
1328
T=T,
1329
H=H,
1330
D=D,
1331
BT=chunk_size,
1332
+ USE_LOWER_BOUND=lower_bound is not None,
1333
)
1334
return y
1335
1305
-
1336
def _chunk_kda_fwd_with_cumulative_g(
1337
q: torch.Tensor,
1338
k: torch.Tensor,
@@ -1424,6 +1454,7 @@ def chunk_kda_with_fused_gate_fwd(
1454
scale: float,
1455
initial_state: torch.Tensor,
1456
output_final_state: bool,
1457
+ lower_bound: float | None = None,
1458
cu_seqlens: torch.Tensor | None = None,
1459
):
1460
chunk_size = FLA_CHUNK_SIZE
@@ -1439,6 +1470,7 @@ def chunk_kda_with_fused_gate_fwd(
1470
cu_seqlens=cu_seqlens,
1471
chunk_indices=chunk_indices,
1472
chunk_size=chunk_size,
1473
+ lower_bound=lower_bound,
1474
)
1475
return _chunk_kda_fwd_with_cumulative_g(
1476
q=q,
@@ -1500,6 +1532,7 @@ def chunk_kda_with_fused_gate(
1532
scale: float | None = None,
1533
initial_state: torch.Tensor | None = None,
1534
output_final_state: bool = False,
1535
+ lower_bound: float | None = None,
1536
use_qk_l2norm_in_kernel: bool = False,
1537
cu_seqlens: torch.Tensor | None = None,
1538
**kwargs,
@@ -1523,6 +1556,7 @@ def chunk_kda_with_fused_gate(
1556
scale=scale,
1557
initial_state=initial_state.contiguous() if initial_state is not None else None,
1558
output_final_state=output_final_state,
1559
+ lower_bound=lower_bound,
1560
cu_seqlens=cu_seqlens,
1561
)
1562
return o, final_state
@@ -1543,6 +1577,7 @@ def kda_gate_fwd_kernel(
1577
A,
1578
y,
1579
g_bias,
1580
+ lower_bound,
1581
beta: tl.constexpr,
1582
threshold: tl.constexpr,
1583
T,
@@ -1551,12 +1586,12 @@ def kda_gate_fwd_kernel(
1586
BT: tl.constexpr,
1587
BD: tl.constexpr,
1588
HAS_BIAS: tl.constexpr,
1589
+ USE_LOWER_BOUND: tl.constexpr,
1590
):
1591
i_t, i_h = tl.program_id(0), tl.program_id(1)
1592
n_t = i_t * BT
1593
1558
- b_a = tl.load(A + i_h).to(tl.float32)
1559
- b_a = -tl.exp(b_a)
1594
+ b_a = tl.exp(tl.load(A + i_h).to(tl.float32))
1595
1596
stride_row = H * D
1597
stride_col = 1
@@ -1589,13 +1624,13 @@ def kda_gate_fwd_kernel(
1624
)
1625
b_g = b_g + b_bias[None, :]
1626
1592
- # softplus(x, beta) = (1/beta) * log(1 + exp(beta * x))
1593
- # When beta * x > threshold, use linear approximation x
1594
- # Use threshold to switch to linear when beta*x > threshold
1595
- g_scaled = b_g * beta
1596
- use_linear = g_scaled > threshold
1597
- sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled)))
1598
- b_y = b_a * sp
1627
+ if USE_LOWER_BOUND:
1628
+ b_y = lower_bound * tl.sigmoid(b_a * b_g)
1629
+ else:
1630
+ g_scaled = b_g * beta
1631
+ use_linear = g_scaled > threshold
1632
+ sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled)))
1633
+ b_y = -b_a * sp
1634
1635
tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1))
1636
@@ -1607,6 +1642,7 @@ def fused_kda_gate(
1642
g_bias: torch.Tensor | None = None,
1643
beta: float = 1.0,
1644
threshold: float = 20.0,
1645
+ lower_bound: float | None = None,
1646
) -> torch.Tensor:
1647
"""
1648
Forward pass for KDA gate:
@@ -1634,6 +1670,7 @@ def fused_kda_gate(
1670
A,
1671
y,
1672
g_bias,
1673
+ lower_bound or 0.0,
1674
beta,
1675
threshold,
1676
T,
@@ -1641,6 +1678,7 @@ def fused_kda_gate(
1678
head_k_dim,
1679
BD=next_power_of_2(head_k_dim),
1680
HAS_BIAS=g_bias is not None,
1681
+ USE_LOWER_BOUND=lower_bound is not None,
1682
)
1683
1684
y = y.view(*orig_shape, H, head_k_dim)
vllm/transformers_utils/configs/kimi_k3.py
new
+139
@@ -0,0 +1,139 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Kimi-K3 multimodal configuration."""
4
+
5
+from transformers.configuration_utils import PretrainedConfig
6
+
7
+from vllm.logger import init_logger
8
+from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
9
+
10
+logger = init_logger(__name__)
11
+
12
+
13
+class KimiK3VisionConfig(PretrainedConfig):
14
+ model_type = "kimi_k3_vision"
15
+
16
+ def __init__(
17
+ self,
18
+ patch_size: int = 14,
19
+ init_pos_emb_height: int = 64,
20
+ init_pos_emb_width: int = 64,
21
+ init_pos_emb_time: int = 4,
22
+ pos_emb_type: str = "divided_fixed",
23
+ vt_num_attention_heads: int = 12,
24
+ vt_num_hidden_layers: int = 27,
25
+ vt_hidden_size: int = 1024,
26
+ vt_intermediate_size: int = 4096,
27
+ merge_kernel_size: tuple[int, int] = (2, 2),
28
+ video_attn_type: str = "spatial_temporal",
29
+ merge_type: str = "sd2_tpool",
30
+ _attn_implementation: str = "flash_attention_2",
31
+ mm_projector_type: str = "patchmergerv2",
32
+ mm_hidden_size: int | None = None,
33
+ projector_hidden_act: str = "gelu",
34
+ projector_ln_eps: float = 1e-5,
35
+ qkv_hidden_size: int = 1536,
36
+ norm_type: str = "rmsnorm",
37
+ attn_bias: bool = False,
38
+ patch_embed_proj_bias: bool = False,
39
+ mlp_type: str = "mlp2",
40
+ linear_bias: bool = False,
41
+ activation_func: str = "gelu_pytorch_tanh",
42
+ pos_emb_interpolation_mode: str = "bilinear",
43
+ text_hidden_size: int = 2304,
44
+ **kwargs,
45
+ ):
46
+ super().__init__(**kwargs)
47
+
48
+ self.patch_size = patch_size
49
+ self.init_pos_emb_height = init_pos_emb_height
50
+ self.init_pos_emb_width = init_pos_emb_width
51
+ self.init_pos_emb_time = init_pos_emb_time
52
+ self.pos_emb_type = pos_emb_type
53
+ self.vt_num_attention_heads = vt_num_attention_heads
54
+ self.vt_num_hidden_layers = vt_num_hidden_layers
55
+ self.vt_hidden_size = vt_hidden_size
56
+ self.vt_intermediate_size = vt_intermediate_size
57
+ self.merge_kernel_size = tuple(merge_kernel_size)
58
+ self.video_attn_type = video_attn_type
59
+ self.merge_type = merge_type
60
+ self._attn_implementation = _attn_implementation
61
+
62
+ self.mm_projector_type = mm_projector_type
63
+ self.mm_hidden_size = (
64
+ mm_hidden_size if mm_hidden_size is not None else vt_hidden_size
65
+ )
66
+ self.projector_hidden_act = projector_hidden_act
67
+ self.projector_ln_eps = projector_ln_eps
68
+ self.text_hidden_size = text_hidden_size
69
+
70
+ self.qkv_hidden_size = qkv_hidden_size
71
+ self.norm_type = norm_type
72
+ self.attn_bias = attn_bias
73
+ self.patch_embed_proj_bias = patch_embed_proj_bias
74
+ self.mlp_type = mlp_type
75
+ self.linear_bias = linear_bias
76
+ self.activation_func = activation_func
77
+ self.pos_emb_interpolation_mode = pos_emb_interpolation_mode
78
+
79
+ # Aliases consumed by the Kimi-K2.5 vision implementation.
80
+ self.num_attention_heads = vt_num_attention_heads
81
+ self.num_hidden_layers = vt_num_hidden_layers
82
+ self.hidden_size = vt_hidden_size
83
+ self.intermediate_size = vt_intermediate_size
84
+
85
+
86
+class KimiK3Config(PretrainedConfig):
87
+ model_type = "kimi_k3"
88
+
89
+ def __init__(
90
+ self,
91
+ text_config: dict | KimiLinearConfig | None = None,
92
+ vision_config: dict | KimiK3VisionConfig | None = None,
93
+ ignore_index: int = -100,
94
+ media_placeholder_token_id: int = 163605,
95
+ pad_token_id: int = 0,
96
+ image_placeholder: str = "<|kimi_image_placeholder|>",
97
+ **kwargs,
98
+ ):
99
+ if text_config is None:
100
+ self.text_config = KimiLinearConfig()
101
+ elif isinstance(text_config, dict):
102
+ self.text_config = KimiLinearConfig(**text_config)
103
+ else:
104
+ self.text_config = text_config
105
+
106
+ if vision_config is None:
107
+ self.vision_config = KimiK3VisionConfig()
108
+ elif isinstance(vision_config, dict):
109
+ self.vision_config = KimiK3VisionConfig(**vision_config)
110
+ else:
111
+ self.vision_config = vision_config
112
+
113
+ # K3's vision projector output must match the text model's hidden size.
114
+ # Override any value provided in vision_config for safety.
115
+ if self.vision_config.text_hidden_size != self.text_config.hidden_size:
116
+ logger.info(
117
+ "Overriding vision_config.text_hidden_size from %s to %s "
118
+ "to match text_config.hidden_size",
119
+ self.vision_config.text_hidden_size,
120
+ self.text_config.hidden_size,
121
+ )
122
+ self.vision_config.text_hidden_size = self.text_config.hidden_size
123
+
124
+ self.ignore_index = ignore_index
125
+ self.media_placeholder_token_id = media_placeholder_token_id
126
+ self.image_placeholder = image_placeholder
127
+
128
+ if getattr(self.text_config, "quantization_config", None) is not None:
129
+ self.quantization_config = self.text_config.quantization_config
130
+
131
+ super().__init__(pad_token_id=pad_token_id, **kwargs)
132
+
133
+ @property
134
+ def hidden_size(self) -> int:
135
+ return self.text_config.hidden_size
136
+
137
+ @property
138
+ def vocab_size(self) -> int:
139
+ return self.text_config.vocab_size
vllm/transformers_utils/processors/kimi_k3.py
new
+61
@@ -0,0 +1,61 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from transformers import BaseImageProcessor, BatchFeature, TensorType
4
+from transformers.processing_utils import ProcessorMixin
5
+
6
+from vllm.tokenizers.hf import HfTokenizer
7
+
8
+
9
+class KimiK3Processor(ProcessorMixin):
10
+ """HF-style processor wrapper for the image-only Kimi-K3 model.
11
+
12
+ K3 exposes the standard ``image`` modality, so vLLM calls this processor
13
+ with ``images=[PIL, ...]``. The underlying checkpoint image processor
14
+ (``KimiK3VisionProcessor``) works on ``{"type": "image", "image": PIL}``
15
+ media dicts, so this wrapper adapts bare PIL images into that shape before
16
+ delegating to ``preprocess``.
17
+
18
+ Text is only tokenized here; the single ``<|kimi_image_placeholder|>``
19
+ token per image is expanded into the resolution-aware media block by the
20
+ model's ``_get_prompt_updates`` on the vLLM side.
21
+ """
22
+
23
+ attributes = ["image_processor", "tokenizer"]
24
+
25
+ def __init__(
26
+ self,
27
+ image_processor: BaseImageProcessor,
28
+ tokenizer: HfTokenizer,
29
+ ) -> None:
30
+ self.image_processor = image_processor
31
+ self.tokenizer = tokenizer
32
+
33
+ def __call__(
34
+ self,
35
+ text: str | list[str] | None = None,
36
+ images: object | list[object] | None = None,
37
+ return_tensors: str | TensorType | None = None,
38
+ **kwargs,
39
+ ) -> BatchFeature:
40
+ if images is not None:
41
+ if not isinstance(images, list):
42
+ images = [images]
43
+ medias = [{"type": "image", "image": image} for image in images]
44
+ mm_inputs = self.image_processor.preprocess(
45
+ medias,
46
+ return_tensors=return_tensors,
47
+ )
48
+ else:
49
+ mm_inputs = {}
50
+
51
+ if text is not None:
52
+ if not isinstance(text, list):
53
+ text = [text]
54
+ text_inputs = self.tokenizer(text)
55
+ else:
56
+ text_inputs = {}
57
+
58
+ return BatchFeature(
59
+ data={**text_inputs, **mm_inputs},
60
+ tensor_type=return_tensors,
61
+ )
vllm/v1/attention/ops/triton_merge_attn_states.py
+48
-6
@@ -142,6 +142,15 @@ def merge_attn_states(
142
# attention backend.
143
prefix_head_stride = prefix_output.stride(1)
144
output_head_stride = output.stride(1)
145
+ # lse tensors are [NUM_HEADS, NUM_TOKENS] but may be non-contiguous views
146
+ # (e.g. a transpose of a backend's [NUM_TOKENS, NUM_HEADS] output), so index
147
+ # them by their actual strides rather than assuming a contiguous layout.
148
+ prefix_lse_head_stride = prefix_lse.stride(0)
149
+ prefix_lse_token_stride = prefix_lse.stride(1)
150
+ suffix_lse_head_stride = suffix_lse.stride(0)
151
+ suffix_lse_token_stride = suffix_lse.stride(1)
152
+ output_lse_head_stride = output_lse.stride(0) if output_lse is not None else 0
153
+ output_lse_token_stride = output_lse.stride(1) if output_lse is not None else 0
154
155
# If prefill_tokens_with_context is None, all tokens should use prefix context
156
if prefill_tokens_with_context is None:
@@ -157,6 +166,12 @@ def merge_attn_states(
166
suffix_lse,
167
prefix_head_stride,
168
output_head_stride,
169
+ prefix_lse_head_stride,
170
+ prefix_lse_token_stride,
171
+ suffix_lse_head_stride,
172
+ suffix_lse_token_stride,
173
+ output_lse_head_stride,
174
+ output_lse_token_stride,
175
output_scale,
176
prefill_tokens_with_context,
177
head_size,
@@ -176,6 +191,12 @@ def merge_attn_states_kernel(
191
suffix_lse, # [NUM_HEADS, NUM_TOKENS]
192
prefix_head_stride,
193
output_head_stride,
194
+ prefix_lse_head_stride,
195
+ prefix_lse_token_stride,
196
+ suffix_lse_head_stride,
197
+ suffix_lse_token_stride,
198
+ output_lse_head_stride,
199
+ output_lse_token_stride,
200
output_scale, # scale tensor or None
201
prefill_tokens_with_context,
202
HEAD_SIZE: tl.constexpr,
@@ -186,7 +207,6 @@ def merge_attn_states_kernel(
207
FP8_MAX: tl.constexpr = float8_info.max,
208
):
209
token_idx = tl.program_id(0)
189
- num_tokens = tl.num_programs(0)
210
head_idx = tl.program_id(1)
211
num_heads = tl.num_programs(1)
212
@@ -198,9 +218,18 @@ def merge_attn_states_kernel(
218
# For tokens without context (token_idx >= prefill_tokens_with_context),
219
# directly copy from suffix_output
220
if not prefix_mask:
201
- s_lse = tl.load(suffix_lse + head_idx * num_tokens + token_idx)
221
+ s_lse = tl.load(
222
+ suffix_lse
223
+ + head_idx * suffix_lse_head_stride
224
+ + token_idx * suffix_lse_token_stride
225
+ )
226
if OUTPUT_LSE:
203
- tl.store(output_lse + head_idx * num_tokens + token_idx, s_lse)
227
+ tl.store(
228
+ output_lse
229
+ + head_idx * output_lse_head_stride
230
+ + token_idx * output_lse_token_stride,
231
+ s_lse,
232
+ )
233
234
s_out = tl.load(
235
suffix_output
@@ -227,8 +256,16 @@ def merge_attn_states_kernel(
256
257
# For tokens with context (token_idx < prefill_tokens_with_context),
258
# perform normal merge operation
230
- p_lse = tl.load(prefix_lse + head_idx * num_tokens + token_idx)
231
- s_lse = tl.load(suffix_lse + head_idx * num_tokens + token_idx)
259
+ p_lse = tl.load(
260
+ prefix_lse
261
+ + head_idx * prefix_lse_head_stride
262
+ + token_idx * prefix_lse_token_stride
263
+ )
264
+ s_lse = tl.load(
265
+ suffix_lse
266
+ + head_idx * suffix_lse_head_stride
267
+ + token_idx * suffix_lse_token_stride
268
+ )
269
270
# FA2 and FA3 have different behavior for when the sum-exp is 0, this namely
271
# arises with 0 len seqlens. FA3 returns -inf here while FA2 returns inf.
@@ -251,7 +288,12 @@ def merge_attn_states_kernel(
288
# Both sides empty (max_lse == -inf) => undefined merge; keep -inf so
289
# downstream merges continue to treat the token as empty.
290
out_lse = tl.where(max_lse == float("-inf"), float("-inf"), out_lse)
254
- tl.store(output_lse + head_idx * num_tokens + token_idx, out_lse)
291
+ tl.store(
292
+ output_lse
293
+ + head_idx * output_lse_head_stride
294
+ + token_idx * output_lse_token_stride,
295
+ out_lse,
296
+ )
297
298
p_out = tl.load(
299
prefix_output