diff --git a/csrc/kda/cake_kda_packed_t1_binding.cuh b/csrc/kda/cake_kda_packed_t1_binding.cuh new file mode 100644 index 00000000000..dc2d76bff8d --- /dev/null +++ b/csrc/kda/cake_kda_packed_t1_binding.cuh @@ -0,0 +1,281 @@ +/* + * 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 + +#ifndef CAKE_KDA_PACKED_T1_BODY_FILE +#error "CAKE_KDA_PACKED_T1_BODY_FILE must name one frozen generated body" +#endif +#ifndef CAKE_KDA_PACKED_T1_KERNEL +#error "CAKE_KDA_PACKED_T1_KERNEL must name the frozen kernel symbol" +#endif +#ifndef CAKE_KDA_PACKED_T1_VALUE_TILES +#error "CAKE_KDA_PACKED_T1_VALUE_TILES must describe the frozen value tiling" +#endif +#ifndef CAKE_KDA_PACKED_T1_THREADS +#error "CAKE_KDA_PACKED_T1_THREADS must describe the frozen thread count" +#endif +#ifndef CAKE_KDA_PACKED_T1_SMEM_BYTES +#error "CAKE_KDA_PACKED_T1_SMEM_BYTES must describe dynamic shared memory" +#endif +#ifndef CAKE_KDA_PACKED_T1_REQUIRES_AUX_VEC4 +#error "CAKE_KDA_PACKED_T1_REQUIRES_AUX_VEC4 must describe auxiliary alignment" +#endif +#ifndef FLASHINFER_CAKE_KDA_PACKED_T1_TARGET_KIND +#error "FLASHINFER_CAKE_KDA_PACKED_T1_TARGET_KIND must identify the target" +#endif + +#include +#include +#include +#include + +#include +#include +#include + +#include "tvm_ffi_utils.h" + +// Generated bodies carry private fixed-width aliases and a tensor-map stand-in. +// Rename them at the include boundary so they cannot collide with CUDA headers. +#define uint8_t cake_kda_packed_generated_uint8_t +#define uint16_t cake_kda_packed_generated_uint16_t +#define uint32_t cake_kda_packed_generated_uint32_t +#define uint64_t cake_kda_packed_generated_uint64_t +#define int32_t cake_kda_packed_generated_int32_t +#define int16_t cake_kda_packed_generated_int16_t +#define CakeTensorMap cake_kda_packed_generated_CakeTensorMap +#define CakeTensorMapPack cake_kda_packed_generated_CakeTensorMapPack +#define CUtensorMap cake_kda_packed_generated_CUtensorMap +#include CAKE_KDA_PACKED_T1_BODY_FILE +#undef uint8_t +#undef uint16_t +#undef uint32_t +#undef uint64_t +#undef int32_t +#undef int16_t +#undef CakeTensorMap +#undef CakeTensorMapPack +#undef CUtensorMap +#undef THREADS +#undef NUM_MAIN_STAGES +#undef CAKE_INF + +namespace flashinfer { +namespace cake_kda_packed_t1 { + +constexpr int32_t kHeads = 12; +constexpr int32_t kHeadDim = 128; +constexpr int32_t kMixedWidth = 3 * kHeads * kHeadDim; +constexpr int32_t kGateWidth = kHeads * kHeadDim; +constexpr int32_t kTargetFamily = 100; +constexpr int32_t kTargetSM100a = 1000; +constexpr int32_t kTargetKind = FLASHINFER_CAKE_KDA_PACKED_T1_TARGET_KIND; + +static_assert(kTargetKind == kTargetFamily || kTargetKind == kTargetSM100a, + "packed KDA T=1 must be compiled for SM100f or legacy exact SM100a"); +static_assert(CAKE_KDA_PACKED_T1_VALUE_TILES == 1 || CAKE_KDA_PACKED_T1_VALUE_TILES == 2 || + CAKE_KDA_PACKED_T1_VALUE_TILES == 8 || CAKE_KDA_PACKED_T1_VALUE_TILES == 16, + "packed KDA T=1 has an unsupported value tiling"); +static_assert(CAKE_KDA_PACKED_T1_THREADS == 32 || CAKE_KDA_PACKED_T1_THREADS == 128, + "packed KDA T=1 has an unsupported thread count"); +static_assert(CAKE_KDA_PACKED_T1_REQUIRES_AUX_VEC4 == 0 || + CAKE_KDA_PACKED_T1_REQUIRES_AUX_VEC4 == 1, + "packed KDA T=1 auxiliary alignment must be boolean"); + +inline void CheckCuda(cudaError_t status, const char* operation) { + TVM_FFI_ICHECK(status == cudaSuccess) << operation << " failed: " << cudaGetErrorString(status); +} + +inline void CheckTarget(int32_t device_id) { + int major = 0; + int minor = 0; + CheckCuda(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device_id), + "cudaDeviceGetAttribute(major)"); + CheckCuda(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device_id), + "cudaDeviceGetAttribute(minor)"); + if (kTargetKind == kTargetFamily) { + TVM_FFI_ICHECK(major == 10 && (minor == 0 || minor == 3)) + << "this packed KDA T=1 module requires the SM100 family " + "(compute capability 10.0 or 10.3), got " + << major << "." << minor; + } else { + TVM_FFI_ICHECK(major == 10 && minor == 0) + << "this packed KDA T=1 module requires exact compute capability 10.0, got " << major << "." + << minor; + } +} + +inline std::pair TensorByteRange(const TensorView& tensor, const char* name) { + const DLDataType dtype = tensor.dtype(); + const uint64_t bits = static_cast(dtype.bits) * dtype.lanes; + TVM_FFI_ICHECK(bits > 0 && bits % 8 == 0) << name << " has a non-byte dtype"; + uint64_t last_element = 0; + for (int32_t i = 0; i < tensor.ndim(); ++i) { + TVM_FFI_ICHECK(tensor.size(i) >= 0 && tensor.stride(i) >= 0) + << name << " must not have negative shapes or strides"; + if (tensor.size(i) > 0) { + const uint64_t extent = static_cast(tensor.size(i) - 1); + const uint64_t stride = static_cast(tensor.stride(i)); + TVM_FFI_ICHECK(stride == 0 || extent <= std::numeric_limits::max() / stride) + << name << " byte range overflows uint64"; + const uint64_t contribution = extent * stride; + TVM_FFI_ICHECK(last_element <= std::numeric_limits::max() - contribution) + << name << " byte range overflows uint64"; + last_element += contribution; + } + } + const uint64_t elements = tensor.numel() == 0 ? 0 : last_element + 1; + TVM_FFI_ICHECK(elements <= std::numeric_limits::max() / (bits / 8)) + << name << " byte range overflows uint64"; + const uint64_t bytes = elements * (bits / 8); + const uintptr_t begin = reinterpret_cast(tensor.data_ptr()); + TVM_FFI_ICHECK(bytes <= std::numeric_limits::max() - begin) + << name << " byte range overflows uintptr_t"; + return {begin, begin + static_cast(bytes)}; +} + +inline void CheckNoOverlap(const TensorView& lhs, const char* lhs_name, const TensorView& rhs, + const char* rhs_name) { + const auto lhs_range = TensorByteRange(lhs, lhs_name); + const auto rhs_range = TensorByteRange(rhs, rhs_name); + TVM_FFI_ICHECK(lhs_range.first >= rhs_range.second || rhs_range.first >= lhs_range.second) + << lhs_name << " must not overlap " << rhs_name + << ": the frozen kernel uses __restrict__ pointers"; +} + +void Run(TensorView mixed_qkv, TensorView raw_gate, TensorView raw_beta, TensorView A_log, + TensorView dt_bias, TensorView state, TensorView state_indices, TensorView out, + int64_t cuda_stream) { + TVM_FFI_ICHECK(cuda_stream >= 0) << "cuda_stream must be a non-negative stream handle"; + CHECK_CUDA(mixed_qkv); + const int32_t device_id = mixed_qkv.device().device_id; + ffi::CUDADeviceGuard device_guard(device_id); + CheckTarget(device_id); + + CHECK_CUDA(raw_gate); + CHECK_CUDA(raw_beta); + CHECK_CUDA(A_log); + CHECK_CUDA(dt_bias); + CHECK_CUDA(state); + CHECK_CUDA(state_indices); + CHECK_CUDA(out); + CHECK_DEVICE(mixed_qkv, raw_gate); + CHECK_DEVICE(mixed_qkv, raw_beta); + CHECK_DEVICE(mixed_qkv, A_log); + CHECK_DEVICE(mixed_qkv, dt_bias); + CHECK_DEVICE(mixed_qkv, state); + CHECK_DEVICE(mixed_qkv, state_indices); + CHECK_DEVICE(mixed_qkv, out); + + CHECK_INPUT_TYPE(mixed_qkv, dl_bfloat16); + CHECK_INPUT_TYPE(raw_gate, dl_bfloat16); + CHECK_INPUT_TYPE(raw_beta, dl_bfloat16); + CHECK_INPUT_TYPE(A_log, dl_float32); + CHECK_INPUT_TYPE(dt_bias, dl_float32); + CHECK_INPUT_TYPE(state, dl_bfloat16); + CHECK_INPUT_TYPE(state_indices, dl_int32); + CHECK_INPUT_TYPE(out, dl_bfloat16); + + TVM_FFI_ICHECK(mixed_qkv.ndim() == 2 && mixed_qkv.size(0) > 0 && mixed_qkv.size(1) == kMixedWidth) + << "mixed_qkv must have shape [B, " << kMixedWidth << "]"; + const int64_t batch = mixed_qkv.size(0); + TVM_FFI_ICHECK(batch <= 65535) << "batch exceeds the CUDA grid.y limit"; + CHECK_LAST_DIM_CONTIGUOUS(mixed_qkv); + TVM_FFI_ICHECK(mixed_qkv.stride(0) >= kMixedWidth) + << "mixed_qkv must have a compact last dimension and disjoint rows"; + + TVM_FFI_ICHECK(raw_gate.ndim() == 2 && raw_gate.size(0) == batch && + raw_gate.size(1) == kGateWidth && raw_gate.stride(0) >= kGateWidth) + << "raw_gate must have shape [B, " << kGateWidth + << "] with a compact last dimension and disjoint rows"; + CHECK_LAST_DIM_CONTIGUOUS(raw_gate); + TVM_FFI_ICHECK(raw_beta.ndim() == 2 && raw_beta.size(0) == batch && raw_beta.size(1) == kHeads && + raw_beta.stride(0) >= kHeads) + << "raw_beta must have shape [B, " << kHeads + << "] with a compact last dimension and disjoint rows"; + CHECK_LAST_DIM_CONTIGUOUS(raw_beta); + + TVM_FFI_ICHECK(A_log.ndim() == 1 && A_log.numel() == kHeads) + << "A_log must be one-dimensional with " << kHeads << " elements"; + TVM_FFI_ICHECK(dt_bias.ndim() == 1 && dt_bias.numel() == kGateWidth) + << "dt_bias must be one-dimensional with " << kGateWidth << " elements"; + CHECK_CONTIGUOUS(A_log); + CHECK_CONTIGUOUS(dt_bias); + + TVM_FFI_ICHECK(state.ndim() == 4 && state.size(0) > 0 && state.size(1) == kHeads && + state.size(2) == kHeadDim && state.size(3) == kHeadDim) + << "state must have shape [N, " << kHeads << ", " << kHeadDim << ", " << kHeadDim << "]"; + CHECK_LAST_DIM_CONTIGUOUS(state); + TVM_FFI_ICHECK(state.stride(2) == kHeadDim && state.stride(1) == kHeadDim * kHeadDim && + state.stride(0) >= kHeads * kHeadDim * kHeadDim) + << "state must have compact [H,V,K] blocks and a positive, disjoint outer slot stride"; + TVM_FFI_ICHECK(state.stride(0) > 0 && state.stride(0) % 8 == 0) + << "optimized state slot stride must be positive and eight-element aligned"; + TVM_FFI_ICHECK(reinterpret_cast(state.data_ptr()) % 16 == 0) + << "optimized state base must be 16-byte aligned"; + + TVM_FFI_ICHECK(state_indices.ndim() == 1 && state_indices.numel() == batch) + << "state_indices must have shape [B]"; + CHECK_CONTIGUOUS(state_indices); + TVM_FFI_ICHECK(out.ndim() == 3 && out.size(0) == batch && out.size(1) == kHeads && + out.size(2) == kHeadDim) + << "output must have shape [B, " << kHeads << ", " << kHeadDim << "]"; + CHECK_CONTIGUOUS(out); + + if constexpr (CAKE_KDA_PACKED_T1_REQUIRES_AUX_VEC4 != 0) { + const uintptr_t mixed_base = reinterpret_cast(mixed_qkv.data_ptr()); + const uintptr_t k_base = mixed_base + kHeads * kHeadDim * sizeof(__nv_bfloat16); + TVM_FFI_ICHECK(mixed_base % 8 == 0 && k_base % 8 == 0 && + reinterpret_cast(raw_gate.data_ptr()) % 8 == 0 && + mixed_qkv.stride(0) % 4 == 0 && raw_gate.stride(0) % 4 == 0 && + reinterpret_cast(dt_bias.data_ptr()) % 16 == 0) + << "this packed KDA variant requires vec4-aligned Q/K/gate rows and dt_bias"; + } + + const std::pair read_tensors[] = { + {&mixed_qkv, "mixed_qkv"}, {&raw_gate, "raw_gate"}, {&raw_beta, "raw_beta"}, + {&A_log, "A_log"}, {&dt_bias, "dt_bias"}, {&state_indices, "state_indices"}, + }; + CheckNoOverlap(state, "state", out, "output"); + for (const auto& named : read_tensors) { + CheckNoOverlap(state, "state", *named.first, named.second); + CheckNoOverlap(out, "output", *named.first, named.second); + } + + auto* mixed = reinterpret_cast<__nv_bfloat16*>(mixed_qkv.data_ptr()); + auto* q = mixed; + auto* k = mixed + kHeads * kHeadDim; + auto* v = mixed + 2 * kHeads * kHeadDim; + const dim3 grid(kHeads * CAKE_KDA_PACKED_T1_VALUE_TILES, static_cast(batch), 1); + const dim3 block(CAKE_KDA_PACKED_T1_THREADS, 1, 1); + const auto stream = reinterpret_cast(cuda_stream); + CAKE_KDA_PACKED_T1_KERNEL<<>>( + q, k, v, reinterpret_cast<__nv_bfloat16*>(raw_gate.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(raw_beta.data_ptr()), + reinterpret_cast(A_log.data_ptr()), reinterpret_cast(dt_bias.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(state.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), + reinterpret_cast(state_indices.data_ptr()), 0.08838834764831845F, mixed_qkv.stride(0), + mixed_qkv.stride(0), mixed_qkv.stride(0), raw_gate.stride(0), raw_beta.stride(0), + state.stride(0)); + CheckCuda(cudaGetLastError(), "frozen packed KDA T=1 launch"); +} + +} // namespace cake_kda_packed_t1 +} // namespace flashinfer + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, flashinfer::cake_kda_packed_t1::Run); diff --git a/csrc/kda/cake_kda_packed_t1_cpasync_tile128_ilp2.cu b/csrc/kda/cake_kda_packed_t1_cpasync_tile128_ilp2.cu new file mode 100644 index 00000000000..c73091a41da --- /dev/null +++ b/csrc/kda/cake_kda_packed_t1_cpasync_tile128_ilp2.cu @@ -0,0 +1,468 @@ +/* + * 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. + */ + +// clang-format off +// Generated by CAKE from the Kimi K3 serving-native packed T=1 schedule. +// BEGIN FROZEN GENERATED BODY +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +template +struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; + +typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_STATE_SMEM_OFF 0 +#define SMEM_STATE_SMEM_STAGE_BYTES 4096 +#define SMEM_STATE_SMEM_STRIDE 4096 +#define SMEM_V_SMEM_OFF 20224 +#define SMEM_V_SMEM_STAGE_BYTES 256 +#define SMEM_V_SMEM_STRIDE 256 +#define SMEM_TOTAL 20480 +#define THREADS 128 + +#include + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flashinfer_packed_kda_t1_cpasync_tile128(__nv_bfloat16* __restrict__ q, __nv_bfloat16* __restrict__ k, __nv_bfloat16* __restrict__ v, __nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ beta, float* __restrict__ A_log, float* __restrict__ dt_bias, __nv_bfloat16* __restrict__ state, __nv_bfloat16* __restrict__ out, int* __restrict__ state_indices, float scale, long long q_stride_token, long long k_stride_token, long long v_stride_token, long long g_stride_token, long long beta_stride_token, long long state_stride_slot) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* state_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int state_smem_addr = smem + 0; + __nv_bfloat16* v_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 20224); + const int v_smem_addr = smem + 20224; + + // === Task calls (dependency order) === + int tid_0 = tid; + int lane_1 = lane; + int k_lane = tid_0 % 16; + int group = tid_0 / 16; + int copy_row = tid_0 / 8; + int copy_col = tid_0 % 8 * 16; + int work = blockIdx.x; + int value_tile = 0; + int hv = work; + int n = blockIdx.y; + int raw_slot = state_indices[n]; + bool active = raw_slot >= 0; + int output_head_base = (n * 12 + hv) * 128; + int tile_row_base = value_tile * 128; + if (active) { + long long state_head_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)tile_row_base * 128; + unsigned int v_registers[8]; + { + } + #pragma unroll + for (int prefetch_chunk = 0; prefetch_chunk < 4; prefetch_chunk++) { + { + #pragma unroll + for (int private_row = 0; private_row < 2; private_row++) { + int copy_state_row = group * 2 + private_row; + int copy_elem = (prefetch_chunk * 16 + copy_state_row) * 128 + k_lane * 8; + int copy_dst = state_smem_addr + (unsigned int)(prefetch_chunk * 4096) + (unsigned int)((copy_state_row * 128 + k_lane * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst), "l"(state + (state_head_base + (long long)copy_elem))); + } + } + asm volatile("cp.async.commit_group;"); + } + int elem_start = lane_1 * 4; + long long q_base = (long long)n * q_stride_token + (long long)hv * 128 + (long long)elem_start; + long long k_base = (long long)n * k_stride_token + (long long)hv * 128 + (long long)elem_start; + long long gate_base = (long long)n * g_stride_token + (long long)hv * 128 + (long long)elem_start; + float q_src[4]; + float k_src[4]; + float gate_src[4]; + float dt_bias_src[4]; + float q_reg[8]; + float k_reg[8]; + float decay_reg[8]; + unsigned int packed_state_row[4]; + unsigned int lookahead_state[16]; + float state_row[8]; + float state_rows[16]; + unsigned int paired_state_bits[8]; + float paired_state_rows[16]; + { + { + uint2 _vld_0 = *reinterpret_cast(q + q_base); + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&q_src[0 + _pair * 2])[0]), "=f"((&q_src[0 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + { + uint2 _vld_1 = *reinterpret_cast(k + k_base); + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&k_src[0 + _pair * 2])[0]), "=f"((&k_src[0 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + { + uint2 _vld_2 = *reinterpret_cast(g + gate_base); + uint32_t* _vpairs_2 = reinterpret_cast(&_vld_2); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&gate_src[0 + _pair * 2])[0]), "=f"((&gate_src[0 + _pair * 2])[1]) + : "r"(_vpairs_2[_pair])); + } + } + { + float4 _v4 = *reinterpret_cast(dt_bias + hv * 128 + elem_start); + dt_bias_src[0 + 0] = _v4.x; + dt_bias_src[0 + 1] = _v4.y; + dt_bias_src[0 + 2] = _v4.z; + dt_bias_src[0 + 3] = _v4.w; + } + } + float q_sum_sq = q_src[0] * q_src[0] + q_src[1] * q_src[1] + (q_src[2] * q_src[2] + q_src[3] * q_src[3]); + float k_sum_sq = k_src[0] * k_src[0] + k_src[1] * k_src[1] + (k_src[2] * k_src[2] + k_src[3] * k_src[3]); + float _warp_reduce_0 = q_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset); + q_sum_sq = _warp_reduce_0; + float _warp_reduce_1 = k_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + k_sum_sq = _warp_reduce_1; + float _rsqrt_0 = rsqrtf(q_sum_sq + 1e-06f); + float q_scale = _rsqrt_0 * scale; + float _rsqrt_1 = rsqrtf(k_sum_sq + 1e-06f); + float k_scale = _rsqrt_1; + float _expf_0 = __expf(A_log[hv]); + float gate_a = _expf_0; + #pragma unroll + for (int source_value = 0; source_value < 4; source_value++) { + float biased_gate; + { + biased_gate = gate_src[source_value] + dt_bias_src[source_value]; + } + float _expf_1 = __expf((-gate_a) * biased_gate); + float gate_sigmoid = 1.0f / (1.0f + _expf_1); + float _expf_2 = __expf(-5.0f * gate_sigmoid); + gate_src[source_value] = _expf_2; + } + #pragma unroll + for (int i = 0; i < 8; i++) { + int source_lane = 2 * k_lane + i / 4; + float _shfl_0 = __shfl_sync(0xFFFFFFFF, q_src[i % 4], source_lane); + q_reg[i] = _shfl_0 * q_scale; + float _shfl_1 = __shfl_sync(0xFFFFFFFF, k_src[i % 4], source_lane); + k_reg[i] = _shfl_1 * k_scale; + float _shfl_2 = __shfl_sync(0xFFFFFFFF, gate_src[i % 4], source_lane); + decay_reg[i] = _shfl_2; + } + float k_dot_q = 0.0f; + #pragma unroll + for (int i_1 = 0; i_1 < 8; i_1++) { + k_dot_q += k_reg[i_1] * q_reg[i_1]; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 8); + k_dot_q += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 4); + k_dot_q += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 2); + k_dot_q += _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 1); + k_dot_q += _shfl_xor_3; + float beta_value = (float)beta[(long long)n * beta_stride_token + (long long)hv]; + if (beta_value >= 0.0f) { + float _expf_3 = __expf(-beta_value); + beta_value = 1.0f / (1.0f + _expf_3); + } else { + float _expf_4 = __expf(beta_value); + float beta_exp = _expf_4; + beta_value = beta_exp / (1.0f + beta_exp); + } + #pragma unroll + for (int chunk = 0; chunk < 8; chunk++) { + { + { + if (chunk < 5) { + asm volatile("cp.async.wait_group 3;"); + } else if (chunk == 5) { + asm volatile("cp.async.wait_group 2;"); + } else { + if (chunk == 6) { + asm volatile("cp.async.wait_group 1;"); + } else { + asm volatile("cp.async.wait_group 0;"); + } + } + } + } + { + if (chunk + 5 - 1 < 8) { + #pragma unroll + for (int private_row_1 = 0; private_row_1 < 2; private_row_1++) { + int copy_state_row_1 = group * 2 + private_row_1; + int copy_elem_1 = ((chunk + 5 - 1) * 16 + copy_state_row_1) * 128 + k_lane * 8; + int copy_dst_1 = state_smem_addr + (unsigned int)((chunk + 5 - 1) % 5 * 4096) + (unsigned int)((copy_state_row_1 * 128 + k_lane * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst_1), "l"(state + (state_head_base + (long long)copy_elem_1))); + } + asm volatile("cp.async.commit_group;"); + } + } + { + { + #pragma unroll + for (int row_local = 0; row_local < 2; row_local++) { + unsigned int packed_state[4]; + int smem_row = group * 2 + row_local; + int state_smem_addr_0 = 0; + { + state_smem_addr_0 = state_smem_addr + (unsigned int)(chunk % 5 * 4096) + (unsigned int)((smem_row * 128 + k_lane * 8) * 2); + } + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_state[0])), "=r"(*reinterpret_cast(&packed_state[(0) + 1])), "=r"(*reinterpret_cast(&packed_state[(0) + 2])), "=r"(*reinterpret_cast(&packed_state[(0) + 3])) + : "r"(state_smem_addr_0)); + float packed_state_f32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_state_f32[_pair * 2])[0]), "=f"((&packed_state_f32[_pair * 2])[1]) + : "r"(packed_state[_pair])); + } + #pragma unroll + for (int i_2 = 0; i_2 < 8; i_2++) { + state_rows[row_local * 8 + i_2] = packed_state_f32[i_2]; + } + } + } + #pragma unroll + for (int row_local_1 = 0; row_local_1 < 2; row_local_1++) { + int value_row = tile_row_base + chunk * 16 + group * 2 + row_local_1; + float prediction; + float base; + { + prediction = 0.0f; + base = 0.0f; + #pragma unroll + for (int i_3 = 0; i_3 < 8; i_3++) { + float decayed = state_rows[row_local_1 * 8 + i_3] * decay_reg[i_3]; + prediction += decayed * k_reg[i_3]; + base += decayed * q_reg[i_3]; + } + } + float _shfl_xor_20 = __shfl_xor_sync(0xFFFFFFFF, prediction, 8); + prediction += _shfl_xor_20; + float _shfl_xor_21 = __shfl_xor_sync(0xFFFFFFFF, prediction, 4); + prediction += _shfl_xor_21; + float _shfl_xor_22 = __shfl_xor_sync(0xFFFFFFFF, prediction, 2); + prediction += _shfl_xor_22; + float _shfl_xor_23 = __shfl_xor_sync(0xFFFFFFFF, prediction, 1); + prediction += _shfl_xor_23; + float _shfl_xor_24 = __shfl_xor_sync(0xFFFFFFFF, base, 8); + base += _shfl_xor_24; + float _shfl_xor_25 = __shfl_xor_sync(0xFFFFFFFF, base, 4); + base += _shfl_xor_25; + float _shfl_xor_26 = __shfl_xor_sync(0xFFFFFFFF, base, 2); + base += _shfl_xor_26; + float _shfl_xor_27 = __shfl_xor_sync(0xFFFFFFFF, base, 1); + base += _shfl_xor_27; + float v_value = 0.0f; + { + { + if (k_lane == 0) { + v_value = (float)v[(long long)n * v_stride_token + (long long)hv * 128 + (long long)value_row]; + } + } + } + { + float _shfl_5 = __shfl_sync(0xFFFFFFFF, v_value, group % 2 * 16); + v_value = _shfl_5; + } + float delta = (v_value - prediction) * beta_value; + { + #pragma unroll + for (int i_4 = 0; i_4 < 8; i_4++) { + state_rows[row_local_1 * 8 + i_4] = state_rows[row_local_1 * 8 + i_4] * decay_reg[i_4] + delta * k_reg[i_4]; + } + } + long long state_row_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)value_row * 128 + (long long)k_lane * 8; + { + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 0], state_rows[row_local_1 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 2], state_rows[row_local_1 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 4], state_rows[row_local_1 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 6], state_rows[row_local_1 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(state))[state_row_base + 0]) = *reinterpret_cast(&_pk[0]); + } + } + { + if (k_lane == 0) { + out[output_head_base + value_row] = base + delta * k_dot_q; + } + } + } + } + } + } else if (k_lane == 0) { + #pragma unroll + for (int chunk_1 = 0; chunk_1 < 8; chunk_1++) { + #pragma unroll + for (int row_local_2 = 0; row_local_2 < 2; row_local_2++) { + int value_row_1 = tile_row_base + chunk_1 * 16 + group * 2 + row_local_2; + out[output_head_base + value_row_1] = 0.0f; + } + } + } +} + +} // extern "C" + +// END FROZEN GENERATED BODY +// clang-format on diff --git a/csrc/kda/cake_kda_packed_t1_cpasync_tile128_ilp4.cu b/csrc/kda/cake_kda_packed_t1_cpasync_tile128_ilp4.cu new file mode 100644 index 00000000000..6686562f314 --- /dev/null +++ b/csrc/kda/cake_kda_packed_t1_cpasync_tile128_ilp4.cu @@ -0,0 +1,426 @@ +/* + * 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. + */ + +// clang-format off +// Generated by CAKE from the Kimi K3 serving-native packed T=1 schedule. +// BEGIN FROZEN GENERATED BODY +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +template +struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; + +typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_STATE_SMEM_OFF 0 +#define SMEM_STATE_SMEM_STAGE_BYTES 8192 +#define SMEM_STATE_SMEM_STRIDE 8192 +#define SMEM_V_SMEM_OFF 24320 +#define SMEM_V_SMEM_STAGE_BYTES 256 +#define SMEM_V_SMEM_STRIDE 256 +#define SMEM_TOTAL 24576 +#define THREADS 128 + +#include + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flashinfer_packed_kda_t1_cpasync_tile128_ilp4(__nv_bfloat16* __restrict__ q, __nv_bfloat16* __restrict__ k, __nv_bfloat16* __restrict__ v, __nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ beta, float* __restrict__ A_log, float* __restrict__ dt_bias, __nv_bfloat16* __restrict__ state, __nv_bfloat16* __restrict__ out, int* __restrict__ state_indices, float scale, long long q_stride_token, long long k_stride_token, long long v_stride_token, long long g_stride_token, long long beta_stride_token, long long state_stride_slot) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* state_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int state_smem_addr = smem + 0; + __nv_bfloat16* v_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 24320); + const int v_smem_addr = smem + 24320; + + // === Task calls (dependency order) === + int tid_0 = tid; + int lane_1 = lane; + int k_lane = tid_0 % 16; + int group = tid_0 / 16; + int copy_row = tid_0 / 8; + int copy_col = tid_0 % 8 * 16; + int work = blockIdx.x; + int value_tile = 0; + int hv = work; + int n = blockIdx.y; + int raw_slot = state_indices[n]; + bool active = raw_slot >= 0; + int output_head_base = (n * 12 + hv) * 128; + int tile_row_base = value_tile * 128; + if (active) { + long long state_head_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)tile_row_base * 128; + unsigned int v_registers[4]; + { + } + #pragma unroll + for (int prefetch_chunk = 0; prefetch_chunk < 2; prefetch_chunk++) { + { + #pragma unroll + for (int private_row = 0; private_row < 4; private_row++) { + int copy_state_row = group * 4 + private_row; + int copy_elem = (prefetch_chunk * 32 + copy_state_row) * 128 + k_lane * 8; + int copy_dst = state_smem_addr + (unsigned int)(prefetch_chunk * 8192) + (unsigned int)((copy_state_row * 128 + k_lane * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst), "l"(state + (state_head_base + (long long)copy_elem))); + } + } + asm volatile("cp.async.commit_group;"); + } + int elem_start = lane_1 * 4; + long long q_base = (long long)n * q_stride_token + (long long)hv * 128 + (long long)elem_start; + long long k_base = (long long)n * k_stride_token + (long long)hv * 128 + (long long)elem_start; + long long gate_base = (long long)n * g_stride_token + (long long)hv * 128 + (long long)elem_start; + float q_src[4]; + float k_src[4]; + float gate_src[4]; + float dt_bias_src[4]; + float q_reg[8]; + float k_reg[8]; + float decay_reg[8]; + unsigned int packed_state_row[4]; + unsigned int lookahead_state[32]; + float state_row[8]; + float state_rows[32]; + unsigned int paired_state_bits[16]; + float paired_state_rows[32]; + { + #pragma unroll + for (int source_value = 0; source_value < 4; source_value++) { + q_src[source_value] = (float)q[q_base + (long long)source_value]; + k_src[source_value] = (float)k[k_base + (long long)source_value]; + gate_src[source_value] = (float)g[gate_base + (long long)source_value]; + } + } + float q_sum_sq = q_src[0] * q_src[0] + q_src[1] * q_src[1] + (q_src[2] * q_src[2] + q_src[3] * q_src[3]); + float k_sum_sq = k_src[0] * k_src[0] + k_src[1] * k_src[1] + (k_src[2] * k_src[2] + k_src[3] * k_src[3]); + float _warp_reduce_0 = q_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset); + q_sum_sq = _warp_reduce_0; + float _warp_reduce_1 = k_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + k_sum_sq = _warp_reduce_1; + float _rsqrt_0 = rsqrtf(q_sum_sq + 1e-06f); + float q_scale = _rsqrt_0 * scale; + float _rsqrt_1 = rsqrtf(k_sum_sq + 1e-06f); + float k_scale = _rsqrt_1; + float _expf_0 = __expf(A_log[hv]); + float gate_a = _expf_0; + #pragma unroll + for (int source_value_1 = 0; source_value_1 < 4; source_value_1++) { + float biased_gate; + { + int k_idx = elem_start + source_value_1; + biased_gate = gate_src[source_value_1] + dt_bias[hv * 128 + k_idx]; + } + float _expf_1 = __expf((-gate_a) * biased_gate); + float gate_sigmoid = 1.0f / (1.0f + _expf_1); + float _expf_2 = __expf(-5.0f * gate_sigmoid); + gate_src[source_value_1] = _expf_2; + } + #pragma unroll + for (int i = 0; i < 8; i++) { + int source_lane = 2 * k_lane + i / 4; + float _shfl_0 = __shfl_sync(0xFFFFFFFF, q_src[i % 4], source_lane); + q_reg[i] = _shfl_0 * q_scale; + float _shfl_1 = __shfl_sync(0xFFFFFFFF, k_src[i % 4], source_lane); + k_reg[i] = _shfl_1 * k_scale; + float _shfl_2 = __shfl_sync(0xFFFFFFFF, gate_src[i % 4], source_lane); + decay_reg[i] = _shfl_2; + } + float k_dot_q = 0.0f; + #pragma unroll + for (int i_1 = 0; i_1 < 8; i_1++) { + k_dot_q += k_reg[i_1] * q_reg[i_1]; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 8); + k_dot_q += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 4); + k_dot_q += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 2); + k_dot_q += _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 1); + k_dot_q += _shfl_xor_3; + float beta_value = (float)beta[(long long)n * beta_stride_token + (long long)hv]; + if (beta_value >= 0.0f) { + float _expf_3 = __expf(-beta_value); + beta_value = 1.0f / (1.0f + _expf_3); + } else { + float _expf_4 = __expf(beta_value); + float beta_exp = _expf_4; + beta_value = beta_exp / (1.0f + beta_exp); + } + #pragma unroll + for (int chunk = 0; chunk < 4; chunk++) { + { + { + if (chunk < 3) { + asm volatile("cp.async.wait_group 1;"); + } else if (chunk == 1) { + asm volatile("cp.async.wait_group 2;"); + } else { + if (chunk == 2) { + asm volatile("cp.async.wait_group 1;"); + } else { + asm volatile("cp.async.wait_group 0;"); + } + } + } + } + { + if (chunk + 3 - 1 < 4) { + #pragma unroll + for (int private_row_1 = 0; private_row_1 < 4; private_row_1++) { + int copy_state_row_1 = group * 4 + private_row_1; + int copy_elem_1 = ((chunk + 3 - 1) * 32 + copy_state_row_1) * 128 + k_lane * 8; + int copy_dst_1 = state_smem_addr + (unsigned int)((chunk + 3 - 1) % 3 * 8192) + (unsigned int)((copy_state_row_1 * 128 + k_lane * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst_1), "l"(state + (state_head_base + (long long)copy_elem_1))); + } + asm volatile("cp.async.commit_group;"); + } + } + { + { + #pragma unroll + for (int row_local = 0; row_local < 4; row_local++) { + unsigned int packed_state[4]; + int smem_row = group * 4 + row_local; + int state_smem_addr_0 = 0; + { + state_smem_addr_0 = state_smem_addr + (unsigned int)(chunk % 3 * 8192) + (unsigned int)((smem_row * 128 + k_lane * 8) * 2); + } + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_state[0])), "=r"(*reinterpret_cast(&packed_state[(0) + 1])), "=r"(*reinterpret_cast(&packed_state[(0) + 2])), "=r"(*reinterpret_cast(&packed_state[(0) + 3])) + : "r"(state_smem_addr_0)); + float packed_state_f32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_state_f32[_pair * 2])[0]), "=f"((&packed_state_f32[_pair * 2])[1]) + : "r"(packed_state[_pair])); + } + #pragma unroll + for (int i_2 = 0; i_2 < 8; i_2++) { + state_rows[row_local * 8 + i_2] = packed_state_f32[i_2]; + } + } + } + #pragma unroll + for (int row_local_1 = 0; row_local_1 < 4; row_local_1++) { + int value_row = tile_row_base + chunk * 32 + group * 4 + row_local_1; + float prediction; + float base; + { + prediction = 0.0f; + base = 0.0f; + #pragma unroll + for (int i_3 = 0; i_3 < 8; i_3++) { + float decayed = state_rows[row_local_1 * 8 + i_3] * decay_reg[i_3]; + prediction += decayed * k_reg[i_3]; + base += decayed * q_reg[i_3]; + } + } + float _shfl_xor_20 = __shfl_xor_sync(0xFFFFFFFF, prediction, 8); + prediction += _shfl_xor_20; + float _shfl_xor_21 = __shfl_xor_sync(0xFFFFFFFF, prediction, 4); + prediction += _shfl_xor_21; + float _shfl_xor_22 = __shfl_xor_sync(0xFFFFFFFF, prediction, 2); + prediction += _shfl_xor_22; + float _shfl_xor_23 = __shfl_xor_sync(0xFFFFFFFF, prediction, 1); + prediction += _shfl_xor_23; + float _shfl_xor_24 = __shfl_xor_sync(0xFFFFFFFF, base, 8); + base += _shfl_xor_24; + float _shfl_xor_25 = __shfl_xor_sync(0xFFFFFFFF, base, 4); + base += _shfl_xor_25; + float _shfl_xor_26 = __shfl_xor_sync(0xFFFFFFFF, base, 2); + base += _shfl_xor_26; + float _shfl_xor_27 = __shfl_xor_sync(0xFFFFFFFF, base, 1); + base += _shfl_xor_27; + float v_value = 0.0f; + { + { + if (k_lane == 0) { + v_value = (float)v[(long long)n * v_stride_token + (long long)hv * 128 + (long long)value_row]; + } + } + } + { + float _shfl_5 = __shfl_sync(0xFFFFFFFF, v_value, group % 2 * 16); + v_value = _shfl_5; + } + float delta = (v_value - prediction) * beta_value; + { + #pragma unroll + for (int i_4 = 0; i_4 < 8; i_4++) { + state_rows[row_local_1 * 8 + i_4] = state_rows[row_local_1 * 8 + i_4] * decay_reg[i_4] + delta * k_reg[i_4]; + } + } + long long state_row_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)value_row * 128 + (long long)k_lane * 8; + { + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 0], state_rows[row_local_1 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 2], state_rows[row_local_1 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 4], state_rows[row_local_1 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 6], state_rows[row_local_1 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(state))[state_row_base + 0]) = *reinterpret_cast(&_pk[0]); + } + } + { + if (k_lane == 0) { + out[output_head_base + value_row] = base + delta * k_dot_q; + } + } + } + } + } + } else if (k_lane == 0) { + #pragma unroll + for (int chunk_1 = 0; chunk_1 < 4; chunk_1++) { + #pragma unroll + for (int row_local_2 = 0; row_local_2 < 4; row_local_2++) { + int value_row_1 = tile_row_base + chunk_1 * 32 + group * 4 + row_local_2; + out[output_head_base + value_row_1] = 0.0f; + } + } + } +} + +} // extern "C" + +// END FROZEN GENERATED BODY +// clang-format on diff --git a/csrc/kda/cake_kda_packed_t1_cpasync_tile128_packed_state_v_private_prefetch.cu b/csrc/kda/cake_kda_packed_t1_cpasync_tile128_packed_state_v_private_prefetch.cu new file mode 100644 index 00000000000..046f8c3630f --- /dev/null +++ b/csrc/kda/cake_kda_packed_t1_cpasync_tile128_packed_state_v_private_prefetch.cu @@ -0,0 +1,513 @@ +/* + * 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. + */ + +// clang-format off +// Generated by CAKE from the Kimi K3 serving-native packed T=1 schedule. +// BEGIN FROZEN GENERATED BODY +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +template +struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; + +typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_STATE_SMEM_OFF 0 +#define SMEM_STATE_SMEM_STAGE_BYTES 4096 +#define SMEM_STATE_SMEM_STRIDE 4096 +#define SMEM_V_SMEM_OFF 20480 +#define SMEM_V_SMEM_STAGE_BYTES 256 +#define SMEM_V_SMEM_STRIDE 256 +#define SMEM_TOTAL 20736 +#define THREADS 128 + +#include + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flashinfer_packed_kda_t1_cpasync_tile128_packed_state_v_private_prefetch(__nv_bfloat16* __restrict__ q, __nv_bfloat16* __restrict__ k, __nv_bfloat16* __restrict__ v, __nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ beta, float* __restrict__ A_log, float* __restrict__ dt_bias, __nv_bfloat16* __restrict__ state, __nv_bfloat16* __restrict__ out, int* __restrict__ state_indices, float scale, long long q_stride_token, long long k_stride_token, long long v_stride_token, long long g_stride_token, long long beta_stride_token, long long state_stride_slot) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* state_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int state_smem_addr = smem + 0; + __nv_bfloat16* v_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 20480); + const int v_smem_addr = smem + 20480; + + // === Task calls (dependency order) === + int tid_0 = tid; + int lane_1 = lane; + int k_lane = tid_0 % 16; + int group = tid_0 / 16; + int copy_row = tid_0 / 8; + int copy_col = tid_0 % 8 * 16; + int work = blockIdx.x; + int value_tile = 0; + int hv = work; + int n = blockIdx.y; + int raw_slot = state_indices[n]; + bool active = raw_slot >= 0; + int output_head_base = (n * 12 + hv) * 128; + int tile_row_base = value_tile * 128; + if (active) { + long long state_head_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)tile_row_base * 128; + unsigned int v_registers[8]; + { + { + if (k_lane == 0) { + #pragma unroll + for (int v_chunk = 0; v_chunk < 8; v_chunk++) { + int v_row = tile_row_base + v_chunk * 16 + group * 2; + asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 4;" + :: "r"(v_smem_addr + (unsigned int)(v_row * 2)), "l"(v + ((long long)n * v_stride_token + (long long)hv * 128 + (long long)v_row))); + } + } + asm volatile("cp.async.commit_group;"); + } + } + #pragma unroll + for (int prefetch_chunk = 0; prefetch_chunk < 4; prefetch_chunk++) { + { + #pragma unroll + for (int private_row = 0; private_row < 2; private_row++) { + int copy_state_row = group * 2 + private_row; + int copy_elem = (prefetch_chunk * 16 + copy_state_row) * 128 + k_lane * 8; + int copy_dst = state_smem_addr + (unsigned int)(prefetch_chunk * 4096) + (unsigned int)((copy_state_row * 128 + k_lane * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst), "l"(state + (state_head_base + (long long)copy_elem))); + } + } + asm volatile("cp.async.commit_group;"); + } + int elem_start = lane_1 * 4; + long long q_base = (long long)n * q_stride_token + (long long)hv * 128 + (long long)elem_start; + long long k_base = (long long)n * k_stride_token + (long long)hv * 128 + (long long)elem_start; + long long gate_base = (long long)n * g_stride_token + (long long)hv * 128 + (long long)elem_start; + float q_src[4]; + float k_src[4]; + float gate_src[4]; + float dt_bias_src[4]; + float q_reg[8]; + float k_reg[8]; + float decay_reg[8]; + unsigned int packed_state_row[4]; + unsigned int lookahead_state[16]; + float state_row[8]; + float state_rows[16]; + unsigned int paired_state_bits[8]; + float paired_state_rows[16]; + { + { + uint2 _vld_0 = *reinterpret_cast(q + q_base); + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&q_src[0 + _pair * 2])[0]), "=f"((&q_src[0 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + { + uint2 _vld_1 = *reinterpret_cast(k + k_base); + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&k_src[0 + _pair * 2])[0]), "=f"((&k_src[0 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + { + uint2 _vld_2 = *reinterpret_cast(g + gate_base); + uint32_t* _vpairs_2 = reinterpret_cast(&_vld_2); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&gate_src[0 + _pair * 2])[0]), "=f"((&gate_src[0 + _pair * 2])[1]) + : "r"(_vpairs_2[_pair])); + } + } + { + float4 _v4 = *reinterpret_cast(dt_bias + hv * 128 + elem_start); + dt_bias_src[0 + 0] = _v4.x; + dt_bias_src[0 + 1] = _v4.y; + dt_bias_src[0 + 2] = _v4.z; + dt_bias_src[0 + 3] = _v4.w; + } + } + float q_sum_sq = q_src[0] * q_src[0] + q_src[1] * q_src[1] + (q_src[2] * q_src[2] + q_src[3] * q_src[3]); + float k_sum_sq = k_src[0] * k_src[0] + k_src[1] * k_src[1] + (k_src[2] * k_src[2] + k_src[3] * k_src[3]); + float _warp_reduce_0 = q_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset); + q_sum_sq = _warp_reduce_0; + float _warp_reduce_1 = k_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + k_sum_sq = _warp_reduce_1; + float _rsqrt_0 = rsqrtf(q_sum_sq + 1e-06f); + float q_scale = _rsqrt_0 * scale; + float _rsqrt_1 = rsqrtf(k_sum_sq + 1e-06f); + float k_scale = _rsqrt_1; + float _expf_0 = __expf(A_log[hv]); + float gate_a = _expf_0; + #pragma unroll + for (int source_value = 0; source_value < 4; source_value++) { + float biased_gate; + { + biased_gate = gate_src[source_value] + dt_bias_src[source_value]; + } + float _expf_1 = __expf((-gate_a) * biased_gate); + float gate_sigmoid = 1.0f / (1.0f + _expf_1); + float _expf_2 = __expf(-5.0f * gate_sigmoid); + gate_src[source_value] = _expf_2; + } + #pragma unroll + for (int i = 0; i < 8; i++) { + int source_lane = 2 * k_lane + i / 4; + float _shfl_0 = __shfl_sync(0xFFFFFFFF, q_src[i % 4], source_lane); + q_reg[i] = _shfl_0 * q_scale; + float _shfl_1 = __shfl_sync(0xFFFFFFFF, k_src[i % 4], source_lane); + k_reg[i] = _shfl_1 * k_scale; + float _shfl_2 = __shfl_sync(0xFFFFFFFF, gate_src[i % 4], source_lane); + decay_reg[i] = _shfl_2; + } + float k_dot_q = 0.0f; + #pragma unroll + for (int i_1 = 0; i_1 < 8; i_1++) { + k_dot_q += k_reg[i_1] * q_reg[i_1]; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 8); + k_dot_q += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 4); + k_dot_q += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 2); + k_dot_q += _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 1); + k_dot_q += _shfl_xor_3; + float beta_value = (float)beta[(long long)n * beta_stride_token + (long long)hv]; + if (beta_value >= 0.0f) { + float _expf_3 = __expf(-beta_value); + beta_value = 1.0f / (1.0f + _expf_3); + } else { + float _expf_4 = __expf(beta_value); + float beta_exp = _expf_4; + beta_value = beta_exp / (1.0f + beta_exp); + } + { + asm volatile("cp.async.wait_group 4;"); + } + #pragma unroll + for (int chunk = 0; chunk < 8; chunk++) { + { + { + if (chunk < 5) { + asm volatile("cp.async.wait_group 3;"); + } else if (chunk == 5) { + asm volatile("cp.async.wait_group 2;"); + } else { + if (chunk == 6) { + asm volatile("cp.async.wait_group 1;"); + } else { + asm volatile("cp.async.wait_group 0;"); + } + } + } + } + { + if (chunk + 5 - 1 < 8) { + #pragma unroll + for (int private_row_1 = 0; private_row_1 < 2; private_row_1++) { + int copy_state_row_1 = group * 2 + private_row_1; + int copy_elem_1 = ((chunk + 5 - 1) * 16 + copy_state_row_1) * 128 + k_lane * 8; + int copy_dst_1 = state_smem_addr + (unsigned int)((chunk + 5 - 1) % 5 * 4096) + (unsigned int)((copy_state_row_1 * 128 + k_lane * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst_1), "l"(state + (state_head_base + (long long)copy_elem_1))); + } + asm volatile("cp.async.commit_group;"); + } + } + { + #pragma unroll + for (int row_local = 0; row_local < 2; row_local++) { + { + { + int smem_row = group * 2 + row_local; + int state_smem_addr_0 = state_smem_addr + (unsigned int)(chunk % 5 * 4096) + (unsigned int)((smem_row * 128 + k_lane * 8) * 2); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_state_row[0])), "=r"(*reinterpret_cast(&packed_state_row[(0) + 1])), "=r"(*reinterpret_cast(&packed_state_row[(0) + 2])), "=r"(*reinterpret_cast(&packed_state_row[(0) + 3])) + : "r"(state_smem_addr_0)); + } + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_row[_pair * 2])[0]), "=f"((&state_row[_pair * 2])[1]) + : "r"(packed_state_row[_pair])); + } + } + int value_row = tile_row_base + chunk * 16 + group * 2 + row_local; + float prediction; + float base; + { + float2 _f2_19 = make_float2(0.0f, 0.0f); + float2 prediction_pair = _f2_19; + float2 _f2_20 = make_float2(0.0f, 0.0f); + float2 base_pair = _f2_20; + #pragma unroll + for (int i_pair = 0; i_pair < 4; i_pair++) { + float2 _f2_21 = make_float2(state_row[i_pair * 2], state_row[i_pair * 2 + 1]); + float2 _f2_22 = make_float2(decay_reg[i_pair * 2], decay_reg[i_pair * 2 + 1]); + float2 _f2_23 = make_float2(0.0f, 0.0f); + float2 decayed_pair = fma_f32x2(_f2_21, _f2_22, _f2_23); + float2 _f2_24 = make_float2(k_reg[i_pair * 2], k_reg[i_pair * 2 + 1]); + prediction_pair = fma_f32x2(decayed_pair, _f2_24, prediction_pair); + float2 _f2_25 = make_float2(q_reg[i_pair * 2], q_reg[i_pair * 2 + 1]); + base_pair = fma_f32x2(decayed_pair, _f2_25, base_pair); + } + prediction = prediction_pair.x + prediction_pair.y; + base = base_pair.x + base_pair.y; + } + float _shfl_xor_20 = __shfl_xor_sync(0xFFFFFFFF, prediction, 8); + prediction += _shfl_xor_20; + float _shfl_xor_21 = __shfl_xor_sync(0xFFFFFFFF, prediction, 4); + prediction += _shfl_xor_21; + float _shfl_xor_22 = __shfl_xor_sync(0xFFFFFFFF, prediction, 2); + prediction += _shfl_xor_22; + float _shfl_xor_23 = __shfl_xor_sync(0xFFFFFFFF, prediction, 1); + prediction += _shfl_xor_23; + float _shfl_xor_24 = __shfl_xor_sync(0xFFFFFFFF, base, 8); + base += _shfl_xor_24; + float _shfl_xor_25 = __shfl_xor_sync(0xFFFFFFFF, base, 4); + base += _shfl_xor_25; + float _shfl_xor_26 = __shfl_xor_sync(0xFFFFFFFF, base, 2); + base += _shfl_xor_26; + float _shfl_xor_27 = __shfl_xor_sync(0xFFFFFFFF, base, 1); + base += _shfl_xor_27; + float v_value = 0.0f; + { + { + if (k_lane == 0) { + unsigned int packed_v[1]; + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&packed_v[0])) : "r"(v_smem_addr + (unsigned int)((value_row - row_local) * 2))); + float packed_v_f32[2]; + #pragma unroll + for (int _pair = 0; _pair < 1; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_v_f32[_pair * 2])[0]), "=f"((&packed_v_f32[_pair * 2])[1]) + : "r"(packed_v[_pair])); + } + v_value = packed_v_f32[row_local]; + } + } + } + { + float _shfl_5 = __shfl_sync(0xFFFFFFFF, v_value, group % 2 * 16); + v_value = _shfl_5; + } + float delta = (v_value - prediction) * beta_value; + { + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_row[_pair * 2])[0]), "=f"((&state_row[_pair * 2])[1]) + : "r"(packed_state_row[_pair])); + } + #pragma unroll + for (int i_pair_1 = 0; i_pair_1 < 4; i_pair_1++) { + float2 _f2_26 = make_float2(k_reg[i_pair_1 * 2], k_reg[i_pair_1 * 2 + 1]); + float2 _f2_27 = make_float2(delta, delta); + float2 _f2_28 = make_float2(state_row[i_pair_1 * 2], state_row[i_pair_1 * 2 + 1]); + float2 _f2_29 = make_float2(decay_reg[i_pair_1 * 2], decay_reg[i_pair_1 * 2 + 1]); + float2 _f2_30 = make_float2(0.0f, 0.0f); + float2 updated_pair = fma_f32x2(_f2_26, _f2_27, fma_f32x2(_f2_28, _f2_29, _f2_30)); + state_row[i_pair_1 * 2] = updated_pair.x; + state_row[i_pair_1 * 2 + 1] = updated_pair.y; + } + } + long long state_row_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)value_row * 128 + (long long)k_lane * 8; + { + { + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(state_row[0 + 0], state_row[0 + 1]); + _pk[1] = __floats2bfloat162_rn(state_row[0 + 2], state_row[0 + 3]); + _pk[2] = __floats2bfloat162_rn(state_row[0 + 4], state_row[0 + 5]); + _pk[3] = __floats2bfloat162_rn(state_row[0 + 6], state_row[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(state))[state_row_base + 0]) = *reinterpret_cast(&_pk[0]); + } + } + } + { + if (k_lane == 0) { + out[output_head_base + value_row] = base + delta * k_dot_q; + } + } + } + } + } + } else if (k_lane == 0) { + #pragma unroll + for (int chunk_1 = 0; chunk_1 < 8; chunk_1++) { + #pragma unroll + for (int row_local_1 = 0; row_local_1 < 2; row_local_1++) { + int value_row_1 = tile_row_base + chunk_1 * 16 + group * 2 + row_local_1; + out[output_head_base + value_row_1] = 0.0f; + } + } + } +} + +} // extern "C" + +// END FROZEN GENERATED BODY +// clang-format on diff --git a/csrc/kda/cake_kda_packed_t1_cpasync_tile128_paired_row_pipeline.cu b/csrc/kda/cake_kda_packed_t1_cpasync_tile128_paired_row_pipeline.cu new file mode 100644 index 00000000000..8b3c4963b57 --- /dev/null +++ b/csrc/kda/cake_kda_packed_t1_cpasync_tile128_paired_row_pipeline.cu @@ -0,0 +1,531 @@ +/* + * 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. + */ + +// clang-format off +// Generated by CAKE from the Kimi K3 serving-native packed T=1 schedule. +// BEGIN FROZEN GENERATED BODY +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +template +struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; + +typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_STATE_SMEM_OFF 0 +#define SMEM_STATE_SMEM_STAGE_BYTES 4096 +#define SMEM_STATE_SMEM_STRIDE 4096 +#define SMEM_V_SMEM_OFF 20480 +#define SMEM_V_SMEM_STAGE_BYTES 256 +#define SMEM_V_SMEM_STRIDE 256 +#define SMEM_TOTAL 20736 +#define THREADS 128 + +#include + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flashinfer_packed_kda_t1_cpasync_tile128_paired_row_pipeline(__nv_bfloat16* __restrict__ q, __nv_bfloat16* __restrict__ k, __nv_bfloat16* __restrict__ v, __nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ beta, float* __restrict__ A_log, float* __restrict__ dt_bias, __nv_bfloat16* __restrict__ state, __nv_bfloat16* __restrict__ out, int* __restrict__ state_indices, float scale, long long q_stride_token, long long k_stride_token, long long v_stride_token, long long g_stride_token, long long beta_stride_token, long long state_stride_slot) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* state_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int state_smem_addr = smem + 0; + __nv_bfloat16* v_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 20480); + const int v_smem_addr = smem + 20480; + + // === Task calls (dependency order) === + int tid_0 = tid; + int lane_1 = lane; + int k_lane = tid_0 % 16; + int group = tid_0 / 16; + int copy_row = tid_0 / 8; + int copy_col = tid_0 % 8 * 16; + int work = blockIdx.x; + int value_tile = 0; + int hv = work; + int n = blockIdx.y; + int raw_slot = state_indices[n]; + bool active = raw_slot >= 0; + int output_head_base = (n * 12 + hv) * 128; + int tile_row_base = value_tile * 128; + if (active) { + long long state_head_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)tile_row_base * 128; + unsigned int v_registers[8]; + { + { + if (k_lane == 0) { + #pragma unroll + for (int v_chunk = 0; v_chunk < 8; v_chunk++) { + int v_row = tile_row_base + v_chunk * 16 + group * 2; + asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 4;" + :: "r"(v_smem_addr + (unsigned int)(v_row * 2)), "l"(v + ((long long)n * v_stride_token + (long long)hv * 128 + (long long)v_row))); + } + } + asm volatile("cp.async.commit_group;"); + } + } + #pragma unroll + for (int prefetch_chunk = 0; prefetch_chunk < 4; prefetch_chunk++) { + { + #pragma unroll + for (int private_row = 0; private_row < 2; private_row++) { + int copy_state_row = group * 2 + private_row; + int copy_elem = (prefetch_chunk * 16 + copy_state_row) * 128 + k_lane * 8; + int copy_dst = state_smem_addr + (unsigned int)(prefetch_chunk * 4096) + (unsigned int)((copy_state_row * 128 + k_lane * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst), "l"(state + (state_head_base + (long long)copy_elem))); + } + } + asm volatile("cp.async.commit_group;"); + } + int elem_start = lane_1 * 4; + long long q_base = (long long)n * q_stride_token + (long long)hv * 128 + (long long)elem_start; + long long k_base = (long long)n * k_stride_token + (long long)hv * 128 + (long long)elem_start; + long long gate_base = (long long)n * g_stride_token + (long long)hv * 128 + (long long)elem_start; + float q_src[4]; + float k_src[4]; + float gate_src[4]; + float dt_bias_src[4]; + float q_reg[8]; + float k_reg[8]; + float decay_reg[8]; + unsigned int packed_state_row[4]; + unsigned int lookahead_state[16]; + float state_row[8]; + float state_rows[16]; + unsigned int paired_state_bits[8]; + float paired_state_rows[16]; + { + { + uint2 _vld_0 = *reinterpret_cast(q + q_base); + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&q_src[0 + _pair * 2])[0]), "=f"((&q_src[0 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + { + uint2 _vld_1 = *reinterpret_cast(k + k_base); + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&k_src[0 + _pair * 2])[0]), "=f"((&k_src[0 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + { + uint2 _vld_2 = *reinterpret_cast(g + gate_base); + uint32_t* _vpairs_2 = reinterpret_cast(&_vld_2); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&gate_src[0 + _pair * 2])[0]), "=f"((&gate_src[0 + _pair * 2])[1]) + : "r"(_vpairs_2[_pair])); + } + } + { + float4 _v4 = *reinterpret_cast(dt_bias + hv * 128 + elem_start); + dt_bias_src[0 + 0] = _v4.x; + dt_bias_src[0 + 1] = _v4.y; + dt_bias_src[0 + 2] = _v4.z; + dt_bias_src[0 + 3] = _v4.w; + } + } + float q_sum_sq = q_src[0] * q_src[0] + q_src[1] * q_src[1] + (q_src[2] * q_src[2] + q_src[3] * q_src[3]); + float k_sum_sq = k_src[0] * k_src[0] + k_src[1] * k_src[1] + (k_src[2] * k_src[2] + k_src[3] * k_src[3]); + float _warp_reduce_0 = q_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset); + q_sum_sq = _warp_reduce_0; + float _warp_reduce_1 = k_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + k_sum_sq = _warp_reduce_1; + float _rsqrt_0 = rsqrtf(q_sum_sq + 1e-06f); + float q_scale = _rsqrt_0 * scale; + float _rsqrt_1 = rsqrtf(k_sum_sq + 1e-06f); + float k_scale = _rsqrt_1; + float _expf_0 = __expf(A_log[hv]); + float gate_a = _expf_0; + #pragma unroll + for (int source_value = 0; source_value < 4; source_value++) { + float biased_gate; + { + biased_gate = gate_src[source_value] + dt_bias_src[source_value]; + } + float _expf_1 = __expf((-gate_a) * biased_gate); + float gate_sigmoid = 1.0f / (1.0f + _expf_1); + float _expf_2 = __expf(-5.0f * gate_sigmoid); + gate_src[source_value] = _expf_2; + } + #pragma unroll + for (int i = 0; i < 8; i++) { + int source_lane = 2 * k_lane + i / 4; + float _shfl_0 = __shfl_sync(0xFFFFFFFF, q_src[i % 4], source_lane); + q_reg[i] = _shfl_0 * q_scale; + float _shfl_1 = __shfl_sync(0xFFFFFFFF, k_src[i % 4], source_lane); + k_reg[i] = _shfl_1 * k_scale; + float _shfl_2 = __shfl_sync(0xFFFFFFFF, gate_src[i % 4], source_lane); + decay_reg[i] = _shfl_2; + } + float k_dot_q = 0.0f; + #pragma unroll + for (int i_1 = 0; i_1 < 8; i_1++) { + k_dot_q += k_reg[i_1] * q_reg[i_1]; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 8); + k_dot_q += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 4); + k_dot_q += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 2); + k_dot_q += _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 1); + k_dot_q += _shfl_xor_3; + float beta_value = (float)beta[(long long)n * beta_stride_token + (long long)hv]; + if (beta_value >= 0.0f) { + float _expf_3 = __expf(-beta_value); + beta_value = 1.0f / (1.0f + _expf_3); + } else { + float _expf_4 = __expf(beta_value); + float beta_exp = _expf_4; + beta_value = beta_exp / (1.0f + beta_exp); + } + { + asm volatile("cp.async.wait_group 4;"); + } + #pragma unroll + for (int chunk = 0; chunk < 8; chunk++) { + { + { + if (chunk < 5) { + asm volatile("cp.async.wait_group 3;"); + } else if (chunk == 5) { + asm volatile("cp.async.wait_group 2;"); + } else { + if (chunk == 6) { + asm volatile("cp.async.wait_group 1;"); + } else { + asm volatile("cp.async.wait_group 0;"); + } + } + } + } + { + if (chunk + 5 - 1 < 8) { + #pragma unroll + for (int private_row_1 = 0; private_row_1 < 2; private_row_1++) { + int copy_state_row_1 = group * 2 + private_row_1; + int copy_elem_1 = ((chunk + 5 - 1) * 16 + copy_state_row_1) * 128 + k_lane * 8; + int copy_dst_1 = state_smem_addr + (unsigned int)((chunk + 5 - 1) % 5 * 4096) + (unsigned int)((copy_state_row_1 * 128 + k_lane * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst_1), "l"(state + (state_head_base + (long long)copy_elem_1))); + } + asm volatile("cp.async.commit_group;"); + } + } + { + #pragma unroll + for (int paired_row = 0; paired_row < 2; paired_row++) { + int paired_smem_row = group * 2 + paired_row; + int paired_state_smem_addr = state_smem_addr + (unsigned int)(chunk % 5 * 4096) + (unsigned int)((paired_smem_row * 128 + k_lane * 8) * 2); + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&paired_state_bits[paired_row * 4])), "=r"(*reinterpret_cast(&paired_state_bits[(paired_row * 4) + 1])), "=r"(*reinterpret_cast(&paired_state_bits[(paired_row * 4) + 2])), "=r"(*reinterpret_cast(&paired_state_bits[(paired_row * 4) + 3])) + : "r"(paired_state_smem_addr)); + } + #pragma unroll + for (int _pair = 0; _pair < 8; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&paired_state_rows[_pair * 2])[0]), "=f"((&paired_state_rows[_pair * 2])[1]) + : "r"(paired_state_bits[_pair])); + } + float2 _f2_0 = make_float2(0.0f, 0.0f); + float2 paired_prediction0 = _f2_0; + float2 _f2_1 = make_float2(0.0f, 0.0f); + float2 paired_prediction1 = _f2_1; + float2 _f2_2 = make_float2(0.0f, 0.0f); + float2 paired_base0 = _f2_2; + float2 _f2_3 = make_float2(0.0f, 0.0f); + float2 paired_base1 = _f2_3; + #pragma unroll + for (int paired_pair = 0; paired_pair < 4; paired_pair++) { + float2 _f2_4 = make_float2(decay_reg[paired_pair * 2], decay_reg[paired_pair * 2 + 1]); + float2 paired_decay = _f2_4; + float2 _f2_5 = make_float2(k_reg[paired_pair * 2], k_reg[paired_pair * 2 + 1]); + float2 paired_k = _f2_5; + float2 _f2_6 = make_float2(q_reg[paired_pair * 2], q_reg[paired_pair * 2 + 1]); + float2 paired_q = _f2_6; + float2 _f2_7 = make_float2(paired_state_rows[paired_pair * 2], paired_state_rows[paired_pair * 2 + 1]); + float2 _f2_8 = make_float2(0.0f, 0.0f); + float2 paired_decayed0 = fma_f32x2(_f2_7, paired_decay, _f2_8); + float2 _f2_9 = make_float2(paired_state_rows[8 + paired_pair * 2], paired_state_rows[8 + (paired_pair * 2 + 1)]); + float2 _f2_10 = make_float2(0.0f, 0.0f); + float2 paired_decayed1 = fma_f32x2(_f2_9, paired_decay, _f2_10); + paired_prediction0 = fma_f32x2(paired_decayed0, paired_k, paired_prediction0); + paired_prediction1 = fma_f32x2(paired_decayed1, paired_k, paired_prediction1); + paired_base0 = fma_f32x2(paired_decayed0, paired_q, paired_base0); + paired_base1 = fma_f32x2(paired_decayed1, paired_q, paired_base1); + } + float prediction0 = paired_prediction0.x + paired_prediction0.y; + float prediction1 = paired_prediction1.x + paired_prediction1.y; + float base0 = paired_base0.x + paired_base0.y; + float base1 = paired_base1.x + paired_base1.y; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, prediction0, 8); + prediction0 += _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, prediction1, 8); + prediction1 += _shfl_xor_5; + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, base0, 8); + base0 += _shfl_xor_6; + float _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, base1, 8); + base1 += _shfl_xor_7; + float _shfl_xor_8 = __shfl_xor_sync(0xFFFFFFFF, prediction0, 4); + prediction0 += _shfl_xor_8; + float _shfl_xor_9 = __shfl_xor_sync(0xFFFFFFFF, prediction1, 4); + prediction1 += _shfl_xor_9; + float _shfl_xor_10 = __shfl_xor_sync(0xFFFFFFFF, base0, 4); + base0 += _shfl_xor_10; + float _shfl_xor_11 = __shfl_xor_sync(0xFFFFFFFF, base1, 4); + base1 += _shfl_xor_11; + float _shfl_xor_12 = __shfl_xor_sync(0xFFFFFFFF, prediction0, 2); + prediction0 += _shfl_xor_12; + float _shfl_xor_13 = __shfl_xor_sync(0xFFFFFFFF, prediction1, 2); + prediction1 += _shfl_xor_13; + float _shfl_xor_14 = __shfl_xor_sync(0xFFFFFFFF, base0, 2); + base0 += _shfl_xor_14; + float _shfl_xor_15 = __shfl_xor_sync(0xFFFFFFFF, base1, 2); + base1 += _shfl_xor_15; + float _shfl_xor_16 = __shfl_xor_sync(0xFFFFFFFF, prediction0, 1); + prediction0 += _shfl_xor_16; + float _shfl_xor_17 = __shfl_xor_sync(0xFFFFFFFF, prediction1, 1); + prediction1 += _shfl_xor_17; + float _shfl_xor_18 = __shfl_xor_sync(0xFFFFFFFF, base0, 1); + base0 += _shfl_xor_18; + float _shfl_xor_19 = __shfl_xor_sync(0xFFFFFFFF, base1, 1); + base1 += _shfl_xor_19; + int paired_value_row0 = tile_row_base + chunk * 16 + group * 2; + unsigned int paired_v[1]; + float residual0 = -prediction0; + float residual1 = -prediction1; + if (k_lane == 0) { + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&paired_v[0])) : "r"(v_smem_addr + (unsigned int)(paired_value_row0 * 2))); + float _bf16x2_add_f32_0[2]; + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %2;\n\t" + "add.rn.f32.bf16 %0, lo, %3;\n\t" + "add.rn.f32.bf16 %1, hi, %4;\n\t" + "}\n" + : "=f"(_bf16x2_add_f32_0[0]), "=f"(_bf16x2_add_f32_0[1]) : "r"(paired_v[0]), "f"(-prediction0), "f"(-prediction1)); + residual0 = _bf16x2_add_f32_0[0]; + residual1 = _bf16x2_add_f32_0[1]; + } + float _shfl_3 = __shfl_sync(0xFFFFFFFF, residual0, group % 2 * 16); + residual0 = _shfl_3; + float _shfl_4 = __shfl_sync(0xFFFFFFFF, residual1, group % 2 * 16); + residual1 = _shfl_4; + float delta0 = residual0 * beta_value; + float delta1 = residual1 * beta_value; + if (k_lane == 0) { + out[output_head_base + paired_value_row0] = base0 + delta0 * k_dot_q; + out[output_head_base + paired_value_row0 + 1] = base1 + delta1 * k_dot_q; + } + #pragma unroll + for (int paired_pair_1 = 0; paired_pair_1 < 4; paired_pair_1++) { + float2 _f2_11 = make_float2(decay_reg[paired_pair_1 * 2], decay_reg[paired_pair_1 * 2 + 1]); + float2 paired_decay_1 = _f2_11; + float2 _f2_12 = make_float2(k_reg[paired_pair_1 * 2], k_reg[paired_pair_1 * 2 + 1]); + float2 paired_k_1 = _f2_12; + float2 _f2_13 = make_float2(delta0, delta0); + float2 _f2_14 = make_float2(paired_state_rows[paired_pair_1 * 2], paired_state_rows[paired_pair_1 * 2 + 1]); + float2 _f2_15 = make_float2(0.0f, 0.0f); + float2 updated0 = fma_f32x2(paired_k_1, _f2_13, fma_f32x2(_f2_14, paired_decay_1, _f2_15)); + float2 _f2_16 = make_float2(delta1, delta1); + float2 _f2_17 = make_float2(paired_state_rows[8 + paired_pair_1 * 2], paired_state_rows[8 + (paired_pair_1 * 2 + 1)]); + float2 _f2_18 = make_float2(0.0f, 0.0f); + float2 updated1 = fma_f32x2(paired_k_1, _f2_16, fma_f32x2(_f2_17, paired_decay_1, _f2_18)); + paired_state_rows[paired_pair_1 * 2] = updated0.x; + paired_state_rows[paired_pair_1 * 2 + 1] = updated0.y; + paired_state_rows[8 + paired_pair_1 * 2] = updated1.x; + paired_state_rows[8 + (paired_pair_1 * 2 + 1)] = updated1.y; + } + long long paired_state_row_base0 = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)paired_value_row0 * 128 + (long long)k_lane * 8; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(paired_state_rows[0 + 0], paired_state_rows[0 + 1]); + _pk[1] = __floats2bfloat162_rn(paired_state_rows[0 + 2], paired_state_rows[0 + 3]); + _pk[2] = __floats2bfloat162_rn(paired_state_rows[0 + 4], paired_state_rows[0 + 5]); + _pk[3] = __floats2bfloat162_rn(paired_state_rows[0 + 6], paired_state_rows[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(state))[paired_state_row_base0 + 0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(paired_state_rows[8 + 0], paired_state_rows[8 + 1]); + _pk[1] = __floats2bfloat162_rn(paired_state_rows[8 + 2], paired_state_rows[8 + 3]); + _pk[2] = __floats2bfloat162_rn(paired_state_rows[8 + 4], paired_state_rows[8 + 5]); + _pk[3] = __floats2bfloat162_rn(paired_state_rows[8 + 6], paired_state_rows[8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(state))[paired_state_row_base0 + 128 + 0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } else if (k_lane == 0) { + #pragma unroll + for (int chunk_1 = 0; chunk_1 < 8; chunk_1++) { + #pragma unroll + for (int row_local = 0; row_local < 2; row_local++) { + int value_row = tile_row_base + chunk_1 * 16 + group * 2 + row_local; + out[output_head_base + value_row] = 0.0f; + } + } + } +} + +} // extern "C" + +// END FROZEN GENERATED BODY +// clang-format on diff --git a/csrc/kda/cake_kda_packed_t1_cpasync_tile128_register_pipeline.cu b/csrc/kda/cake_kda_packed_t1_cpasync_tile128_register_pipeline.cu new file mode 100644 index 00000000000..2ad6f791296 --- /dev/null +++ b/csrc/kda/cake_kda_packed_t1_cpasync_tile128_register_pipeline.cu @@ -0,0 +1,516 @@ +/* + * 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. + */ + +// clang-format off +// Generated by CAKE from the Kimi K3 serving-native packed T=1 schedule. +// BEGIN FROZEN GENERATED BODY +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +template +struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; + +typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_STATE_SMEM_OFF 0 +#define SMEM_STATE_SMEM_STAGE_BYTES 4096 +#define SMEM_STATE_SMEM_STRIDE 4096 +#define SMEM_V_SMEM_OFF 20224 +#define SMEM_V_SMEM_STAGE_BYTES 256 +#define SMEM_V_SMEM_STRIDE 256 +#define SMEM_TOTAL 20480 +#define THREADS 128 + +#include + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flashinfer_packed_kda_t1_cpasync_tile128_register_pipeline(__nv_bfloat16* __restrict__ q, __nv_bfloat16* __restrict__ k, __nv_bfloat16* __restrict__ v, __nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ beta, float* __restrict__ A_log, float* __restrict__ dt_bias, __nv_bfloat16* __restrict__ state, __nv_bfloat16* __restrict__ out, int* __restrict__ state_indices, float scale, long long q_stride_token, long long k_stride_token, long long v_stride_token, long long g_stride_token, long long beta_stride_token, long long state_stride_slot) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* state_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int state_smem_addr = smem + 0; + __nv_bfloat16* v_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 20224); + const int v_smem_addr = smem + 20224; + + // === Task calls (dependency order) === + int tid_0 = tid; + int lane_1 = lane; + int k_lane = tid_0 % 16; + int group = tid_0 / 16; + int copy_row = tid_0 / 8; + int copy_col = tid_0 % 8 * 16; + int work = blockIdx.x; + int value_tile = 0; + int hv = work; + int n = blockIdx.y; + int raw_slot = state_indices[n]; + bool active = raw_slot >= 0; + int output_head_base = (n * 12 + hv) * 128; + int tile_row_base = value_tile * 128; + if (active) { + long long state_head_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)tile_row_base * 128; + unsigned int v_registers[8]; + { + #pragma unroll + for (int v_chunk = 0; v_chunk < 8; v_chunk++) { + int v_row = tile_row_base + v_chunk * 16 + group * 2; + { + v_registers[v_chunk] = *reinterpret_cast(v + (long long)n * v_stride_token + (long long)hv * 128 + (long long)v_row); + } + } + } + #pragma unroll + for (int prefetch_chunk = 0; prefetch_chunk < 4; prefetch_chunk++) { + { + #pragma unroll + for (int private_row = 0; private_row < 2; private_row++) { + int copy_state_row = group * 2 + private_row; + int copy_elem = (prefetch_chunk * 16 + copy_state_row) * 128 + k_lane * 8; + int copy_dst = state_smem_addr + (unsigned int)(prefetch_chunk * 4096) + (unsigned int)((copy_state_row * 128 + k_lane * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst), "l"(state + (state_head_base + (long long)copy_elem))); + } + } + asm volatile("cp.async.commit_group;"); + } + int elem_start = lane_1 * 4; + long long q_base = (long long)n * q_stride_token + (long long)hv * 128 + (long long)elem_start; + long long k_base = (long long)n * k_stride_token + (long long)hv * 128 + (long long)elem_start; + long long gate_base = (long long)n * g_stride_token + (long long)hv * 128 + (long long)elem_start; + float q_src[4]; + float k_src[4]; + float gate_src[4]; + float dt_bias_src[4]; + float q_reg[8]; + float k_reg[8]; + float decay_reg[8]; + unsigned int packed_state_row[4]; + unsigned int lookahead_state[16]; + float state_row[8]; + float state_rows[16]; + unsigned int paired_state_bits[8]; + float paired_state_rows[16]; + { + { + uint2 _vld_1 = *reinterpret_cast(q + q_base); + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&q_src[0 + _pair * 2])[0]), "=f"((&q_src[0 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + { + uint2 _vld_2 = *reinterpret_cast(k + k_base); + uint32_t* _vpairs_2 = reinterpret_cast(&_vld_2); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&k_src[0 + _pair * 2])[0]), "=f"((&k_src[0 + _pair * 2])[1]) + : "r"(_vpairs_2[_pair])); + } + } + { + uint2 _vld_3 = *reinterpret_cast(g + gate_base); + uint32_t* _vpairs_3 = reinterpret_cast(&_vld_3); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&gate_src[0 + _pair * 2])[0]), "=f"((&gate_src[0 + _pair * 2])[1]) + : "r"(_vpairs_3[_pair])); + } + } + { + float4 _v4 = *reinterpret_cast(dt_bias + hv * 128 + elem_start); + dt_bias_src[0 + 0] = _v4.x; + dt_bias_src[0 + 1] = _v4.y; + dt_bias_src[0 + 2] = _v4.z; + dt_bias_src[0 + 3] = _v4.w; + } + } + float q_sum_sq = q_src[0] * q_src[0] + q_src[1] * q_src[1] + (q_src[2] * q_src[2] + q_src[3] * q_src[3]); + float k_sum_sq = k_src[0] * k_src[0] + k_src[1] * k_src[1] + (k_src[2] * k_src[2] + k_src[3] * k_src[3]); + float _warp_reduce_0 = q_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset); + q_sum_sq = _warp_reduce_0; + float _warp_reduce_1 = k_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + k_sum_sq = _warp_reduce_1; + float _rsqrt_0 = rsqrtf(q_sum_sq + 1e-06f); + float q_scale = _rsqrt_0 * scale; + float _rsqrt_1 = rsqrtf(k_sum_sq + 1e-06f); + float k_scale = _rsqrt_1; + float _expf_0 = __expf(A_log[hv]); + float gate_a = _expf_0; + #pragma unroll + for (int source_value = 0; source_value < 4; source_value++) { + float biased_gate; + { + biased_gate = gate_src[source_value] + dt_bias_src[source_value]; + } + float _expf_1 = __expf((-gate_a) * biased_gate); + float gate_sigmoid = 1.0f / (1.0f + _expf_1); + float _expf_2 = __expf(-5.0f * gate_sigmoid); + gate_src[source_value] = _expf_2; + } + #pragma unroll + for (int i = 0; i < 8; i++) { + int source_lane = 2 * k_lane + i / 4; + float _shfl_0 = __shfl_sync(0xFFFFFFFF, q_src[i % 4], source_lane); + q_reg[i] = _shfl_0 * q_scale; + float _shfl_1 = __shfl_sync(0xFFFFFFFF, k_src[i % 4], source_lane); + k_reg[i] = _shfl_1 * k_scale; + float _shfl_2 = __shfl_sync(0xFFFFFFFF, gate_src[i % 4], source_lane); + decay_reg[i] = _shfl_2; + } + float k_dot_q = 0.0f; + #pragma unroll + for (int i_1 = 0; i_1 < 8; i_1++) { + k_dot_q += k_reg[i_1] * q_reg[i_1]; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 8); + k_dot_q += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 4); + k_dot_q += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 2); + k_dot_q += _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 1); + k_dot_q += _shfl_xor_3; + float beta_value = (float)beta[(long long)n * beta_stride_token + (long long)hv]; + if (beta_value >= 0.0f) { + float _expf_3 = __expf(-beta_value); + beta_value = 1.0f / (1.0f + _expf_3); + } else { + float _expf_4 = __expf(beta_value); + float beta_exp = _expf_4; + beta_value = beta_exp / (1.0f + beta_exp); + } + #pragma unroll + for (int chunk = 0; chunk < 8; chunk++) { + { + if (chunk < 5) { + asm volatile("cp.async.wait_group 2;"); + } else if (chunk == 5) { + asm volatile("cp.async.wait_group 1;"); + } else { + if (chunk == 6) { + asm volatile("cp.async.wait_group 0;"); + } + } + } + { + if (chunk + 5 - 1 < 8) { + #pragma unroll + for (int private_row_1 = 0; private_row_1 < 2; private_row_1++) { + int copy_state_row_1 = group * 2 + private_row_1; + int copy_elem_1 = ((chunk + 5 - 1) * 16 + copy_state_row_1) * 128 + k_lane * 8; + int copy_dst_1 = state_smem_addr + (unsigned int)((chunk + 5 - 1) % 5 * 4096) + (unsigned int)((copy_state_row_1 * 128 + k_lane * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst_1), "l"(state + (state_head_base + (long long)copy_elem_1))); + } + asm volatile("cp.async.commit_group;"); + } + } + { + if (chunk == 0) { + #pragma unroll + for (int preload_row = 0; preload_row < 2; preload_row++) { + int preload_smem_row = group * 2 + preload_row; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&lookahead_state[preload_row * 4])), "=r"(*reinterpret_cast(&lookahead_state[(preload_row * 4) + 1])), "=r"(*reinterpret_cast(&lookahead_state[(preload_row * 4) + 2])), "=r"(*reinterpret_cast(&lookahead_state[(preload_row * 4) + 3])) + : "r"(state_smem_addr + (unsigned int)((preload_smem_row * 128 + k_lane * 8) * 2))); + } + } + if (chunk + 1 < 8) { + int next_buffer = 1 - chunk % 2; + #pragma unroll + for (int preload_row_1 = 0; preload_row_1 < 2; preload_row_1++) { + int preload_smem_row_1 = group * 2 + preload_row_1; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&lookahead_state[(next_buffer * 2 + preload_row_1) * 4])), "=r"(*reinterpret_cast(&lookahead_state[((next_buffer * 2 + preload_row_1) * 4) + 1])), "=r"(*reinterpret_cast(&lookahead_state[((next_buffer * 2 + preload_row_1) * 4) + 2])), "=r"(*reinterpret_cast(&lookahead_state[((next_buffer * 2 + preload_row_1) * 4) + 3])) + : "r"(state_smem_addr + (unsigned int)((chunk + 1) % 5 * 4096) + (unsigned int)((preload_smem_row_1 * 128 + k_lane * 8) * 2))); + } + } + } + { + #pragma unroll + for (int row_local = 0; row_local < 2; row_local++) { + { + { + int current_buffer = chunk % 2; + #pragma unroll + for (int state_pair = 0; state_pair < 4; state_pair++) { + packed_state_row[state_pair] = lookahead_state[(current_buffer * 2 + row_local) * 4 + state_pair]; + } + } + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_row[_pair * 2])[0]), "=f"((&state_row[_pair * 2])[1]) + : "r"(packed_state_row[_pair])); + } + } + int value_row = tile_row_base + chunk * 16 + group * 2 + row_local; + float prediction; + float base; + { + float2 _f2_19 = make_float2(0.0f, 0.0f); + float2 prediction_pair = _f2_19; + float2 _f2_20 = make_float2(0.0f, 0.0f); + float2 base_pair = _f2_20; + #pragma unroll + for (int i_pair = 0; i_pair < 4; i_pair++) { + float2 _f2_21 = make_float2(state_row[i_pair * 2], state_row[i_pair * 2 + 1]); + float2 _f2_22 = make_float2(decay_reg[i_pair * 2], decay_reg[i_pair * 2 + 1]); + float2 _f2_23 = make_float2(0.0f, 0.0f); + float2 decayed_pair = fma_f32x2(_f2_21, _f2_22, _f2_23); + float2 _f2_24 = make_float2(k_reg[i_pair * 2], k_reg[i_pair * 2 + 1]); + prediction_pair = fma_f32x2(decayed_pair, _f2_24, prediction_pair); + float2 _f2_25 = make_float2(q_reg[i_pair * 2], q_reg[i_pair * 2 + 1]); + base_pair = fma_f32x2(decayed_pair, _f2_25, base_pair); + } + prediction = prediction_pair.x + prediction_pair.y; + base = base_pair.x + base_pair.y; + } + float _shfl_xor_20 = __shfl_xor_sync(0xFFFFFFFF, prediction, 8); + prediction += _shfl_xor_20; + float _shfl_xor_21 = __shfl_xor_sync(0xFFFFFFFF, prediction, 4); + prediction += _shfl_xor_21; + float _shfl_xor_22 = __shfl_xor_sync(0xFFFFFFFF, prediction, 2); + prediction += _shfl_xor_22; + float _shfl_xor_23 = __shfl_xor_sync(0xFFFFFFFF, prediction, 1); + prediction += _shfl_xor_23; + float _shfl_xor_24 = __shfl_xor_sync(0xFFFFFFFF, base, 8); + base += _shfl_xor_24; + float _shfl_xor_25 = __shfl_xor_sync(0xFFFFFFFF, base, 4); + base += _shfl_xor_25; + float _shfl_xor_26 = __shfl_xor_sync(0xFFFFFFFF, base, 2); + base += _shfl_xor_26; + float _shfl_xor_27 = __shfl_xor_sync(0xFFFFFFFF, base, 1); + base += _shfl_xor_27; + float v_value = 0.0f; + { + float v_registers_f32[2]; + #pragma unroll + for (int _pair = 0; _pair < 1; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&v_registers_f32[_pair * 2])[0]), "=f"((&v_registers_f32[_pair * 2])[1]) + : "r"((v_registers + chunk)[_pair])); + } + v_value = v_registers_f32[row_local]; + } + float delta = (v_value - prediction) * beta_value; + { + if (k_lane == 0) { + out[output_head_base + value_row] = base + delta * k_dot_q; + } + } + { + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_row[_pair * 2])[0]), "=f"((&state_row[_pair * 2])[1]) + : "r"(packed_state_row[_pair])); + } + #pragma unroll + for (int i_pair_1 = 0; i_pair_1 < 4; i_pair_1++) { + float2 _f2_26 = make_float2(k_reg[i_pair_1 * 2], k_reg[i_pair_1 * 2 + 1]); + float2 _f2_27 = make_float2(delta, delta); + float2 _f2_28 = make_float2(state_row[i_pair_1 * 2], state_row[i_pair_1 * 2 + 1]); + float2 _f2_29 = make_float2(decay_reg[i_pair_1 * 2], decay_reg[i_pair_1 * 2 + 1]); + float2 _f2_30 = make_float2(0.0f, 0.0f); + float2 updated_pair = fma_f32x2(_f2_26, _f2_27, fma_f32x2(_f2_28, _f2_29, _f2_30)); + state_row[i_pair_1 * 2] = updated_pair.x; + state_row[i_pair_1 * 2 + 1] = updated_pair.y; + } + } + long long state_row_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)value_row * 128 + (long long)k_lane * 8; + { + { + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(state_row[0 + 0], state_row[0 + 1]); + _pk[1] = __floats2bfloat162_rn(state_row[0 + 2], state_row[0 + 3]); + _pk[2] = __floats2bfloat162_rn(state_row[0 + 4], state_row[0 + 5]); + _pk[3] = __floats2bfloat162_rn(state_row[0 + 6], state_row[0 + 7]); + uint4 _st_v4 = *reinterpret_cast(&_pk[0]); + asm volatile( + "st.global.L2::cache_hint.v4.b32 [%0], {%1, %2, %3, %4}, %5;" + :: "l"((void*)(&((__nv_bfloat16*)(state))[state_row_base + 0])), "r"(_st_v4.x), "r"(_st_v4.y), "r"(_st_v4.z), "r"(_st_v4.w), "l"(0x14F0000000000000ULL) : "memory"); + } + } + } + } + } + } + } else if (k_lane == 0) { + #pragma unroll + for (int chunk_1 = 0; chunk_1 < 8; chunk_1++) { + #pragma unroll + for (int row_local_1 = 0; row_local_1 < 2; row_local_1++) { + int value_row_1 = tile_row_base + chunk_1 * 16 + group * 2 + row_local_1; + out[output_head_base + value_row_1] = 0.0f; + } + } + } +} + +} // extern "C" + +// END FROZEN GENERATED BODY +// clang-format on diff --git a/csrc/kda/cake_kda_packed_t1_cpasync_tile128_v_private_prefetch.cu b/csrc/kda/cake_kda_packed_t1_cpasync_tile128_v_private_prefetch.cu new file mode 100644 index 00000000000..9baf6d7c920 --- /dev/null +++ b/csrc/kda/cake_kda_packed_t1_cpasync_tile128_v_private_prefetch.cu @@ -0,0 +1,495 @@ +/* + * 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. + */ + +// clang-format off +// Generated by CAKE from the Kimi K3 serving-native packed T=1 schedule. +// BEGIN FROZEN GENERATED BODY +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +template +struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; + +typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_STATE_SMEM_OFF 0 +#define SMEM_STATE_SMEM_STAGE_BYTES 4096 +#define SMEM_STATE_SMEM_STRIDE 4096 +#define SMEM_V_SMEM_OFF 20480 +#define SMEM_V_SMEM_STAGE_BYTES 256 +#define SMEM_V_SMEM_STRIDE 256 +#define SMEM_TOTAL 20736 +#define THREADS 128 + +#include + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flashinfer_packed_kda_t1_cpasync_tile128_v_private_prefetch(__nv_bfloat16* __restrict__ q, __nv_bfloat16* __restrict__ k, __nv_bfloat16* __restrict__ v, __nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ beta, float* __restrict__ A_log, float* __restrict__ dt_bias, __nv_bfloat16* __restrict__ state, __nv_bfloat16* __restrict__ out, int* __restrict__ state_indices, float scale, long long q_stride_token, long long k_stride_token, long long v_stride_token, long long g_stride_token, long long beta_stride_token, long long state_stride_slot) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* state_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int state_smem_addr = smem + 0; + __nv_bfloat16* v_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 20480); + const int v_smem_addr = smem + 20480; + + // === Task calls (dependency order) === + int tid_0 = tid; + int lane_1 = lane; + int k_lane = tid_0 % 16; + int group = tid_0 / 16; + int copy_row = tid_0 / 8; + int copy_col = tid_0 % 8 * 16; + int work = blockIdx.x; + int value_tile = 0; + int hv = work; + int n = blockIdx.y; + int raw_slot = state_indices[n]; + bool active = raw_slot >= 0; + int output_head_base = (n * 12 + hv) * 128; + int tile_row_base = value_tile * 128; + if (active) { + long long state_head_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)tile_row_base * 128; + unsigned int v_registers[8]; + { + { + if (k_lane == 0) { + #pragma unroll + for (int v_chunk = 0; v_chunk < 8; v_chunk++) { + int v_row = tile_row_base + v_chunk * 16 + group * 2; + asm volatile("cp.async.ca.shared::cta.global [%0], [%1], 4;" + :: "r"(v_smem_addr + (unsigned int)(v_row * 2)), "l"(v + ((long long)n * v_stride_token + (long long)hv * 128 + (long long)v_row))); + } + } + asm volatile("cp.async.commit_group;"); + } + } + #pragma unroll + for (int prefetch_chunk = 0; prefetch_chunk < 4; prefetch_chunk++) { + { + #pragma unroll + for (int private_row = 0; private_row < 2; private_row++) { + int copy_state_row = group * 2 + private_row; + int copy_elem = (prefetch_chunk * 16 + copy_state_row) * 128 + k_lane * 8; + int copy_dst = state_smem_addr + (unsigned int)(prefetch_chunk * 4096) + (unsigned int)((copy_state_row * 128 + k_lane * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst), "l"(state + (state_head_base + (long long)copy_elem))); + } + } + asm volatile("cp.async.commit_group;"); + } + int elem_start = lane_1 * 4; + long long q_base = (long long)n * q_stride_token + (long long)hv * 128 + (long long)elem_start; + long long k_base = (long long)n * k_stride_token + (long long)hv * 128 + (long long)elem_start; + long long gate_base = (long long)n * g_stride_token + (long long)hv * 128 + (long long)elem_start; + float q_src[4]; + float k_src[4]; + float gate_src[4]; + float dt_bias_src[4]; + float q_reg[8]; + float k_reg[8]; + float decay_reg[8]; + unsigned int packed_state_row[4]; + unsigned int lookahead_state[16]; + float state_row[8]; + float state_rows[16]; + unsigned int paired_state_bits[8]; + float paired_state_rows[16]; + { + { + uint2 _vld_0 = *reinterpret_cast(q + q_base); + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&q_src[0 + _pair * 2])[0]), "=f"((&q_src[0 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + { + uint2 _vld_1 = *reinterpret_cast(k + k_base); + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&k_src[0 + _pair * 2])[0]), "=f"((&k_src[0 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + { + uint2 _vld_2 = *reinterpret_cast(g + gate_base); + uint32_t* _vpairs_2 = reinterpret_cast(&_vld_2); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&gate_src[0 + _pair * 2])[0]), "=f"((&gate_src[0 + _pair * 2])[1]) + : "r"(_vpairs_2[_pair])); + } + } + { + float4 _v4 = *reinterpret_cast(dt_bias + hv * 128 + elem_start); + dt_bias_src[0 + 0] = _v4.x; + dt_bias_src[0 + 1] = _v4.y; + dt_bias_src[0 + 2] = _v4.z; + dt_bias_src[0 + 3] = _v4.w; + } + } + float q_sum_sq = q_src[0] * q_src[0] + q_src[1] * q_src[1] + (q_src[2] * q_src[2] + q_src[3] * q_src[3]); + float k_sum_sq = k_src[0] * k_src[0] + k_src[1] * k_src[1] + (k_src[2] * k_src[2] + k_src[3] * k_src[3]); + float _warp_reduce_0 = q_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset); + q_sum_sq = _warp_reduce_0; + float _warp_reduce_1 = k_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + k_sum_sq = _warp_reduce_1; + float _rsqrt_0 = rsqrtf(q_sum_sq + 1e-06f); + float q_scale = _rsqrt_0 * scale; + float _rsqrt_1 = rsqrtf(k_sum_sq + 1e-06f); + float k_scale = _rsqrt_1; + float _expf_0 = __expf(A_log[hv]); + float gate_a = _expf_0; + #pragma unroll + for (int source_value = 0; source_value < 4; source_value++) { + float biased_gate; + { + biased_gate = gate_src[source_value] + dt_bias_src[source_value]; + } + float _expf_1 = __expf((-gate_a) * biased_gate); + float gate_sigmoid = 1.0f / (1.0f + _expf_1); + float _expf_2 = __expf(-5.0f * gate_sigmoid); + gate_src[source_value] = _expf_2; + } + #pragma unroll + for (int i = 0; i < 8; i++) { + int source_lane = 2 * k_lane + i / 4; + float _shfl_0 = __shfl_sync(0xFFFFFFFF, q_src[i % 4], source_lane); + q_reg[i] = _shfl_0 * q_scale; + float _shfl_1 = __shfl_sync(0xFFFFFFFF, k_src[i % 4], source_lane); + k_reg[i] = _shfl_1 * k_scale; + float _shfl_2 = __shfl_sync(0xFFFFFFFF, gate_src[i % 4], source_lane); + decay_reg[i] = _shfl_2; + } + float k_dot_q = 0.0f; + #pragma unroll + for (int i_1 = 0; i_1 < 8; i_1++) { + k_dot_q += k_reg[i_1] * q_reg[i_1]; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 8); + k_dot_q += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 4); + k_dot_q += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 2); + k_dot_q += _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 1); + k_dot_q += _shfl_xor_3; + float beta_value = (float)beta[(long long)n * beta_stride_token + (long long)hv]; + if (beta_value >= 0.0f) { + float _expf_3 = __expf(-beta_value); + beta_value = 1.0f / (1.0f + _expf_3); + } else { + float _expf_4 = __expf(beta_value); + float beta_exp = _expf_4; + beta_value = beta_exp / (1.0f + beta_exp); + } + { + asm volatile("cp.async.wait_group 4;"); + } + #pragma unroll + for (int chunk = 0; chunk < 8; chunk++) { + { + { + if (chunk < 5) { + asm volatile("cp.async.wait_group 3;"); + } else if (chunk == 5) { + asm volatile("cp.async.wait_group 2;"); + } else { + if (chunk == 6) { + asm volatile("cp.async.wait_group 1;"); + } else { + asm volatile("cp.async.wait_group 0;"); + } + } + } + } + { + if (chunk + 5 - 1 < 8) { + #pragma unroll + for (int private_row_1 = 0; private_row_1 < 2; private_row_1++) { + int copy_state_row_1 = group * 2 + private_row_1; + int copy_elem_1 = ((chunk + 5 - 1) * 16 + copy_state_row_1) * 128 + k_lane * 8; + int copy_dst_1 = state_smem_addr + (unsigned int)((chunk + 5 - 1) % 5 * 4096) + (unsigned int)((copy_state_row_1 * 128 + k_lane * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst_1), "l"(state + (state_head_base + (long long)copy_elem_1))); + } + asm volatile("cp.async.commit_group;"); + } + } + { + { + #pragma unroll + for (int row_local = 0; row_local < 2; row_local++) { + unsigned int packed_state[4]; + int smem_row = group * 2 + row_local; + int state_smem_addr_0 = 0; + { + state_smem_addr_0 = state_smem_addr + (unsigned int)(chunk % 5 * 4096) + (unsigned int)((smem_row * 128 + k_lane * 8) * 2); + } + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_state[0])), "=r"(*reinterpret_cast(&packed_state[(0) + 1])), "=r"(*reinterpret_cast(&packed_state[(0) + 2])), "=r"(*reinterpret_cast(&packed_state[(0) + 3])) + : "r"(state_smem_addr_0)); + float packed_state_f32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_state_f32[_pair * 2])[0]), "=f"((&packed_state_f32[_pair * 2])[1]) + : "r"(packed_state[_pair])); + } + #pragma unroll + for (int i_2 = 0; i_2 < 8; i_2++) { + state_rows[row_local * 8 + i_2] = packed_state_f32[i_2]; + } + } + } + #pragma unroll + for (int row_local_1 = 0; row_local_1 < 2; row_local_1++) { + int value_row = tile_row_base + chunk * 16 + group * 2 + row_local_1; + float prediction; + float base; + { + prediction = 0.0f; + base = 0.0f; + #pragma unroll + for (int i_3 = 0; i_3 < 8; i_3++) { + float decayed = state_rows[row_local_1 * 8 + i_3] * decay_reg[i_3]; + prediction += decayed * k_reg[i_3]; + base += decayed * q_reg[i_3]; + } + } + float _shfl_xor_20 = __shfl_xor_sync(0xFFFFFFFF, prediction, 8); + prediction += _shfl_xor_20; + float _shfl_xor_21 = __shfl_xor_sync(0xFFFFFFFF, prediction, 4); + prediction += _shfl_xor_21; + float _shfl_xor_22 = __shfl_xor_sync(0xFFFFFFFF, prediction, 2); + prediction += _shfl_xor_22; + float _shfl_xor_23 = __shfl_xor_sync(0xFFFFFFFF, prediction, 1); + prediction += _shfl_xor_23; + float _shfl_xor_24 = __shfl_xor_sync(0xFFFFFFFF, base, 8); + base += _shfl_xor_24; + float _shfl_xor_25 = __shfl_xor_sync(0xFFFFFFFF, base, 4); + base += _shfl_xor_25; + float _shfl_xor_26 = __shfl_xor_sync(0xFFFFFFFF, base, 2); + base += _shfl_xor_26; + float _shfl_xor_27 = __shfl_xor_sync(0xFFFFFFFF, base, 1); + base += _shfl_xor_27; + float v_value = 0.0f; + { + { + if (k_lane == 0) { + unsigned int packed_v[1]; + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&packed_v[0])) : "r"(v_smem_addr + (unsigned int)((value_row - row_local_1) * 2))); + float packed_v_f32[2]; + #pragma unroll + for (int _pair = 0; _pair < 1; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_v_f32[_pair * 2])[0]), "=f"((&packed_v_f32[_pair * 2])[1]) + : "r"(packed_v[_pair])); + } + v_value = packed_v_f32[row_local_1]; + } + } + } + { + float _shfl_5 = __shfl_sync(0xFFFFFFFF, v_value, group % 2 * 16); + v_value = _shfl_5; + } + float delta = (v_value - prediction) * beta_value; + { + #pragma unroll + for (int i_4 = 0; i_4 < 8; i_4++) { + state_rows[row_local_1 * 8 + i_4] = state_rows[row_local_1 * 8 + i_4] * decay_reg[i_4] + delta * k_reg[i_4]; + } + } + long long state_row_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)value_row * 128 + (long long)k_lane * 8; + { + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 0], state_rows[row_local_1 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 2], state_rows[row_local_1 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 4], state_rows[row_local_1 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 6], state_rows[row_local_1 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(state))[state_row_base + 0]) = *reinterpret_cast(&_pk[0]); + } + } + { + if (k_lane == 0) { + out[output_head_base + value_row] = base + delta * k_dot_q; + } + } + } + } + } + } else if (k_lane == 0) { + #pragma unroll + for (int chunk_1 = 0; chunk_1 < 8; chunk_1++) { + #pragma unroll + for (int row_local_2 = 0; row_local_2 < 2; row_local_2++) { + int value_row_1 = tile_row_base + chunk_1 * 16 + group * 2 + row_local_2; + out[output_head_base + value_row_1] = 0.0f; + } + } + } +} + +} // extern "C" + +// END FROZEN GENERATED BODY +// clang-format on diff --git a/csrc/kda/cake_kda_packed_t1_cpasync_tile64.cu b/csrc/kda/cake_kda_packed_t1_cpasync_tile64.cu new file mode 100644 index 00000000000..ec47cce59b5 --- /dev/null +++ b/csrc/kda/cake_kda_packed_t1_cpasync_tile64.cu @@ -0,0 +1,421 @@ +/* + * 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. + */ + +// clang-format off +// Generated by CAKE from the Kimi K3 serving-native packed T=1 schedule. +// BEGIN FROZEN GENERATED BODY +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +template +struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; + +typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_STATE_SMEM_OFF 0 +#define SMEM_STATE_SMEM_STAGE_BYTES 4096 +#define SMEM_STATE_SMEM_STRIDE 4096 +#define SMEM_V_SMEM_OFF 16128 +#define SMEM_V_SMEM_STAGE_BYTES 256 +#define SMEM_V_SMEM_STRIDE 256 +#define SMEM_TOTAL 16384 +#define THREADS 128 + +#include + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flashinfer_packed_kda_t1_cpasync(__nv_bfloat16* __restrict__ q, __nv_bfloat16* __restrict__ k, __nv_bfloat16* __restrict__ v, __nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ beta, float* __restrict__ A_log, float* __restrict__ dt_bias, __nv_bfloat16* __restrict__ state, __nv_bfloat16* __restrict__ out, int* __restrict__ state_indices, float scale, long long q_stride_token, long long k_stride_token, long long v_stride_token, long long g_stride_token, long long beta_stride_token, long long state_stride_slot) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* state_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int state_smem_addr = smem + 0; + __nv_bfloat16* v_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 16128); + const int v_smem_addr = smem + 16128; + + // === Task calls (dependency order) === + int tid_0 = tid; + int lane_1 = lane; + int k_lane = tid_0 % 16; + int group = tid_0 / 16; + int copy_row = tid_0 / 8; + int copy_col = tid_0 % 8 * 16; + int work = blockIdx.x; + int value_tile = work % 2; + int hv = work / 2; + int n = blockIdx.y; + int raw_slot = state_indices[n]; + bool active = raw_slot >= 0; + int output_head_base = (n * 12 + hv) * 128; + int tile_row_base = value_tile * 64; + if (active) { + long long state_head_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)tile_row_base * 128; + unsigned int v_registers[4]; + { + } + #pragma unroll + for (int prefetch_chunk = 0; prefetch_chunk < 3; prefetch_chunk++) { + { + #pragma unroll + for (int copy_vec = 0; copy_vec < 2; copy_vec++) { + int copy_elem = (prefetch_chunk * 16 + copy_row) * 128 + copy_col + copy_vec * 8; + int copy_dst = state_smem_addr + (unsigned int)(prefetch_chunk * 4096) + (unsigned int)((copy_row * 128 + copy_col + copy_vec * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst), "l"(state + (state_head_base + (long long)copy_elem))); + } + } + asm volatile("cp.async.commit_group;"); + } + int elem_start = lane_1 * 4; + long long q_base = (long long)n * q_stride_token + (long long)hv * 128 + (long long)elem_start; + long long k_base = (long long)n * k_stride_token + (long long)hv * 128 + (long long)elem_start; + long long gate_base = (long long)n * g_stride_token + (long long)hv * 128 + (long long)elem_start; + float q_src[4]; + float k_src[4]; + float gate_src[4]; + float dt_bias_src[4]; + float q_reg[8]; + float k_reg[8]; + float decay_reg[8]; + unsigned int packed_state_row[4]; + unsigned int lookahead_state[16]; + float state_row[8]; + float state_rows[16]; + unsigned int paired_state_bits[8]; + float paired_state_rows[16]; + { + #pragma unroll + for (int source_value = 0; source_value < 4; source_value++) { + q_src[source_value] = (float)q[q_base + (long long)source_value]; + k_src[source_value] = (float)k[k_base + (long long)source_value]; + gate_src[source_value] = (float)g[gate_base + (long long)source_value]; + } + } + float q_sum_sq = q_src[0] * q_src[0] + q_src[1] * q_src[1] + (q_src[2] * q_src[2] + q_src[3] * q_src[3]); + float k_sum_sq = k_src[0] * k_src[0] + k_src[1] * k_src[1] + (k_src[2] * k_src[2] + k_src[3] * k_src[3]); + float _warp_reduce_0 = q_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset); + q_sum_sq = _warp_reduce_0; + float _warp_reduce_1 = k_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + k_sum_sq = _warp_reduce_1; + float _rsqrt_0 = rsqrtf(q_sum_sq + 1e-06f); + float q_scale = _rsqrt_0 * scale; + float _rsqrt_1 = rsqrtf(k_sum_sq + 1e-06f); + float k_scale = _rsqrt_1; + float _expf_0 = __expf(A_log[hv]); + float gate_a = _expf_0; + #pragma unroll + for (int source_value_1 = 0; source_value_1 < 4; source_value_1++) { + float biased_gate; + { + int k_idx = elem_start + source_value_1; + biased_gate = gate_src[source_value_1] + dt_bias[hv * 128 + k_idx]; + } + float _expf_1 = __expf((-gate_a) * biased_gate); + float gate_sigmoid = 1.0f / (1.0f + _expf_1); + float _expf_2 = __expf(-5.0f * gate_sigmoid); + gate_src[source_value_1] = _expf_2; + } + #pragma unroll + for (int i = 0; i < 8; i++) { + int source_lane = 2 * k_lane + i / 4; + float _shfl_0 = __shfl_sync(0xFFFFFFFF, q_src[i % 4], source_lane); + q_reg[i] = _shfl_0 * q_scale; + float _shfl_1 = __shfl_sync(0xFFFFFFFF, k_src[i % 4], source_lane); + k_reg[i] = _shfl_1 * k_scale; + float _shfl_2 = __shfl_sync(0xFFFFFFFF, gate_src[i % 4], source_lane); + decay_reg[i] = _shfl_2; + } + float k_dot_q = 0.0f; + #pragma unroll + for (int i_1 = 0; i_1 < 8; i_1++) { + k_dot_q += k_reg[i_1] * q_reg[i_1]; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 8); + k_dot_q += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 4); + k_dot_q += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 2); + k_dot_q += _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 1); + k_dot_q += _shfl_xor_3; + float beta_value = (float)beta[(long long)n * beta_stride_token + (long long)hv]; + if (beta_value >= 0.0f) { + float _expf_3 = __expf(-beta_value); + beta_value = 1.0f / (1.0f + _expf_3); + } else { + float _expf_4 = __expf(beta_value); + float beta_exp = _expf_4; + beta_value = beta_exp / (1.0f + beta_exp); + } + #pragma unroll + for (int chunk = 0; chunk < 4; chunk++) { + { + { + if (chunk < 2) { + asm volatile("cp.async.wait_group 1;"); + } else { + asm volatile("cp.async.wait_group 0;"); + } + } + } + { + __syncthreads(); + } + { + if (chunk == 0) { + #pragma unroll + for (int copy_vec_1 = 0; copy_vec_1 < 2; copy_vec_1++) { + int copy_elem_1 = (48 + copy_row) * 128 + copy_col + copy_vec_1 * 8; + int copy_dst_1 = state_smem_addr + 12288 + (unsigned int)((copy_row * 128 + copy_col + copy_vec_1 * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst_1), "l"(state + (state_head_base + (long long)copy_elem_1))); + } + asm volatile("cp.async.commit_group;"); + } + } + { + { + #pragma unroll + for (int row_local = 0; row_local < 2; row_local++) { + unsigned int packed_state[4]; + int smem_row = group * 2 + row_local; + int state_smem_addr_0 = 0; + { + state_smem_addr_0 = state_smem_addr + (unsigned int)(chunk * 4096) + (unsigned int)((smem_row * 128 + k_lane * 8) * 2); + } + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_state[0])), "=r"(*reinterpret_cast(&packed_state[(0) + 1])), "=r"(*reinterpret_cast(&packed_state[(0) + 2])), "=r"(*reinterpret_cast(&packed_state[(0) + 3])) + : "r"(state_smem_addr_0)); + float packed_state_f32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_state_f32[_pair * 2])[0]), "=f"((&packed_state_f32[_pair * 2])[1]) + : "r"(packed_state[_pair])); + } + #pragma unroll + for (int i_2 = 0; i_2 < 8; i_2++) { + state_rows[row_local * 8 + i_2] = packed_state_f32[i_2]; + } + } + } + #pragma unroll + for (int row_local_1 = 0; row_local_1 < 2; row_local_1++) { + int value_row = tile_row_base + chunk * 16 + group * 2 + row_local_1; + float prediction; + float base; + { + prediction = 0.0f; + base = 0.0f; + #pragma unroll + for (int i_3 = 0; i_3 < 8; i_3++) { + float decayed = state_rows[row_local_1 * 8 + i_3] * decay_reg[i_3]; + prediction += decayed * k_reg[i_3]; + base += decayed * q_reg[i_3]; + } + } + float _shfl_xor_20 = __shfl_xor_sync(0xFFFFFFFF, prediction, 8); + prediction += _shfl_xor_20; + float _shfl_xor_21 = __shfl_xor_sync(0xFFFFFFFF, prediction, 4); + prediction += _shfl_xor_21; + float _shfl_xor_22 = __shfl_xor_sync(0xFFFFFFFF, prediction, 2); + prediction += _shfl_xor_22; + float _shfl_xor_23 = __shfl_xor_sync(0xFFFFFFFF, prediction, 1); + prediction += _shfl_xor_23; + float _shfl_xor_24 = __shfl_xor_sync(0xFFFFFFFF, base, 8); + base += _shfl_xor_24; + float _shfl_xor_25 = __shfl_xor_sync(0xFFFFFFFF, base, 4); + base += _shfl_xor_25; + float _shfl_xor_26 = __shfl_xor_sync(0xFFFFFFFF, base, 2); + base += _shfl_xor_26; + float _shfl_xor_27 = __shfl_xor_sync(0xFFFFFFFF, base, 1); + base += _shfl_xor_27; + float v_value = 0.0f; + { + { + if (k_lane == 0) { + v_value = (float)v[(long long)n * v_stride_token + (long long)hv * 128 + (long long)value_row]; + } + } + } + { + float _shfl_5 = __shfl_sync(0xFFFFFFFF, v_value, group % 2 * 16); + v_value = _shfl_5; + } + float delta = (v_value - prediction) * beta_value; + { + #pragma unroll + for (int i_4 = 0; i_4 < 8; i_4++) { + state_rows[row_local_1 * 8 + i_4] = state_rows[row_local_1 * 8 + i_4] * decay_reg[i_4] + delta * k_reg[i_4]; + } + } + long long state_row_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)value_row * 128 + (long long)k_lane * 8; + { + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 0], state_rows[row_local_1 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 2], state_rows[row_local_1 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 4], state_rows[row_local_1 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 6], state_rows[row_local_1 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(state))[state_row_base + 0]) = *reinterpret_cast(&_pk[0]); + } + } + { + if (k_lane == 0) { + out[output_head_base + value_row] = base + delta * k_dot_q; + } + } + } + } + } + } else if (k_lane == 0) { + #pragma unroll + for (int chunk_1 = 0; chunk_1 < 4; chunk_1++) { + #pragma unroll + for (int row_local_2 = 0; row_local_2 < 2; row_local_2++) { + int value_row_1 = tile_row_base + chunk_1 * 16 + group * 2 + row_local_2; + out[output_head_base + value_row_1] = 0.0f; + } + } + } +} + +} // extern "C" + +// END FROZEN GENERATED BODY +// clang-format on diff --git a/csrc/kda/cake_kda_packed_t1_cpasync_tile64_ilp4.cu b/csrc/kda/cake_kda_packed_t1_cpasync_tile64_ilp4.cu new file mode 100644 index 00000000000..195db6fe59a --- /dev/null +++ b/csrc/kda/cake_kda_packed_t1_cpasync_tile64_ilp4.cu @@ -0,0 +1,426 @@ +/* + * 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. + */ + +// clang-format off +// Generated by CAKE from the Kimi K3 serving-native packed T=1 schedule. +// BEGIN FROZEN GENERATED BODY +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +template +struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; + +typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_STATE_SMEM_OFF 0 +#define SMEM_STATE_SMEM_STAGE_BYTES 8192 +#define SMEM_STATE_SMEM_STRIDE 8192 +#define SMEM_V_SMEM_OFF 24320 +#define SMEM_V_SMEM_STAGE_BYTES 256 +#define SMEM_V_SMEM_STRIDE 256 +#define SMEM_TOTAL 24576 +#define THREADS 128 + +#include + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flashinfer_packed_kda_t1_cpasync_tile64_ilp4(__nv_bfloat16* __restrict__ q, __nv_bfloat16* __restrict__ k, __nv_bfloat16* __restrict__ v, __nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ beta, float* __restrict__ A_log, float* __restrict__ dt_bias, __nv_bfloat16* __restrict__ state, __nv_bfloat16* __restrict__ out, int* __restrict__ state_indices, float scale, long long q_stride_token, long long k_stride_token, long long v_stride_token, long long g_stride_token, long long beta_stride_token, long long state_stride_slot) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* state_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int state_smem_addr = smem + 0; + __nv_bfloat16* v_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 24320); + const int v_smem_addr = smem + 24320; + + // === Task calls (dependency order) === + int tid_0 = tid; + int lane_1 = lane; + int k_lane = tid_0 % 16; + int group = tid_0 / 16; + int copy_row = tid_0 / 8; + int copy_col = tid_0 % 8 * 16; + int work = blockIdx.x; + int value_tile = work % 2; + int hv = work / 2; + int n = blockIdx.y; + int raw_slot = state_indices[n]; + bool active = raw_slot >= 0; + int output_head_base = (n * 12 + hv) * 128; + int tile_row_base = value_tile * 64; + if (active) { + long long state_head_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)tile_row_base * 128; + unsigned int v_registers[2]; + { + } + #pragma unroll + for (int prefetch_chunk = 0; prefetch_chunk < 2; prefetch_chunk++) { + { + #pragma unroll + for (int private_row = 0; private_row < 4; private_row++) { + int copy_state_row = group * 4 + private_row; + int copy_elem = (prefetch_chunk * 32 + copy_state_row) * 128 + k_lane * 8; + int copy_dst = state_smem_addr + (unsigned int)(prefetch_chunk * 8192) + (unsigned int)((copy_state_row * 128 + k_lane * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst), "l"(state + (state_head_base + (long long)copy_elem))); + } + } + asm volatile("cp.async.commit_group;"); + } + int elem_start = lane_1 * 4; + long long q_base = (long long)n * q_stride_token + (long long)hv * 128 + (long long)elem_start; + long long k_base = (long long)n * k_stride_token + (long long)hv * 128 + (long long)elem_start; + long long gate_base = (long long)n * g_stride_token + (long long)hv * 128 + (long long)elem_start; + float q_src[4]; + float k_src[4]; + float gate_src[4]; + float dt_bias_src[4]; + float q_reg[8]; + float k_reg[8]; + float decay_reg[8]; + unsigned int packed_state_row[4]; + unsigned int lookahead_state[32]; + float state_row[8]; + float state_rows[32]; + unsigned int paired_state_bits[16]; + float paired_state_rows[32]; + { + #pragma unroll + for (int source_value = 0; source_value < 4; source_value++) { + q_src[source_value] = (float)q[q_base + (long long)source_value]; + k_src[source_value] = (float)k[k_base + (long long)source_value]; + gate_src[source_value] = (float)g[gate_base + (long long)source_value]; + } + } + float q_sum_sq = q_src[0] * q_src[0] + q_src[1] * q_src[1] + (q_src[2] * q_src[2] + q_src[3] * q_src[3]); + float k_sum_sq = k_src[0] * k_src[0] + k_src[1] * k_src[1] + (k_src[2] * k_src[2] + k_src[3] * k_src[3]); + float _warp_reduce_0 = q_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset); + q_sum_sq = _warp_reduce_0; + float _warp_reduce_1 = k_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + k_sum_sq = _warp_reduce_1; + float _rsqrt_0 = rsqrtf(q_sum_sq + 1e-06f); + float q_scale = _rsqrt_0 * scale; + float _rsqrt_1 = rsqrtf(k_sum_sq + 1e-06f); + float k_scale = _rsqrt_1; + float _expf_0 = __expf(A_log[hv]); + float gate_a = _expf_0; + #pragma unroll + for (int source_value_1 = 0; source_value_1 < 4; source_value_1++) { + float biased_gate; + { + int k_idx = elem_start + source_value_1; + biased_gate = gate_src[source_value_1] + dt_bias[hv * 128 + k_idx]; + } + float _expf_1 = __expf((-gate_a) * biased_gate); + float gate_sigmoid = 1.0f / (1.0f + _expf_1); + float _expf_2 = __expf(-5.0f * gate_sigmoid); + gate_src[source_value_1] = _expf_2; + } + #pragma unroll + for (int i = 0; i < 8; i++) { + int source_lane = 2 * k_lane + i / 4; + float _shfl_0 = __shfl_sync(0xFFFFFFFF, q_src[i % 4], source_lane); + q_reg[i] = _shfl_0 * q_scale; + float _shfl_1 = __shfl_sync(0xFFFFFFFF, k_src[i % 4], source_lane); + k_reg[i] = _shfl_1 * k_scale; + float _shfl_2 = __shfl_sync(0xFFFFFFFF, gate_src[i % 4], source_lane); + decay_reg[i] = _shfl_2; + } + float k_dot_q = 0.0f; + #pragma unroll + for (int i_1 = 0; i_1 < 8; i_1++) { + k_dot_q += k_reg[i_1] * q_reg[i_1]; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 8); + k_dot_q += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 4); + k_dot_q += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 2); + k_dot_q += _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 1); + k_dot_q += _shfl_xor_3; + float beta_value = (float)beta[(long long)n * beta_stride_token + (long long)hv]; + if (beta_value >= 0.0f) { + float _expf_3 = __expf(-beta_value); + beta_value = 1.0f / (1.0f + _expf_3); + } else { + float _expf_4 = __expf(beta_value); + float beta_exp = _expf_4; + beta_value = beta_exp / (1.0f + beta_exp); + } + #pragma unroll + for (int chunk = 0; chunk < 2; chunk++) { + { + { + if (chunk < 1) { + asm volatile("cp.async.wait_group 1;"); + } else if (chunk == -1) { + asm volatile("cp.async.wait_group 2;"); + } else { + if (chunk == 0) { + asm volatile("cp.async.wait_group 1;"); + } else { + asm volatile("cp.async.wait_group 0;"); + } + } + } + } + { + if (chunk + 3 - 1 < 2) { + #pragma unroll + for (int private_row_1 = 0; private_row_1 < 4; private_row_1++) { + int copy_state_row_1 = group * 4 + private_row_1; + int copy_elem_1 = ((chunk + 3 - 1) * 32 + copy_state_row_1) * 128 + k_lane * 8; + int copy_dst_1 = state_smem_addr + (unsigned int)((chunk + 3 - 1) % 3 * 8192) + (unsigned int)((copy_state_row_1 * 128 + k_lane * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst_1), "l"(state + (state_head_base + (long long)copy_elem_1))); + } + asm volatile("cp.async.commit_group;"); + } + } + { + { + #pragma unroll + for (int row_local = 0; row_local < 4; row_local++) { + unsigned int packed_state[4]; + int smem_row = group * 4 + row_local; + int state_smem_addr_0 = 0; + { + state_smem_addr_0 = state_smem_addr + (unsigned int)(chunk % 3 * 8192) + (unsigned int)((smem_row * 128 + k_lane * 8) * 2); + } + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_state[0])), "=r"(*reinterpret_cast(&packed_state[(0) + 1])), "=r"(*reinterpret_cast(&packed_state[(0) + 2])), "=r"(*reinterpret_cast(&packed_state[(0) + 3])) + : "r"(state_smem_addr_0)); + float packed_state_f32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_state_f32[_pair * 2])[0]), "=f"((&packed_state_f32[_pair * 2])[1]) + : "r"(packed_state[_pair])); + } + #pragma unroll + for (int i_2 = 0; i_2 < 8; i_2++) { + state_rows[row_local * 8 + i_2] = packed_state_f32[i_2]; + } + } + } + #pragma unroll + for (int row_local_1 = 0; row_local_1 < 4; row_local_1++) { + int value_row = tile_row_base + chunk * 32 + group * 4 + row_local_1; + float prediction; + float base; + { + prediction = 0.0f; + base = 0.0f; + #pragma unroll + for (int i_3 = 0; i_3 < 8; i_3++) { + float decayed = state_rows[row_local_1 * 8 + i_3] * decay_reg[i_3]; + prediction += decayed * k_reg[i_3]; + base += decayed * q_reg[i_3]; + } + } + float _shfl_xor_20 = __shfl_xor_sync(0xFFFFFFFF, prediction, 8); + prediction += _shfl_xor_20; + float _shfl_xor_21 = __shfl_xor_sync(0xFFFFFFFF, prediction, 4); + prediction += _shfl_xor_21; + float _shfl_xor_22 = __shfl_xor_sync(0xFFFFFFFF, prediction, 2); + prediction += _shfl_xor_22; + float _shfl_xor_23 = __shfl_xor_sync(0xFFFFFFFF, prediction, 1); + prediction += _shfl_xor_23; + float _shfl_xor_24 = __shfl_xor_sync(0xFFFFFFFF, base, 8); + base += _shfl_xor_24; + float _shfl_xor_25 = __shfl_xor_sync(0xFFFFFFFF, base, 4); + base += _shfl_xor_25; + float _shfl_xor_26 = __shfl_xor_sync(0xFFFFFFFF, base, 2); + base += _shfl_xor_26; + float _shfl_xor_27 = __shfl_xor_sync(0xFFFFFFFF, base, 1); + base += _shfl_xor_27; + float v_value = 0.0f; + { + { + if (k_lane == 0) { + v_value = (float)v[(long long)n * v_stride_token + (long long)hv * 128 + (long long)value_row]; + } + } + } + { + float _shfl_5 = __shfl_sync(0xFFFFFFFF, v_value, group % 2 * 16); + v_value = _shfl_5; + } + float delta = (v_value - prediction) * beta_value; + { + #pragma unroll + for (int i_4 = 0; i_4 < 8; i_4++) { + state_rows[row_local_1 * 8 + i_4] = state_rows[row_local_1 * 8 + i_4] * decay_reg[i_4] + delta * k_reg[i_4]; + } + } + long long state_row_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)value_row * 128 + (long long)k_lane * 8; + { + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 0], state_rows[row_local_1 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 2], state_rows[row_local_1 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 4], state_rows[row_local_1 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(state_rows[row_local_1 * 8 + 6], state_rows[row_local_1 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(state))[state_row_base + 0]) = *reinterpret_cast(&_pk[0]); + } + } + { + if (k_lane == 0) { + out[output_head_base + value_row] = base + delta * k_dot_q; + } + } + } + } + } + } else if (k_lane == 0) { + #pragma unroll + for (int chunk_1 = 0; chunk_1 < 2; chunk_1++) { + #pragma unroll + for (int row_local_2 = 0; row_local_2 < 4; row_local_2++) { + int value_row_1 = tile_row_base + chunk_1 * 32 + group * 4 + row_local_2; + out[output_head_base + value_row_1] = 0.0f; + } + } + } +} + +} // extern "C" + +// END FROZEN GENERATED BODY +// clang-format on diff --git a/csrc/kda/cake_kda_packed_t1_cpasync_tile64_register_pipeline.cu b/csrc/kda/cake_kda_packed_t1_cpasync_tile64_register_pipeline.cu new file mode 100644 index 00000000000..ef66279e0af --- /dev/null +++ b/csrc/kda/cake_kda_packed_t1_cpasync_tile64_register_pipeline.cu @@ -0,0 +1,513 @@ +/* + * 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. + */ + +// clang-format off +// Generated by CAKE from the Kimi K3 serving-native packed T=1 schedule. +// BEGIN FROZEN GENERATED BODY +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +template +struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; + +typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_STATE_SMEM_OFF 0 +#define SMEM_STATE_SMEM_STAGE_BYTES 4096 +#define SMEM_STATE_SMEM_STRIDE 4096 +#define SMEM_V_SMEM_OFF 16128 +#define SMEM_V_SMEM_STAGE_BYTES 256 +#define SMEM_V_SMEM_STRIDE 256 +#define SMEM_TOTAL 16384 +#define THREADS 128 + +#include + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flashinfer_packed_kda_t1_cpasync_tile64_register_pipeline(__nv_bfloat16* __restrict__ q, __nv_bfloat16* __restrict__ k, __nv_bfloat16* __restrict__ v, __nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ beta, float* __restrict__ A_log, float* __restrict__ dt_bias, __nv_bfloat16* __restrict__ state, __nv_bfloat16* __restrict__ out, int* __restrict__ state_indices, float scale, long long q_stride_token, long long k_stride_token, long long v_stride_token, long long g_stride_token, long long beta_stride_token, long long state_stride_slot) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* state_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 0); + const int state_smem_addr = smem + 0; + __nv_bfloat16* v_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 16128); + const int v_smem_addr = smem + 16128; + + // === Task calls (dependency order) === + int tid_0 = tid; + int lane_1 = lane; + int k_lane = tid_0 % 16; + int group = tid_0 / 16; + int copy_row = tid_0 / 8; + int copy_col = tid_0 % 8 * 16; + int work = blockIdx.x; + int value_tile = work % 2; + int hv = work / 2; + int n = blockIdx.y; + int raw_slot = state_indices[n]; + bool active = raw_slot >= 0; + int output_head_base = (n * 12 + hv) * 128; + int tile_row_base = value_tile * 64; + if (active) { + long long state_head_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)tile_row_base * 128; + unsigned int v_registers[4]; + { + #pragma unroll + for (int v_chunk = 0; v_chunk < 4; v_chunk++) { + int v_row = tile_row_base + v_chunk * 16 + group * 2; + { + v_registers[v_chunk] = *reinterpret_cast(v + (long long)n * v_stride_token + (long long)hv * 128 + (long long)v_row); + } + } + } + #pragma unroll + for (int prefetch_chunk = 0; prefetch_chunk < 3; prefetch_chunk++) { + { + #pragma unroll + for (int private_row = 0; private_row < 2; private_row++) { + int copy_state_row = group * 2 + private_row; + int copy_elem = (prefetch_chunk * 16 + copy_state_row) * 128 + k_lane * 8; + int copy_dst = state_smem_addr + (unsigned int)(prefetch_chunk * 4096) + (unsigned int)((copy_state_row * 128 + k_lane * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst), "l"(state + (state_head_base + (long long)copy_elem))); + } + } + asm volatile("cp.async.commit_group;"); + } + int elem_start = lane_1 * 4; + long long q_base = (long long)n * q_stride_token + (long long)hv * 128 + (long long)elem_start; + long long k_base = (long long)n * k_stride_token + (long long)hv * 128 + (long long)elem_start; + long long gate_base = (long long)n * g_stride_token + (long long)hv * 128 + (long long)elem_start; + float q_src[4]; + float k_src[4]; + float gate_src[4]; + float dt_bias_src[4]; + float q_reg[8]; + float k_reg[8]; + float decay_reg[8]; + unsigned int packed_state_row[4]; + unsigned int lookahead_state[16]; + float state_row[8]; + float state_rows[16]; + unsigned int paired_state_bits[8]; + float paired_state_rows[16]; + { + { + uint2 _vld_1 = *reinterpret_cast(q + q_base); + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&q_src[0 + _pair * 2])[0]), "=f"((&q_src[0 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + { + uint2 _vld_2 = *reinterpret_cast(k + k_base); + uint32_t* _vpairs_2 = reinterpret_cast(&_vld_2); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&k_src[0 + _pair * 2])[0]), "=f"((&k_src[0 + _pair * 2])[1]) + : "r"(_vpairs_2[_pair])); + } + } + { + uint2 _vld_3 = *reinterpret_cast(g + gate_base); + uint32_t* _vpairs_3 = reinterpret_cast(&_vld_3); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&gate_src[0 + _pair * 2])[0]), "=f"((&gate_src[0 + _pair * 2])[1]) + : "r"(_vpairs_3[_pair])); + } + } + { + float4 _v4 = *reinterpret_cast(dt_bias + hv * 128 + elem_start); + dt_bias_src[0 + 0] = _v4.x; + dt_bias_src[0 + 1] = _v4.y; + dt_bias_src[0 + 2] = _v4.z; + dt_bias_src[0 + 3] = _v4.w; + } + } + float q_sum_sq = q_src[0] * q_src[0] + q_src[1] * q_src[1] + (q_src[2] * q_src[2] + q_src[3] * q_src[3]); + float k_sum_sq = k_src[0] * k_src[0] + k_src[1] * k_src[1] + (k_src[2] * k_src[2] + k_src[3] * k_src[3]); + float _warp_reduce_0 = q_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset); + q_sum_sq = _warp_reduce_0; + float _warp_reduce_1 = k_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + k_sum_sq = _warp_reduce_1; + float _rsqrt_0 = rsqrtf(q_sum_sq + 1e-06f); + float q_scale = _rsqrt_0 * scale; + float _rsqrt_1 = rsqrtf(k_sum_sq + 1e-06f); + float k_scale = _rsqrt_1; + float _expf_0 = __expf(A_log[hv]); + float gate_a = _expf_0; + #pragma unroll + for (int source_value = 0; source_value < 4; source_value++) { + float biased_gate; + { + biased_gate = gate_src[source_value] + dt_bias_src[source_value]; + } + float _expf_1 = __expf((-gate_a) * biased_gate); + float gate_sigmoid = 1.0f / (1.0f + _expf_1); + float _expf_2 = __expf(-5.0f * gate_sigmoid); + gate_src[source_value] = _expf_2; + } + #pragma unroll + for (int i = 0; i < 8; i++) { + int source_lane = 2 * k_lane + i / 4; + float _shfl_0 = __shfl_sync(0xFFFFFFFF, q_src[i % 4], source_lane); + q_reg[i] = _shfl_0 * q_scale; + float _shfl_1 = __shfl_sync(0xFFFFFFFF, k_src[i % 4], source_lane); + k_reg[i] = _shfl_1 * k_scale; + float _shfl_2 = __shfl_sync(0xFFFFFFFF, gate_src[i % 4], source_lane); + decay_reg[i] = _shfl_2; + } + float k_dot_q = 0.0f; + #pragma unroll + for (int i_1 = 0; i_1 < 8; i_1++) { + k_dot_q += k_reg[i_1] * q_reg[i_1]; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 8); + k_dot_q += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 4); + k_dot_q += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 2); + k_dot_q += _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, k_dot_q, 1); + k_dot_q += _shfl_xor_3; + float beta_value = (float)beta[(long long)n * beta_stride_token + (long long)hv]; + if (beta_value >= 0.0f) { + float _expf_3 = __expf(-beta_value); + beta_value = 1.0f / (1.0f + _expf_3); + } else { + float _expf_4 = __expf(beta_value); + float beta_exp = _expf_4; + beta_value = beta_exp / (1.0f + beta_exp); + } + #pragma unroll + for (int chunk = 0; chunk < 4; chunk++) { + { + if (chunk < 1) { + asm volatile("cp.async.wait_group 1;"); + } else if (chunk == 1) { + asm volatile("cp.async.wait_group 1;"); + } else { + if (chunk == 2) { + asm volatile("cp.async.wait_group 0;"); + } + } + } + { + if (chunk + 4 - 1 < 4) { + #pragma unroll + for (int private_row_1 = 0; private_row_1 < 2; private_row_1++) { + int copy_state_row_1 = group * 2 + private_row_1; + int copy_elem_1 = ((chunk + 4 - 1) * 16 + copy_state_row_1) * 128 + k_lane * 8; + int copy_dst_1 = state_smem_addr + (unsigned int)((chunk + 4 - 1) % 4 * 4096) + (unsigned int)((copy_state_row_1 * 128 + k_lane * 8) * 2); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" + :: "r"(copy_dst_1), "l"(state + (state_head_base + (long long)copy_elem_1))); + } + asm volatile("cp.async.commit_group;"); + } + } + { + if (chunk == 0) { + #pragma unroll + for (int preload_row = 0; preload_row < 2; preload_row++) { + int preload_smem_row = group * 2 + preload_row; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&lookahead_state[preload_row * 4])), "=r"(*reinterpret_cast(&lookahead_state[(preload_row * 4) + 1])), "=r"(*reinterpret_cast(&lookahead_state[(preload_row * 4) + 2])), "=r"(*reinterpret_cast(&lookahead_state[(preload_row * 4) + 3])) + : "r"(state_smem_addr + (unsigned int)((preload_smem_row * 128 + k_lane * 8) * 2))); + } + } + if (chunk + 1 < 4) { + int next_buffer = 1 - chunk % 2; + #pragma unroll + for (int preload_row_1 = 0; preload_row_1 < 2; preload_row_1++) { + int preload_smem_row_1 = group * 2 + preload_row_1; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&lookahead_state[(next_buffer * 2 + preload_row_1) * 4])), "=r"(*reinterpret_cast(&lookahead_state[((next_buffer * 2 + preload_row_1) * 4) + 1])), "=r"(*reinterpret_cast(&lookahead_state[((next_buffer * 2 + preload_row_1) * 4) + 2])), "=r"(*reinterpret_cast(&lookahead_state[((next_buffer * 2 + preload_row_1) * 4) + 3])) + : "r"(state_smem_addr + (unsigned int)((chunk + 1) % 4 * 4096) + (unsigned int)((preload_smem_row_1 * 128 + k_lane * 8) * 2))); + } + } + } + { + #pragma unroll + for (int row_local = 0; row_local < 2; row_local++) { + { + { + int current_buffer = chunk % 2; + #pragma unroll + for (int state_pair = 0; state_pair < 4; state_pair++) { + packed_state_row[state_pair] = lookahead_state[(current_buffer * 2 + row_local) * 4 + state_pair]; + } + } + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_row[_pair * 2])[0]), "=f"((&state_row[_pair * 2])[1]) + : "r"(packed_state_row[_pair])); + } + } + int value_row = tile_row_base + chunk * 16 + group * 2 + row_local; + float prediction; + float base; + { + float2 _f2_19 = make_float2(0.0f, 0.0f); + float2 prediction_pair = _f2_19; + float2 _f2_20 = make_float2(0.0f, 0.0f); + float2 base_pair = _f2_20; + #pragma unroll + for (int i_pair = 0; i_pair < 4; i_pair++) { + float2 _f2_21 = make_float2(state_row[i_pair * 2], state_row[i_pair * 2 + 1]); + float2 _f2_22 = make_float2(decay_reg[i_pair * 2], decay_reg[i_pair * 2 + 1]); + float2 _f2_23 = make_float2(0.0f, 0.0f); + float2 decayed_pair = fma_f32x2(_f2_21, _f2_22, _f2_23); + float2 _f2_24 = make_float2(k_reg[i_pair * 2], k_reg[i_pair * 2 + 1]); + prediction_pair = fma_f32x2(decayed_pair, _f2_24, prediction_pair); + float2 _f2_25 = make_float2(q_reg[i_pair * 2], q_reg[i_pair * 2 + 1]); + base_pair = fma_f32x2(decayed_pair, _f2_25, base_pair); + } + prediction = prediction_pair.x + prediction_pair.y; + base = base_pair.x + base_pair.y; + } + float _shfl_xor_20 = __shfl_xor_sync(0xFFFFFFFF, prediction, 8); + prediction += _shfl_xor_20; + float _shfl_xor_21 = __shfl_xor_sync(0xFFFFFFFF, prediction, 4); + prediction += _shfl_xor_21; + float _shfl_xor_22 = __shfl_xor_sync(0xFFFFFFFF, prediction, 2); + prediction += _shfl_xor_22; + float _shfl_xor_23 = __shfl_xor_sync(0xFFFFFFFF, prediction, 1); + prediction += _shfl_xor_23; + float _shfl_xor_24 = __shfl_xor_sync(0xFFFFFFFF, base, 8); + base += _shfl_xor_24; + float _shfl_xor_25 = __shfl_xor_sync(0xFFFFFFFF, base, 4); + base += _shfl_xor_25; + float _shfl_xor_26 = __shfl_xor_sync(0xFFFFFFFF, base, 2); + base += _shfl_xor_26; + float _shfl_xor_27 = __shfl_xor_sync(0xFFFFFFFF, base, 1); + base += _shfl_xor_27; + float v_value = 0.0f; + { + float v_registers_f32[2]; + #pragma unroll + for (int _pair = 0; _pair < 1; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&v_registers_f32[_pair * 2])[0]), "=f"((&v_registers_f32[_pair * 2])[1]) + : "r"((v_registers + chunk)[_pair])); + } + v_value = v_registers_f32[row_local]; + } + float delta = (v_value - prediction) * beta_value; + { + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_row[_pair * 2])[0]), "=f"((&state_row[_pair * 2])[1]) + : "r"(packed_state_row[_pair])); + } + #pragma unroll + for (int i_pair_1 = 0; i_pair_1 < 4; i_pair_1++) { + float2 _f2_26 = make_float2(k_reg[i_pair_1 * 2], k_reg[i_pair_1 * 2 + 1]); + float2 _f2_27 = make_float2(delta, delta); + float2 _f2_28 = make_float2(state_row[i_pair_1 * 2], state_row[i_pair_1 * 2 + 1]); + float2 _f2_29 = make_float2(decay_reg[i_pair_1 * 2], decay_reg[i_pair_1 * 2 + 1]); + float2 _f2_30 = make_float2(0.0f, 0.0f); + float2 updated_pair = fma_f32x2(_f2_26, _f2_27, fma_f32x2(_f2_28, _f2_29, _f2_30)); + state_row[i_pair_1 * 2] = updated_pair.x; + state_row[i_pair_1 * 2 + 1] = updated_pair.y; + } + } + long long state_row_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)value_row * 128 + (long long)k_lane * 8; + { + { + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(state_row[0 + 0], state_row[0 + 1]); + _pk[1] = __floats2bfloat162_rn(state_row[0 + 2], state_row[0 + 3]); + _pk[2] = __floats2bfloat162_rn(state_row[0 + 4], state_row[0 + 5]); + _pk[3] = __floats2bfloat162_rn(state_row[0 + 6], state_row[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(state))[state_row_base + 0]) = *reinterpret_cast(&_pk[0]); + } + } + } + { + if (k_lane == 0) { + out[output_head_base + value_row] = base + delta * k_dot_q; + } + } + } + } + } + } else if (k_lane == 0) { + #pragma unroll + for (int chunk_1 = 0; chunk_1 < 4; chunk_1++) { + #pragma unroll + for (int row_local_1 = 0; row_local_1 < 2; row_local_1++) { + int value_row_1 = tile_row_base + chunk_1 * 16 + group * 2 + row_local_1; + out[output_head_base + value_row_1] = 0.0f; + } + } + } +} + +} // extern "C" + +// END FROZEN GENERATED BODY +// clang-format on diff --git a/csrc/kda/cake_kda_packed_t1_register_tile16.cu b/csrc/kda/cake_kda_packed_t1_register_tile16.cu new file mode 100644 index 00000000000..4631c78512d --- /dev/null +++ b/csrc/kda/cake_kda_packed_t1_register_tile16.cu @@ -0,0 +1,388 @@ +/* + * 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. + */ + +// clang-format off +// Generated by CAKE from the Kimi K3 serving-native packed T=1 schedule. +// BEGIN FROZEN GENERATED BODY +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +template +struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; + +typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 128 + +#include + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flashinfer_packed_kda_t1_register_tile16(__nv_bfloat16* __restrict__ q, __nv_bfloat16* __restrict__ k, __nv_bfloat16* __restrict__ v, __nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ beta, float* __restrict__ A_log, float* __restrict__ dt_bias, __nv_bfloat16* __restrict__ state, __nv_bfloat16* __restrict__ out, int* __restrict__ state_indices, float scale, long long q_stride_token, long long k_stride_token, long long v_stride_token, long long g_stride_token, long long beta_stride_token, long long state_stride_slot) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int tid_0 = tid; + int lane_1 = lane; + int k_lane = tid_0 % 16; + int group = tid_0 / 16; + int work = blockIdx.x; + int value_tile = work % 8; + int hv = work / 8; + int n = blockIdx.y; + int raw_slot = state_indices[n]; + bool active = raw_slot >= 0; + int tile_row_base = value_tile * 16; + int group_row_base = tile_row_base + group * 2; + int output_head_base = (n * 12 + hv) * 128; + if (active) { + unsigned int state_bits[8]; + float v_values[2]; + #pragma unroll + for (int row_local = 0; row_local < 2; row_local++) { + int value_row = group_row_base + row_local; + long long state_row_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)value_row * 128 + (long long)k_lane * 8; + { + const uint4* _vptr_0 = reinterpret_cast(state + state_row_base); + uint4* _vdst_0 = reinterpret_cast(&state_bits[row_local * 4]); + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vdst_0[_blk] = _vptr_0[_blk]; + } + } + v_values[row_local] = 0.0f; + if (k_lane == 0) { + v_values[row_local] = (float)v[(long long)n * v_stride_token + (long long)hv * 128 + (long long)value_row]; + } + } + int elem_start = lane_1 * 4; + long long q_base = (long long)n * q_stride_token + (long long)hv * 128 + (long long)elem_start; + long long k_base = (long long)n * k_stride_token + (long long)hv * 128 + (long long)elem_start; + long long gate_base = (long long)n * g_stride_token + (long long)hv * 128 + (long long)elem_start; + float q_src[4]; + float k_src[4]; + float gate_src[4]; + float dt_bias_src[4]; + float q_reg[8]; + float k_reg[8]; + float decay_reg[8]; + { + uint2 _vld_1 = *reinterpret_cast(q + q_base); + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&q_src[0 + _pair * 2])[0]), "=f"((&q_src[0 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + { + uint2 _vld_2 = *reinterpret_cast(k + k_base); + uint32_t* _vpairs_2 = reinterpret_cast(&_vld_2); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&k_src[0 + _pair * 2])[0]), "=f"((&k_src[0 + _pair * 2])[1]) + : "r"(_vpairs_2[_pair])); + } + } + { + uint2 _vld_3 = *reinterpret_cast(g + gate_base); + uint32_t* _vpairs_3 = reinterpret_cast(&_vld_3); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&gate_src[0 + _pair * 2])[0]), "=f"((&gate_src[0 + _pair * 2])[1]) + : "r"(_vpairs_3[_pair])); + } + } + { + float4 _v4 = *reinterpret_cast(dt_bias + hv * 128 + elem_start); + dt_bias_src[0 + 0] = _v4.x; + dt_bias_src[0 + 1] = _v4.y; + dt_bias_src[0 + 2] = _v4.z; + dt_bias_src[0 + 3] = _v4.w; + } + float q_sum_sq = q_src[0] * q_src[0] + q_src[1] * q_src[1] + (q_src[2] * q_src[2] + q_src[3] * q_src[3]); + float k_sum_sq = k_src[0] * k_src[0] + k_src[1] * k_src[1] + (k_src[2] * k_src[2] + k_src[3] * k_src[3]); + float _warp_reduce_0 = q_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset); + q_sum_sq = _warp_reduce_0; + float _warp_reduce_1 = k_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + k_sum_sq = _warp_reduce_1; + float _rsqrt_0 = rsqrtf(q_sum_sq + 1e-06f); + float q_scale = _rsqrt_0 * scale; + float _rsqrt_1 = rsqrtf(k_sum_sq + 1e-06f); + float k_scale = _rsqrt_1; + float _expf_0 = __expf(A_log[hv]); + float gate_a = _expf_0; + #pragma unroll + for (int source_value = 0; source_value < 4; source_value++) { + float biased_gate = gate_src[source_value] + dt_bias_src[source_value]; + float _expf_1 = __expf((-gate_a) * biased_gate); + float gate_sigmoid = 1.0f / (1.0f + _expf_1); + float _expf_2 = __expf(-5.0f * gate_sigmoid); + gate_src[source_value] = _expf_2; + } + #pragma unroll + for (int i = 0; i < 8; i++) { + int source_lane = 2 * k_lane + i / 4; + float _shfl_0 = __shfl_sync(0xFFFFFFFF, q_src[i % 4], source_lane); + q_reg[i] = _shfl_0 * q_scale; + float _shfl_1 = __shfl_sync(0xFFFFFFFF, k_src[i % 4], source_lane); + k_reg[i] = _shfl_1 * k_scale; + float _shfl_2 = __shfl_sync(0xFFFFFFFF, gate_src[i % 4], source_lane); + decay_reg[i] = _shfl_2; + } + float beta_value = (float)beta[(long long)n * beta_stride_token + (long long)hv]; + if (beta_value >= 0.0f) { + float _expf_3 = __expf(-beta_value); + beta_value = 1.0f / (1.0f + _expf_3); + } else { + float _expf_4 = __expf(beta_value); + float beta_exp = _expf_4; + beta_value = beta_exp / (1.0f + beta_exp); + } + #pragma unroll + for (int row_local_1 = 0; row_local_1 < 2; row_local_1++) { + float2 _f2_0 = make_float2(0.0f, 0.0f); + float2 prediction_pair = _f2_0; + float state_bits_f32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_bits_f32[_pair * 2])[0]), "=f"((&state_bits_f32[_pair * 2])[1]) + : "r"((state_bits + row_local_1 * 4)[_pair])); + } + #pragma unroll + for (int pair = 0; pair < 4; pair++) { + float2 _f2_1 = make_float2(state_bits_f32[pair * 2], state_bits_f32[pair * 2 + 1]); + float2 state_pair = _f2_1; + float2 _f2_2 = make_float2(decay_reg[pair * 2], decay_reg[pair * 2 + 1]); + float2 decay_pair = _f2_2; + float2 _f2_3 = make_float2(k_reg[pair * 2], k_reg[pair * 2 + 1]); + float2 k_pair = _f2_3; + float2 decayed_pair = mul_f32x2(state_pair, decay_pair); + prediction_pair = fma_f32x2(decayed_pair, k_pair, prediction_pair); + } + float prediction = prediction_pair.x + prediction_pair.y; + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, prediction, 8); + prediction += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, prediction, 4); + prediction += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, prediction, 2); + prediction += _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, prediction, 1); + prediction += _shfl_xor_3; + float _shfl_3 = __shfl_sync(0xFFFFFFFF, v_values[row_local_1], group % 2 * 16); + float v_value = _shfl_3; + float delta = (v_value - prediction) * beta_value; + float updated[8]; + float2 _f2_4 = make_float2(0.0f, 0.0f); + float2 base_pair = _f2_4; + float state_bits_f32_0[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_bits_f32_0[_pair * 2])[0]), "=f"((&state_bits_f32_0[_pair * 2])[1]) + : "r"((state_bits + row_local_1 * 4)[_pair])); + } + #pragma unroll + for (int pair_1 = 0; pair_1 < 4; pair_1++) { + float2 _f2_5 = make_float2(state_bits_f32_0[pair_1 * 2], state_bits_f32_0[pair_1 * 2 + 1]); + float2 state_pair_1 = _f2_5; + float2 _f2_6 = make_float2(decay_reg[pair_1 * 2], decay_reg[pair_1 * 2 + 1]); + float2 decay_pair_1 = _f2_6; + float2 _f2_7 = make_float2(k_reg[pair_1 * 2], k_reg[pair_1 * 2 + 1]); + float2 k_pair_1 = _f2_7; + float2 _f2_8 = make_float2(q_reg[pair_1 * 2], q_reg[pair_1 * 2 + 1]); + float2 q_pair = _f2_8; + float2 _f2_9 = make_float2(delta, delta); + float2 delta_pair = _f2_9; + float2 decayed_pair_1 = mul_f32x2(state_pair_1, decay_pair_1); + float2 updated_pair = fma_f32x2(k_pair_1, delta_pair, decayed_pair_1); + base_pair = fma_f32x2(updated_pair, q_pair, base_pair); + updated[pair_1 * 2] = updated_pair.x; + updated[pair_1 * 2 + 1] = updated_pair.y; + } + int value_row_1 = group_row_base + row_local_1; + long long state_row_base_1 = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)value_row_1 * 128 + (long long)k_lane * 8; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(updated[0 + 0], updated[0 + 1]); + _pk[1] = __floats2bfloat162_rn(updated[0 + 2], updated[0 + 3]); + _pk[2] = __floats2bfloat162_rn(updated[0 + 4], updated[0 + 5]); + _pk[3] = __floats2bfloat162_rn(updated[0 + 6], updated[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(state))[state_row_base_1 + 0]) = *reinterpret_cast(&_pk[0]); + } + float base = base_pair.x + base_pair.y; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, base, 8); + base += _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, base, 4); + base += _shfl_xor_5; + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, base, 2); + base += _shfl_xor_6; + float _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, base, 1); + base += _shfl_xor_7; + if (k_lane == 0) { + out[output_head_base + value_row_1] = base; + } + } + } else if (k_lane == 0) { + #pragma unroll + for (int row_local_2 = 0; row_local_2 < 2; row_local_2++) { + out[output_head_base + group_row_base + row_local_2] = 0.0f; + } + } +} + +} // extern "C" + +// END FROZEN GENERATED BODY +// clang-format on diff --git a/csrc/kda/cake_kda_packed_t1_register_tile16_warp.cu b/csrc/kda/cake_kda_packed_t1_register_tile16_warp.cu new file mode 100644 index 00000000000..a347d927756 --- /dev/null +++ b/csrc/kda/cake_kda_packed_t1_register_tile16_warp.cu @@ -0,0 +1,388 @@ +/* + * 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. + */ + +// clang-format off +// Generated by CAKE from the Kimi K3 serving-native packed T=1 schedule. +// BEGIN FROZEN GENERATED BODY +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +template +struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; + +typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 32 + +#include + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +extern "C" { + +__global__ __launch_bounds__(32) void +kernel_flashinfer_packed_kda_t1_register_tile16_warp(__nv_bfloat16* __restrict__ q, __nv_bfloat16* __restrict__ k, __nv_bfloat16* __restrict__ v, __nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ beta, float* __restrict__ A_log, float* __restrict__ dt_bias, __nv_bfloat16* __restrict__ state, __nv_bfloat16* __restrict__ out, int* __restrict__ state_indices, float scale, long long q_stride_token, long long k_stride_token, long long v_stride_token, long long g_stride_token, long long beta_stride_token, long long state_stride_slot) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int tid_0 = tid; + int lane_1 = lane; + int k_lane = tid_0 % 16; + int group = tid_0 / 16; + int work = blockIdx.x; + int value_tile = work % 8; + int hv = work / 8; + int n = blockIdx.y; + int raw_slot = state_indices[n]; + bool active = raw_slot >= 0; + int tile_row_base = value_tile * 16; + int group_row_base = tile_row_base + group * 8; + int output_head_base = (n * 12 + hv) * 128; + if (active) { + unsigned int state_bits[32]; + float v_values[8]; + #pragma unroll + for (int row_local = 0; row_local < 8; row_local++) { + int value_row = group_row_base + row_local; + long long state_row_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)value_row * 128 + (long long)k_lane * 8; + { + const uint4* _vptr_0 = reinterpret_cast(state + state_row_base); + uint4* _vdst_0 = reinterpret_cast(&state_bits[row_local * 4]); + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vdst_0[_blk] = _vptr_0[_blk]; + } + } + v_values[row_local] = 0.0f; + if (k_lane == 0) { + v_values[row_local] = (float)v[(long long)n * v_stride_token + (long long)hv * 128 + (long long)value_row]; + } + } + int elem_start = lane_1 * 4; + long long q_base = (long long)n * q_stride_token + (long long)hv * 128 + (long long)elem_start; + long long k_base = (long long)n * k_stride_token + (long long)hv * 128 + (long long)elem_start; + long long gate_base = (long long)n * g_stride_token + (long long)hv * 128 + (long long)elem_start; + float q_src[4]; + float k_src[4]; + float gate_src[4]; + float dt_bias_src[4]; + float q_reg[8]; + float k_reg[8]; + float decay_reg[8]; + { + uint2 _vld_1 = *reinterpret_cast(q + q_base); + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&q_src[0 + _pair * 2])[0]), "=f"((&q_src[0 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + { + uint2 _vld_2 = *reinterpret_cast(k + k_base); + uint32_t* _vpairs_2 = reinterpret_cast(&_vld_2); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&k_src[0 + _pair * 2])[0]), "=f"((&k_src[0 + _pair * 2])[1]) + : "r"(_vpairs_2[_pair])); + } + } + { + uint2 _vld_3 = *reinterpret_cast(g + gate_base); + uint32_t* _vpairs_3 = reinterpret_cast(&_vld_3); + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&gate_src[0 + _pair * 2])[0]), "=f"((&gate_src[0 + _pair * 2])[1]) + : "r"(_vpairs_3[_pair])); + } + } + { + float4 _v4 = *reinterpret_cast(dt_bias + hv * 128 + elem_start); + dt_bias_src[0 + 0] = _v4.x; + dt_bias_src[0 + 1] = _v4.y; + dt_bias_src[0 + 2] = _v4.z; + dt_bias_src[0 + 3] = _v4.w; + } + float q_sum_sq = q_src[0] * q_src[0] + q_src[1] * q_src[1] + (q_src[2] * q_src[2] + q_src[3] * q_src[3]); + float k_sum_sq = k_src[0] * k_src[0] + k_src[1] * k_src[1] + (k_src[2] * k_src[2] + k_src[3] * k_src[3]); + float _warp_reduce_0 = q_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset); + q_sum_sq = _warp_reduce_0; + float _warp_reduce_1 = k_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + k_sum_sq = _warp_reduce_1; + float _rsqrt_0 = rsqrtf(q_sum_sq + 1e-06f); + float q_scale = _rsqrt_0 * scale; + float _rsqrt_1 = rsqrtf(k_sum_sq + 1e-06f); + float k_scale = _rsqrt_1; + float _expf_0 = __expf(A_log[hv]); + float gate_a = _expf_0; + #pragma unroll + for (int source_value = 0; source_value < 4; source_value++) { + float biased_gate = gate_src[source_value] + dt_bias_src[source_value]; + float _expf_1 = __expf((-gate_a) * biased_gate); + float gate_sigmoid = 1.0f / (1.0f + _expf_1); + float _expf_2 = __expf(-5.0f * gate_sigmoid); + gate_src[source_value] = _expf_2; + } + #pragma unroll + for (int i = 0; i < 8; i++) { + int source_lane = 2 * k_lane + i / 4; + float _shfl_0 = __shfl_sync(0xFFFFFFFF, q_src[i % 4], source_lane); + q_reg[i] = _shfl_0 * q_scale; + float _shfl_1 = __shfl_sync(0xFFFFFFFF, k_src[i % 4], source_lane); + k_reg[i] = _shfl_1 * k_scale; + float _shfl_2 = __shfl_sync(0xFFFFFFFF, gate_src[i % 4], source_lane); + decay_reg[i] = _shfl_2; + } + float beta_value = (float)beta[(long long)n * beta_stride_token + (long long)hv]; + if (beta_value >= 0.0f) { + float _expf_3 = __expf(-beta_value); + beta_value = 1.0f / (1.0f + _expf_3); + } else { + float _expf_4 = __expf(beta_value); + float beta_exp = _expf_4; + beta_value = beta_exp / (1.0f + beta_exp); + } + #pragma unroll + for (int row_local_1 = 0; row_local_1 < 8; row_local_1++) { + float2 _f2_0 = make_float2(0.0f, 0.0f); + float2 prediction_pair = _f2_0; + float state_bits_f32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_bits_f32[_pair * 2])[0]), "=f"((&state_bits_f32[_pair * 2])[1]) + : "r"((state_bits + row_local_1 * 4)[_pair])); + } + #pragma unroll + for (int pair = 0; pair < 4; pair++) { + float2 _f2_1 = make_float2(state_bits_f32[pair * 2], state_bits_f32[pair * 2 + 1]); + float2 state_pair = _f2_1; + float2 _f2_2 = make_float2(decay_reg[pair * 2], decay_reg[pair * 2 + 1]); + float2 decay_pair = _f2_2; + float2 _f2_3 = make_float2(k_reg[pair * 2], k_reg[pair * 2 + 1]); + float2 k_pair = _f2_3; + float2 decayed_pair = mul_f32x2(state_pair, decay_pair); + prediction_pair = fma_f32x2(decayed_pair, k_pair, prediction_pair); + } + float prediction = prediction_pair.x + prediction_pair.y; + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, prediction, 8); + prediction += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, prediction, 4); + prediction += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, prediction, 2); + prediction += _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, prediction, 1); + prediction += _shfl_xor_3; + float _shfl_3 = __shfl_sync(0xFFFFFFFF, v_values[row_local_1], group % 2 * 16); + float v_value = _shfl_3; + float delta = (v_value - prediction) * beta_value; + float updated[8]; + float2 _f2_4 = make_float2(0.0f, 0.0f); + float2 base_pair = _f2_4; + float state_bits_f32_0[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_bits_f32_0[_pair * 2])[0]), "=f"((&state_bits_f32_0[_pair * 2])[1]) + : "r"((state_bits + row_local_1 * 4)[_pair])); + } + #pragma unroll + for (int pair_1 = 0; pair_1 < 4; pair_1++) { + float2 _f2_5 = make_float2(state_bits_f32_0[pair_1 * 2], state_bits_f32_0[pair_1 * 2 + 1]); + float2 state_pair_1 = _f2_5; + float2 _f2_6 = make_float2(decay_reg[pair_1 * 2], decay_reg[pair_1 * 2 + 1]); + float2 decay_pair_1 = _f2_6; + float2 _f2_7 = make_float2(k_reg[pair_1 * 2], k_reg[pair_1 * 2 + 1]); + float2 k_pair_1 = _f2_7; + float2 _f2_8 = make_float2(q_reg[pair_1 * 2], q_reg[pair_1 * 2 + 1]); + float2 q_pair = _f2_8; + float2 _f2_9 = make_float2(delta, delta); + float2 delta_pair = _f2_9; + float2 decayed_pair_1 = mul_f32x2(state_pair_1, decay_pair_1); + float2 updated_pair = fma_f32x2(k_pair_1, delta_pair, decayed_pair_1); + base_pair = fma_f32x2(updated_pair, q_pair, base_pair); + updated[pair_1 * 2] = updated_pair.x; + updated[pair_1 * 2 + 1] = updated_pair.y; + } + int value_row_1 = group_row_base + row_local_1; + long long state_row_base_1 = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)value_row_1 * 128 + (long long)k_lane * 8; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(updated[0 + 0], updated[0 + 1]); + _pk[1] = __floats2bfloat162_rn(updated[0 + 2], updated[0 + 3]); + _pk[2] = __floats2bfloat162_rn(updated[0 + 4], updated[0 + 5]); + _pk[3] = __floats2bfloat162_rn(updated[0 + 6], updated[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(state))[state_row_base_1 + 0]) = *reinterpret_cast(&_pk[0]); + } + float base = base_pair.x + base_pair.y; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, base, 8); + base += _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, base, 4); + base += _shfl_xor_5; + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, base, 2); + base += _shfl_xor_6; + float _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, base, 1); + base += _shfl_xor_7; + if (k_lane == 0) { + out[output_head_base + value_row_1] = base; + } + } + } else if (k_lane == 0) { + #pragma unroll + for (int row_local_2 = 0; row_local_2 < 8; row_local_2++) { + out[output_head_base + group_row_base + row_local_2] = 0.0f; + } + } +} + +} // extern "C" + +// END FROZEN GENERATED BODY +// clang-format on diff --git a/csrc/kda/cake_kda_packed_t1_register_tile8_interleaved.cu b/csrc/kda/cake_kda_packed_t1_register_tile8_interleaved.cu new file mode 100644 index 00000000000..a1aedf1dbb6 --- /dev/null +++ b/csrc/kda/cake_kda_packed_t1_register_tile8_interleaved.cu @@ -0,0 +1,617 @@ +/* + * 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. + */ + +// clang-format off +// Generated by CAKE from the Kimi K3 serving-native packed T=1 schedule. +// BEGIN FROZEN GENERATED BODY +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(128) CakeTensorMap { uint64_t opaque[16]; }; +template +struct __align__(128) CakeTensorMapPack { CakeTensorMap maps[N]; }; + +typedef struct __align__(64) { uint64_t opaque[16]; } CUtensorMap; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 32 + +#include + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 fma_sub_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +extern "C" { + +__global__ __launch_bounds__(32) void +kernel_flashinfer_packed_kda_t1_register_tile8_interleaved(__nv_bfloat16* __restrict__ q, __nv_bfloat16* __restrict__ k, __nv_bfloat16* __restrict__ v, __nv_bfloat16* __restrict__ g, __nv_bfloat16* __restrict__ beta, float* __restrict__ A_log, float* __restrict__ dt_bias, __nv_bfloat16* __restrict__ state, __nv_bfloat16* __restrict__ out, int* __restrict__ state_indices, float scale, long long q_stride_token, long long k_stride_token, long long v_stride_token, long long g_stride_token, long long beta_stride_token, long long state_stride_slot) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int tid_0 = tid; + int lane_1 = lane; + int k_lane = tid_0 % 16; + int group = tid_0 / 16; + int work = blockIdx.x; + int value_tile = work % 16; + int hv = work / 16; + int n = blockIdx.y; + int raw_slot = state_indices[n]; + bool active = raw_slot >= 0; + int tile_row_base = value_tile * 8; + int group_row_base = tile_row_base + group * 4; + int output_head_base = (n * 12 + hv) * 128; + if (active) { + unsigned int state_bits[16]; + unsigned int v_bits[2]; + #pragma unroll + for (int row_local = 0; row_local < 4; row_local++) { + int value_row = group_row_base + row_local; + long long state_row_base = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)value_row * 128 + (long long)k_lane * 8; + { + const uint4* _vptr_0 = reinterpret_cast(state + state_row_base); + uint4* _vdst_0 = reinterpret_cast(&state_bits[row_local * 4]); + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vdst_0[_blk] = _vptr_0[_blk]; + } + } + } + long long v_base = (long long)n * v_stride_token + (long long)hv * 128 + (long long)group_row_base; + { + uint2 _vld_1 = *reinterpret_cast(v + v_base); + uint2* _vdst_1 = reinterpret_cast(&v_bits[0]); + *_vdst_1 = _vld_1; + } + int elem_start = lane_1 * 4; + long long q_base = (long long)n * q_stride_token + (long long)hv * 128 + (long long)elem_start; + long long k_base = (long long)n * k_stride_token + (long long)hv * 128 + (long long)elem_start; + long long gate_base = (long long)n * g_stride_token + (long long)hv * 128 + (long long)elem_start; + float q_src[4]; + float k_src[4]; + unsigned int q_bits[2]; + unsigned int k_bits[2]; + unsigned int gate_bits[2]; + float gate_src[4]; + float dt_bias_src[4]; + float q_reg[8]; + float k_reg[8]; + float decay_reg[8]; + { + uint2 _vld_2 = *reinterpret_cast(q + q_base); + uint2* _vdst_2 = reinterpret_cast(&q_bits[0]); + *_vdst_2 = _vld_2; + } + { + uint2 _vld_3 = *reinterpret_cast(k + k_base); + uint2* _vdst_3 = reinterpret_cast(&k_bits[0]); + *_vdst_3 = _vld_3; + } + { + uint2 _vld_4 = *reinterpret_cast(g + gate_base); + uint2* _vdst_4 = reinterpret_cast(&gate_bits[0]); + *_vdst_4 = _vld_4; + } + { + float4 _v4 = *reinterpret_cast(dt_bias + hv * 128 + elem_start); + dt_bias_src[0 + 0] = _v4.x; + dt_bias_src[0 + 1] = _v4.y; + dt_bias_src[0 + 2] = _v4.z; + dt_bias_src[0 + 3] = _v4.w; + } + float _bf16x2_dot_f32_0; + asm volatile( + "{\n\t" + ".reg .b16 a_lo, a_hi, b_lo, b_hi;\n\t" + "mov.b32 {a_lo, a_hi}, %1;\n\t" + "mov.b32 {b_lo, b_hi}, %2;\n\t" + "fma.rn.f32.bf16 %0, a_lo, b_lo, %3;\n\t" + "fma.rn.f32.bf16 %0, a_hi, b_hi, %0;\n\t" + "}\n" + : "=f"(_bf16x2_dot_f32_0) : "r"(q_bits[0]), "r"(q_bits[0]), "f"(0.0f)); + float q_sum_sq = _bf16x2_dot_f32_0; + float _bf16x2_dot_f32_1; + asm volatile( + "{\n\t" + ".reg .b16 a_lo, a_hi, b_lo, b_hi;\n\t" + "mov.b32 {a_lo, a_hi}, %1;\n\t" + "mov.b32 {b_lo, b_hi}, %2;\n\t" + "fma.rn.f32.bf16 %0, a_lo, b_lo, %3;\n\t" + "fma.rn.f32.bf16 %0, a_hi, b_hi, %0;\n\t" + "}\n" + : "=f"(_bf16x2_dot_f32_1) : "r"(q_bits[1]), "r"(q_bits[1]), "f"(q_sum_sq)); + q_sum_sq = _bf16x2_dot_f32_1; + float _bf16x2_dot_f32_2; + asm volatile( + "{\n\t" + ".reg .b16 a_lo, a_hi, b_lo, b_hi;\n\t" + "mov.b32 {a_lo, a_hi}, %1;\n\t" + "mov.b32 {b_lo, b_hi}, %2;\n\t" + "fma.rn.f32.bf16 %0, a_lo, b_lo, %3;\n\t" + "fma.rn.f32.bf16 %0, a_hi, b_hi, %0;\n\t" + "}\n" + : "=f"(_bf16x2_dot_f32_2) : "r"(k_bits[0]), "r"(k_bits[0]), "f"(0.0f)); + float k_sum_sq = _bf16x2_dot_f32_2; + float _bf16x2_dot_f32_3; + asm volatile( + "{\n\t" + ".reg .b16 a_lo, a_hi, b_lo, b_hi;\n\t" + "mov.b32 {a_lo, a_hi}, %1;\n\t" + "mov.b32 {b_lo, b_hi}, %2;\n\t" + "fma.rn.f32.bf16 %0, a_lo, b_lo, %3;\n\t" + "fma.rn.f32.bf16 %0, a_hi, b_hi, %0;\n\t" + "}\n" + : "=f"(_bf16x2_dot_f32_3) : "r"(k_bits[1]), "r"(k_bits[1]), "f"(k_sum_sq)); + k_sum_sq = _bf16x2_dot_f32_3; + float _warp_reduce_0 = q_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_0 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_0, offset); + q_sum_sq = _warp_reduce_0; + float _warp_reduce_1 = k_sum_sq; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + _warp_reduce_1 += __shfl_xor_sync(0xFFFFFFFF, _warp_reduce_1, offset); + k_sum_sq = _warp_reduce_1; + float _rsqrt_0 = rsqrtf(q_sum_sq + 1e-06f); + float q_scale = _rsqrt_0 * scale; + float _rsqrt_1 = rsqrtf(k_sum_sq + 1e-06f); + float k_scale = _rsqrt_1; + float _expf_0 = __expf(A_log[hv]); + float gate_a = _expf_0; + #pragma unroll + for (int source_pair = 0; source_pair < 2; source_pair++) { + float _bf16x2_add_f32_0[2]; + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %2;\n\t" + "add.rn.f32.bf16 %0, lo, %3;\n\t" + "add.rn.f32.bf16 %1, hi, %4;\n\t" + "}\n" + : "=f"(_bf16x2_add_f32_0[0]), "=f"(_bf16x2_add_f32_0[1]) : "r"(gate_bits[source_pair]), "f"(dt_bias_src[source_pair * 2]), "f"(dt_bias_src[source_pair * 2 + 1])); + gate_src[source_pair * 2] = _bf16x2_add_f32_0[0]; + gate_src[source_pair * 2 + 1] = _bf16x2_add_f32_0[1]; + } + #pragma unroll + for (int source_value = 0; source_value < 4; source_value++) { + float _expf_1 = __expf((-gate_a) * gate_src[source_value]); + float gate_sigmoid = 1.0f / (1.0f + _expf_1); + float _expf_2 = __expf(-5.0f * gate_sigmoid); + gate_src[source_value] = _expf_2; + } + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&q_src[_pair * 2])[0]), "=f"((&q_src[_pair * 2])[1]) + : "r"(q_bits[_pair])); + } + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&k_src[_pair * 2])[0]), "=f"((&k_src[_pair * 2])[1]) + : "r"(k_bits[_pair])); + } + #pragma unroll + for (int i = 0; i < 8; i++) { + int source_lane = 2 * k_lane + i / 4; + float _shfl_0 = __shfl_sync(0xFFFFFFFF, q_src[i % 4], source_lane); + q_reg[i] = _shfl_0 * q_scale; + float _shfl_1 = __shfl_sync(0xFFFFFFFF, k_src[i % 4], source_lane); + k_reg[i] = _shfl_1 * k_scale; + float _shfl_2 = __shfl_sync(0xFFFFFFFF, gate_src[i % 4], source_lane); + decay_reg[i] = _shfl_2; + } + float beta_value = (float)beta[(long long)n * beta_stride_token + (long long)hv]; + if (beta_value >= 0.0f) { + float _expf_3 = __expf(-beta_value); + beta_value = 1.0f / (1.0f + _expf_3); + } else { + float _expf_4 = __expf(beta_value); + float beta_exp = _expf_4; + beta_value = beta_exp / (1.0f + beta_exp); + } + float2 _f2_0 = make_float2(0.0f, 0.0f); + float2 prediction_pair0 = _f2_0; + float2 _f2_1 = make_float2(0.0f, 0.0f); + float2 prediction_pair1 = _f2_1; + float2 _f2_2 = make_float2(0.0f, 0.0f); + float2 prediction_pair2 = _f2_2; + float2 _f2_3 = make_float2(0.0f, 0.0f); + float2 prediction_pair3 = _f2_3; + #pragma unroll + for (int pair = 0; pair < 4; pair++) { + float2 _f2_4 = make_float2(decay_reg[pair * 2], decay_reg[pair * 2 + 1]); + float2 decay_pair = _f2_4; + float2 _f2_5 = make_float2(k_reg[pair * 2], k_reg[pair * 2 + 1]); + float2 k_pair = _f2_5; + float state_bits_f32[2]; + #pragma unroll + for (int _pair = 0; _pair < 1; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_bits_f32[_pair * 2])[0]), "=f"((&state_bits_f32[_pair * 2])[1]) + : "r"((state_bits + pair)[_pair])); + } + float state_bits_f32_0[2]; + #pragma unroll + for (int _pair = 0; _pair < 1; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_bits_f32_0[_pair * 2])[0]), "=f"((&state_bits_f32_0[_pair * 2])[1]) + : "r"((state_bits + 4 + pair)[_pair])); + } + float state_bits_f32_1[2]; + #pragma unroll + for (int _pair = 0; _pair < 1; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_bits_f32_1[_pair * 2])[0]), "=f"((&state_bits_f32_1[_pair * 2])[1]) + : "r"((state_bits + 8 + pair)[_pair])); + } + float state_bits_f32_2[2]; + #pragma unroll + for (int _pair = 0; _pair < 1; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_bits_f32_2[_pair * 2])[0]), "=f"((&state_bits_f32_2[_pair * 2])[1]) + : "r"((state_bits + 12 + pair)[_pair])); + } + float2 _f2_6 = make_float2(state_bits_f32[0], state_bits_f32[1]); + float2 decayed0 = mul_f32x2(_f2_6, decay_pair); + float2 _f2_7 = make_float2(state_bits_f32_0[0], state_bits_f32_0[1]); + float2 decayed1 = mul_f32x2(_f2_7, decay_pair); + float2 _f2_8 = make_float2(state_bits_f32_1[0], state_bits_f32_1[1]); + float2 decayed2 = mul_f32x2(_f2_8, decay_pair); + float2 _f2_9 = make_float2(state_bits_f32_2[0], state_bits_f32_2[1]); + float2 decayed3 = mul_f32x2(_f2_9, decay_pair); + prediction_pair0 = fma_f32x2(decayed0, k_pair, prediction_pair0); + prediction_pair1 = fma_f32x2(decayed1, k_pair, prediction_pair1); + prediction_pair2 = fma_f32x2(decayed2, k_pair, prediction_pair2); + prediction_pair3 = fma_f32x2(decayed3, k_pair, prediction_pair3); + } + float prediction0 = prediction_pair0.x + prediction_pair0.y; + float prediction1 = prediction_pair1.x + prediction_pair1.y; + float prediction2 = prediction_pair2.x + prediction_pair2.y; + float prediction3 = prediction_pair3.x + prediction_pair3.y; + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, prediction0, 8); + prediction0 += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, prediction1, 8); + prediction1 += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, prediction2, 8); + prediction2 += _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, prediction3, 8); + prediction3 += _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, prediction0, 4); + prediction0 += _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, prediction1, 4); + prediction1 += _shfl_xor_5; + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, prediction2, 4); + prediction2 += _shfl_xor_6; + float _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, prediction3, 4); + prediction3 += _shfl_xor_7; + float _shfl_xor_8 = __shfl_xor_sync(0xFFFFFFFF, prediction0, 2); + prediction0 += _shfl_xor_8; + float _shfl_xor_9 = __shfl_xor_sync(0xFFFFFFFF, prediction1, 2); + prediction1 += _shfl_xor_9; + float _shfl_xor_10 = __shfl_xor_sync(0xFFFFFFFF, prediction2, 2); + prediction2 += _shfl_xor_10; + float _shfl_xor_11 = __shfl_xor_sync(0xFFFFFFFF, prediction3, 2); + prediction3 += _shfl_xor_11; + float _shfl_xor_12 = __shfl_xor_sync(0xFFFFFFFF, prediction0, 1); + prediction0 += _shfl_xor_12; + float _shfl_xor_13 = __shfl_xor_sync(0xFFFFFFFF, prediction1, 1); + prediction1 += _shfl_xor_13; + float _shfl_xor_14 = __shfl_xor_sync(0xFFFFFFFF, prediction2, 1); + prediction2 += _shfl_xor_14; + float _shfl_xor_15 = __shfl_xor_sync(0xFFFFFFFF, prediction3, 1); + prediction3 += _shfl_xor_15; + float _bf16x2_add_f32_1[2]; + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %2;\n\t" + "add.rn.f32.bf16 %0, lo, %3;\n\t" + "add.rn.f32.bf16 %1, hi, %4;\n\t" + "}\n" + : "=f"(_bf16x2_add_f32_1[0]), "=f"(_bf16x2_add_f32_1[1]) : "r"(v_bits[0]), "f"(-prediction0), "f"(-prediction1)); + float _bf16x2_add_f32_2[2]; + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %2;\n\t" + "add.rn.f32.bf16 %0, lo, %3;\n\t" + "add.rn.f32.bf16 %1, hi, %4;\n\t" + "}\n" + : "=f"(_bf16x2_add_f32_2[0]), "=f"(_bf16x2_add_f32_2[1]) : "r"(v_bits[1]), "f"(-prediction2), "f"(-prediction3)); + float delta0 = _bf16x2_add_f32_1[0] * beta_value; + float delta1 = _bf16x2_add_f32_1[1] * beta_value; + float delta2 = _bf16x2_add_f32_2[0] * beta_value; + float delta3 = _bf16x2_add_f32_2[1] * beta_value; + float2 _f2_10 = make_float2(0.0f, 0.0f); + float2 base_pair0 = _f2_10; + float2 _f2_11 = make_float2(0.0f, 0.0f); + float2 base_pair1 = _f2_11; + float2 _f2_12 = make_float2(0.0f, 0.0f); + float2 base_pair2 = _f2_12; + float2 _f2_13 = make_float2(0.0f, 0.0f); + float2 base_pair3 = _f2_13; + #pragma unroll + for (int pair_1 = 0; pair_1 < 4; pair_1++) { + float2 _f2_14 = make_float2(decay_reg[pair_1 * 2], decay_reg[pair_1 * 2 + 1]); + float2 decay_pair_1 = _f2_14; + float2 _f2_15 = make_float2(k_reg[pair_1 * 2], k_reg[pair_1 * 2 + 1]); + float2 k_pair_1 = _f2_15; + float2 _f2_16 = make_float2(q_reg[pair_1 * 2], q_reg[pair_1 * 2 + 1]); + float2 q_pair = _f2_16; + float state_bits_f32_3[2]; + #pragma unroll + for (int _pair = 0; _pair < 1; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_bits_f32_3[_pair * 2])[0]), "=f"((&state_bits_f32_3[_pair * 2])[1]) + : "r"((state_bits + pair_1)[_pair])); + } + float state_bits_f32_0_1[2]; + #pragma unroll + for (int _pair = 0; _pair < 1; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_bits_f32_0_1[_pair * 2])[0]), "=f"((&state_bits_f32_0_1[_pair * 2])[1]) + : "r"((state_bits + 4 + pair_1)[_pair])); + } + float state_bits_f32_1_1[2]; + #pragma unroll + for (int _pair = 0; _pair < 1; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_bits_f32_1_1[_pair * 2])[0]), "=f"((&state_bits_f32_1_1[_pair * 2])[1]) + : "r"((state_bits + 8 + pair_1)[_pair])); + } + float state_bits_f32_2_1[2]; + #pragma unroll + for (int _pair = 0; _pair < 1; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_bits_f32_2_1[_pair * 2])[0]), "=f"((&state_bits_f32_2_1[_pair * 2])[1]) + : "r"((state_bits + 12 + pair_1)[_pair])); + } + float2 _f2_17 = make_float2(state_bits_f32_3[0], state_bits_f32_3[1]); + float2 decayed0_1 = mul_f32x2(_f2_17, decay_pair_1); + float2 _f2_18 = make_float2(state_bits_f32_0_1[0], state_bits_f32_0_1[1]); + float2 decayed1_1 = mul_f32x2(_f2_18, decay_pair_1); + float2 _f2_19 = make_float2(state_bits_f32_1_1[0], state_bits_f32_1_1[1]); + float2 decayed2_1 = mul_f32x2(_f2_19, decay_pair_1); + float2 _f2_20 = make_float2(state_bits_f32_2_1[0], state_bits_f32_2_1[1]); + float2 decayed3_1 = mul_f32x2(_f2_20, decay_pair_1); + float2 _f2_21 = make_float2(delta0, delta0); + float2 updated0 = fma_f32x2(k_pair_1, _f2_21, decayed0_1); + float2 _f2_22 = make_float2(delta1, delta1); + float2 updated1 = fma_f32x2(k_pair_1, _f2_22, decayed1_1); + float2 _f2_23 = make_float2(delta2, delta2); + float2 updated2 = fma_f32x2(k_pair_1, _f2_23, decayed2_1); + float2 _f2_24 = make_float2(delta3, delta3); + float2 updated3 = fma_f32x2(k_pair_1, _f2_24, decayed3_1); + base_pair0 = fma_f32x2(updated0, q_pair, base_pair0); + base_pair1 = fma_f32x2(updated1, q_pair, base_pair1); + base_pair2 = fma_f32x2(updated2, q_pair, base_pair2); + base_pair3 = fma_f32x2(updated3, q_pair, base_pair3); + __nv_bfloat162 _bf16x2_0 = __float22bfloat162_rn(make_float2(updated0.x, updated0.y)); + __nv_bfloat162 _bf16x2_1 = __float22bfloat162_rn(make_float2(updated1.x, updated1.y)); + __nv_bfloat162 _bf16x2_2 = __float22bfloat162_rn(make_float2(updated2.x, updated2.y)); + __nv_bfloat162 _bf16x2_3 = __float22bfloat162_rn(make_float2(updated3.x, updated3.y)); + state_bits[pair_1] = reinterpret_cast(&_bf16x2_0)[0]; + state_bits[4 + pair_1] = reinterpret_cast(&_bf16x2_1)[0]; + state_bits[8 + pair_1] = reinterpret_cast(&_bf16x2_2)[0]; + state_bits[12 + pair_1] = reinterpret_cast(&_bf16x2_3)[0]; + } + int value_row0 = group_row_base; + int value_row1 = group_row_base + 1; + int value_row2 = group_row_base + 2; + int value_row3 = group_row_base + 3; + long long state_base_common = (long long)raw_slot * state_stride_slot + (long long)hv * 128 * 128 + (long long)k_lane * 8; + reinterpret_cast(state + (state_base_common + (long long)value_row0 * 128))[0] = reinterpret_cast(state_bits + 0)[0]; + reinterpret_cast(state + (state_base_common + (long long)value_row1 * 128))[0] = reinterpret_cast(state_bits + 4)[0]; + reinterpret_cast(state + (state_base_common + (long long)value_row2 * 128))[0] = reinterpret_cast(state_bits + 8)[0]; + reinterpret_cast(state + (state_base_common + (long long)value_row3 * 128))[0] = reinterpret_cast(state_bits + 12)[0]; + float base0 = base_pair0.x + base_pair0.y; + float base1 = base_pair1.x + base_pair1.y; + float base2 = base_pair2.x + base_pair2.y; + float base3 = base_pair3.x + base_pair3.y; + float _shfl_xor_16 = __shfl_xor_sync(0xFFFFFFFF, base0, 8); + base0 += _shfl_xor_16; + float _shfl_xor_17 = __shfl_xor_sync(0xFFFFFFFF, base1, 8); + base1 += _shfl_xor_17; + float _shfl_xor_18 = __shfl_xor_sync(0xFFFFFFFF, base2, 8); + base2 += _shfl_xor_18; + float _shfl_xor_19 = __shfl_xor_sync(0xFFFFFFFF, base3, 8); + base3 += _shfl_xor_19; + float _shfl_xor_20 = __shfl_xor_sync(0xFFFFFFFF, base0, 4); + base0 += _shfl_xor_20; + float _shfl_xor_21 = __shfl_xor_sync(0xFFFFFFFF, base1, 4); + base1 += _shfl_xor_21; + float _shfl_xor_22 = __shfl_xor_sync(0xFFFFFFFF, base2, 4); + base2 += _shfl_xor_22; + float _shfl_xor_23 = __shfl_xor_sync(0xFFFFFFFF, base3, 4); + base3 += _shfl_xor_23; + float _shfl_xor_24 = __shfl_xor_sync(0xFFFFFFFF, base0, 2); + base0 += _shfl_xor_24; + float _shfl_xor_25 = __shfl_xor_sync(0xFFFFFFFF, base1, 2); + base1 += _shfl_xor_25; + float _shfl_xor_26 = __shfl_xor_sync(0xFFFFFFFF, base2, 2); + base2 += _shfl_xor_26; + float _shfl_xor_27 = __shfl_xor_sync(0xFFFFFFFF, base3, 2); + base3 += _shfl_xor_27; + float _shfl_xor_28 = __shfl_xor_sync(0xFFFFFFFF, base0, 1); + base0 += _shfl_xor_28; + float _shfl_xor_29 = __shfl_xor_sync(0xFFFFFFFF, base1, 1); + base1 += _shfl_xor_29; + float _shfl_xor_30 = __shfl_xor_sync(0xFFFFFFFF, base2, 1); + base2 += _shfl_xor_30; + float _shfl_xor_31 = __shfl_xor_sync(0xFFFFFFFF, base3, 1); + base3 += _shfl_xor_31; + if (k_lane == 0) { + out[output_head_base + value_row0] = base0; + out[output_head_base + value_row1] = base1; + out[output_head_base + value_row2] = base2; + out[output_head_base + value_row3] = base3; + } + } else if (k_lane == 0) { + #pragma unroll + for (int row_local_1 = 0; row_local_1 < 4; row_local_1++) { + out[output_head_base + group_row_base + row_local_1] = 0.0f; + } + } +} + +} // extern "C" + +// END FROZEN GENERATED BODY +// clang-format on diff --git a/flashinfer/aot.py b/flashinfer/aot.py index e49c4567329..ad934f6ad7c 100644 --- a/flashinfer/aot.py +++ b/flashinfer/aot.py @@ -77,6 +77,10 @@ FLASH_KDA_PACKED_T1_VARIANTS, gen_flash_kda_packed_t1_module, ) +from .jit.cake_kda_packed_t1 import ( + CAKE_KDA_PACKED_T1_VARIANTS, + gen_cake_kda_packed_t1_module, +) from .jit.nvfp4_attention_sm120 import gen_nvfp4_attention_sm120_module from .jit.fp8_quantization import gen_mxfp8_quantization_sm100_module from .jit.fused_moe import ( @@ -602,11 +606,19 @@ def gen_all_modules( gen_flash_kda_packed_t1_module(variant, "sm100a") for variant in FLASH_KDA_PACKED_T1_VARIANTS ) + jit_specs.extend( + gen_cake_kda_packed_t1_module(variant, "sm100a") + for variant in CAKE_KDA_PACKED_T1_VARIANTS + ) if has_flash_kda_packed_t1_sm100f: jit_specs.extend( gen_flash_kda_packed_t1_module(variant, "sm100f") for variant in FLASH_KDA_PACKED_T1_VARIANTS ) + jit_specs.extend( + gen_cake_kda_packed_t1_module(variant, "sm100f") + for variant in CAKE_KDA_PACKED_T1_VARIANTS + ) if add_act: for act_name in act_func_def_str: diff --git a/flashinfer/jit/cake_kda_packed_t1.py b/flashinfer/jit/cake_kda_packed_t1.py new file mode 100644 index 00000000000..143ee95a675 --- /dev/null +++ b/flashinfer/jit/cake_kda_packed_t1.py @@ -0,0 +1,347 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +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 functools +from pathlib import Path +from typing import Literal, NamedTuple, Optional + +from . import env as jit_env +from .core import ( + JitSpec, + gen_jit_spec, + logger, + sm100a_nvcc_flags, + sm100f_nvcc_flags, +) +from .utils import write_if_different + +CakeKDAPackedT1Variant = Literal[ + "register_tile16", + "register_tile8_interleaved", + "register_tile16_warp", + "cpasync_tile64_ilp4", + "cpasync_tile64", + "cpasync_tile128_ilp4", + "cpasync_tile64_register_pipeline", + "cpasync_tile128_packed_state_v_private_prefetch", + "cpasync_tile128_v_private_prefetch", + "cpasync_tile128_paired_row_pipeline", + "cpasync_tile128_register_pipeline", + "cpasync_tile128_ilp2", +] +CakeKDAPackedT1Target = Literal["sm100a", "sm100f"] + +CAKE_KDA_PACKED_T1_VARIANTS: tuple[CakeKDAPackedT1Variant, ...] = ( + "register_tile16", + "register_tile8_interleaved", + "register_tile16_warp", + "cpasync_tile64_ilp4", + "cpasync_tile64", + "cpasync_tile128_ilp4", + "cpasync_tile64_register_pipeline", + "cpasync_tile128_packed_state_v_private_prefetch", + "cpasync_tile128_v_private_prefetch", + "cpasync_tile128_paired_row_pipeline", + "cpasync_tile128_register_pipeline", + "cpasync_tile128_ilp2", +) + +_CAKE_KDA_PACKED_T1_NVCC_FLAGS = { + "sm100a": sm100a_nvcc_flags, + "sm100f": sm100f_nvcc_flags, +} +_CAKE_KDA_PACKED_T1_TARGET_KIND = {"sm100a": 1000, "sm100f": 100} + + +class CakeKDAPackedT1VariantMetadata(NamedTuple): + body: str + symbol: str + value_tiles: int + threads: int + smem_bytes: int + requires_aux_vec4: bool + extra_cuda_flags: tuple[str, ...] = ("--maxrregcount=128",) + + +CAKE_KDA_PACKED_T1_VARIANT_METADATA: dict[ + CakeKDAPackedT1Variant, CakeKDAPackedT1VariantMetadata +] = { + "register_tile16": CakeKDAPackedT1VariantMetadata( + body="cake_kda_packed_t1_register_tile16.cu", + symbol="kernel_flashinfer_packed_kda_t1_register_tile16", + value_tiles=8, + threads=128, + smem_bytes=0, + requires_aux_vec4=True, + ), + "register_tile8_interleaved": CakeKDAPackedT1VariantMetadata( + body="cake_kda_packed_t1_register_tile8_interleaved.cu", + symbol="kernel_flashinfer_packed_kda_t1_register_tile8_interleaved", + value_tiles=16, + threads=32, + smem_bytes=0, + requires_aux_vec4=True, + extra_cuda_flags=("--ftz=false", "--maxrregcount=128"), + ), + "register_tile16_warp": CakeKDAPackedT1VariantMetadata( + body="cake_kda_packed_t1_register_tile16_warp.cu", + symbol="kernel_flashinfer_packed_kda_t1_register_tile16_warp", + value_tiles=8, + threads=32, + smem_bytes=0, + requires_aux_vec4=True, + ), + "cpasync_tile64_ilp4": CakeKDAPackedT1VariantMetadata( + body="cake_kda_packed_t1_cpasync_tile64_ilp4.cu", + symbol="kernel_flashinfer_packed_kda_t1_cpasync_tile64_ilp4", + value_tiles=2, + threads=128, + smem_bytes=24576, + requires_aux_vec4=False, + ), + "cpasync_tile64": CakeKDAPackedT1VariantMetadata( + body="cake_kda_packed_t1_cpasync_tile64.cu", + symbol="kernel_flashinfer_packed_kda_t1_cpasync", + value_tiles=2, + threads=128, + smem_bytes=16384, + requires_aux_vec4=False, + ), + "cpasync_tile128_ilp4": CakeKDAPackedT1VariantMetadata( + body="cake_kda_packed_t1_cpasync_tile128_ilp4.cu", + symbol="kernel_flashinfer_packed_kda_t1_cpasync_tile128_ilp4", + value_tiles=1, + threads=128, + smem_bytes=24576, + requires_aux_vec4=False, + ), + "cpasync_tile64_register_pipeline": CakeKDAPackedT1VariantMetadata( + body="cake_kda_packed_t1_cpasync_tile64_register_pipeline.cu", + symbol="kernel_flashinfer_packed_kda_t1_cpasync_tile64_register_pipeline", + value_tiles=2, + threads=128, + smem_bytes=16384, + requires_aux_vec4=True, + ), + "cpasync_tile128_packed_state_v_private_prefetch": CakeKDAPackedT1VariantMetadata( + body="cake_kda_packed_t1_cpasync_tile128_packed_state_v_private_prefetch.cu", + symbol=( + "kernel_flashinfer_packed_kda_t1_cpasync_" + "tile128_packed_state_v_private_prefetch" + ), + value_tiles=1, + threads=128, + smem_bytes=20736, + requires_aux_vec4=True, + ), + "cpasync_tile128_v_private_prefetch": CakeKDAPackedT1VariantMetadata( + body="cake_kda_packed_t1_cpasync_tile128_v_private_prefetch.cu", + symbol="kernel_flashinfer_packed_kda_t1_cpasync_tile128_v_private_prefetch", + value_tiles=1, + threads=128, + smem_bytes=20736, + requires_aux_vec4=True, + ), + "cpasync_tile128_paired_row_pipeline": CakeKDAPackedT1VariantMetadata( + body="cake_kda_packed_t1_cpasync_tile128_paired_row_pipeline.cu", + symbol="kernel_flashinfer_packed_kda_t1_cpasync_tile128_paired_row_pipeline", + value_tiles=1, + threads=128, + smem_bytes=20736, + requires_aux_vec4=True, + ), + "cpasync_tile128_register_pipeline": CakeKDAPackedT1VariantMetadata( + body="cake_kda_packed_t1_cpasync_tile128_register_pipeline.cu", + symbol="kernel_flashinfer_packed_kda_t1_cpasync_tile128_register_pipeline", + value_tiles=1, + threads=128, + smem_bytes=20480, + requires_aux_vec4=True, + ), + "cpasync_tile128_ilp2": CakeKDAPackedT1VariantMetadata( + body="cake_kda_packed_t1_cpasync_tile128_ilp2.cu", + symbol="kernel_flashinfer_packed_kda_t1_cpasync_tile128", + value_tiles=1, + threads=128, + smem_bytes=20480, + requires_aux_vec4=True, + ), +} + + +def select_cake_kda_packed_t1_variant( + batch: int, + *, + state_aligned: bool, + aux_vec4_aligned: bool, +) -> Optional[CakeKDAPackedT1Variant]: + """Return the qualified final selector, or ``None`` for the legacy route.""" + + if batch <= 0: + raise ValueError(f"packed KDA T=1 batch must be positive, got {batch}") + if not state_aligned: + return None + if aux_vec4_aligned: + if batch <= 14: + return "register_tile16" + if batch <= 29: + return "register_tile8_interleaved" + if batch <= 38: + return "register_tile16_warp" + if batch <= 41: + return "cpasync_tile64_register_pipeline" + if batch <= 80: + return "cpasync_tile128_packed_state_v_private_prefetch" + if batch <= 101: + return "cpasync_tile128_v_private_prefetch" + if batch <= 152: + return "cpasync_tile128_paired_row_pipeline" + return "cpasync_tile128_register_pipeline" + if batch <= 24: + return "cpasync_tile64_ilp4" + if batch <= 37: + return "cpasync_tile64" + if batch == 38: + return "cpasync_tile128_ilp4" + return None + + +def _get_csrc_dir() -> Path: + installed = jit_env.FLASHINFER_CSRC_DIR / "kda" + if installed.exists(): + return installed + checkout = Path(__file__).resolve().parents[2] / "csrc" / "kda" + if checkout.exists(): + return checkout + raise FileNotFoundError( + "frozen Cake KDA packed T=1 sources were not found. Checked:\n" + f" - {installed}\n - {checkout}" + ) + + +def _get_include_dir() -> Path: + if jit_env.FLASHINFER_INCLUDE_DIR.exists(): + return jit_env.FLASHINFER_INCLUDE_DIR + checkout = Path(__file__).resolve().parents[2] / "include" + if checkout.exists(): + return checkout + raise FileNotFoundError( + "FlashInfer headers were not found. Checked:\n" + f" - {jit_env.FLASHINFER_INCLUDE_DIR}\n - {checkout}" + ) + + +def get_cake_kda_packed_t1_uri( + variant: CakeKDAPackedT1Variant, + target: CakeKDAPackedT1Target, +) -> str: + if variant not in CAKE_KDA_PACKED_T1_VARIANTS: + raise ValueError(f"unsupported Cake KDA packed T=1 variant: {variant}") + if target not in _CAKE_KDA_PACKED_T1_NVCC_FLAGS: + raise ValueError(f"unsupported Cake KDA packed T=1 target: {target}") + return f"cake_kda_packed_t1_{variant}_{target}" + + +def _get_binding_cu(metadata: CakeKDAPackedT1VariantMetadata) -> str: + return f"""\ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * Licensed under the Apache License, Version 2.0. + */ + +#define CAKE_KDA_PACKED_T1_BODY_FILE "{metadata.body}" +#define CAKE_KDA_PACKED_T1_KERNEL {metadata.symbol} +#define CAKE_KDA_PACKED_T1_VALUE_TILES {metadata.value_tiles} +#define CAKE_KDA_PACKED_T1_THREADS {metadata.threads} +#define CAKE_KDA_PACKED_T1_SMEM_BYTES {metadata.smem_bytes} +#define CAKE_KDA_PACKED_T1_REQUIRES_AUX_VEC4 {int(metadata.requires_aux_vec4)} + +#include "cake_kda_packed_t1_binding.cuh" +""" + + +@functools.cache +def gen_cake_kda_packed_t1_module( + variant: CakeKDAPackedT1Variant, + target: CakeKDAPackedT1Target, +) -> JitSpec: + if variant not in CAKE_KDA_PACKED_T1_VARIANTS: + raise ValueError(f"unsupported Cake KDA packed T=1 variant: {variant}") + if target not in _CAKE_KDA_PACKED_T1_NVCC_FLAGS: + raise ValueError(f"unsupported Cake KDA packed T=1 target: {target}") + + csrc_dir = _get_csrc_dir() + metadata = CAKE_KDA_PACKED_T1_VARIANT_METADATA[variant] + body = csrc_dir / metadata.body + if not body.exists(): + raise FileNotFoundError(f"frozen Cake KDA packed T=1 body not found: {body}") + binding_header = csrc_dir / "cake_kda_packed_t1_binding.cuh" + if not binding_header.exists(): + raise FileNotFoundError( + f"Cake KDA packed T=1 binding header not found: {binding_header}" + ) + + uri = get_cake_kda_packed_t1_uri(variant, target) + binding = jit_env.FLASHINFER_GEN_SRC_DIR / uri / "cake_kda_packed_t1_binding.cu" + write_if_different(binding, _get_binding_cu(metadata)) + spec = gen_jit_spec( + name=uri, + sources=[binding], + extra_cuda_cflags=[ + *_CAKE_KDA_PACKED_T1_NVCC_FLAGS[target], + ( + "-DFLASHINFER_CAKE_KDA_PACKED_T1_TARGET_KIND=" + f"{_CAKE_KDA_PACKED_T1_TARGET_KIND[target]}" + ), + *metadata.extra_cuda_flags, + ], + extra_include_paths=[csrc_dir, csrc_dir.parent, _get_include_dir()], + ) + logger.info( + "Generated Cake KDA packed T=1 %s %s JIT spec: %s", variant, target, spec.name + ) + return spec + + +@functools.cache +def load_cake_kda_packed_t1_module( + variant: CakeKDAPackedT1Variant, + target: CakeKDAPackedT1Target, +): + module = gen_cake_kda_packed_t1_module(variant, target).build_and_load() + logger.info("Loaded Cake KDA packed T=1 %s %s module", variant, target) + return module + + +def get_cake_kda_packed_t1_module( + variant: CakeKDAPackedT1Variant, + target: CakeKDAPackedT1Target, +): + return load_cake_kda_packed_t1_module(variant, target) + + +__all__ = [ + "CAKE_KDA_PACKED_T1_VARIANTS", + "CAKE_KDA_PACKED_T1_VARIANT_METADATA", + "CakeKDAPackedT1Target", + "CakeKDAPackedT1Variant", + "CakeKDAPackedT1VariantMetadata", + "gen_cake_kda_packed_t1_module", + "get_cake_kda_packed_t1_module", + "get_cake_kda_packed_t1_uri", + "load_cake_kda_packed_t1_module", + "select_cake_kda_packed_t1_variant", +] diff --git a/flashinfer/kda_kernels/cake_packed_kda_decode.py b/flashinfer/kda_kernels/cake_packed_kda_decode.py index ae5db35d23b..c9019e0eb50 100644 --- a/flashinfer/kda_kernels/cake_packed_kda_decode.py +++ b/flashinfer/kda_kernels/cake_packed_kda_decode.py @@ -21,9 +21,13 @@ import torch from ..jit.cpp_ext import is_cuda_version_at_least +from ..jit.cake_kda_packed_t1 import ( + CakeKDAPackedT1Target, + get_cake_kda_packed_t1_module, + select_cake_kda_packed_t1_variant, +) from ..jit.cake_flash_kda_packed_t1 import ( - FlashKDAPackedT1Target, - _variant_for_batch, + _variant_for_batch as _legacy_variant_for_batch, get_flash_kda_packed_t1_module, ) from ..utils import get_compute_capability @@ -33,7 +37,7 @@ _MIXED_WIDTH = 3 * _HEADS * _HEAD_DIM -def _target_for_device(device: torch.device) -> FlashKDAPackedT1Target: +def _target_for_device(device: torch.device) -> CakeKDAPackedT1Target: """Select the legacy exact target or the SM100-family target.""" compute_capability = get_compute_capability(device) @@ -56,6 +60,36 @@ def _target_for_device(device: torch.device) -> FlashKDAPackedT1Target: ) +def _optimized_alignment_flags( + mixed_qkv: torch.Tensor, + raw_gate: torch.Tensor, + dt_bias: torch.Tensor, + state: torch.Tensor, +) -> tuple[bool, bool]: + """Resolve only the physical alignment facts used by the final selector.""" + + state_aligned = ( + isinstance(state, torch.Tensor) + and state.ndim >= 1 + and int(state.data_ptr()) % 16 == 0 + and int(state.stride(0)) % 8 == 0 + ) + aux_vec4_aligned = ( + isinstance(raw_gate, torch.Tensor) + and raw_gate.ndim >= 1 + and isinstance(dt_bias, torch.Tensor) + and int(mixed_qkv.data_ptr()) % 8 == 0 + and (int(mixed_qkv.data_ptr()) + _HEADS * _HEAD_DIM * mixed_qkv.element_size()) + % 8 + == 0 + and int(mixed_qkv.stride(0)) % 4 == 0 + and int(raw_gate.data_ptr()) % 8 == 0 + and int(raw_gate.stride(0)) % 4 == 0 + and int(dt_bias.data_ptr()) % 16 == 0 + ) + return state_aligned, aux_vec4_aligned + + def run_packed_kda_decode( mixed_qkv: torch.Tensor, raw_gate: torch.Tensor, @@ -104,9 +138,20 @@ def run_packed_kda_decode( raise ValueError("output must be contiguous with shape [B,1,12,128]") output_view = output.view(batch, _HEADS, _HEAD_DIM) - variant = _variant_for_batch(batch) target = _target_for_device(mixed_qkv.device) - module = get_flash_kda_packed_t1_module(variant, target) + state_aligned, aux_vec4_aligned = _optimized_alignment_flags( + mixed_qkv, raw_gate, dt_bias, state + ) + optimized_variant = select_cake_kda_packed_t1_variant( + batch, + state_aligned=state_aligned, + aux_vec4_aligned=aux_vec4_aligned, + ) + if optimized_variant is None: + legacy_variant = _legacy_variant_for_batch(batch) + module = get_flash_kda_packed_t1_module(legacy_variant, target) + else: + module = get_cake_kda_packed_t1_module(optimized_variant, target) module.run( mixed_qkv, raw_gate, diff --git a/tests/jit/test_cake_kda_packed_t1_jit.py b/tests/jit/test_cake_kda_packed_t1_jit.py new file mode 100644 index 00000000000..7ea84c7d093 --- /dev/null +++ b/tests/jit/test_cake_kda_packed_t1_jit.py @@ -0,0 +1,149 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# 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 + +import pytest + +from flashinfer.jit import cake_kda_packed_t1 +from flashinfer.jit import core as jit_core + + +@pytest.mark.parametrize( + ("batch", "expected"), + [ + (1, "register_tile16"), + (14, "register_tile16"), + (15, "register_tile8_interleaved"), + (29, "register_tile8_interleaved"), + (30, "register_tile16_warp"), + (38, "register_tile16_warp"), + (39, "cpasync_tile64_register_pipeline"), + (41, "cpasync_tile64_register_pipeline"), + (42, "cpasync_tile128_packed_state_v_private_prefetch"), + (80, "cpasync_tile128_packed_state_v_private_prefetch"), + (81, "cpasync_tile128_v_private_prefetch"), + (101, "cpasync_tile128_v_private_prefetch"), + (102, "cpasync_tile128_paired_row_pipeline"), + (152, "cpasync_tile128_paired_row_pipeline"), + (153, "cpasync_tile128_register_pipeline"), + (65535, "cpasync_tile128_register_pipeline"), + ], +) +def test_aligned_selector_matches_qualified_batch_bands(batch, expected): + assert ( + cake_kda_packed_t1.select_cake_kda_packed_t1_variant( + batch, + state_aligned=True, + aux_vec4_aligned=True, + ) + == expected + ) + + +@pytest.mark.parametrize( + ("batch", "expected"), + [ + (1, "cpasync_tile64_ilp4"), + (24, "cpasync_tile64_ilp4"), + (25, "cpasync_tile64"), + (37, "cpasync_tile64"), + (38, "cpasync_tile128_ilp4"), + (39, None), + (65535, None), + ], +) +def test_scalar_aux_selector_fails_closed_outside_qualified_bands(batch, expected): + assert ( + cake_kda_packed_t1.select_cake_kda_packed_t1_variant( + batch, + state_aligned=True, + aux_vec4_aligned=False, + ) + == expected + ) + + +def test_unaligned_state_uses_legacy_route(): + for batch in (1, 38, 512): + assert ( + cake_kda_packed_t1.select_cake_kda_packed_t1_variant( + batch, + state_aligned=False, + aux_vec4_aligned=True, + ) + is None + ) + with pytest.raises(ValueError, match="batch must be positive"): + cake_kda_packed_t1.select_cake_kda_packed_t1_variant( + 0, + state_aligned=True, + aux_vec4_aligned=True, + ) + + +@pytest.mark.parametrize("variant", cake_kda_packed_t1.CAKE_KDA_PACKED_T1_VARIANTS) +def test_jit_specs_bind_frozen_source_and_physical_launch_metadata( + monkeypatch, + tmp_path, + variant, +): + monkeypatch.setattr( + jit_core.current_compilation_context, + "TARGET_CUDA_ARCHS", + {(10, "3a")}, + ) + monkeypatch.setattr( + cake_kda_packed_t1.jit_env, + "FLASHINFER_GEN_SRC_DIR", + tmp_path, + ) + cake_kda_packed_t1.gen_cake_kda_packed_t1_module.cache_clear() + + metadata = cake_kda_packed_t1.CAKE_KDA_PACKED_T1_VARIANT_METADATA[variant] + spec = cake_kda_packed_t1.gen_cake_kda_packed_t1_module(variant, "sm100f") + uri = cake_kda_packed_t1.get_cake_kda_packed_t1_uri(variant, "sm100f") + + assert spec.name == uri + assert spec.sources == [tmp_path / uri / "cake_kda_packed_t1_binding.cu"] + assert "-gencode=arch=compute_100f,code=sm_100f" in spec.extra_cuda_cflags + assert "-DFLASHINFER_CAKE_KDA_PACKED_T1_TARGET_KIND=100" in spec.extra_cuda_cflags + assert "-use_fast_math" in spec.extra_cuda_cflags + assert "--maxrregcount=128" in spec.extra_cuda_cflags + assert ("--ftz=false" in spec.extra_cuda_cflags) == ( + variant == "register_tile8_interleaved" + ) + + source = (cake_kda_packed_t1._get_csrc_dir() / metadata.body).read_text() + assert metadata.symbol in source + binding = spec.sources[0].read_text() + assert f'#define CAKE_KDA_PACKED_T1_BODY_FILE "{metadata.body}"' in binding + assert f"#define CAKE_KDA_PACKED_T1_KERNEL {metadata.symbol}" in binding + assert f"#define CAKE_KDA_PACKED_T1_VALUE_TILES {metadata.value_tiles}" in binding + assert f"#define CAKE_KDA_PACKED_T1_THREADS {metadata.threads}" in binding + assert f"#define CAKE_KDA_PACKED_T1_SMEM_BYTES {metadata.smem_bytes}" in binding + assert ( + "#define CAKE_KDA_PACKED_T1_REQUIRES_AUX_VEC4 " + f"{int(metadata.requires_aux_vec4)}" + ) in binding + assert '#include "cake_kda_packed_t1_binding.cuh"' in binding + cake_kda_packed_t1.gen_cake_kda_packed_t1_module.cache_clear() + + +def test_binding_preserves_stream_stride_index_and_alignment_contracts(): + binding = ( + cake_kda_packed_t1._get_csrc_dir() / "cake_kda_packed_t1_binding.cuh" + ).read_text() + assert "CHECK_INPUT_TYPE(mixed_qkv, dl_bfloat16)" in binding + assert "state_indices must have shape [B]" in binding + assert "state.stride(0) % 8 == 0" in binding + assert "state.data_ptr()) % 16 == 0" in binding + assert "CAKE_KDA_PACKED_T1_REQUIRES_AUX_VEC4" in binding + assert "mixed_qkv.stride(0)" in binding + assert "raw_gate.stride(0)" in binding + assert "raw_beta.stride(0)" in binding + assert "reinterpret_cast(cuda_stream)" in binding + assert "TVM_FFI_DLL_EXPORT_TYPED_FUNC(run" in binding diff --git a/tests/jit/test_flash_kda_packed_t1_jit.py b/tests/jit/test_flash_kda_packed_t1_jit.py index b9171a7b378..d31df3004de 100644 --- a/tests/jit/test_flash_kda_packed_t1_jit.py +++ b/tests/jit/test_flash_kda_packed_t1_jit.py @@ -18,6 +18,7 @@ from packaging.version import Version from flashinfer.jit import core as jit_core +from flashinfer.jit import cake_kda_packed_t1 from flashinfer.jit import cake_flash_kda_packed_t1 as flash_kda_packed_t1 @@ -237,25 +238,14 @@ def get_nvcc_flags_list(self, supported_major_versions=None): @pytest.mark.parametrize( ("capabilities", "expected_targets"), [ - ( - {"flash_kda_packed_t1_sm100a": True}, - [(variant, "sm100a") for variant in ("tile8", "tile16")], - ), - ( - {"flash_kda_packed_t1_sm100f": True}, - [(variant, "sm100f") for variant in ("tile8", "tile16")], - ), + ({"flash_kda_packed_t1_sm100a": True}, ["sm100a"]), + ({"flash_kda_packed_t1_sm100f": True}, ["sm100f"]), ( { "flash_kda_packed_t1_sm100a": True, "flash_kda_packed_t1_sm100f": True, }, - [ - ("tile8", "sm100a"), - ("tile16", "sm100a"), - ("tile8", "sm100f"), - ("tile16", "sm100f"), - ], + ["sm100a", "sm100f"], ), ({"sm103": True}, []), ], @@ -265,17 +255,27 @@ def test_aot_registers_flash_kda_packed_t1_portfolio( ): from flashinfer import aot - calls = [] + legacy_calls = [] + cake_calls = [] def fake_flash_kda_packed_t1(variant, target): - calls.append((variant, target)) + legacy_calls.append((variant, target)) return SimpleNamespace(name=f"flash_kda_packed_t1_{variant}_{target}") + def fake_cake_kda_packed_t1(variant, target): + cake_calls.append((variant, target)) + return SimpleNamespace(name=f"cake_kda_packed_t1_{variant}_{target}") + monkeypatch.setattr( aot, "gen_flash_kda_packed_t1_module", fake_flash_kda_packed_t1, ) + monkeypatch.setattr( + aot, + "gen_cake_kda_packed_t1_module", + fake_cake_kda_packed_t1, + ) monkeypatch.setattr( aot, "gen_spdlog_module", lambda: SimpleNamespace(name="spdlog") ) @@ -301,12 +301,33 @@ def fake_flash_kda_packed_t1(variant, target): False, ) - assert calls == expected_targets + expected_legacy_calls = [ + (variant, target) + for target in expected_targets + for variant in flash_kda_packed_t1.FLASH_KDA_PACKED_T1_VARIANTS + ] + expected_cake_calls = [ + (variant, target) + for target in expected_targets + for variant in cake_kda_packed_t1.CAKE_KDA_PACKED_T1_VARIANTS + ] + assert legacy_calls == expected_legacy_calls + assert cake_calls == expected_cake_calls assert [spec.name for spec in specs] == [ "spdlog", - *( - f"flash_kda_packed_t1_{variant}_{target}" - for variant, target in expected_targets - ), + *[ + name + for target in expected_targets + for name in [ + *( + f"flash_kda_packed_t1_{variant}_{target}" + for variant in flash_kda_packed_t1.FLASH_KDA_PACKED_T1_VARIANTS + ), + *( + f"cake_kda_packed_t1_{variant}_{target}" + for variant in cake_kda_packed_t1.CAKE_KDA_PACKED_T1_VARIANTS + ), + ] + ], "cudnn", ] diff --git a/tests/kda/test_packed_kda_decode.py b/tests/kda/test_packed_kda_decode.py index d6aa04babb7..bc09e71f85c 100644 --- a/tests/kda/test_packed_kda_decode.py +++ b/tests/kda/test_packed_kda_decode.py @@ -353,15 +353,18 @@ def test_packed_kda_decode_all_inactive_is_bitwise_noop(packed_kda_device): @pytest.mark.arch_blackwell -@pytest.mark.parametrize("batch", [8, 64]) -def test_packed_kda_decode_sanitizer_schedules(packed_kda_device, batch): - """Small named entry points for tile8/tile16 sanitizer invocations.""" +@pytest.mark.parametrize( + ("batch", "state_padding"), + [(8, 17), (8, _PRODUCTION_STATE_PADDING), (64, _PRODUCTION_STATE_PADDING)], +) +def test_packed_kda_decode_sanitizer_schedules(packed_kda_device, batch, state_padding): + """Small named entry points for legacy and optimized sanitizer runs.""" case = _make_case( batch, packed_kda_device, seed=20260824 + batch, inactive=False, - state_padding=17, + state_padding=state_padding, ) _call(case) torch.cuda.synchronize(packed_kda_device) @@ -615,7 +618,7 @@ def new_empty(self, shape): ("batch", "target", "variant"), [(31, "sm100a", "tile8"), (32, "sm100f", "tile16")], ) -def test_kernel_facade_selects_variant_and_caller_stream_cpu( +def test_kernel_facade_preserves_legacy_fallback_and_caller_stream_cpu( monkeypatch, batch, target, variant ): packed_module = importlib.import_module( @@ -630,6 +633,11 @@ def test_kernel_facade_selects_variant_and_caller_stream_cpu( ) monkeypatch.setattr(packed_module, "torch", fake_torch) monkeypatch.setattr(packed_module, "_target_for_device", lambda device: target) + monkeypatch.setattr( + packed_module, + "_optimized_alignment_flags", + lambda *args: (False, False), + ) module_calls = [] launch_calls = [] @@ -652,6 +660,48 @@ def get_module(selected_variant, selected_target): assert launch_calls[0][-1] == stream_handle +def test_kernel_facade_selects_optimized_variant_and_caller_stream_cpu(monkeypatch): + packed_module = importlib.import_module( + "flashinfer.kda_kernels.cake_packed_kda_decode" + ) + batch = 29 + target = "sm100f" + stream_handle = 0x12345678 + fake_torch = SimpleNamespace( + Tensor=_FakeCudaTensor, + cuda=SimpleNamespace( + current_stream=lambda device: SimpleNamespace(cuda_stream=stream_handle) + ), + ) + monkeypatch.setattr(packed_module, "torch", fake_torch) + monkeypatch.setattr(packed_module, "_target_for_device", lambda device: target) + monkeypatch.setattr( + packed_module, + "_optimized_alignment_flags", + lambda *args: (True, True), + ) + + module_calls = [] + launch_calls = [] + + def get_module(selected_variant, selected_target): + module_calls.append((selected_variant, selected_target)) + return SimpleNamespace(run=lambda *args: launch_calls.append(args)) + + monkeypatch.setattr(packed_module, "get_cake_kda_packed_t1_module", get_module) + mixed_qkv = _FakeCudaTensor((batch, _MIXED_WIDTH)) + inputs = [mixed_qkv, *(_FakeCudaTensor((1,)) for _ in range(6))] + output = _FakeCudaTensor((batch, 1, _HEADS, _HEAD_DIM)) + + result = packed_module.run_packed_kda_decode(*inputs, output=output) + + assert result is output + assert module_calls == [("register_tile8_interleaved", target)] + assert len(launch_calls) == 1 + assert launch_calls[0][-2] is output.view_result + assert launch_calls[0][-1] == stream_handle + + @pytest.mark.parametrize( ("capability", "cuda_at_least", "expected"), [