From 35f84cbe162d134c2a453d7f5dccc03beedb5b7e Mon Sep 17 00:00:00 2001 From: dawn <90008@klbr.net> Date: Sat, 17 Oct 2026 20:49:41 +0300 Subject: [PATCH] Round 18 blitz: per-stream gr_norm_comb 5.09x (-80%), xprep float4 (-46.8%), moe_combine_sh (-49%), Q6_K moe_gu NE>=3 (-25%), qsa_topk2/qsa_score (-5..9%) -> 40.03 tok/s decode & 445.79 tok/s 10.18K prefill --- audits/argon_r18_qsa2.md | 70 ++++ audits/opus_r18_moe.md | 65 ++++ audits/sol_r18_grcomb.md | 88 +++++ scripts/strata_metal_complete.diff | 550 ++++++++++++++++++++++++++--- 4 files changed, 726 insertions(+), 47 deletions(-) create mode 100644 audits/argon_r18_qsa2.md create mode 100644 audits/opus_r18_moe.md create mode 100644 audits/sol_r18_grcomb.md diff --git a/audits/argon_r18_qsa2.md b/audits/argon_r18_qsa2.md new file mode 100644 index 0000000..c791cee --- /dev/null +++ b/audits/argon_r18_qsa2.md @@ -0,0 +1,70 @@ +# Round 18 Track C: `qsa_topk2` Register-Cached Keys & Barrier Reduction + `qsa_score` Query Hoisting & Unrolling + +## 1. Summary of Changes + +### Target 1: `metal/kernels/qsa2.metal` (`qsa_topk2`, lines 121–238) +1. **Redundant `threadgroup_barrier` Removal (`20 -> 13` radix-pass barriers, `-7` barriers per `qsa_topk2` dispatch = `-84` barriers per 12-layer decode window):** + - Hoisted the initial `if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed);` + `threadgroup_barrier(mem_flags::mem_threadgroup);` once before the `pass = 0..3` loop. + - Inside each radix pass, immediately after `cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed);`, cleared the bin for the next pass via `if (pass < 3) atomic_store_explicit(&hist[255 - tid], 0u, memory_order_relaxed);`. Because two `threadgroup_barrier`s (the `sgs` barrier and `sel` barrier) execute between that store and the next pass's `atomic_fetch_add_explicit`, `hist` is guaranteed zeroed before `pass + 1` accumulates into `hist`. + - Removed the redundant per-pass start barrier and trailing barrier after `need = sel[1];` (`sel[0]` and `sel[1]` are only read at the end of the pass and are not written again until after the histogram and `sgs` barriers of `pass + 1`). +2. **Register-Cached Order Keys (`nb <= 4096`, `n_kv <= 16,384`) with Scoped Live Range & Uniform Early `break`:** + - Scoped `uint keys[4];` inside `if (nb <= 4096)` (covering `n_kv <= 16,384`) so the `nb > 4096` (`pos0 = 30000`) path keeps zero `keys` register live-range across its barriers. + - Preloaded `keys[0..3]` via `#pragma clang loop unroll(full)` before `pass = 0..3` and reused `keys[i]` across all 4 radix passes with an early `if (b >= nb) break;` so `nb <= 1024` (`pos0 = 3000`) exits after `i = 0` (`1` iteration instead of `4`) and `nb = 2049` (`pos0 = 8192`) exits after `i = 1..2`. + - **Evaluated `ekeys[4]` in the emit pass (`E4` vs `M1_break`):** Adding a second `uint ekeys[4]` register array across the two emit-pass `threadgroup_barrier`s increased peak register pressure in the 1024-thread threadgroup and measured slower than `M1_break` at `pos0 = 3000` (`6.52 us` in `M1_break` vs `6.61 us` in `E4`) and `pos0 = 8192` (`7.81 us` in `M1_break` vs `8.07 us` in `E4`). Retaining `M1_break` (without `ekeys[4]`) delivers the fastest `qsa_topk2` across both `pos0 = 3000` (`-3.17%` to `-4.95%`) and `pos0 = 8192` (`-5.60%` to `-5.87%`) while also improving `pos0 = 30000` (`-2.19%` to `-2.34%`). + +### Target 2: `metal/kernels/qsa.metal` (`qsa_score`, lines 199–222) +- Hoisted `device const float* qp = qidx + (t * IH) * ID + lane;`, preloaded `float q[IH][4]` and `float k[4]` with `#pragma clang loop unroll(full)`, and unrolled the `h < IH` and `i < 4` FMA reduction loops so all 16 query loads (`qp[h * ID + 32 * i]`) and 4 key loads (`key[lane + 32 * i]`) use compile-time immediate offsets. +- **Evaluated `threadgroup float q_tg[512]` (`vC`) vs Register Hoisting (`vB`):** Staging `qidx` (`2 KiB`) in `threadgroup float q_tg[512]` required adding a 256-thread cooperative load and `threadgroup_barrier(mem_flags::mem_threadgroup)` to a barrier-free kernel and increased `static_tg_mem` from `0 B` to `2,048 B`, which was slower than register hoisting (`vB`) across every tested configuration (`-0.79%..-5.19%` for `vC` vs `-5.47%..-8.94%` for `vB`). Shipped `vB` (`0 B` shared memory, `0` barriers). + +--- + +## 2. Verification & Bitwise Parity + +1. **Standalone Parity Suites (`/tmp/argon_r18_qsa2/strata/build-metal`):** + - `qsa_parity` (`P = 3000` and `P = 8192`): `ALL PASS` (`0 of 6153 ids differ`, `score rel L2 5.88e-08`, `< 0.6s`). + - `qsa_prompt_parity`: `ALL PASS` (`pos0 = 0, 1984, 2050, 30000`, `< 1.1s`). +2. **Paired A/B Bitwise Parity (`/tmp/argon_r18_qsa2/ab_bench`):** + - **100% bitwise `memcmp == 0`** on both `scores` and `ids` (including `0xa5`-poisoned padding/unused entries) across all 56 boundary & context combinations `pos0 in {0, 1, 3, 2048, 2050, 2051, 2052, 3000, 4092, 8192, 16380, 16384, 20000, 30000}` x `ntok in {1, 2, 3, 4}`, plus adversarial tie-breaking and signed-zero (`+0.0f` / `-0.0f` / all-equal / small-integer ties) tests across `pos0 in {2052, 3000, 8192, 16380, 16384, 30000}`. + +--- + +## 3. Paired Interleaved A/B Benchmark Results (Apple M5) + +Measured on a quiet Apple M5 GPU (`11` interleaved pairs x `32` dispatches per command buffer after `0.30s` P-state warmup; median per-dispatch time in microseconds): + +| `pos0` | `ntok` | Stage | Baseline (us) | Candidate (us) | Speedup | Delta | Bitwise `memcmp` | +|---:|---:|---|---:|---:|---:|---:|:---:| +| 3000 | 1 | `qsa_score` | 2.50 | 2.36 | **1.06x** | **-5.57%** | `PASS (0)` | +| 3000 | 1 | `qsa_topk2` | 6.86 | 6.52 | **1.05x** | **-4.95%** | `PASS (0)` | +| 3000 | 1 | `combined` | 9.52 | 9.05 | **1.05x** | **-4.87%** | `PASS (0)` | +| 3000 | 2 | `qsa_score` | 3.64 | 3.43 | **1.06x** | **-5.90%** | `PASS (0)` | +| 3000 | 2 | `qsa_topk2` | 6.80 | 6.52 | **1.04x** | **-4.08%** | `PASS (0)` | +| 3000 | 2 | `combined` | 10.66 | 10.14 | **1.05x** | **-4.90%** | `PASS (0)` | +| 3000 | 4 | `qsa_score` | 6.23 | 5.76 | **1.08x** | **-7.48%** | `PASS (0)` | +| 3000 | 4 | `qsa_topk2` | 6.77 | 6.56 | **1.03x** | **-3.17%** | `PASS (0)` | +| 3000 | 4 | `combined` | 13.27 | 12.48 | **1.06x** | **-5.99%** | `PASS (0)` | +| 8192 | 1 | `qsa_score` | 4.71 | 4.42 | **1.07x** | **-6.33%** | `PASS (0)` | +| 8192 | 1 | `qsa_topk2` | 8.30 | 7.84 | **1.06x** | **-5.60%** | `PASS (0)` | +| 8192 | 1 | `combined` | 13.16 | 12.40 | **1.06x** | **-5.76%** | `PASS (0)` | +| 8192 | 2 | `qsa_score` | 8.02 | 7.48 | **1.07x** | **-6.80%** | `PASS (0)` | +| 8192 | 2 | `qsa_topk2` | 8.29 | 7.81 | **1.06x** | **-5.72%** | `PASS (0)` | +| 8192 | 2 | `combined` | 16.47 | 15.43 | **1.07x** | **-6.32%** | `PASS (0)` | +| 8192 | 4 | `qsa_score` | 14.92 | 13.78 | **1.08x** | **-7.61%** | `PASS (0)` | +| 8192 | 4 | `qsa_topk2` | 8.30 | 7.81 | **1.06x** | **-5.87%** | `PASS (0)` | +| 8192 | 4 | `combined` | 23.39 | 21.77 | **1.07x** | **-6.92%** | `PASS (0)` | +| 30000 | 1 | `qsa_score` | 25.48 | 25.45 | **1.00x** | **-0.12%** | `PASS (0)` | +| 30000 | 1 | `qsa_topk2` | 14.01 | 13.71 | **1.02x** | **-2.19%** | `PASS (0)` | +| 30000 | 1 | `combined` | 39.96 | 39.65 | **1.01x** | **-0.78%** | `PASS (0)` | +| 30000 | 2 | `qsa_score` | 26.86 | 25.27 | **1.06x** | **-5.94%** | `PASS (0)` | +| 30000 | 2 | `qsa_topk2` | 14.06 | 13.73 | **1.02x** | **-2.32%** | `PASS (0)` | +| 30000 | 2 | `combined` | 41.60 | 39.51 | **1.05x** | **-5.04%** | `PASS (0)` | +| 30000 | 4 | `qsa_score` | 52.52 | 47.82 | **1.10x** | **-8.94%** | `PASS (0)` | +| 30000 | 4 | `qsa_topk2` | 14.12 | 13.79 | **1.02x** | **-2.34%** | `PASS (0)` | +| 30000 | 4 | `combined` | 67.35 | 62.22 | **1.08x** | **-7.61%** | `PASS (0)` | + +--- + +## 4. Artifacts +- Unified patch: `/tmp/argon_r18_qsa2/argon_r18_qsa2.patch` (verified clean with `git -C vendor/strata apply --check /tmp/argon_r18_qsa2/argon_r18_qsa2.patch`) +- Paired A/B benchmark source & binary: `/tmp/argon_r18_qsa2/ab_bench.cpp`, `/tmp/argon_r18_qsa2/ab_bench` +- Scratch workspace & build tree: `/tmp/argon_r18_qsa2/strata/build-metal` diff --git a/audits/opus_r18_moe.md b/audits/opus_r18_moe.md new file mode 100644 index 0000000..a5a0d6c --- /dev/null +++ b/audits/opus_r18_moe.md @@ -0,0 +1,65 @@ +# opus_r18_moe — Round 18 Track A: moe_gu / moe_down / moe_route NE>1 reuse, xprep float4, swiglu_prep, combines + +Patch: `/tmp/opus_r18_moe/opus_r18_moe.patch` (185 lines; `git -C vendor/strata apply --check` = OK). +Touches only `metal/kernels/{moe,units,glue}.metal` and `metal/tests/expert_parity.cpp`. +A/B harness (scratch, not in the patch since it would be a new file): `/tmp/opus_r18_moe/moe_ab.cpp` +(compiles the untouched kernels from `/tmp/opus_r18_moe/base_kernels` as A and the scratch tree as B in one process, +memcmp of every output buffer, paired alternating timing, median of 11 B/A ratios). Raw logs: +`/tmp/opus_r18_moe/ab_logs.txt`, parity log `/tmp/opus_r18_moe/parity_log.txt`. +Device: Apple M5 (the 24 GB Air), live :8080 server up the whole time, so absolute µs drift between runs; the +paired ratio is the number to trust. gu/dn bench: 2560 x 640, one 512-expert arena, 48 windows, each window's NE +tokens routed to the same 10 experts (so every group has exactly NE entries). + +## what landed (all bitwise: memcmp == 0 on every A/B run) + +| change | result | +|---|---| +| `xprep` float4 loads/stores + unroll (units.metal) | **-41.5 % .. -46.8 %** (n 2560 ntok 8: 19.8 → 10.5 µs; ntok 1 -41.5 %; n 640 ntok 80 -44.4 %; n 5120 -45.9 %) | +| `moe_combine_sh` topk==10 unrolled (glue.metal) | **-32.0 %** ntok 8 (17.6 → 12.0 µs), **-45.1 %** ntok 1 | +| `moe_combine` topk==10 unrolled (moe.metal) | **-33.3 %** ntok 8, **-48.7 %** ntok 1 | +| `gu_body` fused gate+up per row, one x load per entry — **only for UNIT==32 formats at NE≥3, IQ4_XS from NE 2** (`gu_fused`) | Q6_K 32,1,4: NE3 -18.5 %, NE4 -20.5 %; Q4_K -17.6/-18.6 %; Q5_K -15/-18 %; Q8_0 -9/-19 %; IQ4_XS NE2/3/4 -31.7/-33.2/-35.8 %; NR=2 tiles up to -25 % | +| production formats IQ2_XS gate/up + Q2_0R down | unchanged code path (A/B ±0.4 %, noise) | + +`gr_norm_comb` (gr2.metal) stays bitwise with `moe_combine_sh`: the unrolled sum keeps k order. fuse_parity confirms. + +## what was tested and rejected (numbers so nobody re-tries them) + +1. **gu NE>1 fusion for the 16-weight units (IQ2_XS, IQ3_S, Q2_0)** — the brief's main target. Regresses where it + matters: IQ2_XS 32,1,4 NE3 **+17.0 %**, 16,1,4 NE3 **+21.4 %**, NE4 +4..+10 %; IQ3_S 32,1,4 NE3 +14.7 %, + 16,1,4 NE3 +17.3 %. Holding two decoded units live costs more than the second x load (which is an L1 hit of 4 + float4s). Only NR=2 tiles gain (IQ3_S 32,2,4 NE4 -9 %, 16,2,4 NE4 -14.5 %; IQ2_XS ≈ ±2 %; Q2_0 NR2 **+4 %**), + and no M5 production tile uses gu NR=2 (window.cpp: gr = 1; gl = 16 for IQ2_XS/IQ3_S at ntok ≥ 3). +2. **"decode every row first" (`qg[NR][UNIT], qu[NR][UNIT]`)** — catastrophic: IQ2_XS/IQ3_S NR=2 **+37..+86 %**, + Q6_K NR=2 **+383..+623 %** (spills). +3. **Hoisting `tok[NE]` / x row pointers before the u loop** — neutral to worse: IQ2_XS +2..+5 %, IQ3_S 16,1,4 + NE3/4 +6 %. 2*NE live 64-bit pointers cost more than the uniform `etok[t]` reload. +4. **`dn_body` Q2_0/Q2_0R on the shared-decode path (`q[NR][UNIT]`)** — `moe_down_142_8_2_4`: NE2 **+5.0 %**, NE3 + **+11.0 %**, NE4 **+5.8 %**; 4,2,4: +7.7/+14.2/+12.1 %; 8,1,4 = 0 (NR 1 is the same code). The existing exclusion is + right; the patch only adds a one-line comment saying so. `ew[NE]` hoist: it is only read in the epilogue, nothing + to hoist. +5. **`swiglu_prep` via `simd_shuffle_down`** (0 B threadgroup memory, no barrier, bitwise) — **+2.0..+2.7 %** at + ntok 8 (n 640 and 5120), -0.9 % at ntok 1, ±0 at ntok 4 (two repeats each). 16 shuffles on every lane cost more + than one lane's 16 threadgroup reads after a barrier in a one-simdgroup threadgroup. Not landed. (The real lever + is the 16-thread threadgroup itself — half a simdgroup — but that is a host change in window.cpp, out of scope.) +6. **`moe_route` unroll pragmas** — 0 % (ntok 1: +0.6 %, ntok 8: +0.4 %): the compiler already unrolls these. + **Replacing `v[ex / 32] = -INFINITY` by a 16-way select** to keep v in registers — **+19 % (ntok 8) .. +56 % + (ntok 1)**. Neither landed; moe_route is untouched. + +## verification + +- A/B memcmp == 0 for every run above (gu 11 formats x tiles x NE 1..4, dn 142 x 3 tiles x NE 1..4, xprep x4 shapes, + swiglu x3, combine/combine_sh ntok 1/8, route ids/w/plan ntok 1/4/8). +- `expert_parity` (vs double CPU ref): 160/160 PASS — Q6_K, Q4_K, Q5_K, Q8_0, IQ4_XS, IQ2_XS, IQ3_S, Q2_0 gate/up x + ntok 1,2,3,4,8 x {random 0.3-share routing, new `--same`} x gu tiles {32,1,4, 32,2,4}; plus IQ2_XS/Q2_0 `--dn-prep`. + `--same` (new) routes every token to the same 10 experts so groups have exactly ntok entries and the NE 2/3/4/8 + bodies (incl. `gu_fused`) are hit deterministically. +- `route_parity` PASS (0 mismatches), `fuse_parity` ALL PASS (gr_norm_comb bitwise), `prompt_parity` ALL PASS. +- No :8080 traffic, no strata_window_cli runs. + +## expected effect on the live model + +Our K512 model is IQ2_XS gate/up + Q2_0 (142) down: its gu/dn kernels are unchanged. What it gets is `xprep` +(~9 µs a call saved at ntok 8 — xprep runs several times a layer) and `moe_combine_sh` (~5.6 µs a layer at ntok 8, +if it is not folded into gr_norm_comb; the folded gr_norm_comb path is a different kernel in gr2.metal that this +round does not own and that likely has the same non-unrolled topk loop — worth a follow-up there). The gu_fused win +lands for the UD-Q4/5/6_K_XL and IQ4_XS packs. diff --git a/audits/sol_r18_grcomb.md b/audits/sol_r18_grcomb.md new file mode 100644 index 0000000..c4e0628 --- /dev/null +++ b/audits/sol_r18_grcomb.md @@ -0,0 +1,88 @@ +# round 18 track b + +## contract / preflight +owner: metal grouped-residual kernels and Window::gr_norm dispatch selection. preserve operation order, generic shape fallback, apply=false semantics and buffer bindings. no live server or full-model runs. +question: do fixed-shape pointer hoisting / unrolling and stream staging improve latency without changing a single output bit? compare baseline 4s, optimized 4s, per-stream grid and staged 4s; compare baseline/optimized gr_down_q8_8_4. +evaluator: deterministic fuse_parity inputs, memcmp R/xn and down/inject for ntok 1..8; ordered alternating gpu timings, 3 exploratory and 5 confirmation pairs, 64 dispatch repetitions each. improvement must exceed 3% median, with bitwise parity mandatory. each invocation timeout 4.9s, no concurrent gpu tests; compile jobs exempt from test invocation limit. live server may introduce noise; report it. +source snapshots: /tmp/sol_r18_grcomb/baseline and candidate scratch tree; no tracked vendor edits. stop if compilation cannot be repaired in two attempts or tests exceed the invocation budget. + +## result / decision + +**ship the hoisted/unrolled `gr_norm_comb` with grid `(4, ntok)`.** keep optimized `gr_norm_comb_4s` available, but do not dispatch it from `Window::gr_norm`. shared-bo staging was slower. do not ship the `gr_down_f` change: it passed bitwise parity but did not establish a repeatable improvement. + +main patch: `/tmp/sol_r18_grcomb/sol_r18_grcomb.patch` (three files; no `gr.metal` edits). rejected/neutral down experiment: `/tmp/sol_r18_grcomb/neutral_down.patch`. neither patch has been applied to tracked vendor sources. + +## execution / evidence + +- machine: Apple M5 (the local 10-core GPU laptop). no server access, full-model benchmarks, or spawned agents. +- vendor input HEAD: `296b59d2f73c8f04a805df29c2fa97655adace78`; input working-tree snapshots are in `/tmp/sol_r18_grcomb/baseline/`, since HEAD alone does not describe this already-dirty development tree. +- scratch baseline commit: `faeb472`; all executed shader fixtures, evaluator source, binary, and input/candidate hashes: `/tmp/sol_r18_grcomb/sha256.json`. +- all narrow GPU invocations completed below 1.6 seconds. first existing full `fuse_parity` invocation was terminated at 4.9s during cold shader compilation. parent then explicitly permitted a 20s timeout for the full test's compile overhead. full suite subsequently passed twice: 5.45s and 13.73s (including compilation), exit 0; final log `/tmp/sol_r18_grcomb/final_full_parity.log`. +- standalone compile/link succeeded; scratch `window.cpp` compiled successfully with the existing engine compiler flags, output `/tmp/sol_r18_grcomb/window.o`. +- final `git -C vendor/strata apply --check /tmp/sol_r18_grcomb/sol_r18_grcomb.patch`: exit 0. scratch `git diff --check`: exit 0. +- `--gr-ab ... 0 norm` rejects invalid token counts with exit 2 before creating a GPU context. + +### correctness + +`memcmp == 0` for both R and xn across ntok 1..8, apply false/true, topk 1/10, for optimized 4s, per-stream grid and staged-bo variants versus the original 4s. 96 final A/B parity cases passed. original full suite also checks both production combine kernels versus `moe_combine_sh` + `gr_norm`, and passed all 32 cases (plus gemv/router tests). down Q8_0 lo and inject are bitwise equal for ntok 1..8. generic n != 2560 behavior is retained but was not newly exercised. + +### timing contract adjustment + +initial 80ms/30ms warm-up left obvious clock ramps. before confirmation, increased warm-up to 300ms baseline + 100ms candidate, with five alternating-order pairs of 64 dispatches each. preserved both exploratory and confirmation logs rather than replacing the original evidence. the live machine can change clocks or have other GPU work; absolute timings are not stable, so a direct optimized-4s-vs-grid confirmation was also run. this is a kernel microbenchmark, **not** an end-to-end speedup estimate. + +### original 4s vs selected per-stream kernel + +medians of five paired batches, microseconds/dispatch: + +| ntok | original 4s | selected grid4 | reduction | +|---:|---:|---:|---:| +| 1 | 22.673 | 4.458 | 80.3% | +| 2 | 22.820 | 4.680 | 79.5% | +| 3 | 22.882 | 5.902 | 74.2% | +| 4 | 22.878 | 6.453 | 71.8% | +| 5 | 54.246 | 16.863 | 68.9% | +| 6 | 22.977 | 8.649 | 62.4% | +| 7 | 43.543 | 13.086 | 69.9% | +| 8 | 22.951 | 10.456 | 54.4% | + +optimized 4s reduced baseline time by 46–49% in this confirmation; staged shared bo reduced it by 43–46%. the selected grid kernel reduced it by 54–80%. + +### direct optimized 4s vs selected grid4 + +these runs happened at much slower clocks for most token counts, but both candidates were measured back-to-back in alternating order: + +| ntok | optimized 4s | selected grid4 | reduction | +|---:|---:|---:|---:| +| 1 | 51.937 | 20.042 | 61.4% | +| 2 | 51.971 | 20.734 | 60.1% | +| 3 | 52.184 | 26.059 | 50.1% | +| 4 | 52.386 | 28.171 | 46.2% | +| 5 | 52.582 | 31.563 | 40.0% | +| 6 | 11.881 | 8.551 | 28.0% | +| 7 | 54.546 | 30.451 | 44.2% | +| 8 | 52.787 | 46.443 | 12.0% | + +all eight direct comparisons favor grid4 (12–61%), exceeding the 3% decision threshold. + +### down experiment (not shipped) + +implemented exactly the requested xnp pointer hoist, UNIT/16 load unroll, and KS reduction unroll. ntok 1..8 all bitwise PASS. confirmation median reductions were 1.30%, 4.40%, -0.14%, 0.13%, -0.18%, -0.03%, 0.11%, 0.12%. ntok2 was clock-noisy, not a broad repeatable win. preserve this as a measured neutral result, not a performance claim. + +## reproduce + +artifacts live under `/tmp/sol_r18_grcomb/`; `norm.metal` contains original 4s plus all three candidates, `down.metal` original/candidate down, `direct.metal` uses optimized 4s as the reference. the focused evaluator is the new `--gr-ab ` mode in `metal/tests/fuse_parity.cpp`. + +from the project root (replace `$PWD` with the checkout path if launched elsewhere): + +```text +clang++ -std=c++20 -O2 -Ivendor/strata/metal/src /tmp/sol_r18_grcomb/strata/metal/tests/fuse_parity.cpp vendor/strata/build-metal/libmtl_rt.a -framework Metal -framework Foundation -o /tmp/sol_r18_grcomb/fuse_parity +STRATA_METAL_KERNELS=/tmp/sol_r18_grcomb/strata/metal/kernels /tmp/sol_r18_grcomb/fuse_parity --gr-ab /tmp/sol_r18_grcomb/norm.metal 1 norm +STRATA_METAL_KERNELS=/tmp/sol_r18_grcomb/strata/metal/kernels /tmp/sol_r18_grcomb/fuse_parity --gr-ab /tmp/sol_r18_grcomb/direct.metal 1 norm +STRATA_METAL_KERNELS=/tmp/sol_r18_grcomb/strata/metal/kernels /tmp/sol_r18_grcomb/fuse_parity --gr-ab /tmp/sol_r18_grcomb/down.metal 1 down +``` + +repeat ntok 1..8, sequentially. paired logs: `confirm_norm_*.log`, `confirm_down_*.log`, `direct_*.log`; additional topk1 regression logs `final_norm_*.log`, post-review endpoint checks `postreview_norm_1.log` / `postreview_norm_8.log`. machine-readable medians: `measurements.json`. + +## review + +reviewed every patch hunk with `view_diff`, helper/Window context with `view_code`, then re-reviewed formatting fixes. no arithmetic reassociation or reduction-order changes; all new pointer accesses stay inside the existing fixed-shape specialization and topk guard, generic fallback remains. no live server settings changed. diff --git a/scripts/strata_metal_complete.diff b/scripts/strata_metal_complete.diff index 57fd56c..212085f 100644 --- a/scripts/strata_metal_complete.diff +++ b/scripts/strata_metal_complete.diff @@ -1017,6 +1017,26 @@ index 6d4ed3a..6540548 100644 const int n = int(a.n_in), s16n = n / 16; threadgroup float* st = xt + NTOK * n; for (int t = 0; t < NTOK; ++t) { +diff --git a/metal/kernels/glue.metal b/metal/kernels/glue.metal +index b7bc8bc..bfb08bc 100644 +--- a/metal/kernels/glue.metal ++++ b/metal/kernels/glue.metal +@@ -40,8 +40,14 @@ kernel void moe_combine_sh(device const float* ye [[buffer(0)]], device const in + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.n_embd, k = gid % a.n_embd; + if (t >= a.n_tok) return; ++ device const int* te = plan + P_TOKENT + t * a.topk; + float s = 0.0f; +- for (uint j = 0; j < a.topk; ++j) s += ye[plan[P_TOKENT + t * a.topk + j] * int(a.n_embd) + k]; ++ if (a.topk == 10) { ++#pragma clang loop unroll(full) ++ for (uint j = 0; j < 10; ++j) s += ye[te[j] * int(a.n_embd) + k]; ++ } else { ++ for (uint j = 0; j < a.topk; ++j) s += ye[te[j] * int(a.n_embd) + k]; ++ } + const float gs = 1.0f / (1.0f + exp(-sg[t])); + bo[gid] = fma(gs, ysh[gid], s); + } diff --git a/metal/kernels/gr.metal b/metal/kernels/gr.metal index c1960fd..cd86797 100644 --- a/metal/kernels/gr.metal @@ -1114,7 +1134,7 @@ index c1960fd..cd86797 100644 +template [[host_name("gr_down_q5_8_8")]] kernel decltype(gr_down_f<8, 8, T_Q5_0, T_Q8_0>) gr_down_f<8, 8, T_Q5_0, T_Q8_0>; +template [[host_name("gr_up_q5")]] kernel decltype(gr_up_f) gr_up_f; diff --git a/metal/kernels/gr2.metal b/metal/kernels/gr2.metal -index 1253a11..17daff0 100644 +index 1253a11..331a63a 100644 --- a/metal/kernels/gr2.metal +++ b/metal/kernels/gr2.metal @@ -138,7 +138,6 @@ kernel void gr_up3(device const uchar* w_up [[buffer(0)]], device const float* D @@ -1144,29 +1164,75 @@ index 1253a11..17daff0 100644 } } -@@ -317,6 +315,7 @@ kernel void gr_norm_comb(device float* R [[buffer(0)]], device const float* ye [ +@@ -317,25 +315,129 @@ kernel void gr_norm_comb(device float* R [[buffer(0)]], device const float* ye [ device float* r = R + (t * a.hc + c) * n; const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; const float gs = 1.0f / (1.0f + exp(-sgv[t])); + float rv[10]; float ss = 0.0f; - for (uint d = tid; d < n; d += 256) { - float v = r[d]; -@@ -327,6 +326,7 @@ kernel void gr_norm_comb(device float* R [[buffer(0)]], device const float* ye [ - v = fma(bo, g, v); - r[d] = v; +- for (uint d = tid; d < n; d += 256) { +- float v = r[d]; ++ if (n == 2560 && topk == 10) { ++ device const float* yep[10]; + if (a.apply) { +- float s = 0.0f; +- for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; +- const float bo = fma(gs, ysh[t * n + d], s); +- v = fma(bo, g, v); +- r[d] = v; ++ #pragma clang loop unroll(full) ++ for (uint j = 0; j < 10; ++j) ++ yep[j] = ye + uint(plan[P_TOKENT + t * 10 + j]) * 2560u + tid; ++ } ++ device const float* yshp = ysh + t * 2560u + tid; ++ device float* rp = r + tid; ++ #pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) { ++ float v = rp[i * 256]; ++ if (a.apply) { ++ float s = 0.0f; ++ #pragma clang loop unroll(full) ++ for (uint j = 0; j < 10; ++j) s += yep[j][i * 256]; ++ v = fma(fma(gs, yshp[i * 256], s), g, v); ++ rp[i * 256] = v; ++ } ++ rv[i] = v; ++ ss = fma(v, v, ss); ++ } ++ } else { ++ for (uint d = tid; d < n; d += 256) { ++ float v = r[d]; ++ if (a.apply) { ++ float s = 0.0f; ++ for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; ++ const float bo = fma(gs, ysh[t * n + d], s); ++ v = fma(bo, g, v); ++ r[d] = v; ++ } ++ if (n <= 2560) rv[d / 256] = v; ++ ss = fma(v, v, ss); } -+ if (n <= 2560) rv[d / 256] = v; - ss = fma(v, v, ss); +- ss = fma(v, v, ss); } ss = simd_sum(ss); -@@ -337,5 +337,54 @@ kernel void gr_norm_comb(device float* R [[buffer(0)]], device const float* ye [ + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; ++ #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; const float rs = rsqrt(tot / float(n) + a.eps); device float* o = xn + (t * a.hc + c) * n; device const float* w = w_norm + c * n; - for (uint d = tid; d < n; d += 256) o[d] = r[d] * w[d] * rs; -+ for (uint d = tid; d < n; d += 256) -+ o[d] = (n <= 2560 ? rv[d / 256] : r[d]) * w[d] * rs; ++ if (n == 2560) { ++ o += tid; ++ w += tid; ++ #pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) o[i * 256] = rv[i] * w[i * 256] * rs; ++ } else { ++ for (uint d = tid; d < n; d += 256) ++ o[d] = (n <= 2560 ? rv[d / 256] : r[d]) * w[d] * rs; ++ } +} + +// Fixed n_embd = 2560, hc = 4; grid (1, n_tok), 256 threads. Form bo once for all streams. @@ -1179,45 +1245,66 @@ index 1253a11..17daff0 100644 + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + const uint t = tg.y, n = 2560; ++ device float* rp = R + t * 4 * n + tid; ++ device float* op = xn + t * 4 * n + tid; ++ device const float* wp = w_norm + tid; + float rv[4][10], ss[4] = {0.0f, 0.0f, 0.0f, 0.0f}, g[4]; ++ #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) + g[c] = a.apply ? 2.0f / (1.0f + exp(-inj[t * 4 + c] / 4.0f)) : 0.0f; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); ++ device const float* yep[10]; ++ if (a.apply && topk == 10) { ++ #pragma clang loop unroll(full) ++ for (uint j = 0; j < 10; ++j) ++ yep[j] = ye + uint(plan[P_TOKENT + t * 10 + j]) * 2560u + tid; ++ } ++ device const float* yshp = ysh + t * n + tid; ++ #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const uint d = tid + i * 256; + float bo = 0.0f; + if (a.apply) { + float s = 0.0f; -+ for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; -+ bo = fma(gs, ysh[t * n + d], s); ++ if (topk == 10) { ++ #pragma clang loop unroll(full) ++ for (uint j = 0; j < 10; ++j) s += yep[j][i * 256]; ++ } else { ++ for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; ++ } ++ bo = fma(gs, yshp[i * 256], s); + } ++ #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) { -+ device float* r = R + (t * 4 + c) * n; -+ float v = r[d]; -+ if (a.apply) { v = fma(bo, g[c], v); r[d] = v; } ++ device float* r = rp + c * n; ++ float v = r[i * 256]; ++ if (a.apply) { v = fma(bo, g[c], v); r[i * 256] = v; } + rv[c][i] = v; + ss[c] = fma(v, v, ss[c]); + } + } ++ #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) { + const float s = simd_sum(ss[c]); + if (lane == 0) part[c][sg] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); ++ #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) { + float tot = 0.0f; ++ #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[c][i]; + const float rs = rsqrt(tot / float(n) + a.eps); -+ device float* o = xn + (t * 4 + c) * n; -+ device const float* w = w_norm + c * n; ++ device float* o = op + c * n; ++ device const float* w = wp + c * n; ++ #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { -+ const uint d = tid + i * 256; -+ o[d] = rv[c][i] * w[d] * rs; ++ o[i * 256] = rv[c][i] * w[i * 256] * rs; + } + } } diff --git a/metal/kernels/moe.metal b/metal/kernels/moe.metal -index bb8fdbb..8317881 100644 +index bb8fdbb..b27024b 100644 --- a/metal/kernels/moe.metal +++ b/metal/kernels/moe.metal @@ -20,7 +20,13 @@ @@ -1434,7 +1521,7 @@ index bb8fdbb..8317881 100644 } for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); ssum = max(ssum, 6.103515625e-5f); -@@ -160,7 +253,7 @@ kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids +@@ -160,12 +253,21 @@ kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids if (lane < 10) w[sg * 10 + lane] = selected * inv; } threadgroup_barrier(mem_flags::mem_device); @@ -1443,7 +1530,21 @@ index bb8fdbb..8317881 100644 } // ---- gate + up + SwiGLU ------------------------------------------------------------------------------------------- -@@ -176,19 +269,36 @@ static inline void gu_body(device const uchar* g0, device const uchar* u0, devic + // One pass over a group's rows for all its NE entries (1..8: an expert appears once per token); NE is uniform per + // threadgroup, so the kernel switches once into a body specialised for it. ++// ++// gu_fused: per row, gate and up decoded together and each entry's x slice loaded once for both (the 32-weight units, ++// whose x load is two runs and their sums: the M5 Air, 2560 x 640, NE 3-4: Q4_K / Q5_K / Q6_K / Q8_0 -8..-25%, ++// IQ4_XS -31..-36% from NE 2). The 16-weight units keep two x loads a row: the two decoded units held at once cost ++// them more (IQ2_XS / IQ3_S at NR 1, NE 3: +15..+21%). The same sums in the same order either way. ++template struct gu_fused { ++ enum { value = fmt::UNIT == 32 && (NE >= 3 || F == T_IQ4_XS) }; ++}; ++ + template + static inline void gu_body(device const uchar* g0, device const uchar* u0, device const float* x, + device const float* xs16, device float* h, device const int* etok, int c, int row0, int li, +@@ -176,19 +278,50 @@ static inline void gu_body(device const uchar* g0, device const uchar* u0, devic for (int t = 0; t < NE; ++t) ag[r][t] = au[r][t] = 0.0f; const int U = int(a.n_embd) / UNIT; for (int u = li; u < U; u += LPR) { @@ -1469,6 +1570,20 @@ index bb8fdbb..8317881 100644 + Dec::run(u0 + r * a.gu_row, u, q, c, tab); + au[r][0] += unit_combine(q, c, xs, xsum); + } ++ } else if (gu_fused::value) { ++ for (int r = 0; r < NR; ++r) { ++ float qg[UNIT], qu[UNIT]; ++ float4 cg, cu; ++ Dec::run(g0 + r * a.gu_row, u, qg, cg, tab); ++ Dec::run(u0 + r * a.gu_row, u, qu, cu, tab); ++ for (int t = 0; t < NE; ++t) { ++ const int tok = etok[t]; ++ float xs[UNIT], xsum[2]; ++ X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); ++ ag[r][t] += unit_combine(qg, cg, xs, xsum); ++ au[r][t] += unit_combine(qu, cu, xs, xsum); ++ } ++ } + } else { + for (int r = 0; r < NR; ++r) { + float q[UNIT]; @@ -1490,7 +1605,7 @@ index bb8fdbb..8317881 100644 } } } -@@ -199,6 +309,11 @@ static inline void gu_body(device const uchar* g0, device const uchar* u0, devic +@@ -199,6 +332,11 @@ static inline void gu_body(device const uchar* g0, device const uchar* u0, devic } } @@ -1502,7 +1617,7 @@ index bb8fdbb..8317881 100644 // bench probe (function constant 22): 1 = no table load (wrong values for the i-quants) constant int MOE_PROBE [[function_constant(22)]]; constant int moe_probe = is_function_constant_defined(MOE_PROBE) ? MOE_PROBE : 0; -@@ -213,18 +328,19 @@ kernel void moe_gu(device const uchar* W [[buffer(0)]], device const float* x [[ +@@ -213,18 +351,19 @@ kernel void moe_gu(device const uchar* W [[buffer(0)]], device const float* x [[ ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { constexpr int RG = 32 / LPR; threadgroup uint tab[fmt::TAB]; @@ -1526,7 +1641,7 @@ index bb8fdbb..8317881 100644 device const uchar* g0 = E + a.gate_off + (ulong)row0 * a.gu_row; device const uchar* u0 = E + a.up_off + (ulong)row0 * a.gu_row; switch (ne) { -@@ -254,13 +370,35 @@ static inline void dn_body(device const uchar* d0, device const float* h, device +@@ -254,13 +393,36 @@ static inline void dn_body(device const uchar* d0, device const float* h, device for (int t = 0; t < NE; ++t) ac[r][t] = 0.0f; const int U = int(a.n_ff) / UNIT; for (int u = li; u < U; u += LPR) { @@ -1546,6 +1661,7 @@ index bb8fdbb..8317881 100644 + ac[r][0] += unit_combine(q, sc, xs, xsum); + } + } else if (F == T_Q2_0 || F == T_Q2_0R) { ++ // one decoded row at a time (all NR held at once: +5..+14% at 8,2,4 / 4,2,4, NE 2-4, the M5 Air) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 sc; @@ -1568,7 +1684,7 @@ index bb8fdbb..8317881 100644 } } for (int r = 0; r < NR; ++r) -@@ -280,15 +418,18 @@ kernel void moe_down(device const uchar* W [[buffer(0)]], device const float* h +@@ -280,15 +442,18 @@ kernel void moe_down(device const uchar* W [[buffer(0)]], device const float* h ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { constexpr int RG = 32 / LPR; threadgroup uint tab[fmt::TAB]; @@ -1592,6 +1708,22 @@ index bb8fdbb..8317881 100644 switch (ne) { case 1: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; case 2: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; +@@ -310,8 +475,14 @@ kernel void moe_combine(device const float* ye [[buffer(0)]], device float* y [[ + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.n_embd, k = gid % a.n_embd; + if (t >= a.n_tok) return; ++ device const int* te = plan + P_TOKENT + t * a.topk; + float s = 0.0f; +- for (uint j = 0; j < a.topk; ++j) s += ye[plan[P_TOKENT + t * a.topk + j] * int(a.n_embd) + k]; ++ if (a.topk == 10) { ++#pragma clang loop unroll(full) ++ for (uint j = 0; j < 10; ++j) s += ye[te[j] * int(a.n_embd) + k]; ++ } else { ++ for (uint j = 0; j < a.topk; ++j) s += ye[te[j] * int(a.n_embd) + k]; ++ } + y[gid] = s; + } + diff --git a/metal/kernels/moe_mm.metal b/metal/kernels/moe_mm.metal index ff28897..c48125d 100644 --- a/metal/kernels/moe_mm.metal @@ -2704,7 +2836,7 @@ index 64833b0..5a7d0a8 100644 +template [[host_name("gemm_swiglu_up_32_64_2")]] kernel decltype(gemm_swiglu_up<32, 64, 2>) gemm_swiglu_up<32, 64, 2>; +template [[host_name("gemm_swiglu_up_64_128_4")]] kernel decltype(gemm_swiglu_up<64, 128, 4>) gemm_swiglu_up<64, 128, 4>; diff --git a/metal/kernels/qsa.metal b/metal/kernels/qsa.metal -index 9642e58..4a9adc5 100644 +index 9642e58..cf12671 100644 --- a/metal/kernels/qsa.metal +++ b/metal/kernels/qsa.metal @@ -35,6 +35,13 @@ constant int QSA_NTOK [[function_constant(3)]]; @@ -2867,7 +2999,31 @@ index 9642e58..4a9adc5 100644 const float o = norm_rope(x, 0.0f, iqn, row, part, tid, sg, lane, ID, a.eps, cs, absp, true, mtab); if (tid < ID) qidx[(t * IH + h) * ID + tid] = o; } -@@ -190,7 +233,8 @@ kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [ +@@ -164,12 +207,21 @@ kernel void qsa_score(device const float* qidx [[buffer(0)]], device const float + const uint b = tg.x * 8 + sg; + if (b > n_bid) return; + device const float* key = b == n_bid ? dead : pooled + b * ID; +- float k[4]; ++ device const float* qp = qidx + (t * IH) * ID + lane; ++ float k[4], q[IH][4]; ++ #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) k[i] = key[lane + 32 * i]; ++ #pragma clang loop unroll(full) ++ for (int h = 0; h < IH; ++h) { ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 4; ++i) q[h][i] = qp[h * ID + 32 * i]; ++ } + float s = 0.0f; ++ #pragma clang loop unroll(full) + for (int h = 0; h < IH; ++h) { + float d = 0.0f; +- for (int i = 0; i < 4; ++i) d = fma(qidx[(t * IH + h) * ID + lane + 32 * i], k[i], d); ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 4; ++i) d = fma(q[h][i], k[i], d); + d = simd_sum(d); + s += max(d, 0.0f); + } +@@ -190,7 +242,8 @@ kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [ ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { threadgroup atomic_uint hist[256]; threadgroup uint sel[4]; @@ -2877,7 +3033,7 @@ index 9642e58..4a9adc5 100644 const int t = int(tg.x); const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2, nb = n_bid + 1; device int* out = ids + t * QBUDGET; -@@ -199,6 +243,10 @@ kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [ +@@ -199,6 +252,10 @@ kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [ for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); return; } @@ -2888,7 +3044,7 @@ index 9642e58..4a9adc5 100644 // radix select of the `need`-th largest cell key uint prefix = 0, pmask = 0, need = QBUDGET; for (int pass = 0; pass < 4; ++pass) { -@@ -206,28 +254,35 @@ kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [ +@@ -206,28 +263,35 @@ kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [ if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); threadgroup_barrier(mem_flags::mem_threadgroup); for (uint b = tid; b < nb; b += 1024) { @@ -2934,7 +3090,7 @@ index 9642e58..4a9adc5 100644 } // T = prefix; take every cell with key > T and the first `need` cells (lowest index) with key == T, ascending. // Pass over blocks in order of index with a prefix scan of each block's selected cell count. -@@ -237,36 +292,35 @@ kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [ +@@ -237,36 +301,35 @@ kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [ const uint b = b0 + tid; uint w = 0, weq = 0; if (b < nb) { @@ -2991,7 +3147,7 @@ index 9642e58..4a9adc5 100644 } } -@@ -293,19 +347,28 @@ kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* +@@ -293,19 +356,28 @@ kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* return; } const int nc = min(QCH, int(width) - c0); @@ -3024,7 +3180,7 @@ index 9642e58..4a9adc5 100644 d = simd_sum(d); if (lane == 0) s[h][c] = d * 0.0625f; } -@@ -329,12 +392,15 @@ kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* +@@ -329,12 +401,15 @@ kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* threadgroup_barrier(mem_flags::mem_threadgroup); // acc[h][j] = sum_c p[h][c] v[c][j]: thread = dim j float acc[HPG]; @@ -3040,7 +3196,7 @@ index 9642e58..4a9adc5 100644 for (int h = 0; h < HPG; ++h) { device float* o = pp + (g * HPG + h) * (QD + 2); o[tid] = acc[h]; -@@ -342,11 +408,30 @@ kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* +@@ -342,11 +417,30 @@ kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* } } @@ -3072,7 +3228,7 @@ index 9642e58..4a9adc5 100644 const int h = int(tg.x), t = int(tg.y); const int nchunks = (QBUDGET + QCH - 1) / QCH; const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); -@@ -364,21 +449,6 @@ kernel void qsa_merge(device const float* part [[buffer(0)]], device const float +@@ -364,21 +458,6 @@ kernel void qsa_merge(device const float* part [[buffer(0)]], device const float o = o / L; const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; const float y = o * (1.0f / (1.0f + exp(-gate))); @@ -3096,6 +3252,114 @@ index 9642e58..4a9adc5 100644 + qsa_write_prepared(y, xp, a, t, h, tid); } + +diff --git a/metal/kernels/qsa2.metal b/metal/kernels/qsa2.metal +index 8ae189b..3417cfe 100644 +--- a/metal/kernels/qsa2.metal ++++ b/metal/kernels/qsa2.metal +@@ -133,37 +133,76 @@ kernel void qsa_topk2(device const float* scores [[buffer(0)]], device int* ids + for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); + return; + } ++ if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); ++ threadgroup_barrier(mem_flags::mem_threadgroup); + uint prefix = 0, pmask = 0, need = QBUDGET; +- for (int pass = 0; pass < 4; ++pass) { +- const int shift = 24 - 8 * pass; +- if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); +- threadgroup_barrier(mem_flags::mem_threadgroup); +- for (uint b = tid; b < nb; b += 1024) { +- const uint k = okey(sc[b]); +- if ((k & pmask) == prefix) { +- const uint w = b < n_bid ? 4u : (n_kv & 3u); +- if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); +- } ++ if (nb <= 4096) { ++ uint keys[4]; ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 4; ++i) { ++ const uint b = tid + 1024u * uint(i); ++ keys[i] = b < nb ? okey(sc[b]) : 0u; + } +- threadgroup_barrier(mem_flags::mem_threadgroup); +- // thread i < 256 holds bin 255 - i: an exclusive prefix over i = the cells of the higher digits +- uint cnt = 0, ex = 0; +- if (tid < 256) { +- cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); +- ex = simd_prefix_exclusive_sum(cnt); +- if (lane == 31) sgs[sg] = ex + cnt; ++ for (int pass = 0; pass < 4; ++pass) { ++ const int shift = 24 - 8 * pass; ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 4; ++i) { ++ const uint b = tid + 1024u * uint(i); ++ if (b >= nb) break; ++ const uint k = keys[i]; ++ if ((k & pmask) == prefix) { ++ const uint w = b < n_bid ? 4u : (n_kv & 3u); ++ if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); ++ } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ uint cnt = 0, ex = 0; ++ if (tid < 256) { ++ cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); ++ if (pass < 3) atomic_store_explicit(&hist[255 - tid], 0u, memory_order_relaxed); ++ ex = simd_prefix_exclusive_sum(cnt); ++ if (lane == 31) sgs[sg] = ex + cnt; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (tid < 256) { ++ uint above = ex; ++ for (int i = 0; i < int(sg); ++i) above += sgs[i]; ++ if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ prefix |= sel[0] << shift; ++ pmask |= 255u << shift; ++ need = sel[1]; + } +- threadgroup_barrier(mem_flags::mem_threadgroup); +- if (tid < 256) { +- uint above = ex; +- for (int i = 0; i < int(sg); ++i) above += sgs[i]; +- if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } ++ } else { ++ for (int pass = 0; pass < 4; ++pass) { ++ const int shift = 24 - 8 * pass; ++ for (uint b = tid; b < nb; b += 1024) { ++ const uint k = okey(sc[b]); ++ if ((k & pmask) == prefix) { ++ const uint w = b < n_bid ? 4u : (n_kv & 3u); ++ if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); ++ } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ uint cnt = 0, ex = 0; ++ if (tid < 256) { ++ cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); ++ if (pass < 3) atomic_store_explicit(&hist[255 - tid], 0u, memory_order_relaxed); ++ ex = simd_prefix_exclusive_sum(cnt); ++ if (lane == 31) sgs[sg] = ex + cnt; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (tid < 256) { ++ uint above = ex; ++ for (int i = 0; i < int(sg); ++i) above += sgs[i]; ++ if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ prefix |= sel[0] << shift; ++ pmask |= 255u << shift; ++ need = sel[1]; + } +- threadgroup_barrier(mem_flags::mem_threadgroup); +- prefix |= sel[0] << shift; +- pmask |= 255u << shift; +- need = sel[1]; +- threadgroup_barrier(mem_flags::mem_threadgroup); + } + const uint T = prefix; + const uint per = (nb + 1023) / 1024, b0 = min(nb, uint(tid) * per), b1 = min(nb, b0 + per); diff --git a/metal/kernels/qsa_prompt.metal b/metal/kernels/qsa_prompt.metal index df7caf7..a9fd0c7 100644 --- a/metal/kernels/qsa_prompt.metal @@ -3644,6 +3908,56 @@ index e5c7527..4571eac 100644 if (tid == 0) atomic_store_explicit(counter, 0u, memory_order_relaxed); } +diff --git a/metal/kernels/units.metal b/metal/kernels/units.metal +index 15f584c..25d0761 100644 +--- a/metal/kernels/units.metal ++++ b/metal/kernels/units.metal +@@ -107,26 +107,33 @@ template <> struct X { + } + }; + +-// xprep: one thread per 16 consecutive values of one token row; XP = [4 kinds][n_tok][n] then S16 [n_tok][n / 16] ++// xprep: one thread per 16 consecutive values of one token row; XP = [4 kinds][n_tok][n] then S16 [n_tok][n / 16]. ++// raw and plane are 16-float aligned: the four planes are written as float4s (mtp.metal's mtp_prepare, bit for bit) + kernel void xprep(device const float* x [[buffer(0)]], device float* xp [[buffer(1)]], constant uint& n [[buffer(2)]], + constant uint& n_tok [[buffer(3)]], uint gid [[thread_position_in_grid]]) { + const uint per = n / 16; + if (gid >= per * n_tok) return; + const uint t = gid / per, i = gid % per; +- device const float* src = x + t * n + 16 * i; ++ device const float4* src = (device const float4*)(x + t * n + 16 * i); + const uint plane = n_tok * n; + device float* raw = xp + t * n + 16 * i; +- float v[16]; +- load16(src, v); ++ device float4* r0 = (device float4*)raw; ++ device float4* r1 = (device float4*)(raw + plane); ++ device float4* r2 = (device float4*)(raw + 2 * plane); ++ device float4* r3 = (device float4*)(raw + 3 * plane); ++ const float4 q2lo = float4(1.0f, 0.25f, 0.0625f, 0.015625f), ++ q2hi = float4(0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f), ++ p4 = float4(1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f); + float s0 = 0, s1 = 0; +- for (int j = 0; j < 16; j += 2) { s0 += v[j]; s1 += v[j + 1]; } +- const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; +- const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; +- for (int j = 0; j < 16; ++j) { +- raw[j] = v[j]; +- raw[plane + j] = v[j] * q2s[j & 7]; +- raw[2 * plane + j] = v[j] * p4s[j & 3]; +- raw[3 * plane + j] = (j & 1) ? v[j] * 0.00390625f : v[j]; ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 4; ++j) { ++ const float4 v = src[j]; ++ s0 += v.x; s0 += v.z; ++ s1 += v.y; s1 += v.w; ++ r0[j] = v; ++ r1[j] = v * ((j & 1) ? q2hi : q2lo); ++ r2[j] = v * p4; ++ r3[j] = float4(v.x, v.y * 0.00390625f, v.z, v.w * 0.00390625f); + } + xp[4 * plane + t * per + i] = s0 + s1; + } diff --git a/metal/src/engine/arena_budget.h b/metal/src/engine/arena_budget.h index 4cc024b..513a944 100644 --- a/metal/src/engine/arena_budget.h @@ -5248,7 +5562,7 @@ index 0000000..edc7cd6 + +} // namespace strata::spec diff --git a/metal/src/engine/window.cpp b/metal/src/engine/window.cpp -index 9a6d1d3..c17ddd0 100644 +index 9a6d1d3..4fa92e9 100644 --- a/metal/src/engine/window.cpp +++ b/metal/src/engine/window.cpp @@ -22,15 +22,21 @@ @@ -5977,16 +6291,16 @@ index 9a6d1d3..c17ddd0 100644 + GrArgs a{N, HC, LR, (uint32_t) ntok, apply ? 1u : 0u, eps_}; + mtl::Dispatch d; + if (comb) { -+ static_assert(N == 2560 && HC == 4, "gr_norm_comb_4s shape"); ++ static_assert(N == 2560 && HC == 4, "gr_norm_comb shape"); + const uint32_t tk = TOPK; -+ d.p = P(lib_, "gr_norm_comb_4s", {}, "gr_norm_comb_4s"); ++ d.p = P(lib_, "gr_norm_comb", {}, "gr_norm_comb"); + d.args = {mtl::Arg::B(R), mtl::Arg::B(ye_), mtl::Arg::B(inj), mtl::Arg::B(wb, w_off), mtl::Arg::B(xn_), + mtl::Arg::V(a), mtl::Arg::B(plan_), mtl::Arg::B(ysh_), mtl::Arg::B(sg_), mtl::Arg::V(tk)}; + } else { + d.p = P(lib_, "gr_norm", {}, "gr_norm"); + d.args = {mtl::Arg::B(R), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wb, w_off), mtl::Arg::B(xn_), mtl::Arg::V(a)}; + } -+ d.grid[0] = comb ? 1 : HC; ++ d.grid[0] = HC; + d.grid[1] = ntok; + d.tg[0] = 256; + push(ds, d); @@ -9253,8 +9567,62 @@ index 33ce450..0e34814 100644 - return 0; + return dq_fails ? 1 : 0; } +diff --git a/metal/tests/expert_parity.cpp b/metal/tests/expert_parity.cpp +index 3bd15ba..2f561f4 100644 +--- a/metal/tests/expert_parity.cpp ++++ b/metal/tests/expert_parity.cpp +@@ -2,6 +2,7 @@ + // reference, and their speed. + // + // expert_parity [--gu Q2_0] [--dn Q2_0] [--ntok 1..8] [--gucfg L,R,S] [--dncfg L,R,S] parity (synthetic) ++// [--same] every token routed to the same 10 experts (each group has ntok entries) + // expert_parity --bench [--gu ..] [--dn ..] [--ntok N] [--layers 8] [--reps 96] [--sweep] speed + // + // Parity: random valid blocks for 64 experts, random activations x ~ U(-1,1), random distinct top-10 per token, +@@ -69,7 +70,7 @@ Layout layout(int gu, int dn) { + } + + void random_routing(ref::Rng& r, int ntok, int n_expert, std::vector& ids, std::vector& w, +- double share = 0.0, const std::vector* prev = nullptr) { ++ double share = 0.0, const std::vector* prev = nullptr, bool same = false) { + ids.assign(ntok * TOPK, 0); + w.assign(ntok * TOPK, 0.0f); + for (int t = 0; t < ntok; ++t) { +@@ -78,6 +79,7 @@ void random_routing(ref::Rng& r, int ntok, int n_expert, std::vector& ids, + for (int k = 0; k < TOPK; ++k) { + int e; + do { ++ if (same && t > 0) { e = ids[(k + t) % TOPK]; break; } // token 0's experts, rotated + // with probability `share`, reuse an expert of the previous token (verify windows share experts) + if (t > 0 && share > 0 && r.uni() < share) e = ids[(t - 1) * TOPK + (r.next() % TOPK)]; + else e = (int)(r.next() % n_expert); +@@ -192,6 +194,7 @@ int main(int argc, char** argv) { + Cfg cgu{32, 2, 4}, cdn{8, 2, 4}; + bool bench = false, sweep = false; + double share = 0.0; ++ bool same = false; + bool cfg_gu_set = false, cfg_dn_set = false; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; +@@ -207,6 +210,7 @@ int main(int argc, char** argv) { + else if (a == "--reps") reps = std::atoi(nx()); + else if (a == "--share") share = std::atof(nx()); + else if (a == "--dn-prep") g_dn_prep = true; ++ else if (a == "--same") same = true; + else { std::fprintf(stderr, "unknown argument %s\n", a.c_str()); return 2; } + } + (void)cfg_gu_set; +@@ -273,7 +277,7 @@ int main(int argc, char** argv) { + } + std::vector ids; + std::vector w; +- random_routing(rng, ntok, NE, ids, w, 0.3); ++ random_routing(rng, ntok, NE, ids, w, 0.3, nullptr, same); + mtl::Buffer bid = cx.alloc(4 * ids.size()), bw = cx.alloc(4 * w.size()), plan = cx.alloc(4 * P_INTS); + std::memcpy(bid.ptr, ids.data(), 4 * ids.size()); + std::memcpy(bw.ptr, w.data(), 4 * w.size()); diff --git a/metal/tests/fuse_parity.cpp b/metal/tests/fuse_parity.cpp -index e67d9e5..1cc4fe1 100644 +index e67d9e5..2b5dbc8 100644 --- a/metal/tests/fuse_parity.cpp +++ b/metal/tests/fuse_parity.cpp @@ -1,4 +1,4 @@ @@ -9263,7 +9631,7 @@ index e67d9e5..1cc4fe1 100644 // tokens: // gemv_cat (gemv.metal) ssm_alpha + ssm_beta (BF16 2560 -> 48 + 48) in one dispatch == two gemv dispatches // router_route the router GEMV + moe_route in one dispatch == gemv + moe_route (logits, ids, w, plan; -@@ -18,7 +18,7 @@ struct MoeArgs { +@@ -18,12 +18,94 @@ struct MoeArgs { uint64_t gate_off, up_off, down_off, gu_stride, d_stride; uint32_t gu_row, d_row, n_embd, n_ff, n_tok, topk, n_expert, gx_gu, gx_dn, pad; }; @@ -9271,8 +9639,96 @@ index e67d9e5..1cc4fe1 100644 +constexpr int P_INTS = 16 + 128 + 128 + 4 + 128 * 3 + 20 + 128, P_TOKENT = 16 + 128 + 128 + 4 + 128 * 2; struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; struct GlueArgs { uint32_t n, n_tok, n_embd, topk; }; ++ ++// The A/B source holds baseline and candidate entry points in one library. ++int gr_ab(const char* source, int nt, bool down) { ++ if (nt < 1 || nt > 8) return 2; ++ mtl::Context cx; ++ std::printf("device %s\n", cx.device_name().c_str()); ++ const int lib = down ? cx.compile({"quant.metal", "units.metal", source}) ++ : cx.compile({source}, "// NO_GGML_COMMON"); ++ constexpr int N = 2560, HC = 4, LR = 320, TK = 10; ++ ref::Rng rng(1802); ++ auto floats = [&](size_t n) { ++ auto b = cx.alloc(4 * n); ++ for (size_t i = 0; i < n; ++i) b.as()[i] = rng.sym(); ++ return b; ++ }; ++ auto R0 = floats(nt * HC * N), R1 = floats(nt * HC * N), R2 = floats(nt * HC * N); ++ auto x1 = floats(nt * HC * N), x2 = floats(nt * HC * N); ++ GrArgs a{N, HC, LR, (uint32_t)nt, 1, 1e-6f}; ++ int fails = 0; ++ auto bench = [&](mtl::Dispatch b, mtl::Dispatch v, const char* name) { ++ cx.warm({b}, 0.3); cx.warm({v}, 0.1); ++ for (int pair = 0; pair < 5; ++pair) { ++ double tb, tv; ++ if (pair % 2) { tv = cx.run({v}, 64); tb = cx.run({b}, 64); } ++ else { tb = cx.run({b}, 64); tv = cx.run({v}, 64); } ++ std::printf("bench %s ntok %d pair %d base_us %.6f variant_us %.6f\n", name, nt, pair, ++ tb * 1e6 / 64, tv * 1e6 / 64); ++ } ++ }; ++ if (down) { ++ const size_t rb = ref::row_bytes(ref::Q8_0, HC * N); ++ auto wd = cx.alloc(rb * LR), wi = cx.alloc(rb * HC); ++ ref::random_blocks(ref::Q8_0, wd.as(), LR, HC * N, 1803, 0.05f); ++ ref::random_blocks(ref::Q8_0, wi.as(), HC, HC * N, 1804, 0.05f); ++ auto lo1 = floats(nt * LR), lo2 = floats(nt * LR), in1 = floats(nt * HC), in2 = floats(nt * HC); ++ auto pb = cx.pipeline(lib, "baseline_down", mtl::Constants().i(1, nt)); ++ auto pv = cx.pipeline(lib, "optimized_down", mtl::Constants().i(1, nt)); ++ mtl::Dispatch b, v; ++ b.p = &pb; v.p = &pv; ++ b.args = {mtl::Arg::B(wd), mtl::Arg::B(wi), mtl::Arg::B(R0), mtl::Arg::B(lo1), mtl::Arg::B(in1), mtl::Arg::V(a)}; ++ v.args = {mtl::Arg::B(wd), mtl::Arg::B(wi), mtl::Arg::B(R0), mtl::Arg::B(lo2), mtl::Arg::B(in2), mtl::Arg::V(a)}; ++ b.grid[0] = v.grid[0] = (LR + HC) / 2; ++ b.tg[0] = v.tg[0] = 256; ++ cx.run({b}); cx.run({v}); ++ const bool ok = std::memcmp(lo1.ptr, lo2.ptr, lo1.bytes) == 0 && ++ std::memcmp(in1.ptr, in2.ptr, in1.bytes) == 0; ++ fails += !ok; ++ std::printf("down ntok %d bitwise %s\n", nt, ok ? "PASS" : "FAIL"); ++ bench(b, v, "down"); ++ } else { ++ auto ye = floats(nt * TK * N), ysh = floats(nt * N), sgv = floats(nt), inj = floats(nt * HC), wn = floats(HC * N); ++ auto plan = cx.alloc(4 * P_INTS); ++ for (int i = 0; i < nt * TK; ++i) plan.as()[P_TOKENT + i] = (i * 7) % (nt * TK); ++ auto pb = cx.pipeline(lib, "baseline_4s"); ++ for (const char* name : {"gr_norm_comb_4s", "gr_norm_comb", "gr_norm_comb_staged"}) { ++ auto pv = cx.pipeline(lib, name); ++ for (uint32_t tk : {1u, 10u}) for (uint32_t apply : {0u, 1u}) { ++ a.apply = apply; ++ mtl::Dispatch b, v; ++ auto args = [&](const mtl::Buffer& R, const mtl::Buffer& x) { ++ return std::vector{mtl::Arg::B(R), mtl::Arg::B(ye), mtl::Arg::B(inj), mtl::Arg::B(wn), ++ mtl::Arg::B(x), mtl::Arg::V(a), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::V(tk)}; ++ }; ++ b.p = &pb; v.p = &pv; b.args = args(R1, x1); v.args = args(R2, x2); ++ b.grid[0] = 1; v.grid[0] = std::strcmp(name, "gr_norm_comb") == 0 ? HC : 1; ++ b.grid[1] = v.grid[1] = nt; b.tg[0] = v.tg[0] = 256; ++ std::memcpy(R1.ptr, R0.ptr, R0.bytes); std::memcpy(R2.ptr, R0.ptr, R0.bytes); ++ cx.run({b}); cx.run({v}); ++ const bool ok = std::memcmp(R1.ptr, R2.ptr, R0.bytes) == 0 && ++ std::memcmp(x1.ptr, x2.ptr, x1.bytes) == 0; ++ fails += !ok; ++ std::printf("%s ntok %d topk %u apply %u R/xn bitwise %s\n", ++ name, nt, tk, apply, ok ? "PASS" : "FAIL"); ++ if (apply && tk == TK) bench(b, v, name); ++ } ++ } ++ } ++ return fails ? 1 : 0; ++} ++ } // namespace -@@ -37,7 +37,7 @@ int main() { + +-int main() { ++int main(int argc, char** argv) { ++ if (argc == 5 && std::string(argv[1]) == "--gr-ab") ++ return gr_ab(argv[2], std::stoi(argv[3]), std::string(argv[4]) == "down"); + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal", "gemv.metal", "router.metal", "gr.metal", + "gr2.metal", "glue.metal"}); +@@ -37,7 +119,7 @@ int main() { ref::random_blocks(ref::BF16, Wb.as(), NO, N, 2, 0.05f); ref::Rng r(3); for (size_t i = 0; i < xp.bytes / 4; ++i) xp.as()[i] = r.sym(); @@ -9281,7 +9737,7 @@ index e67d9e5..1cc4fe1 100644 mtl::Pipeline pg = cx.pipeline(lib, "gemv_30_32_1_4_1", mtl::Constants().i(0, nt)), pc = cx.pipeline(lib, "gemvc_30_32_1_4", mtl::Constants().i(0, nt)); const size_t s16 = (size_t)4 * 4 * MAXT * N; -@@ -75,7 +75,7 @@ int main() { +@@ -75,7 +157,7 @@ int main() { w2 = cx.alloc(4 * MAXT * 10), pl1 = cx.alloc(4 * P_INTS), pl2 = cx.alloc(4 * P_INTS), cnt = cx.alloc(16); ref::random_blocks(rty, Wr.as(), NE, N, 9, 0.05f); mtl::Pipeline pr = cx.pipeline(lib, "moe_route"); @@ -9290,7 +9746,7 @@ index e67d9e5..1cc4fe1 100644 mtl::Pipeline pg = cx.pipeline(lib, "gemv_" + std::to_string(rty) + "_32_1_4_1", mtl::Constants().i(0, nt)), pf = cx.pipeline(lib, "router_route_" + std::to_string(rty), mtl::Constants().i(0, nt)); const size_t s16 = (size_t)4 * 4 * MAXT * N; -@@ -131,12 +131,14 @@ int main() { +@@ -131,12 +213,14 @@ int main() { for (int i = 0; i < MAXT * HC; ++i) inj.as()[i] = r.sym() * 3; for (int i = 0; i < HC * N; ++i) wn.as()[i] = 0.5f + r.uni(); for (int i = 0; i < MAXT * TOPK; ++i) plan.as()[P_TOKENT + i] = (i * 7) % (MAXT * TOPK); @@ -9308,7 +9764,7 @@ index e67d9e5..1cc4fe1 100644 mtl::Dispatch c; c.p = &pcm; c.args = {mtl::Arg::B(ye), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::B(bo), mtl::Arg::V(ga)}; c.grid[0] = (nt * N + 255) / 256; c.tg[0] = 256; -@@ -145,14 +147,15 @@ int main() { +@@ -145,14 +229,15 @@ int main() { n1.grid[0] = HC; n1.grid[1] = nt; n1.tg[0] = 256; cx.run({c, n1}); const uint32_t tk = TOPK; -- knotmirror 0.0.1