Skip to content

Commit 635a884

Browse files
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

0 commit comments

Comments
 (0)