diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu index 36a297bec25f..f1e284118d6f 100644 --- a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu @@ -29,8 +29,10 @@ // Source-integrated from the NVIDIA+Moonshot jointly developed // Attention_residual kernel at e7f934124acc915575f9f7561f9d1e373ab43089. +#include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h" +#include #include #include #include @@ -38,151 +40,154 @@ #include #include #include +#include namespace { using bf16_t = __nv_bfloat16; - static constexpr int ATTN_RES_BLOCK = 256; static constexpr int ATTN_RES_WARPS = ATTN_RES_BLOCK / 32; -__inline__ __device__ float warp_reduce_sum(float val) { - #pragma unroll +__inline__ __device__ float warp_reduce_sum(float val) +{ +#pragma unroll for (int offset = 16; offset > 0; offset >>= 1) val += __shfl_xor_sync(0xffffffff, val, offset); return val; } -__inline__ __device__ float block_reduce_sum(float val, float* ws) { +__inline__ __device__ float block_reduce_sum(float val, float* ws) +{ int lane = threadIdx.x & 31; - int wid = threadIdx.x >> 5; + int wid = threadIdx.x >> 5; val = warp_reduce_sum(val); - if (lane == 0) ws[wid] = val; + if (lane == 0) + ws[wid] = val; __syncthreads(); val = (threadIdx.x < ATTN_RES_WARPS) ? ws[threadIdx.x] : 0.f; - if (wid == 0) val = warp_reduce_sum(val); + if (wid == 0) + val = warp_reduce_sum(val); return val; } -__device__ __forceinline__ -const bf16_t* v_addr(const bf16_t* block_res, const bf16_t* layer_res, - int n, int N, int t, int b, int T, int B, int H) { +__device__ __forceinline__ bf16_t const* v_addr( + bf16_t const* block_res, bf16_t const* layer_res, int n, int N, int t, int b, int T, int B, int H) +{ if (n < N - 1) - return block_res + (((long long)n * T + t) * B + b) * H; - return layer_res + ((long long)t * B + b) * H; + return block_res + (((long long) n * T + t) * B + b) * H; + return layer_res + ((long long) t * B + b) * H; } -namespace sm100 { - +namespace sm100 +{ CUTE_DEVICE -void tcgen05_after_thread_sync() { +void tcgen05_after_thread_sync() +{ asm volatile("tcgen05.fence::after_thread_sync;"); } CUTE_DEVICE -void umma_arrive_noelect(uint64_t& bar_ptr) { +void umma_arrive_noelect(uint64_t& bar_ptr) +{ uint64_t bar_addr = cute::cast_smem_ptr_to_uint(&bar_ptr); - asm volatile( - "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];" - : - : "l"(bar_addr)); + asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];" : : "l"(bar_addr)); } CUTE_DEVICE -float2 float2_sub(const float2& a, const float2& b) { +float2 float2_sub(float2 const& a, float2 const& b) +{ float2 c; - asm volatile( - "sub.f32x2 %0, %1, %2;\n" - : "=l"(reinterpret_cast(c)) - : "l"(reinterpret_cast(a)), - "l"(reinterpret_cast(b))); + asm volatile("sub.f32x2 %0, %1, %2;\n" + : "=l"(reinterpret_cast(c)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b))); return c; } CUTE_DEVICE -float2 float2_mul(const float2& a, const float2& b) { +float2 float2_mul(float2 const& a, float2 const& b) +{ float2 c; - asm volatile( - "mul.f32x2 %0, %1, %2;\n" - : "=l"(reinterpret_cast(c)) - : "l"(reinterpret_cast(a)), - "l"(reinterpret_cast(b))); + asm volatile("mul.f32x2 %0, %1, %2;\n" + : "=l"(reinterpret_cast(c)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b))); return c; } CUTE_DEVICE -float2 float2_fma(const float2& a, const float2& b, const float2& c) { +float2 float2_fma(float2 const& a, float2 const& b, float2 const& c) +{ float2 d; - asm volatile( - "fma.rn.f32x2 %0, %1, %2, %3;\n" - : "=l"(reinterpret_cast(d)) - : "l"(reinterpret_cast(a)), - "l"(reinterpret_cast(b)), - "l"(reinterpret_cast(c))); + asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n" + : "=l"(reinterpret_cast(d)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b)), + "l"(reinterpret_cast(c))); return d; } CUTE_DEVICE -float2 float2_add(const float2& a, const float2& b) { +float2 float2_add(float2 const& a, float2 const& b) +{ float2 c; - asm volatile( - "add.rn.f32x2 %0, %1, %2;\n" - : "=l"(reinterpret_cast(c)) - : "l"(reinterpret_cast(a)), - "l"(reinterpret_cast(b))); + asm volatile("add.rn.f32x2 %0, %1, %2;\n" + : "=l"(reinterpret_cast(c)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b))); return c; } template -CUTE_DEVICE void tmem_ld_32dp32bNx(uint32_t const& src_addr, T* dst_ptr_) { +CUTE_DEVICE void tmem_ld_32dp32bNx(uint32_t const& src_addr, T* dst_ptr_) +{ uint32_t* dst_ptr = reinterpret_cast(dst_ptr_); - if constexpr (N == 8) { + if constexpr (N == 8) + { asm volatile( "tcgen05.ld.sync.aligned.32x32b.x8.b32" "{%0, %1, %2, %3, %4, %5, %6, %7}," "[%8];\n" - : "=r"(dst_ptr[0]), "=r"(dst_ptr[1]), "=r"(dst_ptr[2]), - "=r"(dst_ptr[3]), "=r"(dst_ptr[4]), "=r"(dst_ptr[5]), - "=r"(dst_ptr[6]), "=r"(dst_ptr[7]) + : "=r"(dst_ptr[0]), "=r"(dst_ptr[1]), "=r"(dst_ptr[2]), "=r"(dst_ptr[3]), "=r"(dst_ptr[4]), + "=r"(dst_ptr[5]), "=r"(dst_ptr[6]), "=r"(dst_ptr[7]) : "r"(src_addr)); - } else { + } + else + { static_assert(N == 4, "attn_res TMEM helpers support x4 and x8"); asm volatile( "tcgen05.ld.sync.aligned.32x32b.x4.b32" "{%0, %1, %2, %3}, [%4];\n" - : "=r"(dst_ptr[0]), "=r"(dst_ptr[1]), - "=r"(dst_ptr[2]), "=r"(dst_ptr[3]) + : "=r"(dst_ptr[0]), "=r"(dst_ptr[1]), "=r"(dst_ptr[2]), "=r"(dst_ptr[3]) : "r"(src_addr)); } } template -CUTE_DEVICE void tmem_st_32dp32bNx(uint32_t const& dst_addr, T* src_ptr_) { +CUTE_DEVICE void tmem_st_32dp32bNx(uint32_t const& dst_addr, T* src_ptr_) +{ uint32_t* src_ptr = reinterpret_cast(src_ptr_); - if constexpr (N == 8) { + if constexpr (N == 8) + { asm volatile( "tcgen05.st.sync.aligned.32x32b.x8.b32" "[%8], {%0, %1, %2, %3, %4, %5, %6, %7};\n" : - : "r"(src_ptr[0]), "r"(src_ptr[1]), "r"(src_ptr[2]), - "r"(src_ptr[3]), "r"(src_ptr[4]), "r"(src_ptr[5]), - "r"(src_ptr[6]), "r"(src_ptr[7]), "r"(dst_addr)); - } else { + : "r"(src_ptr[0]), "r"(src_ptr[1]), "r"(src_ptr[2]), "r"(src_ptr[3]), "r"(src_ptr[4]), "r"(src_ptr[5]), + "r"(src_ptr[6]), "r"(src_ptr[7]), "r"(dst_addr)); + } + else + { static_assert(N == 4, "attn_res TMEM helpers support x4 and x8"); asm volatile( "tcgen05.st.sync.aligned.32x32b.x4.b32" "[%4], {%0, %1, %2, %3};\n" : - : "r"(src_ptr[0]), "r"(src_ptr[1]), - "r"(src_ptr[2]), "r"(src_ptr[3]), "r"(dst_addr)); + : "r"(src_ptr[0]), "r"(src_ptr[1]), "r"(src_ptr[2]), "r"(src_ptr[3]), "r"(dst_addr)); } } - -namespace fwd_prod_v2 { +namespace fwd_prod_v2 +{ using namespace cute; @@ -190,16 +195,17 @@ constexpr int K_TILE = 1024; constexpr int N_MAX = 12; constexpr int N_CHUNK_DEFAULT = 4; constexpr int CHUNK_DEPTH = 2; -constexpr int BLK = 288; // 1 producer warp + 8 consumer warps -constexpr int CONSUMER_THREADS = BLK - 32; // 256 +constexpr int BLK = 288; // 1 producer warp + 8 consumer warps +constexpr int CONSUMER_THREADS = BLK - 32; // 256 constexpr int CONSUMER_WARPS = CONSUMER_THREADS / 32; -constexpr int CONSUMER_GROUPS = 2; // two 128-thread consumer groups +constexpr int CONSUMER_GROUPS = 2; // two 128-thread consumer groups constexpr int CONSUMER_THREADS_PER_GROUP = CONSUMER_THREADS / CONSUMER_GROUPS; constexpr int TMEM_Q_COLS_PER_GROUP = 32; constexpr int TMEM_Q_COLS_TOTAL = 2 * TMEM_Q_COLS_PER_GROUP; template -struct FwdSmemPlan { +struct FwdSmemPlan +{ alignas(16) uint64_t bar_ready[CHUNK_DEPTH]; alignas(16) uint64_t bar_consumed[CHUNK_DEPTH]; alignas(16) float2 ws_stats[CONSUMER_WARPS][NC]; @@ -207,28 +213,20 @@ struct FwdSmemPlan { uint32_t tmem_base; }; -__device__ __forceinline__ -void cp_async_bulk(void* smem_dst, const void* gmem_src, int bytes, uint64_t& mbar) { +__device__ __forceinline__ void cp_async_bulk(void* smem_dst, void const* gmem_src, int bytes, uint64_t& mbar) +{ uint32_t s = cute::cast_smem_ptr_to_uint(smem_dst); uint32_t m = cute::cast_smem_ptr_to_uint(&mbar); - asm volatile( - "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n" - :: "r"(s), "l"(gmem_src), "r"(bytes), "r"(m) : "memory"); + asm volatile("cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n" ::"r"(s), + "l"(gmem_src), "r"(bytes), "r"(m) + : "memory"); } -template -__global__ void __launch_bounds__(BLK, 1) -attn_res_fwd_online_v2_kernel( - const bf16_t* __restrict__ block_res, - const bf16_t* __restrict__ layer_res, - const bf16_t* __restrict__ res_w, - const bf16_t* __restrict__ rms_w, - bf16_t* __restrict__ output, - float* __restrict__ rsigma_out, - float* __restrict__ probs_out, - float* __restrict__ logits_out, - int N, int T, int B, float rms_eps) +template +__global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t const* __restrict__ block_res, + bf16_t const* __restrict__ layer_res, bf16_t const* __restrict__ res_w, bf16_t const* __restrict__ rms_w, + bf16_t* __restrict__ output, float* __restrict__ rsigma_out, float* __restrict__ probs_out, + float* __restrict__ logits_out, int N, int T, int B, float rms_eps) { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 constexpr float LOG2_E = 1.4426950408889634f; @@ -237,292 +235,287 @@ attn_res_fwd_online_v2_kernel( constexpr int NHT = H / K_TILE; constexpr int SLICES_PER_GROUP = (NHT + CONSUMER_GROUPS - 1) / CONSUMER_GROUPS; constexpr int VEC = 8; - constexpr int ACC_PER_THREAD = - H == 7168 ? 28 : SLICES_PER_GROUP * VEC; - constexpr int TMEM_V_COLS_PER_GROUP = - SLICES_PER_GROUP * N_CHUNK * VEC; - constexpr int TMEM_COLS_TOTAL = - CONSUMER_GROUPS * TMEM_V_COLS_PER_GROUP; + constexpr int ACC_PER_THREAD = H == 7168 ? 28 : SLICES_PER_GROUP * VEC; + constexpr int TMEM_V_COLS_PER_GROUP = SLICES_PER_GROUP * N_CHUNK * VEC; + constexpr int TMEM_COLS_TOTAL = CONSUMER_GROUPS * TMEM_V_COLS_PER_GROUP; constexpr int TMEM_COLS_ALLOC = 256; static_assert(TMEM_COLS_TOTAL <= TMEM_COLS_ALLOC); static_assert(H >= 4096 && H <= 8192); static_assert(H % K_TILE == 0); - const int tid = threadIdx.x; - const int wid = tid >> 5; - const int lane = tid & 31; - const int TB = FULL_N12 ? 1024 : T; - const int num_ctas = gridDim.x; - const int num_chunks = (N + N_CHUNK - 1) / N_CHUNK; + int const tid = threadIdx.x; + int const wid = tid >> 5; + int const lane = tid & 31; + int const TB = FULL_N12 ? 1024 : T; + int const num_ctas = gridDim.x; + int const num_chunks = (N + N_CHUNK - 1) / N_CHUNK; - const int comp_wid = wid - 1; - const int comp_tid = tid - 32; - const int group = (comp_wid >= 4) ? 1 : 0; - const int ct_in_group = (comp_tid >= 0) ? (comp_tid & (CONSUMER_THREADS_PER_GROUP - 1)) : -1; - const int k_local = ct_in_group * VEC; + int const comp_wid = wid - 1; + int const comp_tid = tid - 32; + int const group = (comp_wid >= 4) ? 1 : 0; + int const ct_in_group = (comp_tid >= 0) ? (comp_tid & (CONSUMER_THREADS_PER_GROUP - 1)) : -1; + int const k_local = ct_in_group * VEC; extern __shared__ char smem_raw[]; - bf16_t* v_bufs = reinterpret_cast(smem_raw); // [NUM_BUFS][H] - constexpr size_t V_BYTES = (size_t)NUM_BUFS * H * sizeof(bf16_t); + bf16_t* v_bufs = reinterpret_cast(smem_raw); // [NUM_BUFS][H] + constexpr size_t V_BYTES = (size_t) NUM_BUFS * H * sizeof(bf16_t); FwdSmemPlan& plan = *reinterpret_cast*>(smem_raw + V_BYTES); - auto slot_of = [](long long gci, int n) { - return (int)(gci % CHUNK_DEPTH) * N_CHUNK + n; - }; - auto phase_of = [](long long gci) { - return (int)((gci / CHUNK_DEPTH) & 1); - }; - auto buf_ptr = [&](int slot) -> bf16_t* { - return v_bufs + slot * H; - }; - - if (wid == 0 && elect_one_sync()) { - #pragma unroll - for (int i = 0; i < CHUNK_DEPTH; i++) { + auto slot_of = [](long long gci, int n) { return (int) (gci % CHUNK_DEPTH) * N_CHUNK + n; }; + auto phase_of = [](long long gci) { return (int) ((gci / CHUNK_DEPTH) & 1); }; + auto buf_ptr = [&](int slot) -> bf16_t* { return v_bufs + slot * H; }; + + if (wid == 0 && elect_one_sync()) + { +#pragma unroll + for (int i = 0; i < CHUNK_DEPTH; i++) + { cute::initialize_barrier(plan.bar_ready[i], 1); cute::initialize_barrier(plan.bar_consumed[i], CONSUMER_WARPS); } cutlass::arch::fence_barrier_init(); } - if (wid == 1) { + if (wid == 1) + { cute::TMEM::Allocator1Sm alloc; alloc.allocate(TMEM_COLS_ALLOC, &plan.tmem_base); - if constexpr (RELEASE_TMEM) { + if constexpr (RELEASE_TMEM) + { alloc.release_allocation_lock(); } } __syncthreads(); - const uint32_t my_v_tmem = (comp_tid >= 0) - ? (plan.tmem_base + group * TMEM_V_COLS_PER_GROUP) - : 0; + const uint32_t my_v_tmem = (comp_tid >= 0) ? (plan.tmem_base + group * TMEM_V_COLS_PER_GROUP) : 0; float q_cache[ACC_PER_THREAD]; - if (comp_tid >= 0) { - #pragma unroll - for (int si = 0; si < SLICES_PER_GROUP; si++) { - if constexpr (H == 7168) { - if (si == SLICES_PER_GROUP - 1) { - int h_base = 6 * K_TILE + group * (K_TILE / 2) + - ct_in_group * 4; - #pragma unroll - for (int j = 0; j < 4; j++) { + if (comp_tid >= 0) + { +#pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) + { + if constexpr (H == 7168) + { + if (si == SLICES_PER_GROUP - 1) + { + int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; +#pragma unroll + for (int j = 0; j < 4; j++) + { int h = h_base + j; - q_cache[si * VEC + j] = - __bfloat162float(rms_w[h]) * - __bfloat162float(res_w[h]); + q_cache[si * VEC + j] = __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]); } continue; } } int dt = si * CONSUMER_GROUPS + group; - if (dt >= NHT) continue; + if (dt >= NHT) + continue; int h_base = dt * K_TILE + k_local; - #pragma unroll - for (int j = 0; j < VEC; j++) { +#pragma unroll + for (int j = 0; j < VEC; j++) + { int h = h_base + j; - q_cache[si * VEC + j] = - __bfloat162float(rms_w[h]) * - __bfloat162float(res_w[h]); + q_cache[si * VEC + j] = __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]); } } } - if (wid == 0) { - if (elect_one_sync()) { + if (wid == 0) + { + if (elect_one_sync()) + { long long gci = 0; - for (int tb = blockIdx.x; tb < TB; tb += num_ctas) { - for (int ci = 0; ci < num_chunks; ci++, gci++) { + for (int tb = blockIdx.x; tb < TB; tb += num_ctas) + { + for (int ci = 0; ci < num_chunks; ci++, gci++) + { int ns = ci * N_CHUNK; int an = FULL_N12 ? N_CHUNK : min(N_CHUNK, N - ns); - int chunk_slot = (int)(gci % CHUNK_DEPTH); + int chunk_slot = (int) (gci % CHUNK_DEPTH); int pc = phase_of(gci); - cute::wait_barrier( - plan.bar_consumed[chunk_slot], pc ^ 1); - cute::set_barrier_transaction_bytes( - plan.bar_ready[chunk_slot], - an * H * (int)sizeof(bf16_t)); - #pragma unroll - for (int n = 0; n < N_CHUNK; n++) { - if constexpr (!FULL_N12) { - if (n >= an) continue; + cute::wait_barrier(plan.bar_consumed[chunk_slot], pc ^ 1); + cute::set_barrier_transaction_bytes(plan.bar_ready[chunk_slot], an * H * (int) sizeof(bf16_t)); +#pragma unroll + for (int n = 0; n < N_CHUNK; n++) + { + if constexpr (!FULL_N12) + { + if (n >= an) + continue; } int slot = slot_of(gci, n); - const int ng = ns + n; - const bf16_t* src = - (ng < (FULL_N12 ? 11 : N - 1)) - ? block_res + - ((long long)ng * T + tb) * H - : layer_res + (long long)tb * H; - cp_async_bulk( - buf_ptr(slot), src, H * sizeof(bf16_t), - plan.bar_ready[chunk_slot]); + int const ng = ns + n; + bf16_t const* src = (ng < (FULL_N12 ? 11 : N - 1)) ? block_res + ((long long) ng * T + tb) * H + : layer_res + (long long) tb * H; + cp_async_bulk(buf_ptr(slot), src, H * sizeof(bf16_t), plan.bar_ready[chunk_slot]); } } } } - } else { + } + else + { float acc32[ACC_PER_THREAD] = {}; float eps_cache; asm volatile("mov.b32 %0, %1;" : "=f"(eps_cache) : "f"(rms_eps)); long long gci = 0; - for (int tb = blockIdx.x; tb < TB; tb += num_ctas) { + for (int tb = blockIdx.x; tb < TB; tb += num_ctas) + { float m_running = -FLT_MAX; float s_running = 0.f; - #pragma unroll - for (int i = 0; i < ACC_PER_THREAD; i++) { +#pragma unroll + for (int i = 0; i < ACC_PER_THREAD; i++) + { acc32[i] = 0.f; } - for (int ci = 0; ci < num_chunks; ci++, gci++) { + for (int ci = 0; ci < num_chunks; ci++, gci++) + { int ns = ci * N_CHUNK; int an = FULL_N12 ? N_CHUNK : min(N_CHUNK, N - ns); - int chunk_slot = (int)(gci % CHUNK_DEPTH); + int chunk_slot = (int) (gci % CHUNK_DEPTH); int pr = phase_of(gci); float2 sq_local[N_CHUNK] = {}; float2 dot_local[N_CHUNK] = {}; cute::wait_barrier(plan.bar_ready[chunk_slot], pr); - auto pass_A_body = [&](auto AN_TOK) { + auto pass_A_body = [&](auto AN_TOK) + { constexpr int AN = decltype(AN_TOK)::value; - #pragma unroll - for (int si = 0; si < SLICES_PER_GROUP; si++) { - if constexpr (H == 7168) { - if (si == SLICES_PER_GROUP - 1) { - int h_base = 6 * K_TILE + - group * (K_TILE / 2) + - ct_in_group * 4; +#pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) + { + if constexpr (H == 7168) + { + if (si == SLICES_PER_GROUP - 1) + { + int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; const float* qv = &q_cache[si * VEC]; - #pragma unroll - for (int n = 0; n < AN; n++) { +#pragma unroll + for (int n = 0; n < AN; n++) + { int slot = slot_of(gci, n); - int2 vp = *reinterpret_cast( - buf_ptr(slot) + h_base); - __nv_bfloat162* v2 = - reinterpret_cast<__nv_bfloat162*>(&vp); - float2 f[2] = { - __bfloat1622float2(v2[0]), - __bfloat1622float2(v2[1])}; - if constexpr (FULL_N12) { - if (n == AN - 1 && lane == 0) { - cute::arrive_barrier( - plan.bar_consumed[chunk_slot]); + int2 vp = *reinterpret_cast(buf_ptr(slot) + h_base); + __nv_bfloat162* v2 = reinterpret_cast<__nv_bfloat162*>(&vp); + float2 f[2] = {__bfloat1622float2(v2[0]), __bfloat1622float2(v2[1])}; + if constexpr (FULL_N12) + { + if (n == AN - 1 && lane == 0) + { + cute::arrive_barrier(plan.bar_consumed[chunk_slot]); } } tmem_st_32dp32bNx<4>( - my_v_tmem + - (si * N_CHUNK + n) * VEC, - reinterpret_cast(f)); - sq_local[n] = - float2_fma(f[0], f[0], sq_local[n]); - sq_local[n] = - float2_fma(f[1], f[1], sq_local[n]); - dot_local[n] = float2_fma( - f[0], make_float2(qv[0], qv[1]), - dot_local[n]); - dot_local[n] = float2_fma( - f[1], make_float2(qv[2], qv[3]), - dot_local[n]); + my_v_tmem + (si * N_CHUNK + n) * VEC, reinterpret_cast(f)); + sq_local[n] = float2_fma(f[0], f[0], sq_local[n]); + sq_local[n] = float2_fma(f[1], f[1], sq_local[n]); + dot_local[n] = float2_fma(f[0], make_float2(qv[0], qv[1]), dot_local[n]); + dot_local[n] = float2_fma(f[1], make_float2(qv[2], qv[3]), dot_local[n]); } continue; } } int dt = si * CONSUMER_GROUPS + group; - if (dt >= NHT) continue; + if (dt >= NHT) + continue; const float* qv = &q_cache[si * VEC]; - #pragma unroll - for (int n = 0; n < AN; n++) { +#pragma unroll + for (int n = 0; n < AN; n++) + { int slot = slot_of(gci, n); - int4 vp = *reinterpret_cast( - buf_ptr(slot) + dt * K_TILE + k_local); + int4 vp = *reinterpret_cast(buf_ptr(slot) + dt * K_TILE + k_local); __nv_bfloat162* v2 = reinterpret_cast<__nv_bfloat162*>(&vp); - float2 f[4] = { - __bfloat1622float2(v2[0]), - __bfloat1622float2(v2[1]), - __bfloat1622float2(v2[2]), - __bfloat1622float2(v2[3])}; - tmem_st_32dp32bNx( - my_v_tmem + - (si * N_CHUNK + n) * VEC, - reinterpret_cast(f)); + float2 f[4] = {__bfloat1622float2(v2[0]), __bfloat1622float2(v2[1]), + __bfloat1622float2(v2[2]), __bfloat1622float2(v2[3])}; + tmem_st_32dp32bNx(my_v_tmem + (si * N_CHUNK + n) * VEC, reinterpret_cast(f)); sq_local[n] = float2_fma(f[0], f[0], sq_local[n]); sq_local[n] = float2_fma(f[1], f[1], sq_local[n]); sq_local[n] = float2_fma(f[2], f[2], sq_local[n]); sq_local[n] = float2_fma(f[3], f[3], sq_local[n]); - dot_local[n] = float2_fma( - f[0], make_float2(qv[0], qv[1]), dot_local[n]); - dot_local[n] = float2_fma( - f[1], make_float2(qv[2], qv[3]), dot_local[n]); - dot_local[n] = float2_fma( - f[2], make_float2(qv[4], qv[5]), dot_local[n]); - dot_local[n] = float2_fma( - f[3], make_float2(qv[6], qv[7]), dot_local[n]); + dot_local[n] = float2_fma(f[0], make_float2(qv[0], qv[1]), dot_local[n]); + dot_local[n] = float2_fma(f[1], make_float2(qv[2], qv[3]), dot_local[n]); + dot_local[n] = float2_fma(f[2], make_float2(qv[4], qv[5]), dot_local[n]); + dot_local[n] = float2_fma(f[3], make_float2(qv[6], qv[7]), dot_local[n]); } } - if constexpr (!FULL_N12) { + if constexpr (!FULL_N12) + { cutlass::arch::fence_view_async_tmem_store(); } }; - if constexpr (FULL_N12) { + if constexpr (FULL_N12) + { pass_A_body(std::integral_constant{}); - } else if constexpr (NC == 4) { - switch (an) { - case 4: pass_A_body(std::integral_constant{}); break; - case 3: pass_A_body(std::integral_constant{}); break; - case 2: pass_A_body(std::integral_constant{}); break; - case 1: pass_A_body(std::integral_constant{}); break; - default: __builtin_unreachable(); + } + else if constexpr (NC == 4) + { + switch (an) + { + case 4: pass_A_body(std::integral_constant{}); break; + case 3: pass_A_body(std::integral_constant{}); break; + case 2: pass_A_body(std::integral_constant{}); break; + case 1: pass_A_body(std::integral_constant{}); break; + default: __builtin_unreachable(); } - } else if constexpr (NC == 3) { - switch (an) { - case 3: pass_A_body(std::integral_constant{}); break; - case 2: pass_A_body(std::integral_constant{}); break; - case 1: pass_A_body(std::integral_constant{}); break; - default: __builtin_unreachable(); + } + else if constexpr (NC == 3) + { + switch (an) + { + case 3: pass_A_body(std::integral_constant{}); break; + case 2: pass_A_body(std::integral_constant{}); break; + case 1: pass_A_body(std::integral_constant{}); break; + default: __builtin_unreachable(); } - } else { + } + else + { static_assert(NC == 2); - switch (an) { - case 2: pass_A_body(std::integral_constant{}); break; - case 1: pass_A_body(std::integral_constant{}); break; - default: __builtin_unreachable(); + switch (an) + { + case 2: pass_A_body(std::integral_constant{}); break; + case 1: pass_A_body(std::integral_constant{}); break; + default: __builtin_unreachable(); } } - if constexpr (!FULL_N12) { - if (lane == 0) { - cute::arrive_barrier( - plan.bar_consumed[chunk_slot]); + if constexpr (!FULL_N12) + { + if (lane == 0) + { + cute::arrive_barrier(plan.bar_consumed[chunk_slot]); } } float2 reduce_pair[N_CHUNK]; - #pragma unroll - for (int n = 0; n < N_CHUNK; n++) { - reduce_pair[n] = make_float2( - sq_local[n].x + sq_local[n].y, - dot_local[n].x + dot_local[n].y); +#pragma unroll + for (int n = 0; n < N_CHUNK; n++) + { + reduce_pair[n] = make_float2(sq_local[n].x + sq_local[n].y, dot_local[n].x + dot_local[n].y); } - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - #pragma unroll - for (int n = 0; n < N_CHUNK; n++) { - uint64_t packed = - reinterpret_cast(reduce_pair[n]); - packed = __shfl_xor_sync( - 0xffffffff, packed, offset); +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + { +#pragma unroll + for (int n = 0; n < N_CHUNK; n++) + { + uint64_t packed = reinterpret_cast(reduce_pair[n]); + packed = __shfl_xor_sync(0xffffffff, packed, offset); float2 other = reinterpret_cast(packed); - reduce_pair[n] = - float2_add(reduce_pair[n], other); + reduce_pair[n] = float2_add(reduce_pair[n], other); } } - if constexpr (FULL_N12) { + if constexpr (FULL_N12) + { cutlass::arch::fence_view_async_tmem_store(); } - if (lane == 0) { - #pragma unroll - for (int n = 0; n < N_CHUNK; n++) { + if (lane == 0) + { +#pragma unroll + for (int n = 0; n < N_CHUNK; n++) + { plan.ws_stats[comp_wid][n] = reduce_pair[n]; } } @@ -530,37 +523,46 @@ attn_res_fwd_online_v2_kernel( float local_rsig = 0.f; float local_logit = 0.f; - auto cross_warp_tail = [&](int n) { + auto cross_warp_tail = [&](int n) + { float2 totals = {}; - #pragma unroll - for (int w = 0; w < CONSUMER_WARPS; w++) { - totals = float2_add( - totals, plan.ws_stats[w][n]); +#pragma unroll + for (int w = 0; w < CONSUMER_WARPS; w++) + { + totals = float2_add(totals, plan.ws_stats[w][n]); } local_rsig = rsqrtf(totals.x / H + eps_cache); local_logit = totals.y * local_rsig; }; - if constexpr (FULL_N12) { + if constexpr (FULL_N12) + { cross_warp_tail(lane & (N_CHUNK - 1)); - } else if (lane < N_CHUNK) { + } + else if (lane < N_CHUNK) + { cross_warp_tail(lane); } float logit_n[N_CHUNK]; - #pragma unroll - for (int n = 0; n < N_CHUNK; n++) { - logit_n[n] = __shfl_sync( - 0xffffffff, local_logit, n); +#pragma unroll + for (int n = 0; n < N_CHUNK; n++) + { + logit_n[n] = __shfl_sync(0xffffffff, local_logit, n); } float m_chunk = -FLT_MAX; - if constexpr (FULL_N12) { + if constexpr (FULL_N12) + { float m01 = fmaxf(logit_n[0], logit_n[1]); float m23 = fmaxf(logit_n[2], logit_n[3]); m_chunk = fmaxf(m01, m23); - } else { - #pragma unroll - for (int n = 0; n < N_CHUNK; n++) { - if (n < an) { + } + else + { +#pragma unroll + for (int n = 0; n < N_CHUNK; n++) + { + if (n < an) + { m_chunk = fmaxf(m_chunk, logit_n[n]); } } @@ -569,62 +571,66 @@ attn_res_fwd_online_v2_kernel( float corr = exp2f((m_running - m_new) * LOG2_E); float w_n[N_CHUNK] = {}; float w_sum = 0.f; - if constexpr (FULL_N12) { - #pragma unroll - for (int n = 0; n < N_CHUNK; n++) { - w_n[n] = exp2f( - (logit_n[n] - m_new) * LOG2_E); + if constexpr (FULL_N12) + { +#pragma unroll + for (int n = 0; n < N_CHUNK; n++) + { + w_n[n] = exp2f((logit_n[n] - m_new) * LOG2_E); } - w_sum = - (w_n[0] + w_n[1]) + (w_n[2] + w_n[3]); - } else { - #pragma unroll - for (int n = 0; n < N_CHUNK; n++) { - if (n < an) { - w_n[n] = exp2f( - (logit_n[n] - m_new) * LOG2_E); + w_sum = (w_n[0] + w_n[1]) + (w_n[2] + w_n[3]); + } + else + { +#pragma unroll + for (int n = 0; n < N_CHUNK; n++) + { + if (n < an) + { + w_n[n] = exp2f((logit_n[n] - m_new) * LOG2_E); w_sum += w_n[n]; } } } - auto pass_B_body = [&](auto AN_TOK) { + auto pass_B_body = [&](auto AN_TOK) + { constexpr int AN = decltype(AN_TOK)::value; - #pragma unroll - for (int si = 0; si < SLICES_PER_GROUP; si++) { - if constexpr (H == 7168) { - if (si == SLICES_PER_GROUP - 1) { - float2 corr2 = - make_float2(corr, corr); +#pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) + { + if constexpr (H == 7168) + { + if (si == SLICES_PER_GROUP - 1) + { + float2 corr2 = make_float2(corr, corr); float2 a[2]; - #pragma unroll - for (int j = 0; j < 2; j++) { - float2 old = make_float2( - acc32[si * VEC + 2 * j], - acc32[si * VEC + 2 * j + 1]); +#pragma unroll + for (int j = 0; j < 2; j++) + { + float2 old = make_float2(acc32[si * VEC + 2 * j], acc32[si * VEC + 2 * j + 1]); a[j] = float2_mul(old, corr2); } float2 f_cache[AN][2]; - #pragma unroll - for (int n = 0; n < AN; n++) { +#pragma unroll + for (int n = 0; n < AN; n++) + { tmem_ld_32dp32bNx<4>( - my_v_tmem + - (si * N_CHUNK + n) * VEC, - reinterpret_cast(f_cache[n])); + my_v_tmem + (si * N_CHUNK + n) * VEC, reinterpret_cast(f_cache[n])); } - #pragma unroll - for (int n = 0; n < AN; n++) { - float2 wn = - make_float2( - w_n[n], w_n[n]); - #pragma unroll - for (int j = 0; j < 2; j++) { - a[j] = float2_fma( - wn, f_cache[n][j], a[j]); +#pragma unroll + for (int n = 0; n < AN; n++) + { + float2 wn = make_float2(w_n[n], w_n[n]); +#pragma unroll + for (int j = 0; j < 2; j++) + { + a[j] = float2_fma(wn, f_cache[n][j], a[j]); } } - #pragma unroll - for (int j = 0; j < 2; j++) { +#pragma unroll + for (int j = 0; j < 2; j++) + { acc32[si * VEC + 2 * j] = a[j].x; acc32[si * VEC + 2 * j + 1] = a[j].y; } @@ -632,139 +638,153 @@ attn_res_fwd_online_v2_kernel( } } int dt = si * CONSUMER_GROUPS + group; - if (dt >= NHT) continue; + if (dt >= NHT) + continue; float2 a[VEC / 2]; - float2 corr2 = - make_float2(corr, corr); - #pragma unroll - for (int j = 0; j < VEC / 2; j++) { - float2 old = make_float2( - acc32[si * VEC + 2 * j], - acc32[si * VEC + 2 * j + 1]); + float2 corr2 = make_float2(corr, corr); +#pragma unroll + for (int j = 0; j < VEC / 2; j++) + { + float2 old = make_float2(acc32[si * VEC + 2 * j], acc32[si * VEC + 2 * j + 1]); a[j] = float2_mul(old, corr2); } float2 f_cache[AN][VEC / 2]; - #pragma unroll - for (int n = 0; n < AN; n++) { +#pragma unroll + for (int n = 0; n < AN; n++) + { tmem_ld_32dp32bNx( - my_v_tmem + - (si * N_CHUNK + n) * VEC, - reinterpret_cast(f_cache[n])); + my_v_tmem + (si * N_CHUNK + n) * VEC, reinterpret_cast(f_cache[n])); } - #pragma unroll - for (int n = 0; n < AN; n++) { +#pragma unroll + for (int n = 0; n < AN; n++) + { float2 wn = make_float2(w_n[n], w_n[n]); - #pragma unroll - for (int j = 0; j < VEC / 2; j++) { - a[j] = float2_fma( - wn, f_cache[n][j], a[j]); +#pragma unroll + for (int j = 0; j < VEC / 2; j++) + { + a[j] = float2_fma(wn, f_cache[n][j], a[j]); } } - #pragma unroll - for (int j = 0; j < VEC / 2; j++) { +#pragma unroll + for (int j = 0; j < VEC / 2; j++) + { acc32[si * VEC + 2 * j] = a[j].x; acc32[si * VEC + 2 * j + 1] = a[j].y; } } }; - if constexpr (FULL_N12) { - pass_B_body( - std::integral_constant{}); - } else if constexpr (NC == 4) { - switch (an) { - case 4: pass_B_body(std::integral_constant{}); break; - case 3: pass_B_body(std::integral_constant{}); break; - case 2: pass_B_body(std::integral_constant{}); break; - case 1: pass_B_body(std::integral_constant{}); break; - default: __builtin_unreachable(); + if constexpr (FULL_N12) + { + pass_B_body(std::integral_constant{}); + } + else if constexpr (NC == 4) + { + switch (an) + { + case 4: pass_B_body(std::integral_constant{}); break; + case 3: pass_B_body(std::integral_constant{}); break; + case 2: pass_B_body(std::integral_constant{}); break; + case 1: pass_B_body(std::integral_constant{}); break; + default: __builtin_unreachable(); } - } else if constexpr (NC == 3) { - switch (an) { - case 3: pass_B_body(std::integral_constant{}); break; - case 2: pass_B_body(std::integral_constant{}); break; - case 1: pass_B_body(std::integral_constant{}); break; - default: __builtin_unreachable(); + } + else if constexpr (NC == 3) + { + switch (an) + { + case 3: pass_B_body(std::integral_constant{}); break; + case 2: pass_B_body(std::integral_constant{}); break; + case 1: pass_B_body(std::integral_constant{}); break; + default: __builtin_unreachable(); } - } else { + } + else + { static_assert(NC == 2); - switch (an) { - case 2: pass_B_body(std::integral_constant{}); break; - case 1: pass_B_body(std::integral_constant{}); break; - default: __builtin_unreachable(); + switch (an) + { + case 2: pass_B_body(std::integral_constant{}); break; + case 1: pass_B_body(std::integral_constant{}); break; + default: __builtin_unreachable(); } } s_running = s_running * corr + w_sum; m_running = m_new; - if (comp_wid == 0 && lane < an) { + if (comp_wid == 0 && lane < an) + { int ng = ns + lane; - rsigma_out[(long long)ng * TB + tb] = local_rsig; + rsigma_out[(long long) ng * TB + tb] = local_rsig; plan.logits_all[ng] = local_logit; } } + // Publish the final chunk's plan.logits_all stores before the + // cross-lane reads in consumer warp 0 below (earlier chunks are + // covered by the NamedBarrier inside the loop). + __syncwarp(); float inv_s = 1.f / s_running; - bf16_t* out_ptr = output + (long long)tb * H; - #pragma unroll - for (int si = 0; si < SLICES_PER_GROUP; si++) { - if constexpr (H == 7168) { - if (si == SLICES_PER_GROUP - 1) { - int h_base = 6 * K_TILE + group * (K_TILE / 2) + - ct_in_group * 4; + bf16_t* out_ptr = output + (long long) tb * H; +#pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) + { + if constexpr (H == 7168) + { + if (si == SLICES_PER_GROUP - 1) + { + int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; uint2 ov; - __nv_bfloat162* ov2 = - reinterpret_cast<__nv_bfloat162*>(&ov); + __nv_bfloat162* ov2 = reinterpret_cast<__nv_bfloat162*>(&ov); float2 inv2 = make_float2(inv_s, inv_s); - #pragma unroll - for (int j = 0; j < 2; j++) { - float2 old = make_float2( - acc32[si * VEC + 2 * j], - acc32[si * VEC + 2 * j + 1]); - ov2[j] = __float22bfloat162_rn( - float2_mul(old, inv2)); +#pragma unroll + for (int j = 0; j < 2; j++) + { + float2 old = make_float2(acc32[si * VEC + 2 * j], acc32[si * VEC + 2 * j + 1]); + ov2[j] = __float22bfloat162_rn(float2_mul(old, inv2)); } *reinterpret_cast(out_ptr + h_base) = ov; continue; } } int dt = si * CONSUMER_GROUPS + group; - if (dt >= NHT) continue; + if (dt >= NHT) + continue; int h_base = dt * K_TILE + k_local; uint4 ov; - __nv_bfloat162* ov2 = - reinterpret_cast<__nv_bfloat162*>(&ov); + __nv_bfloat162* ov2 = reinterpret_cast<__nv_bfloat162*>(&ov); float2 inv2 = make_float2(inv_s, inv_s); - #pragma unroll - for (int j = 0; j < VEC / 2; j++) { - float2 old = make_float2( - acc32[si * VEC + 2 * j], - acc32[si * VEC + 2 * j + 1]); - ov2[j] = __float22bfloat162_rn( - float2_mul(old, inv2)); +#pragma unroll + for (int j = 0; j < VEC / 2; j++) + { + float2 old = make_float2(acc32[si * VEC + 2 * j], acc32[si * VEC + 2 * j + 1]); + ov2[j] = __float22bfloat162_rn(float2_mul(old, inv2)); } *reinterpret_cast(out_ptr + h_base) = ov; } - if (comp_wid == 0 && lane < (FULL_N12 ? 12 : N)) { - long long out_idx = (long long)lane * TB + tb; + if (comp_wid == 0 && lane < (FULL_N12 ? 12 : N)) + { + long long out_idx = (long long) lane * TB + tb; float lg = plan.logits_all[lane]; logits_out[out_idx] = lg; - probs_out[out_idx] = - exp2f((lg - m_running) * LOG2_E) * inv_s; + probs_out[out_idx] = exp2f((lg - m_running) * LOG2_E) * inv_s; } } } - if (wid > 0) { + if (wid > 0) + { cutlass::arch::NamedBarrier::sync(CONSUMER_THREADS, 2); } - if (wid == 1) { + if (wid == 1) + { cute::TMEM::Allocator1Sm alloc; alloc.free(plan.tmem_base, TMEM_COLS_ALLOC); } #else - if (cute::thread0()) printf("attn_res_fwd_online_v2_kernel requires sm_100a\n"); + if (cute::thread0()) + printf("attn_res_fwd_online_v2_kernel requires sm_100a\n"); #endif } @@ -772,15 +792,9 @@ attn_res_fwd_online_v2_kernel( // Tile multiple contiguous TB rows per CTA to reduce cp.async.bulk overhead. template __global__ void __launch_bounds__(BLK, 1) -attn_res_fwd_n1_ttile_kernel( - const bf16_t* __restrict__ layer_res, - const bf16_t* __restrict__ res_w, - const bf16_t* __restrict__ rms_w, - bf16_t* __restrict__ output, - float* __restrict__ rsigma_out, - float* __restrict__ probs_out, - float* __restrict__ logits_out, - int T, int B, float rms_eps) + attn_res_fwd_n1_ttile_kernel(bf16_t const* __restrict__ layer_res, bf16_t const* __restrict__ res_w, + bf16_t const* __restrict__ rms_w, bf16_t* __restrict__ output, float* __restrict__ rsigma_out, + float* __restrict__ probs_out, float* __restrict__ logits_out, int T, int B, float rms_eps) { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 constexpr int NHT = H / K_TILE; @@ -789,117 +803,124 @@ attn_res_fwd_n1_ttile_kernel( constexpr int ACC_PER_THREAD = SLICES_PER_GROUP * VEC; static_assert(H == 4096 || H == 8192); - const int tid = threadIdx.x; - const int wid = tid >> 5; - const int lane = tid & 31; - const int TB = T * B; - const int comp_wid = wid - 1; - const int comp_tid = tid - 32; - const int group = (comp_wid >= 4) ? 1 : 0; - const int ct_in_group = (comp_tid >= 0) ? (comp_tid & (CONSUMER_THREADS_PER_GROUP - 1)) : -1; - const int k_local = ct_in_group * VEC; + int const tid = threadIdx.x; + int const wid = tid >> 5; + int const lane = tid & 31; + int const TB = T * B; + int const comp_wid = wid - 1; + int const comp_tid = tid - 32; + int const group = (comp_wid >= 4) ? 1 : 0; + int const ct_in_group = (comp_tid >= 0) ? (comp_tid & (CONSUMER_THREADS_PER_GROUP - 1)) : -1; + int const k_local = ct_in_group * VEC; extern __shared__ char smem_raw[]; bf16_t* v_tiles = reinterpret_cast(smem_raw); - constexpr size_t V_BYTES = (size_t)CHUNK_DEPTH * TB_TILE * H * sizeof(bf16_t); + constexpr size_t V_BYTES = (size_t) CHUNK_DEPTH * TB_TILE * H * sizeof(bf16_t); FwdSmemPlan<1>& plan = *reinterpret_cast*>(smem_raw + V_BYTES); - auto phase_of = [](long long tile_i) { - return (int)((tile_i / CHUNK_DEPTH) & 1); - }; - auto tile_ptr = [&](int slot, int row) -> bf16_t* { - return v_tiles + ((slot * TB_TILE + row) * H); - }; + auto phase_of = [](long long tile_i) { return (int) ((tile_i / CHUNK_DEPTH) & 1); }; + auto tile_ptr = [&](int slot, int row) -> bf16_t* { return v_tiles + ((slot * TB_TILE + row) * H); }; - if (wid == 0 && elect_one_sync()) { - #pragma unroll - for (int i = 0; i < CHUNK_DEPTH; i++) { + if (wid == 0 && elect_one_sync()) + { +#pragma unroll + for (int i = 0; i < CHUNK_DEPTH; i++) + { cute::initialize_barrier(plan.bar_ready[i], 1); cute::initialize_barrier(plan.bar_consumed[i], CONSUMER_THREADS); } cutlass::arch::fence_barrier_init(); } - if (wid == 1) { + if (wid == 1) + { cute::TMEM::Allocator1Sm alloc; alloc.allocate(TMEM_Q_COLS_TOTAL, &plan.tmem_base); - if constexpr (RELEASE_TMEM) { + if constexpr (RELEASE_TMEM) + { alloc.release_allocation_lock(); } } __syncthreads(); - const uint32_t my_tmem = (comp_tid >= 0) - ? (plan.tmem_base + ((comp_wid >= 4) ? TMEM_Q_COLS_PER_GROUP : 0)) - : 0; + const uint32_t my_tmem = (comp_tid >= 0) ? (plan.tmem_base + ((comp_wid >= 4) ? TMEM_Q_COLS_PER_GROUP : 0)) : 0; - if (comp_tid >= 0) { + if (comp_tid >= 0) + { float q32[ACC_PER_THREAD]; - #pragma unroll - for (int si = 0; si < SLICES_PER_GROUP; si++) { +#pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) + { int dt = si * CONSUMER_GROUPS + group; - if (dt >= NHT) continue; + if (dt >= NHT) + continue; int h_base = dt * K_TILE + k_local; - #pragma unroll - for (int j = 0; j < VEC; j++) { +#pragma unroll + for (int j = 0; j < VEC; j++) + { int h = h_base + j; - q32[si * VEC + j] = - __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]); + q32[si * VEC + j] = __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]); } } - #pragma unroll - for (int si = 0; si < SLICES_PER_GROUP; si++) { +#pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) + { int dt = si * CONSUMER_GROUPS + group; - if (dt >= NHT) continue; + if (dt >= NHT) + continue; tmem_st_32dp32bNx(my_tmem + si * VEC, &q32[si * VEC]); } cutlass::arch::fence_view_async_tmem_store(); } __syncthreads(); - if (wid == 0) { - if (elect_one_sync()) { + if (wid == 0) + { + if (elect_one_sync()) + { long long tile_i = 0; - for (int tb0 = blockIdx.x * TB_TILE; tb0 < TB; - tb0 += gridDim.x * TB_TILE, tile_i++) { + for (int tb0 = blockIdx.x * TB_TILE; tb0 < TB; tb0 += gridDim.x * TB_TILE, tile_i++) + { int rows = min(TB_TILE, TB - tb0); - int slot = (int)(tile_i % CHUNK_DEPTH); + int slot = (int) (tile_i % CHUNK_DEPTH); int pc = phase_of(tile_i); cute::wait_barrier(plan.bar_consumed[slot], pc ^ 1); - cute::set_barrier_transaction_bytes( - plan.bar_ready[slot], rows * H * (int)sizeof(bf16_t)); - cp_async_bulk( - tile_ptr(slot, 0), - layer_res + (long long)tb0 * H, - rows * H * sizeof(bf16_t), + cute::set_barrier_transaction_bytes(plan.bar_ready[slot], rows * H * (int) sizeof(bf16_t)); + cp_async_bulk(tile_ptr(slot, 0), layer_res + (long long) tb0 * H, rows * H * sizeof(bf16_t), plan.bar_ready[slot]); } } - } else { + } + else + { long long tile_i = 0; - for (int tb0 = blockIdx.x * TB_TILE; tb0 < TB; - tb0 += gridDim.x * TB_TILE, tile_i++) { + for (int tb0 = blockIdx.x * TB_TILE; tb0 < TB; tb0 += gridDim.x * TB_TILE, tile_i++) + { int rows = min(TB_TILE, TB - tb0); - int slot = (int)(tile_i % CHUNK_DEPTH); + int slot = (int) (tile_i % CHUNK_DEPTH); int pc = phase_of(tile_i); cute::wait_barrier(plan.bar_ready[slot], pc); - #pragma unroll - for (int r = 0; r < TB_TILE; r++) { - if (r >= rows) continue; +#pragma unroll + for (int r = 0; r < TB_TILE; r++) + { + if (r >= rows) + continue; int tb = tb0 + r; bf16_t* row_ptr = tile_ptr(slot, r); - bf16_t* out_ptr = output + (long long)tb * H; + bf16_t* out_ptr = output + (long long) tb * H; float sq_local = 0.f; float dot_local = 0.f; - #pragma unroll - for (int si = 0; si < SLICES_PER_GROUP; si++) { +#pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) + { int dt = si * CONSUMER_GROUPS + group; - if (dt >= NHT) continue; + if (dt >= NHT) + continue; int h_base = dt * K_TILE + k_local; float qv[VEC]; tmem_ld_32dp32bNx(my_tmem + si * VEC, qv); - int4 vp = *reinterpret_cast(row_ptr + h_base); + int4 vp = *reinterpret_cast(row_ptr + h_base); *reinterpret_cast(out_ptr + h_base) = vp; __nv_bfloat162* v2 = reinterpret_cast<__nv_bfloat162*>(&vp); @@ -907,40 +928,38 @@ attn_res_fwd_n1_ttile_kernel( float2 f1 = __bfloat1622float2(v2[1]); float2 f2 = __bfloat1622float2(v2[2]); float2 f3 = __bfloat1622float2(v2[3]); - sq_local += - f0.x * f0.x + f0.y * f0.y + - f1.x * f1.x + f1.y * f1.y + - f2.x * f2.x + f2.y * f2.y + - f3.x * f3.x + f3.y * f3.y; - dot_local += - f0.x * qv[0] + f0.y * qv[1] + - f1.x * qv[2] + f1.y * qv[3] + - f2.x * qv[4] + f2.y * qv[5] + - f3.x * qv[6] + f3.y * qv[7]; + sq_local += f0.x * f0.x + f0.y * f0.y + f1.x * f1.x + f1.y * f1.y + f2.x * f2.x + f2.y * f2.y + + f3.x * f3.x + f3.y * f3.y; + dot_local += f0.x * qv[0] + f0.y * qv[1] + f1.x * qv[2] + f1.y * qv[3] + f2.x * qv[4] + f2.y * qv[5] + + f3.x * qv[6] + f3.y * qv[7]; } - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + { sq_local += __shfl_xor_sync(0xffffffff, sq_local, offset); dot_local += __shfl_xor_sync(0xffffffff, dot_local, offset); } - if (lane == 0) { - plan.ws_stats[comp_wid][0] = - make_float2(sq_local, dot_local); + if (lane == 0) + { + plan.ws_stats[comp_wid][0] = make_float2(sq_local, dot_local); } cutlass::arch::NamedBarrier::sync(CONSUMER_THREADS, 0); - if (comp_wid == 0 && lane == 0) { + if (comp_wid == 0 && lane == 0) + { float2 totals = {}; - #pragma unroll - for (int w = 0; w < CONSUMER_WARPS; w++) { - totals = float2_add( - totals, plan.ws_stats[w][0]); +#pragma unroll + for (int w = 0; w < CONSUMER_WARPS; w++) + { + totals = float2_add(totals, plan.ws_stats[w][0]); } float rs = rsqrtf(totals.x / H + rms_eps); rsigma_out[tb] = rs; - if (logits_out) logits_out[tb] = totals.y * rs; - if (probs_out) probs_out[tb] = 1.f; + if (logits_out) + logits_out[tb] = totals.y * rs; + if (probs_out) + probs_out[tb] = 1.f; } cutlass::arch::NamedBarrier::sync(CONSUMER_THREADS, 1); } @@ -949,48 +968,46 @@ attn_res_fwd_n1_ttile_kernel( } __syncthreads(); - if (wid == 1) { + if (wid == 1) + { cute::TMEM::Allocator1Sm alloc; alloc.free(plan.tmem_base, TMEM_Q_COLS_TOTAL); } #else - if (cute::thread0()) printf("attn_res_fwd_n1_ttile_kernel requires sm_100a\n"); + if (cute::thread0()) + printf("attn_res_fwd_n1_ttile_kernel requires sm_100a\n"); #endif } -template -static void launch_fwd( - const bf16_t* block_residual, - const bf16_t* layer_residual, - const bf16_t* res_weight, - const bf16_t* rms_weight, - bf16_t* output, - float* rsigma, - float* probs, - float* logits, - int N, int T, int B, - float rms_eps, - int num_sm, - cudaStream_t stream) +template +static void launch_fwd(bf16_t const* block_residual, bf16_t const* layer_residual, bf16_t const* res_weight, + bf16_t const* rms_weight, bf16_t* output, float* rsigma, float* probs, float* logits, int N, int T, int B, + float rms_eps, int num_sm, cudaStream_t stream) { - constexpr size_t smem_size = - ((size_t)CHUNK_DEPTH * NC * H * sizeof(bf16_t) + sizeof(FwdSmemPlan) + 15) & - ~size_t(15); - auto kernel = - &attn_res_fwd_online_v2_kernel; - static bool attrs_set = false; - if (!attrs_set) { - if (smem_size > 48 * 1024) { - cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + constexpr size_t smem_size + = ((size_t) CHUNK_DEPTH * NC * H * sizeof(bf16_t) + sizeof(FwdSmemPlan) + 15) & ~size_t(15); + auto kernel = &attn_res_fwd_online_v2_kernel; + if (smem_size > 48 * 1024) + { + // cudaFuncSetAttribute applies to the current device only; set it + // once per device (per kernel instantiation). + static std::once_flag attrs_set[64]; + int dev = 0; + TLLM_CUDA_CHECK(cudaGetDevice(&dev)); + auto const set_attr = [&] + { TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); }; + if (dev >= 0 && dev < 64) + { + std::call_once(attrs_set[dev], set_attr); + } + else + { + set_attr(); } - attrs_set = true; } int grid = RELEASE_TMEM ? num_sm * 2 : num_sm; kernel<<>>( - block_residual, layer_residual, res_weight, rms_weight, - output, rsigma, probs, logits, N, T, B, rms_eps); + block_residual, layer_residual, res_weight, rms_weight, output, rsigma, probs, logits, N, T, B, rms_eps); } // Small-N counterpart to the Triton one-program topology. One CTA owns the @@ -999,16 +1016,9 @@ static void launch_fwd( // boundary; N=1 can write V directly because its softmax is identically one. template __global__ void __launch_bounds__(256, 1) -attn_res_fwd_s1_single_cta_kernel( - const bf16_t* __restrict__ block_res, - const bf16_t* __restrict__ layer_res, - const bf16_t* __restrict__ res_w, - const bf16_t* __restrict__ rms_w, - bf16_t* __restrict__ output, - float* __restrict__ rsigma_out, - float* __restrict__ probs_out, - float* __restrict__ logits_out, - float rms_eps) + attn_res_fwd_s1_single_cta_kernel(bf16_t const* __restrict__ block_res, bf16_t const* __restrict__ layer_res, + bf16_t const* __restrict__ res_w, bf16_t const* __restrict__ rms_w, bf16_t* __restrict__ output, + float* __restrict__ rsigma_out, float* __restrict__ probs_out, float* __restrict__ logits_out, float rms_eps) { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 constexpr int H = 7168; @@ -1021,95 +1031,109 @@ attn_res_fwd_s1_single_cta_kernel( __shared__ float2 warp_stats[WARPS * N]; __shared__ float weights[N]; - const int tid = threadIdx.x; - const int lane = tid & 31; - const int warp = tid >> 5; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; float2 stats[N] = {}; uint32_t v_cache_bf16[ITEMS][(N + 1) / 2]; - #pragma unroll - for (int item = 0; item < ITEMS; item++) { +#pragma unroll + for (int item = 0; item < ITEMS; item++) + { int h = tid + item * THREADS; - float q = __bfloat162float(res_w[h]) * - __bfloat162float(rms_w[h]); + float q = __bfloat162float(res_w[h]) * __bfloat162float(rms_w[h]); bf16_t item_v[N]; - #pragma unroll - for (int n = 0; n < N; n++) { - const bf16_t* row = n < N - 1 - ? block_res + (size_t)n * H - : layer_res; +#pragma unroll + for (int n = 0; n < N; n++) + { + bf16_t const* row = n < N - 1 ? block_res + (size_t) n * H : layer_res; bf16_t packed_v = row[h]; float v = __bfloat162float(packed_v); - if constexpr (N == 1) { + if constexpr (N == 1) + { output[h] = packed_v; - } else { + } + else + { item_v[n] = packed_v; } - stats[n] = float2_fma( - make_float2(v, v), make_float2(v, q), stats[n]); + stats[n] = float2_fma(make_float2(v, v), make_float2(v, q), stats[n]); } - if constexpr (N > 1) { - #pragma unroll - for (int pair = 0; pair < N / 2; pair++) { - union { + if constexpr (N > 1) + { +#pragma unroll + for (int pair = 0; pair < N / 2; pair++) + { + union + { __nv_bfloat162 bf16x2; uint32_t bits; } packed; - packed.bf16x2 = __halves2bfloat162( - item_v[2 * pair], item_v[2 * pair + 1]); + + packed.bf16x2 = __halves2bfloat162(item_v[2 * pair], item_v[2 * pair + 1]); v_cache_bf16[item][pair] = packed.bits; } } } - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - #pragma unroll - for (int n = 0; n < N; n++) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + { +#pragma unroll + for (int n = 0; n < N; n++) + { uint64_t packed = reinterpret_cast(stats[n]); packed = __shfl_down_sync(0xffffffff, packed, offset); float2 other = reinterpret_cast(packed); stats[n] = float2_add(stats[n], other); } } - if (lane == 0) { - #pragma unroll - for (int n = 0; n < N; n++) { + if (lane == 0) + { +#pragma unroll + for (int n = 0; n < N; n++) + { warp_stats[warp * N + n] = stats[n]; } } __syncthreads(); - if (tid < N) { + if (tid < N) + { float2 total = {}; - #pragma unroll - for (int w = 0; w < WARPS; w++) { +#pragma unroll + for (int w = 0; w < WARPS; w++) + { total = float2_add(total, warp_stats[w * N + tid]); } warp_stats[tid] = total; } __syncthreads(); - if (tid == 0) { + if (tid == 0) + { float local_rsigma[N]; float local_logits[N]; float max_logit = -FLT_MAX; - #pragma unroll - for (int n = 0; n < N; n++) { +#pragma unroll + for (int n = 0; n < N; n++) + { float2 total = warp_stats[n]; local_rsigma[n] = rsqrtf(total.x / H + rms_eps); local_logits[n] = total.y * local_rsigma[n]; max_logit = fmaxf(max_logit, local_logits[n]); } float denominator = 0.0f; - #pragma unroll - for (int n = 0; n < N; n++) { +#pragma unroll + for (int n = 0; n < N; n++) + { weights[n] = exp2f((local_logits[n] - max_logit) * LOG2_E); denominator += weights[n]; } float inv_denominator = 1.0f / denominator; - #pragma unroll - for (int n = 0; n < N; n++) { +#pragma unroll + for (int n = 0; n < N; n++) + { weights[n] *= inv_denominator; rsigma_out[n] = local_rsigma[n]; logits_out[n] = local_logits[n]; @@ -1118,16 +1142,21 @@ attn_res_fwd_s1_single_cta_kernel( } __syncthreads(); - if constexpr (N > 1) { - #pragma unroll - for (int item = 0; item < ITEMS; item++) { + if constexpr (N > 1) + { +#pragma unroll + for (int item = 0; item < ITEMS; item++) + { float value = 0.0f; - #pragma unroll - for (int pair = 0; pair < N / 2; pair++) { - union { +#pragma unroll + for (int pair = 0; pair < N / 2; pair++) + { + union + { __nv_bfloat162 bf16x2; uint32_t bits; } packed; + packed.bits = v_cache_bf16[item][pair]; float2 v = __bfloat1622float2(packed.bf16x2); value = fmaf(weights[2 * pair], v.x, value); @@ -1138,28 +1167,20 @@ attn_res_fwd_s1_single_cta_kernel( } } #else - if (cute::thread0()) { + if (cute::thread0()) + { printf("attn_res_fwd_s1_single_cta_kernel requires sm_100a\n"); } #endif } template -static void launch_s1_single_cta( - const bf16_t* block_residual, - const bf16_t* layer_residual, - const bf16_t* res_weight, - const bf16_t* rms_weight, - bf16_t* output, - float* rsigma, - float* probs, - float* logits, - float rms_eps, +static void launch_s1_single_cta(bf16_t const* block_residual, bf16_t const* layer_residual, bf16_t const* res_weight, + bf16_t const* rms_weight, bf16_t* output, float* rsigma, float* probs, float* logits, float rms_eps, cudaStream_t stream) { attn_res_fwd_s1_single_cta_kernel<<<1, 256, 0, stream>>>( - block_residual, layer_residual, res_weight, rms_weight, - output, rsigma, probs, logits, rms_eps); + block_residual, layer_residual, res_weight, rms_weight, output, rsigma, probs, logits, rms_eps); } // Single-token split-K specialization. The complete grid is one CTA cluster: @@ -1167,16 +1188,9 @@ static void launch_s1_single_cta( // rank-local shared memory, and exchanges only (square, dot) partials via DSM. template __global__ void __launch_bounds__(256, 1) -attn_res_fwd_s1_splitk_kernel( - const bf16_t* __restrict__ block_res, - const bf16_t* __restrict__ layer_res, - const bf16_t* __restrict__ res_w, - const bf16_t* __restrict__ rms_w, - bf16_t* __restrict__ output, - float* __restrict__ rsigma_out, - float* __restrict__ probs_out, - float* __restrict__ logits_out, - float rms_eps) + attn_res_fwd_s1_splitk_kernel(bf16_t const* __restrict__ block_res, bf16_t const* __restrict__ layer_res, + bf16_t const* __restrict__ res_w, bf16_t const* __restrict__ rms_w, bf16_t* __restrict__ output, + float* __restrict__ rsigma_out, float* __restrict__ probs_out, float* __restrict__ logits_out, float rms_eps) { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 namespace cg = cooperative_groups; @@ -1189,56 +1203,60 @@ attn_res_fwd_s1_splitk_kernel( extern __shared__ char smem_raw[]; float* v_cache = reinterpret_cast(smem_raw); - float2* warp_stats = reinterpret_cast( - smem_raw + (size_t)N * K_PER_CTA * sizeof(float)); + float2* warp_stats = reinterpret_cast(smem_raw + (size_t) N * K_PER_CTA * sizeof(float)); float* weights = reinterpret_cast(warp_stats + WARPS * N); - const int tid = threadIdx.x; - const int lane = tid & 31; - const int warp = tid >> 5; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; cg::cluster_group cluster = cg::this_cluster(); - const int group = cluster.block_rank(); - const int h_begin = group * K_PER_CTA; + int const group = cluster.block_rank(); + int const h_begin = group * K_PER_CTA; float sq[N] = {}; float dot[N] = {}; - #pragma unroll - for (int ki = tid; ki < K_PER_CTA; ki += THREADS) { +#pragma unroll + for (int ki = tid; ki < K_PER_CTA; ki += THREADS) + { int h = h_begin + ki; - float q = __bfloat162float(res_w[h]) * - __bfloat162float(rms_w[h]); - #pragma unroll - for (int n = 0; n < N; n++) { - const bf16_t* row = n < N - 1 - ? block_res + (size_t)n * H - : layer_res; + float q = __bfloat162float(res_w[h]) * __bfloat162float(rms_w[h]); +#pragma unroll + for (int n = 0; n < N; n++) + { + bf16_t const* row = n < N - 1 ? block_res + (size_t) n * H : layer_res; float v = __bfloat162float(row[h]); - v_cache[(size_t)n * K_PER_CTA + ki] = v; + v_cache[(size_t) n * K_PER_CTA + ki] = v; sq[n] = fmaf(v, v, sq[n]); dot[n] = fmaf(v, q, dot[n]); } } - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - #pragma unroll - for (int n = 0; n < N; n++) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + { +#pragma unroll + for (int n = 0; n < N; n++) + { sq[n] += __shfl_down_sync(0xffffffff, sq[n], offset); dot[n] += __shfl_down_sync(0xffffffff, dot[n], offset); } } - if (lane == 0) { - #pragma unroll - for (int n = 0; n < N; n++) { + if (lane == 0) + { +#pragma unroll + for (int n = 0; n < N; n++) + { warp_stats[warp * N + n] = make_float2(sq[n], dot[n]); } } __syncthreads(); - if (tid < N) { + if (tid < N) + { float2 total = {}; - #pragma unroll - for (int w = 0; w < WARPS; w++) { +#pragma unroll + for (int w = 0; w < WARPS; w++) + { total = float2_add(total, warp_stats[w * N + tid]); } warp_stats[tid] = total; @@ -1249,40 +1267,46 @@ attn_res_fwd_s1_splitk_kernel( // One thread per candidate reduces across CTA ranks. Parallelizing this // avoids making a single leader issue all GROUPS*N remote DSM reads. - if (tid < N) { + if (tid < N) + { float2 total = {}; - #pragma unroll - for (int g = 0; g < GROUPS; g++) { - const float2* remote_stats = - cluster.map_shared_rank(warp_stats, g); +#pragma unroll + for (int g = 0; g < GROUPS; g++) + { + float2 const* remote_stats = cluster.map_shared_rank(warp_stats, g); total = float2_add(total, remote_stats[tid]); } warp_stats[tid] = total; } __syncthreads(); - if (tid == 0) { + if (tid == 0) + { float local_rsigma[N]; float local_logits[N]; float max_logit = -FLT_MAX; - #pragma unroll - for (int n = 0; n < N; n++) { +#pragma unroll + for (int n = 0; n < N; n++) + { float2 total = warp_stats[n]; local_rsigma[n] = rsqrtf(total.x / H + rms_eps); local_logits[n] = total.y * local_rsigma[n]; max_logit = fmaxf(max_logit, local_logits[n]); } float sum = 0.0f; - #pragma unroll - for (int n = 0; n < N; n++) { +#pragma unroll + for (int n = 0; n < N; n++) + { weights[n] = exp2f((local_logits[n] - max_logit) * LOG2_E); sum += weights[n]; } float inv_sum = 1.0f / sum; - #pragma unroll - for (int n = 0; n < N; n++) { +#pragma unroll + for (int n = 0; n < N; n++) + { weights[n] *= inv_sum; - if (group == 0) { + if (group == 0) + { rsigma_out[n] = local_rsigma[n]; logits_out[n] = local_logits[n]; probs_out[n] = weights[n]; @@ -1292,55 +1316,55 @@ attn_res_fwd_s1_splitk_kernel( cluster.sync(); - #pragma unroll - for (int ki = tid; ki < K_PER_CTA; ki += THREADS) { +#pragma unroll + for (int ki = tid; ki < K_PER_CTA; ki += THREADS) + { float value = 0.0f; - #pragma unroll - for (int n = 0; n < N; n++) { - value = fmaf( - weights[n], v_cache[(size_t)n * K_PER_CTA + ki], value); +#pragma unroll + for (int n = 0; n < N; n++) + { + value = fmaf(weights[n], v_cache[(size_t) n * K_PER_CTA + ki], value); } output[h_begin + ki] = __float2bfloat16_rn(value); } #else - if (cute::thread0()) { + if (cute::thread0()) + { printf("attn_res_fwd_s1_splitk_kernel requires sm_100a\n"); } #endif } template -static void launch_s1_splitk( - const bf16_t* block_residual, - const bf16_t* layer_residual, - const bf16_t* res_weight, - const bf16_t* rms_weight, - bf16_t* output, - float* rsigma, - float* probs, - float* logits, - float rms_eps, +static void launch_s1_splitk(bf16_t const* block_residual, bf16_t const* layer_residual, bf16_t const* res_weight, + bf16_t const* rms_weight, bf16_t* output, float* rsigma, float* probs, float* logits, float rms_eps, cudaStream_t stream) { constexpr int K_PER_CTA = 7168 / GROUPS; constexpr int WARPS = 8; - constexpr size_t smem_size = - (size_t)N * K_PER_CTA * sizeof(float) + - (size_t)WARPS * N * sizeof(float2) + - (size_t)N * sizeof(float); + constexpr size_t smem_size + = (size_t) N * K_PER_CTA * sizeof(float) + (size_t) WARPS * N * sizeof(float2) + (size_t) N * sizeof(float); auto kernel = &attn_res_fwd_s1_splitk_kernel; - static bool attrs_set = false; - if (!attrs_set) { - cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); - attrs_set = true; + { + // cudaFuncSetAttribute applies to the current device only; set it + // once per device (per kernel instantiation). + static std::once_flag attrs_set[64]; + int dev = 0; + TLLM_CUDA_CHECK(cudaGetDevice(&dev)); + auto const set_attr = [&] + { TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); }; + if (dev >= 0 && dev < 64) + { + std::call_once(attrs_set[dev], set_attr); + } + else + { + set_attr(); + } } - void* args[] = { - const_cast(&block_residual), - const_cast(&layer_residual), - const_cast(&res_weight), - const_cast(&rms_weight), - &output, &rsigma, &probs, &logits, &rms_eps}; + void* args[] = {const_cast(&block_residual), const_cast(&layer_residual), + const_cast(&res_weight), const_cast(&rms_weight), &output, &rsigma, &probs, &logits, + &rms_eps}; cudaLaunchConfig_t config{}; config.gridDim = dim3(GROUPS); config.blockDim = dim3(256); @@ -1353,39 +1377,37 @@ static void launch_s1_splitk( attribute.val.clusterDim.z = 1; config.attrs = &attribute; config.numAttrs = 1; - cudaLaunchKernelExC( - &config, reinterpret_cast(kernel), args); + cudaLaunchKernelExC(&config, reinterpret_cast(kernel), args); } template -static void launch_n1_ttile( - const bf16_t* layer_residual, - const bf16_t* res_weight, - const bf16_t* rms_weight, - bf16_t* output, - float* rsigma, - float* probs, - float* logits, - int T, int B, - float rms_eps, - int num_sm, +static void launch_n1_ttile(bf16_t const* layer_residual, bf16_t const* res_weight, bf16_t const* rms_weight, + bf16_t* output, float* rsigma, float* probs, float* logits, int T, int B, float rms_eps, int num_sm, cudaStream_t stream) { - constexpr size_t smem_size = - ((size_t)CHUNK_DEPTH * TB_TILE * H * sizeof(bf16_t) + - sizeof(FwdSmemPlan<1>) + 15) & ~size_t(15); + constexpr size_t smem_size + = ((size_t) CHUNK_DEPTH * TB_TILE * H * sizeof(bf16_t) + sizeof(FwdSmemPlan<1>) + 15) & ~size_t(15); auto kernel = &attn_res_fwd_n1_ttile_kernel; - static bool attrs_set = false; - if (!attrs_set) { - if (smem_size > 48 * 1024) { - cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + if (smem_size > 48 * 1024) + { + // cudaFuncSetAttribute applies to the current device only; set it + // once per device (per kernel instantiation). + static std::once_flag attrs_set[64]; + int dev = 0; + TLLM_CUDA_CHECK(cudaGetDevice(&dev)); + auto const set_attr = [&] + { TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); }; + if (dev >= 0 && dev < 64) + { + std::call_once(attrs_set[dev], set_attr); + } + else + { + set_attr(); } - attrs_set = true; } kernel<<>>( - layer_residual, res_weight, rms_weight, - output, rsigma, probs, logits, T, B, rms_eps); + layer_residual, res_weight, rms_weight, output, rsigma, probs, logits, T, B, rms_eps); } } // namespace fwd_prod_v2 @@ -1435,12 +1457,10 @@ void invokeAttnResFwd(AttnResFwdParams const& params, cudaStream_t stream) float const rms_eps = params.rmsEps; int dev = 0; - cudaGetDevice(&dev); + TLLM_CUDA_CHECK(cudaGetDevice(&dev)); int num_sm = attn_res_fwd_grid_size(dev); - if (num_sm <= 0 || N > N_MAX) - { - return; - } + TLLM_CHECK_WITH_INFO(num_sm > 0, "attn_res_fwd: failed to query the SM count of device %d", dev); + TLLM_CHECK_WITH_INFO(N <= N_MAX, "attn_res_fwd: unsupported N=%d (max %d)", N, N_MAX); if (H == 8192) { @@ -1490,7 +1510,7 @@ void invokeAttnResFwd(AttnResFwdParams const& params, cudaStream_t stream) else if (N == 12 && T == 1024) { launch_fwd<7168, 4, false, true>(block_residual, layer_residual, res_weight, rms_weight, output, rsigma, - probs, logits, N, T, B, rms_eps, num_sm - 1, stream); + probs, logits, N, T, B, rms_eps, std::max(1, num_sm - 1), stream); } else { @@ -1526,6 +1546,10 @@ void invokeAttnResFwd(AttnResFwdParams const& params, cudaStream_t stream) logits, N, T, B, rms_eps, num_sm, stream); } } + else + { + TLLM_CHECK_WITH_INFO(false, "attn_res_fwd: unsupported hidden size H=%d", H); + } } } // namespace kernels::kimiK3AttnRes diff --git a/cpp/tensorrt_llm/thop/attnResOp.cpp b/cpp/tensorrt_llm/thop/attnResOp.cpp index c45201c62ba2..6a7d16cadaf2 100644 --- a/cpp/tensorrt_llm/thop/attnResOp.cpp +++ b/cpp/tensorrt_llm/thop/attnResOp.cpp @@ -75,11 +75,13 @@ std::tuple attn_res_fwd( int const B = static_cast(layer_residual.size(1)); int const H = static_cast(layer_residual.size(2)); int const N = static_cast(block_residual.size(0)) + 1; - check_attn_res_contract(N, T, B, H); - c10::cuda::CUDAGuard device_guard(layer_residual.device()); - TORCH_CHECK(layer_residual.is_cuda() && block_residual.is_cuda() && res_weight.is_cuda() && rms_weight.is_cuda(), "attn_res_fwd: all input tensors must be CUDA tensors"); + // Set the device before check_attn_res_contract: is_sm100_family() reads + // the current device, which must match the tensors' device. + c10::cuda::CUDAGuard device_guard(layer_residual.device()); + check_attn_res_contract(N, T, B, H); + TORCH_CHECK(layer_residual.scalar_type() == at::kBFloat16, "attn_res_fwd: layer_residual must be bf16"); TORCH_CHECK(block_residual.scalar_type() == at::kBFloat16, "attn_res_fwd: block_residual must be bf16"); TORCH_CHECK(res_weight.scalar_type() == at::kBFloat16, "attn_res_fwd: res_weight must be bf16"); diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py index 1ce0861a0294..f34908d8c708 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py @@ -31,9 +31,8 @@ final recurrent state. Intermediate matrices remain private runner workspace. """ -from typing import Optional, Tuple - import weakref +from typing import Optional, Tuple import torch @@ -282,8 +281,11 @@ def _get_side_stream(dev): # Padded-input scratch cache for the eqlen partial-chunk path. Keyed by # (B, T_padded, H, K, dtype_qkv, dtype_g, dtype_beta, device, real_T). # real_T is part of the key so the g sentinel tail [real_T:T_padded] = -1e3 -# is set once and reused across calls with the same shape. +# is set once and reused across calls with the same shape. LRU-bounded like +# _buf_cache: real_T varies per prefill batch, so an unbounded dict would pin +# scratch for every distinct token count forever. _padded_input_cache = {} +_PAD_CACHE_MAX_ENTRIES = 8 # Sentinel-padded g scratch for varlen single-seq Phase 2.1 path. Keyed by # (B, T_padded, H, K, dtype, device, real_T). The tail [real_T:T_padded] is @@ -299,7 +301,12 @@ def _get_g_sentinel_buffer(B, T_padded, H, K, dtype_g, device, real_T): e = torch.zeros(B, T_padded, H, K, dtype=dtype_g, device=device) if real_T < T_padded: e[:, real_T:] = -1000.0 + while len(_g_sentinel_cache) >= _PAD_CACHE_MAX_ENTRIES: + _g_sentinel_cache.pop(next(iter(_g_sentinel_cache))) _g_sentinel_cache[key] = e + else: + # LRU refresh so hot shapes survive eviction. + _g_sentinel_cache[key] = _g_sentinel_cache.pop(key) return e @@ -327,12 +334,16 @@ def _get_padded_input_buffers(B, T_padded, H, K, dtype_qkv, dtype_g, dtype_beta, if real_T < T_padded: g_pad[:, real_T:] = -1000.0 e = (q_pad, k_pad, v_pad, g_pad, beta_pad) + while len(_padded_input_cache) >= _PAD_CACHE_MAX_ENTRIES: + _padded_input_cache.pop(next(iter(_padded_input_cache))) _padded_input_cache[key] = e + else: + # LRU refresh so hot shapes survive eviction. + _padded_input_cache[key] = _padded_input_cache.pop(key) return e -def _get_buffers(dev, dtype_k, B, T, H, K_dim, V_dim, NT, N_seqs, BT, - varlen=False): +def _get_buffers(dev, dtype_k, B, T, H, K_dim, V_dim, NT, N_seqs, BT, varlen=False): """All beta fusion lives in akk_inv kernel epilogue (post-inv column-scale).""" key = (dev.index or 0, B, T, H, K_dim, V_dim, NT, N_seqs, varlen) if key not in _buf_cache: @@ -487,6 +498,12 @@ def _launch_k4_persistent( ): """Launch persistent K4 with cached CuTe wrappers. + Precondition: all sequence lengths in cu_seqlens must be > 0 (a + zero-length sequence deadlocks the kernel's chunk-loop barriers; see the + k4_persistent module docstring). Not validated here: cu_seqlens is on the + GPU and a host-side check would sync the hot path, and the prefill + runtime never emits zero-length sequences. + No fast-launch (args-tuple) cache here: such a cache pins the per-call v/initial-state tensors via their cute wrappers (the wrapper holds the storage, so the keyed object never dies and weakref pruning never @@ -792,12 +809,10 @@ def _launch_fused_k123_inv( stream, ) akk_fn = _akk_inv_cache[akk_cache_key] - akk_args = (akk_in_view, akk_out_view, beta_ct, B, NT, akk_cu_ct, - akk_ci_ct, T_val, stream) + akk_args = (akk_in_view, akk_out_view, beta_ct, B, NT, akk_cu_ct, akk_ci_ct, T_val, stream) akk_fn(*akk_args) - # ========== Fused K1234 compilation cache ========== _fused_k1234_cache = {} _BT = 64 @@ -923,10 +938,8 @@ def _chunk_kda_fwd( # pool the initial state may alias. final_state = initial_state.to(torch.float32).clone() else: - final_state = torch.zeros( - n_seqs, H, K, V_dim, dtype=torch.float32, device=q.device) - return (o, final_state, None, None, None, None, None, None, None, - None, None, initial_state) + final_state = torch.zeros(n_seqs, H, K, V_dim, dtype=torch.float32, device=q.device) + return (o, final_state, None, None, None, None, None, None, None, None, None, initial_state) # ===== Fused K1234 path (eqlen only, single kernel launch) ===== if use_fused_k1234 and not is_varlen: @@ -1028,7 +1041,8 @@ def _chunk_kda_fwd( assert cur_T % BT == 0 and cur_T >= real_T, ( f"varlen single-seq path expects caller-padded input " f"(T={cur_T}, seqlen={real_T}); see " - "KDAKernelDispatch.prefill_chunk_kda") + "KDAKernelDispatch.prefill_chunk_kda" + ) g_pad = _get_g_sentinel_buffer(B, cur_T, H, K, g.dtype, g.device, real_T) g_pad[:, :real_T].copy_(g[:, :real_T]) g = g_pad @@ -1060,7 +1074,8 @@ def _chunk_kda_fwd( # KDAKernelDispatch.prefill_chunk_kda. raise ValueError( f"kda_prefill requires >= 4 total varlen chunks (got {NT}); " - "route small varlen batches to the FLA fallback") + "route small varlen batches to the FLA fallback" + ) N_seqs = len(cu_seqlens) - 1 else: NT = T // BT @@ -1079,8 +1094,7 @@ def _chunk_kda_fwd( cu_eqlen, co_eqlen, cute_wrappers, - ) = _get_buffers(device, k.dtype, B, T, H, K, V_dim, NT, N_seqs, BT, - varlen=is_varlen) + ) = _get_buffers(device, k.dtype, B, T, H, K, V_dim, NT, N_seqs, BT, varlen=is_varlen) # ===== State copy on side stream, parallel with K123 ===== # K4 needs S_out populated with initial_state. By doing this copy on a @@ -1213,6 +1227,12 @@ def forward( raise RuntimeError("Kimi K3 KDA prefill requires NVIDIA CUTLASS DSL") if chunk_size != 64: raise ValueError(f"Kimi K3 KDA prefill requires chunk_size=64, got {chunk_size}") + # Zero-token calls take the early return in _chunk_kda_fwd and + # never touch A_log; the runtime emits such batches (overlap + # scheduler + logprobs flows), so only require A_log when there + # is work to do. + if A_log is None and q.shape[1] != 0: + raise ValueError("Kimi K3 KDA prefill requires A_log") result = _chunk_kda_fwd( q=q, diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py index 21e2b247a898..bbd4ddb9203e 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py @@ -68,7 +68,10 @@ from cutlass.cute.runtime import from_dlpack from ..cute_dsl_kernels.blackwell.kimi_k3_kda.kda_mtp_decode import ( - NUM_THREADS, TILE_K, kda_decode_mtp_kernel) + NUM_THREADS, + TILE_K, + kda_decode_mtp_kernel, + ) else: raise ImportError("Kimi K3 KDA MTP decode requires NVIDIA CUTLASS DSL") @@ -230,16 +233,12 @@ def _require_stride_layout( } for name, tensor in last_dim_tensors.items(): if tensor.stride(-1) != 1: - raise ValueError( - f"Expected {name} to be contiguous in its last dimension.") + raise ValueError(f"Expected {name} to be contiguous in its last dimension.") - if w_q.shape != (H * K, W) or w_k.shape != (H * K, W) or w_v.shape != ( - HV * V, W): - raise ValueError(f"Expected w_q/w_k shape [{H * K}, {W}] and w_v " - f"shape [{HV * V}, {W}].") + if w_q.shape != (H * K, W) or w_k.shape != (H * K, W) or w_v.shape != (HV * V, W): + raise ValueError(f"Expected w_q/w_k shape [{H * K}, {W}] and w_v shape [{HV * V}, {W}].") if w_q.stride(1) != 1 or w_k.stride(1) != 1 or w_v.stride(1) != 1: - raise ValueError( - "Expected w_q/w_k/w_v to be contiguous in the kernel-width axis.") + raise ValueError("Expected w_q/w_k/w_v to be contiguous in the kernel-width axis.") if A_log.ndim != 1 or A_log.shape[0] != H: raise ValueError(f"Expected A_log shape [{H}].") @@ -248,43 +247,41 @@ def _require_stride_layout( state_s = W - 1 + num_spec if cs_q.ndim != 3 or cs_k.ndim != 3 or cs_v.ndim != 3: - raise ValueError("Expected cs_q/cs_k/cs_v to have shape " - "[pool, dim, S].") + raise ValueError("Expected cs_q/cs_k/cs_v to have shape [pool, dim, S].") if cs_q.shape[1] != H * K or cs_k.shape[1] != H * K: raise ValueError(f"Expected cs_q/cs_k shape [pool, {H * K}, S].") if cs_v.shape[1] != HV * V: raise ValueError(f"Expected cs_v shape [pool, {HV * V}, S].") - if cs_q.shape[2] < state_s or cs_k.shape[2] < state_s or \ - cs_v.shape[2] < state_s: - raise ValueError( - f"Expected conv-state S dimension to be at least {state_s}.") + if cs_q.shape[2] < state_s or cs_k.shape[2] < state_s or cs_v.shape[2] < state_s: + raise ValueError(f"Expected conv-state S dimension to be at least {state_s}.") if cs_q.stride(1) != 1 or cs_k.stride(1) != 1 or cs_v.stride(1) != 1: raise ValueError( "Expected cs_q/cs_k/cs_v to use dim-contiguous layout " - "(allocate as [pool, S, dim] and transpose(1, 2)).") + "(allocate as [pool, S, dim] and transpose(1, 2))." + ) pool_size = recurrent_state.shape[0] if recurrent_state.ndim != 4 or recurrent_state.shape[1:] != (HV, V, K): - raise ValueError(f"Expected recurrent_state shape " - f"[pool, {HV}, {V}, {K}] (V-first pool layout).") - if qkg_cache.ndim != 4 or qkg_cache.shape[1:] != (num_spec, 3, H * K): raise ValueError( - f"Expected qkg_cache shape [pool, {num_spec}, 3, {H * K}].") + f"Expected recurrent_state shape [pool, {HV}, {V}, {K}] (V-first pool layout)." + ) + if qkg_cache.ndim != 4 or qkg_cache.shape[1:] != (num_spec, 3, H * K): + raise ValueError(f"Expected qkg_cache shape [pool, {num_spec}, 3, {H * K}].") if v_cache.ndim != 3 or v_cache.shape[1:] != (num_spec, HV * V): - raise ValueError( - f"Expected v_cache shape [pool, {num_spec}, {HV * V}].") + raise ValueError(f"Expected v_cache shape [pool, {num_spec}, {HV * V}].") if beta_cache.ndim != 3 or beta_cache.shape[1:] != (num_spec, HV): - raise ValueError( - f"Expected beta_cache shape [pool, {num_spec}, {HV}].") - if qkg_cache.shape[0] < pool_size or v_cache.shape[0] < pool_size or \ - beta_cache.shape[0] < pool_size: - raise ValueError( - "Expected cache pool dimensions to cover recurrent_state rows.") + raise ValueError(f"Expected beta_cache shape [pool, {num_spec}, {HV}].") + if ( + qkg_cache.shape[0] < pool_size + or v_cache.shape[0] < pool_size + or beta_cache.shape[0] < pool_size + ): + raise ValueError("Expected cache pool dimensions to cover recurrent_state rows.") - if ssm_state_indices.ndim != 1 or cu_seqlens.ndim != 1 or \ - num_accepted_tokens.ndim != 1: - raise ValueError("Expected ssm_state_indices, cu_seqlens, and " - "num_accepted_tokens to be 1D.") + if ssm_state_indices.ndim != 1 or cu_seqlens.ndim != 1 or num_accepted_tokens.ndim != 1: + raise ValueError( + "Expected ssm_state_indices, cu_seqlens, and num_accepted_tokens to be 1D." + ) if cu_seqlens.shape[0] != ssm_state_indices.shape[0] + 1: raise ValueError("Expected cu_seqlens length to be N + 1.") if num_accepted_tokens.shape[0] != ssm_state_indices.shape[0]: @@ -292,8 +289,7 @@ def _require_stride_layout( def _layout_key(tensor: torch.Tensor): - return (tensor.dtype, tuple(tensor.shape), tuple(tensor.stride()), - _fits_32bit_stride(tensor)) + return (tensor.dtype, tuple(tensor.shape), tuple(tensor.stride()), _fits_32bit_stride(tensor)) def _fits_32bit_stride(tensor: torch.Tensor) -> bool: @@ -330,15 +326,13 @@ def _dlpack_arg(tensor: torch.Tensor): _precompute_control_cache = {} -def _precompute_control_tensor(device: torch.device, - enabled: bool) -> torch.Tensor: +def _precompute_control_tensor(device: torch.device, enabled: bool) -> torch.Tensor: dev = torch.device(device) - key = (dev.index - if dev.index is not None else torch.cuda.current_device(), - bool(enabled)) + key = (dev.index if dev.index is not None else torch.cuda.current_device(), bool(enabled)) if key not in _precompute_control_cache: _precompute_control_cache[key] = torch.tensor( - [1 if enabled else 0], dtype=torch.int32, device=dev) + [1 if enabled else 0], dtype=torch.int32, device=dev + ) return _precompute_control_cache[key] @@ -364,10 +358,18 @@ def _try_flatten_args( return True, h0, x_q_flat, x_k_flat, x_v_flat -def _is_benchmark_static_shape(N: int, H: int, HV: int, K: int, V: int, - W: int, num_spec: int) -> bool: - return (K == 128 and V == 128 and W == 4 and num_spec == 2 and H == HV - and N in (32, 128) and H in (2, 12, 32)) +def _is_benchmark_static_shape( + N: int, H: int, HV: int, K: int, V: int, W: int, num_spec: int +) -> bool: + return ( + K == 128 + and V == 128 + and W == 4 + and num_spec == 2 + and H == HV + and N in (32, 128) + and H in (2, 12, 32) + ) # Layout-and-constexpr-keyed compile cache. Compilation is per (N, T_total, @@ -448,12 +450,7 @@ def kda_mtp_decode_impl( N = cu_seqlens.shape[0] - 1 if out is None: - out = torch.zeros(1, - T_total, - HV, - V_dim, - dtype=x_q.dtype, - device=x_q.device) + out = torch.zeros(1, T_total, HV, V_dim, dtype=x_q.dtype, device=x_q.device) if num_accepted_tokens.dtype != torch.int32: num_accepted_tokens = num_accepted_tokens.to(torch.int32) @@ -503,14 +500,27 @@ def kda_mtp_decode_impl( V=V_dim, ) pool_size = h0_arg.shape[0] - is_benchmark_static_shape = _is_benchmark_static_shape( - N, H, HV, K, V_dim, W, num_spec) + is_benchmark_static_shape = _is_benchmark_static_shape(N, H, HV, K, V_dim, W, num_spec) use_setmaxreg = is_benchmark_static_shape use_reg_q_weights = is_benchmark_static_shape use_regular_metadata = bool(regular_metadata_hint) - use_zero_accepted = bool(zero_accepted_hint) + # The kernel's USE_ZERO_ACCEPTED fast path unrolls exactly + # 1 + NUM_SPEC == 3 new tokens, so it is only valid for num_spec == 2. + # For other num_spec fall back to the generic loop (the hint implies + # num_accepted_tokens is all zeros, so the generic path computes the + # same result). + use_zero_accepted = bool(zero_accepted_hint) and num_spec == 2 # stage_timing is unused (PROFILE_STAGES=False); pass `out` as the - # placeholder tensor argument like the drop's runner does. + # placeholder tensor argument like the drop's runner does. The alias is + # only valid while profiling stays off: with PROFILE_STAGES=True the + # kernel writes int64 stage deltas through this tensor, corrupting the + # bf16 output. Enabling profiling requires a dedicated int64 buffer of + # at least HV * N * 4 elements. + profile_stages = False + assert not profile_stages, ( + "stage_timing aliases `out`; allocate a dedicated int64 [HV * N * 4] " + "buffer before enabling PROFILE_STAGES" + ) stage_timing_arg = out key = ( @@ -558,7 +568,8 @@ def kda_mtp_decode_impl( f"kda_mtp_decode: compiling variant N={N} H={HV} T={T_total} " f"num_spec={num_spec} zero_accepted={use_zero_accepted} " f"regular_metadata={use_regular_metadata} " - f"static_shape={is_benchmark_static_shape}") + f"static_shape={is_benchmark_static_shape}" + ) _compiled_cache[key] = cute.compile( _run_kda_decode_mtp, _from_dlpack_arg(h0_arg), @@ -601,7 +612,7 @@ def kda_mtp_decode_impl( USE_ZERO_ACCEPTED=use_zero_accepted, FUSE_PRECOMPUTE=True, RUNTIME_PRECOMPUTE_FLAG=False, - PROFILE_STAGES=False, + PROFILE_STAGES=profile_stages, stream=stream, ) @@ -735,4 +746,4 @@ def _( zero_accepted_hint: bool = False, regular_metadata_hint: bool = False, ) -> torch.Tensor: - return x_v.new_empty(x_v.shape) + return x_q.new_empty(x_v.shape) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py index 771bf43e73f9..befa73ab929a 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py @@ -21,26 +21,26 @@ Grid: (NUM_SMS, 1, 1) — 148 persistent blocks, each loops over work units Total work units = (NT/4) * H * B, distributed round-robin across SMs Block i processes work units i, i+NUM_SMS, i+2*NUM_SMS, ... -Block: 1024 threads (32 warps), warp-specialized with setmaxnreg (all groups 4-aligned): +Block: 992 threads (31 warps), warp-specialized with setmaxnreg: Warps 0-15: TMA+K1 fused (8×2, vec2, prefetch pipeline) – 4 WGs, 56 regs - Warps 16-27: K2 MMA compute (10 active + 2 idle for WG alignment) – 3 WGs, 72 regs - Warps 28-31: Store/Inversion warps – 1 WG, 24 regs + Warps 16-26: K2 MMA compute (10 active + warp 26 as TMA producer) – 72 regs + Warps 27-30: Store/Inversion warps – 24 regs Pipeline (single for_generate, warp groups separated by if-blocks): per work unit: Warps 0-15: prefetch chunk 0→stage 0 (warp 0), then loop: TMA next chunk (warp 0), wait cur chunk, K1 compute, arrive(k1_done) - Warps 16-27: wait(k1_done)+wait(store_done), MMA, arrive(mma_done+stage_reuse) - Warps 28-31: wait(mma_done), store sAqk/sAkk→GMEM, arrive(store_done) + Warps 16-26: wait(k1_done)+wait(store_done), MMA, arrive(mma_done+stage_reuse) + Warps 27-30: wait(mma_done), store sAqk/sAkk→GMEM, arrive(store_done) All warp-group invariants are computed inside each group's if-block (not hoisted) to eliminate cross-group register pressure — same budget as the _all version. Mbarrier phases self-reset after 4 iterations (2 stages × 2 phases). Mbarriers: tma_mbars[2]: count=1, warp 0 lane 0 → K1+MMA wait for TMA data - stage_reuse_mbars[2]: count=384, MMA(12 warps) → warp 0 waits before TMA reuse + stage_reuse_mbars[2]: count=320, MMA(10 warps) → warp 0 waits before TMA reuse k1_done_mbars[2]: count=512, K1(16 warps) → MMA waits for g_cumsum ready - mma_done_mbars[2]: count=384, MMA(12 warps) → Store waits for sAqk/sAkk ready + mma_done_mbars[2]: count=320, MMA(10 warps) → Store waits for sAqk/sAkk ready store_done_mbars[2]: count=128, Store(4 warps) → MMA waits for sAqk/sAkk stage free SMEM: ~215KB (q+k+g × [64,128] bf16 × 2 stages + g_cumsum [64,136] fp32 × 2 stages @@ -85,9 +85,9 @@ NUM_MMA_WARPS = 11 # Warps 16-26: MMA (10 active + 1 TMA producer, dropped idle warp 27) NUM_MMA_ACTIVE = 10 # mma_warp 0..9: actual MMA work TMA_WARP_ID = NUM_K1_TMA_WARPS + NUM_MMA_ACTIVE # warp 26 = dedicated TMA producer -NUM_STORE_WARPS = 4 # Warps 28-31: Store/Inversion (1 warpgroup) -NUM_WARPS = NUM_K1_TMA_WARPS + NUM_MMA_WARPS + NUM_STORE_WARPS # 32 -THREADS = NUM_WARPS * 32 # 1024 +NUM_STORE_WARPS = 4 # Warps 27-30: Store/Inversion (1 warpgroup) +NUM_WARPS = NUM_K1_TMA_WARPS + NUM_MMA_WARPS + NUM_STORE_WARPS # 31 +THREADS = NUM_WARPS * 32 # 992 NUM_SUB_CHUNKS = BT // BC # 4 NUM_TILES = NUM_SUB_CHUNKS * (NUM_SUB_CHUNKS + 1) // 2 # 10 lower-tri tiles @@ -1049,30 +1049,29 @@ def fused_kernel123( # or store warps). Required for downstream row-major store optimizations # — positions outside MMA-written sub-tiles stay at 0. # - # Cooperative pattern (32 warps × 32 lanes = 1024 threads): - # - Each warp owns 2 contiguous rows (warp_id*2, warp_id*2+1) + # Cooperative pattern (31 warps × 32 lanes = 992 threads): + # - Rows strided by warp count: warp w owns rows w, w+31, w+62 (< BT) # - Each lane owns 2 contiguous bf16 cols (lane*2, lane*2+1) - # - Per lane: 2 stages × 2 rows × 2 buffers × 2 cols = 16 bf16 stores # - Adjacent (lane*2, lane*2+1) bf16 pairs are 4-byte aligned → - # ptxas should fuse into STS.32 (8 wide stores per lane). + # ptxas should fuse into STS.32 wide stores. # ===================================================================== - _warp_id_in_cta = tidx >> 5 # tidx // 32, range 0..31 + _warp_id_in_cta = tidx >> 5 # tidx // 32, range 0..30 _lane_id_warp = tidx & 31 # tidx % 32, range 0..31 - _row_base = _warp_id_in_cta * 2 # this warp owns rows [_row_base, _row_base+1] _col_lo = _lane_id_warp * 2 # this lane owns cols [_col_lo, _col_lo+1] _col_hi = _col_lo + 1 for _s in cutlass.range_constexpr(NUM_STAGES): - for _ri in cutlass.range_constexpr(2): - _row = _row_base + _ri - sAqk[_row, _col_lo, _s] = cutlass.BFloat16(0.0) - sAqk[_row, _col_hi, _s] = cutlass.BFloat16(0.0) - sAkk[_row, _col_lo, _s] = cutlass.BFloat16(0.0) - sAkk[_row, _col_hi, _s] = cutlass.BFloat16(0.0) + for _ri in cutlass.range_constexpr((BT + NUM_WARPS - 1) // NUM_WARPS): + _row = _warp_id_in_cta + _ri * NUM_WARPS + if _row < BT: + sAqk[_row, _col_lo, _s] = cutlass.BFloat16(0.0) + sAqk[_row, _col_hi, _s] = cutlass.BFloat16(0.0) + sAkk[_row, _col_lo, _s] = cutlass.BFloat16(0.0) + sAkk[_row, _col_hi, _s] = cutlass.BFloat16(0.0) cute.arch.barrier() # ===================================================================== # Pre-arrive (MMA warps only) - # stage_reuse_mbars: warp 0 waits before MMA arrives → pre-arrive all 12 MMA warps + # stage_reuse_mbars: warp 0 waits before MMA arrives → pre-arrive all 10 MMA warps # store_done_mbars: MMA waits before Store arrives → pre-arrive first 4 MMA warps # ===================================================================== if ( @@ -1441,7 +1440,7 @@ def fused_kernel123( cute.arch.mbarrier_arrive(tma_mbars + next_stage) # ============================================================= - # Warps 16-27 (excluding TMA_WARP_ID=26): K2 MMA Compute + # Warps 16-26 (excluding TMA_WARP_ID=26): K2 MMA Compute # ============================================================= if ( warp_idx >= NUM_K1_TMA_WARPS @@ -1525,20 +1524,20 @@ def fused_kernel123( beta_row1 = _z if IS_VARLEN and not VARLEN_PURE: if chunk_start + q_row_base + row0 < mma_eos: - beta_row0 = mBeta[ - i_b, chunk_start + q_row_base + row0, i_h - ].to(cutlass.Float32) + beta_row0 = mBeta[i_b, chunk_start + q_row_base + row0, i_h].to( + cutlass.Float32 + ) if chunk_start + q_row_base + row1 < mma_eos: - beta_row1 = mBeta[ - i_b, chunk_start + q_row_base + row1, i_h - ].to(cutlass.Float32) + beta_row1 = mBeta[i_b, chunk_start + q_row_base + row1, i_h].to( + cutlass.Float32 + ) else: - beta_row0 = mBeta[ - i_b, chunk_start + q_row_base + row0, i_h - ].to(cutlass.Float32) - beta_row1 = mBeta[ - i_b, chunk_start + q_row_base + row1, i_h - ].to(cutlass.Float32) + beta_row0 = mBeta[i_b, chunk_start + q_row_base + row0, i_h].to( + cutlass.Float32 + ) + beta_row1 = mBeta[i_b, chunk_start + q_row_base + row1, i_h].to( + cutlass.Float32 + ) acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = _z, _z, _z, _z acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = _z, _z, _z, _z @@ -1970,7 +1969,7 @@ def fused_kernel123( cute.arch.mbarrier_arrive(mma_done_mbars + s) # ============================================================= - # Warps 28-31: Store/Inversion warps + # Warps 27-30: Store/Inversion warps # ============================================================= if warp_idx >= NUM_K1_TMA_WARPS + NUM_MMA_WARPS: store_warp = warp_idx - (NUM_K1_TMA_WARPS + NUM_MMA_WARPS) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py index 7c2cadc6d9c8..a71d76f4dc37 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py @@ -18,6 +18,13 @@ for per-tile TMA descriptor updates, and domain_offset+flat_divide for per-chunk addressing. Supports variable-length sequences via cu_seqlens. +Precondition: every sequence in cu_seqlens must have length > 0. A zero-length +sequence yields num_chunks = 0, so its scheduler tile skips the chunk loop and +never commits the per-chunk mbarriers the other warp groups wait on (deadlock). +The prefill runtime never produces zero-length sequences, so this is not +checked on the device or in the launch path (cu_seqlens lives on the GPU; a +host-side check would force a sync on the hot path). + K4 chunk loop with 6 MMAs per chunk: MMA1: W = AB @ KS (K-MN, K=64) MMA2: U = AB @ V (K-MN, K=64) @@ -60,6 +67,10 @@ from cutlass.cutlass_dsl import Int32, T, dsl_user_op from cutlass.utils import TensorMapManager, TensorMapUpdateMode +# Compatibility shim, global effect: flashinfer's gated_delta_net_tile_scheduler +# imports CuteExperimentalDSL from cutlass.cutlass_dsl, which older CUTLASS DSL +# releases do not define. Install an inert placeholder so the import below +# succeeds; other modules in the process see the same attribute. if not hasattr(_dsl_mod, "CuteExperimentalDSL"): class _DummyExperimentalDSL: @@ -425,6 +436,20 @@ def transform_partitioned_tensor_layout(tensor): NUM_REGS_WG2 = 232 MAX_REGS = 168 +try: + from tensorrt_llm.logger import logger as _logger +except ImportError: # standalone kernel use outside the full package + import logging as _logging + + _logger = _logging.getLogger(__name__) + +# Monkey-patch, GLOBAL effect: CuTeDSL._get_pipeline is class-level state, so +# every CuTe DSL kernel compiled in this process after this import (not just +# K4) gets ptxas --uumn (unified uniform register allocation, needed here to +# keep WG0 under its 40-register budget). The DSL exposes no per-compile +# ptx-options hook at the cute.compile call sites, hence the patch. The extra +# flag is benign for the other kernels in this package. Idempotent: skips +# pipelines that already carry ptx-options. try: from cutlass.cutlass_dsl.cutlass import CuTeDSL as _CuTeDSL @@ -439,17 +464,23 @@ def _patched_get_pipeline(self, _pipeline): result = result.replace("cubin-format=bin", "cubin-format=bin ptx-options='--uumn'") _patch_applied = True else: - print( - f" [WARN] monkey-patch: 'cubin-format=bin' not found in pipeline: {result[:200]}" + _logger.warning( + f"k4_persistent ptx-options patch: 'cubin-format=bin' not found " + f"in pipeline: {result[:200]}" ) elif result and "ptx-options=" in result: - print(" [INFO] monkey-patch: ptx-options already present") + _logger.debug("k4_persistent ptx-options patch: ptx-options already present") _patch_applied = True return result _CuTeDSL._get_pipeline = _patched_get_pipeline +# Workaround for a CuTe DSL parser bug: the nvidia-cutlass-dsl 4.5.0 AST +# preprocessor cannot parse tuple except handlers anywhere in a kernel +# module ("'Tuple' object has no attribute 'id'"), which breaks every +# cute.compile of this file. Keep a single bare Exception until the DSL +# pin moves past the bug. except Exception as e: - print(f" [WARN] monkey-patch failed: {e}") + _logger.warning(f"k4_persistent ptx-options patch failed: {e}") _patch_applied = False @@ -1453,7 +1484,6 @@ def host_fn( o_out, tm_workspace, scheduler_params, - ).launch(grid=grid_shape, block=(threads_per_cta, 1, 1), use_pdl=True, - stream=stream) + ).launch(grid=grid_shape, block=(threads_per_cta, 1, 1), use_pdl=True, stream=stream) return host_fn diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py index ca6ae67333ce..e769566218fe 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py @@ -40,8 +40,8 @@ import cutlass import cutlass.cute as cute from cutlass._mlir.dialects import llvm -from cutlass.cutlass_dsl import T, dsl_user_op from cutlass.cute.typing import Int64 +from cutlass.cutlass_dsl import T, dsl_user_op NUM_THREADS = 256 TILE_K = 128 @@ -109,7 +109,13 @@ def kda_decode_mtp_kernel( stage_timing: cute.Tensor, PROFILE_STAGES: cutlass.Constexpr[bool], ): - """KDA MTP decode — SMEM pre-compute + register-resident state.""" + """KDA MTP decode — SMEM pre-compute + register-resident state. + + With ``PROFILE_STAGES=True``, ``stage_timing`` must be an int64 tensor + with at least ``HV * N * 4`` elements, indexed as + ``(i_hv * grid_n + i_n) * 4``. With profiling off it is never accessed + and the host may pass any placeholder tensor. + """ tidx, _, _ = cute.arch.thread_idx() in_warp_tid = tidx % 32 warp_idx = cute.arch.warp_idx() @@ -118,6 +124,10 @@ def kda_decode_mtp_kernel( i_h = i_hv if cutlass.const_expr(PROFILE_STAGES): t_stage0 = read_globaltimer() + # Pre-declare so the dynamic `run_precompute` branch below only + # reassigns (first assignment inside a dynamic branch is untraceable; + # see the module docstring on v_row_a/v_row_b). + t_stage1 = Int64(0) if cutlass.const_expr(USE_REGULAR_METADATA): bos = i_n * (2 * NUM_SPEC + 1) eos = bos + (2 * NUM_SPEC + 1) @@ -133,6 +143,11 @@ def kda_decode_mtp_kernel( commit_len = 0 else: commit_len = num_accepted_tokens[i_n] + # Only NUM_SPEC drafts can be pending from the previous round. Clamp + # so a malformed count cannot drive T_loop past the t_max-sized SMEM + # buffers or the num_spec extents of the replay caches. + if commit_len > NUM_SPEC: + commit_len = cutlass.Int32(NUM_SPEC) if cutlass.const_expr(USE_ZERO_ACCEPTED): T_loop = 1 + NUM_SPEC t_max = 1 + NUM_SPEC @@ -155,10 +170,10 @@ def kda_decode_mtp_kernel( sK = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) sG = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max,)), 16) - # Preserve the original shared-memory offsets after sBeta. The removed - # output-norm path used these 8 floats; shifting later buffers changed - # bank mapping in earlier experiments. - sWarpSum = smem.allocate_tensor(cutlass.Float32, cute.make_layout((8,)), 16) + # Reserve the 8-float scratch region formerly used by the removed + # output-norm reduction. Nothing reads or writes this allocation; it only + # preserves the offsets and bank mapping of the shared-memory buffers below. + smem.allocate_tensor(cutlass.Float32, cute.make_layout((8,)), 16) sVall = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max * V,)), 16) sConvW = smem.allocate_tensor( cutlass.Float32, @@ -169,9 +184,13 @@ def kda_decode_mtp_kernel( r_k = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) r_decay = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) r_bk = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) - r_state = cute.make_rmem_tensor(cute.make_layout((NUM_V_ROWS * vec_size,), stride=(1,)), cutlass.Float32) + r_state = cute.make_rmem_tensor( + cute.make_layout((NUM_V_ROWS * vec_size,), stride=(1,)), cutlass.Float32 + ) if cutlass.const_expr(USE_REG_Q_WEIGHTS): - r_wq = cute.make_rmem_tensor(cute.make_layout((KERNEL_WIDTH * vec_size,), stride=(1,)), cutlass.Float32) + r_wq = cute.make_rmem_tensor( + cute.make_layout((KERNEL_WIDTH * vec_size,), stride=(1,)), cutlass.Float32 + ) r_exp_A = cutlass.Float32(0.0) if cutlass.const_expr(USE_REGULAR_METADATA) or eos > bos: if cutlass.const_expr(FUSE_PRECOMPUTE or RUNTIME_PRECOMPUTE_FLAG): @@ -184,20 +203,28 @@ def kda_decode_mtp_kernel( k_idx = i * 32 + in_warp_tid for w in range(KERNEL_WIDTH - 1): r_state[w * vec_size + i] = cutlass.Float32(cs_q[slot, hk_off + k_idx, w]) - r_state[(KERNEL_WIDTH - 1) * vec_size + w * vec_size + i] = cutlass.Float32(cs_k[slot, hk_off + k_idx, w]) + r_state[(KERNEL_WIDTH - 1) * vec_size + w * vec_size + i] = cutlass.Float32( + cs_k[slot, hk_off + k_idx, w] + ) for w in range(KERNEL_WIDTH): if tidx < K: if cutlass.const_expr(not USE_REG_Q_WEIGHTS): sConvW[w * K + tidx] = cutlass.Float32(w_q[hk_off + tidx, w]) - sConvW[k_weight_base + w * K + tidx] = cutlass.Float32(w_k[hk_off + tidx, w]) + sConvW[k_weight_base + w * K + tidx] = cutlass.Float32( + w_k[hk_off + tidx, w] + ) for ld in range(V * KERNEL_WIDTH // NUM_THREADS): flat = ld * NUM_THREADS + tidx - sConvW[v_weight_base + flat] = cutlass.Float32(w_v[hv_off + flat % V, flat // V]) + sConvW[v_weight_base + flat] = cutlass.Float32( + w_v[hv_off + flat % V, flat // V] + ) if cutlass.const_expr(USE_REG_Q_WEIGHTS): if warp_idx == 0: for _w in range(KERNEL_WIDTH): for _i in range(vec_size): - r_wq[_w * vec_size + _i] = cutlass.Float32(w_q[hk_off + _i * 32 + in_warp_tid, _w]) + r_wq[_w * vec_size + _i] = cutlass.Float32( + w_q[hk_off + _i * 32 + in_warp_tid, _w] + ) cute.arch.barrier() if cutlass.const_expr(USE_SETMAXREG): cute.arch.warpgroup_reg_dealloc(64) @@ -214,31 +241,47 @@ def kda_decode_mtp_kernel( if warp_idx == 0: for i in range(vec_size): k_idx = i * 32 + in_warp_tid - sQ[i_t, k_idx] = cutlass.Float32(qkg_cache[slot, i_t, 0, hk_off + k_idx]) + sQ[i_t, k_idx] = cutlass.Float32( + qkg_cache[slot, i_t, 0, hk_off + k_idx] + ) for i in range(vec_size): k_idx = i * 32 + in_warp_tid - r_xq_raw = cutlass.Float32(cs_q[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + i_t]) + r_xq_raw = cutlass.Float32( + cs_q[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + i_t] + ) for w in range(KERNEL_WIDTH - 2): r_state[w * vec_size + i] = r_state[(w + 1) * vec_size + i] r_state[(KERNEL_WIDTH - 2) * vec_size + i] = r_xq_raw elif warp_idx == 1: for i in range(vec_size): k_idx = i * 32 + in_warp_tid - sK[i_t, k_idx] = cutlass.Float32(qkg_cache[slot, i_t, 1, hk_off + k_idx]) + sK[i_t, k_idx] = cutlass.Float32( + qkg_cache[slot, i_t, 1, hk_off + k_idx] + ) if in_warp_tid == 0: sBeta[i_t] = cutlass.Float32(beta_cache[slot, i_t, i_hv]) for i in range(vec_size): k_idx = i * 32 + in_warp_tid - r_xk_raw = cutlass.Float32(cs_k[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + i_t]) + r_xk_raw = cutlass.Float32( + cs_k[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + i_t] + ) for w in range(KERNEL_WIDTH - 2): - r_state[(KERNEL_WIDTH - 1) * vec_size + w * vec_size + i] = r_state[ + r_state[ + (KERNEL_WIDTH - 1) * vec_size + w * vec_size + i + ] = r_state[ (KERNEL_WIDTH - 1) * vec_size + (w + 1) * vec_size + i ] - r_state[(KERNEL_WIDTH - 1) * vec_size + (KERNEL_WIDTH - 2) * vec_size + i] = r_xk_raw + r_state[ + (KERNEL_WIDTH - 1) * vec_size + + (KERNEL_WIDTH - 2) * vec_size + + i + ] = r_xk_raw else: for i in range(vec_size): k_idx = i * 32 + in_warp_tid - r_gk_c2 = cutlass.Float32(qkg_cache[slot, i_t, 2, hk_off + k_idx]) + r_gk_c2 = cutlass.Float32( + qkg_cache[slot, i_t, 2, hk_off + k_idx] + ) sG[i_t, k_idx] = cute.math.exp(r_gk_c2, fastmath=True) else: token = bos + i_t @@ -281,8 +324,12 @@ def kda_decode_mtp_kernel( _cwq_last_0 = r_wq[(KERNEL_WIDTH - 1) * vec_size + i0] _cwq_last_1 = r_wq[(KERNEL_WIDTH - 1) * vec_size + i1] else: - _cwq_last_0 = sConvW[(KERNEL_WIDTH - 1) * K + i0 * 32 + in_warp_tid] - _cwq_last_1 = sConvW[(KERNEL_WIDTH - 1) * K + i1 * 32 + in_warp_tid] + _cwq_last_0 = sConvW[ + (KERNEL_WIDTH - 1) * K + i0 * 32 + in_warp_tid + ] + _cwq_last_1 = sConvW[ + (KERNEL_WIDTH - 1) * K + i1 * 32 + in_warp_tid + ] r_conv_0 += r_xq_0 * _cwq_last_0 r_conv_1 += r_xq_1 * _cwq_last_1 e0 = cute.math.exp(-r_conv_0, fastmath=True) @@ -301,8 +348,12 @@ def kda_decode_mtp_kernel( for i in range(vec_size): sum_q += r_q[i] * r_q[i] for offset in [16, 8, 4, 2, 1]: - sum_q += cute.arch.shuffle_sync_bfly(sum_q, offset=offset, mask=-1, mask_and_clamp=31) - rnorm_q_scaled = cute.math.rsqrt(sum_q + 1e-06, fastmath=True) * scale + sum_q += cute.arch.shuffle_sync_bfly( + sum_q, offset=offset, mask=-1, mask_and_clamp=31 + ) + rnorm_q_scaled = ( + cute.math.rsqrt(sum_q + 1e-06, fastmath=True) * scale + ) for i in range(vec_size): r_q[i] = r_q[i] * rnorm_q_scaled for i in range(vec_size): @@ -330,25 +381,33 @@ def kda_decode_mtp_kernel( r_xk = cutlass.Float32(x_k[0, token, hk_off + k_idx]) else: r_xk = cutlass.Float32(x_k[0, token, i_h, k_idx]) - r_conv += r_xk * sConvW[ - k_weight_base + (KERNEL_WIDTH - 1) * K + i * 32 + in_warp_tid - ] + r_conv += ( + r_xk + * sConvW[ + k_weight_base + + (KERNEL_WIDTH - 1) * K + + i * 32 + + in_warp_tid + ] + ) r_conv = r_conv * cute.arch.rcp_approx( cutlass.Float32(1.0) + cute.math.exp(-r_conv, fastmath=True) ) r_k[i] = r_conv - r_state[(KERNEL_WIDTH - 1) * vec_size + 0 * vec_size + i] = r_state[ - (KERNEL_WIDTH - 1) * vec_size + 1 * vec_size + i - ] - r_state[(KERNEL_WIDTH - 1) * vec_size + 1 * vec_size + i] = r_state[ - (KERNEL_WIDTH - 1) * vec_size + 2 * vec_size + i - ] + r_state[(KERNEL_WIDTH - 1) * vec_size + 0 * vec_size + i] = ( + r_state[(KERNEL_WIDTH - 1) * vec_size + 1 * vec_size + i] + ) + r_state[(KERNEL_WIDTH - 1) * vec_size + 1 * vec_size + i] = ( + r_state[(KERNEL_WIDTH - 1) * vec_size + 2 * vec_size + i] + ) r_state[(KERNEL_WIDTH - 1) * vec_size + 2 * vec_size + i] = r_xk sum_k = 0.0 for i in range(vec_size): sum_k += r_k[i] * r_k[i] for offset in [16, 8, 4, 2, 1]: - sum_k += cute.arch.shuffle_sync_bfly(sum_k, offset=offset, mask=-1, mask_and_clamp=31) + sum_k += cute.arch.shuffle_sync_bfly( + sum_k, offset=offset, mask=-1, mask_and_clamp=31 + ) rnorm_k = cute.math.rsqrt(sum_k + 1e-06, fastmath=True) for i in range(vec_size): r_k[i] = r_k[i] * rnorm_k @@ -357,7 +416,8 @@ def kda_decode_mtp_kernel( sK[i_t, k_idx] = r_k[i] if in_warp_tid == 0: sBeta[i_t] = cute.arch.rcp_approx( - cutlass.Float32(1.0) + cute.math.exp(-r_b_raw, fastmath=True) + cutlass.Float32(1.0) + + cute.math.exp(-r_b_raw, fastmath=True) ) else: for i in range(vec_size): @@ -366,7 +426,8 @@ def kda_decode_mtp_kernel( r_g_raw = r_g_raw + cutlass.Float32(dt_bias[i_h * K + k_idx]) exp_A_x = r_exp_A * r_g_raw sigmoid_val = cute.arch.rcp_approx( - cutlass.Float32(1.0) + cute.math.exp(-exp_A_x, fastmath=True) + cutlass.Float32(1.0) + + cute.math.exp(-exp_A_x, fastmath=True) ) r_gk = lower_bound * sigmoid_val sG[i_t, k_idx] = cute.math.exp(r_gk, fastmath=True) @@ -381,7 +442,9 @@ def kda_decode_mtp_kernel( for i in range(vec_size): k_idx = i * 32 + in_warp_tid for w in range(KERNEL_WIDTH - 1): - cs_q[slot, hk_off + k_idx, w] = r_state[w * vec_size + i] + cs_q[slot, hk_off + k_idx, w] = r_state[ + w * vec_size + i + ] elif warp_idx == 1: for i in range(vec_size): k_idx = i * 32 + in_warp_tid @@ -394,23 +457,31 @@ def kda_decode_mtp_kernel( if warp_idx == 0: for i in range(vec_size): k_idx = i * 32 + in_warp_tid - qkg_cache[slot, cache_pos, 0, hk_off + k_idx] = sQ[i_t, k_idx] + qkg_cache[slot, cache_pos, 0, hk_off + k_idx] = sQ[ + i_t, k_idx + ] for i in range(vec_size): k_idx = i * 32 + in_warp_tid - cs_q[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + cache_pos] = r_state[ - (KERNEL_WIDTH - 2) * vec_size + i - ] + cs_q[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + cache_pos] = ( + r_state[(KERNEL_WIDTH - 2) * vec_size + i] + ) elif warp_idx == 1: for i in range(vec_size): k_idx = i * 32 + in_warp_tid - qkg_cache[slot, cache_pos, 1, hk_off + k_idx] = sK[i_t, k_idx] + qkg_cache[slot, cache_pos, 1, hk_off + k_idx] = sK[ + i_t, k_idx + ] if in_warp_tid == 0: beta_cache[slot, cache_pos, i_hv] = sBeta[i_t] for i in range(vec_size): k_idx = i * 32 + in_warp_tid - cs_k[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + cache_pos] = r_state[ - (KERNEL_WIDTH - 1) * vec_size + (KERNEL_WIDTH - 2) * vec_size + i - ] + cs_k[slot, hk_off + k_idx, KERNEL_WIDTH - 1 + cache_pos] = ( + r_state[ + (KERNEL_WIDTH - 1) * vec_size + + (KERNEL_WIDTH - 2) * vec_size + + i + ] + ) i_t = i_t + 1 else: _v_idx = tidx - 96 @@ -442,7 +513,9 @@ def kda_decode_mtp_kernel( cs_v[slot, hv_off + _v_idx, 0] = _csv1 cs_v[slot, hv_off + _v_idx, 1] = _csv2 cs_v[slot, hv_off + _v_idx, 2] = _xv0 - _vconv1, _vconv2 = cute.arch.mul_packed_f32x2((_csv1, _csv2), (_wv0, _wv0)) + _vconv1, _vconv2 = cute.arch.mul_packed_f32x2( + (_csv1, _csv2), (_wv0, _wv0) + ) _vconv1, _vconv2 = cute.arch.fma_packed_f32x2( (_csv2, _xv0), (_wv1, _wv1), (_vconv1, _vconv2) ) @@ -468,8 +541,12 @@ def kda_decode_mtp_kernel( _i_t = 0 while _i_t < T_loop: if _i_t < commit_len: - sVall[_i_t * V + _v_idx] = cutlass.Float32(v_cache[slot, _i_t, hv_off + _v_idx]) - _xv_replay = cutlass.Float32(cs_v[slot, hv_off + _v_idx, KERNEL_WIDTH - 1 + _i_t]) + sVall[_i_t * V + _v_idx] = cutlass.Float32( + v_cache[slot, _i_t, hv_off + _v_idx] + ) + _xv_replay = cutlass.Float32( + cs_v[slot, hv_off + _v_idx, KERNEL_WIDTH - 1 + _i_t] + ) _csv0 = _csv1 _csv1 = _csv2 _csv2 = _xv_replay @@ -483,9 +560,13 @@ def kda_decode_mtp_kernel( _xv = cutlass.Float32(x_v[0, _token_v, hv_off + _v_idx]) else: _xv = cutlass.Float32(x_v[0, _token_v, i_hv, _v_idx]) - _v_conv += _xv * sConvW[v_weight_base + (KERNEL_WIDTH - 1) * V + _v_idx] + _v_conv += ( + _xv + * sConvW[v_weight_base + (KERNEL_WIDTH - 1) * V + _v_idx] + ) _v_conv = _v_conv * cute.arch.rcp_approx( - cutlass.Float32(1.0) + cute.math.exp(-_v_conv, fastmath=True) + cutlass.Float32(1.0) + + cute.math.exp(-_v_conv, fastmath=True) ) sVall[_i_t * V + _v_idx] = _v_conv _csv0 = _csv1 @@ -521,9 +602,13 @@ def kda_decode_mtp_kernel( v_row = warp_idx * NUM_V_ROWS + row for i in range(vec_size): if cutlass.const_expr(USE_FLAT_LAYOUT): - r_state[row * vec_size + i] = cutlass.Float32(h0[h0_idx, v_row, i * 32 + in_warp_tid]) + r_state[row * vec_size + i] = cutlass.Float32( + h0[h0_idx, v_row, i * 32 + in_warp_tid] + ) else: - r_state[row * vec_size + i] = cutlass.Float32(h0[slot, i_hv, v_row, i * 32 + in_warp_tid]) + r_state[row * vec_size + i] = cutlass.Float32( + h0[slot, i_hv, v_row, i * 32 + in_warp_tid] + ) if cutlass.const_expr(USE_SETMAXREG): cute.arch.warpgroup_reg_alloc(72) cute.arch.barrier() @@ -595,8 +680,12 @@ def kda_decode_mtp_kernel( shk_a = shk_a1 + shk_a2 shk_b = shk_b1 + shk_b2 for offset in [16, 8, 4, 2, 1]: - shk_a += cute.arch.shuffle_sync_bfly(shk_a, offset=offset, mask=-1, mask_and_clamp=31) - shk_b += cute.arch.shuffle_sync_bfly(shk_b, offset=offset, mask=-1, mask_and_clamp=31) + shk_a += cute.arch.shuffle_sync_bfly( + shk_a, offset=offset, mask=-1, mask_and_clamp=31 + ) + shk_b += cute.arch.shuffle_sync_bfly( + shk_b, offset=offset, mask=-1, mask_and_clamp=31 + ) vn_a = r_va - shk_a vn_b = r_vb - shk_b shq_a1 = 0.0 @@ -611,21 +700,25 @@ def kda_decode_mtp_kernel( vnbk_b0, vnbk_b1 = cute.arch.mul_packed_f32x2( (vn_b, vn_b), (r_bk[_p], r_bk[_p + 1]) ) - r_state[ra * vec_size + _p], r_state[ra * vec_size + _p + 1] = cute.arch.fma_packed_f32x2( - src_a=(r_decay[_p], r_decay[_p + 1]), - src_b=( - r_state[ra * vec_size + _p], - r_state[ra * vec_size + _p + 1], - ), - src_c=(vnbk_a0, vnbk_a1), + r_state[ra * vec_size + _p], r_state[ra * vec_size + _p + 1] = ( + cute.arch.fma_packed_f32x2( + src_a=(r_decay[_p], r_decay[_p + 1]), + src_b=( + r_state[ra * vec_size + _p], + r_state[ra * vec_size + _p + 1], + ), + src_c=(vnbk_a0, vnbk_a1), + ) ) - r_state[rb * vec_size + _p], r_state[rb * vec_size + _p + 1] = cute.arch.fma_packed_f32x2( - src_a=(r_decay[_p], r_decay[_p + 1]), - src_b=( - r_state[rb * vec_size + _p], - r_state[rb * vec_size + _p + 1], - ), - src_c=(vnbk_b0, vnbk_b1), + r_state[rb * vec_size + _p], r_state[rb * vec_size + _p + 1] = ( + cute.arch.fma_packed_f32x2( + src_a=(r_decay[_p], r_decay[_p + 1]), + src_b=( + r_state[rb * vec_size + _p], + r_state[rb * vec_size + _p + 1], + ), + src_c=(vnbk_b0, vnbk_b1), + ) ) shq_a1, shq_a2 = cute.arch.fma_packed_f32x2( src_a=( @@ -646,8 +739,12 @@ def kda_decode_mtp_kernel( shq_a = shq_a1 + shq_a2 shq_b = shq_b1 + shq_b2 for offset in [16, 8, 4, 2, 1]: - shq_a += cute.arch.shuffle_sync_bfly(shq_a, offset=offset, mask=-1, mask_and_clamp=31) - shq_b += cute.arch.shuffle_sync_bfly(shq_b, offset=offset, mask=-1, mask_and_clamp=31) + shq_a += cute.arch.shuffle_sync_bfly( + shq_a, offset=offset, mask=-1, mask_and_clamp=31 + ) + shq_b += cute.arch.shuffle_sync_bfly( + shq_b, offset=offset, mask=-1, mask_and_clamp=31 + ) if in_warp_tid == 0: v_row_a = warp_idx * NUM_V_ROWS + ra v_row_b = warp_idx * NUM_V_ROWS + rb @@ -662,9 +759,13 @@ def kda_decode_mtp_kernel( v_row = warp_idx * NUM_V_ROWS + row for i in range(vec_size): if cutlass.const_expr(USE_FLAT_LAYOUT): - ht[h0_idx, v_base + v_row, i * 32 + in_warp_tid] = r_state[row * vec_size + i] + ht[h0_idx, v_base + v_row, i * 32 + in_warp_tid] = r_state[ + row * vec_size + i + ] else: - ht[slot, i_hv, v_base + v_row, i * 32 + in_warp_tid] = r_state[row * vec_size + i] + ht[slot, i_hv, v_base + v_row, i * 32 + in_warp_tid] = r_state[ + row * vec_size + i + ] i_t = i_t + 1 if cutlass.const_expr(PROFILE_STAGES): cute.arch.barrier()