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

[Kimi-K3][Perf] Make native CUDA AttnRes the SM100 default (#54261)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

Yongye Zhu committed Aug 31, 2026 at 16:18 UTC d6d665854314f0aa90ad6ef32a3382136c71314c
7 files changed +415 -94
benchmarks/kernels/benchmark_kimi_k3_attn_res.py new
+235
@@ -0,0 +1,235 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Benchmark vLLM Kimi-K3 AttnRes kernels on NVIDIA Blackwell.
4 +
5 +The benchmark uses a zero delta so repeated in-place launches remain stable.
6 +
7 +Example:
8 +
9 +.. code-block:: console
10 +
11 + .venv/bin/python benchmarks/kernels/benchmark_kimi_k3_attn_res.py \
12 + --output /tmp/attn_res.json
13 +"""
14 +
15 +from __future__ import annotations
16 +
17 +import argparse
18 +import json
19 +from collections.abc import Callable
20 +from dataclasses import asdict, dataclass
21 +from functools import partial
22 +from pathlib import Path
23 +
24 +import torch
25 +
26 +import vllm._custom_ops # noqa: F401 - loads torch.ops._C
27 +from vllm.models.kimi_k3.nvidia.ops.attn_res import _attn_res_kernel
28 +from vllm.platforms import current_platform
29 +from vllm.triton_utils import triton
30 +
31 +HIDDEN_SIZE = 7168
32 +MAX_BLOCKS = 8
33 +EPS = 1e-5
34 +
35 +
36 +@dataclass
37 +class Inputs:
38 + prefix: torch.Tensor
39 + delta: torch.Tensor
40 + blocks: torch.Tensor
41 + norm_weight: torch.Tensor
42 + qk_weight: torch.Tensor
43 + output_norm_weight: torch.Tensor
44 +
45 +
46 +@dataclass
47 +class Result:
48 + num_tokens: int
49 + num_blocks: int
50 + backend: str
51 + median_us: float
52 + p20_us: float
53 + p80_us: float
54 +
55 +
56 +def parse_args() -> argparse.Namespace:
57 + parser = argparse.ArgumentParser(description=__doc__)
58 + parser.add_argument(
59 + "--num-tokens",
60 + type=int,
61 + nargs="+",
62 + default=[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192],
63 + )
64 + parser.add_argument("--num-blocks", type=int, nargs="+", default=[1, 4, 8])
65 + parser.add_argument("--warmup-ms", type=int, default=50)
66 + parser.add_argument("--repeat-ms", type=int, default=200)
67 + parser.add_argument("--output", type=Path)
68 + return parser.parse_args()
69 +
70 +
71 +def make_inputs(num_tokens: int) -> Inputs:
72 + prefix = torch.randn(num_tokens, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
73 + blocks = torch.randn(
74 + num_tokens,
75 + MAX_BLOCKS,
76 + HIDDEN_SIZE,
77 + device="cuda",
78 + dtype=torch.bfloat16,
79 + )
80 + norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
81 + qk_weight = (
82 + torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5
83 + )
84 + output_norm_weight = 1 + 0.1 * torch.randn(
85 + HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
86 + )
87 + return Inputs(
88 + prefix=prefix,
89 + delta=torch.zeros_like(prefix),
90 + blocks=blocks,
91 + norm_weight=norm_weight,
92 + qk_weight=qk_weight,
93 + output_norm_weight=output_norm_weight,
94 + )
95 +
96 +
97 +def launch_triton(inputs: Inputs, output: torch.Tensor, num_blocks: int) -> None:
98 + num_tokens = inputs.prefix.shape[0]
99 + if num_tokens >= 256 or num_blocks <= 1:
100 + block_l, num_warps = 1, 4
101 + else:
102 + block_l, num_warps = 4, 8
103 + _attn_res_kernel[(num_tokens,)](
104 + inputs.prefix,
105 + inputs.delta,
106 + inputs.blocks,
107 + inputs.norm_weight,
108 + inputs.qk_weight,
109 + inputs.output_norm_weight,
110 + output,
111 + inputs.prefix.stride(0),
112 + inputs.delta.stride(0),
113 + inputs.blocks.stride(0),
114 + inputs.blocks.stride(1),
115 + output.stride(0),
116 + num_blocks,
117 + HIDDEN_SIZE,
118 + -1,
119 + EPS,
120 + EPS,
121 + HAS_DELTA=True,
122 + WRITE_BLOCK=False,
123 + APPLY_OUTPUT_NORM=True,
124 + BLOCK_L=block_l,
125 + BLOCK_D=triton.next_power_of_2(HIDDEN_SIZE),
126 + num_warps=num_warps,
127 + num_stages=2,
128 + launch_pdl=current_platform.is_arch_support_pdl(),
129 + )
130 +
131 +
132 +def launch_cuda(inputs: Inputs, output: torch.Tensor, num_blocks: int) -> None:
133 + torch.ops._C.kimi_k3_attn_res(
134 + inputs.prefix,
135 + inputs.delta,
136 + inputs.blocks,
137 + inputs.norm_weight,
138 + inputs.qk_weight,
139 + inputs.output_norm_weight,
140 + output,
141 + num_blocks,
142 + -1,
143 + EPS,
144 + EPS,
145 + )
146 +
147 +
148 +def benchmark(
149 + function: Callable[[], None], warmup_ms: int, repeat_ms: int
150 +) -> tuple[float, float, float]:
151 + median_ms, p20_ms, p80_ms = triton.testing.do_bench(
152 + function,
153 + warmup=warmup_ms,
154 + rep=repeat_ms,
155 + quantiles=[0.5, 0.2, 0.8],
156 + )
157 + return median_ms * 1000, p20_ms * 1000, p80_ms * 1000
158 +
159 +
160 +def validate(inputs: Inputs, num_blocks: int) -> None:
161 + outputs = [torch.empty_like(inputs.prefix) for _ in range(2)]
162 + launch_triton(inputs, outputs[0], num_blocks)
163 + launch_cuda(inputs, outputs[1], num_blocks)
164 + torch.accelerator.synchronize()
165 + torch.testing.assert_close(outputs[1], outputs[0], atol=8e-2, rtol=3e-2)
166 +
167 +
168 +def main() -> None:
169 + args = parse_args()
170 + if torch.cuda.get_device_capability()[0] != 10:
171 + raise SystemExit("The vLLM CUDA AttnRes kernel requires the SM100 family")
172 + if not hasattr(torch.ops._C, "kimi_k3_attn_res"):
173 + raise SystemExit("This vLLM build does not contain kimi_k3_attn_res")
174 + if any(not 0 <= blocks <= MAX_BLOCKS for blocks in args.num_blocks):
175 + raise SystemExit(f"--num-blocks values must be in [0, {MAX_BLOCKS}]")
176 +
177 + torch.manual_seed(0)
178 + check_inputs = make_inputs(5)
179 + for num_blocks in args.num_blocks:
180 + validate(check_inputs, num_blocks)
181 + print("correctness: all selected block counts passed", flush=True)
182 +
183 + results: list[Result] = []
184 + for num_blocks in args.num_blocks:
185 + for num_tokens in args.num_tokens:
186 + inputs = make_inputs(num_tokens)
187 + output = torch.empty_like(inputs.prefix)
188 + candidates = {
189 + "vllm_triton": partial(launch_triton, inputs, output, num_blocks),
190 + "vllm_cuda": partial(launch_cuda, inputs, output, num_blocks),
191 + }
192 + row: dict[str, float] = {}
193 + for backend, candidate in candidates.items():
194 + median_us, p20_us, p80_us = benchmark(
195 + candidate, args.warmup_ms, args.repeat_ms
196 + )
197 + results.append(
198 + Result(
199 + num_tokens=num_tokens,
200 + num_blocks=num_blocks,
201 + backend=backend,
202 + median_us=median_us,
203 + p20_us=p20_us,
204 + p80_us=p80_us,
205 + )
206 + )
207 + row[backend] = median_us
208 + print(
209 + f"blocks={num_blocks} tokens={num_tokens:5d} "
210 + f"triton={row['vllm_triton']:9.2f} us "
211 + f"cuda={row['vllm_cuda']:9.2f} us",
212 + flush=True,
213 + )
214 +
215 + if args.output is not None:
216 + metadata = {
217 + "device": torch.cuda.get_device_name(),
218 + "capability": torch.cuda.get_device_capability(),
219 + "torch": torch.__version__,
220 + "cuda": torch.version.cuda,
221 + "warmup_ms": args.warmup_ms,
222 + "repeat_ms": args.repeat_ms,
223 + }
224 + args.output.write_text(
225 + json.dumps(
226 + {"metadata": metadata, "results": [asdict(item) for item in results]},
227 + indent=2,
228 + )
229 + + "\n"
230 + )
231 + print(f"wrote {args.output}")
232 +
233 +
234 +if __name__ == "__main__":
235 + main()
csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu
+104 -77
@@ -10,7 +10,7 @@
10 // - Q=res_weight*rms_weight remains in registers across persistent tokens.
11 // - V rows are converted once and cached as FP32 in TMEM between passes.
12 //
13 -// Integration contract: Kimi K3 H=7168, 1<=num_blocks<=8, and token-major
13 +// Integration contract: Kimi K3 H=7168, 0<=num_blocks<=8, and token-major
14 // block residual storage.
15
16 #include "../torch_utils.h"
@@ -230,14 +230,15 @@ __device__ __forceinline__ void cp_async_bulk(void* smem_dst,
230 }
231
232 template <int H, int N, int NC = N_CHUNK_DEFAULT, int B = 1,
233 - bool RELEASE_TMEM = false, bool HAS_DELTA = false,
234 - bool HAS_OUTPUT_NORM = false, bool OUTPUT_NORM_IN_SMEM = false>
233 + bool RELEASE_TMEM = false, bool STATIC_COMMON_PATH = false,
234 + bool HAS_OUTPUT_NORM = false>
235 __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
236 - const bf16_t* __restrict__ block_res, bf16_t* __restrict__ layer_res,
236 + bf16_t* __restrict__ block_res, bf16_t* __restrict__ layer_res,
237 const bf16_t* __restrict__ delta, const bf16_t* __restrict__ res_w,
238 const bf16_t* __restrict__ rms_w, bf16_t* __restrict__ output, int T,
239 int block_stride_m, int block_stride_r, float rms_eps,
240 - const bf16_t* __restrict__ output_norm_weight, float output_norm_eps) {
240 + const bf16_t* __restrict__ output_norm_weight, float output_norm_eps,
241 + int block_write_idx) {
242 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100
243 constexpr float LOG2_E = 1.4426950408889634f;
244 constexpr int N_CHUNK = NC;
@@ -254,12 +255,15 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
255 static_assert(TMEM_V_COLS_TOTAL <= TMEM_COLS_ALLOC);
256 static_assert(H >= 4096 && H <= 8192);
257 static_assert(H % K_TILE == 0);
258 + static_assert(!STATIC_COMMON_PATH || HAS_OUTPUT_NORM);
259
260 const int tid = threadIdx.x;
261 const int wid = tid >> 5;
262 const int lane = tid & 31;
263 const int TB = T * B;
264 const int num_ctas = gridDim.x;
265 + const bool has_delta = STATIC_COMMON_PATH || delta != nullptr;
266 + const int write_idx = STATIC_COMMON_PATH ? -1 : block_write_idx;
267 constexpr int num_chunks = (N + N_CHUNK - 1) / N_CHUNK;
268
269 const int comp_wid = wid - 1;
@@ -270,10 +274,9 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
274 const int k_local = ct_in_group * VEC;
275
276 constexpr size_t V_BYTES = (size_t)NUM_BUFS * H * sizeof(bf16_t);
273 - constexpr size_t DELTA_BYTES =
274 - HAS_DELTA ? (size_t)CHUNK_DEPTH * H * sizeof(bf16_t) : 0;
277 + constexpr size_t DELTA_BYTES = (size_t)CHUNK_DEPTH * H * sizeof(bf16_t);
278 constexpr size_t OUTPUT_NORM_BYTES =
276 - OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0;
279 + HAS_OUTPUT_NORM ? (size_t)H * sizeof(bf16_t) : 0;
280 extern __shared__ __align__(16) char smem_raw[];
281 bf16_t* v_bufs = reinterpret_cast<bf16_t*>(smem_raw); // [NUM_BUFS][H]
282 bf16_t* delta_bufs = reinterpret_cast<bf16_t*>(smem_raw + V_BYTES);
@@ -297,15 +300,12 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
300 mbarrier_init(plan.bar_ready[i], 1);
301 mbarrier_init(plan.bar_consumed[i], CONSUMER_WARPS);
302 }
300 - if constexpr (OUTPUT_NORM_IN_SMEM) {
303 + if constexpr (HAS_OUTPUT_NORM) {
304 mbarrier_init(plan.bar_output_norm_ready, 1);
305 }
306 fence_mbarrier_init();
307 }
308
306 - // gdc wait BEFORE tmem alloc
307 - cudaGridDependencySynchronize();
308 -
309 if (wid == 1) {
310 tmem_allocate(TMEM_COLS_ALLOC, &plan.tmem_base);
311 if constexpr (RELEASE_TMEM) {
@@ -314,7 +314,7 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
314 }
315 __syncthreads();
316
317 - if constexpr (OUTPUT_NORM_IN_SMEM) {
317 + if constexpr (HAS_OUTPUT_NORM) {
318 if (wid == 0 && elect_one_sync()) {
319 mbarrier_expect_tx(plan.bar_output_norm_ready, H * (int)sizeof(bf16_t));
320 cp_async_bulk(output_norm_buf, output_norm_weight, H * sizeof(bf16_t),
@@ -365,6 +365,7 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
365 if (wid == 0) {
366 if (elect_one_sync()) {
367 long long gci = 0;
368 + bool pdl_waited = false;
369 for (int tb = blockIdx.x; tb < TB; tb += num_ctas) {
370 const int t = tb / B;
371 for (int ci = 0; ci < num_chunks; ci++, gci++) {
@@ -374,7 +375,7 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
375 int pc = phase_of(gci);
376 mbarrier_wait(plan.bar_consumed[chunk_slot], pc ^ 1);
377 int transaction_bytes = an * H * (int)sizeof(bf16_t);
377 - if constexpr (HAS_DELTA) {
378 + if (has_delta) {
379 int prefix_n = N - 1 - ns;
380 if (prefix_n >= 0 && prefix_n < an) {
381 transaction_bytes += H * (int)sizeof(bf16_t);
@@ -388,12 +389,20 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
389 const bf16_t* src =
390 residual_addr(block_res, layer_res, ns + n, N, t,
391 block_stride_m, block_stride_r, H);
392 + if (!has_delta && ns + n == N - 1 && !pdl_waited) {
393 + cudaGridDependencySynchronize();
394 + pdl_waited = true;
395 + }
396 cp_async_bulk(buf_ptr(slot), src, H * sizeof(bf16_t),
397 plan.bar_ready[chunk_slot]);
398 }
394 - if constexpr (HAS_DELTA) {
399 + if (has_delta) {
400 int prefix_n = N - 1 - ns;
401 if (prefix_n >= 0 && prefix_n < an) {
402 + if (!pdl_waited) {
403 + cudaGridDependencySynchronize();
404 + pdl_waited = true;
405 + }
406 cp_async_bulk(delta_buf_ptr(chunk_slot),
407 delta + (long long)tb * H, H * sizeof(bf16_t),
408 plan.bar_ready[chunk_slot]);
@@ -401,6 +410,7 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
410 }
411 }
412 }
413 + cudaTriggerProgrammaticLaunchCompletion();
414 }
415 } else {
416 float acc32[ACC_PER_THREAD] = {};
@@ -442,9 +452,9 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
452 int2 vp =
453 *reinterpret_cast<const int2*>(buf_ptr(slot) + h_base);
454 auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp);
445 - if constexpr (HAS_DELTA) {
446 - int prefix_n = N - 1 - ns;
447 - if (n == prefix_n) {
455 + int prefix_n = N - 1 - ns;
456 + if (n == prefix_n && (has_delta || write_idx >= 0)) {
457 + if (has_delta) {
458 const bf16_t* delta_ptr =
459 delta_buf_ptr(chunk_slot) + h_base;
460 #pragma unroll
@@ -456,6 +466,11 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
466 *reinterpret_cast<int2*>(layer_res + (long long)tb * H +
467 h_base) = vp;
468 }
469 + if (write_idx >= 0) {
470 + *reinterpret_cast<int2*>(
471 + block_res + (long long)(tb / B) * block_stride_m +
472 + write_idx * block_stride_r + h_base) = vp;
473 + }
474 }
475 float2 f[2] = {__bfloat1622float2(v2[0]),
476 __bfloat1622float2(v2[1])};
@@ -480,9 +495,9 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
495 int slot = slot_of(gci, n);
496 int4 vp = *reinterpret_cast<const int4*>(buf_ptr(slot) + h_base);
497 auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp);
483 - if constexpr (HAS_DELTA) {
484 - int prefix_n = N - 1 - ns;
485 - if (n == prefix_n) {
498 + int prefix_n = N - 1 - ns;
499 + if (n == prefix_n && (has_delta || write_idx >= 0)) {
500 + if (has_delta) {
501 const bf16_t* delta_ptr = delta_buf_ptr(chunk_slot) + h_base;
502 #pragma unroll
503 for (int j = 0; j < VEC / 2; j++) {
@@ -493,6 +508,11 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
508 *reinterpret_cast<int4*>(layer_res + (long long)tb * H +
509 h_base) = vp;
510 }
511 + if (write_idx >= 0) {
512 + *reinterpret_cast<int4*>(
513 + block_res + (long long)(tb / B) * block_stride_m +
514 + write_idx * block_stride_r + h_base) = vp;
515 + }
516 }
517 float2 f[4] = {
518 __bfloat1622float2(v2[0]), __bfloat1622float2(v2[1]),
@@ -791,11 +811,9 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
811 }
812
813 if constexpr (HAS_OUTPUT_NORM) {
794 - if constexpr (OUTPUT_NORM_IN_SMEM) {
795 - // The immutable weight copy is acquired once, at its first use.
796 - if (tb == blockIdx.x) {
797 - mbarrier_wait(plan.bar_output_norm_ready, 0);
798 - }
814 + // The immutable weight copy is acquired once, at its first use.
815 + if (tb == blockIdx.x) {
816 + mbarrier_wait(plan.bar_output_norm_ready, 0);
817 }
818 float output_sq = output_sq_pair.x + output_sq_pair.y;
819 #pragma unroll
@@ -826,9 +844,7 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
844 auto* values = reinterpret_cast<bf16_t*>(&packed);
845 #pragma unroll
846 for (int j = 0; j < 4; j++) {
829 - const bf16_t* weight_ptr =
830 - OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight;
831 - float weight = __bfloat162float(weight_ptr[h_base + j]);
847 + float weight = __bfloat162float(output_norm_buf[h_base + j]);
848 values[j] = __float2bfloat16(acc32[si * VEC + j] *
849 output_rsigma * weight);
850 }
@@ -843,9 +859,7 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
859 auto* values = reinterpret_cast<bf16_t*>(&packed);
860 #pragma unroll
861 for (int j = 0; j < VEC; j++) {
846 - const bf16_t* weight_ptr =
847 - OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight;
848 - float weight = __bfloat162float(weight_ptr[h_base + j]);
862 + float weight = __bfloat162float(output_norm_buf[h_base + j]);
863 values[j] =
864 __float2bfloat16(acc32[si * VEC + j] * output_rsigma * weight);
865 }
@@ -855,7 +869,6 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
869 }
870 }
871
858 - cudaTriggerProgrammaticLaunchCompletion();
872 __syncthreads();
873 if (wid == 1) {
874 tmem_free(plan.tmem_base, TMEM_COLS_ALLOC);
@@ -868,23 +881,22 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
881 }
882
883 template <int H, int N, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false,
871 - bool HAS_DELTA = false, bool HAS_OUTPUT_NORM = false,
872 - bool OUTPUT_NORM_IN_SMEM = false>
873 -static void launch_fwd(const bf16_t* block_residual, bf16_t* layer_residual,
884 + bool STATIC_COMMON_PATH = false, bool HAS_OUTPUT_NORM = false>
885 +static void launch_fwd(bf16_t* block_residual, bf16_t* layer_residual,
886 const bf16_t* delta, const bf16_t* res_weight,
887 const bf16_t* rms_weight, bf16_t* output, int T,
888 float rms_eps, int num_sm, cudaStream_t stream,
889 const bf16_t* output_norm_weight = nullptr,
890 float output_norm_eps = 0.f, int block_stride_m = 0,
879 - int block_stride_r = 0) {
891 + int block_stride_r = 0, int block_write_idx = -1) {
892 constexpr size_t smem_size =
881 - ((size_t)CHUNK_DEPTH * (NC + (HAS_DELTA ? 1 : 0)) * H * sizeof(bf16_t) +
882 - (OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0) +
893 + ((size_t)(CHUNK_DEPTH * (NC + 1) + (HAS_OUTPUT_NORM ? 1 : 0)) * H *
894 + sizeof(bf16_t) +
895 sizeof(FwdSmemPlan<NC>) + 15) &
896 ~size_t(15);
897 auto kernel =
886 - &attn_res_fwd_online_v2_kernel<H, N, NC, 1, RELEASE_TMEM, HAS_DELTA,
887 - HAS_OUTPUT_NORM, OUTPUT_NORM_IN_SMEM>;
898 + &attn_res_fwd_online_v2_kernel<H, N, NC, 1, RELEASE_TMEM,
899 + STATIC_COMMON_PATH, HAS_OUTPUT_NORM>;
900 static bool attrs_set = false;
901 if (!attrs_set) {
902 if (smem_size > 48 * 1024) {
@@ -907,20 +919,21 @@ static void launch_fwd(const bf16_t* block_residual, bf16_t* layer_residual,
919 cudaLaunchKernelEx(&config, kernel, block_residual, layer_residual, delta,
920 res_weight, rms_weight, output, T, block_stride_m,
921 block_stride_r, rms_eps, output_norm_weight,
910 - output_norm_eps);
922 + output_norm_eps, block_write_idx);
923 }
924
925 } // namespace fwd_prod_v2
926 } // namespace sm100
927
928 void kimi_k3_attn_res(torch::stable::Tensor& prefix,
917 - torch::stable::Tensor const& delta,
918 - torch::stable::Tensor const& blocks,
929 + std::optional<torch::stable::Tensor> delta,
930 + torch::stable::Tensor& blocks,
931 torch::stable::Tensor const& norm_weight,
932 torch::stable::Tensor const& qk_weight,
921 - torch::stable::Tensor const& output_norm_weight,
933 + std::optional<torch::stable::Tensor> output_norm_weight,
934 torch::stable::Tensor& output, int64_t num_blocks,
923 - double eps, double output_norm_eps) {
935 + int64_t block_write_idx, double eps,
936 + double output_norm_eps) {
937 int const num_tokens = static_cast<int>(prefix.size(0));
938 int const device = prefix.get_device_index();
939 torch::stable::accelerator::DeviceGuard const device_guard(device);
@@ -928,49 +941,57 @@ void kimi_k3_attn_res(torch::stable::Tensor& prefix,
941 STD_TORCH_CHECK(properties->major == 10,
942 "Kimi K3 AttnRes requires the SM100 family");
943 int64_t const hidden_size = prefix.size(1);
944 + STD_TORCH_CHECK(hidden_size == 7168,
945 + "Kimi K3 AttnRes requires hidden_size=7168, got ",
946 + hidden_size);
947 STD_TORCH_CHECK(prefix.stride(0) == hidden_size &&
932 - delta.stride(0) == hidden_size &&
948 + (!delta.has_value() || delta->stride(0) == hidden_size) &&
949 output.stride(0) == hidden_size,
950 "Kimi K3 AttnRes requires densely packed rows; got strides ",
935 - prefix.stride(0), ", ", delta.stride(0), ", ",
951 + prefix.stride(0), ", ",
952 + delta.has_value() ? delta->stride(0) : 0, ", ",
953 output.stride(0), " for hidden_size ", hidden_size);
954 + STD_TORCH_CHECK(
955 + num_blocks >= 0 && num_blocks <= 8 && num_blocks <= blocks.size(1),
956 + "Kimi K3 AttnRes: num_blocks must be 0..min(8, blocks.size(1))");
957 + STD_TORCH_CHECK(block_write_idx >= -1 && block_write_idx < blocks.size(1),
958 + "Kimi K3 AttnRes: block_write_idx must be -1 or a valid "
959 + "block row, got ",
960 + block_write_idx);
961 +
962 + if (num_tokens == 0) {
963 + return;
964 + }
965
966 using namespace sm100::fwd_prod_v2;
939 - // Two-source chunks and two resident CTAs are beneficial once setup is
940 - // amortized by the long, full eight-block prefill workload.
941 - if (num_blocks == 8 && num_tokens >= 4096) {
942 - launch_fwd<7168, 9, 2, true, true, true, true>(
943 - static_cast<bf16_t const*>(blocks.data_ptr()),
944 - static_cast<bf16_t*>(prefix.data_ptr()),
945 - static_cast<bf16_t const*>(delta.data_ptr()),
946 - static_cast<bf16_t const*>(qk_weight.data_ptr()),
947 - static_cast<bf16_t const*>(norm_weight.data_ptr()),
948 - static_cast<bf16_t*>(output.data_ptr()), num_tokens,
949 - static_cast<float>(eps), properties->multiProcessorCount,
950 - get_current_cuda_stream(device),
951 - static_cast<bf16_t const*>(output_norm_weight.data_ptr()),
952 - static_cast<float>(output_norm_eps), static_cast<int>(blocks.stride(0)),
953 - static_cast<int>(blocks.stride(1)));
954 - } else {
955 - auto dispatch = [&](auto nsrc_tok) {
956 - constexpr int NSRC = decltype(nsrc_tok)::value;
957 - launch_fwd<7168, NSRC, 3, false, true, true, true>(
958 - static_cast<bf16_t const*>(blocks.data_ptr()),
967 + cudaStream_t stream = get_current_cuda_stream(device);
968 + auto dispatch_blocks = [&](auto static_common_path, auto has_output_norm) {
969 + constexpr bool STATIC_COMMON_PATH = decltype(static_common_path)::value;
970 + constexpr bool HAS_OUTPUT_NORM = decltype(has_output_norm)::value;
971 + auto dispatch = [&](auto source_count) {
972 + constexpr int NUM_SOURCES = decltype(source_count)::value;
973 + launch_fwd<7168, NUM_SOURCES, 3, false, STATIC_COMMON_PATH,
974 + HAS_OUTPUT_NORM>(
975 + static_cast<bf16_t*>(blocks.data_ptr()),
976 static_cast<bf16_t*>(prefix.data_ptr()),
960 - static_cast<bf16_t const*>(delta.data_ptr()),
977 + delta.has_value() ? static_cast<bf16_t const*>(delta->data_ptr())
978 + : nullptr,
979 static_cast<bf16_t const*>(qk_weight.data_ptr()),
980 static_cast<bf16_t const*>(norm_weight.data_ptr()),
981 static_cast<bf16_t*>(output.data_ptr()), num_tokens,
964 - static_cast<float>(eps), properties->multiProcessorCount,
965 - get_current_cuda_stream(device),
966 - static_cast<bf16_t const*>(output_norm_weight.data_ptr()),
982 + static_cast<float>(eps), properties->multiProcessorCount, stream,
983 + output_norm_weight.has_value()
984 + ? static_cast<bf16_t const*>(output_norm_weight->data_ptr())
985 + : nullptr,
986 static_cast<float>(output_norm_eps),
987 static_cast<int>(blocks.stride(0)),
969 - static_cast<int>(blocks.stride(1)));
988 + static_cast<int>(blocks.stride(1)),
989 + static_cast<int>(block_write_idx));
990 };
971 - // The source count is num_blocks + 1; specialising on it makes the chunk
972 - // count, per-chunk source count and prefix-chunk test compile-time.
991 switch (num_blocks) {
992 + case 0:
993 + dispatch(std::integral_constant<int, 1>{});
994 + break;
995 case 1:
996 dispatch(std::integral_constant<int, 2>{});
997 break;
@@ -995,9 +1016,15 @@ void kimi_k3_attn_res(torch::stable::Tensor& prefix,
1016 case 8:
1017 dispatch(std::integral_constant<int, 9>{});
1018 break;
998 - default:
999 - STD_TORCH_CHECK(false, "Kimi K3 AttnRes: num_blocks must be 1..8");
1019 }
1020 + };
1021 + if (delta.has_value() && output_norm_weight.has_value() &&
1022 + block_write_idx < 0) {
1023 + dispatch_blocks(std::true_type{}, std::true_type{});
1024 + } else if (output_norm_weight.has_value()) {
1025 + dispatch_blocks(std::false_type{}, std::true_type{});
1026 + } else {
1027 + dispatch_blocks(std::false_type{}, std::false_type{});
1028 }
1029 cudaError_t const error = cudaGetLastError();
1030 STD_TORCH_CHECK(
csrc/libtorch_stable/ops.h
+5 -4
@@ -411,13 +411,14 @@ void fused_gdn_decode_post_conv_mtp(
411
412 #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
413 void kimi_k3_attn_res(torch::stable::Tensor& prefix,
414 - torch::stable::Tensor const& delta,
415 - torch::stable::Tensor const& blocks,
414 + std::optional<torch::stable::Tensor> delta,
415 + torch::stable::Tensor& blocks,
416 torch::stable::Tensor const& norm_weight,
417 torch::stable::Tensor const& qk_weight,
418 - torch::stable::Tensor const& output_norm_weight,
418 + std::optional<torch::stable::Tensor> output_norm_weight,
419 torch::stable::Tensor& output, int64_t num_blocks,
420 - double eps, double output_norm_eps);
420 + int64_t block_write_idx, double eps,
421 + double output_norm_eps);
422 #endif
423
424 // Sampler kernels (shared CUDA/ROCm)
csrc/libtorch_stable/torch_bindings.cpp
+5 -3
@@ -542,9 +542,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
542 #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
543 ops.def(
544 "kimi_k3_attn_res("
545 - "Tensor! prefix, Tensor delta, Tensor blocks, Tensor norm_weight, "
546 - "Tensor qk_weight, Tensor output_norm_weight, Tensor! output, "
547 - "int num_blocks, float eps, float output_norm_eps) -> ()");
545 + "Tensor! prefix, Tensor? delta, Tensor! blocks, Tensor norm_weight, "
546 + "Tensor qk_weight, Tensor? output_norm_weight, Tensor! output, "
547 + "int num_blocks, "
548 + "int block_write_idx, float eps, "
549 + "float output_norm_eps) -> ()");
550 #endif
551
552 // Apply repetition penalties to logits in-place.
tests/models/kimi_k3/test_attn_res.py
+53
@@ -1,6 +1,8 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import importlib
5 +
6 import pytest
7 import torch
8 import torch.nn.functional as F
@@ -12,6 +14,7 @@ from vllm.platforms import current_platform
14 HIDDEN_SIZE = 7168
15 MAX_BLOCKS = 8
16 EPS = 1e-5
17 +attn_res_module = importlib.import_module("vllm.models.kimi_k3.nvidia.ops.attn_res")
18
19
20 def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor:
@@ -60,6 +63,7 @@ def _reference(
63 pytest.param(17, 5, 7, True, True, "triton", id="triton-write-add"),
64 pytest.param(3, 8, 0, False, False, "triton", id="triton-full"),
65 pytest.param(3, 8, 0, False, True, "triton", id="triton-full-add"),
66 + pytest.param(1, 4, 0, False, True, "nvidia", id="nvidia-single-token"),
67 pytest.param(320, 1, 0, False, True, "nvidia", id="nvidia-1"),
68 pytest.param(320, 4, 0, False, True, "nvidia", id="nvidia-4"),
69 pytest.param(320, 8, 0, False, True, "nvidia", id="nvidia-8"),
@@ -194,6 +198,55 @@ def test_attn_res_without_output_norm():
198 torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
199
200
201 +@pytest.mark.parametrize(
202 + ("num_blocks", "has_delta", "has_output_norm", "write_block"),
203 + [
204 + (0, False, True, True),
205 + (4, True, True, True),
206 + (4, False, True, False),
207 + (8, True, False, False),
208 + ],
209 +)
210 +def test_sm100_variants_do_not_fall_back_to_triton(
211 + monkeypatch,
212 + num_blocks: int,
213 + has_delta: bool,
214 + has_output_norm: bool,
215 + write_block: bool,
216 +):
217 + if not current_platform.is_device_capability_family(100) or not hasattr(
218 + torch.ops._C, "kimi_k3_attn_res"
219 + ):
220 + pytest.skip("native AttnRes requires an SM100 build")
221 +
222 + class FailingKernel:
223 + def __getitem__(self, grid):
224 + raise AssertionError(f"unexpected Triton launch with grid {grid}")
225 +
226 + monkeypatch.setattr(attn_res_module, "_attn_res_kernel", FailingKernel())
227 + prefix = torch.randn(1, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
228 + delta = torch.zeros_like(prefix) if has_delta else None
229 + blocks = torch.randn(
230 + 1, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
231 + )
232 + norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
233 + qk_weight = torch.randn_like(norm_weight) / HIDDEN_SIZE**0.5
234 + output_norm_weight = torch.ones_like(norm_weight) if has_output_norm else None
235 +
236 + attn_res(
237 + prefix,
238 + delta,
239 + blocks,
240 + norm_weight,
241 + qk_weight,
242 + output_norm_weight,
243 + num_blocks,
244 + num_blocks if write_block else -1,
245 + EPS,
246 + EPS,
247 + )
248 +
249 +
250 @pytest.mark.parametrize("num_tokens", [0, 1, 17])
251 def test_fused_mtp_input(num_tokens: int):
252 positions = torch.arange(num_tokens, device="cuda")
vllm/_custom_ops.py
+4 -2
@@ -2904,12 +2904,13 @@ def concat_and_cache_mla_grouped(
2904
2905 def kimi_k3_attn_res(
2906 prefix: torch.Tensor,
2907 - delta: torch.Tensor,
2907 + delta: torch.Tensor | None,
2908 blocks: torch.Tensor,
2909 norm_weight: torch.Tensor,
2910 qk_weight: torch.Tensor,
2911 - output_norm_weight: torch.Tensor,
2911 + output_norm_weight: torch.Tensor | None,
2912 num_blocks: int,
2913 + block_write_idx: int,
2914 eps: float,
2915 output_norm_eps: float,
2916 ) -> torch.Tensor:
@@ -2923,6 +2924,7 @@ def kimi_k3_attn_res(
2924 output_norm_weight,
2925 output,
2926 num_blocks,
2927 + block_write_idx,
2928 eps,
2929 output_norm_eps,
2930 )
vllm/models/kimi_k3/nvidia/ops/attn_res.py
+9 -8
@@ -186,16 +186,16 @@ def attn_res(
186 assert norm_weight.stride(-1) == 1
187 assert qk_weight.stride(-1) == 1
188 assert output_norm_weight is None or output_norm_weight.stride(-1) == 1
189 - # The in-tree NVIDIA kernel covers the common fused-add + output-norm path;
190 - # Triton handles block boundaries and final pre-norm output. The native op
191 - # is only compiled for SM100 under CUDA >= 13, so a device check alone is
192 - # not enough to know it exists.
189 + # The native kernel covers every Kimi-K3 AttnRes variant on dense SM100
190 + # inputs. The op is only compiled under CUDA >= 13, so a device check alone
191 + # is not enough to know it exists.
192 if (
193 hidden_size == 7168
195 - and delta is not None
196 - and output_norm_weight is not None
197 - and num_blocks > 0
198 - and block_write_idx < 0
194 + and prefix.stride(0) == hidden_size
195 + and (delta is None or delta.stride(0) == hidden_size)
196 + and 0 <= num_blocks <= 8
197 + and num_blocks <= blocks.shape[1]
198 + and -1 <= block_write_idx < blocks.shape[1]
199 and current_platform.is_device_capability_family(100)
200 and hasattr(torch.ops._C, "kimi_k3_attn_res")
201 ):
@@ -207,6 +207,7 @@ def attn_res(
207 qk_weight,
208 output_norm_weight,
209 num_blocks,
210 + block_write_idx,
211 eps,
212 output_norm_eps,
213 )