Skip to content

[Kernel][MI350] Add bias, alibi bias and sink to flash attention - #960

Open
amd-nprotaso wants to merge 4 commits into
ROCm:mainfrom
amd-nprotaso:new_bias_alibi_sink_attention
Open

[Kernel][MI350] Add bias, alibi bias and sink to flash attention#960
amd-nprotaso wants to merge 4 commits into
ROCm:mainfrom
amd-nprotaso:new_bias_alibi_sink_attention

Conversation

@amd-nprotaso

@amd-nprotaso amd-nprotaso commented Aug 3, 2026

Copy link
Copy Markdown
Contributor

Add bias, alibi bias and sink to flash attention 950gfx implementation.
performance: D=128 H=32 H_KV=8 dtype=torch.bfloat16
Without bias and alibi

B seq mode causal GFLOP new ms main ms aiter ms new/main new/ait new TF main TF aiter TF
1 512 dense 0 4.3 0.0248 0.0240 0.0355 0.97x 1.43x 172.9 178.8 121.1
1 512 dense 1 2.2 0.0243 0.0237 0.0451 0.98x 1.86x 88.5 90.6 47.7
1 1024 dense 0 17.2 0.0351 0.0352 0.0466 1.00x 1.33x 489.5 488.4 368.6
1 1024 dense 1 8.6 0.0345 0.0355 0.0591 1.03x 1.71x 249.2 242.4 145.5
1 2048 dense 0 68.7 0.0815 0.0822 0.0970 1.01x 1.19x 843.2 835.6 708.2
1 2048 dense 1 34.4 0.0614 0.0631 0.0817 1.03x 1.33x 560.1 544.7 421.0
1 4096 dense 0 274.9 0.2792 0.2804 0.3036 1.00x 1.09x 984.5 980.2 905.3
1 4096 dense 1 137.5 0.1749 0.1765 0.1778 1.01x 1.02x 786.0 778.8 773.2
1 8192 dense 0 1099.5 1.0471 1.0548 1.1178 1.01x 1.07x 1050.1 1042.4 983.7
1 8192 dense 1 549.8 0.5903 0.5958 0.5960 1.01x 1.01x 931.4 922.8 922.5
1 16384 dense 0 4398.0 4.1278 4.1669 4.3691 1.01x 1.06x 1065.5 1055.5 1006.6
1 16384 dense 1 2199.2 2.1678 2.1890 2.2426 1.01x 1.03x 1014.5 1004.6 980.6
1 32768 dense 0 17592.2 16.4058 16.5720 17.2883 1.01x 1.05x 1072.3 1061.6 1017.6
1 32768 dense 1 8796.4 8.4807 8.5568 8.6957 1.01x 1.03x 1037.2 1028.0 1011.6
1 65536 dense 0 70368.7 65.6934 66.0745 68.8566 1.01x 1.05x 1071.2 1065.0 1022.0
1 65536 dense 1 35184.9 33.3594 33.6627 34.2810 1.01x 1.03x 1054.7 1045.2 1026.4

Bias:

B seq mode causal GFLOP new ms main ms aiter ms new/main new/ait new TF main TF aiter TF
1 512 dense 0 4.3 0.0380 0.0240 0.0401 0.63x 1.06x 112.9 178.8 107.0
1 512 dense 1 2.2 0.0390 0.0237 0.0407 0.61x 1.05x 55.2 90.6 52.8
1 1024 dense 0 17.2 0.0665 0.0352 0.0711 0.53x 1.07x 258.3 488.4 241.6
1 1024 dense 1 8.6 0.0702 0.0355 0.0716 0.51x 1.02x 122.5 242.4 120.1
1 2048 dense 0 68.7 0.1321 0.0822 0.2204 0.62x 1.67x 520.2 835.6 311.8
1 2048 dense 1 34.4 0.1310 0.0631 0.1335 0.48x 1.02x 262.5 544.7 257.5
1 4096 dense 0 274.9 0.4838 0.2804 0.8058 0.58x 1.67x 568.2 980.2 341.1
1 4096 dense 1 137.5 0.3750 0.1765 0.4393 0.47x 1.17x 366.6 778.8 312.9
1 8192 dense 0 1099.5 2.0100 1.0548 3.0922 0.52x 1.54x 547.0 1042.4 355.6
1 8192 dense 1 549.8 1.2518 0.5958 1.6288 0.48x 1.30x 439.2 922.8 337.6
1 16384 dense 0 4398.0 7.9285 4.1669 12.1398 0.53x 1.53x 554.7 1055.5 362.3
1 16384 dense 1 2199.2 4.8240 2.1890 6.3136 0.45x 1.31x 455.9 1004.6 348.3
1 32768 dense 0 17592.2 30.6325 16.5720 48.2565 0.54

Alibi:

B seq mode causal GFLOP new ms main ms aiter ms new/main new/ait new TF main TF aiter TF
1 512 dense 0 4.3 0.0301 0.0240 0.0485 0.80x 1.61x 142.9 178.8 88.5
1 512 dense 1 2.2 0.0283 0.0237 0.0368 0.84x 1.30x 76.0 90.6 58.5
1 1024 dense 0 17.2 0.0487 0.0352 0.0547 0.72x 1.12x 352.9 488.4 314.2
1 1024 dense 1 8.6 0.0466 0.0355 0.0570 0.76x 1.22x 184.5 242.4 150.8
1 2048 dense 0 68.7 0.0988 0.0822 0.1236 0.83x 1.25x 695.7 835.6 556.0
1 2048 dense 1 34.4 0.0872 0.0631 0.0971 0.72x 1.11x 394.2 544.7 354.0
1 4096 dense 0 274.9 0.3309 0.2804 0.4089 0.85x 1.24x 830.7 980.2 672.2
1 4096 dense 1 137.5 0.2439 0.1765 0.2511 0.72x 1.03x 563.7 778.8 547.4
1 8192 dense 0 1099.5 1.2388 1.0548 1.5123 0.85x 1.22x 887.6 1042.4 727.0
1 8192 dense 1 549.8 0.8014 0.5958 0.8210 0.74x 1.02x 686.1 922.8 669.7
1 16384 dense 0 4398.0 4.8249 4.1669 5.8702 0.86x 1.22x 911.5 1055.5 749.2
1 16384 dense 1 2199.2 2.9316 2.1890 3.1826 0.75x 1.09x 750.2 1004.6 691.0
1 32768 dense 0 17592.2 19.0073 16.5720 22.9100 0.87x 1.21x 925.6 1061.6 767.9
1 32768 dense 1 8796.4 11.1354 8.5568 12.2412 0.77x 1.10x 789.9 1028.0 718.6
1 65536 dense 0 70368.7 75.9698 66.0745 91.4092 0.87x 1.20x 926.3 1065.0 769.8
1 65536 dense 1 35184.9 43.5166 33.6627 49.1306 0.77x 1.13x 808.5 1045.2 716.2

@coderfeli
coderfeli requested a review from yanguahe August 4, 2026 10:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant