Repository navigation
[AMD] Optimize Qwen3.5 MTP unified attention on gfx950 - #36330
Conversation
ec8ba3d to
2e742d3
Compare
Record in a module docstring which launch parameters this fork changes and why, so the next reader does not have to diff against ROCm/aiter to find out that block_m=32 is tied to the 16:1 GQA ratio rather than to head_dim 256. Derive num_queries_per_kv in the dispatch gate instead of pinning tp_q_head_num, and move the rationale above the condition.
9f9633f to
90d832d
Compare
… counts The dispatch gate already derived num_queries_per_kv, but paired it with an absolute tp_v_head_num == 1 while the wrapper asserted num_query_heads == 16 and num_kv_heads == 1. Only the tp_v_head_num clause kept that assert from firing, so a 16:1 layout with more than one kv head (Qwen3.5-397B-A17B at TP1: 32 q / 2 kv) silently fell back instead of taking the kernel. The kernel is kv-head generic: kv_head_idx = program_id(1) drives the Q head offset and both cache loads, the segment buffers stride on num_query_heads, and the descales are per-tensor scalars. What is fixed is the 16:1 ratio, via block_q = block_m // num_queries_per_kv = 2. Assert the ratio in the wrapper, and gate on tp_k_head_num == tp_v_head_num, which is the wrapper's real precondition -- K and V get separate views but the kernel has a single kv-head grid dim. Covered by a 32 q / 2 kv case in test_aiter_unified_attention_mtp; both it and the existing 16 q / 1 kv case match upstream unified_attention at 1e-2.
75257a6 to
f62ee5d
Compare
| is_gfx95_supported() | ||
| and 1 < self.forward_metadata.max_q_len <= 4 | ||
| and max_kv_len > 512 | ||
| and num_queries_per_kv == 16 |
There was a problem hiding this comment.
Hi, @yichiche
The backend uses num_queries_per_kv == 16, but the wrapper still has assert num_query_heads == 16 and num_kv_heads == 1.
Are these asserts necessary? Can this kernel handle kv heads > 1, or is it tp2-only?
Right now Qwen-3.5 tp1 (32 q_head / 2 kv_head) may enter the kernel path but hits the assert fail.
There was a problem hiding this comment.
Thanks. Comment resolved.
Changed in f62ee5d.
Extended the kernel usage to q/kv ratio = 16 cases. Also added test coverage to q/kv=32/2 situation.
There was a problem hiding this comment.
Update the assert to accommodate the general case where num_queries_per_kv == 16, either with num_query_heads = 32 and num_kv_heads = 2 (TP2) or num_query_heads = 16 and num_kv_heads = 1 (TP1):
assert num_query_heads % num_kv_heads == 0
assert num_queries_per_kv == 16
|
@amd-bot ci-status |
CI Status for PR #36330Merge verdict: Safe to merge from a correctness standpoint — this PR's changed code is exercised and passed. The new gfx950 kernel + its unit test ran on real MI35x hardware ( Note The PR's value (a gfx950-only Triton MTP kernel behind a strict shape gate in Warning PR CI is incomplete. Changed files: Executed CI failure attribution: AMD: 13 executed failures (0 related) · Others: 3 executed failures (0 related) · Extra pipelines: not run (gate-skipped, not counted) AMD Executed Failures
(Other failed small/large shards 1,7,8,10,12,13 share the same Other Executed Failures
Details / what to do before merge
Generated by amd-bot using Claude Code CLI |
The module holds a Qwen3.5 MTP-verify specialization of the Triton kernel_unified_attention_3d launch; name it after the kernel it specializes rather than after AITER.
…_qwen35 Match the module name, so the entry point reads as the qwen35 variant of unified_attention_3d_mtp.
The kernel is selected by shape, not by model, so no identifier should name one: unified_attention_3d_mtp_qwen35 -> unified_attention_3d_mtp_func, use_qwen35_mtp_kernel -> use_unified_attention_3d_mtp, and the Triton kernel/repr names lose their qwen35_mtp prefix. Qwen3.5-397B-A17B at TP1/TP2 being the only config that matches today is now stated in comments, where it belongs.
4 participants
Motivation
Qwen3.5 target verification on MI355X uses a short multi-token query (
max_seqlen_q <= 4) with 16 query heads, one KV head, head size 256, page size 16, and FP8 KV cache. AITER's generic 3D unified-attention configuration processes one query token per workgroup and reloads the same paged KV data for each speculative token, leaving this latency-sensitive path substantially slower than the comparison implementation.Modifications
BLOCK_M=32andBLOCK_Q=2, reusing each loaded KV tile across both tokens.TILE_SIZE=32and select 8-64 segments using both CU parallelism and per-segment KV work.Accuracy Tests
IL70000 / OL300: gsm8k
GPQA (
/raid/models/amd-Qwen3.5-397B-A17B-MXFP4, 198 questions, repeat 8)Benchmarking and Profiling
MI355X microbenchmark configuration: batch 20, query length 4, 16 query heads, one KV head, head size 256, page size 16, BF16 query, and FP8 KV cache. Values are kernel-path latency measured with
triton.testing.do_benchafter warmup.Decode-phase trace at input length 70,000. Before is
kernel_unified_attention_3d(TILE_SIZE=16,NUM_SEGMENTS=128); after isqwen35_mtp_unified_attention_3d(TILE_SIZE=32,NUM_SEGMENTS=64).The 51.5 us and 53.8 us saved per layer, across the 15 attention layers on this path, show up as a 933 us and 1,103 us drop in total decode time.
End-to-end serving on MI355X, TP2, input length 70,000 and output length 300 (
tok/s/gpuistotal tok/sdivided by 2):End-to-end latency, throughput, and TTFT improve at every concurrency level tested. Median TPOT regresses 5.4% at concurrency 16; the cause has not been isolated yet, though median end-to-end time still drops 2.4% there.
The same run reports 368.547 s total latency and 613.599 output token/s.
Checklist
Review Process
/tag-run-ci-label,/rerun-failed-ci,/tag-and-rerun-ciCI States
Latest PR Test (Base): ⏳ Run #33045741639
Latest PR Test (Extra): ❌ Run #33045741325
Latest PR Test (AMD ROCm 7.2): ⏳ Run #33045741649