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

[amd-nprotaso]

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.0288	0.0277	0.0434	0.96x	1.51x	149.4	154.9	98.9
1	512	dense	1	2.2	0.0288	0.0288	0.0546	1.00x	1.89x	74.7	74.8	39.4
1	1024	dense	0	17.2	0.0401	0.0397	0.0551	0.99x	1.37x	427.9	432.9	311.7
1	1024	dense	1	8.6	0.0392	0.0386	0.0642	0.98x	1.64x	219.1	222.6	133.9
1	2048	dense	0	68.7	0.0740	0.0736	0.0918	0.99x	1.24x	928.7	933.8	748.7
1	2048	dense	1	34.4	0.0620	0.0618	0.0867	1.00x	1.40x	554.9	556.3	396.3
1	4096	dense	0	274.9	0.2423	0.2430	0.2656	1.00x	1.10x	1134.4	1131.4	1034.8
1	4096	dense	1	137.5	0.1594	0.1602	0.1580	1.00x	0.99x	862.2	858.2	870.1
1	8192	dense	0	1099.5	0.9563	0.8861	0.9432	0.93x	0.99x	1149.8	1240.9	1165.7
1	8192	dense	1	549.8	0.5263	0.5197	0.5059	0.99x	0.96x	1044.7	1057.9	1086.8
1	16384	dense	0	4398.0	4.0628	3.8351	4.1030	0.94x	1.01x	1082.5	1146.8	1071.9
1	16384	dense	1	2199.2	2.2386	2.1735	2.0957	0.97x	0.94x	982.4	1011.8	1049.3
1	32768	dense	0	17592.2	15.8300	15.7176	16.2748	0.99x	1.03x	1111.3	1119.3	1080.9
1	32768	dense	1	8796.4	8.5478	8.4433	8.1752	0.99x	0.96x	1029.1	1041.8	1076.0
1	65536	dense	0	70368.7	65.0197	64.7224	65.1340	1.00x	1.00x	1082.3	1087.2	1080.4
1	65536	dense	1	35184.9	33.6178	33.6022	32.9892	1.00x	0.98x	1046.6	1047.1	1066.6

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.0339	0.0277	0.0477	0.82x	1.41x	126.8	154.9	90.0
1	512	dense	1	2.2	0.0339	0.0288	0.0493	0.85x	1.45x	63.4	74.8	43.6
1	1024	dense	0	17.2	0.0512	0.0397	0.0820	0.77x	1.60x	335.3	432.9	209.6
1	1024	dense	1	8.6	0.0494	0.0386	0.0786	0.78x	1.59x	174.0	222.6	109.4
1	2048	dense	0	68.7	0.0916	0.0736	0.2363	0.80x	2.58x	750.1	933.8	290.8
1	2048	dense	1	34.4	0.0819	0.0618	0.1411	0.75x	1.72x	419.8	556.3	243.7
1	4096	dense	0	274.9	0.3283	0.2430	0.9359	0.74x	2.85x	837.3	1131.4	293.7
1	4096	dense	1	137.5	0.2135	0.1602	0.4508	0.75x	2.11x	643.8	858.2	305.0
1	8192	dense	0	1099.5	1.5590	0.8861	3.6604	0.57x	2.35x	705.3	1240.9	300.4
1	8192	dense	1	549.8	0.7961	0.5197	1.8233	0.65x	2.29x	690.7	1057.9	301.6
1	16384	dense	0	4398.0	5.6241	3.8351	13.5141	0.68x	2.40x	782.0	1146.8	325.4
1	16384	dense	1	2199.2	3.2136	2.1735	7.0074	0.68x	2.18x	684.3	1011.8	313.8
1	32768	dense	0	17592.2	22.2808	15.7176	52.9751	0.71x	2.38x	789.6	1119.3	332.1
1	32768	dense	1	8796.4	11.6732	8.4433	26.5002	0.72x	2.27x	753.5	1041.8	331.9

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.0343	0.0277	0.0428	0.81x	1.25x	125.3	154.9	100.5
1	512	dense	1	2.2	0.0330	0.0288	0.0449	0.87x	1.36x	65.3	74.8	47.9
1	1024	dense	0	17.2	0.0529	0.0397	0.0638	0.75x	1.20x	324.5	432.9	269.3
1	1024	dense	1	8.6	0.0508	0.0386	0.0648	0.76x	1.28x	169.3	222.6	132.8
1	2048	dense	0	68.7	0.0934	0.0736	0.1255	0.79x	1.34x	736.0	933.8	547.7
1	2048	dense	1	34.4	0.0857	0.0618	0.1017	0.72x	1.19x	401.1	556.3	338.0
1	4096	dense	0	274.9	0.3279	0.2430	0.4364	0.74x	1.33x	838.3	1131.4	629.9
1	4096	dense	1	137.5	0.2234	0.1602	0.2411	0.72x	1.08x	615.3	858.2	570.2
1	8192	dense	0	1099.5	1.4016	0.8861	1.4219	0.63x	1.01x	784.5	1240.9	773.2
1	8192	dense	1	549.8	0.8891	0.5197	0.7837	0.58x	0.88x	618.4	1057.9	701.6
1	16384	dense	0	4398.0	5.1092	3.8351	6.3239	0.75x	1.24x	860.8	1146.8	695.5
1	16384	dense	1	2199.2	3.0911	2.1735	3.3299	0.70x	1.08x	711.4	1011.8	660.4
1	32768	dense	0	17592.2	19.9893	15.7176	24.7110	0.79x	1.24x	880.1	1119.3	711.9
1	32768	dense	1	8796.4	11.8189	8.4433	13.4218	0.71x	1.14x	744.3	1041.8	655.4
1	65536	dense	0	70368.7	82.0275	64.7224	101.1079	0.79x	1.23x	857.9	1087.2	696.0
1	65536	dense	1	35184.9	47.2505	33.6022	55.4363	0.71x	1.17x	744.6	1047.1	634.7

