[OPUS] add ut test for sparse mla kernel

[minmengdie]

Summary

Adds three serving-shape sweep UTs, each with one markdown table per kernel path:

• op_tests/test_mla_gqa_logits.py — persistent LEGACY MLA decode (and a round-robin CP phase) across GQA head counts
• op_tests/test_mha_opus_d192_v128_logits.py — gfx950 OPUS dense MHA prefill, Q/K=192 V=128
• op_tests/test_mha_d256_logits.py — aiter.flash_attn_func dense MHA prefill, BSHD D=256

Motivation

The serving axes that pick a kernel config are M (query rows / decode batch × MTP) and N (KV length), plus head count. In-tree coverage for these three paths is either a single-shape correctness UT or a bench that reports no accuracy:

• MLA persistent decode lives in op_tests/test_mla_persistent.py / test_mla_persistent_round_robin.py, which do not sweep GQA nhead ∈ {32,64,96,128} against the ctx/batch/MTP grid a model actually serves.
• OPUS D=192/V=128 is exercised by one shape in test_mha.py (test_flash_attn_func_opus_d192_v128).
• BSHD D=256 goes through flash_attn_func without a dedicated M×N table.

There was no single place to answer "what does this op do across the shapes a model actually serves, and is it still correct there".

Changes

Three new files. Each path is a different kernel with a different calling convention, so each gets its own @benchmark function and table:

file	kernel	shape mapping
test_mla_gqa_logits.py decode	aiter.mla_decode_fwd (persistent LEGACY, nhead_kv=1)	paged cache, Q [B, decode_qlen, H, D], fragmented randperm page tables
test_mla_gqa_logits.py cp	same kernel per CP rank + online-softmax merge	round-robin g(j)=j*W+r, -cpw ranks, skipped when ctx < W
test_mha_opus_d192_v128_logits.py	fmha_fwd_bf16_opus_fwd	dense BSHD, Q/K last dim 192, V last dim 128, causal prefill
test_mha_d256_logits.py	aiter.flash_attn_func	dense BSHD, D=256, causal prefill

MLA decode reproduces the production paged-KV convention: LEGACY layout, page size 1, each sequence owns its pages, and a randomized (fragmented) page table rather than a sequential one a real KV cache never hands the kernel. Prefill, non-persistent decode, 3BUFFER and DS32_OPUS are out of scope. KV length is uniform across the batch.

Two things worth flagging for review:

• --ref is off by default. The default product is a kernel-timing sweep. Pass --ref to compare against torch_mla_extend / CP merge / sampled attention_ref. The tables below were taken with --no-ref, so the err columns are empty on purpose — they are not a claim that every default shape was golden-checked.
• Prefill's reference is row-sampled, deliberately. The full [B, H, Sq, Sk] fp32 score tensor does not fit at the large shapes (same issue as [Test] Sweep the MQA logits indexer's M and N across prefill and decode #5434). Opus samples query rows (ref_rows); d256 does the same when --ref is on. LSE is compared separately (opus_ref_lse / flash LSE).

TB/s is compulsory traffic only (Q+K+V+O bytes / us). Prefill re-reads KV per query row, so that column is a lower bound there and TFLOPS is the metric to read; MLA decode at large ctx is closer to bandwidth-bound and TB/s is the one that matters.

Default products after trimming shapes that OOM'd or illegal-memory-access'd on MI355X:

• MLA: nhead={32,64,96,128}, ctx={4096,16384,65536,131072}, batch={16,32,64,128}, mtp={1,2,4,8}, dtype=kvtype ∈ {bf16,fp8} (Q dtype must equal KV dtype), cpw={2,3,4,7,8}. The numbers below only cover nhead 32/64 — 96/128 have no a16w16 LEGACY PS .co on this tree and abort if swept.
• Opus: batch={1,16,64,128}, nheads={32,64}, seqlen_q={1024…16384}, seqlen_k={4096,16384,65664}, skip sk < sq. Dropped sk=131072 and batch=256 from defaults after OOM.
• d256: batch={1,16}, nheads={32,64}, same q/k lists. _shape_ok also skips b≥64 and b≥16 & h≥64 & sk≥131072.

How to run it

python op_tests/test_mla_gqa_logits.py                          # decode + cp, no golden
python op_tests/test_mla_gqa_logits.py -p decode -n 32 64 --ref
python op_tests/test_mla_gqa_logits.py -p cp -cpw 4 -n 32 -mtp 1 -c 64 -b 1
python op_tests/test_mha_opus_d192_v128_logits.py                     # 99 shapes, no golden
python op_tests/test_mha_opus_d192_v128_logits.py -n 32 -b 1 -q 1024 -k 4096 --ref
python op_tests/test_mha_d256_logits.py
python op_tests/test_mha_d256_logits.py --ref

Every flag is a swept list (-n, -c/-k, -b, -mtp, -d, -kvd, -cpw, -p, -q), so a different model's shapes go in without touching the file. Prefill skips seqlen_k < seqlen_q. --perf on the opus file is an alias for --no-ref.

Performance

