Commit 635a884
authored
[cuda backend] optimized L_kv threshold for sdpa implementation selection. (pytorch#20142)
Our `triton/replacement_pass.py` adopted a suboptimal threshold to
choose between `splited_sdpa` and `regular_sdpa` due to misbenchmarking:
the benchmark script was a. didn't take cuda graph optimization into
account, and b. overfitted into qwen sdpa configuration.
This PR updates benchmarking script to cover gemma4 case as well as
other representative case. Based on the results shows as below we take
L_kv == 256 as the threshold. With the update, the sliding-window decode
attention in gemma4-31b can be benefit from split-sdpa performance gain
and decode perf raised from 38.5 tok/s to 43.98 token/s, beats
llama.cpp.
### gemma sliding (D=256, CTA=16) [B=1 H_q=32 H_kv=16 D=256]
| L_kv | ET Triton (GQA) (us) | ET Split-K (GQA) (us) |
| --- | --- | --- |
| 64 | 6.5 | 7.0 |
| 128 | 9.7 | 9.0 |
| 256 | 16.1 | 12.7 |
| 512 | 27.6 | 13.1 |
| 1024 | 50.8 | 15.8 |
| 2048 | 102.6 | 26.0 |
| 4096 | 200.8 | 54.8 |
| 8192 | 395.2 | 96.0 |
| 16384 | 784.8 | 180.2 |
### qwen (D=256, CTA=2) [B=1 H_q=16 H_kv=2 D=256]
| L_kv | ET Triton (GQA) (us) | ET Split-K (GQA) (us) |
| --- | --- | --- |
| 64 | 6.4 | 7.2 |
| 128 | 9.8 | 9.0 |
| 256 | 15.5 | 12.4 |
| 512 | 27.0 | 12.6 |
| 1024 | 49.8 | 12.7 |
| 2048 | 95.6 | 13.2 |
| 4096 | 186.6 | 14.2 |
| 8192 | 370.2 | 20.6 |
| 16384 | 767.2 | 33.2 |
### head_dim=128 (D=128, CTA=16) [B=1 H_q=32 H_kv=16 D=128]
| L_kv | ET Triton (GQA) (us) | ET Split-K (GQA) (us) |
| --- | --- | --- |
| 64 | 4.5 | 5.8 |
| 128 | 6.8 | 7.2 |
| 256 | 11.2 | 9.5 |
| 512 | 18.6 | 9.9 |
| 1024 | 33.1 | 10.5 |
| 2048 | 62.4 | 13.9 |
| 4096 | 125.8 | 22.2 |
| 8192 | 247.8 | 55.5 |
| 16384 | 490.5 | 98.1 |
## Next
Instead of hard coded the threshold, we should choose the operator base
on the actual performance during export. Other refactor is on the way.1 parent b47e588 commit 635a884
3 files changed
Lines changed: 225 additions & 205 deletions
File tree
- backends/cuda
- benchmarks
- tests
- triton
0 commit comments