[Bug]: Potential Integer Overflow and Out-of-bounds in selective_scan_fwd.cu

[molly-ting]

Your current environment

The output of python collect_env.py

Your output of `python collect_env.py` here

🐛 Describe the bug

While performing static analysis on CUDA kernels, I identified a potential integer overflow and subsequent out-of-bounds memory access in selective_scan_fwd.cu.

vllm/csrc/mamba/mamba_ssm/selective_scan_fwd.cu

Lines 127 to 128 in a00d625

	input_t *u = reinterpret_cast<input_t *>(params.u_ptr) + sequence_start_index * params.u_batch_stride
	+ dim_id * kNRows * params.u_d_stride;

dim_id * kNRows * params.u_d_stride may overflow.
dim_id * kNRows * params.u_d_stride = blockIdx.y * u.size[1] where u.size[1] = batch_size * seq_len.
Example Scenario:

blockIdx.y = 4018
seq_len = 267264
batch_size = 2

In this case, dim_id * kNRows * params.u_d_stride exceeds the 32-bit integer range, causing integer overflow.
As a result, the computed pointer offset becomes negative, and the dereference of *u leads to out-of-bounds memory access.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.