[jhinpan]

I found one correctness blocker in the new sink integration, plus two smaller actionable issues in bias validation/pipelining. I reproduced the sink/LSE failure on gfx950 (Sq=512, Skv=128); the device CI for this head did not run because prepare-mlir timed out. I also tested odd-pitch and storage-offset bias tensors on gfx950 and did not report alignment as a finding because identical inputs produced identical kernel outputs.

[yanguahe]

Additional review findings. The first two are specific to this PR. The ragged dense-paged issue appears to be pre-existing, but it remains relevant because the new paged-bias API and test currently avoid rather than reject the unsupported case.

[amd-nprotaso]

updated 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.0246	0.0249	0.0359	1.01x	1.46x	174.7	172.3	119.7
1	512	dense	1	2.2	0.0249	0.0240	0.0468	0.96x	1.88x	86.6	89.8	45.9
1	1024	dense	0	17.2	0.0360	0.0356	0.0463	0.99x	1.29x	476.7	482.3	370.8
1	1024	dense	1	8.6	0.0346	0.0343	0.0560	0.99x	1.62x	248.2	250.4	153.4
1	2048	dense	0	68.7	0.0821	0.0822	0.0968	1.00x	1.18x	837.0	836.3	709.7
1	2048	dense	1	34.4	0.0628	0.0630	0.0835	1.00x	1.33x	547.3	546.0	411.5
1	4096	dense	0	274.9	0.2839	0.2834	0.3053	1.00x	1.08x	968.1	970.0	900.2
1	4096	dense	1	137.5	0.1774	0.1800	0.1775	1.01x	1.00x	775.0	763.9	774.3
1	8192	dense	0	1099.5	1.0656	1.0630	1.1182	1.00x	1.05x	1031.8	1034.3	983.3
1	8192	dense	1	549.8	0.6014	0.6021	0.5999	1.00x	1.00x	914.3	913.2	916.5
1	16384	dense	0	4398.0	4.2151	4.2035	4.3781	1.00x	1.04x	1043.4	1046.3	1004.6
1	16384	dense	1	2199.2	2.2182	2.1918	2.2430	0.99x	1.01x	991.4	1003.3	980.5
1	32768	dense	0	17592.2	16.7338	16.7546	17.2803	1.00x	1.03x	1051.3	1050.0	1018.0
1	32768	dense	1	8796.4	8.6496	8.6463	8.7641	1.00x	1.01x	1017.0	1017.3	1003.7
1	65536	dense	0	70368.7	67.3649	66.8795	69.1947	0.99x	1.03x	1044.6	1052.2	1017.0
1	65536	dense	1	35184.9	33.9038	33.8535	34.7213	1.00x	1.02x	1037.8	1039.3	1013.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.0293	0.0249	0.0403	0.85x	1.38x	146.4	172.3	106.4
1	512	dense	1	2.2	0.0288	0.0240	0.0415	0.83x	1.44x	74.8	89.8	51.8
1	1024	dense	0	17.2	0.0448	0.0356	0.0712	0.80x	1.59x	383.5	482.3	241.2
1	1024	dense	1	8.6	0.0444	0.0343	0.0688	0.77x	1.55x	193.6	250.4	125.0
1	2048	dense	0	68.7	0.0970	0.0822	0.2198	0.85x	2.26x	708.2	836.3	312.7
1	2048	dense	1	34.4	0.0808	0.0630	0.1331	0.78x	1.65x	425.2	546.0	258.3
1	4096	dense	0	274.9	0.3462	0.2834	0.7949	0.82x	2.30x	794.1	970.0	345.8
1	4096	dense	1	137.5	0.2308	0.1800	0.4380	0.78x	1.90x	595.5	763.9	313.9
1	8192	dense	0	1099.5	1.3062	1.0630	3.0839	0.81x	2.36x	841.8	1034.3	356.5
1	8192	dense	1	549.8	0.7747	0.6021	1.6237	0.78x	2.10x	709.7	913.2	338.6
1	16384	dense	0	4398.0	5.1691	4.2035	12.1289	0.81x	2.35x	850.8	1046.3	362.6
1	16384	dense	1	2199.2	2.8078	2.1918	6.2640	0.78x	2.23x	783.2	1003.3	351.1
1	32768	dense	0	17592.2	20.6223	16.7546	48.1869	0.81x	2.34x	853.1	1050.0	365.1
1	32768	dense	1	8796.4	10.7157	8.6463	24.2576	0.81x	2.26x	820.9	1017.3	362.6

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.0249	0.0357	0.83x	1.19x	142.8	172.3	120.3
1	512	dense	1	2.2	0.0284	0.0240	0.0377	0.85x	1.33x	75.8	89.8	57.1
1	1024	dense	0	17.2	0.0467	0.0356	0.0539	0.76x	1.15x	368.1	482.3	318.9
1	1024	dense	1	8.6	0.0441	0.0343	0.0543	0.78x	1.23x	194.8	250.4	158.2
1	2048	dense	0	68.7	0.0982	0.0822	0.1239	0.84x	1.26x	699.8	836.3	554.8
1	2048	dense	1	34.4	0.0849	0.0630	0.0968	0.74x	1.14x	404.8	546.0	355.2
1	4096	dense	0	274.9	0.3308	0.2834	0.4093	0.86x	1.24x	830.9	970.0	671.6
1	4096	dense	1	137.5	0.2377	0.1800	0.2509	0.76x	1.06x	578.4	763.9	547.9
1	8192	dense	0	1099.5	1.2303	1.0630	1.5128	0.86x	1.23x	893.7	1034.3	726.8
1	8192	dense	1	549.8	0.7925	0.6021	0.8271	0.76x	1.04x	693.8	913.2	664.8
1	16384	dense	0	4398.0	4.8128	4.2035	5.8535	0.87x	1.22x	913.8	1046.3	751.3
1	16384	dense	1	2199.2	2.8872	2.1918	3.1668	0.76x	1.10x	761.7	1003.3	694.4
1	32768	dense	0	17592.2	18.9699	16.7546	22.9919	0.88x	1.21x	927.4	1050.0	765.1
1	32768	dense	1	8796.4	11.0617	8.6463	12.2793	0.78x	1.11x	795.2	1017.3	716.4
1	65536	dense	0	70368.7	75.6528	66.8795	91.1370	0.88x	1.20x	930.2	1052.2	772.1
1	65536	dense	1	35184.9	43.3252	33.8535	48.9696	0.78x	1.13x	812.1	1039.3	718.5