[None][fix] Kimi K3 attn_res persistent fused kernel: 64-bit candidate offsets

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

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 <int H, int N, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false, boo
 static void launch_fwd(bf16_t const* block_residual, bf16_t const* layer_residual, bf16_t* updated_layer_residual,
     bf16_t const* delta, bf16_t const* res_weight, bf16_t const* rms_weight, bf16_t* output, int T, float rms_eps,
     int num_sm, cudaStream_t stream, bf16_t const* output_norm_weight = nullptr, float output_norm_eps = 0.f,
-    int block_stride_m = 0, int block_stride_r = 0)
+    int block_stride_m = 0, long long block_stride_r = 0)
 {
     constexpr size_t smem_size
         = ((size_t) CHUNK_DEPTH * (NC + (HAS_DELTA ? 1 : 0)) * H * sizeof(bf16_t)
@@ -1071,6 +1073,10 @@ static void launch_fwd(bf16_t const* block_residual, bf16_t const* layer_residua
         }
     }
     int grid = RELEASE_TMEM ? num_sm * 2 : num_sm;
+    // The token index stays int: each CTA steps it by the grid size until it is no longer below T, so its largest
+    // value is T - 1 + grid.
+    TLLM_CHECK_WITH_INFO(static_cast<long long>(T) - 1 + grid <= std::numeric_limits<int>::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<long long>(T) * 7168;
-    TLLM_CHECK_WITH_INFO(block_stride_r <= std::numeric_limits<int>::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<int>::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<int>(block_stride_r));
+                stream, params.outputRmsWeight, params.outputRmsEps, block_stride_m, block_stride_r);
         };
         if (has_delta)
         {

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

cpp/tensorrt_llm/thop/attnResOp.cpp

@@ -24,6 +24,7 @@
 #include <cuda_runtime.h>
 #include <torch/library.h>
 
+#include <limits>
 #include <optional>
 #include <tuple>
 
@@ -296,6 +297,9 @@ std::tuple<at::Tensor, at::Tensor> 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<int>::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

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,