ggml-hrx: kernels for ZAYA's router and grouped conv; ZAYA1-8B decodes entirely on HRX (25.5 → 47.9 tok/s)

[bong-water-water-bong]

ZAYA1-8B decodes almost twice as fast on HRX0: its decode graph now runs entirely on the GPU (641 graph splits per token → 1).

Why it was slow. HRX sent seven ops in every ZAYA layer to the CPU, so each decode token made about three GPU↔CPU round trips per layer. The loader also left the 40 grouped-conv weights on the CPU, because HRX had no batched F16 matmul.

New HRX kernels (loom, with dispatch registrations; each matcher only claims cases the existing kernels do not take):

Kernel	Covers
ggml_grouped_mul_mat_f16_f32	MUL_MAT with a batched F16 weight: ZAYA's grouped conv, 10 groups of 128×128, per tap
ggml_softmax_rows_f32	SOFT_MAX without mask, scale 1 (the router's 17 slots)
ggml_sum_rows_f32	SUM_ROWS
ggml_argsort_rows_f32	ARGSORT for rows ≤ 1024, one rank per element, ties by index
ggml_get_rows_small_f32	GET_ROWS for rows narrower than 4 floats, ids may be strided (top-k views)
ggml_copy_strided_f32	CONT of strided views; REPEAT that only broadcasts (zero strides)

test-backend-ops -b HRX0, against the CPU:

• MUL_MAT: batched F16 cases, including groups of 3, 8 and 1536;
• SOFT_MAX 10/10, SUM_ROWS 6/6, ARGSORT 48/48, GET_ROWS 17/17;
• CONT 2/2, CPY 53/53, REPEAT 5/5.

The AMD files get one-line hookups: CMakeLists, dispatch-common.cpp, and REPEAT in the declared ops.

ZAYA graph changes (src/models/zaya.cpp):

• One token per sequence: the conv input transpose is a reshape, not a copy.
• MEAN over the query-head group is SUM_ROWS + SCALE.
• Qpre/Kpre reshape the contiguous matmul outputs instead of copying them. With the whole block on HRX, the copy made 512-token prefill wrong: perplexity 71.8. It's right with a CPU split after the copy, or without the copy; decode was never affected. HRX gives the same wrong answer whether K is broadcast by REPEAT or by a broadcasting ADD, and serialized commands and forced syncs do not change it.

Measured on Strix Halo, ZAYA1-8B Q4_K_M, GGML_HRX_DISABLE_DISPATCH=decode_split (what 1bit serve --device hrx sets):

	before	after
HRX0 decode, llama-bench tg128	25.5 tok/s (engine docs)	47.9 tok/s
HRX0 decode, llama-server, 128 tokens (warm run)	23.2 / 24.9 / 24.0	34.7 / 44.1 / 45.3
HRX0 prefill, pp512	—	2137 tok/s
Vulkan0 decode, llama-server	77.1 / 82.7 / 83.3	83.8 / 83.1 / 83.6

Correctness:

Check	before	after
Wikitext perplexity, Vulkan0 -b 512 / -b 2048	21.5731 / 21.6692	21.5731 / 21.6692 (identical)
CPU, 1 and 4 sequences per ubatch (8 chunks)	25.1372	25.1372 (identical)
HRX0 perplexity	21.5518	21.6456
Teacher-forced top-1 vs transformers FP32, Q4_K_M: HRX0 / Vulkan0	69/96 / 70/96	69/96 / 70/96
F16 Vulkan0	95/96	95/96

• The HRX perplexity change (+0.4%) is the new kernels' rounding amplified by top-1 routing. With the three numeric kernels disabled the build reproduces the old 25.1265 (8 chunks) exactly. Each one alone moves it by under 0.5%, in both directions.
• test-llama-archs -a zaya now covers HRX0 too: OK (8.57e-06), and on Vulkan and CPU, all with save/reload.

Not in this PR: prefill at 512 tokens still has 161 splits on HRX, and the decode-split flash-attention kernel stays off (ggml-org#140/ggml-org#123).

🤖 Generated with Claude Code