@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 61ac368021

[Kimi-K3] Add AttnRes kernels (#50090)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

Thien Tran committed Jul 28, 2026 at 15:22 UTC 61ac368021033dea54448d36ed98c7426e82cd18
14 files changed +1719
.buildkite/test_areas/models_basic.yaml
+12
@@ -61,6 +61,18 @@ steps:
61 # FA4 kernel tests require SM100; the suite skips them elsewhere.
62 - pytest -v -s models/inkling
63
64 +- label: Kimi K3 Unit Tests (B200)
65 + key: kimi-k3-unit-tests-b200
66 + timeout_in_minutes: 40
67 + device: b200-k8s
68 + source_file_dependencies:
69 + - vllm/models/kimi_k3/
70 + - csrc/libtorch_stable/kimi_k3/
71 + - tests/models/kimi_k3/
72 + commands:
73 + # The native NVIDIA AttnRes kernel requires the SM100 family.
74 + - pytest -v -s models/kimi_k3
75 +
76 - label: Basic Models Test (Other CPU) # 5min
77 key: basic-models-test-other-cpu
78 depends_on:
CMakeLists.txt
+22
@@ -1079,6 +1079,24 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1079 set(MLA_ARCHS)
1080 endif()
1081
1082 + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
1083 + cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS
1084 + "10.0f" "${CUDA_ARCHS}")
1085 + endif()
1086 + if(KIMI_K3_ATTN_RES_ARCHS)
1087 + set(KIMI_K3_ATTN_RES_SRC
1088 + "csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu")
1089 + set_gencode_flags_for_srcs(
1090 + SRCS "${KIMI_K3_ATTN_RES_SRC}"
1091 + CUDA_ARCHS "${KIMI_K3_ATTN_RES_ARCHS}")
1092 + set_property(SOURCE ${KIMI_K3_ATTN_RES_SRC} APPEND PROPERTY
1093 + COMPILE_OPTIONS
1094 + "$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr;--expt-extended-lambda;--use_fast_math>")
1095 + list(APPEND VLLM_STABLE_EXT_SRC "${KIMI_K3_ATTN_RES_SRC}")
1096 + message(STATUS
1097 + "Building Kimi K3 AttnRes for archs: ${KIMI_K3_ATTN_RES_ARCHS}")
1098 + endif()
1099 +
1100 # Hadacore kernels
1101 cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
1102 if(HADACORE_ARCHS)
@@ -1120,6 +1138,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1138 target_compile_definitions(_C_stable_libtorch PRIVATE
1139 VLLM_ENABLE_COOPERATIVE_TOPK=1)
1140 endif()
1141 + if(KIMI_K3_ATTN_RES_ARCHS)
1142 + target_compile_definitions(_C_stable_libtorch PRIVATE
1143 + VLLM_ENABLE_KIMI_K3_ATTN_RES=1)
1144 + endif()
1145 # Needed by CUTLASS kernels
1146 target_compile_definitions(_C_stable_libtorch PRIVATE
1147 CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu new
+954
@@ -0,0 +1,954 @@
1 +/*
2 + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
3 + */
4 +
5 +// Production AttnRes forward for Blackwell (SM100).
6 +//
7 +// Warp-specialized online softmax + residual + RMSNorm:
8 +// - 1 producer warp issues cp.async.bulk row loads into shared memory.
9 +// - 8 consumer warps compute reductions and output.
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
14 +// block residual storage.
15 +
16 +#include "../torch_utils.h"
17 +
18 +#include <cfloat>
19 +#include <cstdint>
20 +#include <cstdio>
21 +#include <cuda_runtime.h>
22 +#include <type_traits>
23 +
24 +using bf16_t = __nv_bfloat16;
25 +
26 +namespace sm100 {
27 +namespace fwd_prod_v2 {
28 +
29 +constexpr int K_TILE = 1024;
30 +constexpr int N_CHUNK_DEFAULT = 4;
31 +constexpr int CHUNK_DEPTH = 2;
32 +constexpr int BLK = 288; // 1 producer warp + 8 consumer warps
33 +constexpr int CONSUMER_THREADS = BLK - 32; // 256
34 +constexpr int CONSUMER_WARPS = CONSUMER_THREADS / 32;
35 +constexpr int CONSUMER_GROUPS = 2; // two 128-thread consumer groups
36 +constexpr int CONSUMER_THREADS_PER_GROUP = CONSUMER_THREADS / CONSUMER_GROUPS;
37 +constexpr int FIRST_USER_NAMED_BARRIER = 8;
38 +
39 +__device__ __forceinline__ const bf16_t* residual_addr(
40 + const bf16_t* block_res, const bf16_t* layer_res, int source, int N,
41 + int token, int block_stride_m, int block_stride_r, int H) {
42 + if (source < N - 1) {
43 + return block_res + static_cast<long long>(token) * block_stride_m +
44 + source * block_stride_r;
45 + }
46 + return layer_res + static_cast<long long>(token) * H;
47 +}
48 +
49 +__device__ __forceinline__ uint32_t elect_one_sync() {
50 + uint32_t pred = 0;
51 + uint32_t laneid = 0;
52 + asm volatile(
53 + "{\n"
54 + ".reg .b32 %%rx;\n"
55 + ".reg .pred %%px;\n"
56 + " elect.sync %%rx|%%px, %2;\n"
57 + "@%%px mov.s32 %1, 1;\n"
58 + " mov.s32 %0, %%rx;\n"
59 + "}\n"
60 + : "+r"(laneid), "+r"(pred)
61 + : "r"(0xffffffff));
62 + return pred;
63 +}
64 +
65 +__device__ __forceinline__ void mbarrier_init(uint64_t& barrier,
66 + int thread_count) {
67 + uint32_t const barrier_addr =
68 + static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
69 + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(barrier_addr),
70 + "r"(thread_count));
71 +}
72 +
73 +__device__ __forceinline__ void mbarrier_expect_tx(uint64_t& barrier,
74 + uint32_t bytes) {
75 + uint32_t const barrier_addr =
76 + static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
77 + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(
78 + barrier_addr),
79 + "r"(bytes));
80 +}
81 +
82 +__device__ __forceinline__ void mbarrier_wait(uint64_t& barrier, int phase) {
83 + uint32_t const barrier_addr =
84 + static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
85 + asm volatile(
86 + "{\n"
87 + ".reg .pred p;\n"
88 + "WAIT:\n"
89 + "mbarrier.try_wait.parity.shared::cta.b64 p, [%0], %1;\n"
90 + "@p bra DONE;\n"
91 + "bra WAIT;\n"
92 + "DONE:\n"
93 + "}\n" ::"r"(barrier_addr),
94 + "r"(phase));
95 +}
96 +
97 +__device__ __forceinline__ void mbarrier_arrive(uint64_t& barrier) {
98 + uint32_t const barrier_addr =
99 + static_cast<uint32_t>(__cvta_generic_to_shared(&barrier));
100 + asm volatile(
101 + "{\n"
102 + ".reg .b64 state;\n"
103 + "mbarrier.arrive.shared::cta.b64 state, [%0];\n"
104 + "}\n" ::"r"(barrier_addr));
105 +}
106 +
107 +__device__ __forceinline__ void fence_mbarrier_init() {
108 + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
109 +}
110 +
111 +__device__ __forceinline__ void named_barrier_sync(uint32_t num_threads,
112 + uint32_t user_barrier_id) {
113 + asm volatile(
114 + "bar.sync %0, %1;" ::"r"(user_barrier_id + FIRST_USER_NAMED_BARRIER),
115 + "r"(num_threads)
116 + : "memory");
117 +}
118 +
119 +__device__ __forceinline__ void tmem_allocate(int num_columns, uint32_t* dst) {
120 + uint32_t const dst_addr =
121 + static_cast<uint32_t>(__cvta_generic_to_shared(dst));
122 + asm volatile(
123 + "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(
124 + dst_addr),
125 + "r"(num_columns));
126 +}
127 +
128 +__device__ __forceinline__ void tmem_free(uint32_t tmem_ptr, int num_columns) {
129 + asm volatile(
130 + "tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(tmem_ptr),
131 + "r"(num_columns));
132 +}
133 +
134 +__device__ __forceinline__ void tmem_release_allocation_lock() {
135 + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;");
136 +}
137 +
138 +__device__ __forceinline__ void tmem_store_wait() {
139 + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory");
140 +}
141 +
142 +template <int N, typename T>
143 +__device__ __forceinline__ void tmem_load(uint32_t src_addr, T* dst) {
144 + uint32_t* values = reinterpret_cast<uint32_t*>(dst);
145 + if constexpr (N == 8) {
146 + asm volatile(
147 + "tcgen05.ld.sync.aligned.32x32b.x8.b32"
148 + "{%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n"
149 + : "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3]),
150 + "=r"(values[4]), "=r"(values[5]), "=r"(values[6]), "=r"(values[7])
151 + : "r"(src_addr));
152 + } else {
153 + static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8");
154 + asm volatile(
155 + "tcgen05.ld.sync.aligned.32x32b.x4.b32"
156 + "{%0, %1, %2, %3}, [%4];\n"
157 + : "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3])
158 + : "r"(src_addr));
159 + }
160 +}
161 +
162 +template <int N, typename T>
163 +__device__ __forceinline__ void tmem_store(uint32_t dst_addr, T* src) {
164 + uint32_t* values = reinterpret_cast<uint32_t*>(src);
165 + if constexpr (N == 8) {
166 + asm volatile(
167 + "tcgen05.st.sync.aligned.32x32b.x8.b32"
168 + "[%8], {%0, %1, %2, %3, %4, %5, %6, %7};\n" ::"r"(values[0]),
169 + "r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(values[4]),
170 + "r"(values[5]), "r"(values[6]), "r"(values[7]), "r"(dst_addr));
171 + } else {
172 + static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8");
173 + asm volatile(
174 + "tcgen05.st.sync.aligned.32x32b.x4.b32"
175 + "[%4], {%0, %1, %2, %3};\n" ::"r"(values[0]),
176 + "r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(dst_addr));
177 + }
178 +}
179 +
180 +__device__ __forceinline__ float2 float2_add(const float2& a, const float2& b) {
181 + float2 result;
182 + asm volatile("add.rn.f32x2 %0, %1, %2;\n"
183 + : "=l"(reinterpret_cast<uint64_t&>(result))
184 + : "l"(reinterpret_cast<uint64_t const&>(a)),
185 + "l"(reinterpret_cast<uint64_t const&>(b)));
186 + return result;
187 +}
188 +
189 +__device__ __forceinline__ float2 float2_mul(const float2& a, const float2& b) {
190 + float2 result;
191 + asm volatile("mul.f32x2 %0, %1, %2;\n"
192 + : "=l"(reinterpret_cast<uint64_t&>(result))
193 + : "l"(reinterpret_cast<uint64_t const&>(a)),
194 + "l"(reinterpret_cast<uint64_t const&>(b)));
195 + return result;
196 +}
197 +
198 +__device__ __forceinline__ float2 float2_fma(const float2& a, const float2& b,
199 + const float2& c) {
200 + float2 result;
201 + asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n"
202 + : "=l"(reinterpret_cast<uint64_t&>(result))
203 + : "l"(reinterpret_cast<uint64_t const&>(a)),
204 + "l"(reinterpret_cast<uint64_t const&>(b)),
205 + "l"(reinterpret_cast<uint64_t const&>(c)));
206 + return result;
207 +}
208 +
209 +template <int NC>
210 +struct FwdSmemPlan {
211 + alignas(16) uint64_t bar_ready[CHUNK_DEPTH];
212 + alignas(16) uint64_t bar_consumed[CHUNK_DEPTH];
213 + alignas(16) uint64_t bar_output_norm_ready;
214 + alignas(16) float2 ws_stats[CONSUMER_WARPS][NC];
215 + uint32_t tmem_base;
216 +};
217 +
218 +__device__ __forceinline__ void cp_async_bulk(void* smem_dst,
219 + const void* gmem_src, int bytes,
220 + uint64_t& mbar) {
221 + uint32_t const s = static_cast<uint32_t>(__cvta_generic_to_shared(smem_dst));
222 + uint32_t const m = static_cast<uint32_t>(__cvta_generic_to_shared(&mbar));
223 + asm volatile(
224 + "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], "
225 + "[%1], %2, [%3];\n" ::"r"(s),
226 + "l"(gmem_src), "r"(bytes), "r"(m)
227 + : "memory");
228 +}
229 +
230 +template <int H, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false,
231 + bool HAS_DELTA = false, bool HAS_OUTPUT_NORM = false,
232 + bool OUTPUT_NORM_IN_SMEM = false>
233 +__global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(
234 + const bf16_t* __restrict__ block_res, bf16_t* __restrict__ layer_res,
235 + const bf16_t* __restrict__ delta, const bf16_t* __restrict__ res_w,
236 + const bf16_t* __restrict__ rms_w, bf16_t* __restrict__ output, int N, int T,
237 + int B, int block_stride_m, int block_stride_r, float rms_eps,
238 + const bf16_t* __restrict__ output_norm_weight, float output_norm_eps) {
239 +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100
240 + constexpr float LOG2_E = 1.4426950408889634f;
241 + constexpr int N_CHUNK = NC;
242 + // The two-source specialization only consumes half of the TMEM columns.
243 + constexpr int TMEM_COLS_ALLOC = NC == 2 ? 128 : 256;
244 + constexpr int NUM_BUFS = CHUNK_DEPTH * NC;
245 + constexpr int NHT = H / K_TILE;
246 + constexpr int SLICES_PER_GROUP =
247 + (NHT + CONSUMER_GROUPS - 1) / CONSUMER_GROUPS;
248 + constexpr int VEC = 8;
249 + constexpr int ACC_PER_THREAD = H == 7168 ? 28 : SLICES_PER_GROUP * VEC;
250 + constexpr int TMEM_V_COLS_PER_GROUP = SLICES_PER_GROUP * N_CHUNK * VEC;
251 + constexpr int TMEM_V_COLS_TOTAL = CONSUMER_GROUPS * TMEM_V_COLS_PER_GROUP;
252 + static_assert(TMEM_V_COLS_TOTAL <= TMEM_COLS_ALLOC);
253 + static_assert(H >= 4096 && H <= 8192);
254 + static_assert(H % K_TILE == 0);
255 +
256 + const int tid = threadIdx.x;
257 + const int wid = tid >> 5;
258 + const int lane = tid & 31;
259 + const int TB = T * B;
260 + const int num_ctas = gridDim.x;
261 + const int num_chunks = (N + N_CHUNK - 1) / N_CHUNK;
262 +
263 + const int comp_wid = wid - 1;
264 + const int comp_tid = tid - 32;
265 + const int group = (comp_wid >= 4) ? 1 : 0;
266 + const int ct_in_group =
267 + (comp_tid >= 0) ? (comp_tid & (CONSUMER_THREADS_PER_GROUP - 1)) : -1;
268 + const int k_local = ct_in_group * VEC;
269 +
270 + constexpr size_t V_BYTES = (size_t)NUM_BUFS * H * sizeof(bf16_t);
271 + constexpr size_t DELTA_BYTES =
272 + HAS_DELTA ? (size_t)CHUNK_DEPTH * H * sizeof(bf16_t) : 0;
273 + constexpr size_t OUTPUT_NORM_BYTES =
274 + OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0;
275 + extern __shared__ __align__(16) char smem_raw[];
276 + bf16_t* v_bufs = reinterpret_cast<bf16_t*>(smem_raw); // [NUM_BUFS][H]
277 + bf16_t* delta_bufs = reinterpret_cast<bf16_t*>(smem_raw + V_BYTES);
278 + bf16_t* output_norm_buf =
279 + reinterpret_cast<bf16_t*>(smem_raw + V_BYTES + DELTA_BYTES);
280 + FwdSmemPlan<NC>& plan = *reinterpret_cast<FwdSmemPlan<NC>*>(
281 + smem_raw + V_BYTES + DELTA_BYTES + OUTPUT_NORM_BYTES);
282 +
283 + auto slot_of = [](long long gci, int n) {
284 + return (int)(gci % CHUNK_DEPTH) * N_CHUNK + n;
285 + };
286 + auto phase_of = [](long long gci) { return (int)((gci / CHUNK_DEPTH) & 1); };
287 + auto buf_ptr = [&](int slot) -> bf16_t* { return v_bufs + slot * H; };
288 + auto delta_buf_ptr = [&](int chunk_slot) -> bf16_t* {
289 + return delta_bufs + chunk_slot * H;
290 + };
291 +
292 + if (wid == 0 && elect_one_sync()) {
293 + #pragma unroll
294 + for (int i = 0; i < CHUNK_DEPTH; i++) {
295 + mbarrier_init(plan.bar_ready[i], 1);
296 + mbarrier_init(plan.bar_consumed[i], CONSUMER_WARPS);
297 + }
298 + if constexpr (OUTPUT_NORM_IN_SMEM) {
299 + mbarrier_init(plan.bar_output_norm_ready, 1);
300 + }
301 + fence_mbarrier_init();
302 + }
303 +
304 + // gdc wait BEFORE tmem alloc
305 + cudaGridDependencySynchronize();
306 +
307 + if (wid == 1) {
308 + tmem_allocate(TMEM_COLS_ALLOC, &plan.tmem_base);
309 + if constexpr (RELEASE_TMEM) {
310 + tmem_release_allocation_lock();
311 + }
312 + }
313 + __syncthreads();
314 +
315 + if constexpr (OUTPUT_NORM_IN_SMEM) {
316 + if (wid == 0 && elect_one_sync()) {
317 + mbarrier_expect_tx(plan.bar_output_norm_ready, H * (int)sizeof(bf16_t));
318 + cp_async_bulk(output_norm_buf, output_norm_weight, H * sizeof(bf16_t),
319 + plan.bar_output_norm_ready);
320 + }
321 + }
322 +
323 + const uint32_t my_v_tmem =
324 + comp_tid >= 0 ? plan.tmem_base + group * TMEM_V_COLS_PER_GROUP : 0;
325 + float q_cache[ACC_PER_THREAD];
326 + if (comp_tid >= 0) {
327 + #pragma unroll
328 + for (int si = 0; si < SLICES_PER_GROUP; si++) {
329 + if constexpr (H == 7168) {
330 + if (si == SLICES_PER_GROUP - 1) {
331 + int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
332 + #pragma unroll
333 + for (int j = 0; j < 4; j++) {
334 + int h = h_base + j;
335 + q_cache[si * VEC + j] =
336 + __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]);
337 + }
338 + continue;
339 + }
340 + }
341 + int dt = si * CONSUMER_GROUPS + group;
342 + if (dt >= NHT) continue;
343 + int h_base = dt * K_TILE + k_local;
344 + #pragma unroll
345 + for (int j = 0; j < VEC; j++) {
346 + int h = h_base + j;
347 + q_cache[si * VEC + j] =
348 + __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]);
349 + }
350 + }
351 + }
352 +
353 + if (wid == 0) {
354 + if (elect_one_sync()) {
355 + long long gci = 0;
356 + for (int tb = blockIdx.x; tb < TB; tb += num_ctas) {
357 + const int t = tb / B;
358 + for (int ci = 0; ci < num_chunks; ci++, gci++) {
359 + int ns = ci * N_CHUNK;
360 + int an = min(N_CHUNK, N - ns);
361 + int chunk_slot = (int)(gci % CHUNK_DEPTH);
362 + int pc = phase_of(gci);
363 + mbarrier_wait(plan.bar_consumed[chunk_slot], pc ^ 1);
364 + int transaction_bytes = an * H * (int)sizeof(bf16_t);
365 + if constexpr (HAS_DELTA) {
366 + int prefix_n = N - 1 - ns;
367 + if (prefix_n >= 0 && prefix_n < an) {
368 + transaction_bytes += H * (int)sizeof(bf16_t);
369 + }
370 + }
371 + mbarrier_expect_tx(plan.bar_ready[chunk_slot], transaction_bytes);
372 + #pragma unroll
373 + for (int n = 0; n < N_CHUNK; n++) {
374 + if (n >= an) continue;
375 + int slot = slot_of(gci, n);
376 + const bf16_t* src =
377 + residual_addr(block_res, layer_res, ns + n, N, t,
378 + block_stride_m, block_stride_r, H);
379 + cp_async_bulk(buf_ptr(slot), src, H * sizeof(bf16_t),
380 + plan.bar_ready[chunk_slot]);
381 + }
382 + if constexpr (HAS_DELTA) {
383 + int prefix_n = N - 1 - ns;
384 + if (prefix_n >= 0 && prefix_n < an) {
385 + cp_async_bulk(delta_buf_ptr(chunk_slot),
386 + delta + (long long)tb * H, H * sizeof(bf16_t),
387 + plan.bar_ready[chunk_slot]);
388 + }
389 + }
390 + }
391 + }
392 + }
393 + } else {
394 + float acc32[ACC_PER_THREAD] = {};
395 + float eps_cache;
396 + asm volatile("mov.b32 %0, %1;" : "=f"(eps_cache) : "f"(rms_eps));
397 +
398 + long long gci = 0;
399 + for (int tb = blockIdx.x; tb < TB; tb += num_ctas) {
400 + float m_running = -FLT_MAX;
401 + float s_running = 0.f;
402 + #pragma unroll
403 + for (int i = 0; i < ACC_PER_THREAD; i++) {
404 + acc32[i] = 0.f;
405 + }
406 +
407 + for (int ci = 0; ci < num_chunks; ci++, gci++) {
408 + int ns = ci * N_CHUNK;
409 + int an = min(N_CHUNK, N - ns);
410 + int chunk_slot = (int)(gci % CHUNK_DEPTH);
411 + int pr = phase_of(gci);
412 + mbarrier_wait(plan.bar_ready[chunk_slot], pr);
413 +
414 + float2 sq_local[N_CHUNK] = {};
415 + float2 dot_local[N_CHUNK] = {};
416 +
417 + auto pass_A_body = [&](auto AN_TOK) {
418 + constexpr int AN = decltype(AN_TOK)::value;
419 + #pragma unroll
420 + for (int si = 0; si < SLICES_PER_GROUP; si++) {
421 + if constexpr (H == 7168) {
422 + if (si == SLICES_PER_GROUP - 1) {
423 + int h_base =
424 + 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
425 + const float* qv = &q_cache[si * VEC];
426 + #pragma unroll
427 + for (int n = 0; n < AN; n++) {
428 + int slot = slot_of(gci, n);
429 + int2 vp =
430 + *reinterpret_cast<const int2*>(buf_ptr(slot) + h_base);
431 + auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp);
432 + if constexpr (HAS_DELTA) {
433 + int prefix_n = N - 1 - ns;
434 + if (n == prefix_n) {
435 + const bf16_t* delta_ptr =
436 + delta_buf_ptr(chunk_slot) + h_base;
437 + #pragma unroll
438 + for (int j = 0; j < 2; j++) {
439 + auto delta2 = *reinterpret_cast<const __nv_bfloat162*>(
440 + delta_ptr + 2 * j);
441 + v2[j] = __hadd2(v2[j], delta2);
442 + }
443 + *reinterpret_cast<int2*>(layer_res + (long long)tb * H +
444 + h_base) = vp;
445 + }
446 + }
447 + float2 f[2] = {__bfloat1622float2(v2[0]),
448 + __bfloat1622float2(v2[1])};
449 + tmem_store<4>(my_v_tmem + (si * N_CHUNK + n) * VEC, f);
450 + sq_local[n] = float2_fma(f[0], f[0], sq_local[n]);
451 + sq_local[n] = float2_fma(f[1], f[1], sq_local[n]);
452 + dot_local[n] =
453 + float2_fma(f[0], make_float2(qv[0], qv[1]), dot_local[n]);
454 + dot_local[n] =
455 + float2_fma(f[1], make_float2(qv[2], qv[3]), dot_local[n]);
456 + }
457 + continue;
458 + }
459 + }
460 + int dt = si * CONSUMER_GROUPS + group;
461 + if (dt >= NHT) continue;
462 + int h_base = dt * K_TILE + k_local;
463 + const float* qv = &q_cache[si * VEC];
464 +
465 + #pragma unroll
466 + for (int n = 0; n < AN; n++) {
467 + int slot = slot_of(gci, n);
468 + int4 vp = *reinterpret_cast<const int4*>(buf_ptr(slot) + h_base);
469 + auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp);
470 + if constexpr (HAS_DELTA) {
471 + int prefix_n = N - 1 - ns;
472 + if (n == prefix_n) {
473 + const bf16_t* delta_ptr = delta_buf_ptr(chunk_slot) + h_base;
474 + #pragma unroll
475 + for (int j = 0; j < VEC / 2; j++) {
476 + auto delta2 = *reinterpret_cast<const __nv_bfloat162*>(
477 + delta_ptr + 2 * j);
478 + v2[j] = __hadd2(v2[j], delta2);
479 + }
480 + *reinterpret_cast<int4*>(layer_res + (long long)tb * H +
481 + h_base) = vp;
482 + }
483 + }
484 + float2 f[4] = {
485 + __bfloat1622float2(v2[0]), __bfloat1622float2(v2[1]),
486 + __bfloat1622float2(v2[2]), __bfloat1622float2(v2[3])};
487 + tmem_store<VEC>(my_v_tmem + (si * N_CHUNK + n) * VEC, f);
488 + #pragma unroll
489 + for (int j = 0; j < VEC / 2; j++) {
490 + sq_local[n] = float2_fma(f[j], f[j], sq_local[n]);
491 + dot_local[n] = float2_fma(
492 + f[j], make_float2(qv[2 * j], qv[2 * j + 1]), dot_local[n]);
493 + }
494 + }
495 + }
496 + };
497 + if constexpr (NC == 4) {
498 + switch (an) {
499 + case 4:
500 + pass_A_body(std::integral_constant<int, 4>{});
501 + break;
502 + case 3:
503 + pass_A_body(std::integral_constant<int, 3>{});
504 + break;
505 + case 2:
506 + pass_A_body(std::integral_constant<int, 2>{});
507 + break;
508 + case 1:
509 + pass_A_body(std::integral_constant<int, 1>{});
510 + break;
511 + default:
512 + __builtin_unreachable();
513 + }
514 + } else if constexpr (NC == 3) {
515 + switch (an) {
516 + case 3:
517 + pass_A_body(std::integral_constant<int, 3>{});
518 + break;
519 + case 2:
520 + pass_A_body(std::integral_constant<int, 2>{});
521 + break;
522 + case 1:
523 + pass_A_body(std::integral_constant<int, 1>{});
524 + break;
525 + default:
526 + __builtin_unreachable();
527 + }
528 + } else {
529 + static_assert(NC == 2);
530 + switch (an) {
531 + case 2:
532 + pass_A_body(std::integral_constant<int, 2>{});
533 + break;
534 + case 1:
535 + pass_A_body(std::integral_constant<int, 1>{});
536 + break;
537 + default:
538 + __builtin_unreachable();
539 + }
540 + }
541 + if (lane == 0) {
542 + mbarrier_arrive(plan.bar_consumed[chunk_slot]);
543 + }
544 + tmem_store_wait();
545 +
546 + float2 reduce_pair[N_CHUNK];
547 + #pragma unroll
548 + for (int n = 0; n < N_CHUNK; n++) {
549 + reduce_pair[n] = make_float2(sq_local[n].x + sq_local[n].y,
550 + dot_local[n].x + dot_local[n].y);
551 + }
552 + #pragma unroll
553 + for (int offset = 16; offset > 0; offset >>= 1) {
554 + #pragma unroll
555 + for (int n = 0; n < N_CHUNK; n++) {
556 + uint64_t packed = reinterpret_cast<uint64_t&>(reduce_pair[n]);
557 + packed = __shfl_xor_sync(0xffffffff, packed, offset);
558 + float2 other = reinterpret_cast<float2&>(packed);
559 + reduce_pair[n] = float2_add(reduce_pair[n], other);
560 + }
561 + }
562 + if (lane == 0) {
563 + #pragma unroll
564 + for (int n = 0; n < N_CHUNK; n++) {
565 + plan.ws_stats[comp_wid][n] = reduce_pair[n];
566 + }
567 + }
568 + named_barrier_sync(CONSUMER_THREADS, 0);
569 +
570 + float local_rsig = 0.f;
571 + float local_logit = 0.f;
572 + int stat_n = lane / CONSUMER_WARPS;
573 + int stat_w = lane % CONSUMER_WARPS;
574 + float2 totals = {};
575 + if (stat_n < N_CHUNK) {
576 + totals = plan.ws_stats[stat_w][stat_n];
577 + }
578 + #pragma unroll
579 + for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) {
580 + totals.x +=
581 + __shfl_down_sync(0xffffffff, totals.x, offset, CONSUMER_WARPS);
582 + totals.y +=
583 + __shfl_down_sync(0xffffffff, totals.y, offset, CONSUMER_WARPS);
584 + }
585 + if (stat_n < N_CHUNK && stat_w == 0) {
586 + local_rsig = rsqrtf(totals.x / H + eps_cache);
587 + local_logit = totals.y * local_rsig;
588 + }
589 + float logit_n[N_CHUNK];
590 + #pragma unroll
591 + for (int n = 0; n < N_CHUNK; n++) {
592 + logit_n[n] = __shfl_sync(0xffffffff, local_logit, n * CONSUMER_WARPS);
593 + }
594 +
595 + float m_chunk = -FLT_MAX;
596 + #pragma unroll
597 + for (int n = 0; n < N_CHUNK; n++) {
598 + if (n < an) m_chunk = fmaxf(m_chunk, logit_n[n]);
599 + }
600 + float m_new = fmaxf(m_running, m_chunk);
601 + float corr = exp2f((m_running - m_new) * LOG2_E);
602 + float w_n[N_CHUNK] = {};
603 + float w_sum = 0.f;
604 + #pragma unroll
605 + for (int n = 0; n < N_CHUNK; n++) {
606 + if (n < an) {
607 + w_n[n] = exp2f((logit_n[n] - m_new) * LOG2_E);
608 + w_sum += w_n[n];
609 + }
610 + }
611 +
612 + auto pass_B_body = [&](auto AN_TOK) {
613 + constexpr int AN = decltype(AN_TOK)::value;
614 + #pragma unroll
615 + for (int si = 0; si < SLICES_PER_GROUP; si++) {
616 + if constexpr (H == 7168) {
617 + if (si == SLICES_PER_GROUP - 1) {
618 + float2 corr2 = make_float2(corr, corr);
619 + float2 a[2];
620 + #pragma unroll
621 + for (int j = 0; j < 2; j++) {
622 + float2 old = make_float2(acc32[si * VEC + 2 * j],
623 + acc32[si * VEC + 2 * j + 1]);
624 + a[j] = float2_mul(old, corr2);
625 + }
626 + float2 f_cache[AN][2];
627 + #pragma unroll
628 + for (int n = 0; n < AN; n++) {
629 + tmem_load<4>(my_v_tmem + (si * N_CHUNK + n) * VEC,
630 + f_cache[n]);
631 + }
632 + #pragma unroll
633 + for (int n = 0; n < AN; n++) {
634 + float2 wn = make_float2(w_n[n], w_n[n]);
635 + #pragma unroll
636 + for (int j = 0; j < 2; j++) {
637 + a[j] = float2_fma(wn, f_cache[n][j], a[j]);
638 + }
639 + }
640 + #pragma unroll
641 + for (int j = 0; j < 2; j++) {
642 + acc32[si * VEC + 2 * j] = a[j].x;
643 + acc32[si * VEC + 2 * j + 1] = a[j].y;
644 + }
645 + continue;
646 + }
647 + }
648 + int dt = si * CONSUMER_GROUPS + group;
649 + if (dt >= NHT) continue;
650 + float2 corr2 = make_float2(corr, corr);
651 + float2 a[VEC / 2];
652 + #pragma unroll
653 + for (int j = 0; j < VEC / 2; j++) {
654 + float2 old = make_float2(acc32[si * VEC + 2 * j],
655 + acc32[si * VEC + 2 * j + 1]);
656 + a[j] = float2_mul(old, corr2);
657 + }
658 + float2 f_cache[AN][VEC / 2];
659 + #pragma unroll
660 + for (int n = 0; n < AN; n++) {
661 + tmem_load<VEC>(my_v_tmem + (si * N_CHUNK + n) * VEC, f_cache[n]);
662 + }
663 + #pragma unroll
664 + for (int n = 0; n < AN; n++) {
665 + float2 wn = make_float2(w_n[n], w_n[n]);
666 + #pragma unroll
667 + for (int j = 0; j < VEC / 2; j++) {
668 + a[j] = float2_fma(wn, f_cache[n][j], a[j]);
669 + }
670 + }
671 + #pragma unroll
672 + for (int j = 0; j < VEC / 2; j++) {
673 + acc32[si * VEC + 2 * j] = a[j].x;
674 + acc32[si * VEC + 2 * j + 1] = a[j].y;
675 + }
676 + }
677 + };
678 + if constexpr (NC == 4) {
679 + switch (an) {
680 + case 4:
681 + pass_B_body(std::integral_constant<int, 4>{});
682 + break;
683 + case 3:
684 + pass_B_body(std::integral_constant<int, 3>{});
685 + break;
686 + case 2:
687 + pass_B_body(std::integral_constant<int, 2>{});
688 + break;
689 + case 1:
690 + pass_B_body(std::integral_constant<int, 1>{});
691 + break;
692 + default:
693 + __builtin_unreachable();
694 + }
695 + } else if constexpr (NC == 3) {
696 + switch (an) {
697 + case 3:
698 + pass_B_body(std::integral_constant<int, 3>{});
699 + break;
700 + case 2:
701 + pass_B_body(std::integral_constant<int, 2>{});
702 + break;
703 + case 1:
704 + pass_B_body(std::integral_constant<int, 1>{});
705 + break;
706 + default:
707 + __builtin_unreachable();
708 + }
709 + } else {
710 + static_assert(NC == 2);
711 + switch (an) {
712 + case 2:
713 + pass_B_body(std::integral_constant<int, 2>{});
714 + break;
715 + case 1:
716 + pass_B_body(std::integral_constant<int, 1>{});
717 + break;
718 + default:
719 + __builtin_unreachable();
720 + }
721 + }
722 +
723 + s_running = s_running * corr + w_sum;
724 + m_running = m_new;
725 + }
726 +
727 + float inv_s = 1.f / s_running;
728 + bf16_t* out_ptr = output + (long long)tb * H;
729 + float2 output_sq_pair = {};
730 + // When output RMSNorm is fused, the softmax denominator cancels:
731 + // (acc / s) * rsqrt(mean((acc / s)^2) + eps)
732 + // = acc * rsqrt(mean(acc^2) + eps * s^2).
733 + #pragma unroll
734 + for (int si = 0; si < SLICES_PER_GROUP; si++) {
735 + if constexpr (H == 7168) {
736 + if (si == SLICES_PER_GROUP - 1) {
737 + int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
738 + uint2 packed;
739 + auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed);
740 + float2 inv2 = make_float2(inv_s, inv_s);
741 + #pragma unroll
742 + for (int j = 0; j < 2; j++) {
743 + float2 old = make_float2(acc32[si * VEC + 2 * j],
744 + acc32[si * VEC + 2 * j + 1]);
745 + if constexpr (HAS_OUTPUT_NORM) {
746 + output_sq_pair = float2_fma(old, old, output_sq_pair);
747 + } else {
748 + float2 mixed = float2_mul(old, inv2);
749 + ov2[j] = __float22bfloat162_rn(mixed);
750 + }
751 + }
752 + if constexpr (!HAS_OUTPUT_NORM) {
753 + *reinterpret_cast<uint2*>(out_ptr + h_base) = packed;
754 + }
755 + continue;
756 + }
757 + }
758 + int dt = si * CONSUMER_GROUPS + group;
759 + if (dt >= NHT) continue;
760 + int h_base = dt * K_TILE + k_local;
761 + uint4 packed;
762 + auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed);
763 + float2 inv2 = make_float2(inv_s, inv_s);
764 + #pragma unroll
765 + for (int j = 0; j < VEC / 2; j++) {
766 + float2 old =
767 + make_float2(acc32[si * VEC + 2 * j], acc32[si * VEC + 2 * j + 1]);
768 + if constexpr (HAS_OUTPUT_NORM) {
769 + output_sq_pair = float2_fma(old, old, output_sq_pair);
770 + } else {
771 + float2 mixed = float2_mul(old, inv2);
772 + ov2[j] = __float22bfloat162_rn(mixed);
773 + }
774 + }
775 + if constexpr (!HAS_OUTPUT_NORM) {
776 + *reinterpret_cast<uint4*>(out_ptr + h_base) = packed;
777 + }
778 + }
779 +
780 + if constexpr (HAS_OUTPUT_NORM) {
781 + if constexpr (OUTPUT_NORM_IN_SMEM) {
782 + // The immutable weight copy is acquired once, at its first use.
783 + if (tb == blockIdx.x) {
784 + mbarrier_wait(plan.bar_output_norm_ready, 0);
785 + }
786 + }
787 + float output_sq = output_sq_pair.x + output_sq_pair.y;
788 + #pragma unroll
789 + for (int offset = 16; offset > 0; offset >>= 1) {
790 + output_sq += __shfl_xor_sync(0xffffffff, output_sq, offset);
791 + }
792 + if (lane == 0) {
793 + plan.ws_stats[comp_wid][0] = make_float2(output_sq, 0.f);
794 + }
795 + named_barrier_sync(CONSUMER_THREADS, 0);
796 + float total_sq = lane < CONSUMER_WARPS ? plan.ws_stats[lane][0].x : 0.f;
797 + #pragma unroll
798 + for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) {
799 + total_sq +=
800 + __shfl_down_sync(0xffffffff, total_sq, offset, CONSUMER_WARPS);
801 + }
802 + if (lane == 0) {
803 + total_sq =
804 + rsqrtf(total_sq / H + output_norm_eps * s_running * s_running);
805 + }
806 + float output_rsigma = __shfl_sync(0xffffffff, total_sq, 0);
807 + #pragma unroll
808 + for (int si = 0; si < SLICES_PER_GROUP; si++) {
809 + if constexpr (H == 7168) {
810 + if (si == SLICES_PER_GROUP - 1) {
811 + int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4;
812 + uint2 packed;
813 + auto* values = reinterpret_cast<bf16_t*>(&packed);
814 + #pragma unroll
815 + for (int j = 0; j < 4; j++) {
816 + const bf16_t* weight_ptr =
817 + OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight;
818 + float weight = __bfloat162float(weight_ptr[h_base + j]);
819 + values[j] = __float2bfloat16(acc32[si * VEC + j] *
820 + output_rsigma * weight);
821 + }
822 + *reinterpret_cast<uint2*>(out_ptr + h_base) = packed;
823 + continue;
824 + }
825 + }
826 + int dt = si * CONSUMER_GROUPS + group;
827 + if (dt >= NHT) continue;
828 + int h_base = dt * K_TILE + k_local;
829 + uint4 packed;
830 + auto* values = reinterpret_cast<bf16_t*>(&packed);
831 + #pragma unroll
832 + for (int j = 0; j < VEC; j++) {
833 + const bf16_t* weight_ptr =
834 + OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight;
835 + float weight = __bfloat162float(weight_ptr[h_base + j]);
836 + values[j] =
837 + __float2bfloat16(acc32[si * VEC + j] * output_rsigma * weight);
838 + }
839 + *reinterpret_cast<uint4*>(out_ptr + h_base) = packed;
840 + }
841 + }
842 + }
843 + }
844 +
845 + cudaTriggerProgrammaticLaunchCompletion();
846 + __syncthreads();
847 + if (wid == 1) {
848 + tmem_free(plan.tmem_base, TMEM_COLS_ALLOC);
849 + }
850 +#else
851 + if (threadIdx.x == 0) {
852 + printf("attn_res_fwd_online_v2_kernel requires sm_10x\n");
853 + }
854 +#endif
855 +}
856 +
857 +template <int H, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false,
858 + bool HAS_DELTA = false, bool HAS_OUTPUT_NORM = false,
859 + bool OUTPUT_NORM_IN_SMEM = false>
860 +static void launch_fwd(const bf16_t* block_residual, bf16_t* layer_residual,
861 + const bf16_t* delta, const bf16_t* res_weight,
862 + const bf16_t* rms_weight, bf16_t* output, int N, int T,
863 + int B, float rms_eps, int num_sm, cudaStream_t stream,
864 + const bf16_t* output_norm_weight = nullptr,
865 + float output_norm_eps = 0.f, int block_stride_m = 0,
866 + int block_stride_r = 0) {
867 + constexpr size_t smem_size =
868 + ((size_t)CHUNK_DEPTH * (NC + (HAS_DELTA ? 1 : 0)) * H * sizeof(bf16_t) +
869 + (OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0) +
870 + sizeof(FwdSmemPlan<NC>) + 15) &
871 + ~size_t(15);
872 + auto kernel =
873 + &attn_res_fwd_online_v2_kernel<H, NC, RELEASE_TMEM, HAS_DELTA,
874 + HAS_OUTPUT_NORM, OUTPUT_NORM_IN_SMEM>;
875 + static bool attrs_set = false;
876 + if (!attrs_set) {
877 + if (smem_size > 48 * 1024) {
878 + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
879 + smem_size);
880 + }
881 + attrs_set = true;
882 + }
883 + int grid = RELEASE_TMEM ? num_sm * 2 : num_sm;
884 + cudaLaunchConfig_t config{};
885 + config.gridDim = grid;
886 + config.blockDim = BLK;
887 + config.dynamicSmemBytes = smem_size;
888 + config.stream = stream;
889 + cudaLaunchAttribute attrs[1];
890 + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
891 + attrs[0].val.programmaticStreamSerializationAllowed = 1;
892 + config.attrs = attrs;
893 + config.numAttrs = 1;
894 + cudaLaunchKernelEx(&config, kernel, block_residual, layer_residual, delta,
895 + res_weight, rms_weight, output, N, T, B, block_stride_m,
896 + block_stride_r, rms_eps, output_norm_weight,
897 + output_norm_eps);
898 +}
899 +
900 +} // namespace fwd_prod_v2
901 +} // namespace sm100
902 +
903 +void kimi_k3_attn_res(torch::stable::Tensor& prefix,
904 + torch::stable::Tensor const& delta,
905 + torch::stable::Tensor const& blocks,
906 + torch::stable::Tensor const& norm_weight,
907 + torch::stable::Tensor const& qk_weight,
908 + torch::stable::Tensor const& output_norm_weight,
909 + torch::stable::Tensor& output, int64_t num_blocks,
910 + double eps, double output_norm_eps) {
911 + int const num_tokens = static_cast<int>(prefix.size(0));
912 + int const device = prefix.get_device_index();
913 + torch::stable::accelerator::DeviceGuard const device_guard(device);
914 + cudaDeviceProp const* properties = get_device_prop();
915 + STD_TORCH_CHECK(properties->major == 10,
916 + "Kimi K3 AttnRes requires the SM100 family");
917 +
918 + using namespace sm100::fwd_prod_v2;
919 + // Two-source chunks and two resident CTAs are beneficial once setup is
920 + // amortized by the long, full eight-block prefill workload.
921 + if (num_blocks == 8 && num_tokens >= 4096) {
922 + launch_fwd<7168, 2, true, true, true, true>(
923 + static_cast<bf16_t const*>(blocks.data_ptr()),
924 + static_cast<bf16_t*>(prefix.data_ptr()),
925 + static_cast<bf16_t const*>(delta.data_ptr()),
926 + static_cast<bf16_t const*>(qk_weight.data_ptr()),
927 + static_cast<bf16_t const*>(norm_weight.data_ptr()),
928 + static_cast<bf16_t*>(output.data_ptr()),
929 + static_cast<int>(num_blocks) + 1, num_tokens, 1,
930 + static_cast<float>(eps), properties->multiProcessorCount,
931 + get_current_cuda_stream(device),
932 + static_cast<bf16_t const*>(output_norm_weight.data_ptr()),
933 + static_cast<float>(output_norm_eps), static_cast<int>(blocks.stride(0)),
934 + static_cast<int>(blocks.stride(1)));
935 + } else {
936 + launch_fwd<7168, 4, false, true, true, true>(
937 + static_cast<bf16_t const*>(blocks.data_ptr()),
938 + static_cast<bf16_t*>(prefix.data_ptr()),
939 + static_cast<bf16_t const*>(delta.data_ptr()),
940 + static_cast<bf16_t const*>(qk_weight.data_ptr()),
941 + static_cast<bf16_t const*>(norm_weight.data_ptr()),
942 + static_cast<bf16_t*>(output.data_ptr()),
943 + static_cast<int>(num_blocks) + 1, num_tokens, 1,
944 + static_cast<float>(eps), properties->multiProcessorCount,
945 + get_current_cuda_stream(device),
946 + static_cast<bf16_t const*>(output_norm_weight.data_ptr()),
947 + static_cast<float>(output_norm_eps), static_cast<int>(blocks.stride(0)),
948 + static_cast<int>(blocks.stride(1)));
949 + }
950 + cudaError_t const error = cudaGetLastError();
951 + STD_TORCH_CHECK(
952 + error == cudaSuccess,
953 + "Kimi K3 AttnRes kernel launch failed: ", cudaGetErrorString(error));
954 +}
csrc/libtorch_stable/ops.h
+11
@@ -315,6 +315,17 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
315 std::optional<torch::stable::Tensor> index_q_out,
316 const std::string& kv_cache_dtype, bool skip_index_branch);
317
318 +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
319 +void kimi_k3_attn_res(torch::stable::Tensor& prefix,
320 + torch::stable::Tensor const& delta,
321 + torch::stable::Tensor const& blocks,
322 + torch::stable::Tensor const& norm_weight,
323 + torch::stable::Tensor const& qk_weight,
324 + torch::stable::Tensor const& output_norm_weight,
325 + torch::stable::Tensor& output, int64_t num_blocks,
326 + double eps, double output_norm_eps);
327 +#endif
328 +
329 // Sampler kernels (shared CUDA/ROCm)
330 void apply_repetition_penalties_(
331 torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
csrc/libtorch_stable/torch_bindings.cpp
+11
@@ -468,6 +468,14 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
468 "int block_size, Tensor!? q_out, Tensor!? index_q_out, "
469 "str kv_cache_dtype, bool skip_index_branch=False) -> ()");
470
471 +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
472 + ops.def(
473 + "kimi_k3_attn_res("
474 + "Tensor! prefix, Tensor delta, Tensor blocks, Tensor norm_weight, "
475 + "Tensor qk_weight, Tensor output_norm_weight, Tensor! output, "
476 + "int num_blocks, float eps, float output_norm_eps) -> ()");
477 +#endif
478 +
479 // Apply repetition penalties to logits in-place.
480 ops.def(
481 "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -693,6 +701,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
701 #endif
702 ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
703 TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert));
704 +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
705 + ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res));
706 +#endif
707
708 // Sampler kernels (shared CUDA/ROCm)
709 ops.impl("apply_repetition_penalties_",
tests/models/kimi_k3/test_amd_attn_res.py new
+102
@@ -0,0 +1,102 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +import torch.nn.functional as F
7 +
8 +from vllm.models.kimi_k3.amd.ops.attn_res import attn_res
9 +from vllm.platforms import current_platform
10 +
11 +pytestmark = pytest.mark.skipif(
12 + not current_platform.is_rocm(),
13 + reason="AMD AttnRes requires ROCm",
14 +)
15 +
16 +
17 +def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor:
18 + storage = torch.randn(
19 + *shape[:-1],
20 + shape[-1] + padding,
21 + device="cuda",
22 + dtype=torch.bfloat16,
23 + )
24 + return storage[..., : shape[-1]]
25 +
26 +
27 +def _reference(
28 + prefix: torch.Tensor,
29 + blocks: torch.Tensor,
30 + norm_weight: torch.Tensor,
31 + qk_weight: torch.Tensor,
32 + num_blocks: int,
33 + eps: float,
34 +) -> torch.Tensor:
35 + hidden_size = prefix.shape[-1]
36 + values = torch.cat((blocks[:, :num_blocks], prefix.unsqueeze(1)), dim=1)
37 + keys = F.rms_norm(values, (hidden_size,), norm_weight, eps)
38 + probs = (keys @ qk_weight).softmax(dim=-1)
39 + return torch.matmul(probs.unsqueeze(1), values).squeeze(1)
40 +
41 +
42 +@pytest.mark.parametrize(
43 + (
44 + "num_tokens",
45 + "num_blocks",
46 + "block_capacity",
47 + "hidden_size",
48 + "row_padding",
49 + ),
50 + [
51 + pytest.param(0, 3, 5, 128, 0, id="empty"),
52 + pytest.param(1, 1, 2, 128, 0, id="decode-single"),
53 + pytest.param(17, 4, 6, 1024, 7, id="decode-padded"),
54 + pytest.param(320, 8, 10, 7168, 0, id="prefill-full"),
55 + ],
56 +)
57 +def test_amd_attn_res_matches_reference(
58 + num_tokens: int,
59 + num_blocks: int,
60 + block_capacity: int,
61 + hidden_size: int,
62 + row_padding: int,
63 +) -> None:
64 + eps = 1e-5
65 + prefix = _randn_with_row_padding(num_tokens, hidden_size, padding=row_padding)
66 + blocks = _randn_with_row_padding(
67 + num_tokens,
68 + block_capacity,
69 + hidden_size,
70 + padding=row_padding,
71 + )
72 + norm_weight = 1 + 0.1 * torch.randn(
73 + hidden_size, device="cuda", dtype=torch.bfloat16
74 + )
75 + qk_weight = (
76 + torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) / hidden_size**0.5
77 + )
78 + expected = _reference(
79 + prefix,
80 + blocks,
81 + norm_weight,
82 + qk_weight,
83 + num_blocks,
84 + eps,
85 + )
86 + original_prefix = prefix.clone()
87 + original_blocks = blocks.clone()
88 +
89 + actual = attn_res(
90 + prefix,
91 + blocks,
92 + norm_weight,
93 + qk_weight,
94 + num_blocks,
95 + eps,
96 + )
97 +
98 + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
99 + torch.testing.assert_close(prefix, original_prefix, atol=0, rtol=0)
100 + torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0)
101 + assert actual.shape == prefix.shape
102 + assert actual.is_contiguous()
tests/models/kimi_k3/test_attn_res.py new
+193
@@ -0,0 +1,193 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +import torch.nn.functional as F
7 +
8 +from vllm.models.kimi_k3.nvidia.ops import attn_res
9 +from vllm.platforms import current_platform
10 +
11 +HIDDEN_SIZE = 7168
12 +MAX_BLOCKS = 8
13 +EPS = 1e-5
14 +
15 +
16 +def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor:
17 + storage = torch.randn(
18 + *shape[:-1],
19 + shape[-1] + padding,
20 + device="cuda",
21 + dtype=torch.bfloat16,
22 + )
23 + return storage[..., : shape[-1]]
24 +
25 +
26 +def _reference(
27 + prefix: torch.Tensor,
28 + delta: torch.Tensor | None,
29 + blocks: torch.Tensor,
30 + norm_weight: torch.Tensor,
31 + qk_weight: torch.Tensor,
32 + output_norm_weight: torch.Tensor | None,
33 + num_blocks: int,
34 +) -> tuple[torch.Tensor, torch.Tensor]:
35 + if delta is not None:
36 + prefix = prefix + delta
37 + values = torch.cat((blocks[:, :num_blocks], prefix.unsqueeze(1)), dim=1)
38 + keys = F.rms_norm(values, (HIDDEN_SIZE,), norm_weight, EPS)
39 + probs = (keys @ qk_weight).softmax(dim=-1)
40 + output = torch.matmul(probs.unsqueeze(1), values).squeeze(1)
41 + if output_norm_weight is not None:
42 + output = F.rms_norm(output, (HIDDEN_SIZE,), output_norm_weight, EPS)
43 + return output, prefix
44 +
45 +
46 +@pytest.mark.parametrize(
47 + (
48 + "num_tokens",
49 + "num_blocks",
50 + "row_padding",
51 + "write_block",
52 + "has_delta",
53 + "backend",
54 + ),
55 + [
56 + pytest.param(1, 0, 0, True, False, "triton", id="triton-empty"),
57 + pytest.param(1, 0, 0, True, True, "triton", id="triton-empty-add"),
58 + pytest.param(17, 5, 7, True, False, "triton", id="triton-write"),
59 + pytest.param(17, 5, 7, True, True, "triton", id="triton-write-add"),
60 + pytest.param(3, 8, 0, False, False, "triton", id="triton-full"),
61 + pytest.param(3, 8, 0, False, True, "triton", id="triton-full-add"),
62 + pytest.param(320, 1, 0, False, True, "nvidia", id="nvidia-1"),
63 + pytest.param(320, 4, 0, False, True, "nvidia", id="nvidia-4"),
64 + pytest.param(320, 8, 0, False, True, "nvidia", id="nvidia-8"),
65 + ],
66 +)
67 +def test_attn_res(
68 + num_tokens: int,
69 + num_blocks: int,
70 + row_padding: int,
71 + write_block: bool,
72 + has_delta: bool,
73 + backend: str,
74 +):
75 + if backend == "nvidia" and not current_platform.is_device_capability_family(100):
76 + pytest.skip("NVIDIA AttnRes requires the SM100 family")
77 +
78 + prefix = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding)
79 + delta = (
80 + _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding)
81 + if has_delta
82 + else None
83 + )
84 + blocks = _randn_with_row_padding(
85 + num_tokens, MAX_BLOCKS, HIDDEN_SIZE, padding=row_padding
86 + )
87 + norm_weight = 1 + 0.1 * torch.randn(
88 + HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
89 + )
90 + qk_weight = (
91 + torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5
92 + )
93 + output_norm_weight = 1 + 0.1 * torch.randn(
94 + HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
95 + )
96 + original_blocks = blocks.clone()
97 + expected, expected_prefix = _reference(
98 + prefix.clone(),
99 + delta,
100 + blocks,
101 + norm_weight,
102 + qk_weight,
103 + output_norm_weight,
104 + num_blocks,
105 + )
106 + block_write_idx = num_blocks if write_block else -1
107 +
108 + actual = attn_res(
109 + prefix,
110 + delta,
111 + blocks,
112 + norm_weight,
113 + qk_weight,
114 + output_norm_weight,
115 + num_blocks,
116 + block_write_idx,
117 + EPS,
118 + EPS,
119 + )
120 +
121 + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
122 + torch.testing.assert_close(prefix, expected_prefix, atol=0, rtol=0)
123 + if write_block:
124 + original_blocks[:, block_write_idx].copy_(expected_prefix)
125 + torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0)
126 + assert actual.is_contiguous()
127 +
128 +
129 +@pytest.mark.parametrize("num_blocks", range(MAX_BLOCKS + 1))
130 +def test_attn_res_block_counts(num_blocks: int):
131 + prefix = torch.randn(1, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
132 + blocks = torch.randn(
133 + 1, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
134 + )
135 + norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
136 + qk_weight = (
137 + torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5
138 + )
139 + output_norm_weight = torch.ones_like(norm_weight)
140 + expected, _ = _reference(
141 + prefix.clone(),
142 + None,
143 + blocks,
144 + norm_weight,
145 + qk_weight,
146 + output_norm_weight,
147 + num_blocks,
148 + )
149 +
150 + actual = attn_res(
151 + prefix,
152 + None,
153 + blocks,
154 + norm_weight,
155 + qk_weight,
156 + output_norm_weight,
157 + num_blocks,
158 + -1,
159 + EPS,
160 + EPS,
161 + )
162 +
163 + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
164 +
165 +
166 +def test_attn_res_without_output_norm():
167 + prefix = torch.randn(7, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
168 + delta = torch.randn_like(prefix)
169 + blocks = torch.randn(
170 + 7, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
171 + )
172 + norm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
173 + qk_weight = (
174 + torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5
175 + )
176 + expected, _ = _reference(
177 + prefix.clone(), delta, blocks, norm_weight, qk_weight, None, MAX_BLOCKS
178 + )
179 +
180 + actual = attn_res(
181 + prefix,
182 + delta,
183 + blocks,
184 + norm_weight,
185 + qk_weight,
186 + None,
187 + MAX_BLOCKS,
188 + -1,
189 + EPS,
190 + 0.0,
191 + )
192 +
193 + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
vllm/_custom_ops.py
+27
@@ -2746,6 +2746,33 @@ def concat_and_cache_mla(
2746 )
2747
2748
2749 +def kimi_k3_attn_res(
2750 + prefix: torch.Tensor,
2751 + delta: torch.Tensor,
2752 + blocks: torch.Tensor,
2753 + norm_weight: torch.Tensor,
2754 + qk_weight: torch.Tensor,
2755 + output_norm_weight: torch.Tensor,
2756 + num_blocks: int,
2757 + eps: float,
2758 + output_norm_eps: float,
2759 +) -> torch.Tensor:
2760 + output = torch.empty_like(prefix)
2761 + torch.ops._C.kimi_k3_attn_res(
2762 + prefix,
2763 + delta,
2764 + blocks,
2765 + norm_weight,
2766 + qk_weight,
2767 + output_norm_weight,
2768 + output,
2769 + num_blocks,
2770 + eps,
2771 + output_norm_eps,
2772 + )
2773 + return output
2774 +
2775 +
2776 def concat_and_cache_mla_rope_fused(
2777 positions: torch.Tensor,
2778 q_pe: torch.Tensor,
vllm/models/kimi_k3/amd/__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/__init__.py
vllm/models/kimi_k3/amd/ops/attn_res.py new
+132
@@ -0,0 +1,132 @@
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 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 +
10 +import torch
11 +
12 +from vllm.triton_utils import tl, triton
13 +
14 +
15 +@triton.jit
16 +def _attn_res_kernel(
17 + prefix_ptr,
18 + blocks_ptr,
19 + norm_weight_ptr,
20 + qk_weight_ptr,
21 + output_ptr,
22 + stride_prefix_m: tl.constexpr,
23 + stride_block_m: tl.constexpr,
24 + stride_block_r: tl.constexpr,
25 + stride_output_m: tl.constexpr,
26 + num_blocks: tl.constexpr,
27 + hidden_size: tl.constexpr,
28 + eps: tl.constexpr,
29 + BLOCK_L: tl.constexpr,
30 + BLOCK_D: tl.constexpr,
31 +):
32 + row_idx = tl.program_id(0).to(tl.int64)
33 + d_offsets = tl.max_contiguous(tl.arange(0, BLOCK_D), BLOCK_D)
34 + d_mask = d_offsets < hidden_size
35 +
36 + prefix = tl.load(
37 + prefix_ptr + row_idx * stride_prefix_m + d_offsets,
38 + mask=d_mask,
39 + other=0.0,
40 + ).to(tl.float32)
41 + input_qk_weight = tl.load(norm_weight_ptr + d_offsets, mask=d_mask, other=0.0).to(
42 + tl.float32
43 + ) * tl.load(qk_weight_ptr + d_offsets, mask=d_mask, other=0.0).to(tl.float32)
44 +
45 + max_logit = tl.full((), -float("inf"), tl.float32)
46 + denominator = tl.zeros((), tl.float32)
47 + mixed = tl.zeros((BLOCK_D,), tl.float32)
48 + num_sources = num_blocks + 1
49 +
50 + for source_tile in range(tl.cdiv(num_sources, BLOCK_L)):
51 + source_offsets = source_tile * BLOCK_L + tl.arange(0, BLOCK_L)
52 + source_mask = source_offsets < num_sources
53 + is_prefix = source_offsets == num_blocks
54 + block_ptrs = (
55 + blocks_ptr
56 + + row_idx * stride_block_m
57 + + source_offsets[:, None] * stride_block_r
58 + + d_offsets[None, :]
59 + )
60 + block_values = tl.load(
61 + block_ptrs,
62 + mask=(source_mask[:, None] & ~is_prefix[:, None] & d_mask[None, :]),
63 + other=0.0,
64 + eviction_policy="evict_first",
65 + ).to(tl.float32)
66 + values = tl.where(is_prefix[:, None], prefix[None, :], block_values)
67 + reciprocal_std = tl.rsqrt(
68 + tl.sum(values * values, axis=1) * (1.0 / hidden_size) + eps
69 + )
70 + logits = tl.sum(values * input_qk_weight[None, :], axis=1) * reciprocal_std
71 + scores = tl.where(source_mask, logits, -float("inf"))
72 +
73 + new_max_logit = tl.maximum(max_logit, tl.max(scores, axis=0))
74 + old_scale = tl.exp(max_logit - new_max_logit)
75 + block_scales = tl.exp(scores - new_max_logit)
76 + denominator = denominator * old_scale + tl.sum(block_scales, axis=0)
77 + mixed = mixed * old_scale + tl.sum(block_scales[:, None] * values, axis=0)
78 + max_logit = new_max_logit
79 +
80 + output = mixed / denominator
81 + tl.store(
82 + output_ptr + row_idx * stride_output_m + d_offsets,
83 + output,
84 + mask=d_mask,
85 + )
86 +
87 +
88 +def attn_res(
89 + prefix: torch.Tensor,
90 + blocks: torch.Tensor,
91 + norm_weight: torch.Tensor,
92 + qk_weight: torch.Tensor,
93 + num_blocks: int,
94 + eps: float,
95 +) -> torch.Tensor:
96 + num_tokens, hidden_size = prefix.shape
97 + assert 0 < num_blocks <= blocks.shape[1]
98 + assert blocks.shape[0] == num_tokens
99 + assert norm_weight.numel() == hidden_size
100 + assert qk_weight.numel() == hidden_size
101 + assert prefix.stride(-1) == 1
102 + assert blocks.stride(-1) == 1
103 + assert norm_weight.stride(-1) == 1
104 + assert qk_weight.stride(-1) == 1
105 +
106 + output = prefix.new_empty(prefix.shape)
107 + if num_tokens == 0:
108 + return output
109 +
110 + if num_tokens >= 256 or num_blocks <= 1:
111 + block_l, num_warps = 1, 4
112 + else:
113 + block_l, num_warps = 4, 8
114 + _attn_res_kernel[(num_tokens,)](
115 + prefix,
116 + blocks,
117 + norm_weight,
118 + qk_weight,
119 + output,
120 + prefix.stride(0),
121 + blocks.stride(0),
122 + blocks.stride(1),
123 + output.stride(0),
124 + num_blocks,
125 + hidden_size,
126 + eps,
127 + BLOCK_L=block_l,
128 + BLOCK_D=triton.next_power_of_2(hidden_size),
129 + num_warps=num_warps,
130 + num_stages=2,
131 + )
132 + return output
vllm/models/kimi_k3/nvidia/__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/__init__.py new
+6
@@ -0,0 +1,6 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from .attn_res import attn_res
5 +
6 +__all__ = ["attn_res"]
vllm/models/kimi_k3/nvidia/ops/attn_res.py new
+245
@@ -0,0 +1,245 @@
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 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 +
10 +
11 +import torch
12 +
13 +from vllm import _custom_ops as ops
14 +from vllm.platforms import current_platform
15 +from vllm.triton_utils import tl, triton
16 +
17 +
18 +# Consumed by kimi_k3_triton_warmup.py during kernel_warmup().
19 +def get_attn_res_triton_warmup_profiles(
20 + max_blocks: int,
21 +) -> tuple[tuple[int, bool, int, bool], ...]:
22 + """Return the small-batch profiles that bypass the native kernel."""
23 + profiles = [
24 + (num_blocks, False, -1, True) for num_blocks in range(2, max_blocks + 1)
25 + ]
26 + profiles.extend(
27 + (block_write_idx, True, block_write_idx, True)
28 + for block_write_idx in range(2, max_blocks)
29 + )
30 + profiles.append((max_blocks, True, -1, False))
31 + return tuple(profiles)
32 +
33 +
34 +@triton.jit
35 +def _attn_res_kernel(
36 + prefix_ptr,
37 + delta_ptr,
38 + blocks_ptr,
39 + norm_weight_ptr,
40 + qk_weight_ptr,
41 + output_norm_weight_ptr,
42 + output_ptr,
43 + stride_prefix_m: tl.constexpr,
44 + stride_delta_m: tl.constexpr,
45 + stride_block_m: tl.constexpr,
46 + stride_block_r: tl.constexpr,
47 + stride_output_m: tl.constexpr,
48 + num_blocks: tl.constexpr,
49 + hidden_size: tl.constexpr,
50 + block_write_idx: tl.constexpr,
51 + eps: tl.constexpr,
52 + output_norm_eps: tl.constexpr,
53 + HAS_DELTA: tl.constexpr,
54 + WRITE_BLOCK: tl.constexpr,
55 + APPLY_OUTPUT_NORM: tl.constexpr,
56 + BLOCK_L: tl.constexpr,
57 + BLOCK_D: tl.constexpr,
58 + launch_pdl: tl.constexpr,
59 +):
60 + row_idx = tl.program_id(0).to(tl.int64)
61 + d_offsets = tl.max_contiguous(tl.arange(0, BLOCK_D), BLOCK_D)
62 + d_mask = d_offsets < hidden_size
63 +
64 + if launch_pdl:
65 + tl.extra.cuda.gdc_wait()
66 +
67 + updated_prefix = tl.load(
68 + prefix_ptr + row_idx * stride_prefix_m + d_offsets,
69 + mask=d_mask,
70 + other=0.0,
71 + ).to(tl.float32)
72 + if HAS_DELTA:
73 + delta = tl.load(
74 + delta_ptr + row_idx * stride_delta_m + d_offsets,
75 + mask=d_mask,
76 + other=0.0,
77 + ).to(tl.float32)
78 + updated_prefix += delta
79 + # Match the BF16 prefix-add result before using it as a residual source.
80 + updated_prefix = updated_prefix.to(prefix_ptr.dtype.element_ty).to(tl.float32)
81 + tl.store(
82 + prefix_ptr + row_idx * stride_prefix_m + d_offsets,
83 + updated_prefix,
84 + mask=d_mask,
85 + )
86 + if WRITE_BLOCK:
87 + tl.store(
88 + blocks_ptr
89 + + row_idx * stride_block_m
90 + + block_write_idx * stride_block_r
91 + + d_offsets,
92 + updated_prefix,
93 + mask=d_mask,
94 + )
95 + # With only the prefix source, the AttnRes softmax is exactly one.
96 + if num_blocks == 0:
97 + mixed = updated_prefix
98 + else:
99 + # Reloading avoids keeping the full prefix vector live across the loop.
100 + if HAS_DELTA:
101 + tl.debug_barrier()
102 + input_qk_weight = tl.load(
103 + norm_weight_ptr + d_offsets, mask=d_mask, other=0.0
104 + ).to(tl.float32) * tl.load(
105 + qk_weight_ptr + d_offsets, mask=d_mask, other=0.0
106 + ).to(tl.float32)
107 + max_logit = tl.full((), -float("inf"), tl.float32)
108 + denominator = tl.zeros((), tl.float32)
109 + mixed = tl.zeros((BLOCK_D,), tl.float32)
110 +
111 + num_sources = num_blocks + 1
112 + for source_tile in range(tl.cdiv(num_sources, BLOCK_L)):
113 + source_offsets = source_tile * BLOCK_L + tl.arange(0, BLOCK_L)
114 + source_mask = source_offsets < num_sources
115 + is_prefix = source_offsets == num_blocks
116 + block_ptrs = (
117 + blocks_ptr
118 + + row_idx * stride_block_m
119 + + source_offsets[:, None] * stride_block_r
120 + + d_offsets[None, :]
121 + )
122 + prefix_ptrs = (
123 + prefix_ptr
124 + + row_idx * stride_prefix_m
125 + + source_offsets[:, None] * 0
126 + + d_offsets[None, :]
127 + )
128 + value_ptrs = tl.where(is_prefix[:, None], prefix_ptrs, block_ptrs)
129 + values = tl.load(
130 + value_ptrs,
131 + mask=source_mask[:, None] & d_mask[None, :],
132 + other=0.0,
133 + eviction_policy="evict_first",
134 + ).to(tl.float32)
135 + reciprocal_std = tl.rsqrt(
136 + tl.sum(values * values, axis=1) * (1.0 / hidden_size) + eps
137 + )
138 + logits = tl.sum(values * input_qk_weight[None, :], axis=1) * reciprocal_std
139 + scores = tl.where(source_mask, logits, -float("inf"))
140 +
141 + new_max_logit = tl.maximum(max_logit, tl.max(scores, axis=0))
142 + old_scale = tl.exp(max_logit - new_max_logit)
143 + block_scales = tl.exp(scores - new_max_logit)
144 + denominator = denominator * old_scale + tl.sum(block_scales, axis=0)
145 + mixed = mixed * old_scale + tl.sum(block_scales[:, None] * values, axis=0)
146 + max_logit = new_max_logit
147 +
148 + mixed /= denominator
149 + output = mixed
150 +
151 + if launch_pdl:
152 + tl.extra.cuda.gdc_launch_dependents()
153 +
154 + if APPLY_OUTPUT_NORM:
155 + output_reciprocal_std = tl.rsqrt(
156 + tl.sum(tl.where(d_mask, mixed * mixed, 0.0), axis=0) * (1.0 / hidden_size)
157 + + output_norm_eps
158 + )
159 + output_norm_weight = tl.load(
160 + output_norm_weight_ptr + d_offsets, mask=d_mask, other=0.0
161 + ).to(tl.float32)
162 + output = mixed * output_reciprocal_std * output_norm_weight
163 + tl.store(
164 + output_ptr + row_idx * stride_output_m + d_offsets,
165 + output,
166 + mask=d_mask,
167 + )
168 +
169 +
170 +def attn_res(
171 + prefix: torch.Tensor,
172 + delta: torch.Tensor | None,
173 + blocks: torch.Tensor,
174 + norm_weight: torch.Tensor,
175 + qk_weight: torch.Tensor,
176 + output_norm_weight: torch.Tensor | None,
177 + num_blocks: int,
178 + block_write_idx: int,
179 + eps: float,
180 + output_norm_eps: float,
181 +) -> torch.Tensor:
182 + num_tokens, hidden_size = prefix.shape
183 + assert prefix.stride(-1) == 1
184 + assert delta is None or delta.stride(-1) == 1
185 + assert blocks.stride(-1) == 1
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.
191 + if (
192 + hidden_size == 7168
193 + and delta is not None
194 + and output_norm_weight is not None
195 + and num_blocks > 0
196 + and block_write_idx < 0
197 + and current_platform.is_device_capability_family(100)
198 + ):
199 + return ops.kimi_k3_attn_res(
200 + prefix,
201 + delta,
202 + blocks,
203 + norm_weight,
204 + qk_weight,
205 + output_norm_weight,
206 + num_blocks,
207 + eps,
208 + output_norm_eps,
209 + )
210 + output = prefix.new_empty(prefix.shape)
211 + # Tuned on GB300: source tiling helps decode, while one-source tiles scale
212 + # better for prefill.
213 + # Keep get_attn_res_triton_warmup_profiles in sync with these fallbacks.
214 + if num_tokens >= 256 or num_blocks <= 1:
215 + block_l, num_warps = 1, 4
216 + else:
217 + block_l, num_warps = 4, 8
218 + _attn_res_kernel[(num_tokens,)](
219 + prefix,
220 + delta,
221 + blocks,
222 + norm_weight,
223 + qk_weight,
224 + output_norm_weight,
225 + output,
226 + prefix.stride(0),
227 + 0 if delta is None else delta.stride(0),
228 + blocks.stride(0),
229 + blocks.stride(1),
230 + output.stride(0),
231 + num_blocks,
232 + hidden_size,
233 + block_write_idx,
234 + eps,
235 + output_norm_eps,
236 + HAS_DELTA=delta is not None,
237 + WRITE_BLOCK=block_write_idx >= 0,
238 + APPLY_OUTPUT_NORM=output_norm_weight is not None,
239 + BLOCK_L=block_l,
240 + BLOCK_D=triton.next_power_of_2(hidden_size),
241 + num_warps=num_warps,
242 + num_stages=2,
243 + launch_pdl=current_platform.is_arch_support_pdl(),
244 + )
245 + return output