Environment: MI355X (gfx950, 256 CU, 288 GB HBM), idle GPUs, this tree, --no-ref. Tables list every timed shape (decode 256, opus 99, d256 45). CP timings omitted. Empty err / always-true columns (gfx=gfx950, opus causal, ref_rows=0, kvtype==dtype) are dropped so the comment fits GitHub's 65535-byte limit. us is kernel time.

MLA decode (256/256, nhead 32/64 × bf16+fp8): peak 1829 TFLOPS at fp8 nhead=64 mtp=8 batch=16 ctx=131072, and peak 6.01 TB/s at bf16 nhead=32 mtp=2 batch=128 ctx=131072. fp8 is the TFLOPS story; bf16 at large batch/ctx is the bandwidth story.

Opus D=192/V=128 (99/99 current defaults): peak 1269 TFLOPS at b=1 h=32 sq=sk=16384. Peak 1.49 TB/s at b=1 h=64 sq=1024 sk=4096 — short-M is more traffic-heavy; long-M is the TFLOPS column. An earlier untrimmed run (180 calling shapes) hit 41 OOMs at sk=131072 / b=256; those are no longer in the default product.

BSHD D=256 (45 timed of 52 defaults): peak 512 TFLOPS at b=1 h=32 sq=sk=4096, peak 0.71 TB/s at b=1 h=64 sq=1024 sk=4096. Seven b=16 h=64 shapes with sq>=4096 never timed (HIP illegal memory access on nearby large-sk cases).

MLA decode — mla_decode_fwd persistent LEGACY, 256 shapes

