From 6643925d68cbe3f565b76a9bd7f1775c88bb0061 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Mon, 3 Aug 2026 14:09:19 -0700 Subject: [PATCH 01/12] [None][feat] Add Kimi K3 fused attention-residual kernel and torch op Wave-A pre-flight (plan item A4), based directly on main. Cherry of feat/kimi_k3 b42090808ec plus its later sm_100-family pinning fix. All-new files plus CMake registrations. Same content as k3/14813-preflight-a4-attn-res-kernel minus dependency-branch context in the shared CMake files. Signed-off-by: Brian Nguyen --- cpp/tensorrt_llm/CMakeLists.txt | 1 + cpp/tensorrt_llm/kernels/CMakeLists.txt | 3 + .../kernels/kimiK3AttnRes/CMakeLists.txt | 31 + .../kernels/kimiK3AttnRes/attnResFwd.cu | 1533 +++++++++++++++++ .../kernels/kimiK3AttnRes/attnResFwd.h | 59 + cpp/tensorrt_llm/thop/CMakeLists.txt | 1 + cpp/tensorrt_llm/thop/attnResOp.cpp | 140 ++ 7 files changed, 1768 insertions(+) create mode 100644 cpp/tensorrt_llm/kernels/kimiK3AttnRes/CMakeLists.txt create mode 100644 cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu create mode 100644 cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h create mode 100644 cpp/tensorrt_llm/thop/attnResOp.cpp diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 20e212149d18..d9a7853bc9d5 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -187,6 +187,7 @@ set(TRTLLM_LINK_LIBS selective_scan_src mamba2_mtp_ssm_cache_src kda_decode_src + kimi_k3_attn_res_src ws_layernorm_src fusedGatedRMSNormQuant_src fpA_intB_gemm_src diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index 45f30d424171..a8dae5de5b3f 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -31,6 +31,7 @@ add_subdirectory(causalConv1d) add_subdirectory(fusedGatedRMSNormQuant) add_subdirectory(mamba2MTPSSMCache) add_subdirectory(kdaDecode) +add_subdirectory(kimiK3AttnRes) add_subdirectory(mhcKernels) add_subdirectory(compressorKernels) @@ -60,6 +61,8 @@ list(FILTER SRC_CU EXCLUDE REGEX "fusedGatedRMSNormQuant/.*") list(FILTER SRC_CU EXCLUDE REGEX "mamba2MTPSSMCache/.*") list(FILTER SRC_CPP EXCLUDE REGEX "kdaDecode/.*") list(FILTER SRC_CU EXCLUDE REGEX "kdaDecode/.*") +list(FILTER SRC_CPP EXCLUDE REGEX "kimiK3AttnRes/.*") +list(FILTER SRC_CU EXCLUDE REGEX "kimiK3AttnRes/.*") list(FILTER SRC_CPP EXCLUDE REGEX "mhcKernels/.*") list(FILTER SRC_CU EXCLUDE REGEX "mhcKernels/.*") list(FILTER SRC_CPP EXCLUDE REGEX "compressorKernels/.*") diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/CMakeLists.txt b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/CMakeLists.txt new file mode 100644 index 000000000000..d282ac9ee3f8 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/CMakeLists.txt @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not +# use this file except in compliance with the License. You may obtain a copy of +# the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations under +# the License. + +add_library(kimi_k3_attn_res_src OBJECT attnResFwd.cu) + +# The kernel is warp-specialized for the SM100 (datacenter Blackwell) family and +# uses tcgen05/TMEM PTX instructions that ptxas rejects for other architectures +# (e.g. sm_120f), so pin the target to the sm_100 family instead of inheriting +# the global architecture list. On builds without any sm_100-family architecture +# the target still compiles (the tcgen05 code paths are guarded by +# __CUDA_ARCH__), and the Torch-op bridge rejects unsupported devices at +# runtime. +set_cuda_architectures(kimi_k3_attn_res_src 100f) + +set_property(TARGET kimi_k3_attn_res_src PROPERTY POSITION_INDEPENDENT_CODE ON) +set_property(TARGET kimi_k3_attn_res_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS + ON) +target_compile_options(kimi_k3_attn_res_src + PRIVATE $<$:--use_fast_math>) diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu new file mode 100644 index 000000000000..36a297bec25f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu @@ -0,0 +1,1533 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Fused Kimi K3 attention-residual forward for Blackwell (sm_100 family). +// +// Warp-specialized online softmax + residual + RMSNorm: +// - 1 producer warp issues cp.async.bulk row loads into shared memory. +// - 8 consumer warps compute reductions and output. +// - Q=res_weight*rms_weight remains in registers across persistent tokens. +// - V rows are converted once and cached as FP32 in TMEM between passes. +// +// Contract: B=1, N<=12, H in [4096,8192] and divisible by 1024; checked at +// the Torch-op bridge. The kernel uses separate layer/block residual inputs +// and does not require a concatenated V tensor. +// +// Source-integrated from the NVIDIA+Moonshot jointly developed +// Attention_residual kernel at e7f934124acc915575f9f7561f9d1e373ab43089. + +#include "tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h" + +#include +#include +#include +#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 + 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) { + int lane = threadIdx.x & 31; + int wid = threadIdx.x >> 5; + val = warp_reduce_sum(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); + 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) { + if (n < N - 1) + return block_res + (((long long)n * T + t) * B + b) * H; + return layer_res + ((long long)t * B + b) * H; +} + +namespace sm100 { + + +CUTE_DEVICE +void tcgen05_after_thread_sync() { + asm volatile("tcgen05.fence::after_thread_sync;"); +} + +CUTE_DEVICE +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)); +} + +CUTE_DEVICE +float2 float2_sub(const float2& a, const float2& b) { + float2 c; + 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 c; + 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 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))); + return d; +} + +CUTE_DEVICE +float2 float2_add(const float2& a, const float2& b) { + float2 c; + 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_) { + uint32_t* dst_ptr = reinterpret_cast(dst_ptr_); + 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"(src_addr)); + } 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"(src_addr)); + } +} + +template +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) { + 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 { + 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)); + } +} + + +namespace fwd_prod_v2 { + +using namespace cute; + +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 CONSUMER_WARPS = CONSUMER_THREADS / 32; +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 { + alignas(16) uint64_t bar_ready[CHUNK_DEPTH]; + alignas(16) uint64_t bar_consumed[CHUNK_DEPTH]; + alignas(16) float2 ws_stats[CONSUMER_WARPS][NC]; + alignas(16) float logits_all[N_MAX]; + uint32_t tmem_base; +}; + +__device__ __forceinline__ +void cp_async_bulk(void* smem_dst, const void* 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"); +} + +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) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + constexpr float LOG2_E = 1.4426950408889634f; + constexpr int N_CHUNK = NC; + constexpr int NUM_BUFS = CHUNK_DEPTH * NC; + 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 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; + + 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; + + 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); + 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++) { + 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) { + cute::TMEM::Allocator1Sm alloc; + alloc.allocate(TMEM_COLS_ALLOC, &plan.tmem_base); + 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; + + 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++) { + int h = h_base + j; + q_cache[si * VEC + j] = + __bfloat162float(rms_w[h]) * + __bfloat162float(res_w[h]); + } + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + int h_base = dt * K_TILE + k_local; + #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]); + } + } + } + + 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++) { + int ns = ci * N_CHUNK; + int an = FULL_N12 ? N_CHUNK : min(N_CHUNK, N - ns); + 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; + } + 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]); + } + } + } + } + } 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) { + float m_running = -FLT_MAX; + float s_running = 0.f; + #pragma unroll + for (int i = 0; i < ACC_PER_THREAD; i++) { + acc32[i] = 0.f; + } + + 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 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) { + 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; + const float* qv = &q_cache[si * VEC]; + #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]); + } + } + 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]); + } + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + const float* qv = &q_cache[si * VEC]; + + #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); + __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)); + 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]); + } + } + if constexpr (!FULL_N12) { + cutlass::arch::fence_view_async_tmem_store(); + } + }; + 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 == 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 { + 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(); + } + } + 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 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); + } + } + if constexpr (FULL_N12) { + cutlass::arch::fence_view_async_tmem_store(); + } + if (lane == 0) { + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + plan.ws_stats[comp_wid][n] = reduce_pair[n]; + } + } + cutlass::arch::NamedBarrier::sync(CONSUMER_THREADS, 0); + + float local_rsig = 0.f; + float local_logit = 0.f; + 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]); + } + local_rsig = rsqrtf(totals.x / H + eps_cache); + local_logit = totals.y * local_rsig; + }; + if constexpr (FULL_N12) { + cross_warp_tail(lane & (N_CHUNK - 1)); + } 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); + } + + float m_chunk = -FLT_MAX; + 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) { + m_chunk = fmaxf(m_chunk, logit_n[n]); + } + } + } + float m_new = fmaxf(m_running, m_chunk); + 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); + } + 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) { + 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); + 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]); + a[j] = float2_mul(old, corr2); + } + float2 f_cache[AN][2]; + #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])); + } + #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++) { + acc32[si * VEC + 2 * j] = a[j].x; + acc32[si * VEC + 2 * j + 1] = a[j].y; + } + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + 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]); + a[j] = float2_mul(old, corr2); + } + float2 f_cache[AN][VEC / 2]; + #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])); + } + #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++) { + 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(); + } + } 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 { + 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(); + } + } + + s_running = s_running * corr + w_sum; + m_running = m_new; + + if (comp_wid == 0 && lane < an) { + int ng = ns + lane; + rsigma_out[(long long)ng * TB + tb] = local_rsig; + plan.logits_all[ng] = local_logit; + } + } + + 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; + uint2 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)); + } + *reinterpret_cast(out_ptr + h_base) = ov; + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + int h_base = dt * K_TILE + k_local; + uint4 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)); + } + *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; + float lg = plan.logits_all[lane]; + logits_out[out_idx] = lg; + probs_out[out_idx] = + exp2f((lg - m_running) * LOG2_E) * inv_s; + } + } + } + + if (wid > 0) { + cutlass::arch::NamedBarrier::sync(CONSUMER_THREADS, 2); + } + 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"); +#endif +} + +// N=1 specialization: softmax is degenerate, so output is the layer row. +// 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) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + 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 = 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; + + 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); + 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); + }; + + 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) { + cute::TMEM::Allocator1Sm alloc; + alloc.allocate(TMEM_Q_COLS_TOTAL, &plan.tmem_base); + 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; + + if (comp_tid >= 0) { + float q32[ACC_PER_THREAD]; + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + int h_base = dt * K_TILE + k_local; + #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]); + } + } + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + int dt = si * CONSUMER_GROUPS + group; + 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()) { + long long tile_i = 0; + 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 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), + plan.bar_ready[slot]); + } + } + } else { + long long tile_i = 0; + 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 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; + int tb = tb0 + r; + bf16_t* row_ptr = tile_ptr(slot, r); + 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++) { + int dt = si * CONSUMER_GROUPS + group; + 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); + *reinterpret_cast(out_ptr + h_base) = vp; + + __nv_bfloat162* v2 = reinterpret_cast<__nv_bfloat162*>(&vp); + float2 f0 = __bfloat1622float2(v2[0]); + 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]; + } + + #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); + } + cutlass::arch::NamedBarrier::sync(CONSUMER_THREADS, 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]); + } + 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; + } + cutlass::arch::NamedBarrier::sync(CONSUMER_THREADS, 1); + } + cute::arrive_barrier(plan.bar_consumed[slot]); + } + } + + __syncthreads(); + 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"); +#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) +{ + 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); + } + 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); +} + +// Small-N counterpart to the Triton one-program topology. One CTA owns the +// complete token, with exactly 28 hidden elements per thread at H=7168. For +// N=2/4, packed BF16 V remains in registers across the statistics/softmax +// 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) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + constexpr int H = 7168; + constexpr int THREADS = 256; + constexpr int WARPS = THREADS / 32; + constexpr int ITEMS = H / THREADS; + constexpr float LOG2_E = 1.4426950408889634f; + static_assert(H % THREADS == 0); + static_assert(N == 1 || N == 2 || N == 4); + + __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; + + float2 stats[N] = {}; + uint32_t v_cache_bf16[ITEMS][(N + 1) / 2]; + #pragma unroll + for (int item = 0; item < ITEMS; item++) { + int h = tid + item * THREADS; + 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; + bf16_t packed_v = row[h]; + float v = __bfloat162float(packed_v); + if constexpr (N == 1) { + output[h] = packed_v; + } else { + item_v[n] = packed_v; + } + 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 { + __nv_bfloat162 bf16x2; + uint32_t bits; + } packed; + 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++) { + 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++) { + warp_stats[warp * N + n] = stats[n]; + } + } + __syncthreads(); + + if (tid < N) { + float2 total = {}; + #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) { + float local_rsigma[N]; + float local_logits[N]; + float max_logit = -FLT_MAX; + #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++) { + 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++) { + weights[n] *= inv_denominator; + rsigma_out[n] = local_rsigma[n]; + logits_out[n] = local_logits[n]; + probs_out[n] = weights[n]; + } + } + __syncthreads(); + + 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 { + __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); + value = fmaf(weights[2 * pair + 1], v.y, value); + } + int h = tid + item * THREADS; + output[h] = __float2bfloat16_rn(value); + } + } +#else + 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, + 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); +} + +// Single-token split-K specialization. The complete grid is one CTA cluster: +// rank g owns a disjoint H/GROUPS slice, keeps that slice of FP32 V in its +// 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) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + namespace cg = cooperative_groups; + constexpr int H = 7168; + constexpr int K_PER_CTA = H / GROUPS; + constexpr int THREADS = 256; + constexpr int WARPS = THREADS / 32; + constexpr float LOG2_E = 1.4426950408889634f; + static_assert(H % GROUPS == 0); + + 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)); + float* weights = reinterpret_cast(warp_stats + WARPS * N); + + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + cg::cluster_group cluster = cg::this_cluster(); + const int group = cluster.block_rank(); + const int h_begin = group * K_PER_CTA; + + float sq[N] = {}; + float dot[N] = {}; + #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 v = __bfloat162float(row[h]); + 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++) { + 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++) { + warp_stats[warp * N + n] = make_float2(sq[n], dot[n]); + } + } + __syncthreads(); + + if (tid < N) { + float2 total = {}; + #pragma unroll + for (int w = 0; w < WARPS; w++) { + total = float2_add(total, warp_stats[w * N + tid]); + } + warp_stats[tid] = total; + } + + // Publish every rank's reduced statistics to distributed shared memory. + cluster.sync(); + + // 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) { + float2 total = {}; + #pragma unroll + for (int g = 0; g < GROUPS; g++) { + const float2* remote_stats = + cluster.map_shared_rank(warp_stats, g); + total = float2_add(total, remote_stats[tid]); + } + warp_stats[tid] = total; + } + __syncthreads(); + + 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++) { + 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++) { + 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++) { + weights[n] *= inv_sum; + if (group == 0) { + rsigma_out[n] = local_rsigma[n]; + logits_out[n] = local_logits[n]; + probs_out[n] = weights[n]; + } + } + } + + cluster.sync(); + + #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); + } + output[h_begin + ki] = __float2bfloat16_rn(value); + } +#else + 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, + 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); + auto kernel = &attn_res_fwd_s1_splitk_kernel; + static bool attrs_set = false; + if (!attrs_set) { + cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + attrs_set = true; + } + 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); + config.dynamicSmemBytes = smem_size; + config.stream = stream; + cudaLaunchAttribute attribute{}; + attribute.id = cudaLaunchAttributeClusterDimension; + attribute.val.clusterDim.x = GROUPS; + attribute.val.clusterDim.y = 1; + attribute.val.clusterDim.z = 1; + config.attrs = &attribute; + config.numAttrs = 1; + 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, + 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); + 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); + } + attrs_set = true; + } + kernel<<>>( + layer_residual, res_weight, rms_weight, + output, rsigma, probs, logits, T, B, rms_eps); +} + +} // namespace fwd_prod_v2 +} // namespace sm100 + +int attn_res_fwd_grid_size(int dev) +{ + static int cached_num_sm[64] = {}; + if (dev >= 0 && dev < 64 && cached_num_sm[dev] > 0) + { + return cached_num_sm[dev]; + } + int n = 0; + cudaError_t err = cudaDeviceGetAttribute(&n, cudaDevAttrMultiProcessorCount, dev); + if (err != cudaSuccess || n <= 0) + return 0; + if (dev >= 0 && dev < 64) + { + cached_num_sm[dev] = n; + } + return n; +} + +} // namespace + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::kimiK3AttnRes +{ + +void invokeAttnResFwd(AttnResFwdParams const& params, cudaStream_t stream) +{ + using namespace sm100::fwd_prod_v2; + + bf16_t const* block_residual = params.blockResidual; + bf16_t const* layer_residual = params.layerResidual; + bf16_t const* res_weight = params.resWeight; + bf16_t const* rms_weight = params.rmsWeight; + bf16_t* output = params.output; + float* rsigma = params.rsigma; + float* probs = params.probs; + float* logits = params.logits; + int const N = params.numCandidates; + int const T = params.seqLen; + int const B = params.batchSize; + int const H = params.hiddenSize; + float const rms_eps = params.rmsEps; + + int dev = 0; + cudaGetDevice(&dev); + int num_sm = attn_res_fwd_grid_size(dev); + if (num_sm <= 0 || N > N_MAX) + { + return; + } + + if (H == 8192) + { + if (N == 1) + { + launch_n1_ttile<8192, 2>( + layer_residual, res_weight, rms_weight, output, rsigma, probs, logits, T, B, rms_eps, num_sm, stream); + } + else if (N <= 2) + { + launch_fwd<8192, 2, true>(block_residual, layer_residual, res_weight, rms_weight, output, rsigma, probs, + logits, N, T, B, rms_eps, num_sm, stream); + } + else + { + launch_fwd<8192, 4, false>(block_residual, layer_residual, res_weight, rms_weight, output, rsigma, probs, + logits, N, T, B, rms_eps, num_sm, stream); + } + } + else if (H == 7168) + { + if (T == 1 && N == 1) + { + launch_s1_single_cta<1>( + block_residual, layer_residual, res_weight, rms_weight, output, rsigma, probs, logits, rms_eps, stream); + } + else if (T == 1 && N == 2) + { + launch_s1_single_cta<2>( + block_residual, layer_residual, res_weight, rms_weight, output, rsigma, probs, logits, rms_eps, stream); + } + else if (T == 1 && N == 4) + { + launch_s1_single_cta<4>( + block_residual, layer_residual, res_weight, rms_weight, output, rsigma, probs, logits, rms_eps, stream); + } + else if (T == 1 && N == 8) + { + launch_s1_splitk<8, 8>( + block_residual, layer_residual, res_weight, rms_weight, output, rsigma, probs, logits, rms_eps, stream); + } + else if (T == 1 && N == 12) + { + launch_s1_splitk<12, 8>( + block_residual, layer_residual, res_weight, rms_weight, output, rsigma, probs, logits, rms_eps, 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); + } + else + { + launch_fwd<7168, 4, false>(block_residual, layer_residual, res_weight, rms_weight, output, rsigma, probs, + logits, N, T, B, rms_eps, num_sm, stream); + } + } + else if (H == 6144) + { + launch_fwd<6144, 4, false>(block_residual, layer_residual, res_weight, rms_weight, output, rsigma, probs, + logits, N, T, B, rms_eps, num_sm, stream); + } + else if (H == 5120) + { + launch_fwd<5120, 4, false>(block_residual, layer_residual, res_weight, rms_weight, output, rsigma, probs, + logits, N, T, B, rms_eps, num_sm, stream); + } + else if (H == 4096) + { + if (N == 1) + { + launch_n1_ttile<4096, 4>( + layer_residual, res_weight, rms_weight, output, rsigma, probs, logits, T, B, rms_eps, num_sm, stream); + } + else if (N <= 4) + { + launch_fwd<4096, 2, true>(block_residual, layer_residual, res_weight, rms_weight, output, rsigma, probs, + logits, N, T, B, rms_eps, num_sm, stream); + } + else + { + launch_fwd<4096, 3, true>(block_residual, layer_residual, res_weight, rms_weight, output, rsigma, probs, + logits, N, T, B, rms_eps, num_sm, stream); + } + } +} + +} // namespace kernels::kimiK3AttnRes + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h new file mode 100644 index 000000000000..65a9913f5afb --- /dev/null +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h @@ -0,0 +1,59 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::kimiK3AttnRes +{ + +//! Parameters for the fused Kimi K3 attention-residual forward kernel +//! (warp-specialised online softmax + residual selection + RMSNorm, +//! Blackwell sm_100/sm_103 only). +//! +//! Contract (checked at the Torch-op bridge): B == 1, N in [1, 12], +//! T in [1, 16384], H a multiple of 1024 in [4096, 8192]; all residual +//! tensors bf16 contiguous, rsigma/probs/logits fp32 [N, T, B]. +//! blockResidual may be nullptr when N == 1. +struct AttnResFwdParams +{ + __nv_bfloat16 const* blockResidual; // [N-1, T, B, H], nullptr iff N == 1 + __nv_bfloat16 const* layerResidual; // [T, B, H] + __nv_bfloat16 const* resWeight; // [H] + __nv_bfloat16 const* rmsWeight; // [H] + __nv_bfloat16* output; // [T, B, H] + float* rsigma; // [N, T, B] + float* probs; // [N, T, B] + float* logits; // [N, T, B] + int numCandidates; // N = K + 1 + int seqLen; // T + int batchSize; // B + int hiddenSize; // H + float rmsEps; +}; + +//! Launches the fused attention-residual forward on the supplied stream. +void invokeAttnResFwd(AttnResFwdParams const& params, cudaStream_t stream); + +} // namespace kernels::kimiK3AttnRes + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index f64b815e4240..7af68a90f1c9 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -76,6 +76,7 @@ add_library( deepseekV4BlockTableOp.cpp inverseRopeFp8QuantOp.cpp kdaDecodeOp.cpp + attnResOp.cpp fusedQKNormRopeOp.cpp fusedAdaptiveLayerNormOp.cpp fusedDiTQKNormRopeOp.cpp diff --git a/cpp/tensorrt_llm/thop/attnResOp.cpp b/cpp/tensorrt_llm/thop/attnResOp.cpp new file mode 100644 index 000000000000..c45201c62ba2 --- /dev/null +++ b/cpp/tensorrt_llm/thop/attnResOp.cpp @@ -0,0 +1,140 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h" + +#include +#include +#include +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +bool is_sm100_family() +{ + int dev = 0; + cudaGetDevice(&dev); + static int cached_state[64] = {}; // 0 unknown, 1 false, 2 true + if (dev >= 0 && dev < 64 && cached_state[dev] != 0) + { + return cached_state[dev] == 2; + } + cudaDeviceProp prop; + cudaGetDeviceProperties(&prop, dev); + // The kernel binary is compiled for the sm_100 family only (it relies on + // tcgen05/TMEM, which later architectures such as sm_120 do not support), + // so require compute capability major == 10 rather than >= 10. + bool const ok = prop.major == 10; + if (dev >= 0 && dev < 64) + { + cached_state[dev] = ok ? 2 : 1; + } + return ok; +} + +void check_attn_res_contract(int N, int T, int B, int H) +{ + TORCH_CHECK(is_sm100_family(), "attn_res_fwd requires an sm_100-family (datacenter Blackwell) GPU"); + TORCH_CHECK(B == 1, "attn_res_fwd: unsupported B=", B, " (only B=1 is supported)"); + TORCH_CHECK(N >= 1 && N <= 12, "attn_res_fwd: unsupported N=", N, " (must be in [1, 12])"); + TORCH_CHECK(T >= 1 && T <= 16384, "attn_res_fwd: unsupported T=", T, " (must be in [1, 16384])"); + TORCH_CHECK(H >= 4096 && H <= 8192 && H % 1024 == 0, "attn_res_fwd: unsupported H=", H, + " (must be a multiple of 1024 in [4096, 8192])"); +} + +std::tuple attn_res_fwd( + at::Tensor layer_residual, at::Tensor block_residual, at::Tensor res_weight, at::Tensor rms_weight, double rms_eps) +{ + TORCH_CHECK(layer_residual.dim() == 3, "attn_res_fwd: layer_residual must be [T, B, H]"); + TORCH_CHECK(block_residual.dim() == 4, "attn_res_fwd: block_residual must be [K, T, B, H]"); + + int const T = static_cast(layer_residual.size(0)); + 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"); + 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"); + TORCH_CHECK(rms_weight.scalar_type() == at::kBFloat16, "attn_res_fwd: rms_weight must be bf16"); + TORCH_CHECK(layer_residual.is_contiguous() && block_residual.is_contiguous() && res_weight.is_contiguous() + && rms_weight.is_contiguous(), + "attn_res_fwd: inputs must be contiguous"); + TORCH_CHECK(block_residual.sizes() == at::IntArrayRef({N - 1, T, B, H}), + "attn_res_fwd: block_residual shape must match layer_residual"); + TORCH_CHECK(res_weight.numel() == H, "attn_res_fwd: res_weight must have H elements"); + TORCH_CHECK(rms_weight.numel() == H, "attn_res_fwd: rms_weight must have H elements"); + + auto output = at::empty_like(layer_residual); + auto float_options = layer_residual.options().dtype(at::kFloat); + auto rsigma = at::empty({N, T, B}, float_options); + auto probs = at::empty({N, T, B}, float_options); + auto logits = at::empty({N, T, B}, float_options); + + kernels::kimiK3AttnRes::AttnResFwdParams params{}; + params.blockResidual = N > 1 ? reinterpret_cast<__nv_bfloat16 const*>(block_residual.const_data_ptr()) : nullptr; + params.layerResidual = reinterpret_cast<__nv_bfloat16 const*>(layer_residual.const_data_ptr()); + params.resWeight = reinterpret_cast<__nv_bfloat16 const*>(res_weight.const_data_ptr()); + params.rmsWeight = reinterpret_cast<__nv_bfloat16 const*>(rms_weight.const_data_ptr()); + params.output = reinterpret_cast<__nv_bfloat16*>(output.data_ptr()); + params.rsigma = rsigma.data_ptr(); + params.probs = probs.data_ptr(); + params.logits = logits.data_ptr(); + params.numCandidates = N; + params.seqLen = T; + params.batchSize = B; + params.hiddenSize = H; + params.rmsEps = static_cast(rms_eps); + + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + kernels::kimiK3AttnRes::invokeAttnResFwd(params, stream); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + return {output, rsigma, probs, logits}; +} + +} // namespace + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "attn_res_fwd(Tensor layer_residual, Tensor block_residual, " + "Tensor res_weight, Tensor rms_weight, float rms_eps) " + "-> (Tensor, Tensor, Tensor, Tensor)"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("attn_res_fwd", &tensorrt_llm::torch_ext::attn_res_fwd); +} From b8260aade97f286230bde780bbf4a78b86ab45ed Mon Sep 17 00:00:00 2001 From: "Brian Nguyen (TensorRT)" Date: Sun, 26 Jul 2026 13:42:20 -0700 Subject: [PATCH 02/12] [None][feat] Add CuTe DSL KDA prefill and MTP decode kernels for Kimi K3 Add Blackwell CuTe DSL kernels for KDA chunked prefill (fused K1/K2/K3/K4 stages, standalone K4 persistent variant, and the A_kk inverse solve) and for KDA MTP speculative decode with replay-cache semantics, plus the custom-op wrappers (kda_prefill, KDA MTP ops) that JIT-compile and dispatch them. Prefill supports varlen batches (including empty and small-tail batches) with shape-independent JIT caching and OOB guards on the beta/state tiles; the op layer validates cu_seqlens/chunk_indices dtype consistency to avoid mixed-dtype capture corruption. Co-authored-by: Tao Li Co-authored-by: Pengbo Wang Signed-off-by: Brian Nguyen --- tensorrt_llm/_torch/custom_ops/__init__.py | 18 + .../custom_ops/cute_dsl_kimi_k3_custom_ops.py | 1331 ++++++++++ .../cute_dsl_kimi_k3_kda_mtp_ops.py | 738 ++++++ .../blackwell/kimi_k3_kda/__init__.py | 15 + .../blackwell/kimi_k3_kda/akk_inverse.py | 1075 ++++++++ .../blackwell/kimi_k3_kda/fused_k123.py | 2306 +++++++++++++++++ .../blackwell/kimi_k3_kda/fused_k1234.py | 2276 ++++++++++++++++ .../blackwell/kimi_k3_kda/k4_persistent.py | 1459 +++++++++++ .../blackwell/kimi_k3_kda/kda_mtp_decode.py | 678 +++++ 9 files changed, 9896 insertions(+) create mode 100644 tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py create mode 100644 tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/__init__.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/akk_inverse.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k1234.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py diff --git a/tensorrt_llm/_torch/custom_ops/__init__.py b/tensorrt_llm/_torch/custom_ops/__init__.py index 600c4f501ac7..de9202100657 100644 --- a/tensorrt_llm/_torch/custom_ops/__init__.py +++ b/tensorrt_llm/_torch/custom_ops/__init__.py @@ -1,3 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import torch from ..cuda_tile_utils import IS_CUDA_TILE_AVAILABLE @@ -70,6 +84,10 @@ def inplace_slice_copy(dest: torch.Tensor, src: torch.Tensor, dim1_start: int, from .cute_dsl_megamoe_custom_op import cute_dsl_megamoe_nvfp4_blackwell __all__ += ['cute_dsl_megamoe_nvfp4_blackwell'] +if IS_CUTLASS_DSL_AVAILABLE and IS_FLASHINFER_AVAILABLE: + from .cute_dsl_kimi_k3_custom_ops import kda_prefill + __all__ += ['kda_prefill'] + if IS_CUDA_TILE_AVAILABLE: from .cuda_tile_custom_ops import (cuda_tile_rms_norm, cuda_tile_rms_norm_fuse_residual_) 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 new file mode 100644 index 000000000000..1ce0861a0294 --- /dev/null +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py @@ -0,0 +1,1331 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CuTe DSL custom op for Kimi K3 KDA prefill. + +The implementation source-integrates the production orchestration from the +KDA ``chunk_kda_fwd`` benchmark. The default pipeline launches fused K123, +Akk inverse, and persistent K4 kernels. Equal-length inputs may instead use +the fused K1234 kernel. + +Supported: + - Equal-length sequences (B ≥ 1, with the existing B=1 padding path) + - Variable-length sequences (B=1 with cu_seqlens) + - safe_gate mode (sigmoid + lower_bound) and softplus mode + - dt_bias + - use_gate_in_kernel=True with A_log + - chunk_size=64 + +The public ``trtllm::kda_prefill`` operator returns only the KDA output and +final recurrent state. Intermediate matrices remain private runner workspace. +""" + +from typing import Optional, Tuple + +import weakref + +import torch + +from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE + +if IS_CUTLASS_DSL_AVAILABLE and IS_FLASHINFER_AVAILABLE: + import cuda.bindings.driver as cuda_driver + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import from_dlpack + + from ..cute_dsl_kernels.blackwell.kimi_k3_kda.akk_inverse import akk_inv_host as _akk_inv_host + from ..cute_dsl_kernels.blackwell.kimi_k3_kda.fused_k123 import ( + make_host_function as _fused_make_host, + ) + from ..cute_dsl_kernels.blackwell.kimi_k3_kda.fused_k1234 import ( + make_host_fn as _fused_k1234_make_host, + ) + from ..cute_dsl_kernels.blackwell.kimi_k3_kda.k4_persistent import ( + BYTES_PER_TENSORMAP as _K4P_BTM, + ) + from ..cute_dsl_kernels.blackwell.kimi_k3_kda.k4_persistent import NUM_TENSORMAPS as _K4P_NTM + from ..cute_dsl_kernels.blackwell.kimi_k3_kda.k4_persistent import ( + make_host_fn as _k4p_make_host, + ) + from ..modules.fla.index import prepare_chunk_indices +else: + raise ImportError("Kimi K3 KDA prefill requires NVIDIA CUTLASS DSL and FlashInfer") + + +def _ct(t, etype): + """Create a CuTe tensor from PyTorch tensor.""" + r = from_dlpack(t, assumed_align=16) + r.element_type = etype + return r + + +def _current_cu_stream(device): + """CUstream handle of torch's CURRENT stream, queried per call. + + The executor runs the model on a dedicated non-blocking + ``torch.cuda.Stream`` (``py_executor.execution_stream``), not the default + stream. Every kernel launch must go to this stream: launching on the DSL + default stream (what ``.launch()`` does without a ``stream`` argument) + races with the projections that produce q/k/v/g/beta and with the + consumers of O/final_state — silent, intermittent corruption in the + runtime while single-stream unit tests pass. This is also why the value + must never be cached alongside the scratch buffers: the buffers may be + created under a different stream (e.g. warmup) than later forwards. + """ + return cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream) + + +# Cached eqlen dummy cu/ci cute wrappers — shared by K123 and akk_inv. +# Avoids per-call torch.empty + from_dlpack overhead (~10-12us each). +_eqlen_dummy_cache = {} + + +def _get_eqlen_dummies(device, idx_dtype=torch.int64): + """Returns cached (cu_ct, ci_ct) cute wrappers for eqlen (B+1=2, NT+1=2).""" + key = (device.index if device.index is not None else 0, idx_dtype) + if key not in _eqlen_dummy_cache: + cu_t = torch.empty(2, dtype=idx_dtype, device=device) + ci_t = torch.empty(1, 2, dtype=idx_dtype, device=device) + cu_etype = cutlass.Int64 if idx_dtype == torch.int64 else cutlass.Int32 + _eqlen_dummy_cache[key] = (_ct(cu_t, cu_etype), _ct(ci_t, cu_etype)) + return _eqlen_dummy_cache[key] + + +def _cute_int_type(dtype): + """Map PyTorch integer dtype to CUTLASS element type.""" + if dtype == torch.int32: + return cutlass.Int32 + elif dtype == torch.int64: + return cutlass.Int64 + else: + raise ValueError(f"Unsupported integer dtype: {dtype}") + + +# ========== Fused K1+K2+K3 compilation cache ========== +_fused_k123_cache = {} +# id(cu_seqlens) -> bool. Skips per-call GPU->CPU sync on subsequent calls +# when the same cu_seqlens tensor is reused (typical training/inference loop). +_varlen_pure_cache = {} +# id(cu_seqlens) -> int seqlen, populated alongside _varlen_pure_cache for +# single-seq cu_seqlens. +_varlen_single_seqlen_cache = {} + +# id(tensor) -> cute_wrapper. The wrappers themselves are stateless views +# over the tensor's storage, so they remain valid as long as the tensor's +# data pointer / shape / strides don't change. Callers that reuse the same +# tensor objects across iterations (typical benchmark pattern) hit the +# cache; per-call fresh tensors (the executor runtime pattern) rebuild. +_input_wrap_cache = {} + + +def _prune_on_gc(cache, key, *keyobjs): + """Drop ``cache[key]`` when any of ``keyobjs`` is garbage-collected. + + The id()-keyed caches in this module are only sound while the keyed + Python objects stay alive — CPython reuses an object's address (its id) + the moment it is freed, so a dead entry could otherwise be served for an + unrelated new tensor with recycled id but different contents/storage. + A finalizer runs at dealloc, before the address can be recycled, so a + pruned entry can never alias a new object. This also bounds cache growth + when callers pass fresh tensors every call. + """ + for o in keyobjs: + weakref.finalize(o, cache.pop, key, None) + + +def _ct_cached(t, etype): + """`_ct(t, etype)` with id(t)-based cache. Returns the same cute wrapper + for repeated calls with the same tensor object, avoiding per-call + `from_dlpack` overhead (~5-10us each). + + ONLY use for tensors with process-long lifetime (module params, the + module-level scratch from ``_get_buffers``): the cached wrapper pins the + tensor's storage, so the weakref pruning never fires for the keyed + object and a per-call activation would be pinned forever (~100MB/call + leak in the executor runtime). Per-call tensors must use plain ``_ct``. + """ + key = (id(t), etype) + w = _input_wrap_cache.get(key) + if w is None: + w = _ct(t, etype) + _input_wrap_cache[key] = w + _prune_on_gc(_input_wrap_cache, key, t) + return w + + +# Cache for dt_bias `.float().contiguous().view(H, K)` + cute wrapper. +# dt_bias is typically a nn.Parameter — same object across iterations. +_dt_bias_cache = {} + + +def _get_dt_bias_ct(dt_bias, H, K): + """Returns cached cute wrapper for dt_bias.float().view(H, K).""" + key = (id(dt_bias), H, K) + entry = _dt_bias_cache.get(key) + if entry is None: + # Copy (never view) the keyed object: a view would pin it, making + # the entry immortal (weakref pruning could never fire). Callers + # pass a fresh .detach() per call, so an aliasing entry would leak. + bias_t = dt_bias.detach().float().reshape(H, K).clone() + entry = (bias_t, _ct(bias_t, cutlass.Float32)) + _dt_bias_cache[key] = entry + _prune_on_gc(_dt_bias_cache, key, dt_bias) + return entry[1] + + +# Cache for the empty 1x1 fp32 bias tensor used when dt_bias is None. +_empty_bias_cache = {} + + +def _get_empty_bias_ct(device): + idx = device.index if device.index is not None else 0 + if idx not in _empty_bias_cache: + t = torch.empty(1, 1, dtype=torch.float32, device=device) + _empty_bias_cache[idx] = _ct(t, cutlass.Float32) + return _empty_bias_cache[idx] + + +# K4 varlen cu_seqlens / chunk_offsets cute wrappers. Caches the int32-cast +# tensor and its mark_layout_dynamic wrapper so they survive multiple calls +# with the same input objects. +_k4_varlen_cu_co_cache = {} + + +def _get_k4_varlen_cu_co(cu_seqlens, chunk_offsets): + key = (id(cu_seqlens), id(chunk_offsets)) + entry = _k4_varlen_cu_co_cache.get(key) + if entry is None: + # Always materialize copies: if the cached value aliased the keyed + # object, the entry would pin it and the weakref pruning could + # never fire (immortal entry). + cu_int32 = cu_seqlens.to(torch.int32).contiguous() + if cu_int32 is cu_seqlens: + cu_int32 = cu_seqlens.clone() + co_int32 = chunk_offsets.to(torch.int32).contiguous() + if co_int32 is chunk_offsets: + co_int32 = chunk_offsets.clone() + cu_ct = from_dlpack(cu_int32, assumed_align=4).mark_layout_dynamic() + cu_ct.element_type = cutlass.Int32 + co_ct = from_dlpack(co_int32, assumed_align=4).mark_layout_dynamic() + co_ct.element_type = cutlass.Int32 + # Hold refs to the int32 tensors so they don't get GC'd and the + # underlying storage remain valid as long as cu_ct / co_ct live. + entry = (cu_int32, co_int32, cu_ct, co_ct) + _k4_varlen_cu_co_cache[key] = entry + _prune_on_gc(_k4_varlen_cu_co_cache, key, cu_seqlens, chunk_offsets) + return entry[2], entry[3] + + +# Cache the (cu_for_k4, chunk_offsets_for_k4) pair keyed by id(cu_seqlens). +# Both tensors are computed on-GPU with no host sync — replaces the previous +# `cu_seqlens.cpu().tolist() + Python loop + torch.tensor` chain that forced +# a 50-200us GPU->CPU stall on the K4 prep path every varlen call. +_varlen_k4_input_cache = {} + + +def _get_varlen_k4_inputs(cu_seqlens, BT): + key = id(cu_seqlens) + entry = _varlen_k4_input_cache.get(key) + if entry is None: + # Copy (never alias) the keyed object — see _get_k4_varlen_cu_co. + cu_int32 = cu_seqlens.to(torch.int32).contiguous() + if cu_int32 is cu_seqlens: + cu_int32 = cu_seqlens.clone() + seq_lens = cu_int32[1:] - cu_int32[:-1] + chunk_counts = (seq_lens + (BT - 1)) // BT + zero = torch.zeros(1, dtype=torch.int32, device=cu_int32.device) + co_int32 = torch.cat([zero, torch.cumsum(chunk_counts, dim=0).to(torch.int32)]) + co_int32 = co_int32.contiguous() + _varlen_k4_input_cache[key] = (cu_int32, co_int32) + _prune_on_gc(_varlen_k4_input_cache, key, cu_seqlens) + return _varlen_k4_input_cache[key] + + +# ========== BF16 akk_inv compilation cache ========== +_akk_inv_cache = {} +# ========== K4 persistent (varlen via cu_seqlens) compilation cache ========== +_k4p_cache = {} +_k4p_tm_ws = {} + +# Side-stream cache for the initial-state copy overlapping K123. +_side_streams = {} + + +def _get_side_stream(dev): + idx = dev.index or 0 + if idx not in _side_streams: + _side_streams[idx] = torch.cuda.Stream(device=dev) + return _side_streams[idx] + + +# Buffer cache: avoid re-allocating ~67us of intermediate tensors per call. +# Also caches cute.Tensor wrappers (saves ~7us each call from from_dlpack). +# LRU-bounded: entries are keyed by (B, T, ...) shapes, and the runtime +# executor calls with a different token count per prefill batch — an +# unbounded dict would pin ~T*150KB of scratch per distinct shape forever. +_buf_cache = {} +_BUF_CACHE_MAX_ENTRIES = 8 + +# 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. +_padded_input_cache = {} + +# 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 +# pre-set to -1e3 once at cache init; subsequent calls only overwrite the +# valid prefix [0, real_T). +_g_sentinel_cache = {} + + +def _get_g_sentinel_buffer(B, T_padded, H, K, dtype_g, device, real_T): + key = (B, T_padded, H, K, dtype_g, device.index if device.index is not None else 0, real_T) + e = _g_sentinel_cache.get(key) + if e is None: + e = torch.zeros(B, T_padded, H, K, dtype=dtype_g, device=device) + if real_T < T_padded: + e[:, real_T:] = -1000.0 + _g_sentinel_cache[key] = e + return e + + +def _get_padded_input_buffers(B, T_padded, H, K, dtype_qkv, dtype_g, dtype_beta, device, real_T): + key = ( + B, + T_padded, + H, + K, + dtype_qkv, + dtype_g, + dtype_beta, + device.index if device.index is not None else 0, + real_T, + ) + e = _padded_input_cache.get(key) + if e is None: + q_pad = torch.zeros(B, T_padded, H, K, dtype=dtype_qkv, device=device) + k_pad = torch.zeros(B, T_padded, H, K, dtype=dtype_qkv, device=device) + v_pad = torch.zeros(B, T_padded, H, K, dtype=dtype_qkv, device=device) + beta_pad = torch.zeros(B, T_padded, H, dtype=dtype_beta, device=device) + # g: zero in [0, real_T), sentinel -1e3 in [real_T, T_padded). Caller's + # data overwrites the prefix each call; the sentinel tail never moves. + g_pad = torch.zeros(B, T_padded, H, K, dtype=dtype_g, device=device) + if real_T < T_padded: + g_pad[:, real_T:] = -1000.0 + e = (q_pad, k_pad, v_pad, g_pad, beta_pad) + _padded_input_cache[key] = e + return e + + +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: + bf16 = cutlass.BFloat16 + fp32 = cutlass.Float32 + # Varlen chunk-tile kernels (akk_inv's Stage-0 cp.async A-tile load, + # K4's A-tile loads) transfer the full BT-row tile of every chunk + # and neutralize invalid rows only after the access; when the + # batch's FINAL chunk is partial they read up to BT-1 rows past the + # logical T. These are driver-owned scratch buffers, so honor the + # kernel's boundary-at-the-data contract by allocating one chunk of + # zeroed slack past T (the eqlen path gets the same guarantee from + # its 256-multiple input padding). The slack rows are never + # consumed — OOB rows are zeroed in SMEM or masked at stores. + # Input tensors (beta) get the opposite treatment: the K123 kernel + # bounds-checks those loads, since the driver doesn't own them. + if varlen: + assert B == 1, f"varlen expects packed B=1 input, got B={B}" + T_alloc = T + BT if varlen else T + + def _with_slack(ctor, *shape, dtype): + full = ctor(*shape, device=dev, dtype=dtype) + return full[:, :T] + + k_scaled = _with_slack(torch.empty, B, T_alloc, H, K_dim, dtype=dtype_k) # raw, no beta + kg = _with_slack(torch.empty, B, T_alloc, H, K_dim, dtype=dtype_k) + q_scaled = _with_slack(torch.empty, B, T_alloc, H, K_dim, dtype=dtype_k) + gk_last_exp = torch.empty(B, NT, H, K_dim, device=dev, dtype=torch.float32) + A_qk = _with_slack(torch.zeros, B, T_alloc, H, BT, dtype=dtype_k) + A_kk = _with_slack(torch.zeros, B, T_alloc, H, BT, dtype=dtype_k) + O_flat = _with_slack(torch.empty, B, T_alloc, H, V_dim, dtype=dtype_k) + # K4 reads initial state from S_out (caller copies it in) and writes + # the final state back into the same buffer — no separate s_4d, no + # extra D2D memcpy at the end of the K4 launcher. Layout matches + # caller's [N_seqs, H, K, V] contig fp32 convention. + S_out = torch.empty(N_seqs, H, K_dim, V_dim, device=dev, dtype=torch.float32) + cu_eqlen = torch.arange(0, (B + 1) * T, T, dtype=torch.int32, device=dev) + # co_eqlen is consumed only on the eqlen path (varlen K4 derives its + # chunk offsets from cu_seqlens). Eqlen inputs are pre-padded to a + # 256-multiple upstream, so T // BT >= 4 there; varlen calls can + # carry T < BT (short-prompt batches), where arange(step=T // BT) + # would raise "step must be nonzero" building this dead buffer. + # Clamp the step so the buffer stays constructible. + nt_eq = max(T // BT, 1) + co_eqlen = torch.arange(0, (B + 1) * nt_eq, nt_eq, dtype=torch.int32, device=dev) + + T_total = B * T + A_kk_flat = A_kk.reshape(T_total, H, BT) + A_qk_flat = A_qk.reshape(T_total, H, BT) + KS_flat = k_scaled.reshape(T_total, H, K_dim) + QS_flat = q_scaled.reshape(T_total, H, K_dim) + KG_flat = kg.reshape(T_total, H, K_dim) + O_token = O_flat.reshape(T_total, H, V_dim) + gk_flat = gk_last_exp.reshape(-1, H, K_dim) + + def _wrap(t, etype): + r = from_dlpack(t, assumed_align=16).mark_layout_dynamic() + r.element_type = etype + return r + + a_ct = _wrap(A_kk_flat, bf16) + aqc_ct = _wrap(A_qk_flat, bf16) + ks_ct = _wrap(KS_flat, bf16) # raw k_scaled (beta absorbed in akk_inv) + qs_ct = _wrap(QS_flat, bf16) + kg_ct = _wrap(KG_flat, bf16) + o_ct = _wrap(O_token, bf16) + gk_ct = _wrap(gk_flat, fp32) + cu_eqlen_ct = from_dlpack(cu_eqlen, assumed_align=4).mark_layout_dynamic() + cu_eqlen_ct.element_type = cutlass.Int32 + co_eqlen_ct = from_dlpack(co_eqlen, assumed_align=4).mark_layout_dynamic() + co_eqlen_ct.element_type = cutlass.Int32 + + # akk_inv views: bf16 storage reinterpreted as fp32 (packed 2x bf16 -> 1x fp32). + # Layout-dynamic so the single compiled akk_inv (shape-independent + # key) accepts them regardless of T; the host rebuilds its own + # runtime-shaped views from the iterator anyway. + akk_in_view = from_dlpack(A_kk, assumed_align=16).mark_layout_dynamic() + akk_in_view.element_type = fp32 + akk_out_view = from_dlpack(A_kk, assumed_align=16).mark_layout_dynamic() + akk_out_view.element_type = fp32 + + # K4 state wrapper points directly at S_out — K4 reads/writes in + # place. Only N_seqs (the outer compact mode) varies with the runtime + # context batch. Keep the inner H/K/V shape and strides static so the + # DSL can prove each (batch, head) state tile remains 16-byte aligned. + s_ct = from_dlpack(S_out, assumed_align=16).mark_compact_shape_dynamic( + mode=0, + stride_order=(0, 1, 2, 3), + ) + s_ct.element_type = fp32 + + # Side stream is process-long, safe to cache. The MAIN stream must + # NOT be cached here: the executor creates the buffers under one + # stream (warmup) and calls under another (execution_stream), so it + # is re-queried per call in _chunk_kda_fwd. + side_stream_cached = _get_side_stream(dev) + + cute_wrappers = dict( + a_ct=a_ct, + aqc_ct=aqc_ct, + ks_ct=ks_ct, + qs_ct=qs_ct, + kg_ct=kg_ct, + o_ct=o_ct, + gk_ct=gk_ct, + cu_eqlen_ct=cu_eqlen_ct, + co_eqlen_ct=co_eqlen_ct, + akk_in_view=akk_in_view, + akk_out_view=akk_out_view, + s_ct=s_ct, + side_stream=side_stream_cached, + # Filled lazily on first launch — saves cache_key tuple build + + # outer dict lookup on subsequent calls. + _k123_fns={}, + _akk_inv_fn=None, + _k4_fn=None, + ) + + while len(_buf_cache) >= _BUF_CACHE_MAX_ENTRIES: + _buf_cache.pop(next(iter(_buf_cache))) + _buf_cache[key] = ( + k_scaled, + kg, + q_scaled, + gk_last_exp, + A_qk, + A_kk, + O_flat, + S_out, + cu_eqlen, + co_eqlen, + cute_wrappers, + ) + else: + # LRU refresh so hot shapes survive eviction. + _buf_cache[key] = _buf_cache.pop(key) + return _buf_cache[key] + + +def _launch_k4_persistent( + cute_wrappers, + v_beta, + S_in, + S_out, + cu_seqlens, + chunk_offsets, + cu_eqlen_passed=False, + num_sm=148, + H=None, + V_dim=None, + use_fast_sync=False, +): + """Launch persistent K4 with cached CuTe wrappers. + + 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 + fires) — a per-call GPU memory leak in the executor runtime. The + per-call cost is re-wrapping a handful of tensors (~10us each). + """ + bf16 = cutlass.BFloat16 + N_seqs = cu_seqlens.shape[0] - 1 + dev = v_beta.device + dev_idx = dev.index or 0 + + s_ct = cute_wrappers["s_ct"] + a_ct = cute_wrappers["a_ct"] + b_ct = cute_wrappers["ks_ct"] # raw k_scaled (beta absorbed in akk_inv) + q_ct = cute_wrappers["qs_ct"] + aqc_ct = cute_wrappers["aqc_ct"] + kg_ct = cute_wrappers["kg_ct"] + o_ct = cute_wrappers["o_ct"] + gk_ct = cute_wrappers["gk_ct"] + + # v is a per-call activation — wrap fresh every call (never cache; see + # _ct_cached docstring). + v_view = v_beta.reshape(-1, H, V_dim) if v_beta.dim() == 4 else v_beta + v_ct = from_dlpack(v_view, assumed_align=16).mark_layout_dynamic() + v_ct.element_type = bf16 + + if cu_eqlen_passed: + cu_ct = cute_wrappers["cu_eqlen_ct"] + co_ct = cute_wrappers["co_eqlen_ct"] + else: + cu_ct, co_ct = _get_k4_varlen_cu_co(cu_seqlens, chunk_offsets) + + # Allocate for the maximum persistent grid once. The scheduler applies + # min(N_seqs * H, num_sm) at runtime, so pre-minimizing here would turn + # every distinct request count into another compiled host function. + tm_key = (dev_idx, num_sm) + if tm_key not in _k4p_tm_ws: + tm_ws_t = torch.zeros( + num_sm * _K4P_NTM * _K4P_BTM, + dtype=torch.uint8, + device=dev, + ) + tm_ct = from_dlpack(tm_ws_t, assumed_align=16) + tm_ct.element_type = cutlass.Uint8 + _k4p_tm_ws[tm_key] = (tm_ws_t, tm_ct) + else: + tm_ws_t, tm_ct = _k4p_tm_ws[tm_key] + + # Launch on torch's CURRENT stream — see _current_cu_stream. + stream = _current_cu_stream(dev) + + k4_fn = cute_wrappers.get("_k4_fn") + if k4_fn is None: + # H and the state K/V dims MUST be in the key: s_ct bakes the inner + # [H, K, V] shape and all strides at compile time + # (mark_compact_shape_dynamic(mode=0) in _get_buffers), so the + # compiled function is head-count-specific. Reusing an H=96-compiled + # kernel for a smaller H misaddresses every (seq, head) state tile + # (baked mode-0 stride H*K*V) — silently corrupting final_state and, + # through the recurrence, O. Observed as the small-H parity + # regression whenever another head count compiled first in the same + # process (isolated processes were clean). N_seqs and the token/chunk + # counts stay runtime — only layout-baked dims belong here. + cache_key = (dev_idx, num_sm, H, S_out.shape[-2], V_dim) + k4_fn = _k4p_cache.get(cache_key) + if k4_fn is None: + host_fn = _k4p_make_host(num_sm=num_sm) + k4_fn = cute.compile( + host_fn, + a_ct, + b_ct, + v_ct, + q_ct, + aqc_ct, + kg_ct, + o_ct, + gk_ct, + s_ct, + cu_ct, + co_ct, + tm_ct, + N_seqs, + stream, + ) + _k4p_cache[cache_key] = k4_fn + cute_wrappers["_k4_fn"] = k4_fn + + args = ( + a_ct, + b_ct, + v_ct, + q_ct, + aqc_ct, + kg_ct, + o_ct, + gk_ct, + s_ct, + cu_ct, + co_ct, + tm_ct, + N_seqs, + stream, + ) + k4_fn(*args) + + +def _launch_fused_k123_inv( + q, + k, + g, + A_log, + beta, + scale, + k_scaled, + kg, + q_scaled, + gk_last_exp, + A_qk, + A_kk_inv, + cu_seqlens, + chunk_indices, + is_varlen, + NT, + dt_bias=None, + safe_gate=False, + lower_bound=None, + akk_in_view=None, + akk_out_view=None, + cute_wrappers=None, + varlen_pure_override=False, +): + """Persistent K1+K2 (writes A_kk in I+L format with diag=1) chained with + BF16 akk_inv (in-place inversion). Final A_kk_inv = (I+L)^-1.""" + + # No fast-launch (args-tuple) cache: it would pin the per-call q/k/g/beta + # activations via their cute wrappers (the wrapper holds the storage, so + # the keyed objects never die and weakref pruning never fires) — a + # per-call GPU memory leak in the executor runtime. Wrapper gathering + # below costs ~10us per tensor. + B, T, H, K = q.shape + BT = 64 + dev = q.device.index or 0 + T_padded = T if is_varlen else None + has_bias = dt_bias is not None + + # Auto-detect VARLEN_PURE eligibility, cached by id(cu_seqlens). + # First call with a given cu_seqlens object pays one GPU->CPU sync; later + # calls with the same tensor object are a dict lookup (~100 ns). + varlen_pure = False + if is_varlen and cu_seqlens is not None: + if varlen_pure_override: + # Caller (Phase 2.1 single-seq path) sentinel-padded the data for + # this call; the cached per-object verdict stays untouched. + varlen_pure = True + else: + _vp_key = id(cu_seqlens) + if _vp_key not in _varlen_pure_cache: + cu_cpu = cu_seqlens.cpu().tolist() + seq_lens = [cu_cpu[i + 1] - cu_cpu[i] for i in range(len(cu_cpu) - 1)] + _varlen_pure_cache[_vp_key] = all((sl % BT) == 0 for sl in seq_lens) + _prune_on_gc(_varlen_pure_cache, _vp_key, cu_seqlens) + varlen_pure = _varlen_pure_cache[_vp_key] + # B/NT/T are runtime args of the compiled host_fn (see make_host_function), + # avoiding shape-dependent grid and view specialization. Varlen remains + # fully shape-independent; eqlen retains T only as a cache discriminator + # for the raw tensor batch strides described below. + # + # cu/ci DTYPE must be part of the key: the kernel reads + # mCuSeqlens/mChunkIndices with the element type baked at compile time + # (int64 elements are addressed with stride 8, int32 with stride 4). + # Reusing an int64-compiled kernel on int32 tensors (or vice versa) + # misaddresses every cu/ci element — garbage seq ids/chunk starts -> + # OOB reads (cudaErrorIllegalAddress) or silent corruption. Observed as + # a crash when a dump-replay batch (int64 cu/ci) preceded a synthetic + # batch (int32 cu/ci) in one process. + cu_ci_dtypes = (cu_seqlens.dtype, chunk_indices.dtype) if is_varlen else None + # Eqlen kernels directly index mBeta/mAqk/mAkk using the input wrappers' + # compile-time batch stride, which depends on T. Varlen fixes B=1, so that + # batch stride is never used and all T shapes can keep sharing one compile. + eqlen_t_key = T if not is_varlen else -1 + cache_key = ( + H, + is_varlen, + dev, + has_bias, + safe_gate, + varlen_pure, + cu_ci_dtypes, + eqlen_t_key, + ) + + # Inputs are guaranteed contiguous by upstream linear projections. + # A_log is fp32 model param; .float() is no-op when dtype already matches. + # q/k/g/beta/cu/ci are per-call activations: plain _ct, never cached + # (see _ct_cached docstring). The _get_buffers scratch below is + # module-persistent, so caching its wrappers is safe and worthwhile. + q_ct = _ct(q, cutlass.BFloat16) + k_ct = _ct(k, cutlass.BFloat16) + g_ct = _ct(g, cutlass.BFloat16) + alog_ct = _ct(A_log if A_log.dtype == torch.float32 else A_log.float(), cutlass.Float32) + beta_ct = _ct(beta, cutlass.BFloat16) + + ks_ct = _ct_cached(k_scaled, cutlass.BFloat16) + kg_ct = _ct_cached(kg, cutlass.BFloat16) + qs_ct = _ct_cached(q_scaled, cutlass.BFloat16) + gk_ct = _ct_cached(gk_last_exp, cutlass.Float32) + aqk_ct = _ct_cached(A_qk, cutlass.BFloat16) + akk_ct = _ct_cached(A_kk_inv, cutlass.BFloat16) + + if is_varlen: + cu_ct = _ct(cu_seqlens, _cute_int_type(cu_seqlens.dtype)) + ci_ct = _ct(chunk_indices, _cute_int_type(chunk_indices.dtype)) + else: + cu_ct, ci_ct = _get_eqlen_dummies(q.device, torch.int64) + + if dt_bias is not None: + bias_ct = _get_dt_bias_ct(dt_bias, H, K) + else: + bias_ct = _get_empty_bias_ct(q.device) + lb_val = float(lower_bound) if lower_bound is not None else 0.0 + + # Launch on torch's CURRENT stream — see _current_cu_stream. + stream = _current_cu_stream(q.device) + + ct_args = ( + q_ct, + k_ct, + g_ct, + alog_ct, + beta_ct, + scale, + ks_ct, + kg_ct, + qs_ct, + gk_ct, + aqk_ct, + akk_ct, + cu_ct, + ci_ct, + bias_ct, + lb_val, + # Runtime shape scalars (rt_nt / rt_b / rt_t_total in host_fn). + NT, + B, + B * (T_padded if is_varlen else NT * BT), + stream, + ) + + if cache_key not in _fused_k123_cache: + host_fn = _fused_make_host( + B, + NT, + H, + is_varlen=is_varlen, + T_padded=T_padded, + has_bias=has_bias, + use_safe_gate=safe_gate, + varlen_pure=varlen_pure, + ) + _fused_k123_cache[cache_key] = cute.compile(host_fn, *ct_args) + k123_fn = _fused_k123_cache[cache_key] + k123_fn(*ct_args) + + # ===== Chained BF16 akk_inv (in-place: A_kk_inv = (I+L)^-1) ===== + # Views are layout-dynamic so one compiled akk_inv serves every batch + # shape (the host builds its own runtime-shaped views from the iterator; + # the wrapper layout only matters for the call signature). + if akk_in_view is None: + akk_in_view = from_dlpack(A_kk_inv, assumed_align=16).mark_layout_dynamic() + akk_in_view.element_type = cutlass.Float32 + akk_out_view = from_dlpack(A_kk_inv, assumed_align=16).mark_layout_dynamic() + akk_out_view.element_type = cutlass.Float32 + + if is_varlen: + # Reuse the cached cute wrappers from K123 (same tensor objects). + akk_cu_ct = cu_ct + akk_ci_ct = ci_ct + is_varlen_int = 1 + T_val = T + else: + akk_cu_ct, akk_ci_ct = cu_ct, ci_ct + is_varlen_int = 0 + T_val = NT * BT + + # B/NT/T_val remain runtime args of akk_inv_host. Eqlen T is still part of + # the cache key because the raw mBeta wrapper bakes its batch stride. + # cu/ci dtype is a specializer for the same reason as in the K123 key above + # (element type baked into the compiled reader). + akk_cache_key = (H, is_varlen, dev, cu_ci_dtypes, eqlen_t_key) + if akk_cache_key not in _akk_inv_cache: + _akk_inv_cache[akk_cache_key] = cute.compile( + _akk_inv_host, + akk_in_view, + akk_out_view, + beta_ct, + B, + NT, + H, + akk_cu_ct, + akk_ci_ct, + is_varlen_int, + T_val, + 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_fn(*akk_args) + + + +# ========== Fused K1234 compilation cache ========== +_fused_k1234_cache = {} +_BT = 64 + + +def _launch_fused_k1234( + q, + k, + v, + g, + A_log, + beta, + scale, + initial_state, + output_final_state, + dt_bias=None, + safe_gate=False, + lower_bound=None, +): + """Launch single fused K1+K2+K3+K4 SMEM-direct kernel (eqlen only).""" + B, T, H, K = q.shape + V_dim = v.shape[-1] + device = q.device + dev = device.index or 0 + NC = T // _BT + has_bias = dt_bias is not None + + # Dynamic shapes: cache only on mode flags, not B/NC/H + cache_key = (dev, has_bias, safe_gate) + + # Buffers + BH = B * H + N_seqs = B + S_fp32 = torch.empty(BH, K, V_dim, dtype=torch.float32, device=device) + if initial_state is None: + S_fp32.zero_() + else: + S_fp32.copy_(initial_state.reshape(BH, K, V_dim)) + O_out = torch.empty(B, T, H, V_dim, dtype=q.dtype, device=device) + clocks = torch.empty(2, BH, dtype=torch.int64, device=device) + + bf16 = cutlass.BFloat16 + fp32 = cutlass.Float32 + if has_bias: + bias_ct = _ct(dt_bias.float().contiguous().view(H, K), fp32) + else: + bias_ct = _ct(torch.empty(1, 1, dtype=torch.float32, device=device), fp32) + lb_val = float(lower_bound) if lower_bound is not None else 0.0 + + ct_args = ( + _ct(q.contiguous(), bf16), + _ct(k.contiguous(), bf16), + _ct(g.contiguous(), bf16), + _ct(A_log.float().contiguous(), fp32), + _ct(beta.contiguous(), bf16), + float(scale), + _ct(v.contiguous(), bf16), + _ct(O_out, bf16), + _ct(S_fp32, fp32), + bias_ct, + lb_val, + _ct(clocks, cutlass.Int64), + NC, + H, + B, + # Launch on torch's CURRENT stream — see _current_cu_stream. + _current_cu_stream(device), + ) + + if cache_key not in _fused_k1234_cache: + host_fn = _fused_k1234_make_host(has_bias=has_bias, use_safe_gate=safe_gate) + _fused_k1234_cache[cache_key] = cute.compile(host_fn, *ct_args) + _fused_k1234_cache[cache_key](*ct_args) + + S_out = S_fp32.reshape(N_seqs, H, K, V_dim) + final_state = S_out if output_final_state else None + return O_out, final_state + + +def _chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: float | None = None, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + disable_recompute: bool = False, + return_intermediate_states: bool = False, + cp_context=None, + use_fused_k1234: bool = False, +): + """KDA forward — optimized, FLA-compatible interface.""" + if safe_gate and lower_bound is None: + lower_bound = -5.0 + + is_varlen = cu_seqlens is not None + + if q.shape[1] == 0: + # Zero-token call: the runtime can emit a context batch whose token + # payload is empty (observed under the overlap scheduler + logprobs + # flows; the FLA fallback tolerates it). No tokens means no output + # rows and an unchanged recurrent state — return before any kernel + # or buffer setup (whose arange(step=T) would raise "step must be + # nonzero"). + B, _, H, K = q.shape + V_dim = v.shape[-1] + n_seqs = B if cu_seqlens is None else cu_seqlens.shape[0] - 1 + o = v.new_empty(B, 0, H, V_dim) + if not output_final_state: + final_state = None + elif initial_state is not None: + # clone: the caller index_copy_s final_state back into the state + # 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) + + # ===== Fused K1234 path (eqlen only, single kernel launch) ===== + if use_fused_k1234 and not is_varlen: + o, final_state = _launch_fused_k1234( + q, + k, + v, + g, + A_log, + beta, + scale, + initial_state, + output_final_state, + dt_bias=dt_bias, + safe_gate=safe_gate, + lower_bound=lower_bound, + ) + return (o, final_state, None, None, None, None, None, None, None, None, None, initial_state) + + B, T, H, K = q.shape + V_dim = v.shape[-1] + device = q.device + BT = 64 + + # Phase 1: handle eqlen with T % 64 != 0 entirely on the eqlen path — + # never borrow varlen's mask code. Pad inputs to a CHUNKS_PER_BLOCK*BT + # = 256 multiple so the persistent scheduler's `cgs_per_head = NT // 4` + # divides cleanly and every chunk has 64 valid rows of (zero-padded / + # sentinel-padded) data. K123 eqlen kernel runs unchanged — it doesn't + # know partial chunks exist, the boundary is handled at the data + # boundary (K4's bounded-TMA principle, just expressed via host pad + # since K123 stores via autovec_copy not TMA). + # + # Varlen with non-aligned seq lengths is a separate problem (Phase 2): + # multi-seq varlen can't be host-padded without repacking memory. + real_T = T + needs_eqlen_pad = (not is_varlen) and (T % BT != 0) + if needs_eqlen_pad: + if B != 1: + raise NotImplementedError( + f"eqlen with B>1 and T % {BT} != 0 not supported (got B={B}, T={T})." + ) + CPB_BT = 4 * BT # CHUNKS_PER_BLOCK * BT, the cgs_per_head divisibility unit + T_padded = ((T + CPB_BT - 1) // CPB_BT) * CPB_BT + # Pre-allocated padded scratch buffers (per (B,T_padded,H,K,dtype) cache + # key). torch.cat would reallocate + copy the full 200MB q tensor every + # call — caching the destination buffer drops that to a single slice + # copy of the valid prefix (caller already lives in our buffer for + # subsequent calls reusing the same id, but we re-copy unconditionally + # since the caller may have updated the data in-place). + q_pad, k_pad, v_pad, g_pad, beta_pad = _get_padded_input_buffers( + B, T_padded, H, K, q.dtype, g.dtype, beta.dtype, q.device, real_T + ) + # q/k/v/beta zero-padded → K1/K2 MMAs naturally produce 0 for OOB rows. + # Tail [real_T:] of q_pad/k_pad/v_pad/beta_pad is pre-zeroed at cache + # init and never written, so we only copy the valid prefix. + q_pad[:, :real_T].copy_(q) + k_pad[:, :real_T].copy_(k) + v_pad[:, :real_T].copy_(v) + beta_pad[:, :real_T].copy_(beta) + # g uses a -1e3 sentinel so the gate activation saturates to 0 for OOB + # rows (both safe_gate sigmoid and softplus paths). Plain g=0 gives + # nonzero activation that would corrupt the cumsum past seq end. The + # tail is set to -1e3 once at cache init; we only copy the valid prefix. + g_pad[:, :real_T].copy_(g) + q, k, v, g, beta = q_pad, k_pad, v_pad, g_pad, beta_pad + T = T_padded # downstream buffer alloc + kernel layout use T_padded + + # Phase 2.1: varlen with a SINGLE non-aligned sequence — caller already + # zero-padded q/k/v/beta to a 64-multiple (FLA convention), but g's tail + # is also zero, which causes the gate activation to be non-zero past seq + # end and corrupts the cumsum / GkLast. We sentinel-pad g (cheap: ~5MB + # copy) and force VARLEN_PURE=1 so all 4 mask sites compile-elide. Same + # K4 "boundary at the data" principle as the eqlen path. + # + # Multi-seq varlen is NOT handled here — its OOB regions overlap with + # adjacent seqs' data, so sentinel-pad on g would corrupt the next seq. + # Multi-seq optimization needs per-seq dynamic tensormap (Phase 2.2). + needs_varlen_single_pad = False + if is_varlen and cu_seqlens is not None and cu_seqlens.shape[0] == 2 and B == 1: + _vl_key = id(cu_seqlens) + if _vl_key not in _varlen_pure_cache: + cu_cpu = cu_seqlens.cpu().tolist() + sl = cu_cpu[1] - cu_cpu[0] + _varlen_pure_cache[_vl_key] = sl % BT == 0 + _varlen_single_seqlen_cache[_vl_key] = sl + _prune_on_gc(_varlen_pure_cache, _vl_key, cu_seqlens) + _prune_on_gc(_varlen_single_seqlen_cache, _vl_key, cu_seqlens) + if not _varlen_pure_cache[_vl_key]: + real_T = _varlen_single_seqlen_cache[_vl_key] + needs_varlen_single_pad = True + varlen_pure_override = False + if needs_varlen_single_pad: + # q/k/v/beta already zero-padded by caller (FLA convention). Re-build g + # with -1000 sentinel in the tail so VARLEN_PURE=1 path is correct. + # Cache the resulting g buffer so repeated calls with same input ids + # don't re-allocate. + cur_T = q.shape[1] # caller's padded T + 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") + 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 + # Force VARLEN_PURE=1 for THIS call only (the data is sentinel-padded + # now). Never write the override into _varlen_pure_cache: the cached + # verdict must keep saying "not pure" so a later call with the SAME + # cu_seqlens object re-runs this sentinel pad — poisoning the cache + # made the second call skip the pad while still compiling the + # mask-free variant (silent tail corruption). + varlen_pure_override = True + + # Phase 2.2 (multi-seq via host repack) was attempted but isn't net positive: + # the scatter/gather memcpy cost (~800us GPU bandwidth per call) exceeds + # the kernel mask-elision savings (~250us). Keeping multi-seq non-pure on + # the original masked path. The right fix is kernel-level dynamic + # tensormap (K4-style per-tile bounded TMA) but that requires a major + # K123 kernel refactor — left as future work. + # multiseq_info = None + + if is_varlen: + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = len(chunk_indices) + if NT < 4: + # The persistent K123 scheduler launches NT // 4 cooperative + # groups per head; fewer than 4 total chunks produces a + # zero-size grid (DSLCudaRuntimeError at launch). Callers must + # route such batches to the FLA path — see + # KDAKernelDispatch.prefill_chunk_kda. + raise ValueError( + f"kda_prefill requires >= 4 total varlen chunks (got {NT}); " + "route small varlen batches to the FLA fallback") + N_seqs = len(cu_seqlens) - 1 + else: + NT = T // BT + N_seqs = B + + # ===== Cached buffers + cute wrappers (avoid alloc + from_dlpack overhead per call) ===== + ( + k_scaled, + kg, + q_scaled, + gk_last_exp, + A_qk, + A_kk, + O_flat, + S_out, + cu_eqlen, + co_eqlen, + cute_wrappers, + ) = _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 + # side stream BEFORE K123 launches, the D2D memcpy overlaps with K123's + # compute. Especially big win for high-N_seqs varlen where state is huge + # (192MB for N=32, ~64us memcpy) — would otherwise serialize before K4. + if initial_state is None: + S_in = torch.zeros(N_seqs, H, K, V_dim, dtype=torch.float32, device=device) + else: + S_in = initial_state + # Current stream per call — never the buffer-creation-time stream (the + # executor creates buffers under warmup's stream and calls under + # execution_stream; syncing against a stale stream un-orders the copy). + main_stream = torch.cuda.current_stream(device) + side_stream = cute_wrappers["side_stream"] + needs_copy = S_in.data_ptr() != S_out.data_ptr() + if needs_copy: + side_stream.wait_stream(main_stream) + with torch.cuda.stream(side_stream): + S_out.copy_(S_in) + + # Beta is fused entirely in akk_inv kernel epilogue (post-inv column-scale). + # No host v*beta and no K1 k_scaled*beta any more. + _launch_fused_k123_inv( + q, + k, + g, + A_log, + beta, + scale, + k_scaled, + kg, + q_scaled, + gk_last_exp, + A_qk, + A_kk, + cu_seqlens, + chunk_indices, + is_varlen, + NT, + dt_bias=dt_bias, + safe_gate=safe_gate, + lower_bound=lower_bound, + akk_in_view=cute_wrappers["akk_in_view"], + akk_out_view=cute_wrappers["akk_out_view"], + cute_wrappers=cute_wrappers, + varlen_pure_override=varlen_pure_override, + ) + + # ===== K4: persistent kernel (eqlen + varlen via cu_seqlens) ===== + if is_varlen: + # GPU-side cumsum + cache by id(cu_seqlens). No host sync. + cu_for_k4, chunk_offsets_for_k4 = _get_varlen_k4_inputs(cu_seqlens, BT) + else: + cu_for_k4 = cu_eqlen + chunk_offsets_for_k4 = co_eqlen + + # Wait for the side-stream state copy to finish before K4 reads S_out. + if needs_copy: + main_stream.wait_stream(side_stream) + + _launch_k4_persistent( + cute_wrappers, + v, + S_in, + S_out, + cu_for_k4, + chunk_offsets_for_k4, + cu_eqlen_passed=(not is_varlen), + H=H, + V_dim=V_dim, + use_fast_sync=(not is_varlen), + ) + + o = O_flat + A_qk_out = A_qk + A_kk_out = A_kk + if needs_eqlen_pad: + # Caller called with original T = real_T; their downstream code expects + # outputs at that shape. Slice the padded scratch tail back off. + o = o[:, :real_T] + A_qk_out = A_qk_out[:, :real_T] + A_kk_out = A_kk_out[:, :real_T] + # multiseq_info is always None here (Phase 2.2 disabled — see above) + final_state = S_out if output_final_state else None + + return ( + o, + final_state, + None, + A_qk_out, + A_kk_out, + None, + None, + None, + None, + None, + None, + initial_state, + ) + + +class KdaPrefillRunner: + """Runs the source-integrated Kimi K3 KDA prefill pipeline.""" + + @staticmethod + def forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: Optional[torch.Tensor], + output_final_state: bool, + cu_seqlens: Optional[torch.Tensor] = None, + chunk_indices: Optional[torch.Tensor] = None, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: Optional[float] = None, + use_gate_in_kernel: bool = False, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + disable_recompute: bool = False, + return_intermediate_states: bool = False, + use_fused_k1234: bool = False, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Run KDA prefill and return output plus final recurrent state.""" + if not IS_CUTLASS_DSL_AVAILABLE: + 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}") + + result = _chunk_kda_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + lower_bound=lower_bound, + use_gate_in_kernel=use_gate_in_kernel, + A_log=A_log, + dt_bias=dt_bias, + disable_recompute=disable_recompute, + return_intermediate_states=return_intermediate_states, + use_fused_k1234=use_fused_k1234, + ) + return result[0], result[1] + + +if IS_CUTLASS_DSL_AVAILABLE: + + @torch.library.custom_op("trtllm::kda_prefill", mutates_args=(), device_types="cuda") + def kda_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: Optional[torch.Tensor], + output_final_state: bool, + cu_seqlens: Optional[torch.Tensor] = None, + chunk_indices: Optional[torch.Tensor] = None, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: Optional[float] = None, + use_gate_in_kernel: bool = False, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + disable_recompute: bool = False, + return_intermediate_states: bool = False, + use_fused_k1234: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Run Kimi K3 KDA chunked prefill on Blackwell GPUs. + + Returns ``(output, final_state)``. ``final_state`` is an empty + tensor when ``output_final_state`` is False — + ``torch.library.custom_op`` cannot infer a schema for + ``Optional[Tensor]`` returns. + """ + output, final_state = KdaPrefillRunner.forward( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + scale=scale, + dt_bias=dt_bias, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + lower_bound=lower_bound, + use_gate_in_kernel=use_gate_in_kernel, + disable_recompute=disable_recompute, + return_intermediate_states=return_intermediate_states, + use_fused_k1234=use_fused_k1234, + ) + if final_state is None: + final_state = q.new_empty(0, dtype=torch.float32) + return output, final_state + + @kda_prefill.register_fake + def _( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: Optional[torch.Tensor], + output_final_state: bool, + cu_seqlens: Optional[torch.Tensor] = None, + chunk_indices: Optional[torch.Tensor] = None, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: Optional[float] = None, + use_gate_in_kernel: bool = False, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + disable_recompute: bool = False, + return_intermediate_states: bool = False, + use_fused_k1234: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + del k, g, beta, A_log, scale, dt_bias, initial_state + del chunk_indices, chunk_size, safe_gate, lower_bound + del use_gate_in_kernel, disable_recompute, return_intermediate_states + del use_fused_k1234 + num_sequences = q.shape[0] if cu_seqlens is None else cu_seqlens.shape[0] - 1 + output = v.new_empty(v.shape) + if output_final_state: + final_state = q.new_empty( + (num_sequences, q.shape[2], q.shape[3], v.shape[3]), dtype=torch.float32 + ) + else: + final_state = q.new_empty(0, dtype=torch.float32) + return output, final_state 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 new file mode 100644 index 000000000000..21e2b247a898 --- /dev/null +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py @@ -0,0 +1,738 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CuTe DSL custom op for Kimi K3 KDA multi-token speculative verify. + +Wraps the source-integrated ``kda_decode_mtp_kernel`` (see +``cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py``) as the +``trtllm::kda_mtp_decode`` operator. One launch fuses, per generation +request: replay of previously-accepted draft tokens from the ``qkg/v/beta`` +caches, causal conv + SiLU, Q/K L2 norm, beta sigmoid, lower-bound gate, and +the KDA delta-rule recurrence over the ``1 + num_spec`` new tokens. The +recurrent state and base conv windows are committed **in place** after the +first new (golden) token; the spec tokens are cached for the next round. + +State-management contract (differs from the legacy intermediate-buffer + +``update_mamba_states`` promotion flow): after this op returns, the pools +hold the state as of the last golden token, with the new spec tokens pending +in the replay caches. The next round passes ``num_accepted_tokens`` (how +many of those pending drafts the sampler accepted) and the kernel replays +them before the new tokens. No host-side promotion of KDA SSM/conv state is +required or allowed. + +Productization deltas vs the drop's host wrapper (``reference.py``): + +* ``zero_accepted_hint`` / ``regular_metadata_hint`` are explicit caller + arguments. The drop derived them by *reading the device tensors* + (``torch.count_nonzero(...).item()`` / ``torch.equal``) behind + ``id()``-keyed caches — a host-device sync per novel tensor object plus a + stale-cache hazard on id reuse. Callers that statically know the pattern + (benchmarks, the first verify round after prefill) may pass the hints; + the runtime default (``False``/``False``) is always correct and never + syncs. +* The compile cache is keyed purely by dtype/shape/stride layouts and + constexpr flags — no ``id()`` or ``data_ptr`` keys. +* Bias / output-norm / ``pad_slot_id`` arguments (unsupported by the + specialized kernel, previously validated-then-rejected) are dropped from + the signature. + +Kernel shape contract: ``K == V == 128``, conv width ``W == 4``, +``HV == H``, TILE_V=64. ``num_spec`` is a compile-time constant per cache +allocation. Benchmark-tuned fast variants exist for ``N in (32, 128), H in +(2, 12, 32), num_spec == 2``; other shapes compile the general variant. +""" + +from typing import Optional, Tuple + +import torch + +from tensorrt_llm.logger import logger + +from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + +if IS_CUTLASS_DSL_AVAILABLE: + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + 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) +else: + raise ImportError("Kimi K3 KDA MTP decode requires NVIDIA CUTLASS DSL") + +_TILE_V = 64 + + +if IS_CUTLASS_DSL_AVAILABLE: + + @cute.jit + def _run_kda_decode_mtp( + h0: cute.Tensor, + x_q: cute.Tensor, + x_k: cute.Tensor, + x_v: cute.Tensor, + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + A_log: cute.Tensor, + g: cute.Tensor, + dt_bias: cute.Tensor, + beta: cute.Tensor, + o: cute.Tensor, + ht: cute.Tensor, + qkg_cache: cute.Tensor, + v_cache: cute.Tensor, + beta_cache: cute.Tensor, + stage_timing: cute.Tensor, + ssm_state_indices: cute.Tensor, + cu_seqlens: cute.Tensor, + num_accepted_tokens: cute.Tensor, + precompute_control: cute.Tensor, + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + N: cutlass.Constexpr[int], + NUM_SPEC: cutlass.Constexpr[int], + TILE_V: cutlass.Constexpr[int], + KERNEL_WIDTH: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + USE_FLAT_LAYOUT: cutlass.Constexpr[bool], + USE_SETMAXREG: cutlass.Constexpr[bool], + USE_REGULAR_METADATA: cutlass.Constexpr[bool], + USE_REG_Q_WEIGHTS: cutlass.Constexpr[bool], + USE_ZERO_ACCEPTED: cutlass.Constexpr[bool], + FUSE_PRECOMPUTE: cutlass.Constexpr[bool], + RUNTIME_PRECOMPUTE_FLAG: cutlass.Constexpr[bool], + PROFILE_STAGES: cutlass.Constexpr[bool], + stream: cuda.CUstream, + ): + if cutlass.const_expr(USE_ZERO_ACCEPTED): + t_max = 1 + NUM_SPEC + else: + t_max = 2 * NUM_SPEC + 1 + smem_qk_layout = cute.make_layout((t_max, K), stride=(K, 1)) + kda_decode_mtp_kernel( + h0, + x_q, + x_k, + x_v, + w_q, + w_k, + w_v, + cs_q, + cs_k, + cs_v, + A_log, + g, + dt_bias, + beta, + o, + ht, + qkg_cache, + v_cache, + beta_cache, + smem_qk_layout, + ssm_state_indices, + cu_seqlens, + num_accepted_tokens, + precompute_control, + TILE_V, + scale, + HV, + K, + V, + NUM_SPEC, + KERNEL_WIDTH, + lower_bound, + USE_FLAT_LAYOUT, + USE_SETMAXREG, + USE_REGULAR_METADATA, + USE_REG_Q_WEIGHTS, + USE_ZERO_ACCEPTED, + FUSE_PRECOMPUTE, + RUNTIME_PRECOMPUTE_FLAG, + stage_timing, + PROFILE_STAGES, + ).launch(grid=(HV, N, 1), block=[NUM_THREADS, 1, 1], stream=stream) + + +def _require_stride_layout( + *, + x_q, + x_k, + x_v, + w_q, + w_k, + w_v, + cs_q, + cs_k, + cs_v, + g, + beta, + A_log, + dt_bias, + recurrent_state, + qkg_cache, + v_cache, + beta_cache, + ssm_state_indices, + cu_seqlens, + num_accepted_tokens, + out, + H, + HV, + K, + V, + W, + num_spec, + T_total, +): + if x_q.ndim != 4 or x_k.ndim != 4 or x_v.ndim != 4: + raise ValueError("Expected x_q/x_k/x_v to have shape [1, T, H, D].") + if x_q.shape != (1, T_total, H, K) or x_k.shape != (1, T_total, H, K): + raise ValueError(f"Expected x_q/x_k shape [1, {T_total}, {H}, {K}].") + if x_v.shape != (1, T_total, HV, V): + raise ValueError(f"Expected x_v shape [1, {T_total}, {HV}, {V}].") + if g.ndim != 4 or g.shape != (1, T_total, HV, K): + raise ValueError(f"Expected g shape [1, {T_total}, {HV}, {K}].") + if beta.ndim != 3 or beta.shape != (1, T_total, HV): + raise ValueError(f"Expected beta shape [1, {T_total}, {HV}].") + if out.ndim != 4 or out.shape != (1, T_total, HV, V): + raise ValueError(f"Expected out shape [1, {T_total}, {HV}, {V}].") + + last_dim_tensors = { + "x_q": x_q, + "x_k": x_k, + "x_v": x_v, + "g": g, + "beta": beta, + "out": out, + "recurrent_state": recurrent_state, + "qkg_cache": qkg_cache, + "v_cache": v_cache, + "beta_cache": beta_cache, + } + 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.") + + 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.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.") + + if A_log.ndim != 1 or A_log.shape[0] != H: + raise ValueError(f"Expected A_log shape [{H}].") + if dt_bias.ndim != 1 or dt_bias.shape[0] != H * K: + raise ValueError(f"Expected dt_bias shape [{H * K}].") + + 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].") + 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.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)).") + + 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}].") + 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}].") + 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.") + + 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]: + raise ValueError("Expected num_accepted_tokens length to match N.") + + +def _layout_key(tensor: torch.Tensor): + return (tensor.dtype, tuple(tensor.shape), tuple(tensor.stride()), + _fits_32bit_stride(tensor)) + + +def _fits_32bit_stride(tensor: torch.Tensor) -> bool: + int32_max = 2**31 - 1 + max_offset = int(tensor.storage_offset()) + for size, stride in zip(tensor.shape, tensor.stride()): + stride = abs(int(stride)) + if stride > int32_max: + return False + if size: + max_offset += (int(size) - 1) * stride + if max_offset > int32_max: + return False + return True + + +def _from_dlpack_arg(tensor: torch.Tensor): + return from_dlpack( + tensor, + assumed_align=16, + use_32bit_stride=_fits_32bit_stride(tensor), + ) + + +def _dlpack_arg(tensor: torch.Tensor): + for dim, stride in enumerate(tensor.stride()): + if stride == 1: + return _from_dlpack_arg(tensor).mark_layout_dynamic(dim) + return _from_dlpack_arg(tensor) + + +# (device_index, enabled) -> persistent int32 [1] control tensor. Keys are +# plain values (not tensor identities), so entries never go stale. +_precompute_control_cache = {} + + +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)) + if key not in _precompute_control_cache: + _precompute_control_cache[key] = torch.tensor( + [1 if enabled else 0], dtype=torch.int32, device=dev) + return _precompute_control_cache[key] + + +def _try_flatten_args( + *, + recurrent_state: torch.Tensor, + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + T_total: int, + H: int, + HV: int, + K: int, + V: int, +) -> Tuple[bool, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + try: + h0 = recurrent_state.view(-1, V, K) + x_q_flat = x_q.view(1, T_total, H * K) + x_k_flat = x_k.view(1, T_total, H * K) + x_v_flat = x_v.view(1, T_total, HV * V) + except RuntimeError: + return False, recurrent_state, x_q, x_k, x_v + 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)) + + +# Layout-and-constexpr-keyed compile cache. Compilation is per (N, T_total, +# pool size, layouts, flags) — a new generation batch size triggers a +# multi-second cute.compile, after which the artifact is reused. +_compiled_cache = {} + + +def kda_mtp_decode_impl( + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + qkg_cache: torch.Tensor, + v_cache: torch.Tensor, + beta_cache: torch.Tensor, + ssm_state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + num_spec: int, + num_accepted_tokens: torch.Tensor, + lower_bound: float, + scale: Optional[float] = None, + out: Optional[torch.Tensor] = None, + zero_accepted_hint: bool = False, + regular_metadata_hint: bool = False, +) -> torch.Tensor: + """Launch the fused KDA MTP verify kernel. See the module docstring. + + Args (device tensors unless noted): + x_q/x_k/x_v: post-projection, pre-conv token states + ``[1, T_total, H, 128]`` bf16 — new tokens only, ``1 + + num_spec`` per request, packed per ``cu_seqlens``. + w_q/w_k/w_v: conv weights ``[H*128, W]`` fp32, width-contiguous. + cs_q/cs_k/cs_v: extended conv caches ``[pool, H*128, >= W-1+M]`` + fp32, dim-contiguous. Columns ``[0, W-1)`` are the committed + window; tail columns hold raw pending-draft inputs. Mutated. + g, beta: raw gate ``[1, T, H, 128]`` and beta ``[1, T, H]`` bf16. + A_log, dt_bias: fp32 ``[H]`` / ``[H*128]``. + recurrent_state: pool ``[pool, H, V, K]`` fp32, **V-first** layout + (matches the executor ssm pool and the single-token decode + kernel). Committed in place. + qkg_cache/v_cache/beta_cache: replay caches ``[pool, M, 3, H*K]`` / + ``[pool, M, H*V]`` / ``[pool, M, H]`` fp32. Mutated. + ssm_state_indices / cu_seqlens / num_accepted_tokens: per-request + slot, token offsets ``[N+1]``, accepted-draft counts ``[N]``. + zero_accepted_hint: caller asserts every ``num_accepted_tokens`` is + zero (compiles the smaller-smem no-replay variant). Wrong hints + produce wrong results — pass True only when statically known. + regular_metadata_hint: caller asserts ``cu_seqlens`` is the uniform + ``arange * (2*num_spec+1)`` pattern and ``ssm_state_indices`` + is ``arange(N)`` (benchmark identity layout). + + Returns the output ``[1, T_total, H, V]`` bf16 (rows for replayed + positions are zero; only new-token rows are written). + """ + _, T_total, _, D = x_q.shape + H = A_log.shape[0] + HV = g.shape[2] if g.ndim == 4 else H + K = D + V_dim = x_v.shape[-1] + W = w_q.shape[1] + if K != TILE_K or V_dim != 128 or W != 4: + raise ValueError("specialized kernel expects K=128, V=128, W=4") + if HV != H: + raise ValueError("specialized kernel expects HV == H") + if scale is None: + scale = K**-0.5 + + 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) + if num_accepted_tokens.dtype != torch.int32: + num_accepted_tokens = num_accepted_tokens.to(torch.int32) + + _require_stride_layout( + x_q=x_q, + x_k=x_k, + x_v=x_v, + w_q=w_q, + w_k=w_k, + w_v=w_v, + cs_q=cs_q, + cs_k=cs_k, + cs_v=cs_v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + qkg_cache=qkg_cache, + v_cache=v_cache, + beta_cache=beta_cache, + ssm_state_indices=ssm_state_indices, + cu_seqlens=cu_seqlens, + num_accepted_tokens=num_accepted_tokens, + out=out, + H=H, + HV=HV, + K=K, + V=V_dim, + W=W, + num_spec=num_spec, + T_total=T_total, + ) + + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + precompute_control = _precompute_control_tensor(x_q.device, True) + + use_flat_layout, h0_arg, x_q_arg, x_k_arg, x_v_arg = _try_flatten_args( + recurrent_state=recurrent_state, + x_q=x_q, + x_k=x_k, + x_v=x_v, + T_total=T_total, + H=H, + HV=HV, + K=K, + 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) + 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) + # stage_timing is unused (PROFILE_STAGES=False); pass `out` as the + # placeholder tensor argument like the drop's runner does. + stage_timing_arg = out + + key = ( + x_q.dtype, + scale, + HV, + K, + V_dim, + num_spec, + W, + N, + T_total, + pool_size, + lower_bound, + use_flat_layout, + _layout_key(h0_arg), + _layout_key(x_q_arg), + _layout_key(x_k_arg), + _layout_key(x_v_arg), + _layout_key(w_q), + _layout_key(w_k), + _layout_key(w_v), + _layout_key(cs_q), + _layout_key(cs_k), + _layout_key(cs_v), + _layout_key(A_log), + _layout_key(g), + _layout_key(dt_bias), + _layout_key(beta), + _layout_key(out), + _layout_key(qkg_cache), + _layout_key(v_cache), + _layout_key(beta_cache), + _layout_key(ssm_state_indices), + _layout_key(cu_seqlens), + _layout_key(num_accepted_tokens), + use_setmaxreg, + use_regular_metadata, + use_reg_q_weights, + use_zero_accepted, + ) + + if key not in _compiled_cache: + logger.info( + 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}") + _compiled_cache[key] = cute.compile( + _run_kda_decode_mtp, + _from_dlpack_arg(h0_arg), + _from_dlpack_arg(x_q_arg), + _from_dlpack_arg(x_k_arg), + _from_dlpack_arg(x_v_arg), + _from_dlpack_arg(w_q), + _from_dlpack_arg(w_k), + _from_dlpack_arg(w_v), + _from_dlpack_arg(cs_q), + _from_dlpack_arg(cs_k), + _from_dlpack_arg(cs_v), + _from_dlpack_arg(A_log), + _from_dlpack_arg(g), + _from_dlpack_arg(dt_bias), + _from_dlpack_arg(beta), + _from_dlpack_arg(out), + _from_dlpack_arg(h0_arg), + _from_dlpack_arg(qkg_cache), + _from_dlpack_arg(v_cache), + _from_dlpack_arg(beta_cache), + _from_dlpack_arg(stage_timing_arg), + _from_dlpack_arg(ssm_state_indices), + _from_dlpack_arg(cu_seqlens), + _from_dlpack_arg(num_accepted_tokens), + _from_dlpack_arg(precompute_control), + scale=scale, + HV=HV, + K=K, + V=V_dim, + N=N, + NUM_SPEC=num_spec, + TILE_V=_TILE_V, + KERNEL_WIDTH=W, + lower_bound=lower_bound, + USE_FLAT_LAYOUT=use_flat_layout, + USE_SETMAXREG=use_setmaxreg, + USE_REGULAR_METADATA=use_regular_metadata, + USE_REG_Q_WEIGHTS=use_reg_q_weights, + USE_ZERO_ACCEPTED=use_zero_accepted, + FUSE_PRECOMPUTE=True, + RUNTIME_PRECOMPUTE_FLAG=False, + PROFILE_STAGES=False, + stream=stream, + ) + + _compiled_cache[key]( + _dlpack_arg(h0_arg), + _dlpack_arg(x_q_arg), + _dlpack_arg(x_k_arg), + _dlpack_arg(x_v_arg), + _dlpack_arg(w_q), + _dlpack_arg(w_k), + _dlpack_arg(w_v), + _dlpack_arg(cs_q), + _dlpack_arg(cs_k), + _dlpack_arg(cs_v), + _dlpack_arg(A_log), + _dlpack_arg(g), + _dlpack_arg(dt_bias), + _dlpack_arg(beta), + _dlpack_arg(out), + _dlpack_arg(h0_arg), + _dlpack_arg(qkg_cache), + _dlpack_arg(v_cache), + _dlpack_arg(beta_cache), + _dlpack_arg(stage_timing_arg), + _dlpack_arg(ssm_state_indices), + _dlpack_arg(cu_seqlens), + _dlpack_arg(num_accepted_tokens), + _dlpack_arg(precompute_control), + stream, + ) + + return out + + +if IS_CUTLASS_DSL_AVAILABLE: + + @torch.library.custom_op( + "trtllm::kda_mtp_decode", + mutates_args=( + "cs_q", + "cs_k", + "cs_v", + "recurrent_state", + "qkg_cache", + "v_cache", + "beta_cache", + ), + device_types="cuda", + ) + def kda_mtp_decode( + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + qkg_cache: torch.Tensor, + v_cache: torch.Tensor, + beta_cache: torch.Tensor, + ssm_state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + num_spec: int, + num_accepted_tokens: torch.Tensor, + lower_bound: float, + scale: Optional[float] = None, + zero_accepted_hint: bool = False, + regular_metadata_hint: bool = False, + ) -> torch.Tensor: + """Fused KDA multi-token verify with in-place state commit.""" + return kda_mtp_decode_impl( + x_q=x_q, + x_k=x_k, + x_v=x_v, + w_q=w_q, + w_k=w_k, + w_v=w_v, + cs_q=cs_q, + cs_k=cs_k, + cs_v=cs_v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + qkg_cache=qkg_cache, + v_cache=v_cache, + beta_cache=beta_cache, + ssm_state_indices=ssm_state_indices, + cu_seqlens=cu_seqlens, + num_spec=num_spec, + num_accepted_tokens=num_accepted_tokens, + lower_bound=lower_bound, + scale=scale, + zero_accepted_hint=zero_accepted_hint, + regular_metadata_hint=regular_metadata_hint, + ) + + @kda_mtp_decode.register_fake + def _( + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + qkg_cache: torch.Tensor, + v_cache: torch.Tensor, + beta_cache: torch.Tensor, + ssm_state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + num_spec: int, + num_accepted_tokens: torch.Tensor, + lower_bound: float, + scale: Optional[float] = None, + zero_accepted_hint: bool = False, + regular_metadata_hint: bool = False, + ) -> torch.Tensor: + return x_v.new_empty(x_v.shape) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/__init__.py new file mode 100644 index 000000000000..2c274d483d1a --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CuTe DSL kernels for Kimi K3 KDA prefill on Blackwell GPUs.""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/akk_inverse.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/akk_inverse.py new file mode 100644 index 000000000000..a8a96aa43bac --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/akk_inverse.py @@ -0,0 +1,1075 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Akk 64×64 Lower Triangular Block Inversion — mixed FP32/BF16. + +BF16 input/output and packed BF16 shared storage. The four 16×16 diagonal +blocks use FP32 forward substitution to avoid unstable BF16 Neumann +cancellation; the cross-block products retain BF16 MMA with FP32 accumulators. + +Architecture: + - 4 warps (128 threads) per CTA, each CTA processes one 64×64 Akk matrix + - sAkk [64,36] fp32 — each FP32 slot holds packed bf16x2 (stride 36) + Total SMEM: 64*36*4 = 9216 bytes (vs 64*72*4 = 18432 in FP32 version) + - sTemp [16, 24, 2] fp32: inter-warp FP32 accumulator communication + - cp.async: BF16 global → packed bf16x2 in FP32 SMEM (raw bytes align) + - FP32 forward substitution for diagonal blocks + - BF16 MMA m16n8k16 with FP32 accumulators for cross-block products + - movmatrix.sync.aligned.m8n8.trans.b16 for A→B layout conversion + - C→A chain: cvt.rn.bf16x2.f32 to pack FP32 accum → bf16x2 A-operand + +Block layout in sAkk (4×4 sub-blocks of 16×16, packed bf16x2): + Upper tri = INPUT, Diagonal = in-place inversion, Lower tri = OUTPUT + +Stages: + 0. cp.async load bf16 64×64 → sAkk (packed bf16x2) + 1. Invert 4 diagonal blocks via FP32 forward substitution (2 warps) + 2. Warps 0-2: Ai10, Ai21, Ai32 via chain MMA (C→A pack) + 3. Warps 0+2 → Ai20, warps 1+3 → Ai31 (parallel pairs, sTemp) + 4. Warps 0+1+2 → Ai30 (sTemp aggregation) + 5. All warps: store bf16x2 SMEM → bf16 global + +Inputs: A_in [B, T, H, BT] bf16 +Outputs: A_out [B, T, H, BT] bf16 +""" + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync +from cutlass.cutlass_dsl import T, dsl_user_op + +# =========================================================================== +# Constants +# =========================================================================== +BS = 64 +SB = 16 +THREADS = 128 +TEMP_PAD = 8 +TEMP_COLS = SB + TEMP_PAD # 24 +NUM_TEMPS = 2 +AKK_PAD = 4 +AKK_STRIDE = BS // 2 + AKK_PAD # 36 (in FP32 units = 72 bf16 elements) + + +# =========================================================================== +# BF16 MMA m16n8k16 with FP32 accumulator +# =========================================================================== +@dsl_user_op +def mma_bf16_m16n8k16( + a0, + a1, + a2, + a3, + b0, + b1, + c0, + c1, + c2, + c3, + *, + loc=None, + ip=None, +): + """mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 + A: 4×i32 (bf16x2 pairs), B: 2×i32, C/D: 4×f32 (FP32 accum).""" + a0b = llvm.bitcast(T.i32(), a0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a1b = llvm.bitcast(T.i32(), a1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a2b = llvm.bitcast(T.i32(), a2.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a3b = llvm.bitcast(T.i32(), a3.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b0b = llvm.bitcast(T.i32(), b0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b1b = llvm.bitcast(T.i32(), b1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [ + a0b, + a1b, + a2b, + a3b, + b0b, + b1b, + c0.ir_value(loc=loc, ip=ip), + c1.ir_value(loc=loc, ip=ip), + c2.ir_value(loc=loc, ip=ip), + c3.ir_value(loc=loc, ip=ip), + ], + """{ + mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 + {$0, $1, $2, $3}, + {$4, $5, $6, $7}, + {$8, $9}, + {$10, $11, $12, $13}; + }""", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +# =========================================================================== +# movmatrix.sync.aligned.m8n8.trans.b16 — hardware A→B transpose (warp-level) +# Works on any 16-bit format including BF16. +# =========================================================================== +@dsl_user_op +def _movmatrix_trans(src, *, loc=None, ip=None): + src_b = llvm.bitcast(T.i32(), src.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [src_b], + "movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;", + "=r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +# =========================================================================== +# BF16x2 pack/unpack helpers +# =========================================================================== +@dsl_user_op +def _pack_bf16x2(lo_f32, hi_f32, *, loc=None, ip=None): + """Pack two FP32 values into one i32 holding two BF16 values. + cvt.rn.bf16x2.f32 converts and packs in one instruction.""" + result = llvm.inline_asm( + T.i32(), + [lo_f32.ir_value(loc=loc, ip=ip), hi_f32.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _mask_packed_ltri(packed, row, pair, *, loc=None, ip=None): + """Apply lower-triangular mask to packed bf16x2. + Zero out elements where row < col. col0=2*pair, col1=2*pair+1.""" + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p, row.ir_value(loc=loc, ip=ip), pair.ir_value(loc=loc, ip=ip)], + """{ + .reg .b32 %c0, %c1, %mlo, %mhi, %mask; + .reg .pred %p0, %p1; + shl.b32 %c0, $3, 1; + add.u32 %c1, %c0, 1; + setp.ge.s32 %p0, $2, %c0; + setp.ge.s32 %p1, $2, %c1; + selp.b32 %mlo, 0xFFFF, 0, %p0; + selp.b32 %mhi, 0xFFFF0000, 0, %p1; + or.b32 %mask, %mlo, %mhi; + and.b32 $0, $1, %mask; + }""", + "=r,r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _unpack_bf16x2_lo(packed, *, loc=None, ip=None): + """Unpack lower BF16 from packed i32 to FP32. + BF16 is upper 16 bits of FP32, so shift left by 16.""" + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p], + """{ + shl.b32 $0, $1, 16; + }""", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _unpack_bf16x2_hi(packed, *, loc=None, ip=None): + """Unpack upper BF16 from packed i32 to FP32. + Upper 16 bits are already in the right position for FP32.""" + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p], + """{ + and.b32 $0, $1, 0xFFFF0000; + }""", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +# =========================================================================== +# Neumann diagonal 16×16 inversion using BF16 MMA m16n8k16 + FP32 accum +# +# (I+L)⁻¹ = (I-L)(I+L²)(I+L⁴)(I+L⁸) +# +# All intermediate values kept in FP32. Pack to bf16x2 only at MMA boundaries. +# sAkk is packed bf16x2: sAkk[row, pair] = {bf16[row, 2*pair], bf16[row, 2*pair+1]} +# =========================================================================== +@dsl_user_op +def _invert_diag_neumann(sAkk: cute.Tensor, block_idx, lane_id, *, loc=None, ip=None): + r_off = block_idx * 16 + c_off = block_idx * 8 # packed bf16x2: 16 cols → 8 pairs + gid = lane_id // 4 + tid = lane_id % 4 + + # --- Load A from packed bf16x2 SMEM → unpack to 8 FP32 values --- + packed0 = cutlass.Float32(sAkk[r_off + gid, c_off + tid]) + packed1 = cutlass.Float32(sAkk[r_off + gid + 8, c_off + tid]) + packed2 = cutlass.Float32(sAkk[r_off + gid, c_off + 4 + tid]) + packed3 = cutlass.Float32(sAkk[r_off + gid + 8, c_off + 4 + tid]) + + A_f0 = _unpack_bf16x2_lo(packed0) # A[gid, 2*tid] + A_f1 = _unpack_bf16x2_hi(packed0) # A[gid, 2*tid+1] + A_f2 = _unpack_bf16x2_lo(packed1) # A[gid+8, 2*tid] + A_f3 = _unpack_bf16x2_hi(packed1) # A[gid+8, 2*tid+1] + A_f4 = _unpack_bf16x2_lo(packed2) # A[gid, 8+2*tid] + A_f5 = _unpack_bf16x2_hi(packed2) # A[gid, 8+2*tid+1] + A_f6 = _unpack_bf16x2_lo(packed3) # A[gid+8, 8+2*tid] + A_f7 = _unpack_bf16x2_hi(packed3) # A[gid+8, 8+2*tid+1] + + # Build identity (FP32) + _one = cutlass.Float32(1.0) + _zero = cutlass.Float32(0.0) + I_f0 = _one * cutlass.Float32(gid == 2 * tid) + _zero * cutlass.Float32(gid != 2 * tid) + I_f1 = _one * cutlass.Float32(gid == 2 * tid + 1) + _zero * cutlass.Float32(gid != 2 * tid + 1) + I_f2 = _one * cutlass.Float32(gid + 8 == 2 * tid) + _zero * cutlass.Float32(gid + 8 != 2 * tid) + I_f3 = _one * cutlass.Float32(gid + 8 == 2 * tid + 1) + _zero * cutlass.Float32( + gid + 8 != 2 * tid + 1 + ) + I_f4 = _one * cutlass.Float32(gid == 8 + 2 * tid) + _zero * cutlass.Float32(gid != 8 + 2 * tid) + I_f5 = _one * cutlass.Float32(gid == 8 + 2 * tid + 1) + _zero * cutlass.Float32( + gid != 8 + 2 * tid + 1 + ) + I_f6 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid) + _zero * cutlass.Float32( + gid + 8 != 8 + 2 * tid + ) + I_f7 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid + 1) + _zero * cutlass.Float32( + gid + 8 != 8 + 2 * tid + 1 + ) + + # L = A - I (FP32) + L_f0 = A_f0 - I_f0 + L_f1 = A_f1 - I_f1 + L_f2 = A_f2 - I_f2 + L_f3 = A_f3 - I_f3 + L_f4 = A_f4 - I_f4 + L_f5 = A_f5 - I_f5 + L_f6 = A_f6 - I_f6 + L_f7 = A_f7 - I_f7 + + # INV = I - L = 2I - A (FP32) + INV_f0 = I_f0 - L_f0 + INV_f1 = I_f1 - L_f1 + INV_f2 = I_f2 - L_f2 + INV_f3 = I_f3 - L_f3 + INV_f4 = I_f4 - L_f4 + INV_f5 = I_f5 - L_f5 + INV_f6 = I_f6 - L_f6 + INV_f7 = I_f7 - L_f7 + + _zf = cutlass.Float32(0.0) + + # Pack L and INV → bf16x2 for MMA + L_a0 = _pack_bf16x2(L_f0, L_f1) + L_a1 = _pack_bf16x2(L_f2, L_f3) + L_a2 = _pack_bf16x2(L_f4, L_f5) + L_a3 = _pack_bf16x2(L_f6, L_f7) + + INV_a0 = _pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _pack_bf16x2(INV_f6, INV_f7) + + # === Iteration 1: L² = L × L, then INV = INV + INV × L² === + L_b0 = _movmatrix_trans(L_a0) + L_b1 = _movmatrix_trans(L_a1) + L_b2 = _movmatrix_trans(L_a2) + L_b3 = _movmatrix_trans(L_a3) + + Lp_c0, Lp_c1, Lp_c2, Lp_c3 = mma_bf16_m16n8k16( + L_a0, L_a1, L_a2, L_a3, L_b0, L_b1, _zf, _zf, _zf, _zf + ) + Lp_c4, Lp_c5, Lp_c6, Lp_c7 = mma_bf16_m16n8k16( + L_a0, L_a1, L_a2, L_a3, L_b2, L_b3, _zf, _zf, _zf, _zf + ) + + Lp_a0 = _pack_bf16x2(Lp_c0, Lp_c1) + Lp_a1 = _pack_bf16x2(Lp_c2, Lp_c3) + Lp_a2 = _pack_bf16x2(Lp_c4, Lp_c5) + Lp_a3 = _pack_bf16x2(Lp_c6, Lp_c7) + + Lp_b0 = _movmatrix_trans(Lp_a0) + Lp_b1 = _movmatrix_trans(Lp_a1) + Lp_b2 = _movmatrix_trans(Lp_a2) + Lp_b3 = _movmatrix_trans(Lp_a3) + + # mm = INV × L² (FP32 accum output) + mm_c0, mm_c1, mm_c2, mm_c3 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf + ) + + # INV += mm (FP32) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + + INV_a0 = _pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _pack_bf16x2(INV_f6, INV_f7) + + # === Iteration 2: L⁴ = L² × L², then INV = INV + INV × L⁴ === + L4_c0, L4_c1, L4_c2, L4_c3 = mma_bf16_m16n8k16( + Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf + ) + L4_c4, L4_c5, L4_c6, L4_c7 = mma_bf16_m16n8k16( + Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf + ) + + L4_a0 = _pack_bf16x2(L4_c0, L4_c1) + L4_a1 = _pack_bf16x2(L4_c2, L4_c3) + L4_a2 = _pack_bf16x2(L4_c4, L4_c5) + L4_a3 = _pack_bf16x2(L4_c6, L4_c7) + + L4_b0 = _movmatrix_trans(L4_a0) + L4_b1 = _movmatrix_trans(L4_a1) + L4_b2 = _movmatrix_trans(L4_a2) + L4_b3 = _movmatrix_trans(L4_a3) + + mm_c0, mm_c1, mm_c2, mm_c3 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf + ) + + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + + INV_a0 = _pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _pack_bf16x2(INV_f6, INV_f7) + + # === Iteration 3: L⁸ = L⁴ × L⁴, then INV = INV + INV × L⁸ === + L8_c0, L8_c1, L8_c2, L8_c3 = mma_bf16_m16n8k16( + L4_a0, L4_a1, L4_a2, L4_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf + ) + L8_c4, L8_c5, L8_c6, L8_c7 = mma_bf16_m16n8k16( + L4_a0, L4_a1, L4_a2, L4_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf + ) + + L8_a0 = _pack_bf16x2(L8_c0, L8_c1) + L8_a1 = _pack_bf16x2(L8_c2, L8_c3) + L8_a2 = _pack_bf16x2(L8_c4, L8_c5) + L8_a3 = _pack_bf16x2(L8_c6, L8_c7) + + L8_b0 = _movmatrix_trans(L8_a0) + L8_b1 = _movmatrix_trans(L8_a1) + L8_b2 = _movmatrix_trans(L8_a2) + L8_b3 = _movmatrix_trans(L8_a3) + + mm_c0, mm_c1, mm_c2, mm_c3 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, L8_b0, L8_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = mma_bf16_m16n8k16( + INV_a0, INV_a1, INV_a2, INV_a3, L8_b2, L8_b3, _zf, _zf, _zf, _zf + ) + + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + + # --- Store INV to packed bf16x2 SMEM --- + sAkk[r_off + gid, c_off + tid] = _pack_bf16x2(INV_f0, INV_f1) + sAkk[r_off + gid + 8, c_off + tid] = _pack_bf16x2(INV_f2, INV_f3) + sAkk[r_off + gid, c_off + 4 + tid] = _pack_bf16x2(INV_f4, INV_f5) + sAkk[r_off + gid + 8, c_off + 4 + tid] = _pack_bf16x2(INV_f6, INV_f7) + + +@dsl_user_op +def _load_packed_bf16(sAkk: cute.Tensor, row, packed_col_base, col, *, loc=None, ip=None): + """Load one logical BF16 value from packed-bf16x2 shared storage.""" + valid_i = cutlass.Int32(col >= 0) + safe_col = col * valid_i + packed = cutlass.Float32(sAkk[row, packed_col_base + safe_col // 2]) + lo = _unpack_bf16x2_lo(packed) + hi = _unpack_bf16x2_hi(packed) + use_hi = cutlass.Float32(safe_col % 2) + value = lo * (cutlass.Float32(1.0) - use_hi) + hi * use_hi + return value * cutlass.Float32(valid_i) + + +@dsl_user_op +def _invert_diag_forward_fp32( + sAkk: cute.Tensor, + block_idx, + lane_id, + *, + loc=None, + ip=None, +): + """FP32 forward substitution with one final BF16 pack per inverse row.""" + row = lane_id % 16 + halfwarp_base = (lane_id // 16) * 16 + r_off = block_idx * 16 + c_off = block_idx * 8 + inv = cute.make_rmem_tensor(cute.make_layout((16,), stride=(1,)), cutlass.Float32) + inv[0] = cutlass.Float32(1.0) + for d in range(1, 16): + inv[d] = cutlass.Float32(0.0) + + for d in range(1, 16): + col_d = row - d + valid = cutlass.Float32(col_d >= 0) + a_val = _load_packed_bf16(sAkk, r_off + row, c_off, col_d) + acc = cutlass.Float32(0.0) + for j in range(1, d): + a_re = _load_packed_bf16(sAkk, r_off + row, c_off, row - (d - j)) + inv_prev = cute.arch.shuffle_sync(inv[j], halfwarp_base + row - d + j) + acc = acc + a_re * inv_prev + inv[d] = (-a_val - acc) * valid + + for pair in range(8): + col0 = pair * 2 + col1 = col0 + 1 + d0 = row - col0 + d1 = row - col1 + v0 = cutlass.Float32(row == col0) + v1 = cutlass.Float32(row == col1) + for d in range(1, 16): + v0 = v0 + inv[d] * cutlass.Float32(d0 == d) + v1 = v1 + inv[d] * cutlass.Float32(d1 == d) + sAkk[r_off + row, c_off + pair] = _pack_bf16x2(v0, v1) + + +# =========================================================================== +# 16×16 matmul: load A & B from packed bf16x2 sAkk, BF16 MMA, return FP32 C +# =========================================================================== +@dsl_user_op +def _matmul_AB(sAkk: cute.Tensor, br_A, bc_A, br_B, bc_B, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 8 # packed: 16 cols → 8 pairs + rB = br_B * 16 + cB = bc_B * 8 + + # Load A-operand (packed bf16x2, direct from SMEM) + a0 = cutlass.Float32(sAkk[rA + gid, cA + tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 4 + tid]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 4 + tid]) + + # Load B-operand: load in A-layout then movmatrix → B-layout + bA0 = cutlass.Float32(sAkk[rB + gid, cB + tid]) + bA1 = cutlass.Float32(sAkk[rB + gid + 8, cB + tid]) + bA2 = cutlass.Float32(sAkk[rB + gid, cB + 4 + tid]) + bA3 = cutlass.Float32(sAkk[rB + gid + 8, cB + 4 + tid]) + b0 = _movmatrix_trans(bA0) + b1 = _movmatrix_trans(bA1) + b2 = _movmatrix_trans(bA2) + b3 = _movmatrix_trans(bA3) + + # 16×16 = 2 × m16n8k16 + cn0_0, cn0_1, cn0_2, cn0_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf) + + # Return 8 FP32 C-registers (C-layout of m16n8k16 FP32 accum) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +# =========================================================================== +# Chain MMA: pre-loaded A (from C→A pack), load B from sAkk (Stage 2) +# A-operand already packed as bf16x2 from previous C result. +# =========================================================================== +@dsl_user_op +def _chain_mma_B(sAkk: cute.Tensor, br_B, bc_B, a0, a1, a2, a3, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rB = br_B * 16 + cB = bc_B * 8 + + bA0 = cutlass.Float32(sAkk[rB + gid, cB + tid]) + bA1 = cutlass.Float32(sAkk[rB + gid + 8, cB + tid]) + bA2 = cutlass.Float32(sAkk[rB + gid, cB + 4 + tid]) + bA3 = cutlass.Float32(sAkk[rB + gid + 8, cB + 4 + tid]) + b0 = _movmatrix_trans(bA0) + b1 = _movmatrix_trans(bA1) + b2 = _movmatrix_trans(bA2) + b3 = _movmatrix_trans(bA3) + + cn0_0, cn0_1, cn0_2, cn0_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf) + + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +# =========================================================================== +# Chain MMA: load A from sAkk, pre-loaded B (from C→B shuffle) (Stages 3-4) +# B-operand from shuffle is still FP32 TF32-layout → need to convert. +# Actually in BF16 version we use movmatrix, so B comes from pack+movmatrix. +# This function takes B already in B-layout (packed bf16x2 after movmatrix). +# =========================================================================== +@dsl_user_op +def _chain_mma_A(sAkk: cute.Tensor, br_A, bc_A, b0, b1, b2, b3, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 8 + + a0 = cutlass.Float32(sAkk[rA + gid, cA + tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 4 + tid]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 4 + tid]) + + cn0_0, cn0_1, cn0_2, cn0_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf) + + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +# =========================================================================== +# Store negated C result (16×16, FP32 accum) to packed bf16x2 sAkk +# C-layout for m16n8k16 FP32 accum: +# cn0_0 = C[gid, 0..7 left half col0], cn0_1 = C[gid+8, left half col0] +# cn0_2 = C[gid, left half col1], cn0_3 = C[gid+8, left half col1] +# cn1_* = right half +# Packing: negate FP32 then pack pairs → bf16x2 +# =========================================================================== +@dsl_user_op +def _store_neg_C( + sAkk: cute.Tensor, br, bc, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None +): + gid = lane_id // 4 + tid = lane_id % 4 + r = br * 16 + c = bc * 8 # packed + + # Pack negated FP32 pairs → bf16x2, then store + sAkk[r + gid, c + tid] = _pack_bf16x2(-c0, -c1) + sAkk[r + gid + 8, c + tid] = _pack_bf16x2(-c2, -c3) + sAkk[r + gid, c + 4 + tid] = _pack_bf16x2(-c4, -c5) + sAkk[r + gid + 8, c + 4 + tid] = _pack_bf16x2(-c6, -c7) + + +# =========================================================================== +# Pack FP32 C-accum → bf16x2 A-operand for C→A chain +# C-layout (FP32 accum, m16n8k16): 8 floats → 4 bf16x2 A-regs +# c0,c1 → a0 (rows gid/gid+8, left-half k0..7) +# c2,c3 → a1 +# c4,c5 → a2 (right-half k8..15) +# c6,c7 → a3 +# =========================================================================== +@dsl_user_op +def _pack_C_to_A(c0, c1, c2, c3, c4, c5, c6, c7, *, loc=None, ip=None): + a0 = _pack_bf16x2(c0, c1) + a1 = _pack_bf16x2(c2, c3) + a2 = _pack_bf16x2(c4, c5) + a3 = _pack_bf16x2(c6, c7) + return a0, a1, a2, a3 + + +# =========================================================================== +# Convert FP32 C-accum → bf16x2 B-operand via pack + movmatrix +# =========================================================================== +@dsl_user_op +def _pack_C_to_B(c0, c1, c2, c3, c4, c5, c6, c7, *, loc=None, ip=None): + a0 = _pack_bf16x2(c0, c1) + a1 = _pack_bf16x2(c2, c3) + a2 = _pack_bf16x2(c4, c5) + a3 = _pack_bf16x2(c6, c7) + b0 = _movmatrix_trans(a0) + b1 = _movmatrix_trans(a1) + b2 = _movmatrix_trans(a2) + b3 = _movmatrix_trans(a3) + return b0, b1, b2, b3 + + +# =========================================================================== +# sTemp helpers (FP32, non-swizzled, for inter-warp accumulator exchange) +# =========================================================================== +@dsl_user_op +def _store_C_temp( + sT: cute.Tensor, + buf, + c0, + c1, + c2, + c3, + c4, + c5, + c6, + c7, + lane_id, + *, + loc=None, + ip=None, +): + gid = lane_id // 4 + tid = lane_id % 4 + sT[gid, 2 * tid, buf] = c0 + sT[gid, 2 * tid + 1, buf] = c1 + sT[gid + 8, 2 * tid, buf] = c2 + sT[gid + 8, 2 * tid + 1, buf] = c3 + sT[gid, 8 + 2 * tid, buf] = c4 + sT[gid, 8 + 2 * tid + 1, buf] = c5 + sT[gid + 8, 8 + 2 * tid, buf] = c6 + sT[gid + 8, 8 + 2 * tid + 1, buf] = c7 + + +@dsl_user_op +def _load_C_temp(sT: cute.Tensor, buf, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + c0 = cutlass.Float32(sT[gid, 2 * tid, buf]) + c1 = cutlass.Float32(sT[gid, 2 * tid + 1, buf]) + c2 = cutlass.Float32(sT[gid + 8, 2 * tid, buf]) + c3 = cutlass.Float32(sT[gid + 8, 2 * tid + 1, buf]) + c4 = cutlass.Float32(sT[gid, 8 + 2 * tid, buf]) + c5 = cutlass.Float32(sT[gid, 8 + 2 * tid + 1, buf]) + c6 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid, buf]) + c7 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid + 1, buf]) + return c0, c1, c2, c3, c4, c5, c6, c7 + + +# =========================================================================== +# Main kernel +# =========================================================================== +@cute.kernel +def akk_inv_kernel( + g2s_copy: cute.TiledCopy, + gA_tensor: cute.Tensor, + mOut: cute.Tensor, + mBeta: cute.Tensor, + akk_smem_layout: cute.Layout, + temp_layout: cute.Layout, + NT: cutlass.Int32, + H: int, + mCuSeqlens: cute.Tensor, + mChunkIndices: cute.Tensor, + IS_VARLEN: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + lane_id = tidx % 32 + h_idx, nt_idx, b_idx = cute.arch.block_idx() + + # ===== SMEM allocation ===== + smem = cutlass.utils.SmemAllocator() + sAkk = smem.allocate_tensor(cutlass.Float32, akk_smem_layout, 128) + sTemp = smem.allocate_tensor(cutlass.Float32, temp_layout, 128) + sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout(BS, stride=1), 128) + + # ===== Stage 0: cp.async load bf16 global → packed bf16x2 sAkk ===== + ld_bnt = nt_idx + if IS_VARLEN: + ld_seq_id = cutlass.Int32(mChunkIndices[nt_idx, 0]) + ld_local = cutlass.Int32(mChunkIndices[nt_idx, 1]) + ld_bos = cutlass.Int32(mCuSeqlens[ld_seq_id]) + ld_bnt = ld_bos + ld_local * BS + gA_batch = gA_tensor[(None, None, h_idx, ld_bnt, b_idx)] + + thr_g2s = g2s_copy.get_slice(tidx) + thr_gSrc = thr_g2s.partition_S(gA_batch) + thr_sDst = thr_g2s.partition_D(sAkk) + cute.copy(g2s_copy, thr_gSrc, thr_sDst) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.barrier() + + # Zero out-of-bounds rows for varlen partial chunks. + if IS_VARLEN: + _z_seq = cutlass.Int32(mChunkIndices[nt_idx, 0]) + _z_local = cutlass.Int32(mChunkIndices[nt_idx, 1]) + _z_bos = cutlass.Int32(mCuSeqlens[_z_seq]) + _z_eos = cutlass.Int32(mCuSeqlens[_z_seq + 1]) + _z_cs = _z_bos + _z_local * BS + _z_vr = _z_eos - _z_cs + _zr_start = warp_idx * SB + for ri in cutlass.range_constexpr(SB): + row = _zr_start + ri + if row >= _z_vr: + # Zero packed bf16x2 slots (each slot = 2 bf16 elements) + c0 = lane_id + if c0 < AKK_STRIDE: + sAkk[row, c0] = cutlass.Float32(0.0) + cute.arch.barrier() + + # ===== Stage 0b: Load 64 per-token betas into sBeta (fp32) ===== + _b_chunk_start = nt_idx * BS + _b_eos = cutlass.Int32(_b_chunk_start + BS) + if IS_VARLEN: + _b_seq_id = cutlass.Int32(mChunkIndices[nt_idx, 0]) + _b_local = cutlass.Int32(mChunkIndices[nt_idx, 1]) + _b_bos = cutlass.Int32(mCuSeqlens[_b_seq_id]) + _b_eos = cutlass.Int32(mCuSeqlens[_b_seq_id + 1]) + _b_chunk_start = _b_bos + _b_local * BS + + if warp_idx == 0: + _bcol_lo = lane_id * 2 + _bcol_hi = _bcol_lo + 1 + if _bcol_lo < BS: + _bt_lo = _b_chunk_start + _bcol_lo + _bt_hi = _b_chunk_start + _bcol_hi + if IS_VARLEN: + if _bt_lo < _b_eos: + sBeta[_bcol_lo] = cutlass.Float32(mBeta[b_idx, _bt_lo, h_idx]) + else: + sBeta[_bcol_lo] = cutlass.Float32(0.0) + if _bt_hi < _b_eos: + sBeta[_bcol_hi] = cutlass.Float32(mBeta[b_idx, _bt_hi, h_idx]) + else: + sBeta[_bcol_hi] = cutlass.Float32(0.0) + else: + sBeta[_bcol_lo] = cutlass.Float32(mBeta[b_idx, _bt_lo, h_idx]) + sBeta[_bcol_hi] = cutlass.Float32(mBeta[b_idx, _bt_hi, h_idx]) + cute.arch.barrier() + + # ===== Stage 1: FP32 forward solve, two diagonal blocks per warp ===== + if warp_idx == 0: + _invert_diag_forward_fp32(sAkk, lane_id // 16, lane_id) + if warp_idx == 1: + _invert_diag_forward_fp32(sAkk, 2 + lane_id // 16, lane_id) + + cute.arch.barrier() + + # ===== Stage 2: First batch — Ai10, Ai21, Ai32 ===== + # C→A chain: pack FP32 C-accum → bf16x2 A-operand via _pack_C_to_A + if warp_idx == 0: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 1, 1, 0, 1, lane_id) + a0, a1, a2, a3 = _pack_C_to_A(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B(sAkk, 0, 0, a0, a1, a2, a3, lane_id) + _store_neg_C(sAkk, 1, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + if warp_idx == 1: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 2, 2, 1, 2, lane_id) + a0, a1, a2, a3 = _pack_C_to_A(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B(sAkk, 1, 1, a0, a1, a2, a3, lane_id) + _store_neg_C(sAkk, 2, 1, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + if warp_idx == 2: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 3, 3, 2, 3, lane_id) + a0, a1, a2, a3 = _pack_C_to_A(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B(sAkk, 2, 2, a0, a1, a2, a3, lane_id) + _store_neg_C(sAkk, 3, 2, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + cute.arch.barrier() + + # ===== Stage 3: Second batch — Ai20, Ai31 (warp pairs via sTemp) ===== + _z = cutlass.Float32(0.0) + t0 = _z + t1 = _z + t2 = _z + t3 = _z + t4 = _z + t5 = _z + t6 = _z + t7 = _z + + # --- Ai20 = -Ai22 @ (Akk20 @ Ai00 + Akk21 @ Ai10) --- + if warp_idx == 0: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 0, 2, 0, 0, lane_id) + + if warp_idx == 2: + s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk, 1, 2, 1, 0, lane_id) + _store_C_temp(sTemp, 0, s0, s1, s2, s3, s4, s5, s6, s7, lane_id) + + # --- Ai31 = -Ai33 @ (Akk31 @ Ai11 + Akk32 @ Ai21) --- + if warp_idx == 1: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 1, 3, 1, 1, lane_id) + + if warp_idx == 3: + s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk, 2, 3, 2, 1, lane_id) + _store_C_temp(sTemp, 1, s0, s1, s2, s3, s4, s5, s6, s7, lane_id) + + cute.arch.barrier() + + # Warp 0: accumulate T1+T2, pack→B, multiply by Ai22 + if warp_idx == 0: + e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 0, lane_id) + t0 = t0 + e0 + t1 = t1 + e1 + t2 = t2 + e2 + t3 = t3 + e3 + t4 = t4 + e4 + t5 = t5 + e5 + t6 = t6 + e6 + t7 = t7 + e7 + b0, b1, b2, b3 = _pack_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A(sAkk, 2, 2, b0, b1, b2, b3, lane_id) + _store_neg_C(sAkk, 2, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + # Warp 1: accumulate T1'+T2', pack→B, multiply by Ai33 + if warp_idx == 1: + e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 1, lane_id) + t0 = t0 + e0 + t1 = t1 + e1 + t2 = t2 + e2 + t3 = t3 + e3 + t4 = t4 + e4 + t5 = t5 + e5 + t6 = t6 + e6 + t7 = t7 + e7 + b0, b1, b2, b3 = _pack_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A(sAkk, 3, 3, b0, b1, b2, b3, lane_id) + _store_neg_C(sAkk, 3, 1, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + cute.arch.barrier() + + # ===== Stage 4: Third batch — Ai30 ===== + t0 = _z + t1 = _z + t2 = _z + t3 = _z + t4 = _z + t5 = _z + t6 = _z + t7 = _z + + if warp_idx == 0: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 0, 3, 0, 0, lane_id) + + if warp_idx == 1: + s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk, 1, 3, 1, 0, lane_id) + _store_C_temp(sTemp, 0, s0, s1, s2, s3, s4, s5, s6, s7, lane_id) + + if warp_idx == 2: + s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk, 2, 3, 2, 0, lane_id) + _store_C_temp(sTemp, 1, s0, s1, s2, s3, s4, s5, s6, s7, lane_id) + + cute.arch.barrier() + + # Warp 0: accumulate all three, pack→B, multiply by Ai33 + if warp_idx == 0: + e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 0, lane_id) + t0 = t0 + e0 + t1 = t1 + e1 + t2 = t2 + e2 + t3 = t3 + e3 + t4 = t4 + e4 + t5 = t5 + e5 + t6 = t6 + e6 + t7 = t7 + e7 + e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 1, lane_id) + t0 = t0 + e0 + t1 = t1 + e1 + t2 = t2 + e2 + t3 = t3 + e3 + t4 = t4 + e4 + t5 = t5 + e5 + t6 = t6 + e6 + t7 = t7 + e7 + b0, b1, b2, b3 = _pack_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A(sAkk, 3, 3, b0, b1, b2, b3, lane_id) + _store_neg_C(sAkk, 3, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + + cute.arch.barrier() + + # ===== Stage 5: Store packed bf16x2 sAkk → global (b32 with direct bit-mask) ===== + vl_chunk_start = nt_idx * BS + vl_eos = cutlass.Int32(vl_chunk_start + BS) + if IS_VARLEN: + vl_seq_id = cutlass.Int32(mChunkIndices[nt_idx, 0]) + vl_local = cutlass.Int32(mChunkIndices[nt_idx, 1]) + vl_bos = cutlass.Int32(mCuSeqlens[vl_seq_id]) + vl_eos = cutlass.Int32(mCuSeqlens[vl_seq_id + 1]) + vl_chunk_start = vl_bos + vl_local * BS + + row_start = warp_idx * SB + + # PDL: hint downstream K4 to pre-launch so its setup work overlaps with our + # gmem writes below. K4 has griddepcontrol_wait after its setup but before + # reading any gmem (TMA loads / gk_last_exp), so ordering is correct. + cute.arch.griddepcontrol_launch_dependents() + + for ri in cutlass.range_constexpr(SB): + row = row_start + ri + pair = lane_id + if pair < BS // 2: + packed = cutlass.Float32(sAkk[row, pair]) + masked = _mask_packed_ltri(packed, cutlass.Int32(row), cutlass.Int32(pair)) + + _beta_lo = cutlass.Float32(sBeta[2 * pair]) + _beta_hi = cutlass.Float32(sBeta[2 * pair + 1]) + _lo_f32 = _unpack_bf16x2_lo(masked) * _beta_lo + _hi_f32 = _unpack_bf16x2_hi(masked) * _beta_hi + final = _pack_bf16x2(_lo_f32, _hi_f32) + + t_row = vl_chunk_start + row + if IS_VARLEN: + if t_row < vl_eos: + mOut[b_idx, t_row, h_idx, pair] = final + else: + mOut[b_idx, t_row, h_idx, pair] = final + + +# =========================================================================== +# Host JIT function +# =========================================================================== +@cute.jit +def akk_inv_host( + A_in: cute.Tensor, + A_out: cute.Tensor, + Beta_in: cute.Tensor, + # B / NT / T_VAL are runtime scalars (shape-independent compile): + # baking them recompiled the kernel per batch shape — at eval traffic + # (unique total T per prefill batch) ~100 s of JIT per batch. Only H + # and IS_VARLEN remain compile-time specializers. + B: cutlass.Int32, + NT: cutlass.Int32, + H: cutlass.Constexpr[int], + mCuSeqlens: cute.Tensor, + mChunkIndices: cute.Tensor, + IS_VARLEN: cutlass.Constexpr[int], + T_VAL: cutlass.Int32, + # Launch stream — runtime argument; launching on the DSL default stream + # races with the executor's non-blocking execution stream. + stream: cuda.CUstream, +): + # BF16 input: view as FP32 (packed bf16x2). Each FP32 element = 2 BF16. + # Original BF16 shape: [B, T, H, BS] with stride per-element in BF16. + # As FP32 packed bf16x2: shape [BS, BS//2, H, dim3, B] + # The row stride in BF16 units is H*BS. In FP32 (bf16x2) units = H*BS/2. + # The col stride in BF16 is 1. In bf16x2 pairs: pair stride = 1 (adjacent pairs in memory). + _dim3_size = IS_VARLEN * T_VAL + (1 - IS_VARLEN) * (T_VAL // BS) + _dim3_stride_bf16 = IS_VARLEN * (H * BS) + (1 - IS_VARLEN) * (BS * H * BS) + # All strides in FP32 units (each FP32 = 2 BF16): + _row_stride = H * BS // 2 # row stride (BF16 row stride = H*BS, /2 for bf16x2) + _col_stride = 1 # adjacent bf16x2 pairs + _h_stride = BS // 2 # head stride (BF16 = BS, /2 for bf16x2) + _dim3_stride = _dim3_stride_bf16 // 2 + # T_VAL is a runtime Int32, so this stride is dynamic. Mathematically it is + # T_VAL * (H * BS // 2) = T_VAL * 32 * H, i.e. always a multiple of 4 fp32 + # elements (16 B). The IR verifier can't derive that on its own, and without + # it the b_idx slice offset breaks the 128-bit alignment proof required by + # the cp.async G2S atom below (ICE: "src ptr alignment (32 bits) does not + # meet requirement (128 bits)"). Assert divisibility explicitly. + _batch_stride = cute.assume(T_VAL * (H * BS // 2), divby=H * BS // 2) + + view_layout = cute.make_layout( + (BS, BS // 2, H, _dim3_size, B), + stride=(_row_stride, _col_stride, _h_stride, _dim3_stride, _batch_stride), + ) + gA_view = cute.make_tensor(A_in.iterator, view_layout) + + # sAkk: packed bf16x2, logical shape (64, 32) with stride 36 for bank-conflict-free access + akk_smem_2d = cute.make_layout((BS, BS // 2), stride=(AKK_STRIDE, 1)) + + # cp.async G→S copy: 128-bit (4×fp32 = 8 bf16) vectorised + copy_atom = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + cutlass.Float32, + num_bits_per_copy=128, + ) + # Source layout: BS rows × BS//2 pairs = 64×32 + # 128 threads, each copies 4 FP32 (128 bits) per iteration + # Need 64*32/128/4 = 4 iterations → but (16,8) × (1,4) = 128 threads × 4 vals = 512 per iter + # 64×32 = 2048 elements / 512 = 4 iterations. But SMEM is 64×36, need to handle padding. + # Actually: the copy maps source shape to SMEM shape. + # Source is (BS, BS//2) = (64, 32), SMEM is (BS, AKK_STRIDE) = (64, 36). + # The copy handles the actual data region (64, 32); padding cols 32-35 stay zero. + g2s_copy = cute.make_tiled_copy_tv( + copy_atom, + thr_layout=cute.make_layout((16, 8), stride=(8, 1)), + val_layout=cute.make_layout((1, 4)), + ) + + # sTemp layout (non-swizzled) + temp_layout = cute.make_layout( + (SB, TEMP_COLS, NUM_TEMPS), stride=(TEMP_COLS, 1, SB * TEMP_COLS) + ) + + smem_bytes = BS * AKK_STRIDE * 4 + SB * TEMP_COLS * NUM_TEMPS * 4 + BS * 4 + 256 + + out_layout = cute.make_layout( + (B, T_VAL, H, BS // 2), stride=(T_VAL * H * BS // 2, H * BS // 2, BS // 2, 1) + ) + out_view = cute.make_tensor(A_out.iterator, out_layout) + + akk_inv_kernel( + g2s_copy, + gA_view, + out_view, + Beta_in, + akk_smem_2d, + temp_layout, + NT, + H, + mCuSeqlens, + mChunkIndices, + IS_VARLEN, + ).launch( + grid=(H, NT, B), + block=(THREADS, 1, 1), + smem=smem_bytes, + stream=stream, + ) 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 new file mode 100644 index 000000000000..771bf43e73f9 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py @@ -0,0 +1,2306 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Persistent Fused K1+K2+K3 Kernel for KDA. + +Fuses gate activation + cumsum + scaling (K1), intra sub-chunk Aqk/Akk (K2), +and inter sub-chunk solve + merged inverse (K3) into a single persistent kernel. + +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): + 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 + +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) + 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 + 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 + 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 + + sPartialLast [8,132] fp32 + sAqk [16,168,2] bf16 + + sAkk [64,72,2] fp32 block-transposed upper-tri layout) + +Inputs: + g [B,T,H,K] bf16 raw gate + k [B,T,H,K] bf16 + q [B,T,H,K] bf16 + A_log [H] fp32 per-head log decay + beta [B,T,H] bf16 used for Akk unit lower triangular + scale fp32 1/sqrt(K) + +Outputs (g_cumsum stays in SMEM, not written to GMEM): + k_scaled [B,T,H,K] bf16 + q_scaled [B,T,H,K] bf16 + kg [B,T,H,K] bf16 + gk_last_exp[B,NT,H,K] fp32 + A_qk [B,T,H,BT] bf16 full merged (diagonal + off-diagonal) + A_kk [B,T,H,BT] fp32 block-transposed upper-tri (input to akk_inv) +""" + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +from cutlass import for_generate, yield_out +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cutlass_dsl import T, dsl_user_op + +BT = 64 +BC = 16 +K_DIM = 128 +K_PAD = 8 +K_STRIDE = K_DIM + K_PAD # 136, padded row stride to avoid bank conflicts +CHUNKS_PER_BLOCK = 4 +NUM_SMS = 148 # Persistent kernel: one resident block per SM + +NUM_K1_TMA_WARPS = 16 # Warps 0-15: K1 compute (4 warpgroups, 8×2) -- TMA offloaded +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_SUB_CHUNKS = BT // BC # 4 +NUM_TILES = NUM_SUB_CHUNKS * (NUM_SUB_CHUNKS + 1) // 2 # 10 lower-tri tiles +MMA_K_TILE = 16 +NUM_MMA_K_TILES = K_DIM // MMA_K_TILE # 8 (bf16 m16n8k16) +AQK_TILE_PAD = 8 +AQK_TILE_STRIDE = BT + AQK_TILE_PAD # 72 — sAqk now 64x72 row-major (same shape as sAkk) + +AKK_PAD = 8 +AKK_STRIDE = BT + AKK_PAD # 72 + +# For fused akk_inv: sAkk_pkd viewed as fp32 packed bf16x2, stride = BT/2 + small pad +AKK_PKD_PAD = 4 +AKK_PKD_STRIDE = BT // 2 + AKK_PKD_PAD # 36 fp32 units = 72 bf16 elements per row +AKK_TEMP_PAD = 8 +AKK_TEMP_COLS = 16 + AKK_TEMP_PAD # 24 +AKK_TEMP_BUFS = 2 + +K1_ROW_GROUPS = 8 +K1_COL_GROUPS = 2 +ROWS_PER_K1_WARP = BT // K1_ROW_GROUPS # 8 +K1_COLS_PER_WARP = K_DIM // K1_COL_GROUPS # 64 +ROWS_PER_STORE_WARP = BT // NUM_STORE_WARPS # 16 + +VEC = K1_COLS_PER_WARP // 32 # 2 +K_VEC = K_DIM // VEC # 64 +NUM_STAGES = 2 +PARTIAL_COLS = K_DIM + 4 # 132 +PARTIAL_COLS_PER_WARP = K_DIM // NUM_K1_TMA_WARPS # 8 + +_TILE_IQ = [0, 1, 1, 2, 2, 2, 3, 3, 3, 3] +_TILE_IK = [0, 0, 1, 0, 1, 2, 0, 1, 2, 3] + +LOG2E = 1.4426950408889634 +LN2 = 0.6931471805599453 +RCP_LN2 = LOG2E + + +@dsl_user_op +def k1_internal_barrier(*, loc=None, ip=None): + """Named barrier for K1+TMA warps (0-15, 512 threads). barrier_id=2.""" + llvm.inline_asm( + T.i32(), + [], + "membar.cta; bar.sync 2, 512; mov.u32 $0, 0;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def pack_bf16x2_f32(hi_f32, lo_f32, *, loc=None, ip=None): + """Pack two fp32 values into a bf16x2 u32 register. + Returns u32 with bf16(hi) in bits [31:16] and bf16(lo) in bits [15:0]. + PTX: cvt.rn.bf16x2.f32 d, a, b -> d[31:16]=bf16(a), d[15:0]=bf16(b) + """ + result = llvm.inline_asm( + T.i32(), + [hi_f32.ir_value(loc=loc, ip=ip), lo_f32.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Int32(result) + + +@dsl_user_op +def mma_bf16_m16n8k16( + a0, + a1, + a2, + a3, # 4 u32 (bf16x2 packed) A operands + b0, + b1, # 2 u32 (bf16x2 packed) B operands + c0, + c1, + c2, + c3, # 4 fp32 accumulators + *, + loc=None, + ip=None, +): + """bf16 MMA with fp32 accumulator, shape m16n8k16. + D_fp32 = A_bf16 * B_bf16 + C_fp32 + """ + # a/b already i32 (from pack_bf16x2_f32) -> no bitcast needed + result = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [ + a0.ir_value(loc=loc, ip=ip), + a1.ir_value(loc=loc, ip=ip), + a2.ir_value(loc=loc, ip=ip), + a3.ir_value(loc=loc, ip=ip), + b0.ir_value(loc=loc, ip=ip), + b1.ir_value(loc=loc, ip=ip), + c0.ir_value(loc=loc, ip=ip), + c1.ir_value(loc=loc, ip=ip), + c2.ir_value(loc=loc, ip=ip), + c3.ir_value(loc=loc, ip=ip), + ], + """{ + mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 + {$0, $1, $2, $3}, + {$4, $5, $6, $7}, + {$8, $9}, + {$10, $11, $12, $13}; + }""", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +SHFL_W8_CLAMP = 0x1800 + + +@dsl_user_op +def fast_rcp(x, *, loc=None, ip=None): + """Hardware fast reciprocal: rcp.approx.ftz.f32 (~2 cycles vs ~20 for div).""" + result = llvm.inline_asm( + T.f32(), + [x.ir_value(loc=loc, ip=ip)], + "rcp.approx.ftz.f32 $0, $1;", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(result) + + +# ============================================================ +# Akk inverse helpers (fused from Akk_inverse_lower_triangle_bf16.py) +# ============================================================ +@dsl_user_op +def store_internal_barrier(*, loc=None, ip=None): + """Named barrier for Store warps (4 warps, 128 threads). barrier_id=3.""" + llvm.inline_asm( + T.i32(), + [], + "membar.cta; bar.sync 3, 128; mov.u32 $0, 0;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ak_pack_bf16x2(lo_f32, hi_f32, *, loc=None, ip=None): + """cvt.rn.bf16x2.f32 -- pack two fp32 into bf16x2 (as fp32 bitcast view).""" + result = llvm.inline_asm( + T.i32(), + [lo_f32.ir_value(loc=loc, ip=ip), hi_f32.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _ak_movmatrix_trans(src, *, loc=None, ip=None): + src_b = llvm.bitcast(T.i32(), src.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [src_b], + "movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;", + "=r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _ak_mask_packed_ltri(packed, row, pair, *, loc=None, ip=None): + """Mask packed bf16x2 to lower-tri: zero elements where row < col.""" + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p, row.ir_value(loc=loc, ip=ip), pair.ir_value(loc=loc, ip=ip)], + """{ + .reg .b32 %c0, %c1, %mlo, %mhi, %mask; + .reg .pred %p0, %p1; + shl.b32 %c0, $3, 1; + add.u32 %c1, %c0, 1; + setp.ge.s32 %p0, $2, %c0; + setp.ge.s32 %p1, $2, %c1; + selp.b32 %mlo, 0xFFFF, 0, %p0; + selp.b32 %mhi, 0xFFFF0000, 0, %p1; + or.b32 %mask, %mlo, %mhi; + and.b32 $0, $1, %mask; + }""", + "=r,r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _ak_unpack_bf16x2_lo(packed, *, loc=None, ip=None): + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p], + "shl.b32 $0, $1, 16;", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _ak_unpack_bf16x2_hi(packed, *, loc=None, ip=None): + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p], + "and.b32 $0, $1, 0xFFFF0000;", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _ak_mma(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None): + """BF16 MMA m16n8k16, args are Float32 (bitcast-viewed as i32 for bf16x2).""" + a0b = llvm.bitcast(T.i32(), a0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a1b = llvm.bitcast(T.i32(), a1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a2b = llvm.bitcast(T.i32(), a2.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a3b = llvm.bitcast(T.i32(), a3.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b0b = llvm.bitcast(T.i32(), b0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b1b = llvm.bitcast(T.i32(), b1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [ + a0b, + a1b, + a2b, + a3b, + b0b, + b1b, + c0.ir_value(loc=loc, ip=ip), + c1.ir_value(loc=loc, ip=ip), + c2.ir_value(loc=loc, ip=ip), + c3.ir_value(loc=loc, ip=ip), + ], + """{ + mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 + {$0, $1, $2, $3}, + {$4, $5, $6, $7}, + {$8, $9}, + {$10, $11, $12, $13}; + }""", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +@dsl_user_op +def _ak_invert_diag_neumann(sAkk, block_idx, lane_id, *, loc=None, ip=None): + """Invert 16x16 diag block via Neumann series (I+L)^-1 = (I-L)(I+L^2)(I+L^4)(I+L^8). + Reads/writes sAkk as fp32 packed bf16x2 (stride 36). Each block handled by 1 warp.""" + r_off = block_idx * 16 + c_off = block_idx * 8 # 16 cols = 8 pairs + gid = lane_id // 4 + tid = lane_id % 4 + + packed0 = cutlass.Float32(sAkk[r_off + gid, c_off + tid]) + packed1 = cutlass.Float32(sAkk[r_off + gid + 8, c_off + tid]) + packed2 = cutlass.Float32(sAkk[r_off + gid, c_off + 4 + tid]) + packed3 = cutlass.Float32(sAkk[r_off + gid + 8, c_off + 4 + tid]) + + A_f0 = _ak_unpack_bf16x2_lo(packed0) + A_f1 = _ak_unpack_bf16x2_hi(packed0) + A_f2 = _ak_unpack_bf16x2_lo(packed1) + A_f3 = _ak_unpack_bf16x2_hi(packed1) + A_f4 = _ak_unpack_bf16x2_lo(packed2) + A_f5 = _ak_unpack_bf16x2_hi(packed2) + A_f6 = _ak_unpack_bf16x2_lo(packed3) + A_f7 = _ak_unpack_bf16x2_hi(packed3) + + _one = cutlass.Float32(1.0) + _zero = cutlass.Float32(0.0) + I_f0 = _one * cutlass.Float32(gid == 2 * tid) + _zero * cutlass.Float32(gid != 2 * tid) + I_f1 = _one * cutlass.Float32(gid == 2 * tid + 1) + _zero * cutlass.Float32(gid != 2 * tid + 1) + I_f2 = _one * cutlass.Float32(gid + 8 == 2 * tid) + _zero * cutlass.Float32(gid + 8 != 2 * tid) + I_f3 = _one * cutlass.Float32(gid + 8 == 2 * tid + 1) + _zero * cutlass.Float32( + gid + 8 != 2 * tid + 1 + ) + I_f4 = _one * cutlass.Float32(gid == 8 + 2 * tid) + _zero * cutlass.Float32(gid != 8 + 2 * tid) + I_f5 = _one * cutlass.Float32(gid == 8 + 2 * tid + 1) + _zero * cutlass.Float32( + gid != 8 + 2 * tid + 1 + ) + I_f6 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid) + _zero * cutlass.Float32( + gid + 8 != 8 + 2 * tid + ) + I_f7 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid + 1) + _zero * cutlass.Float32( + gid + 8 != 8 + 2 * tid + 1 + ) + + L_f0 = A_f0 - I_f0 + L_f1 = A_f1 - I_f1 + L_f2 = A_f2 - I_f2 + L_f3 = A_f3 - I_f3 + L_f4 = A_f4 - I_f4 + L_f5 = A_f5 - I_f5 + L_f6 = A_f6 - I_f6 + L_f7 = A_f7 - I_f7 + INV_f0 = I_f0 - L_f0 + INV_f1 = I_f1 - L_f1 + INV_f2 = I_f2 - L_f2 + INV_f3 = I_f3 - L_f3 + INV_f4 = I_f4 - L_f4 + INV_f5 = I_f5 - L_f5 + INV_f6 = I_f6 - L_f6 + INV_f7 = I_f7 - L_f7 + + _zf = cutlass.Float32(0.0) + L_a0 = _ak_pack_bf16x2(L_f0, L_f1) + L_a1 = _ak_pack_bf16x2(L_f2, L_f3) + L_a2 = _ak_pack_bf16x2(L_f4, L_f5) + L_a3 = _ak_pack_bf16x2(L_f6, L_f7) + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 1: L^2, INV += INV*L^2 + L_b0 = _ak_movmatrix_trans(L_a0) + L_b1 = _ak_movmatrix_trans(L_a1) + L_b2 = _ak_movmatrix_trans(L_a2) + L_b3 = _ak_movmatrix_trans(L_a3) + Lp_c0, Lp_c1, Lp_c2, Lp_c3 = _ak_mma(L_a0, L_a1, L_a2, L_a3, L_b0, L_b1, _zf, _zf, _zf, _zf) + Lp_c4, Lp_c5, Lp_c6, Lp_c7 = _ak_mma(L_a0, L_a1, L_a2, L_a3, L_b2, L_b3, _zf, _zf, _zf, _zf) + Lp_a0 = _ak_pack_bf16x2(Lp_c0, Lp_c1) + Lp_a1 = _ak_pack_bf16x2(Lp_c2, Lp_c3) + Lp_a2 = _ak_pack_bf16x2(Lp_c4, Lp_c5) + Lp_a3 = _ak_pack_bf16x2(Lp_c6, Lp_c7) + Lp_b0 = _ak_movmatrix_trans(Lp_a0) + Lp_b1 = _ak_movmatrix_trans(Lp_a1) + Lp_b2 = _ak_movmatrix_trans(Lp_a2) + Lp_b3 = _ak_movmatrix_trans(Lp_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf + ) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 2: L^4, INV += INV*L^4 + L4_c0, L4_c1, L4_c2, L4_c3 = _ak_mma( + Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf + ) + L4_c4, L4_c5, L4_c6, L4_c7 = _ak_mma( + Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf + ) + L4_a0 = _ak_pack_bf16x2(L4_c0, L4_c1) + L4_a1 = _ak_pack_bf16x2(L4_c2, L4_c3) + L4_a2 = _ak_pack_bf16x2(L4_c4, L4_c5) + L4_a3 = _ak_pack_bf16x2(L4_c6, L4_c7) + L4_b0 = _ak_movmatrix_trans(L4_a0) + L4_b1 = _ak_movmatrix_trans(L4_a1) + L4_b2 = _ak_movmatrix_trans(L4_a2) + L4_b3 = _ak_movmatrix_trans(L4_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf + ) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 3: L^8, INV += INV*L^8 + L8_c0, L8_c1, L8_c2, L8_c3 = _ak_mma( + L4_a0, L4_a1, L4_a2, L4_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf + ) + L8_c4, L8_c5, L8_c6, L8_c7 = _ak_mma( + L4_a0, L4_a1, L4_a2, L4_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf + ) + L8_a0 = _ak_pack_bf16x2(L8_c0, L8_c1) + L8_a1 = _ak_pack_bf16x2(L8_c2, L8_c3) + L8_a2 = _ak_pack_bf16x2(L8_c4, L8_c5) + L8_a3 = _ak_pack_bf16x2(L8_c6, L8_c7) + L8_b0 = _ak_movmatrix_trans(L8_a0) + L8_b1 = _ak_movmatrix_trans(L8_a1) + L8_b2 = _ak_movmatrix_trans(L8_a2) + L8_b3 = _ak_movmatrix_trans(L8_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, L8_b0, L8_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, L8_b2, L8_b3, _zf, _zf, _zf, _zf + ) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + + sAkk[r_off + gid, c_off + tid] = _ak_pack_bf16x2(INV_f0, INV_f1) + sAkk[r_off + gid + 8, c_off + tid] = _ak_pack_bf16x2(INV_f2, INV_f3) + sAkk[r_off + gid, c_off + 4 + tid] = _ak_pack_bf16x2(INV_f4, INV_f5) + sAkk[r_off + gid + 8, c_off + 4 + tid] = _ak_pack_bf16x2(INV_f6, INV_f7) + + +@dsl_user_op +def _ak_invert_diag_neumann_inreg( + sAkk, + block_idx, + lane_id, + raw_f0, + raw_f1, + raw_f2, + raw_f3, + raw_f4, + raw_f5, + raw_f6, + raw_f7, + *, + loc=None, + ip=None, +): + """In-register variant: invert 16x16 diag block from K2 acc fragment values + (no SMEM read). Applies I+L mask, runs Neumann iterations, writes INV to sAkk. + + raw_f0..raw_f7 are K2 fp32 acc*beta values in C-fragment layout for m16n8k16: + raw_f0,raw_f1 = (gid, 2*tid), (gid, 2*tid+1) + raw_f2,raw_f3 = (gid+8, 2*tid), (gid+8, 2*tid+1) + raw_f4,raw_f5 = (gid, 2*tid+8), (gid, 2*tid+9) + raw_f6,raw_f7 = (gid+8, 2*tid+8), (gid+8, 2*tid+9) + """ + r_off = block_idx * 16 + c_off = block_idx * 8 + gid = lane_id // 4 + tid = lane_id % 4 + + _one = cutlass.Float32(1.0) + _zero = cutlass.Float32(0.0) + + # Identity pattern (matches the SMEM-read variant's I_f computation) + I_f0 = _one * cutlass.Float32(gid == 2 * tid) + _zero * cutlass.Float32(gid != 2 * tid) + I_f1 = _one * cutlass.Float32(gid == 2 * tid + 1) + _zero * cutlass.Float32(gid != 2 * tid + 1) + I_f2 = _one * cutlass.Float32(gid + 8 == 2 * tid) + _zero * cutlass.Float32( + gid + 8 != 2 * tid + ) # always 0 + I_f3 = _one * cutlass.Float32(gid + 8 == 2 * tid + 1) + _zero * cutlass.Float32( + gid + 8 != 2 * tid + 1 + ) # always 0 + I_f4 = _one * cutlass.Float32(gid == 8 + 2 * tid) + _zero * cutlass.Float32( + gid != 8 + 2 * tid + ) # always 0 + I_f5 = _one * cutlass.Float32(gid == 8 + 2 * tid + 1) + _zero * cutlass.Float32( + gid != 8 + 2 * tid + 1 + ) # always 0 + I_f6 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid) + _zero * cutlass.Float32( + gid + 8 != 8 + 2 * tid + ) + I_f7 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid + 1) + _zero * cutlass.Float32( + gid + 8 != 8 + 2 * tid + 1 + ) + + # Apply I+L mask to raw K2 values: + # diag (row==col): replace with 1 + # strict upper (rowcol within block): keep raw value (= L value) + m_gt_a = cutlass.Float32(gid > 2 * tid) # for cols 2*tid (with row=gid) + m_gt_b = cutlass.Float32(gid > 2 * tid + 1) # for cols 2*tid+1 + # (gid, 2*tid) and (gid, 2*tid+1): mask conditional + A_f0 = m_gt_a * raw_f0 + I_f0 + A_f1 = m_gt_b * raw_f1 + I_f1 + # (gid+8, 2*tid) and (gid+8, 2*tid+1): always strict lower (gid+8 > 2*tid+1 always) + A_f2 = raw_f2 + A_f3 = raw_f3 + # (gid, 2*tid+8) and (gid, 2*tid+9): always strict upper (gid <= 7 < 8 <= 2*tid+8) + A_f4 = _zero + A_f5 = _zero + # (gid+8, 2*tid+8): row-col offset = gid+8 - (2*tid+8) = gid - 2*tid -> same as A_f0 mask + # (gid+8, 2*tid+9): similar -> same as A_f1 mask + A_f6 = m_gt_a * raw_f6 + I_f6 + A_f7 = m_gt_b * raw_f7 + I_f7 + + # L = A - I, INV = I - L (first two Neumann terms) + L_f0 = A_f0 - I_f0 + L_f1 = A_f1 - I_f1 + L_f2 = A_f2 - I_f2 + L_f3 = A_f3 - I_f3 + L_f4 = A_f4 - I_f4 + L_f5 = A_f5 - I_f5 + L_f6 = A_f6 - I_f6 + L_f7 = A_f7 - I_f7 + INV_f0 = I_f0 - L_f0 + INV_f1 = I_f1 - L_f1 + INV_f2 = I_f2 - L_f2 + INV_f3 = I_f3 - L_f3 + INV_f4 = I_f4 - L_f4 + INV_f5 = I_f5 - L_f5 + INV_f6 = I_f6 - L_f6 + INV_f7 = I_f7 - L_f7 + + _zf = cutlass.Float32(0.0) + L_a0 = _ak_pack_bf16x2(L_f0, L_f1) + L_a1 = _ak_pack_bf16x2(L_f2, L_f3) + L_a2 = _ak_pack_bf16x2(L_f4, L_f5) + L_a3 = _ak_pack_bf16x2(L_f6, L_f7) + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 1: L^2, INV += INV*L^2 + L_b0 = _ak_movmatrix_trans(L_a0) + L_b1 = _ak_movmatrix_trans(L_a1) + L_b2 = _ak_movmatrix_trans(L_a2) + L_b3 = _ak_movmatrix_trans(L_a3) + Lp_c0, Lp_c1, Lp_c2, Lp_c3 = _ak_mma(L_a0, L_a1, L_a2, L_a3, L_b0, L_b1, _zf, _zf, _zf, _zf) + Lp_c4, Lp_c5, Lp_c6, Lp_c7 = _ak_mma(L_a0, L_a1, L_a2, L_a3, L_b2, L_b3, _zf, _zf, _zf, _zf) + Lp_a0 = _ak_pack_bf16x2(Lp_c0, Lp_c1) + Lp_a1 = _ak_pack_bf16x2(Lp_c2, Lp_c3) + Lp_a2 = _ak_pack_bf16x2(Lp_c4, Lp_c5) + Lp_a3 = _ak_pack_bf16x2(Lp_c6, Lp_c7) + Lp_b0 = _ak_movmatrix_trans(Lp_a0) + Lp_b1 = _ak_movmatrix_trans(Lp_a1) + Lp_b2 = _ak_movmatrix_trans(Lp_a2) + Lp_b3 = _ak_movmatrix_trans(Lp_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf + ) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 2: L^4, INV += INV*L^4 + L4_c0, L4_c1, L4_c2, L4_c3 = _ak_mma( + Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf + ) + L4_c4, L4_c5, L4_c6, L4_c7 = _ak_mma( + Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf + ) + L4_a0 = _ak_pack_bf16x2(L4_c0, L4_c1) + L4_a1 = _ak_pack_bf16x2(L4_c2, L4_c3) + L4_a2 = _ak_pack_bf16x2(L4_c4, L4_c5) + L4_a3 = _ak_pack_bf16x2(L4_c6, L4_c7) + L4_b0 = _ak_movmatrix_trans(L4_a0) + L4_b1 = _ak_movmatrix_trans(L4_a1) + L4_b2 = _ak_movmatrix_trans(L4_a2) + L4_b3 = _ak_movmatrix_trans(L4_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf + ) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 3: L^8, INV += INV*L^8 + L8_c0, L8_c1, L8_c2, L8_c3 = _ak_mma( + L4_a0, L4_a1, L4_a2, L4_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf + ) + L8_c4, L8_c5, L8_c6, L8_c7 = _ak_mma( + L4_a0, L4_a1, L4_a2, L4_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf + ) + L8_a0 = _ak_pack_bf16x2(L8_c0, L8_c1) + L8_a1 = _ak_pack_bf16x2(L8_c2, L8_c3) + L8_a2 = _ak_pack_bf16x2(L8_c4, L8_c5) + L8_a3 = _ak_pack_bf16x2(L8_c6, L8_c7) + L8_b0 = _ak_movmatrix_trans(L8_a0) + L8_b1 = _ak_movmatrix_trans(L8_a1) + L8_b2 = _ak_movmatrix_trans(L8_a2) + L8_b3 = _ak_movmatrix_trans(L8_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, L8_b0, L8_b1, _zf, _zf, _zf, _zf + ) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma( + INV_a0, INV_a1, INV_a2, INV_a3, L8_b2, L8_b3, _zf, _zf, _zf, _zf + ) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + + sAkk[r_off + gid, c_off + tid] = _ak_pack_bf16x2(INV_f0, INV_f1) + sAkk[r_off + gid + 8, c_off + tid] = _ak_pack_bf16x2(INV_f2, INV_f3) + sAkk[r_off + gid, c_off + 4 + tid] = _ak_pack_bf16x2(INV_f4, INV_f5) + sAkk[r_off + gid + 8, c_off + 4 + tid] = _ak_pack_bf16x2(INV_f6, INV_f7) + + +@dsl_user_op +def _ak_matmul_AB(sAkk, br_A, bc_A, br_B, bc_B, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 8 + rB = br_B * 16 + cB = bc_B * 8 + a0 = cutlass.Float32(sAkk[rA + gid, cA + tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 4 + tid]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 4 + tid]) + bA0 = cutlass.Float32(sAkk[rB + gid, cB + tid]) + bA1 = cutlass.Float32(sAkk[rB + gid + 8, cB + tid]) + bA2 = cutlass.Float32(sAkk[rB + gid, cB + 4 + tid]) + bA3 = cutlass.Float32(sAkk[rB + gid + 8, cB + 4 + tid]) + b0 = _ak_movmatrix_trans(bA0) + b1 = _ak_movmatrix_trans(bA1) + b2 = _ak_movmatrix_trans(bA2) + b3 = _ak_movmatrix_trans(bA3) + cn0_0, cn0_1, cn0_2, cn0_3 = _ak_mma(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf) + cn1_0, cn1_1, cn1_2, cn1_3 = _ak_mma(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _ak_chain_mma_B(sAkk, br_B, bc_B, a0, a1, a2, a3, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rB = br_B * 16 + cB = bc_B * 8 + bA0 = cutlass.Float32(sAkk[rB + gid, cB + tid]) + bA1 = cutlass.Float32(sAkk[rB + gid + 8, cB + tid]) + bA2 = cutlass.Float32(sAkk[rB + gid, cB + 4 + tid]) + bA3 = cutlass.Float32(sAkk[rB + gid + 8, cB + 4 + tid]) + b0 = _ak_movmatrix_trans(bA0) + b1 = _ak_movmatrix_trans(bA1) + b2 = _ak_movmatrix_trans(bA2) + b3 = _ak_movmatrix_trans(bA3) + cn0_0, cn0_1, cn0_2, cn0_3 = _ak_mma(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf) + cn1_0, cn1_1, cn1_2, cn1_3 = _ak_mma(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _ak_chain_mma_A(sAkk, br_A, bc_A, b0, b1, b2, b3, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 8 + a0 = cutlass.Float32(sAkk[rA + gid, cA + tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 4 + tid]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 4 + tid]) + cn0_0, cn0_1, cn0_2, cn0_3 = _ak_mma(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf) + cn1_0, cn1_1, cn1_2, cn1_3 = _ak_mma(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _ak_store_neg_C(sAkk, br, bc, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + r = br * 16 + c = bc * 8 + sAkk[r + gid, c + tid] = _ak_pack_bf16x2(-c0, -c1) + sAkk[r + gid + 8, c + tid] = _ak_pack_bf16x2(-c2, -c3) + sAkk[r + gid, c + 4 + tid] = _ak_pack_bf16x2(-c4, -c5) + sAkk[r + gid + 8, c + 4 + tid] = _ak_pack_bf16x2(-c6, -c7) + + +@dsl_user_op +def _ak_pack_C_to_A(c0, c1, c2, c3, c4, c5, c6, c7, *, loc=None, ip=None): + a0 = _ak_pack_bf16x2(c0, c1) + a1 = _ak_pack_bf16x2(c2, c3) + a2 = _ak_pack_bf16x2(c4, c5) + a3 = _ak_pack_bf16x2(c6, c7) + return a0, a1, a2, a3 + + +@dsl_user_op +def _ak_pack_C_to_B(c0, c1, c2, c3, c4, c5, c6, c7, *, loc=None, ip=None): + a0 = _ak_pack_bf16x2(c0, c1) + a1 = _ak_pack_bf16x2(c2, c3) + a2 = _ak_pack_bf16x2(c4, c5) + a3 = _ak_pack_bf16x2(c6, c7) + b0 = _ak_movmatrix_trans(a0) + b1 = _ak_movmatrix_trans(a1) + b2 = _ak_movmatrix_trans(a2) + b3 = _ak_movmatrix_trans(a3) + return b0, b1, b2, b3 + + +@dsl_user_op +def _ak_store_C_temp(sT, buf, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + sT[gid, 2 * tid, buf] = c0 + sT[gid, 2 * tid + 1, buf] = c1 + sT[gid + 8, 2 * tid, buf] = c2 + sT[gid + 8, 2 * tid + 1, buf] = c3 + sT[gid, 8 + 2 * tid, buf] = c4 + sT[gid, 8 + 2 * tid + 1, buf] = c5 + sT[gid + 8, 8 + 2 * tid, buf] = c6 + sT[gid + 8, 8 + 2 * tid + 1, buf] = c7 + + +@dsl_user_op +def _ak_load_C_temp(sT, buf, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + c0 = cutlass.Float32(sT[gid, 2 * tid, buf]) + c1 = cutlass.Float32(sT[gid, 2 * tid + 1, buf]) + c2 = cutlass.Float32(sT[gid + 8, 2 * tid, buf]) + c3 = cutlass.Float32(sT[gid + 8, 2 * tid + 1, buf]) + c4 = cutlass.Float32(sT[gid, 8 + 2 * tid, buf]) + c5 = cutlass.Float32(sT[gid, 8 + 2 * tid + 1, buf]) + c6 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid, buf]) + c7 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid + 1, buf]) + return c0, c1, c2, c3, c4, c5, c6, c7 + + +@dsl_user_op +def opaque_zero_from_work_id(*, loc=None, ip=None): + """ + Return 0 via opaque side-effectful asm (no inputs needed). + + Because has_side_effects=True, MLIR LICM treats this as having memory + effects and will NOT hoist it outside the for_generate loop. Any value + computed from the result (_oz) therefore appears loop-variant to LICM, + preventing get_slice() and scalar layout-invariant computations from being + hoisted to the kernel prologue. This keeps prologue register pressure < 64 + and eliminates the 440-byte stack frame that caused ~300-cycle L2 LDL + penalties per iteration. + """ + result = llvm.inline_asm( + T.i32(), + [], + "mov.b32 $0, 0;", # output = 0 (opaque to compiler constant-folding) + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Int32(result) + + +@cute.kernel +def fused_kernel123( + tma_atom_Q: cute.CopyAtom, + tma_tensor_Q: cute.Tensor, + tma_atom_K: cute.CopyAtom, + tma_tensor_K: cute.Tensor, + tma_atom_G: cute.CopyAtom, + tma_tensor_G: cute.Tensor, + mA_log: cute.Tensor, + mBeta: cute.Tensor, + scale: cutlass.Float32, + mKscaled: cute.Tensor, + mKg: cute.Tensor, + mQscaled: cute.Tensor, + mGkLast: cute.Tensor, + mAqk: cute.Tensor, # 4D (B, T, H, BT) — used by per-tile pure path + mAkk: cute.Tensor, # 4D (B, T, H, BT) — used by per-tile pure path + mAqk_v2: cute.Tensor, # 5D (B, T, H, BT/2, 2) — used by vec autovec non-pure path + mAkk_v2: cute.Tensor, # 5D (B, T, H, BT/2, 2) — used by vec autovec non-pure path + tiled_copy_qk_k1, + tiled_mma_k2, + tiled_copy_mma_A, + tiled_copy_mma_B, + tiled_copy_Gcum_norm, + tiled_copy_Gcum_gate, + qk_smem_layout, + g_smem_layout, + g_cumsum_layout, + num_chunks: cutlass.Int32, # runtime — shape-independent compile + num_heads: int, # baked (single variant in practice) + batch_size: cutlass.Int32, # runtime — shape-independent compile + mCuSeqlens: cute.Tensor, + mChunkIndices: cute.Tensor, + IS_VARLEN: cutlass.Constexpr[int], + mDtBias: cute.Tensor, + lower_bound: cutlass.Float32, + HAS_BIAS: cutlass.Constexpr[int], + USE_SAFE_GATE: cutlass.Constexpr[int], + VARLEN_PURE: cutlass.Constexpr[int] = 0, +): + block_id, _, _ = cute.arch.block_idx() + tidx = cute.arch.thread_idx()[0] + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + lane_id = tidx % 32 + + total_cgs_per_head = cutlass.Int32(0) + cgs_per_head = cutlass.Int32(0) + total_cgs = cutlass.Int32(0) + if IS_VARLEN: + total_cgs_per_head = (num_chunks + CHUNKS_PER_BLOCK - 1) // CHUNKS_PER_BLOCK + total_cgs = total_cgs_per_head * num_heads + else: + cgs_per_head = num_chunks // CHUNKS_PER_BLOCK + total_cgs = cgs_per_head * num_heads * batch_size + + # ===================================================================== + # SMEM allocation + # ===================================================================== + smem = cutlass.utils.SmemAllocator() + sQ = smem.allocate_tensor( + cutlass.BFloat16, qk_smem_layout.outer, 128, swizzle=qk_smem_layout.inner + ) + sK = smem.allocate_tensor( + cutlass.BFloat16, qk_smem_layout.outer, 128, swizzle=qk_smem_layout.inner + ) + sG = smem.allocate_tensor(cutlass.BFloat16, g_smem_layout, 128) + sGcum = smem.allocate_tensor(cutlass.Float32, g_cumsum_layout, 128) + partial_last_layout = cute.make_layout((K1_ROW_GROUPS, PARTIAL_COLS), stride=(PARTIAL_COLS, 1)) + sPartialLast = smem.allocate_tensor(cutlass.Float32, partial_last_layout, 128) + + # sAqk: 64x72 row-major (BT rows, BT+pad cols), same shape as sAkk. + # Each sub-tile (i_q, i_k) sits at SMEM rows [i_q*BC, (i_q+1)*BC) cols + # [i_k*BC, (i_k+1)*BC) — directly mirrors the 64x64 attention matrix. + aqk_tile_layout = cute.make_layout( + (BT, AQK_TILE_STRIDE, NUM_STAGES), stride=(AQK_TILE_STRIDE, 1, BT * AQK_TILE_STRIDE) + ) + sAqk = smem.allocate_tensor(cutlass.BFloat16, aqk_tile_layout, 128) + + akk_tile_layout = cute.make_layout( + (BT, AKK_STRIDE, NUM_STAGES), stride=(AKK_STRIDE, 1, BT * AKK_STRIDE) + ) + sAkk = smem.allocate_tensor(cutlass.BFloat16, akk_tile_layout, 128) + # sAkk_pkd / sTemp removed: akk_inv runs as a separate kernel call (chained back-to-back). + + # sBeta staging removed: its only consumer (K1 Pass 2b) moved beta fusion + # into the akk_inv epilogue, and the dead staging loop kept reading + # mBeta[chunk_start .. chunk_start+63] unguarded — OOB past the tensor + # end for the final partial chunk of a varlen batch. + + # ===================================================================== + # Mbarrier allocation & init + # ===================================================================== + tma_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + stage_reuse_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + k1_done_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + mma_done_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + store_done_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + + bytes_per_stage = BT * K_DIM * 2 * 3 + + if tidx == 0: + for s in range(NUM_STAGES): + cute.arch.mbarrier_init(tma_mbars + s, 1) + # TMA warp is waiter (not arriver) on stage_reuse; and it skips mma_done arrive + cute.arch.mbarrier_init(stage_reuse_mbars + s, (NUM_MMA_WARPS - 1) * 32) + cute.arch.mbarrier_init(k1_done_mbars + s, NUM_K1_TMA_WARPS * 32) + cute.arch.mbarrier_init(mma_done_mbars + s, (NUM_MMA_WARPS - 1) * 32) + cute.arch.mbarrier_init(store_done_mbars + s, NUM_STORE_WARPS * 32) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + + # ===================================================================== + # SMEM init: zero out sAqk and sAkk valid 64x64 region (cols 64..71 are + # padding for SMEM bank-conflict avoidance, never read or written by MMA + # 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) + # - 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). + # ===================================================================== + _warp_id_in_cta = tidx >> 5 # tidx // 32, range 0..31 + _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) + cute.arch.barrier() + + # ===================================================================== + # Pre-arrive (MMA warps only) + # stage_reuse_mbars: warp 0 waits before MMA arrives → pre-arrive all 12 MMA warps + # store_done_mbars: MMA waits before Store arrives → pre-arrive first 4 MMA warps + # ===================================================================== + if ( + warp_idx >= NUM_K1_TMA_WARPS + and warp_idx < NUM_K1_TMA_WARPS + NUM_MMA_WARPS + and warp_idx != TMA_WARP_ID + ): + mma_warp_tmp = warp_idx - NUM_K1_TMA_WARPS + for s in range(NUM_STAGES): + cute.arch.mbarrier_arrive(stage_reuse_mbars + s) + if mma_warp_tmp < NUM_STORE_WARPS: + cute.arch.mbarrier_arrive(store_done_mbars + s) + + # ================================================================= + # Persistent outer loop. Single for_generate at top level (required). + # Opaque asm barrier on work_id prevents MLIR LICM from hoisting + # get_slice() and scalar layout invariants to the kernel prologue, + # keeping register pressure < 64 and eliminating prologue spill. + # ================================================================= + for work_id in for_generate(block_id, total_cgs, NUM_SMS): + i_cg = cutlass.Int32(0) + i_h = cutlass.Int32(0) + i_b = cutlass.Int32(0) + chunk_base = cutlass.Int32(0) + if IS_VARLEN: + i_cg = work_id % total_cgs_per_head + i_h = work_id // total_cgs_per_head + i_b = cutlass.Int32(0) + chunk_base = i_cg * CHUNKS_PER_BLOCK + else: + i_cg = work_id % cgs_per_head + i_h = (work_id // cgs_per_head) % num_heads + i_b = work_id // (cgs_per_head * num_heads) + chunk_base = i_cg * CHUNKS_PER_BLOCK + + # Anti-LICM barrier: _oz is always 0 but appears to depend on work_id. + # Because this asm has side_effects=True, it stays inside the loop. + # Any value computed from _oz/_lane/_warp is also loop-variant from + # LICM's perspective → get_slice() and scalar invariants stay in-loop. + _oz = opaque_zero_from_work_id() + _lane = lane_id + _oz + _warp = warp_idx + _oz + + # ============================================================= + # Warps 0-15: Fused TMA + K1 + # ============================================================= + if warp_idx < NUM_K1_TMA_WARPS: + # Warp-layout invariants (scope-local → no cross-group register spill) + k1_warp = _warp + warp_row_group = k1_warp % K1_ROW_GROUPS + warp_col_group = k1_warp // K1_ROW_GROUPS + k1_row_start = warp_row_group * ROWS_PER_K1_WARP + col_base = warp_col_group * K1_COLS_PER_WARP + _lane * VEC + col_vec_idx = warp_col_group * (K1_COLS_PER_WARP // VEC) + _lane + cumsum_scale = cutlass.Float32(RCP_LN2) + thr_copy_k1 = tiled_copy_qk_k1.get_slice(_lane) + + rAcc = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + rPrefix = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + rGkLast = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + rKsOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + rQsOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + rKgOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + rGkOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + + # exp_A depends on i_h (changes per work unit) + exp_A = cute.exp(mA_log[i_h], fastmath=True) + + # Load dt_bias per (head, col) — broadcast across all rows + rBias = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + if HAS_BIAS: + for vi in cutlass.range_constexpr(VEC): + rBias[vi] = mDtBias[i_h, col_base + vi] + else: + for vi in cutlass.range_constexpr(VEC): + rBias[vi] = cutlass.Float32(0.0) + + # 3D TMA head slices (fixed for this work unit's head) + gQ_head = tma_tensor_Q[(None, None, i_h)] + gK_head = tma_tensor_K[(None, None, i_h)] + gG_head = tma_tensor_G[(None, None, i_h)] + + for chunk_iter in cutlass.range_constexpr(CHUNKS_PER_BLOCK): + cur_stage = chunk_iter % NUM_STAGES + cur_phase = chunk_iter // NUM_STAGES % 2 + chunk_idx = chunk_base + chunk_iter + chunk_start = cutlass.Int32(0) + ci_eos = cutlass.Int32(0) + if IS_VARLEN: + if chunk_idx < num_chunks: + _sid = cutlass.Int32(mChunkIndices[chunk_idx, 0]) + chunk_start = ( + cutlass.Int32(mCuSeqlens[_sid]) + + cutlass.Int32(mChunkIndices[chunk_idx, 1]) * BT + ) + ci_eos = cutlass.Int32(mCuSeqlens[_sid + 1]) + else: + chunk_start = chunk_idx * BT + + cute.arch.mbarrier_wait(tma_mbars + cur_stage, cur_phase) + + csG = sG[(None, None, cur_stage)] + csGcum = sGcum[(None, None, cur_stage)] + csQ = sQ[(None, None, cur_stage)] + csK = sK[(None, None, cur_stage)] + rGact = cute.make_rmem_tensor( + cute.make_layout((ROWS_PER_K1_WARP, VEC)), cutlass.Float32 + ) + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = cutlass.Float32(0.0) + + for ri in cutlass.range_constexpr(ROWS_PER_K1_WARP): + row = k1_row_start + ri + for vi in cutlass.range_constexpr(VEC): + c = col_base + vi + g_val = csG[row, c].to(cutlass.Float32) + if HAS_BIAS: + g_val = g_val + rBias[vi] + g_activated = cutlass.Float32(0.0) + if USE_SAFE_GATE: + sigmoid_g = fast_rcp( + cutlass.Float32(1.0) + + cute.exp2(-exp_A * g_val * LOG2E, fastmath=True) + ) + g_activated = lower_bound * sigmoid_g + else: + softplus_g = ( + cute.log2( + cutlass.Float32(1.0) + cute.exp2(g_val * LOG2E, fastmath=True), + fastmath=True, + ) + * LN2 + ) + g_activated = -exp_A * softplus_g + # Varlen: zero gate for out-of-bounds rows so cumsum + # stays flat beyond the last valid position. + # VARLEN_PURE=1 elides this at compile time — caller + # guarantees all seq lengths are multiples of BT so no + # chunk has OOB rows. + if IS_VARLEN and not VARLEN_PURE: + if chunk_start + row >= ci_eos: + g_activated = cutlass.Float32(0.0) + rGact[ri, vi] = g_activated + rAcc[vi] = rAcc[vi] + g_activated + + for vi in cutlass.range_constexpr(VEC): + sPartialLast[warp_row_group, col_base + vi] = rAcc[vi] + + k1_internal_barrier() + + prefix_col_start = k1_warp * PARTIAL_COLS_PER_WARP + row_in_prefix = lane_id % K1_ROW_GROUPS + col_in_group = lane_id // K1_ROW_GROUPS + + for j in cutlass.range_constexpr(PARTIAL_COLS_PER_WARP // 4): + col = prefix_col_start + j * 4 + col_in_group + val = cutlass.Float32(sPartialLast[row_in_prefix, col]) + tmp = cute.arch.shuffle_sync_up(val, 1, mask=-1, mask_and_clamp=SHFL_W8_CLAMP) + if row_in_prefix >= 1: + val = val + tmp + tmp = cute.arch.shuffle_sync_up(val, 2, mask=-1, mask_and_clamp=SHFL_W8_CLAMP) + if row_in_prefix >= 2: + val = val + tmp + tmp = cute.arch.shuffle_sync_up(val, 4, mask=-1, mask_and_clamp=SHFL_W8_CLAMP) + if row_in_prefix >= 4: + val = val + tmp + sPartialLast[row_in_prefix, col] = val + + k1_internal_barrier() + + for vi in cutlass.range_constexpr(VEC): + rGkLast[vi] = sPartialLast[K1_ROW_GROUPS - 1, col_base + vi] + + for vi in cutlass.range_constexpr(VEC): + rPrefix[vi] = cutlass.Float32(0.0) + if warp_row_group > 0: + for vi in cutlass.range_constexpr(VEC): + rPrefix[vi] = sPartialLast[warp_row_group - 1, col_base + vi] + + # ---- Pass 2a: ONLY cumsum + write csGcum (critical path, minimal work) ---- + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = rPrefix[vi] + + for ri in cutlass.range_constexpr(ROWS_PER_K1_WARP): + row = k1_row_start + ri + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = rAcc[vi] + rGact[ri, vi] + csGcum[row, col_base + vi] = rAcc[vi] * cumsum_scale + + # Signal MMA early: csGcum is ready + cute.arch.mbarrier_arrive(k1_done_mbars + cur_stage) + + # ---- Pass 2b: recompute + write GMEM (overlaps with MMA, off critical path) ---- + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = rPrefix[vi] + + for ri in cutlass.range_constexpr(ROWS_PER_K1_WARP): + row = k1_row_start + ri + t = chunk_start + row + + sK_tile = cute.local_tile( + csK, tiler=(1, K1_COLS_PER_WARP), coord=(row, warp_col_group) + ) + tCsK = thr_copy_k1.partition_S(sK_tile) + tCrK = cute.make_fragment_like(tCsK) + cute.copy(tiled_copy_qk_k1, tCsK, thr_copy_k1.retile(tCrK)) + + sQ_tile = cute.local_tile( + csQ, tiler=(1, K1_COLS_PER_WARP), coord=(row, warp_col_group) + ) + tCsQ = thr_copy_k1.partition_S(sQ_tile) + tCrQ = cute.make_fragment_like(tCsQ) + cute.copy(tiled_copy_qk_k1, tCsQ, thr_copy_k1.retile(tCrQ)) + + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = rAcc[vi] + rGact[ri, vi] + cs = rAcc[vi] * cumsum_scale + + k_val = tCrK[vi].to(cutlass.Float32) + q_val = tCrQ[vi].to(cutlass.Float32) + + exp2_cs = cute.exp2(cs, fastmath=True) + gk_last_cs = rGkLast[vi] * cumsum_scale + exp2_kg = cute.exp2(gk_last_cs - cs, fastmath=True) + + rKsOut[vi] = (k_val * exp2_cs).to(cutlass.BFloat16) + rQsOut[vi] = (q_val * exp2_cs * scale).to(cutlass.BFloat16) + rKgOut[vi] = (k_val * exp2_kg).to(cutlass.BFloat16) + + if IS_VARLEN and not VARLEN_PURE: + if t < ci_eos: + cute.autovec_copy(rKsOut, mKscaled[i_b, t, i_h, col_vec_idx, None]) + cute.autovec_copy(rQsOut, mQscaled[i_b, t, i_h, col_vec_idx, None]) + cute.autovec_copy(rKgOut, mKg[i_b, t, i_h, col_vec_idx, None]) + elif IS_VARLEN and VARLEN_PURE: + if chunk_idx < num_chunks: + cute.autovec_copy(rKsOut, mKscaled[i_b, t, i_h, col_vec_idx, None]) + cute.autovec_copy(rQsOut, mQscaled[i_b, t, i_h, col_vec_idx, None]) + cute.autovec_copy(rKgOut, mKg[i_b, t, i_h, col_vec_idx, None]) + else: + cute.autovec_copy(rKsOut, mKscaled[i_b, t, i_h, col_vec_idx, None]) + cute.autovec_copy(rQsOut, mQscaled[i_b, t, i_h, col_vec_idx, None]) + cute.autovec_copy(rKgOut, mKg[i_b, t, i_h, col_vec_idx, None]) + + if warp_row_group == 0: + for vi in cutlass.range_constexpr(VEC): + rGkOut[vi] = cute.exp2(rGkLast[vi] * cumsum_scale, fastmath=True) + if IS_VARLEN: + if ci_eos > cutlass.Int32(0): + cute.autovec_copy( + rGkOut, mGkLast[i_b, chunk_idx, i_h, col_vec_idx, None] + ) + else: + cute.autovec_copy(rGkOut, mGkLast[i_b, chunk_idx, i_h, col_vec_idx, None]) + + # ============================================================= + # Warp 26 (TMA_WARP_ID): dedicated TMA producer. + # Waits stage_reuse (gated by MMA arrives), issues TMA for Q/K/G, + # signals tma_mbar. Decouples MMA -> TMA dependency from K1 compute. + # ============================================================= + if warp_idx == TMA_WARP_ID: + gQ_head = tma_tensor_Q[(None, None, i_h)] + gK_head = tma_tensor_K[(None, None, i_h)] + gG_head = tma_tensor_G[(None, None, i_h)] + + # Prefetch chunk 0 -> stage 0 (stage_reuse[0] pre-arrived) + pf_cs = cutlass.Int32(0) + if IS_VARLEN: + pf_seq_id_0 = cutlass.Int32(mChunkIndices[chunk_base, 0]) + pf_local_0 = cutlass.Int32(mChunkIndices[chunk_base, 1]) + pf_bos_0 = cutlass.Int32(mCuSeqlens[pf_seq_id_0]) + pf_cs = pf_bos_0 + pf_local_0 * BT + else: + pf_cs = i_b * num_chunks * BT + chunk_base * BT + cute.arch.mbarrier_wait(stage_reuse_mbars, 0) + if lane_id == 0: + cute.arch.mbarrier_expect_tx(tma_mbars, bytes_per_stage) + sQ_pf = sQ[(None, None, 0)] + gQ_pf = cute.local_tile(cute.domain_offset((pf_cs, 0), gQ_head), (BT, K_DIM), (0, 0)) + ts_pf, tg_pf = cpasync.tma_partition( + tma_atom_Q, + 0, + cute.make_layout(1), + cute.group_modes(sQ_pf, 0, 2), + cute.group_modes(gQ_pf, 0, 2), + ) + cute.copy(tma_atom_Q, tg_pf, ts_pf, tma_bar_ptr=tma_mbars) + sK_pf = sK[(None, None, 0)] + gK_pf = cute.local_tile(cute.domain_offset((pf_cs, 0), gK_head), (BT, K_DIM), (0, 0)) + ts_pf, tg_pf = cpasync.tma_partition( + tma_atom_K, + 0, + cute.make_layout(1), + cute.group_modes(sK_pf, 0, 2), + cute.group_modes(gK_pf, 0, 2), + ) + cute.copy(tma_atom_K, tg_pf, ts_pf, tma_bar_ptr=tma_mbars) + sG_pf = sG[(None, None, 0)] + gG_pf = cute.local_tile(cute.domain_offset((pf_cs, 0), gG_head), (BT, K_DIM), (0, 0)) + ts_pf, tg_pf = cpasync.tma_partition( + tma_atom_G, + 0, + cute.make_layout(1), + cute.group_modes(sG_pf, 0, 2), + cute.group_modes(gG_pf, 0, 2), + ) + cute.copy(tma_atom_G, tg_pf, ts_pf, tma_bar_ptr=tma_mbars) + if lane_id == 0: + cute.arch.mbarrier_arrive(tma_mbars) + + # Issue TMAs for chunks 1..CHUNKS_PER_BLOCK-1 + for next_i in cutlass.range_constexpr(1, CHUNKS_PER_BLOCK): + next_stage = next_i % NUM_STAGES + next_phase = next_i // NUM_STAGES % 2 + next_cs = cutlass.Int32(0) + if IS_VARLEN: + next_chunk_idx = chunk_base + next_i + if next_chunk_idx < num_chunks: + _nsid = cutlass.Int32(mChunkIndices[next_chunk_idx, 0]) + next_cs = ( + cutlass.Int32(mCuSeqlens[_nsid]) + + cutlass.Int32(mChunkIndices[next_chunk_idx, 1]) * BT + ) + else: + next_cs = i_b * num_chunks * BT + (chunk_base + next_i) * BT + cute.arch.mbarrier_wait(stage_reuse_mbars + next_stage, next_phase) + if lane_id == 0: + cute.arch.mbarrier_expect_tx(tma_mbars + next_stage, bytes_per_stage) + sQ_ns = sQ[(None, None, next_stage)] + gQ_ns = cute.local_tile( + cute.domain_offset((next_cs, 0), gQ_head), (BT, K_DIM), (0, 0) + ) + ts_ns, tg_ns = cpasync.tma_partition( + tma_atom_Q, + 0, + cute.make_layout(1), + cute.group_modes(sQ_ns, 0, 2), + cute.group_modes(gQ_ns, 0, 2), + ) + cute.copy(tma_atom_Q, tg_ns, ts_ns, tma_bar_ptr=tma_mbars + next_stage) + sK_ns = sK[(None, None, next_stage)] + gK_ns = cute.local_tile( + cute.domain_offset((next_cs, 0), gK_head), (BT, K_DIM), (0, 0) + ) + ts_ns, tg_ns = cpasync.tma_partition( + tma_atom_K, + 0, + cute.make_layout(1), + cute.group_modes(sK_ns, 0, 2), + cute.group_modes(gK_ns, 0, 2), + ) + cute.copy(tma_atom_K, tg_ns, ts_ns, tma_bar_ptr=tma_mbars + next_stage) + sG_ns = sG[(None, None, next_stage)] + gG_ns = cute.local_tile( + cute.domain_offset((next_cs, 0), gG_head), (BT, K_DIM), (0, 0) + ) + ts_ns, tg_ns = cpasync.tma_partition( + tma_atom_G, + 0, + cute.make_layout(1), + cute.group_modes(sG_ns, 0, 2), + cute.group_modes(gG_ns, 0, 2), + ) + cute.copy(tma_atom_G, tg_ns, ts_ns, tma_bar_ptr=tma_mbars + next_stage) + if lane_id == 0: + cute.arch.mbarrier_arrive(tma_mbars + next_stage) + + # ============================================================= + # Warps 16-27 (excluding TMA_WARP_ID=26): K2 MMA Compute + # ============================================================= + if ( + warp_idx >= NUM_K1_TMA_WARPS + and warp_idx < NUM_K1_TMA_WARPS + NUM_MMA_WARPS + and warp_idx != TMA_WARP_ID + ): + # Warp-layout invariants (scope-local → no cross-group register spill) + _tid_in_group = _lane % 4 + _group_id = _lane // 4 + mma_warp = _warp - NUM_K1_TMA_WARPS + my_i_q = cutlass.Int32(0) + my_i_k = cutlass.Int32(0) + if mma_warp < 1: + my_i_q = cutlass.Int32(0) + my_i_k = mma_warp + elif mma_warp < 3: + my_i_q = cutlass.Int32(1) + my_i_k = mma_warp - 1 + elif mma_warp < 6: + my_i_q = cutlass.Int32(2) + my_i_k = mma_warp - 3 + elif mma_warp < NUM_MMA_ACTIVE: + my_i_q = cutlass.Int32(3) + my_i_k = mma_warp - 6 + q_row_base = my_i_q * BC + k_row_base = my_i_k * BC + akk_row_base = k_row_base + akk_col_base = q_row_base + norm_row = q_row_base + if my_i_q == my_i_k: + norm_row = q_row_base + cutlass.Int32(BC // 2) + row0 = _group_id + row1 = _group_id + 8 + col0 = _tid_in_group * 2 + col1 = _tid_in_group * 2 + 1 + col2 = 8 + _tid_in_group * 2 + col3 = 8 + _tid_in_group * 2 + 1 + thr_mma = tiled_mma_k2.get_slice(_lane) + thr_copy_A = tiled_copy_mma_A.get_slice(_lane) + thr_copy_B = tiled_copy_mma_B.get_slice(_lane) + thr_copy_Gn = tiled_copy_Gcum_norm.get_slice(_tid_in_group) + thr_copy_Ggate = tiled_copy_Gcum_gate.get_slice(_lane) + + for chunk_iter in cutlass.range_constexpr(CHUNKS_PER_BLOCK): + s = chunk_iter % NUM_STAGES + phase = chunk_iter // NUM_STAGES % 2 + chunk_idx = chunk_base + chunk_iter + chunk_start = cutlass.Int32(0) + mma_eos = cutlass.Int32(0) + if IS_VARLEN: + if chunk_idx < num_chunks: + _sid = cutlass.Int32(mChunkIndices[chunk_idx, 0]) + chunk_start = ( + cutlass.Int32(mCuSeqlens[_sid]) + + cutlass.Int32(mChunkIndices[chunk_idx, 1]) * BT + ) + mma_eos = cutlass.Int32(mCuSeqlens[_sid + 1]) + else: + chunk_start = chunk_idx * BT + + cute.arch.mbarrier_wait(k1_done_mbars + s, phase) + cute.arch.mbarrier_wait(store_done_mbars + s, phase) + + if mma_warp < NUM_MMA_ACTIVE: + csQ = sQ[(None, None, s)] + csK = sK[(None, None, s)] + csGcum = sGcum[(None, None, s)] + csAqk = sAqk[(None, None, s)] + csAkk = sAkk[(None, None, s)] + + _z = cutlass.Float32(0.0) + + # Varlen non-pure: a partial chunk's rows past the + # sequence end must not be read — for the batch's final + # chunk they lie past the end of the beta tensor + # entirely (OOB read; NaN/IMA under memory pressure). + # Their A-row contributions are discarded downstream, so + # beta=0 is safe. Eqlen and VARLEN_PURE inputs are + # padded/aligned upstream — load unconditionally there. + beta_row0 = _z + 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) + 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) + 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) + + 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 + acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = _z, _z, _z, _z + acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = _z, _z, _z, _z + + # bf16 m16n8k16 MMA: each k_block covers k=16. + for k_block in cutlass.range_constexpr(NUM_MMA_K_TILES): + # ---- Load Q/Kq bf16 fragments (16x16, 8 bf16/thread) ---- + sQ_tile = cute.local_tile(csQ, tiler=(16, 16), coord=(my_i_q, k_block)) + tCrQ = tiled_mma_k2.make_fragment_A(thr_mma.partition_A(sQ_tile)) + cute.copy( + tiled_copy_mma_A, + thr_copy_A.partition_S(sQ_tile), + thr_copy_A.retile(tCrQ), + ) + + sKq_tile = cute.local_tile(csK, tiler=(16, 16), coord=(my_i_q, k_block)) + tCrKq = tiled_mma_k2.make_fragment_A(thr_mma.partition_A(sKq_tile)) + cute.copy( + tiled_copy_mma_A, + thr_copy_A.partition_S(sKq_tile), + thr_copy_A.retile(tCrKq), + ) + + # ---- Issue K n0/n1 LDSMs early for better ILP ---- + sK_tile_n0 = cute.local_tile( + csK, tiler=(8, 16), coord=(my_i_k * 2, k_block) + ) + tCrK_n0 = tiled_mma_k2.make_fragment_B(thr_mma.partition_B(sK_tile_n0)) + cute.copy( + tiled_copy_mma_B, + thr_copy_B.partition_S(sK_tile_n0), + thr_copy_B.retile(tCrK_n0), + ) + + sK_tile_n1 = cute.local_tile( + csK, tiler=(8, 16), coord=(my_i_k * 2 + 1, k_block) + ) + tCrK_n1 = tiled_mma_k2.make_fragment_B(thr_mma.partition_B(sK_tile_n1)) + cute.copy( + tiled_copy_mma_B, + thr_copy_B.partition_S(sK_tile_n1), + thr_copy_B.retile(tCrK_n1), + ) + + # ---- Gate norm (2x k=8 covers k=16) ---- + sGn_a = cute.local_tile(csGcum, tiler=(1, 8), coord=(norm_row, k_block * 2)) + tCsGn_a = thr_copy_Gn.partition_S(sGn_a) + tCrGn_a = cute.make_fragment_like(tCsGn_a, cutlass.Float32) + cute.copy(tiled_copy_Gcum_norm, tCsGn_a, thr_copy_Gn.retile(tCrGn_a)) + gn_a0 = tCrGn_a[0] + gn_a1 = tCrGn_a[1] + + sGn_b = cute.local_tile( + csGcum, tiler=(1, 8), coord=(norm_row, k_block * 2 + 1) + ) + tCsGn_b = thr_copy_Gn.partition_S(sGn_b) + tCrGn_b = cute.make_fragment_like(tCsGn_b, cutlass.Float32) + cute.copy(tiled_copy_Gcum_norm, tCsGn_b, thr_copy_Gn.retile(tCrGn_b)) + gn_b0 = tCrGn_b[0] + gn_b1 = tCrGn_b[1] + + # ---- Gate Q (2x (16,8) partition_C covers m=16,k=16) ---- + sGq_a = cute.local_tile(csGcum, tiler=(16, 8), coord=(my_i_q, k_block * 2)) + tCrGq_a = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGq_a)) + cute.copy( + tiled_copy_Gcum_gate, + thr_copy_Ggate.partition_S(sGq_a), + thr_copy_Ggate.retile(tCrGq_a), + ) + + sGq_b = cute.local_tile( + csGcum, tiler=(16, 8), coord=(my_i_q, k_block * 2 + 1) + ) + tCrGq_b = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGq_b)) + cute.copy( + tiled_copy_Gcum_gate, + thr_copy_Ggate.partition_S(sGq_b), + thr_copy_Ggate.retile(tCrGq_b), + ) + + # 8 Q gate values per thread (matching A bf16 m16n8k16 layout): + # first half k=0..7 (tCrGq_a): a0=(r0,c0) a1=(r0,c0+1) a2=(r0+8,c0) a3=(r0+8,c0+1) + # second half k=8..15 (tCrGq_b): a4=(r0,c0+8) a5=(r0,c0+9) a6=(r0+8,c0+8) a7=(r0+8,c0+9) + gate_q_0 = cute.exp2(tCrGq_a[0] - gn_a0, fastmath=True) + gate_q_1 = cute.exp2(tCrGq_a[1] - gn_a1, fastmath=True) + gate_q_2 = cute.exp2(tCrGq_a[2] - gn_a0, fastmath=True) + gate_q_3 = cute.exp2(tCrGq_a[3] - gn_a1, fastmath=True) + gate_q_4 = cute.exp2(tCrGq_b[0] - gn_b0, fastmath=True) + gate_q_5 = cute.exp2(tCrGq_b[1] - gn_b1, fastmath=True) + gate_q_6 = cute.exp2(tCrGq_b[2] - gn_b0, fastmath=True) + gate_q_7 = cute.exp2(tCrGq_b[3] - gn_b1, fastmath=True) + + # qa fp32 = Q*gate (8 per thread). tCrQ indexing assumption: + # [0..3] first k-chunk (k=0..7), [4..7] second k-chunk (k=8..15). + qa0 = tCrQ[0].to(cutlass.Float32) * gate_q_0 + qa1 = tCrQ[1].to(cutlass.Float32) * gate_q_1 + qa2 = tCrQ[2].to(cutlass.Float32) * gate_q_2 + qa3 = tCrQ[3].to(cutlass.Float32) * gate_q_3 + qa4 = tCrQ[4].to(cutlass.Float32) * gate_q_4 + qa5 = tCrQ[5].to(cutlass.Float32) * gate_q_5 + qa6 = tCrQ[6].to(cutlass.Float32) * gate_q_6 + qa7 = tCrQ[7].to(cutlass.Float32) * gate_q_7 + + # Pack fp32 qa -> 4 u32 bf16x2 for MMA A (k-adjacent pairs per u32) + qa_u32_0 = pack_bf16x2_f32(qa1, qa0) # reg0 = [qa0 lo | qa1 hi] + qa_u32_1 = pack_bf16x2_f32(qa3, qa2) + qa_u32_2 = pack_bf16x2_f32(qa5, qa4) + qa_u32_3 = pack_bf16x2_f32(qa7, qa6) + + ka0 = tCrKq[0].to(cutlass.Float32) * gate_q_0 + ka1 = tCrKq[1].to(cutlass.Float32) * gate_q_1 + ka2 = tCrKq[2].to(cutlass.Float32) * gate_q_2 + ka3 = tCrKq[3].to(cutlass.Float32) * gate_q_3 + ka4 = tCrKq[4].to(cutlass.Float32) * gate_q_4 + ka5 = tCrKq[5].to(cutlass.Float32) * gate_q_5 + ka6 = tCrKq[6].to(cutlass.Float32) * gate_q_6 + ka7 = tCrKq[7].to(cutlass.Float32) * gate_q_7 + + ka_u32_0 = pack_bf16x2_f32(ka1, ka0) + ka_u32_1 = pack_bf16x2_f32(ka3, ka2) + ka_u32_2 = pack_bf16x2_f32(ka5, ka4) + ka_u32_3 = pack_bf16x2_f32(ka7, ka6) + + # ---- Gate K (2x (16,8) partition_C covers m=16,k=16) ---- + sGk_a = cute.local_tile(csGcum, tiler=(16, 8), coord=(my_i_k, k_block * 2)) + tCrGk_a = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGk_a)) + cute.copy( + tiled_copy_Gcum_gate, + thr_copy_Ggate.partition_S(sGk_a), + thr_copy_Ggate.retile(tCrGk_a), + ) + + sGk_b = cute.local_tile( + csGcum, tiler=(16, 8), coord=(my_i_k, k_block * 2 + 1) + ) + tCrGk_b = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGk_b)) + cute.copy( + tiled_copy_Gcum_gate, + thr_copy_Ggate.partition_S(sGk_b), + thr_copy_Ggate.retile(tCrGk_b), + ) + + # n0 uses rows 0..7 of (16,*) tile (tCrGk_*[0,1]) + # n1 uses rows 8..15 of (16,*) tile (tCrGk_*[2,3]) + gk_n0_0 = cute.exp2(gn_a0 - tCrGk_a[0], fastmath=True) + gk_n0_1 = cute.exp2(gn_a1 - tCrGk_a[1], fastmath=True) + gk_n0_2 = cute.exp2(gn_b0 - tCrGk_b[0], fastmath=True) + gk_n0_3 = cute.exp2(gn_b1 - tCrGk_b[1], fastmath=True) + + gk_n1_0 = cute.exp2(gn_a0 - tCrGk_a[2], fastmath=True) + gk_n1_1 = cute.exp2(gn_a1 - tCrGk_a[3], fastmath=True) + gk_n1_2 = cute.exp2(gn_b0 - tCrGk_b[2], fastmath=True) + gk_n1_3 = cute.exp2(gn_b1 - tCrGk_b[3], fastmath=True) + + # tCrK_n0 bf16 fragment: 4 elems/thread at (n_row, c0), (n_row, c0+1), + # (n_row, c0+8), (n_row, c0+9) — k-adjacent pairs + k_n0_b0 = tCrK_n0[0].to(cutlass.Float32) * gk_n0_0 + k_n0_b1 = tCrK_n0[1].to(cutlass.Float32) * gk_n0_1 + k_n0_b2 = tCrK_n0[2].to(cutlass.Float32) * gk_n0_2 + k_n0_b3 = tCrK_n0[3].to(cutlass.Float32) * gk_n0_3 + + k_n1_b0 = tCrK_n1[0].to(cutlass.Float32) * gk_n1_0 + k_n1_b1 = tCrK_n1[1].to(cutlass.Float32) * gk_n1_1 + k_n1_b2 = tCrK_n1[2].to(cutlass.Float32) * gk_n1_2 + k_n1_b3 = tCrK_n1[3].to(cutlass.Float32) * gk_n1_3 + + # Pack fp32 k -> 2 u32 bf16x2 for MMA B + k_n0_u32_0 = pack_bf16x2_f32(k_n0_b1, k_n0_b0) + k_n0_u32_1 = pack_bf16x2_f32(k_n0_b3, k_n0_b2) + k_n1_u32_0 = pack_bf16x2_f32(k_n1_b1, k_n1_b0) + k_n1_u32_1 = pack_bf16x2_f32(k_n1_b3, k_n1_b2) + + # ---- 4 bf16 MMA calls ---- + acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = mma_bf16_m16n8k16( + qa_u32_0, + qa_u32_1, + qa_u32_2, + qa_u32_3, + k_n0_u32_0, + k_n0_u32_1, + acc_aqk_n0_0, + acc_aqk_n0_1, + acc_aqk_n0_2, + acc_aqk_n0_3, + ) + acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = mma_bf16_m16n8k16( + qa_u32_0, + qa_u32_1, + qa_u32_2, + qa_u32_3, + k_n1_u32_0, + k_n1_u32_1, + acc_aqk_n1_0, + acc_aqk_n1_1, + acc_aqk_n1_2, + acc_aqk_n1_3, + ) + acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = mma_bf16_m16n8k16( + ka_u32_0, + ka_u32_1, + ka_u32_2, + ka_u32_3, + k_n0_u32_0, + k_n0_u32_1, + acc_akk_n0_0, + acc_akk_n0_1, + acc_akk_n0_2, + acc_akk_n0_3, + ) + acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = mma_bf16_m16n8k16( + ka_u32_0, + ka_u32_1, + ka_u32_2, + ka_u32_3, + k_n1_u32_0, + k_n1_u32_1, + acc_akk_n1_0, + acc_akk_n1_1, + acc_akk_n1_2, + acc_akk_n1_3, + ) + + # sQ/sK/sG reads done, signal TMA before SMEM writes + cute.arch.mbarrier_arrive(stage_reuse_mbars + s) + + # Dual-path MMA write (constexpr-gated): + # - non-pure: apply causal + diag=1 inline so SMEM matches + # final GMEM layout (pairs with row-major vec autovec + # store warp that does no causal). + # - pure: write all 16x16 unconditionally (pairs with + # per-tile store warp that applies causal + diag in + # store; this is the original baseline behavior — no + # extra MMA-write cost). + _z16 = cutlass.BFloat16(0.0) + _one16 = cutlass.BFloat16(1.0) + if IS_VARLEN and not VARLEN_PURE: + if my_i_q == my_i_k: + # Diagonal sub-tile: causal-mask sAqk and write + # diag=1 / strict-lower=MMA*beta / strict-upper=0 + # to sAkk so SMEM is in final form. + _v_q00 = (acc_aqk_n0_0 * scale).to(cutlass.BFloat16) + if row0 < col0: + _v_q00 = _z16 + csAqk[q_row_base + row0, k_row_base + col0] = _v_q00 + _v_q01 = (acc_aqk_n0_1 * scale).to(cutlass.BFloat16) + if row0 < col1: + _v_q01 = _z16 + csAqk[q_row_base + row0, k_row_base + col1] = _v_q01 + _v_q02 = (acc_aqk_n0_2 * scale).to(cutlass.BFloat16) + if row1 < col0: + _v_q02 = _z16 + csAqk[q_row_base + row1, k_row_base + col0] = _v_q02 + _v_q03 = (acc_aqk_n0_3 * scale).to(cutlass.BFloat16) + if row1 < col1: + _v_q03 = _z16 + csAqk[q_row_base + row1, k_row_base + col1] = _v_q03 + _v_q04 = (acc_aqk_n1_0 * scale).to(cutlass.BFloat16) + if row0 < col2: + _v_q04 = _z16 + csAqk[q_row_base + row0, k_row_base + col2] = _v_q04 + _v_q05 = (acc_aqk_n1_1 * scale).to(cutlass.BFloat16) + if row0 < col3: + _v_q05 = _z16 + csAqk[q_row_base + row0, k_row_base + col3] = _v_q05 + _v_q06 = (acc_aqk_n1_2 * scale).to(cutlass.BFloat16) + if row1 < col2: + _v_q06 = _z16 + csAqk[q_row_base + row1, k_row_base + col2] = _v_q06 + _v_q07 = (acc_aqk_n1_3 * scale).to(cutlass.BFloat16) + if row1 < col3: + _v_q07 = _z16 + csAqk[q_row_base + row1, k_row_base + col3] = _v_q07 + + _v_k00 = (acc_akk_n0_0 * beta_row0).to(cutlass.BFloat16) + if row0 == col0: + _v_k00 = _one16 + if row0 < col0: + _v_k00 = _z16 + csAkk[akk_row_base + row0, akk_col_base + col0] = _v_k00 + _v_k01 = (acc_akk_n0_1 * beta_row0).to(cutlass.BFloat16) + if row0 == col1: + _v_k01 = _one16 + if row0 < col1: + _v_k01 = _z16 + csAkk[akk_row_base + row0, akk_col_base + col1] = _v_k01 + _v_k02 = (acc_akk_n0_2 * beta_row1).to(cutlass.BFloat16) + if row1 == col0: + _v_k02 = _one16 + if row1 < col0: + _v_k02 = _z16 + csAkk[akk_row_base + row1, akk_col_base + col0] = _v_k02 + _v_k03 = (acc_akk_n0_3 * beta_row1).to(cutlass.BFloat16) + if row1 == col1: + _v_k03 = _one16 + if row1 < col1: + _v_k03 = _z16 + csAkk[akk_row_base + row1, akk_col_base + col1] = _v_k03 + _v_k04 = (acc_akk_n1_0 * beta_row0).to(cutlass.BFloat16) + if row0 == col2: + _v_k04 = _one16 + if row0 < col2: + _v_k04 = _z16 + csAkk[akk_row_base + row0, akk_col_base + col2] = _v_k04 + _v_k05 = (acc_akk_n1_1 * beta_row0).to(cutlass.BFloat16) + if row0 == col3: + _v_k05 = _one16 + if row0 < col3: + _v_k05 = _z16 + csAkk[akk_row_base + row0, akk_col_base + col3] = _v_k05 + _v_k06 = (acc_akk_n1_2 * beta_row1).to(cutlass.BFloat16) + if row1 == col2: + _v_k06 = _one16 + if row1 < col2: + _v_k06 = _z16 + csAkk[akk_row_base + row1, akk_col_base + col2] = _v_k06 + _v_k07 = (acc_akk_n1_3 * beta_row1).to(cutlass.BFloat16) + if row1 == col3: + _v_k07 = _one16 + if row1 < col3: + _v_k07 = _z16 + csAkk[akk_row_base + row1, akk_col_base + col3] = _v_k07 + else: + # Non-diag (i_q > i_k): write all 16x16 unchanged. + csAqk[q_row_base + row0, k_row_base + col0] = (acc_aqk_n0_0 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row0, k_row_base + col1] = (acc_aqk_n0_1 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row1, k_row_base + col0] = (acc_aqk_n0_2 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row1, k_row_base + col1] = (acc_aqk_n0_3 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row0, k_row_base + col2] = (acc_aqk_n1_0 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row0, k_row_base + col3] = (acc_aqk_n1_1 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row1, k_row_base + col2] = (acc_aqk_n1_2 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row1, k_row_base + col3] = (acc_aqk_n1_3 * scale).to( + cutlass.BFloat16 + ) + csAkk[akk_row_base + row0, akk_col_base + col0] = ( + acc_akk_n0_0 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col1] = ( + acc_akk_n0_1 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col0] = ( + acc_akk_n0_2 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col1] = ( + acc_akk_n0_3 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col2] = ( + acc_akk_n1_0 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col3] = ( + acc_akk_n1_1 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col2] = ( + acc_akk_n1_2 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col3] = ( + acc_akk_n1_3 * beta_row1 + ).to(cutlass.BFloat16) + else: + # PURE: write all 16x16 unconditionally — store warp + # applies causal+diag in its per-tile loop. This is + # the original baseline behavior (no extra MMA cost). + csAqk[q_row_base + row0, k_row_base + col0] = (acc_aqk_n0_0 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row0, k_row_base + col1] = (acc_aqk_n0_1 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row1, k_row_base + col0] = (acc_aqk_n0_2 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row1, k_row_base + col1] = (acc_aqk_n0_3 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row0, k_row_base + col2] = (acc_aqk_n1_0 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row0, k_row_base + col3] = (acc_aqk_n1_1 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row1, k_row_base + col2] = (acc_aqk_n1_2 * scale).to( + cutlass.BFloat16 + ) + csAqk[q_row_base + row1, k_row_base + col3] = (acc_aqk_n1_3 * scale).to( + cutlass.BFloat16 + ) + csAkk[akk_row_base + row0, akk_col_base + col0] = ( + acc_akk_n0_0 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col1] = ( + acc_akk_n0_1 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col0] = ( + acc_akk_n0_2 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col1] = ( + acc_akk_n0_3 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col2] = ( + acc_akk_n1_0 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col3] = ( + acc_akk_n1_1 * beta_row0 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col2] = ( + acc_akk_n1_2 * beta_row1 + ).to(cutlass.BFloat16) + csAkk[akk_row_base + row1, akk_col_base + col3] = ( + acc_akk_n1_3 * beta_row1 + ).to(cutlass.BFloat16) + else: + cute.arch.mbarrier_arrive(stage_reuse_mbars + s) + + cute.arch.mbarrier_arrive(mma_done_mbars + s) + + # ============================================================= + # Warps 28-31: Store/Inversion warps + # ============================================================= + if warp_idx >= NUM_K1_TMA_WARPS + NUM_MMA_WARPS: + store_warp = warp_idx - (NUM_K1_TMA_WARPS + NUM_MMA_WARPS) + for chunk_iter in cutlass.range_constexpr(CHUNKS_PER_BLOCK): + s = chunk_iter % NUM_STAGES + phase = chunk_iter // NUM_STAGES % 2 + chunk_idx = chunk_base + chunk_iter + chunk_start = cutlass.Int32(0) + st_eos = cutlass.Int32(0) + if IS_VARLEN: + if chunk_idx < num_chunks: + _sid = cutlass.Int32(mChunkIndices[chunk_idx, 0]) + chunk_start = ( + cutlass.Int32(mCuSeqlens[_sid]) + + cutlass.Int32(mChunkIndices[chunk_idx, 1]) * BT + ) + st_eos = cutlass.Int32(mCuSeqlens[_sid + 1]) + else: + chunk_start = chunk_idx * BT + + cute.arch.mbarrier_wait(mma_done_mbars + s, phase) + + csAqk = sAqk[(None, None, s)] + csAkk = sAkk[(None, None, s)] + + # Dual-path row-major store. SMEM is in final GMEM layout + # already (causal mask + diag=1 applied at MMA write; upper-tri + # SMEM is zero from CTA-startup init). Both paths use vec2 + # autovec_copy → STG.E.32 (4-byte coalesced). + # + # mAqk_v2 / mAkk_v2 shape (B, T, H, BT/2, 2): last dim is vec2. + # + # NON-PURE: full 64-col row-major + row mask (chunk may + # overflow seq end). 32 lanes/warp × 1 vec2 = full row. + # PURE: reduced cols. Warp s writes (s+1)*16 cols/row (no row + # mask, all rows in seq). Saves ~37% GMEM bandwidth vs full. + col_lo = lane_id * 2 + col_hi = col_lo + 1 + col_vec_idx = lane_id + rAqkOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + rAkkOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + + if IS_VARLEN and not VARLEN_PURE: + # NON-PURE: row-major full-row vec autovec + row mask. SMEM + # upper-tri is zero (MMA-masked at write), so writing all 64 + # cols is correct. STG.E.32 (32 lanes × 4 bytes per row). + row_base_warp = store_warp * (BT // NUM_STORE_WARPS) + for ri in cutlass.range_constexpr(BT // NUM_STORE_WARPS): + local_row = row_base_warp + ri + abs_row = chunk_start + local_row + rAqkOut[0] = csAqk[local_row, col_lo] + rAqkOut[1] = csAqk[local_row, col_hi] + rAkkOut[0] = csAkk[local_row, col_lo] + rAkkOut[1] = csAkk[local_row, col_hi] + if abs_row < st_eos: + cute.autovec_copy( + rAqkOut, mAqk_v2[i_b, abs_row, i_h, col_vec_idx, None] + ) + cute.autovec_copy( + rAkkOut, mAkk_v2[i_b, abs_row, i_h, col_vec_idx, None] + ) + else: + # PURE/eqlen: reduced lower-triangular store. Eqlen folds + # tile_valid to True; pure varlen checks the scheduler's + # potentially padded tail slot once per store warp. + tile_valid = True if not IS_VARLEN else chunk_idx < num_chunks + if tile_valid: + for tile_idx in cutlass.range_constexpr(NUM_TILES): + i_q = _TILE_IQ[tile_idx] + i_k = _TILE_IK[tile_idx] + is_diag = _TILE_IQ[tile_idx] == _TILE_IK[tile_idx] + gmem_aqk_row_base = chunk_start + i_q * BC + gmem_aqk_col_base = i_k * BC + gmem_akk_row_base = chunk_start + i_k * BC + gmem_akk_col_base = i_q * BC + smem_aqk_row_base = i_q * BC + smem_aqk_col_base = i_k * BC + smem_akk_row_base = i_k * BC + smem_akk_col_base = i_q * BC + for ri in cutlass.range_constexpr(BC // NUM_STORE_WARPS): + local_row = store_warp * (BC // NUM_STORE_WARPS) + ri + if lane_id < BC: + local_col = lane_id + aqk_val = csAqk[ + smem_aqk_row_base + local_row, + smem_aqk_col_base + local_col, + ] + akk_val = csAkk[ + smem_akk_row_base + local_row, + smem_akk_col_base + local_col, + ] + if is_diag and local_row < local_col: + aqk_val = cutlass.BFloat16(0.0) + if is_diag and local_row < local_col: + akk_val = cutlass.BFloat16(0.0) + if is_diag and local_row == local_col: + akk_val = cutlass.BFloat16(1.0) + mAqk[ + i_b, + gmem_aqk_row_base + local_row, + i_h, + gmem_aqk_col_base + local_col, + ] = aqk_val + mAkk[ + i_b, + gmem_akk_row_base + local_row, + i_h, + gmem_akk_col_base + local_col, + ] = akk_val + + cute.arch.mbarrier_arrive(store_done_mbars + s) + + yield_out() + + +# ========================================================================= +# Host function +# ========================================================================= +def make_host_function( + B, NT, H, is_varlen=False, T_padded=None, has_bias=False, use_safe_gate=False, varlen_pure=False +): + """ + `varlen_pure=True` asserts that all seq lengths in the batch are multiples + of BT (= 64). Under that assumption every chunk has 64 valid rows so the + four mask sites (K1 row mask, K1 store mask, MMA accumulator zero-fill, + Store row mask) are guaranteed to never fire and are dead-code eliminated + at compile time. Caller's data layout is unchanged — this is a hint only. + """ + _H = H + _IS_VARLEN = 1 if is_varlen else 0 + _HAS_BIAS = 1 if has_bias else 0 + _USE_SAFE_GATE = 1 if use_safe_gate else 0 + _VARLEN_PURE = 1 if (is_varlen and varlen_pure) else 0 + if is_varlen: + assert B == 1, "Varlen requires B=1" + assert T_padded is not None, "T_padded required for varlen" + + # B / NT / T are RUNTIME arguments of host_fn (rt_b / rt_nt / + # rt_t_total below), not closure constants: baking them keyed every + # compiled kernel to the batch shape, and eval traffic (a distinct + # total token count per prefill batch) recompiled ~100 s per batch — + # observed as GSM8K "hangs" (0/1319 responses). Only genuine constexpr + # specializers (H, is_varlen, bias/gate/pure variants) stay baked. + s_row = _H * K_DIM + s_col = 1 + s_h = K_DIM + + @cute.jit + def host_fn( + mQ, + mK, + mG, + mA_log, + mBeta, + scale, + mKscaled, + mKg, + mQscaled, + mGkLast, + mAqk, + mAkk, + mCuSeqlens, + mChunkIndices, + mDtBias, + lower_bound_val, + rt_nt: cutlass.Int32, + rt_b: cutlass.Int32, + rt_t_total: cutlass.Int32, + # Launch stream — a runtime argument (the executor runs the model on + # a dedicated non-blocking torch stream; launching on the DSL default + # stream races with the caller's stream). See _launch_fused_k123_inv. + stream: cuda.CUstream, + ): + # 3D TMA view: (B*T, K_DIM, H). domain_offset in kernel shifts to + # arbitrary chunk_start — no BT alignment required for varlen. + # Dynamic T extent: the TMA tensormap is materialized per call from + # the runtime shape, so one compiled host_fn serves every batch. + view_layout_3d = cute.make_layout( + (rt_t_total, K_DIM, _H), + stride=(s_row, s_col, s_h), + ) + mQ_view = cute.make_tensor(mQ.iterator, view_layout_3d) + mK_view = cute.make_tensor(mK.iterator, view_layout_3d) + mG_view = cute.make_tensor(mG.iterator, view_layout_3d) + + smem_atom_qk = tcgen05.make_smem_layout_atom( + tcgen05.SmemLayoutAtomKind.K_SW128, cutlass.BFloat16 + ) + qk_smem_2d = cute.tile_to_shape(smem_atom_qk, (BT, K_DIM), order=(0, 1)) + qk_smem_3d = cute.tile_to_shape(smem_atom_qk, (BT, K_DIM, NUM_STAGES), order=(0, 1, 2)) + + g_smem_2d = cute.make_layout((BT, K_DIM), stride=(K_DIM, 1)) + g_smem_3d = cute.make_layout((BT, K_DIM, NUM_STAGES), stride=(K_DIM, 1, BT * K_DIM)) + + tma_op = cpasync.CopyBulkTensorTileG2SOp(cpasync.CtaGroup.ONE) + tma_atom_Q, tma_tensor_Q = cpasync.make_tiled_tma_atom( + tma_op, mQ_view, qk_smem_2d, cute.product_each(qk_smem_2d.shape), num_multicast=1 + ) + tma_atom_K, tma_tensor_K = cpasync.make_tiled_tma_atom( + tma_op, mK_view, qk_smem_2d, cute.product_each(qk_smem_2d.shape), num_multicast=1 + ) + tma_atom_G, tma_tensor_G = cpasync.make_tiled_tma_atom( + tma_op, mG_view, g_smem_2d, cute.product_each(g_smem_2d.shape), num_multicast=1 + ) + + g_cumsum_layout = cute.make_layout( + (BT, K_DIM, NUM_STAGES), stride=(K_STRIDE, 1, BT * K_STRIDE) + ) + + copy_atom_qk_k1 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=32 + ) + tiled_copy_qk_k1 = cute.make_tiled_copy_tv( + copy_atom_qk_k1, + thr_layout=cute.make_layout((1, 32)), + val_layout=cute.make_layout((1, 2)), + ) + + # The v2 views are indexed [i_b, per-batch row, ...] by the kernel, + # so their T extent / batch stride must be PER-BATCH T, not the flat + # B*T used by the 3D TMA view above. Using rt_t_total here sent every + # i_b > 0 store rt_b× too far (OOB writes + parity failure for the + # eqlen B=2 path; varlen is B=1 so it never noticed). + rt_t_pb = rt_t_total // rt_b + out_v2_layout = cute.make_layout( + (rt_b, rt_t_pb, _H, K_VEC, VEC), + stride=(rt_t_pb * _H * K_DIM, _H * K_DIM, K_DIM, VEC, 1), + ) + mKscaled_v2 = cute.make_tensor(mKscaled.iterator, out_v2_layout) + mQscaled_v2 = cute.make_tensor(mQscaled.iterator, out_v2_layout) + mKg_v2 = cute.make_tensor(mKg.iterator, out_v2_layout) + + # vec2 views for mAqk / mAkk (BT dimension instead of K_DIM). + # Shape: (B, T, H, BT/2, 2). Each VEC=2 slot = 2 contiguous bf16 = 1 + # fp32-aligned 4-byte unit. Used by store warp's autovec_copy → + # STG.E.32 with 32 lanes coalesced to 1 cache line per row. + BT_VEC = BT // VEC # 32 + akk_v2_layout = cute.make_layout( + (rt_b, rt_t_pb, _H, BT_VEC, VEC), + stride=(rt_t_pb * _H * BT, _H * BT, BT, VEC, 1), + ) + mAqk_v2 = cute.make_tensor(mAqk.iterator, akk_v2_layout) + mAkk_v2 = cute.make_tensor(mAkk.iterator, akk_v2_layout) + + gklast_v2_layout = cute.make_layout( + (rt_b, rt_nt, _H, K_VEC, VEC), + stride=(rt_nt * _H * K_DIM, _H * K_DIM, K_DIM, VEC, 1), + ) + mGkLast_v2 = cute.make_tensor(mGkLast.iterator, gklast_v2_layout) + + mma_op = cute.nvgpu.warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16)) + tiled_mma_k2 = cute.make_tiled_mma( + mma_op, cute.make_layout((1, 1, 1)), permutation_mnk=(16, 8, 16) + ) + + tiled_copy_mma_A = cute.make_tiled_copy_A( + cute.make_copy_atom(cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), cutlass.BFloat16), + tiled_mma_k2, + ) + tiled_copy_mma_B = cute.make_tiled_copy_B( + cute.make_copy_atom(cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 2), cutlass.BFloat16), + tiled_mma_k2, + ) + + copy_atom_Gcum = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=64 + ) + tiled_copy_Gcum_norm = cute.make_tiled_copy_tv( + copy_atom_Gcum, thr_layout=cute.make_layout((1, 4)), val_layout=cute.make_layout((1, 2)) + ) + tiled_copy_Gcum_gate = cute.make_tiled_copy_C(copy_atom_Gcum, tiled_mma_k2) + + smem_size = ( + BT * K_DIM * 2 * 2 * NUM_STAGES + + BT * K_DIM * 2 * NUM_STAGES + + BT * K_STRIDE * 4 * NUM_STAGES + + K1_ROW_GROUPS * PARTIAL_COLS * 4 + + BT * AQK_TILE_STRIDE * 2 * NUM_STAGES # sAqk bf16 (64x72 row-major) + + BT * AKK_STRIDE * 2 * NUM_STAGES # sAkk bf16 + + 512 + ) + + # Constant grid: the persistent for_generate loop bounds work by the + # runtime total_cgs, so blocks past the work count exit immediately. + # (A shape-dependent min() here would re-bake the grid per NT.) + _grid_x = NUM_SMS + + fused_kernel123( + tma_atom_Q, + tma_tensor_Q, + tma_atom_K, + tma_tensor_K, + tma_atom_G, + tma_tensor_G, + mA_log, + mBeta, + scale, + mKscaled_v2, + mKg_v2, + mQscaled_v2, + mGkLast_v2, + mAqk, + mAkk, + mAqk_v2, + mAkk_v2, + tiled_copy_qk_k1, + tiled_mma_k2, + tiled_copy_mma_A, + tiled_copy_mma_B, + tiled_copy_Gcum_norm, + tiled_copy_Gcum_gate, + qk_smem_3d, + g_smem_3d, + g_cumsum_layout, + rt_nt, + _H, + rt_b, + mCuSeqlens, + mChunkIndices, + _IS_VARLEN, + mDtBias, + lower_bound_val, + _HAS_BIAS, + _USE_SAFE_GATE, + _VARLEN_PURE, + ).launch( + grid=(_grid_x, 1, 1), + block=(THREADS, 1, 1), + smem=smem_size, + stream=stream, + ) + + return host_fn diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k1234.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k1234.py new file mode 100644 index 000000000000..f4f617814918 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k1234.py @@ -0,0 +1,2276 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fused K1+K2+K3+K4 single kernel for KDA forward pass. + +Architecture: + 640 threads (20 warps = 5 warpgroups), grid = (B*H, 1, 1) for eqlen. + Per-chunk two-phase execution: + Phase A (K123): K1/K2/K3 warps — gate activation, intra-attention, inversion + Phase B (K4): WG0 + WG1 — 6 MMAs, readout, state update + + K123 writes intermediates to Zone B SMEM + sGkLast. K4 reads from SMEM. + Only O and S (final outputs) touch GMEM. SU and gk_last are SMEM-only. + SMEM aliased between phases via flat pool (~225KB, within B200's 228KB). + +Warp assignment (5 warpgroups, 128 threads each): + WG0 (W0-3): K4 MMA(W0) + TMA(W2) + idle(W1,W3) — idle during K123 + WG1 (W4-7): state readout+decay (K123) AND W/NV/O readout (K4) + WG2 (W8-11): K1 — TMA+gate/cumsum/KG/KS/QS + K3 inversion+store (merged) + WG3 (W12-15): K2 — intra-attention MMA (first 4 warps) + WG4 (W16-19): K2 — intra-attention MMA (second 4 warps) + + Per-warpgroup register reallocation via setmaxnreg (CUDA 13.1 cu13 libs): + WG0: dec(56) donates 5120 regs (MMA+TMA+idle) + WG1: inc(104) claims 1024 (state readout+decay + K4 readout) + WG2: default 96 (K1+K3 stays at base) + WG3+4 (K2): inc(112) claims 4096 (heavy intra-attention, 8 warps total) + Donate 5120 == Claim 5120 (exact balance — INC requires matching DEC). + + WG1 sync (mbarriers, asymmetric): st_ready_mbar(WG1→MMA: sST ready), + state_decayed_mbar(WG1→MMA: TMEM ready), mma6_done_mbar(MMA→WG1), + gk_last_ready_mbar(K1→WG1: gkLast ready) +""" + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.nvgpu.tcgen05 import Field, OperandMajorMode, OperandSource +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass.utils.gemm.sm100 import transform_partitioned_tensor_layout + +# ============================================================================= +# Constants +# ============================================================================= +mma_dtype = cutlass.BFloat16 +acc_dtype = cutlass.Float32 +out_dtype = cutlass.BFloat16 + +BT = 64 # chunk_size +BC = 16 # sub-chunk size +K_DIM = 128 # head dim +K_PAD = 8 +K_STRIDE = K_DIM + K_PAD # 136 + +# K4 tile dimensions +M4 = 64 # chunk_size +N4 = 128 # head dim (K or V) +K4_K = 64 # inner dim for MMA1/2/5 +K4_K3 = 128 # inner dim for MMA3/4 (state) +M6 = 128 # MMA6 M +N6 = 128 # MMA6 N +K6 = 64 # MMA6 K + +# Debug flags (set to 0 for normal operation) +# DEBUG_K4_LEVEL: 0=skip K4, 1=TMA only, 2=TMA+MMA, 3=full K4 +DEBUG_K4_LEVEL = 3 +# When True, skip MMA5 (AQC@NV) to isolate: O = OI = QS @ S only. +# If O becomes correct, bug is in sAQC. If still wrong, bug is in sQS. +DEBUG_SKIP_MMA5 = False + +# Thread/warp counts +THREADS = 640 +WARP_SIZE = 32 +WG_SIZE = 128 +NUM_WARPS = THREADS // WARP_SIZE # 28 + +# K123 warp assignment +NUM_K1_WARPS = 4 # Warps 8-11 (WG2) +NUM_MMA_WARPS = 8 # Warps 12-19 +NUM_MMA_ACTIVE = 8 +NUM_STORE_WARPS = 4 # Warps 20-23 +K1_FIRST_WARP = 8 +K2_FIRST_WARP = 12 +K3_FIRST_WARP = 20 + +# K4 warp assignment +K4_MMA_WARP = 0 +K4_TMA_WARP = 2 +K4_READOUT_WG = 1 # warpgroup index (warps 4-7) + +# K123 sub-parameters +K1_ROW_GROUPS = 4 +K1_COL_GROUPS = 1 +ROWS_PER_K1_WARP = BT // K1_ROW_GROUPS # 16 +K1_COLS_PER_WARP = K_DIM // K1_COL_GROUPS # 128 +ROWS_PER_STORE_WARP = BT // NUM_STORE_WARPS # 16 +VEC = K1_COLS_PER_WARP // 32 # 4 +K_VEC = K_DIM // VEC # 32 + + +NUM_SUB_CHUNKS = BT // BC # 4 +NUM_TILES = NUM_SUB_CHUNKS * (NUM_SUB_CHUNKS + 1) // 2 # 10 +MMA_K_TILE = 8 +NUM_MMA_K_TILES = K_DIM // MMA_K_TILE # 16 +AQK_TILE_COLS = NUM_TILES * BC # 160 +AQK_TILE_PAD = 8 +AQK_TILE_STRIDE = AQK_TILE_COLS + AQK_TILE_PAD # 168 +AKK_PAD = 8 +AKK_STRIDE = BT + AKK_PAD # 72 + +TEMP_PAD = 8 +TEMP_COLS = BC + TEMP_PAD # 24 +NUM_TEMPS = 2 + +_TILE_IQ = [0, 1, 1, 2, 2, 2, 3, 3, 3, 3] +_TILE_IK = [0, 0, 1, 0, 1, 2, 0, 1, 2, 3] + +LOG2E = 1.4426950408889634 +LN2 = 0.6931471805599453 +RCP_LN2 = LOG2E +SHFL_W4_CLAMP = 0x1C00 + +# SMEM layout byte offsets (informational — allocator handles actual placement) +# Persistent region (0-112KB): sAB(8) + sAQC(8) + sKG(16) + sQS(16) + sKS(16) + sST(32) + sV_ext(16) +# sQS/sKS: time-shared with K123 TMA Q/K, overwritten by K4-format after K2 +# sST: persistent — readout WG writes during K123 +# Zone A (112-192KB): sW(16) + sNV(16) + sO(16) + scratch(31KB) +# K123 scratch aliases: sG(16) + sGcum(34) + sAqk(5) + sAkk(18) + sTemp(3) +# Total peak: ~192KB + + +# ============================================================================= +# dsl_user_op helpers (from fuse_kernel123_no_persistent.py) +# ============================================================================= + + +@dsl_user_op +def k1_internal_barrier(*, loc=None, ip=None): + """Named barrier for K1 warps (0-3, 128 threads). barrier_id=2.""" + llvm.inline_asm( + T.i32(), + [], + "membar.cta; bar.sync 2, 128; mov.u32 $0, 0;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def mma_tf32_m16n8k8(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None): + """TF32 MMA: D = A * B + C, shape m16n8k8""" + a0_bits = llvm.bitcast(T.i32(), a0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a1_bits = llvm.bitcast(T.i32(), a1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a2_bits = llvm.bitcast(T.i32(), a2.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a3_bits = llvm.bitcast(T.i32(), a3.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b0_bits = llvm.bitcast(T.i32(), b0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b1_bits = llvm.bitcast(T.i32(), b1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [ + a0_bits, + a1_bits, + a2_bits, + a3_bits, + b0_bits, + b1_bits, + c0.ir_value(loc=loc, ip=ip), + c1.ir_value(loc=loc, ip=ip), + c2.ir_value(loc=loc, ip=ip), + c3.ir_value(loc=loc, ip=ip), + ], + """{ + mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 + {$0, $1, $2, $3}, + {$4, $5, $6, $7}, + {$8, $9}, + {$10, $11, $12, $13}; + }""", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +@dsl_user_op +def read_clock(*, loc=None, ip=None): + """Read globaltimer (ns). Returns i64 as two i32 words packed into fp32 pair.""" + result = llvm.inline_asm( + T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return result + + +@dsl_user_op +def fast_rcp(x, *, loc=None, ip=None): + result = llvm.inline_asm( + T.f32(), + [x.ir_value(loc=loc, ip=ip)], + "rcp.approx.ftz.f32 $0, $1;", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(result) + + +@dsl_user_op +def inv_internal_barrier(*, loc=None, ip=None): + llvm.inline_asm( + T.i32(), + [], + "bar.sync 3, 128; mov.u32 $0, 0;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _invert_diag(sAkk: cute.Tensor, block_rc, lane_id, *, loc=None, ip=None): + my_row = lane_id % 16 + halfwarp_base = (lane_id // 16) * 16 + r_off = block_rc * 16 + c_off = block_rc * 16 + rInv = cute.make_rmem_tensor(cute.make_layout((16,), stride=(1,)), cutlass.Float32) + rInv[0] = cutlass.Float32(1.0) + for x in range(1, 16): + rInv[x] = cutlass.Float32(0.0) + for d in range(1, 16): + col_d = my_row - d + valid = cutlass.Float32(col_d >= 0) + a_val = cutlass.Float32(sAkk[r_off + my_row, c_off + col_d]) * valid + acc = cutlass.Float32(0.0) + for j in range(1, d): + a_re = cutlass.Float32(sAkk[r_off + my_row, c_off + my_row - (d - j)]) + inv_shfl = cute.arch.shuffle_sync(rInv[j], halfwarp_base + my_row - d + j) + acc = acc + a_re * inv_shfl + rInv[d] = (-a_val - acc) * valid + rInv[0] = cutlass.Float32(1.0) + sAkk[r_off + my_row, c_off + my_row] = rInv[0] + for d in range(1, 16): + sAkk[r_off + my_row, c_off + (my_row + 16 - d) % 16] = rInv[d] * cutlass.Float32( + my_row >= d + ) + + +@dsl_user_op +def _matmul_AB(sAkk: cute.Tensor, br_A, bc_A, br_B, bc_B, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _z = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 16 + rB = br_B * 16 + cB = bc_B * 16 + a0 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid + 1]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid + 1]) + b0n0 = cutlass.Float32(sAkk[rB + 2 * tid, cB + gid]) + b1n0 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + gid]) + b0n1 = cutlass.Float32(sAkk[rB + 2 * tid, cB + 8 + gid]) + b1n1 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + 8 + gid]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n0, b1n0, _z, _z, _z, _z) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n1, b1n1, _z, _z, _z, _z) + a0 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid + 1]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid + 1]) + b0n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + gid]) + b1n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + gid]) + b0n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + 8 + gid]) + b1n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + 8 + gid]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8( + a0, a1, a2, a3, b0n0, b1n0, cn0_0, cn0_1, cn0_2, cn0_3 + ) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8( + a0, a1, a2, a3, b0n1, b1n1, cn1_0, cn1_1, cn1_2, cn1_3 + ) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _chain_mma_B( + sAkk: cute.Tensor, + br_B, + bc_B, + a0k0, + a1k0, + a2k0, + a3k0, + a0k1, + a1k1, + a2k1, + a3k1, + lane_id, + *, + loc=None, + ip=None, +): + gid = lane_id // 4 + tid = lane_id % 4 + _z = cutlass.Float32(0.0) + rB = br_B * 16 + cB = bc_B * 16 + b0n0 = cutlass.Float32(sAkk[rB + 2 * tid, cB + gid]) + b1n0 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + gid]) + b0n1 = cutlass.Float32(sAkk[rB + 2 * tid, cB + 8 + gid]) + b1n1 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + 8 + gid]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8( + a0k0, a1k0, a2k0, a3k0, b0n0, b1n0, _z, _z, _z, _z + ) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8( + a0k0, a1k0, a2k0, a3k0, b0n1, b1n1, _z, _z, _z, _z + ) + b0n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + gid]) + b1n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + gid]) + b0n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + 8 + gid]) + b1n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + 8 + gid]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8( + a0k1, a1k1, a2k1, a3k1, b0n0, b1n0, cn0_0, cn0_1, cn0_2, cn0_3 + ) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8( + a0k1, a1k1, a2k1, a3k1, b0n1, b1n1, cn1_0, cn1_1, cn1_2, cn1_3 + ) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _chain_mma_A( + sAkk: cute.Tensor, + br_A, + bc_A, + b0_k0n0, + b1_k0n0, + b0_k0n1, + b1_k0n1, + b0_k1n0, + b1_k1n0, + b0_k1n1, + b1_k1n1, + lane_id, + *, + loc=None, + ip=None, +): + gid = lane_id // 4 + tid = lane_id % 4 + _z = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 16 + a0 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid + 1]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid + 1]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0_k0n0, b1_k0n0, _z, _z, _z, _z) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0_k0n1, b1_k0n1, _z, _z, _z, _z) + a0 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid + 1]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid + 1]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8( + a0, a1, a2, a3, b0_k1n0, b1_k1n0, cn0_0, cn0_1, cn0_2, cn0_3 + ) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8( + a0, a1, a2, a3, b0_k1n1, b1_k1n1, cn1_0, cn1_1, cn1_2, cn1_3 + ) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _store_neg_C( + sAkk: cute.Tensor, br, bc, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None +): + gid = lane_id // 4 + tid = lane_id % 4 + r = br * 16 + c = bc * 16 + sAkk[r + gid, c + 2 * tid] = -c0 + sAkk[r + gid, c + 2 * tid + 1] = -c1 + sAkk[r + gid + 8, c + 2 * tid] = -c2 + sAkk[r + gid + 8, c + 2 * tid + 1] = -c3 + sAkk[r + gid, c + 8 + 2 * tid] = -c4 + sAkk[r + gid, c + 8 + 2 * tid + 1] = -c5 + sAkk[r + gid + 8, c + 8 + 2 * tid] = -c6 + sAkk[r + gid + 8, c + 8 + 2 * tid + 1] = -c7 + + +@dsl_user_op +def _shuffle_C_to_B(c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + src_a = 8 * tid + gid // 2 + src_b = src_a + 4 + f_odd = cutlass.Float32(gid % 2) + f_even = cutlass.Float32(1) - f_odd + c0_a = cute.arch.shuffle_sync(c0, src_a) + c1_a = cute.arch.shuffle_sync(c1, src_a) + c2_a = cute.arch.shuffle_sync(c2, src_a) + c3_a = cute.arch.shuffle_sync(c3, src_a) + c4_a = cute.arch.shuffle_sync(c4, src_a) + c5_a = cute.arch.shuffle_sync(c5, src_a) + c6_a = cute.arch.shuffle_sync(c6, src_a) + c7_a = cute.arch.shuffle_sync(c7, src_a) + c0_b = cute.arch.shuffle_sync(c0, src_b) + c1_b = cute.arch.shuffle_sync(c1, src_b) + c2_b = cute.arch.shuffle_sync(c2, src_b) + c3_b = cute.arch.shuffle_sync(c3, src_b) + c4_b = cute.arch.shuffle_sync(c4, src_b) + c5_b = cute.arch.shuffle_sync(c5, src_b) + c6_b = cute.arch.shuffle_sync(c6, src_b) + c7_b = cute.arch.shuffle_sync(c7, src_b) + return ( + c0_a * f_even + c1_a * f_odd, + c0_b * f_even + c1_b * f_odd, + c2_a * f_even + c3_a * f_odd, + c2_b * f_even + c3_b * f_odd, + c4_a * f_even + c5_a * f_odd, + c4_b * f_even + c5_b * f_odd, + c6_a * f_even + c7_a * f_odd, + c6_b * f_even + c7_b * f_odd, + ) + + +@dsl_user_op +def tma_store_fence(*, loc=None, ip=None): + """Fence: wait for all outstanding TMA S2G stores to finish reading SMEM. + Must be called after TMA store and before SMEM is reused (e.g., by K123 phase).""" + llvm.inline_asm( + T.i32(), + [], + "cp.async.bulk.commit_group; cp.async.bulk.wait_group.read 0; mov.u32 $0, 0;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def threadfence_gl(*, loc=None, ip=None): + """Global memory fence — flush L1 writes to L2 for TMA visibility.""" + llvm.inline_asm( + T.i32(), + [], + "membar.gl; mov.u32 $0, 0;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _store_C_temp( + sT: cute.Tensor, buf, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None +): + gid = lane_id // 4 + tid = lane_id % 4 + sT[gid, 2 * tid, buf] = c0 + sT[gid, 2 * tid + 1, buf] = c1 + sT[gid + 8, 2 * tid, buf] = c2 + sT[gid + 8, 2 * tid + 1, buf] = c3 + sT[gid, 8 + 2 * tid, buf] = c4 + sT[gid, 8 + 2 * tid + 1, buf] = c5 + sT[gid + 8, 8 + 2 * tid, buf] = c6 + sT[gid + 8, 8 + 2 * tid + 1, buf] = c7 + + +@dsl_user_op +def _load_C_temp(sT: cute.Tensor, buf, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + return ( + cutlass.Float32(sT[gid, 2 * tid, buf]), + cutlass.Float32(sT[gid, 2 * tid + 1, buf]), + cutlass.Float32(sT[gid + 8, 2 * tid, buf]), + cutlass.Float32(sT[gid + 8, 2 * tid + 1, buf]), + cutlass.Float32(sT[gid, 8 + 2 * tid, buf]), + cutlass.Float32(sT[gid, 8 + 2 * tid + 1, buf]), + cutlass.Float32(sT[gid + 8, 8 + 2 * tid, buf]), + cutlass.Float32(sT[gid + 8, 8 + 2 * tid + 1, buf]), + ) + + +# ============================================================================= +# Fused K1234 Kernel +# ============================================================================= + + +@cute.kernel +def fused_k1234_kernel( + # --- K4 MMA descriptors --- + tiled_mma_kmn: cute.TiledMma, + tiled_mma_mn_mn: cute.TiledMma, + # --- K4 TMA load atoms (V only; S goes TMEM→SMEM via readout WG) --- + tma_atom_v: cute.CopyAtom, + mV_nkl: cute.Tensor, + v_sl: cute.ComposedLayout, + s_sl: cute.ComposedLayout, # sST SMEM layout (filled from TMEM, not TMA) + # --- K4 Zone B swizzle layouts (for SMEM allocation + write views) --- + ab_sl: cute.ComposedLayout, + ks_sl: cute.ComposedLayout, + qs_sl: cute.ComposedLayout, + aqc_sl: cute.ComposedLayout, + kg_sl: cute.ComposedLayout, + # --- K4 TMA store --- + tma_atom_o_st: cute.CopyAtom, + mOo: cute.Tensor, + store_sl: cute.ComposedLayout, + # --- K4 readout/reinterpret layouts --- + readout_k_sl: cute.ComposedLayout, + nv_b_sl: cute.ComposedLayout, + nv_a_sl: cute.ComposedLayout, + # --- K123 TMA load atoms + tensors --- + tma_atom_q_k123: cute.CopyAtom, + tma_tensor_q_k123: cute.Tensor, + tma_atom_k_k123: cute.CopyAtom, + tma_tensor_k_k123: cute.Tensor, + tma_atom_g_k123: cute.CopyAtom, + tma_tensor_g_k123: cute.Tensor, + # --- K123 SMEM layouts (passed from host) --- + qk_smem_layout: cute.ComposedLayout, + g_smem_layout, + g_cumsum_layout, + # --- K123 tiled copies --- + tiled_copy_qk_k1, + tiled_mma_k2, + tiled_copy_mma_A, + tiled_copy_mma_B, + tiled_copy_Gcum_norm, + tiled_copy_Gcum_gate, + # --- K123 GMEM tensors (K123 outputs now go to Zone B SMEM, not GMEM) --- + mA_log: cute.Tensor, + mBeta: cute.Tensor, + scale: cutlass.Float32, + # --- K4 GMEM (state in TMEM; GMEM only for init/final store) --- + mS_fp32: cute.Tensor, # [V, K, B*H] fp32 — initial state + final state output + # --- dt_bias / safe_gate --- + mDtBias: cute.Tensor, + lower_bound: cutlass.Float32, + HAS_BIAS: cutlass.Constexpr[int], + USE_SAFE_GATE: cutlass.Constexpr[int], + # --- Clock profiling (optional) --- + mClocks: cute.Tensor, # [16, B*H] i64 — per-role clock profiling + PROFILE_CLOCKS: cutlass.Constexpr[int], + # --- Dimensions (dynamic — one compilation serves all shapes) --- + num_chunks: cutlass.Int32, + num_heads: cutlass.Int32, + batch_size: cutlass.Int32, +): + # === Thread indices === + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + warpgroup_idx = cute.arch.make_warp_uniform(tidx // WG_SIZE) + warpgroup_tidx = tidx % WG_SIZE + lane_id = tidx % WARP_SIZE + + # ======================================================================== + # Per-warpgroup register reallocation (Blackwell SM100, CUDA 13.1). + # + # CRITICAL: setmaxnreg.inc is effectively BLOCKING on SM100. If CTAPOOL + # doesn't have the full requested regs, INC spins forever (hang at + # torch.cuda.synchronize). Total INC claims MUST be <= total DEC donations. + # + # Requirements: + # (1) nvidia-cutlass-dsl-libs-cu13==4.4.2 installed (start_env.sh) + # (2) CUDA_HOME = CUDA 13.1 toolkit (so 13.1 ptxas is used) + # (3) DEC target >= 40 (else ptxas rejects with C7507) + # (4) sum(DEC donations) >= sum(INC claims) + # + # Budget (20 warps, 640 threads, compiler default 96/thread): + # WG0 (W0-3, 128t): dec(56) donates (96-56)*128 = 5120 regs + # WG1 (W4-7, 128t): inc(104) claims (104-96)*128 = 1024 regs + # WG2 (W8-11, 128t): default 96, no call + # WG3+4 (W12-19, 256t): inc(112) claims (112-96)*256 = 4096 regs + # Donate 5120 == Claim 5120 (exact balance). + # ======================================================================== + if warpgroup_idx == 0: + cute.arch.setmaxregister_decrease(56) + elif warpgroup_idx == 1: + cute.arch.setmaxregister_increase(104) + elif warpgroup_idx >= 3: + cute.arch.setmaxregister_increase(112) + + bid = cute.arch.block_idx()[0] + i_b = bid // num_heads + i_h = bid % num_heads + + # === K4 thr slices (created once, used every K4 phase) === + thr_kmn = tiled_mma_kmn.get_slice(0) + thr_mn = tiled_mma_mn_mn.get_slice(0) + dice = (None, None, None) + + # ======================================================================== + # SMEM ALLOCATION — Persistent + Aliased Layout (~224KB) + # Persistent (0-112KB): K4 operands + sST + sV_ext + # sQS/sKS: K4-only, NOT aliased by K123 (separate sQ_k123/sK_k123 in Zone A) + # sST: persistent — K1 warps write decayed state bf16 during K123. + # Zone A (112-224KB): K123 scratch / K4 readout (aliased) + # During K123: sQ, sK, sG, sGcum, sAqk, sAkk, sTemp (~109KB) + # During K4: sW, sNV, sO readout buffers (48KB) + # sGkLast (512B): separate allocation, persists from K1 to K4 + # State S[V,K] fp32 lives in TMEM offset 256-383 (persistent across chunks). + # ======================================================================== + smem = cutlass.utils.SmemAllocator() + AL = 128 + + # --- Persistent region (written by K123, read by K4) --- + sAB = smem.allocate_tensor(mma_dtype, ab_sl.outer, AL, ab_sl.inner) # 8KB A_kk_inv + sAQC = smem.allocate_tensor(mma_dtype, aqc_sl.outer, AL, aqc_sl.inner) # 8KB A_qk + sKG = smem.allocate_tensor(mma_dtype, kg_sl.outer, AL, kg_sl.inner) # 16KB kg + sQS = smem.allocate_tensor(mma_dtype, qs_sl.outer, AL, qs_sl.inner) # 16KB q_scaled (K4 only) + sKS = smem.allocate_tensor(mma_dtype, ks_sl.outer, AL, ks_sl.inner) # 16KB k_scaled (K4 only) + sST = smem.allocate_tensor( + mma_dtype, s_sl.outer, AL, s_sl.inner + ) # 32KB state bf16 (persistent) + sV_ext = smem.allocate_tensor(mma_dtype, v_sl.outer, AL, v_sl.inner) # 16KB V (non-aliased) + # Persistent total: 112KB + + # --- Zone A: K4 readout (aliases K123 scratch) --- + sW = smem.allocate_tensor(mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner) # 16KB readout + sNV = smem.allocate_tensor( + mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner + ) # 16KB readout + sO = smem.allocate_tensor(mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner) # 16KB readout + # K4 readout: 48KB. K123 scratch needs ~111KB, extra ~64KB allocated below. + _smem_extra = smem.allocate_array(cutlass.Float32, 16384) # ~64KB for K123 overflow + + # --- K123 alias: sQ_k123/sK_k123 in Zone A (separate from sQS/sKS) --- + # sQ_k123 aliases sW, sK_k123 aliases sNV during K123. + # TMA loads Q/K here in K_SW128 swizzle layout. K1 writes KS/QS to persistent sQS/sKS. + _za = sW.iterator # Zone A base (bf16 pointer) + sQ_k123 = cute.make_tensor( + cute.recast_ptr(_za, qk_smem_layout.inner, cutlass.BFloat16), qk_smem_layout.outer + ) + sK_k123 = cute.make_tensor( + cute.recast_ptr(_za + 8192, qk_smem_layout.inner, cutlass.BFloat16), qk_smem_layout.outer + ) + + # --- K123 alias: scratch tensors (overlap Zone A, after sQ/sK) --- + # sG_k123 [64,128] bf16 at +32KB (+16384 bf16), aliases sO during K4 + sG_k123 = cute.make_tensor(cute.recast_ptr(_za + 16384, dtype=cutlass.BFloat16), g_smem_layout) + # sGcum [64,136] fp32 at +48KB (+24576 bf16), 34816B + sGcum = cute.make_tensor(cute.recast_ptr(_za + 24576, dtype=cutlass.Float32), g_cumsum_layout) + # sPartialLast removed — 1-col-per-thread K1 doesn't need cross-warp shuffle + # sAqk_k123 [16,168] bf16 at +85KB (+43040 bf16), 5376B + aqk_tile_layout = cute.make_layout((BC, AQK_TILE_STRIDE), stride=(AQK_TILE_STRIDE, 1)) + sAqk_k123 = cute.make_tensor( + cute.recast_ptr(_za + 43040, dtype=cutlass.BFloat16), aqk_tile_layout + ) + # sAkk_k123 [64,72] fp32 at +91KB (+45728 bf16), 18432B + akk_tile_layout = cute.make_layout((BT, AKK_STRIDE), stride=(AKK_STRIDE, 1)) + sAkk_k123 = cute.make_tensor( + cute.recast_ptr(_za + 45728, dtype=cutlass.Float32), akk_tile_layout + ) + # sTemp [16,24,2] fp32 at +110KB (+54944 bf16), 3072B + temp_layout = cute.make_layout( + (BC, TEMP_COLS, NUM_TEMPS), stride=(TEMP_COLS, 1, BC * TEMP_COLS) + ) + sTemp = cute.make_tensor(cute.recast_ptr(_za + 54944, dtype=cutlass.Float32), temp_layout) + + # --- K4 reinterpret views --- + sNV_b = cute.make_tensor(cute.recast_ptr(sNV.iterator, nv_b_sl.inner, mma_dtype), nv_b_sl.outer) + sNV_a = cute.make_tensor(cute.recast_ptr(sNV.iterator, nv_a_sl.inner, mma_dtype), nv_a_sl.outer) + sO_st = cute.make_tensor( + cute.recast_ptr(sO.iterator, store_sl.inner, out_dtype), store_sl.outer + ) + + # sSU_fp32 removed — SU stays in TMEM/RMEM, no SMEM copy needed + + # gk_last: 128 fp32 per (b,h) chunk — persists from K1 to K4 within same chunk + _sGkLast_buf = smem.allocate_array(cutlass.Float32, K_DIM) # 128 fp32 = 512 bytes + sGkLast = cute.make_tensor(_sGkLast_buf, cute.make_layout((K_DIM,), stride=(1,))) + + # beta: 64 bf16 per chunk — loaded once in K1, read by K2/K3 + _sBeta_buf = smem.allocate_array(cutlass.BFloat16, BT) # 64 bf16 = 128 bytes + sBeta = cute.make_tensor(_sBeta_buf, cute.make_layout((BT,), stride=(1,))) + + # ======================================================================== + # TMEM ALLOCATION (K4, warp 0 only) + # ======================================================================== + tmem_smem = smem.allocate_array(cutlass.Int32, 1) + if warp_idx == 0: + cute.arch.alloc_tmem(512, tmem_smem) + + # ======================================================================== + # K123 MBARRIERS (single-stage, no double-buffering) + # ======================================================================== + k123_tma_mbar = smem.allocate_array(cutlass.Int64, 1) + k123_k1_done_mbar = smem.allocate_array(cutlass.Int64, 1) + k123_mma_done_mbar = smem.allocate_array(cutlass.Int64, 1) + # Asymmetric mbarriers for WG1 async (arrive = non-blocking signal, wait = blocking) + mma6_done_mbar = smem.allocate_array(cutlass.Int64, 1) # MMA→WG1: MMA6 complete + st_ready_mbar = smem.allocate_array(cutlass.Int64, 1) # WG1→MMA: sST bf16 ready + gk_last_ready_mbar = smem.allocate_array(cutlass.Int64, 1) # K1→WG1: gkLast ready + state_decayed_mbar = smem.allocate_array(cutlass.Int64, 1) # WG1→MMA: decayed state ready + final_state_done_mbar = smem.allocate_array(cutlass.Int64, 1) # WG1→MMA: final state stored + + # K4 named barriers (symmetric — both sides arrive and wait) + sW_ready_nbar = pipeline.NamedBarrier(4, WG_SIZE + WARP_SIZE) + sNV_ready_nbar = pipeline.NamedBarrier(5, WG_SIZE + WARP_SIZE) + store_nbar = pipeline.NamedBarrier(6, WG_SIZE + 2 * WARP_SIZE) + phase_nbar = pipeline.NamedBarrier(7, THREADS) # all threads (WG1 now participates) + + elect_one = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + wg_coop = pipeline.CooperativeGroup(pipeline.Agent.Thread, WG_SIZE) + + # K4 TMA pipelines + def _make_tma_pipe(byte_count): + ptr = smem.allocate_array(cutlass.Int64, 2) + return pipeline.PipelineTmaUmma.create( + barrier_storage=ptr, + num_stages=1, + producer_group=elect_one, + consumer_group=elect_one, + tx_count=byte_count, + defer_sync=True, + ).make_participants() + + # Only V still uses TMA pipeline (KS/QS now in persistent SMEM) + v_bytes = cute.size_in_bytes(mma_dtype, v_sl) + v_prod, v_cons = _make_tma_pipe(v_bytes) + + # K4 UMMA pipelines (MMA → readout) + def _make_umma_pipe(): + ptr = smem.allocate_array(cutlass.Int64, 2) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=ptr, + num_stages=1, + producer_group=elect_one, + consumer_group=wg_coop, + defer_sync=True, + ).make_participants() + + w_prod, w_cons = _make_umma_pipe() + nv_prod, nv_cons = _make_umma_pipe() + o_prod, o_cons = _make_umma_pipe() + # kv_acc_pipeline removed: barrier() between K123→K4 ensures state readout/decay done + + # ======================================================================== + # INIT + # ======================================================================== + k123_tma_bytes = BT * K_DIM * 2 * 3 # q+k+g bf16 + + if tidx == 0: + cute.arch.mbarrier_init(k123_tma_mbar, 1) + cute.arch.mbarrier_init(k123_k1_done_mbar, NUM_K1_WARPS * WARP_SIZE) + cute.arch.mbarrier_init(k123_mma_done_mbar, NUM_MMA_WARPS * WARP_SIZE) + cute.arch.mbarrier_init(mma6_done_mbar, WARP_SIZE) # MMA warp (W0, 32t) arrives + cute.arch.mbarrier_init(st_ready_mbar, WG_SIZE) # WG1 (128t) arrives + cute.arch.mbarrier_init(gk_last_ready_mbar, NUM_K1_WARPS * WARP_SIZE) # K1 (128t) arrives + cute.arch.mbarrier_init(state_decayed_mbar, WG_SIZE) # WG1 (128t) arrives + cute.arch.mbarrier_init(final_state_done_mbar, WG_SIZE) # WG1 (128t) arrives + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + + # ======================================================================== + # TMEM tensors (K4, created once) + # ======================================================================== + tmem_ptr = cute.arch.retrieve_tmem_ptr(cutlass.Int32, 16, tmem_smem) + + tCtW_shape = tiled_mma_kmn.partition_shape_C((M4, N4)) + tCtW_fake = tiled_mma_kmn.make_fragment_C(tCtW_shape) + tCtW = cute.make_tensor(cute.recast_ptr(tmem_ptr + 0, dtype=acc_dtype), tCtW_fake.layout) + + tCtNV_shape = tiled_mma_kmn.partition_shape_C((M4, N4)) + tCtNV_fake = tiled_mma_kmn.make_fragment_C(tCtNV_shape) + tCtNV = cute.make_tensor(cute.recast_ptr(tmem_ptr + 128, dtype=acc_dtype), tCtNV_fake.layout) + tCtO = cute.make_tensor(cute.recast_ptr(tmem_ptr + 384, dtype=acc_dtype), tCtNV_fake.layout) + + tCtS_shape = tiled_mma_mn_mn.partition_shape_C((M6, N6)) + tCtS_fake = tiled_mma_mn_mn.make_fragment_C(tCtS_shape) + # tCtS = cute.make_tensor(cute.recast_ptr(tmem_ptr + 128, dtype=acc_dtype), tCtS_fake.layout) + + # State [V=128, K=128] fp32 at TMEM offset 256 — persistent across chunks (GDN pattern) + # Same MN-MN layout as SU (tCtS) so corresponding elements align for state update + tCtState = cute.make_tensor(cute.recast_ptr(tmem_ptr + 256, dtype=acc_dtype), tCtS_fake.layout) + + # ======================================================================== + # WARP-SPECIALIZED INDEPENDENT LOOPS (GDN pattern) + # Each role: per-role setup -> own chunk loop -> cleanup + # (setmaxnreg deferred to Phase 2) + # ======================================================================== + + # ============================================================== + # WG1 (warps 4-7): state readout + decay (K123) + W/NV/O readout (K4) + # Both tasks in one WG — they never overlap (K123 vs K4 phases). + # ============================================================== + if warpgroup_idx == K4_READOUT_WG: + # --- W/NV/O readout setup (M=64: Ld16x256bOp) --- + tCtW_mn = transform_partitioned_tensor_layout(tCtW) + tCtNV_mn = transform_partitioned_tensor_layout(tCtNV) + tCtO_mn = transform_partitioned_tensor_layout(tCtO) + + atom_t2r = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition(1)), acc_dtype) + tiled_t2r = tcgen05.make_tmem_copy(atom_t2r, tCtW[(None, None), 0, 0]) + thr_t2r = tiled_t2r.get_slice(warpgroup_tidx) + + tTR_W = thr_t2r.partition_S(tCtW_mn) + tTR_NV = thr_t2r.partition_S(tCtNV_mn) + tTR_O = thr_t2r.partition_S(tCtO_mn) + + atom_r2s_k = sm100_utils.get_smem_store_op( + utils.LayoutEnum.ROW_MAJOR, mma_dtype, acc_dtype, tiled_t2r + ) + tiled_r2s_k = cute.make_tiled_copy_D(atom_r2s_k, tiled_t2r) + thr_r2s_k = tiled_r2s_k.get_slice(warpgroup_tidx) + tCsW = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sW)) + tCsO = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sO)) + tCsNV = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sNV)) + + cId = cute.make_identity_tensor((M4, N4)) + tTR_cId = thr_t2r.partition_D(cId) + + # --- State TMEM setup (M=128, GDN pattern) --- + cId_128 = cute.make_identity_tensor((M6, N6)) + + # --- State TMEM copy atoms (M=128, GDN pattern) --- + tCtState_mn = transform_partitioned_tensor_layout(tCtState) + + # TMEM->RMEM (Ld32x32bOp for M=128) + atom_state_t2r = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), acc_dtype) + tiled_state_t2r = tcgen05.make_tmem_copy(atom_state_t2r, tCtState[(None, None), 0, 0]) + thr_state_t2r = tiled_state_t2r.get_slice(warpgroup_tidx) + tTR_tCtState = thr_state_t2r.partition_S(tCtState_mn) + tTR_tCcState = thr_state_t2r.partition_D(cId_128) + tRrState = cute.make_rmem_tensor_like(tTR_tCcState, acc_dtype) # fp32 RMEM + + # RMEM->TMEM (St32x32bOp for M=128, state write-back) + atom_state_r2t = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition(32)), acc_dtype) + tiled_state_r2t = tcgen05.make_tmem_copy(atom_state_r2t, tCtState[(None, None), 0, 0]) + thr_state_r2t = tiled_state_r2t.get_slice(warpgroup_tidx) + tRT_tCtState = thr_state_r2t.partition_D(tCtState_mn) + + # RMEM bf16 -> swizzled sST SMEM (CopyUniversalOp, NO domain transpose) + atom_state_r2s = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), mma_dtype, num_bits_per_copy=16 + ) + tiled_state_r2s = cute.make_tiled_copy_D(atom_state_r2s, tiled_state_t2r) + thr_state_r2s = tiled_state_r2s.get_slice(warpgroup_tidx) + sST_mn_view = transform_partitioned_tensor_layout(sST) + tCsState_inp = thr_state_r2s.partition_D(sST_mn_view) + tRrState_bf16 = cute.make_rmem_tensor_like(tTR_tCcState, mma_dtype) # bf16 RMEM + tCrState_bf16 = tiled_state_r2s.retile(tRrState_bf16) + + # GMEM fp32 -> RMEM fp32 (for initial/final state, runs once) + atom_state_g2r = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), acc_dtype, num_bits_per_copy=32 + ) + tiled_state_g2r = cute.make_tiled_copy_S(atom_state_g2r, tiled_state_r2t) + thr_state_g2r = tiled_state_g2r.get_slice(warpgroup_tidx) + + # RMEM fp32 -> GMEM fp32 (for final state store, runs once) + atom_state_r2g = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), acc_dtype, num_bits_per_copy=32 + ) + tiled_state_r2g = cute.make_tiled_copy_D(atom_state_r2g, tiled_state_t2r) + thr_state_r2g = tiled_state_r2g.get_slice(warpgroup_tidx) + + # --- Initial state load: GMEM fp32 -> RMEM -> state TMEM --- + gS_init = mS_fp32[(None, None, bid)] # [V=128, K=128] fp32 + tGR_tCgState_in = thr_state_g2r.partition_S(gS_init) + tGR_tCrState_in = thr_state_g2r.retile(tRrState) + # 1. GMEM fp32 -> RMEM fp32 (coalesced CopyUniversalOp) + cute.copy(tiled_state_g2r, tGR_tCgState_in, tGR_tCrState_in) + # 2. RMEM fp32 -> state TMEM (St32x32bOp per sub-tile) + num_state_subs = tRrState.shape[2] + for sub in cutlass.range(num_state_subs): + cute.copy(tiled_state_r2t, tRrState[None, 0, sub], tRT_tCtState[None, 0, sub]) + cute.arch.fence_view_async_tmem_store() + + # --- WG1 chunk loop --- + clk_wg1_start = cutlass.Int64(0) + if PROFILE_CLOCKS: + clk_wg1_start = read_clock() + for chunk_c in cutlass.range(num_chunks): + num_state_subs = tRrState.shape[2] + _sub_tile_size = cute.size(tRrState.shape[0]) + + # ---- Part 1: State readout + decay ---- + # Wait previous MMA6 (chunk > 0). Chunk 0: state from init. + if chunk_c > 0: + cute.arch.mbarrier_wait(mma6_done_mbar, phase=(chunk_c - 1) % 2) + + # Step A: State TMEM -> RMEM(fp32) -> bf16 -> sST SMEM + # Per-sub streaming: only 1 sub-tile in RMEM at a time (32 regs, not 128). + # tRrState[sub] is overwritten each iteration — NOT held across subs. + for sub in cutlass.range(num_state_subs): + cute.copy(tiled_state_t2r, tTR_tCtState[None, 0, sub], tRrState[None, 0, sub]) + tRrState_bf16[None, 0, sub].store(tRrState[None, 0, sub].load().to(mma_dtype)) + cute.copy( + tiled_state_r2s, tCrState_bf16[None, 0, sub], tCsState_inp[None, 0, sub, 0] + ) + cute.arch.fence_view_async_shared() + # Signal MMA: sST bf16 ready for MMA4/MMA3 (non-blocking) + cute.arch.mbarrier_arrive(st_ready_mbar) + + # Step B: Wait K1 to finish gk_last (current chunk) + cute.arch.mbarrier_wait(gk_last_ready_mbar, phase=chunk_c % 2) + + # Step C+D: Re-read state from TMEM, decay, write back. + # tRrState was overwritten per-sub in Step A, so re-read is needed. + # Cost: 4 extra TMEM reads (~40 cycles total) vs saving 96 spilled regs. + for sub in cutlass.range(num_state_subs): + # Re-read this sub-tile from TMEM (only 32 regs live) + cute.copy(tiled_state_t2r, tTR_tCtState[None, 0, sub], tRrState[None, 0, sub]) + # Decay by gk_last + for i in cutlass.range(_sub_tile_size): + coord = tTR_tCcState[i, 0, sub] + k_idx = coord[1] # N dimension = K in (V,K) convention + gk_val = cutlass.Float32(sGkLast[k_idx]) + tRrState[i, 0, sub] = tRrState[i, 0, sub] * gk_val + # Write decayed sub-tile back to TMEM immediately + cute.copy(tiled_state_r2t, tRrState[None, 0, sub], tRT_tCtState[None, 0, sub]) + cute.arch.fence_view_async_tmem_store() + # Signal MMA: decayed state in TMEM ready for MMA6 (non-blocking) + cute.arch.mbarrier_arrive(state_decayed_mbar) + + # ---- K4 phase: W/NV/O readout ---- + phase_nbar.arrive_and_wait() # K123->K4 + + tRrR = cute.make_rmem_tensor_like(tTR_cId, acc_dtype) + tRrR_out = cute.make_rmem_tensor_like(tRrR, mma_dtype) + tCrR_k = tiled_r2s_k.retile(tRrR_out) + num_subs = tRrR.shape[2] + + # W readout + wh = w_cons.wait_and_advance() + for sub in cutlass.range(num_subs): + cute.copy(tiled_t2r, tTR_W[None, 0, sub], tRrR[None, 0, sub]) + tRrR_out[None, 0, sub].store(tRrR[None, 0, sub].load().to(mma_dtype)) + cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsW[None, 0, sub, 0]) + cute.arch.fence_view_async_tmem_load() + wh.release() + cute.arch.fence_view_async_shared() + sW_ready_nbar.arrive_and_wait() + + # NV readout + nvh = nv_cons.wait_and_advance() + for sub in cutlass.range(num_subs): + cute.copy(tiled_t2r, tTR_NV[None, 0, sub], tRrR[None, 0, sub]) + tRrR_out[None, 0, sub].store(tRrR[None, 0, sub].load().to(mma_dtype)) + cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsNV[None, 0, sub, 0]) + cute.arch.fence_view_async_tmem_load() + nvh.release() + cute.arch.fence_view_async_shared() + sNV_ready_nbar.arrive_and_wait() + + # O readout + oh = o_cons.wait_and_advance() + for sub in cutlass.range(num_subs): + cute.copy(tiled_t2r, tTR_O[None, 0, sub], tRrR[None, 0, sub]) + tRrR_out[None, 0, sub].store(tRrR[None, 0, sub].load().to(mma_dtype)) + cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsO[None, 0, sub, 0]) + cute.arch.fence_view_async_tmem_load() + oh.release() + cute.arch.fence_view_async_shared() + store_nbar.arrive_and_wait() + + phase_nbar.arrive_and_wait() # K4->next + + # --- Wait last MMA6, final state store --- + cute.arch.mbarrier_wait(mma6_done_mbar, phase=(num_chunks - 1) % 2) + + # TMEM -> RMEM fp32 + gS_out = mS_fp32[(None, None, bid)] # [V=128, K=128] fp32 + tGR_tCgState_out = thr_state_r2g.partition_D(gS_out) + tGR_tCrState_out = thr_state_r2g.retile(tRrState) + num_state_subs_final = tRrState.shape[2] + for sub in cutlass.range(num_state_subs_final): + cute.copy(tiled_state_t2r, tTR_tCtState[None, 0, sub], tRrState[None, 0, sub]) + cute.arch.fence_view_async_tmem_load() + # RMEM fp32 -> GMEM fp32 (coalesced CopyUniversalOp) + for sub in cutlass.range(num_state_subs_final): + cute.copy( + tiled_state_r2g, tGR_tCrState_out[None, 0, sub], tGR_tCgState_out[None, 0, sub] + ) + + # Signal MMA warp: safe to dealloc TMEM + cute.arch.mbarrier_arrive(final_state_done_mbar) + + if PROFILE_CLOCKS: + clk_wg1_end = read_clock() + if warpgroup_tidx == 0: + mClocks[8, bid] = clk_wg1_end - clk_wg1_start + + # ============================================================== + # K1 (warps 8-11): TMA + gate activation + cumsum + KG/KS/QS + # ============================================================== + elif warp_idx >= K1_FIRST_WARP and warp_idx < K1_FIRST_WARP + NUM_K1_WARPS: + # K1 per-warp setup — each thread handles 1 column, all 64 rows sequentially + clk_k1_start = cutlass.Int64(0) + if PROFILE_CLOCKS: + clk_k1_start = read_clock() + k1_warp = warp_idx - K1_FIRST_WARP + my_col = k1_warp * WARP_SIZE + lane_id # 0-127, one column per thread + + # Per-head constants (invariant across chunks) + exp_A = cute.exp(mA_log[i_h], fastmath=True) + cumsum_scale = cutlass.Float32(RCP_LN2) + rBias_val = cutlass.Float32(0.0) + if HAS_BIAS: + rBias_val = mDtBias[i_h, my_col].to(cutlass.Float32) + + # Zero sAQC once (upper triangle stays zero across all chunks) + _aqc_zero = cutlass.BFloat16(0.0) + for ri in cutlass.range(BC): + aqc_row = k1_warp * BC + ri + c0 = lane_id * 2 + c1 = lane_id * 2 + 1 + sAQC[aqc_row + (c0 % 16) * 64, 0, c0 // 16, 0] = _aqc_zero + sAQC[aqc_row + (c1 % 16) * 64, 0, c1 // 16, 0] = _aqc_zero + inv_internal_barrier() + + for chunk_c in cutlass.range(num_chunks): + c_phase = chunk_c % 2 + chunk_start = i_b * num_chunks * BT + chunk_c * BT + chunk_start_local = chunk_c * BT + + # TMA load q, k, g for this chunk + if warp_idx == K1_FIRST_WARP: + gQ_head = tma_tensor_q_k123[(None, None, i_h)] + gK_head = tma_tensor_k_k123[(None, None, i_h)] + gG_head = tma_tensor_g_k123[(None, None, i_h)] + + if lane_id == 0: + cute.arch.mbarrier_expect_tx(k123_tma_mbar, k123_tma_bytes) + + sQ_s = sQ_k123 + gQ_s = cute.local_tile( + cute.domain_offset((chunk_start, 0), gQ_head), (BT, K_DIM), (0, 0) + ) + ts_q, tg_q = cpasync.tma_partition( + tma_atom_q_k123, + 0, + cute.make_layout(1), + cute.group_modes(sQ_s, 0, 2), + cute.group_modes(gQ_s, 0, 2), + ) + cute.copy(tma_atom_q_k123, tg_q, ts_q, tma_bar_ptr=k123_tma_mbar) + + sK_s = sK_k123 + gK_s = cute.local_tile( + cute.domain_offset((chunk_start, 0), gK_head), (BT, K_DIM), (0, 0) + ) + ts_k, tg_k = cpasync.tma_partition( + tma_atom_k_k123, + 0, + cute.make_layout(1), + cute.group_modes(sK_s, 0, 2), + cute.group_modes(gK_s, 0, 2), + ) + cute.copy(tma_atom_k_k123, tg_k, ts_k, tma_bar_ptr=k123_tma_mbar) + + sG_s = sG_k123 + gG_s = cute.local_tile( + cute.domain_offset((chunk_start, 0), gG_head), (BT, K_DIM), (0, 0) + ) + ts_g, tg_g = cpasync.tma_partition( + tma_atom_g_k123, + 0, + cute.make_layout(1), + cute.group_modes(sG_s, 0, 2), + cute.group_modes(gG_s, 0, 2), + ) + cute.copy(tma_atom_g_k123, tg_g, ts_g, tma_bar_ptr=k123_tma_mbar) + + if lane_id == 0: + cute.arch.mbarrier_arrive(k123_tma_mbar) + + # Wait TMA + cute.arch.mbarrier_wait(k123_tma_mbar, phase=c_phase) + + # Load beta[chunk] into SMEM (first K1 warp, 32 threads load 64 bf16 in 2 iters) + if warp_idx == K1_FIRST_WARP: + for bi in cutlass.range_constexpr(2): + beta_t = bi * WARP_SIZE + lane_id # 0..63 + sBeta[beta_t] = mBeta[i_b, chunk_start_local + beta_t, i_h] + + # ============================================================ + # Fused gate activation + cumsum + KG/KS/QS scaling + # Each thread processes 1 column across all 64 rows sequentially. + # No rGact, no sPartialLast, no cross-warp shuffle, no barriers. + # ============================================================ + + # Fused Pass: gate activation → cumsum → sGcum (single pass, 64 rows) + # 16 dynamic iterations × 4 constexpr unrolled = 64 rows + running_sum = cutlass.Float32(0.0) + for row_base in cutlass.range(BT // 4): + for ri in cutlass.range_constexpr(4): + row = row_base * 4 + ri + g_val = sG_k123[row, my_col].to(cutlass.Float32) + if HAS_BIAS: + g_val = g_val + rBias_val + g_activated = cutlass.Float32(0.0) + if USE_SAFE_GATE: + sigmoid_g = fast_rcp( + cutlass.Float32(1.0) + cute.exp2(-exp_A * g_val * LOG2E, fastmath=True) + ) + g_activated = lower_bound * sigmoid_g + else: + softplus_g = ( + cute.log2( + cutlass.Float32(1.0) + cute.exp2(g_val * LOG2E, fastmath=True), + fastmath=True, + ) + * LN2 + ) + g_activated = -exp_A * softplus_g + running_sum = running_sum + g_activated + sGcum[row, my_col] = running_sum * cumsum_scale + + # Signal K2: sGcum ready (all 64 rows written) + cute.arch.mbarrier_arrive(k123_k1_done_mbar) + + # KG/KS/QS scaling (reads sGcum + sK + sQ, writes Zone B) + # 16 dynamic × 4 constexpr = 64 rows + gk_last_cs = running_sum * cumsum_scale # total cumsum for this column + for row_base in cutlass.range(BT // 4): + for ri in cutlass.range_constexpr(4): + row = row_base * 4 + ri + cs = sGcum[row, my_col] + k_val = sK_k123[row, my_col].to(cutlass.Float32) + q_val = sQ_k123[row, my_col].to(cutlass.Float32) + exp2_cs = cute.exp2(cs, fastmath=True) + exp2_kg = cute.exp2(gk_last_cs - cs, fastmath=True) + # Write directly to Zone B swizzled layout + sKG[my_col + (row % 16) * 128, 0, row // 16, 0] = (k_val * exp2_kg).to( + cutlass.BFloat16 + ) + sKS[my_col + (row % 16) * 128, 0, row // 16, 0] = ( + cutlass.Float32(-1.0) * k_val * exp2_cs + ).to(cutlass.BFloat16) + sQS[row + (my_col % 16) * 64, 0, my_col // 16, 0] = ( + q_val * exp2_cs * scale + ).to(cutlass.BFloat16) + + # gkLast: all threads write their own column + sGkLast[my_col] = cute.exp2(gk_last_cs, fastmath=True) + + # Signal readout WG: gk_last is in SMEM (K1 arrive, non-blocking) + cute.arch.mbarrier_arrive(gk_last_ready_mbar) + + # ============================================================ + # K3 phase: wait K2, then inversion + store to Zone B + # (merged into K1 warps — K1 is idle here, K3 reuses same 4 warps) + # ============================================================ + store_warp = k1_warp # k1_warp 0-3 maps directly to store_warp 0-3 + + cute.arch.mbarrier_wait(k123_mma_done_mbar, phase=c_phase) + + # Write lower-triangle tiles with actual A_qk values + # (sAQC zeroed once before chunk loop — upper triangle stays zero) + for tile_idx in cutlass.range_constexpr(NUM_TILES): + i_q = _TILE_IQ[tile_idx] + i_k = _TILE_IK[tile_idx] + is_diag = _TILE_IQ[tile_idx] == _TILE_IK[tile_idx] + aqk_col_base = tile_idx * BC + + for ri in cutlass.range(BC // NUM_STORE_WARPS): + local_row = store_warp * (BC // NUM_STORE_WARPS) + ri + if lane_id < BC: + local_col = lane_id + aqk_val = sAqk_k123[local_row, aqk_col_base + local_col] + if is_diag and local_row < local_col: + aqk_val = cutlass.BFloat16(0.0) + aqc_m = i_q * BC + local_row + aqc_k = i_k * BC + local_col + sAQC[aqc_m + (aqc_k % 16) * 64, 0, aqc_k // 16, 0] = aqk_val + + # Akk inversion: 4 stages + if store_warp == 0: + _invert_diag(sAkk_k123, lane_id // 16, lane_id) + if store_warp == 1: + _invert_diag(sAkk_k123, 2 + lane_id // 16, lane_id) + inv_internal_barrier() + + # Stage 2: Ai10, Ai21, Ai32 + if store_warp == 0: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk_k123, 1, 1, 0, 1, lane_id) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B( + sAkk_k123, 0, 0, t0, t2, t1, t3, t4, t6, t5, t7, lane_id + ) + _store_neg_C(sAkk_k123, 1, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + if store_warp == 1: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk_k123, 2, 2, 1, 2, lane_id) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B( + sAkk_k123, 1, 1, t0, t2, t1, t3, t4, t6, t5, t7, lane_id + ) + _store_neg_C(sAkk_k123, 2, 1, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + if store_warp == 2: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk_k123, 3, 3, 2, 3, lane_id) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B( + sAkk_k123, 2, 2, t0, t2, t1, t3, t4, t6, t5, t7, lane_id + ) + _store_neg_C(sAkk_k123, 3, 2, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + inv_internal_barrier() + + # Stage 3: Ai20, Ai31 + _zz = cutlass.Float32(0.0) + t0 = _zz + t1 = _zz + t2 = _zz + t3 = _zz + t4 = _zz + t5 = _zz + t6 = _zz + t7 = _zz + if store_warp == 0: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk_k123, 0, 2, 0, 0, lane_id) + if store_warp == 2: + s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk_k123, 1, 2, 1, 0, lane_id) + _store_C_temp(sTemp, 0, s0, s1, s2, s3, s4, s5, s6, s7, lane_id) + if store_warp == 1: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk_k123, 1, 3, 1, 1, lane_id) + if store_warp == 3: + s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk_k123, 2, 3, 2, 1, lane_id) + _store_C_temp(sTemp, 1, s0, s1, s2, s3, s4, s5, s6, s7, lane_id) + inv_internal_barrier() + if store_warp == 0: + e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 0, lane_id) + t0 = t0 + e0 + t1 = t1 + e1 + t2 = t2 + e2 + t3 = t3 + e3 + t4 = t4 + e4 + t5 = t5 + e5 + t6 = t6 + e6 + t7 = t7 + e7 + sb = _shuffle_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7, lane_id) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A( + sAkk_k123, 2, 2, sb[0], sb[1], sb[4], sb[5], sb[2], sb[3], sb[6], sb[7], lane_id + ) + _store_neg_C(sAkk_k123, 2, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + if store_warp == 1: + e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 1, lane_id) + t0 = t0 + e0 + t1 = t1 + e1 + t2 = t2 + e2 + t3 = t3 + e3 + t4 = t4 + e4 + t5 = t5 + e5 + t6 = t6 + e6 + t7 = t7 + e7 + sb = _shuffle_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7, lane_id) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A( + sAkk_k123, 3, 3, sb[0], sb[1], sb[4], sb[5], sb[2], sb[3], sb[6], sb[7], lane_id + ) + _store_neg_C(sAkk_k123, 3, 1, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + inv_internal_barrier() + + # Stage 4: Ai30 + t0 = _zz + t1 = _zz + t2 = _zz + t3 = _zz + t4 = _zz + t5 = _zz + t6 = _zz + t7 = _zz + if store_warp == 0: + t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk_k123, 0, 3, 0, 0, lane_id) + if store_warp == 1: + s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk_k123, 1, 3, 1, 0, lane_id) + _store_C_temp(sTemp, 0, s0, s1, s2, s3, s4, s5, s6, s7, lane_id) + if store_warp == 2: + s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk_k123, 2, 3, 2, 0, lane_id) + _store_C_temp(sTemp, 1, s0, s1, s2, s3, s4, s5, s6, s7, lane_id) + inv_internal_barrier() + if store_warp == 0: + e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 0, lane_id) + t0 = t0 + e0 + t1 = t1 + e1 + t2 = t2 + e2 + t3 = t3 + e3 + t4 = t4 + e4 + t5 = t5 + e5 + t6 = t6 + e6 + t7 = t7 + e7 + e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 1, lane_id) + t0 = t0 + e0 + t1 = t1 + e1 + t2 = t2 + e2 + t3 = t3 + e3 + t4 = t4 + e4 + t5 = t5 + e5 + t6 = t6 + e6 + t7 = t7 + e7 + sb = _shuffle_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7, lane_id) + r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A( + sAkk_k123, 3, 3, sb[0], sb[1], sb[4], sb[5], sb[2], sb[3], sb[6], sb[7], lane_id + ) + _store_neg_C(sAkk_k123, 3, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + inv_internal_barrier() + + # Write inverted Akk * beta to K4's swizzled SMEM (sAB in Zone B) + inv_row_start = store_warp * 16 + for ri in cutlass.range(BC): + inv_row = inv_row_start + ri + c0 = lane_id * 2 + c1 = lane_id * 2 + 1 + beta_c0 = sBeta[c0].to(cutlass.Float32) + beta_c1 = sBeta[c1].to(cutlass.Float32) + v0 = ( + cutlass.Float32(sAkk_k123[inv_row, c0]) + * cutlass.Float32(inv_row >= c0) + * beta_c0 + ) + v1 = ( + cutlass.Float32(sAkk_k123[inv_row, c1]) + * cutlass.Float32(inv_row >= c1) + * beta_c1 + ) + sAB[inv_row + (c0 % 16) * 64, 0, c0 // 16, 0] = v0.to(cutlass.BFloat16) + sAB[inv_row + (c1 % 16) * 64, 0, c1 // 16, 0] = v1.to(cutlass.BFloat16) + + # K1+K3 done, idle during K4 phase + phase_nbar.arrive_and_wait() # K123->K4 + phase_nbar.arrive_and_wait() # K4->next + + if PROFILE_CLOCKS: + clk_k1_end = read_clock() + if k1_warp == 0 and lane_id == 0: + mClocks[9, bid] = clk_k1_end - clk_k1_start + + # ============================================================== + # K2 (warps 12-19): intra sub-chunk attention MMA + # 8 warps: warps 0-5 do 1 tile each, warps 6-7 do 2 tiles each (row 3) + # ============================================================== + elif warp_idx >= K2_FIRST_WARP and warp_idx < K2_FIRST_WARP + NUM_MMA_WARPS: + # K2 warp decode + clk_k2_start = cutlass.Int64(0) + if PROFILE_CLOCKS: + clk_k2_start = read_clock() + mma_warp = warp_idx - K2_FIRST_WARP + # Tile A assignment (all 8 warps) + my_i_q = cutlass.Int32(0) + my_i_k = cutlass.Int32(0) + if mma_warp < 1: + my_i_q = cutlass.Int32(0) + my_i_k = mma_warp + elif mma_warp < 3: + my_i_q = cutlass.Int32(1) + my_i_k = mma_warp - 1 + elif mma_warp < 6: + my_i_q = cutlass.Int32(2) + my_i_k = mma_warp - 3 + else: + # Warps 6-7: first tile of row 3 pair + # W6: (3,0), W7: (3,2) + my_i_q = cutlass.Int32(3) + my_i_k = (mma_warp - 6) * 2 + # Tile B assignment (warps 6-7 only): i_k_b = i_k_a + 1 + # W6: (3,1), W7: (3,3) + my_i_k_b = my_i_k + cutlass.Int32(1) + + for chunk_c in cutlass.range(num_chunks): + c_phase = chunk_c % 2 + + cute.arch.mbarrier_wait(k123_k1_done_mbar, phase=c_phase) + + # --- Tile A (all 8 warps) --- + q_row_base = my_i_q * BC + k_row_base = my_i_k * BC + # tile_col_base maps to sAqk slot index: + # warps 0-5: 1:1 mapping (tile slot = mma_warp) + # warps 6-7: tile slot = 6+(mma_warp-6)*2 = 6 or 8 + tile_col_base = mma_warp * BC + if mma_warp >= 6: + tile_col_base = (6 + (mma_warp - 6) * 2) * BC + akk_row_base = k_row_base + akk_col_base = q_row_base + norm_row = q_row_base + if my_i_q == my_i_k: + norm_row = q_row_base + cutlass.Int32(BC // 2) + + group_id = lane_id // 4 + tid_in_group = lane_id % 4 + row0, row1 = group_id, group_id + 8 + + thr_mma = tiled_mma_k2.get_slice(lane_id) + thr_copy_A = tiled_copy_mma_A.get_slice(lane_id) + thr_copy_B = tiled_copy_mma_B.get_slice(lane_id) + thr_copy_Gn = tiled_copy_Gcum_norm.get_slice(tid_in_group) + thr_copy_Ggate = tiled_copy_Gcum_gate.get_slice(lane_id) + + beta_row0 = sBeta[q_row_base + row0].to(cutlass.Float32) + beta_row1 = sBeta[q_row_base + row1].to(cutlass.Float32) + + _z = cutlass.Float32(0.0) + 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 + acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = _z, _z, _z, _z + acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = _z, _z, _z, _z + + for k_block in cutlass.range_constexpr(NUM_MMA_K_TILES): + sQ_tile = cute.local_tile(sQ_k123, tiler=(16, 8), coord=(my_i_q, k_block)) + tCrQ = tiled_mma_k2.make_fragment_A(thr_mma.partition_A(sQ_tile)) + cute.copy( + tiled_copy_mma_A, thr_copy_A.partition_S(sQ_tile), thr_copy_A.retile(tCrQ) + ) + + sKq_tile = cute.local_tile(sK_k123, tiler=(16, 8), coord=(my_i_q, k_block)) + tCrKq = tiled_mma_k2.make_fragment_A(thr_mma.partition_A(sKq_tile)) + cute.copy( + tiled_copy_mma_A, thr_copy_A.partition_S(sKq_tile), thr_copy_A.retile(tCrKq) + ) + + sGn_tile = cute.local_tile(sGcum, tiler=(1, 8), coord=(norm_row, k_block)) + tCsGn = thr_copy_Gn.partition_S(sGn_tile) + tCrGn = cute.make_fragment_like(tCsGn, cutlass.Float32) + cute.copy(tiled_copy_Gcum_norm, tCsGn, thr_copy_Gn.retile(tCrGn)) + g_norm_0 = tCrGn[0] + g_norm_1 = tCrGn[1] + + sGq_tile = cute.local_tile(sGcum, tiler=(16, 8), coord=(my_i_q, k_block)) + tCrGq = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGq_tile)) + cute.copy( + tiled_copy_Gcum_gate, + thr_copy_Ggate.partition_S(sGq_tile), + thr_copy_Ggate.retile(tCrGq), + ) + gate_q_0 = cute.exp2(tCrGq[0] - g_norm_0, fastmath=True) + gate_q_1 = cute.exp2(tCrGq[1] - g_norm_1, fastmath=True) + gate_q_2 = cute.exp2(tCrGq[2] - g_norm_0, fastmath=True) + gate_q_3 = cute.exp2(tCrGq[3] - g_norm_1, fastmath=True) + + qa0 = tCrQ[0].to(cutlass.Float32) * gate_q_0 + qa1 = tCrQ[2].to(cutlass.Float32) * gate_q_2 + qa2 = tCrQ[1].to(cutlass.Float32) * gate_q_1 + qa3 = tCrQ[3].to(cutlass.Float32) * gate_q_3 + ka0 = tCrKq[0].to(cutlass.Float32) * gate_q_0 + ka1 = tCrKq[2].to(cutlass.Float32) * gate_q_2 + ka2 = tCrKq[1].to(cutlass.Float32) * gate_q_1 + ka3 = tCrKq[3].to(cutlass.Float32) * gate_q_3 + + sK_tile_n0 = cute.local_tile(sK_k123, tiler=(8, 8), coord=(my_i_k * 2, k_block)) + tCrK_n0 = tiled_mma_k2.make_fragment_B(thr_mma.partition_B(sK_tile_n0)) + cute.copy( + tiled_copy_mma_B, thr_copy_B.partition_S(sK_tile_n0), thr_copy_B.retile(tCrK_n0) + ) + + sK_tile_n1 = cute.local_tile(sK_k123, tiler=(8, 8), coord=(my_i_k * 2 + 1, k_block)) + tCrK_n1 = tiled_mma_k2.make_fragment_B(thr_mma.partition_B(sK_tile_n1)) + cute.copy( + tiled_copy_mma_B, thr_copy_B.partition_S(sK_tile_n1), thr_copy_B.retile(tCrK_n1) + ) + + sGk_tile = cute.local_tile(sGcum, tiler=(16, 8), coord=(my_i_k, k_block)) + tCrGk = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGk_tile)) + cute.copy( + tiled_copy_Gcum_gate, + thr_copy_Ggate.partition_S(sGk_tile), + thr_copy_Ggate.retile(tCrGk), + ) + gk_n0_0 = cute.exp2(g_norm_0 - tCrGk[0], fastmath=True) + gk_n0_1 = cute.exp2(g_norm_1 - tCrGk[1], fastmath=True) + k_n0_b0 = tCrK_n0[0].to(cutlass.Float32) * gk_n0_0 + k_n0_b1 = tCrK_n0[1].to(cutlass.Float32) * gk_n0_1 + gk_n1_0 = cute.exp2(g_norm_0 - tCrGk[2], fastmath=True) + gk_n1_1 = cute.exp2(g_norm_1 - tCrGk[3], fastmath=True) + k_n1_b0 = tCrK_n1[0].to(cutlass.Float32) * gk_n1_0 + k_n1_b1 = tCrK_n1[1].to(cutlass.Float32) * gk_n1_1 + + acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = mma_tf32_m16n8k8( + qa0, + qa1, + qa2, + qa3, + k_n0_b0, + k_n0_b1, + acc_aqk_n0_0, + acc_aqk_n0_1, + acc_aqk_n0_2, + acc_aqk_n0_3, + ) + acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = mma_tf32_m16n8k8( + qa0, + qa1, + qa2, + qa3, + k_n1_b0, + k_n1_b1, + acc_aqk_n1_0, + acc_aqk_n1_1, + acc_aqk_n1_2, + acc_aqk_n1_3, + ) + acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = mma_tf32_m16n8k8( + ka0, + ka1, + ka2, + ka3, + k_n0_b0, + k_n0_b1, + acc_akk_n0_0, + acc_akk_n0_1, + acc_akk_n0_2, + acc_akk_n0_3, + ) + acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = mma_tf32_m16n8k8( + ka0, + ka1, + ka2, + ka3, + k_n1_b0, + k_n1_b1, + acc_akk_n1_0, + acc_akk_n1_1, + acc_akk_n1_2, + acc_akk_n1_3, + ) + + # Store tile A + col0, col1 = tid_in_group * 2, tid_in_group * 2 + 1 + col2, col3 = 8 + tid_in_group * 2, 8 + tid_in_group * 2 + 1 + + sAqk_k123[row0, tile_col_base + col0] = (acc_aqk_n0_0 * scale).to(cutlass.BFloat16) + sAqk_k123[row0, tile_col_base + col1] = (acc_aqk_n0_1 * scale).to(cutlass.BFloat16) + sAqk_k123[row1, tile_col_base + col0] = (acc_aqk_n0_2 * scale).to(cutlass.BFloat16) + sAqk_k123[row1, tile_col_base + col1] = (acc_aqk_n0_3 * scale).to(cutlass.BFloat16) + sAqk_k123[row0, tile_col_base + col2] = (acc_aqk_n1_0 * scale).to(cutlass.BFloat16) + sAqk_k123[row0, tile_col_base + col3] = (acc_aqk_n1_1 * scale).to(cutlass.BFloat16) + sAqk_k123[row1, tile_col_base + col2] = (acc_aqk_n1_2 * scale).to(cutlass.BFloat16) + sAqk_k123[row1, tile_col_base + col3] = (acc_aqk_n1_3 * scale).to(cutlass.BFloat16) + + sAkk_k123[akk_row_base + row0, akk_col_base + col0] = acc_akk_n0_0 * beta_row0 + sAkk_k123[akk_row_base + row0, akk_col_base + col1] = acc_akk_n0_1 * beta_row0 + sAkk_k123[akk_row_base + row1, akk_col_base + col0] = acc_akk_n0_2 * beta_row1 + sAkk_k123[akk_row_base + row1, akk_col_base + col1] = acc_akk_n0_3 * beta_row1 + sAkk_k123[akk_row_base + row0, akk_col_base + col2] = acc_akk_n1_0 * beta_row0 + sAkk_k123[akk_row_base + row0, akk_col_base + col3] = acc_akk_n1_1 * beta_row0 + sAkk_k123[akk_row_base + row1, akk_col_base + col2] = acc_akk_n1_2 * beta_row1 + sAkk_k123[akk_row_base + row1, akk_col_base + col3] = acc_akk_n1_3 * beta_row1 + + # --- Tile B (warps 6-7 only): second tile of row 3 pair --- + # W6: (3,1), W7: (3,3). Accumulators reused (zero extra reg pressure). + if mma_warp >= 6: + k_row_base_b = my_i_k_b * BC + tile_col_base_b = tile_col_base + BC # tile B = tile A + 1 slot + akk_row_base_b = k_row_base_b + norm_row_b = q_row_base + if my_i_q == my_i_k_b: + norm_row_b = q_row_base + cutlass.Int32(BC // 2) + + 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 + acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = _z, _z, _z, _z + acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = _z, _z, _z, _z + + for k_block in cutlass.range_constexpr(NUM_MMA_K_TILES): + sQ_tile = cute.local_tile(sQ_k123, tiler=(16, 8), coord=(my_i_q, k_block)) + tCrQ = tiled_mma_k2.make_fragment_A(thr_mma.partition_A(sQ_tile)) + cute.copy( + tiled_copy_mma_A, thr_copy_A.partition_S(sQ_tile), thr_copy_A.retile(tCrQ) + ) + + sKq_tile = cute.local_tile(sK_k123, tiler=(16, 8), coord=(my_i_q, k_block)) + tCrKq = tiled_mma_k2.make_fragment_A(thr_mma.partition_A(sKq_tile)) + cute.copy( + tiled_copy_mma_A, thr_copy_A.partition_S(sKq_tile), thr_copy_A.retile(tCrKq) + ) + + sGn_tile = cute.local_tile(sGcum, tiler=(1, 8), coord=(norm_row_b, k_block)) + tCsGn = thr_copy_Gn.partition_S(sGn_tile) + tCrGn = cute.make_fragment_like(tCsGn, cutlass.Float32) + cute.copy(tiled_copy_Gcum_norm, tCsGn, thr_copy_Gn.retile(tCrGn)) + g_norm_0 = tCrGn[0] + g_norm_1 = tCrGn[1] + + sGq_tile = cute.local_tile(sGcum, tiler=(16, 8), coord=(my_i_q, k_block)) + tCrGq = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGq_tile)) + cute.copy( + tiled_copy_Gcum_gate, + thr_copy_Ggate.partition_S(sGq_tile), + thr_copy_Ggate.retile(tCrGq), + ) + gate_q_0 = cute.exp2(tCrGq[0] - g_norm_0, fastmath=True) + gate_q_1 = cute.exp2(tCrGq[1] - g_norm_1, fastmath=True) + gate_q_2 = cute.exp2(tCrGq[2] - g_norm_0, fastmath=True) + gate_q_3 = cute.exp2(tCrGq[3] - g_norm_1, fastmath=True) + + qa0 = tCrQ[0].to(cutlass.Float32) * gate_q_0 + qa1 = tCrQ[2].to(cutlass.Float32) * gate_q_2 + qa2 = tCrQ[1].to(cutlass.Float32) * gate_q_1 + qa3 = tCrQ[3].to(cutlass.Float32) * gate_q_3 + ka0 = tCrKq[0].to(cutlass.Float32) * gate_q_0 + ka1 = tCrKq[2].to(cutlass.Float32) * gate_q_2 + ka2 = tCrKq[1].to(cutlass.Float32) * gate_q_1 + ka3 = tCrKq[3].to(cutlass.Float32) * gate_q_3 + + sK_tile_n0 = cute.local_tile( + sK_k123, tiler=(8, 8), coord=(my_i_k_b * 2, k_block) + ) + tCrK_n0 = tiled_mma_k2.make_fragment_B(thr_mma.partition_B(sK_tile_n0)) + cute.copy( + tiled_copy_mma_B, + thr_copy_B.partition_S(sK_tile_n0), + thr_copy_B.retile(tCrK_n0), + ) + + sK_tile_n1 = cute.local_tile( + sK_k123, tiler=(8, 8), coord=(my_i_k_b * 2 + 1, k_block) + ) + tCrK_n1 = tiled_mma_k2.make_fragment_B(thr_mma.partition_B(sK_tile_n1)) + cute.copy( + tiled_copy_mma_B, + thr_copy_B.partition_S(sK_tile_n1), + thr_copy_B.retile(tCrK_n1), + ) + + sGk_tile = cute.local_tile(sGcum, tiler=(16, 8), coord=(my_i_k_b, k_block)) + tCrGk = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGk_tile)) + cute.copy( + tiled_copy_Gcum_gate, + thr_copy_Ggate.partition_S(sGk_tile), + thr_copy_Ggate.retile(tCrGk), + ) + gk_n0_0 = cute.exp2(g_norm_0 - tCrGk[0], fastmath=True) + gk_n0_1 = cute.exp2(g_norm_1 - tCrGk[1], fastmath=True) + k_n0_b0 = tCrK_n0[0].to(cutlass.Float32) * gk_n0_0 + k_n0_b1 = tCrK_n0[1].to(cutlass.Float32) * gk_n0_1 + gk_n1_0 = cute.exp2(g_norm_0 - tCrGk[2], fastmath=True) + gk_n1_1 = cute.exp2(g_norm_1 - tCrGk[3], fastmath=True) + k_n1_b0 = tCrK_n1[0].to(cutlass.Float32) * gk_n1_0 + k_n1_b1 = tCrK_n1[1].to(cutlass.Float32) * gk_n1_1 + + acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = mma_tf32_m16n8k8( + qa0, + qa1, + qa2, + qa3, + k_n0_b0, + k_n0_b1, + acc_aqk_n0_0, + acc_aqk_n0_1, + acc_aqk_n0_2, + acc_aqk_n0_3, + ) + acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = mma_tf32_m16n8k8( + qa0, + qa1, + qa2, + qa3, + k_n1_b0, + k_n1_b1, + acc_aqk_n1_0, + acc_aqk_n1_1, + acc_aqk_n1_2, + acc_aqk_n1_3, + ) + acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = mma_tf32_m16n8k8( + ka0, + ka1, + ka2, + ka3, + k_n0_b0, + k_n0_b1, + acc_akk_n0_0, + acc_akk_n0_1, + acc_akk_n0_2, + acc_akk_n0_3, + ) + acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = mma_tf32_m16n8k8( + ka0, + ka1, + ka2, + ka3, + k_n1_b0, + k_n1_b1, + acc_akk_n1_0, + acc_akk_n1_1, + acc_akk_n1_2, + acc_akk_n1_3, + ) + + # Store tile B + sAqk_k123[row0, tile_col_base_b + col0] = (acc_aqk_n0_0 * scale).to( + cutlass.BFloat16 + ) + sAqk_k123[row0, tile_col_base_b + col1] = (acc_aqk_n0_1 * scale).to( + cutlass.BFloat16 + ) + sAqk_k123[row1, tile_col_base_b + col0] = (acc_aqk_n0_2 * scale).to( + cutlass.BFloat16 + ) + sAqk_k123[row1, tile_col_base_b + col1] = (acc_aqk_n0_3 * scale).to( + cutlass.BFloat16 + ) + sAqk_k123[row0, tile_col_base_b + col2] = (acc_aqk_n1_0 * scale).to( + cutlass.BFloat16 + ) + sAqk_k123[row0, tile_col_base_b + col3] = (acc_aqk_n1_1 * scale).to( + cutlass.BFloat16 + ) + sAqk_k123[row1, tile_col_base_b + col2] = (acc_aqk_n1_2 * scale).to( + cutlass.BFloat16 + ) + sAqk_k123[row1, tile_col_base_b + col3] = (acc_aqk_n1_3 * scale).to( + cutlass.BFloat16 + ) + + sAkk_k123[akk_row_base_b + row0, akk_col_base + col0] = acc_akk_n0_0 * beta_row0 + sAkk_k123[akk_row_base_b + row0, akk_col_base + col1] = acc_akk_n0_1 * beta_row0 + sAkk_k123[akk_row_base_b + row1, akk_col_base + col0] = acc_akk_n0_2 * beta_row1 + sAkk_k123[akk_row_base_b + row1, akk_col_base + col1] = acc_akk_n0_3 * beta_row1 + sAkk_k123[akk_row_base_b + row0, akk_col_base + col2] = acc_akk_n1_0 * beta_row0 + sAkk_k123[akk_row_base_b + row0, akk_col_base + col3] = acc_akk_n1_1 * beta_row0 + sAkk_k123[akk_row_base_b + row1, akk_col_base + col2] = acc_akk_n1_2 * beta_row1 + sAkk_k123[akk_row_base_b + row1, akk_col_base + col3] = acc_akk_n1_3 * beta_row1 + + cute.arch.mbarrier_arrive(k123_mma_done_mbar) + + # K2 idle during K4 phase + phase_nbar.arrive_and_wait() # K123->K4 + phase_nbar.arrive_and_wait() # K4->next + + if PROFILE_CLOCKS: + clk_k2_end = read_clock() + if mma_warp == 0 and lane_id == 0: + mClocks[10, bid] = clk_k2_end - clk_k2_start + + # ============================================================== + # MMA warp (warp 0): 6 MMAs per chunk + # ============================================================== + elif warp_idx == K4_MMA_WARP: + mc = (0, 0, 0, 0) + ml = cute.make_layout((1, 1, 1, 1)) + + # Clock accumulators (thread 0 only, i64) + clk_k123_acc = cutlass.Int64(0) + clk_k4_acc = cutlass.Int64(0) + clk0 = cutlass.Int64(0) + clk1 = cutlass.Int64(0) + clk2 = cutlass.Int64(0) + clk_mma1_acc = cutlass.Int64(0) + clk_mma2_acc = cutlass.Int64(0) + clk_wait_s_mma4_acc = cutlass.Int64(0) + clk_wait_w_mma3_acc = cutlass.Int64(0) + clk_wait_nv_mma5_acc = cutlass.Int64(0) + clk_mma6_drain_acc = cutlass.Int64(0) + clk_k4a = cutlass.Int64(0) + clk_k4b = cutlass.Int64(0) + + for chunk_c in cutlass.range(num_chunks): + chunk_start = i_b * num_chunks * BT + chunk_c * BT + + if PROFILE_CLOCKS: + clk0 = read_clock() + + # Wait for K123 to finish + phase_nbar.arrive_and_wait() # K123->K4 + + if PROFILE_CLOCKS: + clk1 = read_clock() + clk_k123_acc = clk_k123_acc + (clk1 - clk0) + + if DEBUG_K4_LEVEL == 0: + pass # Skip K4 entirely for debugging + else: + fA_kmn = thr_kmn.make_fragment_A(sAB) + fB_ks = thr_kmn.make_fragment_B(sKS) + fB_v = thr_kmn.make_fragment_B(sV_ext) + fA_w = thr_kmn.make_fragment_A(sW) + fB_s = thr_kmn.make_fragment_B(sST) + fA_q = thr_kmn.make_fragment_A(sQS) + fA_aqc = thr_kmn.make_fragment_A(sAQC) + fB_nv = thr_kmn.make_fragment_B(sNV_b) + fA_nv_mn = thr_mn.make_fragment_A(sNV_a) + fB_kg = thr_mn.make_fragment_B(sKG) + + if PROFILE_CLOCKS: + clk_k4a = read_clock() + + # MMA1: W = AB @ KS + w_h = w_prod.acquire_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, False) + for k in cutlass.range_constexpr(cute.size(sKS.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtW, + fA_kmn[dice + (0,)][None, None, k], + fB_ks[dice + (0,)][None, None, k], + tCtW, + ) + if k == 0: + tiled_mma_kmn.set(Field.ACCUMULATE, True) + w_h.commit() + + if PROFILE_CLOCKS: + clk_k4b = read_clock() + clk_mma1_acc = clk_mma1_acc + (clk_k4b - clk_k4a) + clk_k4a = clk_k4b + + # MMA2: U = AB @ V + vh = v_cons.wait_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, False) + for k in cutlass.range_constexpr(cute.size(sV_ext.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtNV, + fA_kmn[dice + (0,)][None, None, k], + fB_v[dice + (vh.index,)][None, None, k], + tCtNV, + ) + if k == 0: + tiled_mma_kmn.set(Field.ACCUMULATE, True) + vh.release() + + if PROFILE_CLOCKS: + clk_k4b = read_clock() + clk_mma2_acc = clk_mma2_acc + (clk_k4b - clk_k4a) + clk_k4a = clk_k4b + + # Wait WG1: sST bf16 ready in SMEM + cute.arch.mbarrier_wait(st_ready_mbar, phase=chunk_c % 2) + + # MMA4: OI = QS @ S + tiled_mma_kmn.set(Field.ACCUMULATE, False) + for k in cutlass.range_constexpr(cute.size(sST.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtO, + fA_q[dice + (0,)][None, None, k], + fB_s[dice + (0,)][None, None, k], + tCtO, + ) + if k == 0: + tiled_mma_kmn.set(Field.ACCUMULATE, True) + + if PROFILE_CLOCKS: + clk_k4b = read_clock() + clk_wait_s_mma4_acc = clk_wait_s_mma4_acc + (clk_k4b - clk_k4a) + clk_k4a = clk_k4b + + # Wait W readout + sW_ready_nbar.arrive_and_wait() + + # MMA3: NV += sW @ S (accumulate into U) + nv_h = nv_prod.acquire_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, True) + for k in cutlass.range_constexpr(cute.size(sST.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtNV, + fA_w[dice + (0,)][None, None, k], + fB_s[dice + (0,)][None, None, k], + tCtNV, + ) + nv_h.commit() + + if PROFILE_CLOCKS: + clk_k4b = read_clock() + clk_wait_w_mma3_acc = clk_wait_w_mma3_acc + (clk_k4b - clk_k4a) + clk_k4a = clk_k4b + + # Wait NV readout + sNV_ready_nbar.arrive_and_wait() + + # MMA5: O += AQC @ NV + o_h = o_prod.acquire_and_advance() + if not DEBUG_SKIP_MMA5: + tiled_mma_kmn.set(Field.ACCUMULATE, True) + for k in cutlass.range_constexpr(cute.size(sNV_b.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtO, + fA_aqc[dice + (0,)][None, None, k], + fB_nv[dice + (0,)][None, None, k], + tCtO, + ) + o_h.commit() + + if PROFILE_CLOCKS: + clk_k4b = read_clock() + clk_wait_nv_mma5_acc = clk_wait_nv_mma5_acc + (clk_k4b - clk_k4a) + clk_k4a = clk_k4b + + # Wait WG1: decayed state written back to TMEM + cute.arch.mbarrier_wait(state_decayed_mbar, phase=chunk_c % 2) + + # MMA6: accumulate SU = NV^T @ KG onto decayed state in TMEM@256 + tiled_mma_mn_mn.set(Field.ACCUMULATE, True) + for k in cutlass.range_constexpr(cute.size(sKG.shape[2])): + cute.gemm( + tiled_mma_mn_mn, + tCtState, + fA_nv_mn[dice + (0,)][None, None, k], + fB_kg[dice + (0,)][None, None, k], + tCtState, + ) + + # Drain W/NV/O UMMA pipelines, then sync with TMA+Readout + w_prod.tail() + nv_prod.tail() + o_prod.tail() + store_nbar.arrive_and_wait() + + # Signal WG1: MMA6 done (all chunks, including last) + cute.arch.mbarrier_arrive(mma6_done_mbar) + + if PROFILE_CLOCKS: + clk_k4b = read_clock() + clk_mma6_drain_acc = clk_mma6_drain_acc + (clk_k4b - clk_k4a) + + # K4->next chunk + phase_nbar.arrive_and_wait() + + if PROFILE_CLOCKS: + clk2 = read_clock() + clk_k4_acc = clk_k4_acc + (clk2 - clk1) + + # Store accumulated clocks (thread 0 = warp 0 = MMA warp) + if PROFILE_CLOCKS: + if tidx == 0: + mClocks[0, bid] = clk_k123_acc + mClocks[1, bid] = clk_k4_acc + mClocks[2, bid] = clk_mma1_acc + mClocks[3, bid] = clk_mma2_acc + mClocks[4, bid] = clk_wait_s_mma4_acc + mClocks[5, bid] = clk_wait_w_mma3_acc + mClocks[6, bid] = clk_wait_nv_mma5_acc + mClocks[7, bid] = clk_mma6_drain_acc + + # Wait for WG1 final state store, then dealloc TMEM + cute.arch.mbarrier_wait(final_state_done_mbar, phase=0) + cute.arch.relinquish_tmem_alloc_permit() + cute.arch.dealloc_tmem(tmem_ptr, 512) + + # ============================================================== + # TMA warp (warp 2): V early load + O store + # ============================================================== + elif warp_idx == K4_TMA_WARP: + mc = (0, 0, 0, 0) + ml = cute.make_layout((1, 1, 1, 1)) + clk_tma_start = cutlass.Int64(0) + if PROFILE_CLOCKS: + clk_tma_start = read_clock() + + for chunk_c in cutlass.range(num_chunks): + chunk_start = i_b * num_chunks * BT + chunk_c * BT + + # V early load (overlaps K123 phase) + gV_h = mV_nkl[(None, None, i_h)] + gV = cute.local_tile( + cute.domain_offset((0, chunk_start), gV_h), tiler=(N4, K4_K), coord=(0, 0) + ) + tBsV, tBgV = cpasync.tma_partition( + tma_atom_v, + mc, + ml, + cute.group_modes(sV_ext, 0, 3), + cute.group_modes(thr_kmn.partition_B(gV), 0, 3), + ) + vh = v_prod.acquire_and_advance() + cute.copy(tma_atom_v, tBgV, tBsV[None, vh.index], tma_bar_ptr=vh.barrier) + + # Wait for K123 to finish + phase_nbar.arrive_and_wait() # K123->K4 + + # Store O[c] to GMEM + gO_h = mOo[(None, None, i_h)] # [T_total, V_dim] + + # Wait for MMA + Readout to finish (O is in SMEM) + store_nbar.arrive_and_wait() + + gOo_c = cute.local_tile( + cute.domain_offset((chunk_start, 0), gO_h), tiler=(M4, N4), coord=(0, 0) + ) + sOt, gOt = cpasync.tma_partition( + tma_atom_o_st, mc, ml, cute.group_modes(sO_st, 0, 2), cute.group_modes(gOo_c, 0, 2) + ) + cute.copy(tma_atom_o_st, sOt, gOt) + + # Fence: wait for TMA to finish reading sO from SMEM + tma_store_fence() + + # K4->next chunk + phase_nbar.arrive_and_wait() + + # Drain V pipeline + v_prod.tail() + + if PROFILE_CLOCKS: + clk_tma_end = read_clock() + if lane_id == 0: + mClocks[12, bid] = clk_tma_end - clk_tma_start + + # ============================================================== + # Idle warps (1, 3): phase barriers only + # ============================================================== + else: + for chunk_c in cutlass.range(num_chunks): + phase_nbar.arrive_and_wait() # K123->K4 + phase_nbar.arrive_and_wait() # K4->next + + +# ============================================================================= +# Host Function +# ============================================================================= + + +def make_host_fn(has_bias=False, use_safe_gate=False, profile_clocks=False): + """Create the host function for the fused K1234 kernel. + + B, H, NC are all dynamic — one compilation serves all shapes. + Only mode flags (has_bias, use_safe_gate, profile_clocks) specialize at compile time. + """ + _HAS_BIAS = 1 if has_bias else 0 + _USE_SAFE_GATE = 1 if use_safe_gate else 0 + _PROFILE_CLOCKS = 1 if profile_clocks else 0 + + tile1 = (M4, N4, K4_K) + tile3 = (M4, N4, K4_K3) + tile6 = (M6, N6, K6) + + @cute.jit + def host_fn( + # K123 raw inputs + mQ, + mK, + mG, + mA_log, + mBeta, + scale_val, + # K4 inputs (raw pointers — host_fn creates 3D views internally) + mV_in, # [B, T, H, V_dim] bf16 + mO_out, # [B, T, H, V_dim] bf16 output + mS_fp32, # [B*H, K, V] fp32 state (init + final output) + # dt_bias / safe_gate + mDtBias, # [H, K] fp32 (or dummy [1,1] if no bias) + lower_bound_val, # float (0.0 if unused) + # Clock profiling + mClocks, # [2, B*H] i64 (or dummy if not profiling) + # Runtime dimensions (dynamic — no recompilation needed) + num_chunks: cutlass.Int32, + num_heads: cutlass.Int32, + batch_size: cutlass.Int32, + # Launch stream — runtime argument; launching on the DSL default + # stream races with the executor's non-blocking execution stream. + stream: cuda.CUstream, + ): + # Derive GMEM shapes from runtime dimensions + T_total = batch_size * num_chunks * BT + s_row = num_heads * K_DIM + s_col = 1 + s_h = K_DIM + BH = batch_size * num_heads + + # --- Common 3D layout for [T_total, K_DIM, H] tensors --- + view_layout_3d = cute.make_layout((T_total, K_DIM, num_heads), stride=(s_row, s_col, s_h)) + + # --- K4 MMA setup --- + mma_kmn = sm100_utils.make_trivial_tiled_mma( + mma_dtype, + OperandMajorMode.K, + OperandMajorMode.MN, + acc_dtype, + tcgen05.CtaGroup.ONE, + (M4, N4), + OperandSource.SMEM, + ) + mma_mn_mn = sm100_utils.make_trivial_tiled_mma( + mma_dtype, + OperandMajorMode.MN, + OperandMajorMode.MN, + acc_dtype, + tcgen05.CtaGroup.ONE, + (M6, N6), + OperandSource.SMEM, + ) + + sl_ab = sm100_utils.make_smem_layout_a(mma_kmn, tile1, mma_dtype, 1) + sl_ks = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1) + sl_v = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1) + sl_s = sm100_utils.make_smem_layout_b(mma_kmn, tile3, mma_dtype, 1) + sl_qs = sm100_utils.make_smem_layout_a(mma_kmn, tile3, mma_dtype, 1) + sl_aqc = sm100_utils.make_smem_layout_a(mma_kmn, tile1, mma_dtype, 1) + sl_kg = sm100_utils.make_smem_layout_b(mma_mn_mn, tile6, mma_dtype, 1) + + sl_readout_k = sm100_utils.make_smem_layout_a(mma_kmn, tile3, mma_dtype, 1) + sl_nv_b = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1) + sl_nv_a = sm100_utils.make_smem_layout_a(mma_mn_mn, tile6, mma_dtype, 1) + + # --- K4 3D GMEM views (V, O, S_fp32 — AB/KS/QS/AQC/KG in Zone B SMEM) --- + b_view_k4 = cute.make_layout( + (K_DIM, T_total, num_heads), stride=(1, num_heads * K_DIM, K_DIM) + ) + mV_k4 = cute.make_tensor(mV_in.iterator, b_view_k4) + + # O output: [T_total, V_dim, H] strides (H*V_dim, 1, V_dim) — same as A K_DIM view + mO_k4 = cute.make_tensor(mO_out.iterator, view_layout_3d) + + # State S_fp32: [V, K, B*H] view — mode 0=V(stride 1), mode 1=K(stride V_dim) + # Matches TMEM state [V,K] convention (M=V, N=K from tiled_mma_mn_mn) + s_fp32_vk_view = cute.make_layout((K_DIM, K_DIM, BH), stride=(1, K_DIM, K_DIM * K_DIM)) + mS_fp32_vk = cute.make_tensor(mS_fp32.iterator, s_fp32_vk_view) + + tma_ld = cpasync.CopyBulkTensorTileG2SOp() + # Only V loaded via TMA (S now in TMEM, not TMA) + ta_v, mVt = cute.nvgpu.make_tiled_tma_atom_B( + tma_ld, mV_k4, cute.select(sl_v, mode=[0, 1, 2]), tile1, mma_kmn + ) + + # TMA store for O (3D view) + sk = sm100_utils.get_smem_layout_atom_ab(OperandMajorMode.K, out_dtype, (M4, N4)) + sl_store = cute.tile_to_shape( + sm100_utils.make_smem_layout_atom(sk, out_dtype), (M4, N4), order=(0, 1) + ) + tma_st = cpasync.CopyBulkTensorTileS2GOp() + ta_o, mOo = cpasync.make_tiled_tma_atom(tma_st, mO_k4, sl_store, (M4, N4)) + + # --- K123 TMA setup --- + mQ_view = cute.make_tensor(mQ.iterator, view_layout_3d) + mK_view = cute.make_tensor(mK.iterator, view_layout_3d) + mG_view = cute.make_tensor(mG.iterator, view_layout_3d) + + smem_atom_qk = tcgen05.make_smem_layout_atom( + tcgen05.SmemLayoutAtomKind.K_SW128, cutlass.BFloat16 + ) + qk_smem_2d = cute.tile_to_shape(smem_atom_qk, (BT, K_DIM), order=(0, 1)) + + g_smem_2d = cute.make_layout((BT, K_DIM), stride=(K_DIM, 1)) + + tma_op_k123 = cpasync.CopyBulkTensorTileG2SOp(cpasync.CtaGroup.ONE) + ta_q_k123, tt_q = cpasync.make_tiled_tma_atom( + tma_op_k123, mQ_view, qk_smem_2d, cute.product_each(qk_smem_2d.shape), num_multicast=1 + ) + ta_k_k123, tt_k = cpasync.make_tiled_tma_atom( + tma_op_k123, mK_view, qk_smem_2d, cute.product_each(qk_smem_2d.shape), num_multicast=1 + ) + ta_g_k123, tt_g = cpasync.make_tiled_tma_atom( + tma_op_k123, mG_view, g_smem_2d, cute.product_each(g_smem_2d.shape), num_multicast=1 + ) + + g_cumsum_layout = cute.make_layout((BT, K_STRIDE), stride=(K_STRIDE, 1)) + + # K123 tiled copies + copy_atom_qk_k1 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=64 + ) + tiled_copy_qk_k1 = cute.make_tiled_copy_tv( + copy_atom_qk_k1, + thr_layout=cute.make_layout((1, 32)), + val_layout=cute.make_layout((1, 4)), + ) + + # K123 output v2 views removed — k_scaled/q_scaled/kg/gk_last now write to SMEM + + mma_op_k2 = cute.nvgpu.warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 8)) + tiled_mma_k2 = cute.make_tiled_mma( + mma_op_k2, cute.make_layout((1, 1, 1)), permutation_mnk=(16, 8, 8) + ) + tiled_copy_mma_A = cute.make_tiled_copy_A( + cute.make_copy_atom(cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 2), cutlass.BFloat16), + tiled_mma_k2, + ) + tiled_copy_mma_B = cute.make_tiled_copy_B( + cute.make_copy_atom(cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 1), cutlass.BFloat16), + tiled_mma_k2, + ) + copy_atom_Gcum = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=64 + ) + tiled_copy_Gcum_norm = cute.make_tiled_copy_tv( + copy_atom_Gcum, thr_layout=cute.make_layout((1, 4)), val_layout=cute.make_layout((1, 2)) + ) + tiled_copy_Gcum_gate = cute.make_tiled_copy_C(copy_atom_Gcum, tiled_mma_k2) + + fused_k1234_kernel( + mma_kmn, + mma_mn_mn, + # K4 TMA: V only (S in TMEM, not TMA) + ta_v, + mVt, + sl_v, + sl_s, # sST SMEM layout (filled from TMEM, not TMA) + # Zone B swizzle layouts + sl_ab, + sl_ks, + sl_qs, + sl_aqc, + sl_kg, + # K4 TMA store + readout + ta_o, + mOo, + sl_store, + sl_readout_k, + sl_nv_b, + sl_nv_a, + # K123 TMA + ta_q_k123, + tt_q, + ta_k_k123, + tt_k, + ta_g_k123, + tt_g, + qk_smem_2d, + g_smem_2d, + g_cumsum_layout, + tiled_copy_qk_k1, + tiled_mma_k2, + tiled_copy_mma_A, + tiled_copy_mma_B, + tiled_copy_Gcum_norm, + tiled_copy_Gcum_gate, + # K123 GMEM + mA_log, + mBeta, + scale_val, + # K4 GMEM: S_fp32 [V,K,BH] for initial/final state + mS_fp32_vk, + mDtBias, + lower_bound_val, + _HAS_BIAS, + _USE_SAFE_GATE, + mClocks, + _PROFILE_CLOCKS, + num_chunks, + num_heads, + batch_size, + ).launch( + grid=(BH, 1, 1), block=(THREADS, 1, 1), smem=225 * 1024, + stream=stream, + ) # Force high SMEM to prevent >1 block per SM (TMEM conflict) + + return host_fn 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 new file mode 100644 index 000000000000..7c2cadc6d9c8 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py @@ -0,0 +1,1459 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""K4-only fused KDA kernel: persistent variant with varlen support. + +Persistent scheduler (GDNTileScheduler) with 3D tensor layout, TensorMapManager +for per-tile TMA descriptor updates, and domain_offset+flat_divide for per-chunk +addressing. Supports variable-length sequences via cu_seqlens. + +K4 chunk loop with 6 MMAs per chunk: + MMA1: W = AB @ KS (K-MN, K=64) + MMA2: U = AB @ V (K-MN, K=64) + MMA3: NV = U - W_bf16 @ S (SS-mode: A=SMEM sO, B=SMEM sST, K=128, accumulate; + W is negated in its readout, FLA: b_v = u - w @ h) + MMA4: OI = QS @ S (K-MN, K=128) + MMA5: O = OI + AQC @ NV (K-MN, K=64, accumulate) + MMA6: State += NV^T @ KG (MN-MN, K=64, accumulate on decayed state) + +Execution order: MMA1->MMA2->MMA4->MMA3->MMA5->MMA6 + +Warp assignment (3 warpgroups, 12 warps, 384 threads): + WG0 (W0-3): W0=MMA issue, W2=TMA load/store, W1/W3=idle + WG1 (W4-7): State 2-pass: TMEM->bf16->sST, gk(SMEM)->decay->TMEM + WG2 (W8-11): GDN readout: W/NV/O TMEM->bf16->SMEM + +TMA pipelines (all PipelineTmaUmma, 1-stage): + - KS, V, QS: prefetch c+1 while MMA uses c + - AB, AQC, KG: 1-stage with explicit consumer release (no prefetch) + - TMA warp decoupled from O readout (no store_nbar wait) + +State management: + - tCtState [128,128] fp32 @ TMEM offset 256: persistent + - 2-pass per chunk: bf16->sST early, gk decay->TMEM late + - gk_last preloaded to SMEM (coalesced), read in Pass 2 + - single state_ready_mbar after both passes (TMEM no-overlap constraint) +""" + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.cutlass_dsl as _dsl_mod +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.nvgpu.tcgen05 import Field, OperandMajorMode, OperandSource +from cutlass.cutlass_dsl import Int32, T, dsl_user_op +from cutlass.utils import TensorMapManager, TensorMapUpdateMode + +if not hasattr(_dsl_mod, "CuteExperimentalDSL"): + + class _DummyExperimentalDSL: + jit = None + kernel = None + compile = None + + _dsl_mod.CuteExperimentalDSL = _DummyExperimentalDSL + +from flashinfer.gdn_kernels.blackwell.gated_delta_net_tile_scheduler import ( + GDNTileScheduler, + GDNTileSchedulerParams, +) + +SB = 16 +AKK_PAD = 8 +AKK_STRIDE = 64 + AKK_PAD # 72 +TEMP_COLS = SB + AKK_PAD # 24 +NUM_TEMPS = 2 + + +# =========================================================================== +# Inverse dsl_user_op functions (TF32 MMA m16n8k8, barrier 7) +# =========================================================================== +@dsl_user_op +def mma_tf32_m16n8k8(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None): + a0b = llvm.bitcast(T.i32(), a0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a1b = llvm.bitcast(T.i32(), a1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a2b = llvm.bitcast(T.i32(), a2.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a3b = llvm.bitcast(T.i32(), a3.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b0b = llvm.bitcast(T.i32(), b0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b1b = llvm.bitcast(T.i32(), b1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [ + a0b, + a1b, + a2b, + a3b, + b0b, + b1b, + c0.ir_value(loc=loc, ip=ip), + c1.ir_value(loc=loc, ip=ip), + c2.ir_value(loc=loc, ip=ip), + c3.ir_value(loc=loc, ip=ip), + ], + """{ + mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 + {$0, $1, $2, $3}, + {$4, $5, $6, $7}, + {$8, $9}, + {$10, $11, $12, $13}; + }""", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +@dsl_user_op +def inv_internal_barrier(*, loc=None, ip=None): + llvm.inline_asm( + T.i32(), + [], + "bar.sync 7, 128; mov.u32 $0, 0;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _invert_diag(sAkk: cute.Tensor, block_rc, lane_id, *, loc=None, ip=None): + my_row = lane_id % 16 + halfwarp_base = (lane_id // 16) * 16 + r_off = block_rc * 16 + c_off = block_rc * 16 + rInv = cute.make_rmem_tensor(cute.make_layout((16,), stride=(1,)), cutlass.Float32) + rInv[0] = cutlass.Float32(1.0) + for x in range(1, 16): + rInv[x] = cutlass.Float32(0.0) + for d in range(1, 16): + col_d = my_row - d + valid = cutlass.Float32(col_d >= 0) + a_val = cutlass.Float32(sAkk[r_off + my_row, c_off + col_d]) * valid + acc = cutlass.Float32(0.0) + for j in range(1, d): + a_re = cutlass.Float32(sAkk[r_off + my_row, c_off + my_row - (d - j)]) + inv_shfl = cute.arch.shuffle_sync(rInv[j], halfwarp_base + my_row - d + j) + acc = acc + a_re * inv_shfl + rInv[d] = (-a_val - acc) * valid + rInv[0] = cutlass.Float32(1.0) + sAkk[r_off + my_row, c_off + my_row] = rInv[0] + for d in range(1, 16): + sAkk[r_off + my_row, c_off + (my_row + 16 - d) % 16] = rInv[d] * cutlass.Float32( + my_row >= d + ) + + +@dsl_user_op +def _matmul_AB_inv(sAkk: cute.Tensor, br_A, bc_A, br_B, bc_B, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _z = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 16 + rB = br_B * 16 + cB = bc_B * 16 + a0 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid + 1]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid + 1]) + b0n0 = cutlass.Float32(sAkk[rB + 2 * tid, cB + gid]) + b1n0 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + gid]) + b0n1 = cutlass.Float32(sAkk[rB + 2 * tid, cB + 8 + gid]) + b1n1 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + 8 + gid]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n0, b1n0, _z, _z, _z, _z) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n1, b1n1, _z, _z, _z, _z) + a0 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid + 1]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid + 1]) + b0n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + gid]) + b1n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + gid]) + b0n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + 8 + gid]) + b1n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + 8 + gid]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8( + a0, a1, a2, a3, b0n0, b1n0, cn0_0, cn0_1, cn0_2, cn0_3 + ) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8( + a0, a1, a2, a3, b0n1, b1n1, cn1_0, cn1_1, cn1_2, cn1_3 + ) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _chain_mma_B_inv( + sAkk: cute.Tensor, + br_B, + bc_B, + a0k0, + a1k0, + a2k0, + a3k0, + a0k1, + a1k1, + a2k1, + a3k1, + lane_id, + *, + loc=None, + ip=None, +): + gid = lane_id // 4 + tid = lane_id % 4 + _z = cutlass.Float32(0.0) + rB = br_B * 16 + cB = bc_B * 16 + b0n0 = cutlass.Float32(sAkk[rB + 2 * tid, cB + gid]) + b1n0 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + gid]) + b0n1 = cutlass.Float32(sAkk[rB + 2 * tid, cB + 8 + gid]) + b1n1 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + 8 + gid]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8( + a0k0, a1k0, a2k0, a3k0, b0n0, b1n0, _z, _z, _z, _z + ) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8( + a0k0, a1k0, a2k0, a3k0, b0n1, b1n1, _z, _z, _z, _z + ) + b0n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + gid]) + b1n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + gid]) + b0n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + 8 + gid]) + b1n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + 8 + gid]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8( + a0k1, a1k1, a2k1, a3k1, b0n0, b1n0, cn0_0, cn0_1, cn0_2, cn0_3 + ) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8( + a0k1, a1k1, a2k1, a3k1, b0n1, b1n1, cn1_0, cn1_1, cn1_2, cn1_3 + ) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _chain_mma_A_inv( + sAkk: cute.Tensor, + br_A, + bc_A, + b0_k0n0, + b1_k0n0, + b0_k0n1, + b1_k0n1, + b0_k1n0, + b1_k1n0, + b0_k1n1, + b1_k1n1, + lane_id, + *, + loc=None, + ip=None, +): + gid = lane_id // 4 + tid = lane_id % 4 + _z = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 16 + a0 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid + 1]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid + 1]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0_k0n0, b1_k0n0, _z, _z, _z, _z) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0_k0n1, b1_k0n1, _z, _z, _z, _z) + a0 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid + 1]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid + 1]) + cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8( + a0, a1, a2, a3, b0_k1n0, b1_k1n0, cn0_0, cn0_1, cn0_2, cn0_3 + ) + cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8( + a0, a1, a2, a3, b0_k1n1, b1_k1n1, cn1_0, cn1_1, cn1_2, cn1_3 + ) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _store_neg_C_inv( + sAkk: cute.Tensor, br, bc, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None +): + gid = lane_id // 4 + tid = lane_id % 4 + r = br * 16 + c = bc * 16 + sAkk[r + gid, c + 2 * tid] = -c0 + sAkk[r + gid, c + 2 * tid + 1] = -c1 + sAkk[r + gid + 8, c + 2 * tid] = -c2 + sAkk[r + gid + 8, c + 2 * tid + 1] = -c3 + sAkk[r + gid, c + 8 + 2 * tid] = -c4 + sAkk[r + gid, c + 8 + 2 * tid + 1] = -c5 + sAkk[r + gid + 8, c + 8 + 2 * tid] = -c6 + sAkk[r + gid + 8, c + 8 + 2 * tid + 1] = -c7 + + +@dsl_user_op +def _shuffle_C_to_B_inv(c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + src_a = 8 * tid + gid // 2 + src_b = src_a + 4 + f_odd = cutlass.Float32(gid % 2) + f_even = cutlass.Float32(1) - f_odd + c0_a = cute.arch.shuffle_sync(c0, src_a) + c1_a = cute.arch.shuffle_sync(c1, src_a) + c2_a = cute.arch.shuffle_sync(c2, src_a) + c3_a = cute.arch.shuffle_sync(c3, src_a) + c4_a = cute.arch.shuffle_sync(c4, src_a) + c5_a = cute.arch.shuffle_sync(c5, src_a) + c6_a = cute.arch.shuffle_sync(c6, src_a) + c7_a = cute.arch.shuffle_sync(c7, src_a) + c0_b = cute.arch.shuffle_sync(c0, src_b) + c1_b = cute.arch.shuffle_sync(c1, src_b) + c2_b = cute.arch.shuffle_sync(c2, src_b) + c3_b = cute.arch.shuffle_sync(c3, src_b) + c4_b = cute.arch.shuffle_sync(c4, src_b) + c5_b = cute.arch.shuffle_sync(c5, src_b) + c6_b = cute.arch.shuffle_sync(c6, src_b) + c7_b = cute.arch.shuffle_sync(c7, src_b) + b0_00 = c0_a * f_even + c1_a * f_odd + b1_00 = c0_b * f_even + c1_b * f_odd + b0_10 = c2_a * f_even + c3_a * f_odd + b1_10 = c2_b * f_even + c3_b * f_odd + b0_01 = c4_a * f_even + c5_a * f_odd + b1_01 = c4_b * f_even + c5_b * f_odd + b0_11 = c6_a * f_even + c7_a * f_odd + b1_11 = c6_b * f_even + c7_b * f_odd + return b0_00, b1_00, b0_10, b1_10, b0_01, b1_01, b0_11, b1_11 + + +@dsl_user_op +def _store_C_temp_inv( + sT: cute.Tensor, buf, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None +): + gid = lane_id // 4 + tid = lane_id % 4 + sT[gid, 2 * tid, buf] = c0 + sT[gid, 2 * tid + 1, buf] = c1 + sT[gid + 8, 2 * tid, buf] = c2 + sT[gid + 8, 2 * tid + 1, buf] = c3 + sT[gid, 8 + 2 * tid, buf] = c4 + sT[gid, 8 + 2 * tid + 1, buf] = c5 + sT[gid + 8, 8 + 2 * tid, buf] = c6 + sT[gid + 8, 8 + 2 * tid + 1, buf] = c7 + + +@dsl_user_op +def _load_C_temp_inv(sT: cute.Tensor, buf, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + c0 = cutlass.Float32(sT[gid, 2 * tid, buf]) + c1 = cutlass.Float32(sT[gid, 2 * tid + 1, buf]) + c2 = cutlass.Float32(sT[gid + 8, 2 * tid, buf]) + c3 = cutlass.Float32(sT[gid + 8, 2 * tid + 1, buf]) + c4 = cutlass.Float32(sT[gid, 8 + 2 * tid, buf]) + c5 = cutlass.Float32(sT[gid, 8 + 2 * tid + 1, buf]) + c6 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid, buf]) + c7 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid + 1, buf]) + return c0, c1, c2, c3, c4, c5, c6, c7 + + +def transform_partitioned_tensor_layout(tensor): + layout = tensor.layout + stored_layout = layout + if isinstance(stored_layout, cute.ComposedLayout): + layout = layout.outer + shape = layout.shape + stride = layout.stride + new_shape = ((shape[0][0], shape[1]), (shape[0][1], shape[2]), *shape[3:]) + new_stride = ((stride[0][0], stride[1]), (stride[0][1], stride[2]), *stride[3:]) + new_layout = cute.make_layout(shape=new_shape, stride=new_stride) + if isinstance(stored_layout, cute.ComposedLayout): + new_layout = cute.make_composed_layout( + stored_layout.inner, stored_layout.offset, new_layout + ) + return cute.make_tensor(tensor.iterator, new_layout) + + +mma_dtype = cutlass.BFloat16 +acc_dtype = cutlass.Float32 +out_dtype = cutlass.BFloat16 + +M = 64 +N = 128 +K = 64 +K3 = 128 +M6 = 128 +N6 = 128 +K6 = 64 + +threads_per_cta = 384 +warp_threads = 32 +warpgroup_threads = 128 + +BYTES_PER_TENSORMAP = 128 +NUM_TENSORMAPS = 7 # a, b, v, q, aqc, kg, o + +MMA_WARP = 0 +O_STORE_WARP = 1 +TMA_WARP = 2 +STATE_WG = 1 +READOUT_WG = 2 + +NUM_REGS_WG0 = 40 +NUM_REGS_WG1 = 232 +NUM_REGS_WG2 = 232 +MAX_REGS = 168 + +try: + from cutlass.cutlass_dsl.cutlass import CuTeDSL as _CuTeDSL + + _orig_get_pipeline = _CuTeDSL._get_pipeline + _patch_applied = False + + def _patched_get_pipeline(self, _pipeline): + global _patch_applied + result = _orig_get_pipeline(self, _pipeline) + if result and "ptx-options=" not in result: + if "cubin-format=bin" in result: + 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]}" + ) + elif result and "ptx-options=" in result: + print(" [INFO] monkey-patch: ptx-options already present") + _patch_applied = True + return result + + _CuTeDSL._get_pipeline = _patched_get_pipeline +except Exception as e: + print(f" [WARN] monkey-patch failed: {e}") + _patch_applied = False + + +@cute.kernel +def k4_persistent_kernel( + tiled_mma_kmn: cute.TiledMma, + tiled_mma_mn_mn: cute.TiledMma, + tma_a, + a_sl: cute.ComposedLayout, + tma_b, + b_sl: cute.ComposedLayout, + tma_v, + v_sl: cute.ComposedLayout, + s_sl: cute.ComposedLayout, + tma_q, + q_sl: cute.ComposedLayout, + tma_aqc, + aqc_sl: cute.ComposedLayout, + tma_kg, + kg_sl: cute.ComposedLayout, + tma_o, + store_sl: cute.ComposedLayout, + readout_k_sl: cute.ComposedLayout, + nv_b_sl: cute.ComposedLayout, + nv_a_sl: cute.ComposedLayout, + kg_a_sl: cute.ComposedLayout, + nv_b_mn_sl: cute.ComposedLayout, + mGkLastExp: cute.Tensor, + mS_fp32: cute.Tensor, + cu_seqlens: cute.Tensor, + chunk_offsets: cute.Tensor, + mA: cute.Tensor, + mB: cute.Tensor, + mV_g: cute.Tensor, + mQ: cute.Tensor, + mAQC: cute.Tensor, + mKG: cute.Tensor, + mO: cute.Tensor, + tensormap_workspace: cute.Tensor, + scheduler_params: GDNTileSchedulerParams, +): + bidx, bidy, bidz = cute.arch.block_idx() + grid_dim = cute.arch.grid_dim() + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(tidx // warp_threads) + warpgroup_idx = cute.arch.make_warp_uniform(tidx // warpgroup_threads) + warpgroup_tidx = tidx % warpgroup_threads + thr_kmn = tiled_mma_kmn.get_slice(0) + thr_mn = tiled_mma_mn_mn.get_slice(0) + dice = (None, None, None) + + cta_linear_idx = bidz * grid_dim[1] * grid_dim[0] + bidy * grid_dim[0] + bidx + tensormap_manager = TensorMapManager(TensorMapUpdateMode.GMEM, BYTES_PER_TENSORMAP) + tm_ws = cute.make_tensor( + tensormap_workspace.iterator, + cute.make_layout( + (grid_dim[0] * grid_dim[1] * grid_dim[2], NUM_TENSORMAPS, BYTES_PER_TENSORMAP), + stride=(NUM_TENSORMAPS * BYTES_PER_TENSORMAP, BYTES_PER_TENSORMAP, 1), + ), + ) + tm_a_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 0, None)].iterator) + tm_b_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 1, None)].iterator) + tm_v_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 2, None)].iterator) + tm_q_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 3, None)].iterator) + tm_aqc_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 4, None)].iterator) + tm_kg_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 5, None)].iterator) + tm_o_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 6, None)].iterator) + + smem = cutlass.utils.SmemAllocator() + AL = 128 + sA = smem.allocate_tensor(mma_dtype, a_sl.outer, AL, a_sl.inner) + sB = smem.allocate_tensor(mma_dtype, b_sl.outer, AL, b_sl.inner) + sV = smem.allocate_tensor(mma_dtype, v_sl.outer, AL, v_sl.inner) + sST = smem.allocate_tensor(mma_dtype, s_sl.outer, AL, s_sl.inner) + sQ = smem.allocate_tensor(mma_dtype, q_sl.outer, AL, q_sl.inner) + sAQC = smem.allocate_tensor(mma_dtype, aqc_sl.outer, AL, aqc_sl.inner) + sKG = smem.allocate_tensor(mma_dtype, kg_sl.outer, AL, kg_sl.inner) + sNV = smem.allocate_tensor(mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner) + sO = smem.allocate_tensor(mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner) + sO_out = smem.allocate_tensor(mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner) + sGk_buf = smem.allocate_array(acc_dtype, N) + + sNV_b = cute.make_tensor(cute.recast_ptr(sNV.iterator, nv_b_sl.inner, mma_dtype), nv_b_sl.outer) + # sNV_a = cute.make_tensor(cute.recast_ptr(sNV.iterator, nv_a_sl.inner, mma_dtype), nv_a_sl.outer) + sKG_a = cute.make_tensor(cute.recast_ptr(sKG.iterator, kg_a_sl.inner, mma_dtype), kg_a_sl.outer) + sNV_b_mn = cute.make_tensor( + cute.recast_ptr(sNV.iterator, nv_b_mn_sl.inner, mma_dtype), nv_b_mn_sl.outer + ) + sO_out_st = cute.make_tensor( + cute.recast_ptr(sO_out.iterator, store_sl.inner, out_dtype), store_sl.outer + ) + + tmem_smem = smem.allocate_array(cutlass.Int32, 1) + if warp_idx == 0: + cute.arch.alloc_tmem(512, tmem_smem) + + sNV_ready_nbar = pipeline.NamedBarrier(3, warpgroup_threads + warp_threads) + sW_ready_nbar = pipeline.NamedBarrier(2, warpgroup_threads + warp_threads) + gk_load_nbar = pipeline.NamedBarrier(5, warpgroup_threads) + + elect_one = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + wg_coop = pipeline.CooperativeGroup(pipeline.Agent.Thread, warpgroup_threads) + warp_coop = pipeline.CooperativeGroup(pipeline.Agent.Thread, warp_threads) + + mma6_done_mbar = smem.allocate_array(cutlass.Int64, 1) + gmem_done_mbar = smem.allocate_array(cutlass.Int64, 1) + state_ready_mbar = smem.allocate_array(cutlass.Int64, 1) + final_state_done_mbar = smem.allocate_array(cutlass.Int64, 1) + + if warp_idx == 0: + with cute.arch.elect_one(): + cute.arch.mbarrier_init(mma6_done_mbar, warp_threads) + cute.arch.mbarrier_init(gmem_done_mbar, warpgroup_threads) + cute.arch.mbarrier_init(state_ready_mbar, warpgroup_threads) + cute.arch.mbarrier_init(final_state_done_mbar, warpgroup_threads) + + def _make_tma_pipe(total_byte_count, num_stages=1): + ptr = smem.allocate_array(cutlass.Int64, 2 * num_stages) + return pipeline.PipelineTmaUmma.create( + barrier_storage=ptr, + num_stages=num_stages, + producer_group=elect_one, + consumer_group=elect_one, + tx_count=total_byte_count // num_stages, + defer_sync=True, + ).make_participants() + + b_prod, b_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, b_sl), num_stages=1) + v_prod, v_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, v_sl), num_stages=1) + q_prod, q_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, q_sl), num_stages=1) + a_prod, a_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, a_sl), num_stages=1) + aqc_prod, aqc_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, aqc_sl), num_stages=1) + kg_prod, kg_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, kg_sl), num_stages=1) + + def _make_umma_pipe(): + ptr = smem.allocate_array(cutlass.Int64, 2) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=ptr, + num_stages=1, + producer_group=elect_one, + consumer_group=wg_coop, + defer_sync=True, + ).make_participants() + + w_prod, w_cons = _make_umma_pipe() + nv_prod, nv_cons = _make_umma_pipe() + o_prod, o_cons = _make_umma_pipe() + + o_store_mbar = smem.allocate_array(cutlass.Int64, 2) + o_store_prod, o_store_cons = pipeline.PipelineAsync.create( + barrier_storage=o_store_mbar, + num_stages=1, + producer_group=wg_coop, + consumer_group=warp_coop, + defer_sync=True, + ).make_participants() + + cute.arch.sync_threads() + + tmem_ptr = cute.arch.retrieve_tmem_ptr(cutlass.Int32, 16, tmem_smem) + + tCtW_shape = tiled_mma_kmn.partition_shape_C((M, N)) + tCtW_fake = tiled_mma_kmn.make_fragment_C(tCtW_shape) + tCtW = cute.make_tensor(cute.recast_ptr(tmem_ptr + 0, dtype=acc_dtype), tCtW_fake.layout) + + tCtNV_shape = tiled_mma_kmn.partition_shape_C((M, N)) + tCtNV_fake = tiled_mma_kmn.make_fragment_C(tCtNV_shape) + tCtNV = cute.make_tensor(cute.recast_ptr(tmem_ptr + 128, dtype=acc_dtype), tCtNV_fake.layout) + tCtO = cute.make_tensor(cute.recast_ptr(tmem_ptr + 384, dtype=acc_dtype), tCtNV_fake.layout) + + tCtS_shape = tiled_mma_mn_mn.partition_shape_C((M6, N6)) + tCtS_fake = tiled_mma_mn_mn.make_fragment_C(tCtS_shape) + tCtState = cute.make_tensor(cute.recast_ptr(tmem_ptr + 256, dtype=acc_dtype), tCtS_fake.layout) + + # PDL: setup is done; now wait for upstream akk_inv to commit before reading + # gmem inputs (TMA loads, gk_last_exp reads). PDL allows the entire setup phase + # above (~tmem alloc, smem layout, pipeline init) to overlap with akk_inv's tail. + cute.arch.griddepcontrol_wait() + + # ==== WG1: State readout + decay ==== + if warpgroup_idx == STATE_WG: + cId_128 = cute.make_identity_tensor((M6, N6)) + tCtState_mn = transform_partitioned_tensor_layout(tCtState) + + atom_state_t2r = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), acc_dtype) + tiled_state_t2r = tcgen05.make_tmem_copy(atom_state_t2r, tCtState[(None, None), 0, 0]) + thr_state_t2r = tiled_state_t2r.get_slice(warpgroup_tidx) + tTR_tCtState = thr_state_t2r.partition_S(tCtState_mn) + tTR_tCcState = thr_state_t2r.partition_D(cId_128) + tRrState = cute.make_rmem_tensor_like(tTR_tCcState, acc_dtype) + + atom_state_r2t = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition(32)), acc_dtype) + tiled_state_r2t = tcgen05.make_tmem_copy(atom_state_r2t, tCtState[(None, None), 0, 0]) + thr_state_r2t = tiled_state_r2t.get_slice(warpgroup_tidx) + tRT_tCtState = thr_state_r2t.partition_D(tCtState_mn) + + atom_state_g2r = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), acc_dtype, num_bits_per_copy=128 + ) + tiled_state_g2r = cute.make_tiled_copy_S(atom_state_g2r, tiled_state_r2t) + thr_state_g2r = tiled_state_g2r.get_slice(warpgroup_tidx) + + atom_state_r2g = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), acc_dtype, num_bits_per_copy=128 + ) + tiled_state_r2g = cute.make_tiled_copy_D(atom_state_r2g, tiled_state_t2r) + thr_state_r2g = tiled_state_r2g.get_slice(warpgroup_tidx) + + atom_state_r2s = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), mma_dtype, num_bits_per_copy=128 + ) + tiled_state_r2s = cute.make_tiled_copy_D(atom_state_r2s, tiled_state_t2r) + thr_state_r2s = tiled_state_r2s.get_slice(warpgroup_tidx) + sST_vk_view = transform_partitioned_tensor_layout(sST) + sST_kv_view = cute.make_tensor( + sST.iterator, cute.select(sST_vk_view.layout, mode=[1, 0, 2]) + ) + tCsState_inp = thr_state_r2s.partition_D(sST_kv_view) + tRrState_bf16 = cute.make_rmem_tensor_like(tTR_tCcState, mma_dtype) + tCrState_bf16 = tiled_state_r2s.retile(tRrState_bf16) + + sGk = cute.make_tensor(sGk_buf, cute.make_layout(((N,), (1,)))) + + scheduler = GDNTileScheduler.create(scheduler_params, (bidx, bidy, bidz), grid_dim) + work = scheduler.initial_work_tile_info() + global_chunk = Int32(0) + + while work.is_valid_tile: + batch_idx, head_idx, _ = work.tile_idx + batch_start = cu_seqlens[batch_idx] + batch_end = cu_seqlens[batch_idx + 1] + num_chunks = cute.ceil_div(batch_end - batch_start, M) + # chunk_base from cumulative chunk-offsets array (handles non-64-aligned varlen). + # batch_start // M only works when all seq lengths are multiples of M. + chunk_base = chunk_offsets[batch_idx] + + gS_init = cute.flat_divide(mS_fp32[batch_idx, head_idx, None, None], (M6, N6))[ + None, None, 0, 0 + ] + tGR_tCgState_in = thr_state_g2r.partition_S(gS_init) + tGR_tCrState_in = thr_state_g2r.retile(tRrState) + cute.copy(tiled_state_g2r, tGR_tCgState_in, tGR_tCrState_in) + num_state_subs = tRrState.shape[2] + for sub in cutlass.range(num_state_subs): + cute.copy(tiled_state_r2t, tRrState[None, 0, sub], tRT_tCtState[None, 0, sub]) + cute.arch.fence_view_async_tmem_store() + + sGk[warpgroup_tidx] = cutlass.Float32(mGkLastExp[chunk_base, head_idx, warpgroup_tidx]) + gk_load_nbar.arrive_and_wait() + + for chunk_c in cutlass.range(num_chunks): + num_state_subs = tRrState.shape[2] + _sub_tile_size = cute.size(tRrState.shape[0]) + + if chunk_c > 0: + cute.arch.mbarrier_wait(mma6_done_mbar, phase=(global_chunk - 1) % 2) + gk_load_nbar.arrive_and_wait() + + cute.copy(tiled_state_t2r, tTR_tCtState, tRrState) + tRrState_bf16.store(tRrState.load().to(mma_dtype)) + cute.copy(tiled_state_r2s, tCrState_bf16, tCsState_inp[None, None, None, 0]) + cute.arch.fence_view_async_shared() + cute.arch.mbarrier_arrive(gmem_done_mbar) + + for sub in cutlass.range(num_state_subs): + cute.copy(tiled_state_t2r, tTR_tCtState[None, 0, sub], tRrState[None, 0, sub]) + cute.arch.fence_view_async_tmem_load() + for sub in cutlass.range(num_state_subs): + for i in cutlass.range(_sub_tile_size): + coord = tTR_tCcState[i, 0, sub] + # FIX: probe shows kernel state TMEM[m=K, n=V] (m corresponds to K-dim); + # OLD K4 / fla wants K-axis decay (state[k, v] *= gk[k]) → use coord[0]=m=K-idx. + k_idx = coord[0] + gk_val = cutlass.Float32(sGk[k_idx]) + tRrState[i, 0, sub] = tRrState[i, 0, sub] * gk_val + cute.copy(tiled_state_r2t, tRrState[None, 0, sub], tRT_tCtState[None, 0, sub]) + cute.arch.fence_view_async_tmem_store() + cute.arch.mbarrier_arrive(state_ready_mbar) + + if chunk_c + 1 < num_chunks: + sGk[warpgroup_tidx] = cutlass.Float32( + mGkLastExp[chunk_base + chunk_c + 1, head_idx, warpgroup_tidx] + ) + global_chunk = global_chunk + 1 + + cute.arch.mbarrier_wait(mma6_done_mbar, phase=(global_chunk - 1) % 2) + gS_out = cute.flat_divide(mS_fp32[batch_idx, head_idx, None, None], (M6, N6))[ + None, None, 0, 0 + ] + tGR_tCgState_out = thr_state_r2g.partition_D(gS_out) + tGR_tCrState_out = thr_state_r2g.retile(tRrState) + num_state_subs_final = tRrState.shape[2] + for sub in cutlass.range(num_state_subs_final): + cute.copy(tiled_state_t2r, tTR_tCtState[None, 0, sub], tRrState[None, 0, sub]) + cute.arch.fence_view_async_tmem_load() + for sub in cutlass.range(num_state_subs_final): + cute.copy( + tiled_state_r2g, tGR_tCrState_out[None, 0, sub], tGR_tCgState_out[None, 0, sub] + ) + + cute.arch.mbarrier_arrive(final_state_done_mbar) + + scheduler.advance_to_next_work() + work = scheduler.get_current_work() + + # ==== MMA warp (warp 0): 6 MMAs ==== + elif warp_idx == MMA_WARP: + fA_kmn = thr_kmn.make_fragment_A(sA) + fB_ks = thr_kmn.make_fragment_B(sB) + fB_v = thr_kmn.make_fragment_B(sV) + fB_s = thr_kmn.make_fragment_B(sST) + fA_q = thr_kmn.make_fragment_A(sQ) + fA_aqc = thr_kmn.make_fragment_A(sAQC) + fB_nv = thr_kmn.make_fragment_B(sNV_b) + fA_kg = thr_mn.make_fragment_A(sKG_a) + fB_nv_mn = thr_mn.make_fragment_B(sNV_b_mn) + fA_w = thr_kmn.make_fragment_A(sO) + + scheduler = GDNTileScheduler.create(scheduler_params, (bidx, bidy, bidz), grid_dim) + work = scheduler.initial_work_tile_info() + global_chunk = Int32(0) + tile_count = Int32(0) + + while work.is_valid_tile: + batch_idx_m, head_idx_m, _ = work.tile_idx + num_chunks_m = cute.ceil_div(cu_seqlens[batch_idx_m + 1] - cu_seqlens[batch_idx_m], M) + + for chunk_c in cutlass.range(num_chunks_m): + c_phase = global_chunk % 2 + + ah = a_cons.wait_and_advance() + bh = b_cons.wait_and_advance() + w_h = w_prod.acquire_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, False) + for k in cutlass.range_constexpr(cute.size(sB.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtW, + fA_kmn[dice + (0,)][None, None, k], + fB_ks[dice + (bh.index,)][None, None, k], + tCtW, + ) + if k == 0: + tiled_mma_kmn.set(Field.ACCUMULATE, True) + bh.release() + w_h.commit() + + vh = v_cons.wait_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, False) + for k in cutlass.range_constexpr(cute.size(sV.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtNV, + fA_kmn[dice + (0,)][None, None, k], + fB_v[dice + (vh.index,)][None, None, k], + tCtNV, + ) + if k == 0: + tiled_mma_kmn.set(Field.ACCUMULATE, True) + vh.release() + ah.release() + + cute.arch.mbarrier_wait(gmem_done_mbar, phase=c_phase) + sW_ready_nbar.arrive_and_wait() + + nv_h = nv_prod.acquire_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, True) + for k in cutlass.range_constexpr(cute.size(sST.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtNV, + fA_w[dice + (0,)][None, None, k], + fB_s[dice + (0,)][None, None, k], + tCtNV, + ) + nv_h.commit() + + qh = q_cons.wait_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, False) + for k in cutlass.range_constexpr(cute.size(sST.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtO, + fA_q[dice + (qh.index,)][None, None, k], + fB_s[dice + (0,)][None, None, k], + tCtO, + ) + if k == 0: + tiled_mma_kmn.set(Field.ACCUMULATE, True) + qh.release() + + sNV_ready_nbar.arrive_and_wait() + + aqch = aqc_cons.wait_and_advance() + o_h = o_prod.acquire_and_advance() + tiled_mma_kmn.set(Field.ACCUMULATE, True) + for k in cutlass.range_constexpr(cute.size(sNV_b.shape[2])): + cute.gemm( + tiled_mma_kmn, + tCtO, + fA_aqc[dice + (0,)][None, None, k], + fB_nv[dice + (0,)][None, None, k], + tCtO, + ) + o_h.commit() + aqch.release() + + cute.arch.mbarrier_wait(state_ready_mbar, phase=c_phase) + kgh = kg_cons.wait_and_advance() + tiled_mma_mn_mn.set(Field.ACCUMULATE, True) + for k in cutlass.range_constexpr(cute.size(sKG_a.shape[2])): + cute.gemm( + tiled_mma_mn_mn, + tCtState, + fA_kg[dice + (0,)][None, None, k], + fB_nv_mn[dice + (0,)][None, None, k], + tCtState, + ) + kgh.release() + + w_prod.tail() + nv_prod.tail() + o_prod.tail() + tcgen05.commit(mma6_done_mbar) + global_chunk = global_chunk + 1 + + cute.arch.mbarrier_wait(final_state_done_mbar, phase=tile_count % 2) + tile_count = tile_count + 1 + + scheduler.advance_to_next_work() + work = scheduler.get_current_work() + + cute.arch.relinquish_tmem_alloc_permit() + cute.arch.dealloc_tmem(tmem_ptr, 512) + + # ==== TMA warp (warp 2) ==== + elif warp_idx == TMA_WARP: + cta_layout = cute.make_layout(1) + + scheduler = GDNTileScheduler.create(scheduler_params, (bidx, bidy, bidz), grid_dim) + work = scheduler.initial_work_tile_info() + + if work.is_valid_tile: + tensormap_manager.init_tensormap_from_atom(tma_a[0], tm_a_ptr, TMA_WARP) + tensormap_manager.init_tensormap_from_atom(tma_b[0], tm_b_ptr, TMA_WARP) + tensormap_manager.init_tensormap_from_atom(tma_v[0], tm_v_ptr, TMA_WARP) + tensormap_manager.init_tensormap_from_atom(tma_q[0], tm_q_ptr, TMA_WARP) + tensormap_manager.init_tensormap_from_atom(tma_aqc[0], tm_aqc_ptr, TMA_WARP) + tensormap_manager.init_tensormap_from_atom(tma_kg[0], tm_kg_ptr, TMA_WARP) + tensormap_manager.fence_tensormap_initialization() + + while work.is_valid_tile: + batch_idx_t, head_idx_t, _ = work.tile_idx + batch_start_t = cu_seqlens[batch_idx_t] + batch_end_t = cu_seqlens[batch_idx_t + 1] + num_chunks_t = cute.ceil_div(batch_end_t - batch_start_t, M) + + bounded_a = cute.make_tensor( + mA.iterator, + cute.make_layout( + (batch_end_t, mA.shape[1], mA.shape[2]), + stride=(mA.stride[0], mA.stride[1], mA.stride[2]), + ), + ) + bounded_b = cute.make_tensor( + mB.iterator, + cute.make_layout( + (mB.shape[0], batch_end_t, mB.shape[2]), + stride=(mB.stride[0], mB.stride[1], mB.stride[2]), + ), + ) + bounded_v = cute.make_tensor( + mV_g.iterator, + cute.make_layout( + (mV_g.shape[0], batch_end_t, mV_g.shape[2]), + stride=(mV_g.stride[0], mV_g.stride[1], mV_g.stride[2]), + ), + ) + bounded_q = cute.make_tensor( + mQ.iterator, + cute.make_layout( + (batch_end_t, mQ.shape[1], mQ.shape[2]), + stride=(mQ.stride[0], mQ.stride[1], mQ.stride[2]), + ), + ) + bounded_aqc = cute.make_tensor( + mAQC.iterator, + cute.make_layout( + (batch_end_t, mAQC.shape[1], mAQC.shape[2]), + stride=(mAQC.stride[0], mAQC.stride[1], mAQC.stride[2]), + ), + ) + bounded_kg = cute.make_tensor( + mKG.iterator, + cute.make_layout( + (mKG.shape[0], batch_end_t, mKG.shape[2]), + stride=(mKG.stride[0], mKG.stride[1], mKG.stride[2]), + ), + ) + tensormap_manager.update_tensormap( + (bounded_a, bounded_b, bounded_v, bounded_q, bounded_aqc, bounded_kg), + (tma_a[0], tma_b[0], tma_v[0], tma_q[0], tma_aqc[0], tma_kg[0]), + (tm_a_ptr, tm_b_ptr, tm_v_ptr, tm_q_ptr, tm_aqc_ptr, tm_kg_ptr), + TMA_WARP, + (None, None, None, None, None, None), + ) + + for chunk_c in cutlass.range(num_chunks_t): + chunk_offset = batch_start_t + chunk_c * M + + # AB: A-operand (T, K, H) → domain_offset on (tokens, 0) + mA_c = cute.domain_offset( + (chunk_offset, Int32(0)), tma_a[1][None, None, head_idx_t] + ) + gA = cute.flat_divide(mA_c, (M, K)) + tCgA = thr_kmn.partition_A(gA) + tAsA, tAgA = cpasync.tma_partition( + tma_a[0], + 0, + cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + ah = a_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_a_ptr) + cute.copy( + tma_a[0], + tAgA[(None, 0, 0)], + tAsA[(None, ah.index)], + tma_bar_ptr=ah.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_a_ptr, cute.AddressSpace.generic + ), + ) + + # KS: B-operand (N, T, H) → domain_offset on (0, tokens) + mB_c = cute.domain_offset( + (Int32(0), chunk_offset), tma_b[1][None, None, head_idx_t] + ) + gB_c = cute.flat_divide(mB_c, (N, K)) + tCgB = thr_kmn.partition_B(gB_c) + tBsB, tBgB = cpasync.tma_partition( + tma_b[0], + 0, + cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + bh = b_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_b_ptr) + cute.copy( + tma_b[0], + tBgB[(None, 0, 0)], + tBsB[(None, bh.index)], + tma_bar_ptr=bh.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_b_ptr, cute.AddressSpace.generic + ), + ) + + # V: B-operand (N, T, H) → domain_offset on (0, tokens) + mV_c = cute.domain_offset( + (Int32(0), chunk_offset), tma_v[1][None, None, head_idx_t] + ) + gV_c = cute.flat_divide(mV_c, (N, K)) + tCgV = thr_kmn.partition_B(gV_c) + tBsV, tBgV = cpasync.tma_partition( + tma_v[0], + 0, + cta_layout, + cute.group_modes(sV, 0, 3), + cute.group_modes(tCgV, 0, 3), + ) + vh = v_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_v_ptr) + cute.copy( + tma_v[0], + tBgV[(None, 0, 0)], + tBsV[(None, vh.index)], + tma_bar_ptr=vh.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_v_ptr, cute.AddressSpace.generic + ), + ) + + # QS: A-operand (T, N, H) → domain_offset on (tokens, 0) + mQ_c = cute.domain_offset( + (chunk_offset, Int32(0)), tma_q[1][None, None, head_idx_t] + ) + gQ_c = cute.flat_divide(mQ_c, (M, K3)) + tCgQ = thr_kmn.partition_A(gQ_c) + tAsQ, tAgQ = cpasync.tma_partition( + tma_q[0], + 0, + cta_layout, + cute.group_modes(sQ, 0, 3), + cute.group_modes(tCgQ, 0, 3), + ) + qh = q_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_q_ptr) + cute.copy( + tma_q[0], + tAgQ[(None, 0, 0)], + tAsQ[(None, qh.index)], + tma_bar_ptr=qh.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_q_ptr, cute.AddressSpace.generic + ), + ) + + # AQC: A-operand (T, K, H) → domain_offset on (tokens, 0) + mAQC_c = cute.domain_offset( + (chunk_offset, Int32(0)), tma_aqc[1][None, None, head_idx_t] + ) + gAQC_c = cute.flat_divide(mAQC_c, (M, K)) + tCgAQC = thr_kmn.partition_A(gAQC_c) + tAsAQC, tAgAQC = cpasync.tma_partition( + tma_aqc[0], + 0, + cta_layout, + cute.group_modes(sAQC, 0, 3), + cute.group_modes(tCgAQC, 0, 3), + ) + aqch = aqc_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_aqc_ptr) + cute.copy( + tma_aqc[0], + tAgAQC[(None, 0, 0)], + tAsAQC[(None, aqch.index)], + tma_bar_ptr=aqch.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_aqc_ptr, cute.AddressSpace.generic + ), + ) + + # KG: B-operand (N6, T, H) → domain_offset on (0, tokens) + mKG_c = cute.domain_offset( + (Int32(0), chunk_offset), tma_kg[1][None, None, head_idx_t] + ) + gKG_c = cute.flat_divide(mKG_c, (N6, K6)) + tCgKG = thr_mn.partition_B(gKG_c) + tBsKG, tBgKG = cpasync.tma_partition( + tma_kg[0], + 0, + cta_layout, + cute.group_modes(sKG, 0, 3), + cute.group_modes(tCgKG, 0, 3), + ) + kgh = kg_prod.acquire_and_advance() + if chunk_c == 0: + tensormap_manager.fence_tensormap_update(tm_kg_ptr) + cute.copy( + tma_kg[0], + tBgKG[(None, 0, 0)], + tBsKG[(None, kgh.index)], + tma_bar_ptr=kgh.barrier, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_kg_ptr, cute.AddressSpace.generic + ), + ) + + scheduler.advance_to_next_work() + work = scheduler.get_current_work() + + # ==== Warp 1: O TMA store ==== + elif warp_idx == O_STORE_WARP: + cta_layout_o = cute.make_layout(1) + + scheduler = GDNTileScheduler.create(scheduler_params, (bidx, bidy, bidz), grid_dim) + work = scheduler.initial_work_tile_info() + + if work.is_valid_tile: + tensormap_manager.init_tensormap_from_atom(tma_o[0], tm_o_ptr, O_STORE_WARP) + tensormap_manager.fence_tensormap_initialization() + + while work.is_valid_tile: + batch_idx_o, head_idx_o, _ = work.tile_idx + batch_start_o = cu_seqlens[batch_idx_o] + batch_end_o = cu_seqlens[batch_idx_o + 1] + num_chunks_o = cute.ceil_div(batch_end_o - batch_start_o, M) + + bounded_o = cute.make_tensor( + mO.iterator, + cute.make_layout( + (batch_end_o, mO.shape[1], mO.shape[2]), + stride=(mO.stride[0], mO.stride[1], mO.stride[2]), + ), + ) + tensormap_manager.update_tensormap( + (bounded_o,), (tma_o[0],), (tm_o_ptr,), O_STORE_WARP, (None,) + ) + tensormap_manager.fence_tensormap_update(tm_o_ptr) + + for chunk_c in cutlass.range(num_chunks_o): + os_h = o_store_cons.wait_and_advance() + chunk_offset_o = batch_start_o + chunk_c * M + mO_c = cute.domain_offset( + (chunk_offset_o, Int32(0)), tma_o[1][None, None, head_idx_o] + ) + gOo = cute.flat_divide(mO_c, (M, N)) + sOt, gOt = cpasync.tma_partition( + tma_o[0], + 0, + cta_layout_o, + cute.group_modes(sO_out_st, 0, 2), + cute.group_modes(gOo, 0, 2), + ) + cute.copy( + tma_o[0], + sOt[None], + gOt[(None, 0, 0)], + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tm_o_ptr, cute.AddressSpace.generic + ), + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + os_h.release() + + scheduler.advance_to_next_work() + work = scheduler.get_current_work() + + # ==== WG2: W/NV/O readout (SS-mode, no Phase 2) ==== + elif warpgroup_idx == READOUT_WG: + tCtW_mn = transform_partitioned_tensor_layout(tCtW) + tCtNV_mn = transform_partitioned_tensor_layout(tCtNV) + tCtO_mn = transform_partitioned_tensor_layout(tCtO) + + atom_t2r = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition(1)), acc_dtype) + tiled_t2r = tcgen05.make_tmem_copy(atom_t2r, tCtW[(None, None), 0, 0]) + thr_t2r = tiled_t2r.get_slice(warpgroup_tidx) + + tTR_W = thr_t2r.partition_S(tCtW_mn) + tTR_NV = thr_t2r.partition_S(tCtNV_mn) + tTR_O = thr_t2r.partition_S(tCtO_mn) + + atom_r2s_k = sm100_utils.get_smem_store_op( + utils.LayoutEnum.ROW_MAJOR, mma_dtype, acc_dtype, tiled_t2r + ) + tiled_r2s_k = cute.make_tiled_copy_D(atom_r2s_k, tiled_t2r) + thr_r2s_k = tiled_r2s_k.get_slice(warpgroup_tidx) + tCsO = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sO)) + tCsO_out = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sO_out)) + tCsNV = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sNV)) + + cId = cute.make_identity_tensor((M, N)) + tTR_cId = thr_t2r.partition_D(cId) + + scheduler = GDNTileScheduler.create(scheduler_params, (bidx, bidy, bidz), grid_dim) + work = scheduler.initial_work_tile_info() + global_chunk = Int32(0) + + while work.is_valid_tile: + batch_idx_r, head_idx_r, _ = work.tile_idx + num_chunks_r = cute.ceil_div(cu_seqlens[batch_idx_r + 1] - cu_seqlens[batch_idx_r], M) + + for chunk_c in cutlass.range(num_chunks_r): + c_phase = global_chunk % 2 + tRrR = cute.make_rmem_tensor_like(tTR_cId, acc_dtype) + tRrR_out = cute.make_rmem_tensor_like(tRrR, mma_dtype) + tCrR_k = tiled_r2s_k.retile(tRrR_out) + num_subs = tRrR.shape[2] + + wh = w_cons.wait_and_advance() + for sub in cutlass.range(num_subs): + cute.copy(tiled_t2r, tTR_W[None, 0, sub], tRrR[None, 0, sub]) + # Negate W in its readout so MMA3 accumulates NV = U - W @ S + # (FLA: b_v = u - w @ h). Folded into the bf16 cast — free. + tRrR_out[None, 0, sub].store((-(tRrR[None, 0, sub].load())).to(mma_dtype)) + cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsO[None, 0, sub, 0]) + cute.arch.fence_view_async_tmem_load() + wh.release() + cute.arch.fence_view_async_shared() + sW_ready_nbar.arrive_and_wait() + + nvh = nv_cons.wait_and_advance() + for sub in cutlass.range(num_subs): + cute.copy(tiled_t2r, tTR_NV[None, 0, sub], tRrR[None, 0, sub]) + tRrR_out[None, 0, sub].store(tRrR[None, 0, sub].load().to(mma_dtype)) + cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsNV[None, 0, sub, 0]) + cute.arch.fence_view_async_tmem_load() + nvh.release() + cute.arch.fence_view_async_shared() + sNV_ready_nbar.arrive_and_wait() + + os_h = o_store_prod.acquire_and_advance() + + oh = o_cons.wait_and_advance() + for sub in cutlass.range(num_subs): + cute.copy(tiled_t2r, tTR_O[None, 0, sub], tRrR[None, 0, sub]) + tRrR_out[None, 0, sub].store(tRrR[None, 0, sub].load().to(mma_dtype)) + cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsO_out[None, 0, sub, 0]) + cute.arch.fence_view_async_tmem_load() + oh.release() + cute.arch.fence_view_async_shared() + os_h.commit() + global_chunk = global_chunk + 1 + + scheduler.advance_to_next_work() + work = scheduler.get_current_work() + o_store_prod.tail() + + +def make_host_fn(num_sm=148): + """Create one K4 host function for all runtime context batch sizes. + + ``num_seqs`` is a runtime scalar. ``H`` and the token extent come from + dynamic tensor layouts, while ``num_sm`` is the fixed persistent-grid + ceiling. The scheduler computes ``min(num_seqs * H, num_sm)`` at launch. + """ + _num_sm = num_sm + + @cute.jit + def host_fn( + a_raw: cute.Tensor, + b_raw: cute.Tensor, + v_raw: cute.Tensor, + q_raw: cute.Tensor, + aqc_raw: cute.Tensor, + kg_raw: cute.Tensor, + o_raw: cute.Tensor, + gk_last_exp: cute.Tensor, + s_fp32: cute.Tensor, + cu_seqlens: cute.Tensor, + chunk_offsets: cute.Tensor, + tm_workspace: cute.Tensor, + num_seqs: cutlass.Int32, + # Launch stream — runtime argument; launching on the DSL default + # stream races with the executor's non-blocking execution stream. + # Same stream as akk_inv, so the PDL chain stays intact. + stream: cuda.CUstream, + ): + # Raw tensors from from_dlpack have PyTorch layout (T, H, dim) + # stride (H*dim, dim, 1). Reshape to 3D for TMA: + # A-operands: (T, dim, H), stride (dim*H, 1, dim) + # B-operands: (dim, T, H), stride (1, dim*H, dim) + T_tok = a_raw.shape[0] + H = a_raw.shape[1] + + a = cute.make_tensor( + a_raw.iterator, + cute.make_layout( + (T_tok, a_raw.shape[2], H), + stride=(a_raw.stride[0], a_raw.stride[2], a_raw.stride[1]), + ), + ) + q = cute.make_tensor( + q_raw.iterator, + cute.make_layout( + (T_tok, q_raw.shape[2], H), + stride=(q_raw.stride[0], q_raw.stride[2], q_raw.stride[1]), + ), + ) + aqc = cute.make_tensor( + aqc_raw.iterator, + cute.make_layout( + (T_tok, aqc_raw.shape[2], H), + stride=(aqc_raw.stride[0], aqc_raw.stride[2], aqc_raw.stride[1]), + ), + ) + o_out = cute.make_tensor( + o_raw.iterator, + cute.make_layout( + (T_tok, o_raw.shape[2], H), + stride=(o_raw.stride[0], o_raw.stride[2], o_raw.stride[1]), + ), + ) + + b = cute.make_tensor( + b_raw.iterator, + cute.make_layout( + (b_raw.shape[2], T_tok, H), + stride=(b_raw.stride[2], b_raw.stride[0], b_raw.stride[1]), + ), + ) + v = cute.make_tensor( + v_raw.iterator, + cute.make_layout( + (v_raw.shape[2], T_tok, H), + stride=(v_raw.stride[2], v_raw.stride[0], v_raw.stride[1]), + ), + ) + kg = cute.make_tensor( + kg_raw.iterator, + cute.make_layout( + (kg_raw.shape[2], T_tok, H), + stride=(kg_raw.stride[2], kg_raw.stride[0], kg_raw.stride[1]), + ), + ) + + tile1 = (M, N, K) + tile3 = (M, N, K3) + tile6 = (M6, N6, K6) + + mma_kmn = sm100_utils.make_trivial_tiled_mma( + mma_dtype, + OperandMajorMode.K, + OperandMajorMode.MN, + acc_dtype, + tcgen05.CtaGroup.ONE, + (M, N), + OperandSource.SMEM, + ) + mma_mn_mn = sm100_utils.make_trivial_tiled_mma( + mma_dtype, + OperandMajorMode.MN, + OperandMajorMode.MN, + acc_dtype, + tcgen05.CtaGroup.ONE, + (M6, N6), + OperandSource.SMEM, + ) + sl_a = sm100_utils.make_smem_layout_a(mma_kmn, tile1, mma_dtype, 1) + sl_b = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1) + sl_v = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1) + sl_s = sm100_utils.make_smem_layout_b(mma_kmn, tile3, mma_dtype, 1) + sl_q = sm100_utils.make_smem_layout_a(mma_kmn, tile3, mma_dtype, 1) + sl_aqc = sm100_utils.make_smem_layout_a(mma_kmn, tile1, mma_dtype, 1) + sl_kg = sm100_utils.make_smem_layout_b(mma_mn_mn, tile6, mma_dtype, 1) + + sl_readout_k = sm100_utils.make_smem_layout_a(mma_kmn, tile3, mma_dtype, 1) + sl_nv_b = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1) + sl_nv_a = sm100_utils.make_smem_layout_a(mma_mn_mn, tile6, mma_dtype, 1) + sl_kg_a = sm100_utils.make_smem_layout_a(mma_mn_mn, tile6, mma_dtype, 1) + sl_nv_b_mn = sm100_utils.make_smem_layout_b(mma_mn_mn, tile6, mma_dtype, 1) + + tma_ld = cpasync.CopyBulkTensorTileG2SOp() + + ta_a = cute.nvgpu.make_tiled_tma_atom_A( + tma_ld, a, cute.select(sl_a, mode=[0, 1, 2]), tile1, mma_kmn + ) + ta_b = cute.nvgpu.make_tiled_tma_atom_B( + tma_ld, b, cute.select(sl_b, mode=[0, 1, 2]), tile1, mma_kmn + ) + ta_v = cute.nvgpu.make_tiled_tma_atom_B( + tma_ld, v, cute.select(sl_v, mode=[0, 1, 2]), tile1, mma_kmn + ) + ta_q = cute.nvgpu.make_tiled_tma_atom_A( + tma_ld, q, cute.select(sl_q, mode=[0, 1, 2]), tile3, mma_kmn + ) + ta_aqc = cute.nvgpu.make_tiled_tma_atom_A( + tma_ld, aqc, cute.select(sl_aqc, mode=[0, 1, 2]), tile1, mma_kmn + ) + ta_kg = cute.nvgpu.make_tiled_tma_atom_B( + tma_ld, kg, cute.select(sl_kg, mode=[0, 1, 2]), tile6, mma_mn_mn + ) + sk = sm100_utils.get_smem_layout_atom_ab(OperandMajorMode.K, out_dtype, (M, N)) + sl_store = cute.tile_to_shape( + sm100_utils.make_smem_layout_atom(sk, out_dtype), (M, N), order=(0, 1) + ) + tma_st = cpasync.CopyBulkTensorTileS2GOp() + ta_o = cpasync.make_tiled_tma_atom(tma_st, o_out, sl_store, (M, N)) + + scheduler_params = GDNTileSchedulerParams( + num_seqs=num_seqs, + num_q_heads=H, + num_v_heads=H, + is_GQA=False, + is_persistent=True, + ) + grid_shape = GDNTileScheduler.get_grid_shape(scheduler_params, _num_sm) + + k4_persistent_kernel( + mma_kmn, + mma_mn_mn, + ta_a, + sl_a, + ta_b, + sl_b, + ta_v, + sl_v, + sl_s, + ta_q, + sl_q, + ta_aqc, + sl_aqc, + ta_kg, + sl_kg, + ta_o, + sl_store, + sl_readout_k, + sl_nv_b, + sl_nv_a, + sl_kg_a, + sl_nv_b_mn, + gk_last_exp, + s_fp32, + cu_seqlens, + chunk_offsets, + a, + b, + v, + q, + aqc, + kg, + o_out, + tm_workspace, + scheduler_params, + ).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 new file mode 100644 index 000000000000..ca6ae67333ce --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py @@ -0,0 +1,678 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CuTe DSL device kernel for KDA multi-token speculative verify (conv-MTP). + +Source-integrated from the ``KDA_decode_mtp`` kernel drop ("readable KDA +conv-MTP"). Specialized for the benchmarked contract: conv enabled, no bias, +no output norm, lower_bound gate, Q/K L2 norm, beta sigmoid, TILE_V=64, +ILP=2, W=4, K == V == 128, HV == H. + +Per request ``n`` the kernel processes ``num_accepted_tokens[n] + 1 + +NUM_SPEC`` steps: it first *replays* the accepted draft tokens from the +``qkg/v/beta`` caches (raw conv inputs from the extended tail columns of +``cs_q/cs_k/cs_v``), then processes the ``1 + NUM_SPEC`` new tokens. The +recurrent state and the base conv windows are committed in place after the +first new (golden) token; the new spec tokens are cached for the next +round's replay. The pool invariant is therefore "state after the last +golden token; accepted drafts pending in the replay caches". + +Vendoring delta vs the drop: the ``v_row_a``/``v_row_b`` readout indices in +the output stage are hoisted above the ``if/elif`` on ``USE_ZERO_ACCEPTED`` +/ ``i_t >= commit_len``. In the drop they were first assigned inside the +*dynamic* ``elif`` branch, which the CuTe DSL cannot trace (names must +pre-exist before dynamic assignment), so every non-``USE_ZERO_ACCEPTED`` +compile — i.e. any replay round, and any non-benchmark shape — failed with +``DSLRuntimeError: v_row_a is None``. The hoist is semantics-preserving. +""" + +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 + +NUM_THREADS = 256 +TILE_K = 128 + + +@dsl_user_op +def read_globaltimer(*, loc=None, ip=None) -> Int64: + """Read the SM global timer for optional in-kernel stage profiling.""" + return Int64( + llvm.inline_asm( + T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.kernel +def kda_decode_mtp_kernel( + h0: cute.Tensor, + x_q: cute.Tensor, + x_k: cute.Tensor, + x_v: cute.Tensor, + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + A_log: cute.Tensor, + g: cute.Tensor, + dt_bias: cute.Tensor, + beta: cute.Tensor, + o: cute.Tensor, + ht: cute.Tensor, + qkg_cache: cute.Tensor, + v_cache: cute.Tensor, + beta_cache: cute.Tensor, + smem_qk_layout: cute.Layout, + ssm_state_indices: cute.Tensor, + cu_seqlens: cute.Tensor, + num_accepted_tokens: cute.Tensor, + precompute_control: cute.Tensor, + TILE_V: cutlass.Constexpr[int], + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + NUM_SPEC: cutlass.Constexpr[int], + KERNEL_WIDTH: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + USE_FLAT_LAYOUT: cutlass.Constexpr[bool], + USE_SETMAXREG: cutlass.Constexpr[bool], + USE_REGULAR_METADATA: cutlass.Constexpr[bool], + USE_REG_Q_WEIGHTS: cutlass.Constexpr[bool], + USE_ZERO_ACCEPTED: cutlass.Constexpr[bool], + FUSE_PRECOMPUTE: cutlass.Constexpr[bool], + RUNTIME_PRECOMPUTE_FLAG: cutlass.Constexpr[bool], + stage_timing: cute.Tensor, + PROFILE_STAGES: cutlass.Constexpr[bool], +): + """KDA MTP decode — SMEM pre-compute + register-resident state.""" + tidx, _, _ = cute.arch.thread_idx() + in_warp_tid = tidx % 32 + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + i_hv, i_n, _ = cute.arch.block_idx() + i_h = i_hv + if cutlass.const_expr(PROFILE_STAGES): + t_stage0 = read_globaltimer() + if cutlass.const_expr(USE_REGULAR_METADATA): + bos = i_n * (2 * NUM_SPEC + 1) + eos = bos + (2 * NUM_SPEC + 1) + slot = i_n + else: + bos = cu_seqlens[i_n] + eos = cu_seqlens[i_n + 1] + slot = ssm_state_indices[i_n] + h0_idx = slot * HV + i_hv + hk_off = i_h * K + hv_off = i_hv * V + if cutlass.const_expr(USE_ZERO_ACCEPTED): + commit_len = 0 + else: + commit_len = num_accepted_tokens[i_n] + if cutlass.const_expr(USE_ZERO_ACCEPTED): + T_loop = 1 + NUM_SPEC + t_max = 1 + NUM_SPEC + else: + T_loop = commit_len + 1 + NUM_SPEC + t_max = 2 * NUM_SPEC + 1 + vec_size = TILE_K // 32 + num_v_tiles = V // TILE_V + NUM_V_ROWS = TILE_V // (NUM_THREADS // 32) + if cutlass.const_expr(USE_REG_Q_WEIGHTS): + q_weight_elems = 0 + else: + q_weight_elems = KERNEL_WIDTH * K + v_weight_elems = KERNEL_WIDTH * V + k_weight_base = q_weight_elems + v_weight_base = q_weight_elems + KERNEL_WIDTH * K + conv_weight_elems = q_weight_elems + KERNEL_WIDTH * K + v_weight_elems + smem = cutlass.utils.SmemAllocator() + sQ = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) + 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) + sVall = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max * V,)), 16) + sConvW = smem.allocate_tensor( + cutlass.Float32, + cute.make_layout((conv_weight_elems,)), + 16, + ) + r_q = cute.make_rmem_tensor(cute.make_layout((vec_size,), stride=(1,)), cutlass.Float32) + 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) + 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_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): + if cutlass.const_expr(RUNTIME_PRECOMPUTE_FLAG): + run_precompute = precompute_control[0] != 0 + else: + run_precompute = True + if run_precompute: + for i in range(vec_size): + 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]) + 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]) + 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]) + 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]) + cute.arch.barrier() + if cutlass.const_expr(USE_SETMAXREG): + cute.arch.warpgroup_reg_dealloc(64) + if warp_idx < 3: + if warp_idx == 2: + r_exp_A = cute.math.exp(cutlass.Float32(A_log[i_h]), fastmath=True) + i_t = 0 + while i_t < T_loop: + if cutlass.const_expr(USE_ZERO_ACCEPTED): + replay_from_cache = False + else: + replay_from_cache = i_t < commit_len + if replay_from_cache: + 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]) + 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]) + 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]) + 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]) + for w in range(KERNEL_WIDTH - 2): + 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 + 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]) + sG[i_t, k_idx] = cute.math.exp(r_gk_c2, fastmath=True) + else: + token = bos + i_t + if warp_idx == 0: + for i_pair in range(vec_size // 2): + i0 = i_pair * 2 + i1 = i_pair * 2 + 1 + k_idx0 = i0 * 32 + in_warp_tid + k_idx1 = i1 * 32 + in_warp_tid + r_conv_0 = 0.0 + r_conv_1 = 0.0 + if cutlass.const_expr(USE_REG_Q_WEIGHTS): + _cwq_0 = r_wq[0 * vec_size + i0] + _cwq_1 = r_wq[0 * vec_size + i1] + r_conv_0 += r_state[0 * vec_size + i0] * _cwq_0 + r_conv_1 += r_state[0 * vec_size + i1] * _cwq_1 + _cwq_0 = r_wq[1 * vec_size + i0] + _cwq_1 = r_wq[1 * vec_size + i1] + r_conv_0 += r_state[1 * vec_size + i0] * _cwq_0 + r_conv_1 += r_state[1 * vec_size + i1] * _cwq_1 + _cwq_0 = r_wq[2 * vec_size + i0] + _cwq_1 = r_wq[2 * vec_size + i1] + r_conv_0 += r_state[2 * vec_size + i0] * _cwq_0 + r_conv_1 += r_state[2 * vec_size + i1] * _cwq_1 + else: + _cwq_0 = cutlass.Float32(0.0) + _cwq_1 = cutlass.Float32(0.0) + for w in range(KERNEL_WIDTH - 1): + _cwq_0 = sConvW[w * K + i0 * 32 + in_warp_tid] + _cwq_1 = sConvW[w * K + i1 * 32 + in_warp_tid] + r_conv_0 += r_state[w * vec_size + i0] * _cwq_0 + r_conv_1 += r_state[w * vec_size + i1] * _cwq_1 + if cutlass.const_expr(USE_FLAT_LAYOUT): + r_xq_0 = cutlass.Float32(x_q[0, token, hk_off + k_idx0]) + r_xq_1 = cutlass.Float32(x_q[0, token, hk_off + k_idx1]) + else: + r_xq_0 = cutlass.Float32(x_q[0, token, i_h, k_idx0]) + r_xq_1 = cutlass.Float32(x_q[0, token, i_h, k_idx1]) + if cutlass.const_expr(USE_REG_Q_WEIGHTS): + _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] + 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) + e1 = cute.math.exp(-r_conv_1, fastmath=True) + sig_0 = cute.arch.rcp_approx(cutlass.Float32(1.0) + e0) + sig_1 = cute.arch.rcp_approx(cutlass.Float32(1.0) + e1) + r_q[i0] = r_conv_0 * sig_0 + r_q[i1] = r_conv_1 * sig_1 + r_state[0 * vec_size + i0] = r_state[1 * vec_size + i0] + r_state[0 * vec_size + i1] = r_state[1 * vec_size + i1] + r_state[1 * vec_size + i0] = r_state[2 * vec_size + i0] + r_state[1 * vec_size + i1] = r_state[2 * vec_size + i1] + r_state[2 * vec_size + i0] = r_xq_0 + r_state[2 * vec_size + i1] = r_xq_1 + sum_q = 0.0 + 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 + for i in range(vec_size): + r_q[i] = r_q[i] * rnorm_q_scaled + for i in range(vec_size): + k_idx = i * 32 + in_warp_tid + sQ[i_t, k_idx] = r_q[i] + elif warp_idx == 1: + r_b_raw = cutlass.Float32(0.0) + if in_warp_tid == 0: + r_b_raw = cutlass.Float32(beta[0, token, i_hv]) + for i in range(vec_size): + k_idx = i * 32 + in_warp_tid + r_conv = ( + r_state[(KERNEL_WIDTH - 1) * vec_size + 0 * vec_size + i] + * sConvW[k_weight_base + 0 * K + i * 32 + in_warp_tid] + ) + r_conv += ( + r_state[(KERNEL_WIDTH - 1) * vec_size + 1 * vec_size + i] + * sConvW[k_weight_base + 1 * K + i * 32 + in_warp_tid] + ) + r_conv += ( + r_state[(KERNEL_WIDTH - 1) * vec_size + 2 * vec_size + i] + * sConvW[k_weight_base + 2 * K + i * 32 + in_warp_tid] + ) + if cutlass.const_expr(USE_FLAT_LAYOUT): + 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_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 + 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) + 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 + for i in range(vec_size): + k_idx = i * 32 + in_warp_tid + 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) + ) + else: + for i in range(vec_size): + k_idx = i * 32 + in_warp_tid + r_g_raw = cutlass.Float32(g[0, token, i_hv, k_idx]) + 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) + ) + r_gk = lower_bound * sigmoid_val + sG[i_t, k_idx] = cute.math.exp(r_gk, fastmath=True) + r_decay[i] = r_gk + if i_t > commit_len: + cache_pos = i_t - commit_len - 1 + for i in range(vec_size): + k_idx = i * 32 + in_warp_tid + qkg_cache[slot, cache_pos, 2, hk_off + k_idx] = r_decay[i] + if i_t == commit_len: + if warp_idx == 0: + 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] + elif warp_idx == 1: + for i in range(vec_size): + k_idx = i * 32 + in_warp_tid + for w in range(KERNEL_WIDTH - 1): + cs_k[slot, hk_off + k_idx, w] = r_state[ + (KERNEL_WIDTH - 1) * vec_size + w * vec_size + i + ] + if i_t > commit_len: + cache_pos = i_t - commit_len - 1 + 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] + 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 + ] + 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] + 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 + ] + i_t = i_t + 1 + else: + _v_idx = tidx - 96 + if _v_idx < V: + _csv0 = cutlass.Float32(cs_v[slot, hv_off + _v_idx, 0]) + _csv1 = cutlass.Float32(cs_v[slot, hv_off + _v_idx, 1]) + _csv2 = cutlass.Float32(cs_v[slot, hv_off + _v_idx, 2]) + if cutlass.const_expr(USE_ZERO_ACCEPTED): + _wv0 = sConvW[v_weight_base + 0 * V + _v_idx] + _wv1 = sConvW[v_weight_base + 1 * V + _v_idx] + _wv2 = sConvW[v_weight_base + 2 * V + _v_idx] + _wv3 = sConvW[v_weight_base + (KERNEL_WIDTH - 1) * V + _v_idx] + if cutlass.const_expr(USE_FLAT_LAYOUT): + _xv0 = cutlass.Float32(x_v[0, bos + 0, hv_off + _v_idx]) + _xv1 = cutlass.Float32(x_v[0, bos + 1, hv_off + _v_idx]) + _xv2 = cutlass.Float32(x_v[0, bos + 2, hv_off + _v_idx]) + else: + _xv0 = cutlass.Float32(x_v[0, bos + 0, i_hv, _v_idx]) + _xv1 = cutlass.Float32(x_v[0, bos + 1, i_hv, _v_idx]) + _xv2 = cutlass.Float32(x_v[0, bos + 2, i_hv, _v_idx]) + _vconv0 = _csv0 * _wv0 + _vconv0 += _csv1 * _wv1 + _vconv0 += _csv2 * _wv2 + _vconv0 += _xv0 * _wv3 + _vconv0 = _vconv0 * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-_vconv0, fastmath=True) + ) + sVall[0 * V + _v_idx] = _vconv0 + 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.fma_packed_f32x2( + (_csv2, _xv0), (_wv1, _wv1), (_vconv1, _vconv2) + ) + _vconv1, _vconv2 = cute.arch.fma_packed_f32x2( + (_xv0, _xv1), (_wv2, _wv2), (_vconv1, _vconv2) + ) + _vconv1, _vconv2 = cute.arch.fma_packed_f32x2( + (_xv1, _xv2), (_wv3, _wv3), (_vconv1, _vconv2) + ) + _vconv1 = _vconv1 * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-_vconv1, fastmath=True) + ) + _vconv2 = _vconv2 * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-_vconv2, fastmath=True) + ) + sVall[1 * V + _v_idx] = _vconv1 + v_cache[slot, 0, hv_off + _v_idx] = _vconv1 + cs_v[slot, hv_off + _v_idx, KERNEL_WIDTH - 1] = _xv1 + sVall[2 * V + _v_idx] = _vconv2 + v_cache[slot, 1, hv_off + _v_idx] = _vconv2 + cs_v[slot, hv_off + _v_idx, KERNEL_WIDTH] = _xv2 + else: + _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]) + _csv0 = _csv1 + _csv1 = _csv2 + _csv2 = _xv_replay + else: + _token_v = bos + _i_t + _v_conv = 0.0 + _v_conv += _csv0 * sConvW[v_weight_base + 0 * V + _v_idx] + _v_conv += _csv1 * sConvW[v_weight_base + 1 * V + _v_idx] + _v_conv += _csv2 * sConvW[v_weight_base + 2 * V + _v_idx] + if cutlass.const_expr(USE_FLAT_LAYOUT): + _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 = _v_conv * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-_v_conv, fastmath=True) + ) + sVall[_i_t * V + _v_idx] = _v_conv + _csv0 = _csv1 + _csv1 = _csv2 + _csv2 = _xv + if _i_t == commit_len: + cs_v[slot, hv_off + _v_idx, 0] = _csv0 + cs_v[slot, hv_off + _v_idx, 1] = _csv1 + cs_v[slot, hv_off + _v_idx, 2] = _csv2 + if _i_t > commit_len: + _cp = _i_t - commit_len - 1 + v_cache[slot, _cp, hv_off + _v_idx] = _v_conv + cs_v[slot, hv_off + _v_idx, KERNEL_WIDTH - 1 + _cp] = _xv + _i_t = _i_t + 1 + if cutlass.const_expr(PROFILE_STAGES): + cute.arch.barrier() + t_stage1 = read_globaltimer() + else: + cute.arch.barrier() + if cutlass.const_expr(USE_SETMAXREG): + cute.arch.warpgroup_reg_dealloc(64) + if cutlass.const_expr(PROFILE_STAGES): + cute.arch.barrier() + t_stage1 = read_globaltimer() + else: + cute.arch.barrier() + if cutlass.const_expr(USE_SETMAXREG): + cute.arch.warpgroup_reg_dealloc(64) + if cutlass.const_expr(PROFILE_STAGES): + cute.arch.barrier() + t_stage1 = read_globaltimer() + for row in range(NUM_V_ROWS): + 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]) + else: + 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() + for i_v in range(num_v_tiles): + v_base = i_v * TILE_V + if i_v > 0: + for row in range(NUM_V_ROWS): + 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_base + v_row, i * 32 + in_warp_tid] + ) + else: + r_state[row * vec_size + i] = cutlass.Float32( + h0[slot, i_hv, v_base + v_row, i * 32 + in_warp_tid] + ) + r_v_val = cutlass.Float32(0.0) + i_t = 0 + while i_t < T_loop: + if in_warp_tid < NUM_V_ROWS: + v_idx = v_base + warp_idx * NUM_V_ROWS + in_warp_tid + r_v_val = sVall[i_t * V + v_idx] + r_beta_val = sBeta[i_t] + for i_pair in range(vec_size // 2): + i0 = i_pair * 2 + i1 = i_pair * 2 + 1 + k_idx0 = i0 * 32 + in_warp_tid + k_idx1 = i1 * 32 + in_warp_tid + r_q[i0] = sQ[i_t, k_idx0] + r_q[i1] = sQ[i_t, k_idx1] + _k0 = sK[i_t, k_idx0] + _k1 = sK[i_t, k_idx1] + r_decay[i0] = sG[i_t, k_idx0] + r_decay[i1] = sG[i_t, k_idx1] + r_bk[i0], r_bk[i1] = cute.arch.mul_packed_f32x2( + (r_beta_val, r_beta_val), (_k0, _k1) + ) + r_k[i0], r_k[i1] = cute.arch.mul_packed_f32x2( + (r_decay[i0], r_decay[i1]), (_k0, _k1) + ) + for row_pair in range(NUM_V_ROWS // 2): + ra = row_pair * 2 + rb = row_pair * 2 + 1 + r_va = cute.arch.shuffle_sync(r_v_val, ra) + r_vb = cute.arch.shuffle_sync(r_v_val, rb) + shk_a1 = 0.0 + shk_a2 = 0.0 + shk_b1 = 0.0 + shk_b2 = 0.0 + for _pi in range(vec_size // 2): + _p = _pi * 2 + shk_a1, shk_a2 = cute.arch.fma_packed_f32x2( + src_a=( + r_state[ra * vec_size + _p], + r_state[ra * vec_size + _p + 1], + ), + src_b=(r_k[_p], r_k[_p + 1]), + src_c=(shk_a1, shk_a2), + ) + shk_b1, shk_b2 = cute.arch.fma_packed_f32x2( + src_a=( + r_state[rb * vec_size + _p], + r_state[rb * vec_size + _p + 1], + ), + src_b=(r_k[_p], r_k[_p + 1]), + src_c=(shk_b1, shk_b2), + ) + 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) + vn_a = r_va - shk_a + vn_b = r_vb - shk_b + shq_a1 = 0.0 + shq_a2 = 0.0 + shq_b1 = 0.0 + shq_b2 = 0.0 + for _pi in range(vec_size // 2): + _p = _pi * 2 + vnbk_a0, vnbk_a1 = cute.arch.mul_packed_f32x2( + (vn_a, vn_a), (r_bk[_p], r_bk[_p + 1]) + ) + 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[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=( + r_state[ra * vec_size + _p], + r_state[ra * vec_size + _p + 1], + ), + src_b=(r_q[_p], r_q[_p + 1]), + src_c=(shq_a1, shq_a2), + ) + shq_b1, shq_b2 = cute.arch.fma_packed_f32x2( + src_a=( + r_state[rb * vec_size + _p], + r_state[rb * vec_size + _p + 1], + ), + src_b=(r_q[_p], r_q[_p + 1]), + src_c=(shq_b1, shq_b2), + ) + 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) + if in_warp_tid == 0: + v_row_a = warp_idx * NUM_V_ROWS + ra + v_row_b = warp_idx * NUM_V_ROWS + rb + if cutlass.const_expr(USE_ZERO_ACCEPTED): + o[0, bos + i_t, i_hv, v_base + v_row_a] = cutlass.BFloat16(shq_a) + o[0, bos + i_t, i_hv, v_base + v_row_b] = cutlass.BFloat16(shq_b) + elif i_t >= commit_len: + o[0, bos + i_t, i_hv, v_base + v_row_a] = cutlass.BFloat16(shq_a) + o[0, bos + i_t, i_hv, v_base + v_row_b] = cutlass.BFloat16(shq_b) + if i_t == commit_len: + for row in range(NUM_V_ROWS): + 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] + else: + 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() + t_stage2 = read_globaltimer() + if tidx == 0: + _, grid_n, _ = cute.arch.grid_dim() + timing_base = (i_hv * grid_n + i_n) * 4 + stage_timing[timing_base + 0] = t_stage1 - t_stage0 + stage_timing[timing_base + 1] = t_stage2 - t_stage1 + stage_timing[timing_base + 2] = t_stage2 - t_stage0 + stage_timing[timing_base + 3] = t_stage0 From ae37fbead655dd8549c1f0575892ae86b7d3910b Mon Sep 17 00:00:00 2001 From: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com> Date: Mon, 3 Aug 2026 06:29:23 -0700 Subject: [PATCH 03/12] fix cache key for fused prefill Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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..46fb320771aa 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 @@ -825,8 +825,9 @@ def _launch_fused_k1234( NC = T // _BT has_bias = dt_bias is not None - # Dynamic shapes: cache only on mode flags, not B/NC/H - cache_key = (dev, has_bias, safe_gate) + # B only changes the outer extent and launch grid. T, H, and K change + # compiled tensor layouts, so they must select distinct artifacts. + cache_key = (dev, has_bias, safe_gate, T, H, K) # Buffers BH = B * H From 7f373e5f4a249c7e968f34a059ff5b13d337a8f2 Mon Sep 17 00:00:00 2001 From: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com> Date: Mon, 3 Aug 2026 06:29:46 -0700 Subject: [PATCH 04/12] fix cache key of batchsize for decode mtp Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com> --- .../cute_dsl_kimi_k3_kda_mtp_ops.py | 67 ++++++++++--------- 1 file changed, 36 insertions(+), 31 deletions(-) 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..334eefb957fc 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 @@ -107,7 +107,7 @@ def _run_kda_decode_mtp( HV: cutlass.Constexpr[int], K: cutlass.Constexpr[int], V: cutlass.Constexpr[int], - N: cutlass.Constexpr[int], + N: cutlass.Int32, NUM_SPEC: cutlass.Constexpr[int], TILE_V: cutlass.Constexpr[int], KERNEL_WIDTH: cutlass.Constexpr[int], @@ -291,11 +291,6 @@ def _require_stride_layout( raise ValueError("Expected num_accepted_tokens length to match N.") -def _layout_key(tensor: torch.Tensor): - return (tensor.dtype, tuple(tensor.shape), tuple(tensor.stride()), - _fits_32bit_stride(tensor)) - - def _fits_32bit_stride(tensor: torch.Tensor) -> bool: int32_max = 2**31 - 1 max_offset = int(tensor.storage_offset()) @@ -322,7 +317,18 @@ def _dlpack_arg(tensor: torch.Tensor): for dim, stride in enumerate(tensor.stride()): if stride == 1: return _from_dlpack_arg(tensor).mark_layout_dynamic(dim) - return _from_dlpack_arg(tensor) + return _from_dlpack_arg(tensor).mark_layout_dynamic() + + +def _layout_key(tensor: torch.Tensor, dynamic_layout: bool = False): + arg = _dlpack_arg(tensor) if dynamic_layout else _from_dlpack_arg(tensor) + shape_mask = arg.dynamic_shapes_mask + stride_mask = arg.dynamic_strides_mask + shape = tuple(None if dynamic else size + for size, dynamic in zip(tensor.shape, shape_mask)) + stride = tuple(None if dynamic else value + for value, dynamic in zip(tensor.stride(), stride_mask)) + return (tensor.dtype, shape, stride, _fits_32bit_stride(tensor)) # (device_index, enabled) -> persistent int32 [1] control tensor. Keys are @@ -370,9 +376,9 @@ def _is_benchmark_static_shape(N: int, H: int, HV: int, K: int, V: int, and N in (32, 128) and H in (2, 12, 32)) -# Layout-and-constexpr-keyed compile cache. Compilation is per (N, T_total, -# pool size, layouts, flags) — a new generation batch size triggers a -# multi-second cute.compile, after which the artifact is reused. +# Layout-and-constexpr-keyed compile cache. Request count and packed-token +# length are dynamic; batches sharing the same kernel variant reuse one +# artifact even when their launch grid and token-buffer extents differ. _compiled_cache = {} @@ -521,15 +527,13 @@ def kda_mtp_decode_impl( V_dim, num_spec, W, - N, - T_total, pool_size, lower_bound, use_flat_layout, _layout_key(h0_arg), - _layout_key(x_q_arg), - _layout_key(x_k_arg), - _layout_key(x_v_arg), + _layout_key(x_q_arg, dynamic_layout=True), + _layout_key(x_k_arg, dynamic_layout=True), + _layout_key(x_v_arg, dynamic_layout=True), _layout_key(w_q), _layout_key(w_k), _layout_key(w_v), @@ -537,16 +541,16 @@ def kda_mtp_decode_impl( _layout_key(cs_k), _layout_key(cs_v), _layout_key(A_log), - _layout_key(g), + _layout_key(g, dynamic_layout=True), _layout_key(dt_bias), - _layout_key(beta), - _layout_key(out), + _layout_key(beta, dynamic_layout=True), + _layout_key(out, dynamic_layout=True), _layout_key(qkg_cache), _layout_key(v_cache), _layout_key(beta_cache), - _layout_key(ssm_state_indices), - _layout_key(cu_seqlens), - _layout_key(num_accepted_tokens), + _layout_key(ssm_state_indices, dynamic_layout=True), + _layout_key(cu_seqlens, dynamic_layout=True), + _layout_key(num_accepted_tokens, dynamic_layout=True), use_setmaxreg, use_regular_metadata, use_reg_q_weights, @@ -562,9 +566,9 @@ def kda_mtp_decode_impl( _compiled_cache[key] = cute.compile( _run_kda_decode_mtp, _from_dlpack_arg(h0_arg), - _from_dlpack_arg(x_q_arg), - _from_dlpack_arg(x_k_arg), - _from_dlpack_arg(x_v_arg), + _dlpack_arg(x_q_arg), + _dlpack_arg(x_k_arg), + _dlpack_arg(x_v_arg), _from_dlpack_arg(w_q), _from_dlpack_arg(w_k), _from_dlpack_arg(w_v), @@ -572,18 +576,18 @@ def kda_mtp_decode_impl( _from_dlpack_arg(cs_k), _from_dlpack_arg(cs_v), _from_dlpack_arg(A_log), - _from_dlpack_arg(g), + _dlpack_arg(g), _from_dlpack_arg(dt_bias), - _from_dlpack_arg(beta), - _from_dlpack_arg(out), + _dlpack_arg(beta), + _dlpack_arg(out), _from_dlpack_arg(h0_arg), _from_dlpack_arg(qkg_cache), _from_dlpack_arg(v_cache), _from_dlpack_arg(beta_cache), - _from_dlpack_arg(stage_timing_arg), - _from_dlpack_arg(ssm_state_indices), - _from_dlpack_arg(cu_seqlens), - _from_dlpack_arg(num_accepted_tokens), + _dlpack_arg(stage_timing_arg), + _dlpack_arg(ssm_state_indices), + _dlpack_arg(cu_seqlens), + _dlpack_arg(num_accepted_tokens), _from_dlpack_arg(precompute_control), scale=scale, HV=HV, @@ -630,6 +634,7 @@ def kda_mtp_decode_impl( _dlpack_arg(cu_seqlens), _dlpack_arg(num_accepted_tokens), _dlpack_arg(precompute_control), + N, stream, ) From 3415e6a4128b254e9f5d0b734c2d498777721178 Mon Sep 17 00:00:00 2001 From: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com> Date: Mon, 3 Aug 2026 08:43:17 -0700 Subject: [PATCH 05/12] chore: apply pre-commit fixes Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com> --- .../custom_ops/cute_dsl_kimi_k3_custom_ops.py | 25 +- .../cute_dsl_kimi_k3_kda_mtp_ops.py | 108 ++++----- .../blackwell/kimi_k3_kda/fused_k123.py | 24 +- .../blackwell/kimi_k3_kda/fused_k1234.py | 4 +- .../blackwell/kimi_k3_kda/k4_persistent.py | 3 +- .../blackwell/kimi_k3_kda/kda_mtp_decode.py | 226 ++++++++++++------ 6 files changed, 235 insertions(+), 155 deletions(-) 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 46fb320771aa..82d0e8ef82af 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 @@ -331,8 +330,7 @@ def _get_padded_input_buffers(B, T_padded, H, K, dtype_qkv, dtype_g, dtype_beta, 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: @@ -792,12 +790,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 @@ -924,10 +920,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: @@ -1029,7 +1023,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 @@ -1061,7 +1056,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 @@ -1080,8 +1076,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 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 334eefb957fc..4c0afabab520 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]: @@ -324,10 +321,10 @@ def _layout_key(tensor: torch.Tensor, dynamic_layout: bool = False): arg = _dlpack_arg(tensor) if dynamic_layout else _from_dlpack_arg(tensor) shape_mask = arg.dynamic_shapes_mask stride_mask = arg.dynamic_strides_mask - shape = tuple(None if dynamic else size - for size, dynamic in zip(tensor.shape, shape_mask)) - stride = tuple(None if dynamic else value - for value, dynamic in zip(tensor.stride(), stride_mask)) + shape = tuple(None if dynamic else size for size, dynamic in zip(tensor.shape, shape_mask)) + stride = tuple( + None if dynamic else value for value, dynamic in zip(tensor.stride(), stride_mask) + ) return (tensor.dtype, shape, stride, _fits_32bit_stride(tensor)) @@ -336,15 +333,13 @@ def _layout_key(tensor: torch.Tensor, dynamic_layout: bool = False): _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] @@ -370,10 +365,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. Request count and packed-token @@ -454,12 +457,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) @@ -509,8 +507,7 @@ 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) @@ -562,7 +559,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), 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..adc25e90ad65 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 @@ -1525,20 +1525,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 diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k1234.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k1234.py index f4f617814918..32dbdb2297c7 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k1234.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k1234.py @@ -2269,7 +2269,9 @@ def host_fn( num_heads, batch_size, ).launch( - grid=(BH, 1, 1), block=(THREADS, 1, 1), smem=225 * 1024, + grid=(BH, 1, 1), + block=(THREADS, 1, 1), + smem=225 * 1024, stream=stream, ) # Force high SMEM to prevent >1 block per SM (TMEM conflict) 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..fcab7bd6b673 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 @@ -1453,7 +1453,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..8aef40782c8b 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 @@ -155,10 +155,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 +169,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 +188,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 +226,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 +309,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 +333,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 +366,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 +401,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 +411,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 +427,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 +442,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 +498,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 +526,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 +545,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 +587,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 +665,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 +685,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 +724,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 +744,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() From 57b38ed0357b96887335fafad0b1678daf253de3 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Mon, 3 Aug 2026 14:28:29 -0700 Subject: [PATCH 06/12] [None][chore] Apply clang-format to attnResFwd.cu Signed-off-by: Brian Nguyen --- .../kernels/kimiK3AttnRes/attnResFwd.cu | 1364 ++++++++--------- 1 file changed, 675 insertions(+), 689 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu index 36a297bec25f..8e1e1b391e8b 100644 --- a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu @@ -44,145 +44,147 @@ 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 +192,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 +210,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 +232,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 +520,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 +568,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 +635,149 @@ 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; } } 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 +785,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 +796,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 +921,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 +961,37 @@ 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; + 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); + if (!attrs_set) + { + if (smem_size > 48 * 1024) + { + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); } 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 +1000,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 +1015,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 +1126,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 +1151,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 +1172,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 +1187,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 +1251,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 +1300,44 @@ 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); + if (!attrs_set) + { + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); attrs_set = true; } - 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 +1350,28 @@ 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 (!attrs_set) + { + if (smem_size > 48 * 1024) + { + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); } 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 From d75d3090b2b081cdbd397386176cf8bdaa63ec5a Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Mon, 3 Aug 2026 16:01:37 -0700 Subject: [PATCH 07/12] [None][fix] Address review feedback on KDA prefill/MTP kernels and attn-res op - attnResFwd.cu: publish plan.logits_all with __syncwarp before the tail cross-lane reads; per-device cudaFuncSetAttribute via std::once_flag in all three launchers; replace silent returns with TLLM_CHECK_WITH_INFO (SM query, N > N_MAX, unsupported H); clamp num_sm - 1 grid to >= 1 - attnResOp.cpp: check CUDA placement, then set the device guard, then run the contract check so is_sm100_family() sees the tensors' device - cute_dsl_kimi_k3_custom_ops.py: LRU-bound the padded-input and g-sentinel scratch caches; validate A_log; document the K4 positive sequence-length precondition - cute_dsl_kimi_k3_kda_mtp_ops.py: gate the zero-accepted fast path on num_spec == 2; assert profiling stays off while stage_timing aliases out; fake kernel output dtype from x_q - kda_mtp_decode.py: clamp commit_len to NUM_SPEC on device; pre-declare t_stage1 for the profiling path; document stage_timing requirements - fused_k123.py: cover all 64 SMEM rows in the A_qk/A_kk zero-init (31 warps left rows 62-63 stale for the transposed A_kk store path); fix stale warp/thread counts in comments - k4_persistent.py: narrow monkey-patch exception handling, route messages through tensorrt_llm.logger, document the global ptx-options effect Signed-off-by: Brian Nguyen --- .../kernels/kimiK3AttnRes/attnResFwd.cu | 78 ++++++++++++++----- cpp/tensorrt_llm/thop/attnResOp.cpp | 8 +- .../custom_ops/cute_dsl_kimi_k3_custom_ops.py | 23 +++++- .../cute_dsl_kimi_k3_kda_mtp_ops.py | 22 +++++- .../blackwell/kimi_k3_kda/fused_k123.py | 49 ++++++------ .../blackwell/kimi_k3_kda/k4_persistent.py | 36 +++++++-- .../blackwell/kimi_k3_kda/kda_mtp_decode.py | 17 +++- 7 files changed, 174 insertions(+), 59 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu index 8e1e1b391e8b..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,6 +40,7 @@ #include #include #include +#include namespace { @@ -716,6 +719,10 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t c 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; @@ -980,14 +987,23 @@ static void launch_fwd(bf16_t const* block_residual, bf16_t const* layer_residua 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) { - 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) { - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + std::call_once(attrs_set[dev], set_attr); + } + else + { + set_attr(); } - attrs_set = true; } int grid = RELEASE_TMEM ? num_sm * 2 : num_sm; kernel<<>>( @@ -1329,11 +1345,22 @@ static void launch_s1_splitk(bf16_t const* block_residual, bf16_t const* layer_r 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, @@ -1361,14 +1388,23 @@ static void launch_n1_ttile(bf16_t const* layer_residual, bf16_t const* res_weig 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) { - 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) { - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + 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); @@ -1421,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) { @@ -1476,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 { @@ -1512,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 82d0e8ef82af..63df6fa20ab3 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 @@ -281,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 @@ -298,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 @@ -326,7 +334,12 @@ 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 @@ -485,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 @@ -1209,6 +1228,8 @@ 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}") + if A_log is None: + 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 4c0afabab520..94b76ee5b161 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 @@ -511,9 +511,23 @@ def kda_mtp_decode_impl( 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 = ( @@ -603,7 +617,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, ) @@ -738,4 +752,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 adc25e90ad65..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 @@ -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 fcab7bd6b673..460a7ddfb8cf 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,18 @@ 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 -except Exception as e: - print(f" [WARN] monkey-patch failed: {e}") +except (AttributeError, ImportError) as e: + _logger.warning(f"k4_persistent ptx-options patch failed: {e}") _patch_applied = False 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 8aef40782c8b..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 @@ -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 From e7e277e5133a2156ad59f71b1d18b6a8bd8a2539 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Tue, 4 Aug 2026 05:01:22 -0700 Subject: [PATCH 08/12] [None][fix] Revert tuple except handler that breaks CuTe DSL compilation The review-round change narrowing the k4_persistent monkey-patch handler to `except (AttributeError, ImportError)` breaks every cute.compile of the module: the nvidia-cutlass-dsl 4.5.0 AST preprocessor cannot parse tuple except handlers ("'Tuple' object has no attribute 'id'"). Revert to a single bare Exception with a comment explaining the constraint. Also gate the new A_log validation on a non-empty token batch: zero-token calls take the early return in _chunk_kda_fwd and never touch A_log, and the runtime emits such batches under the overlap scheduler + logprobs flows. Signed-off-by: Brian Nguyen --- .../_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py | 6 +++++- .../blackwell/kimi_k3_kda/k4_persistent.py | 7 ++++++- 2 files changed, 11 insertions(+), 2 deletions(-) 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 63df6fa20ab3..5ff484f2d999 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 @@ -1228,7 +1228,11 @@ 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}") - if A_log is None: + # 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( 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 460a7ddfb8cf..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 @@ -474,7 +474,12 @@ def _patched_get_pipeline(self, _pipeline): return result _CuTeDSL._get_pipeline = _patched_get_pipeline -except (AttributeError, ImportError) as e: +# 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: _logger.warning(f"k4_persistent ptx-options patch failed: {e}") _patch_applied = False From 9b35e959dc535de515eb559fac1c00d7110adc70 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Tue, 4 Aug 2026 11:03:25 -0700 Subject: [PATCH 09/12] [None][fix] Re-raise unexpected errors in the k4 ptx-options patch Review follow-up: the patch's bare `except Exception` swallowed every failure kind, so an unexpected error (e.g. a changed CUTLASS API) would leave the module silently unpatched and surface later as a ptxas register-budget failure at the first K4 compile, with no trail back to the patch. Re-raise anything that is not the expected AttributeError / ImportError. The handler clause itself stays a single `except Exception` with an isinstance check in the body: the nvidia-cutlass-dsl 4.5.0 AST preprocessor cannot parse tuple except handlers anywhere in a kernel module (see the previous commit's revert). Add CPU-only contract tests pinning both constraints (no tuple except handlers in the module; benign kinds soft-continue while unexpected errors fail the import loudly) and wire them into the CPU pre-merge lists. Nothing else in CI compiles these kernels yet, so the structural check is the only automated guard against reintroducing the parser breakage. Signed-off-by: Brian Nguyen --- .../blackwell/kimi_k3_kda/k4_persistent.py | 11 +- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../cute_dsl/test_kimi_k3_kda_ptx_patch.py | 168 ++++++++++++++++++ 3 files changed, 178 insertions(+), 2 deletions(-) create mode 100644 tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py 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 a71d76f4dc37..7c118a3051b6 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 @@ -477,9 +477,16 @@ def _patched_get_pipeline(self, _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. +# cute.compile of this file. Keep a single Exception clause until the DSL +# pin moves past the bug; the isinstance re-raise below narrows it to the +# expected failure kinds without putting a tuple in the except clause. except Exception as e: + if not isinstance(e, (AttributeError, ImportError)): + # Not a known-benign "DSL not present / API surface moved" case: + # swallowing it would defer the blow-up to the first K4 compile + # (ptxas over WG0's 40-register budget without --uumn) with no + # trail back to this patch. Fail here, where the cause is visible. + raise _logger.warning(f"k4_persistent ptx-options patch failed: {e}") _patch_applied = False diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 5a1a08f2e19b..bc2d8c0dea3d 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -12,6 +12,7 @@ l0_cpu: backend: generic orchestrator: mpi tests: + - unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py - unittest/_torch/distributed - unittest/_torch/executor - unittest/_torch/lora diff --git a/tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py b/tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py new file mode 100644 index 000000000000..a3a0dfed9116 --- /dev/null +++ b/tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py @@ -0,0 +1,168 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contract tests for the k4_persistent ptx-options monkey-patch. + +The patch's exception handling encodes two constraints that are invisible in +its behavior and have already round-tripped through review once: + + * The handler must stay a single ``except Exception`` clause: the + nvidia-cutlass-dsl 4.5.0 AST preprocessor rejects tuple except handlers + anywhere in a kernel module ("'Tuple' object has no attribute 'id'"), + breaking every ``cute.compile`` of the file. A narrowing to + ``except (AttributeError, ImportError)`` had to be reverted for exactly + this (commit f79c6ef3af6). + * Only the expected "DSL absent / API surface moved" kinds + (``AttributeError``, ``ImportError``) may be swallowed. Anything else + must fail the import loudly: silently skipping the patch defers the + blow-up to the first K4 compile (ptxas over WG0's 40-register budget + without ``--uumn``) with no trail back to the patch. + +CPU-only: the import scenarios run the module top level in subprocesses (no +kernel is compiled); the parser constraint is checked structurally on the AST. +""" + +import ast +import importlib.util +import subprocess +import sys +from pathlib import Path + +import pytest + +try: + import cutlass # noqa: F401 + import flashinfer.gdn_kernels # noqa: F401 + + _KERNEL_DEPS_AVAILABLE = True +except Exception: + _KERNEL_DEPS_AVAILABLE = False + +needs_kernel_deps = pytest.mark.skipif( + not _KERNEL_DEPS_AVAILABLE, reason="requires nvidia-cutlass-dsl and flashinfer (gdn_kernels)" +) + + +def _k4_path() -> Path: + """The k4_persistent.py that is actually in use (installed package if + importable, else the source tree this test file lives in).""" + try: + spec = importlib.util.find_spec( + "tensorrt_llm._torch.cute_dsl_kernels.blackwell.kimi_k3_kda.k4_persistent" + ) + if spec is not None and spec.origin: + return Path(spec.origin) + except Exception: + pass + return ( + Path(__file__).resolve().parents[4] + / "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/" + "k4_persistent.py" + ) + + +def test_no_tuple_except_handlers(): + """Guard against reintroducing ``except (A, B):`` anywhere in the module. + + Nothing else in CI compiles these kernels, so without this check the + breakage only shows up as a runtime cute.compile failure on Blackwell. + """ + tree = ast.parse(_k4_path().read_text()) + bad = [ + handler.lineno + for node in ast.walk(tree) + for handler in getattr(node, "handlers", []) + if isinstance(handler.type, ast.Tuple) + ] + assert not bad, ( + f"tuple except handler(s) at line(s) {bad} of k4_persistent.py: the " + f"nvidia-cutlass-dsl 4.5.0 AST preprocessor cannot parse these and " + f"every cute.compile of the module fails. Use a single `except " + f"Exception` with an isinstance re-raise in the body instead." + ) + + +_LOAD_K4 = """ +import importlib.util, sys +spec = importlib.util.spec_from_file_location("k4_under_test", {k4!r}) +mod = importlib.util.module_from_spec(spec) +spec.loader.exec_module(mod) +""" + +_FAKE_DSL = """ +import sys, types +import cutlass.cutlass_dsl # materialize the real parent package first +fake = types.ModuleType("cutlass.cutlass_dsl.cutlass") +{fake_body} +sys.modules["cutlass.cutlass_dsl.cutlass"] = fake +""" + + +def _run_import(inject: str = "") -> subprocess.CompletedProcess: + code = _FAKE_DSL.format(fake_body=inject) if inject else "" + code += _LOAD_K4.format(k4=str(_k4_path())) + return subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=600) + + +@needs_kernel_deps +def test_import_applies_patch(): + proc = subprocess.run( + [ + sys.executable, + "-c", + _LOAD_K4.format(k4=str(_k4_path())) + + """ +from cutlass.cutlass_dsl.cutlass import CuTeDSL +assert CuTeDSL._get_pipeline.__name__ == "_patched_get_pipeline", \\ + f"ptx-options patch not applied: {CuTeDSL._get_pipeline}" +""", + ], + capture_output=True, + text=True, + timeout=600, + ) + assert proc.returncode == 0, proc.stderr[-2000:] + + +@needs_kernel_deps +def test_expected_patch_failure_is_soft(): + """AttributeError (API surface moved) logs a warning and continues.""" + proc = _run_import( + inject=""" +class CuTeDSL: # no _get_pipeline -> AttributeError in the patch block + pass +fake.CuTeDSL = CuTeDSL +""" + ) + assert proc.returncode == 0, proc.stderr[-2000:] + + +@needs_kernel_deps +def test_unexpected_patch_failure_raises(): + """Anything but AttributeError/ImportError must fail the import.""" + proc = _run_import( + inject=""" +class _Meta(type): + def __getattr__(cls, name): + raise RuntimeError("simulated unexpected patch failure") +class CuTeDSL(metaclass=_Meta): + pass +fake.CuTeDSL = CuTeDSL +""" + ) + assert proc.returncode != 0, ( + "an unexpected exception in the ptx-options patch was swallowed at " + "import; it must propagate (a silently skipped patch surfaces later " + "as an unexplained ptxas register-budget failure)" + ) + assert "simulated unexpected patch failure" in proc.stderr From ec96f8bfde47a94c19aaad48a8a425b65121ef0a Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Tue, 4 Aug 2026 12:14:42 -0700 Subject: [PATCH 10/12] Address trivial review comments Signed-off-by: Brian Nguyen --- tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py b/tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py index a3a0dfed9116..292835875484 100644 --- a/tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py +++ b/tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py @@ -45,7 +45,7 @@ import flashinfer.gdn_kernels # noqa: F401 _KERNEL_DEPS_AVAILABLE = True -except Exception: +except ImportError: _KERNEL_DEPS_AVAILABLE = False needs_kernel_deps = pytest.mark.skipif( @@ -62,7 +62,7 @@ def _k4_path() -> Path: ) if spec is not None and spec.origin: return Path(spec.origin) - except Exception: + except ModuleNotFoundError: pass return ( Path(__file__).resolve().parents[4] From 666c9a36fc86742a5a91d5afadfdc12c01ab4605 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Tue, 4 Aug 2026 18:49:01 -0700 Subject: [PATCH 11/12] [None][fix] Mark k4 ptx-patch contract tests cpu_only for the L0 CPU stages The L0 CPU-Generic stages run the inner pytest with -m cpu_only, so without the marker all four tests are deselected and pytest exits 5 (no tests collected), which test_unittests_v2 reports as a failure on both x86_64 and SBSA. Signed-off-by: Brian Nguyen --- .../unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py b/tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py index 292835875484..3dd42fb498ca 100644 --- a/tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py +++ b/tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py @@ -40,6 +40,12 @@ import pytest +# The L0 CPU-Generic stages run `pytest -m cpu_only`, and their conftest only +# collects files containing the literal string "pytest.mark.cpu_only"; without +# this marker every test here is deselected and pytest exits 5 (no tests +# collected), which the test_unittests_v2 wrapper reports as a failure. +pytestmark = pytest.mark.cpu_only + try: import cutlass # noqa: F401 import flashinfer.gdn_kernels # noqa: F401 From 0434374502eb88559658b1320ca20accdab20469 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Wed, 5 Aug 2026 12:56:28 -0700 Subject: [PATCH 12/12] [None][fix] Add a fake impl for the trtllm::attn_res_fwd custom op test_custom_ops.py::test_register_fake requires every trtllm:: op to register a fake (meta) kernel. Register one for attn_res_fwd matching the CUDA impl: output is empty_like(layer_residual) and rsigma, probs, and logits are [N, T, B] fp32 with N = block_residual.size(0) + 1. Signed-off-by: Brian Nguyen --- .../_torch/custom_ops/cpp_custom_ops.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 1cd00b3e9b6c..7a5408b77b79 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -140,6 +140,23 @@ def _(q, k, norm_weight_q, norm_weight_k, workspace, rank, nranks, eps, def _(q: torch.Tensor, num_heads: int, head_dim: int, eps: float): return torch.empty_like(q) + @torch.library.register_fake("trtllm::attn_res_fwd") + def _( + layer_residual: torch.Tensor, block_residual: torch.Tensor, + res_weight: torch.Tensor, rms_weight: torch.Tensor, rms_eps: float + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + # layer_residual: [T, B, H] bf16; block_residual: [N - 1, T, B, H]. + num_candidates = block_residual.shape[0] + 1 + seq_len, batch_size = layer_residual.shape[0], layer_residual.shape[1] + output = torch.empty_like(layer_residual) + rsigma = layer_residual.new_empty((num_candidates, seq_len, batch_size), + dtype=torch.float32) + probs = layer_residual.new_empty((num_candidates, seq_len, batch_size), + dtype=torch.float32) + logits = layer_residual.new_empty((num_candidates, seq_len, batch_size), + dtype=torch.float32) + return output, rsigma, probs, logits + @torch.library.register_fake("trtllm::fused_inv_rope_fp8_quant_vllm_port") def _(o: torch.Tensor, positions: torch.Tensor, cos_sin_cache: torch.Tensor, n_groups: int, heads_per_group: int, nope_dim: int, rope_dim: int,