mamba-ssm: replace deprecated torch.get_autocast_gpu_dtype() calls

[albertvillanova]

Problem

Running the mamba-ssm kernel under CUDA autocast on a recent PyTorch (2.13+) emits a DeprecationWarning:

mamba-ssm/torch-ext/mamba_ssm/ops/triton/ssd_combined.py:1511: DeprecationWarning: torch.get_autocast_gpu_dtype() is deprecated. Please use torch.get_autocast_dtype('cuda') instead. (Triggered internally at .../torch/csrc/autograd/init.cpp:1007.)
    dtype = torch.get_autocast_gpu_dtype()

We hit this via huggingface/transformers' NemotronH (Nemotron 3) Mamba2 mixer, which fetches this kernel through kernels's lazy_load_kernel("mamba-ssm"), and it surfaces as CI noise downstream — see huggingface/trl#6555, and huggingface/transformers#47577 for the related discussion on when this kernel path gets engaged.

Occurrences found

All under mamba-ssm/torch-ext/mamba_ssm/:

• ops/triton/ssd_combined.py:1511
• ops/selective_scan_interface.py:264,266,268,270
• ops/triton/layer_norm.py:1023,1028
• distributed/tensor_parallel.py:36,45,46

Suggested fix

Replace torch.get_autocast_gpu_dtype() with torch.get_autocast_dtype('cuda'), per PyTorch's own deprecation message (the device-agnostic autocast API landed in pytorch/pytorch#124359). Since mamba-ssm only targets CUDA today, this should be a mechanical drop-in replacement at all the call sites above.

[ccyyy1023]

Hi, I’d like to work on this. I’ll replace the remaining torch.get_autocast_gpu_dtype() calls in the mamba-ssm kernel with torch.get_autocast_dtype("cuda"), add a focused regression check to prevent the deprecated API from being reintroduced, and run the relevant local checks. Please let me know if this scope matches your expectations.

[sayakpaul]

Cc: @vasqu

[sayakpaul]

@vasqu a gentle ping