nhead	mtp	batch	ctx	dtype	us	TFLOPS	TB/s
32	1	16	4096	bf16	34.7793	131.21	2.20279
32	1	32	4096	bf16	47.2596	193.121	3.24216
32	1	64	4096	bf16	75.7614	240.936	4.04489
32	1	128	4096	bf16	130.537	279.67	4.69518
32	1	16	16384	bf16	73.4773	248.425	4.12514
32	1	32	16384	bf16	126.233	289.206	4.80231
32	1	64	16384	bf16	230.744	316.43	5.25438
32	1	128	16384	bf16	428.799	340.553	5.65494
32	1	16	65536	bf16	229.433	318.239	5.26983
32	1	32	65536	bf16	431.547	338.384	5.60343
32	1	64	65536	bf16	844.862	345.687	5.72436
32	1	128	65536	bf16	1637.81	356.645	5.90581
32	1	16	131072	bf16	433.825	336.608	5.57145
32	1	32	131072	bf16	837.91	348.555	5.7692
32	1	64	131072	bf16	1639.74	356.224	5.89613
32	1	128	131072	bf16	3246.21	359.875	5.95656
32	1	16	4096	fp8	26.8914	169.697	1.43421
32	1	32	4096	fp8	35.2533	258.892	2.18805
32	1	64	4096	fp8	48.9394	372.984	3.1523
32	1	128	4096	fp8	78.1666	467.044	3.94725
32	1	16	16384	fp8	49.038	372.234	3.09585
32	1	32	16384	fp8	75.8823	481.104	4.00131
32	1	64	16384	fp8	128.989	566.051	4.70781
32	1	128	16384	fp8	245.223	595.493	4.95268
32	1	16	65536	fp8	129.788	562.566	4.65989
32	1	32	65536	fp8	241.599	604.427	5.00664
32	1	64	65536	fp8	462.273	631.786	5.23326
32	1	128	65536	fp8	914.998	638.379	5.28787
32	1	16	131072	fp8	241.978	603.48	4.99541
32	1	32	131072	fp8	463.991	629.448	5.21036
32	1	64	131072	fp8	918.196	636.156	5.26589
32	1	128	131072	fp8	1780.46	656.139	5.4313
32	2	16	4096	bf16	38.2909	238.354	2.02987
32	2	32	4096	bf16	53.0181	344.29	2.93205
32	2	64	4096	bf16	82.4391	442.839	3.7713
32	2	128	4096	bf16	136.596	534.529	4.55216
32	2	16	16384	bf16	82.395	443.076	3.69219
32	2	32	16384	bf16	132.368	551.601	4.59654
32	2	64	16384	bf16	236.292	618.002	5.14987
32	2	128	16384	bf16	448.001	651.913	5.43245
32	2	16	65536	bf16	235.325	620.542	5.14263
32	2	32	65536	bf16	435.731	670.271	5.55475
32	2	64	65536	bf16	842.67	693.173	5.74454
32	2	128	65536	bf16	1639.54	712.535	5.905
32	2	16	131072	bf16	436.677	668.819	5.53761
32	2	32	131072	bf16	843.301	692.654	5.73496
32	2	64	131072	bf16	1649.45	708.254	5.86412
32	2	128	131072	bf16	3216.47	726.405	6.01441
32	2	16	4096	fp8	33.0489	276.161	1.19178
32	2	32	4096	fp8	40.7217	448.252	1.93445
32	2	64	4096	fp8	58.0653	628.727	2.7133
32	2	128	4096	fp8	91.1447	801.083	3.45711
32	2	16	16384	fp8	54.1605	674.057	2.81817
32	2	32	16384	fp8	86.4295	844.786	3.53197
32	2	64	16384	fp8	147.616	989.248	4.13596
32	2	128	16384	fp8	262.675	1111.86	4.64858
32	2	16	65536	fp8	138.624	1053.42	4.36878
32	2	32	65536	fp8	251.87	1159.56	4.80898
32	2	64	65536	fp8	480.496	1215.65	5.04161
32	2	128	65536	fp8	1006.46	1160.73	4.81384
32	2	16	131072	fp8	259.972	1123.42	4.65281
32	2	32	131072	fp8	488.779	1195.05	4.94947
32	2	64	131072	fp8	1002.92	1164.82	4.82428
32	2	128	131072	fp8	2092.31	1116.69	4.62493
32	4	16	4096	bf16	63.862	285.829	1.25198
32	4	32	4096	bf16	82.1515	444.389	1.9465
32	4	64	4096	bf16	123.077	593.242	2.5985
32	4	128	4096	bf16	205.209	711.61	3.11697
32	4	16	16384	bf16	121.014	603.356	2.53233
32	4	32	16384	bf16	199.826	730.78	3.06713
32	4	64	16384	bf16	369.929	789.497	3.31357
32	4	128	16384	bf16	692.574	843.398	3.5398
32	4	16	65536	bf16	370.401	788.49	3.27325
32	4	32	65536	bf16	696.273	838.917	3.48259
32	4	64	65536	bf16	1355.88	861.605	3.57677
32	4	128	65536	bf16	2629.18	888.667	3.68911
32	4	16	131072	bf16	696.651	838.463	3.4743
32	4	32	131072	bf16	1346.79	867.416	3.59428
32	4	64	131072	bf16	2629.71	888.488	3.68159
32	4	128	131072	bf16	5213.95	896.235	3.71369
32	4	16	4096	fp8	51.6132	353.661	0.794864
32	4	32	4096	fp8	61.1803	596.715	1.34113
32	4	64	4096	fp8	83.2575	876.971	1.97102
32	4	128	4096	fp8	128.204	1139.04	2.56002
32	4	16	16384	fp8	84.2623	866.514	1.83085
32	4	32	16384	fp8	126.737	1152.22	2.43452
32	4	64	16384	fp8	215.19	1357.21	2.86764
32	4	128	16384	fp8	396.835	1471.93	3.11004
32	4	16	65536	fp8	213.291	1369.29	2.84708
32	4	32	65536	fp8	382.614	1526.64	3.17425
32	4	64	65536	fp8	771.643	1513.95	3.14786
32	4	128	65536	fp8	1541.98	1515.24	3.15054
32	4	16	131072	fp8	383.199	1524.31	3.16086
32	4	32	131072	fp8	772.668	1511.94	3.1352
32	4	64	131072	fp8	1552.06	1505.39	3.12162
32	4	128	131072	fp8	3077.7	1518.32	3.14842
32	8	16	4096	bf16	79.1359	461.323	1.06665
32	8	32	4096	bf16	124.942	584.386	1.35119
32	8	64	4096	bf16	196.461	743.296	1.71862
32	8	128	4096	bf16	345.463	845.41	1.95472
32	8	16	16384	bf16	192.233	759.645	1.61732
32	8	32	16384	bf16	342.533	852.642	1.81532
32	8	64	16384	bf16	678.941	860.333	1.83169
32	8	128	16384	bf16	1295.36	901.861	1.92011
32	8	16	65536	bf16	646.642	903.306	1.88183
32	8	32	65536	bf16	1325.47	881.371	1.83614
32	8	64	65536	bf16	2662.09	877.678	1.82844
32	8	128	65536	bf16	5120	912.681	1.90136
32	8	16	131072	bf16	1269.49	920.236	1.91008
32	8	32	131072	bf16	2641.63	884.477	1.83586
32	8	64	131072	bf16	5109.76	914.509	1.8982
32	8	128	131072	bf16	10239	912.768	1.89458
32	8	16	4096	fp8	60.9189	599.276	0.727235
32	8	32	4096	fp8	83.4999	874.426	1.06114
32	8	64	4096	fp8	127.189	1148.12	1.39328
32	8	128	4096	fp8	191.769	1522.96	1.84815
32	8	16	16384	fp8	122.082	1196.15	1.29051
32	8	32	16384	fp8	210.325	1388.6	1.49814
32	8	64	16384	fp8	363.776	1605.7	1.73237
32	8	128	16384	fp8	674.28	1732.56	1.86924
32	8	16	65536	fp8	361.349	1616.49	1.68959
32	8	32	65536	fp8	690.181	1692.64	1.7692
32	8	64	65536	fp8	1458.78	1601.65	1.67409
32	8	128	65536	fp8	2863.17	1632.08	1.70589
32	8	16	131072	fp8	673.824	1733.73	1.80242
32	8	32	131072	fp8	1344.24	1738.13	1.80699
32	8	64	131072	fp8	2878.35	1623.47	1.68779
32	8	128	131072	fp8	5941.21	1573.05	1.63537
64	1	16	4096	bf16	39.0198	233.902	1.99196
64	1	32	4096	bf16	53.0744	343.925	2.92893
64	1	64	4096	bf16	83.2075	438.749	3.73648
64	1	128	4096	bf16	138.456	527.348	4.491
64	1	16	16384	bf16	81.4343	448.303	3.73575
64	1	32	16384	bf16	131.421	555.579	4.62969
64	1	64	16384	bf16	238.051	613.435	5.11181
64	1	128	16384	bf16	441.995	660.772	5.50627
64	1	16	65536	bf16	232.839	627.165	5.19752
64	1	32	65536	bf16	437.378	667.747	5.53383
64	1	64	65536	bf16	847.118	689.532	5.71437
64	1	128	65536	bf16	1657.22	704.933	5.842
64	1	16	131072	bf16	439.472	664.564	5.50239
64	1	32	131072	bf16	844.628	691.565	5.72595
64	1	64	131072	bf16	1651.29	707.465	5.85759
64	1	128	131072	bf16	3245.69	719.865	5.96026
64	1	16	4096	fp8	30.9089	295.28	1.2743
64	1	32	4096	fp8	39.8408	458.163	1.97722
64	1	64	4096	fp8	57.0921	639.444	2.75955
64	1	128	4096	fp8	92.461	789.678	3.40789
64	1	16	16384	fp8	54.48	670.103	2.80164
64	1	32	16384	fp8	85.0937	858.048	3.58742
64	1	64	16384	fp8	145.634	1002.71	4.19224
64	1	128	16384	fp8	262.814	1111.27	4.64613
64	1	16	65536	fp8	144.747	1008.86	4.18398
64	1	32	65536	fp8	252.447	1156.91	4.79798
64	1	64	65536	fp8	491.621	1188.14	4.92753
64	1	128	65536	fp8	1009.58	1157.14	4.79896
64	1	16	131072	fp8	258.411	1130.2	4.6809
64	1	32	131072	fp8	485.366	1203.45	4.98427
64	1	64	131072	fp8	1002.6	1165.21	4.82587
64	1	128	131072	fp8	2101.4	1111.86	4.60492
64	2	16	4096	bf16	63.8034	286.091	1.25313
64	2	32	4096	bf16	82.4539	442.759	1.93936
64	2	64	4096	bf16	123.983	588.905	2.5795
64	2	128	4096	bf16	205.498	710.611	3.1126
64	2	16	16384	bf16	122.489	596.092	2.50183
64	2	32	16384	bf16	201.782	723.697	3.0374
64	2	64	16384	bf16	373.514	781.919	3.28177
64	2	128	16384	bf16	687.433	849.705	3.56627
64	2	16	65536	bf16	369.556	790.294	3.28074
64	2	32	65536	bf16	696.774	838.315	3.48009
64	2	64	65536	bf16	1355.08	862.109	3.57886
64	2	128	65536	bf16	2627.98	889.071	3.69079
64	2	16	131072	bf16	698.253	836.538	3.46633
64	2	32	131072	bf16	1352.6	863.69	3.57884
64	2	64	131072	bf16	2632.64	887.499	3.67749
64	2	128	131072	bf16	5217.34	895.652	3.71128
64	2	16	4096	fp8	46.5388	392.223	0.881534
64	2	32	4096	fp8	60.9209	599.256	1.34685
64	2	64	4096	fp8	85.3387	855.584	1.92295
64	2	128	4096	fp8	132.968	1098.22	2.46829
64	2	16	16384	fp8	84.5371	863.697	1.8249
64	2	32	16384	fp8	128.424	1137.08	2.40253
64	2	64	16384	fp8	215.039	1358.16	2.86965
64	2	128	16384	fp8	402.681	1450.57	3.06489
64	2	16	65536	fp8	214.256	1363.13	2.83426
64	2	32	65536	fp8	387.568	1507.13	3.13368
64	2	64	65536	fp8	775.609	1506.21	3.13176
64	2	128	65536	fp8	1546.72	1510.59	3.14088
64	2	16	131072	fp8	404.773	1443.07	2.99239
64	2	32	131072	fp8	765.629	1525.84	3.16403
64	2	64	131072	fp8	1566.13	1491.87	3.09357
64	2	128	131072	fp8	3066.71	1523.76	3.15971
64	4	16	4096	bf16	79.3763	459.926	1.06342
64	4	32	4096	bf16	120.617	605.343	1.39965
64	4	64	4096	bf16	204.091	715.509	1.65437
64	4	128	4096	bf16	335.296	871.043	2.01399
64	4	16	16384	bf16	190.351	767.157	1.63331
64	4	32	16384	bf16	343.698	849.752	1.80916
64	4	64	16384	bf16	680.332	858.574	1.82795
64	4	128	16384	bf16	1291.52	904.539	1.92581
64	4	16	65536	bf16	631.855	924.445	1.92587
64	4	32	65536	bf16	1327.18	880.238	1.83378
64	4	64	65536	bf16	2611.96	894.524	1.86354
64	4	128	65536	bf16	5095.8	917.015	1.91039
64	4	16	131072	bf16	1271.46	918.812	1.90713
64	4	32	131072	bf16	2674.98	873.451	1.81297
64	4	64	131072	bf16	5102.18	915.867	1.90101
64	4	128	131072	bf16	10258.6	911.026	1.89097
64	4	16	4096	fp8	58.7965	620.908	0.753486
64	4	32	4096	fp8	84.6019	863.035	1.04731
64	4	64	4096	fp8	132.382	1103.09	1.33862
64	4	128	4096	fp8	191.501	1525.1	1.85074
64	4	16	16384	fp8	119.971	1217.21	1.31323
64	4	32	16384	fp8	209.011	1397.33	1.50757
64	4	64	16384	fp8	370.064	1578.42	1.70293
64	4	128	16384	fp8	703.93	1659.58	1.7905
64	4	16	65536	fp8	352.976	1654.83	1.72967
64	4	32	65536	fp8	726.485	1608.06	1.68079
64	4	64	65536	fp8	1473.94	1585.18	1.65687
64	4	128	65536	fp8	2881.16	1621.89	1.69524
64	4	16	131072	fp8	671.929	1738.62	1.8075
64	4	32	131072	fp8	1477.03	1581.87	1.64454
64	4	64	131072	fp8	2873.87	1626	1.69042
64	4	128	131072	fp8	5893.42	1585.81	1.64864
64	8	16	4096	bf16	116.815	625.045	0.798901
64	8	32	4096	bf16	197.653	738.813	0.944313
64	8	64	4096	bf16	326.052	895.741	1.14489
64	8	128	4096	bf16	663.448	880.424	1.12531
64	8	16	16384	bf16	329.644	885.978	0.970184
64	8	32	16384	bf16	613.439	952.198	1.0427
64	8	64	16384	bf16	1266.36	922.508	1.01018
64	8	128	16384	bf16	2660.62	878.164	0.961627
64	8	16	65536	bf16	1168.66	999.636	1.04888
64	8	32	65536	bf16	2576.16	906.954	0.951636
64	8	64	65536	bf16	4938.57	946.21	0.992826
64	8	128	65536	bf16	10254.7	911.375	0.956275
64	8	16	131072	bf16	2369.04	986.249	1.02731
64	8	32	131072	bf16	4949.72	944.079	0.983387
64	8	64	131072	bf16	10093.8	925.897	0.964448
64	8	128	131072	bf16	20707.5	902.653	0.940236
64	8	16	4096	fp8	81.8203	892.376	0.621556
64	8	32	4096	fp8	129.694	1125.95	0.784243
64	8	64	4096	fp8	191.274	1526.91	1.06352
64	8	128	4096	fp8	369.498	1580.83	1.10108
64	8	16	16384	fp8	200.102	1459.55	0.820094
64	8	32	16384	fp8	362.12	1613.04	0.906341
64	8	64	16384	fp8	659.491	1771.41	0.995325
64	8	128	16384	fp8	1438.75	1623.95	0.91247
64	8	16	65536	fp8	656.582	1779.26	0.939847
64	8	32	65536	fp8	1310.14	1783.36	0.942015
64	8	64	65536	fp8	2575.13	1814.63	0.958532
64	8	128	65536	fp8	6359.94	1469.49	0.776217
64	8	16	131072	fp8	1277.51	1828.92	0.955821
64	8	32	131072	fp8	2585.41	1807.42	0.944581
64	8	64	131072	fp8	5600.68	1668.7	0.872085
64	8	128	131072	fp8	13126.3	1423.99	0.744195

