diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu index f1e284118d6f..b6eee63b6100 100644 --- a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu @@ -30,6 +30,7 @@ // Attention_residual kernel at e7f934124acc915575f9f7561f9d1e373ab43089. #include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h" #include @@ -72,6 +73,14 @@ __inline__ __device__ float block_reduce_sum(float val, float* ws) return val; } +__device__ __forceinline__ bf16_t apply_output_rms_norm(bf16_t value, float rsigma, bf16_t weight) +{ + // Preserve KimiK3RMSNorm semantics: normalize in FP32, round to the + // activation dtype, then apply the BF16 weight. + bf16_t const normalized = __float2bfloat16_rn(__bfloat162float(value) * rsigma); + return __float2bfloat16_rn(__bfloat162float(normalized) * __bfloat162float(weight)); +} + __device__ __forceinline__ bf16_t const* v_addr( bf16_t const* block_res, bf16_t const* layer_res, int n, int N, int t, int b, int T, int B, int H) { @@ -719,10 +728,6 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t c plan.logits_all[ng] = local_logit; } } - // Publish the final chunk's plan.logits_all stores before the - // cross-lane reads in consumer warp 0 below (earlier chunks are - // covered by the NamedBarrier inside the loop). - __syncwarp(); float inv_s = 1.f / s_running; bf16_t* out_ptr = output + (long long) tb * H; @@ -1012,22 +1017,29 @@ static void launch_fwd(bf16_t const* block_residual, bf16_t const* layer_residua // Small-N counterpart to the Triton one-program topology. One CTA owns the // complete token, with exactly 28 hidden elements per thread at H=7168. For -// N=2/4, packed BF16 V remains in registers across the statistics/softmax -// boundary; N=1 can write V directly because its softmax is identically one. -template -__global__ void __launch_bounds__(256, 1) - attn_res_fwd_s1_single_cta_kernel(bf16_t const* __restrict__ block_res, bf16_t const* __restrict__ layer_res, - bf16_t const* __restrict__ res_w, bf16_t const* __restrict__ rms_w, bf16_t* __restrict__ output, - float* __restrict__ rsigma_out, float* __restrict__ probs_out, float* __restrict__ logits_out, float rms_eps) +// N=2/3/4, packed BF16 V remains in registers across the statistics/softmax +// boundary. The fused output is retained in registers for the trailing +// RMSNorm, preserving the BF16 boundary without a shared-memory round trip. +template +__global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_single_cta_kernel(bf16_t const* __restrict__ block_res, + bf16_t const* __restrict__ layer_res, bf16_t const* __restrict__ layer_res_add, bf16_t const* __restrict__ res_w, + bf16_t const* __restrict__ rms_w, bf16_t const* __restrict__ output_rms_w, bf16_t* __restrict__ updated_layer_res, + bf16_t* __restrict__ output, float* __restrict__ rsigma_out, float* __restrict__ probs_out, + float* __restrict__ logits_out, float rms_eps, float output_rms_eps) { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + if constexpr (ENABLE_PDL) + { + cudaGridDependencySynchronize(); + } + constexpr int H = 7168; constexpr int THREADS = 256; constexpr int WARPS = THREADS / 32; constexpr int ITEMS = H / THREADS; constexpr float LOG2_E = 1.4426950408889634f; static_assert(H % THREADS == 0); - static_assert(N == 1 || N == 2 || N == 4); + static_assert(N >= 1 && N <= 4); __shared__ float2 warp_stats[WARPS * N]; __shared__ float weights[N]; @@ -1036,7 +1048,9 @@ __global__ void __launch_bounds__(256, 1) int const warp = tid >> 5; float2 stats[N] = {}; + float output_sq_local = 0.f; uint32_t v_cache_bf16[ITEMS][(N + 1) / 2]; + bf16_t mixed_cache[ITEMS]; #pragma unroll for (int item = 0; item < ITEMS; item++) { @@ -1048,10 +1062,26 @@ __global__ void __launch_bounds__(256, 1) { bf16_t const* row = n < N - 1 ? block_res + (size_t) n * H : layer_res; bf16_t packed_v = row[h]; + if constexpr (FUSE_LAYER_ADD) + { + if (n == N - 1) + { + packed_v = __float2bfloat16_rn(__bfloat162float(packed_v) + __bfloat162float(layer_res_add[h])); + updated_layer_res[h] = packed_v; + } + } float v = __bfloat162float(packed_v); if constexpr (N == 1) { - output[h] = packed_v; + if constexpr (FUSE_OUTPUT_RMS_NORM) + { + mixed_cache[item] = packed_v; + output_sq_local = fmaf(v, v, output_sq_local); + } + else + { + output[h] = packed_v; + } } else { @@ -1073,6 +1103,17 @@ __global__ void __launch_bounds__(256, 1) packed.bf16x2 = __halves2bfloat162(item_v[2 * pair], item_v[2 * pair + 1]); v_cache_bf16[item][pair] = packed.bits; } + if constexpr (N % 2 == 1) + { + union + { + __nv_bfloat162 bf16x2; + uint32_t bits; + } packed; + + packed.bf16x2 = __halves2bfloat162(item_v[N - 1], __float2bfloat16_rn(0.f)); + v_cache_bf16[item][N / 2] = packed.bits; + } } } @@ -1135,9 +1176,12 @@ __global__ void __launch_bounds__(256, 1) for (int n = 0; n < N; n++) { weights[n] *= inv_denominator; - rsigma_out[n] = local_rsigma[n]; - logits_out[n] = local_logits[n]; - probs_out[n] = weights[n]; + if (rsigma_out) + rsigma_out[n] = local_rsigma[n]; + if (logits_out) + logits_out[n] = local_logits[n]; + if (probs_out) + probs_out[n] = weights[n]; } } __syncthreads(); @@ -1162,10 +1206,66 @@ __global__ void __launch_bounds__(256, 1) value = fmaf(weights[2 * pair], v.x, value); value = fmaf(weights[2 * pair + 1], v.y, value); } + if constexpr (N % 2 == 1) + { + union + { + __nv_bfloat162 bf16x2; + uint32_t bits; + } packed; + + packed.bits = v_cache_bf16[item][N / 2]; + float2 v = __bfloat1622float2(packed.bf16x2); + value = fmaf(weights[N - 1], v.x, value); + } int h = tid + item * THREADS; - output[h] = __float2bfloat16_rn(value); + bf16_t const mixed = __float2bfloat16_rn(value); + if constexpr (FUSE_OUTPUT_RMS_NORM) + { + mixed_cache[item] = mixed; + float const mixed_float = __bfloat162float(mixed); + output_sq_local = fmaf(mixed_float, mixed_float, output_sq_local); + } + else + { + output[h] = mixed; + } + } + } + + if constexpr (FUSE_OUTPUT_RMS_NORM) + { + output_sq_local = warp_reduce_sum(output_sq_local); + if (lane == 0) + { + warp_stats[warp].x = output_sq_local; + } + __syncthreads(); + if (tid == 0) + { + float output_sq = 0.f; +#pragma unroll + for (int w = 0; w < WARPS; w++) + { + output_sq += warp_stats[w].x; + } + weights[0] = rsqrtf(output_sq / H + output_rms_eps); + } + __syncthreads(); + + float const output_rsigma = weights[0]; +#pragma unroll + for (int item = 0; item < ITEMS; item++) + { + int h = tid + item * THREADS; + output[h] = apply_output_rms_norm(mixed_cache[item], output_rsigma, output_rms_w[h]); } } + + if constexpr (ENABLE_PDL) + { + cudaTriggerProgrammaticLaunchCompletion(); + } #else if (cute::thread0()) { @@ -1174,30 +1274,67 @@ __global__ void __launch_bounds__(256, 1) #endif } +template +static void launch_s1_single_cta(bf16_t const* block_residual, bf16_t const* layer_residual, + bf16_t const* layer_residual_add, bf16_t const* res_weight, bf16_t const* rms_weight, + bf16_t const* output_rms_weight, bf16_t* updated_layer_residual, bf16_t* output, float* rsigma, float* probs, + float* logits, float rms_eps, float output_rms_eps, cudaStream_t stream) +{ + if (tensorrt_llm::common::getEnvEnablePDL()) + { + auto kernel = &attn_res_fwd_s1_single_cta_kernel; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(1); + config.blockDim = dim3(256); + config.stream = stream; + cudaLaunchAttribute attribute{}; + attribute.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute.val.programmaticStreamSerializationAllowed = 1; + config.attrs = &attribute; + config.numAttrs = 1; + cudaLaunchKernelEx(&config, kernel, block_residual, layer_residual, layer_residual_add, res_weight, rms_weight, + output_rms_weight, updated_layer_residual, output, rsigma, probs, logits, rms_eps, output_rms_eps); + } + else + { + attn_res_fwd_s1_single_cta_kernel + <<<1, 256, 0, stream>>>(block_residual, layer_residual, layer_residual_add, res_weight, rms_weight, + output_rms_weight, updated_layer_residual, output, rsigma, probs, logits, rms_eps, output_rms_eps); + } +} + template static void launch_s1_single_cta(bf16_t const* block_residual, bf16_t const* layer_residual, bf16_t const* res_weight, bf16_t const* rms_weight, bf16_t* output, float* rsigma, float* probs, float* logits, float rms_eps, cudaStream_t stream) { - attn_res_fwd_s1_single_cta_kernel<<<1, 256, 0, stream>>>( - block_residual, layer_residual, res_weight, rms_weight, output, rsigma, probs, logits, rms_eps); + launch_s1_single_cta(block_residual, layer_residual, nullptr, res_weight, rms_weight, nullptr, + nullptr, output, rsigma, probs, logits, rms_eps, 0.f, stream); } // Single-token split-K specialization. The complete grid is one CTA cluster: // rank g owns a disjoint H/GROUPS slice, keeps that slice of FP32 V in its // rank-local shared memory, and exchanges only (square, dot) partials via DSM. -template -__global__ void __launch_bounds__(256, 1) - attn_res_fwd_s1_splitk_kernel(bf16_t const* __restrict__ block_res, bf16_t const* __restrict__ layer_res, - bf16_t const* __restrict__ res_w, bf16_t const* __restrict__ rms_w, bf16_t* __restrict__ output, - float* __restrict__ rsigma_out, float* __restrict__ probs_out, float* __restrict__ logits_out, float rms_eps) +template +__global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_splitk_kernel(bf16_t const* __restrict__ block_res, + bf16_t const* __restrict__ layer_res, bf16_t const* __restrict__ layer_res_add, bf16_t const* __restrict__ res_w, + bf16_t const* __restrict__ rms_w, bf16_t const* __restrict__ output_rms_w, bf16_t* __restrict__ updated_layer_res, + bf16_t* __restrict__ output, float* __restrict__ rsigma_out, float* __restrict__ probs_out, + float* __restrict__ logits_out, float rms_eps, float output_rms_eps) { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + if constexpr (ENABLE_PDL) + { + cudaGridDependencySynchronize(); + } + namespace cg = cooperative_groups; constexpr int H = 7168; constexpr int K_PER_CTA = H / GROUPS; constexpr int THREADS = 256; constexpr int WARPS = THREADS / 32; + constexpr int ITEMS = (K_PER_CTA + THREADS - 1) / THREADS; constexpr float LOG2_E = 1.4426950408889634f; static_assert(H % GROUPS == 0); @@ -1224,7 +1361,16 @@ __global__ void __launch_bounds__(256, 1) for (int n = 0; n < N; n++) { bf16_t const* row = n < N - 1 ? block_res + (size_t) n * H : layer_res; - float v = __bfloat162float(row[h]); + bf16_t packed_v = row[h]; + if constexpr (FUSE_LAYER_ADD) + { + if (n == N - 1) + { + packed_v = __float2bfloat16_rn(__bfloat162float(packed_v) + __bfloat162float(layer_res_add[h])); + updated_layer_res[h] = packed_v; + } + } + float v = __bfloat162float(packed_v); v_cache[(size_t) n * K_PER_CTA + ki] = v; sq[n] = fmaf(v, v, sq[n]); dot[n] = fmaf(v, q, dot[n]); @@ -1267,16 +1413,26 @@ __global__ void __launch_bounds__(256, 1) // One thread per candidate reduces across CTA ranks. Parallelizing this // avoids making a single leader issue all GROUPS*N remote DSM reads. + float2 cluster_total = {}; if (tid < N) { - float2 total = {}; #pragma unroll for (int g = 0; g < GROUPS; g++) { float2 const* remote_stats = cluster.map_shared_rank(warp_stats, g); - total = float2_add(total, remote_stats[tid]); + cluster_total = float2_add(cluster_total, remote_stats[tid]); } - warp_stats[tid] = total; + } + + // No rank may overwrite its published DSM partial until every other rank + // has finished reading it. This second cluster barrier is required even + // though every rank traverses the same remote-rank loop: CTAs can make + // progress independently, especially at the higher-register N=9/12 + // specializations. + cluster.sync(); + if (tid < N) + { + warp_stats[tid] = cluster_total; } __syncthreads(); @@ -1307,17 +1463,22 @@ __global__ void __launch_bounds__(256, 1) weights[n] *= inv_sum; if (group == 0) { - rsigma_out[n] = local_rsigma[n]; - logits_out[n] = local_logits[n]; - probs_out[n] = weights[n]; + if (rsigma_out) + rsigma_out[n] = local_rsigma[n]; + if (logits_out) + logits_out[n] = local_logits[n]; + if (probs_out) + probs_out[n] = weights[n]; } } } cluster.sync(); + float output_sq_local = 0.f; + bf16_t mixed_cache[ITEMS]; #pragma unroll - for (int ki = tid; ki < K_PER_CTA; ki += THREADS) + for (int ki = tid, item = 0; ki < K_PER_CTA; ki += THREADS, item++) { float value = 0.0f; #pragma unroll @@ -1325,7 +1486,70 @@ __global__ void __launch_bounds__(256, 1) { value = fmaf(weights[n], v_cache[(size_t) n * K_PER_CTA + ki], value); } - output[h_begin + ki] = __float2bfloat16_rn(value); + bf16_t const mixed = __float2bfloat16_rn(value); + if constexpr (FUSE_OUTPUT_RMS_NORM) + { + float const mixed_float = __bfloat162float(mixed); + mixed_cache[item] = mixed; + output_sq_local = fmaf(mixed_float, mixed_float, output_sq_local); + } + else + { + output[h_begin + ki] = mixed; + } + } + + if constexpr (FUSE_OUTPUT_RMS_NORM) + { + output_sq_local = warp_reduce_sum(output_sq_local); + if (lane == 0) + { + warp_stats[warp].x = output_sq_local; + } + __syncthreads(); + if (tid == 0) + { + float output_sq = 0.f; +#pragma unroll + for (int w = 0; w < WARPS; w++) + { + output_sq += warp_stats[w].x; + } + warp_stats[0].x = output_sq; + } + + cluster.sync(); + + if (tid == 0) + { + float output_sq = 0.f; +#pragma unroll + for (int g = 0; g < GROUPS; g++) + { + float2 const* remote_stats = cluster.map_shared_rank(warp_stats, g); + output_sq += remote_stats[0].x; + } + weights[0] = rsqrtf(output_sq / H + output_rms_eps); + } + + // Keep every source CTA alive until all remote DSM reads above have + // completed. A block-local barrier is insufficient: a faster CTA + // could otherwise leave the cluster while a peer still reads its + // rank-local output-square partial. + cluster.sync(); + + float const output_rsigma = weights[0]; +#pragma unroll + for (int ki = tid, item = 0; ki < K_PER_CTA; ki += THREADS, item++) + { + int const h = h_begin + ki; + output[h] = apply_output_rms_norm(mixed_cache[item], output_rsigma, output_rms_w[h]); + } + } + + if constexpr (ENABLE_PDL) + { + cudaTriggerProgrammaticLaunchCompletion(); } #else if (cute::thread0()) @@ -1335,24 +1559,32 @@ __global__ void __launch_bounds__(256, 1) #endif } -template -static void launch_s1_splitk(bf16_t const* block_residual, bf16_t const* layer_residual, bf16_t const* res_weight, - bf16_t const* rms_weight, bf16_t* output, float* rsigma, float* probs, float* logits, float rms_eps, - cudaStream_t stream) +template +static void launch_s1_splitk(bf16_t const* block_residual, bf16_t const* layer_residual, + bf16_t const* layer_residual_add, bf16_t const* res_weight, bf16_t const* rms_weight, + bf16_t const* output_rms_weight, bf16_t* updated_layer_residual, bf16_t* output, float* rsigma, float* probs, + float* logits, float rms_eps, float output_rms_eps, cudaStream_t stream) { constexpr int K_PER_CTA = 7168 / GROUPS; constexpr int WARPS = 8; constexpr size_t smem_size = (size_t) N * K_PER_CTA * sizeof(float) + (size_t) WARPS * N * sizeof(float2) + (size_t) N * sizeof(float); - auto kernel = &attn_res_fwd_s1_splitk_kernel; + bool const enable_pdl = tensorrt_llm::common::getEnvEnablePDL(); + auto kernel_pdl = &attn_res_fwd_s1_splitk_kernel; + auto kernel_nopdl = &attn_res_fwd_s1_splitk_kernel; + auto kernel = enable_pdl ? kernel_pdl : kernel_nopdl; { // cudaFuncSetAttribute applies to the current device only; set it - // once per device (per kernel instantiation). + // once per device (per kernel instantiation). Both PDL variants are + // registered so a later flip of getEnvEnablePDL() is still valid. static std::once_flag attrs_set[64]; int dev = 0; TLLM_CUDA_CHECK(cudaGetDevice(&dev)); auto const set_attr = [&] - { TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); }; + { + TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel_pdl, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + TLLM_CUDA_CHECK(cudaFuncSetAttribute(kernel_nopdl, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + }; if (dev >= 0 && dev < 64) { std::call_once(attrs_set[dev], set_attr); @@ -1363,23 +1595,38 @@ static void launch_s1_splitk(bf16_t const* block_residual, bf16_t const* layer_r } } void* args[] = {const_cast(&block_residual), const_cast(&layer_residual), - const_cast(&res_weight), const_cast(&rms_weight), &output, &rsigma, &probs, &logits, - &rms_eps}; + const_cast(&layer_residual_add), const_cast(&res_weight), const_cast(&rms_weight), + const_cast(&output_rms_weight), &updated_layer_residual, &output, &rsigma, &probs, &logits, &rms_eps, + &output_rms_eps}; cudaLaunchConfig_t config{}; config.gridDim = dim3(GROUPS); config.blockDim = dim3(256); config.dynamicSmemBytes = smem_size; config.stream = stream; - cudaLaunchAttribute attribute{}; - attribute.id = cudaLaunchAttributeClusterDimension; - attribute.val.clusterDim.x = GROUPS; - attribute.val.clusterDim.y = 1; - attribute.val.clusterDim.z = 1; - config.attrs = &attribute; - config.numAttrs = 1; + cudaLaunchAttribute attributes[2]{}; + attributes[0].id = cudaLaunchAttributeClusterDimension; + attributes[0].val.clusterDim.x = GROUPS; + attributes[0].val.clusterDim.y = 1; + attributes[0].val.clusterDim.z = 1; + if (enable_pdl) + { + attributes[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[1].val.programmaticStreamSerializationAllowed = 1; + } + config.attrs = attributes; + config.numAttrs = enable_pdl ? 2 : 1; cudaLaunchKernelExC(&config, reinterpret_cast(kernel), args); } +template +static void launch_s1_splitk(bf16_t const* block_residual, bf16_t const* layer_residual, bf16_t const* res_weight, + bf16_t const* rms_weight, bf16_t* output, float* rsigma, float* probs, float* logits, float rms_eps, + cudaStream_t stream) +{ + launch_s1_splitk(block_residual, layer_residual, nullptr, res_weight, rms_weight, nullptr, + nullptr, output, rsigma, probs, logits, rms_eps, 0.f, stream); +} + template static void launch_n1_ttile(bf16_t const* layer_residual, bf16_t const* res_weight, bf16_t const* rms_weight, bf16_t* output, float* rsigma, float* probs, float* logits, int T, int B, float rms_eps, int num_sm, @@ -1552,6 +1799,67 @@ void invokeAttnResFwd(AttnResFwdParams const& params, cudaStream_t stream) } } +template +static void launchAttnResDecodeRmsNorm(AttnResFwdParams const& params, cudaStream_t stream) +{ + using namespace sm100::fwd_prod_v2; + + auto const* layer_residual_add = FUSE_LAYER_ADD ? params.layerResidualAdd : nullptr; + auto* updated_layer_residual = FUSE_LAYER_ADD ? params.updatedLayerResidual : nullptr; + + if constexpr (N <= 4) + { + launch_s1_single_cta(params.blockResidual, params.layerResidual, layer_residual_add, + params.resWeight, params.rmsWeight, params.outputRmsWeight, updated_layer_residual, params.output, nullptr, + nullptr, nullptr, params.rmsEps, params.outputRmsEps, stream); + } + else + { + launch_s1_splitk(params.blockResidual, params.layerResidual, layer_residual_add, + params.resWeight, params.rmsWeight, params.outputRmsWeight, updated_layer_residual, params.output, nullptr, + nullptr, nullptr, params.rmsEps, params.outputRmsEps, stream); + } +} + +template +static void invokeAttnResDecodeRmsNorm(AttnResFwdParams const& params, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(params.seqLen == 1 && params.batchSize == 1 && params.hiddenSize == 7168, + "attn_res decode RMSNorm supports T=B=1, H=7168 only"); + TLLM_CHECK_WITH_INFO(params.outputRmsWeight != nullptr, "attn_res decode RMSNorm requires outputRmsWeight"); + if constexpr (FUSE_LAYER_ADD) + { + TLLM_CHECK_WITH_INFO(params.layerResidualAdd != nullptr && params.updatedLayerResidual != nullptr, + "attn_res decode add+RMSNorm requires layerResidualAdd and updatedLayerResidual"); + } + switch (params.numCandidates) + { + case 1: launchAttnResDecodeRmsNorm<1, FUSE_LAYER_ADD>(params, stream); break; + case 2: launchAttnResDecodeRmsNorm<2, FUSE_LAYER_ADD>(params, stream); break; + case 3: launchAttnResDecodeRmsNorm<3, FUSE_LAYER_ADD>(params, stream); break; + case 4: launchAttnResDecodeRmsNorm<4, FUSE_LAYER_ADD>(params, stream); break; + case 5: launchAttnResDecodeRmsNorm<5, FUSE_LAYER_ADD>(params, stream); break; + case 6: launchAttnResDecodeRmsNorm<6, FUSE_LAYER_ADD>(params, stream); break; + case 7: launchAttnResDecodeRmsNorm<7, FUSE_LAYER_ADD>(params, stream); break; + case 8: launchAttnResDecodeRmsNorm<8, FUSE_LAYER_ADD>(params, stream); break; + case 9: launchAttnResDecodeRmsNorm<9, FUSE_LAYER_ADD>(params, stream); break; + case 12: launchAttnResDecodeRmsNorm<12, FUSE_LAYER_ADD>(params, stream); break; + default: + TLLM_CHECK_WITH_INFO(false, "attn_res decode RMSNorm: unsupported numCandidates=%d (expected [1, 9] or 12)", + params.numCandidates); + } +} + +void invokeAttnResRmsNormFwd(AttnResFwdParams const& params, cudaStream_t stream) +{ + invokeAttnResDecodeRmsNorm(params, stream); +} + +void invokeAttnResAddRmsNormFwd(AttnResFwdParams const& params, cudaStream_t stream) +{ + invokeAttnResDecodeRmsNorm(params, stream); +} + } // namespace kernels::kimiK3AttnRes TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h index 65a9913f5afb..4eba27e585bd 100644 --- a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h @@ -32,28 +32,44 @@ namespace kernels::kimiK3AttnRes //! //! Contract (checked at the Torch-op bridge): B == 1, N in [1, 12], //! T in [1, 16384], H a multiple of 1024 in [4096, 8192]; all residual -//! tensors bf16 contiguous, rsigma/probs/logits fp32 [N, T, B]. +//! tensors bf16 contiguous, rsigma/probs/logits fp32 [N, T, B] when requested. //! blockResidual may be nullptr when N == 1. struct AttnResFwdParams { - __nv_bfloat16 const* blockResidual; // [N-1, T, B, H], nullptr iff N == 1 - __nv_bfloat16 const* layerResidual; // [T, B, H] - __nv_bfloat16 const* resWeight; // [H] - __nv_bfloat16 const* rmsWeight; // [H] - __nv_bfloat16* output; // [T, B, H] - float* rsigma; // [N, T, B] - float* probs; // [N, T, B] - float* logits; // [N, T, B] - int numCandidates; // N = K + 1 - int seqLen; // T - int batchSize; // B - int hiddenSize; // H + __nv_bfloat16 const* blockResidual; // [N-1, T, B, H], nullptr iff N == 1 + __nv_bfloat16 const* layerResidual; // [T, B, H] + __nv_bfloat16 const* layerResidualAdd; // [T, B, H], optional fused addend + __nv_bfloat16 const* resWeight; // [H] + __nv_bfloat16 const* rmsWeight; // [H] + __nv_bfloat16 const* outputRmsWeight; // [H], nullptr unless trailing RMSNorm is fused + __nv_bfloat16* updatedLayerResidual; // [T, B, H], optional fused-add output + __nv_bfloat16* output; // [T, B, H] + float* rsigma; // [N, T, B], optional + float* probs; // [N, T, B], optional + float* logits; // [N, T, B], optional + int numCandidates; // N = K + 1 + int seqLen; // T + int batchSize; // B + int hiddenSize; // H float rmsEps; + float outputRmsEps; }; //! Launches the fused attention-residual forward on the supplied stream. void invokeAttnResFwd(AttnResFwdParams const& params, cudaStream_t stream); +//! Launches attention-residual selection followed by the next RMSNorm in the +//! same kernel. The BF16 attention-residual output rounding boundary is +//! preserved before applying outputRmsWeight. Supported only for T=B=1, +//! H=7168. +void invokeAttnResRmsNormFwd(AttnResFwdParams const& params, cudaStream_t stream); + +//! Launches the production decode specialization with +//! updatedLayerResidual = bf16(layerResidual + layerResidualAdd), then uses +//! that rounded value as the final attention-residual candidate and fuses the +//! immediately following RMSNorm. Supported only for T=B=1, H=7168. +void invokeAttnResAddRmsNormFwd(AttnResFwdParams const& params, cudaStream_t stream); + } // namespace kernels::kimiK3AttnRes TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/attnResOp.cpp b/cpp/tensorrt_llm/thop/attnResOp.cpp index 6a7d16cadaf2..2c540afb67c3 100644 --- a/cpp/tensorrt_llm/thop/attnResOp.cpp +++ b/cpp/tensorrt_llm/thop/attnResOp.cpp @@ -122,6 +122,133 @@ std::tuple attn_res_fwd( return {output, rsigma, probs, logits}; } +at::Tensor attn_res_rmsnorm_fwd(at::Tensor layer_residual, at::Tensor block_residual, at::Tensor res_weight, + at::Tensor rms_weight, at::Tensor output_rms_weight, double rms_eps, double output_rms_eps) +{ + TORCH_CHECK(layer_residual.dim() == 3, "attn_res_rmsnorm_fwd: layer_residual must be [T, B, H]"); + TORCH_CHECK(block_residual.dim() == 4, "attn_res_rmsnorm_fwd: block_residual must be [K, T, B, H]"); + + int const T = static_cast(layer_residual.size(0)); + int const B = static_cast(layer_residual.size(1)); + int const H = static_cast(layer_residual.size(2)); + int const N = static_cast(block_residual.size(0)) + 1; + + TORCH_CHECK(layer_residual.is_cuda() && block_residual.is_cuda() && res_weight.is_cuda() && rms_weight.is_cuda() + && output_rms_weight.is_cuda(), + "attn_res_rmsnorm_fwd: all input tensors must be CUDA tensors"); + TORCH_CHECK(block_residual.device() == layer_residual.device() && res_weight.device() == layer_residual.device() + && rms_weight.device() == layer_residual.device() && output_rms_weight.device() == layer_residual.device(), + "attn_res_rmsnorm_fwd: all input tensors must be on the same CUDA device"); + c10::cuda::CUDAGuard device_guard(layer_residual.device()); + check_attn_res_contract(N, T, B, H); + TORCH_CHECK( + T == 1 && B == 1 && H == 7168, "attn_res_rmsnorm_fwd: only production decode shape T=B=1, H=7168 is supported"); + TORCH_CHECK((N >= 1 && N <= 9) || N == 12, "attn_res_rmsnorm_fwd: supported N values are [1, 9] and 12"); + + TORCH_CHECK(layer_residual.scalar_type() == at::kBFloat16, "attn_res_rmsnorm_fwd: layer_residual must be bf16"); + TORCH_CHECK(block_residual.scalar_type() == at::kBFloat16, "attn_res_rmsnorm_fwd: block_residual must be bf16"); + TORCH_CHECK(res_weight.scalar_type() == at::kBFloat16, "attn_res_rmsnorm_fwd: res_weight must be bf16"); + TORCH_CHECK(rms_weight.scalar_type() == at::kBFloat16, "attn_res_rmsnorm_fwd: rms_weight must be bf16"); + TORCH_CHECK( + output_rms_weight.scalar_type() == at::kBFloat16, "attn_res_rmsnorm_fwd: output_rms_weight must be bf16"); + TORCH_CHECK(layer_residual.is_contiguous() && block_residual.is_contiguous() && res_weight.is_contiguous() + && rms_weight.is_contiguous() && output_rms_weight.is_contiguous(), + "attn_res_rmsnorm_fwd: inputs must be contiguous"); + TORCH_CHECK(block_residual.sizes() == at::IntArrayRef({N - 1, T, B, H}), + "attn_res_rmsnorm_fwd: block_residual shape must match layer_residual"); + TORCH_CHECK(res_weight.numel() == H, "attn_res_rmsnorm_fwd: res_weight must have H elements"); + TORCH_CHECK(rms_weight.numel() == H, "attn_res_rmsnorm_fwd: rms_weight must have H elements"); + TORCH_CHECK(output_rms_weight.numel() == H, "attn_res_rmsnorm_fwd: output_rms_weight must have H elements"); + + auto output = at::empty_like(layer_residual); + kernels::kimiK3AttnRes::AttnResFwdParams params{}; + params.blockResidual = N > 1 ? reinterpret_cast<__nv_bfloat16 const*>(block_residual.const_data_ptr()) : nullptr; + params.layerResidual = reinterpret_cast<__nv_bfloat16 const*>(layer_residual.const_data_ptr()); + params.resWeight = reinterpret_cast<__nv_bfloat16 const*>(res_weight.const_data_ptr()); + params.rmsWeight = reinterpret_cast<__nv_bfloat16 const*>(rms_weight.const_data_ptr()); + params.outputRmsWeight = reinterpret_cast<__nv_bfloat16 const*>(output_rms_weight.const_data_ptr()); + params.output = reinterpret_cast<__nv_bfloat16*>(output.data_ptr()); + params.numCandidates = N; + params.seqLen = T; + params.batchSize = B; + params.hiddenSize = H; + params.rmsEps = static_cast(rms_eps); + params.outputRmsEps = static_cast(output_rms_eps); + + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + kernels::kimiK3AttnRes::invokeAttnResRmsNormFwd(params, stream); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return output; +} + +std::tuple attn_res_add_rmsnorm_fwd(at::Tensor layer_residual, at::Tensor layer_residual_add, + at::Tensor block_residual, at::Tensor res_weight, at::Tensor rms_weight, at::Tensor output_rms_weight, + double rms_eps, double output_rms_eps) +{ + TORCH_CHECK(layer_residual.dim() == 3, "attn_res_add_rmsnorm_fwd: layer_residual must be [T, B, H]"); + TORCH_CHECK(layer_residual_add.sizes() == layer_residual.sizes(), + "attn_res_add_rmsnorm_fwd: layer_residual_add must match layer_residual"); + TORCH_CHECK(block_residual.dim() == 4, "attn_res_add_rmsnorm_fwd: block_residual must be [K, T, B, H]"); + + int const T = static_cast(layer_residual.size(0)); + int const B = static_cast(layer_residual.size(1)); + int const H = static_cast(layer_residual.size(2)); + int const N = static_cast(block_residual.size(0)) + 1; + + TORCH_CHECK(layer_residual.is_cuda() && layer_residual_add.is_cuda() && block_residual.is_cuda() + && res_weight.is_cuda() && rms_weight.is_cuda() && output_rms_weight.is_cuda(), + "attn_res_add_rmsnorm_fwd: all input tensors must be CUDA tensors"); + TORCH_CHECK(layer_residual_add.device() == layer_residual.device() + && block_residual.device() == layer_residual.device() && res_weight.device() == layer_residual.device() + && rms_weight.device() == layer_residual.device() && output_rms_weight.device() == layer_residual.device(), + "attn_res_add_rmsnorm_fwd: all input tensors must be on the same CUDA device"); + c10::cuda::CUDAGuard device_guard(layer_residual.device()); + check_attn_res_contract(N, T, B, H); + TORCH_CHECK(T == 1 && B == 1 && H == 7168, + "attn_res_add_rmsnorm_fwd: only production decode shape T=B=1, H=7168 is supported"); + TORCH_CHECK((N >= 1 && N <= 9) || N == 12, "attn_res_add_rmsnorm_fwd: supported N values are [1, 9] and 12"); + + TORCH_CHECK(layer_residual.scalar_type() == at::kBFloat16, "attn_res_add_rmsnorm_fwd: layer_residual must be bf16"); + TORCH_CHECK( + layer_residual_add.scalar_type() == at::kBFloat16, "attn_res_add_rmsnorm_fwd: layer_residual_add must be bf16"); + TORCH_CHECK(block_residual.scalar_type() == at::kBFloat16, "attn_res_add_rmsnorm_fwd: block_residual must be bf16"); + TORCH_CHECK(res_weight.scalar_type() == at::kBFloat16, "attn_res_add_rmsnorm_fwd: res_weight must be bf16"); + TORCH_CHECK(rms_weight.scalar_type() == at::kBFloat16, "attn_res_add_rmsnorm_fwd: rms_weight must be bf16"); + TORCH_CHECK( + output_rms_weight.scalar_type() == at::kBFloat16, "attn_res_add_rmsnorm_fwd: output_rms_weight must be bf16"); + TORCH_CHECK(layer_residual.is_contiguous() && layer_residual_add.is_contiguous() && block_residual.is_contiguous() + && res_weight.is_contiguous() && rms_weight.is_contiguous() && output_rms_weight.is_contiguous(), + "attn_res_add_rmsnorm_fwd: inputs must be contiguous"); + TORCH_CHECK(block_residual.sizes() == at::IntArrayRef({N - 1, T, B, H}), + "attn_res_add_rmsnorm_fwd: block_residual shape must match layer_residual"); + TORCH_CHECK(res_weight.numel() == H, "attn_res_add_rmsnorm_fwd: res_weight must have H elements"); + TORCH_CHECK(rms_weight.numel() == H, "attn_res_add_rmsnorm_fwd: rms_weight must have H elements"); + TORCH_CHECK(output_rms_weight.numel() == H, "attn_res_add_rmsnorm_fwd: output_rms_weight must have H elements"); + + auto updated_layer_residual = at::empty_like(layer_residual); + auto output = at::empty_like(layer_residual); + kernels::kimiK3AttnRes::AttnResFwdParams params{}; + params.blockResidual = N > 1 ? reinterpret_cast<__nv_bfloat16 const*>(block_residual.const_data_ptr()) : nullptr; + params.layerResidual = reinterpret_cast<__nv_bfloat16 const*>(layer_residual.const_data_ptr()); + params.layerResidualAdd = reinterpret_cast<__nv_bfloat16 const*>(layer_residual_add.const_data_ptr()); + params.resWeight = reinterpret_cast<__nv_bfloat16 const*>(res_weight.const_data_ptr()); + params.rmsWeight = reinterpret_cast<__nv_bfloat16 const*>(rms_weight.const_data_ptr()); + params.outputRmsWeight = reinterpret_cast<__nv_bfloat16 const*>(output_rms_weight.const_data_ptr()); + params.updatedLayerResidual = reinterpret_cast<__nv_bfloat16*>(updated_layer_residual.data_ptr()); + params.output = reinterpret_cast<__nv_bfloat16*>(output.data_ptr()); + params.numCandidates = N; + params.seqLen = T; + params.batchSize = B; + params.hiddenSize = H; + params.rmsEps = static_cast(rms_eps); + params.outputRmsEps = static_cast(output_rms_eps); + + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + kernels::kimiK3AttnRes::invokeAttnResAddRmsNormFwd(params, stream); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {updated_layer_residual, output}; +} + } // namespace } // namespace torch_ext @@ -134,9 +261,19 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "attn_res_fwd(Tensor layer_residual, Tensor block_residual, " "Tensor res_weight, Tensor rms_weight, float rms_eps) " "-> (Tensor, Tensor, Tensor, Tensor)"); + m.def( + "attn_res_rmsnorm_fwd(Tensor layer_residual, Tensor block_residual, " + "Tensor res_weight, Tensor rms_weight, Tensor output_rms_weight, " + "float rms_eps, float output_rms_eps) -> Tensor"); + m.def( + "attn_res_add_rmsnorm_fwd(Tensor layer_residual, Tensor layer_residual_add, " + "Tensor block_residual, Tensor res_weight, Tensor rms_weight, Tensor output_rms_weight, " + "float rms_eps, float output_rms_eps) -> (Tensor, Tensor)"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("attn_res_fwd", &tensorrt_llm::torch_ext::attn_res_fwd); + m.impl("attn_res_rmsnorm_fwd", &tensorrt_llm::torch_ext::attn_res_rmsnorm_fwd); + m.impl("attn_res_add_rmsnorm_fwd", &tensorrt_llm::torch_ext::attn_res_add_rmsnorm_fwd); } diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index f1f104fd8851..3cb321aa2a1b 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -340,6 +340,19 @@ def _is_mla_layer(cfg, layer_idx: int) -> bool: _FUSED_ATTN_RES_ENABLED = os.environ.get(KIMI_K3_FUSED_ATTN_RES_ENV, "1") == "1" +KIMI_K3_FUSED_ATTN_RES_NORM_ENV = "KIMI_K3_FUSED_ATTN_RES_NORM" +"""Set to ``0`` to keep the trailing RMSNorm out of the attention-residual +kernel, i.e. ``trtllm::attn_res_fwd`` followed by the production RMSNorm +module. + +This is the *only* knob that isolates the norm-fusing ops +(``attn_res_rmsnorm_fwd`` / ``attn_res_add_rmsnorm_fwd``) from the +attention-residual selection itself. ``KIMI_K3_FUSED_ATTN_RES=0`` disables +both and falls all the way back to the exact fp32 reference, so it cannot be +used to A/B the norm fusion against the unfused-norm path.""" + +_FUSED_ATTN_RES_NORM_ENABLED = os.environ.get(KIMI_K3_FUSED_ATTN_RES_NORM_ENV, "1") == "1" + def _apply_attn_res_fused( prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm @@ -352,7 +365,11 @@ def _apply_attn_res_fused( layout. Candidate order matches the reference: snapshots first, the running prefix sum last. """ - if prefix_sum.dtype is not torch.bfloat16: + if ( + prefix_sum.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not block_residual.is_cuda + ): return None M, H = prefix_sum.shape K = int(block_residual.shape[0]) @@ -374,6 +391,123 @@ def _apply_attn_res_fused( return output.reshape(M, H) +def _rms_norm_eps(norm: nn.Module) -> float: + if hasattr(norm, "eps"): + return float(norm.eps) + return float(norm.variance_epsilon) + + +def _note_attn_res_fusion(site: str, fused: bool, M: int, H: int, N: int) -> None: + """Report whether the fused path was actually reached, once per shape. + + ``_FUSED_ATTN_RES_ENABLED`` only says the feature is switched on. It does + not say the shape gate below let the call through, and a rejected call + looks exactly like a disabled one in the logs. Under attention-DP the + per-rank token count decides it, so one job can fuse at low concurrency and + fall back at high concurrency -- without this line, a benchmark that shows + no change is indistinguishable from one that never ran the kernel. + """ + logger.info_once( + f"Kimi K3 attn-res fusion [{site}]: " + f"{'FUSED' if fused else 'fallback'} (M={M}, H={H}, N={N})", + key=f"kimi_k3_attn_res_fusion_{site}_{fused}_{M}_{H}_{N}", + ) + + +def _apply_attn_res_rmsnorm_fused( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, +) -> Optional[torch.Tensor]: + """Fuse attention-residual mixing with its immediately following norm.""" + if ( + prefix_sum.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not block_residual.is_cuda + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + N = K + 1 + # The fused topology is beneficial for the production decode shape only. + # Keep prefill on attn_res_fwd + the production RMSNorm, which exposes + # independent work across tokens and was 41-108% faster in GB300 tests. + if M != 1 or H != 7168 or (N > 9 and N != 12): + _note_attn_res_fusion("attn_res+norm", False, M, H, N) + return None + try: + attn_res_rmsnorm_op = torch.ops.trtllm.attn_res_rmsnorm_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + output = attn_res_rmsnorm_op( + layer_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + ) + _note_attn_res_fusion("attn_res+norm", True, M, H, N) + return output.reshape(M, H) + + +def _apply_attn_res_add_rmsnorm_fused( + prefix_sum: torch.Tensor, + addend: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, +) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """Fuse ``prefix_sum + addend``, attention-residual, and trailing norm. + + The production residual add produces a BF16 tensor that remains live + across the following MLP. The kernel therefore returns that materialized, + BF16-rounded prefix sum alongside the normalized attention-residual + output, while avoiding a separate add launch and a re-read of the + intermediate by attention-residual selection. + """ + if ( + prefix_sum.dtype is not torch.bfloat16 + or addend.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not addend.is_cuda + or not block_residual.is_cuda + or prefix_sum.shape != addend.shape + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + N = K + 1 + if M != 1 or H != 7168 or (N > 9 and N != 12): + _note_attn_res_fusion("add+attn_res+norm", False, M, H, N) + return None + try: + attn_res_add_rmsnorm_op = torch.ops.trtllm.attn_res_add_rmsnorm_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + addend_kernel = addend.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + updated_prefix_sum, output = attn_res_add_rmsnorm_op( + layer_kernel, + addend_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + ) + _note_attn_res_fusion("add+attn_res+norm", True, M, H, N) + return updated_prefix_sum.reshape(M, H), output.reshape(M, H) + + def _apply_attn_res( prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm ) -> torch.Tensor: @@ -402,6 +536,42 @@ def _apply_attn_res( return hidden_states.to(v.dtype) +def _apply_attn_res_and_rmsnorm( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, +) -> torch.Tensor: + """Apply attention-residual selection and the next RMSNorm.""" + if _FUSED_ATTN_RES_ENABLED and _FUSED_ATTN_RES_NORM_ENABLED: + fused = _apply_attn_res_rmsnorm_fused(prefix_sum, block_residual, proj, norm, output_norm) + if fused is not None: + return fused + return output_norm(_apply_attn_res(prefix_sum, block_residual, proj, norm)) + + +def _apply_attn_res_add_and_rmsnorm( + prefix_sum: torch.Tensor, + addend: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Add an attention output to the running residual, then select and norm.""" + if _FUSED_ATTN_RES_ENABLED and _FUSED_ATTN_RES_NORM_ENABLED: + fused = _apply_attn_res_add_rmsnorm_fused( + prefix_sum, addend, block_residual, proj, norm, output_norm + ) + if fused is not None: + return fused + updated_prefix_sum = prefix_sum + addend + return updated_prefix_sum, _apply_attn_res_and_rmsnorm( + updated_prefix_sum, block_residual, proj, norm, output_norm + ) + + # --------------------------------------------------------------------------- # Dense / shared-expert MLP: fused [gate | up] layout (``GatedMLP``). # @@ -2294,12 +2464,15 @@ def forward( valid_block_residual = block_residual[:num_snapshots] if num_snapshots > 0: - hidden_states = _apply_attn_res( + hidden_states = _apply_attn_res_and_rmsnorm( prefix_sum, valid_block_residual, self.self_attention_res_proj, self.self_attention_res_norm, + self.input_layernorm, ) + else: + hidden_states = self.input_layernorm(hidden_states) if self.layer_idx % self.attn_res_block_size == 0: block_residual[num_snapshots].copy_(prefix_sum) @@ -2307,19 +2480,26 @@ def forward( valid_block_residual = block_residual[:num_snapshots] prefix_sum = None - hidden_states = self.input_layernorm(hidden_states) hidden_states = self.self_attn(hidden_states, attn_metadata) - if prefix_sum is not None: - prefix_sum = prefix_sum + hidden_states - else: + if prefix_sum is None: prefix_sum = hidden_states - - hidden_states = _apply_attn_res( - prefix_sum, valid_block_residual, self.mlp_res_proj, self.mlp_res_norm - ) - - hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + hidden_states, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) if self.is_moe: hidden_states = self.block_sparse_moe( hidden_states, getattr(attn_metadata, "all_rank_num_tokens", None) @@ -2368,6 +2548,12 @@ def __init__(self, model_config: ModelConfig): cfg.num_hidden_layers + cfg.attn_res_block_size - 1 ) // cfg.attn_res_block_size + logger.info_once( + f"Kimi K3 attention-residual kernels: fused={_FUSED_ATTN_RES_ENABLED}, " + f"fused_norm={_FUSED_ATTN_RES_NORM_ENABLED}", + key="kimi_k3_attn_res_fusion", + ) + def forward( self, attn_metadata: AttentionMetadata, @@ -2411,13 +2597,13 @@ def forward( # before real weights are used. spec_metadata.maybe_capture_hidden_states(layer.layer_idx, hidden_states, None) - hidden_states = _apply_attn_res( + return _apply_attn_res_and_rmsnorm( hidden_states, block_residual[:num_snapshots], self.output_attn_res_proj, self.output_attn_res_norm, + self.norm, ) - return self.norm(hidden_states) # --------------------------------------------------------------------------- diff --git a/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py b/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py index cea663614621..108561e2659a 100644 --- a/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py +++ b/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py @@ -2,12 +2,25 @@ # SPDX-License-Identifier: Apache-2.0 """Parity tests for the fused Kimi K3 attention-residual op.""" +from unittest import mock + import pytest import torch from torch import nn -from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3RMSNorm, _apply_attn_res_fused +from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE +from tensorrt_llm._torch.models import modeling_kimi_linear +from tensorrt_llm._torch.models.modeling_kimi_linear import ( + KimiK3RMSNorm, + _apply_attn_res, + _apply_attn_res_add_and_rmsnorm, + _apply_attn_res_add_rmsnorm_fused, + _apply_attn_res_and_rmsnorm, + _apply_attn_res_fused, + _apply_attn_res_rmsnorm_fused, +) from tensorrt_llm._torch.modules.kimi_k3_attn_res import apply_attn_res_reference +from tensorrt_llm._torch.modules.rms_norm import RMSNorm HIDDEN_SIZE = 7168 RMS_EPS = 1e-6 @@ -84,3 +97,157 @@ def test_fused_attn_res_matches_torch_reference(num_tokens: int, num_snapshots: cosine, relative_l2 = _similarity(actual, expected) assert cosine > 0.999 assert relative_l2 < 3e-2 + + +OUTPUT_RMS_EPS = 1e-6 +# Production decode: N=4 single-CTA and N=8 split-K. Other legal N share those two topologies. +_DECODE_SNAPSHOTS = (3, 7) + + +def _production_rms_norm( + hidden_states: torch.Tensor, weight: torch.Tensor, eps: float +) -> torch.Tensor: + if IS_FLASHINFER_AVAILABLE: + from tensorrt_llm._torch.custom_ops import flashinfer_rmsnorm + + return flashinfer_rmsnorm(hidden_states.contiguous(), weight, eps) + hidden_float = hidden_states.float() + variance = hidden_float.square().mean(dim=-1, keepdim=True) + return weight * (hidden_float * torch.rsqrt(variance + eps)).to(hidden_states.dtype) + + +def _make_decode_case(num_snapshots: int): + torch.manual_seed(0) + prefix_sum = torch.randn(1, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05 + addend = torch.randn(1, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05 + block_residual = ( + torch.randn(num_snapshots, 1, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05 + ) + projection = nn.Linear(HIDDEN_SIZE, 1, bias=False, dtype=torch.bfloat16, device="cuda") + score_norm = KimiK3RMSNorm(HIDDEN_SIZE, eps=RMS_EPS).to(device="cuda", dtype=torch.bfloat16) + output_norm = RMSNorm( + hidden_size=HIDDEN_SIZE, + eps=OUTPUT_RMS_EPS, + dtype=torch.bfloat16, + device=torch.device("cuda"), + ) + projection.weight.mul_(0.02) + return prefix_sum, addend, block_residual, projection, score_norm, output_norm + + +@pytest.mark.parametrize("num_snapshots", _DECODE_SNAPSHOTS) +@torch.no_grad() +def test_decode_rmsnorm_fusion_matches_unfused(num_snapshots: int) -> None: + prefix_sum, _addend, block_residual, projection, score_norm, output_norm = _make_decode_case( + num_snapshots + ) + expected = output_norm(_apply_attn_res(prefix_sum, block_residual, projection, score_norm)) + actual = _apply_attn_res_and_rmsnorm( + prefix_sum, block_residual, projection, score_norm, output_norm + ) + cosine, relative_l2 = _similarity(actual, expected) + assert cosine > 0.9999 + assert relative_l2 < 5e-3 + + +@pytest.mark.parametrize("num_snapshots", _DECODE_SNAPSHOTS) +@torch.no_grad() +def test_decode_add_rmsnorm_fusion_matches_separate_add(num_snapshots: int) -> None: + prefix_sum, addend, block_residual, projection, score_norm, output_norm = _make_decode_case( + num_snapshots + ) + expected_updated = prefix_sum + addend + expected_output = output_norm( + _apply_attn_res(expected_updated, block_residual, projection, score_norm) + ) + actual_updated, actual_output = _apply_attn_res_add_and_rmsnorm( + prefix_sum, addend, block_residual, projection, score_norm, output_norm + ) + assert torch.equal(actual_updated, expected_updated) + cosine, relative_l2 = _similarity(actual_output, expected_output) + assert cosine > 0.9999 + assert relative_l2 < 5e-3 + + +@torch.no_grad() +def test_decode_fusion_gate_skips_prefill() -> None: + prefix_sum, addend, block_residual, projection, score_norm, output_norm = _make_decode_case(3) + prefix_sum = prefix_sum.expand(64, -1).contiguous() + addend = addend.expand(64, -1).contiguous() + block_residual = block_residual.expand(-1, 64, -1).contiguous() + assert ( + _apply_attn_res_rmsnorm_fused( + prefix_sum, block_residual, projection, score_norm, output_norm + ) + is None + ) + assert ( + _apply_attn_res_add_rmsnorm_fused( + prefix_sum, addend, block_residual, projection, score_norm, output_norm + ) + is None + ) + + +@torch.no_grad() +def test_decode_norm_flag_keeps_unfused_path(monkeypatch: pytest.MonkeyPatch) -> None: + prefix_sum, addend, block_residual, projection, score_norm, output_norm = _make_decode_case(3) + unexpected = mock.Mock(side_effect=AssertionError("norm flag off still reached fused op")) + monkeypatch.setattr(modeling_kimi_linear, "_FUSED_ATTN_RES_NORM_ENABLED", False) + monkeypatch.setattr(modeling_kimi_linear, "_apply_attn_res_rmsnorm_fused", unexpected) + monkeypatch.setattr(modeling_kimi_linear, "_apply_attn_res_add_rmsnorm_fused", unexpected) + expected_updated = prefix_sum + addend + expected_output = output_norm( + _apply_attn_res(expected_updated, block_residual, projection, score_norm) + ) + actual_updated, actual_output = _apply_attn_res_add_and_rmsnorm( + prefix_sum, addend, block_residual, projection, score_norm, output_norm + ) + unexpected.assert_not_called() + assert torch.equal(actual_updated, expected_updated) + cosine, relative_l2 = _similarity(actual_output, expected_output) + assert cosine > 0.9999 + assert relative_l2 < 5e-3 + + +@torch.no_grad() +def test_decode_add_rmsnorm_cuda_graph_replay() -> None: + prefix_sum, addend, block_residual, projection, score_norm, output_norm = _make_decode_case(3) + layer = prefix_sum.reshape(1, 1, HIDDEN_SIZE).contiguous() + addend_b = addend.reshape(1, 1, HIDDEN_SIZE).contiguous() + block = block_residual.reshape(block_residual.shape[0], 1, 1, HIDDEN_SIZE).contiguous() + expected_updated = layer + addend_b + mixed, *_ = torch.ops.trtllm.attn_res_fwd( + expected_updated, block, projection.weight.reshape(-1), score_norm.weight, RMS_EPS + ) + expected_output = _production_rms_norm(mixed, output_norm.weight, OUTPUT_RMS_EPS) + + torch.ops.trtllm.attn_res_add_rmsnorm_fwd( + layer, + addend_b, + block, + projection.weight.reshape(-1), + score_norm.weight, + output_norm.weight, + RMS_EPS, + OUTPUT_RMS_EPS, + ) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual_updated, actual_output = torch.ops.trtllm.attn_res_add_rmsnorm_fwd( + layer, + addend_b, + block, + projection.weight.reshape(-1), + score_norm.weight, + output_norm.weight, + RMS_EPS, + OUTPUT_RMS_EPS, + ) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(actual_updated, expected_updated) + cosine, relative_l2 = _similarity(actual_output, expected_output) + assert cosine > 0.9999 + assert relative_l2 < 5e-3