[Kernel][MI350] Add 8 wave GQA sliding-window attention kernel

[amd-nprotaso]

Summary

Adds kernels/attention/swa_gfx950.py, a hand-scheduled GQA (grouped-query) forward attention kernel targeting CDNA4 (gfx950 / MI350). It supports both dense and sliding-window attention with aiter (LEFT, RIGHT) window semantics, bf16 I/O, and head dim 64/128.

What it does

• Work decomposition: 8 waves per CTA, each wave independently computes attention for its own 32-row Q tile (Q_BLOCK_SIZE=32, KV_BLOCK_SIZE=64, NUM_WARPS=8). Grid is (ATTN_H, ceil(seq_q / (32*8)), ATTN_B).
• GQA: GROUP_SIZE = ATTN_H // ATTN_H_KV query heads share each KV head; head remap (bx % H_KV) * GROUP + (bx // H_KV).
• Sliding window: when sliding_window=(LEFT, RIGHT), each CTA computes only the in-band KV tile range (base_kv_row / nt_ct), and mask_band() applies the per-query [i-LEFT, i+RIGHT] mask via a single unsigned-compare trick. Dense path when sliding_window=None.
• Softmax: online flash-style with a lazy-threshold deferred rescale (rescale_defer, ballot-voted RESCALE_THRESHOLD=8.0) in the steady-state loop, and unconditional rescale in the epilogue. Cross-32-lane reductions via permlane32_swap.
• MMA: MFMA 32x32x16 bf16 atoms for both QK and OV.

Performance engineering

• Double-buffered K/V staged through LDS (k0/k1/v0/v1), Gâ��LDS via BufferCopyLDS128b DMA atoms with swizzled offsets for K; CDNA4 ds_read_tr16 transposing LDS reads for V.
• Manual software pipeline: 2Ã� unrolled steady-state scf.for with loop-carried register state (_flatten/_unflatten phis) to avoid spills, plus a fully unrolled 12-block epilogue.
• Explicit scheduling control: sched_group_barrier MFMA/VALU/EXP interleaving, s_setprio, nop_anchor, ROCDL s_waitcnt with hand-computed vmcnt/lgkmcnt bitfields, and provenance-preserving LDS reads (getelementptr, not ptrtointâ��inttoptr) so alias analysis drops conservative vmcnt(0) drains.

B	seq	win(L,R)	flydsl ms	aiter ms	flydsl/ait	flydsl TF	aiter TF
1	1024	(512,0)	0.0354	0.0482	1.36x	182.2	133.9
1	2048	(512,0)	0.0424	0.0625	1.48x	355.5	240.9
1	2048	(1024,0)	0.0583	0.0800	1.37x	442.4	322.4
1	4096	(512,0)	0.0823	0.1030	1.25x	392.2	313.2
1	4096	(1024,0)	0.1148	0.1386	1.21x	524.3	434.3
1	8192	(512,0)	0.1592	0.1772	1.11x	419.0	376.4
1	8192	(1024,0)	0.2155	0.2513	1.17x	598.5	513.2
1	8192	(4096,0)	0.5412	0.5957	1.10x	762.0	692.4
1	16384	(512,0)	0.3116	0.3235	1.04x	435.1	419.0
1	16384	(1024,0)	0.4527	0.4893	1.08x	588.8	544.7
1	16384	(4096,0)	1.2312	1.3945	1.13x	781.6	690.1
1	16384	(8192,0)	2.0835	2.2723	1.09x	791.7	725.9
1	32768	(512,0)	0.6208	0.6497	1.05x	440.2	420.6
1	32768	(1024,0)	0.9133	1.0199	1.12x	593.1	531.1
1	32768	(4096,0)	2.6153	3.0386	1.16x	788.5	678.6
1	32768	(8192,0)	4.8045	5.4283	1.13x	801.1	709.0
1	65536	(512,0)	1.2541	1.3152	1.05x	437.5	417.2
1	65536	(1024,0)	1.8437	2.0522	1.11x	592.3	532.1
1	65536	(4096,0)	5.3787	6.3342	1.18x	792.3	672.8
1	65536	(8192,0)	10.2574	11.6588	1.14x	804.0	707.4
1	131072	(512,0)	2.5106	2.6223	1.04x	437.9	419.3
1	131072	(1024,0)	3.7263	4.1360	1.11x	588.4	530.1
1	131072	(4096,0)	10.8920	12.9600	1.19x	795.1	668.3
1	131072	(8192,0)	21.2167	24.8864	1.17x	803.4	684.9

API

launch = build_gqa_attn(ATTN_B=16, ATTN_H=64, ATTN_H_KV=8, ATTN_D=128,
sliding_window=(LEFT, RIGHT)) # or None for dense
launch(Q, K, V, O, Q_stride1, K_stride1, V_stride1, O_stride1, seq_len_q, seq_len_kv)

[yanguahe]

Why are there no unit tests for the newly added kernels/attention/swa_gfx950.py kernel?