diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 5f5e37836a05..20e212149d18 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -186,6 +186,7 @@ set(TRTLLM_LINK_LIBS trtllm_gen_batched_gemm selective_scan_src mamba2_mtp_ssm_cache_src + kda_decode_src ws_layernorm_src fusedGatedRMSNormQuant_src fpA_intB_gemm_src diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index fcea8829442b..ac5306af882c 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -30,6 +30,7 @@ add_subdirectory(dsv3MinLatencyKernels) add_subdirectory(causalConv1d) add_subdirectory(fusedGatedRMSNormQuant) add_subdirectory(mamba2MTPSSMCache) +add_subdirectory(kdaDecode) add_subdirectory(mhcKernels) add_subdirectory(compressorKernels) @@ -57,6 +58,8 @@ list(FILTER SRC_CU EXCLUDE REGEX "userbuffers/.*") list(FILTER SRC_CU EXCLUDE REGEX "fusedLayernormKernels/.*") list(FILTER SRC_CU EXCLUDE REGEX "fusedGatedRMSNormQuant/.*") list(FILTER SRC_CU EXCLUDE REGEX "mamba2MTPSSMCache/.*") +list(FILTER SRC_CPP EXCLUDE REGEX "kdaDecode/.*") +list(FILTER SRC_CU EXCLUDE REGEX "kdaDecode/.*") list(FILTER SRC_CPP EXCLUDE REGEX "mhcKernels/.*") list(FILTER SRC_CU EXCLUDE REGEX "mhcKernels/.*") list(FILTER SRC_CPP EXCLUDE REGEX "compressorKernels/.*") diff --git a/cpp/tensorrt_llm/kernels/kdaDecode/CMakeLists.txt b/cpp/tensorrt_llm/kernels/kdaDecode/CMakeLists.txt new file mode 100644 index 000000000000..81560913c019 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/kdaDecode/CMakeLists.txt @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not +# use this file except in compliance with the License. You may obtain a copy of +# the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations under +# the License. + +add_library(kda_decode_src OBJECT kdaDecode.cu) + +set_property(TARGET kda_decode_src PROPERTY POSITION_INDEPENDENT_CODE ON) +set_property(TARGET kda_decode_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) +target_compile_options(kda_decode_src + PRIVATE $<$:--use_fast_math>) diff --git a/cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.cu b/cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.cu new file mode 100644 index 000000000000..1249779454ad --- /dev/null +++ b/cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.cu @@ -0,0 +1,1611 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/kdaDecode/kdaDecode.h" + +#include "tensorrt_llm/common/cudaUtils.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::kdaDecode +{ + +namespace +{ + +constexpr int kDimK = 128; +constexpr int kDimV = 128; +constexpr int kKernelWidth = 4; +constexpr int kConvStateWidth = kKernelWidth - 1; +constexpr int kThreads = 256; +constexpr int kWarps = kThreads / 32; +constexpr int kChunkV = 32; +constexpr int kNumChunks = kDimV / kChunkV; +constexpr int kRowsPerWarp = kChunkV / kWarps; + +template +__device__ __forceinline__ float bf16_load(__nv_bfloat16 const* ptr, Index idx) +{ + return __bfloat162float(ptr[idx]); +} + +__device__ __forceinline__ __nv_bfloat16 bf16_store(float value) +{ + return __float2bfloat16(value); +} + +template +__device__ __forceinline__ void store_state_float4(float* ptr, float4 value) +{ + if constexpr (kUseCacheGlobalStore) + { + __stcg(reinterpret_cast(ptr), value); + } + else + { + *reinterpret_cast(ptr) = value; + } +} + +__device__ __forceinline__ float sigmoid_fast(float x) +{ + return 1.0f / (1.0f + __expf(-x)); +} + +__device__ __forceinline__ float silu_fast(float x) +{ + return x * sigmoid_fast(x); +} + +__device__ __forceinline__ float softplus_fast(float x) +{ + return x > 20.0f ? x : log1pf(__expf(x)); +} + +__device__ __forceinline__ float warp_reduce_sum(float value) +{ +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + { + value += __shfl_xor_sync(0xffffffffu, value, offset); + } + return value; +} + +// cp.async requires sm_80+; pre-SM80 falls back to a synchronous copy, so the +// commit/wait helpers become no-ops there. +__device__ __forceinline__ void cp_async_cg_16b(float* smem_ptr, float const* gmem_ptr) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + uint32_t smem_addr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" : : "r"(smem_addr), "l"(gmem_ptr)); +#else + *reinterpret_cast(smem_ptr) = *reinterpret_cast(gmem_ptr); +#endif +} + +__device__ __forceinline__ void cp_async_commit() +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.commit_group;\n" ::); +#endif +} + +__device__ __forceinline__ void cp_async_wait_all() +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.wait_all;\n" ::); +#endif +} + +__device__ __forceinline__ void cp_async_wait_group_0() +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.wait_group 0;\n" ::); +#endif +} + +__device__ __forceinline__ void cp_async_wait_group_1() +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.wait_group 1;\n" ::); +#endif +} + +__device__ __forceinline__ void cp_async_wait_group_2() +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.wait_group 2;\n" ::); +#endif +} + +template +__device__ __forceinline__ void cp_async_state_chunk_stage( + float* s_state, float const* state, int slot, int i_hv, int64_t state_slot_stride, int chunk, int stage) +{ + constexpr int kFloat4PerChunk = kStageChunkV * kDimK / 4; + int const tid = threadIdx.x; + int const v_base = chunk * kStageChunkV; + for (int linear4 = tid; linear4 < kFloat4PerChunk; linear4 += kThreads) + { + int const elem = linear4 * 4; + int const row = elem / kDimK; + int const k = elem - row * kDimK; + float* dst = s_state + (stage * kStageChunkV + row) * kDimK + k; + int64_t const state_offset = static_cast(slot) * state_slot_stride + + (static_cast(i_hv) * kDimV + v_base + row) * kDimK + k; + float const* src = state + state_offset; + cp_async_cg_16b(dst, src); + } + cp_async_commit(); +} + +template +__device__ __forceinline__ void cp_async_state_chunk_for( + float* s_state, float const* state, int slot, int i_hv, int64_t state_slot_stride, int chunk) +{ + constexpr int kFloat4PerChunk = kChunkV * kDimK / 4; + int const tid = threadIdx.x; + int const stage = chunk & 1; + int const v_base = chunk * kChunkV; + for (int linear4 = tid; linear4 < kFloat4PerChunk; linear4 += kCopyThreads) + { + int const elem = linear4 * 4; + int const row = elem / kDimK; + int const k = elem - row * kDimK; + float* dst = s_state + (stage * kChunkV + row) * kDimK + k; + int64_t const state_offset = static_cast(slot) * state_slot_stride + + (static_cast(i_hv) * kDimV + v_base + row) * kDimK + k; + float const* src = state + state_offset; + cp_async_cg_16b(dst, src); + } + cp_async_commit(); +} + +__device__ __forceinline__ void cp_async_state_chunk( + float* s_state, float const* state, int slot, int i_hv, int64_t state_slot_stride, int chunk) +{ + cp_async_state_chunk_for(s_state, state, slot, i_hv, state_slot_stride, chunk); +} + +__device__ __forceinline__ float block_reduce_sum(float value, float* scratch) +{ + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + + float warp_total = warp_reduce_sum(value); + if (lane == 0) + { + scratch[warp] = warp_total; + } + __syncthreads(); + + float block_total = 0.0f; + if (warp == 0) + { + block_total = lane < kWarps ? scratch[lane] : 0.0f; + block_total = warp_reduce_sum(block_total); + if (lane == 0) + { + scratch[0] = block_total; + } + } + __syncthreads(); + return scratch[0]; +} + +struct Sum2 +{ + float x; + float y; +}; + +__device__ __forceinline__ Sum2 warp_reduce_sum_pair(float x, float y) +{ +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + { + x += __shfl_xor_sync(0xffffffffu, x, offset); + y += __shfl_xor_sync(0xffffffffu, y, offset); + } + return {x, y}; +} + +struct Sum4 +{ + float a; + float b; + float c; + float d; +}; + +__device__ __forceinline__ Sum4 warp_reduce_sum4(float a, float b, float c, float d) +{ +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + { + a += __shfl_xor_sync(0xffffffffu, a, offset); + b += __shfl_xor_sync(0xffffffffu, b, offset); + c += __shfl_xor_sync(0xffffffffu, c, offset); + d += __shfl_xor_sync(0xffffffffu, d, offset); + } + return {a, b, c, d}; +} + +template +__device__ __forceinline__ Sum2 block_reduce_sum2_for(float x, float y, float* scratch) +{ + int const lane = threadIdx.x & 31; + int const warp = threadIdx.x >> 5; + + float const warp_x = warp_reduce_sum(x); + float const warp_y = warp_reduce_sum(y); + if (lane == 0) + { + scratch[warp] = warp_x; + scratch[kReduceWarps + warp] = warp_y; + } + __syncthreads(); + + float block_x = 0.0f; + float block_y = 0.0f; + if (warp == 0) + { + block_x = lane < kReduceWarps ? scratch[lane] : 0.0f; + block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f; + block_x = warp_reduce_sum(block_x); + block_y = warp_reduce_sum(block_y); + if (lane == 0) + { + scratch[0] = block_x; + scratch[1] = block_y; + } + } + __syncthreads(); + return {scratch[0], scratch[1]}; +} + +__device__ __forceinline__ Sum2 block_reduce_sum2(float x, float y, float* scratch) +{ + return block_reduce_sum2_for(x, y, scratch); +} + +template +__device__ __forceinline__ float block_reduce_sum_active_for(float value, float* scratch) +{ + int const lane = threadIdx.x & 31; + int const warp = threadIdx.x >> 5; + + float warp_total = 0.0f; + if (warp < kReduceWarps) + { + warp_total = warp_reduce_sum(value); + } + if (lane == 0 && warp < kReduceWarps) + { + scratch[warp] = warp_total; + } + __syncthreads(); + + float block_total = 0.0f; + if (warp == 0) + { + block_total = lane < kReduceWarps ? scratch[lane] : 0.0f; + block_total = warp_reduce_sum(block_total); + if (lane == 0) + { + scratch[0] = block_total; + } + } + __syncthreads(); + return scratch[0]; +} + +template +__device__ __forceinline__ Sum2 block_reduce_sum2_active_for(float x, float y, float* scratch) +{ + int const lane = threadIdx.x & 31; + int const warp = threadIdx.x >> 5; + + float warp_x = 0.0f; + float warp_y = 0.0f; + if (warp < kReduceWarps) + { + warp_x = warp_reduce_sum(x); + warp_y = warp_reduce_sum(y); + } + if (lane == 0 && warp < kReduceWarps) + { + scratch[warp] = warp_x; + scratch[kReduceWarps + warp] = warp_y; + } + __syncthreads(); + + float block_x = 0.0f; + float block_y = 0.0f; + if (warp == 0) + { + block_x = lane < kReduceWarps ? scratch[lane] : 0.0f; + block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f; + block_x = warp_reduce_sum(block_x); + block_y = warp_reduce_sum(block_y); + if (lane == 0) + { + scratch[0] = block_x; + scratch[1] = block_y; + } + } + __syncthreads(); + return {scratch[0], scratch[1]}; +} + +template +__global__ __launch_bounds__(kThreads, 2) void kda_decode_fusion_compact_heads_kernel( + __nv_bfloat16 const* __restrict__ x_q, __nv_bfloat16 const* __restrict__ x_k, __nv_bfloat16 const* __restrict__ x_v, + __nv_bfloat16 const* __restrict__ w_q_t, __nv_bfloat16 const* __restrict__ w_k_t, + __nv_bfloat16 const* __restrict__ w_v_t, __nv_bfloat16 const* __restrict__ bias_q, + __nv_bfloat16 const* __restrict__ bias_k, __nv_bfloat16 const* __restrict__ bias_v, + __nv_bfloat16* __restrict__ cs_q, __nv_bfloat16* __restrict__ cs_k, __nv_bfloat16* __restrict__ cs_v, + float const* __restrict__ a_log, __nv_bfloat16 const* __restrict__ g, float const* __restrict__ dt_bias, + __nv_bfloat16 const* __restrict__ beta, __nv_bfloat16 const* __restrict__ onorm_g, + float const* __restrict__ onorm_weight, int const* __restrict__ ssm_state_indices, + int const* __restrict__ cu_seqlens, float* __restrict__ state, int64_t state_slot_stride, + __nv_bfloat16* __restrict__ out, int B, int H, int HV, float lower_bound, float scale, float onorm_eps) +{ + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int i_n; + int i_hv; + int i_h; + int bos; + int slot; + if constexpr (kUseStaticDecodeLayout) + { + if constexpr (kUseHeadGrid) + { + i_n = blockIdx.x; + i_hv = blockIdx.y; + } + else + { + int const nhv = blockIdx.x; + i_n = nhv / kFixedValueHeads; + i_hv = nhv - i_n * kFixedValueHeads; + } + i_h = i_hv; + bos = i_n; + slot = i_n; + } + else + { + int const nhv = blockIdx.x; + i_n = nhv / HV; + i_hv = nhv - i_n * HV; + int const hv_per_h = HV / H; + i_h = i_hv / hv_per_h; + + bos = cu_seqlens == nullptr ? i_n : cu_seqlens[i_n]; + int const eos = cu_seqlens == nullptr ? i_n + 1 : cu_seqlens[i_n + 1]; + if (eos <= bos) + { + return; + } + slot = ssm_state_indices[i_n]; + } + + int const hk_off = i_h * kDimK; + int const hv_off = i_hv * kDimV; + int const h_count = kUseStaticDecodeLayout ? kFixedHeads : H; + int const hv_count = kUseStaticDecodeLayout ? kFixedValueHeads : HV; + int const hkv_dim = h_count * kDimK; + int const hvv_dim = hv_count * kDimV; + int const conv_slot = kUpdateConvState ? slot : i_n; + int64_t const qk_input_head = static_cast(bos) * h_count + i_h; + int64_t const value_input_head = static_cast(bos) * hv_count + i_hv; + int64_t const output_head = static_cast(i_n) * hv_count + i_hv; + + constexpr int kStageChunkV = 32; + constexpr int kStageCount = 3; + constexpr int kStageNumChunks = kDimV / kStageChunkV; + extern __shared__ float s_state[]; + __shared__ float s_q[kDimK]; + __shared__ float s_k[kDimK]; + __shared__ float s_decay[kDimK]; + __shared__ float s_v[kDimV]; + __shared__ float s_o[kDimV]; + __shared__ float s_reduce[kThreads]; + __shared__ float s_beta; + float pre_onorm_gate = 0.0f; + float pre_onorm_weight = 0.0f; + + cp_async_state_chunk_stage(s_state, state, slot, i_hv, state_slot_stride, 0, 0); + cp_async_state_chunk_stage(s_state, state, slot, i_hv, state_slot_stride, 1, 1); + + if constexpr (kUpdateConvState) + { + if (tid < kDimK) + { + int const k = tid; + int const hk = hk_off + k; + int64_t const cs_base = static_cast(slot) * hkv_dim * kConvStateWidth + hk; + int64_t const xq_idx = qk_input_head * kDimK + k; + float const exp_a = __shfl_sync(0xffffffffu, lane == 0 ? __expf(a_log[i_h]) : 0.0f, 0); + + float q_acc = bf16_load(bias_q, hk); + float k_acc = bf16_load(bias_k, hk); + __nv_bfloat16 q_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 q_shift1 = __float2bfloat16(0.0f); + __nv_bfloat16 k_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 k_shift1 = __float2bfloat16(0.0f); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) + { + const __nv_bfloat16 q_state = cs_q[cs_base + w * hkv_dim]; + const __nv_bfloat16 k_state = cs_k[cs_base + w * hkv_dim]; + q_acc += __bfloat162float(q_state) * bf16_load(w_q_t, w * hkv_dim + hk); + k_acc += __bfloat162float(k_state) * bf16_load(w_k_t, w * hkv_dim + hk); + if (w == 1) + { + q_shift0 = q_state; + k_shift0 = k_state; + } + else if (w == 2) + { + q_shift1 = q_state; + k_shift1 = k_state; + } + } + const __nv_bfloat16 q_new = x_q[xq_idx]; + const __nv_bfloat16 k_new = x_k[xq_idx]; + q_acc += __bfloat162float(q_new) * bf16_load(w_q_t, (kKernelWidth - 1) * hkv_dim + hk); + k_acc += __bfloat162float(k_new) * bf16_load(w_k_t, (kKernelWidth - 1) * hkv_dim + hk); + + cs_q[cs_base + 0] = q_shift0; + cs_q[cs_base + hkv_dim] = q_shift1; + cs_q[cs_base + 2 * hkv_dim] = q_new; + cs_k[cs_base + 0] = k_shift0; + cs_k[cs_base + hkv_dim] = k_shift1; + cs_k[cs_base + 2 * hkv_dim] = k_new; + + s_q[k] = silu_fast(q_acc); + s_k[k] = silu_fast(k_acc); + + float const g_raw = bf16_load(g, value_input_head * kDimK + k) + dt_bias[hk]; + if constexpr (kUseLowerBound) + { + s_decay[k] = __expf(lower_bound * sigmoid_fast(exp_a * g_raw)); + } + else + { + s_decay[k] = __expf(-exp_a * softplus_fast(g_raw)); + } + } + } + else + { + if (tid < kDimK) + { + int const k = tid; + int const hk = hk_off + k; + float const exp_a = __shfl_sync(0xffffffffu, lane == 0 ? __expf(a_log[i_h]) : 0.0f, 0); + + float q_acc = bf16_load(bias_q, hk); + float k_acc = bf16_load(bias_k, hk); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) + { + int64_t const cs_idx = (static_cast(conv_slot) * hkv_dim + hk) * kConvStateWidth + w; + q_acc += bf16_load(cs_q, cs_idx) * bf16_load(w_q_t, w * hkv_dim + hk); + k_acc += bf16_load(cs_k, cs_idx) * bf16_load(w_k_t, w * hkv_dim + hk); + } + q_acc += bf16_load(x_q, qk_input_head * kDimK + k) * bf16_load(w_q_t, (kKernelWidth - 1) * hkv_dim + hk); + k_acc += bf16_load(x_k, qk_input_head * kDimK + k) * bf16_load(w_k_t, (kKernelWidth - 1) * hkv_dim + hk); + + s_q[k] = silu_fast(q_acc); + s_k[k] = silu_fast(k_acc); + + float const g_raw = bf16_load(g, value_input_head * kDimK + k) + dt_bias[hk]; + if constexpr (kUseLowerBound) + { + s_decay[k] = __expf(lower_bound * sigmoid_fast(exp_a * g_raw)); + } + else + { + s_decay[k] = __expf(-exp_a * softplus_fast(g_raw)); + } + } + } + + if constexpr (kUpdateConvState) + { + if (tid < kDimV) + { + int const v = tid; + int const hvv = hv_off + v; + int64_t const cs_base = static_cast(slot) * hvv_dim * kConvStateWidth + hvv; + int64_t const xv_idx = value_input_head * kDimV + v; + + float v_acc = bf16_load(bias_v, hvv); + __nv_bfloat16 v_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 v_shift1 = __float2bfloat16(0.0f); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) + { + const __nv_bfloat16 v_state = cs_v[cs_base + w * hvv_dim]; + v_acc += __bfloat162float(v_state) * bf16_load(w_v_t, w * hvv_dim + hvv); + if (w == 1) + { + v_shift0 = v_state; + } + else if (w == 2) + { + v_shift1 = v_state; + } + } + const __nv_bfloat16 v_new = x_v[xv_idx]; + v_acc += __bfloat162float(v_new) * bf16_load(w_v_t, (kKernelWidth - 1) * hvv_dim + hvv); + cs_v[cs_base + 0] = v_shift0; + cs_v[cs_base + hvv_dim] = v_shift1; + cs_v[cs_base + 2 * hvv_dim] = v_new; + s_v[v] = silu_fast(v_acc); + + if constexpr (kApplyOnorm && kPreloadOnormParams) + { + int64_t const out_idx = output_head * kDimV + v; + pre_onorm_gate = sigmoid_fast(bf16_load(onorm_g, out_idx)); + pre_onorm_weight = onorm_weight[v]; + } + } + } + else + { + if (tid < kDimV) + { + int const v = tid; + int const hvv = hv_off + v; + + float v_acc = bf16_load(bias_v, hvv); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) + { + int64_t const cs_idx = (static_cast(conv_slot) * hvv_dim + hvv) * kConvStateWidth + w; + v_acc += bf16_load(cs_v, cs_idx) * bf16_load(w_v_t, w * hvv_dim + hvv); + } + v_acc + += bf16_load(x_v, value_input_head * kDimV + v) * bf16_load(w_v_t, (kKernelWidth - 1) * hvv_dim + hvv); + s_v[v] = silu_fast(v_acc); + + if constexpr (kApplyOnorm && kPreloadOnormParams) + { + int64_t const out_idx = output_head * kDimV + v; + pre_onorm_gate = sigmoid_fast(bf16_load(onorm_g, out_idx)); + pre_onorm_weight = onorm_weight[v]; + } + } + } + + if (tid == 0) + { + float const beta_raw = bf16_load(beta, value_input_head); + if constexpr (kApplyBetaSigmoid) + { + s_beta = sigmoid_fast(beta_raw); + } + else + { + s_beta = beta_raw; + } + } + __syncthreads(); + + if constexpr (kIssueThirdStatePrefetchEarly && kStageNumChunks > 2) + { + cp_async_state_chunk_stage(s_state, state, slot, i_hv, state_slot_stride, 2, 2); + } + + float const q_sq = tid < kDimK ? s_q[tid] * s_q[tid] : 0.0f; + float const k_sq = tid < kDimK ? s_k[tid] * s_k[tid] : 0.0f; + Sum2 qk_sum; + if constexpr (kUseActiveQkReduction) + { + qk_sum = block_reduce_sum2_active_for(q_sq, k_sq, s_reduce); + } + else + { + qk_sum = block_reduce_sum2(q_sq, k_sq, s_reduce); + } + if (tid < kDimK) + { + s_q[tid] *= rsqrtf(qk_sum.x + 1.0e-6f) * scale; + s_k[tid] *= rsqrtf(qk_sum.y + 1.0e-6f); + } + __syncthreads(); + + int const k_base = lane * 4; + const float4 q4 = *reinterpret_cast(s_q + k_base); + const float4 k4 = *reinterpret_cast(s_k + k_base); + const float4 decay4 = *reinterpret_cast(s_decay + k_base); + float r_q[4] = {q4.x, q4.y, q4.z, q4.w}; + float r_k[4] = {k4.x, k4.y, k4.z, k4.w}; + float r_decay[4] = {decay4.x, decay4.y, decay4.z, decay4.w}; + float o_sumsq = 0.0f; + +#pragma unroll + for (int chunk = 0; chunk < kStageNumChunks; ++chunk) + { + if constexpr (kIssueThirdStatePrefetchEarly) + { + if (chunk == 0) + { + cp_async_wait_group_2(); + } + else if (chunk + 1 < kStageNumChunks) + { + cp_async_wait_group_1(); + } + else + { + cp_async_wait_group_0(); + } + } + else + { + if (chunk + 1 < kStageNumChunks) + { + cp_async_wait_group_1(); + } + else + { + cp_async_wait_group_0(); + } + } + __syncwarp(); + + int const prefetch = chunk + 2; + if constexpr (kIssueThirdStatePrefetchEarly) + { + if (prefetch < kStageNumChunks && prefetch >= 3) + { + cp_async_state_chunk_stage( + s_state, state, slot, i_hv, state_slot_stride, prefetch, prefetch % kStageCount); + } + } + else + { + if (prefetch < kStageNumChunks) + { + cp_async_state_chunk_stage( + s_state, state, slot, i_hv, state_slot_stride, prefetch, prefetch % kStageCount); + } + } + + int const v_row_a = warp; + int const v_row_b = warp + kWarps; + int const v_row_c = warp + 2 * kWarps; + int const v_row_d = warp + 3 * kWarps; + int const v0 = chunk * kChunkV + v_row_a; + int const v1 = chunk * kChunkV + v_row_b; + int const v2 = chunk * kChunkV + v_row_c; + int const v3 = chunk * kChunkV + v_row_d; + float h_a_vals[4]; + float h_b_vals[4]; + float h_c_vals[4]; + float h_d_vals[4]; + float dot_hk_a = 0.0f; + float dot_hk_b = 0.0f; + float dot_hk_c = 0.0f; + float dot_hk_d = 0.0f; + + float const* state_stage = s_state + (chunk % kStageCount) * kStageChunkV * kDimK; + const float4 raw_h_a = *reinterpret_cast(state_stage + v_row_a * kDimK + k_base); + const float4 raw_h_b = *reinterpret_cast(state_stage + v_row_b * kDimK + k_base); + const float4 raw_h_c = *reinterpret_cast(state_stage + v_row_c * kDimK + k_base); + const float4 raw_h_d = *reinterpret_cast(state_stage + v_row_d * kDimK + k_base); + h_a_vals[0] = raw_h_a.x * r_decay[0]; + h_a_vals[1] = raw_h_a.y * r_decay[1]; + h_a_vals[2] = raw_h_a.z * r_decay[2]; + h_a_vals[3] = raw_h_a.w * r_decay[3]; + h_b_vals[0] = raw_h_b.x * r_decay[0]; + h_b_vals[1] = raw_h_b.y * r_decay[1]; + h_b_vals[2] = raw_h_b.z * r_decay[2]; + h_b_vals[3] = raw_h_b.w * r_decay[3]; + h_c_vals[0] = raw_h_c.x * r_decay[0]; + h_c_vals[1] = raw_h_c.y * r_decay[1]; + h_c_vals[2] = raw_h_c.z * r_decay[2]; + h_c_vals[3] = raw_h_c.w * r_decay[3]; + h_d_vals[0] = raw_h_d.x * r_decay[0]; + h_d_vals[1] = raw_h_d.y * r_decay[1]; + h_d_vals[2] = raw_h_d.z * r_decay[2]; + h_d_vals[3] = raw_h_d.w * r_decay[3]; + dot_hk_a = h_a_vals[0] * r_k[0] + h_a_vals[1] * r_k[1] + h_a_vals[2] * r_k[2] + h_a_vals[3] * r_k[3]; + dot_hk_b = h_b_vals[0] * r_k[0] + h_b_vals[1] * r_k[1] + h_b_vals[2] * r_k[2] + h_b_vals[3] * r_k[3]; + dot_hk_c = h_c_vals[0] * r_k[0] + h_c_vals[1] * r_k[1] + h_c_vals[2] * r_k[2] + h_c_vals[3] * r_k[3]; + dot_hk_d = h_d_vals[0] * r_k[0] + h_d_vals[1] * r_k[1] + h_d_vals[2] * r_k[2] + h_d_vals[3] * r_k[3]; + + const Sum4 dot_hk = warp_reduce_sum4(dot_hk_a, dot_hk_b, dot_hk_c, dot_hk_d); + float const v_new0 = (s_v[v0] - dot_hk.a) * s_beta; + float const v_new1 = (s_v[v1] - dot_hk.b) * s_beta; + float const v_new2 = (s_v[v2] - dot_hk.c) * s_beta; + float const v_new3 = (s_v[v3] - dot_hk.d) * s_beta; + + float dot_hq_a = 0.0f; + float dot_hq_b = 0.0f; + float dot_hq_c = 0.0f; + float dot_hq_d = 0.0f; + int64_t const state_head_offset + = static_cast(slot) * state_slot_stride + static_cast(i_hv) * kDimV * kDimK; + int64_t const state_idx_a = state_head_offset + static_cast(v0) * kDimK + k_base; + int64_t const state_idx_b = state_head_offset + static_cast(v1) * kDimK + k_base; + int64_t const state_idx_c = state_head_offset + static_cast(v2) * kDimK + k_base; + int64_t const state_idx_d = state_head_offset + static_cast(v3) * kDimK + k_base; + float const h_a_0 = h_a_vals[0] + r_k[0] * v_new0; + float const h_a_1 = h_a_vals[1] + r_k[1] * v_new0; + float const h_a_2 = h_a_vals[2] + r_k[2] * v_new0; + float const h_a_3 = h_a_vals[3] + r_k[3] * v_new0; + float const h_b_0 = h_b_vals[0] + r_k[0] * v_new1; + float const h_b_1 = h_b_vals[1] + r_k[1] * v_new1; + float const h_b_2 = h_b_vals[2] + r_k[2] * v_new1; + float const h_b_3 = h_b_vals[3] + r_k[3] * v_new1; + float const h_c_0 = h_c_vals[0] + r_k[0] * v_new2; + float const h_c_1 = h_c_vals[1] + r_k[1] * v_new2; + float const h_c_2 = h_c_vals[2] + r_k[2] * v_new2; + float const h_c_3 = h_c_vals[3] + r_k[3] * v_new2; + float const h_d_0 = h_d_vals[0] + r_k[0] * v_new3; + float const h_d_1 = h_d_vals[1] + r_k[1] * v_new3; + float const h_d_2 = h_d_vals[2] + r_k[2] * v_new3; + float const h_d_3 = h_d_vals[3] + r_k[3] * v_new3; + if constexpr (kComputeOutputBeforeStore) + { + dot_hq_a = h_a_0 * r_q[0] + h_a_1 * r_q[1] + h_a_2 * r_q[2] + h_a_3 * r_q[3]; + dot_hq_b = h_b_0 * r_q[0] + h_b_1 * r_q[1] + h_b_2 * r_q[2] + h_b_3 * r_q[3]; + dot_hq_c = h_c_0 * r_q[0] + h_c_1 * r_q[1] + h_c_2 * r_q[2] + h_c_3 * r_q[3]; + dot_hq_d = h_d_0 * r_q[0] + h_d_1 * r_q[1] + h_d_2 * r_q[2] + h_d_3 * r_q[3]; + store_state_float4(state + state_idx_a, make_float4(h_a_0, h_a_1, h_a_2, h_a_3)); + store_state_float4(state + state_idx_b, make_float4(h_b_0, h_b_1, h_b_2, h_b_3)); + store_state_float4(state + state_idx_c, make_float4(h_c_0, h_c_1, h_c_2, h_c_3)); + store_state_float4(state + state_idx_d, make_float4(h_d_0, h_d_1, h_d_2, h_d_3)); + } + else + { + store_state_float4(state + state_idx_a, make_float4(h_a_0, h_a_1, h_a_2, h_a_3)); + store_state_float4(state + state_idx_b, make_float4(h_b_0, h_b_1, h_b_2, h_b_3)); + store_state_float4(state + state_idx_c, make_float4(h_c_0, h_c_1, h_c_2, h_c_3)); + store_state_float4(state + state_idx_d, make_float4(h_d_0, h_d_1, h_d_2, h_d_3)); + dot_hq_a = h_a_0 * r_q[0] + h_a_1 * r_q[1] + h_a_2 * r_q[2] + h_a_3 * r_q[3]; + dot_hq_b = h_b_0 * r_q[0] + h_b_1 * r_q[1] + h_b_2 * r_q[2] + h_b_3 * r_q[3]; + dot_hq_c = h_c_0 * r_q[0] + h_c_1 * r_q[1] + h_c_2 * r_q[2] + h_c_3 * r_q[3]; + dot_hq_d = h_d_0 * r_q[0] + h_d_1 * r_q[1] + h_d_2 * r_q[2] + h_d_3 * r_q[3]; + } + + const Sum4 dot_hq = warp_reduce_sum4(dot_hq_a, dot_hq_b, dot_hq_c, dot_hq_d); + if (lane == 0) + { + s_o[v0] = dot_hq.a; + s_o[v1] = dot_hq.b; + s_o[v2] = dot_hq.c; + s_o[v3] = dot_hq.d; + if constexpr (kApplyOnorm && kAccumulateOnormSumsq) + { + o_sumsq += dot_hq.a * dot_hq.a + dot_hq.b * dot_hq.b + dot_hq.c * dot_hq.c + dot_hq.d * dot_hq.d; + } + } + } + __syncthreads(); + + if constexpr (kApplyOnorm) + { + if constexpr (kAccumulateOnormSumsq) + { + if (lane == 0) + { + s_reduce[warp] = o_sumsq; + } + __syncthreads(); + + float total_sumsq = 0.0f; + if (warp == 0) + { + total_sumsq = lane < kWarps ? s_reduce[lane] : 0.0f; + total_sumsq = warp_reduce_sum(total_sumsq); + if (lane == 0) + { + s_reduce[0] = total_sumsq; + } + } + __syncthreads(); + + if (tid < kDimV) + { + int64_t const out_idx = output_head * kDimV + tid; + float const raw_o = s_o[tid]; + float const rstd = rsqrtf(s_reduce[0] / static_cast(kDimV) + onorm_eps); + float gate; + float weight; + if constexpr (kPreloadOnormParams) + { + gate = pre_onorm_gate; + weight = pre_onorm_weight; + } + else + { + gate = sigmoid_fast(bf16_load(onorm_g, out_idx)); + weight = onorm_weight[tid]; + } + float const y = raw_o * rstd * weight * gate; + out[out_idx] = bf16_store(y); + } + } + else + { + float const raw_o = tid < kDimV ? s_o[tid] : 0.0f; + float const o_sq = raw_o * raw_o; + float const sumsq = block_reduce_sum(o_sq, s_reduce); + + if (tid < kDimV) + { + int64_t const out_idx = output_head * kDimV + tid; + float const rstd = rsqrtf(sumsq / static_cast(kDimV) + onorm_eps); + float gate; + float weight; + if constexpr (kPreloadOnormParams) + { + gate = pre_onorm_gate; + weight = pre_onorm_weight; + } + else + { + gate = sigmoid_fast(bf16_load(onorm_g, out_idx)); + weight = onorm_weight[tid]; + } + float const y = raw_o * rstd * weight * gate; + out[out_idx] = bf16_store(y); + } + } + } + else + { + if (tid < kDimV) + { + int64_t const out_idx = output_head * kDimV + tid; + out[out_idx] = bf16_store(s_o[tid]); + } + } +} + +template +__global__ __launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel( + __nv_bfloat16 const* __restrict__ x_q, __nv_bfloat16 const* __restrict__ x_k, __nv_bfloat16 const* __restrict__ x_v, + __nv_bfloat16 const* __restrict__ w_q_t, __nv_bfloat16 const* __restrict__ w_k_t, + __nv_bfloat16 const* __restrict__ w_v_t, __nv_bfloat16 const* __restrict__ bias_q, + __nv_bfloat16 const* __restrict__ bias_k, __nv_bfloat16 const* __restrict__ bias_v, + __nv_bfloat16* __restrict__ cs_q, __nv_bfloat16* __restrict__ cs_k, __nv_bfloat16* __restrict__ cs_v, + float const* __restrict__ a_log, __nv_bfloat16 const* __restrict__ g, float const* __restrict__ dt_bias, + __nv_bfloat16 const* __restrict__ beta, __nv_bfloat16 const* __restrict__ onorm_g, + float const* __restrict__ onorm_weight, int const* __restrict__ ssm_state_indices, + int const* __restrict__ cu_seqlens, float* __restrict__ state, int64_t state_slot_stride, + __nv_bfloat16* __restrict__ out, int B, int H, int HV, float lower_bound, float scale, float onorm_eps) +{ + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int i_n; + int i_hv; + int i_h; + int bos; + int slot; + if constexpr (kUseStaticDecodeLayout) + { + if constexpr (kUseHeadGrid) + { + i_n = blockIdx.x; + i_hv = blockIdx.y; + } + else + { + int const nhv = blockIdx.x; + i_n = nhv / kFixedValueHeads; + i_hv = nhv - i_n * kFixedValueHeads; + } + i_h = i_hv; + bos = i_n; + slot = i_n; + } + else + { + int const nhv = blockIdx.x; + i_n = nhv / HV; + i_hv = nhv - i_n * HV; + int const hv_per_h = HV / H; + i_h = i_hv / hv_per_h; + + bos = cu_seqlens == nullptr ? i_n : cu_seqlens[i_n]; + int const eos = cu_seqlens == nullptr ? i_n + 1 : cu_seqlens[i_n + 1]; + if (eos <= bos) + { + return; + } + slot = ssm_state_indices[i_n]; + } + + int const hk_off = i_h * kDimK; + int const hv_off = i_hv * kDimV; + int const h_count = kUseStaticDecodeLayout ? kFixedHeads : H; + int const hv_count = kUseStaticDecodeLayout ? kFixedValueHeads : HV; + int const hkv_dim = h_count * kDimK; + int const hvv_dim = hv_count * kDimV; + int const conv_slot = kUpdateConvState ? slot : i_n; + int64_t const qk_input_head = static_cast(bos) * h_count + i_h; + int64_t const value_input_head = static_cast(bos) * hv_count + i_hv; + int64_t const output_head = static_cast(i_n) * hv_count + i_hv; + + __shared__ float s_state[2][kChunkV][kDimK]; + __shared__ float s_q[kDimK]; + __shared__ float s_k[kDimK]; + __shared__ float s_decay[kDimK]; + __shared__ float s_v[kDimV]; + __shared__ float s_o[kDimV]; + __shared__ float s_reduce[kThreads]; + __shared__ float s_beta; + float pre_onorm_gate = 0.0f; + float pre_onorm_weight = 0.0f; + + cp_async_state_chunk(&s_state[0][0][0], state, slot, i_hv, state_slot_stride, 0); + + if constexpr (kUpdateConvState) + { + if (tid < kDimK) + { + int const k = tid; + int const hk = hk_off + k; + int64_t const cs_base = static_cast(slot) * hkv_dim * kConvStateWidth + hk; + int64_t const xq_idx = qk_input_head * kDimK + k; + float const exp_a = __shfl_sync(0xffffffffu, lane == 0 ? __expf(a_log[i_h]) : 0.0f, 0); + + float q_acc = bf16_load(bias_q, hk); + float k_acc = bf16_load(bias_k, hk); + __nv_bfloat16 q_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 q_shift1 = __float2bfloat16(0.0f); + __nv_bfloat16 k_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 k_shift1 = __float2bfloat16(0.0f); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) + { + const __nv_bfloat16 q_state = cs_q[cs_base + w * hkv_dim]; + const __nv_bfloat16 k_state = cs_k[cs_base + w * hkv_dim]; + q_acc += __bfloat162float(q_state) * bf16_load(w_q_t, w * hkv_dim + hk); + k_acc += __bfloat162float(k_state) * bf16_load(w_k_t, w * hkv_dim + hk); + if (w == 1) + { + q_shift0 = q_state; + k_shift0 = k_state; + } + else if (w == 2) + { + q_shift1 = q_state; + k_shift1 = k_state; + } + } + const __nv_bfloat16 q_new = x_q[xq_idx]; + const __nv_bfloat16 k_new = x_k[xq_idx]; + q_acc += __bfloat162float(q_new) * bf16_load(w_q_t, (kKernelWidth - 1) * hkv_dim + hk); + k_acc += __bfloat162float(k_new) * bf16_load(w_k_t, (kKernelWidth - 1) * hkv_dim + hk); + + cs_q[cs_base + 0] = q_shift0; + cs_q[cs_base + hkv_dim] = q_shift1; + cs_q[cs_base + 2 * hkv_dim] = q_new; + cs_k[cs_base + 0] = k_shift0; + cs_k[cs_base + hkv_dim] = k_shift1; + cs_k[cs_base + 2 * hkv_dim] = k_new; + + s_q[k] = silu_fast(q_acc); + s_k[k] = silu_fast(k_acc); + + float const g_raw = bf16_load(g, value_input_head * kDimK + k) + dt_bias[hk]; + if constexpr (kUseLowerBound) + { + s_decay[k] = __expf(lower_bound * sigmoid_fast(exp_a * g_raw)); + } + else + { + s_decay[k] = __expf(-exp_a * softplus_fast(g_raw)); + } + } + } + else + { + if (tid < kDimK) + { + int const k = tid; + int const hk = hk_off + k; + float const exp_a = __shfl_sync(0xffffffffu, lane == 0 ? __expf(a_log[i_h]) : 0.0f, 0); + + float q_acc = bf16_load(bias_q, hk); + float k_acc = bf16_load(bias_k, hk); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) + { + int64_t const cs_idx = (static_cast(conv_slot) * hkv_dim + hk) * kConvStateWidth + w; + q_acc += bf16_load(cs_q, cs_idx) * bf16_load(w_q_t, w * hkv_dim + hk); + k_acc += bf16_load(cs_k, cs_idx) * bf16_load(w_k_t, w * hkv_dim + hk); + } + q_acc += bf16_load(x_q, qk_input_head * kDimK + k) * bf16_load(w_q_t, (kKernelWidth - 1) * hkv_dim + hk); + k_acc += bf16_load(x_k, qk_input_head * kDimK + k) * bf16_load(w_k_t, (kKernelWidth - 1) * hkv_dim + hk); + + s_q[k] = silu_fast(q_acc); + s_k[k] = silu_fast(k_acc); + + float const g_raw = bf16_load(g, value_input_head * kDimK + k) + dt_bias[hk]; + if constexpr (kUseLowerBound) + { + s_decay[k] = __expf(lower_bound * sigmoid_fast(exp_a * g_raw)); + } + else + { + s_decay[k] = __expf(-exp_a * softplus_fast(g_raw)); + } + } + } + + if constexpr (kUpdateConvState) + { + if (tid < kDimV) + { + int const v = tid; + int const hvv = hv_off + v; + int64_t const cs_base = static_cast(slot) * hvv_dim * kConvStateWidth + hvv; + int64_t const xv_idx = value_input_head * kDimV + v; + + float v_acc = bf16_load(bias_v, hvv); + __nv_bfloat16 v_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 v_shift1 = __float2bfloat16(0.0f); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) + { + const __nv_bfloat16 v_state = cs_v[cs_base + w * hvv_dim]; + v_acc += __bfloat162float(v_state) * bf16_load(w_v_t, w * hvv_dim + hvv); + if (w == 1) + { + v_shift0 = v_state; + } + else if (w == 2) + { + v_shift1 = v_state; + } + } + const __nv_bfloat16 v_new = x_v[xv_idx]; + v_acc += __bfloat162float(v_new) * bf16_load(w_v_t, (kKernelWidth - 1) * hvv_dim + hvv); + cs_v[cs_base + 0] = v_shift0; + cs_v[cs_base + hvv_dim] = v_shift1; + cs_v[cs_base + 2 * hvv_dim] = v_new; + s_v[v] = silu_fast(v_acc); + + if constexpr (kApplyOnorm && kPreloadOnormParams) + { + int64_t const out_idx = output_head * kDimV + v; + pre_onorm_gate = sigmoid_fast(bf16_load(onorm_g, out_idx)); + pre_onorm_weight = onorm_weight[v]; + } + } + } + else + { + if (tid < kDimV) + { + int const v = tid; + int const hvv = hv_off + v; + + float v_acc = bf16_load(bias_v, hvv); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) + { + int64_t const cs_idx = (static_cast(conv_slot) * hvv_dim + hvv) * kConvStateWidth + w; + v_acc += bf16_load(cs_v, cs_idx) * bf16_load(w_v_t, w * hvv_dim + hvv); + } + v_acc + += bf16_load(x_v, value_input_head * kDimV + v) * bf16_load(w_v_t, (kKernelWidth - 1) * hvv_dim + hvv); + s_v[v] = silu_fast(v_acc); + + if constexpr (kApplyOnorm && kPreloadOnormParams) + { + int64_t const out_idx = output_head * kDimV + v; + pre_onorm_gate = sigmoid_fast(bf16_load(onorm_g, out_idx)); + pre_onorm_weight = onorm_weight[v]; + } + } + } + + if (tid == 0) + { + float const beta_raw = bf16_load(beta, value_input_head); + if constexpr (kApplyBetaSigmoid) + { + s_beta = sigmoid_fast(beta_raw); + } + else + { + s_beta = beta_raw; + } + } + __syncthreads(); + + if constexpr (kPrefetchNextStateChunk && kNumChunks > 1) + { + cp_async_state_chunk(&s_state[0][0][0], state, slot, i_hv, state_slot_stride, 1); + } + + float const q_sq = tid < kDimK ? s_q[tid] * s_q[tid] : 0.0f; + float const k_sq = tid < kDimK ? s_k[tid] * s_k[tid] : 0.0f; + Sum2 qk_sum; + if constexpr (kUseActiveQkReduction) + { + qk_sum = block_reduce_sum2_active_for(q_sq, k_sq, s_reduce); + } + else + { + qk_sum = block_reduce_sum2(q_sq, k_sq, s_reduce); + } + if (tid < kDimK) + { + s_q[tid] *= rsqrtf(qk_sum.x + 1.0e-6f) * scale; + s_k[tid] *= rsqrtf(qk_sum.y + 1.0e-6f); + } + __syncthreads(); + + int const k_base = lane * 4; + const float4 q4 = *reinterpret_cast(s_q + k_base); + const float4 k4 = *reinterpret_cast(s_k + k_base); + const float4 decay4 = *reinterpret_cast(s_decay + k_base); + float r_q[4] = {q4.x, q4.y, q4.z, q4.w}; + float r_k[4] = {k4.x, k4.y, k4.z, k4.w}; + float r_decay[4] = {decay4.x, decay4.y, decay4.z, decay4.w}; + float o_sumsq = 0.0f; + +#pragma unroll + for (int chunk = 0; chunk < kNumChunks; ++chunk) + { + if constexpr (kPrefetchNextStateChunk && kNumChunks > 1) + { + if (chunk + 1 < kNumChunks) + { + cp_async_wait_group_1(); + } + else + { + cp_async_wait_all(); + } + } + else + { + cp_async_wait_all(); + } + if constexpr (!kSkipWarpSync) + { + __syncwarp(); + } + + if constexpr (!kPrefetchNextStateChunk) + { + if (chunk + 1 < kNumChunks) + { + cp_async_state_chunk(&s_state[0][0][0], state, slot, i_hv, state_slot_stride, chunk + 1); + } + } + +#pragma unroll + for (int row = 0; row < kRowsPerWarp; row += 2) + { + int const v_row_a = warp + row * kWarps; + int const v_row_b = warp + (row + 1) * kWarps; + int const v0 = chunk * kChunkV + v_row_a; + int const v1 = chunk * kChunkV + v_row_b; + float h_a_vals[4]; + float h_b_vals[4]; + float dot_hk_a = 0.0f; + float dot_hk_b = 0.0f; + + const float4 raw_h_a = *reinterpret_cast(&s_state[chunk & 1][v_row_a][k_base]); + const float4 raw_h_b = *reinterpret_cast(&s_state[chunk & 1][v_row_b][k_base]); + h_a_vals[0] = raw_h_a.x * r_decay[0]; + h_a_vals[1] = raw_h_a.y * r_decay[1]; + h_a_vals[2] = raw_h_a.z * r_decay[2]; + h_a_vals[3] = raw_h_a.w * r_decay[3]; + h_b_vals[0] = raw_h_b.x * r_decay[0]; + h_b_vals[1] = raw_h_b.y * r_decay[1]; + h_b_vals[2] = raw_h_b.z * r_decay[2]; + h_b_vals[3] = raw_h_b.w * r_decay[3]; + dot_hk_a = h_a_vals[0] * r_k[0] + h_a_vals[1] * r_k[1] + h_a_vals[2] * r_k[2] + h_a_vals[3] * r_k[3]; + dot_hk_b = h_b_vals[0] * r_k[0] + h_b_vals[1] * r_k[1] + h_b_vals[2] * r_k[2] + h_b_vals[3] * r_k[3]; + + const Sum2 dot_hk = warp_reduce_sum_pair(dot_hk_a, dot_hk_b); + float const v_new0 = (s_v[v0] - dot_hk.x) * s_beta; + float const v_new1 = (s_v[v1] - dot_hk.y) * s_beta; + + float dot_hq_a = 0.0f; + float dot_hq_b = 0.0f; + int64_t const state_head_offset + = static_cast(slot) * state_slot_stride + static_cast(i_hv) * kDimV * kDimK; + int64_t const state_idx_a = state_head_offset + static_cast(v0) * kDimK + k_base; + int64_t const state_idx_b = state_head_offset + static_cast(v1) * kDimK + k_base; + float const h_a_0 = h_a_vals[0] + r_k[0] * v_new0; + float const h_a_1 = h_a_vals[1] + r_k[1] * v_new0; + float const h_a_2 = h_a_vals[2] + r_k[2] * v_new0; + float const h_a_3 = h_a_vals[3] + r_k[3] * v_new0; + float const h_b_0 = h_b_vals[0] + r_k[0] * v_new1; + float const h_b_1 = h_b_vals[1] + r_k[1] * v_new1; + float const h_b_2 = h_b_vals[2] + r_k[2] * v_new1; + float const h_b_3 = h_b_vals[3] + r_k[3] * v_new1; + if constexpr (kComputeOutputBeforeStore) + { + dot_hq_a = h_a_0 * r_q[0] + h_a_1 * r_q[1] + h_a_2 * r_q[2] + h_a_3 * r_q[3]; + dot_hq_b = h_b_0 * r_q[0] + h_b_1 * r_q[1] + h_b_2 * r_q[2] + h_b_3 * r_q[3]; + store_state_float4(state + state_idx_a, make_float4(h_a_0, h_a_1, h_a_2, h_a_3)); + store_state_float4(state + state_idx_b, make_float4(h_b_0, h_b_1, h_b_2, h_b_3)); + } + else + { + store_state_float4(state + state_idx_a, make_float4(h_a_0, h_a_1, h_a_2, h_a_3)); + store_state_float4(state + state_idx_b, make_float4(h_b_0, h_b_1, h_b_2, h_b_3)); + dot_hq_a = h_a_0 * r_q[0] + h_a_1 * r_q[1] + h_a_2 * r_q[2] + h_a_3 * r_q[3]; + dot_hq_b = h_b_0 * r_q[0] + h_b_1 * r_q[1] + h_b_2 * r_q[2] + h_b_3 * r_q[3]; + } + + const Sum2 dot_hq = warp_reduce_sum_pair(dot_hq_a, dot_hq_b); + if (lane == 0) + { + s_o[v0] = dot_hq.x; + s_o[v1] = dot_hq.y; + if constexpr (kApplyOnorm && kAccumulateOnormSumsq) + { + o_sumsq += dot_hq.x * dot_hq.x + dot_hq.y * dot_hq.y; + } + } + } + + if constexpr (kPrefetchNextStateChunk) + { + if (chunk + 2 < kNumChunks) + { + cp_async_state_chunk(&s_state[0][0][0], state, slot, i_hv, state_slot_stride, chunk + 2); + } + } + } + __syncthreads(); + + if constexpr (kApplyOnorm) + { + if constexpr (kAccumulateOnormSumsq) + { + if (lane == 0) + { + s_reduce[warp] = o_sumsq; + } + __syncthreads(); + + float total_sumsq = 0.0f; + if (warp == 0) + { + total_sumsq = lane < kWarps ? s_reduce[lane] : 0.0f; + total_sumsq = warp_reduce_sum(total_sumsq); + if (lane == 0) + { + s_reduce[0] = total_sumsq; + } + } + __syncthreads(); + + if (tid < kDimV) + { + int64_t const out_idx = output_head * kDimV + tid; + float const raw_o = s_o[tid]; + float const rstd = rsqrtf(s_reduce[0] / static_cast(kDimV) + onorm_eps); + float gate; + float weight; + if constexpr (kPreloadOnormParams) + { + gate = pre_onorm_gate; + weight = pre_onorm_weight; + } + else + { + gate = sigmoid_fast(bf16_load(onorm_g, out_idx)); + weight = onorm_weight[tid]; + } + float const y = raw_o * rstd * weight * gate; + out[out_idx] = bf16_store(y); + } + } + else + { + float const raw_o = tid < kDimV ? s_o[tid] : 0.0f; + float const o_sq = raw_o * raw_o; + float sumsq; + if constexpr (kUseActiveOnormReduction || kUseActiveQkReduction) + { + sumsq = block_reduce_sum_active_for(o_sq, s_reduce); + } + else + { + sumsq = block_reduce_sum(o_sq, s_reduce); + } + + if (tid < kDimV) + { + int64_t const out_idx = output_head * kDimV + tid; + float const rstd = rsqrtf(sumsq / static_cast(kDimV) + onorm_eps); + float gate; + float weight; + if constexpr (kPreloadOnormParams) + { + gate = pre_onorm_gate; + weight = pre_onorm_weight; + } + else + { + gate = sigmoid_fast(bf16_load(onorm_g, out_idx)); + weight = onorm_weight[tid]; + } + float const y = raw_o * rstd * weight * gate; + out[out_idx] = bf16_store(y); + } + } + } + else + { + if (tid < kDimV) + { + int64_t const out_idx = output_head * kDimV + tid; + out[out_idx] = bf16_store(s_o[tid]); + } + } +} + +struct KdaDecodeLaunchParams +{ + void const* x_q; + void const* x_k; + void const* x_v; + void const* w_q_t; + void const* w_k_t; + void const* w_v_t; + void const* bias_q; + void const* bias_k; + void const* bias_v; + void* cs_q; + void* cs_k; + void* cs_v; + float const* a_log; + void const* g; + float const* dt_bias; + void const* beta; + void const* onorm_g; + float const* onorm_weight; + int const* ssm_state_indices; + int const* cu_seqlens; + float* state; + int64_t state_slot_stride; + void* out; + int B; + int H; + int HV; + bool apply_onorm; + bool update_conv_cache; + bool use_lower_bound; + bool apply_beta_sigmoid; + float lower_bound; + float scale; + float onorm_eps; + cudaStream_t stream; +}; + +template +void launch_kda_decode_compact_heads_raw(KdaDecodeLaunchParams const& p) +{ + constexpr int kStageDynamicSmemBytes = 3 * 32 * kDimK * static_cast(sizeof(float)); + TLLM_CUDA_CHECK(cudaFuncSetAttribute( + kda_decode_fusion_compact_heads_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, kStageDynamicSmemBytes)); + kda_decode_fusion_compact_heads_kernel + <<>>( + reinterpret_cast<__nv_bfloat16 const*>(p.x_q), reinterpret_cast<__nv_bfloat16 const*>(p.x_k), + reinterpret_cast<__nv_bfloat16 const*>(p.x_v), reinterpret_cast<__nv_bfloat16 const*>(p.w_q_t), + reinterpret_cast<__nv_bfloat16 const*>(p.w_k_t), reinterpret_cast<__nv_bfloat16 const*>(p.w_v_t), + reinterpret_cast<__nv_bfloat16 const*>(p.bias_q), reinterpret_cast<__nv_bfloat16 const*>(p.bias_k), + reinterpret_cast<__nv_bfloat16 const*>(p.bias_v), reinterpret_cast<__nv_bfloat16*>(p.cs_q), + reinterpret_cast<__nv_bfloat16*>(p.cs_k), reinterpret_cast<__nv_bfloat16*>(p.cs_v), p.a_log, + reinterpret_cast<__nv_bfloat16 const*>(p.g), p.dt_bias, reinterpret_cast<__nv_bfloat16 const*>(p.beta), + reinterpret_cast<__nv_bfloat16 const*>(p.onorm_g), p.onorm_weight, p.ssm_state_indices, p.cu_seqlens, + p.state, p.state_slot_stride, reinterpret_cast<__nv_bfloat16*>(p.out), p.B, p.H, p.HV, p.lower_bound, + p.scale, p.onorm_eps); +} + +template +void launch_kda_decode_many_heads_raw(KdaDecodeLaunchParams const& p) +{ + constexpr bool kUseHeadGrid = kUseStaticDecodeLayout; + dim3 const grid = kUseStaticDecodeLayout ? dim3(p.B, p.HV) : dim3(p.B * p.HV); + kda_decode_fusion_many_heads_kernel + <<>>(reinterpret_cast<__nv_bfloat16 const*>(p.x_q), + reinterpret_cast<__nv_bfloat16 const*>(p.x_k), reinterpret_cast<__nv_bfloat16 const*>(p.x_v), + reinterpret_cast<__nv_bfloat16 const*>(p.w_q_t), reinterpret_cast<__nv_bfloat16 const*>(p.w_k_t), + reinterpret_cast<__nv_bfloat16 const*>(p.w_v_t), reinterpret_cast<__nv_bfloat16 const*>(p.bias_q), + reinterpret_cast<__nv_bfloat16 const*>(p.bias_k), reinterpret_cast<__nv_bfloat16 const*>(p.bias_v), + reinterpret_cast<__nv_bfloat16*>(p.cs_q), reinterpret_cast<__nv_bfloat16*>(p.cs_k), + reinterpret_cast<__nv_bfloat16*>(p.cs_v), p.a_log, reinterpret_cast<__nv_bfloat16 const*>(p.g), p.dt_bias, + reinterpret_cast<__nv_bfloat16 const*>(p.beta), reinterpret_cast<__nv_bfloat16 const*>(p.onorm_g), + p.onorm_weight, p.ssm_state_indices, p.cu_seqlens, p.state, p.state_slot_stride, + reinterpret_cast<__nv_bfloat16*>(p.out), p.B, p.H, p.HV, p.lower_bound, p.scale, p.onorm_eps); +} + +template +void launch_kda_decode_raw(KdaDecodeLaunchParams const& p) +{ + if constexpr (kCompact) + { + launch_kda_decode_compact_heads_raw(p); + } + else + { + launch_kda_decode_many_heads_raw(p); + } +} + +template +void launch_kda_decode_selected_backend(KdaDecodeLaunchParams const& p) +{ + if (p.update_conv_cache) + { + launch_kda_decode_raw(p); + } + else + { + launch_kda_decode_raw(p); + } +} + +template +void dispatch_kda_decode_beta(KdaDecodeLaunchParams const& p) +{ + if (p.apply_beta_sigmoid) + { + launch_kda_decode_selected_backend( + p); + } + else + { + launch_kda_decode_selected_backend(p); + } +} + +template +void dispatch_kda_decode_decay(KdaDecodeLaunchParams const& p) +{ + if (p.use_lower_bound) + { + dispatch_kda_decode_beta(p); + } + else + { + dispatch_kda_decode_beta(p); + } +} + +template +void dispatch_kda_decode_features(KdaDecodeLaunchParams const& p) +{ + if (p.apply_onorm) + { + dispatch_kda_decode_decay(p); + } + else + { + dispatch_kda_decode_decay(p); + } +} + +template +void dispatch_kda_decode_layout(KdaDecodeLaunchParams const& p) +{ + if (p.ssm_state_indices == nullptr) + { + dispatch_kda_decode_features(p); + } + else + { + dispatch_kda_decode_features(p); + } +} + +template +void dispatch_kda_decode_heads(KdaDecodeLaunchParams const& p) +{ + switch (p.H) + { + case 1: dispatch_kda_decode_layout(p); break; + case 2: dispatch_kda_decode_layout(p); break; + case 3: dispatch_kda_decode_layout(p); break; + case 4: dispatch_kda_decode_layout(p); break; + case 6: dispatch_kda_decode_layout(p); break; + case 8: dispatch_kda_decode_layout(p); break; + case 12: dispatch_kda_decode_layout(p); break; + case 16: dispatch_kda_decode_layout(p); break; + case 24: dispatch_kda_decode_layout(p); break; + case 32: dispatch_kda_decode_layout(p); break; + case 48: dispatch_kda_decode_layout(p); break; + case 96: dispatch_kda_decode_layout(p); break; + default: + if constexpr (kCompact) + { + TLLM_CHECK_WITH_INFO(false, "KDA compact-heads decode does not support numHeads=%d", p.H); + } + else + { + TLLM_CHECK_WITH_INFO(false, "KDA decode does not support numHeads=%d", p.H); + } + } +} + +} // namespace + +void invokeKdaDecode(KdaDecodeParams const& params, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(params.numHeads == params.numValueHeads, "KDA decode requires numHeads == numValueHeads"); + bool const useCompactHeads = shouldUseCompactHeads(params.batchSize, params.numHeads, params.numValueHeads); + KdaDecodeLaunchParams const launchParams{params.xQ, params.xK, params.xV, params.wQT, params.wKT, params.wVT, + params.biasQ, params.biasK, params.biasV, params.convStateQ, params.convStateK, params.convStateV, params.logA, + params.gate, params.dtBias, params.beta, params.outputNormGate, params.outputNormWeight, params.ssmStateIndices, + params.cuSeqlens, params.state, params.stateSlotStride, params.output, params.batchSize, params.numHeads, + params.numValueHeads, params.applyOutputNorm, params.updateConvCache, params.useLowerBound, + params.applyBetaSigmoid, params.lowerBound, params.scale, params.outputNormEps, stream}; + if (useCompactHeads) + { + dispatch_kda_decode_heads(launchParams); + } + else + { + dispatch_kda_decode_heads(launchParams); + } + TLLM_CUDA_CHECK(cudaGetLastError()); +} + +} // namespace kernels::kdaDecode + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.h b/cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.h new file mode 100644 index 000000000000..4d030ce8e4f5 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/kdaDecode/kdaDecode.h @@ -0,0 +1,89 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::kdaDecode +{ + +constexpr int kCompactHeadsWorkThreshold = 144; + +constexpr bool isSupportedHeadCount(int numHeads) +{ + return numHeads == 1 || numHeads == 2 || numHeads == 3 || numHeads == 4 || numHeads == 6 || numHeads == 8 + || numHeads == 12 || numHeads == 16 || numHeads == 24 || numHeads == 32 || numHeads == 48 || numHeads == 96; +} + +//! Select the compact-head kernel within the measured KDA decode work threshold. +//! Division keeps the B*H threshold overflow-safe. +constexpr bool shouldUseCompactHeads(int batchSize, int numHeads, int numValueHeads) +{ + return batchSize > 0 && numHeads == numValueHeads && isSupportedHeadCount(numHeads) + && batchSize <= kCompactHeadsWorkThreshold / numHeads; +} + +//! Parameters for the fused, single-token KDA decode kernel. +struct KdaDecodeParams +{ + void const* xQ; + void const* xK; + void const* xV; + void const* wQT; + void const* wKT; + void const* wVT; + void const* biasQ; + void const* biasK; + void const* biasV; + void* convStateQ; + void* convStateK; + void* convStateV; + float const* logA; // Named differently from a_log to satisfy codespell. + void const* gate; + float const* dtBias; + void const* beta; + void const* outputNormGate; + float const* outputNormWeight; + int const* ssmStateIndices; + //! Must be arange(batchSize + 1): the kernel only advances each state by one token. + int const* cuSeqlens; + float* state; + int64_t stateSlotStride; + void* output; + int batchSize; + int numHeads; + int numValueHeads; + bool applyOutputNorm; + bool updateConvCache; + bool useLowerBound; + bool applyBetaSigmoid; + float lowerBound; + float scale; + float outputNormEps; +}; + +//! Launches the tuned KDA decode kernel on the supplied CUDA stream. +void invokeKdaDecode(KdaDecodeParams const& params, cudaStream_t stream); + +} // namespace kernels::kdaDecode + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 4ba8cd88bd9c..f64b815e4240 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -75,6 +75,7 @@ add_library( deepseekV4QNormOp.cpp deepseekV4BlockTableOp.cpp inverseRopeFp8QuantOp.cpp + kdaDecodeOp.cpp fusedQKNormRopeOp.cpp fusedAdaptiveLayerNormOp.cpp fusedDiTQKNormRopeOp.cpp diff --git a/cpp/tensorrt_llm/thop/kdaDecodeOp.cpp b/cpp/tensorrt_llm/thop/kdaDecodeOp.cpp new file mode 100644 index 000000000000..081f0eed9945 --- /dev/null +++ b/cpp/tensorrt_llm/thop/kdaDecodeOp.cpp @@ -0,0 +1,245 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/kdaDecode/kdaDecode.h" + +#include +#include +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int kDimK = 128; +constexpr int kDimV = 128; +constexpr int kKernelWidth = 4; + +void validate_kda_decode_fusion_inputs(at::Tensor x_q, at::Tensor x_k, at::Tensor x_v, at::Tensor w_q_t, + at::Tensor w_k_t, at::Tensor w_v_t, at::Tensor bias_q, at::Tensor bias_k, at::Tensor bias_v, at::Tensor cs_q, + at::Tensor cs_k, at::Tensor cs_v, at::Tensor a_log, at::Tensor g, at::Tensor dt_bias, at::Tensor beta, + at::Tensor onorm_g, at::Tensor onorm_weight, std::optional const& ssm_state_indices, + at::Tensor cu_seqlens, at::Tensor state, bool apply_onorm, bool update_conv_cache) +{ + TORCH_CHECK(x_q.is_cuda() && x_q.scalar_type() == at::kBFloat16, "x_q must be a CUDA bfloat16 tensor"); + TORCH_CHECK(x_k.is_cuda() && x_k.scalar_type() == at::kBFloat16, "x_k must be a CUDA bfloat16 tensor"); + TORCH_CHECK(x_v.is_cuda() && x_v.scalar_type() == at::kBFloat16, "x_v must be a CUDA bfloat16 tensor"); + TORCH_CHECK(w_q_t.is_cuda() && w_q_t.scalar_type() == at::kBFloat16, "w_q_t must be a CUDA bfloat16 tensor"); + TORCH_CHECK(w_k_t.is_cuda() && w_k_t.scalar_type() == at::kBFloat16, "w_k_t must be a CUDA bfloat16 tensor"); + TORCH_CHECK(w_v_t.is_cuda() && w_v_t.scalar_type() == at::kBFloat16, "w_v_t must be a CUDA bfloat16 tensor"); + TORCH_CHECK(bias_q.is_cuda() && bias_q.scalar_type() == at::kBFloat16, "bias_q must be a CUDA bfloat16 tensor"); + TORCH_CHECK(bias_k.is_cuda() && bias_k.scalar_type() == at::kBFloat16, "bias_k must be a CUDA bfloat16 tensor"); + TORCH_CHECK(bias_v.is_cuda() && bias_v.scalar_type() == at::kBFloat16, "bias_v must be a CUDA bfloat16 tensor"); + TORCH_CHECK(cs_q.is_cuda() && cs_q.scalar_type() == at::kBFloat16, "cs_q must be a CUDA bfloat16 tensor"); + TORCH_CHECK(cs_k.is_cuda() && cs_k.scalar_type() == at::kBFloat16, "cs_k must be a CUDA bfloat16 tensor"); + TORCH_CHECK(cs_v.is_cuda() && cs_v.scalar_type() == at::kBFloat16, "cs_v must be a CUDA bfloat16 tensor"); + TORCH_CHECK(g.is_cuda() && g.scalar_type() == at::kBFloat16, "g must be a CUDA bfloat16 tensor"); + TORCH_CHECK(beta.is_cuda() && beta.scalar_type() == at::kBFloat16, "beta must be a CUDA bfloat16 tensor"); + TORCH_CHECK(onorm_g.is_cuda() && onorm_g.scalar_type() == at::kBFloat16, "onorm_g must be a CUDA bfloat16 tensor"); + TORCH_CHECK(a_log.is_cuda() && a_log.scalar_type() == at::kFloat, "a_log must be a CUDA float32 tensor"); + TORCH_CHECK(dt_bias.is_cuda() && dt_bias.scalar_type() == at::kFloat, "dt_bias must be a CUDA float32 tensor"); + TORCH_CHECK(onorm_weight.is_cuda() && onorm_weight.scalar_type() == at::kFloat, + "onorm_weight must be a CUDA float32 tensor"); + TORCH_CHECK(state.is_cuda() && state.scalar_type() == at::kFloat, "state must be a CUDA float32 tensor"); + + TORCH_CHECK(x_q.dim() == 4 && x_k.dim() == 4 && x_v.dim() == 4, "x_q, x_k, and x_v must be rank-4 tensors"); + TORCH_CHECK(x_q.size(0) == 1 && x_k.size(0) == 1 && x_v.size(0) == 1, "only T=1 decode inputs are supported"); + TORCH_CHECK(x_q.size(3) == kDimK && x_k.size(3) == kDimK, "only K=128 is supported"); + TORCH_CHECK(x_v.size(3) == kDimV, "only V=128 is supported"); + TORCH_CHECK( + x_q.is_contiguous() && x_k.is_contiguous() && x_v.is_contiguous(), "x_q, x_k, and x_v must be contiguous"); + TORCH_CHECK(w_q_t.dim() == 2 && w_k_t.dim() == 2 && w_v_t.dim() == 2, "w_q_t, w_k_t, and w_v_t must be rank-2"); + TORCH_CHECK(w_q_t.size(0) == kKernelWidth && w_k_t.size(0) == kKernelWidth && w_v_t.size(0) == kKernelWidth, + "only convolution width 4 is supported"); + TORCH_CHECK(w_q_t.is_contiguous() && w_k_t.is_contiguous() && w_v_t.is_contiguous(), + "w_q_t, w_k_t, and w_v_t must be contiguous [4, dim] tensors"); + + int const B = static_cast(x_q.size(1)); + int const H = static_cast(x_q.size(2)); + int const HV = static_cast(x_v.size(2)); + TORCH_CHECK(B > 0, "KDA decode requires a non-empty batch"); + bool const supportedHeads = H == 1 || H == 2 || H == 3 || H == 4 || H == 6 || H == 8 || H == 12 || H == 16 + || H == 24 || H == 32 || H == 48 || H == 96; + TORCH_CHECK( + H == HV && supportedHeads, "KDA decode fusion CUDA supports H == HV in {1,2,3,4,6,8,12,16,24,32,48,96}"); + TORCH_CHECK(x_k.size(1) == B && x_k.size(2) == H && x_v.size(1) == B, + "x_q, x_k, and x_v batch/head dimensions are inconsistent"); + TORCH_CHECK(HV % H == 0, "HV must be divisible by H"); + + int64_t const qk_dim = static_cast(H) * kDimK; + int64_t const v_dim = static_cast(HV) * kDimV; + TORCH_CHECK(w_q_t.size(1) == qk_dim && w_k_t.size(1) == qk_dim && w_v_t.size(1) == v_dim, + "w_q_t and w_k_t must be [4, H*128], w_v_t must be [4, HV*128]"); + TORCH_CHECK(bias_q.is_contiguous() && bias_k.is_contiguous() && bias_v.is_contiguous(), + "bias_q, bias_k, and bias_v must be contiguous"); + TORCH_CHECK(bias_q.numel() == qk_dim && bias_k.numel() == qk_dim && bias_v.numel() == v_dim, + "bias_q and bias_k must hold H*128 elements, bias_v must hold HV*128 elements"); + TORCH_CHECK(a_log.is_contiguous() && a_log.numel() == H, "a_log must be contiguous with H elements"); + TORCH_CHECK(dt_bias.is_contiguous() && dt_bias.numel() == qk_dim, "dt_bias must be contiguous with H*128 elements"); + TORCH_CHECK( + g.is_contiguous() && g.dim() == 4 && g.size(0) == 1 && g.size(1) == B && g.size(2) == HV && g.size(3) == kDimK, + "g must be a contiguous [1, B, HV, 128] tensor"); + TORCH_CHECK(beta.is_contiguous() && beta.dim() == 3 && beta.size(0) == 1 && beta.size(1) == B && beta.size(2) == HV, + "beta must be a contiguous [1, B, HV] tensor"); + if (apply_onorm) + { + TORCH_CHECK(onorm_g.is_contiguous() && onorm_g.dim() == 4 && onorm_g.size(0) == 1 && onorm_g.size(1) == B + && onorm_g.size(2) == HV && onorm_g.size(3) == kDimV, + "onorm_g must be a contiguous [1, B, HV, 128] tensor when apply_onorm is set"); + TORCH_CHECK(onorm_weight.is_contiguous() && onorm_weight.numel() == kDimV, + "onorm_weight must be contiguous with 128 elements when apply_onorm is set"); + } + + TORCH_CHECK(state.dim() == 4 && state.size(0) >= B && state.size(1) == HV && state.size(2) == kDimV + && state.size(3) == kDimK, + "state must have shape [slots, HV, 128, 128] with slots >= B"); + TORCH_CHECK(state.stride(3) == 1 && state.stride(2) == kDimK && state.stride(1) == kDimV * kDimK + && state.stride(0) >= HV * kDimV * kDimK, + "state must be contiguous within each [HV, 128, 128] slot and have a non-overlapping slot stride"); + // The kernel moves recurrent state with 16B cp.async loads and float4 stores at element + // offsets of `slot * stride(0) + `, so both the slot stride and the base + // pointer have to keep those accesses 16B aligned. + TORCH_CHECK(state.stride(0) % 4 == 0, + "state slot stride must be a multiple of 4 floats so that per-slot float4 accesses stay 16B aligned, got ", + state.stride(0)); + TORCH_CHECK(reinterpret_cast(state.data_ptr()) % 16 == 0, + "state must start at a 16B-aligned address (check the storage offset of the view passed in)"); + if (ssm_state_indices.has_value()) + { + TORCH_CHECK(ssm_state_indices->is_cuda() && ssm_state_indices->scalar_type() == at::kInt, + "ssm_state_indices must be a CUDA int32 tensor"); + TORCH_CHECK( + ssm_state_indices->is_contiguous() && ssm_state_indices->dim() == 1 && ssm_state_indices->size(0) == B, + "ssm_state_indices must be contiguous with shape [B]"); + } + TORCH_CHECK(cu_seqlens.is_cuda() && cu_seqlens.scalar_type() == at::kInt, "cu_seqlens must be a CUDA int32 tensor"); + TORCH_CHECK(cu_seqlens.is_contiguous() && cu_seqlens.dim() == 1 && cu_seqlens.size(0) == B + 1, + "cu_seqlens must be contiguous with shape [B + 1]"); + + if (update_conv_cache) + { + TORCH_CHECK(H == HV, "conv state update currently assumes H == HV"); + TORCH_CHECK( + cs_q.dim() == 3 && cs_k.dim() == 3 && cs_v.dim() == 3, "update_conv_cache expects rank-3 conv-state pools"); + TORCH_CHECK(cs_q.size(0) >= state.size(0) && cs_k.size(0) >= state.size(0) && cs_v.size(0) >= state.size(0), + "conv-state pools must cover every recurrent-state slot"); + TORCH_CHECK(cs_q.size(1) == H * kDimK && cs_k.size(1) == H * kDimK && cs_v.size(1) == HV * kDimV + && cs_q.size(2) == kKernelWidth - 1 && cs_k.size(2) == kKernelWidth - 1 + && cs_v.size(2) == kKernelWidth - 1, + "update_conv_cache expects [slots, dim, 3] conv-state pools"); + TORCH_CHECK( + cs_q.stride(0) == 3 * H * kDimK && cs_k.stride(0) == 3 * H * kDimK && cs_v.stride(0) == 3 * HV * kDimV, + "update_conv_cache expects densely packed conv-state slots"); + TORCH_CHECK(cs_q.stride(1) == 1 && cs_k.stride(1) == 1, + "update_conv_cache expects cs_q/cs_k transposed layout with " + "contiguous dim axis"); + TORCH_CHECK(cs_q.stride(2) == H * kDimK && cs_k.stride(2) == H * kDimK, + "update_conv_cache expects cs_q/cs_k token stride H*K"); + TORCH_CHECK(cs_v.stride(1) == 1, + "update_conv_cache expects cs_v transposed layout with " + "contiguous dim axis"); + TORCH_CHECK(cs_v.stride(2) == HV * kDimV, "update_conv_cache expects cs_v token stride HV*V"); + } + else + { + TORCH_CHECK(cs_q.dim() == 3 && cs_k.dim() == 3 && cs_v.dim() == 3, "batch-local conv states must be rank-3"); + TORCH_CHECK(cs_q.size(0) == B && cs_k.size(0) == B && cs_v.size(0) == B && cs_q.size(1) == H * kDimK + && cs_k.size(1) == H * kDimK && cs_v.size(1) == HV * kDimV && cs_q.size(2) == kKernelWidth - 1 + && cs_k.size(2) == kKernelWidth - 1 && cs_v.size(2) == kKernelWidth - 1, + "batch-local conv states must have shapes [B, H*128, 3], " + "[B, H*128, 3], and [B, HV*128, 3]"); + TORCH_CHECK(cs_q.is_contiguous(), "cs_q must be contiguous [B, H*K, 3]"); + TORCH_CHECK(cs_k.is_contiguous(), "cs_k must be contiguous [B, H*K, 3]"); + TORCH_CHECK(cs_v.is_contiguous(), "cs_v must be contiguous [B, HV*V, 3]"); + } +} + +void launch_selected_kernel(at::Tensor x_q, at::Tensor x_k, at::Tensor x_v, at::Tensor w_q_t, at::Tensor w_k_t, + at::Tensor w_v_t, at::Tensor bias_q, at::Tensor bias_k, at::Tensor bias_v, at::Tensor cs_q, at::Tensor cs_k, + at::Tensor cs_v, at::Tensor a_log, at::Tensor g, at::Tensor dt_bias, at::Tensor beta, at::Tensor onorm_g, + at::Tensor onorm_weight, std::optional const& ssm_state_indices, at::Tensor cu_seqlens, + at::Tensor state, at::Tensor out, bool apply_onorm, bool update_conv_cache, bool use_lower_bound, + bool apply_beta_sigmoid, double lower_bound, double scale, double onorm_eps) +{ + int const B = static_cast(x_q.size(1)); + int const H = static_cast(x_q.size(2)); + int const HV = static_cast(x_v.size(2)); + + tensorrt_llm::kernels::kdaDecode::KdaDecodeParams const params{x_q.data_ptr(), x_k.data_ptr(), x_v.data_ptr(), + w_q_t.data_ptr(), w_k_t.data_ptr(), w_v_t.data_ptr(), bias_q.data_ptr(), bias_k.data_ptr(), bias_v.data_ptr(), + cs_q.data_ptr(), cs_k.data_ptr(), cs_v.data_ptr(), a_log.data_ptr(), g.data_ptr(), + dt_bias.data_ptr(), beta.data_ptr(), onorm_g.data_ptr(), onorm_weight.data_ptr(), + ssm_state_indices.has_value() ? ssm_state_indices->data_ptr() : nullptr, cu_seqlens.data_ptr(), + state.data_ptr(), state.stride(0), out.data_ptr(), B, H, HV, apply_onorm, update_conv_cache, + use_lower_bound, apply_beta_sigmoid, static_cast(lower_bound), static_cast(scale), + static_cast(onorm_eps)}; + + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + tensorrt_llm::kernels::kdaDecode::invokeKdaDecode(params, stream); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +at::Tensor kda_decode_fusion_forward(at::Tensor x_q, at::Tensor x_k, at::Tensor x_v, at::Tensor w_q_t, at::Tensor w_k_t, + at::Tensor w_v_t, at::Tensor bias_q, at::Tensor bias_k, at::Tensor bias_v, at::Tensor cs_q, at::Tensor cs_k, + at::Tensor cs_v, at::Tensor a_log, at::Tensor g, at::Tensor dt_bias, at::Tensor beta, at::Tensor onorm_g, + at::Tensor onorm_weight, std::optional ssm_state_indices, at::Tensor cu_seqlens, at::Tensor state, + bool apply_onorm, bool update_conv_cache, bool use_lower_bound, bool apply_beta_sigmoid, double lower_bound, + double scale, double onorm_eps) +{ + validate_kda_decode_fusion_inputs(x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, cs_k, cs_v, + a_log, g, dt_bias, beta, onorm_g, onorm_weight, ssm_state_indices, cu_seqlens, state, apply_onorm, + update_conv_cache); + int const B = static_cast(x_q.size(1)); + int const HV = static_cast(x_v.size(2)); + auto out = at::empty({B, 1, HV, kDimV}, x_q.options()); + launch_selected_kernel(x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, cs_k, cs_v, a_log, g, + dt_bias, beta, onorm_g, onorm_weight, ssm_state_indices, cu_seqlens, state, out, apply_onorm, update_conv_cache, + use_lower_bound, apply_beta_sigmoid, lower_bound, scale, onorm_eps); + return out; +} + +} // namespace + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "kda_decode(Tensor x_q, Tensor x_k, Tensor x_v, Tensor w_q_t, " + "Tensor w_k_t, Tensor w_v_t, Tensor bias_q, Tensor bias_k, " + "Tensor bias_v, Tensor(a!) conv_state_q, Tensor(b!) conv_state_k, " + "Tensor(c!) conv_state_v, Tensor a_log, Tensor g, Tensor dt_bias, " + "Tensor beta, Tensor onorm_g, Tensor onorm_weight, " + "Tensor? ssm_state_indices, Tensor cu_seqlens, Tensor(d!) state, " + "bool apply_onorm, bool update_conv_cache, bool use_lower_bound, " + "bool apply_beta_sigmoid, float lower_bound, float scale, " + "float onorm_eps) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("kda_decode", &tensorrt_llm::torch_ext::kda_decode_fusion_forward); +} diff --git a/requirements-dev.txt b/requirements-dev.txt index 57d64fb0436c..bd98a550e1a2 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -4,6 +4,8 @@ diffusers==0.38.0 boto3 einops +# Reference implementation for KDA kernel parity tests. +flash-linear-attention==0.5.2 lpips graphviz mypy==1.19.1 diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index a9ebd7674314..1cd00b3e9b6c 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -276,6 +276,15 @@ def _(logits, # In-place operation, no return value (void function) pass + @torch.library.register_fake("trtllm::kda_decode") + def _(x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, + conv_state_q, conv_state_k, conv_state_v, a_log, g, dt_bias, beta, + onorm_g, onorm_weight, ssm_state_indices, cu_seqlens, state, + apply_onorm, update_conv_cache, use_lower_bound, apply_beta_sigmoid, + lower_bound, scale, onorm_eps): + # x_q is [1, tokens, H, 128]; the kernel emits one row per token. + return x_q.new_empty((x_q.size(1), 1, x_v.size(2), x_v.size(3))) + @torch.library.register_fake("trtllm::userbuffers_allreduce_finalize") def _(input, force_applying_finalize): return torch.empty_like(input) diff --git a/tests/unittest/_torch/thop/parallel/test_kda_decode.py b/tests/unittest/_torch/thop/parallel/test_kda_decode.py new file mode 100644 index 000000000000..35a85523000d --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/test_kda_decode.py @@ -0,0 +1,457 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FLA parity tests for the fused KDA single-token decode kernel.""" + +from dataclasses import dataclass + +import pytest +import torch +import torch.nn.functional as F + +import tensorrt_llm # noqa: F401 + +CONV_WIDTH = 4 +HEAD_DIM = 128 +NUM_CACHE_SLOTS = 514 +OUTPUT_NORM_EPS = 1e-5 + + +@dataclass(frozen=True) +class KdaInputs: + num_cache_slots: int + x_q: torch.Tensor + x_k: torch.Tensor + x_v: torch.Tensor + w_q_t: torch.Tensor + w_k_t: torch.Tensor + w_v_t: torch.Tensor + bias_q: torch.Tensor + bias_k: torch.Tensor + bias_v: torch.Tensor + conv_state_q: torch.Tensor + conv_state_k: torch.Tensor + conv_state_v: torch.Tensor + a_log: torch.Tensor + g: torch.Tensor + dt_bias: torch.Tensor + beta: torch.Tensor + output_norm_gate: torch.Tensor + output_norm_weight: torch.Tensor + state_storage: torch.Tensor + state_indices: torch.Tensor | None + cu_seqlens: torch.Tensor + + +def _state_view( + storage: torch.Tensor, + num_cache_slots: int, + num_heads: int, + head_dim: int, + slot_gap: int | None, +) -> torch.Tensor: + dense_slot_stride = num_heads * head_dim * head_dim + slot_gap = 0 if slot_gap is None else slot_gap + assert (dense_slot_stride + slot_gap) % 4 == 0, ( + "KDA state slots must remain aligned for float4 state accesses" + ) + return storage.as_strided( + (num_cache_slots, num_heads, head_dim, head_dim), + (dense_slot_stride + slot_gap, head_dim * head_dim, head_dim, 1), + ) + + +def _make_inputs( + *, + batch_size: int, + num_heads: int, + head_dim: int, + use_state_indices: bool, + update_conv_cache: bool, + state_slot_gap: int | None, + seed: int, +) -> KdaInputs: + torch.manual_seed(seed) + projection_size = num_heads * head_dim + num_cache_slots = NUM_CACHE_SLOTS if use_state_indices else batch_size + conv_slots = num_cache_slots if update_conv_cache else batch_size + + def make_conv_state() -> torch.Tensor: + if update_conv_cache: + return torch.empty_strided( + (conv_slots, projection_size, CONV_WIDTH - 1), + ( + projection_size * (CONV_WIDTH - 1), + 1, + projection_size, + ), + device="cuda", + dtype=torch.bfloat16, + ).normal_() + return torch.randn( + (conv_slots, projection_size, CONV_WIDTH - 1), + device="cuda", + dtype=torch.bfloat16, + ) + + dense_slot_stride = num_heads * head_dim * head_dim + slot_gap = 0 if state_slot_gap is None else state_slot_gap + state_storage = torch.randn( + (num_cache_slots * (dense_slot_stride + slot_gap),), + device="cuda", + dtype=torch.float32, + ) + state_indices = None + if use_state_indices: + state_indices = torch.randperm( + num_cache_slots, + device="cuda", + dtype=torch.int32, + )[:batch_size] + + return KdaInputs( + num_cache_slots=num_cache_slots, + x_q=torch.randn( + (1, batch_size, num_heads, head_dim), + device="cuda", + dtype=torch.bfloat16, + ), + x_k=torch.randn( + (1, batch_size, num_heads, head_dim), + device="cuda", + dtype=torch.bfloat16, + ), + x_v=torch.randn( + (1, batch_size, num_heads, head_dim), + device="cuda", + dtype=torch.bfloat16, + ), + w_q_t=torch.randn( + (CONV_WIDTH, projection_size), + device="cuda", + dtype=torch.bfloat16, + ), + w_k_t=torch.randn( + (CONV_WIDTH, projection_size), + device="cuda", + dtype=torch.bfloat16, + ), + w_v_t=torch.randn( + (CONV_WIDTH, projection_size), + device="cuda", + dtype=torch.bfloat16, + ), + bias_q=torch.randn( + (projection_size,), + device="cuda", + dtype=torch.bfloat16, + ), + bias_k=torch.randn( + (projection_size,), + device="cuda", + dtype=torch.bfloat16, + ), + bias_v=torch.randn( + (projection_size,), + device="cuda", + dtype=torch.bfloat16, + ), + conv_state_q=make_conv_state(), + conv_state_k=make_conv_state(), + conv_state_v=make_conv_state(), + a_log=torch.empty(num_heads, device="cuda", dtype=torch.float32).uniform_(1.0, 16.0).log_(), + g=torch.randn( + (1, batch_size, num_heads, head_dim), + device="cuda", + dtype=torch.bfloat16, + ), + dt_bias=torch.empty(projection_size, device="cuda", dtype=torch.float32).uniform_( + -4.0, -2.0 + ), + beta=torch.randn( + (1, batch_size, num_heads), + device="cuda", + dtype=torch.bfloat16, + ), + output_norm_gate=torch.randn( + (1, batch_size, num_heads, head_dim), + device="cuda", + dtype=torch.bfloat16, + ), + output_norm_weight=torch.empty(head_dim, device="cuda", dtype=torch.float32).uniform_( + 0.5, 1.5 + ), + state_storage=state_storage, + state_indices=state_indices, + cu_seqlens=torch.arange(batch_size + 1, device="cuda", dtype=torch.int32), + ) + + +def _conv_reference( + x: torch.Tensor, + conv_state: torch.Tensor, + weight_t: torch.Tensor, + bias: torch.Tensor, + num_heads: int, + head_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + batch_size = x.shape[1] + x_flat = x.transpose(0, 1).reshape(batch_size, num_heads * head_dim) + window = torch.cat((conv_state.float(), x_flat.float().unsqueeze(-1)), dim=-1) + output = bias.float() + (window * weight_t.transpose(0, 1).float()).sum(dim=-1) + output = F.silu(output).reshape(batch_size, 1, num_heads, head_dim).to(torch.bfloat16) + updated_state = torch.cat((conv_state[:, :, 1:], x_flat.unsqueeze(-1)), dim=-1) + return output, updated_state + + +def _fla_reference( + inputs: KdaInputs, + initial_state: torch.Tensor, + *, + num_heads: int, + head_dim: int, + apply_output_norm: bool, + update_conv_cache: bool, + apply_beta_sigmoid: bool, + gate_lower_bound: float | None, +) -> tuple[torch.Tensor, torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + from fla.ops.kda import fused_recurrent_kda + + if update_conv_cache: + assert inputs.state_indices is not None + state_indices = inputs.state_indices.long() + conv_q = inputs.conv_state_q.index_select(0, state_indices) + conv_k = inputs.conv_state_k.index_select(0, state_indices) + conv_v = inputs.conv_state_v.index_select(0, state_indices) + else: + conv_q = inputs.conv_state_q + conv_k = inputs.conv_state_k + conv_v = inputs.conv_state_v + + q, updated_conv_q = _conv_reference( + inputs.x_q, + conv_q, + inputs.w_q_t, + inputs.bias_q, + num_heads, + head_dim, + ) + k, updated_conv_k = _conv_reference( + inputs.x_k, + conv_k, + inputs.w_k_t, + inputs.bias_k, + num_heads, + head_dim, + ) + v, updated_conv_v = _conv_reference( + inputs.x_v, + conv_v, + inputs.w_v_t, + inputs.bias_v, + num_heads, + head_dim, + ) + + output, final_state = fused_recurrent_kda( + q=q, + k=k, + v=v, + g=inputs.g.transpose(0, 1), + beta=inputs.beta.transpose(0, 1).float(), + A_log=inputs.a_log, + dt_bias=inputs.dt_bias, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=apply_beta_sigmoid, + lower_bound=gate_lower_bound, + state_v_first=True, + ) + + if apply_output_norm: + output_float = output.float() + rstd = torch.rsqrt(output_float.square().mean(dim=-1, keepdim=True) + OUTPUT_NORM_EPS) + output = ( + output_float + * rstd + * inputs.output_norm_weight.view(1, 1, 1, head_dim) + * torch.sigmoid(inputs.output_norm_gate.transpose(0, 1).float()) + ) + return ( + output.to(torch.bfloat16), + final_state.float(), + (updated_conv_q, updated_conv_k, updated_conv_v), + ) + + +def _assert_parity( + name: str, + actual: torch.Tensor, + expected: torch.Tensor, + *, + min_cosine: float = 0.9999, + max_relative_l2: float = 1e-2, +) -> None: + actual_float = actual.float().flatten() + expected_float = expected.float().flatten() + cosine = F.cosine_similarity(actual_float, expected_float, dim=0).item() + relative_l2 = ( + (actual_float - expected_float).norm() / expected_float.norm().clamp_min(1e-12) + ).item() + max_abs = (actual_float - expected_float).abs().max().item() + assert cosine > min_cosine, f"{name}: cosine={cosine:.8f}, max_abs={max_abs:.6g}" + assert relative_l2 < max_relative_l2, ( + f"{name}: relative_l2={relative_l2:.8f}, max_abs={max_abs:.6g}" + ) + + +@torch.no_grad() +@pytest.mark.parametrize("batch_size", [1, 17, 32]) +@pytest.mark.parametrize("num_heads", [2, 3, 4, 6, 12, 96]) +@pytest.mark.parametrize( + ( + "use_state_indices," + "update_conv_cache,state_slot_gap,apply_output_norm,apply_beta_sigmoid,gate_lower_bound" + ), + [ + pytest.param(True, False, None, True, True, -5.0, id="indexed"), + pytest.param(False, False, None, False, False, None, id="batch-local-softplus-decay"), + pytest.param(True, True, 73728, True, True, -5.0, id="indexed-conv-strided"), + ], +) +def test_kda_decode_matches_fla( + batch_size: int, + num_heads: int, + use_state_indices: bool, + update_conv_cache: bool, + state_slot_gap: int | None, + apply_output_norm: bool, + apply_beta_sigmoid: bool, + gate_lower_bound: float | None, +) -> None: + head_dim = HEAD_DIM + inputs = _make_inputs( + batch_size=batch_size, + num_heads=num_heads, + head_dim=head_dim, + use_state_indices=use_state_indices, + update_conv_cache=update_conv_cache, + state_slot_gap=state_slot_gap, + seed=2026 + batch_size + num_heads, + ) + actual_state = _state_view( + inputs.state_storage, + inputs.num_cache_slots, + num_heads, + head_dim, + state_slot_gap, + ) + state_gap = None + state_gap_before = None + if state_slot_gap is not None: + dense_slot_stride = num_heads * head_dim * head_dim + state_gap = inputs.state_storage.view( + inputs.num_cache_slots, dense_slot_stride + state_slot_gap + )[:, dense_slot_stride:] + state_gap_before = state_gap.clone() + + if inputs.state_indices is None: + initial_selected_state = actual_state.clone() + state_before = None + else: + initial_selected_state = actual_state.index_select(0, inputs.state_indices.long()).clone() + state_before = actual_state.clone() + + actual_conv_q = inputs.conv_state_q.clone(memory_format=torch.preserve_format) + actual_conv_k = inputs.conv_state_k.clone(memory_format=torch.preserve_format) + actual_conv_v = inputs.conv_state_v.clone(memory_format=torch.preserve_format) + conv_before = (actual_conv_q.clone(), actual_conv_k.clone(), actual_conv_v.clone()) + if update_conv_cache: + projection_size = num_heads * head_dim + expected_conv_stride = ( + projection_size * (CONV_WIDTH - 1), + 1, + projection_size, + ) + for conv_state in (actual_conv_q, actual_conv_k, actual_conv_v): + assert conv_state.stride() == expected_conv_stride + + expected_output, expected_state, expected_conv = _fla_reference( + inputs, + initial_selected_state, + num_heads=num_heads, + head_dim=head_dim, + apply_output_norm=apply_output_norm, + update_conv_cache=update_conv_cache, + apply_beta_sigmoid=apply_beta_sigmoid, + gate_lower_bound=gate_lower_bound, + ) + + actual_output = torch.ops.trtllm.kda_decode( + inputs.x_q, + inputs.x_k, + inputs.x_v, + inputs.w_q_t, + inputs.w_k_t, + inputs.w_v_t, + inputs.bias_q, + inputs.bias_k, + inputs.bias_v, + actual_conv_q, + actual_conv_k, + actual_conv_v, + inputs.a_log, + inputs.g, + inputs.dt_bias, + inputs.beta, + inputs.output_norm_gate, + inputs.output_norm_weight, + inputs.state_indices, + inputs.cu_seqlens, + actual_state, + apply_output_norm, + update_conv_cache, + gate_lower_bound is not None, + apply_beta_sigmoid, + 0.0 if gate_lower_bound is None else gate_lower_bound, + head_dim**-0.5, + OUTPUT_NORM_EPS, + ) + + _assert_parity("output", actual_output, expected_output) + actual_selected_state = ( + actual_state + if inputs.state_indices is None + else actual_state.index_select(0, inputs.state_indices.long()) + ) + _assert_parity( + "recurrent state", + actual_selected_state, + expected_state, + ) + if state_before is not None: + state_before.index_copy_(0, inputs.state_indices.long(), actual_selected_state) + torch.testing.assert_close(actual_state, state_before, rtol=0, atol=0) + if state_gap is not None: + torch.testing.assert_close(state_gap, state_gap_before, rtol=0, atol=0) + + if update_conv_cache: + assert inputs.state_indices is not None + state_indices = inputs.state_indices.long() + for actual, expected, before in zip( + (actual_conv_q, actual_conv_k, actual_conv_v), + expected_conv, + conv_before, + strict=True, + ): + expected_pool = before.clone(memory_format=torch.preserve_format) + expected_pool.index_copy_(0, state_indices, expected) + torch.testing.assert_close(actual, expected_pool, rtol=0, atol=0) + else: + for actual, before in zip( + (actual_conv_q, actual_conv_k, actual_conv_v), + conv_before, + strict=True, + ): + torch.testing.assert_close(actual, before, rtol=0, atol=0)