[XPU][B580] `flex-attn-causal` performance drop with `last_level_cache_size` cache clear between runs
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
This PR https://github.com/pytorch/pytorch/pull/184567 introduced regression in a few cases for `flex-attn-causal` with triton, for example:
`H_q=16, H_kv=2, N_CTX_q=1, N_CTX_kv=1024, D_HEAD_qk=128, D_HEAD_v=128 case dropped from 0.81 -> 0.41 TFLOPS`
More details: https://github.com/intel/intel-xpu-backend-for-triton/issues/7302
TL;DR;
`18MB` of cache clear between runs seems like not enough and it's leading to not choosing optimal config when auto tuning triton kernel. Before that change it had `256MB` as a fallback which confirmed by experiment produced more stable results out of autotune, also per experiments 3x L2 cache size clearing between runs seems to help.
Run [repro.py](https://github.com/user-attachments/files/29742915/repro.py) with `TORCHINDUCTOR_USE_EXPERIMENTAL_BENCHMARKER=0/1` and expects that with 1 (default) results are unstable + picking worse performing kernel comparing to 0 (0 uses triton's default `do_bench` so 256MB cache clear between runs)
It is proved by experiments that it's only about `InductorBenchmarker.get_device_cache_size`
Due to that I'm proposing to increase this amount in `InductorBenchmarker` by at least 3x or even 4x or something like `max(N * cache_size, 256MB)` to be more general not tied to B580
```python
(triton) gta@DUT2130BMGFRD:~/workspace$ TORCHINDUCTOR_USE_EXPERIMENTAL_BENCHMARKER=1 python repro.py
# repro #7302 on BMG | TORCHINDUCTOR_USE_EXPERIMENTAL_BENCHMARKER=1 | 4 fresh runs
run 1: 0.439 TFLOPs
run 2: 0.435 TFLOPs
run 3: 0.499 TFLOPs
run 4: 0.439 TFLOPs
mean=0.453 min=0.435 max=0.499 spread=1.15x slow(<0.62)=4/4
VERDICT: BIMODAL / REGRESSED (repro)
(triton) gta@DUT2130BMGFRD:~/workspace$ TORCHINDUCTOR_USE_EXPERIMENTAL_BENCHMARKER=0 python repro.py
# repro #7302 on BMG | TORCHINDUCTOR_USE_EXPERIMENTAL_BENCHMARKER=0 | 4 fresh runs
run 1: 0.728 TFLOPs
run 2: 0.753 TFLOPs
run 3: 0.754 TFLOPs
run 4: 0.751 TFLOPs
mean=0.746 min=0.728 max=0.754 spread=1.04x slow(<0.62)=0/4
VERDICT: STABLE (no regression)
```
-----------------------
### For more reliable repro recommended to narrow gpu clock
Before test:
`sudo xpu-smi config -d 0 -t 0 --frequencyrange 2800,2850`
Recover back after test:
`sudo xpu-smi config -d 0 -t 0 --frequencyrange 1200,2850`
@guangyey
cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo @gujinghui @fengyuan14 @guangyey @Chillee @drisspg @yanboliang @BoyuanFeng @liangel-02 @howardzhang-cv
Contributor guide
Assessment
This issue has not been assessed yet.