Opus prefill — fmha_fwd_bf16_opus_fwd, 99 shapes (current defaults)

batch	nheads	seqlen_q	seqlen_k	us	TFLOPS	TB/s
1	32	1024	4096	101.156	424.589	1.0366
1	32	1024	16384	381.186	450.695	0.93528
1	32	1024	65664	1506.04	457.185	0.906865
1	32	2048	4096	135.483	634.022	0.928743
1	32	2048	16384	537.336	639.446	0.702516
1	32	2048	65664	2190.4	628.687	0.633101
1	32	4096	4096	169.784	1011.87	0.988151
1	32	4096	16384	945.995	726.425	0.443375
1	32	4096	65664	4530.09	607.968	0.315377
1	32	8192	16384	1630.7	842.821	0.30865
1	32	8192	65664	8004.41	688.157	0.188967
1	32	16384	16384	2166.04	1269.03	0.309823
1	32	16384	65664	14924.9	738.137	0.112587
1	64	1024	4096	140.52	611.296	1.49242
1	64	1024	16384	571.271	601.461	1.24815
1	64	1024	65664	2443.08	563.664	1.11807
1	64	2048	4096	235.627	729.113	1.06804
1	64	2048	16384	1098.14	625.782	0.687505
1	64	2048	65664	5002.6	550.543	0.554408
1	64	4096	4096	325.248	1056.42	1.03166
1	64	4096	16384	1988.28	691.244	0.421902
1	64	4096	65664	9502.18	579.687	0.300707
1	64	8192	16384	3247.82	846.347	0.309942
1	64	8192	65664	17738.2	621.067	0.170544
1	64	16384	16384	4338.74	1267.09	0.309347
1	64	16384	65664	29890.9	737.121	0.112432
16	32	1024	4096	1169.96	587.368	1.434
16	32	1024	16384	5151.25	533.614	1.10735
16	32	1024	65664	20196.9	545.459	1.08196
16	32	2048	4096	1958.7	701.683	1.02786
16	32	2048	16384	9451.79	581.642	0.639011
16	32	2048	65664	38972.7	565.349	0.569318
16	32	4096	4096	2533.4	1085.02	1.05959
16	32	4096	16384	17259.8	637.036	0.388816
16	32	4096	65664	76946.6	572.687	0.297075
16	32	8192	16384	29662.5	741.348	0.27149
16	32	8192	65664	146545	601.403	0.165144
16	32	16384	16384	35187.7	1249.88	0.305147
16	32	16384	65664	272976	645.717	0.0984902
16	64	1024	4096	2378.13	577.929	1.41096
16	64	1024	16384	9814.22	560.163	1.16245
16	64	1024	65664	40713.7	541.174	1.07346
16	64	2048	4096	4097.44	670.852	0.982694
16	64	2048	16384	18973.8	579.49	0.636647
16	64	2048	65664	78091.2	564.293	0.568255
16	64	4096	4096	5458.35	1007.18	0.983577
16	64	4096	16384	35025.5	627.834	0.3832
16	64	4096	65664	156694	562.452	0.291766
16	64	8192	16384	59146.4	743.587	0.27231
16	64	8192	65664	296155	595.179	0.163435
16	64	16384	16384	75495.2	1165.12	0.284453
16	64	16384	65664	540480	652.255	0.0994874
64	32	1024	4096	4878.4	563.459	1.37563
64	32	1024	16384	20083.9	547.46	1.13609
64	32	1024	65664	80082.5	550.262	1.09149
64	32	2048	4096	8267.75	664.94	0.974033
64	32	2048	16384	37723.4	582.933	0.640429
64	32	2048	65664	158969	554.402	0.558295
64	32	4096	4096	11208.3	980.983	0.957991
64	32	4096	16384	71085.7	618.696	0.377622
64	32	4096	65664	313350	562.519	0.291801
64	32	8192	16384	120596	729.384	0.267108
64	32	8192	65664	591094	596.404	0.163772
64	32	16384	16384	162396	1083.29	0.264475
64	32	16384	65664	1.05675e+06	667.199	0.101767
64	64	1024	4096	9488.56	579.388	1.41452
64	64	1024	16384	39898.4	551.156	1.14376
64	64	1024	65664	161874	544.451	1.07996
64	64	2048	4096	16559.2	663.987	0.972637
64	64	2048	16384	75479.6	582.68	0.640151
64	64	2048	65664	321635	548.029	0.551876
64	64	4096	4096	23361.8	941.289	0.919227
64	64	4096	16384	145759	603.469	0.368328
64	64	4096	65664	624763	564.264	0.292706
64	64	8192	16384	243952	721.132	0.264086
64	64	8192	65664	1.17059e+06	602.313	0.165394
64	64	16384	16384	326267	1078.39	0.26328
64	64	16384	65664	2.12845e+06	662.512	0.101052
128	32	1024	4096	9390.16	585.459	1.42934
128	32	1024	16384	39546.8	556.056	1.15393
128	32	1024	65664	163902	537.714	1.0666
128	32	2048	4096	16827	653.421	0.957159
128	32	2048	16384	75478.5	582.689	0.640161
128	32	2048	65664	322929	545.834	0.549666
128	32	4096	4096	23046	954.187	0.931824
128	32	4096	16384	148813	591.083	0.360768
128	32	4096	65664	630250	559.351	0.290157
128	32	8192	16384	244691	718.954	0.263289
128	32	8192	65664	1.17883e+06	598.102	0.164238
128	32	16384	16384	326922	1076.23	0.262752
128	32	16384	65664	2.14328e+06	657.928	0.100353
128	64	1024	4096	19123.4	574.957	1.4037
128	64	1024	16384	78521.7	560.106	1.16233
128	64	2048	4096	32889.8	668.604	0.9794
128	64	2048	16384	150464	584.597	0.642258
128	64	4096	4096	47022.8	935.301	0.91338
128	64	4096	16384	288723	609.31	0.371893
128	64	8192	16384	489017	719.493	0.263486
128	64	16384	16384	653059	1077.53	0.263068

