Make tensor of indices a LongTensor

[hirschsn]

Indices get multiplied with strides in operations like aten::index_put. A 32-bit index can lead to a silent overflow, since the Torch Inductor does not do any overflow checking on its own.

This commit fixes GitHub issue #2430

[tianyu-l]

sounds reasonable to me, but why it only happens for compile, not eager?

[hirschsn]

The kernel that gets launched by index_put in eager mode is (according to a Torch profile) void at::native::index_elementwise_kernel</* ... */, at::native::gpu_index_kernel<at::native::index_put_kernel_impl //... .
The kernel that does the indexing is gpu_index_kernel (aten/src/ATen/native/cuda/IndexKernel.cu:101ff) and it receives the indices as 64-bit integers.

I think they are converted in index_put_with_sort_kernel (which seems to be the dispatch target for CUDA), there the indices go through makeLinearIndex (aten/src/ATen/native/cuda/Indexing.cu:615), which in turn does:

  for (auto & i : indices) {
    if (i.defined() && i.dtype() == at::kInt) {
      i = i.to(at::kLong);
    }
  }

[tianyu-l]

have you measured the perf/memory impact of this change? sounds necessary to me anyway

[hirschsn]

We measured on a Qwen3 MoE model with moderate local batch size (16), performance impact was negligible (<1% TPS).