[Perf] pa_decode_sparse: BLOCK_K=64 is suboptimal above the CU count — up to +43 % on gfx950 with a load-dependent choice

pa_decode_sparse hardcodes BLOCK_K = 64 (attention/pa_decode_sparse.py, gfx950 gluon driver). On MI355X this is only the right choice while the launch fits on the CUs — above that, BLOCK_K = 32 is faster by up to 43 %, with identical numerics.

Measurements

MI355X (gfx950), 64 heads, head_dim 512 (448 NoPE + 64 RoPE), packed fp8 cache, top-k 2048 per query, has_invalid=False. Times are the minimum of 5 rounds × 15 iterations.

Max abs deviation against a torch reference is 0.000061 for both variants — this is purely a scheduling effect, not a numerics tradeoff.

Why it is not monotonic in num_queries

The switch tracks the number of CTAs, not the query count:

ctas = num_queries * num_splits * heads_blocks

num_splits comes from _decode_num_splits, so e.g. 32 queries can launch fewer CTAs than 34. Sorting the same data by ctas makes it monotonic, and the threshold lands exactly on the CU count (get_num_sms() = 256 here):

BLOCK_K = 32 if ctas > get_num_sms() else 64

This predicts 21 of 21 measured load points correctly. Intuition: while every CTA gets its own CU, the larger block (4 warps, 64 KB LDS) does more work per CU; once multiple waves are needed, the smaller block (2 warps, 32 KB) packs better and hides the gather latency.

Suggested change

In the gfx950 driver, replace the constant with the rule above. Wired in, it stays within ±0.6 % of the per-point optimum from 24 queries upward:

Note that num_splits must be computed before BLOCK_K (it is currently derived after), and get_num_sms() should be hoisted to module scope — at 62 µs kernel time a per-call lookup is measurable in microbenchmarks, though not in production where the decode runs inside a CUDA graph.

Things that did not help (same setup)

• GSPT (bytes per thread in the gather layout): 16 is optimal; 8 costs 14 %, 32 costs 45 % at K=64.
• waves_per_eu 1 / 2 / 4: −4 % / −4 % / −69 %.
• Cache modifier .cs instead of .cg: does not compile.
• Forcing kv_splits: the automatic choice is already optimal at 64+ queries; more splits cost up to 45 %.

Happy to open a PR if the rule looks right to you — and glad to re-measure on other shapes (different head counts, top-k, or a bf16 cache) if that would help.

Start in attention/pa_decode_sparse.py and inspect the gfx950 gluon driver, especially _decode_num_splits and the current BLOCK_K selection. Verify the CTA-based choice using get_num_sms() across the measured query sizes, then compare timings and numerical output against the reported torch reference. Done means the load-dependent selection reproduces the reported performance without changing numerics.