[Bug] fp8 scaled_mm_v2 TMA template: NoValidChoicesError on XPU (tensorwise/rowwise scaling)
- Dominant language
- Python
- Stars
- 113
- Forks
- 128
- Avg merge
- 5d 13h
- Merged PRs (30d)
- 107
Description
## Bug Description
`torch.nn.functional.scaled_mm` (fp8 `aten._scaled_mm_v2`) lowering fails to
compile via `torch.compile`/Inductor's max-autotune GEMM path on XPU for both
tensorwise and rowwise scaling with TMA templates. Inductor's autotuner finds
"No choices exist for backend" for `aten._scaled_mm_v2.default` on XPU,
because XPU has no TMA-based (or ATEN fallback) scaled_mm choice registered
for the `max_autotune_gemm_backends` selection, causing a hard
`NoValidChoicesError` instead of falling back to an ATEN implementation.
## Affected Tests
Cases:
op_ut,test/inductor/test_fp8.py,TestFP8LoweringXPU.test_tensorwise_scaling_tma_template_bfloat16_shape_1024,1024,512_use_fast_accum_False_xpu_bfloat16
op_ut,test/inductor/test_fp8.py,TestFP8LoweringXPU.test_tensorwise_scaling_tma_template_bfloat16_shape_1024,1024,512_use_fast_accum_True_xpu_bfloat16
op_ut,test/inductor/test_fp8.py,TestFP8LoweringXPU.test_tensorwise_scaling_tma_template_bfloat16_shape_16,32,32_use_fast_accum_False_xpu_bfloat16
op_ut,test/inductor/test_fp8.py,TestFP8LoweringXPU.test_tensorwise_scaling_tma_template_bfloat16_shape_16,32,32_use_fast_accum_True_xpu_bfloat16
op_ut,test/inductor/test_fp8.py,TestFP8LoweringXPU.test_tensorwise_scaling_tma_template_float32_shape_1024,1024,512_use_fast_accum_False_xpu_float32
op_ut,test/inductor/test_fp8.py,TestFP8LoweringXPU.test_tensorwise_scaling_tma_template_float32_shape_1024,1024,512_use_fast_accum_True_xpu_float32
op_ut,test/inductor/test_fp8.py,TestFP8LoweringXPU.test_tensorwise_scaling_tma_template_float32_shape_16,32,32_use_fast_accum_False_xpu_float32
op_ut,test/inductor/test_fp8.py,TestFP8LoweringXPU.test_tensorwise_scaling_tma_template_float32_shape_16,32,32_use_fast_accum_True_xpu_float32
op_ut,test/inductor/test_fp8.py,TestFP8LoweringXPU.test_rowwise_scaling_tma_template_shape_1024,1024,512_use_fast_accum_False_xpu
op_ut,test/inductor/test_fp8.py,TestFP8LoweringXPU.test_rowwise_scaling_tma_template_shape_1024,1024,512_use_fast_accum_True_xpu
op_ut,test/inductor/test_fp8.py,TestFP8LoweringXPU.test_rowwise_scaling_tma_template_shape_16,32,32_use_fast_accum_False_xpu
op_ut,test/inductor/test_fp8.py,TestFP8LoweringXPU.test_rowwise_scaling_tma_template_shape_16,32,32_use_fast_accum_True_xpu
## Error Message
```
torch._inductor.exc.InductorError: LoweringException: NoValidChoicesError: No choices to select. Provided reason: No choices exist for backend. please consider adding ATEN into max_autotune_gemm_backends config (defined in torch/_inductor/config.py) to allow at least one choice.
target: aten._scaled_mm_v2.default
args[0]: TensorBox(StorageBox(
InputBuffer(name='arg2_1', layout=FixedLayout('xpu:0', torch.float8_e4m3fn, size=[1024, 1024], stride=[1024, 1]))
))
args[1]: TensorBox(StorageBox(
InputBuffer(name='arg1_1', layout=FixedLayout('xpu:0', torch.float8_e4m3fn, size=[1024, 512], stride=[1, 1024]))
))
...
args[9]: torch.bfloat16
kwargs: {'contraction_dim': [], 'use_fast_accum': False}
```
Additionally, a `UserWarning` is emitted for every case using `use_fast_accum=True`:
```
UserWarning: scaled_mm: fast_accum is not supported in XPU for now. It would
silently set use_fast_accum to false. (Triggered internally at
/__w/pytorch/pytorch/aten/src/ATen/native/mkldnn/xpu/ScaledBlas.cpp:218.)
```
## Traceback
```
pytest_command:
pytest test/inductor/test_fp8.py -k "test_tensorwise_scaling_tma_template or test_rowwise_scaling_tma_template" -v
Traceback:
File "test/inductor/test_fp8.py", line 919, in test_tensorwise_scaling_tma_template
y_compiled, code = run_and_get_code(
File ".../torch/_inductor/utils.py", line 2734, in run_and_get_code
result = fn(*args, **kwargs)
File ".../torch/_dynamo/eval_frame.py", line 1278, in compile_wrapper
raise e.remove_dynamo_frames() from None
File ".../torch/_inductor/compile_fx.py", line 1192, in _compile_fx_inner
raise InductorError(e, currentframe()).with_traceback(
File ".../torch/_inductor/graph.py", line 1567, in call_function
raise LoweringException(
File ".../torch/_inductor/lowering.py", line 524, in wrapped
out = decomp_fn(*args, **kwargs)
File ".../torch/_inductor/kernel/mm.py", line 1170, in tuned_scaled_mm_v2
node, _ = autotune_select_algorithm(name, choices, kernel_inputs.nodes(), layout)
File ".../torch/_inductor/select_algorithm.py", line 4063, in __call__
raise self.create_no_valid_choices(name, "No choices exist for backend.")
torch._inductor.exc.InductorError: LoweringException: NoValidChoicesError: No
choices to select. Provided reason: No choices exist for backend. please
consider adding ATEN into max_autotune_gemm_backends config (defined in
torch/_inductor/config.py) to allow at least one choice.
target: aten._scaled_mm_v2.default
Found from :
File "test/inductor/test_fp8.py", line 888, in linear
y = torch.nn.functional.scaled_mm(
File ".../torch/nn/functional.py", line 7236, in scaled_mm
out = torch._scaled_mm_v2(
Same failure pattern (different lowering line numbers: line 1080/1049 for
test_rowwise_scaling_tma_template) for the rowwise-scaling variant.
```
## Root Cause Analysis
`tuned_scaled_mm_v2` in `torch/_inductor/kernel/mm.py` builds a list of
`choices` for `aten._scaled_mm_v2.default` (TMA-based Triton templates and
ATEN fallback) gated by `max_autotune_gemm_backends`. On XPU, the TMA
template path used for tensorwise/rowwise fp8 scaling produces zero valid
choices (either because XPU is not a supported device for the TMA-based
Triton persistent template, or the ATEN choice is not being added/enabled
for XPU when TMA is selected), so `autotune_select_algorithm` raises
`NoValidChoicesError` instead of falling back to an eager/ATEN
`_scaled_mm_v2` implementation. This means fp8 scaled_mm cannot currently be
compiled via Inductor on XPU at all for these configurations, even though
the eager fp8 `scaled_mm` op itself works (as shown by the `UserWarning`
about `use_fast_accum` being silently disabled on XPU, which comes from the
eager `ScaledBlas.cpp` kernel executing normally). The separate
`use_fast_accum` warning is expected/known XPU behavior (fast_accum
unsupported) and is not itself the cause of the failure.
## Related PyTorch Issues
- None found matching `NoValidChoicesError` for `_scaled_mm_v2` on XPU.
## Related Intel/torch-xpu-ops Issues
- None found (searched "scaled_mm_v2 NoValidChoicesError", "tma_template scaled_mm").
## Versions
Branch: daisyden/inductor (repo: daisyden/pytorch)
Commit: 8f036e818c26e60f0a28fb51e89a5010315a27c3
Platform: linux, Python 3.10.20
Contributor guide
Assessment
This issue has not been assessed yet.