CUDA: Follow up of #25635, refactoring FA shared smem swizzle

[ynankani]

Overview

This is follow-up of #25635, refactor FA shared smem swizzle.
This is WIP, I am working on isolating on different hardware to find which cases swizzle helps best on granular level.
For starters I have disabled swizzling for MHA

Any design related suggestion or review are welcome.

Additional information

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: Yes, for testing and refactoring

[ynankani]

Hi @JohannesGaessler

I was trying to repro the numbers posted on the comment #25635 (comment) to gate regressing configs. I observed below, just to confirm you were running with https://huggingface.co/unsloth/Llama-3.1-8B-Instruct-GGUF/tree/main? I see net gain in prefill for llama3.1-8b-instruct

Performance

From the comment
RTX 3090	llama 8B Q4_0	layer	1	pp1024@d32768	92.84	92.31	0.99
RTX 3090	llama 8B Q4_0	layer	2	pp1024@d32768	178.41	176.74	0.99
RTX 3090	llama 8B Q4_0	layer	4	pp1024@d32768	297.95	269.61	0.9
RTX 3090	llama 8B Q4_0	layer	8	pp1024@d32768	398.07	370.72	0.93
RTX 3090	llama 8B Q4_0	layer	16	pp1024@d32768	838.28	785.35	0.94
RTX 3090	llama 8B Q4_0	layer	32	pp1024@d32768	1217.3	1190.95	0.98
RTX 3090	llama 8B Q4_0	layer	64	pp1024@d32768	1592.73	1365.97	0.86
RTX 3090	llama 8B Q4_0	layer	128	pp1024@d32768	1494.87	1458.58	0.98
RTX 3090	llama 8B Q4_0	layer	256	pp1024@d32768	1581.03	1570.17	0.99
RTX 3090	llama 8B Q4_0	layer	512	pp1024@d32768	2076.8	1882.33	0.91
Re-run locally					off (swz=0)	on (swz=1)	Observed
RTX 3090	llama 8B Q4_0	layer	1	pp1024@d32768	93.74	94.21	1.005
RTX 3090	llama 8B Q4_0	layer	2	pp1024@d32768	181.93	182.8	1.005
RTX 3090	llama 8B Q4_0	layer	4	pp1024@d32768	332.46	336.06	1.011
RTX 3090	llama 8B Q4_0	layer	8	pp1024@d32768	520.85	523.54	1.005
RTX 3090	llama 8B Q4_0	layer	16	pp1024@d32768	1005.82	1008.99	1.003
RTX 3090	llama 8B Q4_0	layer	32	pp1024@d32768	1535.79	1559.7	1.016
RTX 3090	llama 8B Q4_0	layer	64	pp1024@d32768	1945.26	2005.82	1.031
RTX 3090	llama 8B Q4_0	layer	128	pp1024@d32768	1808.57	1808.58	1
RTX 3090	llama 8B Q4_0	layer	256	pp1024@d32768	2173.82	2263.12	1.041
RTX 3090	llama 8B Q4_0	layer	512	pp1024@d32768	2204.96	2338.54	1.061

[JohannesGaessler]

Ill re-check the numbers. The RTX 3090 is in a server that I'm sharing with other devs so feasibly there may have been interference and I didn't notice.

[JohannesGaessler]

I re-ran the benchmarks and the numbers I get match yours within 1%. Notably the numbers that I posted in the previous thread were worse than my current ones. So presumably my numbers indicating a regression were simply wrong. Since there seems to be no case where swizzling hurts performance if it can be used I would say to keep the logic equivalent to what we have on master: a function to check whether swizzling can be used but no explicit value in the configs.

[ynankani]

I re-ran the benchmarks and the numbers I get match yours within 1%. Notably the numbers that I posted in the previous thread were worse than my current ones. So presumably my numbers indicating a regression were simply wrong. Since there seems to be no case where swizzling hurts performance if it can be used I would say to keep the logic equivalent to what we have on master: a function to check whether swizzling can be used but no explicit value in the configs.

Thank you for the re-run.

Actually I was also re-running microbenchmark perf test for the enabled swizzle configs and observed regression on some configs. I also checked ncu traces for one the config and saw higher bank-conflicts with swizzle on vs off. So maybe we need more granular control per config or need to handle tilling in some way for these. One approach can be just disable swizzle for configs which are regressing. I am not yet able to find the pattern why these configs regress. Any suggestions?

Performance

4-reps

RTX3090
DKQ	DV	nc1	nc2	mean%	min%	max%
512	512	4	8	−9.99	−23.68	−0.73
256	256	2	8	−3.15	−4.24	−1.08
512	512	8	2	−0.64	−3.87	2.96
512	512	4	4	−0.44	−3.51	3.09
RTX4090
DKQ	DV	nc1	nc2	mean%	min%	max%
512	512	4	4	−13.29	−37.96	−0.31
512	512	8	2	−13.24	−37.77	−0.33
512	512	4	8	−5.78	−7.80	−2.17
256	256	2	8	−4.36	−10.31	−0.99
RTX PRO BLACKWELL 6000
DKQ	DV	nc1	nc2	mean%	min%	max%
512	512	4	8	-12.35	-17.46	-3.72
256	256	2	8	-2.13	-7.91	6.89
512	512	8	2	-1.56	-5.66	6.22
512	512	2	8	-0.76	-3.35	3.62

[JohannesGaessler]

For context, do you mean you are seeing more bank conflicts with the version in this PR?

[ynankani]

For context, do you mean you are seeing more bank conflicts with the version in this PR?

I meant swizzling has more bank conflicts compared to padding for few configs which are reported in previous comment rest configs swizzling is better compared to padding. It is inherent to swizzling not related to refactoring. I mean for these configs padding has lesser bank conflicts compared to both swizzling implementation (merged PR and refactor PR).