From f617d69c60661322cde8fd1a600f87f73edc5bf4 Mon Sep 17 00:00:00 2001 From: Taksh Date: Wed, 24 Jun 2026 16:43:10 +0530 Subject: [PATCH] fix: cast x buffer offset to int64 to prevent int32 overflow Fixes #884 --- csrc/selective_scan/selective_scan_bwd_kernel.cuh | 2 +- csrc/selective_scan/selective_scan_fwd_kernel.cuh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/csrc/selective_scan/selective_scan_bwd_kernel.cuh b/csrc/selective_scan/selective_scan_bwd_kernel.cuh index c720ba28c..26f2be3ee 100755 --- a/csrc/selective_scan/selective_scan_bwd_kernel.cuh +++ b/csrc/selective_scan/selective_scan_bwd_kernel.cuh @@ -139,7 +139,7 @@ void selective_scan_bwd_kernel(SSMParamsBwd params) { float delta_bias = params.delta_bias_ptr == nullptr ? 0 : reinterpret_cast(params.delta_bias_ptr)[dim_id]; scan_t *x = params.x_ptr == nullptr ? nullptr - : reinterpret_cast(params.x_ptr) + (batch_id * params.dim + dim_id) * (params.n_chunks) * params.dstate; + : reinterpret_cast(params.x_ptr) + (int64_t)(batch_id * params.dim + dim_id) * params.n_chunks * params.dstate; float dD_val = 0; float ddelta_bias_val = 0; diff --git a/csrc/selective_scan/selective_scan_fwd_kernel.cuh b/csrc/selective_scan/selective_scan_fwd_kernel.cuh index c64098a3a..4cb671cef 100755 --- a/csrc/selective_scan/selective_scan_fwd_kernel.cuh +++ b/csrc/selective_scan/selective_scan_fwd_kernel.cuh @@ -111,7 +111,7 @@ void selective_scan_fwd_kernel(SSMParamsBase params) { input_t *Bvar = reinterpret_cast(params.B_ptr) + batch_id * params.B_batch_stride + group_id * params.B_group_stride; weight_t *C = reinterpret_cast(params.C_ptr) + dim_id * kNRows * params.C_d_stride; input_t *Cvar = reinterpret_cast(params.C_ptr) + batch_id * params.C_batch_stride + group_id * params.C_group_stride; - scan_t *x = reinterpret_cast(params.x_ptr) + (batch_id * params.dim + dim_id * kNRows) * params.n_chunks * params.dstate; + scan_t *x = reinterpret_cast(params.x_ptr) + (int64_t)(batch_id * params.dim + dim_id * kNRows) * params.n_chunks * params.dstate; float D_val[kNRows] = {0}; if (params.D_ptr != nullptr) {