mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-06-29 19:28:33 +00:00
d77f0bea63
After moving kBlockQ to runtime in the previous commit, the static
NumQPerKV in `variant_config<V>` and the runtime-vs-static assert in
the kernel became the only things still tying a compiled binary to a
specific num_queries_per_kv. Drop both and the existing instances now
serve every num_qpkv that divides kBlockM evenly.
Concretely:
* `variant_config<V>` -- remove the NumQPerKV field from every
specialization.
* `unified_attention_kernel_traits` -- remove the `num_queries_per_kv`
/ `kBlockQ = kBlockM / num_qpkv` derivation. The BlockTile's 2nd
entry (the static `kBlockQ` exposed via UnifiedAttentionShape) is
anchored at kBlockM so it describes the "num_qpkv == 1" fallback;
the actual kBlockQ is always the runtime value.
* `unified_attention_kernel_launch` -- recompute kBlockQ at host time
from `args.num_queries_per_kv` for the total_num_q_blocks math.
* `unified_attention_kernel.hpp` -- drop the
`assert(kBlockQ_dyn == kBlockQ)` (it enforced the very coupling we
just removed).
* `unified_attention.cpp::select_config` -- collapse the two
per-num_qpkv code paths into a single (head_dim, avg_rows,
max_rows) ladder, where avg_rows = avg_q * num_qpkv.
Variant renames (8 variants):
prefill_d128_mha -> prefill_d128
decode_d128_mha_m128 -> decode_d128_m128
decode_d128_mha_m32 -> decode_d128_m32
decode_d128_mha_m16 -> decode_d128_m16
prefill_d64_gqa8 -> prefill_d64
decode_d64_gqa8_m128 -> decode_d64_m128
decode_d64_gqa8_m64 -> decode_d64_m64
decode_d64_gqa8_m16 -> decode_d64_m16
The 16 d=64 instance files lose their `_gqa8` infix to match the
d=128 naming (file count unchanged: 16 dtypes x mask combos per
head_dim).
Validation:
* Correctness suite: 241/245 (same 4 pre-existing int32-overflow
failures in the prefill rebased-pointer path).
* d=128 GQA-8 (a NEW combo we never had a binary for) -- runs
correctly on the existing decode_d128_m* binaries with num_qpkv=8
at runtime. max abs diff <= 1e-2 vs the torch reference at ql in
{1, 4, 16}.
* d=64 MHA (also a new combo) -- runs correctly on the existing
decode_d64_m* binaries with num_qpkv=1. Same tolerance.
* Perf sweep (b=4..256, sk=120000, MI300):
d=64 GQA-8: speedups 1.28x..1.84x vs Triton (within 0.6%
of baseline).
d=128 MHA: speedups 0.98x..1.14x vs Triton (within 0.3%
of baseline).
Unlocked: adding new (head_dim, num_qpkv) combos no longer requires
new kernel binaries -- just a host-side heuristic update mapping the
combo to the appropriate (kBlockM, BlockWarps) ladder.
Co-authored-by: Cursor <cursoragent@cursor.com>