diff --git a/benchmarks/NOTE-hrx-124-multipass-output-2026-09-27.md b/benchmarks/NOTE-hrx-124-multipass-output-2026-09-27.md new file mode 100644 index 000000000000..2b34b42c68bc --- /dev/null +++ b/benchmarks/NOTE-hrx-124-multipass-output-2026-09-27.md @@ -0,0 +1,91 @@ +# HRX decode-split multipass output pass: drop the 128x redundant expf and the scalar block loop (engine#124) + +Worktree: `~/1bit-engine-176/third_party/llama.cpp`, branch `1bit/hrx-124-output-pass` (base `00adc2b`). +Box: gfx1151 (Strix Halo), `Qwen3-Coder-30B-A3B-Instruct-Q4_K_M`, `-dev HRX0`. + +## What was wrong + +`@ggml.flash_attention.decode_split.reduce_completed.multipass` +(`ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom`) +finished with a scalar output pass: `workitem -> one output channel`, then a serial +`scf.for %block = [%c0 to %active_block_count]` over **every** KV block, run by all 256 +workitems. Only 128 are live at `value_head_size = 128`, and each one recomputed +`scale = expf(partial_max[block] - maximum)` for every block, so `expf` ran once per +`(block, output element)` instead of once per `(block)`, and every element walked the +block dimension alone. That is the residual ~16% gap above capacity 2048 (issue #124). + +## The change (bit-exact) + +1. **LDS scale stage.** The lane-strided sum pass already computes + `scale = expf(partial_max[block] - maximum)` for every block; store it into a + per-row workgroup stage (`scale_stage_view[row, block]`) and also publish the row + `sum` (`sum_stage_view[row]`). The output pass then just reloads the exact f32 + value. Removes ~128x redundant `expf` and the redundant `partial_max` reloads. +2. **Vectorised all-rows output pass.** After the per-row max/sum phase, one pass + covers all query rows at once: workitem `w` owns a 4-channel *quad* + (`quad = tile*256 + w`, `row = quad / quads_per_row`, `channel = (quad % + quads_per_row)*4`), loads `vector<4xf16>` from `partial_output`, and accumulates + `vector<4xf32>` over blocks with `unroll(%c4) schedule(interleaved)`. For the + production GQA shape (8 query heads per KV head, `value_head_size = 128`) there are + exactly 256 quads, so the whole workgroup (all four subgroups) is busy and each + channel still accumulates its blocks in the same order as before. + +Because the per-channel accumulation order is unchanged and the vector ops are +elementwise, the reduce output is **bit-identical** to the previous multipass reducer +(verified end-to-end below). + +## Result + +`llama-bench -p 0 -n 8 -r 5 -d 1900,2000,2100,3000,4800`, 6 interleaved A/B rounds, +median of the 6 run means: + +| depth | capacity | blocks | base t/s | patched t/s | delta | path | +|---|---:|---:|---:|---:|---:|---| +| 1900 | 1920 | 30 | 70.72 | 71.48 | +1.1% | cooperative (untouched control) | +| 2000 | 2048 | 32 | 66.06 | 68.82 | +4.2% | cooperative (untouched control) | +| 2100 | 2112 | 33 | 58.19 | 67.15 | **+15.4%** | multipass | +| 3000 | 3008 | 47 | 50.84 | 60.89 | **+19.8%** | multipass | +| 4800 | 4864 | 76 | 41.83 | 51.02 | **+22.0%** | multipass | + +The boundary cliff is gone: `d2100/d2000` moves **0.881 -> 0.976** (and +`d4800/d2000` 0.633 -> 0.741). The two depths the reducer does not touch moved by ++1.1% / +4.2%, i.e. run-to-run noise. + +A throwaway diagnostic that deleted the output block loop entirely (keeping the +max/sum passes) measured **+21.6% / +50.0% / +30.5%** at d2100 / d3000 / d4800, which +is how the output pass was identified as the whole gap rather than the reduction math +(issue #124's diagnosis). + +## Evidence + +- **Bit-exact.** `llama-perplexity -c 2049 -b 1 -f <6600-token text> --save-all-logits` + (token-by-token decode, so the decode-split kernel is dispatched; `GGML_HRX_LOG_DISPATCH=1` + confirms `flash_attention_decode_split`). The 1,244,725,284-byte logits dumps for base + and patched are **byte-identical** (`cmp`), PPL 21.9760 +/- 0.81813 on both. +- **Correctness.** Buried code word `ZX-4718-QQ` at 4700 prompt tokens (capacity 4864, + multipass), `llama-server`, temperature 0, seed 42, `cache_prompt:false` -> answer + `ZX-4718-QQ`. PASS. +- **Faults.** 6 A/B rounds x 5-rep sweeps at all five depths on the patched build: + 0 `HSA_STATUS_ERROR_MEMORY_FAULT`, 0 `res = -3`. The cooperative/direct paths are + untouched, so <= 2048 (d1900/d2000) shows no tok/s or fault regression. + +## Reproduce + +``` +# build (engine worktree ~/1bit-engine-176, HRX llama.cpp build dir) +cmake --build build/hrx/llama --target llama-bench llama-perplexity -j 16 + +# dispatch needs libhsa from the HRX toolchain +export IREE_HAL_AMDGPU_LIBHSA_PATH=/opt/rocm-therock/lib/python3.14/site-packages/_rocm_sdk_core/lib/libhsa-runtime64.so.1 + +# performance +build/hrx/llama/bin/llama-bench -m models/Qwen3-Coder-30B-A3B-Instruct-Q4_K_M.gguf \ + -dev HRX0 -ngl 99 -p 0 -n 8 -r 5 -d 1900,2000,2100,3000,4800 + +# bit-exact reference vs the previous build +build/hrx/llama/bin/llama-perplexity -m models/Qwen3-Coder-30B-A3B-Instruct-Q4_K_M.gguf \ + -f ref-text.txt -c 2049 -b 1 -dev HRX0 -ngl 99 --save-all-logits logits.bin +``` + +Raw runs and the A/B driver used here live under the session scratch dir +(`ab5-*.json`, `ab5-summary.txt`, `diag-noloop.json`). diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom index 526b231b10f4..4f55c76bc658 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom @@ -855,8 +855,10 @@ template.def<@ggml.flash_attention.decode_split.reduce_completed.multipass> devi %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index - %padded_value_head_size = index.add %value_head_size, %c255 : index - %output_tile_count = index.div %padded_value_head_size, %c256 : index + %quads_per_row = index.div %value_head_size, %c4 : index + %total_quad_count = index.mul %query_heads_per_key_value_head, %quads_per_row : index + %padded_quad_count = index.add %total_quad_count, %c255 : index + %output_tile_count = index.div %padded_quad_count, %c256 : index %is_first_subgroup = index.cmp eq, %subgroup, %c0 : index %workitem_is_zero = index.cmp eq, %workitem, %c0 : index %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias, %output_noalias = buffer.assume.noalias %partial_max, %partial_sum, %partial_output, %output : buffer, buffer, buffer, buffer @@ -870,6 +872,18 @@ template.def<@ggml.flash_attention.decode_split.reduce_completed.multipass> devi %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x[%value_head_size]xf32> %reduction_stage = buffer.alloca align(8) %reduction_stage_bytes : buffer %reduction_stage_view = buffer.view %reduction_stage[%c0_offset] : buffer -> view<2xf32> + // Per-row, per-block normalisation scales and per-row sums, computed once by + // the lane-strided max/sum passes and reused by the vectorised output pass, + // so expf runs once per (row, block) instead of once per output element. The + // per-element accumulation order over blocks is unchanged. + %f32_bytes = index.constant 4 : offset + %scale_stage_elements = index.mul %query_heads_per_key_value_head, %active_block_count : index + %scale_stage_bytes = index.scale %scale_stage_elements, %f32_bytes : index, offset -> offset + %scale_stage = buffer.alloca align(16) %scale_stage_bytes : buffer + %scale_stage_view = buffer.view %scale_stage[%c0_offset] : buffer -> view<[%query_heads_per_key_value_head]x[%active_block_count]xf32> + %sum_stage_bytes = index.scale %query_heads_per_key_value_head, %f32_bytes : index, offset -> offset + %sum_stage = buffer.alloca align(16) %sum_stage_bytes : buffer + %sum_stage_view = buffer.view %sum_stage[%c0_offset] : buffer -> view<[%query_heads_per_key_value_head]xf32> scf.for %query_row = [%c0 to %query_heads_per_key_value_head step %c1] { %query_head = index.add %query_head_base, %query_row : index %lane_maximum = scf.if %is_first_subgroup -> (f32) { @@ -893,6 +907,7 @@ template.def<@ggml.flash_attention.decode_split.reduce_completed.multipass> devi %block_maximum = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 %delta = scalar.subf %block_maximum, %maximum : f32 %scale = scalar.expf %delta : f32 + view.store %scale, %scale_stage_view[%query_row, %block] : f32, view<[%query_heads_per_key_value_head]x[%active_block_count]xf32> %block_sum = view.load %partial_sum_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 %scaled_sum = scalar.mulf %block_sum, %scale : f32 %next_sum = scalar.addf %running_sum, %scaled_sum : f32 @@ -908,26 +923,39 @@ template.def<@ggml.flash_attention.decode_split.reduce_completed.multipass> devi } kernel.barrier scope(workgroup) ordering(acq_rel) %sum = view.load %reduction_stage_view[%c1] : view<2xf32> -> f32 - scf.for %output_tile = [%c0 to %output_tile_count step %c1] { - %tile_channel = index.mul %output_tile, %c256 : index - %output_channel = index.add %tile_channel, %workitem : index - %output_channel_valid = index.cmp ult, %output_channel, %value_head_size : index - scf.if %output_channel_valid { - %unnormalized_output = scf.for %block = [%c0 to %active_block_count step %c1](%running_output = %c0_f32 : f32) -> (f32) unroll(%c4) schedule(interleaved) { - %block_maximum = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 - %delta = scalar.subf %block_maximum, %maximum : f32 - %scale = scalar.expf %delta : f32 - %block_output_f16 = view.load %partial_output_view[%key_value_head, %block, %query_row, %output_channel] : view<[%key_value_head_count]x[%partial_block_capacity]x16x[%value_head_size]xf16> -> f16 - %block_output = scalar.extf %block_output_f16 : f16 to f32 - %scaled_output = scalar.mulf %scale, %block_output : f32 - %next_output = scalar.addf %running_output, %scaled_output : f32 - scf.yield %next_output : f32 - } - %normalized_output = scalar.divf %unnormalized_output, %sum : f32 - view.store %normalized_output, %output_view[%query_head, %output_channel] : f32, view<[%query_head_count]x[%value_head_size]xf32> + view.store %sum, %sum_stage_view[%query_row] : f32, view<[%query_heads_per_key_value_head]xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + } + // Vectorised output pass: one 4-channel quad per workitem, spanning every + // query row so all four subgroups stay busy. Each channel accumulates its + // blocks in the same order as the scalar pass, so the values are identical. + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + scf.for %output_tile = [%c0 to %output_tile_count step %c1] { + %output_tile_base = index.mul %output_tile, %c256 : index + %quad = index.add %output_tile_base, %workitem : index + %quad_valid = index.cmp ult, %quad, %total_quad_count : index + scf.if %quad_valid { + %output_row0 = index.div %quad, %quads_per_row : index + %output_row = index.assume %output_row0 [range(%output_row0, 0, 16)] : index + %output_row_base = index.mul %output_row, %quads_per_row : index + %quad_in_row = index.sub %quad, %output_row_base : index + %output_channel0 = index.mul %quad_in_row, %c4 : index + %output_channel = index.assume %output_channel0 [range(%output_channel0, 0, 512)] : index + %query_head = index.add %query_head_base, %output_row : index + %row_sum = view.load %sum_stage_view[%output_row] : view<[%query_heads_per_key_value_head]xf32> -> f32 + %row_sum_vector = vector.splat %row_sum : vector<4xf32> + %unnormalized_output = scf.for %block = [%c0 to %active_block_count step %c1](%running_output = %c0_f32x4 : vector<4xf32>) -> (vector<4xf32>) unroll(%c4) schedule(interleaved) { + %scale = view.load %scale_stage_view[%output_row, %block] : view<[%query_heads_per_key_value_head]x[%active_block_count]xf32> -> f32 + %scale_vector = vector.splat %scale : vector<4xf32> + %block_output_f16 = vector.load %partial_output_view[%key_value_head, %block, %output_row, %output_channel] : view<[%key_value_head_count]x[%partial_block_capacity]x16x[%value_head_size]xf16> -> vector<4xf16> + %block_output = vector.extf %block_output_f16 : vector<4xf16> to vector<4xf32> + %scaled_output = vector.mulf %scale_vector, %block_output : vector<4xf32> + %next_output = vector.addf %running_output, %scaled_output : vector<4xf32> + scf.yield %next_output : vector<4xf32> } + %normalized_output = vector.divf %unnormalized_output, %row_sum_vector : vector<4xf32> + vector.store %normalized_output, %output_view[%query_head, %output_channel] : vector<4xf32>, view<[%query_head_count]x[%value_head_size]xf32> } - kernel.barrier scope(workgroup) ordering(acq_rel) } template.return }