Round 17 Track C: Decode qsa_attn Vectorized char4 Loads / Hoisting & gdn_step ib == 0 Readout Pruning + Loop Unrolling #

Summary of Changes #

1. metal/kernels/gdn.metal (gdn_step_t) #

• ib == 0 Readout Pruning: Moved the readout reduction for (int k = 0; k < IB; ++k) o += red2[k][j], scaling o *= scale, simd_sum(o * o), and part2[cb] = osq write inside if (ib == 0) (tid < 128), keeping the surrounding threadgroup_barrier(mem_flags::mem_threadgroup) uniform across all threads. Threads ib = 1, 2, 3 (tid = 128..511, 75% of the threadgroup) never read o or part2 after red2[ib][j] = o, so pruning their redundant readout eliminates 1,536 shared-memory loads and 384 simd_sums per token per value head (73,728 shared-memory loads per GDN layer).
• State Base Pointer Hoisting & #pragma clang loop unroll(full): Hoisted device float* st_base = state + (i0 * hv + h) * S + j with stride st_stride = hv * S across the initial state load and commit writeback, and added #pragma clang loop unroll(full) to all fixed-trip loops (k < 3, k < 4, r < RPT, k < IB, k < 16; k += 2).
• Threadgroup Occupancy Preserved: All GDN_NTOK = 1..8 specializations compile with maxTotalThreadsPerThreadgroup = 512 (gdn_step) on Apple M5.

2. metal/kernels/qsa.metal (qsa_attn, lines 327–398) #

• Vectorized char4 Key Loads & Query Pointer Hoisting: Replaced the 8 scalar char loads from kq + rowi * QD + lane * 8 (which is 8-byte aligned since QD = 256 and lane * 8 is a multiple of 8) with two 32-bit char4 loads (k0, k1), hoisted threadgroup const float* q_lane = &qs[0][lane * 8] outside the cell loop with stride QD per query head, and added #pragma clang loop unroll(full) on the HPG = 12 and i < 8 loops in both the score and value-accumulation phases.

Verification & Bitwise Parity #

1. Standalone Test Suite (/tmp/argon_r17_gdn/strata/build-metal):
   • gdn_parity: ALL PASS (all GDN_NTOK = 1..8 pipelines report max 512 threads per threadgroup; ntok = 1..4 verify/commit/partial-commit checks pass).
   • gdn_prompt_parity: ALL PASS (T = 1..2048 single-pass vs TB=16 and CB scan checks pass).
   • qsa_parity: ALL PASS (prep, score, top-k, attention + gate pass).
   • qsa_prompt_parity: ALL PASS (pos0 = 0..30000 qsa_attn_mm / qsa_flash / qsa_prep checks pass).
2. Paired A/B Bitwise Parity (/tmp/argon_r17_gdn/ab_bench):
   • gdn_step: 100% bitwise memcmp == 0 across state, hist_out, and y (including all 4 prepared xprep planes + per-16 sums when plane = MAXT * 6144 = 49152) for all ntok in {1, 2, 3, 4}, commit in {0, 1}, plane in {0, 49152}.
   • qsa_attn: 100% bitwise memcmp == 0 across part ([n_tok][chunks][QH][QD + 2]) for all pos0 in {256, 2048, 8192}, ntok in {1, 2, 3, 4}.

Paired Interleaved A/B Benchmark Results (Apple M5) #

1. gdn_step (36 layers rotated, 113 MB state, median of 9 interleaved runs) #

2. qsa_attn (24 dispatches per timing batch, median of 11 interleaved runs) #

Artifacts #

• Unified patch: /tmp/argon_r17_gdn/argon_r17_gdn.patch (verified with git -C vendor/strata apply --check /tmp/argon_r17_gdn/argon_r17_gdn.patch)
• Paired A/B benchmark source & binary: /tmp/argon_r17_gdn/ab_bench.cpp, /tmp/argon_r17_gdn/ab_bench
• Scratch build & test tree: /tmp/argon_r17_gdn/strata/build-metal