Optimize ROCM Prefill and Decode - #18
Merged
Merged
Conversation
Tune MMQ and MMVQ scheduling, add routed expert kernels, share Q8_1 activation quantization, and add measured four-column kernels for quant formats that outperform the generic path on gfx1151. Assisted-by: OpenCode
Add gfx1151 paths for paired matmuls, SwiGLU quantization, expert aggregation, contiguous binary operations, transposed concat, normalization, and conversion operations. Assisted-by: OpenCode
Dequantize Q8 KV data directly in tiles, specialize GQA decode grouping, reduce D256 register pressure, and reuse masks across grouped heads on gfx1151. Assisted-by: OpenCode
Increase gfx1151 column parallelism, tile KDA inputs in LDS, preserve recurrent graph ordering, and schedule compatible Qwen projections together. Assisted-by: OpenCode
Use full MTP rollback checkpoints and retain speculative sequence state on the device to avoid unnecessary host transfers. Assisted-by: OpenCode
Permit backends to allocate quantized CONCAT outputs with row padding under the expanded allocation-size invariant. Assisted-by: OpenCode
Add exact-batch, quantized matmul, routed expert, Flash Attention, concat, normalization, and Gated DeltaNet regression coverage for the optimized gfx1151 paths. Assisted-by: OpenCode
Keep batch-one MMVQ on one wave for all quant types. Wider wave counts changed FP32 reduction order, causing non-identical logits for Q2-Q5 and deterministic decode divergence in the four-corpus correctness matrix. Assisted-by: OpenCode
|
Thanks for putting this together — the RDNA3.5 work here is nice, and the We had a Strix Halo box free tonight, so we ran an independent A/B to help fill in the Measurements section. Sharing the raw numbers in case they're useful — happy to re-run anything with different settings. SetupAll four trees were built and run in the same session, back to back, with nothing else on the GPU. Prompt processing (tok/s, ± stddev over 3 reps)
Token generation (tok/s)
Correctness
We also read through the FA Observations
Two small questionsBoth look unrelated to prefill and we didn't exercise them, so flagging rather than asserting anything:
What we did not verify: speculative/MTP decode paths, the server changes, Vulkan, any non-gfx1151 device, and perplexity (we ran op-level correctness only, not end-to-end output quality). Disclosure: these benchmarks were run and this comment was written by Claude Code, on hardware and at the request of @dzannotti. Numbers are machine-generated but reproducible with the command line above; happy to share the raw JSON. |
Use 32 waves for the production non-KDA H32/S128 Gated DeltaNet shape on gfx1151. This reduces kernel latency by 1.34% and improves Qwen3.6 IQ3_XXS PP2048 by up to 0.97% while preserving exact four-corpus logits and decode output. Assisted-by: OpenCode
With 64 KiB of LDS per CU only 1-2 MMQ blocks are resident on gfx1151, so the global-load latency of every 256-wide K iteration is exposed (PMC: VALU busy 6%, 53% of wave-cycles waiting on data). Stage the next iteration's activation tile (and, for Q8_0, the weight tile) in registers while the current one is consumed. Same loads, same LDS layout, same arithmetic: outputs are bit-identical to the previous kernels. Whitelisted to the specializations where the extra registers do not spill: Q8_0 J48/J128, Q6_K J32, Q5_K J32, Q4_K J48, IQ2_S/IQ3_XXS J128 (all others keep the original loop; no spill regression in any kernel). test-backend-ops perf, MUL_MAT_ID 256 experts top-8, 2048 tokens (includes the quantize/ids overhead): Q8_0 gate/up 2717 -> 2551 us, Q6_K 3600 -> 3481, Q5_K 2571 -> 2510, Q4_K 2776 -> 2701, IQ2_S 3830 -> 3743, IQ3_XXS 3544 -> 3457. Also adds these production shapes to the perf test list.
Two exact-preserving changes to the gated delta net recurrence: * warp_reduce_sum compiled to 5 dependent ds_bpermute round trips per reduction (10 per token); on RDNA3/4 use v_permlanex16 + DPP row_xmask with the same xor-16/8/4/2/1 pairing, so results are bit-identical. * For the non-KDA S=128 prompt-processing path add a tiled kernel: q/k/g/ beta/v of a 16-token tile are loaded cooperatively into LDS and each warp owns 4 state columns, so k/q are read once per token for 4 columns and 4x fewer waves compete for issue slots (the kernel was issue-bound with ~50 VALU per token per wave for 4 state elements). FMA contractions are spelled out to match the original kernel's code generation. Bit-identical to the previous kernels on H=32/64, T=1/17/64/2048 (op-level dump comparison); test-backend-ops GATED_DELTA_NET 43/43. gfx1151, test-backend-ops perf: H32 T2048 5029 -> 2350 us, H64 T2048 9289 -> 4533 us, H32 T64 158 -> 93 us, T=1 13.3 -> 12.6 us. Together with the MMQ prefetch, Qwen3.6-35B-A3B PP2048 (llama-bench -b 2048 -ub 2048): Q8_0 +10.9%, UD-Q6_K_XL +9.4%, UD-Q4_K_XL +9.2%, UD-IQ3_XXS +7.6%.
The fused weighted expert sum used a volatile product to prevent FMA contraction and preserve the unfused graph's rounding. On AMD this spills every product to scratch. Use explicit v_mul_f32/v_add_f32 instructions to retain the exact operation order without scratch traffic, and process four adjacent channels per thread to amortize token indexing and weight loads. The fusion is restricted to contiguous 2048-wide Qwen tensors, so all float4 accesses are aligned. A direct old/new kernel comparison across 1,048,576 random outputs is bit-identical; MUL and ADD backend tests pass. gfx1151 weighted_expert_sum_f32<8>: 813 -> 689 us (-15.3%), scratch 8 -> 0 bytes. Qwen3.6-35B-A3B PP2048 improves by 0.35-1.36% across Q8_0, Q6_K_XL, Q4_K_XL, and IQ3_XXS.
… RDNA3.5 Token generation on gfx1151 launches ~1480 kernels per token with a ~2 us gap between each, so the quantize_q8_1 kernel in front of every MMVQ costs about as much as the quantization itself. Add mul_mat_vec_q_fq for the single-column Q8_0 path: a block of 16 waves (one output row per wave) quantizes the activation vector once into shared memory, replicating quantize_q8_1 exactly (same amax, IEEE division and roundf), and then runs the per-wave K loop of mul_mat_vec_q in the same order against it, with two K iterations of weight loads issued before the quantization and barrier. The two-kernel path is kept for other types, batched decode, and unaligned activations. Dense and MUL_MAT_ID (fused SWIGLU and down) outputs at all Qwen3.6 decode shapes are byte-identical to the previous path; MUL_MAT and MUL_MAT_ID backend tests pass. 16 waves per block and a prefetch depth of 2 were the fastest of the measured variants (8/32 waves, depth 1/4/8/16). Qwen3.6-35B-A3B TG128 on gfx1151: Q8_0 48.83 -> 50.99 t/s (+4.4%) at depth 0 and 42.88 -> 44.47 (+3.7%) at 32K; UD-Q8_K_XL 48.04 -> 49.61 (+3.3%) and 42.24 -> 43.38 (+2.7%). Adds Qwen3.6 batch-1 perf cases to test-backend-ops. Assisted-by: OpenCode
For a single decode token, every linear-attention layer runs 14 graph nodes between the qkv projection and the output gate: concat with the conv state, conv state shift into the cache, state gather, causal conv, SiLU, two L2 norms, alpha bias/softplus/A gate, beta sigmoid, the recurrence, the state copy into the cache, and the gated RMS norm. On gfx1151 these are 11 small kernels (~32 us) separated by ~2 us launch gaps, about 50 us per layer. Match this chain in the CUDA backend and run it as two kernels: one block per value head computes the conv + SiLU of its q/k/v channels (shifting the conv state straight into the cache), the L2 norms, gate and beta, and the recurrence with the state read from and written to the cache in place; a second tiny kernel applies the gated RMS norm (its destination may alias inputs of the first kernel that other blocks are still reading). Every op is spelled with the operation order and FMA contraction of the kernel it replaces, so the outputs, conv state and recurrent state are byte-identical to the unfused graph (checked in place and across cache rows). The fusion requires n_tokens == n_seqs == 1 and S == 128 and falls back otherwise, including the first token of a sequence where the state reset sits inside the chain. Qwen3.6-35B-A3B TG128 on gfx1151 (with the fused activation quantization): Q8_0 48.74 -> 54.18 t/s (+11.2%) at depth 0 and 42.62 -> 46.62 (+9.4%) at 32K; UD-Q8_K_XL 48.02 -> 52.69 (+9.7%) and 42.23 -> 45.71 (+8.2%). Assisted-by: OpenCode
ggml_cuda_op_weighted_expert_sum always staged its result in a pool buffer and copied it into the destination, because the graph allocator places the output over the expert or routing-weight buffers that die inside the fused range (in Qwen3.6 decode the eight routing weights sit inside the output). That copy is one more launch per layer in a decode step that is dominated by launch gaps. When the whole sum fits in one block (n_embd/4 * n_tokens <= 1024, i.e. decode), launch a single block and place a barrier between the reads and the store, so the output can be written in place whatever it aliases. The per-thread arithmetic is unchanged; larger grids keep the staging copy. Greedy 120-token generations with Qwen3.6-35B-A3B are identical. Qwen3.6-35B-A3B UD-Q8_K_XL TG128 on gfx1151: 52.69 -> 52.91 t/s. Assisted-by: OpenCode
The shared expert output is gated by sigmoid(w . y) with a 2048-wide f32 weight vector: a one-row mul_mat_vec_f, a sigmoid and the already fused mul + add + add, i.e. three launches per layer in decode. Match the five nodes and run them in one single-block kernel: the dot product repeats the per-thread fma chain and two-level warp reduction of mul_mat_vec_f<256>, the sigmoid is the same expression, and the elementwise tail uses explicit v_mul_f32/v_add_f32 like shared_mul_add_f32 (plain __fmul_rn/__fadd_rn are contracted into an fma by the HIP compiler). One block reads the whole activation vector before the reduction barrier, so the output may alias it. Outputs are byte-identical to the unfused graph. Qwen3.6-35B-A3B UD-Q8_K_XL TG128 on gfx1151: 52.91 -> 53.42 t/s. Assisted-by: OpenCode
mul_mat_vec_f loops over K with a runtime trip count, so the compiler keeps one 8-byte load per lane in flight and waits for it every iteration. The f32 MoE router (2048 x 256) of Qwen3.6 ran at 131 GB/s on gfx1151 because of this. For the single-column, unfused f32 path on HIP, load four K iterations of weights and activations before consuming them, in the same accumulation order; the remainder loop is unchanged. Outputs are byte-identical. gfx1151 Qwen3.6-35B-A3B decode: router matvec 16.0 -> 10.5 us, alpha/beta projections 3.2 -> 2.6 us; UD-Q8_K_XL TG128 53.42 -> 54.16 t/s. Assisted-by: OpenCode
mul_mat_vec_f picks a 256-thread block for K=512, so every wave issues a single 4-byte bf16 load and then reduces; on gfx1151 the bf16 MoE down projections of Unsloth UD quants (8 x 2048 rows of K=512) run at ~145 GB/s. Add a wave-per-row kernel that emulates the 256-thread block: lane l accumulates the eight virtual warps' lane l in the same fma order, reduces each with the same butterfly and combines them exactly like block_reduce over lanes 0..7 (including the zero-lane adds), with eight loads in flight per lane. Used on RDNA for unfused single-column bf16 with K <= 512; for longer rows the lower wave count loses, so the standard kernel stays. Outputs are byte-identical (dense and MUL_MAT_ID, odd shapes included). gfx1151 Qwen3.6-35B-A3B UD-Q8_K_XL: bf16 down projection 93 -> 67 us, TG128 54.16 -> 54.54 t/s. Assisted-by: OpenCode
Extend mul_mat_vec_q_fq beyond Q8_0 with a generic path that runs the per-wave vec_dot loop of mul_mat_vec_q<type, 1> against the shared-memory activations, and replicate the Q8_1 block sum of quantize_q8_1 in the exact butterfly order for types that read it. Q6_K (scale only) is enabled: it is byte-identical (dense and MUL_MAT_ID) and measured +0.8% on Qwen3.6-35B-A3B UD-Q6_K_XL TG128 (56.64 -> 57.09 t/s). Q4_K/Q5_K are exact as well but neutral without weight prefetch, so they stay on the two-kernel path. Assisted-by: OpenCode
…NA3.5 For mul_mat_id the MMQ tile width J is chosen from ncols_max, which is the total number of token-expert pairs, so the generic path always picks J=128. With 256 experts and 8 active the per-expert row count is only n_tokens/32 (16/32/64 rows for ubatch 512/1024/2048), so 50-87% of every tile was padding. The RDNA3.5 specializations that fixed this were gated on exact shapes and on 32/64 rows per expert, i.e. only ubatch 1024/2048, and only Q8_0/Q6_K used the routed-compact kernel. Replace the shape-specific paths with a per-type rule on rows per expert, and route all of them through mul_mat_q_routed_compact: Q8_0: 16/32/64 rows -> J=48 Q4_K/Q5_K/Q6_K: 16 rows -> J=32, 32 rows -> J=48, 64 rows -> J=32 Values measured on gfx1151 for the 2048x512 and 3072x1024 expert shapes; the previous choices (Q8_0 J48, Q6_K J32 at 64 rows) are unchanged, Q4_K at 64 rows moves from tiled J48 to compact J32 (-7% up-proj, -31% down-proj), Q5_K at 64 rows from tiled J32 to compact J32 (-31% down-proj). MUL_MAT_ID kernel time, 256 experts, 512x2048 / 2048x512, n=512 tokens: Q8_0 2347/2490 us -> 1517/1518 us Q6_K 3918/3931 us -> 1344/1300 us Q5_K 2606/2607 us -> 1087/1084 us Q4_K 2561/2538 us -> 1026/ 991 us n=1024: Q6_K -54%, Q5_K -50%, Q4_K -54% (Q8_0 unchanged). Qwen3.6-35B-A3B on gfx1151, llama-bench -b 2048 -ub 2048: UD-Q8_K_XL PP512 1437 -> 1722 t/s (+19.8%); PP1024/2048/4096 unchanged UD-Q4_K_XL PP512 1439 -> 2086 t/s (+44.9%), PP1024 +25.5%, PP2048 +3.5% UD-Q8_K_XL default ubatch 512: PP512 +22.1%, PP2048 +22.4% TG unchanged. test-backend-ops MUL_MAT_ID: 906/906 passed (adds 16/32 rows-per-expert eval cases and qwen3.6 perf cases).
The previous rule only matched exactly 16/32/64 rows per expert (ubatch
512/1024/2048 with 8 active experts); any other ubatch size, and the
remainder ubatch of every prompt, still fell back to the generic J=128 tiled
grid, which launches ncols_max/J*n_experts mostly-empty blocks.
Use ranges instead, and allow the routed-compact kernel for every selected J:
Q8_0: rows <= 12 -> J16, <= 64 -> J48, else J128
Q4_K/Q5_K/Q6_K: rows <= 12 -> J16, <= 24 -> J32, <= 48 -> J48,
<= 64 -> J32, else J128
The 16/32/64 choices are unchanged from the previous commit. Ranges were
picked from a J sweep over 64..4096 tokens on gfx1151 and validated
end-to-end.
Qwen3.6-35B-A3B on gfx1151, -b 2048 -ub 2048 unless noted, vs cutoff:
PP128 PP256 PP384 PP768 PP1536 PP4096@ub4096
UD-Q8_K_XL +39.7% +61.1% +21.5% +17.2% +12.4% +3.9%
UD-Q4_K_XL +93.5% +75.9% +60.1% +39.1% +18.4% +2.4%
(PP512/1024/2048 as in the previous commit.)
test-backend-ops MUL_MAT_ID: 914/914 passed (adds 4- and 128-row cases).
On gfx1151 the D=256 flash-attention path always fell back to the tile kernel, which is FMA-bound (~8 TFLOPS) and dominates prefill time at deep context (68% of PP time at 32k, 81% at 64k for Qwen3.5-0.8B). - Enable AMD WMMA device code for DKQ=256 (previously DKQ>128 was excluded). - Add a RDNA ncols=64 D=256 config tuned for the 256-VGPR budget: nwarps=8, nstages=1, nbatch_fa=32, nbatch_K2=128, nbatch_V2=64, nbatch_combine=64. Naive settings spill heavily and lose to the tile kernel. - Dispatch to MMA only when Q->ne[1]*gqa_ratio_eff > 32; the tile kernel is still faster for small batches (nb 4/8). - Fix a latent tile_Q reuse race in the combine loop: when DV/2 > nbatch_combine, warps read tile_Q columns written by other warps and the buffer is rewritten by the next batch without a barrier. Qwen3.5-0.8B UD-Q8_K_XL, gfx1151, PP512: depth 0: 9361 -> 9473 t/s (+1.2%) depth 32768: 3211 -> 4105 t/s (+27.9%) depth 65536: 1927 -> 2632 t/s (+36.6%) TG unchanged. test-backend-ops FLASH_ATTN_EXT: 2939/2939 passed.
Constrain architecture-specific paths, validate fused GDN use counts, narrow graph synchronization, remove disabled experiments, and strengthen fusion tests. Assisted-by: OpenCode
gaetan-puleo
marked this pull request as ready for review
September 5, 2026 10:00
gaetan-puleo
changed the title
gaetan-puleo
changed the title
Author
This was a fix for correctness on MTP Nathan uses it in the Vulkan Performance toolbox too |
|
ship this! this is awesome |
dzannotti
approved these changes
Sep 5, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds guarded ROCm/HIP inference optimizations for RDNA3.5 and Strix Halo systems, including the Ryzen AI Max+ 395 with Radeon 8060S graphics.
Main improvements:
MUL_MAT_IDexecution with quant- and geometry-specific tile selectionArchitecture, shape, layout, aliasing, and graph-closure guards preserve existing behavior on unsupported devices and tensor configurations.
Correctness
The complete PR was reviewed by five independent agents, followed by a separate source-level verification of every reported finding.
Confirmed issues were fixed, including:
Final validation:
Performance
Representative measured improvements include:
8704.46 -> 9561.11tok/s, +9.8%130.67 -> 131.85tok/s, +0.9%Benchmark Coverage
Benchmarked on gfx1151 with ROCm 7.14:
-ub 512-ub 2048-b 2048 -fa 1 -ngl 999 --load-mode mmap -r 4Interactive and raw results:
https://gaetan-puleo.github.io/strix-halo-pull-18-results/
The branch includes latest master
b212548e0.Final head:
756f0d443