[Bugfix][Kernel] Fix integer overflow in libtorch_stable/activation_kernels.cu (#49660)
Signed-off-by: molly-ting <molly.cn.mail@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Molly committed
Jul 29, 2026 at 20:50 UTC
451227cb3ff07989698fed982c2d3e4300257924
1 file changed
+9
-6
csrc/libtorch_stable/activation_kernels.cu
+9
-6
@@ -103,9 +103,10 @@ __global__ void act_and_mul_kernel(
103
scalar_t* __restrict__ out, // [..., d]
104
const scalar_t* __restrict__ input, // [..., 2, d]
105
const int d, const float limit, const float alpha, const float beta) {
106
- const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
106
+ const int64_t token_idx = blockIdx.x;
107
+ const scalar_t* x_ptr = input + token_idx * 2 * d;
108
const scalar_t* y_ptr = x_ptr + d;
108
- scalar_t* out_ptr = out + blockIdx.x * d;
109
+ scalar_t* out_ptr = out + token_idx * d;
110
111
if constexpr (use_vec) {
112
using cuda_t = typename CUDATypeConverter<scalar_t>::Type;
@@ -353,9 +354,10 @@ template <typename scalar_t, typename packed_t,
354
__global__ void act_and_mul_kernel_with_param(
355
scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const int d,
356
const float param) {
356
- const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
357
+ const int64_t token_idx = blockIdx.x;
358
+ const scalar_t* x_ptr = input + token_idx * 2 * d;
359
const scalar_t* y_ptr = x_ptr + d;
358
- scalar_t* out_ptr = out + blockIdx.x * d;
360
+ scalar_t* out_ptr = out + token_idx * d;
361
362
if constexpr (use_vec) {
363
using cuda_t = typename CUDATypeConverter<scalar_t>::Type;
@@ -670,8 +672,9 @@ __global__ void activation_kernel(
672
scalar_t* __restrict__ out, // [..., d]
673
const scalar_t* __restrict__ input, // [..., d]
674
const int d) {
673
- const scalar_t* in_ptr = input + blockIdx.x * d;
674
- scalar_t* out_ptr = out + blockIdx.x * d;
675
+ const int64_t token_idx = blockIdx.x;
676
+ const scalar_t* in_ptr = input + token_idx * d;
677
+ scalar_t* out_ptr = out + token_idx * d;
678
679
if constexpr (use_vec) {
680
// Fast path: 128-bit/256-bit vectorized loop