CUDA: XOR swizzle flash attn K,V smem fp16 tiles

[ynankani]

Overview

This PR adds XOR swizzling for K/V shared-memory tiles in the CUDA flash-attention fp16 MMA kernel (fattn-mma-f16.cuh), replacing row padding (+4) as the primary fix for shared-memory bank conflicts on cp.async stores and ldmatrix loads. On Turing+ (Ampere/Ada/Blackwell), when the K/V tile stride is a power-of-two ≥ 32, the kernel uses a pow2 stride with per-row XOR address remapping instead of padded row stride.

Additional information

Followed resource : https://lubits.ch/flash/Part-4

This change helped in improving perf for very high context, collected data for 65K depth

Performance DGX -SPARK

Model	Metric	ToT	Swizzle	Δ
Gemma-4-26B-A4B-NVFP4	pp2048 @ d65536	1616.61 ± 2.94	1644.18 ± 1.80	1.71%
Gemma-4-26B-A4B-NVFP4	tg128 @ d65536	29.30 ± 0.04	30.32 ± 0.03	3.48%
Qwen3.6-27B-Q4_K_M	pp2048 @ d65536	574.91 ± 1.19	583.02 ± 1.44	1.41%
Qwen3.6-27B-Q4_K_M	tg128 @ d65536	9.43 ± 0.01	9.76 ± 0.01	3.50%
Qwen3.6-27B-NVFP4	pp2048 @ d65536	648.32 ± 0.77	665.23 ± 0.60	2.61%
Qwen3.6-27B-NVFP4	tg128 @ d65536	8.29 ± 0.00	8.53 ± 0.01	2.89%
Qwen3.6-35B-A3B-NVFP4	pp2048 @ d65536	1729.35 ± 8.33	1758.57 ± 3.19	1.69%
Qwen3.6-35B-A3B-NVFP4	tg128 @ d65536	49.22 ± 0.20	52.39 ± 0.15	6.44%
Qwen3.6-35B-A3B-UD-Q4_K_M	pp2048 @ d65536	1662.27 ± 4.71	1680.99 ± 3.85	1.13%
Qwen3.6-35B-A3B-UD-Q4_K_M	tg128 @ d65536	46.24 ± 0.20	48.48 ± 0.63	4.84%

Performance RTX6000-PRO-BLACKWELL

Model	Metric	ToT (t/s)	Swizzle (t/s)	Δ
Qwen3.6-35B-A3B NVFP4 (W4A16)	pp2048 @ d65536	5105.62 ± 8.83	5213.21 ± 7.96	2.11%
Qwen3.6-35B-A3B NVFP4 (W4A16)	tg128 @ d65536	191.42 ± 3.47	193.00 ± 2.71	0.83%
Qwen3.6-35B-A3B-UD-Q4_K_M	pp2048 @ d65536	4757.63 ± 13.50	4947.80 ± 9.65	4.00%
Qwen3.6-35B-A3B-UD-Q4_K_M	tg128 @ d65536	197.44 ± 2.57	197.45 ± 2.75	0.01%
Qwen3.6-27B-NVFP4	pp2048 @ d65536	1805.03 ± 4.16	1868.25 ± 2.47	3.50%
Qwen3.6-27B-NVFP4	tg128 @ d65536	53.91 ± 0.20	53.71 ± 0.23	−0.37%
Qwen3.6-27B-Q4_K_M	pp2048 @ d65536	1518.74 ± 1.59	1630.54 ± 5.56	7.36%
Qwen3.6-27B-Q4_K_M	tg128 @ d65536	54.70 ± 0.21	55.91 ± 0.15	2.21%
Gemma-4-26B-A4B-NVFP4	pp2048 @ d65536	4997.70 ± 36.84	5213.21 ± 7.96	4.31%
Gemma-4-26B-A4B-NVFP4	tg128 @ d65536	144.94 ± 0.51	143.93 ± 0.87	−0.70%

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: Yes, paired with Claude on this