| 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, 0<=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 | : "memory"); |
| 96 | } |
| 97 | |
| 98 | __device__ __forceinline__ void mbarrier_arrive(uint64_t& barrier) { |
| 99 | uint32_t const barrier_addr = |
| 100 | static_cast<uint32_t>(__cvta_generic_to_shared(&barrier)); |
| 101 | asm volatile( |
| 102 | "{\n" |
| 103 | ".reg .b64 state;\n" |
| 104 | "mbarrier.arrive.shared::cta.b64 state, [%0];\n" |
| 105 | "}\n" ::"r"(barrier_addr) |
| 106 | : "memory"); |
| 107 | } |
| 108 | |
| 109 | __device__ __forceinline__ void fence_mbarrier_init() { |
| 110 | asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); |
| 111 | } |
| 112 | |
| 113 | __device__ __forceinline__ void named_barrier_sync(uint32_t num_threads, |
| 114 | uint32_t user_barrier_id) { |
| 115 | asm volatile( |
| 116 | "bar.sync %0, %1;" ::"r"(user_barrier_id + FIRST_USER_NAMED_BARRIER), |
| 117 | "r"(num_threads) |
| 118 | : "memory"); |
| 119 | } |
| 120 | |
| 121 | __device__ __forceinline__ void tmem_allocate(int num_columns, uint32_t* dst) { |
| 122 | uint32_t const dst_addr = |
| 123 | static_cast<uint32_t>(__cvta_generic_to_shared(dst)); |
| 124 | asm volatile( |
| 125 | "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"( |
| 126 | dst_addr), |
| 127 | "r"(num_columns)); |
| 128 | } |
| 129 | |
| 130 | __device__ __forceinline__ void tmem_free(uint32_t tmem_ptr, int num_columns) { |
| 131 | asm volatile( |
| 132 | "tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(tmem_ptr), |
| 133 | "r"(num_columns)); |
| 134 | } |
| 135 | |
| 136 | __device__ __forceinline__ void tmem_release_allocation_lock() { |
| 137 | asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); |
| 138 | } |
| 139 | |
| 140 | __device__ __forceinline__ void tmem_store_wait() { |
| 141 | asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); |
| 142 | } |
| 143 | |
| 144 | template <int N, typename T> |
| 145 | __device__ __forceinline__ void tmem_load(uint32_t src_addr, T* dst) { |
| 146 | uint32_t* values = reinterpret_cast<uint32_t*>(dst); |
| 147 | if constexpr (N == 8) { |
| 148 | asm volatile( |
| 149 | "tcgen05.ld.sync.aligned.32x32b.x8.b32" |
| 150 | "{%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n" |
| 151 | : "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3]), |
| 152 | "=r"(values[4]), "=r"(values[5]), "=r"(values[6]), "=r"(values[7]) |
| 153 | : "r"(src_addr)); |
| 154 | } else { |
| 155 | static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8"); |
| 156 | asm volatile( |
| 157 | "tcgen05.ld.sync.aligned.32x32b.x4.b32" |
| 158 | "{%0, %1, %2, %3}, [%4];\n" |
| 159 | : "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3]) |
| 160 | : "r"(src_addr)); |
| 161 | } |
| 162 | } |
| 163 | |
| 164 | template <int N, typename T> |
| 165 | __device__ __forceinline__ void tmem_store(uint32_t dst_addr, T* src) { |
| 166 | uint32_t* values = reinterpret_cast<uint32_t*>(src); |
| 167 | if constexpr (N == 8) { |
| 168 | asm volatile( |
| 169 | "tcgen05.st.sync.aligned.32x32b.x8.b32" |
| 170 | "[%8], {%0, %1, %2, %3, %4, %5, %6, %7};\n" ::"r"(values[0]), |
| 171 | "r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(values[4]), |
| 172 | "r"(values[5]), "r"(values[6]), "r"(values[7]), "r"(dst_addr)); |
| 173 | } else { |
| 174 | static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8"); |
| 175 | asm volatile( |
| 176 | "tcgen05.st.sync.aligned.32x32b.x4.b32" |
| 177 | "[%4], {%0, %1, %2, %3};\n" ::"r"(values[0]), |
| 178 | "r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(dst_addr)); |
| 179 | } |
| 180 | } |
| 181 | |
| 182 | __device__ __forceinline__ float2 float2_add(const float2& a, const float2& b) { |
| 183 | float2 result; |
| 184 | asm volatile("add.rn.f32x2 %0, %1, %2;\n" |
| 185 | : "=l"(reinterpret_cast<uint64_t&>(result)) |
| 186 | : "l"(reinterpret_cast<uint64_t const&>(a)), |
| 187 | "l"(reinterpret_cast<uint64_t const&>(b))); |
| 188 | return result; |
| 189 | } |
| 190 | |
| 191 | __device__ __forceinline__ float2 float2_mul(const float2& a, const float2& b) { |
| 192 | float2 result; |
| 193 | asm volatile("mul.f32x2 %0, %1, %2;\n" |
| 194 | : "=l"(reinterpret_cast<uint64_t&>(result)) |
| 195 | : "l"(reinterpret_cast<uint64_t const&>(a)), |
| 196 | "l"(reinterpret_cast<uint64_t const&>(b))); |
| 197 | return result; |
| 198 | } |
| 199 | |
| 200 | __device__ __forceinline__ float2 float2_fma(const float2& a, const float2& b, |
| 201 | const float2& c) { |
| 202 | float2 result; |
| 203 | asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n" |
| 204 | : "=l"(reinterpret_cast<uint64_t&>(result)) |
| 205 | : "l"(reinterpret_cast<uint64_t const&>(a)), |
| 206 | "l"(reinterpret_cast<uint64_t const&>(b)), |
| 207 | "l"(reinterpret_cast<uint64_t const&>(c))); |
| 208 | return result; |
| 209 | } |
| 210 | |
| 211 | template <int NC> |
| 212 | struct FwdSmemPlan { |
| 213 | alignas(16) uint64_t bar_ready[CHUNK_DEPTH]; |
| 214 | alignas(16) uint64_t bar_consumed[CHUNK_DEPTH]; |
| 215 | alignas(16) uint64_t bar_output_norm_ready; |
| 216 | alignas(16) float2 ws_stats[CONSUMER_WARPS][NC]; |
| 217 | uint32_t tmem_base; |
| 218 | }; |
| 219 | |
| 220 | __device__ __forceinline__ void cp_async_bulk(void* smem_dst, |
| 221 | const void* gmem_src, int bytes, |
| 222 | uint64_t& mbar) { |
| 223 | uint32_t const s = static_cast<uint32_t>(__cvta_generic_to_shared(smem_dst)); |
| 224 | uint32_t const m = static_cast<uint32_t>(__cvta_generic_to_shared(&mbar)); |
| 225 | asm volatile( |
| 226 | "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], " |
| 227 | "[%1], %2, [%3];\n" ::"r"(s), |
| 228 | "l"(gmem_src), "r"(bytes), "r"(m) |
| 229 | : "memory"); |
| 230 | } |
| 231 | |
| 232 | template <int H, int N, int NC = N_CHUNK_DEFAULT, int B = 1, |
| 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 | 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, |
| 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; |
| 245 | // The two-source specialization only consumes half of the TMEM columns. |
| 246 | constexpr int TMEM_COLS_ALLOC = NC == 2 ? 128 : 256; |
| 247 | constexpr int NUM_BUFS = CHUNK_DEPTH * NC; |
| 248 | constexpr int NHT = H / K_TILE; |
| 249 | constexpr int SLICES_PER_GROUP = |
| 250 | (NHT + CONSUMER_GROUPS - 1) / CONSUMER_GROUPS; |
| 251 | constexpr int VEC = 8; |
| 252 | constexpr int ACC_PER_THREAD = H == 7168 ? 28 : SLICES_PER_GROUP * VEC; |
| 253 | constexpr int TMEM_V_COLS_PER_GROUP = SLICES_PER_GROUP * N_CHUNK * VEC; |
| 254 | constexpr int TMEM_V_COLS_TOTAL = CONSUMER_GROUPS * TMEM_V_COLS_PER_GROUP; |
| 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 | const int comp_tid = tid - 32; |
| 271 | const int group = (comp_wid >= 4) ? 1 : 0; |
| 272 | const int ct_in_group = |
| 273 | (comp_tid >= 0) ? (comp_tid & (CONSUMER_THREADS_PER_GROUP - 1)) : -1; |
| 274 | const int k_local = ct_in_group * VEC; |
| 275 | |
| 276 | constexpr size_t V_BYTES = (size_t)NUM_BUFS * H * sizeof(bf16_t); |
| 277 | constexpr size_t DELTA_BYTES = (size_t)CHUNK_DEPTH * H * sizeof(bf16_t); |
| 278 | constexpr size_t OUTPUT_NORM_BYTES = |
| 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); |
| 283 | bf16_t* output_norm_buf = |
| 284 | reinterpret_cast<bf16_t*>(smem_raw + V_BYTES + DELTA_BYTES); |
| 285 | FwdSmemPlan<NC>& plan = *reinterpret_cast<FwdSmemPlan<NC>*>( |
| 286 | smem_raw + V_BYTES + DELTA_BYTES + OUTPUT_NORM_BYTES); |
| 287 | |
| 288 | auto slot_of = [](long long gci, int n) { |
| 289 | return (int)(gci % CHUNK_DEPTH) * N_CHUNK + n; |
| 290 | }; |
| 291 | auto phase_of = [](long long gci) { return (int)((gci / CHUNK_DEPTH) & 1); }; |
| 292 | auto buf_ptr = [&](int slot) -> bf16_t* { return v_bufs + slot * H; }; |
| 293 | auto delta_buf_ptr = [&](int chunk_slot) -> bf16_t* { |
| 294 | return delta_bufs + chunk_slot * H; |
| 295 | }; |
| 296 | |
| 297 | if (wid == 0 && elect_one_sync()) { |
| 298 | #pragma unroll |
| 299 | for (int i = 0; i < CHUNK_DEPTH; i++) { |
| 300 | mbarrier_init(plan.bar_ready[i], 1); |
| 301 | mbarrier_init(plan.bar_consumed[i], CONSUMER_WARPS); |
| 302 | } |
| 303 | if constexpr (HAS_OUTPUT_NORM) { |
| 304 | mbarrier_init(plan.bar_output_norm_ready, 1); |
| 305 | } |
| 306 | fence_mbarrier_init(); |
| 307 | } |
| 308 | |
| 309 | if (wid == 1) { |
| 310 | tmem_allocate(TMEM_COLS_ALLOC, &plan.tmem_base); |
| 311 | if constexpr (RELEASE_TMEM) { |
| 312 | tmem_release_allocation_lock(); |
| 313 | } |
| 314 | } |
| 315 | __syncthreads(); |
| 316 | |
| 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), |
| 321 | plan.bar_output_norm_ready); |
| 322 | } |
| 323 | } |
| 324 | |
| 325 | const uint32_t my_v_tmem = |
| 326 | comp_tid >= 0 ? plan.tmem_base + group * TMEM_V_COLS_PER_GROUP : 0; |
| 327 | float q_cache[ACC_PER_THREAD]; |
| 328 | if (comp_tid >= 0) { |
| 329 | #pragma unroll |
| 330 | for (int si = 0; si < SLICES_PER_GROUP; si++) { |
| 331 | if constexpr (H == 7168) { |
| 332 | if (si == SLICES_PER_GROUP - 1) { |
| 333 | int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; |
| 334 | int2 rms_v = *reinterpret_cast<const int2*>(rms_w + h_base); |
| 335 | int2 res_v = *reinterpret_cast<const int2*>(res_w + h_base); |
| 336 | auto* rms2 = reinterpret_cast<__nv_bfloat162*>(&rms_v); |
| 337 | auto* res2 = reinterpret_cast<__nv_bfloat162*>(&res_v); |
| 338 | #pragma unroll |
| 339 | for (int k = 0; k < 2; k++) { |
| 340 | float2 rf = __bfloat1622float2(rms2[k]); |
| 341 | float2 sf = __bfloat1622float2(res2[k]); |
| 342 | q_cache[si * VEC + 2 * k] = rf.x * sf.x; |
| 343 | q_cache[si * VEC + 2 * k + 1] = rf.y * sf.y; |
| 344 | } |
| 345 | continue; |
| 346 | } |
| 347 | } |
| 348 | int dt = si * CONSUMER_GROUPS + group; |
| 349 | if (dt >= NHT) continue; |
| 350 | int h_base = dt * K_TILE + k_local; |
| 351 | int4 rms_v = *reinterpret_cast<const int4*>(rms_w + h_base); |
| 352 | int4 res_v = *reinterpret_cast<const int4*>(res_w + h_base); |
| 353 | auto* rms2 = reinterpret_cast<__nv_bfloat162*>(&rms_v); |
| 354 | auto* res2 = reinterpret_cast<__nv_bfloat162*>(&res_v); |
| 355 | #pragma unroll |
| 356 | for (int k = 0; k < 4; k++) { |
| 357 | float2 rf = __bfloat1622float2(rms2[k]); |
| 358 | float2 sf = __bfloat1622float2(res2[k]); |
| 359 | q_cache[si * VEC + 2 * k] = rf.x * sf.x; |
| 360 | q_cache[si * VEC + 2 * k + 1] = rf.y * sf.y; |
| 361 | } |
| 362 | } |
| 363 | } |
| 364 | |
| 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++) { |
| 372 | int ns = ci * N_CHUNK; |
| 373 | int an = min(N_CHUNK, N - ns); |
| 374 | int chunk_slot = (int)(gci % CHUNK_DEPTH); |
| 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); |
| 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); |
| 382 | } |
| 383 | } |
| 384 | mbarrier_expect_tx(plan.bar_ready[chunk_slot], transaction_bytes); |
| 385 | #pragma unroll |
| 386 | for (int n = 0; n < N_CHUNK; n++) { |
| 387 | if (n >= an) continue; |
| 388 | int slot = slot_of(gci, n); |
| 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 | } |
| 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]); |
| 409 | } |
| 410 | } |
| 411 | } |
| 412 | } |
| 413 | cudaTriggerProgrammaticLaunchCompletion(); |
| 414 | } |
| 415 | } else { |
| 416 | float acc32[ACC_PER_THREAD] = {}; |
| 417 | float eps_cache; |
| 418 | asm volatile("mov.b32 %0, %1;" : "=f"(eps_cache) : "f"(rms_eps)); |
| 419 | |
| 420 | long long gci = 0; |
| 421 | for (int tb = blockIdx.x; tb < TB; tb += num_ctas) { |
| 422 | float m_running = -FLT_MAX; |
| 423 | float s_running = 0.f; |
| 424 | #pragma unroll |
| 425 | for (int i = 0; i < ACC_PER_THREAD; i++) { |
| 426 | acc32[i] = 0.f; |
| 427 | } |
| 428 | |
| 429 | #pragma unroll |
| 430 | for (int ci = 0; ci < num_chunks; ci++, gci++) { |
| 431 | int ns = ci * N_CHUNK; |
| 432 | int an = min(N_CHUNK, N - ns); |
| 433 | int chunk_slot = (int)(gci % CHUNK_DEPTH); |
| 434 | int pr = phase_of(gci); |
| 435 | mbarrier_wait(plan.bar_ready[chunk_slot], pr); |
| 436 | |
| 437 | float2 sq_local[N_CHUNK] = {}; |
| 438 | float2 dot_local[N_CHUNK] = {}; |
| 439 | |
| 440 | auto pass_A_body = [&](auto AN_TOK) { |
| 441 | constexpr int AN = decltype(AN_TOK)::value; |
| 442 | #pragma unroll |
| 443 | for (int si = 0; si < SLICES_PER_GROUP; si++) { |
| 444 | if constexpr (H == 7168) { |
| 445 | if (si == SLICES_PER_GROUP - 1) { |
| 446 | int h_base = |
| 447 | 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; |
| 448 | const float* qv = &q_cache[si * VEC]; |
| 449 | #pragma unroll |
| 450 | for (int n = 0; n < AN; n++) { |
| 451 | int slot = slot_of(gci, n); |
| 452 | int2 vp = |
| 453 | *reinterpret_cast<const int2*>(buf_ptr(slot) + h_base); |
| 454 | auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp); |
| 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 |
| 461 | for (int j = 0; j < 2; j++) { |
| 462 | auto delta2 = *reinterpret_cast<const __nv_bfloat162*>( |
| 463 | delta_ptr + 2 * j); |
| 464 | v2[j] = __hadd2(v2[j], delta2); |
| 465 | } |
| 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])}; |
| 477 | tmem_store<4>(my_v_tmem + (si * N_CHUNK + n) * VEC, f); |
| 478 | sq_local[n] = float2_fma(f[0], f[0], sq_local[n]); |
| 479 | sq_local[n] = float2_fma(f[1], f[1], sq_local[n]); |
| 480 | dot_local[n] = |
| 481 | float2_fma(f[0], make_float2(qv[0], qv[1]), dot_local[n]); |
| 482 | dot_local[n] = |
| 483 | float2_fma(f[1], make_float2(qv[2], qv[3]), dot_local[n]); |
| 484 | } |
| 485 | continue; |
| 486 | } |
| 487 | } |
| 488 | int dt = si * CONSUMER_GROUPS + group; |
| 489 | if (dt >= NHT) continue; |
| 490 | int h_base = dt * K_TILE + k_local; |
| 491 | const float* qv = &q_cache[si * VEC]; |
| 492 | |
| 493 | #pragma unroll |
| 494 | for (int n = 0; n < AN; n++) { |
| 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); |
| 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++) { |
| 504 | auto delta2 = *reinterpret_cast<const __nv_bfloat162*>( |
| 505 | delta_ptr + 2 * j); |
| 506 | v2[j] = __hadd2(v2[j], delta2); |
| 507 | } |
| 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]), |
| 519 | __bfloat1622float2(v2[2]), __bfloat1622float2(v2[3])}; |
| 520 | tmem_store<VEC>(my_v_tmem + (si * N_CHUNK + n) * VEC, f); |
| 521 | #pragma unroll |
| 522 | for (int j = 0; j < VEC / 2; j++) { |
| 523 | sq_local[n] = float2_fma(f[j], f[j], sq_local[n]); |
| 524 | dot_local[n] = float2_fma( |
| 525 | f[j], make_float2(qv[2 * j], qv[2 * j + 1]), dot_local[n]); |
| 526 | } |
| 527 | } |
| 528 | } |
| 529 | }; |
| 530 | if constexpr (NC == 4) { |
| 531 | switch (an) { |
| 532 | case 4: |
| 533 | pass_A_body(std::integral_constant<int, 4>{}); |
| 534 | break; |
| 535 | case 3: |
| 536 | pass_A_body(std::integral_constant<int, 3>{}); |
| 537 | break; |
| 538 | case 2: |
| 539 | pass_A_body(std::integral_constant<int, 2>{}); |
| 540 | break; |
| 541 | case 1: |
| 542 | pass_A_body(std::integral_constant<int, 1>{}); |
| 543 | break; |
| 544 | default: |
| 545 | __builtin_unreachable(); |
| 546 | } |
| 547 | } else if constexpr (NC == 3) { |
| 548 | switch (an) { |
| 549 | case 3: |
| 550 | pass_A_body(std::integral_constant<int, 3>{}); |
| 551 | break; |
| 552 | case 2: |
| 553 | pass_A_body(std::integral_constant<int, 2>{}); |
| 554 | break; |
| 555 | case 1: |
| 556 | pass_A_body(std::integral_constant<int, 1>{}); |
| 557 | break; |
| 558 | default: |
| 559 | __builtin_unreachable(); |
| 560 | } |
| 561 | } else { |
| 562 | static_assert(NC == 2); |
| 563 | switch (an) { |
| 564 | case 2: |
| 565 | pass_A_body(std::integral_constant<int, 2>{}); |
| 566 | break; |
| 567 | case 1: |
| 568 | pass_A_body(std::integral_constant<int, 1>{}); |
| 569 | break; |
| 570 | default: |
| 571 | __builtin_unreachable(); |
| 572 | } |
| 573 | } |
| 574 | if (lane == 0) { |
| 575 | mbarrier_arrive(plan.bar_consumed[chunk_slot]); |
| 576 | } |
| 577 | tmem_store_wait(); |
| 578 | |
| 579 | float2 reduce_pair[N_CHUNK]; |
| 580 | #pragma unroll |
| 581 | for (int n = 0; n < N_CHUNK; n++) { |
| 582 | reduce_pair[n] = make_float2(sq_local[n].x + sq_local[n].y, |
| 583 | dot_local[n].x + dot_local[n].y); |
| 584 | } |
| 585 | #pragma unroll |
| 586 | for (int offset = 16; offset > 0; offset >>= 1) { |
| 587 | #pragma unroll |
| 588 | for (int n = 0; n < N_CHUNK; n++) { |
| 589 | uint64_t packed = reinterpret_cast<uint64_t&>(reduce_pair[n]); |
| 590 | packed = __shfl_xor_sync(0xffffffff, packed, offset); |
| 591 | float2 other = reinterpret_cast<float2&>(packed); |
| 592 | reduce_pair[n] = float2_add(reduce_pair[n], other); |
| 593 | } |
| 594 | } |
| 595 | if (lane == 0) { |
| 596 | #pragma unroll |
| 597 | for (int n = 0; n < N_CHUNK; n++) { |
| 598 | plan.ws_stats[comp_wid][n] = reduce_pair[n]; |
| 599 | } |
| 600 | } |
| 601 | named_barrier_sync(CONSUMER_THREADS, 0); |
| 602 | |
| 603 | float local_rsig = 0.f; |
| 604 | float local_logit = 0.f; |
| 605 | int stat_n = lane / CONSUMER_WARPS; |
| 606 | int stat_w = lane % CONSUMER_WARPS; |
| 607 | float2 totals = {}; |
| 608 | if (stat_n < N_CHUNK) { |
| 609 | totals = plan.ws_stats[stat_w][stat_n]; |
| 610 | } |
| 611 | #pragma unroll |
| 612 | for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) { |
| 613 | totals.x += |
| 614 | __shfl_down_sync(0xffffffff, totals.x, offset, CONSUMER_WARPS); |
| 615 | totals.y += |
| 616 | __shfl_down_sync(0xffffffff, totals.y, offset, CONSUMER_WARPS); |
| 617 | } |
| 618 | if (stat_n < N_CHUNK && stat_w == 0) { |
| 619 | local_rsig = rsqrtf(totals.x / H + eps_cache); |
| 620 | local_logit = totals.y * local_rsig; |
| 621 | } |
| 622 | float logit_n[N_CHUNK]; |
| 623 | #pragma unroll |
| 624 | for (int n = 0; n < N_CHUNK; n++) { |
| 625 | logit_n[n] = __shfl_sync(0xffffffff, local_logit, n * CONSUMER_WARPS); |
| 626 | } |
| 627 | |
| 628 | float m_chunk = -FLT_MAX; |
| 629 | #pragma unroll |
| 630 | for (int n = 0; n < N_CHUNK; n++) { |
| 631 | if (n < an) m_chunk = fmaxf(m_chunk, logit_n[n]); |
| 632 | } |
| 633 | float m_new = fmaxf(m_running, m_chunk); |
| 634 | float corr = exp2f((m_running - m_new) * LOG2_E); |
| 635 | float w_n[N_CHUNK] = {}; |
| 636 | float w_sum = 0.f; |
| 637 | #pragma unroll |
| 638 | for (int n = 0; n < N_CHUNK; n++) { |
| 639 | if (n < an) { |
| 640 | w_n[n] = exp2f((logit_n[n] - m_new) * LOG2_E); |
| 641 | w_sum += w_n[n]; |
| 642 | } |
| 643 | } |
| 644 | |
| 645 | auto pass_B_body = [&](auto AN_TOK) { |
| 646 | constexpr int AN = decltype(AN_TOK)::value; |
| 647 | #pragma unroll |
| 648 | for (int si = 0; si < SLICES_PER_GROUP; si++) { |
| 649 | if constexpr (H == 7168) { |
| 650 | if (si == SLICES_PER_GROUP - 1) { |
| 651 | float2 corr2 = make_float2(corr, corr); |
| 652 | float2 a[2]; |
| 653 | #pragma unroll |
| 654 | for (int j = 0; j < 2; j++) { |
| 655 | float2 old = make_float2(acc32[si * VEC + 2 * j], |
| 656 | acc32[si * VEC + 2 * j + 1]); |
| 657 | a[j] = float2_mul(old, corr2); |
| 658 | } |
| 659 | float2 f_cache[AN][2]; |
| 660 | #pragma unroll |
| 661 | for (int n = 0; n < AN; n++) { |
| 662 | tmem_load<4>(my_v_tmem + (si * N_CHUNK + n) * VEC, |
| 663 | f_cache[n]); |
| 664 | } |
| 665 | #pragma unroll |
| 666 | for (int n = 0; n < AN; n++) { |
| 667 | float2 wn = make_float2(w_n[n], w_n[n]); |
| 668 | #pragma unroll |
| 669 | for (int j = 0; j < 2; j++) { |
| 670 | a[j] = float2_fma(wn, f_cache[n][j], a[j]); |
| 671 | } |
| 672 | } |
| 673 | #pragma unroll |
| 674 | for (int j = 0; j < 2; j++) { |
| 675 | acc32[si * VEC + 2 * j] = a[j].x; |
| 676 | acc32[si * VEC + 2 * j + 1] = a[j].y; |
| 677 | } |
| 678 | continue; |
| 679 | } |
| 680 | } |
| 681 | int dt = si * CONSUMER_GROUPS + group; |
| 682 | if (dt >= NHT) continue; |
| 683 | float2 corr2 = make_float2(corr, corr); |
| 684 | float2 a[VEC / 2]; |
| 685 | #pragma unroll |
| 686 | for (int j = 0; j < VEC / 2; j++) { |
| 687 | float2 old = make_float2(acc32[si * VEC + 2 * j], |
| 688 | acc32[si * VEC + 2 * j + 1]); |
| 689 | a[j] = float2_mul(old, corr2); |
| 690 | } |
| 691 | float2 f_cache[AN][VEC / 2]; |
| 692 | #pragma unroll |
| 693 | for (int n = 0; n < AN; n++) { |
| 694 | tmem_load<VEC>(my_v_tmem + (si * N_CHUNK + n) * VEC, f_cache[n]); |
| 695 | } |
| 696 | #pragma unroll |
| 697 | for (int n = 0; n < AN; n++) { |
| 698 | float2 wn = make_float2(w_n[n], w_n[n]); |
| 699 | #pragma unroll |
| 700 | for (int j = 0; j < VEC / 2; j++) { |
| 701 | a[j] = float2_fma(wn, f_cache[n][j], a[j]); |
| 702 | } |
| 703 | } |
| 704 | #pragma unroll |
| 705 | for (int j = 0; j < VEC / 2; j++) { |
| 706 | acc32[si * VEC + 2 * j] = a[j].x; |
| 707 | acc32[si * VEC + 2 * j + 1] = a[j].y; |
| 708 | } |
| 709 | } |
| 710 | }; |
| 711 | if constexpr (NC == 4) { |
| 712 | switch (an) { |
| 713 | case 4: |
| 714 | pass_B_body(std::integral_constant<int, 4>{}); |
| 715 | break; |
| 716 | case 3: |
| 717 | pass_B_body(std::integral_constant<int, 3>{}); |
| 718 | break; |
| 719 | case 2: |
| 720 | pass_B_body(std::integral_constant<int, 2>{}); |
| 721 | break; |
| 722 | case 1: |
| 723 | pass_B_body(std::integral_constant<int, 1>{}); |
| 724 | break; |
| 725 | default: |
| 726 | __builtin_unreachable(); |
| 727 | } |
| 728 | } else if constexpr (NC == 3) { |
| 729 | switch (an) { |
| 730 | case 3: |
| 731 | pass_B_body(std::integral_constant<int, 3>{}); |
| 732 | break; |
| 733 | case 2: |
| 734 | pass_B_body(std::integral_constant<int, 2>{}); |
| 735 | break; |
| 736 | case 1: |
| 737 | pass_B_body(std::integral_constant<int, 1>{}); |
| 738 | break; |
| 739 | default: |
| 740 | __builtin_unreachable(); |
| 741 | } |
| 742 | } else { |
| 743 | static_assert(NC == 2); |
| 744 | switch (an) { |
| 745 | case 2: |
| 746 | pass_B_body(std::integral_constant<int, 2>{}); |
| 747 | break; |
| 748 | case 1: |
| 749 | pass_B_body(std::integral_constant<int, 1>{}); |
| 750 | break; |
| 751 | default: |
| 752 | __builtin_unreachable(); |
| 753 | } |
| 754 | } |
| 755 | |
| 756 | s_running = s_running * corr + w_sum; |
| 757 | m_running = m_new; |
| 758 | } |
| 759 | |
| 760 | float inv_s = 1.f / s_running; |
| 761 | bf16_t* out_ptr = output + (long long)tb * H; |
| 762 | float2 output_sq_pair = {}; |
| 763 | // When output RMSNorm is fused, the softmax denominator cancels: |
| 764 | // (acc / s) * rsqrt(mean((acc / s)^2) + eps) |
| 765 | // = acc * rsqrt(mean(acc^2) + eps * s^2). |
| 766 | #pragma unroll |
| 767 | for (int si = 0; si < SLICES_PER_GROUP; si++) { |
| 768 | if constexpr (H == 7168) { |
| 769 | if (si == SLICES_PER_GROUP - 1) { |
| 770 | int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; |
| 771 | uint2 packed; |
| 772 | auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed); |
| 773 | float2 inv2 = make_float2(inv_s, inv_s); |
| 774 | #pragma unroll |
| 775 | for (int j = 0; j < 2; j++) { |
| 776 | float2 old = make_float2(acc32[si * VEC + 2 * j], |
| 777 | acc32[si * VEC + 2 * j + 1]); |
| 778 | if constexpr (HAS_OUTPUT_NORM) { |
| 779 | output_sq_pair = float2_fma(old, old, output_sq_pair); |
| 780 | } else { |
| 781 | float2 mixed = float2_mul(old, inv2); |
| 782 | ov2[j] = __float22bfloat162_rn(mixed); |
| 783 | } |
| 784 | } |
| 785 | if constexpr (!HAS_OUTPUT_NORM) { |
| 786 | *reinterpret_cast<uint2*>(out_ptr + h_base) = packed; |
| 787 | } |
| 788 | continue; |
| 789 | } |
| 790 | } |
| 791 | int dt = si * CONSUMER_GROUPS + group; |
| 792 | if (dt >= NHT) continue; |
| 793 | int h_base = dt * K_TILE + k_local; |
| 794 | uint4 packed; |
| 795 | auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed); |
| 796 | float2 inv2 = make_float2(inv_s, inv_s); |
| 797 | #pragma unroll |
| 798 | for (int j = 0; j < VEC / 2; j++) { |
| 799 | float2 old = |
| 800 | make_float2(acc32[si * VEC + 2 * j], acc32[si * VEC + 2 * j + 1]); |
| 801 | if constexpr (HAS_OUTPUT_NORM) { |
| 802 | output_sq_pair = float2_fma(old, old, output_sq_pair); |
| 803 | } else { |
| 804 | float2 mixed = float2_mul(old, inv2); |
| 805 | ov2[j] = __float22bfloat162_rn(mixed); |
| 806 | } |
| 807 | } |
| 808 | if constexpr (!HAS_OUTPUT_NORM) { |
| 809 | *reinterpret_cast<uint4*>(out_ptr + h_base) = packed; |
| 810 | } |
| 811 | } |
| 812 | |
| 813 | if constexpr (HAS_OUTPUT_NORM) { |
| 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 |
| 820 | for (int offset = 16; offset > 0; offset >>= 1) { |
| 821 | output_sq += __shfl_xor_sync(0xffffffff, output_sq, offset); |
| 822 | } |
| 823 | if (lane == 0) { |
| 824 | plan.ws_stats[comp_wid][0] = make_float2(output_sq, 0.f); |
| 825 | } |
| 826 | named_barrier_sync(CONSUMER_THREADS, 0); |
| 827 | float total_sq = lane < CONSUMER_WARPS ? plan.ws_stats[lane][0].x : 0.f; |
| 828 | #pragma unroll |
| 829 | for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) { |
| 830 | total_sq += |
| 831 | __shfl_down_sync(0xffffffff, total_sq, offset, CONSUMER_WARPS); |
| 832 | } |
| 833 | if (lane == 0) { |
| 834 | total_sq = |
| 835 | rsqrtf(total_sq / H + output_norm_eps * s_running * s_running); |
| 836 | } |
| 837 | float output_rsigma = __shfl_sync(0xffffffff, total_sq, 0); |
| 838 | #pragma unroll |
| 839 | for (int si = 0; si < SLICES_PER_GROUP; si++) { |
| 840 | if constexpr (H == 7168) { |
| 841 | if (si == SLICES_PER_GROUP - 1) { |
| 842 | int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; |
| 843 | uint2 packed; |
| 844 | auto* values = reinterpret_cast<bf16_t*>(&packed); |
| 845 | #pragma unroll |
| 846 | for (int j = 0; j < 4; j++) { |
| 847 | float weight = __bfloat162float(output_norm_buf[h_base + j]); |
| 848 | values[j] = __float2bfloat16(acc32[si * VEC + j] * |
| 849 | output_rsigma * weight); |
| 850 | } |
| 851 | *reinterpret_cast<uint2*>(out_ptr + h_base) = packed; |
| 852 | continue; |
| 853 | } |
| 854 | } |
| 855 | int dt = si * CONSUMER_GROUPS + group; |
| 856 | if (dt >= NHT) continue; |
| 857 | int h_base = dt * K_TILE + k_local; |
| 858 | uint4 packed; |
| 859 | auto* values = reinterpret_cast<bf16_t*>(&packed); |
| 860 | #pragma unroll |
| 861 | for (int j = 0; j < VEC; j++) { |
| 862 | float weight = __bfloat162float(output_norm_buf[h_base + j]); |
| 863 | values[j] = |
| 864 | __float2bfloat16(acc32[si * VEC + j] * output_rsigma * weight); |
| 865 | } |
| 866 | *reinterpret_cast<uint4*>(out_ptr + h_base) = packed; |
| 867 | } |
| 868 | } |
| 869 | } |
| 870 | } |
| 871 | |
| 872 | __syncthreads(); |
| 873 | if (wid == 1) { |
| 874 | tmem_free(plan.tmem_base, TMEM_COLS_ALLOC); |
| 875 | } |
| 876 | #else |
| 877 | if (threadIdx.x == 0) { |
| 878 | printf("attn_res_fwd_online_v2_kernel requires sm_10x\n"); |
| 879 | } |
| 880 | #endif |
| 881 | } |
| 882 | |
| 883 | template <int H, int N, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false, |
| 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, |
| 891 | int block_stride_r = 0, int block_write_idx = -1) { |
| 892 | constexpr size_t smem_size = |
| 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 = |
| 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) { |
| 903 | cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, |
| 904 | smem_size); |
| 905 | } |
| 906 | attrs_set = true; |
| 907 | } |
| 908 | int grid = RELEASE_TMEM ? num_sm * 2 : num_sm; |
| 909 | cudaLaunchConfig_t config{}; |
| 910 | config.gridDim = grid; |
| 911 | config.blockDim = BLK; |
| 912 | config.dynamicSmemBytes = smem_size; |
| 913 | config.stream = stream; |
| 914 | cudaLaunchAttribute attrs[1]; |
| 915 | attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; |
| 916 | attrs[0].val.programmaticStreamSerializationAllowed = 1; |
| 917 | config.attrs = attrs; |
| 918 | config.numAttrs = 1; |
| 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, |
| 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, |
| 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, |
| 933 | std::optional<torch::stable::Tensor> output_norm_weight, |
| 934 | torch::stable::Tensor& output, int64_t num_blocks, |
| 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); |
| 940 | cudaDeviceProp const* properties = get_device_prop(); |
| 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 && |
| 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 ", |
| 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; |
| 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()), |
| 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, |
| 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)), |
| 988 | static_cast<int>(blocks.stride(1)), |
| 989 | static_cast<int>(block_write_idx)); |
| 990 | }; |
| 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; |
| 998 | case 2: |
| 999 | dispatch(std::integral_constant<int, 3>{}); |
| 1000 | break; |
| 1001 | case 3: |
| 1002 | dispatch(std::integral_constant<int, 4>{}); |
| 1003 | break; |
| 1004 | case 4: |
| 1005 | dispatch(std::integral_constant<int, 5>{}); |
| 1006 | break; |
| 1007 | case 5: |
| 1008 | dispatch(std::integral_constant<int, 6>{}); |
| 1009 | break; |
| 1010 | case 6: |
| 1011 | dispatch(std::integral_constant<int, 7>{}); |
| 1012 | break; |
| 1013 | case 7: |
| 1014 | dispatch(std::integral_constant<int, 8>{}); |
| 1015 | break; |
| 1016 | case 8: |
| 1017 | dispatch(std::integral_constant<int, 9>{}); |
| 1018 | break; |
| 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( |
| 1031 | error == cudaSuccess, |
| 1032 | "Kimi K3 AttnRes kernel launch failed: ", cudaGetErrorString(error)); |
| 1033 | } |