pytorch / pytorch/pytorch

Inductor autotune cache calls legacy get_float32_matmul_precision() and hard-fails under fp32_precision="bfx9"

Open
#196,728 1 comment 0 reactions 0 assignees View on GitHub
bot-triaged module: flex attention module: inductor module: tf32 oncall: pt2 release triage triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

## Summary

Setting the new per-backend API

```python
torch.backends.cuda.matmul.fp32_precision = "bfx9"
```

makes every `torch.compile` region that reaches Inductor's autotune cache raise:

```
torch._inductor.exc.InductorError: LoweringException: RuntimeError: PyTorch is checking the matmul precision without a specific backend name,Current status indicate that you have used mix of the legacy and new APIs to set the matmul precision. We suggest only using the new API for matmul precision. See also: https://pytorch.org/docs/main/notes/cuda.html#tensorfloat-32-tf32-on-ampere-and-later-devices
```

`PersistentCache.lookup` still reads the legacy global `torch.get_float32_matmul_precision()`, and `at::Context::float32MatmulPrecision()` rejects that read once the new per-backend API has been used. The user never touches the legacy API — selecting `bfx9` alone is sufficient to trip the guard.

This reproduces with unmodified `main` of both PyTorch and TorchTitan. TorchTitan's `enable_fp32_matmul_emulation_with_bf16x9()` (`torchtitan/distributed/utils.py`) sets `fp32_precision = "bfx9"` unconditionally when `torch.cuda.get_device_capability() >= (10, 0)`, called as the first statement of `init_distributed()`. All TorchTitan training on Blackwell therefore fails at the first step.

## Environment

- `torch` `2.15.0a0+gita56a3a18` (built from `main` after #195301)
- `torchtitan` `0.2.2+gitc215a12` (`main`)
- `triton` `3.8.0`
- NVIDIA GB200, compute capability 10.0, CUDA 13.x, aarch64

## Reproduction

Observed in nightly TorchTitan DeepSeek-V3 16B pretraining on 8x GB200 (2 nodes x 4 GPUs), crashing at training step 1. Note this happens with no `--compile.enable`: TorchTitan's `set_determinism()` already wraps `flex_attention` in `torch.compile`.

Traceback (module-relative paths, single rank, intermediate frames elided):

```
File "torchtitan/models/common/attention.py", line 362, in forward
out_1HTV, aux = FlexAttention.compiled_flex_attn(
File "torchtitan/models/common/attention.py", line 304, in compiled_flex_attn
out, aux = FlexAttention._compiled_flex_attn(
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 1355, in _compile_fx_inner
raise InductorError(e, currentframe()).with_traceback(
...
File "torch/_inductor/graph.py", line 1551, in call_function
out = lowerings[target](*args, **kwargs)
File "torch/_inductor/lowering.py", line 526, in wrapped
out = decomp_fn(*args, **kwargs)
File "torch/_inductor/kernel/flex/flex_attention.py", line 586, in flex_attention
out, _ = autotune_select_algorithm(
File "torch/_inductor/select_algorithm.py", line 6267, in autotune_select_algorithm
return cache(*args, **kwargs)
File "torch/_inductor/select_algorithm.py", line 4222, in __call__
precompile_fn = self.make_precompile_fn(
File "torch/_inductor/select_algorithm.py", line 4817, in make_precompile_fn
timings = self.lookup(
File "torch/_inductor/codecache.py", line 469, in lookup
precision = torch.get_float32_matmul_precision()
File "torch/__init__.py", line 2005, in get_float32_matmul_precision
return _C._get_float32_matmul_precision()
torch._inductor.exc.InductorError: LoweringException: RuntimeError: PyTorch is checking the matmul precision without a specific backend name,Current status indicate that you have used mix of the legacy and new APIs to set the matmul precision. ...
target: flex_attention
```

The FlexAttention `q`/`k`/`v` in that graph are `torch.bfloat16`. The autotune cache lookup is unconditional, so this is not limited to FP32 matmuls.

A minimal standalone reduction should be the following. It is reduced from the traceback above and has **not** been executed standalone — the reporter has no sm_100 access outside CI, so the evidence above is the observed failure:

```python
import torch
from torch.nn.attention.flex_attention import flex_attention

torch.backends.cuda.matmul.fp32_precision = "bfx9"
q, k, v = (torch.randn(1, 16, 4096, 128, device="cuda", dtype=torch.bfloat16) for _ in range(3))
torch.compile(flex_attention)(q, k, v)
```

## Expected vs actual

Expected: selecting `bfx9` through the documented per-backend API is honoured, and Inductor's autotune cache keys on it, as it already does for other cache and guard state.

Actual: the autotune cache's legacy read trips the mixed-API guard and compilation fails hard. A user cannot avoid this while `bfx9` is set.

## Analysis

The migration to the per-backend API inside `torch/_inductor` appears incomplete. #195301 migrated several cache and guard sites (`torch/_inductor/runtime/caching/encoders.py`, `torch/_inductor/fx_passes/pad_mm.py`, `torch/csrc/dynamo/guards.cpp`) and added a correct `bfx9`-aware read in `torch/_inductor/kernel/flex/flex_attention.py`:

```python
precision = torch.backends.cuda.matmul.fp32_precision
if precision == "none":
precision = (
"ieee" if torch.get_float32_matmul_precision() == "highest" else "tf32"
)
if dtype == torch.float32 and precision == "bfx9":
# See Note [BF16x9 precision] in torch/_inductor/utils.py.
```

but it did not change these two, which still call the legacy getter directly:

- `torch/_inductor/codecache.py`, `PersistentCache.lookup` — the raising frame
- `torch/_inductor/select_algorithm.py`, autotune cache key construction

The PR description states that persistent cache keys now carry the full device-specific precision value, so these two sites look like an oversight rather than a deliberate exclusion. Separately, the guard added in `flex/flex_attention.py` only covers the FP32 lowering decision, while the autotune cache lookup runs for every dtype — which is why a bfloat16 FlexAttention call still hits this.

## Suggested fix

Key the autotune cache on the per-backend precision, reusing the pattern already present in `torch/_inductor/kernel/flex/flex_attention.py`: read `torch.backends.cuda.matmul.fp32_precision` and fall back to the legacy getter only when it is `"none"`. That keeps existing IEEE/TF32 cache entries valid while making `bfx9` a distinct cache key — which is also the correct behaviour, since `bfx9` selects different GEMM arithmetic and should not reuse IEEE or TF32 autotune results.

cc @zasdfgbnm @ptrblck @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo @Chillee @drisspg @yanboliang @BoyuanFeng @liangel-02 @howardzhang-cv

---

_This issue was drafted with assistance from the `opus` AI model._

Contributor guide

Open the contributing guide

Research direction

Start with PersistentCache.lookup in torch/_inductor/codecache.py and the autotune cache key construction in torch/_inductor/select_algorithm.py, then compare their precision handling with torch/_inductor/kernel/flex/flex_attention.py. Run the provided FlexAttention reduction on Blackwell or CI; done means bfx9 compilation no longer raises, bfx9 has a distinct cache key, and the existing legacy fallback remains valid.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers, machine-learning, performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
72/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.