[bong-water-water-bong]

Fork half of 1bit-MONSTER/engine#1926. An engine PR that moves the pin will carry Fixes #1926 once this merges.

The fault

A one-token decode fails on any model whose attention head size is not a power of two and whose Q weights are not Q4_K/Q6_K:

HRX Loom JIT Loom compilation failed
compile_kernel: compile gfx1151|679a0aa7…|llm_attention_q_matmul_rope_decode_f32_f32|recipe=direct|token_count=1|…|llm.attention_qkv.head_count=16|llm.attention_qkv.head_size=160|…
llama_decode: failed to decode, ret = -3

The runtime hides the compiler's reason. Compiling the same sources offline (loom-link --mode=merge, then loom-compile --target=amdgpu:gfx1151 --format=amdgpu-hsaco with the config keys from the line above) prints it:

ops/llm_attention_q_matmul_rope_decode_f32_f32.loom:55:5: error [SUBRANGE/010]: vector.store full-vector footprint upper bound
  is not proven on view axis 2 (vector axis 0); origin %78 must satisfy origin + vector_extent <= view_bound for view %88;
  vector_extent is 2, view_bound is 160, and the maximum legal origin is 158

The kernel stores each rotated pair, two values wide, at head_channel = channel rem head_size. channel is always the even row of a pair, so the store is in bounds, but the compiler can only bound the remainder by head_size - 2 when the head size is a power of two.

head_size	before	after
64, 128, 256	compiles	compiles
48, 80, 96, 160	SUBRANGE/010	compiles

The weight format does not matter (format 4 at 160 fails the same way). Q4_K/Q6_K models did not hit this because the matcher sends them to the packed low-token kernel when output_size % 64 == 0.

The fix

Index by rotation pair instead of by channel: pair = channel / 2, head = pair / (head_size / 2), theta_channel = pair rem (head_size / 2), head_channel = theta_channel * 2. Same addresses for every even channel; the remainder is now bounded by head_size / 2, so the two-wide store is provably in bounds at every head size. The K and V decode kernels store on the output_size axis and are unchanged.

Proof (strixhalo, gfx1151, power mode performance, under gpu-hold)

New test tests/test-hrx-attention-q-decode-head-size.cpp: MUL_MAT + head layout + ROPE at one token on HRX0, must plan on this kernel, compared with the CPU backend (normalized MSE ≤ 5e-4). It skips without an HRX0 device, like the other device tests, so CI without a GPU does not exercise it.

Before (kernel as at 2d62cc7af, same test binary otherwise): 15 cases, 9 failures

bf16  head_size= 48 heads= 6  planned, but the compute failed
bf16  head_size= 96 heads= 8  planned, but the compute failed
bf16  head_size=160 heads=16  planned, but the compute failed
bf16  head_size= 64 heads=16  loom_libs:llm_attention_q_matmul_rope_decode_f32_f32 nmse=2e-06 OK
bf16  head_size=128 heads= 8  loom_libs:llm_attention_q_matmul_rope_decode_f32_f32 nmse=2.37e-06 OK
…(f16 and q8_0: the same three head sizes fail)

After: 15 cases, 0 failures

bf16  head_size= 48 heads= 6  loom_libs:llm_attention_q_matmul_rope_decode_f32_f32 nmse=2.46e-06 OK
bf16  head_size= 96 heads= 8  loom_libs:llm_attention_q_matmul_rope_decode_f32_f32 nmse=2.02e-06 OK
bf16  head_size=160 heads=16  loom_libs:llm_attention_q_matmul_rope_decode_f32_f32 nmse=2.15e-06 OK
bf16  head_size= 64 heads=16  loom_libs:llm_attention_q_matmul_rope_decode_f32_f32 nmse=2e-06 OK
bf16  head_size=128 heads= 8  loom_libs:llm_attention_q_matmul_rope_decode_f32_f32 nmse=2.37e-06 OK
f16   head_size= 48 heads= 6  … nmse=3.31e-08 OK
f16   head_size= 96 heads= 8  … nmse=3.36e-08 OK
f16   head_size=160 heads=16  … nmse=3.53e-08 OK
f16   head_size= 64 heads=16  … nmse=3.31e-08 OK
f16   head_size=128 heads= 8  … nmse=3.38e-08 OK
q8_0  head_size= 48 heads= 6  … nmse=1.64e-05 OK
q8_0  head_size= 96 heads= 8  … nmse=1.38e-05 OK
q8_0  head_size=160 heads=16  … nmse=1.35e-05 OK
q8_0  head_size= 64 heads=16  … nmse=1.57e-05 OK
q8_0  head_size=128 heads= 8  … nmse=1.4e-05 OK