BSHD D=256 — flash_attn_func, 45 timed of 52 default shapes

batch	nheads	seqlen_q	seqlen_k	us	TFLOPS	TB/s
1	32	1024	4096	259.800	264.509	0.6458
1	32	1024	16384	1263.310	217.585	0.4515
1	32	1024	65664	5016.530	219.606	0.4356
1	32	2048	4096	395.850	347.200	0.5086
1	32	2048	16384	2378.550	231.131	0.2539
1	32	2048	65664	9899.020	222.579	0.2241
1	32	4096	4096	537.250	511.639	0.4996
1	32	4096	16384	4368.630	251.683	0.1536
1	32	4096	65664	19404.020	227.099	0.1178
1	32	8192	16384	7404.120	297.000	0.1088
1	32	8192	65664	37509.830	234.959	0.0645
1	32	16384	16384	9478.150	464.020	0.1133
1	32	16384	65664	69891.440	252.199	0.0385
1	64	1024	4096	474.860	289.430	0.7066
1	64	1024	16384	2490.890	220.707	0.4580
1	64	1024	65664	10108.830	217.960	0.4323
1	64	2048	4096	809.560	339.540	0.4974
1	64	2048	16384	4782.050	229.925	0.2526
1	64	2048	65664	20031.000	219.991	0.2215
1	64	4096	4096	1122.150	489.913	0.4784
1	64	4096	16384	8903.590	246.982	0.1507
1	64	4096	65664	39394.380	223.719	0.1161
1	64	8192	16384	15233.610	288.707	0.1057
1	64	8192	65664	76120.980	231.560	0.0636
1	64	16384	16384	20092.910	437.771	0.1069
1	64	16384	65664	142074.660	248.131	0.0378
16	32	1024	4096	4148.400	265.045	0.6471
16	32	1024	16384	19126.790	229.942	0.4772
16	32	1024	65664	79288.090	222.310	0.4410
16	32	2048	4096	7043.820	312.192	0.4573
16	32	2048	16384	36949.570	238.057	0.2615
16	32	2048	65664	157184.960	224.278	0.2259
16	32	4096	4096	9348.930	470.433	0.4594
16	32	4096	16384	68780.090	255.774	0.1561
16	32	4096	65664	309336.050	227.927	0.1182
16	32	8192	16384	117240.390	300.105	0.1099
16	32	8192	65664	598336.900	235.674	0.0647
16	32	16384	16384	151913.070	463.217	0.1131
16	32	16384	65664	1116481.250	252.601	0.0385
16	64	1024	4096	8776.010	250.572	0.6117
16	64	1024	16384	39152.430	224.663	0.4662
16	64	1024	65664	161396.290	218.426	0.4333
16	64	2048	4096	14187.840	309.987	0.4541
16	64	2048	16384	75985.980	231.519	0.2544
16	64	2048	65664	319301.950	220.814	0.2224

