intel / intel/torch-xpu-ops

[Bug] fp8 scaled_mm_v2 TMA template: NoValidChoicesError on XPU (tensorwise/rowwise scaling)

Open
#4,852 0 comments 0 reactions 1 assignee Claimed by @jkosnox View on GitHub
module: inductor test: ut
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.