From 15ed27d7d415e9f205d98524000c06790da556f1 Mon Sep 17 00:00:00 2001
From: bingps <46775742+bingps@users.noreply.github.com>
Date: Mon, 3 Nov 2025 03:52:06 +0800
Subject: [PATCH] [Fix] `concat_mla_absorb_q_kernel` fails for long inputs
 (#12453)

---
 sgl-kernel/csrc/elementwise/concat_mla.cu | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/sgl-kernel/csrc/elementwise/concat_mla.cu b/sgl-kernel/csrc/elementwise/concat_mla.cu
index 0335dc724..7d5b8595c 100644
--- a/sgl-kernel/csrc/elementwise/concat_mla.cu
+++ b/sgl-kernel/csrc/elementwise/concat_mla.cu
@@ -18,11 +18,11 @@ __global__ void concat_mla_k_kernel(
     const nv_bfloat16* __restrict__ k_nope,
     const nv_bfloat16* __restrict__ k_rope,
     const int num_tokens,
-    const int k_stride_0,
+    const int64_t k_stride_0,
     const int k_stride_1,
-    const int k_nope_stride_0,
+    const int64_t k_nope_stride_0,
     const int k_nope_stride_1,
-    const int k_rope_stride_0) {
+    const int64_t k_rope_stride_0) {
   const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
   const int token_id = flat_warp_id / NUM_HEAD_CHUNKS;
   const int head_chunk_id = flat_warp_id % NUM_HEAD_CHUNKS;
@@ -126,11 +126,11 @@ __global__ void concat_mla_absorb_q_kernel(
     nv_bfloat16* out,
     const int num_items,
     const int dim_1,
-    const int a_stride_0,
+    const int64_t a_stride_0,
     const int a_stride_1,
-    const int b_stride_0,
+    const int64_t b_stride_0,
     const int b_stride_1,
-    const int out_stride_0,
+    const int64_t out_stride_0,
     const int out_stride_1) {
   const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
   const int lane_id = get_lane_id();