Testing

• Unit tests added (these three files)
• Performance benchmarks run on MI355X (gfx950); full timed tables above (CP omitted)
• Accuracy smoke (--ref) on GPU, small shapes only:
  • MLA decode bf16: err=0 vs torch_mla_extend
  • MLA decode fp8 vs bf16 golden: failed! (max abs ~0.07–0.08) — expected quantization, not a kernel bug; vs fp8 golden: warning!, cal_diff did not abort
  • MLA CP bf16: checked shapes pass; fp8 CP: no cprr heuristic kernel, skipped
  • Opus --ref b=1 h∈{32,64} q=1024 k=4096: output pass, LSE max diff ~2e-6
  • d256 --ref same small shapes: output pass; LSE pass with derived atol=0.125
• Default --no-ref sweeps: MLA decode 256/256, opus 99/99 (after OOM trim), d256 45 timed then IAE on large b/h/sk
• Not tested on gfx942 — SUPPORTED_GFX gates it; MLA also skips bf16 + nhead=32 on gfx942. Opus is gfx950-only.
• nhead={96,128} MLA not in the perf tables: no a16w16 LEGACY PS .co on this tree

Follow-ups (not in this PR)

• Ship or generate a16w16 LEGACY PS kernels for nhead 96/128 so the GQA default list is actually runnable.
• fp8 round-robin CP (cprr) heuristic / kernel — currently decode-only for fp8.
• d256 HIP illegal memory access at b=16 h=64 and sk=131072 (also a few sq>=4096 at that batch/head that never timed). Either fix the kernel path or keep those shapes out of defaults after a clean repro.
• Opus OOM at sk=131072 and batch=256 on MI355X — left out of defaults rather than shipped unrun.
• Variable context lengths — all three sweeps run uniform KV length, so varlen / ragged scheduling is never exercised.
• Full-grid --ref on the default product. Golden is sampled and expensive; the published tables are kernel-only.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5456 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.

