[xpu][test] Enable int8 mixed-precision training tests (test/prototype/test_quantized_training.py) and register scaled_int8_mm for XPU - #4630

What

The PR consists of two parts.

1. Make test_int8_mixed_precision_training and test_bitnet_training (in test/prototype/test_quantized_training.py) accelerator-generic:

• @pytest.mark.skipif(not torch.cuda.is_available(), ...) →
  @pytest.mark.skipif(not torch.accelerator.is_available(), reason="GPU not available")
• device = "cuda" → device = get_current_accelerator_device()

The tests use scaled_int8_mm, which need to be properly registered for XPU.

2. Register the existing Triton kernel under the XPU dispatch key in
torchao/prototype/quantized_training/int8_mm.py:

 @torch.library.impl(lib, "scaled_int8_mm", "CUDA")
-def scaled_int8_mm_cuda(A: Tensor, B: Tensor, row_scale: Tensor, col_scale: Tensor):
+@torch.library.impl(lib, "scaled_int8_mm", "XPU")
+def scaled_int8_mm_gpu(A: Tensor, B: Tensor, row_scale: Tensor, col_scale: Tensor):

Why

Both callers select the Triton path on any backend whenever torch.utils._triton.has_triton()
is true — which includes XPU, because the CI image installs pytorch-triton-xpu:

# torchao/prototype/quantized_training/int8_mixed_precision.py (and bitnet.py)
if has_triton():
    from .int8_mm import scaled_int8_mm
else:
    ...  # torch._int_mm fallback

So on XPU the op is called but has no XPU dispatch key — it raises NotImplementedError
instead of falling back to torch._int_mm.

Neither part is independently valid (kernel registration and UT enablement), hence one PR:

• Kernel registration alone — nothing exercises it on XPU (tests are still CUDA-gated), so there is no evidence it works.
• Test enablement alone — the test hits torch.ops.torchao.scaled_int8_mm and errors out, turning a silent skip into a red XPU failure.

Validation

Environment: Intel XPU node, Python 3.12, torch 2.14.0.dev20260729+xpu, pytorch-triton-xpu present, single card via ZE_AFFINITY_MASK=0.

Quick smoke test

python -c "
import torch, torchao
from torchao.prototype.quantized_training.int8_mm import scaled_int8_mm
A = torch.randint(-8, 8, (64, 64), dtype=torch.int8, device='xpu')
B = torch.randint(-8, 8, (64, 64), dtype=torch.int8, device='xpu')
s = torch.ones(64, device='xpu')
print(scaled_int8_mm(A, B, s, s).shape)"

Before (main):

NotImplementedError: Could not run 'torchao::scaled_int8_mm' with arguments from the 'XPU' backend

After (PR branch):

torch.Size([64, 64])

Full test — pytest -rs test/prototype/test_quantized_training.py

16 test IDs move SKIPPED → PASSED: the full parametrization of test_int8_mixed_precision_training (compile × 4 Int8MixedPrecisionTrainingConfigs × module_swap), including the compile=True half, i.e. inductor-xpu. No test moves in the other direction.

CPU-only collection check

CPU-only collection is unaffected — checked explicitly:

CUDA_VISIBLE_DEVICES="" ZE_AFFINITY_MASK=99 pytest -q --collect-only test/prototype/test_quantized_training.py

Collection succeeds with the same test count as main. No regressions.

Notes

• test_bitnet_training keeps @pytest.mark.skip("Flaky on CI") — this PR only removes the CUDA gate; it stays skipped on all platforms.
• No behavioural change on CUDA: the CUDA registration is untouched.
• This PR is stacked on [xpu][test] Use device-agnostic device list in test_quantized_training #4628; the first commit is [xpu][test] Use device-agnostic device list in test_quantized_training #4628's. Review the second commit.
  Will rebase once [xpu][test] Use device-agnostic device list in test_quantized_training #4628 lands.