diff --git a/audits/argon_r8_dq.md b/audits/argon_r8_dq.md new file mode 100644 index 0000000..1fed2e8 --- /dev/null +++ b/audits/argon_r8_dq.md @@ -0,0 +1,43 @@ +# Round 8 Track C: Barrier-Free Constant-Table `dq_f16` + Vectorized `act_f16` + +## 1. Summary of Changes (`metal/kernels/dense_mm.metal` & `metal/tests/dense_mm_parity.cpp`) + +1. **Zero-Threadgroup-Memory, Barrier-Free `dq_f16` (`DWC`):** + - Previously, `dq_f16` allocated `threadgroup uint tab[fmt::TAB]`, invoked `load_tab(tab, tid, 64)`, and executed `threadgroup_barrier(mem_flags::mem_threadgroup)` in **every single threadgroup** (`(n_in / UNIT / 64) * n_out` threadgroups per matrix, e.g., `30,720` threadgroups for a `2560 x 10240` `IQ3_S` matrix), even though each thread decodes at most **one** unit (`16` or `32` weights) before exiting: + - For `T_IQ3_S` (`TAB = 544`, `2,176` bytes static threadgroup memory), each thread loaded `8` `uint`s from `iq3s_grid` + `2` `uchar`s from `ksigns_iq2xs` (which `deq` does not even use!) into threadgroup memory and waited on a threadgroup barrier just to perform `4` grid lookups in `deq`. + - For `T_IQ4_NL` and `T_IQ4_XS` (`TAB = 256`, `1,024` bytes static threadgroup memory), each thread constructed `4` `half2` codebook entries in threadgroup memory and waited on a barrier. + - For all other 10 formats (`TAB = 1`), `tab` was unused (`static_tg_mem = 0`), yet every threadgroup still executed `threadgroup_barrier(mem_flags::mem_threadgroup)`. + - Added `DWC` in `metal/kernels/dense_mm.metal` with direct constant-memory specializations for `T_IQ3_S` (`iq3s_grid`), `T_IQ4_NL` (`kvalues_iq4nl`), and `T_IQ4_XS` (`kvalues_iq4nl`), forwarding the remaining 10 formats to `DW::run(..., nullptr)`. + - Removed `threadgroup uint tab[fmt::TAB]`, `load_tab`, and `threadgroup_barrier` from `dq_f16`. Static threadgroup memory is now **0 bytes** across all 13 `dq_f16` specializations. + +2. **Vectorized `half4` Output Stores in `dq_f16`:** + - Since `H = UNIT / 2` is always a multiple of `4` (`H = 8` for `UNIT = 16`, `H = 16` for `UNIT = 32`) and chunk offsets `p0`, `p1` are always 4-element (`8`-byte) aligned, replaced the scalar `half` store loop with `half4` vector stores (`*(device half4*)(o + p0 + j)` and `*(device half4*)(o + p1 + j)`), cutting store instructions by `4x` (`4` `half4` stores instead of `16` scalar stores for `UNIT = 16`; `8` `half4` stores instead of `32` scalar stores for `UNIT = 32`). + +3. **Vectorized `float4 -> half4` `act_f16`:** + - Updated `act_f16` to convert 4 elements per thread via `((device half4*)xh)[gid] = half4(((device const float4*)x)[gid])` with a tail loop at `gid == (n >> 2)` for any remainder elements. This remains backward-compatible with existing host dispatch grids `(n + 255) / 256` while allowing `(n + 1023) / 1024` dispatches. + +## 2. Bitwise Parity & Benchmark Results on Apple M5 GPU + +All 13 formats (`F32`, `BF16`, `Q8_0`, `Q4_0`, `Q5_0`, `IQ4_NL`, `IQ4_XS`, `Q2_0`, `IQ3_S`, `Q3_K`, `Q4_K`, `Q5_K`, `Q6_K`) were verified against `ref::dequant` + `ref::f32_to_f16` with `std::memcmp` in both `dense_mm_parity` (`3.65 s`) and `dense_mm_parity --bench` (`1.48 s`): +- **100% bitwise identity (`memcmp == 0`)** across all 13 formats and all tested matrix shapes (`2560 x 10240`, `2560 x 12288`, `2560 x 6144`, `10240 x 324`, plus `2560/640/6144 -> 100/48`). + +| Format | `UNIT` | `static_tg_mem` (Old -> New) | `2560 x 10240` (`us`) | `2560 x 12288` (`us`) | `2560 x 6144` (`us`) | `10240 x 324` (`us`) | `memcmp` | +|---|---:|---:|---:|---:|---:|---:|---:| +| `F32` | 16 | `0 B -> 0 B` | 1191.7 | 1436.9 | 677.9 | 124.2 | `0` (PASS) | +| `BF16` | 16 | `0 B -> 0 B` | 777.6 | 936.1 | 447.0 | 85.7 | `0` (PASS) | +| `Q8_0` | 32 | `0 B -> 0 B` | 582.1 | 699.9 | 324.8 | 62.9 | `0` (PASS) | +| `Q4_0` | 32 | `0 B -> 0 B` | 521.3 | 625.4 | 269.5 | 48.4 | `0` (PASS) | +| `Q5_0` | 32 | `0 B -> 0 B` | 533.3 | 640.8 | 290.4 | 51.6 | `0` (PASS) | +| `IQ4_NL` | 32 | `1024 B -> 0 B` | 520.4 | 624.3 | 268.0 | 49.3 | `0` (PASS) | +| `IQ4_XS` | 32 | `1024 B -> 0 B` | 516.1 | 626.3 | 262.6 | 47.4 | `0` (PASS) | +| `Q2_0` | 16 | `0 B -> 0 B` | 460.0 | 571.6 | 208.3 | 41.1 | `0` (PASS) | +| `IQ3_S` | 16 | `2176 B -> 0 B` | 495.6 | 602.1 | 237.0 | 43.2 | `0` (PASS) | +| `Q3_K` | 32 | `0 B -> 0 B` | 506.7 | 615.3 | 243.4 | 44.8 | `0` (PASS) | +| `Q4_K` | 32 | `0 B -> 0 B` | 529.5 | 637.2 | 274.6 | 46.9 | `0` (PASS) | +| `Q5_K` | 32 | `0 B -> 0 B` | 542.1 | 649.9 | 291.0 | 49.0 | `0` (PASS) | +| `Q6_K` | 32 | `0 B -> 0 B` | 547.7 | 659.1 | 297.8 | 56.0 | `0` (PASS) | + +- **Note on DRAM bandwidth saturation on M5 (`~100–130 GB/s`):** For large matrices (`2560 x 10240` writes `52.4 MB` of `fp16` output and reads `14.1 MB` of `IQ3_S` weights = `66.5 MB` per matrix), `495 us` corresponds to **`134 GB/s` effective DRAM bandwidth**, which is at the M5 hardware memory bandwidth ceiling. On smaller shapes (`2560 x 48` `ssm_alpha`/`ssm_beta`), removing the `2,176`-byte threadgroup table load + barrier in `dq_f16_21` (`IQ3_S`) speeds up each dispatch by **`1.56x`** (`3.3 us -> 2.1 us`). + +## 3. Artifacts +- Clean unified diff: `/tmp/argon_r8_dq/argon_r8_dq.patch` (verified with `git -C vendor/strata apply --check`). diff --git a/audits/opus_r8_moemm.md b/audits/opus_r8_moemm.md new file mode 100644 index 0000000..90a3c9a --- /dev/null +++ b/audits/opus_r8_moemm.md @@ -0,0 +1,36 @@ +# opus r8 track A: mm_gu3 register epilogue + W8 decoder micro-opts + +patch: `/tmp/opus_r8_moemm/opus_r8_moemm.patch` (applies with `git -C vendor/strata apply`, `--check` clean). touches only +`metal/kernels/moe_mm.metal` and `metal/tests/moe_mm_parity.cpp`. + +## method +- scratch copy `/tmp/opus_r8_moemm/strata`, own build-metal; kernel variants as separate kernel dirs (`/tmp/opus_r8_moemm/k_*`) selected with `STRATA_METAL_KERNELS`. +- `moe_mm_parity` gains `--dump ` (HS + y bytes, skips the slow CPU reference -> bitwise A/B of two kernel versions in <0.1 s) and `--gu-bench` (times only the gate/up GEMM, ~2 s an invocation). +- all runs: `--v3 --tile 16,32,32,4` (the production mm_gu3 tile), `--dn Q2_0`, M5. gate/up time a layer, us, min of 3 interleaved runs (other agents were on the GPU at times; medians noisy). + +## Target 1: drop `threadgroup float ot[SG][TM*TN]` (8 KiB) from mm_gu3 -> REJECTED +- direct register stores via `accg.get_multidimensional_index(i)` (scalar half) and a coop_k4 half4 variant: both **bitwise identical** to baseline for IQ2_XS/IQ2_S/IQ2_XXS/IQ3_S/IQ3_XXS (T=256, HS+y). the half4 one confirms the M5 dst layout: elements i..i+3 = 4 consecutive f at one m, 8-byte aligned. +- but slower: IQ2_XS T=2048 base 23973 / scalar 24131 / half4 24136 us (+0.7%, clean GPU, 3 reps, spread <0.1%). +- other tiles with half4: 16,64,32,4 +10%; 16,64,32,2 -7%; 32,32,32,2 -3%; 16,32,32,5 / 16,16,32,8 / 16,16,32,4 within ±0.4%. none beats 16,32,32,4 baseline. +- reading: mm_gu3 is decode-bound and threadgroup memory is not the binding occupancy limit (registers likely are); extra TGs/core where it does change do not help (16,64,32,4 gets worse, probably cache contention). not included in the patch. + +## Target 2: W8 micro-opts +| change | result | kept | +|---|---|---| +| W8: scales folded (l<2: sc0 = scb&15, sc1 = scb>>4) | T=2048 23963 -> 23343 (-2.6%) | yes | +| W8: same fold | 29208 -> 29281 (+0.25%, noise) | no | +| W8: popcount signs instead of ksg tab | erases the fold gain (23947) | no | +| W8: popcount signs instead of ksg tab | 24704 -> 24489 (-0.9%), sweep below -1.8..-2.3% | yes | +| W8: existing MM_PROBE=4 popcount | inconclusive under GPU contention (min -2.3%, medians ±5%) | no change | + +all kept changes bitwise identical (all 5 IQ formats, HS + y dumps). + +## final sweep (gate/up us, min of 3, base -> patched) +| fmt | T=512 | T=1024 | T=2048 | +|---|---|---|---| +| IQ2_XS | 8375 -> 8383 (+0.1%) | 14526 -> 14058 (-3.2%) | 24166 -> 23629 (-2.2%) | +| IQ2_XXS | 8954 -> 8935 (-0.2%) | 15134 -> 14786 (-2.3%) | 26058 -> 25598 (-1.8%) | +| IQ2_S (unchanged code, control) | -0.1% | +0.0% | +0.1% | +| IQ3_S (unchanged code, control) | +0.1% | +0.3% | +2.3% (contended) | + +expected prefill effect: ~2% of mm_gu3 IQ2_XS (806.7 ms) + ~2% of IQ2_XXS (67.9 ms) ~= 17-19 ms at T=1983 (~0.4% of total prefill). small but free and bitwise. diff --git a/audits/sol_r8_grvec.md b/audits/sol_r8_grvec.md new file mode 100644 index 0000000..4141988 --- /dev/null +++ b/audits/sol_r8_grvec.md @@ -0,0 +1,61 @@ +# round 8 track b: gr address/memory experiment + +plan: compare current gr kernels with precomputed 32-bit expert offsets, ordered slot compaction, and residual reread versus register caching. owner: prompt.metal gr helpers only; no host dispatch/interface changes. fma and reduction order must remain unchanged, dropped slots must skip fma entirely. uint offsets are safe for the supported chunk geometry (maximum 230,686,720 elements). + +evaluator: comb_fold_parity short gr-only mode, fixed seed 23+T, N=2560 hc=4 topk=10 with dropped slots, T=1,7,64,512,1024,2048; memcmp residual and half normalization output against original dispatch chains. timing four dispatches per sample, three samples; confirmation interleaves baseline/candidate processes. select >3% repeatable wins; reject any bit mismatch. each standalone process capped at 5s, no full model or server use. archive source variants and logs under /tmp/sol_r8_grvec. no tracked vendor edits. + +## implementation and evidence + +scratch input snapshot: git `81b76e4` in `/tmp/sol_r8_grvec/strata` (only the two owned files tracked). this snapshots the actual dirty working inputs, not vendor HEAD. source variants are archived as `/tmp/sol_r8_grvec/{offset,compact,reread,both,stream,stream_orig,unroll,unrolld,candidate}.metal`; baseline is `/tmp/sol_r8_grvec/baseline/prompt.metal`. all measurements use apple m5, runtime safe-math default, existing `build-metal/libmtl_rt.a` and clang++ -std=c++20 -O2. + +candidate keeps `bv[10]` for bo, then updates/reduces/emits one residual stream at a time with `rv[10]`: nominal live cached float payload 20 rather than 40, plus fewer stream accumulators/gates. this is a source-level register-pressure argument, not a measured hardware register count. reduction order within each stream and per-element fma order stay unchanged. the shared new tail avoids repeating the stream normalization in both kernels. `moe_comb`, `moe_comb_slots`, `gr_apply_moec`, `gr_norm_h_moe`, host dispatch and gemm_hc_up are unchanged in the candidate. + +### rejected or inconclusive variants + +- precomputed uint row offsets: initial separate-process apparent ~5% gain disappeared after warm-up. same-process paired T2048 median ms4: 13.8539 candidate vs13.8124 baseline. not retained. +- order-preserving compaction via first-simdgroup exclusive prefix and sum: median 13.8535 vs13.8479; no win, not retained. dropped slots removed without adding zero-weight fmas. +- residual reread instead of rv: moec essentially flat, h4 apply1 ~1% slower, apply0 ~5% slower. not retained. +- topk10 unroll: no clear improvement over stream candidate, not retained. +- fully unrolled d traversal: contention and no demonstrated incremental win, not retained. + +### exploratory stream result before final helper cleanup + +paired timing format is candidate/baseline milliseconds for four dispatches; three samples alternate AB/BA. median values below, not kprof or whole-model performance: + +| source/log | T | moec | h4 apply1 | h4 apply0 | +|---|---:|---:|---:|---:| +| paired_stream_2048.log (stream+uint) |2048|13.5182/13.8329 (-2.3%)|6.9600/7.1963 (-3.3%)|3.7201/4.0299 (-7.7%)| +| confirm_stream_1024.log (stream+uint)|1024|6.8074/6.9979 (-2.7%)|3.5236/3.6321 (-3.0%)|1.8177/2.0026 (-9.2%)| +| paired_stream_orig_512.log (stream, no uint)|512|4.2696/4.2976 (-0.7%)|1.7997/1.8695 (-3.7%)|0.8299/1.0031 (-17.3%)| + +these are leads, not isolated confirmation. the moec change does not meet the planned >3% threshold; h4 does in these exploratory runs. candidate intentionally drops the neutral uint/compaction change. `gr_apply_moec` shows no stable improvement and is unchanged. + +### correctness and process budget + +final source: `candidate.metal`, current scratch `metal/kernels/prompt.metal`. `/tmp/sol_r8_grvec/final_{1,7,64,512,1024,2048}.log`: all six processes exited 0, all four comparisons bitwise PASS (R and fp16 output, h4 apply0 and apply1), memcmp against original dispatch chains; random scattered routes include pos=-1. final process durations 0.084–1.022s, below5s. every standalone run in this investigation was capped at4.9s and completed below2s. no server access or full-model command used. + +all final timing samples are contaminated: unchanged `gr_apply_moec` jumps from ~12 to ~300ms4 across the investigation and final1024/final512 show broad slowdowns. do not use final timing logs to claim a speedup. parent and the two active sibling tracks were asked for a short exclusive slot; clean final confirmation pending. + +reproduction: + +``` +clang++ -std=c++20 -O2 /tmp/sol_r8_grvec/strata/metal/tests/comb_fold_parity.cpp \ + -I/tmp/sol_r8_grvec/strata/metal/src -I/tmp/sol_r8_grvec/strata/metal/tests \ + vendor/strata/build-metal/libmtl_rt.a -framework Metal -framework Foundation \ + -o /tmp/sol_r8_grvec/comb_fold_parity +STRATA_METAL_KERNELS=/tmp/sol_r8_grvec/strata/metal/kernels \ + /tmp/sol_r8_grvec/comb_fold_parity 2048 grvec /tmp/sol_r8_grvec/baseline/prompt.metal +``` + +patch: `/tmp/sol_r8_grvec/sol_r8_grvec.patch`; `git -C vendor/strata apply --check` passed. scratch `git diff --check` passed. final visual diff/context review completed; removed redundant braces and duplicate comments and centralized the normalization tail. no tracked vendor files edited. + +## final short confirmation (after parent authorization) + +`confirmation_2048.log` and `confirmation_1024.log`, exit0, wall0.689s and0.323s; all bitwise checks pass again. opus track explicitly paused. final T1024 paired medians (candidate/baseline ms4): + +- moec: 6.7921/6.9860 = -2.8% (below planned3% acceptance threshold; marginal). +- h4 apply1: 3.4993/3.6128 = -3.1%. +- h4 apply0: 1.8099/2.0053 = -9.7%. +- unchanged apply_moec: 6.0877/6.1005 = -0.2% (noise control). + +T2048 h4 apply1 6.9608/7.1618 = -2.8%, apply0 3.7630/4.0195 = -6.4%; moec has warm-up/contention drift (17.7 ->14.1ms4) and is inconclusive there. no model-wide extrapolation claimed. parent plans fully serialized isolated confirmation once all patches arrive. recommend testing this candidate, not declaring a completed production speedup; h4 is the stronger lead, moec remains marginal and may be dropped independently. diff --git a/scripts/strata_metal_complete.diff b/scripts/strata_metal_complete.diff index 36d065d..4efe1c3 100644 --- a/scripts/strata_metal_complete.diff +++ b/scripts/strata_metal_complete.diff @@ -248,10 +248,50 @@ index d23a26f..d1c5be4 100644 +target_compile_definitions(heal_k256_q2_0 PRIVATE ACCELERATE_NEW_LAPACK) +target_link_libraries(heal_k256_q2_0 PRIVATE mtl_rt "-framework Accelerate") diff --git a/metal/kernels/dense_mm.metal b/metal/kernels/dense_mm.metal -index 45e29ac..11b0827 100644 +index 45e29ac..15c8584 100644 --- a/metal/kernels/dense_mm.metal +++ b/metal/kernels/dense_mm.metal -@@ -67,3 +67,6 @@ template [[host_name("gemm_f16_64_64_4")]] kernel decltype(gemm_f16<64, 64, 4>) +@@ -8,7 +8,12 @@ using namespace mpp::tensor_ops; + + kernel void act_f16(device const float* x [[buffer(0)]], device half* xh [[buffer(1)]], constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { +- if (gid < n) xh[gid] = half(x[gid]); ++ const uint n4 = n >> 2; ++ if (gid < n4) { ++ ((device half4*)xh)[gid] = half4(((device const float4*)x)[gid]); ++ } else if (gid == n4) { ++ for (uint i = n4 << 2; i < n; ++i) xh[i] = half(x[i]); ++ } + } + + struct DqArgs { +@@ -25,8 +30,10 @@ kernel void dq_f16(device const uchar* W [[buffer(0)]], device half* Wh [[buffer + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, H = UNIT / 2; + threadgroup uint tab[fmt::TAB]; +- load_tab(tab, tid, 64); +- threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (fmt::TAB > 1) { ++ load_tab(tab, tid, 64); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } + const int u = int(tg.x) * 64 + tid, r = int(tg.y); + if (u * UNIT >= int(a.n_in)) return; + // the unit's chunk origin: a 128-wide chunk holds whole units of every format +@@ -37,9 +44,9 @@ kernel void dq_f16(device const uchar* W [[buffer(0)]], device half* Wh [[buffer + const int q = a.perm ? int(a.n_out / a.perm) : 1; + const int ro = a.perm ? (r % q) * int(a.perm) + r / q : r; + device half* o = Wh + ulong(ro) * a.n_in + k0; +- for (int j = 0; j < H; ++j) { +- o[p0 + j] = half(w[j]); +- o[p1 + j] = half(w[H + j]); ++ for (int j = 0; j < H; j += 4) { ++ *(device half4*)(o + p0 + j) = half4(half(w[j]), half(w[j + 1]), half(w[j + 2]), half(w[j + 3])); ++ *(device half4*)(o + p1 + j) = half4(half(w[H + j]), half(w[H + j + 1]), half(w[H + j + 2]), half(w[H + j + 3])); + } + } + +@@ -67,3 +74,6 @@ template [[host_name("gemm_f16_64_64_4")]] kernel decltype(gemm_f16<64, 64, 4>) template [[host_name("gemm_f16_128_64_4")]] kernel decltype(gemm_f16<128, 64, 4>) gemm_f16<128, 64, 4>; template [[host_name("gemm_f16_64_128_4")]] kernel decltype(gemm_f16<64, 128, 4>) gemm_f16<64, 128, 4>; template [[host_name("gemm_f16_64_32_4")]] kernel decltype(gemm_f16<64, 32, 4>) gemm_f16<64, 32, 4>; @@ -1238,7 +1278,7 @@ index bb8fdbb..2943861 100644 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; diff --git a/metal/kernels/moe_mm.metal b/metal/kernels/moe_mm.metal -index ff28897..4602f6d 100644 +index ff28897..7791c5d 100644 --- a/metal/kernels/moe_mm.metal +++ b/metal/kernels/moe_mm.metal @@ -35,7 +35,7 @@ kernel void mm_hist(device const int* ids [[buffer(0)]], device uint* bcnt [[buf @@ -1419,7 +1459,7 @@ index ff28897..4602f6d 100644 } }; template <> struct W4 { // 256 a block: d, scales_h, scales_l[4], qs[128]; 32-groups as IQ4_NL's -@@ -629,16 +678,16 @@ template <> struct W8 { +@@ -629,16 +678,17 @@ template <> struct W8 { } }; template <> struct W8 { @@ -1436,31 +1476,34 @@ index ff28897..4602f6d 100644 - const uint s0 = ks[(a1 >> (7 * l)) & 127], s1 = ks[(a1 >> (7 * (l + 2))) & 127]; - a = w4_signed(tab[2 * i0 + hh], hh ? s0 >> 4 : s0, sc); - b = w4_signed(tab[2 * i1 + hh], hh ? s1 >> 4 : s1, sc); -+ const uint s0 = ksg(tab, 512, (a1 >> (7 * l)) & 127), s1 = ksg(tab, 512, (a1 >> (7 * (l + 2))) & 127); ++ const uint j0 = (a1 >> (7 * l)) & 127, j1 = (a1 >> (7 * (l + 2))) & 127; // the signs computed, as W8's probe 4 ++ const uint s0 = j0 | ((popcount(j0) & 1u) << 7), s1 = j1 | ((popcount(j1) & 1u) << 7); + a = w4_signed(g_iq2xxs(tab, i0, hh), hh ? s0 >> 4 : s0, sc); + b = w4_signed(g_iq2xxs(tab, i1, hh), hh ? s1 >> 4 : s1, sc); } }; template <> struct W8 { -@@ -647,14 +696,13 @@ template <> struct W8 { +@@ -647,14 +697,13 @@ template <> struct W8 { const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; device const ushort* q = (device const ushort*)(bl + 2) + 4 * ib; const uint v0 = q[l], v1 = q[l + 2]; - threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 1024); const uint scb = bl[66 + ib]; const float d = rdh(bl); - const float sc0 = d * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f; - const float sc1 = d * (0.5f + float((scb >> (4 * ((l + 2) >> 1))) & 15)) * 0.25f; +- const float sc0 = d * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f; +- const float sc1 = d * (0.5f + float((scb >> (4 * ((l + 2) >> 1))) & 15)) * 0.25f; - const uint s0 = ks[v0 >> 9], s1 = ks[v1 >> 9]; - a = w4_signed(tab[2 * (v0 & 511) + hh], hh ? s0 >> 4 : s0, sc0); - b = w4_signed(tab[2 * (v1 & 511) + hh], hh ? s1 >> 4 : s1, sc1); ++ const float sc0 = d * (0.5f + float(scb & 15)) * 0.25f; // l < 2 (k % 32 < 16): runs l, l + 2 take nibbles 0, 1 ++ const float sc1 = d * (0.5f + float(scb >> 4)) * 0.25f; + const uint s0 = ksg(tab, 1024, v0 >> 9), s1 = ksg(tab, 1024, v1 >> 9); + a = w4_signed(g_iq2xs(tab, v0 & 511, hh), hh ? s0 >> 4 : s0, sc0); + b = w4_signed(g_iq2xs(tab, v1 & 511, hh), hh ? s1 >> 4 : s1, sc1); } }; template <> struct W8 { -@@ -798,54 +846,72 @@ kernel void mm_gu3(device const uchar* W [[buffer(0)]], device const half* XS [[ +@@ -798,54 +847,72 @@ kernel void mm_gu3(device const uchar* W [[buffer(0)]], device const half* XS [[ constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); matmul2d op; tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); @@ -1570,7 +1613,7 @@ index ff28897..4602f6d 100644 } for (uint16_t i = 0; i < accg.get_capacity(); ++i) if (accg.is_valid_element(i)) { -@@ -1016,6 +1082,8 @@ MM4_GU(T_Q2_0) MM4_GU(T_IQ3_XXS) MM4_GU(T_IQ2_XXS) MM4_GU(T_IQ2_XS) MM4_GU(T_IQ2 +@@ -1016,6 +1083,8 @@ MM4_GU(T_Q2_0) MM4_GU(T_IQ3_XXS) MM4_GU(T_IQ2_XXS) MM4_GU(T_IQ2_XS) MM4_GU(T_IQ2 template kernel void mm_gu5(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], @@ -1579,7 +1622,7 @@ index ff28897..4602f6d 100644 uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { constexpr int TM = 16, TK = 32, BB = fmt::BB, ROWS = SG * TN, RU = BB / 2; -@@ -1025,8 +1093,9 @@ kernel void mm_gu5(device const uchar* W [[buffer(0)]], device const half* XS [[ +@@ -1025,8 +1094,9 @@ kernel void mm_gu5(device const uchar* W [[buffer(0)]], device const half* XS [[ load_tab(tab, tid, SG * 32); const int2 tl = tiles[tg.y]; const int e = tl.x, row0 = tl.y, nt0 = int(tg.x) * ROWS, n0 = nt0 + sg * TN; @@ -1591,14 +1634,14 @@ index ff28897..4602f6d 100644 constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); matmul2d op; tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); -@@ -1086,4 +1155,4 @@ kernel void mm_gu5(device const uchar* W [[buffer(0)]], device const half* XS [[ +@@ -1086,4 +1156,4 @@ kernel void mm_gu5(device const uchar* W [[buffer(0)]], device const half* XS [[ #define MMGU5_(F, N, S) template [[host_name("mm_gu5_" #F "_" #N "_" #S)]] kernel decltype(mm_gu5) mm_gu5; #define MMGU5(F, N, S) MMGU5_(F, N, S) MMGU5(T_IQ3_S, 32, 2) MMGU5(T_IQ3_S, 32, 1) MMGU5(T_IQ3_S, 16, 4) MMGU5(T_IQ3_XXS, 32, 2) MMGU5(T_IQ3_XXS, 32, 1) -MMGU5(T_IQ3_XXS, 16, 4) MMGU5(T_IQ2_S, 32, 2) MMGU5(T_IQ2_S, 32, 1) MMGU5(T_IQ2_S, 16, 4) +MMGU5(T_IQ3_XXS, 16, 4) MMGU5(T_IQ2_S, 32, 2) MMGU5(T_IQ2_S, 32, 1) MMGU5(T_IQ2_S, 16, 4) MMGU5(T_IQ2_XXS, 32, 2) MMGU5(T_IQ2_XXS, 16, 4) diff --git a/metal/kernels/prompt.metal b/metal/kernels/prompt.metal -index 64833b0..8a1c6b4 100644 +index 64833b0..841148b 100644 --- a/metal/kernels/prompt.metal +++ b/metal/kernels/prompt.metal @@ -7,33 +7,64 @@ @@ -1732,7 +1775,7 @@ index 64833b0..8a1c6b4 100644 ss = fma(v, v, ss); } ss = simd_sum(ss); -@@ -138,40 +195,223 @@ kernel void gr_norm_h(device float* R [[buffer(0)]], device const float* bo [[bu +@@ -138,40 +195,237 @@ kernel void gr_norm_h(device float* R [[buffer(0)]], device const float* bo [[bu const float rs = rsqrt(tot / float(n) + a.eps); device half* o = xnh + (t * a.hc + c) * n; device const float* w = w_norm + c * n; @@ -1759,6 +1802,18 @@ index 64833b0..8a1c6b4 100644 + } +} + ++// stream-wise reduction keeps one residual stream plus bo live, rather than all four residual streams. ++inline void gr_norm_h_stream_out(thread const float* rv, device const float* w, device half* o, uint n, float eps, ++ float ss, threadgroup float* part, ushort tid, ushort sg, ushort lane) { ++ ss = simd_sum(ss); ++ if (lane == 0) part[sg] = ss; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ float tot = 0.0f; ++ for (int i = 0; i < 8; ++i) tot += part[i]; ++ const float rs = rsqrt(tot / float(n) + eps); ++ for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[i] * w[d] * rs); ++} ++ +// gr_norm_h4: gr_norm_h with one threadgroup per token over all hc streams, so bo is read once a token rather than +// hc times; R', xnh bitwise gr_norm_h's. grid (1, T), 256 threads; hc <= 4 +kernel void gr_norm_h4(device float* R [[buffer(0)]], device const float* bo [[buffer(1)]], @@ -1769,21 +1824,22 @@ index 64833b0..8a1c6b4 100644 + threadgroup float part[4][8]; + const uint t = tg.y, n = a.n_embd, hc = a.hc; + device float* r = R + t * hc * n; -+ float rv[4][10], g[4], ss[4] = {0.0f, 0.0f, 0.0f, 0.0f}; -+ for (uint c = 0; c < hc; ++c) g[c] = a.apply ? 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))) : 0.0f; -+ for (uint d = tid, i = 0; d < n; d += 256, ++i) { -+ const float b = a.apply ? bo[t * n + d] : 0.0f; -+ for (uint c = 0; c < hc; ++c) { ++ float bv[10]; ++ for (uint d = tid, i = 0; d < n; d += 256, ++i) bv[i] = a.apply ? bo[t * n + d] : 0.0f; ++ for (uint c = 0; c < hc; ++c) { ++ const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))) : 0.0f; ++ float rv[10], ss = 0.0f; ++ for (uint d = tid, i = 0; d < n; d += 256, ++i) { + float v = r[c * n + d]; + if (a.apply) { -+ v = fma(b, g[c], v); ++ v = fma(bv[i], g, v); + r[c * n + d] = v; + } -+ rv[c][i] = v; -+ ss[c] = fma(v, v, ss[c]); ++ rv[i] = v; ++ ss = fma(v, v, ss); + } ++ gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); + } -+ gr_norm_h_out(rv, w_norm, xnh, t, hc, n, a.eps, ss, part, tid, sg, lane); +} + +// gr_norm_h_moe: gr_norm_h with the sublayer output bo formed on the fly as moe_out - bo = fma(sigmoid(sg[t]), ys, yr) @@ -1849,19 +1905,20 @@ index 64833b0..8a1c6b4 100644 + const uint t = tg.y, n = a.n_embd, hc = a.hc; + moe_comb_slots(pos, srt_w, t, topk, rk, wk, tid); + device float* r = R + t * hc * n; -+ float rv[4][10], g[4], ss[4] = {0.0f, 0.0f, 0.0f, 0.0f}; -+ for (uint c = 0; c < hc; ++c) g[c] = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); ++ float bv[10]; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); -+ for (uint d = tid, i = 0; d < n; d += 256, ++i) { -+ const float bo = fma(gs, ys[t * n + d], moe_comb(YE, rk, wk, topk, n, d)); -+ for (uint c = 0; c < hc; ++c) { -+ const float v = fma(bo, g[c], r[c * n + d]); ++ for (uint d = tid, i = 0; d < n; d += 256, ++i) bv[i] = fma(gs, ys[t * n + d], moe_comb(YE, rk, wk, topk, n, d)); ++ for (uint c = 0; c < hc; ++c) { ++ const float g = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); ++ float rv[10], ss = 0.0f; ++ for (uint d = tid, i = 0; d < n; d += 256, ++i) { ++ const float v = fma(bv[i], g, r[c * n + d]); + r[c * n + d] = v; -+ rv[c][i] = v; -+ ss[c] = fma(v, v, ss[c]); ++ rv[i] = v; ++ ss = fma(v, v, ss); + } ++ gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); + } -+ gr_norm_h_out(rv, w_norm, xnh, t, hc, n, a.eps, ss, part, tid, sg, lane); +} + +// gr_apply_moe: the last layer's MoE write applied to R alone (R' of moe_out + gr_norm; its xn has no reader). @@ -3836,7 +3893,7 @@ index 0000000..edc7cd6 + +} // namespace strata::spec diff --git a/metal/src/engine/window.cpp b/metal/src/engine/window.cpp -index 9a6d1d3..ca3b4c7 100644 +index 9a6d1d3..71b8fd8 100644 --- a/metal/src/engine/window.cpp +++ b/metal/src/engine/window.cpp @@ -22,15 +22,21 @@ @@ -5200,9 +5257,9 @@ index 9a6d1d3..ca3b4c7 100644 + std::memcpy(mtail_[0].ptr, c.tails.data() + tail_.size() * ts, ts); + std::memcpy(mdead_[0].ptr, c.dead.data() + dead_.size() * dsz, dsz); + } -+ return true; -+} -+ + return true; + } + +namespace { +// one contiguous range of a buffer in the session file +struct SessionSec { @@ -5346,9 +5403,9 @@ index 9a6d1d3..ca3b4c7 100644 + ss_->ple_prev[0] = h.ple_prev[0]; + ss_->ple_prev[1] = h.ple_prev[1]; + n_cells = h.n_cells; - return true; - } - ++ return true; ++} ++ +bool Window::scr_relocate(int64_t c_src, int64_t c_dst, int64_t n, std::string& err) { + if (mrope_on()) { err = "metal window: SCR of M-RoPE (vision) cells is not supported"; return false; } + if (n < 0 || c_src < 0 || c_dst < 0 || c_src + n >= max_cells_ || c_dst + n >= max_cells_) { err = "metal window: SCR span out of range"; return false; } @@ -6194,7 +6251,7 @@ index 9a6d1d3..ca3b4c7 100644 // ---- the embeddings into every stream, the PLE rows { const gguf::Tensor& te = TN("token_embd.weight"); -@@ -2133,17 +3256,34 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri +@@ -2133,19 +3256,36 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri for (int c = 1; c < HC; ++c) std::memcpy(r + (size_t) c * N, r, (size_t) N * 4); } } @@ -6230,8 +6287,11 @@ index 9a6d1d3..ca3b4c7 100644 } + const double tb1 = now_ms(); std::vector ds; - std::vector cuts; +- std::vector cuts; ++ std::vector cuts, sigs; nb_ = false; + auto D = [&](const mtl::Pipeline* pp, std::vector a, uint32_t gx, uint32_t gy, uint32_t tg, uint32_t gz = 1) { + mtl::Dispatch d; @@ -2159,10 +3299,16 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri return P(lib, fn, c, "pg:" + (key.empty() ? fn : key)); }; @@ -6252,7 +6312,7 @@ index 9a6d1d3..ca3b4c7 100644 }; auto dqb = [&](const mtl::Buffer& src, size_t off, int type, int n_in, int n_out, const mtl::Buffer& out, size_t out_off = 0, uint32_t perm = 0) { -@@ -2172,7 +3318,7 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri +@@ -2172,97 +3318,149 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri }; auto dq = [&](const gguf::Tensor& t, const mtl::Buffer& out, size_t out_off = 0, uint32_t perm = 0) { const int n_in = (int) t.shape[0], n_out = t.shape.size() > 1 ? (int) t.shape[1] : 1; @@ -6261,8 +6321,9 @@ index 9a6d1d3..ca3b4c7 100644 else dqb(W_, t.offset, t.type, n_in, n_out, out, out_off, perm); }; auto act16 = [&](const mtl::Buffer& x, const mtl::Buffer& xh, size_t n) { -@@ -2180,89 +3326,141 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri - D(PP(p.lib, "act_f16"), {mtl::Arg::B(x), mtl::Arg::B(xh), mtl::Arg::V(nn)}, (nn + 255) / 256, 1, 256); + const uint32_t nn = (uint32_t) n; +- D(PP(p.lib, "act_f16"), {mtl::Arg::B(x), mtl::Arg::B(xh), mtl::Arg::V(nn)}, (nn + 255) / 256, 1, 256); ++ D(PP(p.lib, "act_f16"), {mtl::Arg::B(x), mtl::Arg::B(xh), mtl::Arg::V(nn)}, (nn + 1023) / 1024, 1, 256); }; int ci = 0; - auto hc_read = [&](const std::string& pre, bool apply) { @@ -6422,7 +6483,7 @@ index 9a6d1d3..ca3b4c7 100644 D(PP(p.lqsa, "qsa_flash", mtl::Constants().b(7, p.k4), "qsa_flash"), {mtl::Arg::B(p.qcurh), mtl::Arg::B(kq_[(size_t) qi]), mtl::Arg::B(vq_[(size_t) qi]), mtl::Arg::B(ksc_[(size_t) qi]), mtl::Arg::B(vsc_[(size_t) qi]), mtl::Arg::B(p.qids), mtl::Arg::B(p.qfull), mtl::Arg::B(p.xqh), mtl::Arg::V(a)}, -@@ -2287,26 +3485,45 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri +@@ -2287,26 +3485,46 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri // layer 0's attention write applied to R happens in the FFN read below; the PLE block comes after the FFN } // ---- the MoE @@ -6463,6 +6524,7 @@ index 9a6d1d3..ca3b4c7 100644 + mtl::Arg::B(st_->smap(), (size_t) ln * 512 * 4), mtl::Arg::V(p_moce)}, (T + 7) / 8, 1, 256); + } + } ++ if (stream_ && comb_fold) sigs.push_back((int) ds.size()); const int nent = T * TOPK, nblk = (nent + 1023) / 1024; - PgPlan pa{(uint32_t) nent, TOPK, NEXP, (uint32_t) p.TM, (uint32_t) nblk}; - const uint32_t nclr = (uint32_t) p.maxrows; @@ -6476,7 +6538,7 @@ index 9a6d1d3..ca3b4c7 100644 D(PP(p.lib, "mm_place"), {mtl::Arg::B(p.ids), mtl::Arg::B(p.wts), mtl::Arg::B(p.boff), mtl::Arg::B(p.eoff), mtl::Arg::B(p.stok), mtl::Arg::B(p.sw), mtl::Arg::B(p.pos), mtl::Arg::V(pa)}, nblk, 1, 1024); PgMm ma{eg.offset, TN(b + "ffn_up_exps.weight").offset, ed.offset, 0, 0, (uint32_t) ref::row_bytes(eg.type, N), -@@ -2317,8 +3534,8 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri +@@ -2317,13 +3535,24 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri const uint32_t ne_ = N; { mtl::Dispatch g; @@ -6487,27 +6549,61 @@ index 9a6d1d3..ca3b4c7 100644 g.tg[0] = 256; g.indirect = &p.meta; g.ind_off = 8 * 4; -@@ -2371,75 +3588,86 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri + push(ds, g); + } ++ if (comb_fold) { ++ nb_ = true; ++ gemm(p.mixh, p.Wshg, p.gsh, T, NFF, N); ++ nb_ = true; ++ gemm(p.mixh, p.Wshu, p.ush, T, NFF, N); ++ nb_ = true; ++ gemm(p.mixh, p.Wsg, p.sgv, T, 1, N); ++ const uint32_t nh = (uint32_t) T * NFF; ++ D(PP(p.lib, "swiglu_h"), {mtl::Arg::B(p.gsh), mtl::Arg::B(p.ush), mtl::Arg::B(p.hh), mtl::Arg::V(nh)}, (nh + 255) / 256, 1, 256); ++ gemm(p.hh, p.Wshd, p.ys, T, N, NFF); ++ } + const mtl::Buffer& EW = stream_ ? st_->chunk(0) : W_; + auto arena_args = [&](mtl::Dispatch& d) { + if (!stream_) return; +@@ -2349,7 +3578,10 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri + arena_args(gu); + gu.indirect = &p.meta; + gu.ind_off = 2 * 4; +- if (stream_) cuts.push_back((int) ds.size()); // the host's turn: this layer's experts put in place ++ if (stream_) { ++ if (!comb_fold) sigs.push_back((int) ds.size()); ++ cuts.push_back((int) ds.size()); // the host's turn: this layer's experts put in place ++ } + push(ds, gu); + mtl::Dispatch dn; + if (p.dn_v3) { +@@ -2371,75 +3603,87 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri dn.ind_off = 5 * 4; push(ds, dn); const uint32_t ca[4] = {(uint32_t) T, N, TOPK, 0}; - D(PP(p.lib, "mm_combine"), {mtl::Arg::B(p.YE), mtl::Arg::B(p.pos), mtl::Arg::B(p.sw), mtl::Arg::B(p.yr), mtl::Arg::V(ca)}, - (T * N + 255) / 256, 1, 256); -+ if (!comb_fold) -+ D(PP(p.lib, "mm_combine"), {mtl::Arg::B(p.YE), mtl::Arg::B(p.pos), mtl::Arg::B(p.sw), mtl::Arg::B(p.yr), mtl::Arg::V(ca)}, -+ (T * N + 255) / 256, 1, 256); - gemm(p.mixh, p.Wshg, p.gsh, T, NFF, N); -+ nb_ = true; - gemm(p.mixh, p.Wshu, p.ush, T, NFF, N); -+ nb_ = true; -+ gemm(p.mixh, p.Wsg, p.sgv, T, 1, N); - const uint32_t nh = (uint32_t) T * NFF; - D(PP(p.lib, "swiglu_h"), {mtl::Arg::B(p.gsh), mtl::Arg::B(p.ush), mtl::Arg::B(p.hh), mtl::Arg::V(nh)}, (nh + 255) / 256, 1, 256); - gemm(p.hh, p.Wshd, p.ys, T, N, NFF); +- gemm(p.mixh, p.Wshg, p.gsh, T, NFF, N); +- gemm(p.mixh, p.Wshu, p.ush, T, NFF, N); +- const uint32_t nh = (uint32_t) T * NFF; +- D(PP(p.lib, "swiglu_h"), {mtl::Arg::B(p.gsh), mtl::Arg::B(p.ush), mtl::Arg::B(p.hh), mtl::Arg::V(nh)}, (nh + 255) / 256, 1, 256); +- gemm(p.hh, p.Wshd, p.ys, T, N, NFF); - gemm(p.mixh, p.Wsg, p.sgv, T, 1, N); - const uint32_t oa[2] = {(uint32_t) T, N}; - D(PP(p.lib, "moe_out"), {mtl::Arg::B(p.yr), mtl::Arg::B(p.ys), mtl::Arg::B(p.sgv), mtl::Arg::B(p.bo), mtl::Arg::V(oa)}, - (T * N + 255) / 256, 1, 256); ++ if (!comb_fold) { ++ D(PP(p.lib, "mm_combine"), {mtl::Arg::B(p.YE), mtl::Arg::B(p.pos), mtl::Arg::B(p.sw), mtl::Arg::B(p.yr), mtl::Arg::V(ca)}, ++ (T * N + 255) / 256, 1, 256); ++ gemm(p.mixh, p.Wshg, p.gsh, T, NFF, N); ++ nb_ = true; ++ gemm(p.mixh, p.Wshu, p.ush, T, NFF, N); ++ nb_ = true; ++ gemm(p.mixh, p.Wsg, p.sgv, T, 1, N); ++ const uint32_t nh = (uint32_t) T * NFF; ++ D(PP(p.lib, "swiglu_h"), {mtl::Arg::B(p.gsh), mtl::Arg::B(p.ush), mtl::Arg::B(p.hh), mtl::Arg::V(nh)}, (nh + 255) / 256, 1, 256); ++ gemm(p.hh, p.Wshd, p.ys, T, N, NFF); ++ } + // no moe_out: the next read forms the MoE write from yr (or YE), ys, sgv (gr_norm_h_moe(c), ple_moe_gate_norm(_c), + // gr_apply_moe(c)) if (l == 0 && ple_on_) { @@ -6570,9 +6666,10 @@ index 9a6d1d3..ca3b4c7 100644 if (stream_) { seg_T_ = T; cut_ids_ = p.ids.as(); +- cx_->run_graph_seg(ds, cuts, [](void* u, int s) { ((Window*) u)->at_cut(s); }, this); + st_->set_prompt_pred(pred_on ? p.ids2.as() : nullptr, TOPK, + pred_on && p.ids3.ptr ? p.ids3.as() : nullptr); - cx_->run_graph_seg(ds, cuts, [](void* u, int s) { ((Window*) u)->at_cut(s); }, this); ++ cx_->run_graph_seg(ds, cuts, sigs, [](void* u, int s) { ((Window*) u)->at_cut(s); }, nullptr, this); + st_->set_prompt_pred(nullptr, 0); + if (ple_th.joinable()) ple_th.join(); // only if the graph stopped before layer 0's handshake + ple_join_ = nullptr; @@ -6614,18 +6711,20 @@ index 9a6d1d3..ca3b4c7 100644 mtp_src_ = &p.R; // the draft layer's prompt K/V read this chunk's final rows return true; } -@@ -2479,6 +3707,10 @@ bool window_init(void* w, const strata::core::ModelGeometry& g, strata::core::Se +@@ -2479,6 +3723,12 @@ bool window_init(void* w, const strata::core::ModelGeometry& g, strata::core::Se if (!x->init(g, ss, g_shard1, MAXT, max_cells, err)) return false; window_register(w); if (!mtp_load_pending(x, err)) return false; + if (x->streaming() && !(std::getenv("STRATA_METAL_LAZY_ARENA") && std::atoi(std::getenv("STRATA_METAL_LAZY_ARENA")) != 0)) { ++ const char* esep = std::getenv("STRATA_METAL_SEPARATE_PROMPT_BUF"); ++ const bool sep_pg = !(esep && std::string(esep) == "0"); + if (!x->prompt_gemm_init(err)) return false; -+ x->prompt_region(false); ++ if (!sep_pg) x->prompt_region(false); + } } x->set_max_t(max_t); *n_vocab = x->vocab(); -@@ -2494,6 +3726,7 @@ bool window_run(void* w, int T, const int32_t* tokens, int64_t pos0, bool head, +@@ -2494,6 +3744,7 @@ bool window_run(void* w, int T, const int32_t* tokens, int64_t pos0, bool head, static const bool logpos = std::getenv("STRATA_LOGPOS") != nullptr; return ((Window*) w)->run(T, tokens, pos0, head || logpos, out, err); } @@ -6633,7 +6732,7 @@ index 9a6d1d3..ca3b4c7 100644 void window_reset() { if (g_the_window != nullptr && g_the_window->ready()) g_the_window->reset(); } -@@ -2505,6 +3738,18 @@ bool window_checkpoint_restore(const strata::core::ConversationCheckpoint& c, st +@@ -2505,6 +3756,18 @@ bool window_checkpoint_restore(const strata::core::ConversationCheckpoint& c, st if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; } return g_the_window->checkpoint_restore(c, err); } @@ -6652,7 +6751,7 @@ index 9a6d1d3..ca3b4c7 100644 bool window_commit(void* w, int n_keep, std::string& err) { const double t0 = now_ms(); const bool ok = ((Window*) w)->commit(n_keep, err); -@@ -2555,13 +3800,14 @@ bool prompt_run(const int64_t* tokens, int64_t n, int64_t pos0, const float* con +@@ -2555,13 +3818,14 @@ bool prompt_run(const int64_t* tokens, int64_t n, int64_t pos0, const float* con bool prompt_run_rows(const int64_t* tokens, int64_t n, int64_t pos0, const std::function& on_chunk, std::string& err) { if (g_window == nullptr) { err = "metal prompt: the decode window is not up yet"; return false; } @@ -6668,7 +6767,7 @@ index 9a6d1d3..ca3b4c7 100644 if (!(e && std::string(e) == "0") && n >= mn) { if (!g_window->prompt_gemm_init(err)) return false; const int C = g_window->prompt_chunk_rows(); -@@ -2571,7 +3817,9 @@ bool prompt_run_rows(const int64_t* tokens, int64_t n, int64_t pos0, +@@ -2571,7 +3835,9 @@ bool prompt_run_rows(const int64_t* tokens, int64_t n, int64_t pos0, for (int64_t c = 0; c < n && ok; c += C) { const int k = (int) std::min(C, n - c); for (int i = 0; i < k; ++i) tk32[(size_t) i] = (int32_t) tokens[c + i]; @@ -6775,10 +6874,10 @@ index 4c6045b..106657d 100644 [cb waitUntilCompleted]; diff --git a/metal/tests/comb_fold_parity.cpp b/metal/tests/comb_fold_parity.cpp new file mode 100644 -index 0000000..d2f8492 +index 0000000..9f6dcd5 --- /dev/null +++ b/metal/tests/comb_fold_parity.cpp -@@ -0,0 +1,248 @@ +@@ -0,0 +1,299 @@ +// metal/tests/comb_fold_parity.cpp - Round 5 Track A: the prompt chunk's folded mm_combine and the per-token HC read +// against the dispatch chains they replace, BITWISE; `comb_fold_parity [T] bench` adds their GPU times at T (default +// 1024), 48 reps = a chunk's layers: @@ -6786,6 +6885,7 @@ index 0000000..d2f8492 +// gr_apply_moec == mm_combine + gr_apply_moe (R') +// gr_norm_h4 == gr_norm_h (R', xnh; apply 0 and 1) +// gemm_hc_up, mixed unbound == gemm_hc_up's mixh ++// `comb_fold_parity T grvec [baseline-prompt.metal]` checks/times GR only in a short process. +// `comb_fold_parity T hcup [bench|confirm]` checks HC-up tiles alone (both outputs, null mixed too). +#include "mtl_rt.h" +#include "ref_quant.h" @@ -6904,9 +7004,28 @@ index 0000000..d2f8492 + if (argc > 2 && std::strcmp(argv[2], "hcup") == 0) + return hc_up_parity(uint32_t(std::atoi(argv[1])), argc > 3 && std::strcmp(argv[3], "bench") == 0, + argc > 3 && std::strcmp(argv[3], "confirm") == 0); ++ const bool gr_only = argc > 2 && std::strcmp(argv[2], "grvec") == 0; + mtl::Context cx; -+ const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "gr.metal", "prompt.metal"}); -+ const int lmm = cx.compile({"quant.metal", "units.metal", "mma.metal", "moe_mm.metal"}); ++ const std::string source = gr_only ? mtl::read_file(cx.kernel_dir + "/prompt.metal") : ""; ++ const std::string prelude = "#define NO_GGML_COMMON 1\nusing namespace metal;\n" ++ "struct GrArgs { uint n_embd, hc, hc_lr, n_tok, apply; float eps; };\n"; ++ const auto begin = source.find("// ---- the HC read"); ++ const auto end = source.find("// gemm_hc_up:"); ++ if (gr_only && (begin == std::string::npos || end == std::string::npos || end <= begin)) { ++ std::fprintf(stderr, "GR source markers missing\n"); ++ return 1; ++ } ++ const int lib = gr_only ? cx.compile({}, prelude + source.substr(begin, end - begin)) : ++ cx.compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "gr.metal", "prompt.metal"}); ++ const std::string mm = gr_only ? mtl::read_file(cx.kernel_dir + "/moe_mm.metal") : ""; ++ const auto mb = mm.find("kernel void mm_combine("); ++ const auto me = mm.find("\n}", mb == std::string::npos ? 0 : mb); ++ if (gr_only && (mb == std::string::npos || me == std::string::npos)) { ++ std::fprintf(stderr, "mm_combine source marker missing\n"); ++ return 1; ++ } ++ const int lmm = gr_only ? cx.compile({}, prelude + mm.substr(mb, me + 2 - mb)) : ++ cx.compile({"quant.metal", "units.metal", "mma.metal", "moe_mm.metal"}); + std::printf("device %s\n", cx.device_name().c_str()); + const uint32_t N = 2560, HC = 4, LR = 320, TOPK = 10, NEXP = 512, TMR = 16; + const size_t RN = (size_t) HC * N; @@ -6915,7 +7034,16 @@ index 0000000..d2f8492 + mtl::Pipeline pcomb = cx.pipeline(lmm, "mm_combine"), pnhm = cx.pipeline(lib, "gr_norm_h_moe"), + pnhc = cx.pipeline(lib, "gr_norm_h_moec"), pam = cx.pipeline(lib, "gr_apply_moe"), + pamc = cx.pipeline(lib, "gr_apply_moec"), pnh = cx.pipeline(lib, "gr_norm_h"), -+ pnh4 = cx.pipeline(lib, "gr_norm_h4"), pup = cx.pipeline(lib, "gemm_hc_up_64_64_4"); ++ pnh4 = cx.pipeline(lib, "gr_norm_h4"), pup = gr_only ? mtl::Pipeline{} : cx.pipeline(lib, "gemm_hc_up_64_64_4"); ++ std::vector baseline; ++ if (gr_only && argc > 3) { ++ const std::string ref = mtl::read_file(argv[3]); ++ const auto rb = ref.find("// ---- the HC read"), re = ref.find("// gemm_hc_up:"); ++ if (rb == std::string::npos || re == std::string::npos || re <= rb) return 1; ++ const int rl = cx.compile({}, prelude + ref.substr(rb, re - rb)); ++ for (const char* name : {"gr_norm_h_moec", "gr_apply_moec", "gr_norm_h4", "gr_norm_h4"}) ++ baseline.push_back(cx.pipeline(rl, name)); ++ } + auto D = [](const mtl::Pipeline& p, std::vector a, uint32_t gx, uint32_t gy, uint32_t th = 256) { + mtl::Dispatch d; + d.p = &p; @@ -6926,7 +7054,7 @@ index 0000000..d2f8492 + return d; + }; + int fails = 0; -+ for (uint32_t T : {1u, 7u, 64u, TB}) { ++ for (uint32_t T : gr_only ? std::vector{TB} : std::vector{1u, 7u, 64u, TB}) { + const size_t mr = (size_t) T * TOPK + NEXP * TMR; + mtl::Buffer R0 = cx.alloc(4 * T * RN), Ra = cx.alloc(4 * T * RN), Rb = cx.alloc(4 * T * RN), YE = cx.alloc(4 * mr * N), + pos = cx.alloc(4 * T * TOPK), sw = cx.alloc(4 * mr), yr = cx.alloc(4 * T * N), ys = cx.alloc(4 * T * N), @@ -7000,6 +7128,28 @@ index 0000000..d2f8492 + cx.run({nh(pnh4, Rb, xb, *g, 1)}); + check(g->apply ? "gr_norm_h4 (apply 1)" : "gr_norm_h4 (apply 0)", same(Ra, Rb, Ra.bytes) && same(xa, xb, xa.bytes)); + } ++ if (gr_only) { ++ const std::vector ds = { ++ nhc(Rb, xb), amc(Rb), nh(pnh4, Rb, xb, a, 1), nh(pnh4, Rb, xb, a0, 1)}; ++ for (size_t i = 0; i < ds.size(); ++i) { ++ const auto& d = ds[i]; ++ auto ref = d; ++ if (!baseline.empty()) ref.p = &baseline[i]; ++ cx.run({d}); ++ if (!baseline.empty()) cx.run({ref}); ++ std::printf("grvec %s apply %u ms4", d.p->name.c_str(), i == 3 ? 0u : 1u); ++ for (int sample = 0; sample < 3; ++sample) { ++ double ms, rms = 0.0; ++ if (!baseline.empty() && sample % 2 == 0) rms = cx.run({ref}, 4) * 1e3; ++ ms = cx.run({d}, 4) * 1e3; ++ if (!baseline.empty() && sample % 2 != 0) rms = cx.run({ref}, 4) * 1e3; ++ std::printf(" %.4f", ms); ++ if (!baseline.empty()) std::printf("/%.4f", rms); ++ } ++ std::printf("\n"); ++ } ++ continue; ++ } + cx.run({up(mixed, mha)}); + cx.run({up(nil, mhb)}); + check("gemm_hc_up, mixed unbound", same(mha, mhb, mha.bytes)); @@ -7027,6 +7177,116 @@ index 0000000..d2f8492 + std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); + return fails ? 1 : 0; +} +diff --git a/metal/tests/dense_mm_parity.cpp b/metal/tests/dense_mm_parity.cpp +index 33ce450..6f89541 100644 +--- a/metal/tests/dense_mm_parity.cpp ++++ b/metal/tests/dense_mm_parity.cpp +@@ -45,7 +45,7 @@ int main(int argc, char** argv) { + if (which & 1) { + const uint32_t n = (uint32_t)T * n_in; + mtl::Dispatch d; d.p = &pact; d.args = {mtl::Arg::B(x), mtl::Arg::B(xh), mtl::Arg::V(n)}; +- d.grid[0] = (n + 255) / 256; d.tg[0] = 256; ds.push_back(d); ++ d.grid[0] = (n + 1023) / 1024; d.tg[0] = 256; ds.push_back(d); + } + if (which & 2) { + DqArgs a{(uint32_t)n_in, (uint32_t)n_out, (uint32_t)ref::row_bytes(ty, n_in)}; +@@ -78,6 +78,11 @@ int main(int argc, char** argv) { + cx.run(ds); + std::vector wf((size_t)n_in * n_out); + ref::dequant(ty, W.ptr, wf.data(), wf.size()); ++ std::vector wh_ref(wf.size()), xh_ref((size_t)T * n_in); ++ for (size_t i = 0; i < wf.size(); ++i) wh_ref[i] = ref::f32_to_f16(wf[i]); ++ for (size_t i = 0; i < xh_ref.size(); ++i) xh_ref[i] = ref::f32_to_f16(x.as()[i]); ++ const int wh_cmp = std::memcmp(wh.ptr, wh_ref.data(), wh_ref.size() * 2); ++ const int xh_cmp = std::memcmp(xh.ptr, xh_ref.data(), xh_ref.size() * 2); + double se16 = 0, se = 0, sr = 0; + for (int t = 0; t < T; ++t) + for (int o = 0; o < n_out; ++o) { +@@ -93,47 +98,49 @@ int main(int argc, char** argv) { + sr += ae * ae; + } + const double r16 = std::sqrt(se16 / sr), re = std::sqrt(se / sr); +- const bool ok = r16 < 1e-5; ++ const bool ok = r16 < 1e-5 && wh_cmp == 0 && xh_cmp == 0; + fails += !ok; +- std::printf("dense mm %-7s %5d -> %4d, T %d: vs fp16 operands %.2e, vs exact %.2e %s\n", +- ref::type_name(ty), n_in, n_out, T, r16, re, ok ? "PASS" : "FAIL"); ++ std::printf("dense mm %-7s %5d -> %4d, T %d: vs fp16 operands %.2e, vs exact %.2e, memcmp(wh)=%d, memcmp(xh)=%d %s\n", ++ ref::type_name(ty), n_in, n_out, T, r16, re, wh_cmp, xh_cmp, ok ? "PASS" : "FAIL"); + } + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; + } +- struct Shape { const char* name; int type, n_in, n_out; }; +- const std::vector shapes = {{"gdn qkv", ref::IQ4_XS, 2560, 10240}, {"gdn qkv", ref::Q3_K, 2560, 10240}, +- {"gdn gate", ref::Q4_K, 2560, 6144}, {"gdn out", ref::Q6_K, 6144, 2560}, +- {"qsa q", ref::Q2_0, 2560, 12288}, {"hc down", ref::BF16, 10240, 324}, +- {"hc up", ref::BF16, 320, 10240}, {"router", ref::BF16, 2560, 512}, +- {"alpha", ref::BF16, 2560, 48}}; +- const int T = 2048; +- for (const Shape& s : shapes) { +- if (s.n_in % 128) { std::printf("%-9s %-6s: n_in %d not a multiple of 128 (needs the 32-chunk path)\n", s.name, ref::type_name(s.type), s.n_in); continue; } +- const size_t rb = ref::row_bytes(s.type, s.n_in); +- mtl::Buffer W = cx.alloc(rb * s.n_out), x = cx.alloc(4 * (size_t)T * s.n_in), xh = cx.alloc(2 * (size_t)T * s.n_in), +- wh = cx.alloc(2 * (size_t)s.n_out * s.n_in), y = cx.alloc(4 * (size_t)T * s.n_out); +- ref::random_blocks(s.type, W.as(), s.n_out, s.n_in, 3, 0.02f); +- auto timeit = [&](int which, const std::string& gk) { +- std::vector ds; +- for (int i = 0; i < 4; ++i) run_all(s.type, W, 0, s.n_in, s.n_out, T, x, xh, wh, y, gk, ds, which); +- cx.warm(ds, 0.2); +- std::vector t; +- for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds) / 4); +- return median(t); +- }; +- const double tdq = timeit(2, "gemm_f16_64_64_4"), tact = timeit(1, "gemm_f16_64_64_4"); +- double best = 1e9; +- std::string bk; +- for (const char* gk : {"gemm_f16_64_64_4", "gemm_f16_128_64_4", "gemm_f16_64_128_4", "gemm_f16_64_32_4"}) { +- const double tg = timeit(4, gk); +- if (tg < best) { best = tg; bk = gk; } ++ const size_t max_w = (size_t)4 * 2560 * 12288; ++ const size_t max_wh = (size_t)2 * 2560 * 12288; ++ mtl::Buffer W = cx.alloc(max_w), wh = cx.alloc(max_wh); ++ std::vector wf(2560 * 12288); ++ std::vector wh_ref(2560 * 12288); ++ struct DqShape { const char* name; int n_in, n_out; }; ++ const std::vector dq_shapes = {{"2560x10240", 2560, 10240}, {"2560x12288", 2560, 12288}, ++ {"2560x6144", 2560, 6144}, {"10240x324", 10240, 324}}; ++ int dq_fails = 0; ++ for (const auto& ds_sh : dq_shapes) { ++ std::printf("--- dq_f16 all 13 formats on %s (%d x %d): bitwise memcmp + speed ---\n", ds_sh.name, ds_sh.n_in, ds_sh.n_out); ++ for (int ty : types) { ++ const size_t rb = ref::row_bytes(ty, ds_sh.n_in); ++ ref::random_blocks(ty, W.as(), std::min(ds_sh.n_out, 64), ds_sh.n_in, 17 + ty, 0.02f); ++ for (int row = 64; row < ds_sh.n_out; row += 64) { ++ const int ncopy = std::min(64, ds_sh.n_out - row); ++ std::memcpy(W.as() + (size_t)row * rb, W.as(), (size_t)ncopy * rb); ++ } ++ mtl::Pipeline pdq = cx.pipeline(lib, "dq_f16_" + std::to_string(ty)); ++ DqArgs a{(uint32_t)ds_sh.n_in, (uint32_t)ds_sh.n_out, (uint32_t)rb, 0}; ++ mtl::Dispatch d; d.p = &pdq; d.args = {mtl::Arg::B(W), mtl::Arg::B(wh), mtl::Arg::V(a)}; ++ d.grid[0] = (ds_sh.n_in / unit_of(ty) + 63) / 64; d.grid[1] = ds_sh.n_out; d.tg[0] = 64; d.barrier = true; ++ std::vector ds(4, d); ++ cx.run_graph(ds); ++ const size_t nelem = (size_t)ds_sh.n_in * ds_sh.n_out; ++ ref::dequant(ty, W.ptr, wf.data(), nelem); ++ for (size_t i = 0; i < nelem; ++i) wh_ref[i] = ref::f32_to_f16(wf[i]); ++ const int cmp = std::memcmp(wh.ptr, wh_ref.data(), nelem * 2); ++ dq_fails += (cmp != 0); ++ std::vector ts; ++ for (int i = 0; i < 5; ++i) ts.push_back(cx.run_graph(ds) / 4); ++ std::printf("dq_f16 %-7s %-10s: %6.1f us, memcmp=%d %s\n", ++ ref::type_name(ty), ds_sh.name, median(ts) * 1e6, cmp, cmp == 0 ? "PASS" : "FAIL"); + } +- const double fl = 2.0 * T * s.n_in * s.n_out; +- std::printf("%-9s %-6s %5d -> %5d, T %d: dequant %7.1f us, act %6.1f us, GEMM %7.1f us (%5.1f TFLOP/s, %s) -> %.2f us a token\n", +- s.name, ref::type_name(s.type), s.n_in, s.n_out, T, tdq * 1e6, tact * 1e6, best * 1e6, fl / best / 1e12, +- bk.c_str(), (tdq + best) / T * 1e6); + } +- return 0; ++ return dq_fails ? 1 : 0; + } diff --git a/metal/tests/fuse_parity.cpp b/metal/tests/fuse_parity.cpp index e67d9e5..1cc4fe1 100644 --- a/metal/tests/fuse_parity.cpp @@ -7265,10 +7525,32 @@ index 24cb85b..432bbeb 100644 return 0; } diff --git a/metal/tests/moe_mm_parity.cpp b/metal/tests/moe_mm_parity.cpp -index 6f51f1e..224cb15 100644 +index 6f51f1e..63c4e80 100644 --- a/metal/tests/moe_mm_parity.cpp +++ b/metal/tests/moe_mm_parity.cpp -@@ -67,14 +67,18 @@ int main(int argc, char** argv) { +@@ -34,6 +34,8 @@ int main(int argc, char** argv) { + int gu = ref::Q2_0, dn = ref::Q2_0, T = 256, TM = 16, TN = 64, TK = 16, SG = 2, DTN = -1, DTK = -1, DSG = -1; + bool bench = false, v2 = false, v3 = false, v4 = false, v5 = false; + int MS = 1; ++ bool gu_only = false; ++ const char* dump = nullptr; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + if (a == "--gu") gu = type_of(argv[++i]); +@@ -42,10 +44,12 @@ int main(int argc, char** argv) { + else if (a == "--tile") std::sscanf(argv[++i], "%d,%d,%d,%d", &TM, &TN, &TK, &SG); + else if (a == "--dtile") std::sscanf(argv[++i], "%d,%d,%d", &DTN, &DTK, &DSG); // the down kernel's TN,TK,SG + else if (a == "--bench") bench = true; ++ else if (a == "--gu-bench") bench = gu_only = true; // --bench timing the gate/up GEMM alone + else if (a == "--v2") v2 = true; + else if (a == "--v3") v2 = v3 = true; // v2's gather + the register-operand GEMMs (TN per simdgroup) + else if (a == "--v4") { v2 = v3 = v4 = true; MS = std::atoi(argv[++i]); } // mm_gu4 with MS 16-row subtiles + else if (a == "--v5") { v2 = v3 = true; v5 = true; } // mm_gu5 (weight blocks staged in threadgroup memory) ++ else if (a == "--dump") dump = argv[++i]; // HS and y to a file: bitwise comparison of two kernel versions + } + if (bench && T == 256) T = 2048; + if (v4) TM = 16 * MS; // the plan's tile height; mm_gu4's name carries MS instead +@@ -67,14 +71,18 @@ int main(int argc, char** argv) { } const int nent = T * TOPK, nblk = (nent + 1023) / 1024; const int maxrows = nent + NE * TM; @@ -7289,7 +7571,7 @@ index 6f51f1e..224cb15 100644 for (int t = 0; t < T; ++t) { std::set used; float s = 0; -@@ -88,7 +92,7 @@ int main(int argc, char** argv) { +@@ -88,7 +96,7 @@ int main(int argc, char** argv) { for (int k = 0; k < TOPK; ++k) w.as()[t * TOPK + k] /= s; } mtl::Pipeline phist = cx.pipeline(lib, "mm_hist"), pscan = cx.pipeline(lib, "mm_scan"), pplace = cx.pipeline(lib, "mm_place"), @@ -7298,7 +7580,7 @@ index 6f51f1e..224cb15 100644 if (DTN < 0) { DTN = TN; DTK = TK; DSG = SG; } if (dn == ref::IQ4_NL && DTK < 32) DTK = 32; const std::string tl = std::to_string(TM) + "_" + std::to_string(TN) + "_" + std::to_string(TK) + "_" + std::to_string(SG); -@@ -105,15 +109,13 @@ int main(int argc, char** argv) { +@@ -105,15 +113,13 @@ int main(int argc, char** argv) { : cx.pipeline(lib, (v2 ? "mm_gu2_" : "mm_gu_") + std::to_string(gu) + "_" + tl, probe), pdn = v3 ? cx.pipeline(lib, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4)) : cx.pipeline(lib, (v2 ? "mm_down2_" : "mm_down_") + std::to_string(dn) + "_" + (v2 ? dtl : tl)), @@ -7316,7 +7598,7 @@ index 6f51f1e..224cb15 100644 mtl::Dispatch pl; pl.p = &pplace; pl.args = {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok), mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}; pl.grid[0] = nblk; pl.tg[0] = 1024; ds.push_back(pl); } if (gemm) { -@@ -121,7 +123,7 @@ int main(int argc, char** argv) { +@@ -121,7 +127,7 @@ int main(int argc, char** argv) { blob, blob, (uint32_t)gu_row, (uint32_t)d_row, N_EMBD, N_FF}; if (v2) { const uint32_t ne = N_EMBD; @@ -7325,6 +7607,32 @@ index 6f51f1e..224cb15 100644 ga.grid[0] = meta.as()[1]; ga.tg[0] = 256; ds.push_back(ga); } mtl::Dispatch g; g.p = &pgu; +@@ -155,6 +161,13 @@ int main(int argc, char** argv) { + std::vector ds; + encode(0, ds, false, true); + cx.run(ds); ++ if (dump) { ++ FILE* f = std::fopen(dump, "wb"); ++ if (v2) std::fwrite(HS.as(), 2, (size_t)meta.as()[1] * N_FF, f); ++ std::fwrite(y.as(), 4, (size_t)T * N_EMBD, f); ++ std::fclose(f); ++ return 0; ++ } + // reference: fp16 operands (x, weights, h), fp32+ accumulation; and the exact one + double se16 = 0, se = 0, sr = 0; + std::vector wg((size_t)N_FF * N_EMBD), wu(wg.size()), wd((size_t)N_EMBD * N_FF); +@@ -229,6 +242,11 @@ int main(int argc, char** argv) { + return median(t); + }; + const int o = v2 ? 1 : 0; ++ if (gu_only) { ++ const double tg = timeit(false, true, o); ++ std::printf("gate/up %.1f us (%.2f TFLOP/s useful)\n", tg * 1e6, 2.0 * nent * 2.0 * N_FF * N_EMBD / tg / 1e12); ++ return 0; ++ } + const double tp = timeit(true, false, -1), tg = timeit(false, true, o), td = timeit(false, true, o + 1), + tc = timeit(false, true, o + 2), ta = timeit(false, true, -1); + if (v2) std::printf("gather %.1f us\n", timeit(false, true, 0) * 1e6); diff --git a/metal/tests/ple_moe_parity.cpp b/metal/tests/ple_moe_parity.cpp new file mode 100644 index 0000000..f8a0267