[Copilot]

🟡 Changes recommended

Unresolved critical and moderate correctness and reliability issues remain.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds GPU shape-sweep benchmarks and optional accuracy checks for MLA GQA, OPUS D192/V128, and dense MHA D256 kernels.

Changes:

• Adds configurable MLA decode and context-parallel sweeps.
• Adds OPUS and D256 prefill sweeps with sampled references.
• Reports performance metrics and optional correctness results.

File summaries

File	Review findings
op_tests/test_mla_gqa_logits.py	Moderate (2 votes): FP8 LSE is compared against the BF16 reference instead of _lse_ref_fp8.
op_tests/test_mha_opus_d192_v128_logits.py	Moderate (1 vote): Causal FLOPs are incorrectly halved for rectangular bottom-right masks.
op_tests/test_mha_d256_logits.py	Critical (3 votes): Default failing IAE shapes are not excluded or handled. Moderate (2 votes): Missing LSE is silently accepted. Moderate (1 vote): Rectangular causal FLOPs are incorrectly halved.

Review details

Suppressed comments (2)

op_tests/test_mha_d256_logits.py:204

• [verified] The same causal-halving error is present here: the public wrapper documents a bottom-right mask, so for rectangular inputs the visible-token count is not seqlen_q * seqlen_k / 2. For default sq=1024, sk=65664, this makes the table underreport TFLOPS by nearly 2x. Author must derive causal FLOPs from the number of visible key tokens rather than unconditionally dividing the full-matrix count by two.

        "ref_rows": ref_n,
        "fwd us": us,

