fix(mamba): widen causal_conv1d token offsets to int64 (#33665)

python/sglang/kernels/ops/mamba/causal_conv1d_triton.py

@@ -91,7 +91,9 @@ def _causal_conv1d_fwd_kernel(  # continuous batching
 
     # base of the sequence
     x_base = (
-        x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim
+        x_ptr
+        + sequence_start_index.to(tl.int64) * stride_x_token
+        + idx_feats * stride_x_dim
     )  # [BLOCK_N,]
 
     if IS_CONTINUOUS_BATCHING:
@@ -173,7 +175,10 @@ def _causal_conv1d_fwd_kernel(  # continuous batching
             )  # [BLOCK_M]
             x_ptrs = (
                 x_ptr
-                + ((sequence_start_index + idx_tokens_last) * stride_x_token)[:, None]
+                + (
+                    (sequence_start_index + idx_tokens_last).to(tl.int64)
+                    * stride_x_token
+                )[:, None]
                 + (idx_feats * stride_x_dim)[None, :]
             )  # [BLOCK_M,BLOCK_N,]
             mask_x = (
@@ -268,7 +273,7 @@ def _causal_conv1d_fwd_kernel(  # continuous batching
     else:  # chunk_offset > 0
         # read prior-token data from `x`
         load_init_state = True
-        prior_tokens = x_base + (token_offset - 1) * stride_x_token
+        prior_tokens = x_base + (token_offset - 1).to(tl.int64) * stride_x_token
         mask_w = idx_feats < dim
         if KERNEL_WIDTH == 2:
             conv_states_ptrs = prior_tokens  # [BLOCK_N]
@@ -305,7 +310,7 @@ def _causal_conv1d_fwd_kernel(  # continuous batching
     else:
         acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32)
 
-    x_base_1d = x_base + token_offset * stride_x_token  # starting of chunk
+    x_base_1d = x_base + token_offset.to(tl.int64) * stride_x_token  # starting of chunk
 
     # PRE-LOAD WEIGHTS
     mask_w = idx_feats < dim
@@ -372,7 +377,8 @@ def _causal_conv1d_fwd_kernel(  # continuous batching
         )  # token-index  # feature-index
         o_ptrs = (
             o_ptr
-            + (sequence_start_index + token_offset + idx_token) * stride_o_token
+            + (sequence_start_index + token_offset + idx_token).to(tl.int64)
+            * stride_o_token
             + (idx_feats * stride_o_dim)
         )