From 48d059e10f5748525e46e8bdb3ee1391db0fcc08 Mon Sep 17 00:00:00 2001 From: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> Date: Tue, 27 Oct 2026 07:01:20 -0700 Subject: [PATCH] [None][fix] Kimi K3 attn_res persistent fused kernel: 64-bit candidate offsets The persistent attention-residual kernel formed source * block_stride_r in int. block_stride_r is T * 7168 and source reaches N - 2, so with N == 9 any T above 42,799 tokens overflowed; the launcher refused those shapes, which made max_num_tokens 49152 fail in warmup. The kernel now takes the candidate stride as long long, so that product is formed in 64 bits, and T = 49152 and 65536 work at N = 7, 8 and 9. In the PTX, only the producer warp's address arithmetic changes; registers, shared memory and the launch configuration stay the same. launch_fwd checks the remaining bound, the int token index (T - 1 + grid), and the op checks T before narrowing it into the params. Tests cover N = 2..9, both sides of the former int limit and the production token counts, against the unfused reference; without the residual add they also check that the op returns layer_residual itself. Signed-off-by: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> --- .../kernels/kimiK3AttnRes/attnResFwd.h | 4 +- .../attnResFwdPersistentFused.cu | 34 +++-- .../kimiK3AttnRes/attnResFwdPersistentFused.h | 27 ++-- cpp/tensorrt_llm/thop/attnResOp.cpp | 4 + .../test_attn_res_rmsnorm_op.py | 121 ++++++++++++++++++ 5 files changed, 161 insertions(+), 29 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h index feadfe63c8e5..38edfae19481 100644 --- a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h @@ -33,7 +33,9 @@ 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] when requested. -//! blockResidual may be nullptr when N == 1. +//! blockResidual may be nullptr when N == 1. That is the contract of the +//! per-token entry points declared here; invokeAttnResPersistentFusedFwd +//! (attnResFwdPersistentFused.h) takes the same struct under its own contract. struct AttnResFwdParams { __nv_bfloat16 const* blockResidual; // [N-1, T, B, H], nullptr iff N == 1 diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu index ee734add0661..cf3e615a8ad7 100644 --- a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu @@ -97,8 +97,10 @@ constexpr int CONSUMER_GROUPS = 2; // two 128-thread consumer groups constexpr int CONSUMER_THREADS_PER_GROUP = CONSUMER_THREADS / CONSUMER_GROUPS; constexpr int FIRST_USER_NAMED_BARRIER = 8; +// Both terms are formed in 64 bits. The candidate term, source * block_stride_r, reaches (N - 2) * T * H elements, +// which passes INT_MAX from T = 42800 at N = 9. __device__ __forceinline__ bf16_t const* residual_addr(bf16_t const* block_res, bf16_t const* layer_res, int source, - int N, int token, int block_stride_m, int block_stride_r, int H) + int N, int token, int block_stride_m, long long block_stride_r, int H) { if (source < N - 1) { @@ -319,8 +321,8 @@ __global__ void __launch_bounds__(BLK, 1) // consumed value comes from the TMA'd shared-memory copy rather than a // reload of layer_res removes one hazard, not the no-alias contract. bf16_t* __restrict__ updated_layer_res, bf16_t const* __restrict__ delta, bf16_t const* __restrict__ res_w, - bf16_t const* __restrict__ rms_w, bf16_t* __restrict__ output, int T, int block_stride_m, int block_stride_r, - float rms_eps, bf16_t const* __restrict__ output_norm_weight, float output_norm_eps) + bf16_t const* __restrict__ rms_w, bf16_t* __restrict__ output, int T, int block_stride_m, + long long block_stride_r, float rms_eps, bf16_t const* __restrict__ output_norm_weight, float output_norm_eps) { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100 constexpr float LOG2_E = 1.4426950408889634f; @@ -1042,7 +1044,7 @@ template (T) - 1 + grid <= std::numeric_limits::max(), + "attn_res persistent fused: T=%d overflows the kernel's int token index with a grid of %d CTAs", T, grid); cudaLaunchConfig_t config{}; config.gridDim = grid; config.blockDim = BLK; @@ -1139,22 +1145,14 @@ void invokeAttnResPersistentFusedFwd(AttnResFwdParams const& params, cudaStream_ // candidate step is T * H. The source copy carried a token-major bank and // therefore takes both strides as arguments; passing ours keeps the kernel // body unchanged. + // + // residual_addr scales the candidate step by sources up to N - 2, so the + // largest candidate offset is (N - 2) * T * H elements; the kernel takes + // the step as long long and forms that product in 64 bits. launch_fwd + // checks the remaining bound on T, the int token index. int const T = params.seqLen; int const block_stride_m = 7168; long long const block_stride_r = static_cast(T) * 7168; - TLLM_CHECK_WITH_INFO(block_stride_r <= std::numeric_limits::max(), - "attn_res persistent fused: candidate stride %lld overflows the kernel's int stride", block_stride_r); - // One candidate step is not the largest address the kernel forms. - // ``residual_addr`` widens the token term but evaluates - // ``source * block_stride_r`` in int arithmetic, and ``source`` reaches - // ``numCandidates - 2`` (the ``source < N - 1`` branch). With N == 9 that - // product overflows from T == 42801 upward and the load leaves - // blockResidual. Persistent topology accepts any token count, so nothing - // upstream bounds T. - long long const max_candidate_offset = block_stride_r * (params.numCandidates - 2); - TLLM_CHECK_WITH_INFO(max_candidate_offset <= std::numeric_limits::max(), - "attn_res persistent fused: candidate offset %lld (T=%d, N=%d) overflows the kernel's int stride", - max_candidate_offset, T, params.numCandidates); bool const has_delta = params.layerResidualAdd != nullptr; auto launch = [&](auto nsrc_tok, auto nc_tok, auto release_tmem_tok) @@ -1168,7 +1166,7 @@ void invokeAttnResPersistentFusedFwd(AttnResFwdParams const& params, cudaStream_ launch_fwd<7168, NSRC, NC, RELEASE_TMEM, HAS_DELTA, /*HAS_OUTPUT_NORM=*/true, /*OUTPUT_NORM_IN_SMEM=*/true>(params.blockResidual, params.layerResidual, params.updatedLayerResidual, params.layerResidualAdd, params.resWeight, params.rmsWeight, params.output, T, params.rmsEps, num_sm, - stream, params.outputRmsWeight, params.outputRmsEps, block_stride_m, static_cast(block_stride_r)); + stream, params.outputRmsWeight, params.outputRmsEps, block_stride_m, block_stride_r); }; if (has_delta) { diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.h b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.h index 3ef52863143d..65fbb01b34a3 100644 --- a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.h +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.h @@ -25,15 +25,19 @@ namespace kernels::kimiK3AttnRes //! Persistent-grid counterpart to invokeAttnResAddRmsNormFwd. //! -//! Same contract and same AttnResFwdParams, different topology: one CTA per SM -//! looping over tokens, instead of one CTA (or one cluster) per token. That -//! makes it the shape to reach for when there are far more tokens than SMs -- -//! prefill -- where the per-token kernels stop being able to amortize the score -//! projection load. +//! Same AttnResFwdParams, different topology: one CTA per SM looping over +//! tokens, instead of one CTA (or one cluster) per token. That makes it the +//! shape to reach for when there are far more tokens than SMs -- prefill -- +//! where the per-token kernels stop being able to amortize the score +//! projection load. The shape contract is below. //! -//! Requires H == 7168 and N in [2, 9]. Unlike the per-token kernels there is no -//! N == 1 case: the persistent kernel always has at least one snapshot plus the -//! layer residual. +//! Requires B == 1, H == 7168, N in [2, 9] and bf16 contiguous tensors. Unlike +//! the per-token kernels there is no N == 1 case: the persistent kernel always +//! has at least one snapshot plus the layer residual, and it produces no +//! rsigma/probs/logits. T is bounded only by the int token index: offsets that +//! grow with T are formed in 64 bits wherever they can pass INT_MAX, and the +//! token index needs T - 1 + grid <= INT_MAX, where grid is the persistent grid +//! size; the launch checks that bound. //! //! outputRmsWeight is required -- this entry point exists for the fully fused //! form. layerResidualAdd is optional: shapes with no residual add to fold @@ -47,8 +51,11 @@ namespace kernels::kimiK3AttnRes //! aliased form so a caller that stops doing so fails loudly. void invokeAttnResPersistentFusedFwd(AttnResFwdParams const& params, cudaStream_t stream); -//! True when invokeAttnResPersistentFusedFwd can serve this shape, so callers -//! can route without duplicating the constraints above. +//! True when B, H, N, T >= 1, the required tensors and the aliasing rule above +//! all hold, so callers can route without duplicating those checks. +//! invokeAttnResPersistentFusedFwd also rejects non-null rsigma/probs/logits +//! and, at launch, a T past the token index bound, which depends on the grid +//! size. bool attnResPersistentFusedSupported(AttnResFwdParams const& params); } // namespace kernels::kimiK3AttnRes diff --git a/cpp/tensorrt_llm/thop/attnResOp.cpp b/cpp/tensorrt_llm/thop/attnResOp.cpp index a1add45502db..68df195c8042 100644 --- a/cpp/tensorrt_llm/thop/attnResOp.cpp +++ b/cpp/tensorrt_llm/thop/attnResOp.cpp @@ -24,6 +24,7 @@ #include #include +#include #include #include @@ -296,6 +297,9 @@ std::tuple attn_res_add_rmsnorm_persistent_fwd(at::Tenso // Unlike the per-token kernels there is no N == 1 form: the persistent // kernel always mixes at least one snapshot with the layer residual. TORCH_CHECK(N >= 2 && N <= 9, "attn_res_add_rmsnorm_persistent_fwd: supported N values are [2, 9], got ", N); + // AttnResFwdParams carries T as int; the launcher bounds it further by its grid size. + TORCH_CHECK(T <= std::numeric_limits::max(), "attn_res_add_rmsnorm_persistent_fwd: T=", T, + " does not fit the kernel's int token count"); TORCH_CHECK(layer_residual.scalar_type() == at::kBFloat16 && (!has_add || layer_residual_add->scalar_type() == at::kBFloat16) && block_residual.scalar_type() == at::kBFloat16 && res_weight.scalar_type() == at::kBFloat16 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 33afef2de92e..f66d4282177c 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 @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Parity tests for the fused Kimi K3 attention-residual + RMSNorm op.""" +from typing import Optional from unittest import mock import pytest @@ -23,6 +24,8 @@ HIDDEN_SIZE = 7168 ATTN_RES_RMS_EPS = 1e-6 OUTPUT_RMS_EPS = 1e-6 +# Largest token count trtllm::attn_res_fwd accepts (check_attn_res_contract in attnResOp.cpp). +ATTN_RES_FWD_MAX_TOKENS = 16384 def _has_supported_gpu() -> bool: @@ -149,6 +152,26 @@ def _fused_add( ) +def _persistent( + layer_residual: torch.Tensor, + layer_residual_add: Optional[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_persistent_fwd( + layer_residual, + layer_residual_add, + block_residual, + res_weight, + score_rms_weight, + output_rms_weight, + ATTN_RES_RMS_EPS, + OUTPUT_RMS_EPS, + ) + + def _above_ceiling_tokens() -> int: """One more token than the fused epilogue will take. @@ -326,6 +349,104 @@ def test_attn_res_add_rmsnorm_cuda_graph_replay( assert relative_l2 < 5e-3 +def _assert_persistent_matches_unfused( + num_tokens: int, + num_candidates: int, + with_add: bool, +) -> None: + """Compare the persistent op with the unfused reference slice by slice. + + The reference takes at most ATTN_RES_FWD_MAX_TOKENS tokens per call, and + every output row depends only on its own token, so each token slice is + compared on its own; an error confined to late tokens is not averaged away. + """ + ( + layer_residual, + block_residual, + res_weight, + score_rms_weight, + output_rms_weight, + ) = _make_inputs(num_tokens, num_candidates - 1) + layer_residual_add = ( + (torch.randn_like(layer_residual) * 0.05).contiguous() if with_add else None + ) + + actual_updated, actual_output = _persistent( + layer_residual, + layer_residual_add, + block_residual, + res_weight, + score_rms_weight, + output_rms_weight, + ) + + expected_updated = layer_residual + if with_add: + expected_updated = layer_residual + layer_residual_add + assert torch.equal(actual_updated, expected_updated) + else: + assert actual_updated.data_ptr() == layer_residual.data_ptr() + for start in range(0, num_tokens, ATTN_RES_FWD_MAX_TOKENS): + tokens = slice(start, start + ATTN_RES_FWD_MAX_TOKENS) + expected_output = _unfused_reference( + expected_updated[tokens], + block_residual[:, tokens].contiguous(), + res_weight, + score_rms_weight, + output_rms_weight, + ) + cosine, relative_l2 = _similarity(actual_output[tokens], expected_output) + assert cosine > 0.9999, f"tokens from {start}: cosine {cosine}" + assert relative_l2 < 5e-3, f"tokens from {start}: relative L2 {relative_l2}" + + +@pytest.mark.parametrize("with_add", [False, True], ids=["no_add", "add"]) +@pytest.mark.parametrize("num_candidates", range(2, 10)) +@pytest.mark.parametrize("num_tokens", [1, 213]) +@torch.no_grad() +def test_attn_res_persistent_matches_unfused( + num_tokens: int, + num_candidates: int, + with_add: bool, +) -> None: + """Every supported N at T=1 and T=213. + + T=1 leaves most CTAs without a token, and 213 is not a multiple of the CTA + count, so CTAs take different numbers of tokens. + """ + _assert_persistent_matches_unfused(num_tokens, num_candidates, with_add) + + +@pytest.mark.parametrize("with_add", [False, True], ids=["no_add", "add"]) +@pytest.mark.parametrize( + ("num_tokens", "num_candidates"), + [ + (16384, 9), # every candidate offset inside int32 + (42799, 9), # the largest T whose deepest candidate offset fits int32 + (42800, 9), # the smallest T whose deepest candidate offset does not + # The production max_num_tokens candidates; at 65536 the deepest + # candidate offset passes INT32_MAX for N=7 and N=8 as well. + (49152, 9), + (65536, 7), + (65536, 8), + (65536, 9), + ], +) +@torch.no_grad() +def test_attn_res_persistent_matches_unfused_at_prefill_token_counts( + num_tokens: int, + num_candidates: int, + with_add: bool, +) -> None: + """Prefill token counts on both sides of the int32 candidate offset. + + The deepest snapshot, source N-2, starts (N-2) * T * H elements into + block_residual. That offset passes INT32_MAX at N=9 from T=42800, so the + kernel forms it in 64 bits. + """ + _assert_persistent_matches_unfused(num_tokens, num_candidates, with_add) + + @torch.no_grad() def test_model_helper_dispatches_fused_attn_res_add_rmsnorm( monkeypatch: pytest.MonkeyPatch,