The 64 and 128 controls give the same error before and after, and the new head sizes sit at the controls' level.

The two models from the issue, llama-bench --device HRX0 -ngl 99 -p 512 -n 128 -r 3 with the no-SVM trio, one run each:

model	head_size	before	after: pp512 t/s (sd)	after: tg128 t/s (sd)
stories15M Q8_0	48	decode fails, ret = -3	58,882.86 (5,845.70)	545.39 (35.65)
Sarashina2.2-3B bf16	160	decode fails, ret = -3	518.72 (12.18)	26.36 (0.22)

Limits of this proof

• Built with hrx: read the host term before classifying it, so a clang build opens the device again (engine #1942) #279 on top. At the tip a clang build registers no HRX0 device (engine#1942), so every device run above used a local branch of this commit plus hrx: read the host term before classifying it, so a clang build opens the device again (engine #1942) #279's commit. This PR does not contain hrx: read the host term before classifying it, so a clang build opens the device again (engine #1942) #279.
• Greedy text is not identical to the CPU's. stories15M, 96 tokens at --temp 0: HRX and CPU agree for the first 167 bytes, then diverge into different coherent sentences. I did not compare logits per step; the per-kernel comparison above is the numeric evidence.
• One bench run per model, no repeat, and the pp512 figure for Sarashina differs from the 31.11 t/s in the issue. I did not investigate that; it is a different build.
• Head sizes 80 and 256 are covered by the offline compile only, not by the device test.

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at f33978d

[bong-water-water-bong]

VERDICT: READY at f33978d

The bug is a power-of-two assumption in an index expression, and the fix is the weaker, correct constraint:

// %channel is the even row of a rotation pair, so index by pair: the remainder is then bounded by
// head_size / 2 and the two-wide store below is in bounds for every head size, not only for a power of two
// (engine #1926: 48, 96 and 160 did not compile).
%pair_channel = index.div %channel, %c2 : index

Dividing by two before the modulo is what makes the remainder bounded by head_size/2 for any even head size — the old form needed head_size to divide
the channel range evenly, which only holds for powers of two. That is a one-line change with a stated invariant instead of a special case, which is the shape I
want for this class.

And the arm tests exactly the sizes that failed, plus the control:

cases.push_back({ type,  48,  6 });   // stories15M
cases.push_back({ type,  96,  8 });
cases.push_back({ type, 160, 16 });   // Sarashina2.2-3B
cases.push_back({ type, 128,  8 });   // control

Three real models with their own head counts and the previously-working 128 as the control — so the arm proves the fix at the boundaries and that the
power-of-two case still behaves. Nine REQUIREs, with the dispatch resolution, an imported buffer and a per-case NMSE check (nmse=%.3g), so the comparison is
numeric rather than "it compiled".

That last point is the important one for this bug: 48/96/160 did not compile, so the cheapest possible test would be "it builds now". The PR instead
runs the kernel and compares against a reference with a printed NMSE, which is the difference between a compile fix and a correctness claim.

Not built: no llama.cpp build tree here and -DGGML_HRX=ON needs the private runtime, so this verdict rests on the diff, the invariant, and the arm's shape —
not on a run. Labelled rather than implied.

One note: the same power-of-two assumption is worth a grep across the other kernel-index expressions (% head_size, & (head_size-1) shapes); if it hid here
it can hide in a formula that only runs for a model nobody has loaded yet.

[bong-water-water-bong]

VERDICT at f33978d: READY

A fused Q decode kernel that only compiles at some head sizes is a kernel that fails at configuration time for the models nobody tested — the worst place to find out, because it looks like a build break rather than a missing instantiation. Making it compile at every head size, with the sizes handled uniformly, removes a class of failure rather than one instance.

CI green. I read the diff and CI; I did not compile this fork.