5.5 kB 67 lines
f3384f69
Round 16 blitz: fused QSA Wk+Wv+Wik(+Wiq) single GEMM (-24..-36 dispatches, -5.6 MB scratch), decode resident-MoE + shared-expert wave zip (-48 barriers), gemv TAB>1 guards, ple_moe_gate_norm_c SIMD compaction + ple_conv half4/t>=9
dawn
⋯
audits/opus_r16_kvik.md
sourcerendered
download raw history
at main Markdown

opus_r16_kvik — round 16 track A: QSA Wk | Wv | Wik (| Wiq) as one GEMM #

verdict: bitwise-clean, -24/-36 dispatches and -5.6 MB of prompt scratch per chunk, but speed-neutral at C = 2048. the premise that the 512-column GEMMs are bandwidth-bound turned out wrong: the separate GEMMs already run concurrently (nb_) and reach ~13.7 TFLOPS fp16 together (2·2048·1664·2560 flop in 1.27 ms), so they're compute-bound. dropping the extra p.mixh reads and halving the weight re-loads saves nothing you can measure. the patch is safe to merge for the memory and dispatch savings; don't expect tok/s from it.

patch: /tmp/opus_r16_kvik/opus_r16_kvik.patch (git -C vendor/strata apply --check passes) test: /tmp/opus_r16_kvik/kvik_test.cpp (standalone, built in the scratch tree as metal/tests/kvik_test.cpp; NOT part of the patch) raw: /tmp/opus_r16_kvik/{gemm_runs,tile_sweep,prep_run}.txt

the change (window.cpp 3070–3530, qsa.metal 90–200 only) #

  • prompt_chunk QSA branch: dequantize attn_k → Wk rows 0..511, attn_v → 512..1023, indexer.k_proj → 1024..1151, indexer.q_proj (only when pos0 + T > 2051) → 1152..1663, all into p.Wk (= p.Wz, 6144 rows). one gemm(p.mixh, p.Wk, p.kcur, T, 1152|1664, N) into p.kcur (= p.gqn, MT·2048 floats). qsa_prep gets kcur, kcur+512, kcur+1152 (iqraw), kcur+1024 (ikraw).
  • prompt_gemm_init / prompt_gemm_need_bytes: p.iqr (MT·512 f32), p.ikr (MT·128 f32), p.Wik (128·N f16) are no longer allocated (4 MiB + 1 MiB + 640 KiB at C = 2048, N = 2560); p.Wv / p.Wiq / p.vcur aliases dropped. grep: nothing else in window.cpp reads them (prompt_chain.inc has its own b.* struct, untouched).
  • gemm() tile: new bucket Nn <= 1664 (only the fused shape lives in 641..1664): M < 768 ? 32_128_2 : 64_128_4. the default it would otherwise fall into (M < 1536 ? 128_64_4) is 10% slower at M = 1024 (0.89x, +0.64/+0.95 ms/chunk).
  • qsa_prep: i deviated from the brief's pointer-difference detection (vcur == kcur + 512) and used a function constant QSA_KV_STRIDE (26, unused in qsa.metal + qsa_prompt.metal) instead. reasons: the pointer test can false-positive whenever a caller passes K/V as slices of one buffer (or adjacent 1-token allocations), and it would add a runtime branch to every decode/MTP/test caller. with the constant undefined the strides fold to the old 512 / 128 / 512 at compile time, so every other caller's pipeline is unchanged. prompt_chunk keys the pipeline ...#kv1152 / #kv1664 (two extra pipelines, compiled on first use).

evidence (M5 Air; the :8080 server was live, so expect some GPU contention noise) #

GEMM parity: fused vs the 3/4 separate GEMMs (each with prompt_chunk's tile), every K / V / ik / iq sub-slice of every row memcmp == 0 at M = 256, 512, 768, 1024, 1536, 1983, 2048 × Nn = 1152, 1664, and for every fused tile in the sweep (128_64_4, 64_128_4, 64_64_4, 128_128_4, 64_128_2, 32_128_2).

GEMM speed (paired, 15 alternating rounds × 16 reps, run 1; run 2 agreed within ±0.06 ms; a third run had contention spikes and was dropped):

M 1152: sep → fused ms/chunk (12 layers) 1664: sep → fused ms/chunk
256 122.9 → 114.8 us (1.07x) -0.10 179.5 → 174.0 (1.03x) -0.07
512 236.9 → 221.3 (1.07x) -0.19 333.5 → 318.9 (1.05x) -0.17
768 333.4 → 341.8 (0.98x) +0.10 488.8 → 504.8 (0.97x) +0.19
1024 438.5 → 457.6 (0.96x) +0.23 643.7 → 659.3 (0.98x) +0.19
1536 667.4 → 663.0 (1.01x) -0.05 961.1 → 948.7 (1.01x) -0.15
1983 859.3 → 825.8 (1.04x) -0.40 1230.9 → 1227.2 (1.00x) -0.04
2048 874.2 → 882.3 (0.99x) +0.10 1275.0 → 1273.9 (1.00x) -0.01

no fused tile beats this at M ≥ 768 (tile_sweep.txt): 128_128_4 runs at 0.5x, 64_128_2 at 0.8x, the rest at 0.9–1.04x. so the fused GEMM gets small wins only for short tail chunks (≤ 512 rows); full 2048-row chunks are ±0.1 ms.

qsa_prep parity: at stride 1152 and 1664, with fp32 and qf16 (fc 25), 12 cases (T = 2, 3, 5, 517, 1029, 1983, 2048; pos0 = 0, 1, 1022, 2047, 2048, 2050, 4093, 6143, which covers the t < 3 boundary, the tail, pos0 % 4 ≠ 0 and pooling past the budget), all 9 outputs (qcur, qidx, kq, vq, ksc, vsc, tail, dead, pooled) memcmp == 0 against the separate-buffer reference. the reference is compiled from the unpatched vendor qsa.metal (KVIK_BASE), so the default path is also checked bitwise against the old kernel. negative control (KVIK_NEG=1: fused rows read with the default stride): every case DIFFs, so the test does catch a wrong stride. qsa_prep timing is in the noise (single cold dispatches, contention).

regression: scratch qsa_parity ALL PASS, qsa_prompt_parity ALL PASS (incl. qsa_prep qf16 memcmp), and strata_metal_kernels (window.cpp) builds clean.

not done (per the rules) #

  • no full-model run: the need != prompt_gemm_need_bytes() runtime check and end-to-end prompt token parity weren't exercised. by inspection the two removals match (F32(MT*512), F32(MT*128), F16(128*N) dropped on both sides). run a prompt-parity / strata_window_cli check on a > 2051-token prompt once :8080 is free.
  • tests/dense_mm_parity.cpp was left alone (outside my lines); the bench lives in kvik_test.cpp. if the fused tile rule lands, new_tile() there should mirror it.