Skip to content

HRX: decode-split multipass reduce output pass is O(blocks) per element, leaving a ~16% decode gap above 2048 #124

Description

@1bit-traffic-bot

The multi-pass reduction in 1bit-MONSTER/llama.cpp#13 restores flash_attention.decode_split above key_value_token_capacity 2048, but materially slower than the <=2048 path, so the decode cliff is only partly removed.

Measured (gfx1151, Qwen3-Coder-30B-A3B-Instruct-Q4_K_M, -dev HRX0, llama-bench -p 0 -n 8 -r 3):

depth capacity blocks t/s path
1470 1536 24 73.33 cooperative
1900 1920 30 64.65 cooperative
2000 2048 32 64.49 cooperative
2100 2112 33 53.88 multipass
4800 4864 76 37.04 multipass

The split now beats the fallback (d2100 53.88 vs 46.4 t/s; d4800 37.04 vs 32.5 t/s), so the old -32% cliff is reduced to about -16% at d2100 - but decode does not recover toward the <=2048 corridor.

Cause

reduce_completed.multipass (flash_attention_decode_split_f32_f16_wmma.loom) does its final pass as a serial loop per output element:

scf.for %output_tile = [%c0 to %output_tile_count step %c1] {
  ...
  %unnormalized_output = scf.for %block = [%c0 to %active_block_count step %c1] {
    %scale = view.load %partial_max_view[%key_value_head, %block, %query_row]   // per-block scale
    %block_output_f16 = view.load %partial_output_view[..., %block, %query_row, %output_channel]
    ...
  }
}

That is O(blocks) per output element and runs in all 256 workitems, whereas the max/sum passes above it lane-stride the block dimension (scf.for %block = [%lane to %active_block_count step %c64]). The per-block normalisation scale is already written back into partial_max, so the output pass could lane-stride the block dimension too (with a subgroup/workgroup reduction per channel) instead of iterating every block serially.

Impact

The goal for this work is that decode at 2100/3000/4800 "recovers toward the <=2048 tok/s (no sharp boundary cliff)". Today it recovers only about half of the loss. Restructuring the output pass is expected to close most of the gap, and it is also the prime suspect in the residual GPU fault filed separately.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions