Preserve device index in safe_int_mm fallback

[fallintoplace]

Summary

• Return the CPU fallback result to the full input device instead of just the device type.
• Add a regression test for a two-CUDA-device setup where the current CUDA device differs from the input tensor device.

Why

The fallback path computes on CPU and then used .to(input.device.type), which preserves cuda but drops the device index. If safe_int_mm receives tensors on cuda:1 while the current default CUDA device is cuda:0, the fallback output can land on the wrong GPU.

Validation

• python3 -m py_compile torchao/kernel/intmm.py test/kernel/test_autotuner.py
• uv run --isolated --no-project --with ruff==0.11.6 ruff format --check torchao/kernel/intmm.py test/kernel/test_autotuner.py
• uv run --isolated --no-project --with ruff==0.11.6 ruff check torchao/kernel/intmm.py test/kernel/test_autotuner.py
• uv run --isolated --no-project --with torch --with parameterized python -m unittest test.kernel.test_autotuner.TestQuantFlow.test_safe_int_mm_fallback_preserves_cuda_device_index (skips locally without two CUDA devices)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4520

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [ROCm] MI350 CI jobs will have longer queue times due to CI migration

This comment was automatically generated by Dr. CI and updates every 15 minutes.