From c4b7ded450f768f620846efc3120de7f08569940 Mon Sep 17 00:00:00 2001 From: Xin Guan <294044352+xguannv@users.noreply.github.com> Date: Mon, 3 Aug 2026 13:45:03 +0800 Subject: [PATCH 1/5] [TRTLLM][Kimi K3] Fuse decode attention residual tail Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com> --- .../kernels/kimiK3AttnRes/attnResFwd.cu | 405 +++++++++++-- .../kernels/kimiK3AttnRes/attnResFwd.h | 42 +- cpp/tensorrt_llm/thop/attnResOp.cpp | 137 +++++ .../_torch/models/modeling_kimi_linear.py | 179 +++++- .../kimi_k3_attn_res_add_rmsnorm.py | 331 +++++++++++ .../test_attn_res_rmsnorm_op.py | 536 ++++++++++++++++++ 6 files changed, 1548 insertions(+), 82 deletions(-) create mode 100644 tests/microbenchmarks/kimi_k3_attn_res_add_rmsnorm.py create mode 100644 tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu index f1e284118d6f..d13a42eac6ed 100644 --- a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu @@ -30,9 +30,9 @@ // 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 #include #include #include @@ -72,6 +72,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 +727,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 +1016,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 +1047,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 +1061,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 +1102,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 +1175,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 +1205,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 +1273,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 +1360,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 +1412,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 +1462,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 +1485,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 +1558,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 +1594,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, @@ -1457,10 +1703,12 @@ void invokeAttnResFwd(AttnResFwdParams const& params, cudaStream_t stream) float const rms_eps = params.rmsEps; int dev = 0; - TLLM_CUDA_CHECK(cudaGetDevice(&dev)); + cudaGetDevice(&dev); int num_sm = attn_res_fwd_grid_size(dev); - TLLM_CHECK_WITH_INFO(num_sm > 0, "attn_res_fwd: failed to query the SM count of device %d", dev); - TLLM_CHECK_WITH_INFO(N <= N_MAX, "attn_res_fwd: unsupported N=%d (max %d)", N, N_MAX); + if (num_sm <= 0 || N > N_MAX) + { + return; + } if (H == 8192) { @@ -1510,7 +1758,7 @@ void invokeAttnResFwd(AttnResFwdParams const& params, cudaStream_t stream) else if (N == 12 && T == 1024) { launch_fwd<7168, 4, false, true>(block_residual, layer_residual, res_weight, rms_weight, output, rsigma, - probs, logits, N, T, B, rms_eps, std::max(1, num_sm - 1), stream); + probs, logits, N, T, B, rms_eps, num_sm - 1, stream); } else { @@ -1546,12 +1794,59 @@ void invokeAttnResFwd(AttnResFwdParams const& params, cudaStream_t stream) logits, N, T, B, rms_eps, num_sm, 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 { - TLLM_CHECK_WITH_INFO(false, "attn_res_fwd: unsupported hidden size H=%d", H); + 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) +{ + 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: break; + } +} + +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..5b07c1029912 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -352,7 +352,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 +378,102 @@ 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 _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): + 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), + ) + 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): + 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), + ) + 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 +502,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: + 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: + 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 +2430,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 +2446,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 +2514,11 @@ 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}", + key="kimi_k3_attn_res_fusion", + ) + def forward( self, attn_metadata: AttentionMetadata, @@ -2411,13 +2562,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/microbenchmarks/kimi_k3_attn_res_add_rmsnorm.py b/tests/microbenchmarks/kimi_k3_attn_res_add_rmsnorm.py new file mode 100644 index 000000000000..f43c352e71ef --- /dev/null +++ b/tests/microbenchmarks/kimi_k3_attn_res_add_rmsnorm.py @@ -0,0 +1,331 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Microbenchmark the incremental Kimi K3 attention-output add fusion. + +The production baseline after the trailing-RMSNorm fusion is two kernels: + + updated_prefix = prefix_sum + attention_output + trtllm::attn_res_rmsnorm_fwd(updated_prefix, ...) + +The new path is one kernel and still materializes ``updated_prefix`` for the +MLP residual that follows: + + updated_prefix, output = + trtllm::attn_res_add_rmsnorm_fwd(prefix_sum, attention_output, ...) + +CUDA-graph replay timings remove Python and dispatcher overhead. ``--profile`` +emits eager calls inside NVTX ranges for Nsys attribution. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from typing import Any + +import torch + +from tensorrt_llm._torch.custom_ops import flashinfer_rmsnorm + +HIDDEN_SIZE = 7168 +RMS_EPS = 1e-6 + + +@dataclass +class CaseInputs: + prefix_sum: torch.Tensor + attention_output: torch.Tensor + block_residual: torch.Tensor + res_weight: torch.Tensor + score_rms_weight: torch.Tensor + output_rms_weight: torch.Tensor + + +def _parse_candidates(value: str) -> list[int]: + candidates = [int(item) for item in value.split(",")] + if any(candidate not in {*range(1, 10), 12} for candidate in candidates): + raise argparse.ArgumentTypeError("candidate counts must be in [1, 9] or equal to 12") + return candidates + + +def _make_inputs(num_candidates: int) -> CaseInputs: + device = torch.device("cuda") + shape = (1, 1, HIDDEN_SIZE) + prefix_sum = torch.empty(shape, dtype=torch.bfloat16, device=device).uniform_(-0.05, 0.05) + attention_output = torch.empty(shape, dtype=torch.bfloat16, device=device).uniform_(-0.05, 0.05) + block_residual = torch.empty( + (num_candidates - 1, *shape), + dtype=torch.bfloat16, + device=device, + ).uniform_(-0.05, 0.05) + res_weight = torch.empty(HIDDEN_SIZE, dtype=torch.bfloat16, device=device).uniform_(-0.02, 0.02) + score_rms_weight = torch.empty(HIDDEN_SIZE, dtype=torch.bfloat16, device=device).uniform_( + 0.98, 1.02 + ) + output_rms_weight = torch.empty(HIDDEN_SIZE, dtype=torch.bfloat16, device=device).uniform_( + 0.98, 1.02 + ) + return CaseInputs( + prefix_sum=prefix_sum, + attention_output=attention_output, + block_residual=block_residual, + res_weight=res_weight, + score_rms_weight=score_rms_weight, + output_rms_weight=output_rms_weight, + ) + + +def _attn_res(inputs: CaseInputs, updated_prefix: torch.Tensor) -> torch.Tensor: + output, _rsigma, _probs, _logits = torch.ops.trtllm.attn_res_fwd( + updated_prefix, + inputs.block_residual, + inputs.res_weight, + inputs.score_rms_weight, + RMS_EPS, + ) + return output + + +def _three_kernel(inputs: CaseInputs) -> tuple[torch.Tensor, torch.Tensor]: + updated_prefix = inputs.prefix_sum + inputs.attention_output + mixed = _attn_res(inputs, updated_prefix) + output = flashinfer_rmsnorm(mixed, inputs.output_rms_weight, RMS_EPS) + return updated_prefix, output + + +def _two_kernel(inputs: CaseInputs) -> tuple[torch.Tensor, torch.Tensor]: + updated_prefix = inputs.prefix_sum + inputs.attention_output + output = torch.ops.trtllm.attn_res_rmsnorm_fwd( + updated_prefix, + inputs.block_residual, + inputs.res_weight, + inputs.score_rms_weight, + inputs.output_rms_weight, + RMS_EPS, + RMS_EPS, + ) + return updated_prefix, output + + +def _fused(inputs: CaseInputs) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.trtllm.attn_res_add_rmsnorm_fwd( + inputs.prefix_sum, + inputs.attention_output, + inputs.block_residual, + inputs.res_weight, + inputs.score_rms_weight, + inputs.output_rms_weight, + RMS_EPS, + RMS_EPS, + ) + + +def _capture( + fn: Callable[[], Any], +) -> tuple[torch.cuda.CUDAGraph, Any]: + fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = fn() + graph.replay() + torch.cuda.synchronize() + return graph, output + + +def _time_graph( + fn: Callable[[], Any], + iterations: int, + samples: int, + chain_length: int, +) -> tuple[float, float, float]: + def chained_fn() -> Any: + output = None + for _ in range(chain_length): + output = fn() + return output + + graph, output = _capture(chained_fn) + del output + for _ in range(20): + graph.replay() + torch.cuda.synchronize() + + timings = [] + for _ in range(samples): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + graph.replay() + end.record() + end.synchronize() + timings.append(start.elapsed_time(end) * 1000.0 / iterations / chain_length) + return statistics.median(timings), min(timings), max(timings) + + +def _similarity( + actual: torch.Tensor, + expected: torch.Tensor, +) -> tuple[float, float]: + actual_float = actual.float().flatten() + expected_float = expected.float().flatten() + cosine = torch.nn.functional.cosine_similarity(actual_float, expected_float, dim=0).item() + relative_l2 = ((actual_float - expected_float).norm() / (expected_float.norm() + 1e-12)).item() + return cosine, relative_l2 + + +def _benchmark_case( + num_candidates: int, + iterations: int, + samples: int, + chain_length: int, +) -> dict[str, float | int]: + inputs = _make_inputs(num_candidates) + expected_prefix, expected_output = _two_kernel(inputs) + actual_prefix, actual_output = _fused(inputs) + torch.cuda.synchronize() + if not torch.equal(actual_prefix, expected_prefix): + raise AssertionError(f"N={num_candidates}: fused updated prefix is not exact") + cosine, relative_l2 = _similarity(actual_output, expected_output) + if cosine <= 0.9999 or relative_l2 >= 5e-3: + raise AssertionError(f"N={num_candidates}: cosine={cosine}, relative_l2={relative_l2}") + + add_us, add_min_us, add_max_us = _time_graph( + lambda: inputs.prefix_sum + inputs.attention_output, + iterations, + samples, + chain_length, + ) + three_us, three_min_us, three_max_us = _time_graph( + lambda: _three_kernel(inputs), iterations, samples, chain_length + ) + two_us, two_min_us, two_max_us = _time_graph( + lambda: _two_kernel(inputs), iterations, samples, chain_length + ) + fused_us, fused_min_us, fused_max_us = _time_graph( + lambda: _fused(inputs), iterations, samples, chain_length + ) + + return { + "num_tokens": 1, + "num_candidates": num_candidates, + "iterations": iterations, + "samples": samples, + "chain_length": chain_length, + "cosine": cosine, + "relative_l2": relative_l2, + "add_us": add_us, + "add_min_us": add_min_us, + "add_max_us": add_max_us, + "three_kernel_us": three_us, + "three_kernel_min_us": three_min_us, + "three_kernel_max_us": three_max_us, + "two_kernel_us": two_us, + "two_kernel_min_us": two_min_us, + "two_kernel_max_us": two_max_us, + "fused_us": fused_us, + "fused_min_us": fused_min_us, + "fused_max_us": fused_max_us, + "fused_vs_two_kernel_pct": (fused_us / two_us - 1.0) * 100.0, + "saved_vs_two_kernel_us": two_us - fused_us, + "fused_vs_three_kernel_pct": (fused_us / three_us - 1.0) * 100.0, + "saved_vs_three_kernel_us": three_us - fused_us, + } + + +def _profile_case( + num_candidates: int, + iterations: int, +) -> None: + inputs = _make_inputs(num_candidates) + modes: Sequence[tuple[str, Callable[[], Any]]] = ( + ("add", lambda: inputs.prefix_sum + inputs.attention_output), + ("three_kernel", lambda: _three_kernel(inputs)), + ("two_kernel", lambda: _two_kernel(inputs)), + ("fused", lambda: _fused(inputs)), + ) + for _name, fn in modes: + for _ in range(10): + fn() + torch.cuda.synchronize() + + for name, fn in modes: + range_name = f"attn_res_add|T=1|N={num_candidates}|mode={name}" + torch.cuda.nvtx.range_push(range_name) + for _ in range(iterations): + fn() + torch.cuda.synchronize() + torch.cuda.nvtx.range_pop() + print( + json.dumps( + {"profile_range": range_name, "iterations": iterations}, + sort_keys=True, + ), + flush=True, + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--candidates", + type=_parse_candidates, + default=_parse_candidates("1,2,3,4,5,6,7,8,9,12"), + ) + parser.add_argument("--iterations", type=int, default=2000) + parser.add_argument("--samples", type=int, default=7) + parser.add_argument( + "--chain-length", + type=int, + default=1, + help="Capture this many copies of each mode in one CUDA graph.", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Emit eager kernels in NVTX ranges for Nsys instead of timing.", + ) + args = parser.parse_args() + + if args.chain_length < 1: + parser.error("--chain-length must be positive") + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + capability = torch.cuda.get_device_capability() + if capability not in {(10, 0), (10, 3)}: + raise RuntimeError(f"SM100/SM103 is required, got {capability}") + torch.manual_seed(0) + print( + json.dumps( + { + "device": torch.cuda.get_device_name(), + "capability": capability, + "profile": args.profile, + }, + sort_keys=True, + ), + flush=True, + ) + + for num_candidates in args.candidates: + if args.profile: + _profile_case(num_candidates, min(args.iterations, 100)) + else: + print( + json.dumps( + _benchmark_case( + num_candidates, args.iterations, args.samples, args.chain_length + ), + sort_keys=True, + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py b/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py new file mode 100644 index 000000000000..7108be67093d --- /dev/null +++ b/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py @@ -0,0 +1,536 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Parity tests for the fused Kimi K3 attention-residual + RMSNorm op.""" + +from unittest import mock + +import pytest +import torch +from torch import nn + +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_rmsnorm_fused, +) +from tensorrt_llm._torch.modules.rms_norm import RMSNorm + +HIDDEN_SIZE = 7168 +ATTN_RES_RMS_EPS = 1e-6 +OUTPUT_RMS_EPS = 1e-6 + + +def _has_supported_gpu() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability(0) in { + (10, 0), + (10, 3), + } + + +pytestmark = pytest.mark.skipif( + not _has_supported_gpu(), + reason="Kimi K3 attention-residual kernels require SM100/SM103", +) + + +def _production_rms_norm( + hidden_states: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Apply the unfused production RMSNorm that the new op replaces.""" + 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) + normalized = hidden_float * torch.rsqrt(variance + eps) + return weight * normalized.to(hidden_states.dtype) + + +def _make_inputs( + num_tokens: int, + num_snapshots: int, +) -> tuple[torch.Tensor, ...]: + torch.manual_seed(0) + layer_residual = ( + torch.randn( + num_tokens, + 1, + HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + block_residual = ( + torch.randn( + num_snapshots, + num_tokens, + 1, + HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + res_weight = torch.randn(HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.02 + score_rms_weight = 1 + torch.randn(HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.02 + output_rms_weight = 1 + torch.randn(HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.02 + return ( + layer_residual, + block_residual, + res_weight.contiguous(), + score_rms_weight.contiguous(), + output_rms_weight.contiguous(), + ) + + +def _unfused_reference( + layer_residual: torch.Tensor, + block_residual: torch.Tensor, + res_weight: torch.Tensor, + score_rms_weight: torch.Tensor, + output_rms_weight: torch.Tensor, +) -> torch.Tensor: + mixed, _rsigma, _probs, _logits = torch.ops.trtllm.attn_res_fwd( + layer_residual, + block_residual, + res_weight, + score_rms_weight, + ATTN_RES_RMS_EPS, + ) + return _production_rms_norm(mixed, output_rms_weight, OUTPUT_RMS_EPS) + + +def _fused( + layer_residual: torch.Tensor, + block_residual: torch.Tensor, + res_weight: torch.Tensor, + score_rms_weight: torch.Tensor, + output_rms_weight: torch.Tensor, +) -> torch.Tensor: + return torch.ops.trtllm.attn_res_rmsnorm_fwd( + layer_residual, + block_residual, + res_weight, + score_rms_weight, + output_rms_weight, + ATTN_RES_RMS_EPS, + OUTPUT_RMS_EPS, + ) + + +def _fused_add( + layer_residual: torch.Tensor, + layer_residual_add: torch.Tensor, + block_residual: torch.Tensor, + res_weight: torch.Tensor, + score_rms_weight: torch.Tensor, + output_rms_weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.trtllm.attn_res_add_rmsnorm_fwd( + layer_residual, + layer_residual_add, + block_residual, + res_weight, + score_rms_weight, + output_rms_weight, + ATTN_RES_RMS_EPS, + OUTPUT_RMS_EPS, + ) + + +def _similarity( + actual: torch.Tensor, + expected: torch.Tensor, +) -> tuple[float, float]: + actual_float = actual.float() + expected_float = expected.float() + cosine = torch.nn.functional.cosine_similarity( + actual_float.flatten(), + expected_float.flatten(), + dim=0, + ).item() + relative_l2 = ((actual_float - expected_float).norm() / (expected_float.norm() + 1e-12)).item() + return cosine, relative_l2 + + +@pytest.mark.parametrize( + ("num_tokens", "num_snapshots"), + [ + (1, 0), # N=1 single-CTA decode + (1, 1), # N=2 single-CTA decode + (1, 2), # N=3 single-CTA decode + (1, 3), # N=4 single-CTA decode + (1, 4), # N=5 CTA-cluster decode + (1, 5), # N=6 CTA-cluster decode + (1, 6), # N=7 CTA-cluster decode + (1, 7), # N=8 CTA-cluster decode + (1, 8), # N=9 CTA-cluster decode + (1, 11), # N=12 CTA-cluster decode + ], +) +@torch.no_grad() +def test_attn_res_rmsnorm_matches_unfused( + num_tokens: int, + num_snapshots: int, +) -> None: + inputs = _make_inputs(num_tokens, num_snapshots) + expected = _unfused_reference(*inputs) + actual = _fused(*inputs) + + assert actual.shape == expected.shape + assert actual.dtype == torch.bfloat16 + cosine, relative_l2 = _similarity(actual, expected) + assert cosine > 0.9999 + assert relative_l2 < 5e-3 + + +@torch.no_grad() +def test_attn_res_rmsnorm_op_rejects_multi_token() -> None: + inputs = _make_inputs(num_tokens=2, num_snapshots=3) + with pytest.raises(RuntimeError, match="only production decode shape"): + _fused(*inputs) + + +@torch.no_grad() +def test_attn_res_rmsnorm_cuda_graph_replay() -> None: + inputs = _make_inputs(num_tokens=1, num_snapshots=3) + expected = _unfused_reference(*inputs) + + _fused(*inputs) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = _fused(*inputs) + graph.replay() + torch.cuda.synchronize() + + cosine, relative_l2 = _similarity(actual, expected) + assert cosine > 0.9999 + assert relative_l2 < 5e-3 + + +@pytest.mark.parametrize("num_snapshots", list(range(9)) + [11]) +@torch.no_grad() +def test_attn_res_add_rmsnorm_matches_separate_add( + num_snapshots: int, +) -> None: + inputs = _make_inputs(num_tokens=1, num_snapshots=num_snapshots) + layer_residual_add = (torch.randn_like(inputs[0]) * 0.05).contiguous() + expected_updated = inputs[0] + layer_residual_add + expected_output = _unfused_reference( + expected_updated, + inputs[1], + inputs[2], + inputs[3], + inputs[4], + ) + actual_updated, actual_output = _fused_add( + inputs[0], + layer_residual_add, + inputs[1], + inputs[2], + inputs[3], + inputs[4], + ) + + assert torch.equal(actual_updated, expected_updated) + cosine, relative_l2 = _similarity(actual_output, expected_output) + assert cosine > 0.9999 + assert relative_l2 < 5e-3 + + +@pytest.mark.parametrize("num_snapshots", [3, 7, 8]) +@torch.no_grad() +def test_attn_res_add_rmsnorm_cuda_graph_replay( + num_snapshots: int, +) -> None: + inputs = _make_inputs(num_tokens=1, num_snapshots=num_snapshots) + layer_residual_add = (torch.randn_like(inputs[0]) * 0.05).contiguous() + expected_updated = inputs[0] + layer_residual_add + expected_output = _unfused_reference( + expected_updated, + inputs[1], + inputs[2], + inputs[3], + inputs[4], + ) + + _fused_add(inputs[0], layer_residual_add, *inputs[1:]) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual_updated, actual_output = _fused_add(inputs[0], layer_residual_add, *inputs[1:]) + for _ in range(100): + 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 + + +@torch.no_grad() +def test_model_helper_dispatches_fused_attn_res_add_rmsnorm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + ( + layer_residual, + block_residual, + res_weight, + score_rms_weight, + output_rms_weight, + ) = _make_inputs(num_tokens=1, num_snapshots=3) + layer_residual_add = (torch.randn_like(layer_residual) * 0.05).contiguous() + projection = nn.Linear( + HIDDEN_SIZE, + 1, + bias=False, + dtype=torch.bfloat16, + device="cuda", + ) + score_norm = KimiK3RMSNorm( + HIDDEN_SIZE, + eps=ATTN_RES_RMS_EPS, + dtype=torch.bfloat16, + device=torch.device("cuda"), + ) + output_norm = RMSNorm( + hidden_size=HIDDEN_SIZE, + eps=OUTPUT_RMS_EPS, + dtype=torch.bfloat16, + device=torch.device("cuda"), + ) + projection.weight.copy_(res_weight.reshape(1, -1)) + score_norm.weight.copy_(score_rms_weight) + output_norm.weight.copy_(output_rms_weight) + + prefix_sum = layer_residual[:, 0, :] + addend = layer_residual_add[:, 0, :] + block_kernel_layout = block_residual[:, :, 0, :] + expected_updated = prefix_sum + addend + expected_output = output_norm( + _apply_attn_res( + expected_updated, + block_kernel_layout, + projection, + score_norm, + ) + ) + unexpected_fallback = mock.Mock( + side_effect=AssertionError("model helper did not dispatch fused add") + ) + monkeypatch.setattr( + modeling_kimi_linear, + "_apply_attn_res_and_rmsnorm", + unexpected_fallback, + ) + actual_updated, actual_output = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + addend, + block_kernel_layout, + projection, + score_norm, + output_norm, + ) + + unexpected_fallback.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_model_helper_rejects_multi_token_fused_add() -> None: + ( + layer_residual, + block_residual, + res_weight, + score_rms_weight, + output_rms_weight, + ) = _make_inputs(num_tokens=64, num_snapshots=5) + layer_residual_add = (torch.randn_like(layer_residual) * 0.05).contiguous() + projection = nn.Linear( + HIDDEN_SIZE, + 1, + bias=False, + dtype=torch.bfloat16, + device="cuda", + ) + score_norm = KimiK3RMSNorm( + HIDDEN_SIZE, + eps=ATTN_RES_RMS_EPS, + dtype=torch.bfloat16, + device=torch.device("cuda"), + ) + output_norm = RMSNorm( + hidden_size=HIDDEN_SIZE, + eps=OUTPUT_RMS_EPS, + dtype=torch.bfloat16, + device=torch.device("cuda"), + ) + projection.weight.copy_(res_weight.reshape(1, -1)) + score_norm.weight.copy_(score_rms_weight) + output_norm.weight.copy_(output_rms_weight) + + assert ( + _apply_attn_res_add_rmsnorm_fused( + layer_residual[:, 0, :], + layer_residual_add[:, 0, :], + block_residual[:, :, 0, :], + projection, + score_norm, + output_norm, + ) + is None + ) + + +@pytest.mark.parametrize( + ("num_tokens", "num_snapshots"), + [ + (1, 3), + ], +) +@torch.no_grad() +def test_model_helper_dispatches_fused_attn_res_rmsnorm( + num_tokens: int, + num_snapshots: int, + monkeypatch: pytest.MonkeyPatch, +) -> None: + ( + layer_residual, + block_residual, + res_weight, + score_rms_weight, + output_rms_weight, + ) = _make_inputs(num_tokens, num_snapshots) + projection = nn.Linear( + HIDDEN_SIZE, + 1, + bias=False, + dtype=torch.bfloat16, + device="cuda", + ) + score_norm = KimiK3RMSNorm( + HIDDEN_SIZE, + eps=ATTN_RES_RMS_EPS, + dtype=torch.bfloat16, + device=torch.device("cuda"), + ) + output_norm = RMSNorm( + hidden_size=HIDDEN_SIZE, + eps=OUTPUT_RMS_EPS, + dtype=torch.bfloat16, + device=torch.device("cuda"), + ) + projection.weight.copy_(res_weight.reshape(1, -1)) + score_norm.weight.copy_(score_rms_weight) + output_norm.weight.copy_(output_rms_weight) + + prefix_sum = layer_residual[:, 0, :] + block_kernel_layout = block_residual[:, :, 0, :] + expected = output_norm( + _apply_attn_res( + prefix_sum, + block_kernel_layout, + projection, + score_norm, + ) + ) + unexpected_fallback = mock.Mock( + side_effect=AssertionError("model helper did not dispatch the fused op") + ) + monkeypatch.setattr( + modeling_kimi_linear, + "_apply_attn_res", + unexpected_fallback, + ) + actual = _apply_attn_res_and_rmsnorm( + prefix_sum, + block_kernel_layout, + projection, + score_norm, + output_norm, + ) + + cosine, relative_l2 = _similarity(actual, expected) + unexpected_fallback.assert_not_called() + assert cosine > 0.9999 + assert relative_l2 < 5e-3 + + +@torch.no_grad() +def test_model_helper_keeps_multi_token_rmsnorm_split() -> None: + ( + layer_residual, + block_residual, + res_weight, + score_rms_weight, + output_rms_weight, + ) = _make_inputs(num_tokens=64, num_snapshots=5) + projection = nn.Linear( + HIDDEN_SIZE, + 1, + bias=False, + dtype=torch.bfloat16, + device="cuda", + ) + score_norm = KimiK3RMSNorm( + HIDDEN_SIZE, + eps=ATTN_RES_RMS_EPS, + dtype=torch.bfloat16, + device=torch.device("cuda"), + ) + output_norm = RMSNorm( + hidden_size=HIDDEN_SIZE, + eps=OUTPUT_RMS_EPS, + dtype=torch.bfloat16, + device=torch.device("cuda"), + ) + projection.weight.copy_(res_weight.reshape(1, -1)) + score_norm.weight.copy_(score_rms_weight) + output_norm.weight.copy_(output_rms_weight) + + prefix_sum = layer_residual[:, 0, :] + block_kernel_layout = block_residual[:, :, 0, :] + assert ( + _apply_attn_res_rmsnorm_fused( + prefix_sum, + block_kernel_layout, + projection, + score_norm, + output_norm, + ) + is None + ) + + expected = output_norm( + _apply_attn_res( + prefix_sum, + block_kernel_layout, + projection, + score_norm, + ) + ) + actual = _apply_attn_res_and_rmsnorm( + prefix_sum, + block_kernel_layout, + projection, + score_norm, + output_norm, + ) + cosine, relative_l2 = _similarity(actual, expected) + assert cosine > 0.9999 + assert relative_l2 < 5e-3 From 8263ed59bcfaf79c0cd51fe15d8ce8bb062df065 Mon Sep 17 00:00:00 2001 From: Xin Guan <294044352+xguannv@users.noreply.github.com> Date: Tue, 18 Aug 2026 01:48:33 +0000 Subject: [PATCH 2/5] [TRTLLM][Kimi K3] Isolate attn-res norm fusion behind its own A/B switch KIMI_K3_FUSED_ATTN_RES=0 falls back to the fp32 reference, so it cannot A/B this fusion against the pre-port path. Log whether the shape gate actually fired, and compare microbench numerics to the three-kernel baseline. Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com> --- .../_torch/models/modeling_kimi_linear.py | 41 ++++++++++- .../kimi_k3_attn_res_add_rmsnorm.py | 7 ++ .../test_attn_res_rmsnorm_op.py | 71 +++++++++++++++++++ 3 files changed, 116 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 5b07c1029912..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 @@ -384,6 +397,23 @@ def _rms_norm_eps(norm: nn.Module) -> float: 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, @@ -405,6 +435,7 @@ def _apply_attn_res_rmsnorm_fused( # 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 @@ -421,6 +452,7 @@ def _apply_attn_res_rmsnorm_fused( float(norm.eps), _rms_norm_eps(output_norm), ) + _note_attn_res_fusion("attn_res+norm", True, M, H, N) return output.reshape(M, H) @@ -453,6 +485,7 @@ def _apply_attn_res_add_rmsnorm_fused( 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 @@ -471,6 +504,7 @@ def _apply_attn_res_add_rmsnorm_fused( 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) @@ -510,7 +544,7 @@ def _apply_attn_res_and_rmsnorm( output_norm: nn.Module, ) -> torch.Tensor: """Apply attention-residual selection and the next RMSNorm.""" - if _FUSED_ATTN_RES_ENABLED: + 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 @@ -526,7 +560,7 @@ def _apply_attn_res_add_and_rmsnorm( 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: + 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 ) @@ -2515,7 +2549,8 @@ def __init__(self, model_config: ModelConfig): ) // cfg.attn_res_block_size logger.info_once( - f"Kimi K3 attention-residual kernels: fused={_FUSED_ATTN_RES_ENABLED}", + 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", ) diff --git a/tests/microbenchmarks/kimi_k3_attn_res_add_rmsnorm.py b/tests/microbenchmarks/kimi_k3_attn_res_add_rmsnorm.py index f43c352e71ef..ef0e6bb1b3d5 100644 --- a/tests/microbenchmarks/kimi_k3_attn_res_add_rmsnorm.py +++ b/tests/microbenchmarks/kimi_k3_attn_res_add_rmsnorm.py @@ -187,12 +187,17 @@ def _benchmark_case( inputs = _make_inputs(num_candidates) expected_prefix, expected_output = _two_kernel(inputs) actual_prefix, actual_output = _fused(inputs) + _, three_kernel_output = _three_kernel(inputs) torch.cuda.synchronize() if not torch.equal(actual_prefix, expected_prefix): raise AssertionError(f"N={num_candidates}: fused updated prefix is not exact") cosine, relative_l2 = _similarity(actual_output, expected_output) if cosine <= 0.9999 or relative_l2 >= 5e-3: raise AssertionError(f"N={num_candidates}: cosine={cosine}, relative_l2={relative_l2}") + # attn_res_fwd + the production RMSNorm is what shipped before the trailing + # norm was folded into the kernel, so this -- not two_kernel, which is also + # a norm-fusing op -- is the baseline the fusion has to be judged against. + three_cosine, three_relative_l2 = _similarity(actual_output, three_kernel_output) add_us, add_min_us, add_max_us = _time_graph( lambda: inputs.prefix_sum + inputs.attention_output, @@ -218,6 +223,8 @@ def _benchmark_case( "chain_length": chain_length, "cosine": cosine, "relative_l2": relative_l2, + "cosine_vs_three_kernel": three_cosine, + "relative_l2_vs_three_kernel": three_relative_l2, "add_us": add_us, "add_min_us": add_min_us, "add_max_us": add_max_us, diff --git a/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py b/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py index 7108be67093d..6f845b808f33 100644 --- a/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py +++ b/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py @@ -534,3 +534,74 @@ def test_model_helper_keeps_multi_token_rmsnorm_split() -> None: cosine, relative_l2 = _similarity(actual, expected) assert cosine > 0.9999 assert relative_l2 < 5e-3 + + +@torch.no_grad() +def test_model_helper_norm_flag_keeps_unfused_norm_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """KIMI_K3_FUSED_ATTN_RES_NORM=0 must keep attn_res_fwd + production RMSNorm. + + KIMI_K3_FUSED_ATTN_RES=0 is the wrong A/B knob: it drops all the way to + the fp32 reference and skips the pre-port path. + """ + ( + layer_residual, + block_residual, + res_weight, + score_rms_weight, + output_rms_weight, + ) = _make_inputs(num_tokens=1, num_snapshots=3) + projection = nn.Linear( + HIDDEN_SIZE, + 1, + bias=False, + dtype=torch.bfloat16, + device="cuda", + ) + score_norm = KimiK3RMSNorm( + HIDDEN_SIZE, + eps=ATTN_RES_RMS_EPS, + dtype=torch.bfloat16, + device=torch.device("cuda"), + ) + output_norm = RMSNorm( + hidden_size=HIDDEN_SIZE, + eps=OUTPUT_RMS_EPS, + dtype=torch.bfloat16, + device=torch.device("cuda"), + ) + projection.weight.copy_(res_weight.reshape(1, -1)) + score_norm.weight.copy_(score_rms_weight) + output_norm.weight.copy_(output_rms_weight) + + prefix_sum = layer_residual[:, 0, :] + block_kernel_layout = block_residual[:, :, 0, :] + unexpected_fused = mock.Mock( + side_effect=AssertionError("norm flag off still reached the fused norm op") + ) + monkeypatch.setattr(modeling_kimi_linear, "_FUSED_ATTN_RES_NORM_ENABLED", False) + monkeypatch.setattr( + modeling_kimi_linear, + "_apply_attn_res_rmsnorm_fused", + unexpected_fused, + ) + expected = output_norm( + _apply_attn_res( + prefix_sum, + block_kernel_layout, + projection, + score_norm, + ) + ) + actual = _apply_attn_res_and_rmsnorm( + prefix_sum, + block_kernel_layout, + projection, + score_norm, + output_norm, + ) + unexpected_fused.assert_not_called() + cosine, relative_l2 = _similarity(actual, expected) + assert cosine > 0.9999 + assert relative_l2 < 5e-3 From 3c6560d9ec24fa8c8d93b737ac18c80e1312f035 Mon Sep 17 00:00:00 2001 From: Xin Guan <294044352+xguannv@users.noreply.github.com> Date: Wed, 19 Aug 2026 03:45:36 +0000 Subject: [PATCH 3/5] [TRTLLM][Kimi K3] Fail loudly on invalid attn-res launch contracts Restore main's invokeAttnResFwd SM/N/H checks and validate the decode RMSNorm entry instead of silently returning or launching the wrong shape. Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com> --- .../kernels/kimiK3AttnRes/attnResFwd.cu | 27 ++++++++++++++----- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu index d13a42eac6ed..b6eee63b6100 100644 --- a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu @@ -33,6 +33,7 @@ #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h" +#include #include #include #include @@ -1703,12 +1704,10 @@ void invokeAttnResFwd(AttnResFwdParams const& params, cudaStream_t stream) float const rms_eps = params.rmsEps; int dev = 0; - cudaGetDevice(&dev); + TLLM_CUDA_CHECK(cudaGetDevice(&dev)); int num_sm = attn_res_fwd_grid_size(dev); - if (num_sm <= 0 || N > N_MAX) - { - return; - } + TLLM_CHECK_WITH_INFO(num_sm > 0, "attn_res_fwd: failed to query the SM count of device %d", dev); + TLLM_CHECK_WITH_INFO(N <= N_MAX, "attn_res_fwd: unsupported N=%d (max %d)", N, N_MAX); if (H == 8192) { @@ -1758,7 +1757,7 @@ void invokeAttnResFwd(AttnResFwdParams const& params, cudaStream_t stream) else if (N == 12 && T == 1024) { launch_fwd<7168, 4, false, true>(block_residual, layer_residual, res_weight, rms_weight, output, rsigma, - probs, logits, N, T, B, rms_eps, num_sm - 1, stream); + probs, logits, N, T, B, rms_eps, std::max(1, num_sm - 1), stream); } else { @@ -1794,6 +1793,10 @@ void invokeAttnResFwd(AttnResFwdParams const& params, cudaStream_t stream) logits, N, T, B, rms_eps, num_sm, stream); } } + else + { + TLLM_CHECK_WITH_INFO(false, "attn_res_fwd: unsupported hidden size H=%d", H); + } } template @@ -1821,6 +1824,14 @@ static void launchAttnResDecodeRmsNorm(AttnResFwdParams const& params, cudaStrea 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; @@ -1833,7 +1844,9 @@ static void invokeAttnResDecodeRmsNorm(AttnResFwdParams const& params, cudaStrea 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: break; + default: + TLLM_CHECK_WITH_INFO(false, "attn_res decode RMSNorm: unsupported numCandidates=%d (expected [1, 9] or 12)", + params.numCandidates); } } From 91849554b0e15d80c89981671e3cd479e9b4921c Mon Sep 17 00:00:00 2001 From: Xin Guan <294044352+xguannv@users.noreply.github.com> Date: Wed, 19 Aug 2026 13:49:53 +0000 Subject: [PATCH 4/5] [TRTLLM][Kimi K3] Fold decode-norm fusion tests into the existing attn-res file Keep N=4/N=8 parity, the prefill gate, the A/B flag, and one CUDA graph replay. Drop the per-N sweep so L0 only covers the new decode path. Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com> --- .../kimi_k3_attn_res/test_attn_res_op.py | 169 ++++- .../test_attn_res_rmsnorm_op.py | 607 ------------------ 2 files changed, 168 insertions(+), 608 deletions(-) delete mode 100644 tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py 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 diff --git a/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py b/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py deleted file mode 100644 index 6f845b808f33..000000000000 --- a/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py +++ /dev/null @@ -1,607 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Parity tests for the fused Kimi K3 attention-residual + RMSNorm op.""" - -from unittest import mock - -import pytest -import torch -from torch import nn - -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_rmsnorm_fused, -) -from tensorrt_llm._torch.modules.rms_norm import RMSNorm - -HIDDEN_SIZE = 7168 -ATTN_RES_RMS_EPS = 1e-6 -OUTPUT_RMS_EPS = 1e-6 - - -def _has_supported_gpu() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability(0) in { - (10, 0), - (10, 3), - } - - -pytestmark = pytest.mark.skipif( - not _has_supported_gpu(), - reason="Kimi K3 attention-residual kernels require SM100/SM103", -) - - -def _production_rms_norm( - hidden_states: torch.Tensor, - weight: torch.Tensor, - eps: float, -) -> torch.Tensor: - """Apply the unfused production RMSNorm that the new op replaces.""" - 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) - normalized = hidden_float * torch.rsqrt(variance + eps) - return weight * normalized.to(hidden_states.dtype) - - -def _make_inputs( - num_tokens: int, - num_snapshots: int, -) -> tuple[torch.Tensor, ...]: - torch.manual_seed(0) - layer_residual = ( - torch.randn( - num_tokens, - 1, - HIDDEN_SIZE, - dtype=torch.bfloat16, - device="cuda", - ) - * 0.05 - ) - block_residual = ( - torch.randn( - num_snapshots, - num_tokens, - 1, - HIDDEN_SIZE, - dtype=torch.bfloat16, - device="cuda", - ) - * 0.05 - ) - res_weight = torch.randn(HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.02 - score_rms_weight = 1 + torch.randn(HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.02 - output_rms_weight = 1 + torch.randn(HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.02 - return ( - layer_residual, - block_residual, - res_weight.contiguous(), - score_rms_weight.contiguous(), - output_rms_weight.contiguous(), - ) - - -def _unfused_reference( - layer_residual: torch.Tensor, - block_residual: torch.Tensor, - res_weight: torch.Tensor, - score_rms_weight: torch.Tensor, - output_rms_weight: torch.Tensor, -) -> torch.Tensor: - mixed, _rsigma, _probs, _logits = torch.ops.trtllm.attn_res_fwd( - layer_residual, - block_residual, - res_weight, - score_rms_weight, - ATTN_RES_RMS_EPS, - ) - return _production_rms_norm(mixed, output_rms_weight, OUTPUT_RMS_EPS) - - -def _fused( - layer_residual: torch.Tensor, - block_residual: torch.Tensor, - res_weight: torch.Tensor, - score_rms_weight: torch.Tensor, - output_rms_weight: torch.Tensor, -) -> torch.Tensor: - return torch.ops.trtllm.attn_res_rmsnorm_fwd( - layer_residual, - block_residual, - res_weight, - score_rms_weight, - output_rms_weight, - ATTN_RES_RMS_EPS, - OUTPUT_RMS_EPS, - ) - - -def _fused_add( - layer_residual: torch.Tensor, - layer_residual_add: torch.Tensor, - block_residual: torch.Tensor, - res_weight: torch.Tensor, - score_rms_weight: torch.Tensor, - output_rms_weight: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - return torch.ops.trtllm.attn_res_add_rmsnorm_fwd( - layer_residual, - layer_residual_add, - block_residual, - res_weight, - score_rms_weight, - output_rms_weight, - ATTN_RES_RMS_EPS, - OUTPUT_RMS_EPS, - ) - - -def _similarity( - actual: torch.Tensor, - expected: torch.Tensor, -) -> tuple[float, float]: - actual_float = actual.float() - expected_float = expected.float() - cosine = torch.nn.functional.cosine_similarity( - actual_float.flatten(), - expected_float.flatten(), - dim=0, - ).item() - relative_l2 = ((actual_float - expected_float).norm() / (expected_float.norm() + 1e-12)).item() - return cosine, relative_l2 - - -@pytest.mark.parametrize( - ("num_tokens", "num_snapshots"), - [ - (1, 0), # N=1 single-CTA decode - (1, 1), # N=2 single-CTA decode - (1, 2), # N=3 single-CTA decode - (1, 3), # N=4 single-CTA decode - (1, 4), # N=5 CTA-cluster decode - (1, 5), # N=6 CTA-cluster decode - (1, 6), # N=7 CTA-cluster decode - (1, 7), # N=8 CTA-cluster decode - (1, 8), # N=9 CTA-cluster decode - (1, 11), # N=12 CTA-cluster decode - ], -) -@torch.no_grad() -def test_attn_res_rmsnorm_matches_unfused( - num_tokens: int, - num_snapshots: int, -) -> None: - inputs = _make_inputs(num_tokens, num_snapshots) - expected = _unfused_reference(*inputs) - actual = _fused(*inputs) - - assert actual.shape == expected.shape - assert actual.dtype == torch.bfloat16 - cosine, relative_l2 = _similarity(actual, expected) - assert cosine > 0.9999 - assert relative_l2 < 5e-3 - - -@torch.no_grad() -def test_attn_res_rmsnorm_op_rejects_multi_token() -> None: - inputs = _make_inputs(num_tokens=2, num_snapshots=3) - with pytest.raises(RuntimeError, match="only production decode shape"): - _fused(*inputs) - - -@torch.no_grad() -def test_attn_res_rmsnorm_cuda_graph_replay() -> None: - inputs = _make_inputs(num_tokens=1, num_snapshots=3) - expected = _unfused_reference(*inputs) - - _fused(*inputs) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual = _fused(*inputs) - graph.replay() - torch.cuda.synchronize() - - cosine, relative_l2 = _similarity(actual, expected) - assert cosine > 0.9999 - assert relative_l2 < 5e-3 - - -@pytest.mark.parametrize("num_snapshots", list(range(9)) + [11]) -@torch.no_grad() -def test_attn_res_add_rmsnorm_matches_separate_add( - num_snapshots: int, -) -> None: - inputs = _make_inputs(num_tokens=1, num_snapshots=num_snapshots) - layer_residual_add = (torch.randn_like(inputs[0]) * 0.05).contiguous() - expected_updated = inputs[0] + layer_residual_add - expected_output = _unfused_reference( - expected_updated, - inputs[1], - inputs[2], - inputs[3], - inputs[4], - ) - actual_updated, actual_output = _fused_add( - inputs[0], - layer_residual_add, - inputs[1], - inputs[2], - inputs[3], - inputs[4], - ) - - assert torch.equal(actual_updated, expected_updated) - cosine, relative_l2 = _similarity(actual_output, expected_output) - assert cosine > 0.9999 - assert relative_l2 < 5e-3 - - -@pytest.mark.parametrize("num_snapshots", [3, 7, 8]) -@torch.no_grad() -def test_attn_res_add_rmsnorm_cuda_graph_replay( - num_snapshots: int, -) -> None: - inputs = _make_inputs(num_tokens=1, num_snapshots=num_snapshots) - layer_residual_add = (torch.randn_like(inputs[0]) * 0.05).contiguous() - expected_updated = inputs[0] + layer_residual_add - expected_output = _unfused_reference( - expected_updated, - inputs[1], - inputs[2], - inputs[3], - inputs[4], - ) - - _fused_add(inputs[0], layer_residual_add, *inputs[1:]) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual_updated, actual_output = _fused_add(inputs[0], layer_residual_add, *inputs[1:]) - for _ in range(100): - 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 - - -@torch.no_grad() -def test_model_helper_dispatches_fused_attn_res_add_rmsnorm( - monkeypatch: pytest.MonkeyPatch, -) -> None: - ( - layer_residual, - block_residual, - res_weight, - score_rms_weight, - output_rms_weight, - ) = _make_inputs(num_tokens=1, num_snapshots=3) - layer_residual_add = (torch.randn_like(layer_residual) * 0.05).contiguous() - projection = nn.Linear( - HIDDEN_SIZE, - 1, - bias=False, - dtype=torch.bfloat16, - device="cuda", - ) - score_norm = KimiK3RMSNorm( - HIDDEN_SIZE, - eps=ATTN_RES_RMS_EPS, - dtype=torch.bfloat16, - device=torch.device("cuda"), - ) - output_norm = RMSNorm( - hidden_size=HIDDEN_SIZE, - eps=OUTPUT_RMS_EPS, - dtype=torch.bfloat16, - device=torch.device("cuda"), - ) - projection.weight.copy_(res_weight.reshape(1, -1)) - score_norm.weight.copy_(score_rms_weight) - output_norm.weight.copy_(output_rms_weight) - - prefix_sum = layer_residual[:, 0, :] - addend = layer_residual_add[:, 0, :] - block_kernel_layout = block_residual[:, :, 0, :] - expected_updated = prefix_sum + addend - expected_output = output_norm( - _apply_attn_res( - expected_updated, - block_kernel_layout, - projection, - score_norm, - ) - ) - unexpected_fallback = mock.Mock( - side_effect=AssertionError("model helper did not dispatch fused add") - ) - monkeypatch.setattr( - modeling_kimi_linear, - "_apply_attn_res_and_rmsnorm", - unexpected_fallback, - ) - actual_updated, actual_output = _apply_attn_res_add_and_rmsnorm( - prefix_sum, - addend, - block_kernel_layout, - projection, - score_norm, - output_norm, - ) - - unexpected_fallback.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_model_helper_rejects_multi_token_fused_add() -> None: - ( - layer_residual, - block_residual, - res_weight, - score_rms_weight, - output_rms_weight, - ) = _make_inputs(num_tokens=64, num_snapshots=5) - layer_residual_add = (torch.randn_like(layer_residual) * 0.05).contiguous() - projection = nn.Linear( - HIDDEN_SIZE, - 1, - bias=False, - dtype=torch.bfloat16, - device="cuda", - ) - score_norm = KimiK3RMSNorm( - HIDDEN_SIZE, - eps=ATTN_RES_RMS_EPS, - dtype=torch.bfloat16, - device=torch.device("cuda"), - ) - output_norm = RMSNorm( - hidden_size=HIDDEN_SIZE, - eps=OUTPUT_RMS_EPS, - dtype=torch.bfloat16, - device=torch.device("cuda"), - ) - projection.weight.copy_(res_weight.reshape(1, -1)) - score_norm.weight.copy_(score_rms_weight) - output_norm.weight.copy_(output_rms_weight) - - assert ( - _apply_attn_res_add_rmsnorm_fused( - layer_residual[:, 0, :], - layer_residual_add[:, 0, :], - block_residual[:, :, 0, :], - projection, - score_norm, - output_norm, - ) - is None - ) - - -@pytest.mark.parametrize( - ("num_tokens", "num_snapshots"), - [ - (1, 3), - ], -) -@torch.no_grad() -def test_model_helper_dispatches_fused_attn_res_rmsnorm( - num_tokens: int, - num_snapshots: int, - monkeypatch: pytest.MonkeyPatch, -) -> None: - ( - layer_residual, - block_residual, - res_weight, - score_rms_weight, - output_rms_weight, - ) = _make_inputs(num_tokens, num_snapshots) - projection = nn.Linear( - HIDDEN_SIZE, - 1, - bias=False, - dtype=torch.bfloat16, - device="cuda", - ) - score_norm = KimiK3RMSNorm( - HIDDEN_SIZE, - eps=ATTN_RES_RMS_EPS, - dtype=torch.bfloat16, - device=torch.device("cuda"), - ) - output_norm = RMSNorm( - hidden_size=HIDDEN_SIZE, - eps=OUTPUT_RMS_EPS, - dtype=torch.bfloat16, - device=torch.device("cuda"), - ) - projection.weight.copy_(res_weight.reshape(1, -1)) - score_norm.weight.copy_(score_rms_weight) - output_norm.weight.copy_(output_rms_weight) - - prefix_sum = layer_residual[:, 0, :] - block_kernel_layout = block_residual[:, :, 0, :] - expected = output_norm( - _apply_attn_res( - prefix_sum, - block_kernel_layout, - projection, - score_norm, - ) - ) - unexpected_fallback = mock.Mock( - side_effect=AssertionError("model helper did not dispatch the fused op") - ) - monkeypatch.setattr( - modeling_kimi_linear, - "_apply_attn_res", - unexpected_fallback, - ) - actual = _apply_attn_res_and_rmsnorm( - prefix_sum, - block_kernel_layout, - projection, - score_norm, - output_norm, - ) - - cosine, relative_l2 = _similarity(actual, expected) - unexpected_fallback.assert_not_called() - assert cosine > 0.9999 - assert relative_l2 < 5e-3 - - -@torch.no_grad() -def test_model_helper_keeps_multi_token_rmsnorm_split() -> None: - ( - layer_residual, - block_residual, - res_weight, - score_rms_weight, - output_rms_weight, - ) = _make_inputs(num_tokens=64, num_snapshots=5) - projection = nn.Linear( - HIDDEN_SIZE, - 1, - bias=False, - dtype=torch.bfloat16, - device="cuda", - ) - score_norm = KimiK3RMSNorm( - HIDDEN_SIZE, - eps=ATTN_RES_RMS_EPS, - dtype=torch.bfloat16, - device=torch.device("cuda"), - ) - output_norm = RMSNorm( - hidden_size=HIDDEN_SIZE, - eps=OUTPUT_RMS_EPS, - dtype=torch.bfloat16, - device=torch.device("cuda"), - ) - projection.weight.copy_(res_weight.reshape(1, -1)) - score_norm.weight.copy_(score_rms_weight) - output_norm.weight.copy_(output_rms_weight) - - prefix_sum = layer_residual[:, 0, :] - block_kernel_layout = block_residual[:, :, 0, :] - assert ( - _apply_attn_res_rmsnorm_fused( - prefix_sum, - block_kernel_layout, - projection, - score_norm, - output_norm, - ) - is None - ) - - expected = output_norm( - _apply_attn_res( - prefix_sum, - block_kernel_layout, - projection, - score_norm, - ) - ) - actual = _apply_attn_res_and_rmsnorm( - prefix_sum, - block_kernel_layout, - projection, - score_norm, - output_norm, - ) - cosine, relative_l2 = _similarity(actual, expected) - assert cosine > 0.9999 - assert relative_l2 < 5e-3 - - -@torch.no_grad() -def test_model_helper_norm_flag_keeps_unfused_norm_path( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """KIMI_K3_FUSED_ATTN_RES_NORM=0 must keep attn_res_fwd + production RMSNorm. - - KIMI_K3_FUSED_ATTN_RES=0 is the wrong A/B knob: it drops all the way to - the fp32 reference and skips the pre-port path. - """ - ( - layer_residual, - block_residual, - res_weight, - score_rms_weight, - output_rms_weight, - ) = _make_inputs(num_tokens=1, num_snapshots=3) - projection = nn.Linear( - HIDDEN_SIZE, - 1, - bias=False, - dtype=torch.bfloat16, - device="cuda", - ) - score_norm = KimiK3RMSNorm( - HIDDEN_SIZE, - eps=ATTN_RES_RMS_EPS, - dtype=torch.bfloat16, - device=torch.device("cuda"), - ) - output_norm = RMSNorm( - hidden_size=HIDDEN_SIZE, - eps=OUTPUT_RMS_EPS, - dtype=torch.bfloat16, - device=torch.device("cuda"), - ) - projection.weight.copy_(res_weight.reshape(1, -1)) - score_norm.weight.copy_(score_rms_weight) - output_norm.weight.copy_(output_rms_weight) - - prefix_sum = layer_residual[:, 0, :] - block_kernel_layout = block_residual[:, :, 0, :] - unexpected_fused = mock.Mock( - side_effect=AssertionError("norm flag off still reached the fused norm op") - ) - monkeypatch.setattr(modeling_kimi_linear, "_FUSED_ATTN_RES_NORM_ENABLED", False) - monkeypatch.setattr( - modeling_kimi_linear, - "_apply_attn_res_rmsnorm_fused", - unexpected_fused, - ) - expected = output_norm( - _apply_attn_res( - prefix_sum, - block_kernel_layout, - projection, - score_norm, - ) - ) - actual = _apply_attn_res_and_rmsnorm( - prefix_sum, - block_kernel_layout, - projection, - score_norm, - output_norm, - ) - unexpected_fused.assert_not_called() - cosine, relative_l2 = _similarity(actual, expected) - assert cosine > 0.9999 - assert relative_l2 < 5e-3 From 3cea2632ba997cb165ebeec251250b2fbc9f5dbc Mon Sep 17 00:00:00 2001 From: Xin Guan <294044352+xguannv@users.noreply.github.com> Date: Thu, 20 Aug 2026 03:09:34 +0000 Subject: [PATCH 5/5] [TRTLLM][Kimi K3] Drop the one-off attn-res add+rmsnorm microbenchmark Correctness is covered by the existing unittest and e2e. Kernel timing belongs in the layer-wise harness from #17804 rather than a new script. Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com> --- .../kimi_k3_attn_res_add_rmsnorm.py | 338 ------------------ 1 file changed, 338 deletions(-) delete mode 100644 tests/microbenchmarks/kimi_k3_attn_res_add_rmsnorm.py diff --git a/tests/microbenchmarks/kimi_k3_attn_res_add_rmsnorm.py b/tests/microbenchmarks/kimi_k3_attn_res_add_rmsnorm.py deleted file mode 100644 index ef0e6bb1b3d5..000000000000 --- a/tests/microbenchmarks/kimi_k3_attn_res_add_rmsnorm.py +++ /dev/null @@ -1,338 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Microbenchmark the incremental Kimi K3 attention-output add fusion. - -The production baseline after the trailing-RMSNorm fusion is two kernels: - - updated_prefix = prefix_sum + attention_output - trtllm::attn_res_rmsnorm_fwd(updated_prefix, ...) - -The new path is one kernel and still materializes ``updated_prefix`` for the -MLP residual that follows: - - updated_prefix, output = - trtllm::attn_res_add_rmsnorm_fwd(prefix_sum, attention_output, ...) - -CUDA-graph replay timings remove Python and dispatcher overhead. ``--profile`` -emits eager calls inside NVTX ranges for Nsys attribution. -""" - -from __future__ import annotations - -import argparse -import json -import statistics -from collections.abc import Callable, Sequence -from dataclasses import dataclass -from typing import Any - -import torch - -from tensorrt_llm._torch.custom_ops import flashinfer_rmsnorm - -HIDDEN_SIZE = 7168 -RMS_EPS = 1e-6 - - -@dataclass -class CaseInputs: - prefix_sum: torch.Tensor - attention_output: torch.Tensor - block_residual: torch.Tensor - res_weight: torch.Tensor - score_rms_weight: torch.Tensor - output_rms_weight: torch.Tensor - - -def _parse_candidates(value: str) -> list[int]: - candidates = [int(item) for item in value.split(",")] - if any(candidate not in {*range(1, 10), 12} for candidate in candidates): - raise argparse.ArgumentTypeError("candidate counts must be in [1, 9] or equal to 12") - return candidates - - -def _make_inputs(num_candidates: int) -> CaseInputs: - device = torch.device("cuda") - shape = (1, 1, HIDDEN_SIZE) - prefix_sum = torch.empty(shape, dtype=torch.bfloat16, device=device).uniform_(-0.05, 0.05) - attention_output = torch.empty(shape, dtype=torch.bfloat16, device=device).uniform_(-0.05, 0.05) - block_residual = torch.empty( - (num_candidates - 1, *shape), - dtype=torch.bfloat16, - device=device, - ).uniform_(-0.05, 0.05) - res_weight = torch.empty(HIDDEN_SIZE, dtype=torch.bfloat16, device=device).uniform_(-0.02, 0.02) - score_rms_weight = torch.empty(HIDDEN_SIZE, dtype=torch.bfloat16, device=device).uniform_( - 0.98, 1.02 - ) - output_rms_weight = torch.empty(HIDDEN_SIZE, dtype=torch.bfloat16, device=device).uniform_( - 0.98, 1.02 - ) - return CaseInputs( - prefix_sum=prefix_sum, - attention_output=attention_output, - block_residual=block_residual, - res_weight=res_weight, - score_rms_weight=score_rms_weight, - output_rms_weight=output_rms_weight, - ) - - -def _attn_res(inputs: CaseInputs, updated_prefix: torch.Tensor) -> torch.Tensor: - output, _rsigma, _probs, _logits = torch.ops.trtllm.attn_res_fwd( - updated_prefix, - inputs.block_residual, - inputs.res_weight, - inputs.score_rms_weight, - RMS_EPS, - ) - return output - - -def _three_kernel(inputs: CaseInputs) -> tuple[torch.Tensor, torch.Tensor]: - updated_prefix = inputs.prefix_sum + inputs.attention_output - mixed = _attn_res(inputs, updated_prefix) - output = flashinfer_rmsnorm(mixed, inputs.output_rms_weight, RMS_EPS) - return updated_prefix, output - - -def _two_kernel(inputs: CaseInputs) -> tuple[torch.Tensor, torch.Tensor]: - updated_prefix = inputs.prefix_sum + inputs.attention_output - output = torch.ops.trtllm.attn_res_rmsnorm_fwd( - updated_prefix, - inputs.block_residual, - inputs.res_weight, - inputs.score_rms_weight, - inputs.output_rms_weight, - RMS_EPS, - RMS_EPS, - ) - return updated_prefix, output - - -def _fused(inputs: CaseInputs) -> tuple[torch.Tensor, torch.Tensor]: - return torch.ops.trtllm.attn_res_add_rmsnorm_fwd( - inputs.prefix_sum, - inputs.attention_output, - inputs.block_residual, - inputs.res_weight, - inputs.score_rms_weight, - inputs.output_rms_weight, - RMS_EPS, - RMS_EPS, - ) - - -def _capture( - fn: Callable[[], Any], -) -> tuple[torch.cuda.CUDAGraph, Any]: - fn() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - output = fn() - graph.replay() - torch.cuda.synchronize() - return graph, output - - -def _time_graph( - fn: Callable[[], Any], - iterations: int, - samples: int, - chain_length: int, -) -> tuple[float, float, float]: - def chained_fn() -> Any: - output = None - for _ in range(chain_length): - output = fn() - return output - - graph, output = _capture(chained_fn) - del output - for _ in range(20): - graph.replay() - torch.cuda.synchronize() - - timings = [] - for _ in range(samples): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - for _ in range(iterations): - graph.replay() - end.record() - end.synchronize() - timings.append(start.elapsed_time(end) * 1000.0 / iterations / chain_length) - return statistics.median(timings), min(timings), max(timings) - - -def _similarity( - actual: torch.Tensor, - expected: torch.Tensor, -) -> tuple[float, float]: - actual_float = actual.float().flatten() - expected_float = expected.float().flatten() - cosine = torch.nn.functional.cosine_similarity(actual_float, expected_float, dim=0).item() - relative_l2 = ((actual_float - expected_float).norm() / (expected_float.norm() + 1e-12)).item() - return cosine, relative_l2 - - -def _benchmark_case( - num_candidates: int, - iterations: int, - samples: int, - chain_length: int, -) -> dict[str, float | int]: - inputs = _make_inputs(num_candidates) - expected_prefix, expected_output = _two_kernel(inputs) - actual_prefix, actual_output = _fused(inputs) - _, three_kernel_output = _three_kernel(inputs) - torch.cuda.synchronize() - if not torch.equal(actual_prefix, expected_prefix): - raise AssertionError(f"N={num_candidates}: fused updated prefix is not exact") - cosine, relative_l2 = _similarity(actual_output, expected_output) - if cosine <= 0.9999 or relative_l2 >= 5e-3: - raise AssertionError(f"N={num_candidates}: cosine={cosine}, relative_l2={relative_l2}") - # attn_res_fwd + the production RMSNorm is what shipped before the trailing - # norm was folded into the kernel, so this -- not two_kernel, which is also - # a norm-fusing op -- is the baseline the fusion has to be judged against. - three_cosine, three_relative_l2 = _similarity(actual_output, three_kernel_output) - - add_us, add_min_us, add_max_us = _time_graph( - lambda: inputs.prefix_sum + inputs.attention_output, - iterations, - samples, - chain_length, - ) - three_us, three_min_us, three_max_us = _time_graph( - lambda: _three_kernel(inputs), iterations, samples, chain_length - ) - two_us, two_min_us, two_max_us = _time_graph( - lambda: _two_kernel(inputs), iterations, samples, chain_length - ) - fused_us, fused_min_us, fused_max_us = _time_graph( - lambda: _fused(inputs), iterations, samples, chain_length - ) - - return { - "num_tokens": 1, - "num_candidates": num_candidates, - "iterations": iterations, - "samples": samples, - "chain_length": chain_length, - "cosine": cosine, - "relative_l2": relative_l2, - "cosine_vs_three_kernel": three_cosine, - "relative_l2_vs_three_kernel": three_relative_l2, - "add_us": add_us, - "add_min_us": add_min_us, - "add_max_us": add_max_us, - "three_kernel_us": three_us, - "three_kernel_min_us": three_min_us, - "three_kernel_max_us": three_max_us, - "two_kernel_us": two_us, - "two_kernel_min_us": two_min_us, - "two_kernel_max_us": two_max_us, - "fused_us": fused_us, - "fused_min_us": fused_min_us, - "fused_max_us": fused_max_us, - "fused_vs_two_kernel_pct": (fused_us / two_us - 1.0) * 100.0, - "saved_vs_two_kernel_us": two_us - fused_us, - "fused_vs_three_kernel_pct": (fused_us / three_us - 1.0) * 100.0, - "saved_vs_three_kernel_us": three_us - fused_us, - } - - -def _profile_case( - num_candidates: int, - iterations: int, -) -> None: - inputs = _make_inputs(num_candidates) - modes: Sequence[tuple[str, Callable[[], Any]]] = ( - ("add", lambda: inputs.prefix_sum + inputs.attention_output), - ("three_kernel", lambda: _three_kernel(inputs)), - ("two_kernel", lambda: _two_kernel(inputs)), - ("fused", lambda: _fused(inputs)), - ) - for _name, fn in modes: - for _ in range(10): - fn() - torch.cuda.synchronize() - - for name, fn in modes: - range_name = f"attn_res_add|T=1|N={num_candidates}|mode={name}" - torch.cuda.nvtx.range_push(range_name) - for _ in range(iterations): - fn() - torch.cuda.synchronize() - torch.cuda.nvtx.range_pop() - print( - json.dumps( - {"profile_range": range_name, "iterations": iterations}, - sort_keys=True, - ), - flush=True, - ) - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument( - "--candidates", - type=_parse_candidates, - default=_parse_candidates("1,2,3,4,5,6,7,8,9,12"), - ) - parser.add_argument("--iterations", type=int, default=2000) - parser.add_argument("--samples", type=int, default=7) - parser.add_argument( - "--chain-length", - type=int, - default=1, - help="Capture this many copies of each mode in one CUDA graph.", - ) - parser.add_argument( - "--profile", - action="store_true", - help="Emit eager kernels in NVTX ranges for Nsys instead of timing.", - ) - args = parser.parse_args() - - if args.chain_length < 1: - parser.error("--chain-length must be positive") - - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required") - capability = torch.cuda.get_device_capability() - if capability not in {(10, 0), (10, 3)}: - raise RuntimeError(f"SM100/SM103 is required, got {capability}") - torch.manual_seed(0) - print( - json.dumps( - { - "device": torch.cuda.get_device_name(), - "capability": capability, - "profile": args.profile, - }, - sort_keys=True, - ), - flush=True, - ) - - for num_candidates in args.candidates: - if args.profile: - _profile_case(num_candidates, min(args.iterations, 100)) - else: - print( - json.dumps( - _benchmark_case( - num_candidates, args.iterations, args.samples, args.chain_length - ), - sort_keys=True, - ), - flush=True, - ) - - -if __name__ == "__main__": - main()