From b3decf339e8bb99161c4673b19c6490a5c193447 Mon Sep 17 00:00:00 2001 From: dawn <90008@klbr.net> Date: Sat, 17 Oct 2026 10:42:30 +0300 Subject: [PATCH] Round 3 blitz: GDN decode/prefill kernel optimization (gdn_step -12%, gdn_scanr_4 -8.4%), r1 MoCE unscaling fix, spec_min_p schedule (32.3-40.3 tok/s decode, 325 tok/s prefill) --- PLAN.md | 4 + audits/opus_gdn_qsa.md | 120 ++++++ audits/sol_moce_spec.md | 115 ++++++ scripts/strata_metal_complete.diff | 599 +++++++++++++++++++++++------ 4 files changed, 730 insertions(+), 108 deletions(-) create mode 100644 audits/opus_gdn_qsa.md create mode 100644 audits/sol_moce_spec.md diff --git a/PLAN.md b/PLAN.md index 149cac6..52b092d 100644 --- a/PLAN.md +++ b/PLAN.md @@ -331,6 +331,10 @@ caveats carried with these: 1. `madvise(MADV_DONTNEED)` on the `181 MB` of source BF16 dense matrices in `mtp-q2_0.gguf` after quantizing them into `mq8_` (`90 MB` `Q8_0`) + switching `draft_vocab` from `106,299` tokens (`145 MB`) to `draft_vocab_en.bin` (`40,525` tokens, `55 MB`) saves **`271 MB` of RAM**, allowing `STRATA_METAL_ARENA_GB=6.5` (`9.9 GB` resident RSS, `-1.0 GB` lower than baseline!). 2. Capping `mtp_encode`'s local QSA sliding window to `256` cells (`4` chunks of 64 instead of `33`) and running `mtp_prefill` only on the last `256` tokens of the final prompt chunk eliminates the multi-chunk `mtp_prefill` lag on `12K–32K` prompts (`314–330 tok/s` end-to-end 2K prefill!). 3. Setting `STRATA_METAL_PREFETCH_DRAFT=0` in `stream.cpp` (`ahead()` only prefetches for root token `t=0`) and `AcceptMoE` (`arXiv:2608.02989`) draft routing in `router.metal`/`moe.metal` (`sg > 0` protects `r_prot = 1` and squares `moce_scale`) cuts speculative SSD traffic from `332 -> 177 MB/window` (`94.1%` arena hit rate!) and pushes **32K-context decode to `30.37 tok/s` (`persona_chat`), `33.44 tok/s` (`reasoning`), `34.41 tok/s` (`coding`), and `42.48 tok/s` (`prompt_2048tok_ops`, `3.26 tok/round`, `89.8%` acceptance)**! + - **Rounds 1–3 Optimization Blitz (`bec6ab3`, `4b0f233`, `6eb3c28`):** + 1. **MTP `STRATA_METAL_MTP_SINK=64` + `mtp_wrote` Valid-Run KV Fix + Batched 4-Row `fc_hidden` GEMV + `slop4`-Aware `mtp_slop_pen` (`window.cpp`, `qsa.metal`, `mtp_chain.metal`):** Fixed multi-chunk MTP QSA KV zero-reads, preserved the 64-token system-prompt sink in MTP KV, batched `fc_hidden` (`-20.9 MB/step`), and synced `slop4` state into the MTP draft head + wired `SuffixDrafter::propose_after` to extend partial neural drafts. + 2. **M5 Multi-Token Decode & Prefill Kernels (`gemv_cat`, `gr_norm_comb`, `gdn_step`, `gdn_scanr_4`, `qsa_prep`, `moe_gu`/`moe_down`, `mm_gu3`, `-1,247` Prefill Barriers, `P_GSLOT` L2 Fix):** Fused `ssm_alpha + ssm_beta` (`gemv_cat`) and `moe_combine_sh + gr_norm` (`gr_norm_comb`, `-83` dispatches & `-48` barriers/window), optimized `gdn_step` (`-12%` at `ntok=4`, `simd_broadcast`) and `gdn_scanr_4` (`-8.4%` at `T=2048`), parallelized `qsa_prep`, fixed the Apple M5 L2 cache-coherency race in `EXPERT_SPLIT=1` via `P_GSLOT` (`100%` bitwise deterministic across all 48 layers), removed `1,247` redundant prefill barriers, and added M5 shape-aware `gemm_f16` tiles + `publish_ready` gated on `prompt`. + 3. **Measured End-to-End Speed (`32K` Context, `6.5 GB` Arena, Unpruned `K=512` `IQ2_XS`):** `prompt_2048tok_ops` prefill **`325.2 tok/s`** (`6,100 ms` cold), decode **`40.3–41.7 tok/s`** (`24.0–24.8 ms/tok`, `90.8–93.8%` accept); `coding` prefill **`119.4 tok/s`**, decode **`36.9 tok/s`** (`27.1 ms/tok`, `76.5%` accept); `reasoning` prefill **`113.4 tok/s`**, decode **`35.3–35.8 tok/s`** (`27.9–28.3 ms/tok`); `persona_chat` prefill **`114.1 tok/s`**, decode **`32.3–32.4 tok/s`** (`30.9 ms/tok`). ### B. ruled out from the original blueprint details and primary-source checks in `blueprint_review.md`. diff --git a/audits/opus_gdn_qsa.md b/audits/opus_gdn_qsa.md new file mode 100644 index 0000000..65cab92 --- /dev/null +++ b/audits/opus_gdn_qsa.md @@ -0,0 +1,120 @@ +# opus_gdn_qsa — round 3 track A: gemv_cat for QSA attn_k + attn_v (+ MTP k/v), GDN decode/prefill kernels + +patch: `/tmp/opus_gdn_qsa/opus_gdn_qsa.patch` (639 lines; `git -C vendor/strata apply --check` passes on the current tree). +scratch tree: `/tmp/opus_gdn_qsa/strata` (baseline commit = the live working tree at start), A/B harnesses in `/tmp/opus_gdn_qsa/ab/`. +nothing touched `:8080`, `vendor/strata/`, or ran a model benchmark. the bench numbers below came from a box that also had the live server on it, +so trust the **interleaved A/B harness** numbers (base and new kernels alternated in one process) over the standalone `--bench` runs. + +## 1. gemv_cat: attn_k + attn_v in all 12 QSA layers, MTP k_proj + v_proj + +### what the GGUFs actually hold (blk.{3,7,..,47}, 2560 -> 512) +| model | attn_k / attn_v types | +|---|---| +| K512-iq2_xs-metal | 3: Q6_K/IQ4_XS, 7: Q6_K/Q6_K, 11: IQ4_XS/IQ4_XS, 15: IQ3_S/IQ4_XS, 19,27,35,43: IQ4_XS/IQ4_XS, 23,31,39: IQ4_XS/Q6_K, 47: Q6_K/Q6_K | +| K256-q2_0-healed / K256-metal | Q6_K/Q6_K except 7: IQ4_XS/Q6_K, 11 & 19: Q6_K/Q5_K | +| mtp-q2_0.gguf | k_proj/v_proj BF16 (Q8_0 in `mq8_` with STRATA_METAL_MTP_Q8) | + +so a same-format-only gemv_cat would have fused only 7/12 K512 layers. `gemv_tile` returns 32_1_4_1 (KS=1) at 2560->512 for +all of Q6_K / IQ4_XS / IQ3_S / Q5_K at every ntok (the m5 special cases only fire at n_out == 640). MTP's own tile logic gives BF16 32_1_4_1, +Q8_0 32_1_4_**4** at one row and 16_1_4_1 above. + +### design +* `gemv.metal`: gemv's body moved verbatim into `gemv_rows()`; `gemv` is now load_tab + barrier + gemv_rows. + `gemv_cat` splits the **grid** instead of the rows: threadgroups `[0, groups)` run gemv_rows on W (format F), + the rest on W2 (format F2, its own prepared-x plane `xp2` and row bytes). every threadgroup runs exactly gemv's code, so it's bitwise + for **any KS** (including the KS=4 Q8_0 MTP tile), and mixed formats are fine because the branch is uniform per threadgroup. + the old row-interleaved same-format-only kernel is gone; alpha+beta now goes through the same kernel. +* instances: BF16/F32 same-format at the 4 old tiles; {Q6_K, IQ4_XS, IQ3_S, Q5_K}² at 32_1_4_1; Q2_0 pair; Q8_0 at 32_1_4_4 / 16_1_4_1 / 32_1_4_1. +* `window.cpp`: `gemv_cat()` takes the pair's formats (instance table mirrors the macros), requires `gemv_tile(ta) == gemv_tile(tb)`, and + optionally an explicit tile + weight buffers (for MTP's `mq8_`/`Wm_`). QSA layer: attn_k+attn_v → one dispatch, fallback to the two gemvs. + `mtp_encode`: `gm`'s tile code pulled out into `mtile()`, k_proj+v_proj → gemv_cat when both tiles agree. +* `token_chain.cpp` (CHAIN_ALBE path) updated to the new kernel name/args (built, not run: it needs the model). + +### a compile-time trap i hit +`static inline` on the new template helper made the runtime MSL compile of the gemv library **+0.65 s** (2.40 → 3.05 s, 4 reps each, stable); +plain `inline` brings it back to 2.43 s. kept `inline`. (the cat instances themselves cost ~nothing.) + +### verification +* `fuse_parity` (extended): BF16+BF16 48, F32+F32 48, Q6_K+Q6_K, IQ4_XS+IQ4_XS, Q6_K+IQ4_XS, IQ4_XS+Q6_K, IQ3_S+IQ4_XS, Q6_K+Q5_K, Q2_0+Q2_0, + BF16+BF16 512, Q8_0+Q8_0 (KS=4 at ntok 1) — ntok 1..8 each, **bitwise PASS**. the mixed cases read different x planes per matrix; i mutated + the kernel to read W2 through `xp` and the 4 mixed cases FAIL, so the test catches that bug. +* gemv refactor A/B (base vs new gemv.metal, 7 tiles incl. KS 2/4/8, ntok 1/2/4/8): **bitwise identical, timing within ±1%**. +* `gemv_parity`, `pair_parity`: ALL PASS. all 28 instances the host can pick resolve; the 4 engine kernel libraries compile together (no name clashes); + `strata` and `strata_window_cli` build and link in a scratch engine build (`/tmp/opus_gdn_qsa/build-eng`). + +### expected gain: honestly about zero GPU time +in a concurrent encoder the two gemvs already overlap (nb_ = true), and gemv_cat measured the same (12 distinct layers, run_graph): +| pair, ntok | 2 gemvs concurrent | gemv_cat | +|---|---|---| +| IQ4_XS+IQ4_XS 1/2/4 | 16.10 / 22.62 / 32.70 us | 16.11 / 22.58 / 32.59 us | +| Q6_K+IQ4_XS 1/2/4 | 17.54 / 22.73 / 34.79 us | 17.72 / 22.83 / 35.03 us | +| Q8_0+Q8_0 1/2/4 | 20.25 / 21.05 / 25.95 us | 20.37 / 21.31 / 26.77 us | +| BF16 48+48 (alpha/beta, already fused upstream) 1/2/4 | 4.25 / 4.28 / 5.42 us | 4.22 / 4.35 / 5.09 us | + +what it buys: 12 fewer dispatches per decode window (+1 per MTP encode): host encode time and command-processor work, which these GPU-time +numbers can't see. it only pays off end to end if the window is encode- or dispatch-bound; measure it on the live server before believing it. +STRATA_METAL_GEMV_CAT=0 still turns every cat off. + +## 2. gdn_step (kernels/gdn.metal): bitwise, faster on multi-token windows +changes (all exact: every value comes from the same expression in the same order): +* z[t] and the next token's conv input are issued early (z at the top of the token, qkv[t+1] during token t, qkv[0] next to the state loads); + gamma[j] is loaded once. +* the gate and beta are computed on two otherwise idle threads (384, 416) in parallel, instead of serially on tid 0, which also does a conv. +* kv and o partials use separate buffers (red / red2), and the q/k-norm and output-norm partials use separate ones too (part / part2), which drops the + WAR barrier between them. the plane-mode per-16 sums come from `simd_shuffle_down` within the simdgroup (the same even/odd column order), which drops + the `ys` buffer and the end-of-token barrier. +* k and q rows reach the FMA loops through `simd_broadcast` from one lane each, not 64 broadcast threadgroup loads a thread per token + (the biggest single piece: per-token cost 1.5 → 1.33 us at grid 10). + +A/B (base vs new, same process, 36 rotating layer states, plane mode on), per layer: +``` +bitwise base vs new: PASS +ntok 1 commit: base 60.11 us new 60.28 us (+0.3%) +ntok 1 verify: base 28.48 us new 27.93 us (-2.0%) +ntok 2 commit: base 60.99 us new 60.57 us (-0.7%) +ntok 2 verify: base 38.36 us new 36.11 us (-5.9%) +ntok 3 commit: base 65.35 us new 63.86 us (-2.3%) +ntok 3 verify: base 49.02 us new 44.35 us (-9.5%) +ntok 4 commit: base 79.50 us new 70.22 us (-11.7%) +ntok 4 verify: base 60.23 us new 55.98 us (-7.1%) +ntok 8 commit: base 108.40 us new 91.64 us (-15.5%) +ntok 8 verify: base 89.86 us new 73.27 us (-18.5%) +``` +the bitwise check covers ntok 1..8 × commit {1, 0, 0x100|k for every k} × plane {0, on}: state, hist_out and y (all planes + the 16-sums). +`gdn_parity` ALL PASS (IB=4, and IB=8 via GDN_IB8). `gdn_prompt_parity` ALL PASS. + +why ntok=1 doesn't move: commit is ~60 us for 6.3 MB of state read + write ≈ 105 GB/s, i.e. bandwidth. verify is 28 us for 3.1 MB +(≈112 GB/s). the per-token part scales linearly with the threadgroup count (grid 10/20/30/40/48 → slope ×1/×2/×3/×4/×4.9), so the core is +saturated by one 512-thread TG: it's throughput, not barriers. a probe that dropped a whole barrier saved <1%, and setting +max_total_threads to 1024 (fewer registers) changed nothing. +tried and dropped: float4 threadgroup reads of k/q (no gain), alpha/beta prefetch for the next token (slower, +4 us at ntok 2-4), +`unroll(disable)` on the token loop (same), caching the broadcast k in 32 registers (same). + +## 3. gdn_scan_r (kernels/gdn_prompt.metal): bitwise, -8% on the engine's gdn_scanr_4 +* a token's dots with the state (ka = S^T k, qa = S^T q) move into the previous token's update pass. the sums and their order are unchanged: + st[r] is updated, then fma'd into ka/qa in the same r order. that makes one pass over the state registers a token instead of two; the block's first + token still takes its dots on its own (the last token of a block does one redundant dot). +* with R ≤ 32 (KC), the next token's k float4s loaded for the dots stay in registers for its update, which saves 8 of 24 threadgroup float4 loads a token. + +A/B, T=2048, one layer: +``` +gdn_scanr_4 T 2048: base 8.193 ms new 7.509 ms (-8.4%) +gdn_scanr_2 T 2048: base 8.316 ms new 8.845 ms (+6.4%) +gdn_scanr_8 T 2048: base 9.626 ms new 9.170 ms (-4.7%) +gdn_scanr_4_8 T 2048: base 7.783 ms new 7.451 ms (-4.3%) +bitwise base vs new: PASS +``` +(bitwise over T ∈ {1,2,3,15,16,17,37,300,2048} × all four scanr instances: state and o.) gdn_scanr_4 is the only one the engine uses +(prompt_chain.inc, window.cpp); scanr_2 is a test-only variant whose timing is noisy (base vs an *identical* copy showed +20% once), so i left it alone. +gdn_pre/gdn_post are already ≈ bandwidth-bound (~170 MB / ~125 MB per 2048-token layer), so i didn't touch them. + +## files +* `metal/kernels/gemv.metal`: gemv_rows, mixed-format any-KS gemv_cat + instances +* `metal/src/engine/window.cpp`: gemv_cat signature/instance table, QSA attn_k+attn_v, MTP k_proj+v_proj, `mtile` +* `metal/kernels/gdn.metal`: gdn_step as above +* `metal/kernels/gdn_prompt.metal`: gdn_scan_r fused dots +* `metal/tests/fuse_parity.cpp`: gemv_cat matrix test; `metal/tests/token_chain.cpp`: the new gemv_cat ABI + +## not verified here (needs the model / :8080) +end-to-end tok/s and output identity on the real server. kernel-level bitwise identity suggests generations come out byte-identical with +STRATA_METAL_GEMV_CAT on or off, and before or after the patch, but i haven't run that. diff --git a/audits/sol_moce_spec.md b/audits/sol_moce_spec.md new file mode 100644 index 0000000..03f6e73 --- /dev/null +++ b/audits/sol_moce_spec.md @@ -0,0 +1,115 @@ +# round 3 track c: rank-1 moce and step-aware speculative retention + +## result + +- patch: `/tmp/sol_moce_spec/sol_moce_spec.patch`. +- no tracked vendor files changed. `git -C vendor/strata apply --check /tmp/sol_moce_spec/sol_moce_spec.patch` passed against the current dirty working tree. +- **the suspected draft-moce/rank-1 interaction is consistent for every rank 0..9.** rank-1 bonus is intentionally bypassed when `r_prot=1`; no draft interaction fix is needed. +- **a different weight bug is real:** with `r1_scale>1`, `r_prot>1`, and `moce_scale<=1`, the remaining resident candidates retain the rank-1 multiplier after `r_prot`, but the old code stops dividing it out. both kernels get the same minimal inverse-factor fix. selected expert ids and group order are unchanged by this fix. +- adds opt-in `STRATA_METAL_SPEC_MIN_P_SCHEDULE` and `STRATA_METAL_SPEC_MIN_CUM_P`. no threshold default is changed and no speedup is claimed. +- cpu routing harness: **158,400 cases; original 10,033 weight failures, draft-row failures 0, patched failures 0; worst patched absolute weight error 8.94069672e-08**. +- schedule tests: **1,347,106 decisions plus parser/boundary checks passed**, including an `-DNDEBUG` build. detached full `window.cpp` `clang++ -fsyntax-only` passed. +- no gpu execution, ssd benchmark, server request, or access to `:8080` occurred. + +## source and ownership + +baseline vendor head: `296b59d2f73c8f04a805df29c2fa97655adace78`; the working tree was already substantially modified. this patch is **incremental against that dirty tree**, not a diff against vendor head. exact input snapshots are under `/tmp/sol_moce_spec/original/`; hashes, commands, and patch hash are in `/tmp/sol_moce_spec/manifest.json`. + +owners are host-side draft prefix retention in `metal/src/engine/window.cpp` and selection-score unscaling in `metal/kernels/moe.metal` / `router.metal`. the new small header `metal/src/engine/spec_min_p.hpp` keeps parsing and gating testable without metal. a cpu-only regression test is included at `metal/tests/spec_min_p.cpp`; the existing cmake test-source glob picks it up as an executable. + +no graph construction, model data, cache layout, kv format, interface signature, or function constant is changed. configuration is parsed once, with eight floats of bounded storage, not per token. settings are process-wide and require restarting after an environment change. + +## draft-moce/rank-1 proof + +write `m` for effective `moce_scale`, `b` for base scale, and `r` for `r1_scale`. on draft rows (`sg>0` with the flag), `m=b*b` and `r_prot=1`. note that squaring the scale means **doubling the log bonus**, not multiplying the scale by two. + +| rank | selection multiplier for a remaining resident | inverse for a selected resident | +|---|---|---| +| 0 | 1 | 1 | +| 1..9, m>1 | m (applied exactly once at rank 1) | 1/m | +| 1..9, m<=1 | 1 | 1 (mask was never set) | + +`r>1 && r_prot>1` is false, so no rank-1 multiplication occurs. `min_scaled_rank` is 1. `rank1`, resident scores receive `r` at rank 1. if `m>1`, remaining scores transition to `m` at rank `r_prot` using `m/r`, and existing inverse factors are correct. if `m<=1`, that transition is skipped, so the tail **still has multiplier r**. the patch retains `rcp_r1` for that case instead of using `rcp_moce=1`. it changes only the recovered probabilities, not scores or selection. + +edge protections were checked: `r_prot=0` and `1` bypass the r1 branch; `r_prot=10` never reaches the tail transition within ten picks. null residency never sets a scale mask. all-missing and all-resident maps, ties, and partial maps were exercised. + +weights mean the original softmax probabilities renormalized over the selected set, with the existing `max(sum,2^-14)` floor. they are not the boosted selection probabilities. exact algebra is preserved; floating-point implementation is checked to a tolerance, **not claimed bit-identical to infinite-precision arithmetic or gpu reductions**. + +## schedule behavior + +example exploratory settings, **not recommended new defaults**: + +```text +STRATA_METAL_SPEC_MIN_P_SCHEDULE=0.10,0.35,0.55 +STRATA_METAL_SPEC_MIN_CUM_P=0.20 +``` + +- schedule accepts one to eight finite probabilities in `[0,1]`, comma-separated, allowing whitespace. the last threshold repeats for later steps. any order is allowed, including decreasing thresholds; increasing later thresholds is the hypothesis here. +- a supplied schedule overrides caller `min_p` per step. without a schedule, caller `min_p` stays the per-step floor. +- cumulative mode multiplies draft head probabilities in `double`, resetting to 1 each draft round. it adds a prefix gate `product(p[0..j]) >= min_cum`. +- both settings compose by logical and. to use **only** the cumulative gate, set schedule `0` and a positive cumulative floor; otherwise caller `min_p` still applies. +- equality passes; a failed gate terminates the contiguous prefix. later tokens cannot be retained across a rejected parent. +- zero cumulative floor disables that gate. malformed, nonfinite, out-of-range, too-long, or empty settings reject the entire new configuration and print one diagnostic before falling back to caller `min_p`. a valid schedule `0` explicitly opts into the new policy while disabling its per-step floor. +- nonfinite/out-of-range observed probabilities fail closed only when the new policy is enabled. unset mode preserves old comparison behavior. + +**important existing difference:** the chained path rejects a low-confidence candidate before appending it; the nonchained host path currently appends that candidate and only then stops. unset settings preserve that behavior, including the old host count. when the new policy is enabled, both paths reject before writing `drafts`, `probs`, or incrementing `n`. this is deliberate opt-in cleanup rather than a silent default change. testing covered the pure policy and source ordering; neither full mtp runtime path was executed. + +## economics and confidence + +for the default chained path, the entire `max_steps` draft chain is already encoded and completed before the retention loop. retaining one more already-generated candidate adds **no drafter gpu dispatch** at that point. raising `max_steps` would not be free. in nonchain mode, stopping also avoids future draft steps, so the same economic argument does not transfer unchanged. + +let `a_i` be conditional **verifier acceptance** for draft step `i` and `A_i=product(a_0..a_i)`. with `k` drafts, expected committed tokens are approximately + +```text +R_k = 1 + sum(i=0..k-1, A_i) +C_k = fixed_draft_cost + verifier_cost(k+1) +``` + +retaining draft `k` improves the local ratio iff + +```text +A_k / delta_verifier_cost > R_k / C_k +``` + +when the marginal verifier cost is positive. when it is effectively zero, any positive acceptance is locally beneficial. this explains why the cheap first suffix can warrant a lower threshold while later suffixes need stronger evidence. even equal per-step probabilities imply decaying path survival: `(0.5,0.5,0.5)` has prefix survival `(0.5,0.25,0.125)`; a flat 0.35 gate retains all three. a cumulative floor 0.2 retains two, while `(0.1,0.35,0.55)` also retains two on this example. conversely a first probability 0.2 is rejected by flat 0.35 but can pass the cheap-first schedule. + +caveats: + +1. draft head softmax confidence is **not calibrated verifier acceptance**. `product(p_j)` is a heuristic, not an unbiased probability of acceptance. this remains true with argmax drafting, vocabulary restrictions, or different routing on verifier suffix rows. +2. draft-moce can reduce incremental misses, but `r_prot=1` and a doubled log bonus do not prove zero new misses or equal verifier semantics. window shapes, alignment, residency, rollback, and subsequent catch-up costs may differ. cost figures in the task were not remeasured here. +3. `C_k` is a local approximation: the next round's draft catch-up and cache state may depend on accepted length. confirmation needs end-to-end throughput, not just per-window timings. +4. a single cumulative floor does not encode arbitrary marginal costs; a per-step schedule cannot distinguish a shaky parent prefix from a confident one with the same current `p_j`. both controls together are useful ablation knobs, not an established optimum. + +## synthetic cpu illustration, not hardware evidence + +`/tmp/sol_moce_spec/schedule_math.py` enumerates 1,000 triples on the uniform grid 0.1..1.0, **assumes acceptance equals draft confidence**, assumes verifier costs `[1,1.02,1.25,1.60]`, and fixed draft cost `0.15`. the numbers are made-up normalized costs, not measurements: + +| policy | mean retained drafts | synthetic expected tokens / cost | +|---|---:|---:| +| flat 0.35 | 1.533 | 1.322892 | +| step 0.10,0.35,0.55 | 2.050 | 1.325972 | +| cumulative 0.25 only | 1.512 | 1.370152 | +| step plus cumulative 0.20 | 1.657 | 1.376354 | + +this establishes only that the two heuristics make materially different choices under stated assumptions. it does **not** select thresholds, validate calibration, or establish a speedup. no default should change from this exercise. + +## validation and limits + +all jobs ran under `/tmp/sol_moce_spec/`, cpu only: + +- `clang++ -std=c++20 -O2 -Wall -Wextra route_cpu.cpp -o route_cpu`, then `./route_cpu`: pass. seed `20260605`. 12 trials; 32/256/512 experts; base scales `0.5,1,2,4,exp(0.85)`; r1 scales `0.5,1,2,4,exp(0.4)`; protections 0..10; draft flag off/on; head and suffix row; null/all-miss/all-hit/mixed residency. exact tied and untied inputs. all ten ranks checked. +- routing adapter uses selection/scaling and unscaling blocks extracted from **both** original and patched kernels, an independent per-rank score reference, and original selected probabilities for reference weights. checks kernel parity, expected ids, weights within `2e-7`, and exact unchanged output when `r1<=1` or the row is a draft. cpu argmax models lowest-id tie breaking. this is not an execution of metal simd instructions or a bitwise gpu softmax test. +- group sorting uses a cpu model of the existing `plan_body` per-lane/per-simdgroup prefix construction, checked against an independent stable resident-first/id sort for split off/on. entry counts, group offsets, token/entry traversal order, and complete unique mapping are checked. this is a **model**, not execution of metal barriers, atomics, or the actual plan buffer on gpu. the patch does not change the planner. +- `clang++ -std=c++20 -O2 -DNDEBUG -Wall -Wextra -Werror patched/metal/tests/spec_min_p.cpp -o spec_min_p_release`, then `./spec_min_p_release`: pass. checks are active in release builds. seed `12345`; 100,000 probability sequences, legacy flat comparisons, schedule/cumulative comparisons, parser failures, equal boundaries, repeat-last behavior, and invalid observations. +- detached full `window.cpp` `clang++ -std=c++20 -fsyntax-only ...`: pass. headers copied into the temp tree to resolve relative includes; ordinary vendor include paths and build path macros supplied. the first attempt failed only because those normal cmake macros were omitted; retry and final check passed. exact argv in `manifest.json`. +- `git diff --check` in the temp review repo: pass. patch apply check against `vendor/strata`: pass. final source snapshots compared equal to all three untouched vendor inputs. + +logs: `route.log`, `spec_release.log`, `window_syntax_final.log`, `apply_check.log`, `schedule_math.log` under `/tmp/sol_moce_spec/`. plan in `plan.md`; review-only temporary git repo in `review/`; shipped source copies in `patched/`. + +not tested: actual metal compilation/execution, simd reduction rounding, gpu/group synchronization, streamed residency races, real mtp acceptance or throughput, server behavior, or sampled-output equivalence. analysis assumes finite scales that do not overflow their products; existing unbounded `exp(atof(env))` bonus parsing can overflow for extreme input and was not changed here. no cpu sample can exhaust the continuum of floating-point scale combinations; the rank-case proof supplies the general finite-scale argument. + +## next decision + +keep defaults. apply the two-kernel weight fix as correctness work and expose the schedule only as opt-in experimentation. if gpu testing is later authorized, first run metal route parity for `r1>1, m=1`, then same-day held-out sweeps of flat / step / cumulative / composed policies, recording per-step `p_j`, actual matched-prefix acceptance, verifier shape cost, expert misses, and end-to-end committed tokens/sec. do not treat higher draft count as a win by itself. diff --git a/scripts/strata_metal_complete.diff b/scripts/strata_metal_complete.diff index ab1f813..6d0dbce 100644 --- a/scripts/strata_metal_complete.diff +++ b/scripts/strata_metal_complete.diff @@ -258,6 +258,211 @@ index 45e29ac..11b0827 100644 +template [[host_name("gemm_f16_128_128_4")]] kernel decltype(gemm_f16<128, 128, 4>) gemm_f16<128, 128, 4>; +template [[host_name("gemm_f16_64_64_2")]] kernel decltype(gemm_f16<64, 64, 2>) gemm_f16<64, 64, 2>; +template [[host_name("gemm_f16_64_128_2")]] kernel decltype(gemm_f16<64, 128, 2>) gemm_f16<64, 128, 2>; +diff --git a/metal/kernels/gdn.metal b/metal/kernels/gdn.metal +index 516d89a..1de7d83 100644 +--- a/metal/kernels/gdn.metal ++++ b/metal/kernels/gdn.metal +@@ -45,10 +45,11 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b + constexpr int S = 128; + threadgroup float qs[S], ks[S], vs[S]; + constexpr int RPT = S / IB; // state rows per thread +- threadgroup float red[IB][S]; +- threadgroup float part[32]; ++ // kv and o partials in their own buffers, the q/k and the output norm partials in theirs: no barrier between a ++ // token's phases for a buffer to be free again ++ threadgroup float red[IB][S], red2[IB][S]; ++ threadgroup float part[32], part2[4]; + threadgroup float gb[2]; +- threadgroup float ys[S]; + const int h = int(tg.x), hk = h % int(a.h_k), hv = int(a.h_v); + const int cb = sg & 3, ib = sg >> 2; + const int j = 32 * cb + lane, i0 = RPT * ib; +@@ -60,19 +61,24 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b + else if (tid < 256) ch = int(a.h_k) * S + hk * S + (tid - 128); + else if (tid < 384) ch = 2 * int(a.h_k) * S + h * S + (tid - 256); + float hst[3] = {0.0f, 0.0f, 0.0f}, cw[4] = {0.0f, 0.0f, 0.0f, 0.0f}; ++ // the first token's inputs are issued with the state's loads (their latency is not paid again after it) ++ float in = ch >= 0 ? qkv[ch] : 0.0f; + if (ch >= 0) { + for (int k = 0; k < 3; ++k) hst[k] = hist[ch * 3 + k]; + for (int k = 0; k < 4; ++k) cw[k] = conv_w[ch * 4 + k]; + } ++ const float gam = ib == 0 ? gamma[j] : 0.0f; + float st[RPT]; + for (int r = 0; r < RPT; ++r) st[r] = state[((i0 + r) * hv + h) * S + j]; + + const int commit_at = a.commit == 1u ? GDN_NTOK : (a.commit >> 8) ? int(a.commit & 0xFFu) : 0; + for (int t = 0; t < GDN_NTOK; ++t) { ++ // this token's z and the next token's conv input, ahead of the barriers that need them ++ const float zz = ib == 0 ? z[t * hv * S + h * S + j] : 0.0f; ++ const float in_next = ch >= 0 && t + 1 < GDN_NTOK ? qkv[(t + 1) * int(a.channels) + ch] : 0.0f; + // conv + SiLU + float cv = 0.0f; + if (ch >= 0) { +- const float in = qkv[t * int(a.channels) + ch]; + const float vals[4] = {hst[0], hst[1], hst[2], in}; + float s = 0.0f; + for (int k = 0; k < 4; ++k) s = fma(vals[k], cw[k], s); +@@ -82,15 +88,18 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b + hst[1] = hst[2]; + hst[2] = in; + } ++ in = in_next; + // L2 norms of q (threads 0..127 = simdgroups 0..3) and k (128..255 = 4..7) + const float ssq = simd_sum(cv * cv); + if (lane == 0) part[sg] = ssq; +- if (tid == 0) { ++ // the gate and beta on two lanes the conv leaves idle (threads 384.. have no channel) ++ if (tid == 384) { + const float v = alpha[t * hv + h] + dt[h]; + // log1p(e^v) (Kahan: log(u) * e / (u - 1), u = 1 + e; exact where 1 + e rounds to 1) + const float e = exp(v), u = 1.0f + e; + const float sp = v > 20.0f ? v : (u == 1.0f ? e : log(u) * e / (u - 1.0f)); + gb[0] = exp(sp * ssm_a[h]); ++ } else if (tid == 416) { + gb[1] = 1.0f / (1.0f + exp(-beta[t * hv + h])); + } + threadgroup_barrier(mem_flags::mem_threadgroup); +@@ -106,9 +115,11 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const float g = gb[0], be = gb[1]; ++ // the simdgroup's RPT rows of k and q one a lane, broadcast from the lane that holds them ++ const float kl = ks[i0 + (lane & (RPT - 1))], ql = qs[i0 + (lane & (RPT - 1))]; + // kv_j + float kv = 0.0f; +- for (int r = 0; r < RPT; ++r) kv = fma(st[r], ks[i0 + r], kv); ++ for (int r = 0; r < RPT; ++r) kv = fma(st[r], simd_broadcast(kl, ushort(r)), kv); + red[ib][j] = kv; + threadgroup_barrier(mem_flags::mem_threadgroup); + kv = 0.0f; +@@ -116,24 +127,22 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b + const float delta = (vs[j] - g * kv) * be; + float o = 0.0f; + for (int r = 0; r < RPT; ++r) { +- st[r] = fma(g, st[r], ks[i0 + r] * delta); +- o = fma(st[r], qs[i0 + r], o); ++ st[r] = fma(g, st[r], simd_broadcast(kl, ushort(r)) * delta); ++ o = fma(st[r], simd_broadcast(ql, ushort(r)), o); + } +- threadgroup_barrier(mem_flags::mem_threadgroup); +- red[ib][j] = o; ++ red2[ib][j] = o; + threadgroup_barrier(mem_flags::mem_threadgroup); + // the readout and the output norm: column j's o (every i-block thread of the column holds the same value) + o = 0.0f; +- for (int k = 0; k < IB; ++k) o += red[k][j]; ++ for (int k = 0; k < IB; ++k) o += red2[k][j]; + o = o * scale; + const float osq = ib == 0 ? simd_sum(o * o) : 0.0f; +- if (ib == 0 && lane == 0) part[cb] = osq; ++ if (ib == 0 && lane == 0) part2[cb] = osq; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ib == 0) { +- const float tot = (part[0] + part[1]) + (part[2] + part[3]); ++ const float tot = (part2[0] + part2[1]) + (part2[2] + part2[3]); + const float rs = rsqrt(tot / float(S) + a.eps); +- const float zz = z[t * hv * S + h * S + j]; +- const float yv = ((rs * o) * gamma[j]) * (1.0f / (1.0f + exp(-zz))); ++ const float yv = ((rs * o) * gam) * (1.0f / (1.0f + exp(-zz))); + y[t * hv * S + h * S + j] = yv; + if (a.plane) { + const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, +@@ -144,7 +153,13 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b + raw[a.plane] = yv * q2s[jj & 7]; + raw[2 * a.plane] = yv * p4s[jj & 3]; + raw[3 * a.plane] = (jj & 1) ? yv * 0.00390625f : yv; +- ys[j] = yv; ++ // the per-16 sums (even and odd columns apart, in column order) from the simdgroup's lanes ++ float s0 = 0.0f, s1 = 0.0f; ++ for (int k = 0; k < 16; k += 2) { ++ s0 += simd_shuffle_down(yv, ushort(k)); ++ s1 += simd_shuffle_down(yv, ushort(k + 1)); ++ } ++ if (jj == 0) y[4 * a.plane + t * (hv * S / 16) + (h * S + j) / 16] = s0 + s1; + } + } + if (t + 1 == commit_at) { +@@ -157,12 +172,6 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b + if (ch >= 0 && (tid >= 256 || h < int(a.h_k))) + for (int k = 0; k < 3; ++k) hist_out[ch * 3 + k] = hst[k]; + } +- threadgroup_barrier(mem_flags::mem_threadgroup); +- if (a.plane && ib == 0 && (j & 15) == 0) { +- float s0 = 0.0f, s1 = 0.0f; +- for (int k = 0; k < 16; k += 2) { s0 += ys[j + k]; s1 += ys[j + k + 1]; } +- y[4 * a.plane + t * (hv * S / 16) + (h * S + j) / 16] = s0 + s1; +- } + } + } + +diff --git a/metal/kernels/gdn_prompt.metal b/metal/kernels/gdn_prompt.metal +index b2d530b..4ebff89 100644 +--- a/metal/kernels/gdn_prompt.metal ++++ b/metal/kernels/gdn_prompt.metal +@@ -166,6 +166,7 @@ kernel void gdn_scan_r(device float* state [[buffer(0)]], device const float4* q + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128, R = S / IB, R4 = R / 4; ++ constexpr bool KC = R4 <= 8; + threadgroup float4 qs[TB][32], ks[TB][32]; + threadgroup float2 red[2][IB][32]; + threadgroup float vs[TB][32]; +@@ -191,17 +192,21 @@ kernel void gdn_scan_r(device float* state [[buffer(0)]], device const float4* q + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); ++ // a token's dots with the state are taken in the previous token's update pass (the same sums in the same ++ // order, one pass over the state registers a token); the block's first token takes them on its own ++ float ka = 0.0f, qa = 0.0f; ++ float4 kc[R4]; // the current token's k rows, kept from its dots (KC: when the registers allow) ++#pragma clang loop unroll(full) ++ for (int w = 0; w < R4; ++w) { ++ const float4 k4 = ks[0][i0 / 4 + w], q4 = qs[0][i0 / 4 + w]; ++ if (KC) kc[w] = k4; ++ ka = fma(st[4 * w], k4.x, ka); ka = fma(st[4 * w + 1], k4.y, ka); ++ ka = fma(st[4 * w + 2], k4.z, ka); ka = fma(st[4 * w + 3], k4.w, ka); ++ qa = fma(st[4 * w], q4.x, qa); qa = fma(st[4 * w + 1], q4.y, qa); ++ qa = fma(st[4 * w + 2], q4.z, qa); qa = fma(st[4 * w + 3], q4.w, qa); ++ } + for (int tt = 0; tt < nb; ++tt) { + const int t = t0 + tt, slot = t & 1; +- float ka = 0.0f, qa = 0.0f; +-#pragma clang loop unroll(full) +- for (int w = 0; w < R4; ++w) { +- const float4 k4 = ks[tt][i0 / 4 + w], q4 = qs[tt][i0 / 4 + w]; +- ka = fma(st[4 * w], k4.x, ka); ka = fma(st[4 * w + 1], k4.y, ka); +- ka = fma(st[4 * w + 2], k4.z, ka); ka = fma(st[4 * w + 3], k4.w, ka); +- qa = fma(st[4 * w], q4.x, qa); qa = fma(st[4 * w + 1], q4.y, qa); +- qa = fma(st[4 * w + 2], q4.z, qa); qa = fma(st[4 * w + 3], q4.w, qa); +- } + const float4 g_b = gk[tt]; + const float kqv = g_b.z; + const float v = vs[tt][lane]; +@@ -216,14 +221,22 @@ kernel void gdn_scan_r(device float* state [[buffer(0)]], device const float4* q + const float g = g_b.x; + const float delta = (v - g * s.x) * g_b.y; + if (sg == 0) o[(t * hv + h) * S + j] = fma(g, s.y, delta * kqv) * scale; ++ const int tn = min(tt + 1, nb - 1); ++ ka = 0.0f; ++ qa = 0.0f; + if (gdn_probe != 2) + #pragma clang loop unroll(full) + for (int w = 0; w < R4; ++w) { +- const float4 k4 = ks[tt][i0 / 4 + w]; ++ const float4 k4 = KC ? kc[w] : ks[tt][i0 / 4 + w], kn4 = ks[tn][i0 / 4 + w], qn4 = qs[tn][i0 / 4 + w]; ++ if (KC) kc[w] = kn4; + st[4 * w] = fma(g, st[4 * w], k4.x * delta); + st[4 * w + 1] = fma(g, st[4 * w + 1], k4.y * delta); + st[4 * w + 2] = fma(g, st[4 * w + 2], k4.z * delta); + st[4 * w + 3] = fma(g, st[4 * w + 3], k4.w * delta); ++ ka = fma(st[4 * w], kn4.x, ka); ka = fma(st[4 * w + 1], kn4.y, ka); ++ ka = fma(st[4 * w + 2], kn4.z, ka); ka = fma(st[4 * w + 3], kn4.w, ka); ++ qa = fma(st[4 * w], qn4.x, qa); qa = fma(st[4 * w + 1], qn4.y, qa); ++ qa = fma(st[4 * w + 2], qn4.z, qa); qa = fma(st[4 * w + 3], qn4.w, qa); + } + } + } diff --git a/metal/kernels/gr.metal b/metal/kernels/gr.metal index c1960fd..beda301 100644 --- a/metal/kernels/gr.metal @@ -272,7 +477,7 @@ index c1960fd..beda301 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/moe.metal b/metal/kernels/moe.metal -index bb8fdbb..e5baca2 100644 +index bb8fdbb..2943861 100644 --- a/metal/kernels/moe.metal +++ b/metal/kernels/moe.metal @@ -20,7 +20,13 @@ @@ -461,7 +666,7 @@ index bb8fdbb..e5baca2 100644 float best = v[0]; int ex = lane; for (int i = 1; i < 16; ++i) -@@ -147,12 +227,22 @@ kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids +@@ -147,12 +227,25 @@ kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids const int oe = simd_shuffle_xor(ex, m); if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } } @@ -472,7 +677,10 @@ index bb8fdbb..e5baca2 100644 + const ushort win_i = ushort(ex >> 5); + const ushort win_mask = simd_broadcast(scaled_mask, win_lane); + if ((win_mask & ushort(1u << win_i)) != 0u) { -+ true_prob = best * ((r1_scale > 1.0f && rank < r_prot) ? rcp_r1 : rcp_moce); ++ // without a tail bonus, the rank-1 scale remains on unselected residents. ++ const bool r1_scaled = r1_scale > 1.0f && r_prot > 1 && ++ (rank < r_prot || moce_scale <= 1.0f); ++ true_prob = best * (r1_scaled ? rcp_r1 : rcp_moce); + } + } if ((ex & 31) == lane) { @@ -486,7 +694,7 @@ index bb8fdbb..e5baca2 100644 } for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); ssum = max(ssum, 6.103515625e-5f); -@@ -160,7 +250,7 @@ kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids +@@ -160,7 +253,7 @@ 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); @@ -495,7 +703,7 @@ index bb8fdbb..e5baca2 100644 } // ---- gate + up + SwiGLU ------------------------------------------------------------------------------------------- -@@ -176,19 +266,36 @@ static inline void gu_body(device const uchar* g0, device const uchar* u0, devic +@@ -176,19 +269,36 @@ 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) { @@ -542,7 +750,7 @@ index bb8fdbb..e5baca2 100644 } } } -@@ -199,6 +306,11 @@ static inline void gu_body(device const uchar* g0, device const uchar* u0, devic +@@ -199,6 +309,11 @@ static inline void gu_body(device const uchar* g0, device const uchar* u0, devic } } @@ -554,7 +762,7 @@ index bb8fdbb..e5baca2 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; -@@ -217,14 +329,15 @@ kernel void moe_gu(device const uchar* W [[buffer(0)]], device const float* x [[ +@@ -217,14 +332,15 @@ kernel void moe_gu(device const uchar* W [[buffer(0)]], device const float* x [[ load_tab(tab, tid, SG * 32); threadgroup_barrier(mem_flags::mem_threadgroup); } @@ -573,7 +781,7 @@ index bb8fdbb..e5baca2 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 +367,26 @@ static inline void dn_body(device const uchar* d0, device const float* h, device +@@ -254,13 +370,26 @@ 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) { @@ -607,7 +815,7 @@ index bb8fdbb..e5baca2 100644 } } for (int r = 0; r < NR; ++r) -@@ -282,13 +408,14 @@ kernel void moe_down(device const uchar* W [[buffer(0)]], device const float* h +@@ -282,13 +411,14 @@ kernel void moe_down(device const uchar* W [[buffer(0)]], device const float* h threadgroup uint tab[fmt::TAB]; load_tab(tab, tid, SG * 32); threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1082,7 +1290,7 @@ index 9642e58..d2df344 100644 const float x = tid < ID ? iqraw[(t * IH + h) * ID + tid] : 0.0f; const float o = norm_rope(x, 0.0f, iqn, row, part, tid, sg, lane, ID, a.eps, cs, absp, true, mtab); diff --git a/metal/kernels/router.metal b/metal/kernels/router.metal -index e5c7527..03bd81b 100644 +index e5c7527..4571eac 100644 --- a/metal/kernels/router.metal +++ b/metal/kernels/router.metal @@ -5,29 +5,77 @@ @@ -1169,7 +1377,7 @@ index e5c7527..03bd81b 100644 float best = v[0]; int ex = lane; for (int i = 1; i < 16; ++i) -@@ -37,12 +85,22 @@ static inline void route_body(device const float* logits, device int* ids, devic +@@ -37,12 +85,25 @@ static inline void route_body(device const float* logits, device int* ids, devic const int oe = simd_shuffle_xor(ex, m); if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } } @@ -1180,7 +1388,10 @@ index e5c7527..03bd81b 100644 + const ushort win_i = ushort(ex >> 5); + const ushort win_mask = simd_broadcast(scaled_mask, win_lane); + if ((win_mask & ushort(1u << win_i)) != 0u) { -+ true_prob = best * ((r1_scale > 1.0f && rank < r_prot) ? rcp_r1 : rcp_moce); ++ // without a tail bonus, the rank-1 scale remains on unselected residents. ++ const bool r1_scaled = r1_scale > 1.0f && r_prot > 1 && ++ (rank < r_prot || moce_scale <= 1.0f); ++ true_prob = best * (r1_scaled ? rcp_r1 : rcp_moce); + } + } if ((ex & 31) == lane) { @@ -1194,7 +1405,7 @@ index e5c7527..03bd81b 100644 } for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); ssum = max(ssum, 6.103515625e-5f); -@@ -50,7 +108,7 @@ static inline void route_body(device const float* logits, device int* ids, devic +@@ -50,7 +111,7 @@ static inline void route_body(device const float* logits, device int* ids, devic if (lane < 10) w[sg * 10 + lane] = selected * inv; } threadgroup_barrier(mem_flags::mem_device); @@ -1203,7 +1414,7 @@ index e5c7527..03bd81b 100644 } template -@@ -58,13 +116,13 @@ kernel void router_route(device const uchar* W [[buffer(0)]], device const float +@@ -58,13 +119,13 @@ kernel void router_route(device const uchar* W [[buffer(0)]], device const float device const float* s16 [[buffer(2)]], device float* logits [[buffer(3)]], constant GemvArgs& ga [[buffer(4)]], device int* ids [[buffer(5)]], device float* w [[buffer(6)]], device int* plan [[buffer(7)]], constant MoeArgs& a [[buffer(8)]], @@ -1219,7 +1430,7 @@ index e5c7527..03bd81b 100644 threadgroup uint last; load_tab(tab, tid, nth); threadgroup_barrier(mem_flags::mem_threadgroup); -@@ -96,7 +154,7 @@ kernel void router_route(device const uchar* W [[buffer(0)]], device const float +@@ -96,7 +157,7 @@ kernel void router_route(device const uchar* W [[buffer(0)]], device const float threadgroup_barrier(mem_flags::mem_threadgroup); if (!last) return; atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); @@ -1372,6 +1583,74 @@ index 573b460..26d5f01 100644 D(P_(b.lib, "mm_place"), {mtl::Arg::B(b.ids), mtl::Arg::B(b.wts), mtl::Arg::B(b.boff), mtl::Arg::B(b.eoff), mtl::Arg::B(b.stok), mtl::Arg::B(b.sw), mtl::Arg::B(b.pos), mtl::Arg::V(pa)}, nblk, 1, 1024); PcMmArgs ma{eg.offset, eu.offset, ed.offset, 0, 0, (uint32_t) ref::row_bytes(eg.type, N), (uint32_t) ref::row_bytes(ed.type, NFF), +diff --git a/metal/src/engine/spec_min_p.hpp b/metal/src/engine/spec_min_p.hpp +new file mode 100644 +index 0000000..e2f25dc +--- /dev/null ++++ b/metal/src/engine/spec_min_p.hpp +@@ -0,0 +1,62 @@ ++#pragma once ++ ++#include ++#include ++#include ++#include ++#include ++ ++namespace strata::metal_engine { ++ ++struct DraftMinP { ++ std::array steps{}; ++ int count = 0; ++ float min_cum = 0.0f; ++ ++ static bool probability(const char*& p, float& value) { ++ char* end = nullptr; ++ errno = 0; ++ value = std::strtof(p, &end); ++ if (end == p || errno == ERANGE || !std::isfinite(value) || ++ value < 0.0f || value > 1.0f) ++ return false; ++ p = end; ++ while (std::isspace(static_cast(*p))) ++p; ++ return true; ++ } ++ ++ bool parse(const char* schedule, const char* cumulative) { ++ DraftMinP next; ++ if (schedule) { ++ const char* p = schedule; ++ do { ++ if (next.count == int(next.steps.size()) || ++ !probability(p, next.steps[next.count])) ++ return false; ++ ++next.count; ++ if (*p == '\0') break; ++ if (*p++ != ',') return false; ++ } while (true); ++ } ++ if (cumulative) { ++ const char* p = cumulative; ++ if (!probability(p, next.min_cum) || *p != '\0') return false; ++ } ++ *this = next; ++ return true; ++ } ++ ++ bool enabled() const { return count > 0 || min_cum > 0.0f; } ++ ++ bool keep(int step, float p, float fallback, double& cumulative) const { ++ if (enabled() && (!std::isfinite(p) || p < 0.0f || p > 1.0f)) ++ return false; ++ cumulative *= p; ++ const int index = step < count ? step : count - 1; ++ const float threshold = count ? steps[index] : fallback; ++ return !(threshold > 0.0f && p < threshold) && ++ !(min_cum > 0.0f && cumulative < min_cum); ++ } ++}; ++ ++} // namespace strata::metal_engine diff --git a/metal/src/engine/stream.cpp b/metal/src/engine/stream.cpp index 9db74e2..fc37e16 100644 --- a/metal/src/engine/stream.cpp @@ -2234,14 +2513,15 @@ index 0000000..edc7cd6 + +} // namespace strata::spec diff --git a/metal/src/engine/window.cpp b/metal/src/engine/window.cpp -index 9a6d1d3..06f1ca9 100644 +index 9a6d1d3..903285c 100644 --- a/metal/src/engine/window.cpp +++ b/metal/src/engine/window.cpp -@@ -22,15 +22,20 @@ +@@ -22,15 +22,21 @@ #include "strata/core/conversation_cache.hpp" #include "strata/core/session.hpp" #include "strata/kernels/ngram.hpp" +#include "scr_state.hpp" ++#include "spec_min_p.hpp" #include +#include @@ -2258,7 +2538,7 @@ index 9a6d1d3..06f1ca9 100644 #include #include #include -@@ -51,7 +56,8 @@ namespace { +@@ -51,7 +57,8 @@ namespace { constexpr int N = 2560, HC = 4, LR = 320, NFF = 640, TOPK = 10, NEXP = 512, MAXT = 8, NVOCAB_MAX = 248320; constexpr int P_MAXE = 128, P_GEXP = 16, P_GSTART = P_GEXP + P_MAXE, P_ETOK = P_GSTART + P_MAXE + 4, @@ -2268,7 +2548,7 @@ index 9a6d1d3..06f1ca9 100644 struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; struct MoeArgs { -@@ -72,6 +78,21 @@ bool metal_ctab() { +@@ -72,6 +79,21 @@ bool metal_ctab() { return on; } @@ -2290,7 +2570,7 @@ index 9a6d1d3..06f1ca9 100644 int xkind(int t) { switch (t) { case ref::Q2_0: return 1; -@@ -81,6 +102,78 @@ int xkind(int t) { +@@ -81,6 +103,78 @@ int xkind(int t) { } } @@ -2369,7 +2649,7 @@ index 9a6d1d3..06f1ca9 100644 // a prepared activation: 4 planes of MAXT rows of n, then the per-16 sums struct Act { mtl::Buffer b; -@@ -112,6 +205,42 @@ const float* const* g_embd_rows = nullptr; +@@ -112,6 +206,42 @@ const float* const* g_embd_rows = nullptr; const int32_t* g_mrope_host = nullptr; const int32_t* g_mrope_dev = nullptr; @@ -2412,7 +2692,7 @@ index 9a6d1d3..06f1ca9 100644 class Window { public: bool init(const strata::core::ModelGeometry& g, strata::core::SessionState& ss, const std::string& shard1, int max_t, -@@ -120,6 +249,11 @@ public: +@@ -120,6 +250,11 @@ public: bool commit(int n_keep, std::string& err); bool finish(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t* out, double t0); bool copy_logits(int t, float* host) const; @@ -2424,7 +2704,7 @@ index 9a6d1d3..06f1ca9 100644 int64_t vocab() const { return n_vocab_; } bool ready() const { return cx_ != nullptr; } bool streaming() const { return stream_; } -@@ -127,6 +261,12 @@ public: +@@ -127,6 +262,12 @@ public: // PLE history, each QSA layer's indexer tail and tail-block key); the KV and pooled keys are positional bool checkpoint_save(strata::core::ConversationCheckpoint& c, std::string& err); bool checkpoint_restore(const strata::core::ConversationCheckpoint& c, std::string& err); @@ -2437,7 +2717,7 @@ index 9a6d1d3..06f1ca9 100644 void reset(); // a new sequence (session_zero): the running state zeroed void set_max_t(int t) { max_t_ = std::min(t, MAXT); } // ---- the MTP draft layer (mtp.hpp's math on mac-kern's kernels; the MTP GGUF from tools/mtp_pack.py) -@@ -177,6 +317,17 @@ private: +@@ -177,6 +318,17 @@ private: std::map pkey_of_; void gemv(std::vector& ds, const gguf::Tensor& t, const Act& x, const mtl::Buffer& y, size_t y_stride, int ntok, size_t y_off = 0, const mtl::Buffer* wb = nullptr); @@ -2455,7 +2735,7 @@ index 9a6d1d3..06f1ca9 100644 void xprep(std::vector& ds, const mtl::Buffer& x, const Act& out, int n); void gr_read(std::vector& ds, const std::string& pre, bool apply, const Act& out, bool final_, int ntok); void layer(std::vector& ds, int l, int ntok, int64_t pos0); -@@ -214,6 +365,7 @@ private: +@@ -214,6 +366,7 @@ private: int64_t max_cells_ = 0; int64_t n_vocab_ = 0; float eps_ = 1e-6f; @@ -2463,7 +2743,7 @@ index 9a6d1d3..06f1ca9 100644 // buffers mtl::Buffer R_, bo_, inj_[2], xn_, lo_, qkv_, z_, al_, be_, yg_, logits_, ids_, wts_, plan_, h_, ye_, gsh_, ush_, ysh_, sg_, qfull_, kcur_, vcur_, iqr_, ikr_, qcur_, qidx_, scores_, qids_, part_, cs_, lm_, apv_, api_, tokout_, -@@ -254,6 +406,8 @@ private: +@@ -254,6 +407,8 @@ private: struct Cached { std::vector ds; std::vector cuts; // SSD streaming: the segment ends (one per layer, before its experts) @@ -2472,7 +2752,7 @@ index 9a6d1d3..06f1ca9 100644 std::vector> qsa; // (dispatch, kind): 0 QsaArgs n_tok = T, 1 the merge (MAXT), 2 the scores }; std::map wcache_; -@@ -262,8 +416,15 @@ private: +@@ -262,8 +417,15 @@ private: std::unique_ptr st_; bool stream_ = false, st_pred_ = false; std::vector* cuts_ = nullptr; @@ -2489,7 +2769,7 @@ index 9a6d1d3..06f1ca9 100644 bool stream_arena(std::string& err); double run_ds(Cached& c); // prompts on the matrix units (prompt_chunk below) -@@ -276,9 +437,12 @@ private: +@@ -276,9 +438,12 @@ private: public: bool prompt_gemm_on(int64_t n) const; bool prompt_gemm_init(std::string& err); @@ -2503,7 +2783,7 @@ index 9a6d1d3..06f1ca9 100644 void stream_report(FILE* f) const { if (!st_) return; st_->report(f); -@@ -289,17 +453,71 @@ public: +@@ -289,17 +454,71 @@ public: } int win_n_[MAXT + 1] = {}; double win_ms_[MAXT + 1] = {}, win_gpu_[MAXT + 1] = {}; @@ -2575,7 +2855,7 @@ index 9a6d1d3..06f1ca9 100644 std::vector last_norm_; // the window's normalized PLE rows (T x hc_dim), for commit // STRATA_METAL_LOGPOS= (test tooling, one-token windows): per decode position the log-probability of the // token that came next (the next window's input: with --spec-follow, the forced text) and the top 64, in the -@@ -321,13 +539,26 @@ public: +@@ -321,13 +540,26 @@ public: ~Window(); void prof_add(bool draft, double ms) { Prof& x = draft ? dprof_ : cprof_; ++x.n; x.all += ms; } private: @@ -2604,7 +2884,7 @@ index 9a6d1d3..06f1ca9 100644 gguf::File fm_; mtl::Buffer Wm_, mnorm_, mR_, memb_, men_, me2_, mhn_, mh2_, midx_, mbo_, minj_[2], mqids_, mlm_; std::map mnorm_off_; -@@ -381,7 +612,16 @@ private: +@@ -381,7 +613,16 @@ private: // per step with the host's pick and embedding (as before) int lchain_ = -1; bool chain_ = false; @@ -2622,7 +2902,7 @@ index 9a6d1d3..06f1ca9 100644 }; namespace { -@@ -411,10 +651,21 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta +@@ -411,10 +652,21 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta ss_ = &ss; max_t_ = std::min(max_t, MAXT); max_cells_ = max_cells; @@ -2645,7 +2925,7 @@ index 9a6d1d3..06f1ca9 100644 eps_ = 1e-6f; const double t0 = now_ms(); if (const long rp = repack_pid_running(shard1); rp > 0 && !(std::getenv("STRATA_METAL_IGNORE_REPACK") && std::string(std::getenv("STRATA_METAL_IGNORE_REPACK")) == "1")) { -@@ -423,6 +674,9 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta +@@ -423,6 +675,9 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta return false; } if (!f_.open(shard1, err)) return false; @@ -2655,7 +2935,7 @@ index 9a6d1d3..06f1ca9 100644 cx_ = std::make_unique(); stream_ = ExpertStream::wanted(f_, *cx_); if (stream_) { -@@ -467,6 +721,7 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta +@@ -467,6 +722,7 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta XA_ = act(N); XF_ = act(N); XY_ = act(6144); XQ_ = act(6144); XH_ = act(NFF); XE_ = act(N); // RoPE: cos/sin [pos][32] for the first 64 dims, NEOX pairs (i, i + 32), in double as build_rope_table const double base = f_.meta_num.count("qwen4exp.rope.freq_base") ? f_.meta_num["qwen4exp.rope.freq_base"] : 1e7; @@ -2663,7 +2943,7 @@ index 9a6d1d3..06f1ca9 100644 for (int p = 0; p < MAXC + 16; ++p) for (int i = 0; i < 32; ++i) { const double th = (double) p * std::pow(base, -2.0 * i / 64.0); -@@ -495,19 +750,37 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta +@@ -495,19 +751,37 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta { const char* e = std::getenv("STRATA_HC_Q8"); hc_q8_ = !(e && std::string(e) == "0"); @@ -2702,7 +2982,7 @@ index 9a6d1d3..06f1ca9 100644 } Q8_ = cx_->alloc((tot + 4095) / 4096 * 4096 + 4096); // + a page: a file with Q8_0 HC already has no copies std::vector th; -@@ -528,12 +801,18 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta +@@ -528,12 +802,18 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta src = raw.data(); } ref::dequant(ref::BF16, src, tmp.data(), ne); @@ -2724,7 +3004,7 @@ index 9a6d1d3..06f1ca9 100644 } if (stream_) { std::vector skip; // the BF16 originals of the Q8_0 HC copies are not read again -@@ -543,9 +822,51 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta +@@ -543,9 +823,51 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta const char* e = std::getenv("STRATA_METAL_PREFETCH"); st_pred_ = !(e && std::atoi(e) <= 0); if (st_pred_) predL_ = A(4 * (size_t) 48 * MAXT * 512); @@ -2777,7 +3057,7 @@ index 9a6d1d3..06f1ca9 100644 if (ple_on_ && T("blk.1.ple_conv1d.weight").type == ref::F32) { const gguf::Tensor& k = T("blk.1.ple_conv1d.weight"); const size_t ne = k.bytes / 4; -@@ -563,14 +884,60 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta +@@ -563,14 +885,60 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta return true; } @@ -2845,7 +3125,7 @@ index 9a6d1d3..06f1ca9 100644 { // windows of several tokens: NR rows a lane share each x load (the window's cost is its x traffic). The sums // run in the same order for any NR, so the answers do not change. STRATA_METAL_GEMV_NR (1, 2): the tile. -@@ -581,6 +948,14 @@ void Window::gemv(std::vector& ds, const gguf::Tensor& t, const A +@@ -581,6 +949,14 @@ void Window::gemv(std::vector& ds, const gguf::Tensor& t, const A if (K == 1 || (L == 32 && K == 2)) Rr = 2; } } @@ -2860,7 +3140,7 @@ index 9a6d1d3..06f1ca9 100644 const std::string fn = "gemv_" + std::to_string(ty) + "_" + std::to_string(L) + "_" + std::to_string(Rr) + "_" + std::to_string(S) + "_" + std::to_string(K); const mtl::Pipeline* p = P(lib_, fn, mtl::Constants().i(0, ntok), fn + "#" + std::to_string(ntok)); -@@ -596,6 +971,36 @@ void Window::gemv(std::vector& ds, const gguf::Tensor& t, const A +@@ -596,6 +972,36 @@ void Window::gemv(std::vector& ds, const gguf::Tensor& t, const A push(ds, d); } @@ -2897,7 +3177,7 @@ index 9a6d1d3..06f1ca9 100644 void Window::xprep(std::vector& ds, const mtl::Buffer& x, const Act& out, int n) { mtl::Dispatch d; d.p = P(lib_, "xprep", {}, "xprep"); -@@ -606,6 +1011,57 @@ void Window::xprep(std::vector& ds, const mtl::Buffer& x, const A +@@ -606,6 +1012,57 @@ void Window::xprep(std::vector& ds, const mtl::Buffer& x, const A push(ds, d); } @@ -2955,7 +3235,7 @@ index 9a6d1d3..06f1ca9 100644 // the hyper-connection read (gr.metal, 3 dispatches): the previous sublayer's write folded in when `apply` void Window::gr_read(std::vector& ds, const std::string& pre, bool apply, const Act& out, bool final_, int ntok) { -@@ -614,14 +1070,7 @@ void Window::gr_read(std::vector& ds, const std::string& pre, boo +@@ -614,14 +1071,7 @@ void Window::gr_read(std::vector& ds, const std::string& pre, boo const gguf::Tensor& tup = T(pre + "_up.weight"); const gguf::Tensor& tin = final_ ? tdn : T(pre + "_inject.weight"); GrArgs a{N, HC, LR, (uint32_t) ntok, apply ? 1u : 0u, eps_}; @@ -2971,7 +3251,7 @@ index 9a6d1d3..06f1ca9 100644 // each matrix: its load-time Q8_0 copy (q8off_), else the file's own tensor (Unsloth's are Q8_0 already, their // inject rows F32) auto src = [&](const gguf::Tensor& t) -> std::pair { -@@ -633,17 +1082,22 @@ void Window::gr_read(std::vector& ds, const std::string& pre, boo +@@ -633,17 +1083,22 @@ void Window::gr_read(std::vector& ds, const std::string& pre, boo const auto [bi, oi] = src(tin); const auto [bu, ou] = src(tup); const bool q8 = is_q8(tdn) && is_q8(tup); @@ -2997,7 +3277,7 @@ index 9a6d1d3..06f1ca9 100644 d3.p = P(lib_, up, mtl::Constants().i(1, ntok), std::string(up) + "#" + std::to_string(ntok)); GrArgs au = a; au.n_tok = MAXT; // the plane stride of the prepared activation -@@ -669,9 +1123,12 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 +@@ -669,9 +1124,12 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 nb_ = true; gemv(ds, T(b + "attn_gate.weight"), XA_, zL_, 6144, ntok, o_z); nb_ = true; @@ -3013,7 +3293,7 @@ index 9a6d1d3..06f1ca9 100644 GdnArgs a{16, 48, 10240, cm, eps_, fuse_ ? (uint32_t) (MAXT * 6144) : 0u}; mtl::Dispatch d; d.p = P(lgdn_, "gdn_step", mtl::Constants().i(2, ntok), "gdn_step#" + std::to_string(ntok)); -@@ -690,13 +1147,16 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 +@@ -690,13 +1148,16 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 gemv(ds, T(b + "ssm_out.weight"), XY_, bo_, N, ntok); } else { const int qi = l / 4; @@ -3032,7 +3312,7 @@ index 9a6d1d3..06f1ca9 100644 nb_ = true; gemv(ds, T(b + "indexer.k_proj.weight"), XA_, ikr_, 128, ntok); QsaArgs a{(uint32_t) pos0, 0, (uint32_t) max_cells_, (uint32_t) ntok, eps_}; -@@ -721,7 +1181,6 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 +@@ -721,7 +1182,6 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 s.grid[0] = (uint32_t) (((pos0 + ntok) / 4 + 1 + 7) / 8); s.grid[1] = ntok; s.tg[0] = 256; @@ -3040,7 +3320,7 @@ index 9a6d1d3..06f1ca9 100644 if (sel) { push(ds, s); if (rec_) rec_->push_back({(int) ds.size() - 1, 2}); -@@ -739,12 +1198,12 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 +@@ -739,12 +1199,12 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 at.p = P(lqsa_, "qsa_attn", {}, "qsa_attn"); at.args = {mtl::Arg::B(qcur_), 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(qids_), mtl::Arg::B(part_), mtl::Arg::V(a)}; @@ -3055,7 +3335,7 @@ index 9a6d1d3..06f1ca9 100644 mtl::Dispatch m; m.p = P(lqsa_, "qsa_merge", {}, "qsa_merge"); QsaArgs am = a; -@@ -759,28 +1218,41 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 +@@ -759,28 +1219,41 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 } // ---- the MoE gr_read(ds, b + "hc_ffn", true, XF_, false, ntok); @@ -3106,7 +3386,7 @@ index 9a6d1d3..06f1ca9 100644 { MoeArgs a{}; const gguf::Tensor &tg_ = T(b + "ffn_gate_exps.weight"), &tu = T(b + "ffn_up_exps.weight"), -@@ -801,33 +1273,64 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 +@@ -801,33 +1274,64 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 a.n_ff = NFF; a.n_tok = ntok; a.topk = TOPK; @@ -3185,7 +3465,7 @@ index 9a6d1d3..06f1ca9 100644 { GlueArgs ga{NFF, MAXT, N, TOPK}; mtl::Dispatch d; -@@ -836,53 +1339,57 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 +@@ -836,53 +1340,57 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 d.grid[0] = NFF / 16; d.grid[1] = ntok; d.tg[0] = 16; @@ -3282,7 +3562,7 @@ index 9a6d1d3..06f1ca9 100644 } namespace { -@@ -952,6 +1459,12 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t +@@ -952,6 +1460,12 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t if (Tn < 1 || Tn > MAXT || (head && Tn > max_t_)) { err = "metal window: window size out of range"; return false; } if (pos0 + Tn + 1 >= max_cells_) { err = "metal window: the window runs past the context"; return false; } const double t0 = now_ms(); @@ -3295,7 +3575,7 @@ index 9a6d1d3..06f1ca9 100644 // ---- the embeddings (CPU dequant of the token rows straight into the shared buffer), broadcast to the streams const gguf::Tensor& te = T("token_embd.weight"); const size_t rb = ref::row_bytes(te.type, N); -@@ -1000,6 +1513,7 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t +@@ -1000,6 +1514,7 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t if (!g.bytes.empty()) qa = &g; std::memcpy(qa->bytes.data(), &a, sizeof a); if (kind == 2) d.grid[0] = (uint32_t) (((pos0 + Tn) / 4 + 1 + 7) / 8); @@ -3303,7 +3583,7 @@ index 9a6d1d3..06f1ca9 100644 } seg_T_ = Tn; const double r0 = now_ms(); -@@ -1007,7 +1521,7 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t +@@ -1007,7 +1522,7 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t { static const int kt = std::getenv("STRATA_METAL_KPROF") ? std::atoi(std::getenv("STRATA_METAL_KPROF")) : 0; static int seen = 0; @@ -3312,7 +3592,7 @@ index 9a6d1d3..06f1ca9 100644 } tp_rg_ = now_ms() - r0; last_n_disp_ = (int) c.ds.size(); -@@ -1017,6 +1531,12 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t +@@ -1017,6 +1532,12 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t Cached& cw = wcache_[key]; rec_ = &cw.qsa; cuts_ = stream_ ? &cw.cuts : nullptr; @@ -3325,7 +3605,7 @@ index 9a6d1d3..06f1ca9 100644 std::vector& ds = cw.ds; nb_ = false; if (fold_n > 0) fold(ds, fold_n, 1 - hp_); // the last window's commit (hp_ already turned) -@@ -1026,32 +1546,27 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t +@@ -1026,32 +1547,27 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t nb_ = true; gemv(ds, T("blk.1.ple_value.weight"), XE_, pval_, N, Tn); } @@ -3372,7 +3652,7 @@ index 9a6d1d3..06f1ca9 100644 push(ds, g2); g3.p = P(lple_, "ple_gate_norm", {}, "ple_gate_norm"); g3.args = {mtl::Arg::B(keyn_), mtl::Arg::B(qn_), mtl::Arg::B(pval_), mtl::Arg::B(W_, off("blk.1.ple_norm_conv.weight")), -@@ -1085,6 +1600,7 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t +@@ -1085,6 +1601,7 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t } rec_ = nullptr; cuts_ = nullptr; @@ -3380,7 +3660,7 @@ index 9a6d1d3..06f1ca9 100644 seg_T_ = Tn; const double r0 = now_ms(); last_gpu_ms_ = 1000.0 * run_ds(cw); -@@ -1095,7 +1611,9 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t +@@ -1095,7 +1612,9 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t double Window::run_ds(Cached& c) { if (!stream_) return cx_->run_graph(c.ds); @@ -3391,7 +3671,7 @@ index 9a6d1d3..06f1ca9 100644 } namespace { -@@ -1164,6 +1682,7 @@ bool Window::stream_arena(std::string& err) { +@@ -1164,6 +1683,7 @@ bool Window::stream_arena(std::string& err) { if (const char* e = std::getenv("STRATA_METAL_KEEP_FREE_GB")) am.keep_gb = std::atof(e); if (const char* e = std::getenv("STRATA_METAL_ARENA_GB")) am.arena_gb = std::atof(e); if (const char* e = std::getenv("STRATA_METAL_LEAVE_SHARE")) am.leave_share = std::atof(e); // for measuring; the default is kLeaveShare @@ -3399,7 +3679,7 @@ index 9a6d1d3..06f1ca9 100644 const double keep = arena_keep(am); const double avail = am.available; double budget = arena_budget(am); -@@ -1179,25 +1698,31 @@ bool Window::stream_arena(std::string& err) { +@@ -1179,25 +1699,31 @@ bool Window::stream_arena(std::string& err) { std::fprintf(stderr, "strata metal stream: %.2f GB allocated before the arena (the process holds %.2f GB, %.2f GB resident); " "RAM %.1f GB (%.1f GB available, %.1f GB kept for macOS), GPU working set %.1f GB\n", used / gb, pm0.footprint / gb, pm0.rss / gb, phys / gb, avail / gb, keep / gb, cx_->recommended_working_set() / gb); @@ -3441,7 +3721,7 @@ index 9a6d1d3..06f1ca9 100644 [&](uint64_t step) { return arena_grow(am, avail_now(), (double) step); })) break; if (st_->shrink_to() == 0 || tries == 4) return false; -@@ -1206,7 +1731,7 @@ bool Window::stream_arena(std::string& err) { +@@ -1206,7 +1732,7 @@ bool Window::stream_arena(std::string& err) { budget / gb, st_->shrink_to() / gb); budget = (double) st_->shrink_to(); } @@ -3450,7 +3730,7 @@ index 9a6d1d3..06f1ca9 100644 { // what the process really holds now, against the plan and the cap (a number that disagrees is a bug in the count) const ProcMem pm1 = proc_mem(); const double cap = footprint_cap(am); -@@ -1300,6 +1825,8 @@ void Window::kprof(const std::vector& ds, int Tn) { +@@ -1300,6 +1826,8 @@ void Window::kprof(const std::vector& ds, int Tn) { } Window::~Window() { @@ -3459,7 +3739,7 @@ index 9a6d1d3..06f1ca9 100644 if (logpos_) std::fclose(logpos_); if (std::getenv("STRATA_METAL_PROFILE") == nullptr) return; auto line = [](const char* what, const Prof& x) { -@@ -1327,6 +1854,53 @@ bool Window::finish(int Tn, const int32_t* tokens, int64_t pos0, bool head, int3 +@@ -1327,6 +1855,53 @@ bool Window::finish(int Tn, const int32_t* tokens, int64_t pos0, bool head, int3 ++x.n; x.pre += tp_pre_ - t0; x.rg += tp_rg_; x.gpu += last_gpu_ms_; x.fin += t1 - tp_pre_ - tp_rg_; x.all += t1 - t0; } @@ -3513,7 +3793,7 @@ index 9a6d1d3..06f1ca9 100644 logpos(Tn, tokens, pos0, head); if (Tn == 1 || partial_) hp_ = 1 - hp_; // the window committed its first row (or all of them) in the step last_head_ = head; -@@ -1388,12 +1962,17 @@ bool Window::checkpoint_save(strata::core::ConversationCheckpoint& c, std::strin +@@ -1388,12 +1963,17 @@ bool Window::checkpoint_save(strata::core::ConversationCheckpoint& c, std::strin std::memcpy(c.gdn.data() + i * (gs + hs) + gs, ghist_[hp_][i].ptr, hs); } c.ple.assign((const uint8_t*) pleHist_.ptr, (const uint8_t*) pleHist_.ptr + 4 * 9 * 10240); @@ -3533,7 +3813,7 @@ index 9a6d1d3..06f1ca9 100644 c.block_pos.clear(); return true; } -@@ -1412,6 +1991,7 @@ void Window::reset() { +@@ -1412,6 +1992,7 @@ void Window::reset() { zero(tail_[j], 4 * 4 * 128); zero(dead_[j], 4 * 128); } @@ -3541,7 +3821,7 @@ index 9a6d1d3..06f1ca9 100644 for (size_t j = 0; j < mtail_.size(); ++j) { zero(mtail_[j], 4 * 4 * 128); zero(mdead_[j], 4 * 128); -@@ -1420,8 +2000,9 @@ void Window::reset() { +@@ -1420,8 +2001,9 @@ void Window::reset() { bool Window::checkpoint_restore(const strata::core::ConversationCheckpoint& c, std::string& err) { const size_t gs = 4 * 128 * 48 * 128, hs = 4 * 10240 * 3, ts = 4 * 4 * 128, dsz = 4 * 128; @@ -3552,7 +3832,7 @@ index 9a6d1d3..06f1ca9 100644 err = "metal window: the checkpoint is not this engine's (sizes differ)"; return false; } -@@ -1436,12 +2017,198 @@ bool Window::checkpoint_restore(const strata::core::ConversationCheckpoint& c, s +@@ -1436,12 +2018,198 @@ bool Window::checkpoint_restore(const strata::core::ConversationCheckpoint& c, s std::memcpy(tail_[j].ptr, c.tails.data() + j * ts, ts); std::memcpy(dead_[j].ptr, c.dead.data() + j * dsz, dsz); } @@ -3751,7 +4031,7 @@ index 9a6d1d3..06f1ca9 100644 if (last_T_ > 1) { // the GDN layers: the kept prefix folded into the state (gdn_step over n_keep tokens from the window's own // inputs, with commit; bitwise the window's trajectory). K/V and indexer cells past the prefix are rewritten -@@ -1507,6 +2274,17 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) +@@ -1507,6 +2275,17 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) { const char* e = std::getenv("STRATA_METAL_MTP_Q8"); if (!(e && std::string(e) == "0")) { @@ -3769,7 +4049,7 @@ index 9a6d1d3..06f1ca9 100644 size_t tot = 0; std::vector names; for (const auto& kv : fm_.tensors) { -@@ -1514,10 +2292,11 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) +@@ -1514,10 +2293,11 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) uint64_t ne = 1; for (uint64_t d : x.shape) ne *= d; if (x.type != ref::BF16 || x.shape.size() < 2 || x.shape[0] % 32 != 0) continue; @@ -3783,7 +4063,7 @@ index 9a6d1d3..06f1ca9 100644 mq8t_[x.name] = q; names.push_back(x.name); tot += (q.bytes + 255) / 256 * 256; -@@ -1533,10 +2312,20 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) +@@ -1533,10 +2313,20 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) const uint64_t ne = x.bytes / 2; tmp.resize(ne); ref::dequant(ref::BF16, (const uint8_t*) fm_.map + x.offset, tmp.data(), ne); @@ -3805,7 +4085,7 @@ index 9a6d1d3..06f1ca9 100644 } } auto A = [&](size_t bytes) { return cx_->alloc(bytes); }; -@@ -1575,6 +2364,25 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) +@@ -1575,6 +2365,25 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) mhead_t_.bytes = dvocab_.size() * rb; } } @@ -3831,7 +4111,7 @@ index 9a6d1d3..06f1ca9 100644 mtp_on_ = true; { const int ty = T("token_embd.weight").type; -@@ -1582,7 +2390,8 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) +@@ -1582,7 +2391,8 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) const char* e = std::getenv("STRATA_METAL_MTP_CHAIN"); chain_ = !(e && std::string(e) == "0") && std::find(std::begin(ok_types), std::end(ok_types), ty) != std::end(ok_types); if (chain_) { @@ -3841,7 +4121,7 @@ index 9a6d1d3..06f1ca9 100644 dout_ = A(4 * MAXT); dprob_ = A(4 * MAXT); posv_ = A(4 * N); for (int i = 0; i < N; ++i) posv_.as()[i] = (float) i; dmap_ = A(4 * std::max(1, dvocab_.size())); -@@ -1630,12 +2439,7 @@ void Window::gr_read_m(std::vector& ds, const std::string& pre, m +@@ -1630,12 +2440,7 @@ void Window::gr_read_m(std::vector& ds, const std::string& pre, m const bool q8 = tdn.type == ref::Q8_0 && tup.type == ref::Q8_0 && tin.type == ref::Q8_0; const mtl::Buffer& WB = q8 ? mq8_ : Wm_; GrArgs a{N, HC, LR, (uint32_t) ntok, apply ? 1u : 0u, eps_}; @@ -3855,7 +4135,7 @@ index 9a6d1d3..06f1ca9 100644 mtl::Dispatch d2; const char* dn = q8 ? "gr_down_q8_8_8" : "gr_down_8_8"; d2.p = P(lib_, dn, mtl::Constants().i(1, ntok), std::string(dn) + "#" + std::to_string(ntok)); -@@ -1686,10 +2490,35 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, +@@ -1686,10 +2491,35 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, // a GEMV on the MTP model's tensor (the window's gemv() reads the main model's buffer) const gguf::Tensor& tt = TQ(name); const int n_in = (int) tt.shape[0], n_out = tt.shape.size() > 1 ? (int) tt.shape[1] : 1, ty = tt.type; @@ -3894,7 +4174,7 @@ index 9a6d1d3..06f1ca9 100644 const std::string fn = "gemv_" + std::to_string(ty) + "_" + std::to_string(Lq) + "_" + std::to_string(Rr) + "_" + std::to_string(S) + "_" + std::to_string(K); GemvArgs a{(uint32_t) n_in, (uint32_t) n_out, (uint32_t) rows, (uint32_t) ref::row_bytes(ty, n_in), (uint32_t) n_in, -@@ -1703,11 +2532,22 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, +@@ -1703,11 +2533,22 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, push(ds, d); }; // ---- e = fc_embedding(rms(embed(tok)) * (1 + w)), h = fc_hidden per stream (rms_10240(R) * (1 + w)), R = h + e @@ -3919,7 +4199,7 @@ index 9a6d1d3..06f1ca9 100644 xp(mhn_, (size_t) c * N * 4, XHn_, N, n, HC * N); gm("mtp.fc_hidden.weight", XHn_, mh2_, N, n, (size_t) c * MAXT * N * 4); } -@@ -1719,21 +2559,21 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, +@@ -1719,21 +2560,21 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, d.grid[0] = (n * HC * N + 255) / 256; d.tg[0] = 256; push(ds, d); } @@ -3945,7 +4225,7 @@ index 9a6d1d3..06f1ca9 100644 d.args = {mtl::Arg::B(qfull_), mtl::Arg::B(kcur_), mtl::Arg::B(vcur_), mtl::Arg::B(midx_), mtl::Arg::B(midx_, 512 * 4), mtl::Arg::B(mnorm_, NM(L + "self_attn.q_norm.weight")), mtl::Arg::B(mnorm_, NM(L + "self_attn.k_norm.weight")), mtl::Arg::B(mnorm_, NM(L + "self_attn.indexer.q_layernorm.weight")), -@@ -1741,24 +2581,32 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, +@@ -1741,24 +2582,32 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, mtl::Arg::B(cs_), mtl::Arg::B(qcur_), mtl::Arg::B(qidx_), mtl::Arg::B(mkq_[0]), mtl::Arg::B(mvq_[0]), mtl::Arg::B(mksc_[0]), mtl::Arg::B(mvsc_[0]), mtl::Arg::B(mtail_[0]), mtl::Arg::B(mdead_[0]), mtl::Arg::B(mpooled_[0]), mtl::Arg::V(qa), mtl::Arg::B(mrope_)}; @@ -3985,7 +4265,7 @@ index 9a6d1d3..06f1ca9 100644 am.n_tok = MAXT; m.args = {mtl::Arg::B(part_), mtl::Arg::B(qfull_), mtl::Arg::B(XQ_.b), mtl::Arg::V(am)}; m.grid[0] = 24; m.grid[1] = 1; m.tg[0] = 256; -@@ -1790,13 +2638,13 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, +@@ -1790,13 +2639,13 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, // the down tile: Q2_0 8,1; else 4,2 - but on the M4 (no M5 layout) Q5_1 4,1 and Q8_0 8,1 (4,2 spills there: // 24-36 GB/s, mac-kern 513fc615) const bool m4 = !m5_layout(); @@ -4002,7 +4282,7 @@ index 9a6d1d3..06f1ca9 100644 r.tg[0] = 512; push(ds, r); { -@@ -1829,18 +2677,14 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, +@@ -1829,18 +2678,14 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, dd.indirect = &plan_; dd.ind_off = P_IDN * 4; push(ds, dd); @@ -4023,7 +4303,7 @@ index 9a6d1d3..06f1ca9 100644 { ArgArgs a{(uint32_t) nv, 120}; mtl::Dispatch d1, d2; -@@ -1858,26 +2702,86 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, +@@ -1858,26 +2703,86 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, // the prompt's K/V: cell cell0 + i pairs the main window's final residual row i with the token at cell0 + i + 1 bool Window::mtp_prefill(const int32_t* next_tokens, int64_t n, int64_t cell0, std::string& err) { if (!mtp_on_) { err = "metal mtp: not loaded"; return false; } @@ -4118,7 +4398,7 @@ index 9a6d1d3..06f1ca9 100644 void Window::mtp_pick(std::vector& ds, int j, bool embed) { struct PickArgs { uint32_t nv, j, use_map; }; const int64_t nv = dvocab_.empty() ? n_vocab_ : (int64_t) dvocab_.size(); -@@ -1907,7 +2811,16 @@ bool Window::mtp_draft(int Tn, const int32_t* tokens, int64_t p, int a, int max_ +@@ -1907,7 +2812,25 @@ bool Window::mtp_draft(int Tn, const int32_t* tokens, int64_t p, int a, int max_ float min_p, int* n_drafts, std::string& err) { if (!mtp_on_) { err = "metal mtp: not loaded"; return false; } (void) Tn; @@ -4132,14 +4412,23 @@ index 9a6d1d3..06f1ca9 100644 mrope_sync(p, a + max_steps + 1); + // the catch-up rows' cells and draft 0's hold committed tokens; the later drafts' do not + mtp_wrote(p, max_steps > 0 ? a + 1 : 0); ++ static const DraftMinP policy = [] { ++ DraftMinP config; ++ if (!config.parse(std::getenv("STRATA_METAL_SPEC_MIN_P_SCHEDULE"), ++ std::getenv("STRATA_METAL_SPEC_MIN_CUM_P"))) { ++ std::fprintf(stderr, "metal mtp: invalid min-p schedule; using caller min_p\n"); ++ } ++ return config; ++ }(); ++ double cumulative = 1.0; int n = 0; int32_t tok = tokens[a]; if (chain_ && max_steps > 0) { -@@ -1923,10 +2836,10 @@ bool Window::mtp_draft(int Tn, const int32_t* tokens, int64_t p, int a, int max_ +@@ -1923,10 +2846,10 @@ bool Window::mtp_draft(int Tn, const int32_t* tokens, int64_t p, int a, int max_ cx_->run_graph(ds); for (int j = 0; j < max_steps; ++j) { const float pj = dprob_.as()[j]; -+ if (min_p > 0.0f && pj < min_p) break; ++ if (!policy.keep(j, pj, min_p, cumulative)) break; drafts[j] = dout_.as()[j]; if (probs) probs[j] = pj; ++n; @@ -4147,7 +4436,7 @@ index 9a6d1d3..06f1ca9 100644 } if (n_drafts) *n_drafts = n; return true; -@@ -1941,6 +2854,7 @@ bool Window::mtp_draft(int Tn, const int32_t* tokens, int64_t p, int a, int max_ +@@ -1941,6 +2864,7 @@ bool Window::mtp_draft(int Tn, const int32_t* tokens, int64_t p, int a, int max_ mtp_encode(ds, mR_, 0, 1, &tok, p + a + j, true); } cx_->run_graph(ds); @@ -4155,15 +4444,22 @@ index 9a6d1d3..06f1ca9 100644 const int32_t di = tokout_.as()[0]; const int32_t d = dvocab_.empty() ? di : dvocab_[(size_t) di]; // the draft's probability: softmax of the draft head's logits at the pick -@@ -1953,6 +2867,7 @@ bool Window::mtp_draft(int Tn, const int32_t* tokens, int64_t p, int a, int max_ +@@ -1950,11 +2874,13 @@ bool Window::mtp_draft(int Tn, const int32_t* tokens, int64_t p, int a, int max_ + double s = 0.0; + for (int64_t i = 0; i < nv; ++i) s += std::exp((double) (lg[i] - mx)); + const float pj = (float) (1.0 / s); ++ if (policy.enabled() && !policy.keep(j, pj, min_p, cumulative)) break; drafts[j] = d; if (probs) probs[j] = pj; tok = d; + if (draft_slop_) draft_state = slop4::step(draft_state, d); ++n; - if (min_p > 0.0f && pj < min_p) break; +- if (min_p > 0.0f && pj < min_p) break; ++ if (!policy.enabled() && min_p > 0.0f && pj < min_p) break; } -@@ -1999,10 +2914,90 @@ struct Window::PromptGemm { + if (n_drafts) *n_drafts = n; + return true; +@@ -1999,10 +2925,90 @@ struct Window::PromptGemm { mtl::Buffer R, bo, inj[2], xnh, V, loh, mixed, mixh, qkv, z, al, be, ygh, gqn, gkn, gvc, ggb, gkq, go, yg, qfull, kcur, vcur, iqr, ikr, qcur, qcurh, qidx, scores, qids, part, XQ, xqh, logits, ids, wts, bcnt, boff, eoff, tiles, meta, stok, sw, pos, XS, HS, YE, yr, gsh, ush, hh, ys, sgv, emb, embh, pkey, pval, keyn, qn, gated, normW; @@ -4254,7 +4550,7 @@ index 9a6d1d3..06f1ca9 100644 bool Window::prompt_gemm_on(int64_t n) const { static const bool on = [] { -@@ -2018,7 +3013,7 @@ bool Window::prompt_gemm_on(int64_t n) const { +@@ -2018,7 +3024,7 @@ bool Window::prompt_gemm_on(int64_t n) const { bool Window::prompt_gemm_init(std::string& err) { if (pg_) return true; @@ -4263,7 +4559,7 @@ index 9a6d1d3..06f1ca9 100644 const char* ec = std::getenv("STRATA_METAL_PROMPT_CHUNK"); // Streaming: every chunk makes one pass over the layer's experts that are not in the arena, so a longer chunk // streams each expert fewer times. M5 Max, UD-Q5, warm: 1024 -> 700 tok/s at 4K, 2048 -> 920, 4096 -> 1,100 (32K the -@@ -2033,10 +3028,15 @@ bool Window::prompt_gemm_init(std::string& err) { +@@ -2033,10 +3039,15 @@ bool Window::prompt_gemm_init(std::string& err) { PromptGemm& p = *pg; p.MT = C; p.QS = std::min(C, 256); @@ -4283,7 +4579,7 @@ index 9a6d1d3..06f1ca9 100644 p.k4 = ::coop_k4(*cx_); const char* eg = std::getenv("STRATA_METAL_PROMPT_EXPERTS"); // v2 | v3 (default: v3 with the M5's layout) p.gu_v3 = p.dn_v3 = eg ? std::string(eg) == "v3" : p.k4; -@@ -2054,8 +3054,13 @@ bool Window::prompt_gemm_init(std::string& err) { +@@ -2054,8 +3065,13 @@ bool Window::prompt_gemm_init(std::string& err) { auto AB = [&](size_t bytes) -> mtl::Buffer { const size_t b = (bytes + 16383) / 16384 * 16384; if (measure) { need += b; return mtl::Buffer{}; } @@ -4299,7 +4595,7 @@ index 9a6d1d3..06f1ca9 100644 region += b; return r; }; -@@ -2074,6 +3079,12 @@ bool Window::prompt_gemm_init(std::string& err) { +@@ -2074,6 +3090,12 @@ bool Window::prompt_gemm_init(std::string& err) { p.part = p.flash ? F32(16) : F32((size_t) p.QS * 33 * 24 * 258); p.XQ = F32(4 * MT * 6144 + 3 * (size_t) p.QS * 6144 + MT * 384); p.xqh = F16(MT * 6144); p.logits = F32(MT * NEXP); p.ids = AB(4 * MT * TOPK); p.wts = F32(MT * TOPK); @@ -4312,7 +4608,7 @@ index 9a6d1d3..06f1ca9 100644 p.bcnt = AB(4 * (size_t) p.nblk_max * NEXP); p.boff = AB(4 * (size_t) p.nblk_max * NEXP); p.eoff = AB(4 * NEXP); p.tiles = AB(8 * (mr / p.TM + 1)); p.meta = AB(64); p.stok = AB(4 * mr); p.sw = F32(mr); p.pos = AB(4 * MT * TOPK); p.XS = F16(mr * N); p.HS = F16(mr * NFF); -@@ -2097,14 +3108,99 @@ bool Window::prompt_gemm_init(std::string& err) { +@@ -2097,14 +3119,99 @@ bool Window::prompt_gemm_init(std::string& err) { if (region == nullptr) { err = "metal prompt: no arena region for the prompt buffers"; return false; } } make(); @@ -4413,7 +4709,7 @@ index 9a6d1d3..06f1ca9 100644 bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::string& err) { PromptGemm& p = *pg_; p.meta.as()[2] = p.gu_v3 ? NFF / (32 * 4) : NFF / 64; -@@ -2121,6 +3217,7 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri +@@ -2121,6 +3228,7 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri pend_fold_ = 0; } const int T = Tn; @@ -4421,7 +4717,7 @@ index 9a6d1d3..06f1ca9 100644 // ---- the embeddings into every stream, the PLE rows { const gguf::Tensor& te = TN("token_embd.weight"); -@@ -2133,17 +3230,34 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri +@@ -2133,17 +3241,34 @@ 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); } } @@ -4459,7 +4755,7 @@ index 9a6d1d3..06f1ca9 100644 std::vector ds; std::vector cuts; nb_ = false; -@@ -2159,10 +3273,16 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri +@@ -2159,10 +3284,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)); }; auto gemm = [&](const mtl::Buffer& A, const mtl::Buffer& B, const mtl::Buffer& Y, int M, int Nn, int K) { @@ -4479,7 +4775,7 @@ index 9a6d1d3..06f1ca9 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 +3292,7 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri +@@ -2172,7 +3303,7 @@ 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; @@ -4488,7 +4784,7 @@ index 9a6d1d3..06f1ca9 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) { -@@ -2184,10 +3304,13 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri +@@ -2184,10 +3315,13 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri const gguf::Tensor &tn = TN(pre + "_norm.weight"), &tdn = TN(pre + "_down.weight"), &tin = TN(pre + "_inject.weight"), &tup = TN(pre + "_up.weight"); dq(tdn, p.WD); @@ -4502,7 +4798,7 @@ index 9a6d1d3..06f1ca9 100644 D(PP(p.lib, "gr_norm_h"), {mtl::Arg::B(p.R), mtl::Arg::B(p.bo), mtl::Arg::B(p.inj[ci]), mtl::Arg::B(W_, tn.offset), mtl::Arg::B(p.xnh), mtl::Arg::V(a)}, HC, T, 256); gemm(p.xnh, p.WD, p.V, T, LR + HC, HC * N); -@@ -2199,28 +3322,42 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri +@@ -2199,28 +3333,42 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri ci = 1 - ci; }; // the PLE projections (the window's gemv pair as GEMMs) @@ -4547,7 +4843,7 @@ index 9a6d1d3..06f1ca9 100644 gemm(p.mixh, p.Wbe, p.be, T, 48, N); PgGdnP ga{(uint32_t) T, 16, 48, 10240, eps_, 0}; D(PP(p.lgdn, "gdn_pre"), -@@ -2236,16 +3373,30 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri +@@ -2236,16 +3384,30 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri gemm(p.ygh, p.Wout, p.bo, T, N, 6144); } else { const int qi = l / 4; @@ -4580,7 +4876,7 @@ index 9a6d1d3..06f1ca9 100644 gemm(p.mixh, p.Wik, p.ikr, T, 128, N); QsaArgs a{(uint32_t) pos0, 0, (uint32_t) max_cells_, (uint32_t) T, eps_}; D(PP(p.lqsa, "qsa_prep", qsa_c(), qsa_k()), -@@ -2262,6 +3413,7 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri +@@ -2262,6 +3424,7 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri D(PP(p.lqsa, "qsa_topk"), {mtl::Arg::B(p.scores), mtl::Arg::B(p.qids), mtl::Arg::V(a)}, T, 1, 1024); if (p.flash) { const uint32_t nq = (uint32_t) T * 24 * 256; @@ -4588,7 +4884,7 @@ index 9a6d1d3..06f1ca9 100644 D(PP(p.lqsa, "qsa_q_f16"), {mtl::Arg::B(p.qcur), mtl::Arg::B(p.qcurh), mtl::Arg::V(nq)}, (nq + 255) / 256, 1, 256); 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]), -@@ -2292,21 +3444,43 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri +@@ -2292,21 +3455,43 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri &tg_ = TN(b + "ffn_gate_shexp.weight"), &tu_ = TN(b + "ffn_up_shexp.weight"), &td_ = TN(b + "ffn_down_shexp.weight"), &eg = TN(b + "ffn_gate_exps.weight"), &ed = TN(b + "ffn_down_exps.weight"); @@ -4637,7 +4933,7 @@ index 9a6d1d3..06f1ca9 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), -@@ -2374,14 +3548,17 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri +@@ -2374,14 +3559,17 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri 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); @@ -4656,7 +4952,7 @@ index 9a6d1d3..06f1ca9 100644 if (l == 0 && ple_on_) { // layer 0's FFN write applied to R, then the PLE block (as the window, for T rows) const gguf::Tensor& tn = TN("blk.1.hc_attn_norm.weight"); -@@ -2414,15 +3591,25 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri +@@ -2414,15 +3602,25 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri if (stream_) { seg_T_ = T; cut_ids_ = p.ids.as(); @@ -4683,7 +4979,7 @@ index 9a6d1d3..06f1ca9 100644 // the window's state after the chunk: the conv history turned, the PLE history and token window advanced hp_ = 1 - hp_; if (ple_on_) { -@@ -2440,6 +3627,8 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri +@@ -2440,6 +3638,8 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri } if (T >= 2) { ss.ple_prev[0] = tokens[T - 2]; ss.ple_prev[1] = tokens[T - 1]; } else { ss.ple_prev[0] = ss.ple_prev[1]; ss.ple_prev[1] = tokens[0]; } @@ -4692,7 +4988,7 @@ index 9a6d1d3..06f1ca9 100644 mtp_src_ = &p.R; // the draft layer's prompt K/V read this chunk's final rows return true; } -@@ -2479,6 +3668,10 @@ bool window_init(void* w, const strata::core::ModelGeometry& g, strata::core::Se +@@ -2479,6 +3679,10 @@ 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; @@ -4703,7 +4999,7 @@ index 9a6d1d3..06f1ca9 100644 } x->set_max_t(max_t); *n_vocab = x->vocab(); -@@ -2494,6 +3687,7 @@ bool window_run(void* w, int T, const int32_t* tokens, int64_t pos0, bool head, +@@ -2494,6 +3698,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); } @@ -4711,7 +5007,7 @@ index 9a6d1d3..06f1ca9 100644 void window_reset() { if (g_the_window != nullptr && g_the_window->ready()) g_the_window->reset(); } -@@ -2505,6 +3699,18 @@ bool window_checkpoint_restore(const strata::core::ConversationCheckpoint& c, st +@@ -2505,6 +3710,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); } @@ -4730,7 +5026,7 @@ index 9a6d1d3..06f1ca9 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 +3761,14 @@ bool prompt_run(const int64_t* tokens, int64_t n, int64_t pos0, const float* con +@@ -2555,13 +3772,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; } @@ -4746,7 +5042,7 @@ index 9a6d1d3..06f1ca9 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 +3778,9 @@ bool prompt_run_rows(const int64_t* tokens, int64_t n, int64_t pos0, +@@ -2571,7 +3789,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]; @@ -5075,6 +5371,93 @@ index 2f72329..ab7fc40 100644 + fails_256_moce, worst_w_256_moce, total_fails == 0 && worst_w < 2e-6 ? "PASS" : "FAIL"); + return total_fails == 0 && worst_w < 2e-6 ? 0 : 1; } +diff --git a/metal/tests/spec_min_p.cpp b/metal/tests/spec_min_p.cpp +new file mode 100644 +index 0000000..ec5ef0c +--- /dev/null ++++ b/metal/tests/spec_min_p.cpp +@@ -0,0 +1,81 @@ ++#include "../src/engine/spec_min_p.hpp" ++ ++#include ++#include ++#include ++ ++#define CHECK(condition) do { \ ++ if (!(condition)) { \ ++ std::fprintf(stderr, "spec_min_p:%d: %s\n", __LINE__, #condition); \ ++ return 1; \ ++ } \ ++} while (false) ++ ++using strata::metal_engine::DraftMinP; ++ ++int main() { ++ DraftMinP policy; ++ CHECK(policy.parse(nullptr, nullptr) && !policy.enabled()); ++ CHECK(policy.parse(" 0.1 , .35, .55 ", nullptr)); ++ CHECK(policy.count == 3 && policy.enabled()); ++ const char* bad_schedules[] = { ++ "", " ", ",", "0,", "0,,1", "-0.1", "1.01", "nan", "inf", "1x", "1;0", ++ "0,0,0,0,0,0,0,0,0", "1e-999" ++ }; ++ for (const char* bad : bad_schedules) { ++ DraftMinP invalid; ++ CHECK(!invalid.parse(bad, nullptr) && !invalid.enabled()); ++ } ++ for (const char* bad : {"", "nan", "inf", "-1", "1.01", "0.1,0.2", "0.1junk"}) { ++ DraftMinP invalid; ++ CHECK(!invalid.parse("0.1,0.2", bad) && !invalid.enabled()); ++ } ++ CHECK(policy.parse("0,1,0,1,0,1,0,1", ".25")); ++ double cumulative = 1; ++ CHECK(policy.keep(0, .5f, .35f, cumulative)); ++ CHECK(!policy.keep(1, .5f, .35f, cumulative)); ++ CHECK(policy.parse("0,0,0", ".25")); ++ cumulative = 1; ++ CHECK(policy.keep(0, .5f, .35f, cumulative)); ++ CHECK(policy.keep(1, .5f, .35f, cumulative)); ++ CHECK(!policy.keep(2, .5f, .35f, cumulative)); ++ CHECK(policy.parse(".1,.2,.8", nullptr)); ++ cumulative = 1; ++ CHECK(policy.keep(0, .1f, .35f, cumulative)); ++ CHECK(policy.keep(1, .2f, .35f, cumulative)); ++ CHECK(policy.keep(2, .8f, .35f, cumulative)); ++ CHECK(!policy.keep(3, .7f, .35f, cumulative)); ++ const float bad_probs[] = { ++ -1.f, 2.f, std::numeric_limits::quiet_NaN(), std::numeric_limits::infinity() ++ }; ++ for (float bad : bad_probs) { ++ cumulative = 1; ++ CHECK(!policy.keep(0, bad, 0, cumulative)); ++ } ++ std::mt19937 rng(12345); ++ int checks = 0; ++ for (int trial = 0; trial < 100000; ++trial) { ++ float probs[8]; ++ for (float& v : probs) v = float(rng() % 101) / 100; ++ CHECK(policy.parse(nullptr, nullptr)); ++ cumulative = 1; ++ for (int j = 0; j < 8; ++j) { ++ CHECK(policy.keep(j, probs[j], .35f, cumulative) == !(probs[j] < .35f)); ++ ++checks; ++ } ++ for (bool cumulative_mode : {false, true}) { ++ CHECK(policy.parse(".1,.35,.55", cumulative_mode ? ".2" : nullptr)); ++ cumulative = 1; ++ double expected = 1; ++ for (int j = 0; j < 8; ++j) { ++ expected *= probs[j]; ++ const float gate = j == 0 ? .1f : j == 1 ? .35f : .55f; ++ const bool keep = probs[j] >= gate && (!cumulative_mode || expected >= .2f); ++ CHECK(policy.keep(j, probs[j], .35f, cumulative) == keep); ++ ++checks; ++ if (!keep) break; ++ } ++ } ++ } ++ std::printf("schedule_checks=%d parser/boundary/invalid-observation checks=pass\n", checks); ++} diff --git a/metal/tests/token_chain.cpp b/metal/tests/token_chain.cpp index a47c763..a0add9b 100644 --- a/metal/tests/token_chain.cpp -- knotmirror 0.0.1