op_tests/test_mha_opus_d192_v128_logits.py:190

• [verified] With seqlen_k > seqlen_q, flash_attn_func uses a bottom-right causal mask, so row i sees seqlen_k - seqlen_q + i + 1 keys; only square shapes are approximately half of the full QK/PV work. Dividing by two here underreports TFLOPS for rectangular defaults such as sq=1024, sk=65664 by almost 2x, making the reported metric misleading. Author must compute the visible-token count from (seqlen_q, seqlen_k) before deriving causal FLOPs.

• Files reviewed: 3/3 changed files
• Comments generated: 3
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

Unresolved moderate issues affect correctness, benchmark accuracy, and default sweep reliability.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (7)

Previously missed (3) — in code that hasn't changed since the last review.

op_tests/test_mha_d256_logits.py:193

• [verified] For bottom-right causal attention with seqlen_k >= seqlen_q, the matrix is not half full: the visible pair count is seqlen_q * seqlen_k - seqlen_q * (seqlen_q - 1) // 2. Dividing by two therefore under-reports TFLOPS by almost 2x for shapes such as sq=1024, sk=65664, so this sweep's reported performance is inaccurate. Author must compute FLOPs from the actual causal visible-pair count.
  op_tests/test_mha_opus_d192_v128_logits.py:173
• [verified] For bottom-right causal attention with seqlen_k >= seqlen_q, the matrix is not half full: the visible pair count is seqlen_q * seqlen_k - seqlen_q * (seqlen_q - 1) // 2. Dividing by two therefore under-reports TFLOPS by almost 2x for shapes such as sq=1024, sk=65664, making the M×N performance table inaccurate. Author must compute FLOPs from the actual causal visible-pair count.
  op_tests/test_mla_gqa_logits.py:294
• [verified] Despite the PR title saying this adds a sparse MLA UT, the MLA metadata is requested with is_sparse=False here (and again in the normal decode path), while the reference consumes the full page table; this exercises dense persistent GQA, not the sparse kernel/metadata path. The existing sparse test uses is_sparse=True. Author must either add sparse top-k inputs and a matching reference or retitle/describe this as a dense GQA sweep.

op_tests/test_mha_d256_logits.py:301

• [verified] The loop only filters seqlen_k < seqlen_q; it has no _shape_ok for the known HIP illegal-access cases. The documented default therefore still launches b=16, h=64, sq>=4096 cases, and the except below re-raises anything that is not OOM, so the default sweep aborts at the first IAE instead of producing the claimed 45-row table; custom large-batch/large-K cases are also unguarded. Author must add the documented shape filter before allocation/launch rather than catching a poisoned HIP IAE and continuing.

        if seqlen_k < seqlen_q:
            continue

op_tests/test_mha_d256_logits.py:98

• [verified] run_flash passes 0.0 as softmax_scale, but aiter.flash_attn_func treats only None as the default and forwards an explicit zero to the kernel (aiter/ops/mha.py:2619-2620). The benchmark therefore computes unscaled, uniform attention while the reference uses 1/sqrt(256), so --ref is validating a different operation. Author must pass None or `HEAD_DIM**-0.5 here.

        0.0,

op_tests/test_mha_d256_logits.py:166

• [verified] With --ref, requesting return_lse=True is not enough to test the LSE contract because a None result simply skips this check. A dispatch regression that stops returning LSE would still report a successful output comparison, contrary to the stated LSE coverage. Author must assert that softmax_lse is present before slicing and comparing it.

        lse_got = softmax_lse[:, :, rows] if softmax_lse is not None else None
        if lse_got is not None:

op_tests/test_mla_gqa_logits.py:928

• [verified] The no-argument sweep defaults to nhead=96 and 128, but the PR notes that these LEGACY page-size-1 cases have no a16w16 PS code object and abort when swept. check_support only filters dtype/arch, so the advertised default command still reaches those unsupported heads instead of producing a complete table. Author must remove unsupported heads from the defaults or add an explicit capability gate that skips them.

        default=[32, 64, 96, 128],

• Files reviewed: 3/3 changed files
• Comments generated: 1
• Review effort level: Lite