fix: cast x buffer offset to int64 to prevent int32 overflow

[Chessing234]

Summary

• Cast the x buffer pointer offset to int64_t in the forward and backward selective scan CUDA kernels before multiplying by n_chunks * dstate.
• Prevents signed int32 overflow when batch * dim * n_chunks * dstate exceeds INT32_MAX, which previously wrapped to a negative offset and silently corrupted memory or caused illegal CUDA accesses.

Fixes #884

Test plan

• Build CUDA extensions with MAMBA_FORCE_BUILD=TRUE pip install --no-cache-dir --force-reinstall . --no-build-isolation
• Run selective scan forward/backward with large batch, dim, n_chunks, and dstate where the product exceeds 2^31-1
• Verify gradients match a smaller configuration that does not overflow
• Confirm no regression on default/smaller model configs

Made with Cursor