diff --git a/PLAN.md b/PLAN.md index 6f8b31b..4ef4bf3 100644 --- a/PLAN.md +++ b/PLAN.md @@ -369,7 +369,8 @@ caveats carried with these: 35. **Round 36 Decode `predL_` / `predL2_` Token-0 Fast Path When `STRATA_METAL_PREFETCH_DRAFT <= 0` (`-16.0%` Hot / `-6.8%` Cold Across `93` Router GEMVs/Window) + Guarded `TM > 32` Row-Tail Trimming in `gemm_hc_down` & `gemm_swiglu_up` (`audits/opus_r36_decode.md`, `audits/sol_r36_prompt.md`, `audits/argon_r36_qsa.md`):** Added `tile_ntok` to `Window::gemv` (`window.cpp`) so when `STRATA_METAL_PREFETCH_DRAFT <= 0` (`= 0` in production, where `ExpertStream::ahead` only reads `t = 0`), the `93` router-guess GEMVs (`47` `predL_` + `46` `predL2_`, `Q8_0` `2560 -> 512`) run at `ntok = 1` on the window's tile (`gemv_8_16_1_4_1#1`), producing **100% bit-for-bit-identical `t = 0` predictions** while cutting their GPU time by **`-16.0%` hot (`21.78 -> 18.30 us`) and `-6.8%` cold (`23.91 -> 22.29 us`)**; added `if constexpr (TM > 32)` row-tail (`TM = 16 / 32`) trimming to `gemm_hc_down` (**`-4.71%` at `T=1025`, `-3.73%` at `T=1041`**) and `gemm_swiglu_up` (**`-1.42%` at `T=1025`, `-1.14%` at `T=1041`, `-0.78%` at `T=1987`**) while preserving `TM = 32` (`SG = 2`) kernels byte-for-byte (`252/252` `epi` bitwise checks `PASS`); and confirmed `qsa_prep` / `qsa_score_mm` are at their M5 register/barrier optimum. 36. **Round 37 Multi-Chunk Prompt Hot-Expert LRU & Scratch Pool Preservation (`stream.h`, `stream.cpp`, `window.cpp`):** Passed `multi_chunk = (n > C)` into `prompt_region(true, n > C)` -> `hold_region(true, multi_chunk)` so 1-chunk prompts (`<= 2,048` tokens) run 100% identically (`42.49 tok/s` / `2,989.17 ms` `@2K` decode — new all-time record, `88.9%` accept), while multi-chunk prompts (`> 2,048` tokens, e.g., `prompt_albedo_11k.txt` at `10,180` tokens) keep `!x.hot` scratch slots at `lru_back(s)` across chunks, mark `pass == 1` demotions `x.hot = false` in `victim()` so `pass == 0` reuses them instead of evicting more `hot` decode experts, and restrict `allow_hot` lookahead to `distance == 1` — boosting `@10.2K` decode from `25.26 tok/s` (`1,227.18 ms`, `45.2%` accept, `2,713` SSD reads) to **`34.01–34.74 tok/s` (`892.47–911.51 ms`, `+34.6–37.5%` faster, `78.3%` accept, `2,073` SSD reads / `-23.6%` fewer SSD reads)** and `@10.2K` prefill from `423.48 tok/s` (`24,038.85 ms`) to **`458.34–472.89 tok/s` (`21,527–22,210 ms`, `-1.83s to -2.51s` faster)**, with pure GPU `T=1983` prefill (`kprof`) hitting **`3,131.29 ms` (`633.28 tok/s` — new all-time record)**. 37. **Round 38 `hold_region` Layer-0 `allow_hot = false` + Prompt/Decode `bytes_read_`/`reads_` Split + `mtp_prefill` `tail_cap` Sync with `STRATA_METAL_MTP_WIN` (`stream.h`, `stream.cpp`, `window.cpp`, `audits/astra_r38_qsa.md`):** Passed `allow_hot = false` to `victim()` during `hold_region(true)`'s initial Layer-0 prefetch (`stream.cpp:1080`) and fallback static lookahead (`stream.cpp:890`) so `hold_region` only fills the `251` `!hot` scratch slots of `cls_[0]` (`IQ2_XXS`) instead of evicting `~180` `hot` decode experts for tail Layer-0 experts that `at_layer(0)` cancels 5 ms later; split `p_bytes_read_`/`p_reads_` from decode `bytes_read_`/`reads_` and zeroed `bytes_read_`/`reads_` after the `6.5 GB` startup arena fill (`fixing the 1-window prefill tail report from 116,128.9 MB / 79,058 reads to the exact 375.2 MB / 258 reads`); synced `Window::mtp_prefill`'s default `tail_cap` with `STRATA_METAL_MTP_WIN` clamped to `[64, 2050]`; and verified in `qsa_parity --ab` (`audits/astra_r38_qsa.md`) that caching `keys[8]` across `qsa_topk2`'s 1024-thread barriers drops M5 occupancy (`+3.5–12.9%` slower at `pos0 = 10180`), pushing `@10.2K` prefill to **`473.80 tok/s` (`21,485.94 ms` — new all-time record!)**, `@10.2K` decode to **`35.21–35.47 tok/s` (`873.92–880.47 ms`, `mtp_ms = 106.34 ms` — new all-time record!)**, `@2K` `prompt_2048tok_ops` decode to **`42.57 tok/s` (`2,982.99 ms`, `mtp_ms = 194.09 ms` — new all-time record!)**, and `reasoning` decode to **`37.20 tok/s` (`5,886.81 ms` — new all-time record!)**. - 38. **Measured End-to-End Speed (`32K` Context, `6.5 GB` Arena, Unpruned `K=512` `IQ2_XS`, `call_3817687`, `call_2483052`, `call_3102851`, `call_2898988`, `call_4110683`, `call_4178375`, `call_3948986`, `call_4500449`, `call_3444931`, `call_3116575`, `call_4373130`, `call_2769318`, `call_3781775`, `call_4410602`, `call_3511955`, `call_4432071`, `call_2918543`, `call_1514614`, `call_1173886`, `call_2456120`, `call_5024768`, `call_4223761`, `call_3543222`, `call_1624568`, `call_3678353`, `call_4256391`, `call_3418394` & `call_3706633`):** `prompt_2048tok_ops` prefill **`400.88–432.78 tok/s` (`4,584.29–4,949.14 ms` end-to-end — all-time record `4,584.29 ms` / `432.78 tok/s` — `-1,728 ms` / `+37.7%` faster end-to-end, **`3,131.29 ms` / `633.28 tok/s` pure GPU `kprof` — `3,267.43 ms` serialized sum!**, `2,047` dispatches (`-47` dispatches/chunk!), `1.12 GB` transient scratch!)**, decode **`41.53–42.27 tok/s`** (`23.66–24.08 ms/tok`, **`3,004.68–3,005.02 ms` decode, `mtp_ms = 190.62–197.99 ms` / `1.63 ms/draft` — new all-time MTP record `190.62 ms`!**); `prompt_albedo_11k` (`10,180` tokens) prefill **`421.49–446.59 tok/s` (`22.79–24.15 s` vs `28.25 s`, `-5.46 s` faster — all-time record; Chunk 0 = `4,613–4,719 ms`, Chunk 1 = `501.5 tok/s` / `4,084–4,572 ms`, Chunk 2 = `4,362–4,736 ms`, Chunk 3 = `4,610–4,911 ms`, Chunk 4 = `4,685–4,898 ms`, `mtp_ms = 121.54–125.14 ms`, and `@10.2K` decode hit **`25.00–25.26 tok/s` / `39.59–40.01 ms/tok` / `1,227.18–1,240.21 ms` — all-time record `25.26 tok/s` at `10.2K` context!** vs `22.06 tok/s` / `1,405.37 ms`!)**; `coding` prefill **`120.61–126.78 tok/s` (`2,413–2,537 ms`, `+32.2%`)**, decode **`35.46–37.1 tok/s`** (**`6,115.57–6,176.28 ms` decode — all-time record `6,115.57 ms`, `mtp_ms = 374.22–380.29 ms`**); `reasoning` prefill **`121.57–122.92 tok/s` (`2,416–2,458 ms`, `+33.3%`)**, decode **`34.92–36.74 tok/s`** (**`5,961.16–6,270.75 ms` decode, `mtp_ms = 322.48–370.78 ms` — all-time record `322.48 ms`!**); `persona_chat` prefill **`124.62–128.14 tok/s` (`2,255.32–2,443 ms`, `+38.5%` — all-time record `2,255.32 ms` / `128.14 tok/s`!)**, decode **`28.50–33.2 tok/s`** (**`6,628.00–7,685.17 ms` decode, `mtp_ms = 374.58–424.81 ms` — all-time record `374.58 ms`!**). + 38. **round 39 — selected strata v0.1.45 ports and checked local fixes:** incremental prompt tokenization, bounded parser scans, content-keyed logit-bias caching, mac telemetry, thinking budgets/answer reservation, and PLE prefetch-key reuse (`36381540`). local review removed upstream's stale identity/length bias shortcut, capped prompt-cache entries at 16, and cached/released the telemetry device query. preserved the text-only image-omission switch. metal prompt load safety and the default moe parity tile are fixed; the local `qsa_merge2` batching experiment is bitwise exact (40 edge cases), but **only tests dispatch it; live engine decode remains on `qsa_merge`**. rejected the cuda score butterfly (memcmp failure) and `qsa_merge` batching (repeatable nt=3 slowdown). active-chunk bounds were already present. standalone CLI smoke: 42.46 tok/s at 1,984 prompt tokens and 34.23 tok/s at 10,180; these are **not paired speedup evidence**. a later alternating before/after check produced identical token IDs in all five pairs, but substantial temporal drift in both builds left performance inconclusive (2k candidate decode time +1.5–1.7%; 10k mixed -3.6% to +2.8%). no zero-regression claim. serving now enables upstream `answer_room` by default (for thinking with output caps >=1,024 and no explicit budget), which can inject a wrap-up instruction; disable it explicitly for parity comparisons to prior service behavior. exact verification and known baseline failures: `audits/r39_integration.md`, `audits/r39_serve.md`, `audits/r39_opus_review.md`; evidence: `results/round39/`. + 39. **Measured End-to-End Speed (`32K` Context, `6.5 GB` Arena, Unpruned `K=512` `IQ2_XS`, `call_3817687`, `call_2483052`, `call_3102851`, `call_2898988`, `call_4110683`, `call_4178375`, `call_3948986`, `call_4500449`, `call_3444931`, `call_3116575`, `call_4373130`, `call_2769318`, `call_3781775`, `call_4410602`, `call_3511955`, `call_4432071`, `call_2918543`, `call_1514614`, `call_1173886`, `call_2456120`, `call_5024768`, `call_4223761`, `call_3543222`, `call_1624568`, `call_3678353`, `call_4256391`, `call_3418394` & `call_3706633`):** `prompt_2048tok_ops` prefill **`400.88–432.78 tok/s` (`4,584.29–4,949.14 ms` end-to-end — all-time record `4,584.29 ms` / `432.78 tok/s` — `-1,728 ms` / `+37.7%` faster end-to-end, **`3,131.29 ms` / `633.28 tok/s` pure GPU `kprof` — `3,267.43 ms` serialized sum!**, `2,047` dispatches (`-47` dispatches/chunk!), `1.12 GB` transient scratch!)**, decode **`41.53–42.27 tok/s`** (`23.66–24.08 ms/tok`, **`3,004.68–3,005.02 ms` decode, `mtp_ms = 190.62–197.99 ms` / `1.63 ms/draft` — new all-time MTP record `190.62 ms`!**); `prompt_albedo_11k` (`10,180` tokens) prefill **`421.49–446.59 tok/s` (`22.79–24.15 s` vs `28.25 s`, `-5.46 s` faster — all-time record; Chunk 0 = `4,613–4,719 ms`, Chunk 1 = `501.5 tok/s` / `4,084–4,572 ms`, Chunk 2 = `4,362–4,736 ms`, Chunk 3 = `4,610–4,911 ms`, Chunk 4 = `4,685–4,898 ms`, `mtp_ms = 121.54–125.14 ms`, and `@10.2K` decode hit **`25.00–25.26 tok/s` / `39.59–40.01 ms/tok` / `1,227.18–1,240.21 ms` — all-time record `25.26 tok/s` at `10.2K` context!** vs `22.06 tok/s` / `1,405.37 ms`!)**; `coding` prefill **`120.61–126.78 tok/s` (`2,413–2,537 ms`, `+32.2%`)**, decode **`35.46–37.1 tok/s`** (**`6,115.57–6,176.28 ms` decode — all-time record `6,115.57 ms`, `mtp_ms = 374.22–380.29 ms`**); `reasoning` prefill **`121.57–122.92 tok/s` (`2,416–2,458 ms`, `+33.3%`)**, decode **`34.92–36.74 tok/s`** (**`5,961.16–6,270.75 ms` decode, `mtp_ms = 322.48–370.78 ms` — all-time record `322.48 ms`!**); `persona_chat` prefill **`124.62–128.14 tok/s` (`2,255.32–2,443 ms`, `+38.5%` — all-time record `2,255.32 ms` / `128.14 tok/s`!)**, decode **`28.50–33.2 tok/s`** (**`6,628.00–7,685.17 ms` decode, `mtp_ms = 374.58–424.81 ms` — all-time record `374.58 ms`!**). ### B. ruled out from the original blueprint diff --git a/audits/astra_r39_qsa.md b/audits/astra_r39_qsa.md new file mode 100644 index 0000000..56ba472 --- /dev/null +++ b/audits/astra_r39_qsa.md @@ -0,0 +1,71 @@ +# round 39 track a — qsa upstream fixes and local merge experiment + +## decision + +- keep `831ba47a`'s clamped-row/unconditional-load/value-select fix in the four non-cooperative prompt fills. preserve our staged scales and cooperative paths. +- keep `ad734531`'s `TK = 64` default and explicit `v2 || v3` gather/HS/binding conditions. this tree already sets `v2 = v3 = true` for `--v3`, so the binding change is documentary/defensive here, not a newly repaired v3 failure. existing half-input gather and `--gu-bench`, `--dn-bench`, `--ab`, `--dtile` remain intact. +- keep the **local** four-partial-load batching experiment in `qsa_merge2`. this is inspired by `731528be`, not a literal upstream transplant. active-chunk bounds already existed and were not changed. +- reject the exact requested six-shuffle `qsa_score` butterfly: `scores FAIL`, exit 1, on its first `--ab` case at p=3000. no speed claim for a numerically invalid candidate. cuda's reduction-order argument does not establish equivalence to metal `simd_sum`; the mismatch is measured, its compiler-level cause is not proven. +- reject four-partial batching in `qsa_merge`: bitwise correct, but nt=3 regressed repeatedly by 4.8–6.6% across all three contexts. `qsa.metal` is entirely unchanged in the delivered patch. + +## artifacts and reproducibility + +- patch: `/tmp/astra_r39_qsa/astra_r39_qsa.patch` +- durable copy, source snapshots, evaluators, metadata, commands, exit statuses and full logs: `results/round39/astra_r39_qsa/` +- exact baseline: the working-tree snapshot, not merely git head; `metadata.json` records its head and four owned-file sha256 values. `baseline-kernels/` and `baseline-moe_mm_parity.cpp` preserve inputs. `rejected-score.metal` and `rejected-merge1.metal` preserve rejected experiments; `final/` preserves the delivered sources. +- all editing/building ran in `/tmp/astra_r39_qsa/strata`; no tracked vendor files, live service or full-model tests were touched. +- device: apple m5; system version recorded in `metadata.json`. apple clang 21.0.0, release `-O2`. build used `cmake -S metal -B build-metal -G Ninja` followed by `cmake --build build-metal --target qsa_parity qsa_prompt_parity moe_mm_parity -j 3`. +- tests were serialized on the gpu. initial moe runs exceeded the original 15s/20s limits and were terminated; after parent explicitly extended CPU-reference allowance to 180s, both full tests completed. all unsuccessful attempts remain in the logs. +- `git -C vendor/strata apply --check /tmp/astra_r39_qsa/astra_r39_qsa.patch`: exit 0. scratch `git diff --check`: exit 0. final visual diff/context review completed. + +## correctness and test results + +all commands below run from the scratch tree; `BASE=/tmp/astra_r39_qsa/baseline`. + +| test | result | +|---|---| +| baseline `qsa_parity --ab $BASE 3000` | exit 0, all memcmps pass | +| baseline `moe_mm_parity` | exit 3, missing `mm_gu_42_16_64_16_2` | +| final `moe_mm_parity` (no args) | exit 0, 23.961s, `PASS`; fp16-reference rel l2 2.03e-06, exact-reference 2.37e-04 | +| final `moe_mm_parity --v3 --gu IQ2_S --dn Q2_0 --tile 16,32,32,4 --dtile 64,32,2` | exit 0, 20.011s, `PASS`; fp16-reference rel l2 1.60e-05, exact-reference 4.38e-04 | +| same v3 command plus `--ab $BASE` | exit 0, HS/YE/Y memcmp pass, t=256 | +| default and specified v3 plus `--T 1` | exit 0, CPU-reference parity pass | +| v3 plus `--T 1 --gu-bench --ab $BASE`, and `--dn-bench` | both exit 0; HS/YE/Y memcmp pass; flags preserved | +| `qsa_prompt_parity` | exit 0, `ALL PASS`, default cooperative layout | +| `COOP_K4=0 qsa_prompt_parity` | exit 0, `ALL PASS`, edited non-cooperative layout | +| `qsa_prompt_parity --mm-compare attn 96 P $BASE/qsa_prompt.metal 3` | p=0 and 30000, both COOP_K4=0 and default: exit 0; output, post-timing and merge+prepared memcmp pass | +| `qsa_prompt_parity --flash-compare 96 P $BASE/qsa_prompt.metal 1 3` | same four combinations: exit 0, memcmp pass | +| `QSA_ATTN=2 qsa_parity --ab $BASE P` | p=3000,10180,30000: all full-buffer memcmps pass; nt=1,2,3,4,8 plus prompt prep nt=1984,2048 | +| `QSA_ATTN=2 qsa_parity --bitwise $BASE/qsa2.metal` | exit 0; zero of 40 cases differ, including short contexts/remainders | + +important distinction: moe's no-argument and specified-v3 CPU checks are tolerance checks and print `PASS`, not `ALL PASS` or bitwise equality to CPU. the separate v3 A/B command provides bitwise equality to the baseline kernels. prompt's default suite likewise includes tolerance checks; the comparison modes provide bitwise evidence for both edited kernels. no unsupported claim of IEEE equality to a differently accumulated CPU reference is made. + +## merge2 paired timings + +three sweeps: initial, confirmation, then final with rejected merge1 restored. all pass. the final sweep below uses only retained changes. measurements are median GPU µs; paired ratio is median(candidate/baseline), not necessarily the quotient of medians. inspected evaluator code uses **five paired rounds**, 16 dispatch repetitions for decode (the file header's 12-repetition description is stale). each stage is warmed before interleaved A/B timing. + +| p | nt | baseline µs | candidate µs | paired candidate/baseline | +|---:|---:|---:|---:|---:| +| 3000 | 1 | 15.143 | 9.273 | 0.6122 | +| 3000 | 2 | 21.924 | 21.516 | 0.9812 | +| 3000 | 3 | 33.688 | 32.969 | 0.9791 | +| 3000 | 4 | 45.458 | 43.227 | 0.9522 | +| 3000 | 8 | 91.102 | 86.547 | 0.9501 | +| 10180 | 1 | 10.961 | 6.844 | 0.6240 | +| 10180 | 2 | 21.977 | 21.589 | 0.9834 | +| 10180 | 3 | 33.732 | 32.961 | 0.9772 | +| 10180 | 4 | 45.344 | 43.313 | 0.9546 | +| 10180 | 8 | 91.393 | 86.607 | 0.9493 | +| 30000 | 1 | 17.180 | 10.807 | 0.6295 | +| 30000 | 2 | 22.125 | 21.547 | 0.9739 | +| 30000 | 3 | 33.500 | 33.471 | 0.9955 | +| 30000 | 4 | 45.469 | 43.544 | 0.9577 | +| 30000 | 8 | 91.594 | 88.021 | 0.9556 | + +nt=1 improves 35–41% across the earlier sweeps and 37–39% in the final sweep; other cases are neutral to about 5% faster. one initial p=30000, nt=3 ratio was 1.0008 (0.08% noise); subsequent ratios were 0.9902 and 0.9955. no repeatable regression. absolute clocks/times drift (even unchanged stages), so use paired ratios, not cross-run raw times. this is a standalone-kernel result, not evidence of whole-model throughput improvement. + +for the rejected merge1, confirmation nt=3 ratios were 1.0551, 1.0660, 1.0563 at p=3000,10180,30000. despite wins elsewhere, these violate the requested neutral-to-faster-across-suite rule. + +## prompt safety fix timing observation + +with `COOP_K4=0`, p=30000/t=96 non-cooperative attn median fell from 4.954 to 1.537ms, flash from 5.685 to 1.828ms, with memcmp pass. p=0 also improved. these three-pair observations are secondary to the safety fix, not broad performance claims. cooperative mode remains bitwise identical. actual m1 hardware was not available; the m1 fault rationale comes from upstream `831ba47a`, while our runtime evidence is m5 with both constant settings. diff --git a/audits/r39_integration.md b/audits/r39_integration.md new file mode 100644 index 0000000..1484b13 --- /dev/null +++ b/audits/r39_integration.md @@ -0,0 +1,63 @@ +# round 39 integration + +## scope + +selected upstream changes from strata `296b59d2..477ac13c`; not an engine-wide upgrade. serve/tokenizer files carry the range plus local correctness hardening. PLE uses only `36381540`, not the later upstream mutex change. metal adaptations are documented in `astra_r39_qsa.md`. + +- no cuda butterfly: score memcmp failed. +- no batching in production `qsa_merge`: repeatable slowdown. +- retained `qsa_merge2` batching is test-only; the live engine does not dispatch it. qsa2.metal still belongs to the runtime-compiled shader library; the hardened service startup exercised that compile. +- active-chunk bounds already existed. no credit assigned to this round. +- prompt safety edits affect non-cooperative branches, not the M5 default path. +- local image omission, model files, and sampling config unchanged. +- upstream answer reservation remains enabled as approved: large thinking requests without a budget reserve answer space and may receive an injected wrap-up sentence in its reasoning ("I have thought about this long enough; time to give my answer."). the policy also logs its budget per request for caps <=32K. this is a behavior change; service persona evaluations need an explicit `answer_room` setting for comparability. CLI smoke does not run that service policy. + +## parent verification + +logs are under `results/round39/`. + +- integration ninja build: exit 0 (`integration-build.log`). +- `QSA_ATTN=2 qsa_parity --bitwise /qsa2.metal`: exit 0, all 40 cases bitwise equal. +- `qsa_prompt_parity`, default and `COOP_K4=0`: exit 0, ALL PASS (mixed tolerance and explicit memcmp checks). +- `moe_mm_parity` no args: exit 0, fp16-reference relative L2 2.03e-06. +- requested IQ2_S/Q2_0 `--v3` configuration: exit 0, fp16-reference relative L2 1.60e-05. +- real local tokenizer: 18 tests, exit 0, one skip for unavailable external llama.cpp oracle vectors. +- bias/image compatibility: 17 tests, exit 0, including same-length mutation and JSON-key equivalence. +- final telemetry/cache hardening: 39 tests, exit 0. real M5 ctypes/ioreg probe also exits 0. +- PLE new reuse fixture: all 10 cases x 2 cache settings x 2 worker settings pass, targeted runner exit 0. full baseline and candidate selftests both exit 1 on the same pre-existing keepalive timing assertions; thresholds unchanged. do not describe the full C++ suite as green. +- parent full serve suite: 852 tests in 517.636s, exit 0, 11 skips (`parent-serve-full.log`). imports predate the last telemetry/cache-cap hardening; the 39-test final-source rerun above covers those changes. the serve owner’s separate baseline (785 tests) and candidate (852 tests) runs both hit the same existing keepawake acquire/release scheduling race; no production code or timing threshold was changed to suppress it. +- complete `scripts/strata_metal_complete.diff`, including eight new serve tests: applies cleanly to isolated vendor HEAD (`complete-patch-check.log`). + +## integration smoke, not a speed comparison + +one model process at a time, existing service stopped for standalone runs and restarted afterwards. same configured 6.5 GB arena and K512 model. CPU-only fake-engine suites also running, so these measurements are not clean performance comparisons. + +| prompt | prefill | decode | exit | +|---|---:|---:|---:| +| 1,984 tokens | 4,941.54 ms | 127 tokens / 2,991.30 ms (42.46 tok/s) | 0 | +| 10,180 tokens | 23,498.33 ms | 31 tokens / 905.63 ms (34.23 tok/s) | 0 | + +`model-smoke-plan.json` records arguments and environment; full logs retained. no whole-model speedup claimed. + +live HTTP roundtrip returned `ready`, then `still ready`; second request reported 18 cached tokens. final hardened service `/health` reports loaded=true on localhost:8080. telemetry reports the driver’s in-use memory statistic, **not established free unified memory**; do not enable the optional `min_free_vram_mib` gate based on this round. the per-sample device query is gone, but the 1 Hz ioreg subprocess remains; its serving overhead was not A/B measured. + +## review and limits + +parent reviewed production diff/context, metal patch, PLE fixture and local regressions; serve owner reviewed all tracked and new test files. opus independently fuzzed 38,033 real-tokenizer cases (PromptEncoder vs full encode directly), with zero mismatches; that tokenizer source is unchanged since the fuzz run. no CUDA/HIP/SYCL or M1 hardware verification this round. full evidence and source fingerprints: `results/round39/`; individual agent audits distinguish measured results from inferred portability. + +complete patch replay reproduces all 27 changed source files exactly. vendor `git diff --check` flags only the pre-existing blank line at qsa.metal EOF; that file is byte-identical to the pre-round snapshot. + +non-empty request logit biases remain unsupported by the Metal engine (no logit_bias=1 INFO capability); that cache was checked through unit tests, not claimed exercised by the live chat smoke. + +## quiet paired CLI followup + +same-day baseline reconstruction independently checked by opus: only the PLE object differs in the host archive; all other linked objects match. runtime shader directories restore the exact prior versions. both arms run the same inputs/env, with service stopped and other agent suites complete. source/build provenance, frozen shaders, commands, full logs and summary: `results/round39/model-ab/`. + +excluded one warmup per arm, then 10k order A/B/B/A/A/B and 2k A/B/B/A. every run exits 0 and **all five pairs produce identical generated token IDs**. + +| case | candidate/baseline decode time, adjacent pairs | candidate/baseline prefill time | +|---|---|---| +| 10k | 1.0137, 0.9643, 1.0280 | 1.0284, 0.9103, 1.0663 | +| 2k | 1.0146, 1.0170 | 1.1043, 0.9660 | + +**performance remains inconclusive, not a zero-regression certificate.** no consistent 10k slowdown; both 2k comparisons are ~1.5–1.7% slower, below the declared 2% material-slowdown threshold but not evidence of exact neutrality. strong temporal drift affects both arms: baseline 10k prefill rises from 23.4 to 28.5s and decode from 908 to 997ms; candidate drifts too. thermal throttling is plausible on the fanless laptop but was not measured, so no causal claim. do not treat the median 10k advantage as an improvement. the short 31-token decode further limits precision. no more repeated hot-device runs; a cooled/stationary experiment would be needed to resolve small differences. server restarted and health confirmed loaded=true. diff --git a/audits/r39_opus_review.md b/audits/r39_opus_review.md new file mode 100644 index 0000000..6bf1801 --- /dev/null +++ b/audits/r39_opus_review.md @@ -0,0 +1,190 @@ +# round 39 opus review — v0.1.45 serve/tokenizer/ple ports + astra qsa patch + +read-only review. no source edits, gpu runs or service actions by me. scope: prompt cache, logit_bias cache, +telemetry fallback, ple dedup, astra qsa patch. paths are relative to `vendor/strata/` unless noted. + +## what i compared + +- `/tmp/v0145_serve_ple.patch` applied to a clean `git archive HEAD` copy (`/tmp/r39_audit_tree`, exit 0). the patch + does **not** apply to the dirty tree (`serve/server.py:3852` conflicts with the local `STRATA_OMIT_IMAGES_WHEN_NO_VISION` + block); that full `git apply` failed atomically, so nothing in it landed. the serve/tokenizer parts went in + through the serve owner (working-tree server.py = the clean-applied copy plus the local omission block). PLE went + in as **only upstream commit 36381540** (`results/round39/ple.patch`: ngram.cpp + ple_reader_test.cpp, no mutex). + the later upstream commit that adds the `std::recursive_mutex` was **not imported**; the mutex was never in our + tree, so nothing was dropped or reverted. +- `/tmp/astra_r39_qsa/astra_r39_qsa.patch` against current `metal/`. + +## verdict + +no remaining correctness blocker once the logit_bias `_last` removal (already done by the serve owner) is in. +three should-fix items before benchmark numbers or a persona eval are attributed to this build: telemetry +cost on darwin (S1), answer_room default (S2), prompt-cache cap (S3). + +## should-fix + +### S1. telemetry runs ioreg and creates a Metal device in the server every second (darwin) + +- `serve/server.py:6485` `serve()` always calls `start_telemetry()`; on darwin `gpu_reader` now returns `_Metal` + (`serve/telemetry.py` `gpu_reader`), so the sampler thread runs `ioreg -l -w0 -r -c IOAccelerator` every second + (measured here on the M5: 39 ms wall, 57 KB of output; the patch says 22 ms) and + `_Metal.read()` calls `metal_working_set()` on **every sample**, which does `MTLCreateSystemDefaultDevice()` + (create rule: +1 retain that is never released) and the objc lookups each time. before the port, darwin had no + gpu reader, so this is new work in the serving process during decode. +- fix: compute the working set once (`_Metal.__init__`) and reuse it in `read()` and in `Telemetry.static`. +- test: patch `telemetry.metal_working_set` with a counter, build `_Metal(run=fixture)`, call `read()` 3x, assert + 1 call. before citing any tok/s from this build: paired decode A/B with the sampler on vs `gpu_reader` forced to + `_Nvml` (same engine, same prompts). +- semantics (not a blocker; it only matters for opt-in `min_free_vram_mib` / image policy): on this M5 at idle + `"In use system memory"=15007744` (15 MB) vs `"Alloc system memory"=15473033216` (15.5 GB). free = + working set − in-use probably reports far more free than there is while the wired arena is up. needs one check + with the engine loaded before anyone turns on `min_free_vram_mib` on the mac. + +### S2. `answer_room` is on by default and changes every thinking request on our config + +- `serve/server.py:4206-4214`: with no `reasoning_budget_tokens`, every thinking request with `max_new >= 1024` + gets budget = max_new − max(512, max_new/4). our config has `fit_max_tokens: true` and 32K context, so a request + without max_tokens gets max_new = room and a budget of 75% of it, plus a `print` on **every** request + (`max_new <= 32768` is always true at 32K ctx). +- when it fires, `REASONING_WRAP_UP` (`"I have thought about this long enough; time to give my answer."`) is injected + into the reasoning stream, and the generation splits into a second engine pass. `frontend.py` additionally turns + Anthropic `thinking.budget_tokens` into a hard budget. +- effort `low` keeps this rare, but persona/antislop evals before and after the port are not comparable unless + this is pinned. decide explicitly: `"answer_room": false` in `models/K512_iq2_xs/strata-k512-iq2xs.json` for + parity runs, or accept and note it. the per-request print should go behind a once-flag either way. +- model-card sampling (`model_card_sampling`): checked: a no-op for our config, because the `sampling` block sets + temperature, top_p, top_k and presence_penalty, which covers every key in both qwen4exp sets. + +### S3. PromptEncoder `keep` has no upper bound + +- `serve/server.py` Service init: `keep = conversation_slots(engine)` is taken from `--conversation-cache-slots` with no clamp; + the engine accepts up to INT32_MAX (`src/program/generate.cpp:2113`). each entry holds the full prompt str + (1–4 B/char: kaomoji/CJK push the str to UCS-2/4), its ids list (8 B/token) and its mark lists, so ~2–5 MB at 100K + tokens. our config passes no slots, so keep = 4 and this is latent. +- fix: `keep=min(slots, 16)` (or a total-chars budget). test: `conversation_slots` stub returning 10_000 → assert + `svc.prompts.keep <= 16`; feed 20 unrelated prompts → `len(entries) <= keep`. + +## checked and fine + +### prompt cache exactness / mutation / concurrency (`tools/strata_tokenizer.py` PromptEncoder) + +- the margin argument holds for `_encode_matching`: a match decision at q depends only on `text[q:q+max_special_len]` + and on whether q is in a plain span; the finditer state before a cut c (end of a taken match) is identical when + `c + max_special_len - 1 <= same`. `call_end`-style lookbehind does not exist in the literal alternation. +- fuzz (`/tmp/r39_fuzz.py`, `/tmp/r39_fuzz2.py`, real Swift tokenizer from `models/K256_reap_l2_sum/tokenizer`, 33 + specials, max len 21): conversations with appends, mid-inserts, truncations, specials split across appends + (`<|im_sta` + `rt|>`), random plain spans, `store=False` mixed in, keep 1/2/4. 38,033 cases, 8,529 of them reusing + a prefix, **0 mismatches** vs `tok.encode(text, parse_special=True, plain=...)`. +- mutation: `encode(store=True)` returns `list(ids)`, `store=False` returns a list that is not stored, entries are + replaced and never changed in place; `prepare` never mutates `ids` in place (images build `out`, force builds a new + list), so `forget_prompt(encoded_ids)` matches. the local omission block runs before `encode_prompt` + (`serve/test_image_omission.py` covers it). +- concurrency: snapshot under the lock, BPE outside it, replace under the lock against the current list: a race + loses reuse at worst, never data. +- `pick_slot` `TURN_ID = 248045` matches `<|im_start|>` in our vocab. +- parser `_wrapped_call_end` fast path: sound, because `call_end` only ever returns positions of a literal + ``, and `_call_end_checked` is reset on every path that leaves the `call` state. + +### logit_bias (`serve/logit_bias.py`) + +- `_last` identity+len staleness: found by the parent, already removed by the serve owner. +- content key `json.dumps`: no collisions that change the result: `{1: x}` and `{"1": x}` serialize alike and + normalize alike; float/bool keys, NaN and inf never normalize, so they are never stored; `False` vs `0` and + `1` vs `"1"` serialize differently. results are copied (`normalize` → `dict(...)`), and the engine key is an + immutable str. the lazy `entry[1]` write is a benign race (same string). +- cost note only: without `_last`, a 65K-entry list pays `json.dumps` on each of the ~5 calls per request. + compute the key once in the route if that shows up in profiles. +- note: the Metal engine does not report `logit_bias=1`, so `validate_request` refuses non-empty biases on our + backend anyway. + +### PLE dedup (`src/kernels/ngram.cpp` `gather_batch`) + +- any-slot matching plus `matched[t]` reuse is correct (keys are row ids; the table is immutable). dropping + `n_prefetch >= n_tokens` is correct now that duplicates map to one slot. draining every ticket before returning + an error fixes the old early return, which left tickets in flight writing into slots that would be reused. +- Metal uses Direct mode (`metal/src/strata_window_cli.cpp:945`), so this path is live: MTP `prefetch_rows` + (`metal/src/engine/window.cpp:2972`) feeds the verify window's `gather_batch` (T = spec+1 ≤ 16). prompt chunks (2048) + take the slow path, which drops the prefetches; that is fine. +- the later upstream `recursive_mutex` commit was not imported (only 36381540 was transplanted), so the table's + threading is unchanged from our baseline. on Metal that baseline is fine: `ple_th` (`window.cpp:3434`) is joined + inside the forward, and the MTP prefetch runs after it, so they never overlap. other backends (the desktop HIP build + shares this file; upstream's mutex comment is about prompt-path / decode-window overlap there) carry the same + unchanged baseline hazard and are untested in this round. record it in the port notes as "not imported". even with the mutex, `issue()` copies into `impl_->rows` before + its `pending` check (pre-existing), so a second `issue()` while one is in flight corrupts the first ticket's keys: + move the `pending` check first if this ever becomes multi-threaded. + +### astra qsa patch + +- `qsa_merge2` is dispatched only by `metal/tests/token_chain.cpp:249` and `metal/tests/qsa_parity.cpp:49,159,254`. + the engine dispatches `qsa_merge` (v1) at `metal/src/engine/window.cpp:1285,2737,3662` and + `metal/src/engine/prompt_chain.inc:262`. `qsa2.metal` is still compiled into `lqsa_` (`window.cpp:763`), so the change is + compile-surface only: **zero runtime effect, no end-to-end attribution**. the unroll keeps the fma accumulation order + (it only hoists loads), which matches astra's bitwise result. +- `qsa_prompt.metal` clamped-row edits are in the `!coop_k4` branches only; on the M5, `Window::coop_k4()` probes true, + so they run only with `COOP_K4=0` or on M1/M4. it is a safety port, with no M5 default-path effect. row 0 / scale 0 for + empty cells is in bounds whenever the cache holds ≥1 cell. +- `moe_mm_parity` `TK=64` / `v2||v3`: test-only. the baseline default was already broken (missing + `mm_gu_42_16_64_16_2`, exit 3). + +## tests i ran (cpu only, `PYTHONDONTWRITEBYTECODE=1`, from `vendor/strata`) + +- `python -m unittest serve.test_prompt_encoder serve.test_logit_bias serve.test_image_omission serve.test_model_card_sampling serve.test_admit_order`: 33 tests OK + (before the serve owner's `_last` fix landed). +- `python -m unittest serve.test_server.MacTelemetry serve.test_server.AmdTelemetry`: 7 OK. +- `python tools/test_strata_tokenizer.py`: 11 OK, 3 skipped. +- PromptEncoder fuzz above: 0 mismatches. +- not run: anything on the gpu, ple_reader_test (the parent ran it), the full serve suite. + +## followup v2 — sdk ownership check for S1 + +- `MacOSX.sdk/System/Library/Frameworks/Metal.framework/Headers/MTLDevice.h` (CommandLineTools sdk on this machine): + `MTL_EXTERN id __nullable MTLCreateSystemDefaultDevice(void) API_AVAILABLE(macos(10.11), ios(8.0)) NS_RETURNS_RETAINED;` + so the caller owns +1, and each uncached `metal_working_set()` call leaked one retain on the device singleton. that + is a refcount, not memory per call, and the cost is the per-sample framework call. with the working set cached + once, a single unreleased +1 at startup is acceptable. an optional `objc_msgSend(dev, sel_registerName("release"))` + after reading would make it exact, but it is not needed. +- the same header says a call switches a dual-gpu intel mac to the high-power gpu. it does not apply to apple + silicon, but it is one more reason to call it once and not every second. +- answer_room: parent approved keeping the live default; it stays disclosed in S2, with no comparable eval numbers + claimed. + +## followup v3 — hardening review (final sources) + +- the fingerprints in `results/round39/integration-source-sha256.json` match the working tree for all 27 files. +- telemetry: `_Metal.__init__` reads the working set once (`self.working_set`), and both `read()` and `Telemetry.static` + reuse it. `metal_working_set()` releases the device in a `finally` (`objc_release`, matching the + `NS_RETURNS_RETAINED` annotation); the early `if not dev` return happens before the try, so a null is never released. covered by + `test_working_set_is_read_once_for_samples_and_static` plus the release assertion. what remains per sample: the 1 Hz `ioreg` + subprocess (39 ms wall measured earlier); its cost on decode is still not A/B measured. `Service.free_vram_mib()` still + builds a fresh reader per call (2 ioreg runs + 1 working-set read), but only on the opt-in and image paths. +- prompt cache: `PROMPT_ENCODER_KEEP_MAX = 16` caps slots from args and from INFO, and the default stays 4. covered by + `IncrementalPrompts` (asserts 10**9 → 16 on both paths). `tools/strata_tokenizer.py` is byte-identical to the + fuzzed version, so the 38,033-case result still applies to the final source. +- logit_bias: `_last` is fully removed and only the bounded JSON-keyed cache remains; nothing else changed. +- reran on the final sources (cpu): test_prompt_encoder, test_logit_bias, test_image_omission, MacTelemetry, AmdTelemetry: + 33 OK; IncrementalPrompts: 9 OK. +- verdict: S1 (per-sample device creation) and S3 (cap) are resolved. S2 is accepted and disclosed by the parent. + +## followup v4 — provenance of `results/round39/model-ab/` (read-only; no processes launched) + +- `ngram-baseline.cpp` is byte-identical to `git show HEAD:src/kernels/ngram.cpp` (ngram.cpp was clean before the + round). against the current file the diff is 15+/11-, i.e. the 36381540 dedup only. +- host archive: 34 members, same names and order as `build-metal/libstrata_metal_host.a`. extracted and hashed: only + `ngram.cpp.o` and the ranlib `__.SYMDEF` index differ, and the swapped member equals `model-ab/ngram.cpp.o`. no other + linked archive defines `PleTable::gather_batch` (`libstrata_metal_kernels.a` only references it), so link order + cannot pick a stale copy. the disassembly of `PleTable::gather_batch` differs between `baseline-cli` (340 lines) + and `candidate-cli` (391 lines). +- `candidate-cli` is byte-identical to `build-metal/strata_window_cli`, built at 11:59:22, after the last round39 + source write (qsa files 11:59:21, ngram.cpp 11:56:32). +- shaders: the runtime directory comes from `STRATA_METAL_KERNELS` (`metal/src/mtl_rt.mm:51`), with no on-disk library + cache in mtl_rt. `candidate-kernels/` equals the current `metal/kernels/` file for file (22/22). + `baseline-kernels/` differs from it only in `qsa2.metal` and `qsa_prompt.metal`, matches `/tmp/astra_r39_qsa/baseline` + exactly, and matches the sha256 values in astra's `metadata.json` for qsa/qsa2/qsa_prompt. + the engine-private msl (`mtp.metal`, `mtp_chain.metal`, `ple.metal`, `slop4.h`) loads from the compiled-in absolute + path for both arms; all four were last written before the round (02:57–08:11, and slop4.h on 10-17), so this is shared and + unchanged. +- consequence for interpreting results: on the M5 the shader delta is inert (`qsa_merge2` is not dispatched, and the + non-coop branches are not taken), so the arms differ in effect only by the PLE gather_batch dedup in the MTP verify path. +- suggestions, not blockers: (1) both arms should emit identical token ids, because the dedup returns the same bytes; + compare the generated ids per pair as a correctness gate. (2) the 10K case decodes ~31 tokens in ~0.9 s, which is too + short a window to resolve a 2% decode delta with 3 pairs; call that case's decode inconclusive unless the spread is + clearly under 2%, and read prefill separately. diff --git a/audits/r39_serve.md b/audits/r39_serve.md new file mode 100644 index 0000000..f701347 --- /dev/null +++ b/audits/r39_serve.md @@ -0,0 +1,69 @@ +# round39 serve + tokenizer port + +## scope and source + +- upstream range: `296b59d2f73c8f04a805df29c2fa97655adace78..477ac13cbe8e582c20484871ce459d20d82df09d`. +- read `vendor/strata/AGENTS.md`, `docs/DETAILS.md` thinking/API sections, `docs/LOGIT_BIAS.md`, design-preflight and orchestrating-agents. +- owner boundaries: tokenizer owns exact incremental encoding; service owns request preparation/cache admission, thinking budgets and API defaults; telemetry owns fail-open platform readers. no engine/model/GPU execution or running service changes. +- all 21 upstream serve/tokenizer changed files ported. 16 match upstream exactly; exceptions: server.py retains the pre-existing local image omission hunk unchanged and caps prompt retention; logit_bias.py removes unsafe identity caching; test_logit_bias.py adds regression coverage and removes obsolete _last resets; telemetry.py caches/releases the working-set device; test_server.py covers those hardening changes. one local test file added. +- complete owned patch, including untracked upstream tests: `results/round39/serve-owned-final.patch`. +- sha256s of every owned final file and upstream counterpart: `results/round39/serve-source-hashes.json`. + +## behavior and deviations + +ports prompt encoder resume points, plain-span-aware prefix reuse and rejected-prompt forgetting; parser wrapped-call scan cursor and bounded suffix checks; bounded content-keyed logit-bias LRU; apple IOKit telemetry; anthropic hard thinking budgets and automatic answer reservation. also carries upstream admission ordering (opt-in), turn-aware slot choice, sampling/model-card/web defaults, scrape-window throughput gauges, config diagnostics, and prompt telemetry fields. + +local image omission still mutates content lists before image processing and encoding, exactly as before. new tests prove retained text-part identity, omitted images cached as text, full/incremental token parity, repeated fresh/mutated-object requests, disabled-switch error, and no mutation when vision exists. + +parent review identified upstream's `_last` identity+length fast path returning stale values after same-length mutation. removed it rather than asserting request immutability. every large input now receives a JSON content key; four-entry bounded LRU and cached engine-key strings remain. new tests cover nested list/dict edits to valid/invalid values and JSON-equivalent inputs/duplicate IDs. red-on-upstream evidence: `serve-logit-bias-upstream-repro.log` (one test, two failing subtests); fixed bias suite: 14 passed. + +hardening requested after opus review: service retention is capped at the explicit `PROMPT_ENCODER_KEEP_MAX = 16`, including huge engine INFO and lazy-start argument values. apple working-set budget is read once per `_Metal` instance and reused by `read` and static telemetry. the local SDK's `/Library/Developer/CommandLineTools/SDKs/MacOSX.sdk/System/Library/Frameworks/Metal.framework/Headers/MTLDevice.h:130` declares `MTLCreateSystemDefaultDevice` `NS_RETURNS_RETAINED`; the helper now releases it in `finally`, including unsupported-selector returns. mock tests count one helper call through three reads and Telemetry static initialization, and verify device release on success/unsupported selectors. + +## validation + +python: `/tmp/r39-serve-venv/bin/python` (3.14.7); dependencies: `results/round39/serve-python-packages.txt`. initial system python lacked jinja2/regex; initial import-failure logs retained separately. baseline tree copied to `/tmp/r39-serve-baseline`, restoring only owned paths to source baseline plus original local image hunk; shared checkout never reverted. + +- tokenizer baseline: 6 passed. +- tokenizer default candidate: 11 run, 3 skips (real-vocabulary classes/oracle unavailable by default). +- tokenizer with `STRATA_TOKENIZER=$PWD/models/K256_reap_l2_sum/tokenizer`: 18 run, 1 skip (external llama.cpp oracle vectors absent), no failures; reads vocabulary only, no model execution. +- image compatibility: 3 passed. +- final logit bias: 14 passed. +- hardening affected suites: 39 passed in 3.786s (`serve-hardening-affected.log`), with real vocabulary enabled; rerun after final readability adjustment also passes 39 (`serve-hardening-final.log`). +- first full serve baseline/candidate attempts hit 300s deadlines; partial logs retained as `serve-baseline.log` / `serve-post.log`. no completed suite count claimed for these. +- completed full baseline: 785 tests in 494.236s, 1 failure, 11 skips (`serve-baseline-full.log`). +- completed full candidate (pre-hardening imports, before one extra bias-key test): 852 tests in 516.658s, 1 failure, 6 skips (`serve-final-full.log`). +- the sole failure is identical on baseline/candidate: `test_keepawake.KeepAwakeTests.test_service_holds_for_the_whole_request_including_early_close`, line 87. this is an existing asynchronous-observation race, not a port regression: the mock request can acquire and release its hold before the keep-awake daemon acquires `_cond`; `_loop` then sees `_count == 0` with `held == False` and waits without emitting either transition. the test assumes the daemon observed the transient busy state. deterministic baseline/candidate reproductions block that condition lock through the request and both record `count=0, calls=[]`, then only idle reset on close (`keepawake-mechanism-{baseline,candidate}.log`). no timeout increase, expectedFailure, or unrelated product/test edit applied. +- final hardening is checked by the separate 39-test affected-suite run, not falsely folded into the older full-run count. + +## changed files + +- `vendor/strata/serve/frontend.py` +- `vendor/strata/serve/logit_bias.py` +- `vendor/strata/serve/runconfig.py` +- `vendor/strata/serve/server.py` +- `vendor/strata/serve/telemetry.py` +- `vendor/strata/serve/test_admit_order.py` +- `vendor/strata/serve/test_effort_paths.py` +- `vendor/strata/serve/test_logit_bias.py` +- `vendor/strata/serve/test_model_card_sampling.py` +- `vendor/strata/serve/test_parallel.py` +- `vendor/strata/serve/test_parser_call_scan.py` +- `vendor/strata/serve/test_parser_hold.py` +- `vendor/strata/serve/test_prompt_encoder.py` +- `vendor/strata/serve/test_q36_config.py` +- `vendor/strata/serve/test_reasoning_loop_recovery.py` +- `vendor/strata/serve/test_server.py` +- `vendor/strata/serve/test_slot_alloc.py` +- `vendor/strata/serve/test_web_sampling.py` +- `vendor/strata/serve/web/app.js` +- `vendor/strata/tools/strata_tokenizer.py` +- `vendor/strata/tools/test_strata_tokenizer.py` +- `vendor/strata/serve/test_image_omission.py` + +## final review + +view_diff reviewed every tracked owned diff; view_code reviewed every new upstream/local test file and final local mutation/key-equivalence tests. tokenizer reuse helpers are shared with service rather than duplicated. upstream comments/style retained for port fidelity. no assertion or test removed except obsolete `_last` reset setup lines; upstream loop-recovery tests explicitly disable automatic answer_room to continue testing the original loop-guard path. + +## parent independent verification + +parent reports an independent full-suite run passed: 852 tests in 517.636s, 11 skips, exit 0; its independent hardening selection also passed 39/39. these are parent-reported results, not this agent's own runs. the baseline-shared keepawake race above did not reproduce in the parent full run. no additional suites launched after parent began quiet paired CLI benchmarking. diff --git a/results/round39/astra_r39_qsa/astra_r39_qsa.patch b/results/round39/astra_r39_qsa/astra_r39_qsa.patch new file mode 100644 index 0000000..f2b59e6 --- /dev/null +++ b/results/round39/astra_r39_qsa/astra_r39_qsa.patch @@ -0,0 +1,108 @@ +diff --git a/metal/kernels/qsa2.metal b/metal/kernels/qsa2.metal +index d5b594c..30439a1 100644 +--- a/metal/kernels/qsa2.metal ++++ b/metal/kernels/qsa2.metal +@@ -120,7 +120,15 @@ kernel void qsa_merge2(device const float* part [[buffer(0)]], device const floa + } + float o = 0.0f; + device const float* q = pp + tid; +- for (int k = 0; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); ++ int k = 0; ++ for (; k + 3 < nc; k += 4, q += 4 * STRIDE) { ++ const float p0 = q[0], p1 = q[STRIDE], p2 = q[2 * STRIDE], p3 = q[3 * STRIDE]; ++ o = fma(p0, ek[k], o); ++ o = fma(p1, ek[k + 1], o); ++ o = fma(p2, ek[k + 2], o); ++ o = fma(p3, ek[k + 3], o); ++ } ++ for (; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); + threadgroup_barrier(mem_flags::mem_threadgroup); + o = o / L_sh; + const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; +diff --git a/metal/kernels/qsa_prompt.metal b/metal/kernels/qsa_prompt.metal +index 444e0c4..2ee4a4b 100644 +--- a/metal/kernels/qsa_prompt.metal ++++ b/metal/kernels/qsa_prompt.metal +@@ -82,8 +82,11 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* + for (uint16_t i = 0; i < rk.get_capacity(); ++i) { + const auto ix = rk.get_multidimensional_index(i); + const int k = k0 + ix[0], id = cid[ix[1]]; +- const uint row = uint(id) * QKV + g; +- rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)])) : half(0.0f); ++ // the row is clamped and the select is on the value: a select with the device load in the arm ++ // loses the threadgroup on the M1 (macOS 26, measured); the discarded value cannot matter ++ const uint row = uint(max(id, 0)) * QKV + g; ++ const float v = float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)]); ++ rk[i] = id >= 0 ? half(v) : half(0.0f); + } + } else + #pragma clang loop unroll(full) +@@ -136,7 +139,9 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* + for (uint16_t i = 0; i < rv.get_capacity(); ++i) { + const auto ix = rv.get_multidimensional_index(i); // (cell, dim) + const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; +- rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vs[c][d >> 6])) : half(0.0f); ++ const uint row = uint(max(id, 0)) * QKV + g; ++ const float v = float(vq[row * QD + d]) * float(vs[c][d >> 6]); ++ rv[i] = id >= 0 ? half(v) : half(0.0f); + } + } else { + uint vo[8]; +@@ -332,8 +337,9 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq + for (uint16_t i = 0; i < rk.get_capacity(); ++i) { + const auto ix = rk.get_multidimensional_index(i); + const int k = k0 + ix[0], id = cid[ix[1]]; +- const uint row = uint(id) * QKV + g; +- rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ks_tg[ix[1]][k >> 6])) : half(0.0f); ++ const uint row = uint(max(id, 0)) * QKV + g; ++ const float v = float(kq[row * QD + k]) * float(ks_tg[ix[1]][k >> 6]); ++ rk[i] = id >= 0 ? half(v) : half(0.0f); + } + } else + #pragma clang loop unroll(full) +@@ -398,7 +404,9 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq + for (uint16_t i = 0; i < rv.get_capacity(); ++i) { + const auto ix = rv.get_multidimensional_index(i); // (cell, dim) + const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; +- rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vs_tg[c][sg])) : half(0.0f); ++ const uint row = uint(max(id, 0)) * QKV + g; ++ const float v = float(vq[row * QD + d]) * float(vs_tg[c][sg]); ++ rv[i] = id >= 0 ? half(v) : half(0.0f); + } + } else { + // a lane holds the same eight cells (two runs of four, 16 apart) at each of its dims: as for rk, row +diff --git a/metal/tests/moe_mm_parity.cpp b/metal/tests/moe_mm_parity.cpp +index 8e06a77..a44397c 100644 +--- a/metal/tests/moe_mm_parity.cpp ++++ b/metal/tests/moe_mm_parity.cpp +@@ -176,7 +176,7 @@ int main(int argc, char** argv) { + argc > 3 ? std::atoi(argv[3]) : 512, + argc > 4 ? std::atoi(argv[4]) : N_EMBD); + } +- int gu = ref::Q2_0, dn = ref::Q2_0, T = 256, TM = 16, TN = 64, TK = 16, SG = 2, DTN = -1, DTK = -1, DSG = -1; ++ int gu = ref::Q2_0, dn = ref::Q2_0, T = 256, TM = 16, TN = 64, TK = 64, SG = 2, DTN = -1, DTK = -1, DSG = -1; + bool bench = false, v2 = false, v3 = false, v4 = false, v5 = false; + int MS = 1; + bool gu_only = false, dn_only = false; +@@ -283,17 +283,17 @@ int main(int argc, char** argv) { + if (gemm) { + MmArgs ma{(uint64_t)layer * NE * blob, (uint64_t)layer * NE * blob + up_off, (uint64_t)layer * NE * blob + down_off, + blob, blob, (uint32_t)gu_row, (uint32_t)d_row, N_EMBD, N_FF}; +- if (v2) { ++ if (v2 || v3) { + const uint32_t ne = N_EMBD; + mtl::Dispatch ga; ga.p = &pgat; ga.args = {mtl::Arg::B(xh), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; + ga.grid[0] = meta.as()[1]; ga.tg[0] = 256; ds.push_back(ga); + } + mtl::Dispatch g; g.p = &pgu; +- if (v2) g.args = {mtl::Arg::B(W), mtl::Arg::B(XS), mtl::Arg::B(HS), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; ++ if (v2 || v3) g.args = {mtl::Arg::B(W), mtl::Arg::B(XS), mtl::Arg::B(HS), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + else g.args = {mtl::Arg::B(W), mtl::Arg::B(x), mtl::Arg::B(H), mtl::Arg::B(stok), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + g.grid[0] = N_FF / (v3 ? TN * SG : TN); g.grid[1] = meta.as()[0]; g.tg[0] = 32 * SG; ds.push_back(g); +- mtl::Dispatch d; d.p = &pdn; d.args = {mtl::Arg::B(W), mtl::Arg::B(v2 ? HS : H), mtl::Arg::B(YE), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; +- d.grid[0] = N_EMBD / (v3 ? DTN * DSG : v2 ? DTN : TN); d.grid[1] = meta.as()[0]; d.tg[0] = 32 * (v2 ? DSG : SG); ds.push_back(d); ++ mtl::Dispatch d; d.p = &pdn; d.args = {mtl::Arg::B(W), mtl::Arg::B(v2 || v3 ? HS : H), mtl::Arg::B(YE), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; ++ d.grid[0] = N_EMBD / (v3 ? DTN * DSG : v2 ? DTN : TN); d.grid[1] = meta.as()[0]; d.tg[0] = 32 * ((v2 || v3) ? DSG : SG); ds.push_back(d); + const uint32_t ca[4] = {(uint32_t)T, N_EMBD, TOPK, 0}; + mtl::Dispatch cb; cb.p = &pcomb; cb.args = {mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(y), mtl::Arg::V(ca)}; + cb.grid[0] = (T * N_EMBD + 255) / 256; cb.tg[0] = 256; ds.push_back(cb); diff --git a/results/round39/astra_r39_qsa/baseline-kernels/alu_bench.metal b/results/round39/astra_r39_qsa/baseline-kernels/alu_bench.metal new file mode 100644 index 0000000..e34890a --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/alu_bench.metal @@ -0,0 +1,181 @@ +// metal/kernels/alu_bench.metal - instruction-rate probes (alu_bench.cpp): each thread runs ITER iterations of 8 +// independent chains of one operation; the result is written so nothing is eliminated. +#define ITER 1024 + +kernel void a_fma32(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + float a0 = gid, a1 = gid + 1, a2 = gid + 2, a3 = gid + 3, a4 = gid + 4, a5 = gid + 5, a6 = gid + 6, a7 = gid + 7; + for (int i = 0; i < ITER; ++i) { + a0 = fma(a0, k, 0.5f); a1 = fma(a1, k, 0.5f); a2 = fma(a2, k, 0.5f); a3 = fma(a3, k, 0.5f); + a4 = fma(a4, k, 0.5f); a5 = fma(a5, k, 0.5f); a6 = fma(a6, k, 0.5f); a7 = fma(a7, k, 0.5f); + } + out[gid] = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7; +} + +kernel void a_fma16(device float* out [[buffer(0)]], constant float& kf [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + const half k = half(kf); + half a0 = gid & 7, a1 = 1, a2 = 2, a3 = 3, a4 = 4, a5 = 5, a6 = 6, a7 = 7; + for (int i = 0; i < ITER; ++i) { + a0 = fma(a0, k, half(0.5)); a1 = fma(a1, k, half(0.5)); a2 = fma(a2, k, half(0.5)); a3 = fma(a3, k, half(0.5)); + a4 = fma(a4, k, half(0.5)); a5 = fma(a5, k, half(0.5)); a6 = fma(a6, k, half(0.5)); a7 = fma(a7, k, half(0.5)); + } + out[gid] = float(a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7); +} + +// float accumulator, half operands (does the conversion come free?) +kernel void a_fmamix(device float* out [[buffer(0)]], constant float& kf [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + half h0 = half(gid & 7), h1 = h0 + half(1); + float a0 = gid, a1 = 1, a2 = 2, a3 = 3, a4 = 4, a5 = 5, a6 = 6, a7 = 7; + for (int i = 0; i < ITER; ++i) { + a0 = fma(float(h0), kf, a0); a1 = fma(float(h1), kf, a1); a2 = fma(float(h0), kf, a2); a3 = fma(float(h1), kf, a3); + a4 = fma(float(h0), kf, a4); a5 = fma(float(h1), kf, a5); a6 = fma(float(h0), kf, a6); a7 = fma(float(h1), kf, a7); + h0 = h0 * half(0.999); h1 = h1 * half(0.999); + } + out[gid] = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7; +} + +kernel void a_iand(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint a0 = gid, a1 = gid + 1, a2 = gid + 2, a3 = gid + 3, a4 = gid + 4, a5 = gid + 5, a6 = gid + 6, a7 = gid + 7; + for (int i = 0; i < ITER; ++i) { + a0 = (a0 ^ k) + 1; a1 = (a1 ^ k) + 1; a2 = (a2 ^ k) + 1; a3 = (a3 ^ k) + 1; + a4 = (a4 ^ k) + 1; a5 = (a5 ^ k) + 1; a6 = (a6 ^ k) + 1; a7 = (a7 ^ k) + 1; + } + out[gid] = float(a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7); +} + +kernel void a_shift(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint a0 = gid, a1 = gid + 1, a2 = gid + 2, a3 = gid + 3, a4 = gid + 4, a5 = gid + 5, a6 = gid + 6, a7 = gid + 7; + for (int i = 0; i < ITER; ++i) { + a0 = (a0 >> k) | k; a1 = (a1 >> k) | k; a2 = (a2 >> k) | k; a3 = (a3 >> k) | k; + a4 = (a4 >> k) | k; a5 = (a5 >> k) | k; a6 = (a6 >> k) | k; a7 = (a7 >> k) | k; + } + out[gid] = float(a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7); +} + +// int -> float conversion + fma +kernel void a_cvt(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint u = gid; + float a0 = 0, a1 = 0, a2 = 0, a3 = 0, a4 = 0, a5 = 0, a6 = 0, a7 = 0; + for (int i = 0; i < ITER; ++i) { + u += k; + a0 += float(u & 3); a1 += float(u & 12); a2 += float(u & 48); a3 += float(u & 192); + a4 += float(u & 768); a5 += float(u & 3072); a6 += float(u & 12288); a7 += float(u & 49152); + } + out[gid] = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7; +} + +// threadgroup memory loads: random-ish 8-byte reads from a 4 KB table +kernel void a_tgload(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint2 tab[512]; + for (int i = tid; i < 512; i += 256) tab[i] = uint2(i, i * 3); + threadgroup_barrier(mem_flags::mem_threadgroup); + uint u = gid * 2654435761u; + uint2 s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += tab[(u >> 3) & 511]; s1 += tab[(u >> 7) & 511]; s2 += tab[(u >> 11) & 511]; s3 += tab[(u >> 15) & 511]; + u += k; + } + out[gid] = float(s0.x + s1.x + s2.x + s3.x + s0.y + s1.y + s2.y + s3.y); +} + +// select (sign application) +kernel void a_sel(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint u = gid; + float a0 = 1, a1 = 2, a2 = 3, a3 = 4, a4 = 5, a5 = 6, a6 = 7, a7 = 8; + for (int i = 0; i < ITER; ++i) { + u ^= k; + a0 = (u & 1) ? -a0 : a0; a1 = (u & 2) ? -a1 : a1; a2 = (u & 4) ? -a2 : a2; a3 = (u & 8) ? -a3 : a3; + a4 = (u & 16) ? -a4 : a4; a5 = (u & 32) ? -a5 : a5; a6 = (u & 64) ? -a6 : a6; a7 = (u & 128) ? -a7 : a7; + } + out[gid] = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7; +} + +// dot-shaped: acc_i += q_i * x_i with loop-carried accumulators only (q, x vary slowly) +kernel void a_dot32(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + float q[8], x[8], a[8]; + for (int j = 0; j < 8; ++j) { q[j] = float(gid + j); x[j] = k * j; a[j] = 0; } + for (int i = 0; i < ITER; ++i) { + for (int j = 0; j < 8; ++j) a[j] = fma(q[j], x[j], a[j]); + q[0] += 1.0f; + } + float s = 0; + for (int j = 0; j < 8; ++j) s += a[j]; + out[gid] = s; +} +kernel void a_dot16(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + half q[8], x[8], a[8]; + for (int j = 0; j < 8; ++j) { q[j] = half(gid + j); x[j] = half(k * j); a[j] = 0; } + for (int i = 0; i < ITER; ++i) { + for (int j = 0; j < 8; ++j) a[j] = fma(q[j], x[j], a[j]); + q[0] += half(1.0f); + } + float s = 0; + for (int j = 0; j < 8; ++j) s += float(a[j]); + out[gid] = s; +} +kernel void a_dotmix(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + half q[8]; + float x[8], a[8]; + for (int j = 0; j < 8; ++j) { q[j] = half(gid + j); x[j] = k * j; a[j] = 0; } + for (int i = 0; i < ITER; ++i) { + for (int j = 0; j < 8; ++j) a[j] = fma(float(q[j]), x[j], a[j]); + q[0] += half(1.0f); + } + float s = 0; + for (int j = 0; j < 8; ++j) s += a[j]; + out[gid] = s; +} +kernel void a_fma32x16(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + float a[16]; + for (int j = 0; j < 16; ++j) a[j] = gid + j; + for (int i = 0; i < ITER / 2; ++i) + for (int j = 0; j < 16; ++j) a[j] = fma(a[j], k, 0.5f); + float s = 0; + for (int j = 0; j < 16; ++j) s += a[j]; + out[gid] = s; +} + +// 8-byte gathers from a 4 KB table in device memory (L1-resident) and in the constant address space +kernel void a_devload(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], device const uint2* tab [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + uint u = gid * 2654435761u; + uint2 s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += tab[(u >> 3) & 511]; s1 += tab[(u >> 7) & 511]; s2 += tab[(u >> 11) & 511]; s3 += tab[(u >> 15) & 511]; + u += k; + } + out[gid] = float(s0.x + s1.x + s2.x + s3.x + s0.y + s1.y + s2.y + s3.y); +} +kernel void a_cstload(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], constant uint2* tab [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + uint u = gid * 2654435761u; + uint2 s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += tab[(u >> 3) & 511]; s1 += tab[(u >> 7) & 511]; s2 += tab[(u >> 11) & 511]; s3 += tab[(u >> 15) & 511]; + u += k; + } + out[gid] = float(s0.x + s1.x + s2.x + s3.x + s0.y + s1.y + s2.y + s3.y); +} +// 4-byte gathers from a 1 KB threadgroup table (IQ3_XXS's grid) +kernel void a_tgload4(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint tab[256]; + for (int i = tid; i < 256; i += 256) tab[i] = i * 3; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint u = gid * 2654435761u; + uint s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += tab[(u >> 3) & 255]; s1 += tab[(u >> 7) & 255]; s2 += tab[(u >> 11) & 255]; s3 += tab[(u >> 15) & 255]; + u += k; + } + out[gid] = float(s0 + s1 + s2 + s3); +} +// popcount parity (computed ksigns) +kernel void a_popc(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint u = gid, s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += popcount(u & 127); s1 += popcount((u >> 7) & 127); s2 += popcount((u >> 14) & 127); s3 += popcount((u >> 21) & 127); + u += k; + } + out[gid] = float(s0 + s1 + s2 + s3); +} diff --git a/results/round39/astra_r39_qsa/baseline-kernels/dense_mm.metal b/results/round39/astra_r39_qsa/baseline-kernels/dense_mm.metal new file mode 100644 index 0000000..ccf06f7 --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/dense_mm.metal @@ -0,0 +1,101 @@ +// metal/kernels/dense_mm.metal - dense projections for a PROMPT chunk (after quant/units/mma.metal): +// act_f16 X[t][k] -> half (the matrix units' operand) +// dq_f16 a quantized matrix -> half rows (ref_quant's values rounded to fp16), once per chunk +// gemm_f16 Y[t][n] = Xh[t] . Wh[n] (fp32) by MetalPerformancePrimitives matmul2d, TM x TN tiles, K dynamic; +// rows of Y past n_out / t past n_tok are not written (the tensors carry the extents) +#include +using namespace mpp::tensor_ops; + +kernel void act_f16(device const float* x [[buffer(0)]], device half* xh [[buffer(1)]], constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + const uint n4 = n >> 2; + if (gid < n4) { + ((device half4*)xh)[gid] = half4(((device const float4*)x)[gid]); + } else if (gid == n4) { + for (uint i = n4 << 2; i < n; ++i) xh[i] = half(x[i]); + } +} + +struct DqArgs { + uint n_in; // values per row + uint n_out; // rows + uint row_bytes; + uint perm; // 0: row r -> r; p > 0: row r = c * (n_out / p) + d -> d * p + c (the HC up matrix with the p + // streams of an output interleaved, for gemm_hc_up's epilogue) +}; + +// one thread per unit: grid ((units per row + nth.x - 1) / nth.x, n_out), nth.x threads +template +kernel void dq_f16(device const uchar* W [[buffer(0)]], device half* Wh [[buffer(1)]], constant DqArgs& a [[buffer(2)]], + uint3 gid [[thread_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort3 nth [[threads_per_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, H = UNIT / 2; + threadgroup uint tab[fmt::TAB]; + if (F == T_IQ4_XS || F == T_IQ4_NL) { + if (tid < 64) { + const uint r0 = tid & 15u, sh = (tid >> 4) * 8u; + const half h0 = half(as_type(uchar(((r0 & 8u) ? 0x7159453526190d01ul : 0xf6eaddcfbfad9881ul) >> ((r0 & 7u) * 8u)))); + tab[tid] = as_type(half2(h0, half(as_type(uchar(0xbfad9881u >> sh))))); + tab[tid + 64] = as_type(half2(h0, half(as_type(uchar(0xf6eaddcfu >> sh))))); + tab[tid + 128] = as_type(half2(h0, half(as_type(uchar(0x26190d01u >> sh))))); + tab[tid + 192] = as_type(half2(h0, half(as_type(uchar(0x71594535u >> sh))))); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } else if (fmt::TAB > 1) { + load_tab(tab, tid, nth.x); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int u = int(gid.x); + if (u * UNIT >= int(a.n_in)) return; + uint r = gid.y, ro = r; + if (a.perm) { + if (fmt::BB * 16 < fmt::QK * 17) { + r = a.perm == 4 ? (ro & 3u) * (a.n_out >> 2) + (ro >> 2) : (ro % a.perm) * (a.n_out / a.perm) + ro / a.perm; + } else { + const uint q = a.n_out / a.perm; + ro = (r % q) * a.perm + r / q; + } + } + float w[UNIT]; + int p0, p1; + DW::run(W + ulong(r) * a.row_bytes, u, 0, w, p0, p1, tab); + device half* o = Wh + ulong(ro) * a.n_in; +#pragma clang loop unroll(full) + for (int j = 0; j < H; j += 4) { + *(device half4*)(o + p0 + j) = half4(half(w[j]), half(w[j + 1]), half(w[j + 2]), half(w[j + 3])); + *(device half4*)(o + p1 + j) = half4(half(w[H + j]), half(w[H + j + 1]), half(w[H + j + 2]), half(w[H + j + 3])); + } +} + +struct GemmArgs { int M, N, K; int pad; }; // Y is M rows of N floats + +template +kernel void gemm_f16(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device float* Y [[buffer(2)]], + constant GemmArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]]) { + tensor, tensor_inline> tA(A, dextents(a.K, a.M)); + tensor, tensor_inline> tB(B, dextents(a.K, a.N)); + tensor, tensor_inline> tY(Y, dextents(a.N, a.M)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, int(tg.y) * TM); + auto mB = tB.slice(0, int(tg.x) * TN); + auto mY = tY.slice(int(tg.x) * TN, int(tg.y) * TM); + op.run(mA, mB, mY); +} + +#define DQ(F) template [[host_name("dq_f16_" #F)]] kernel decltype(dq_f16) dq_f16; +#define DQ_(F) DQ(F) +DQ_(T_F32) DQ_(T_BF16) DQ_(T_Q8_0) DQ_(T_Q4_0) DQ_(T_Q5_0) DQ_(T_Q5_1) DQ_(T_IQ4_NL) DQ_(T_IQ4_XS) DQ_(T_Q2_0) DQ_(T_IQ3_S) DQ_(T_Q3_K) +DQ_(T_Q4_K) DQ_(T_Q5_K) DQ_(T_Q6_K) +template [[host_name("gemm_f16_64_64_4")]] kernel decltype(gemm_f16<64, 64, 4>) gemm_f16<64, 64, 4>; +template [[host_name("gemm_f16_128_64_4")]] kernel decltype(gemm_f16<128, 64, 4>) gemm_f16<128, 64, 4>; +template [[host_name("gemm_f16_64_128_4")]] kernel decltype(gemm_f16<64, 128, 4>) gemm_f16<64, 128, 4>; +template [[host_name("gemm_f16_64_32_4")]] kernel decltype(gemm_f16<64, 32, 4>) gemm_f16<64, 32, 4>; +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>; +template [[host_name("gemm_f16_16_64_2")]] kernel decltype(gemm_f16<16, 64, 2>) gemm_f16<16, 64, 2>; +template [[host_name("gemm_f16_32_64_2")]] kernel decltype(gemm_f16<32, 64, 2>) gemm_f16<32, 64, 2>; +template [[host_name("gemm_f16_32_128_2")]] kernel decltype(gemm_f16<32, 128, 2>) gemm_f16<32, 128, 2>; +template [[host_name("gemm_f16_64_32_2")]] kernel decltype(gemm_f16<64, 32, 2>) gemm_f16<64, 32, 2>; +template [[host_name("gemm_f16_16_96_2")]] kernel decltype(gemm_f16<16, 96, 2>) gemm_f16<16, 96, 2>; diff --git a/results/round39/astra_r39_qsa/baseline-kernels/gdn.metal b/results/round39/astra_r39_qsa/baseline-kernels/gdn.metal new file mode 100644 index 0000000..c057067 --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/gdn.metal @@ -0,0 +1,197 @@ +// metal/kernels/gdn.metal - one GDN (gated delta-net) layer's recurrence for a window of tokens, fused end to end: +// +// conv (4 taps, history [channels][3] oldest first) -> SiLU, on q (16 heads), k (16), v (48 heads) of 128 +// q, k: L2 norm as native_gdn_l2_norm (rsqrt(sum/128 + eps/128), then x 1/sqrt(128), two roundings) +// gate = exp(softplus(alpha + dt) * a) (softplus = v > 20 ? v : log1p(exp(v))), beta = sigmoid(beta) +// kv_j = sum_i S[i][j] k[i]; delta_j = (v_j - gate * kv_j) * beta; S = gate * S + k[i] delta_j +// o_j = (sum_i S[i][j] q[i]) / sqrt(128); y_j = o_j * rsqrt(mean(o^2) + eps) * gamma_j * sigmoid(z_j) +// +// State layout as Strata's: S[(i * h_v + h) * 128 + j]. Value head h reads q/k head h % h_k. +// One threadgroup per value head, 512 threads = 16 simdgroups: simdgroup s holds columns j = 32 (s % 4) + lane and +// rows i = 32 (s / 4) .. + 32 of the head's 128 x 128 state in registers. GDN_NTOK tokens run in order; with +// `commit` the state and the conv history are written back after the last one (a verify window runs without commit, +// the fold of m accepted tokens is the same kernel with GDN_NTOK = m and commit). The conv history is read from +// `hist` and the committed one written to `hist_out` (ping-pong: the q/k channels are shared by three threadgroups). + +constant int GDN_NTOK [[function_constant(2)]]; + +struct GdnArgs { + uint h_k; // 16 + uint h_v; // 48 + uint channels; // 10240 = (2 h_k + h_v) * 128 + uint commit; // 0: verify (no write); 1: commit after the last token; 0x100 | k (1 <= k <= GDN_NTOK): commit + // the state and the conv history after the first k tokens, in place, while every token's y is + // still computed (a window whose first k rows are kept) - bitwise the commit of a k-token step + float eps; // 1e-6 + uint plane; // 0: y only; else y is the raw plane of prepared activations (xprep layout, n = h_v * 128, + // `plane` floats between planes) and the other planes and the per-16 sums are written too +}; + +constant float gdn_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, + 0.00006103515625f}; +constant float gdn_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + +// The kernel is launched with IB * 128 threads (512 for gdn_step, 1024 for gdn_step8). Without the attribute the compiler +// is free to use as many registers as the 128-row state needs and a pipeline can end up with +// maxTotalThreadsPerThreadgroup below the launch size: an M2 Max gave 448 for the GDN_NTOK >= 2 specializations, and +// Metal then skips the dispatch without an error (#2323). The attribute makes the compiler fit the launch size. +template +[[max_total_threads_per_threadgroup(IB * 128)]] +kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[buffer(1)]], + device const float* qkv [[buffer(2)]], device const float* alpha [[buffer(3)]], + device const float* beta [[buffer(4)]], device const float* z [[buffer(5)]], + device const float* conv_w [[buffer(6)]], device const float* dt [[buffer(7)]], + device const float* ssm_a [[buffer(8)]], device const float* gamma [[buffer(9)]], + device float* y [[buffer(10)]], constant GdnArgs& a [[buffer(11)]], + device float* hist_out [[buffer(12)]], + 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; + threadgroup float qs[S], ks[S], vs[S]; + constexpr int RPT = S / IB; // state rows per thread + // 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 - 1][S]; + threadgroup float part[8], part2[4]; + threadgroup float gb[2]; + 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; + const float scale = 1.0f / sqrt(float(S)); + + // the channel this thread convolves: 0..127 q of head hk, 128..255 k of head hk, 256..383 v of head h + int ch = -1; + if (tid < 128) ch = hk * S + tid; + 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) { + #pragma clang loop unroll(full) + for (int k = 0; k < 3; ++k) hst[k] = hist[ch * 3 + k]; + const float4 cwv = *(device const float4*)(conv_w + ch * 4); + cw[0] = cwv.x; cw[1] = cwv.y; cw[2] = cwv.z; cw[3] = cwv.w; + } + const float gam = ib == 0 ? gamma[j] : 0.0f; + const float dth = tid == 384 ? dt[h] : 0.0f; + const float ssm_ah = tid == 384 ? ssm_a[h] : 0.0f; + float st[RPT]; + const int st_stride = hv * S; + const int st_off = (i0 * hv + h) * S + j; + #pragma clang loop unroll(full) + for (int r = 0; r < RPT; ++r) st[r] = (state + st_off)[r * st_stride]; + + 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 vals[4] = {hst[0], hst[1], hst[2], in}; + float s = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 4; ++k) s = fma(vals[k], cw[k], s); + s = s + 0.0f; + cv = s / (1.0f + exp(-s)); + hst[0] = hst[1]; + hst[1] = hst[2]; + hst[2] = in; + if (t + 1 == commit_at && (tid >= 256 || h < int(a.h_k))) { + #pragma clang loop unroll(full) + for (int k = 0; k < 3; ++k) hist_out[ch * 3 + k] = hst[k]; + } + } + in = in_next; + // raw q/k/v to shared memory + L2 norm partials of q (sg 0..3) and k (sg 4..7); gate and beta on idle lanes + if (tid < 256) { + const float ssq = simd_sum(cv * cv); + if (lane == 0) part[sg] = ssq; + if (tid < 128) qs[tid] = cv; + else ks[tid - 128] = cv; + } else if (tid < 384) { + vs[tid - 256] = cv; + } else if (tid == 384) { + const float v = alpha[t * hv + h] + dth; + // 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_ah); + } else if (tid == 416) { + gb[1] = 1.0f / (1.0f + exp(-beta[t * hv + h])); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const float tot_k = (part[4] + part[5]) + (part[6] + part[7]); + const float rsk = rsqrt(tot_k / float(S) + a.eps / float(S)); + // the simdgroup's RPT rows of k and q one a lane, normalized on the fly and broadcast from the lane that holds them + const int qk_idx = i0 + (lane & (RPT - 1)); + const float kl = (rsk * ks[qk_idx]) * scale; + // kv_j + float kv = 0.0f; + #pragma clang loop unroll(full) + for (int r = 0; r < RPT; ++r) kv = fma(st[r], simd_broadcast(kl, ushort(r)), kv); + red[ib][j] = kv; + const float tot_q = (part[0] + part[1]) + (part[2] + part[3]); + const float rsq = rsqrt(tot_q / float(S) + a.eps / float(S)); + const float ql = (rsq * qs[qk_idx]) * scale; + const float g = gb[0], be = gb[1], vj = vs[j]; + threadgroup_barrier(mem_flags::mem_threadgroup); + kv = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < IB; ++k) kv += red[k][j]; + const float delta = (vj - g * kv) * be; + float o = 0.0f; + #pragma clang loop unroll(full) + for (int r = 0; r < RPT; ++r) { + st[r] = fma(g, st[r], simd_broadcast(kl, ushort(r)) * delta); + o = fma(st[r], simd_broadcast(ql, ushort(r)), o); + } + if (ib > 0) red2[ib - 1][j] = o; + if (t + 1 == commit_at) { + // the state and the conv history after token t (registers: st after this token's update, hst after its + // shift); the state is read only before the loop, so in place is safe. q/k history: the threadgroups + // h < h_k own their q/k head (h == hk); v history: every head its own - to hist_out, NOT in place: a q/k + // channel is read by the three threadgroups h % h_k at their start, and a threadgroup may start after its + // owner has finished (seen on the 10-core M4); the caller swaps the two + device float* st_base = state + st_off; + #pragma clang loop unroll(full) + for (int r = 0; r < RPT; ++r) st_base[r * st_stride] = st[r]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // the readout and the output norm: column j's o (only ib == 0 computes readout and writes y) + if (ib == 0) { + #pragma clang loop unroll(full) + for (int k = 0; k < IB - 1; ++k) o += red2[k][j]; + o = o * scale; + const float osq = simd_sum(o * o); + if (lane == 0) part2[cb] = osq; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ib == 0) { + const float tot = (part2[0] + part2[1]) + (part2[2] + part2[3]); + const float rs = rsqrt(tot / float(S) + a.eps); + const float yv = ((rs * o) * gam) * (1.0f / (1.0f + exp(-zz))); + y[t * hv * S + h * S + j] = yv; + if (a.plane) { + device float* raw = y + t * hv * S + h * S + j; + const int jj = j & 15; + raw[a.plane] = yv * gdn_q2s[jj & 7]; + raw[2 * a.plane] = yv * gdn_p4s[jj & 3]; + raw[3 * a.plane] = (jj & 1) ? yv * 0.00390625f : 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; + #pragma clang loop unroll(full) + 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; + } + } + } +} + +template [[host_name("gdn_step")]] kernel decltype(gdn_step_t<4>) gdn_step_t<4>; +template [[host_name("gdn_step8")]] kernel decltype(gdn_step_t<8>) gdn_step_t<8>; diff --git a/results/round39/astra_r39_qsa/baseline-kernels/gdn_prompt.metal b/results/round39/astra_r39_qsa/baseline-kernels/gdn_prompt.metal new file mode 100644 index 0000000..c2c4020 --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/gdn_prompt.metal @@ -0,0 +1,703 @@ +// metal/kernels/gdn_prompt.metal - a GDN layer's recurrence for a PROMPT chunk of T tokens, in three passes: the work +// that does not touch the state runs parallel over the tokens; only the state update is a scan, and the columns of a +// head's state are independent in it (column j needs k, q, v_j, the gate and beta), so a lane owns one column. +// +// gdn_pre per (token, head): conv (4 taps; history oldest first) -> SiLU; q, k L2 norm * 1/sqrt(128) (the decode +// step's arithmetic); kq = k . q per key head; gate = exp(softplus(alpha + dt) * a), beta = sigmoid(beta); +// the chunk's last 3 conv inputs -> hist_out +// gdn_scan per (value head, 32 columns x CB simdgroups): the state column in registers (128 floats a lane); +// per token a = S^T k, b = S^T q (old state), delta = (v - g a) beta, o = (g b + delta kq) / sqrt(128) +// (= S_new^T q / sqrt(128)), S = g S + k delta^T; q and k staged per TB tokens in threadgroup memory +// gdn_post per (token, head): y = o * rsqrt(mean(o^2) + eps) * gamma * sigmoid(z) -> fp16 (the out GEMM's operand) +// and, with `f32`, fp32 +// +// State layout as Strata's: S[(i * h_v + h) * 128 + j]. Value head h reads q/k head h % h_k. +// The kernels are new entries (nothing of gdn.metal changes). + +struct GdnPArgs { + uint n_tok; + uint h_k; // 16 + uint h_v; // 48 + uint channels; // 10240 + float eps; // 1e-6 + uint f32; // gdn_post: also write fp32 y +}; + +static inline float gp_conv(device const float* qkv, device const float* hist, device const float* cw, int ch, int t, + int channels) { + const float4 w = *(device const float4*)(cw + ch * 4); + float s = 0.0f; +#pragma clang loop unroll(full) + for (int k = 0; k < 4; ++k) { + const int src = t - 3 + k; + const float x = src >= 0 ? qkv[src * channels + ch] : hist[ch * 3 + (src + 3)]; + s = fma(x, w[k], s); + } + s = s + 0.0f; + return s / (1.0f + exp(-s)); +} + +// grid ((h_k + h_v) / 8, T) = (8, T) or (h_k + h_v / 2, T) = (40, T), 256 threads: when grid_x == (h_k + h_v) / 8, +// each simdgroup (32 lanes) owns one key head (x < h_k / 8: q, k and kq in registers) or one value head (else: vc and +// gb), with zero threadgroup memory and zero barriers; grid_x == h_k + h_v / 2 keeps the legacy grid. +kernel void gdn_pre(device const float* qkv [[buffer(0)]], device const float* hist [[buffer(1)]], + device const float* conv_w [[buffer(2)]], device const float* alpha [[buffer(3)]], + device const float* beta [[buffer(4)]], device const float* dt [[buffer(5)]], + device const float* ssm_a [[buffer(6)]], device float* qn [[buffer(7)]], device float* kn [[buffer(8)]], + device float* vc [[buffer(9)]], device float2* gb [[buffer(10)]], device float* kq [[buffer(11)]], + device float* hist_out [[buffer(12)]], constant GdnPArgs& a [[buffer(13)]], + uint3 tg [[threadgroup_position_in_grid]], uint3 tgs [[threadgroups_per_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128; + const int x = int(tg.x), t = int(tg.y), hk_n = int(a.h_k), hv = int(a.h_v), C = int(a.channels), T = int(a.n_tok); + const bool packed = int(tgs.x) * 8 == hk_n + hv; + // alpha and beta as the halves of one [T][2 h_v] GEMM (beta == alpha + h_v), else each [T][h_v] + const int ab = beta == alpha + hv ? 2 * hv : hv; + const int k_tgs = packed ? hk_n / 8 : hk_n; + if (x < k_tgs) { + const int kh = packed ? x * 8 + int(sg) : (sg == 0 ? x : hk_n); + if (kh >= hk_n) return; + float cv_q[4], cv_k[4]; + if (t >= 3) { + device const float* pq = qkv + (t - 3) * C + kh * S + int(lane); + device const float* pk = qkv + (t - 3) * C + (hk_n + kh) * S + int(lane); +#pragma clang loop unroll(full) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; + const float4 wq = *(device const float4*)(conv_w + ch_q * 4); + const float4 wk = *(device const float4*)(conv_w + ch_k * 4); + device const float* pqm = pq + m * 32; + device const float* pkm = pk + m * 32; + float sq = fma(pqm[0], wq.x, 0.0f); + sq = fma(pqm[C], wq.y, sq); + sq = fma(pqm[2 * C], wq.z, sq); + sq = fma(pqm[3 * C], wq.w, sq) + 0.0f; + cv_q[m] = sq / (1.0f + exp(-sq)); + float sk = fma(pkm[0], wk.x, 0.0f); + sk = fma(pkm[C], wk.y, sk); + sk = fma(pkm[2 * C], wk.z, sk); + sk = fma(pkm[3 * C], wk.w, sk) + 0.0f; + cv_k[m] = sk / (1.0f + exp(-sk)); + } + if (t == T - 1) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; + for (int k = 0; k < 3; ++k) { + hist_out[ch_q * 3 + k] = pq[m * 32 + (k + 1) * C]; + hist_out[ch_k * 3 + k] = pk[m * 32 + (k + 1) * C]; + } + } + } else { +#pragma clang loop unroll(full) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; + cv_q[m] = gp_conv(qkv, hist, conv_w, ch_q, t, C); + cv_k[m] = gp_conv(qkv, hist, conv_w, ch_k, t, C); + } + if (t == T - 1) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; + for (int k = 0; k < 3; ++k) { + const int src = T - 3 + k; + hist_out[ch_q * 3 + k] = src >= 0 ? qkv[src * C + ch_q] : hist[ch_q * 3 + (src + 3)]; + hist_out[ch_k * 3 + k] = src >= 0 ? qkv[src * C + ch_k] : hist[ch_k * 3 + (src + 3)]; + } + } + } + const float4 sq = simd_sum(float4(cv_q[0] * cv_q[0], cv_q[1] * cv_q[1], cv_q[2] * cv_q[2], cv_q[3] * cv_q[3])); + const float4 sk = simd_sum(float4(cv_k[0] * cv_k[0], cv_k[1] * cv_k[1], cv_k[2] * cv_k[2], cv_k[3] * cv_k[3])); + const float tot_q = (sq.x + sq.y) + (sq.z + sq.w), tot_k = (sk.x + sk.y) + (sk.z + sk.w); + const float rs_q = rsqrt(tot_q / float(S) + a.eps / float(S)); + const float rs_k = rsqrt(tot_k / float(S) + a.eps / float(S)); + const float scale = 1.0f / sqrt(float(S)); + const int base = (t * hk_n + kh) * S + int(lane); + float4 prod; +#pragma clang loop unroll(full) + for (int m = 0; m < 4; ++m) { + const float nq = (rs_q * cv_q[m]) * scale, nk = (rs_k * cv_k[m]) * scale; + qn[base + m * 32] = nq; + kn[base + m * 32] = nk; + prod[m] = nk * nq; + } + const float4 pv = simd_sum(prod); + if (lane == 0) kq[t * hk_n + kh] = (pv.x + pv.y) + (pv.z + pv.w); + } else if (packed) { + const int h = (x - k_tgs) * 8 + int(sg); + if (t >= 3) { + const int d = int(lane) * 4, ch = (2 * hk_n + h) * S + d; + device const float4* p = (device const float4*)(qkv + (t - 3) * C + ch); + device const float4* w = (device const float4*)(conv_w + ch * 4); + const float4 w0 = w[0], w1 = w[1], w2 = w[2], w3 = w[3]; + float4 cv = fma(p[0], float4(w0.x, w1.x, w2.x, w3.x), float4(0.0f)); + cv = fma(p[C / 4], float4(w0.y, w1.y, w2.y, w3.y), cv); + cv = fma(p[C / 2], float4(w0.z, w1.z, w2.z, w3.z), cv); + cv = fma(p[3 * C / 4], float4(w0.w, w1.w, w2.w, w3.w), cv); + cv = cv + 0.0f; + cv = cv / (1.0f + exp(-cv)); + *((device float4*)(vc + (t * hv + h) * S + d)) = cv; + if (t == T - 1) + for (int k = 0; k < 3; ++k) { + const float4 v = p[(k + 1) * C / 4]; + for (int m = 0; m < 4; ++m) hist_out[(ch + m) * 3 + k] = v[m]; + } + } else { +#pragma clang loop unroll(full) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch = 2 * hk_n * S + h * S + d; + const float cv = gp_conv(qkv, hist, conv_w, ch, t, C); + if (t == T - 1) + for (int k = 0; k < 3; ++k) { + const int src = T - 3 + k; + hist_out[ch * 3 + k] = src >= 0 ? qkv[src * C + ch] : hist[ch * 3 + (src + 3)]; + } + vc[(t * hv + h) * S + d] = cv; + } + } + if (lane == 0) { + const float v = alpha[t * ab + h] + dt[h]; + 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[t * hv + h] = float2(exp(sp * ssm_a[h]), 1.0f / (1.0f + exp(-beta[t * ab + h]))); + } + } else { + const int h = 2 * (x - k_tgs) + (tid >> 7), d = tid & 127; + const int ch = 2 * hk_n * S + h * S + d; + const float cv = gp_conv(qkv, hist, conv_w, ch, t, C); + if (t == T - 1) + for (int k = 0; k < 3; ++k) { + const int src = T - 3 + k; + hist_out[ch * 3 + k] = src >= 0 ? qkv[src * C + ch] : hist[ch * 3 + (src + 3)]; + } + vc[(t * hv + h) * S + d] = cv; + if (d == 0) { + const float v = alpha[t * ab + h] + dt[h]; + 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[t * hv + h] = float2(exp(sp * ssm_a[h]), 1.0f / (1.0f + exp(-beta[t * ab + h]))); + } + } +} + +// grid (h_v * 4 / CB), CB * 32 threads; TB tokens of q and k staged at a time +template +kernel void gdn_scan_t(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]], + device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]], + device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]], + device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]], + 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, NB = 4 / CB; + threadgroup float4 qs[TB][32], ks[TB][32]; + const int h = int(tg.x) / NB, hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k); + const int j = ((int(tg.x) % NB) * CB + sg) * 32 + lane; + const int T = int(a.n_tok); + const float scale = 1.0f / sqrt(float(S)); + float st[S]; +#pragma clang loop unroll(full) + for (int i = 0; i < S; ++i) st[i] = state[(i * hv + h) * S + j]; + for (int t0 = 0; t0 < T; t0 += TB) { + const int nb = min(TB, T - t0); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < nb * 32; idx += CB * 32) { + const int tt = idx >> 5, w = idx & 31; + qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w]; + ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int tt = 0; tt < nb; ++tt) { + const int t = t0 + tt; + const float2 g_b = gb[t * hv + h]; + const float kqv = kq[t * hkn + hk]; + const float v = vc[(t * hv + h) * S + j]; + float ka = 0.0f, kb = 0.0f, qa = 0.0f, qb = 0.0f; +#pragma clang loop unroll(full) + for (int w = 0; w < 32; w += 2) { + const float4 k0 = ks[tt][w], k1 = ks[tt][w + 1], q0 = qs[tt][w], q1 = qs[tt][w + 1]; + ka = fma(st[4 * w], k0.x, ka); ka = fma(st[4 * w + 1], k0.y, ka); + ka = fma(st[4 * w + 2], k0.z, ka); ka = fma(st[4 * w + 3], k0.w, ka); + kb = fma(st[4 * w + 4], k1.x, kb); kb = fma(st[4 * w + 5], k1.y, kb); + kb = fma(st[4 * w + 6], k1.z, kb); kb = fma(st[4 * w + 7], k1.w, kb); + qa = fma(st[4 * w], q0.x, qa); qa = fma(st[4 * w + 1], q0.y, qa); + qa = fma(st[4 * w + 2], q0.z, qa); qa = fma(st[4 * w + 3], q0.w, qa); + qb = fma(st[4 * w + 4], q1.x, qb); qb = fma(st[4 * w + 5], q1.y, qb); + qb = fma(st[4 * w + 6], q1.z, qb); qb = fma(st[4 * w + 7], q1.w, qb); + } + const float g = g_b.x, kv = ka + kb, qv = qa + qb; + const float delta = (v - g * kv) * g_b.y; + o[(t * hv + h) * S + j] = fma(g, qv, delta * kqv) * scale; +#pragma clang loop unroll(full) + for (int w = 0; w < 32; ++w) { + const float4 k4 = ks[tt][w]; + 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); + } + } + } +#pragma clang loop unroll(full) + for (int i = 0; i < S; ++i) state[(i * hv + h) * S + j] = st[i]; +} + +// a probe for the bench (function constant 6): 1 = no per-token exchange (wrong results), 2 = no state update +constant int GDN_PROBE [[function_constant(6)]]; +constant int gdn_probe = is_function_constant_defined(GDN_PROBE) ? GDN_PROBE : 0; + +// the rows of a column split over IB simdgroups (128 / IB rows a lane, fewer registers): grid (h_v * 4), IB * 32 +// threads; lane = column (32 columns a threadgroup); when TB == 0 (single-pass), each simdgroup loads its R rows of q +// and k across its lanes and broadcasts them with simd_broadcast (one initial dot pass at t = 0, one barrier a token, +// only red[2][IB][32] in threadgroup memory); TB > 0 keeps the TB-token threadgroup-staged reference. +template +kernel void gdn_scan_r(device float* state [[buffer(0)]], device const float* qn [[buffer(1)]], + device const float* kn [[buffer(2)]], device const float* vc [[buffer(3)]], + device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]], + device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]], + 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; + threadgroup float2 red[2][IB][32]; + const int g_idx = int(tg.x) / 4, hv = int(a.h_v), hkn = int(a.h_k); + const int hk = (hv == 3 * hkn) ? (g_idx / 3) : (g_idx % hkn); + const int h = (hv == 3 * hkn) ? (hk + (g_idx % 3) * hkn) : g_idx; + const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * R; + const int T = int(a.n_tok); + const float scale = 1.0f / sqrt(float(S)); + const int rs = hv * S, s0 = (i0 * hv + h) * S + j; + float st[R]; +#pragma clang loop unroll(full) + for (int i = 0; i < R; ++i) st[i] = (state + s0)[i * rs]; + if (TB == 0) { + constexpr int C = (R + 31) / 32, R1 = R / C; + const ushort lr = lane & (R1 - 1); + const int ks = hkn * S, k0 = hk * S + i0 + int(lr), v0 = h * S + j; + float kl[C], ql[C]; +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) { + kl[c] = kn[k0 + c * 32]; + ql[c] = qn[k0 + c * 32]; + } + float2 g_b = gb[h]; + float v = vc[v0]; + float ka = 0.0f, qa = 0.0f; +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) +#pragma clang loop unroll(full) + for (int r = 0; r < R1; ++r) { + ka = fma(st[c * 32 + r], simd_broadcast(kl[c], ushort(r)), ka); + qa = fma(st[c * 32 + r], simd_broadcast(ql[c], ushort(r)), qa); + } + for (int t = 0; t < T - 1; ++t) { + const int slot = t & 1, tn = t + 1; + const int kn_off = tn * ks + k0, vn_off = tn * rs + v0; + float kn_l[C], qn_l[C]; +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) { + kn_l[c] = kn[kn_off + c * 32]; + qn_l[c] = qn[kn_off + c * 32]; + } + const float2 g_b_n = gb[tn * hv + h]; + const float v_n = vc[vn_off]; + const float kqv = sg == 0 ? kq[t * hkn + hk] : 0.0f; + float2 s = float2(ka, qa); + if (gdn_probe != 1) { + red[slot][sg][lane] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + s = red[slot][0][lane]; +#pragma clang loop unroll(full) + for (int b = 1; b < IB; ++b) s += red[slot][b][lane]; + } + const float g = g_b.x; + const float delta = (v - g * s.x) * g_b.y; + if (sg == 0) o[t * rs + v0] = fma(g, s.y, delta * kqv) * scale; + ka = 0.0f; + qa = 0.0f; + if (gdn_probe != 2) +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) +#pragma clang loop unroll(full) + for (int r = 0; r < R1; ++r) { + const float kr = simd_broadcast(kl[c], ushort(r)); + const float knr = simd_broadcast(kn_l[c], ushort(r)), qnr = simd_broadcast(qn_l[c], ushort(r)); + st[c * 32 + r] = fma(g, st[c * 32 + r], kr * delta); + ka = fma(st[c * 32 + r], knr, ka); + qa = fma(st[c * 32 + r], qnr, qa); + } +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) kl[c] = kn_l[c]; + g_b = g_b_n; + v = v_n; + } + if (T > 0) { + const int t = T - 1, slot = t & 1; + const float kqv = sg == 0 ? kq[t * hkn + hk] : 0.0f; + float2 s = float2(ka, qa); + if (gdn_probe != 1) { + red[slot][sg][lane] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + s = red[slot][0][lane]; +#pragma clang loop unroll(full) + for (int b = 1; b < IB; ++b) s += red[slot][b][lane]; + } + const float g = g_b.x; + const float delta = (v - g * s.x) * g_b.y; + if (sg == 0) o[t * rs + v0] = fma(g, s.y, delta * kqv) * scale; + if (gdn_probe != 2) +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) +#pragma clang loop unroll(full) + for (int r = 0; r < R1; ++r) { + const float kr = simd_broadcast(kl[c], ushort(r)); + st[c * 32 + r] = fma(g, st[c * 32 + r], kr * delta); + } + } + } else { + constexpr int R4 = R / 4, TBS = TB > 0 ? TB : 1; + constexpr bool KC = R4 <= 8; + threadgroup float4 qs[TBS][32], ks[TBS][32]; + threadgroup float vs[TBS][32]; + threadgroup float4 gk[TBS]; + device const float4* qn4_p = (device const float4*)qn; + device const float4* kn4_p = (device const float4*)kn; + for (int t0 = 0; t0 < T; t0 += TBS) { + const int nb = min(TBS, T - t0); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < nb * 32; idx += IB * 32) { + const int tt = idx >> 5, w = idx & 31; + qs[tt][w] = qn4_p[((t0 + tt) * hkn + hk) * 32 + w]; + ks[tt][w] = kn4_p[((t0 + tt) * hkn + hk) * 32 + w]; + vs[tt][w] = vc[((t0 + tt) * hv + h) * S + (j - lane) + w]; + if (w == 0) { + const float2 g_b = gb[(t0 + tt) * hv + h]; + gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float ka = 0.0f, qa = 0.0f; + float4 kc[R4]; +#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; + const float4 g_b = gk[tt]; + const float kqv = g_b.z; + const float v = vs[tt][lane]; + float2 s = float2(ka, qa); + if (gdn_probe != 1) { + red[slot][sg][lane] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + s = red[slot][0][lane]; +#pragma clang loop unroll(full) + for (int b = 1; b < IB; ++b) s += red[slot][b][lane]; + } + 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 = 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); + } + } + } + } + device float* sp = state + s0; +#pragma clang loop unroll(full) + for (int i = 0; i < R; ++i) sp[i * rs] = st[i]; +} + +// gdn_scan_r<8, 16> with the 16 state rows of a lane in four float4 variables (no array): a register-allocation check +kernel void gdn_scan_v(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]], + device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]], + device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]], + device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]], + 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, IB = 8, TB = 16; + threadgroup float4 qs[TB][32], ks[TB][32]; + threadgroup float2 red[2][IB][32]; + threadgroup float vs[TB][32]; + threadgroup float4 gk[TB]; + const int h = int(tg.x) / 4, hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k); + const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * 16, w0 = sg * 4; + const int T = int(a.n_tok); + const float scale = 1.0f / sqrt(float(S)); + device float* sp = state + (i0 * hv + h) * S + j; + const int rs = hv * S; + float4 s0 = float4(sp[0], sp[rs], sp[2 * rs], sp[3 * rs]); + float4 s1 = float4(sp[4 * rs], sp[5 * rs], sp[6 * rs], sp[7 * rs]); + float4 s2 = float4(sp[8 * rs], sp[9 * rs], sp[10 * rs], sp[11 * rs]); + float4 s3 = float4(sp[12 * rs], sp[13 * rs], sp[14 * rs], sp[15 * rs]); + for (int t0 = 0; t0 < T; t0 += TB) { + const int nb = min(TB, T - t0); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < nb * 32; idx += IB * 32) { + const int tt = idx >> 5, w = idx & 31; + qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w]; + ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w]; + vs[tt][w] = vc[((t0 + tt) * hv + h) * S + (j - lane) + w]; + if (w == 0) { + const float2 g_b = gb[(t0 + tt) * hv + h]; + gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int tt = 0; tt < nb; ++tt) { + const int t = t0 + tt, slot = t & 1; + const float4 k0 = ks[tt][w0], k1 = ks[tt][w0 + 1], k2 = ks[tt][w0 + 2], k3 = ks[tt][w0 + 3]; + const float4 q0 = qs[tt][w0], q1 = qs[tt][w0 + 1], q2 = qs[tt][w0 + 2], q3 = qs[tt][w0 + 3]; + const float ka = (dot(s0, k0) + dot(s1, k1)) + (dot(s2, k2) + dot(s3, k3)); + const float qa = (dot(s0, q0) + dot(s1, q1)) + (dot(s2, q2) + dot(s3, q3)); + red[slot][sg][lane] = float2(ka, qa); + const float4 g_b = gk[tt]; + const float v = vs[tt][lane]; + threadgroup_barrier(mem_flags::mem_threadgroup); + float2 r = red[slot][0][lane]; + for (int b = 1; b < IB; ++b) r += red[slot][b][lane]; + const float g = g_b.x; + const float delta = (v - g * r.x) * g_b.y; + if (sg == 0) o[(t * hv + h) * S + j] = fma(g, r.y, delta * g_b.z) * scale; + s0 = fma(float4(g), s0, k0 * delta); + s1 = fma(float4(g), s1, k1 * delta); + s2 = fma(float4(g), s2, k2 * delta); + s3 = fma(float4(g), s3, k3 * delta); + } + } + sp[0] = s0.x; sp[rs] = s0.y; sp[2 * rs] = s0.z; sp[3 * rs] = s0.w; + sp[4 * rs] = s1.x; sp[5 * rs] = s1.y; sp[6 * rs] = s1.z; sp[7 * rs] = s1.w; + sp[8 * rs] = s2.x; sp[9 * rs] = s2.y; sp[10 * rs] = s2.z; sp[11 * rs] = s2.w; + sp[12 * rs] = s3.x; sp[13 * rs] = s3.y; sp[14 * rs] = s3.z; sp[15 * rs] = s3.w; +} + +// grid (h_v / 4, T) or (h_v, T), 128 threads: one simdgroup (32 lanes) per head h computes the RMSNorm and gated SiLU +// over the 128 channels (4 channels a lane) with zero threadgroup memory and zero barriers. +kernel void gdn_post(device const float* o [[buffer(0)]], device const float* z [[buffer(1)]], + device const float* gamma [[buffer(2)]], device half* yh [[buffer(3)]], device float* y [[buffer(4)]], + constant GdnPArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], uint3 tgs [[threadgroups_per_grid]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const int hv = int(a.h_v), t = int(tg.y); + const int h = int(tgs.x) == hv ? (sg == 0 ? int(tg.x) : hv) : int(tg.x) * 4 + int(sg); + if (h >= hv) return; + const uint base = uint(t * hv + h) * 128 + lane; + device const float* o_row = o + base; + device const float* z_row = z + base; + device const float* g_row = gamma + lane; + device half* yh_row = yh + base; + device float* y_row = y + base; + const float ov0 = o_row[0], ov1 = o_row[32], ov2 = o_row[64], ov3 = o_row[96]; + // the gate and gamma do not depend on the norm: computed while the simd_sums run + float sig[4]; +#pragma clang loop unroll(full) + for (uint k = 0; k < 4; ++k) sig[k] = 1.0f / (1.0f + exp(-z_row[k * 32])); + const float gv[4] = {g_row[0], g_row[32], g_row[64], g_row[96]}; + const float4 pv = simd_sum(float4(ov0 * ov0, ov1 * ov1, ov2 * ov2, ov3 * ov3)); + const float tot = (pv.x + pv.y) + (pv.z + pv.w); + const float rs = rsqrt(tot / 128.0f + a.eps); + const float ov[4] = {ov0, ov1, ov2, ov3}; + float yv[4]; +#pragma clang loop unroll(full) + for (uint k = 0; k < 4; ++k) { + yv[k] = ((rs * ov[k]) * gv[k]) * sig[k]; + yh_row[k * 32] = half(yv[k]); + } + if (a.f32) { +#pragma clang loop unroll(full) + for (uint k = 0; k < 4; ++k) y_row[k * 32] = yv[k]; + } +} + +template [[host_name("gdn_scan_1")]] kernel decltype(gdn_scan_t<1, 16>) gdn_scan_t<1, 16>; +template [[host_name("gdn_scan_2")]] kernel decltype(gdn_scan_t<2, 16>) gdn_scan_t<2, 16>; +template [[host_name("gdn_scan_4")]] kernel decltype(gdn_scan_t<4, 16>) gdn_scan_t<4, 16>; +template [[host_name("gdn_scanr_2")]] kernel decltype(gdn_scan_r<2, 0>) gdn_scan_r<2, 0>; +template [[host_name("gdn_scanr_4")]] kernel decltype(gdn_scan_r<4, 0>) gdn_scan_r<4, 0>; +template [[host_name("gdn_scanr_8")]] kernel decltype(gdn_scan_r<8, 0>) gdn_scan_r<8, 0>; +template [[host_name("gdn_scanr_4_8")]] kernel decltype(gdn_scan_r<4, 8>) gdn_scan_r<4, 8>; +template [[host_name("gdn_scanr_4_16")]] kernel decltype(gdn_scan_r<4, 16>) gdn_scan_r<4, 16>; + +// ---- the chunked form (chunks of GC = 32 tokens), the WY / UT form of the gated delta rule. With G_t the product +// of the chunk's gates up to t (log space: lg_t) and S0 the state entering the chunk: +// (I + A) D = B, A_ts = beta_t (k_t . k_s) G_t / G_s (s < t), B = [beta v | beta G k] -> D = [U | W] +// delta = U - W S0, o_t = G_t S0^T q_t + sum_{s <= t} (q_t . k_s) G_t / G_s delta_s, +// S_L = G_L S0 + sum_s G_L / G_s k_s delta_s^T +// (the sequential recurrence's algebra; the rounding differs). Three kernels: +// gdn_cprep per (head, chunk): lg, A, the forward substitution for U and W (a thread per column of B) +// gdn_cstate per (head, 32 state columns): the chunks in order - delta, the inter-chunk output, the state update +// gdn_cout per (head, chunk): o += the intra-chunk part, times 1/sqrt(128) +constant constexpr int GC = 32; + +// grid (h_v, chunks), 256 threads. U, W: [h][chunk][GC][128]; lgo: [h][chunk][GC] +kernel void gdn_cprep(device const float* qn [[buffer(0)]], device const float* kn [[buffer(1)]], + device const float* vc [[buffer(2)]], device const float2* gb [[buffer(3)]], + device float* U [[buffer(4)]], device float* W [[buffer(5)]], device float* lgo [[buffer(6)]], + constant GdnPArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int S = 128; + threadgroup float A[GC][GC]; + threadgroup float lg[GC], bt[GC]; + const int h = int(tg.x), c = int(tg.y), hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok); + const int NC = (T + GC - 1) / GC, t0 = c * GC, nb = min(GC, T - t0); + if (tid == 0) { + float s = 0.0f; + for (int t = 0; t < GC; ++t) { + if (t < nb) s += max(log(gb[(t0 + t) * hv + h].x), -80.0f); + lg[t] = s; + bt[t] = t < nb ? gb[(t0 + t) * hv + h].y : 0.0f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < GC * GC; idx += 256) { + const int t = idx / GC, s = idx % GC; + float v = 0.0f; + if (s < t && t < nb) { + device const float4* kt = (device const float4*)(kn + ((t0 + t) * hkn + hk) * S); + device const float4* ks = (device const float4*)(kn + ((t0 + s) * hkn + hk) * S); + float d = 0.0f; + for (int i = 0; i < S / 4; ++i) d += dot(kt[i], ks[i]); + v = bt[t] * d * exp(lg[t] - lg[s]); + } + A[t][s] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // forward substitution, column j of B a thread (threads 0..127: beta v; 128..255: beta G k) + const int j = tid & 127; + const bool isw = tid >= 128; + float x[GC]; + for (int t = 0; t < GC; ++t) { + float b = 0.0f; + if (t < nb) b = isw ? bt[t] * exp(lg[t]) * kn[((t0 + t) * hkn + hk) * S + j] : bt[t] * vc[((t0 + t) * hv + h) * S + j]; + for (int s = 0; s < t; ++s) b = fma(-A[t][s], x[s], b); + x[t] = b; + } + device float* o = (isw ? W : U) + (ulong(h) * NC + c) * GC * S + j; + for (int t = 0; t < GC; ++t) o[t * S] = x[t]; + if (tid < GC) lgo[(h * NC + c) * GC + tid] = lg[tid]; +} + +// grid (h_v * 4), 256 threads = 8 row blocks of 16 x 32 columns; the state's 128 x 32 block in registers (16 a thread) +kernel void gdn_cstate(device float* state [[buffer(0)]], device const float* qn [[buffer(1)]], + device const float* kn [[buffer(2)]], device const float* U [[buffer(3)]], + device const float* W [[buffer(4)]], device const float* lgo [[buffer(5)]], + device float* delta [[buffer(6)]], device float* o [[buffer(7)]], constant GdnPArgs& a [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128, RB = 16; + threadgroup float red[8][16][32]; + threadgroup float dl[GC][32]; + const int h = int(tg.x) / 4, hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok); + const int NC = (T + GC - 1) / GC; + const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * RB; + float st[RB]; + for (int r = 0; r < RB; ++r) st[r] = state[((i0 + r) * hv + h) * S + j]; + for (int c = 0; c < NC; ++c) { + const int t0 = c * GC, nb = min(GC, T - t0); + device const float* Uc = U + (ulong(h) * NC + c) * GC * S; + device const float* Wc = W + (ulong(h) * NC + c) * GC * S; + device const float* lg = lgo + (h * NC + c) * GC; + // pass 0: delta = U - W S0; pass 1: o_inter = G_t S0^T q_t; 16 tokens at a time + for (int pass = 0; pass < 2; ++pass) { + for (int hf = 0; hf < 2; ++hf) { + float p[16]; + for (int u = 0; u < 16; ++u) { + const int t = hf * 16 + u; + float s = 0.0f; + if (t < nb) { + device const float* v = pass == 0 ? Wc + t * S + i0 : qn + ((t0 + t) * hkn + hk) * S + i0; + for (int r = 0; r < RB; ++r) s = fma(v[r], st[r], s); + } + p[u] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int u = 0; u < 16; ++u) red[sg][u][lane] = p[u]; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int u = sg * 2; u < sg * 2 + 2; ++u) { + const int t = hf * 16 + u; + float s = 0.0f; + for (int b = 0; b < 8; ++b) s += red[b][u][lane]; + if (t < nb) { + if (pass == 0) { + const float d = Uc[t * S + j] - s; + dl[t][lane] = d; + delta[((ulong(h) * NC + c) * GC + t) * S + j] = d; + } else { + o[((t0 + t) * hv + h) * S + j] = exp(lg[t]) * s; + } + } else if (pass == 0) { + dl[t][lane] = 0.0f; + } + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // S = G_L S0 + sum_s G_L / G_s k_s delta_s^T + const float lgl = lg[nb - 1], gl = exp(lgl); + for (int r = 0; r < RB; ++r) st[r] *= gl; + for (int s = 0; s < nb; ++s) { + const float f = exp(lgl - lg[s]) * dl[s][lane]; + device const float* ks = kn + ((t0 + s) * hkn + hk) * S + i0; + for (int r = 0; r < RB; ++r) st[r] = fma(ks[r], f, st[r]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (int r = 0; r < RB; ++r) state[((i0 + r) * hv + h) * S + j] = st[r]; +} + +// grid (h_v, chunks), 128 threads (column j): o_t = (o_inter_t + sum_{s <= t} (q_t . k_s) G_t / G_s delta_s) / sqrt(128) +kernel void gdn_cout(device const float* qn [[buffer(0)]], device const float* kn [[buffer(1)]], + device const float* delta [[buffer(2)]], device const float* lgo [[buffer(3)]], + device float* o [[buffer(4)]], constant GdnPArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int S = 128; + threadgroup float M[GC][GC]; + const int h = int(tg.x), c = int(tg.y), hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok); + const int NC = (T + GC - 1) / GC, t0 = c * GC, nb = min(GC, T - t0); + device const float* lg = lgo + (h * NC + c) * GC; + for (int idx = tid; idx < GC * GC; idx += 128) { + const int t = idx / GC, s = idx % GC; + float v = 0.0f; + if (s <= t && t < nb) { + device const float4* qt = (device const float4*)(qn + ((t0 + t) * hkn + hk) * S); + device const float4* ks = (device const float4*)(kn + ((t0 + s) * hkn + hk) * S); + float d = 0.0f; + for (int i = 0; i < S / 4; ++i) d += dot(qt[i], ks[i]); + v = d * exp(lg[t] - lg[s]); + } + M[t][s] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int j = tid; + device const float* dc = delta + (ulong(h) * NC + c) * GC * S + j; + float dv[GC]; + for (int s = 0; s < GC; ++s) dv[s] = s < nb ? dc[s * S] : 0.0f; + const float scale = 1.0f / sqrt(float(S)); + for (int t = 0; t < nb; ++t) { + float s = o[((t0 + t) * hv + h) * S + j]; + for (int u = 0; u <= t; ++u) s = fma(M[t][u], dv[u], s); + o[((t0 + t) * hv + h) * S + j] = s * scale; + } +} diff --git a/results/round39/astra_r39_qsa/baseline-kernels/gemv.metal b/results/round39/astra_r39_qsa/baseline-kernels/gemv.metal new file mode 100644 index 0000000..49b7b62 --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/gemv.metal @@ -0,0 +1,500 @@ +// metal/kernels/gemv.metal - y[t][r] = W[r] . x[t] for a window of 1..8 tokens (dense projections, the shared +// expert, the LM head), any format of quant.metal/units.metal. +// +// Shape of a threadgroup: SG simdgroups; KS of them split the row's units (KS = 1: no split), the other SG/KS cover +// different rows. In a simdgroup, LPR lanes share a row and 32/LPR rows run side by side; each lane holds NR rows. +// The weights of a unit are decoded once and applied to every token. Partial sums of a K split are added in a +// fixed order (deterministic). + +struct GemvArgs { + uint n_in; + uint n_out; + uint n_tok; + uint row_bytes; + uint x_stride; // floats between tokens' x rows + uint y_stride; // floats between tokens' y rows +}; + +constant int NTOK [[function_constant(0)]]; // tokens of the window (1..8): the token loops are specialised +// a bench probe (function constant 9): 1 = tokens past the first reuse token 0's x (no load; wrong values), +// 2 = tokens past the first skip the combine (wrong values) +constant int GEMV_PROBE [[function_constant(9)]]; +constant int gemv_probe = is_function_constant_defined(GEMV_PROBE) ? GEMV_PROBE : 0; + +// Wider raw-dot windows and prepared Q2 at NR=1 did not improve on M5. +template +static inline bool gemv_direct_x() { + return F == T_Q8_0 && (NTOK == 1 || (NTOK == 2 && LPR == 16 && NR == 1)); +} + +// Keep unit_combine's four FMA chains without materialising the raw x slice. +template +static inline void gemv_raw_combine(device const float* xp, thread const float (&q)[NR][UNIT], + thread const float4* sc, thread float (&acc)[NR][8], int t) { + constexpr int H = UNIT / 2; + device const float4* x = (device const float4*)xp; + float s[NR][4]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) s[r][0] = s[r][1] = s[r][2] = s[r][3] = 0.0f; +#pragma clang loop unroll(full) + for (int j = 0; j < H; j += 4) { + const float4 a = x[j / 4], b = x[(H + j) / 4]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + s[r][0] = fma(q[r][j], a.x, s[r][0]); + s[r][1] = fma(q[r][j + 1], a.y, s[r][1]); + s[r][2] = fma(q[r][H + j], b.x, s[r][2]); + s[r][3] = fma(q[r][H + j + 1], b.y, s[r][3]); + s[r][0] = fma(q[r][j + 2], a.z, s[r][0]); + s[r][1] = fma(q[r][j + 3], a.w, s[r][1]); + s[r][2] = fma(q[r][H + j + 2], b.z, s[r][2]); + s[r][3] = fma(q[r][H + j + 3], b.w, s[r][3]); + } + } +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + acc[r][t] += sc[r].x * (s[r][0] + s[r][1]) + sc[r].y * (s[r][2] + s[r][3]) - + (sc[r].z * 0.0f + sc[r].w * 0.0f); +} + +static inline float gemv_q2_raw_combine(ushort2 pk, float d, thread const float* xs) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; j += 2) { + const uint2 q = uint2(pk) >> (2 * j); + const uint2 bits = ((q | (q << 14)) & 0x00030003u) | 0x64006400u; + const float4 v = float4(as_type(bits) - 1025.0h); + s0 = fma(v.x, xs[j], s0); + s1 = fma(v.y, xs[j + 1], s1); + s2 = fma(v.z, xs[8 + j], s2); + s3 = fma(v.w, xs[8 + j + 1], s3); + } + return d * (s0 + s1) + d * (s2 + s3) - 0.0f; +} + +// IQ4_XS kept as its 16 code bytes and scale (5 registers a row instead of 36), decoded inside the combine. M5, weights +// from DRAM: 2560 -> 640 at 16,2,4,1 -1.5..-5% (1 token), -2..-3% (2..4); 3..4 tokens -1..-2% on 2560 -> 6144..12288 +// (but 6144 -> 2560 at 16,2,4,1 3 tokens +1.4%); LPR 32 under 4 tokens +0.4..0.7%, so those keep the decoded unit. +// IQ4_NL (18-byte blocks, ushort reads) measured no change and keeps the decoded unit. +template +static inline bool gemv_iq4_packed() { return LPR <= 16 || NTOK >= 4; } + +static inline uint4 gemv_iq4xs_codes(device const uchar* row, int u, thread float& d) { + device const uchar* b = row + (u >> 3) * 136; + d = iq4xs_scale(b, u & 7); + return uint4(*(device const packed_uint4*)(b + 8 + 16 * (u & 7))); +} + +// unit_combine<32> of deq's q, decoding the codes as it goes: the same four FMA chains in the same order +static inline float gemv_iq4_combine(uint4 pk, float d, thread const float* xs, threadgroup const uint* tab) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const uint v = pk[j >> 1] >> (16 * (j & 1)); + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[(v >> 8) & 255]); + s0 = fma(float(w0.x), xs[2 * j], s0); + s1 = fma(float(w1.x), xs[2 * j + 1], s1); + s2 = fma(float(w0.y), xs[16 + 2 * j], s2); + s3 = fma(float(w1.y), xs[16 + 2 * j + 1], s3); + } + return d * (s0 + s1) + d * (s2 + s3) - 0.0f; +} + +template +kernel void gemv(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RS = SG / KS, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + threadgroup float red[KS > 1 ? KS * RS * RG * NR * 8 : 1]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int ks = sg / RS, rs = sg % RS; + const int li = lane % LPR; + const int lrow0 = (rs * RG + lane / LPR) * NR; + const int row0 = int(tg.x) * RS * RG * NR + lrow0; + const int U = int(a.n_in) / UNIT; + const int u0 = ks * U / KS, u1 = (ks + 1) * U / KS; + float acc[NR][8]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) rp[r] = W + (ulong)min(row0 + r, int(a.n_out) - 1) * a.row_bytes; + const int s16n = int(a.n_in) / 16; + if (gemv_probe == 3) { + // rows outer, tokens inner: one row's decoded unit in registers at a time (x re-read per row from L1) + for (int u = u0 + li; u < u1; u += LPR) +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 sc; + Dec::run(rp[r], u, q, sc, tab); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + acc[r][t] += unit_combine(q, sc, xs, xsum); + } + } + } else + for (int u = u0 + li; u < u1; u += LPR) { + if (NTOK == 1 && F == T_Q2_0R) { + ushort2 pk[NR]; + float d[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + device const uchar* b = rp[r] + (u >> 2) * 18; + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + pk[r] = ushort2(p[0], p[1]); + d[r] = rdh(b); + } + float xs[16]; + load16(xp + 16 * u, xs); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][0] += gemv_q2_raw_combine(pk[r], d[r], xs); + } else if (F == T_IQ4_XS && gemv_iq4_packed()) { + uint4 pk[NR]; + float d[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) pk[r] = gemv_iq4xs_codes(rp[r], u, d[r]); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[32]; + load16(xp + t * int(a.x_stride) + 32 * u, xs); + load16(xp + t * int(a.x_stride) + 32 * u + 16, xs + 16); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += gemv_iq4_combine(pk[r], d[r], xs, tab); + } + } else { + float q[NR][UNIT]; + float4 sc[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); + float xs[UNIT], xsum[2]; +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + if (gemv_direct_x()) { + if (gemv_probe != 2 || t == 0) + gemv_raw_combine(xp + (gemv_probe == 1 ? 0 : t) * int(a.x_stride) + u * UNIT, + q, sc, acc, t); + } else { + if (gemv_probe != 1 || t == 0) X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + if (gemv_probe != 2 || t == 0) +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + } + } + constexpr int RPT = RS * RG * NR; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + if (KS == 1) { + if (li == 0 && row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = v; + } else if (li == 0) { + red[(ks * RPT + lrow0 + r) * 8 + t] = v; + } + } + if (KS > 1) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ks == 0 && li == 0) +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + if (row0 + r < int(a.n_out)) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float v = 0.0f; +#pragma clang loop unroll(full) + for (int k = 0; k < KS; ++k) v += red[(k * RPT + lrow0 + r) * 8 + t]; + y[t * int(a.y_stride) + row0 + r] = v; + } + } +} + +#define GV(F, L, R, S, K) template [[host_name("gemv_" #F "_" #L "_" #R "_" #S "_" #K)]] kernel decltype(gemv) gemv; +#define GV_SET(F) GV(F, 32, 1, 4, 1) GV(F, 32, 2, 4, 1) GV(F, 16, 1, 4, 1) GV(F, 16, 2, 4, 1) GV(F, 8, 1, 4, 1) \ + GV(F, 8, 2, 4, 1) GV(F, 4, 2, 4, 1) GV(F, 4, 1, 4, 1) GV(F, 32, 1, 4, 4) GV(F, 32, 2, 4, 2) \ + GV(F, 32, 1, 8, 8) GV(F, 16, 1, 8, 4) GV(F, 32, 1, 8, 2) GV(F, 8, 1, 8, 1) GV(F, 1, 1, 4, 1) GV(F, 1, 2, 4, 1) GV(F, 2, 1, 4, 1) GV(F, 2, 2, 4, 1) GV(F, 1, 1, 2, 1) GV(F, 4, 4, 4, 1) +#ifndef GEMV_TYPES +#define GEMV_TYPES GV_SET(T_F32) GV_SET(T_BF16) GV_SET(T_Q8_0) GV_SET(T_Q4_0) GV_SET(T_Q5_0) GV_SET(T_IQ4_NL) GV_SET(T_IQ4_XS) \ + GV_SET(T_Q2_0) GV_SET(T_IQ3_S) GV_SET(T_Q3_K) GV_SET(T_Q4_K) GV_SET(T_Q5_K) GV_SET(T_Q6_K) +#endif +GEMV_TYPES + +// The shared expert's gate and up in one dispatch: h[t][r] = silu(Wg[r] . x[t]) * (Wu[r] . x[t]), each with its own +// format (and so its own prepared activation kind); h is written raw (its consumer reads raw activations). +// KS = 1; tiles as gemv. Buffers: Wg, xg (Wg's kind), Wu, xu (Wu's kind), s16, h, args (row_bytes = Wg's), rb_up. +template +kernel void gemv_pair(device const uchar* Wg [[buffer(0)]], device const float* xg [[buffer(1)]], + device const uchar* Wu [[buffer(2)]], device const float* xu [[buffer(3)]], + device const float* s16 [[buffer(4)]], device float* h [[buffer(5)]], + constant GemvArgs& a [[buffer(6)]], constant uint& rb_up [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int U1 = fmt::UNIT, U2 = fmt::UNIT, RG = 32 / LPR; + constexpr int TB = fmt::TAB > fmt::TAB ? fmt::TAB : fmt::TAB; + threadgroup uint tab1[fmt::TAB]; + threadgroup uint tab2[fmt::TAB]; + if (fmt::TAB > 1) load_tab(tab1, tid, SG * 32); + if (fmt::TAB > 1) load_tab(tab2, tid, SG * 32); + if (TB > 1) threadgroup_barrier(mem_flags::mem_threadgroup); + const int li = lane % LPR; + const int row = (int(tg.x) * SG + sg) * RG + lane / LPR; + const int rr = min(row, int(a.n_out) - 1); + device const uchar* rg = Wg + ulong(rr) * a.row_bytes; + device const uchar* ru = Wu + ulong(rr) * rb_up; + const int s16n = int(a.n_in) / 16; + float ag[8], au[8]; +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) ag[t] = au[t] = 0.0f; + for (int u = li; u < int(a.n_in) / U1; u += LPR) { + float q[U1]; + float4 sc; + Dec::run(rg, u, q, sc, tab1); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[U1], xsum[2]; + X::load(xg + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + ag[t] += unit_combine(q, sc, xs, xsum); + } + } + for (int u = li; u < int(a.n_in) / U2; u += LPR) { + float q[U2]; + float4 sc; + Dec::run(ru, u, q, sc, tab2); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[U2], xsum[2]; + X::load(xu + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + au[t] += unit_combine(q, sc, xs, xsum); + } + } +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float gv = lane_sum(ag[t]), uv = lane_sum(au[t]); + if (li == 0 && row < int(a.n_out)) h[t * int(a.y_stride) + row] = gv / (1.0f + exp(-gv)) * uv; + } +} + +#define GP_(F1, F2) template [[host_name("gemv_pair_" #F1 "_" #F2)]] kernel decltype(gemv_pair) gemv_pair; +#define GP(F1, F2) GP_(F1, F2) +#define GP_ROW(F1) GP(F1, T_IQ3_S) GP(F1, T_IQ4_XS) GP(F1, T_Q2_0) GP(F1, T_Q3_K) GP(F1, T_Q4_K) GP(F1, T_Q5_K) GP(F1, T_Q6_K) +GP_ROW(T_IQ3_S) GP_ROW(T_IQ4_XS) GP_ROW(T_Q2_0) GP_ROW(T_Q3_K) GP_ROW(T_Q4_K) GP_ROW(T_Q5_K) GP_ROW(T_Q6_K) +// Q2_0 with raw activations (the shared expert's down reading the pair's raw h) +GV_SET(T_Q2_0R) + +// gemv_cat: two matrices of the same format and n_in over the same x in one dispatch - rows [0, n_split) are W's +// (into y), rows [n_split, n_out) are W2's (into y2, from its row 0). Each row's arithmetic is gemv's (KS = 1), so the +// outputs are bitwise the two separate GEMVs'. (GDN's ssm_alpha + ssm_beta: 48 + 48 rows.) +// Buffers: W, xp, s16, y, args (n_out = both), W2, y2, n_split. +template +kernel void gemv_cat(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], + device const uchar* W2 [[buffer(5)]], device float* y2 [[buffer(6)]], constant uint& n_split [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int li = lane % LPR; + const int row0 = int(tg.x) * SG * RG * NR + (sg * RG + lane / LPR) * NR; + const int U = int(a.n_in) / UNIT, sp = int(n_split); + float acc[NR][8]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + const int row = min(row0 + r, int(a.n_out) - 1); + rp[r] = row < sp ? W + (ulong)row * a.row_bytes : W2 + (ulong)(row - sp) * a.row_bytes; + } + const int s16n = int(a.n_in) / 16; + for (int u = li; u < U; u += LPR) { + float q[NR][UNIT]; + float4 sc[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + if (gemv_direct_x()) { + gemv_raw_combine(xp + t * int(a.x_stride) + u * UNIT, q, sc, acc, t); + } else { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + } +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + const int row = row0 + r; + if (li == 0 && row < int(a.n_out)) { + if (row < sp) y[t * int(a.y_stride) + row] = v; + else y2[t * int(a.y_stride) + row - sp] = v; + } + } +} +#define GVC_(F, L, R, S) template [[host_name("gemvc_" #F "_" #L "_" #R "_" #S)]] kernel decltype(gemv_cat) gemv_cat; +#define GVC(F, L, R, S) GVC_(F, L, R, S) +GVC(T_BF16, 32, 1, 4) GVC(T_BF16, 32, 1, 8) GVC(T_BF16, 16, 1, 4) GVC(T_BF16, 32, 2, 4) +GVC(T_F32, 32, 1, 4) GVC(T_F32, 32, 1, 8) GVC(T_F32, 16, 1, 4) GVC(T_F32, 32, 2, 4) +GVC(T_Q8_0, 32, 1, 4) GVC(T_Q8_0, 32, 1, 8) GVC(T_Q8_0, 16, 1, 4) GVC(T_Q8_0, 32, 2, 4) GVC(T_Q8_0, 8, 1, 8) + +// gemv_catk: gemv_cat with gemv's K split (KS simdgroups share a row, partial sums added in gemv's fixed order), so +// the outputs are bitwise gemv_____'s. (The MTP draft step's k + v and shared gate + up: Q8_0 +// 2560 -> 512 + 512 / 640 + 640 at 32,1,4,4 on the M5.) Buffers as gemv_cat. +template +kernel void gemv_catk(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], + device const uchar* W2 [[buffer(5)]], device float* y2 [[buffer(6)]], constant uint& n_split [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RS = SG / KS, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + threadgroup float red[KS * RS * RG * NR * 8]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int ks = sg / RS, rs = sg % RS; + const int li = lane % LPR; + const int lrow0 = (rs * RG + lane / LPR) * NR; + const int row0 = int(tg.x) * RS * RG * NR + lrow0; + const int U = int(a.n_in) / UNIT, sp = int(n_split); + const int u0 = ks * U / KS, u1 = (ks + 1) * U / KS; + float acc[NR][8]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + const int row = min(row0 + r, int(a.n_out) - 1); + rp[r] = row < sp ? W + (ulong)row * a.row_bytes : W2 + (ulong)(row - sp) * a.row_bytes; + } + const int s16n = int(a.n_in) / 16; + for (int u = u0 + li; u < u1; u += LPR) { + float q[NR][UNIT]; + float4 sc[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + if (gemv_direct_x()) { + gemv_raw_combine(xp + t * int(a.x_stride) + u * UNIT, q, sc, acc, t); + } else { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + } + constexpr int RPT = RS * RG * NR; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + if (li == 0) red[(ks * RPT + lrow0 + r) * 8 + t] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ks == 0 && li == 0) +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + const int row = row0 + r; + if (row < int(a.n_out)) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float v = 0.0f; +#pragma clang loop unroll(full) + for (int k = 0; k < KS; ++k) v += red[(k * RPT + lrow0 + r) * 8 + t]; + if (row < sp) y[t * int(a.y_stride) + row] = v; + else y2[t * int(a.y_stride) + row - sp] = v; + } + } +} +#define GVCK_(F, L, R, S, K) template [[host_name("gemvck_" #F "_" #L "_" #R "_" #S "_" #K)]] kernel decltype(gemv_catk) gemv_catk; +#define GVCK(F, L, R, S, K) GVCK_(F, L, R, S, K) +GVCK(T_Q8_0, 32, 1, 4, 4) GVCK(T_Q8_0, 32, 2, 4, 2) GVCK(T_Q8_0, 32, 1, 8, 8) + +// gemv_x: gemv (KS = 1) with the window's prepared x (the format's plane, NTOK rows, and its per-16 sums) staged in +// threadgroup memory once per threadgroup, so the per-unit x reads of every token come from threadgroup memory +// instead of L1 (on the M4 the x traffic through L1 - 4 B a weight a token, shared by the rows in a simdgroup - is +// what a window token costs). Same tile meaning and the same arithmetic as gemv_____1: bitwise. +// Dynamic threadgroup memory 0: 4 * NTOK * (n_in + n_in / 16) bytes. +template +kernel void gemv_x(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], + threadgroup float* xt [[threadgroup(0)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1) load_tab(tab, tid, SG * 32); + const int n = int(a.n_in), s16n = n / 16; + threadgroup float* st = xt + NTOK * n; +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + device const float4* src = (device const float4*)(xp + t * int(a.x_stride)); + threadgroup float4* dst = (threadgroup float4*)(xt + t * n); + for (int i = tid; i < n / 4; i += SG * 32) dst[i] = src[i]; + for (int i = tid; i < s16n; i += SG * 32) st[t * s16n + i] = s16[t * s16n + i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int li = lane % LPR; + const int row0 = int(tg.x) * SG * RG * NR + (sg * RG + lane / LPR) * NR; + const int U = n / UNIT; + float acc[NR][8]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) rp[r] = W + (ulong)min(row0 + r, int(a.n_out) - 1) * a.row_bytes; + for (int u = li; u < U; u += LPR) { + float q[NR][UNIT]; + float4 sc[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[UNIT], xsum[2]; + X::load((threadgroup const float*)(xt + t * n), (threadgroup const float*)(st + t * s16n), u, xs, xsum); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + if (li == 0 && row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = v; + } +} +#define GVX(F, L, R, S) template [[host_name("gemvx_" #F "_" #L "_" #R "_" #S)]] kernel decltype(gemv_x) gemv_x; +#define GVX_SET(F) GVX(F, 32, 1, 4) GVX(F, 32, 2, 4) GVX(F, 16, 1, 4) GVX(F, 16, 2, 4) GVX(F, 8, 1, 4) GVX(F, 8, 2, 4) \ + GVX(F, 4, 2, 4) GVX(F, 16, 2, 8) GVX(F, 8, 2, 8) GVX(F, 16, 4, 4) GVX(F, 32, 2, 8) GVX(F, 2, 2, 4) +GVX_SET(T_BF16) GVX_SET(T_Q8_0) GVX_SET(T_Q4_0) GVX_SET(T_Q5_0) GVX_SET(T_IQ4_NL) GVX_SET(T_IQ4_XS) GVX_SET(T_Q2_0) +GVX_SET(T_IQ3_S) GVX_SET(T_Q3_K) GVX_SET(T_Q4_K) GVX_SET(T_Q5_K) GVX_SET(T_Q6_K) diff --git a/results/round39/astra_r39_qsa/baseline-kernels/glue.metal b/results/round39/astra_r39_qsa/baseline-kernels/glue.metal new file mode 100644 index 0000000..b34bb90 --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/glue.metal @@ -0,0 +1,72 @@ +// metal/kernels/glue.metal - small elementwise steps between the big kernels (after quant/units/moe.metal). +// +// swiglu_prep h = silu(g) * u for the shared expert, written as prepared activations (xprep layout) +// moe_combine_sh bo[t] = sum_k ye[entry(t, k)] + sigmoid(sg[t]) * ysh[t] (k in order) + +struct GlueArgs { + uint n; // vector width + uint n_tok; + uint n_embd; + uint topk; +}; + +constant float glue_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, + 0.00006103515625f}; +constant float glue_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + +// grid (n / 16, n_tok), 16 threads: one group of 16 values +kernel void swiglu_prep(device const float* g [[buffer(0)]], device const float* u [[buffer(1)]], + device float* xp [[buffer(2)]], constant GlueArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort j [[thread_index_in_threadgroup]]) { + const uint t = tg.y, d = tg.x * 16 + j, idx = t * a.n + d, plane = a.n_tok * a.n; + const float gv = g[idx]; + const float h = gv / (1.0f + exp(-gv)) * u[idx]; + device float* raw = xp + idx; + raw[0] = h; + raw[plane] = h * glue_q2s[j & 7]; + raw[2 * plane] = h * glue_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? h * 0.00390625f : h; + float s0 = 0.0f, s1 = 0.0f; +#pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(h, ushort(k)); + s1 += simd_shuffle_down(h, ushort(k + 1)); + } + if (j == 0) xp[4 * plane + t * (a.n / 16) + tg.x] = s0 + s1; +} + +kernel void moe_combine_sh(device const float* ye [[buffer(0)]], device const int* plan [[buffer(1)]], + device const float* ysh [[buffer(2)]], device const float* sg [[buffer(3)]], + device float* bo [[buffer(4)]], constant GlueArgs& a [[buffer(5)]], + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.n_embd, k = gid % a.n_embd; + if (t >= a.n_tok) return; + const int row_base = int(t * a.topk); + device const int* te = plan + P_TOKENT + row_base; + device const float* yep = ye + k; + float s = 0.0f; + if (a.topk == 10) { +#pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += yep[te[j] * int(a.n_embd)]; + } else { + for (uint j = 0; j < a.topk; ++j) s += yep[te[j] * int(a.n_embd)]; + } + const float gs = 1.0f / (1.0f + exp(-sg[t])); + bo[gid] = fma(gs, ysh[gid], s); +} + +// slc_touch: reads n16 16-byte words (a weight range the next kernels will need) so they sit in the SLC; the xor is +// written only if it equals a value it never takes (keeps the loads alive). grid (ceil(n16 / 1024)), 256 threads, +// 4 words a thread. +kernel void slc_touch(device const uint4* src [[buffer(0)]], device uint* sink [[buffer(1)]], constant uint& n16 [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + uint s = 0; + for (uint k = 0; k < 4; ++k) { + const uint j = (gid >> 8) * 1024 + k * 256 + (gid & 255); + if (j < n16) { + const uint4 v = src[j]; + s ^= v.x ^ v.y ^ v.z ^ v.w; + } + } + if (s == 0x9e3779b9u) sink[0] = s; +} diff --git a/results/round39/astra_r39_qsa/baseline-kernels/gr.metal b/results/round39/astra_r39_qsa/baseline-kernels/gr.metal new file mode 100644 index 0000000..c18d2ac --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/gr.metal @@ -0,0 +1,307 @@ +// metal/kernels/gr.metal - the hyper-connection (gated residual) read for a window of tokens, with the previous +// sublayer's write folded in. The math is fused_gr.hpp's (native contract: FP32 activations, BF16 weights): +// +// gr_norm R'[c] = R[c] + bo * 2 sigmoid(inj_prev[c] / hc) (when apply; written back in place) +// xn[c] = R'[c] * w_norm[c] * rsqrt(mean(R'[c]^2) + eps) +// gr_down lo[k] = silu((w_down[k] . xn) / hc) k < hc_lr; inject[c] = w_inject[c] . xn +// gr_up mixed[d] = mean_c xn[c][d] * sigmoid(w_up[c n + d] . lo), written as prepared activations (xprep's +// planes and per-16 sums) for the GEMVs that read it +// +// Shapes: n_embd 2560, hc 4, hc_lr 320. Per token t: R[t] (hc * n), bo[t] (n), inj_prev[t] (hc), xn[t] (hc * n), +// lo[t] (hc_lr), inject[t] (hc), XP planes of n_tok rows. + +constant int GR_NTOK [[function_constant(1)]]; + +// the prepared activations' per-position plane scales (units.metal's xprep): q2 planes 4^-(j%8), p4 planes 256^-(j%4) +constant float gr_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, + 0.00006103515625f}; +constant float gr_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + +struct GrArgs { + uint n_embd; + uint hc; + uint hc_lr; + uint n_tok; + uint apply; + float eps; +}; + +// one threadgroup per (stream c, token t), 256 threads +kernel void gr_norm(device float* R [[buffer(0)]], device const float* bo [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device float* xn [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint c = tg.x, t = tg.y, n = a.n_embd; + device float* r = R + (t * a.hc + c) * n; + const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; + float rv[10]; + float ss = 0.0f; + if (n == 2560) { + device float* rp = r + tid; + if (a.apply) { + device const float* bp = bo + t * 2560u + tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const float v = fma(bp[i * 256], g, rp[i * 256]); + rp[i * 256] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + rv[i] = rp[i * 256]; + ss = fma(rv[i], rv[i], ss); + } + } + } else { + for (uint d = tid; d < n; d += 256) { + float v = r[d]; + if (a.apply) { + v = fma(bo[t * n + d], g, v); + r[d] = v; + } + if (n <= 2560) rv[d / 256] = v; + ss = fma(v, v, ss); + } + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device float* o = xn + (t * a.hc + c) * n; + device const float* w = w_norm + c * n; + if (n == 2560) { + o += tid; + w += tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) o[i * 256] = rv[i] * w[i * 256] * rs; + } else { + for (uint d = tid; d < n; d += 256) + o[d] = (n <= 2560 ? rv[d / 256] : r[d]) * w[d] * rs; + } +} + +// rows: 0 .. hc_lr - 1 = w_down, hc_lr .. hc_lr + hc - 1 = w_inject (both bf16 rows of hc * n_embd). +// SG simdgroups split the row's units KS ways; LPR = 32 (a simdgroup per row); RS = SG / KS simdgroups per +// threadgroup, NR rows each. +// one token's dot with a decoded unit (gr_up): x is the token's slice of lo +template +static inline float gr_dot(device const float* x, thread const float* q, float4 sc) { + float xs[UNIT], xsum[2] = {0.0f, 0.0f}; + #pragma clang loop unroll(full) + for (int k = 0; k < UNIT; k += 16) load16(x + k, xs + k); + return unit_combine(q, sc, xs, xsum); +} + +// a unit's decoded rows for grd_acc, read four codes at a time +template +struct GrdUnit { + float q[NR][fmt::UNIT]; + float4 sc[NR]; + void load(device const uchar* w, uint rb, int u, threadgroup const uint* notab) thread { + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(w + ulong(r) * rb, u, q[r], sc[r], notab); + } + float4 v(int r, int i) const thread { return float4(q[r][4 * i], q[r][4 * i + 1], q[r][4 * i + 2], q[r][4 * i + 3]); } +}; + +// Q8_0 two rows a simdgroup: the codes stay packed (8 registers a row, not 32 floats) and are widened per token. +// Measured 10-16% faster past 6 tokens (where the token loop is rolled), no faster below, and slower for one row. +template <> +struct GrdUnit { + char4 q[2][8]; + float4 sc[2]; + void load(device const uchar* w, uint rb, int u, threadgroup const uint*) thread { + #pragma clang loop unroll(full) + for (int r = 0; r < 2; ++r) { + device const uchar* b = w + ulong(r) * rb + 34 * u; + device const ushort* p = (device const ushort*)(b + 2); + #pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) q[r][k] = as_type(uint(p[2 * k]) | (uint(p[2 * k + 1]) << 16)); + const float d = rdh(b); + sc[r] = float4(d, d, 0.0f, 0.0f); + } + } + float4 v(int r, int i) const thread { return float4(q[r][i]); } +}; + +// NR consecutive rows' K-split partials over xn (units of format F; the rows rb bytes apart), acc[r * 8 + t]. +// unit_combine's sums in its order, with each xn float4 read straight from memory once for all NR rows: xn is +// 15-30x the weights' bytes at 4-8 tokens, so NR = 2 halves the kernel's traffic. GR_NTOK is a function constant, +// which the compiler does not unroll over by itself: the pragma keeps acc in registers instead of the stack. Past 6 +// tokens the unrolled loop measured slower than the rolled one, so it stays rolled there; ROLL keeps it rolled at +// every count (the packed Q8_0 pair at KS = 4: 1-3% faster at 4-6 tokens, where KS = 8 measured 2-6% slower). +template +static inline void grd_acc(device const uchar* w, uint rb, int nin, int ks, int KS, ushort lane, device const float* xn, + thread float* acc, threadgroup const uint* notab) { + constexpr int UNIT = fmt::UNIT, H = UNIT / 2; + const int U = nin / UNIT, u0 = ks * U / KS, u1 = (ks + 1) * U / KS; + for (int u = u0 + lane; u < u1; u += 32) { + GrdUnit d; + d.load(w, rb, u, notab); + auto tok = [&](int t) { + device const float4* x = (device const float4*)(xn + t * nin + UNIT * u); + float s[NR][4]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) s[r][0] = s[r][1] = s[r][2] = s[r][3] = 0.0f; + #pragma clang loop unroll(full) + for (int j = 0; j < H; j += 4) { + const float4 a = x[j / 4], b = x[(H + j) / 4]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + const float4 qa = d.v(r, j / 4), qb = d.v(r, (H + j) / 4); + s[r][0] = fma(qa.x, a.x, s[r][0]); + s[r][1] = fma(qa.y, a.y, s[r][1]); + s[r][2] = fma(qb.x, b.x, s[r][2]); + s[r][3] = fma(qb.y, b.y, s[r][3]); + s[r][0] = fma(qa.z, a.z, s[r][0]); + s[r][1] = fma(qa.w, a.w, s[r][1]); + s[r][2] = fma(qb.z, b.z, s[r][2]); + s[r][3] = fma(qb.w, b.w, s[r][3]); + } + } + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + acc[r * 8 + t] += d.sc[r].x * (s[r][0] + s[r][1]) + d.sc[r].y * (s[r][2] + s[r][3]) - + (d.sc[r].z * 0.0f + d.sc[r].w * 0.0f); + }; + if (GR_NTOK <= 6 && !ROLL) { + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) tok(t); + } else { + for (int t = 0; t < GR_NTOK; ++t) tok(t); + } + } +} + +// FI: the inject rows' format (the Unsloth files: Q8_0 down, F32 inject); FI = F is the original kernel. NR must +// divide hc_lr and hc, so that a simdgroup's rows are all w_down or all w_inject. The sums do not depend on SG or NR: +// every SG / NR with the same KS is bitwise the same kernel. +template +kernel void gr_down_f(device const uchar* w_down [[buffer(0)]], device const uchar* w_inject [[buffer(1)]], + device const float* xn [[buffer(2)]], device float* lo [[buffer(3)]], + device float* inject [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + constexpr int RS = SG / KS; + threadgroup float red[KS * RS * NR * 8]; + threadgroup uint notab[1]; + const int ks = sg / RS, rs = sg % RS; + const int row0 = (int(tg.x) * RS + rs) * NR; + const int nrows = int(a.hc_lr + a.hc); + const int rr = min(row0, nrows - NR); + const int nin = int(a.hc * a.n_embd); + const uint rb = uint(nin / fmt::QK * fmt::BB), rbi = uint(nin / fmt::QK * fmt::BB); + float acc[NR * 8]; + #pragma clang loop unroll(full) + for (int t = 0; t < NR * 8; ++t) acc[t] = 0.0f; + constexpr bool roll = F == T_Q8_0 && NR == 2 && KS == 4; + if (rr < int(a.hc_lr)) grd_acc(w_down + ulong(rr) * rb, rb, nin, ks, KS, lane, xn, acc, notab); + else grd_acc(w_inject + ulong(rr - int(a.hc_lr)) * rbi, rbi, nin, ks, KS, lane, xn, acc, notab); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) { + const float v = simd_sum(acc[r * 8 + t]); + if (lane == 0) red[((ks * RS + rs) * NR + r) * 8 + t] = v; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ks == 0 && lane < NR * 8) { + const int r = lane / 8, t = lane % 8, row = row0 + r; + if (t < GR_NTOK && row < nrows) { + float v = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < KS; ++k) v += red[((k * RS + rs) * NR + r) * 8 + t]; + if (row < int(a.hc_lr)) { + const float z = v / float(a.hc); + lo[t * int(a.hc_lr) + row] = z / (1.0f + exp(-z)); + } else { + inject[t * int(a.hc) + row - int(a.hc_lr)] = v; + } + } + } +} + +// A threadgroup owns 16 consecutive d for all hc = 4 streams: 8 simdgroups, simdgroup s -> stream s / 2 and +// d0 + 8 (s % 2) .. + 8; LPR = 4 lanes per row (320 = 20 units of 16, 5 per lane). +// XP: [4 kinds][n_tok][n] then S16 [n_tok][n / 16] (units.metal's xprep layout). +template +kernel void gr_up_f(device const uchar* w_up [[buffer(0)]], device const float* lo [[buffer(1)]], + device const float* xn [[buffer(2)]], device float* xp [[buffer(3)]], + constant GrArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float gate[8][4][16]; + threadgroup uint notab[1]; + const int n = int(a.n_embd), c = sg >> 1, dl = 8 * (sg & 1) + lane / 4, li = lane & 3; + const int d = int(tg.x) * 16 + dl; + constexpr int UNIT = fmt::UNIT; + device const uchar* w = w_up + ulong(c * n + d) * (a.hc_lr / fmt::QK * fmt::BB); + const int U = int(a.hc_lr) / UNIT; + float acc[8]; + #pragma clang loop unroll(full) + for (int t = 0; t < 8; ++t) acc[t] = 0.0f; + for (int u = li; u < U; u += 4) { + float q[UNIT]; + float4 sc; + Dec::run(w, u, q, sc, notab); + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) acc[t] += gr_dot(lo + t * int(a.hc_lr) + UNIT * u, q, sc); + } + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) { + const float v = lane_sum<4>(acc[t]); + if (li == 0) gate[t][c][dl] = 1.0f / (1.0f + exp(-v)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // 16 threads per token: thread j -> d0 + j + const int t = tid / 16, j = tid & 15; + if (t < GR_NTOK) { + const int dd = int(tg.x) * 16 + j; + const int hc = int(a.hc); + float m = 0.0f; + if (hc == 4) { + #pragma clang loop unroll(full) + for (int cc = 0; cc < 4; ++cc) m = fma(xn[(t * 4 + cc) * n + dd], gate[t][cc][j], m); + } else { + for (int cc = 0; cc < hc; ++cc) m = fma(xn[(t * hc + cc) * n + dd], gate[t][cc][j], m); + } + m = m / float(hc); + const uint plane = a.n_tok * uint(n); + device float* raw = xp + t * n + dd; + raw[0] = m; + raw[plane] = m * gr_q2s[j & 7]; + raw[2 * plane] = m * gr_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? m * 0.00390625f : m; + float s0 = 0.0f, s1 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(m, ushort(k)); + s1 += simd_shuffle_down(m, ushort(k + 1)); + } + if (j == 0) xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1; + } +} + +#define GRD1(S, K, R, SFX) \ + template [[host_name("gr_down_" #S "_" #K SFX)]] kernel decltype(gr_down_f) gr_down_f; \ + template [[host_name("gr_down_q8_" #S "_" #K SFX)]] kernel decltype(gr_down_f) gr_down_f; \ + template [[host_name("gr_down_q8f_" #S "_" #K SFX)]] kernel decltype(gr_down_f) gr_down_f; +#define GRD(S, K) GRD1(S, K, 1, "") +GRD(8, 8) GRD(8, 4) GRD(4, 4) GRD(4, 2) GRD(16, 16) GRD(16, 8) +// two rows a simdgroup (gr_read past 4 tokens) +GRD1(8, 8, 2, "_r2") GRD1(4, 4, 2, "_r2") +template [[host_name("gr_up")]] kernel decltype(gr_up_f) gr_up_f; +template [[host_name("gr_up_q8")]] kernel decltype(gr_up_f) gr_up_f; + +template [[host_name("gr_down_q4_8_8")]] kernel decltype(gr_down_f<8, 8, 1, T_Q4_0, T_Q8_0>) gr_down_f<8, 8, 1, T_Q4_0, T_Q8_0>; +template [[host_name("gr_up_q4")]] kernel decltype(gr_up_f) gr_up_f; +template [[host_name("gr_down_q5_8_8")]] kernel decltype(gr_down_f<8, 8, 1, T_Q5_0, T_Q8_0>) gr_down_f<8, 8, 1, T_Q5_0, T_Q8_0>; +template [[host_name("gr_up_q5")]] kernel decltype(gr_up_f) gr_up_f; diff --git a/results/round39/astra_r39_qsa/baseline-kernels/gr2.metal b/results/round39/astra_r39_qsa/baseline-kernels/gr2.metal new file mode 100644 index 0000000..58342f0 --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/gr2.metal @@ -0,0 +1,484 @@ +// metal/kernels/gr2.metal - the hyper-connection read in TWO dispatches (after quant/units/gr.metal), for BF16 HC +// weights (the exact path) or a Q8_0 copy of them (the STRATA_HC_Q8 option): +// +// gr_down3 grid (row blocks, hc streams). The RMSNorm is per stream, so a threadgroup owns one stream c: it forms +// R'_c = R_c + bo * 2 sigmoid(inj_prev_c / hc) (10 KB a token), its sum of squares and rs_c itself, +// stages xn_c = (R'_c * w_norm_c) * rs_c in threadgroup memory once, and every simdgroup dots its rows' +// stream-c quarter with it: D[t][row][c]. Row block 0 also writes R'_c to R_out (ping-pong, never in +// place) and rs_c. Rows 0..hc_lr-1 are w_down, hc_lr.. the inject rows. +// gr_up3 lo[k] = silu(sum_c D[t][k][c] / hc) formed in threadgroup memory, then as gr_up with xn = R' * w_norm +// * rs on the fly, writing the prepared activations. +// inj_prev of the NEXT read = sum_c D[t][hc_lr + c'][c] of this one (gr_down3 sums it, in c order). +// +// Per element the arithmetic is gr_norm's (fma(bo, g, R), then (R' * w) * rs); the K order differs (per stream, then +// the 4 streams in order). + +#ifndef P_TOKENT +#define P_TOKENT 532 // moe.metal's plan layout +#endif + +struct GrArgs3 { + uint n_embd; + uint hc; + uint hc_lr; + uint n_tok; // the prepared planes' row count (the plane stride) + uint apply; + float eps; + uint rb_down; // bytes per row of w_down / w_inject + uint rb_up; // bytes per row of w_up + uint comb; // 1: bo is not given; it is the MoE combine (moe_combine_sh's arithmetic) of buffers 10..13 + uint topk; +}; + +// bo[t][d], given or the MoE combine: sum_k ye[entry(t, k)][d] (k in order) + sigmoid(sg[t]) * ysh[t][d] +static inline float bo_at(device const float* bo, device const int* plan, device const float* ye, + device const float* ysh, device const float* sg, constant GrArgs3& a, int t, int d) { + const int n = int(a.n_embd); + if (!a.comb) return bo[t * n + d]; + float s = 0.0f; + for (uint j = 0; j < a.topk; ++j) s += ye[plan[P_TOKENT + t * int(a.topk) + j] * n + d]; + const float gs = 1.0f / (1.0f + exp(-sg[t])); + return fma(gs, ysh[t * n + d], s); +} + +template +kernel void gr_down3(device const uchar* w_down [[buffer(0)]], device const uchar* w_inject [[buffer(1)]], + device const float* Rin [[buffer(2)]], device float* Rout [[buffer(3)]], + device const float* bo [[buffer(4)]], device const float* Dprev [[buffer(5)]], + device const float* wn [[buffer(6)]], device float* D [[buffer(7)]], + device float* rs_out [[buffer(8)]], constant GrArgs3& a [[buffer(9)]], + device const int* plan [[buffer(10)]], device const float* ye [[buffer(11)]], + device const float* ysh [[buffer(12)]], device const float* sgv [[buffer(13)]], + 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 UNIT = fmt::UNIT; + threadgroup float xs_tg[TT][2560]; + threadgroup float part[8][TT]; + threadgroup float g_tg[TT], rs_tg[TT]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, 256); + const int n = int(a.n_embd), hc = int(a.hc), NIN = hc * n, c = int(tg.y); + const int nrows = int(a.hc_lr) + hc; + for (int t0 = 0; t0 < GR_NTOK; t0 += TT) { + const int T = min(TT, GR_NTOK - t0); + if (tid < T) { + float g = 0.0f; + if (a.apply) { + device const float* dp = Dprev + ((t0 + tid) * nrows + int(a.hc_lr) + c) * hc; + const float ip = ((dp[0] + dp[1]) + dp[2]) + dp[3]; + g = 2.0f / (1.0f + exp(-ip / float(hc))); + } + g_tg[tid] = g; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int t = 0; t < T; ++t) { + device const float* rt = Rin + (t0 + t) * NIN + c * n; + float ss = 0.0f; + for (int d = tid; d < n; d += 256) { + float v = rt[d]; + if (a.apply) v = fma(bo_at(bo, plan, ye, ysh, sgv, a, t0 + t, d), g_tg[t], v); + xs_tg[t][d] = v; + ss = fma(v, v, ss); + if (tg.x == 0) Rout[(t0 + t) * NIN + c * n + d] = v; + } + ss = simd_sum(ss); + if (lane == 0) part[sg][t] = ss; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < T) { + float tot = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) tot += part[k][tid]; + const float rsv = rsqrt(tot / float(n) + a.eps); + rs_tg[tid] = rsv; + if (tg.x == 0) rs_out[(t0 + tid) * hc + c] = rsv; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int t = 0; t < T; ++t) + for (int d = tid; d < n; d += 256) xs_tg[t][d] = xs_tg[t][d] * wn[c * n + d] * rs_tg[t]; + threadgroup_barrier(mem_flags::mem_threadgroup); + // the rows: simdgroup sg holds RPS rows; lanes split the stream's units + int rows[RPS]; + device const uchar* wr[RPS]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + rows[k] = (int(tg.x) * 8 + sg) * RPS + k; + const int rr = min(rows[k], nrows - 1); + wr[k] = rr < int(a.hc_lr) ? w_down + ulong(rr) * a.rb_down : w_inject + ulong(rr - int(a.hc_lr)) * a.rb_down; + } + float acc[RPS][TT]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) acc[k][t] = 0.0f; + } + const int U = n / UNIT, u0 = c * U; + for (int u = lane; u < U; u += 32) { + float q[RPS][UNIT]; + float4 sc[RPS]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) Dec::run(wr[k], u0 + u, q[k], sc[k], tab); + // TT-bounded with a T guard, so acc[][t] is indexed by a constant (registers, not the stack) + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) { + if (t >= T) break; + float xs[UNIT], xsum[2] = {0.0f, 0.0f}; + #pragma clang loop unroll(full) + for (int e = 0; e < UNIT; ++e) xs[e] = xs_tg[t][u * UNIT + e]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) acc[k][t] += unit_combine(q[k], sc[k], xs, xsum); + } + } + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) { + if (t >= T) break; + const float v = simd_sum(acc[k][t]); + if (lane == 0 && rows[k] < nrows) D[((t0 + t) * nrows + rows[k]) * hc + c] = v; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +// grid (n_embd / 16), 256 threads. +template +kernel void gr_up3(device const uchar* w_up [[buffer(0)]], device const float* D [[buffer(1)]], + device const float* Rn [[buffer(2)]], device float* xp [[buffer(3)]], + constant GrArgs3& a [[buffer(4)]], device const float* wn [[buffer(5)]], + device const float* rs [[buffer(6)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT; + threadgroup float lo_tg[8][320]; + threadgroup float gate[8][4][16]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, 256); + const int n = int(a.n_embd), hc = int(a.hc), lr = int(a.hc_lr), nrows = lr + hc; + for (int e = tid; e < GR_NTOK * lr; e += 256) { + const int t = e / lr, k = e - t * lr; + device const float* dp = D + (t * nrows + k) * hc; + const float z = (((dp[0] + dp[1]) + dp[2]) + dp[3]) / float(hc); + lo_tg[t][k] = z / (1.0f + exp(-z)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int c = sg >> 1, dl = 8 * (sg & 1) + lane / 4, li = lane & 3; + const int d = int(tg.x) * 16 + dl; + device const uchar* w = w_up + ulong(c * n + d) * a.rb_up; + const int U = lr / UNIT; + float acc[8]; + #pragma clang loop unroll(full) + for (int t = 0; t < 8; ++t) acc[t] = 0.0f; + for (int u = li; u < U; u += 4) { + float q[UNIT]; + float4 sc; + Dec::run(w, u, q, sc, tab); + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) { + float xs[UNIT], xsum[2] = {0.0f, 0.0f}; + #pragma clang loop unroll(full) + for (int e = 0; e < UNIT; ++e) xs[e] = lo_tg[t][UNIT * u + e]; + acc[t] += unit_combine(q, sc, xs, xsum); + } + } + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) { + const float v = lane_sum<4>(acc[t]); + if (li == 0) gate[t][c][dl] = 1.0f / (1.0f + exp(-v)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int t = tid / 16, j = tid & 15; + if (t < GR_NTOK) { + const int dd = int(tg.x) * 16 + j; + float m = 0.0f; + if (hc == 4) { + #pragma clang loop unroll(full) + for (int cc = 0; cc < 4; ++cc) { + const float xn = Rn[(t * 4 + cc) * n + dd] * wn[cc * n + dd] * rs[t * 4 + cc]; + m = fma(xn, gate[t][cc][j], m); + } + } else { + for (int cc = 0; cc < hc; ++cc) { + const float xn = Rn[(t * hc + cc) * n + dd] * wn[cc * n + dd] * rs[t * hc + cc]; + m = fma(xn, gate[t][cc][j], m); + } + } + m = m / float(hc); + const uint plane = a.n_tok * uint(n); + device float* raw = xp + t * n + dd; + raw[0] = m; + raw[plane] = m * gr_q2s[j & 7]; + raw[2 * plane] = m * gr_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? m * 0.00390625f : m; + float s0 = 0.0f, s1 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(m, ushort(k)); + s1 += simd_shuffle_down(m, ushort(k + 1)); + } + if (j == 0) xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1; + } +} + +#define GD3(F, T, R) template [[host_name("gr_down3_" #F "_" #T "_" #R)]] kernel decltype(gr_down3) gr_down3; +GD3(30, 1, 1) GD3(30, 2, 1) GD3(30, 3, 1) GD3(30, 1, 2) GD3(30, 2, 2) GD3(30, 3, 2) +GD3(8, 1, 1) GD3(8, 2, 1) GD3(8, 3, 1) GD3(8, 1, 2) GD3(8, 2, 2) GD3(8, 3, 2) +template [[host_name("gr_up3_30")]] kernel decltype(gr_up3) gr_up3; +template [[host_name("gr_up3_8")]] kernel decltype(gr_up3) gr_up3; + +// BF16 only: as gr_down3, but each lane loads its rows' weights (5 units of 16 per row and lane: 2560 = 160 units) +// into registers before the norm phase, so the weight stream is in flight while the threadgroup forms xn. +template +kernel void gr_down3b(device const uchar* w_down [[buffer(0)]], device const uchar* w_inject [[buffer(1)]], + device const float* Rin [[buffer(2)]], device float* Rout [[buffer(3)]], + device const float* bo [[buffer(4)]], device const float* Dprev [[buffer(5)]], + device const float* wn [[buffer(6)]], device float* D [[buffer(7)]], + device float* rs_out [[buffer(8)]], constant GrArgs3& a [[buffer(9)]], + device const int* plan [[buffer(10)]], device const float* ye [[buffer(11)]], + device const float* ysh [[buffer(12)]], device const float* sgv [[buffer(13)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float xs_tg[TT][2560]; + threadgroup float part[8][TT]; + threadgroup float g_tg[TT], rs_tg[TT]; + const int n = int(a.n_embd), hc = int(a.hc), NIN = hc * n, c = int(tg.y); + const int nrows = int(a.hc_lr) + hc; + int rows[RPS]; + uint4 wr[RPS][5][2]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + rows[k] = (int(tg.x) * 8 + sg) * RPS + k; + const int rr = min(rows[k], nrows - 1); + device const uchar* w = rr < int(a.hc_lr) ? w_down + ulong(rr) * a.rb_down : w_inject + ulong(rr - int(a.hc_lr)) * a.rb_down; + device const uint4* p = (device const uint4*)(w + 2 * (c * n)); + #pragma clang loop unroll(full) + for (int i = 0; i < 5; ++i) { + wr[k][i][0] = p[2 * (lane + 32 * i)]; + wr[k][i][1] = p[2 * (lane + 32 * i) + 1]; + } + } + for (int t0 = 0; t0 < GR_NTOK; t0 += TT) { + const int T = min(TT, GR_NTOK - t0); + if (tid < T) { + float g = 0.0f; + if (a.apply) { + device const float* dp = Dprev + ((t0 + tid) * nrows + int(a.hc_lr) + c) * hc; + const float ip = ((dp[0] + dp[1]) + dp[2]) + dp[3]; + g = 2.0f / (1.0f + exp(-ip / float(hc))); + } + g_tg[tid] = g; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int t = 0; t < T; ++t) { + device const float* rt = Rin + (t0 + t) * NIN + c * n; + float ss = 0.0f; + for (int d = tid; d < n; d += 256) { + float v = rt[d]; + if (a.apply) v = fma(bo_at(bo, plan, ye, ysh, sgv, a, t0 + t, d), g_tg[t], v); + xs_tg[t][d] = v; + ss = fma(v, v, ss); + if (tg.x == 0) Rout[(t0 + t) * NIN + c * n + d] = v; + } + ss = simd_sum(ss); + if (lane == 0) part[sg][t] = ss; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < T) { + float tot = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) tot += part[k][tid]; + const float rsv = rsqrt(tot / float(n) + a.eps); + rs_tg[tid] = rsv; + if (tg.x == 0) rs_out[(t0 + tid) * hc + c] = rsv; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int t = 0; t < T; ++t) + for (int d = tid; d < n; d += 256) xs_tg[t][d] = xs_tg[t][d] * wn[c * n + d] * rs_tg[t]; + threadgroup_barrier(mem_flags::mem_threadgroup); + float acc[RPS][TT]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) acc[k][t] = 0.0f; + } + #pragma clang loop unroll(full) + for (int i = 0; i < 5; ++i) { + const int u = lane + 32 * i; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + const uint v[8] = {wr[k][i][0].x, wr[k][i][0].y, wr[k][i][0].z, wr[k][i][0].w, + wr[k][i][1].x, wr[k][i][1].y, wr[k][i][1].z, wr[k][i][1].w}; + float q[16]; + #pragma clang loop unroll(full) + for (int e = 0; e < 8; ++e) { + q[2 * e] = as_type(v[e] << 16); + q[2 * e + 1] = as_type(v[e] & 0xFFFF0000u); + } + const float4 sc = float4(1.0f, 1.0f, 0.0f, 0.0f); + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) { + if (t >= T) break; + float xs[16], xsum[2] = {0.0f, 0.0f}; + #pragma clang loop unroll(full) + for (int e = 0; e < 16; ++e) xs[e] = xs_tg[t][u * 16 + e]; + acc[k][t] += unit_combine<16>(q, sc, xs, xsum); + } + } + } + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) { + if (t >= T) break; + const float v = simd_sum(acc[k][t]); + if (lane == 0 && rows[k] < nrows) D[((t0 + t) * nrows + rows[k]) * hc + c] = v; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} +#define GD3B(T, R) template [[host_name("gr_down3b_" #T "_" #R)]] kernel decltype(gr_down3b) gr_down3b; +GD3B(1, 1) GD3B(2, 1) GD3B(3, 1) GD3B(1, 2) GD3B(2, 2) GD3B(3, 2) + +// gr_norm_comb: gr_norm (gr.metal) with the sublayer output bo formed on the fly as the MoE combine - bo[t][d] = +// sum_k ye[entry(t, k)][d] (k in order) + sigmoid(sg[t]) * ysh[t][d], moe_combine_sh's arithmetic - so the combine's +// dispatch goes away; R', xn bitwise gr_norm's after moe_combine_sh. One threadgroup per (stream c, token t). +kernel void gr_norm_comb(device float* R [[buffer(0)]], device const float* ye [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device float* xn [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + device const int* plan [[buffer(6)]], device const float* ysh [[buffer(7)]], + device const float* sgv [[buffer(8)]], constant uint& topk [[buffer(9)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint c = tg.x, t = tg.y, n = a.n_embd; + device float* r = R + (t * a.hc + c) * n; + const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + float rv[10]; + float ss = 0.0f; + if (n == 2560 && topk == 10) { + device const float* yep[10]; + if (a.apply) { + #pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) + yep[j] = ye + uint(plan[P_TOKENT + t * 10 + j]) * 2560u + tid; + } + device const float* yshp = ysh + t * 2560u + tid; + device float* rp = r + tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + float v = rp[i * 256]; + if (a.apply) { + float s = 0.0f; + #pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += yep[j][i * 256]; + v = fma(fma(gs, yshp[i * 256], s), g, v); + rp[i * 256] = v; + } + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid; d < n; d += 256) { + float v = r[d]; + if (a.apply) { + float s = 0.0f; + for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; + const float bo = fma(gs, ysh[t * n + d], s); + v = fma(bo, g, v); + r[d] = v; + } + if (n <= 2560) rv[d / 256] = v; + ss = fma(v, v, ss); + } + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device float* o = xn + (t * a.hc + c) * n; + device const float* w = w_norm + c * n; + if (n == 2560) { + o += tid; + w += tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) o[i * 256] = rv[i] * w[i * 256] * rs; + } else { + for (uint d = tid; d < n; d += 256) + o[d] = (n <= 2560 ? rv[d / 256] : r[d]) * w[d] * rs; + } +} + +// Fixed n_embd = 2560, hc = 4; grid (1, n_tok), 256 threads. Form bo once for all streams. +kernel void gr_norm_comb_4s(device float* R [[buffer(0)]], device const float* ye [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device float* xn [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + device const int* plan [[buffer(6)]], device const float* ysh [[buffer(7)]], + device const float* sgv [[buffer(8)]], constant uint& topk [[buffer(9)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + const uint t = tg.y, n = 2560; + device float* rp = R + t * 4 * n + tid; + device float* op = xn + t * 4 * n + tid; + device const float* wp = w_norm + tid; + float rv[4][10], ss[4] = {0.0f, 0.0f, 0.0f, 0.0f}, g[4]; + #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) + g[c] = a.apply ? 2.0f / (1.0f + exp(-inj[t * 4 + c] / 4.0f)) : 0.0f; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + device const float* yep[10]; + if (a.apply && topk == 10) { + #pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) + yep[j] = ye + uint(plan[P_TOKENT + t * 10 + j]) * 2560u + tid; + } + device const float* yshp = ysh + t * n + tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const uint d = tid + i * 256; + float bo = 0.0f; + if (a.apply) { + float s = 0.0f; + if (topk == 10) { + #pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += yep[j][i * 256]; + } else { + for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; + } + bo = fma(gs, yshp[i * 256], s); + } + #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) { + device float* r = rp + c * n; + float v = r[i * 256]; + if (a.apply) { v = fma(bo, g[c], v); r[i * 256] = v; } + rv[c][i] = v; + ss[c] = fma(v, v, ss[c]); + } + } + #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) { + const float s = simd_sum(ss[c]); + if (lane == 0) part[c][sg] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) { + float tot = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[c][i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device float* o = op + c * n; + device const float* w = wp + c * n; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + o[i * 256] = rv[c][i] * w[i * 256] * rs; + } + } +} diff --git a/results/round39/astra_r39_qsa/baseline-kernels/gridbar.metal b/results/round39/astra_r39_qsa/baseline-kernels/gridbar.metal new file mode 100644 index 0000000..5cd9a0f --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/gridbar.metal @@ -0,0 +1,54 @@ +// metal/kernels/gridbar.metal - feasibility probe for a persistent kernel: phases separated by an in-kernel grid +// barrier (a device atomic counter) instead of dispatch boundaries. Metal does not promise that all threadgroups of +// a grid are resident at once; the probe launches at most one threadgroup per GPU core. +// +// Each phase p: every thread streams its share of chunk p (n uint4), and threadgroup g writes a token value +// (p * 1000 + g) to flag[p][g]; after the barrier, phase p + 1 checks every flag of phase p (written by the OTHER +// threadgroups) and counts mismatches into bad[0]. + +static inline void grid_barrier(device atomic_uint* bar, uint target, ushort tid) { + threadgroup_barrier(mem_flags::mem_device); +#ifdef GB_FENCE + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); +#endif + if (tid == 0) { + atomic_fetch_add_explicit(bar, 1u, memory_order_relaxed); + while (atomic_load_explicit(bar, memory_order_relaxed) < target) {} + } + threadgroup_barrier(mem_flags::mem_device); +#ifdef GB_FENCE + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); +#endif +} + +kernel void k_phases(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + device atomic_uint* bar [[buffer(2)]], device volatile uint* flag [[buffer(3)]], + device atomic_uint* bad [[buffer(4)]], constant uint& n [[buffer(5)]], + constant uint& phases [[buffer(6)]], + uint tgpx [[threadgroup_position_in_grid]], uint ntgx [[threadgroups_per_grid]], + ushort tid [[thread_index_in_threadgroup]], uint gid [[thread_position_in_grid]], + uint nt [[threads_per_grid]]) { + const uint g = tgpx, ntg = ntgx; + uint4 acc = 0; + for (uint p = 0; p < phases; ++p) { + if (p > 0 && tid < ntg) { + const uint v = flag[(p - 1) * ntg + tid]; + if (v != (p - 1) * 1000 + tid) atomic_fetch_add_explicit(bad, 1u, memory_order_relaxed); + } + device const uint4* s = src + ulong(p) * n; + for (uint i = gid; i < n; i += nt) acc += s[i]; + if (tid == 0) flag[p * ntg + g] = p * 1000 + g; + grid_barrier(bar, (p + 1) * ntg, tid); + } + out[gid] = acc.x ^ acc.y ^ acc.z ^ acc.w; +} + +// one phase per dispatch, for comparison +kernel void k_phase1(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + constant uint& n [[buffer(2)]], constant uint& p [[buffer(3)]], + uint gid [[thread_position_in_grid]], uint nt [[threads_per_grid]]) { + uint4 acc = 0; + device const uint4* s = src + ulong(p) * n; + for (uint i = gid; i < n; i += nt) acc += s[i]; + out[gid] = acc.x ^ acc.y ^ acc.z ^ acc.w; +} diff --git a/results/round39/astra_r39_qsa/baseline-kernels/mm_bench.metal b/results/round39/astra_r39_qsa/baseline-kernels/mm_bench.metal new file mode 100644 index 0000000..c04659f --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/mm_bench.metal @@ -0,0 +1,85 @@ +// metal/kernels/mm_bench.metal - matrix throughput probes (mm_bench.cpp): C[M][N] = A[M][K] . B[N][K]^T, fp32 out. +// mpp_mm Metal 4 MetalPerformancePrimitives matmul2d (the M5's neural accelerators), a TM x TN +// tile of C per threadgroup, K dynamic, A and B straight from device memory +// sg_mm simdgroup_matrix 8x8, a 32 x 32 tile per simdgroup (4 x 4 accumulators), 4 simdgroups +#include +using namespace mpp::tensor_ops; + +struct MmArgs { int M, N, K; }; + +template +kernel void mpp_mm(device T* A [[buffer(0)]], device T* B [[buffer(1)]], device float* C [[buffer(2)]], + constant MmArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]]) { + // extents are (inner, outer): A is M rows of K, B is N rows of K, C is M rows of N + tensor, tensor_inline> tA(A, dextents(a.K, a.M)); + tensor, tensor_inline> tB(B, dextents(a.K, a.N)); + tensor, tensor_inline> tC(C, dextents(a.N, a.M)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, int(tg.y) * TM); + auto mB = tB.slice(0, int(tg.x) * TN); + auto mC = tC.slice(int(tg.x) * TN, int(tg.y) * TM); + op.run(mA, mB, mC); +} + +#define MPP(T, M, N, S) template [[host_name("mpp_mm_" #T "_" #M "_" #N "_" #S)]] kernel decltype(mpp_mm) mpp_mm; +MPP(half, 64, 32, 4) MPP(half, 64, 64, 4) MPP(half, 128, 64, 4) MPP(half, 32, 64, 4) MPP(half, 64, 128, 4) +MPP(half, 16, 64, 2) MPP(half, 128, 128, 4) +MPP(half, 64, 16, 2) MPP(half, 64, 16, 4) MPP(half, 128, 16, 4) MPP(half, 128, 32, 4) MPP(half, 64, 32, 2) MPP(half, 32, 16, 2) +MPP(half, 32, 32, 2) MPP(half, 16, 32, 1) MPP(half, 16, 16, 1) MPP(half, 128, 48, 4) MPP(half, 64, 48, 4) +MPP(float, 64, 32, 4) MPP(float, 64, 64, 4) MPP(float, 128, 64, 4) +MPP(bfloat, 64, 64, 4) MPP(bfloat, 128, 64, 4) + +// simdgroup_matrix: threadgroup tile 64 x 64 (4 simdgroups x 32 x 32), K in steps of 8, operands from device memory +template +kernel void sg_mm(device const T* A [[buffer(0)]], device const T* B [[buffer(1)]], device float* C [[buffer(2)]], + constant MmArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + const int m0 = int(tg.y) * 64 + (sg >> 1) * 32, n0 = int(tg.x) * 64 + (sg & 1) * 32; + simdgroup_matrix acc[4][4]; + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) acc[i][j] = make_filled_simdgroup_matrix(0.0f); + for (int k = 0; k < a.K; k += 8) { + simdgroup_matrix ma[4], mb[4]; + for (int i = 0; i < 4; ++i) simdgroup_load(ma[i], A + (m0 + 8 * i) * a.K + k, a.K); + for (int j = 0; j < 4; ++j) simdgroup_load(mb[j], B + (n0 + 8 * j) * a.K + k, a.K, ulong2(0, 0), true); + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) simdgroup_multiply_accumulate(acc[i][j], ma[i], mb[j], acc[i][j]); + } + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) simdgroup_store(acc[i][j], C + (m0 + 8 * i) * a.N + n0 + 8 * j, a.N); +} +template [[host_name("sg_mm_half")]] kernel decltype(sg_mm) sg_mm; +template [[host_name("sg_mm_float")]] kernel decltype(sg_mm) sg_mm; + + +// the element layout of single-simdgroup matmul2d cooperative tensors (right input, left input, destination) for a +// TM x TN x TK descriptor: out[((which * 32 + lane) * 128 + i) * 4 + (0, 1)] = the two indices, [2] = capacity +template +kernel void coop_layout(device int* out [[buffer(0)]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + auto r = op.template get_right_input_cooperative_tensor(); + auto l = op.template get_left_input_cooperative_tensor(); + auto dd = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < dd.get_capacity() && i < 128; ++i) { + auto ix = dd.get_multidimensional_index(i); + device int* o = out + ((2 * 32 + lane) * 128 + i) * 4; + o[0] = ix[0]; o[1] = ix[1]; o[2] = dd.get_capacity(); + } + for (uint16_t i = 0; i < r.get_capacity() && i < 128; ++i) { + auto ix = r.get_multidimensional_index(i); + device int* o = out + ((0 * 32 + lane) * 128 + i) * 4; + o[0] = ix[0]; o[1] = ix[1]; o[2] = r.get_capacity(); + } + for (uint16_t i = 0; i < l.get_capacity() && i < 128; ++i) { + auto ix = l.get_multidimensional_index(i); + device int* o = out + ((1 * 32 + lane) * 128 + i) * 4; + o[0] = ix[0]; o[1] = ix[1]; o[2] = l.get_capacity(); + } +} +template [[host_name("coop_layout_16_32_16")]] kernel decltype(coop_layout<16, 32, 16>) coop_layout<16, 32, 16>; +template [[host_name("coop_layout_32_32_16")]] kernel decltype(coop_layout<32, 32, 16>) coop_layout<32, 32, 16>; +template [[host_name("coop_layout_16_64_32")]] kernel decltype(coop_layout<16, 64, 32>) coop_layout<16, 64, 32>; +template [[host_name("coop_layout_16_32_32")]] kernel decltype(coop_layout<16, 32, 32>) coop_layout<16, 32, 32>; diff --git a/results/round39/astra_r39_qsa/baseline-kernels/mma.metal b/results/round39/astra_r39_qsa/baseline-kernels/mma.metal new file mode 100644 index 0000000..8fcc0ab --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/mma.metal @@ -0,0 +1,279 @@ +// metal/kernels/mma.metal - a GEMV for windows of tokens on the simdgroup matrix units (after quant/units.metal). +// +// y[t][r] = W[r] . x[t] for NTOK (1..8) tokens: each simdgroup owns 8 rows; the K axis runs in chunks of 128. Per +// chunk the simdgroup decodes its 8 x 128 weights - with ref_quant's arithmetic, so a tile element is exactly the +// dequantized weight - into threadgroup memory, then 16 simdgroup_multiply_accumulate steps take A = the 8 x 8 weight +// tile and B = the 8 (k) x 8 (token) slice of x, loaded transposed straight from device memory (x has 8 rows: the +// tokens past the window are padding, and an output column depends only on its own token). +// +// Activations are the raw plane (no pre-scaling); n_in must be a multiple of 128. Every format's units sit inside +// one 128-wide chunk: Q4_K/Q5_K units span 64 positions, Q6_K units 128, the others are contiguous. + +// ---- decoders to actual weights: w[UNIT] and the chunk offsets of its two halves (p0: w[0..H), p1: w[H..UNIT)) ---- +template struct DW { + // generic: w = s * q (quant.metal's deq is ref_quant's order: the scale times the code), contiguous + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT, UPB = fmt::QK / UNIT; + float q[UNIT]; + const float s = deq(row + (u / UPB) * fmt::BB, u % UPB, q, tab); +#pragma clang loop unroll(full) + for (int j = 0; j < UNIT; ++j) w[j] = q[j] * s; + p0 = u * UNIT - k0; + p1 = p0 + UNIT / 2; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint* tab) { + device const uchar* b = row + u * 18; + device const ushort* p = (device const ushort*)(b + 2); +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); + w[2 * j] = float(w0.x); + w[2 * j + 16] = float(w0.y); + w[2 * j + 1] = float(w1.x); + w[2 * j + 17] = float(w1.y); + } + const float d = rdh(b); +#pragma clang loop unroll(full) + for (int j = 0; j < 32; ++j) w[j] *= d; + p0 = 32 * u - k0; + p1 = p0 + 16; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint* tab) { + device const uchar* b = row + (u >> 3) * 136; + const int ui = u & 7; + device const ushort* p = (device const ushort*)(b + 8 + 16 * ui); +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); + w[2 * j] = float(w0.x); + w[2 * j + 16] = float(w0.y); + w[2 * j + 1] = float(w1.x); + w[2 * j + 17] = float(w1.y); + } + const uint sh = *(device const ushort*)(b + 2); + const int ls = int((b[4 + (ui >> 1)] >> (4 * (ui & 1))) & 15) | int(((sh >> (2 * ui)) & 3) << 4); + const float d = rdh(b) * float(ls - 32); +#pragma clang loop unroll(full) + for (int j = 0; j < 32; ++j) w[j] *= d; + p0 = 32 * u - k0; + p1 = p0 + 16; + } +}; +// Q2_0: (code - 1) * d +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 2) * 18; + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + const ushort a = p[0], c = p[1]; + const float d = rdh(b); +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + w[j] = float(int((a >> (2 * j)) & 3) - 1) * d; + w[8 + j] = float(int((c >> (2 * j)) & 3) - 1) * d; + } + p0 = 16 * u - k0; + p1 = p0 + 8; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint* tab) { + device const uchar* b = row + u * 24; + float q[32]; + const float d = deq(b, 0, q, tab), m = rdh(b + 2); +#pragma clang loop unroll(full) + for (int j = 0; j < 32; ++j) w[j] = q[j] * d + m; + p0 = 32 * u - k0; + p1 = p0 + 16; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const float* p = (device const float*)(row + 64 * u); +#pragma clang loop unroll(full) + for (int j = 0; j < 16; ++j) w[j] = p[j]; + p0 = 16 * u - k0; + p1 = p0 + 8; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uint4* p = (device const uint4*)(row + 32 * u); + const uint4 a = p[0], c = p[1]; + const uint v[8] = {a.x, a.y, a.z, a.w, c.x, c.y, c.z, c.w}; +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) { + w[2 * k] = as_type(v[k] << 16); + w[2 * k + 1] = as_type(v[k] & 0xFFFF0000u); + } + p0 = 16 * u - k0; + p1 = p0 + 8; + } +}; +// Q4_K / Q5_K: d1 * code - m1 (ref: d * sc, dmin * m), runs at 256b + 64c + 16m and + 32 +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 144; + const int c = (u >> 1) & 3, m = u & 1; + const float d = rdh(b), dm = rdh(b + 2); + float s1, m1, s2, m2; + scale_min_k4(2 * c, b + 4, s1, m1); + scale_min_k4(2 * c + 1, b + 4, s2, m2); + const float d1 = d * s1, mm1 = dm * m1, d2 = d * s2, mm2 = dm * m2; + device const uchar* q = b + 16 + 32 * c + 16 * m; +#pragma clang loop unroll(full) + for (int j = 0; j < 16; ++j) { + w[j] = d1 * float(q[j] & 15) - mm1; + w[16 + j] = d2 * float(q[j] >> 4) - mm2; + } + p0 = 256 * (u >> 3) + 64 * c + 16 * m - k0; + p1 = p0 + 32; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 176; + const int c = (u >> 1) & 3, m = u & 1; + const float d = rdh(b), dm = rdh(b + 2); + float s1, m1, s2, m2; + scale_min_k4(2 * c, b + 4, s1, m1); + scale_min_k4(2 * c + 1, b + 4, s2, m2); + const float d1 = d * s1, mm1 = dm * m1, d2 = d * s2, mm2 = dm * m2; + device const uchar* ql = b + 48 + 32 * c + 16 * m; + device const uchar* qh = b + 16 + 16 * m; + const uchar u1 = uchar(1 << (2 * c)), u2 = uchar(2 << (2 * c)); +#pragma clang loop unroll(full) + for (int j = 0; j < 16; ++j) { + w[j] = d1 * float((ql[j] & 15) + ((qh[j] & u1) ? 16 : 0)) - mm1; + w[16 + j] = d2 * float((ql[j] >> 4) + ((qh[j] & u2) ? 16 : 0)) - mm2; + } + p0 = 256 * (u >> 3) + 64 * c + 16 * m - k0; + p1 = p0 + 32; + } +}; +// Q6_K: d * sc * (code - 32), runs at 256b + 128n + 32p + 16m and + 64 +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 210; + const int n = (u >> 2) & 1, m = (u >> 1) & 1, p = u & 1; + device const uchar* ql = b + 64 * n + 32 * p + 16 * m; + device const uchar* qh = b + 128 + 32 * n + 16 * m; + device const char* sc = (device const char*)(b + 192); + const float d = rdh(b + 208); + const float ds1 = d * float(sc[8 * n + m + 2 * p]), ds2 = d * float(sc[8 * n + m + 2 * p + 4]); + const uint sh = uint(2 * p), sh4 = sh + 4u; +#pragma clang loop unroll(full) + for (int j = 0; j < 16; j += 4) { + const uint la = rd32(ql + j), ha = rd32(qh + j); + const uchar4 c1 = as_type((la & 0x0F0F0F0Fu) | (((ha >> sh) & 0x03030303u) << 4)); + const uchar4 c2 = as_type(((la >> 4) & 0x0F0F0F0Fu) | (((ha >> sh4) & 0x03030303u) << 4)); + w[j] = ds1 * float(int(c1.x) - 32); + w[j + 1] = ds1 * float(int(c1.y) - 32); + w[j + 2] = ds1 * float(int(c1.z) - 32); + w[j + 3] = ds1 * float(int(c1.w) - 32); + w[16 + j] = ds2 * float(int(c2.x) - 32); + w[16 + j + 1] = ds2 * float(int(c2.y) - 32); + w[16 + j + 2] = ds2 * float(int(c2.z) - 32); + w[16 + j + 3] = ds2 * float(int(c2.w) - 32); + } + p0 = 256 * (u >> 3) + 128 * n + 32 * p + 16 * m - k0; + p1 = p0 + 64; + } +}; +// Q3_K: d_all * (sc - 32) * (code - (hbit ? 0 : 4)), contiguous 32 +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 110; + const int n = (u >> 2) & 1, j = u & 3; + device const uchar* q = b + 32 + 32 * n; + device const uchar* hm = b; + device const uchar* s = b + 96; + const int i1 = 8 * n + 2 * j, i2 = i1 + 1; + const int l1 = (i1 < 8 ? s[i1] & 0xF : s[i1 - 8] >> 4) | (((s[8 + (i1 & 3)] >> (2 * (i1 >> 2))) & 3) << 4); + const int l2 = (i2 < 8 ? s[i2] & 0xF : s[i2 - 8] >> 4) | (((s[8 + (i2 & 3)] >> (2 * (i2 >> 2))) & 3) << 4); + const float d = rdh(b + 108); + const float dl1 = d * float(l1 - 32), dl2 = d * float(l2 - 32); + const uchar hb = uchar(1 << (4 * n + j)); +#pragma clang loop unroll(full) + for (int l = 0; l < 16; ++l) { + w[l] = dl1 * float(int((q[l] >> (2 * j)) & 3) - ((hm[l] & hb) ? 0 : 4)); + w[16 + l] = dl2 * float(int((q[l + 16] >> (2 * j)) & 3) - ((hm[l + 16] & hb) ? 0 : 4)); + } + p0 = 32 * u - k0; + p1 = p0 + 16; + } +}; + +struct MmaArgs { + uint n_in; + uint n_out; + uint row_bytes; + uint x_stride; // floats between the 8 token rows of x + uint y_stride; +}; +constant int MMA_NTOK [[function_constant(4)]]; + +template +kernel void gemv_mma(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]], + device float* y [[buffer(2)]], constant MmaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, UPC = 128 / UNIT, H = UNIT / 2; + threadgroup float wt[SG][8 * 128]; + threadgroup float ct[SG][64]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int row0 = (int(tg.x) * SG + sg) * 8; + simdgroup_float8x8 C = make_filled_simdgroup_matrix(0.0f); + threadgroup float* w8 = wt[sg]; + for (int k0 = 0; k0 < int(a.n_in); k0 += 128) { + for (int it = lane; it < 8 * UPC; it += 32) { + const int r = it / UPC, uu = it - r * UPC; + const int rr = min(row0 + r, int(a.n_out) - 1); + float w[UNIT]; + int p0, p1; + DW::run(W + ulong(rr) * a.row_bytes, k0 / UNIT + uu, k0, w, p0, p1, tab); +#pragma clang loop unroll(full) + for (int j = 0; j < H; ++j) { + w8[r * 128 + p0 + j] = w[j]; + w8[r * 128 + p1 + j] = w[H + j]; + } + } + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int k8 = 0; k8 < 128; k8 += 8) { + simdgroup_float8x8 A, B; + simdgroup_load(A, w8 + k8, 128); + simdgroup_load(B, x + k0 + k8, a.x_stride, ulong2(0, 0), true); + simdgroup_multiply_accumulate(C, A, B, C); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + } + simdgroup_store(C, ct[sg], 8); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int e = lane; e < 64; e += 32) { + const int r = e >> 3, t = e & 7; + if (t < MMA_NTOK && row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = ct[sg][e]; + } +} + +#define MM(F, S) template [[host_name("gemv_mma_" #F "_" #S)]] kernel decltype(gemv_mma) gemv_mma; +#define MM_SET(F) MM(F, 1) MM(F, 2) MM(F, 4) +MM_SET(T_BF16) MM_SET(T_Q8_0) MM_SET(T_Q4_0) MM_SET(T_Q5_0) MM_SET(T_IQ4_NL) MM_SET(T_IQ4_XS) MM_SET(T_Q2_0) +MM_SET(T_IQ3_S) MM_SET(T_Q3_K) MM_SET(T_Q4_K) MM_SET(T_Q5_K) MM_SET(T_Q6_K) diff --git a/results/round39/astra_r39_qsa/baseline-kernels/moe.metal b/results/round39/astra_r39_qsa/baseline-kernels/moe.metal new file mode 100644 index 0000000..878252e --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/moe.metal @@ -0,0 +1,765 @@ +// metal/kernels/moe.metal - routed experts of one layer for a verify window of 1..8 tokens, grouped by expert: +// +// moe_plan one threadgroup: the window's top-k ids -> groups (distinct experts, by id) with their entries +// (token, weight), each (token, k)'s entry index, and the indirect grid sizes of the next two kernels. +// Deterministic (no atomics). +// moe_gu per group, a block of FFN rows: gate and up rows read once, applied to every entry of the group +// (the shared input x is read from cache), h = silu(gate) * up per entry. +// moe_down per group, a block of output rows: ye[entry] = weight * (down . h[entry]). +// moe_combine y[token] = sum over k of ye[entry(token, k)], in k order. +// +// The plan is one int buffer (see P_* below). Expert addressing: MoeArgs (a pack blob or the GGUF tensors): +// [gate rows | up rows | down rows] in GGUF blocks (tools/iq_pack.py's experts.bin layout). + +#define P_NG 0 +#define P_IGU 4 +#define P_IDN 8 +#define P_MAXE 128 +#define P_GEXP 16 +#define P_GSTART (P_GEXP + P_MAXE) +#define P_ETOK (P_GSTART + P_MAXE + 4) +#define P_EW (P_ETOK + P_MAXE) +#define P_TOKENT (P_EW + P_MAXE) +#define P_SPLIT (P_TOKENT + P_MAXE) // function constant 26 only: [0] resident groups, then the indirect grids +#define P_IGU_R (P_SPLIT + 4) // ... of the resident groups (the GPU computes them during the host's turn) +#define P_IDN_R (P_SPLIT + 8) +#define P_IGU_M (P_SPLIT + 12) // ... and of the missing ones (after it) +#define P_IDN_M (P_SPLIT + 16) +#define P_GSLOT (P_SPLIT + 20) +#define P_INTS (P_GSLOT + P_MAXE) + +// Expert e's matrices: gate at W + gate_off + e * gu_stride, up at W + up_off + e * gu_stride, down at +// W + down_off + e * d_stride (a pack's blobs: gate_off 0, strides = the blob; a GGUF: the three tensors). +struct MoeArgs { + ulong gate_off; + ulong up_off; + ulong down_off; + ulong gu_stride; + ulong d_stride; + uint gu_row; // bytes per gate/up row + uint d_row; // bytes per down row + uint n_embd; + uint n_ff; + uint n_tok; + uint topk; + uint n_expert; + uint gx_gu; // threadgroups along the rows for moe_gu (n_ff / rows per threadgroup) + uint gx_dn; // ... for moe_down (n_embd / rows per threadgroup) + uint pad; +}; + +// The SSD-streaming arena (metal/src/engine/stream.cpp): with function constant 20 set, expert e of the layer lies in +// an arena slot, smap[e] = (chunk << 28) | its 16 KiB page in that chunk; the chunks are W (buffer 0) and buffers 7..9 +// (gate_off / up_off / down_off are then offsets inside the slot). The same block is in moe_mm.metal. +#ifndef STRATA_ARENA_METAL +#define STRATA_ARENA_METAL +constant bool kArenaFc [[function_constant(20)]]; +constant bool kArena = is_function_constant_defined(kArenaFc) && kArenaFc; +static inline device const uchar* arena_slot(uint v, device const uchar* W, device const uchar* W1, device const uchar* W2, + device const uchar* W3) { + const uint c = v >> 28; + device const uchar* b = c == 0 ? W : (c == 1 ? W1 : (c == 2 ? W2 : W3)); + return b + (ulong)(v & 0x0FFFFFFFu) * 16384ul; +} +#endif + +// Function constant 26 (the SSD-streaming decode window, with smap): the groups whose expert is in the arena when the +// routing is made come first, the missing ones after them, each part with its own indirect grids (P_SPLIT): the +// resident part runs while the host reads the missing experts. The entries keep their (token, k) mapping, so the +// combine (and every number) is the same; without it the plan is the old one, bit for bit. +constant bool kSplitResFc [[function_constant(26)]]; +constant bool kSplitRes = is_function_constant_defined(kSplitResFc) && kSplitResFc; + +// s_grp, s_ent: 64 each (the upper halves for the missing groups) +static inline void plan_body(device const int* ids, device const float* w, device int* plan, constant MoeArgs& a, + ushort tid, ushort nth, ushort sg, ushort lane, threadgroup int* s_grp, + threadgroup int* s_ent, device const uint* smap = nullptr) { + const int ne = int(a.n_tok * a.topk); + // one thread per expert id (nth == n_expert, at most 1024) + const int e = tid; + int cnt = 0; + if (e < int(a.n_expert)) { + #pragma clang loop unroll_count(10) + for (int i = 0; i < ne; ++i) cnt += ids[i] == e; + } + const int has = cnt > 0; + const bool split = kSplitRes && smap != nullptr; + const int miss = (split && has && smap[e] == 0xFFFFFFFFu) ? 1 : 0; + const int h0 = has - miss, c0 = miss ? 0 : cnt, c1 = miss ? cnt : 0; + const int pg = simd_prefix_exclusive_sum(h0), pe = simd_prefix_exclusive_sum(c0); + const int pg1 = split ? simd_prefix_exclusive_sum(miss) : 0; + const int pe1 = split ? simd_prefix_exclusive_sum(c1) : 0; + if (lane == 31) { + s_grp[sg] = pg + h0; s_ent[sg] = pe + c0; + if (split) { s_grp[32 + sg] = pg1 + miss; s_ent[32 + sg] = pe1 + c1; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + const int nsg = (nth + 31) / 32; + const int g0 = lane < nsg ? s_grp[lane] : 0, n0 = lane < nsg ? s_ent[lane] : 0; + const int og0 = simd_prefix_exclusive_sum(g0), oe0 = simd_prefix_exclusive_sum(n0); + if (lane < nsg) { s_grp[lane] = og0; s_ent[lane] = oe0; } + const int ng0 = simd_broadcast(og0 + g0, 31), ne0 = simd_broadcast(oe0 + n0, 31); + int og = ng0, oe = ne0; + if (split) { + const int g1 = lane < nsg ? s_grp[32 + lane] : 0, n1 = lane < nsg ? s_ent[32 + lane] : 0; + const int og1 = ng0 + simd_prefix_exclusive_sum(g1), oe1 = ne0 + simd_prefix_exclusive_sum(n1); + if (lane < nsg) { s_grp[32 + lane] = og1; s_ent[32 + lane] = oe1; } + og = simd_broadcast(og1 + g1, 31); + oe = simd_broadcast(oe1 + n1, 31); + } + if (lane == 0) { + plan[P_NG] = og; + plan[P_IGU] = int(a.gx_gu); plan[P_IGU + 1] = og; plan[P_IGU + 2] = 1; + plan[P_IDN] = int(a.gx_dn); plan[P_IDN + 1] = og; plan[P_IDN + 2] = 1; + plan[P_GSTART + og] = oe; + if (split) { + plan[P_SPLIT] = ng0; + plan[P_IGU_R] = int(a.gx_gu); plan[P_IGU_R + 1] = ng0; plan[P_IGU_R + 2] = 1; + plan[P_IDN_R] = int(a.gx_dn); plan[P_IDN_R + 1] = ng0; plan[P_IDN_R + 2] = 1; + plan[P_IGU_M] = int(a.gx_gu); plan[P_IGU_M + 1] = og - ng0; plan[P_IGU_M + 2] = 1; + plan[P_IDN_M] = int(a.gx_dn); plan[P_IDN_M + 1] = og - ng0; plan[P_IDN_M + 2] = 1; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (has) { + const int gi = (miss ? s_grp[32 + sg] + pg1 : s_grp[sg] + pg), off = miss ? s_ent[32 + sg] + pe1 : s_ent[sg] + pe; + plan[P_GEXP + gi] = e; + plan[P_GSTART + gi] = off; + if (split && !miss) plan[P_GSLOT + gi] = int(smap[e]); + const bool topk10 = (a.topk == 10u); + int j = 0; + #pragma clang loop unroll_count(10) + for (int i = 0; i < ne; ++i) + if (ids[i] == e) { + plan[P_ETOK + off + j] = topk10 ? ((i * 205) >> 11) : (i / int(a.topk)); + plan[P_EW + off + j] = as_type(w[i]); + plan[P_TOKENT + i] = off + j; + ++j; + } + } +} + +kernel void moe_plan(device const int* ids [[buffer(0)]], device const float* w [[buffer(1)]], + device int* plan [[buffer(2)]], constant MoeArgs& a [[buffer(3)]], + ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup int s_grp[64], s_ent[64]; + plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); +} + +// The router's tail for the window, then the plan, in one threadgroup of n_expert (512) threads: simdgroup t routes +// token t exactly as native_router.cu (pinned llama.cpp topk-moe): softmax over the 512 logits, top 10 by +// probability (equal values: the lower expert index), weights = p / max(sum of the selected, 2^-14), the sum taken +// by winning lane in selection order and reduced with the same xor tree. +#ifndef STRATA_MOCE_FC_METAL +#define STRATA_MOCE_FC_METAL +constant int kMoceProtectFc [[function_constant(24)]]; +constant int kMoceProtect = is_function_constant_defined(kMoceProtectFc) ? kMoceProtectFc : 3; +constant uint kMoceScaleR1BitsFc [[function_constant(25)]]; +constant float kMoceScaleR1 = is_function_constant_defined(kMoceScaleR1BitsFc) ? as_type(kMoceScaleR1BitsFc) : 1.0f; +#endif +kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]], + device float* w [[buffer(2)]], device int* plan [[buffer(3)]], constant MoeArgs& a [[buffer(4)]], + device const uint* smap [[buffer(5)]], + ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup int s_grp[64], s_ent[64]; + if (sg < a.n_tok) { + device const float* l = logits + sg * 512; + const bool e512 = (a.n_expert == 512u); + float v[16]; + float mx = -INFINITY; + float sum = 0.0f; + if (e512) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] = l[lane + 32 * i]; + mx = max(mx, v[i]); + } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] = exp(v[i] - mx); + sum += v[i]; + } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + } else { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + v[i] = idx < int(a.n_expert) ? l[idx] : -INFINITY; + mx = max(mx, v[i]); + } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] = (lane + 32 * i < int(a.n_expert)) ? exp(v[i] - mx) : 0.0f; + sum += v[i]; + } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + if (lane + 32 * i >= int(a.n_expert)) { + v[i] = -INFINITY; + } else { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + } + } + const bool draft_moce = (a.pad & 0x80000000u) != 0u && sg > 0; + const float base_moce = (a.pad & 0x7FFFFFFFu) != 0u ? as_type(a.pad & 0x7FFFFFFFu) : 1.0f; + const float moce_scale = draft_moce ? (base_moce * base_moce) : base_moce; + const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f; + const int r_prot = draft_moce ? 1 : kMoceProtect; + const float r1_scale = kMoceScaleR1; + const float rcp_r1 = r1_scale > 1.0f ? (1.0f / r1_scale) : 1.0f; + const float r_step_scale = (r1_scale > 1.0f && moce_scale > 1.0f) ? (moce_scale * rcp_r1) : moce_scale; + const bool has_r1 = (r1_scale > 1.0f && r_prot > 1); + const bool has_moce = (moce_scale > 1.0f || r1_scale > 1.0f) && smap != nullptr; + const int min_scaled_rank = has_r1 ? 1 : r_prot; + ushort res_mask = 0; + if (has_moce) { + if (e512) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) + if (smap[lane + 32 * i] != 0xFFFFFFFFu) res_mask |= ushort(1u << i); + } else { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if (idx < int(a.n_expert) && smap[idx] != 0xFFFFFFFFu) res_mask |= ushort(1u << i); + } + } + } + ushort scaled_mask = 0; + float selected = 0.0f, ssum = 0.0f; + if (!has_moce) { + for (int rank = 0; rank < 10; ++rank) { + float best = v[0]; + int ex = lane; + #pragma clang loop unroll(full) + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) v[i] = (ex == lane + 32 * i) ? -INFINITY : v[i]; + if ((ex & 31) == lane) { + ids[sg * 10 + rank] = ex; + ssum += best; + } + if (rank == lane) selected = best; + } + } else { + for (int rank = 0; rank < 10; ++rank) { + if (has_r1) { + if (rank == 1) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + if ((res_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) { + v[i] *= r1_scale; + scaled_mask |= ushort(1u << i); + } + } + } else if (rank == r_prot && moce_scale > 1.0f) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + if ((scaled_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) { + v[i] *= r_step_scale; + } + } + } + } else if (rank == r_prot && moce_scale > 1.0f) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + if ((res_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) { + v[i] *= moce_scale; + scaled_mask |= ushort(1u << i); + } + } + } + float best = v[0]; + int ex = lane; + #pragma clang loop unroll(full) + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + float true_prob = best; + if (rank >= min_scaled_rank && best > 0.0f) { + const ushort win_lane = ushort(ex & 31); + 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) { + // without a tail bonus, the rank-1 scale remains on unselected residents. + const bool r1_scaled = has_r1 && (rank < r_prot || moce_scale <= 1.0f); + true_prob = best * (r1_scaled ? rcp_r1 : rcp_moce); + } + } + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) v[i] = (ex == lane + 32 * i) ? -INFINITY : v[i]; + if ((ex & 31) == lane) { + ids[sg * 10 + rank] = ex; + ssum += true_prob; + } + if (rank == lane) selected = true_prob; + } + } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + ssum = max(ssum, 6.103515625e-5f); + const float inv = 1.0f / ssum; + if (lane < 10) w[sg * 10 + lane] = selected * inv; + } + threadgroup_barrier(mem_flags::mem_device); + plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent, smap); +} + +// ---- gate + up + SwiGLU ------------------------------------------------------------------------------------------- +// One pass over a group's rows for all its NE entries (1..8: an expert appears once per token); NE is uniform per +// threadgroup, so the kernel switches once into a body specialised for it. +// +// gu_fused: per row, gate and up decoded together and each entry's x slice loaded once for both (the 32-weight units, +// whose x load is two runs and their sums: the M5 Air, 2560 x 640, NE 3-4: Q4_K / Q5_K / Q6_K / Q8_0 -8..-25%, +// IQ4_XS -31..-36% from NE 2). The 16-weight units keep two x loads a row: the two decoded units held at once cost +// them more (IQ2_XS / IQ3_S at NR 1, NE 3: +15..+21%). The same sums in the same order either way. +template struct gu_fused { + enum { value = fmt::UNIT == 32 && (NE >= 3 || F == T_IQ4_XS) }; +}; + +template +static inline float q2_combine(ushort2 pk, float d, thread const float* xs, float xsum0) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; + #pragma clang loop unroll(full) + for (int j = 0; j < 8; j += 2) { + if (F == T_Q2_0R) { + const uint2 q = uint2(pk) >> (2 * j); + const uint2 bits = ((q | (q << 14)) & 0x00030003u) | 0x64006400u; + const float4 v = float4(as_type(bits) - 1025.0h); + s0 = fma(v.x, xs[j], s0); + s1 = fma(v.y, xs[j + 1], s1); + s2 = fma(v.z, xs[8 + j], s2); + s3 = fma(v.w, xs[8 + j + 1], s3); + } else { + s0 = fma(float(pk.x & ushort(3u << (2 * j))), xs[j], s0); + s1 = fma(float(pk.x & ushort(3u << (2 * j + 2))), xs[j + 1], s1); + s2 = fma(float(pk.y & ushort(3u << (2 * j))), xs[8 + j], s2); + s3 = fma(float(pk.y & ushort(3u << (2 * j + 2))), xs[8 + j + 1], s3); + } + } + const float dot = d * (s0 + s1) + d * (s2 + s3); + return F == T_Q2_0R ? (dot - 0.0f) : (dot - (d * xsum0 + 0.0f)); +} + +static inline float iq_combine(uint4 g, uint s16, float sc, thread const float* xs) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 2; ++k) { + const uchar4 v0 = as_type(g[k]), v1 = as_type(g[2 + k]); + const uint m0 = s16 >> (4 * k), m1 = s16 >> (8 + 4 * k); + #pragma clang loop unroll(full) + for (int j = 0; j < 4; j += 2) { + s0 = fma((m0 & (1u << j)) ? -float(v0[j]) : float(v0[j]), xs[4 * k + j], s0); + s1 = fma((m0 & (2u << j)) ? -float(v0[j + 1]) : float(v0[j + 1]), xs[4 * k + j + 1], s1); + s2 = fma((m1 & (1u << j)) ? -float(v1[j]) : float(v1[j]), xs[8 + 4 * k + j], s2); + s3 = fma((m1 & (2u << j)) ? -float(v1[j + 1]) : float(v1[j + 1]), xs[8 + 4 * k + j + 1], s3); + } + } + return sc * (s0 + s1) + sc * (s2 + s3) - 0.0f; +} + +template +static inline float deq_iq_packed(device const uchar* b, int ui, thread uint4& g, thread uint& s16, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + if (F == T_IQ2_S) { + const uint q2 = *(device const ushort*)(b + 2 + 2 * ui); + s16 = *(device const ushort*)(b + 34 + 2 * ui); + const uint qh = uint(b[66 + ib]) >> (4 * h); + const uint i0 = (q2 & 255) | ((qh << 8) & 0x300); + const uint i1 = (q2 >> 8) | ((qh << 6) & 0x300); + const uint2 g0 = ctab ? as_type(iq2s_grid[i0]) : ((threadgroup const uint2*)tab)[i0]; + const uint2 g1 = ctab ? as_type(iq2s_grid[i1]) : ((threadgroup const uint2*)tab)[i1]; + g = uint4(g0.x, g0.y, g1.x, g1.y); + const uint sc = b[74 + ib]; + return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); + } else if (F == T_IQ2_XXS) { + device const uchar* qs = b + 2 + 4 * ib + 2 * ui; + const uint a1 = rd32(b + 2 + 8 * ib + 4), s0 = a1 >> (14 * h); + const uint sa = s0 & 127, sb = (s0 >> 7) & 127; + s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12); + const uint2 g0 = ctab ? as_type(iq2xxs_grid[qs[0]]) : ((threadgroup const uint2*)tab)[qs[0]]; + const uint2 g1 = ctab ? as_type(iq2xxs_grid[qs[1]]) : ((threadgroup const uint2*)tab)[qs[1]]; + g = uint4(g0.x, g0.y, g1.x, g1.y); + return rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f); + } else if (F == T_IQ2_XS) { + device const ushort* qs = (device const ushort*)(b + 2) + 2 * ui; + const uint v0 = qs[0], v1 = qs[1]; + const uint sa = v0 >> 9, sb = v1 >> 9; + s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12); + const uint2 g0 = ctab ? as_type(iq2xs_grid[v0 & 511]) : ((threadgroup const uint2*)tab)[v0 & 511]; + const uint2 g1 = ctab ? as_type(iq2xs_grid[v1 & 511]) : ((threadgroup const uint2*)tab)[v1 & 511]; + g = uint4(g0.x, g0.y, g1.x, g1.y); + const uint sc = b[66 + ib]; + return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); + } else if (F == T_IQ3_S) { + const int l0 = 2 * h, l1 = 2 * h + 1; + const uint qh = b[66 + ib]; + s16 = *(device const ushort*)(b + 74 + 2 * ui); + const uint i0 = uint(b[2 + 8 * ib + 2 * l0]) | ((qh << (8 - 2 * l0)) & 256); + const uint i1 = uint(b[2 + 8 * ib + 2 * l0 + 1]) | ((qh << (7 - 2 * l0)) & 256); + const uint i2 = uint(b[2 + 8 * ib + 2 * l1]) | ((qh << (8 - 2 * l1)) & 256); + const uint i3 = uint(b[2 + 8 * ib + 2 * l1 + 1]) | ((qh << (7 - 2 * l1)) & 256); + g = uint4(g_iq3s(tab, i0), g_iq3s(tab, i1), g_iq3s(tab, i2), g_iq3s(tab, i3)); + const uint sc = b[106 + (ib >> 1)]; + return rdh(b) * fma(float((sc >> (4 * (ib & 1))) & 15), 2.0f, 1.0f); + } else { + const uint aux = rd32(b + 2 + 64 + 4 * ib), s0 = aux >> (14 * h); + const uint sa = s0 & 127, sb = (s0 >> 7) & 127; + s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12); + device const uchar* qs = b + 2 + 4 * ui; + g = uint4(g_iq3xxs(tab, qs[0]), g_iq3xxs(tab, qs[1]), g_iq3xxs(tab, qs[2]), g_iq3xxs(tab, qs[3])); + return rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f); + } +} + +template +static inline void gu_body(device const uchar* g0, device const uchar* u0, device const float* x, + device const float* xs16, device float* h, device const int* etok, int c, int row0, int li, + constant MoeArgs& a, threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT; + float ag[NR][NE], au[NR][NE]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) ag[r][t] = au[r][t] = 0.0f; + const int U = int(a.n_embd) / UNIT; + if (NE == 1) { + const int tok = etok[0]; + x += tok * int(a.n_embd); + xs16 += tok * int(a.n_embd / 16); + } + for (int u = li; u < U; u += LPR) { + if (NE == 1 && UNIT == 16 && F != T_Q2_0 && F != T_Q2_0R) { + float xs[UNIT], xsum[2]; + X::load(x, xs16, u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 c; + Dec::run(g0 + r * a.gu_row, u, q, c, tab); + ag[r][0] += unit_combine(q, c, xs, xsum); + Dec::run(u0 + r * a.gu_row, u, q, c, tab); + au[r][0] += unit_combine(q, c, xs, xsum); + } + } else if (NE == 1) { + float xs[UNIT], xsum[2]; + X::load(x, xs16, u, xs, xsum); + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 c; + Dec::run(g0 + r * a.gu_row, u, q, c, tab); + ag[r][0] += unit_combine(q, c, xs, xsum); + Dec::run(u0 + r * a.gu_row, u, q, c, tab); + au[r][0] += unit_combine(q, c, xs, xsum); + } + } else if (F == T_Q2_0 || F == T_Q2_0R) { + const int boff = (u >> 2) * 18, qoff = 2 + 4 * (u & 3); + ushort2 pkg[NR], pku[NR]; + float dg[NR], du[NR]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + device const uchar* bg = g0 + r * a.gu_row + boff; + device const uchar* bu = u0 + r * a.gu_row + boff; + device const ushort* pg = (device const ushort*)(bg + qoff); + device const ushort* pu = (device const ushort*)(bu + qoff); + pkg[r] = ushort2(pg[0], pg[1]); + pku[r] = ushort2(pu[0], pu[1]); + dg[r] = rdh(bg); + du[r] = rdh(bu); + } + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + ag[r][t] += q2_combine(pkg[r], dg[r], xs, xsum[0]); + au[r][t] += q2_combine(pku[r], du[r], xs, xsum[0]); + } + } + } else if (gu_fused::value) { + for (int r = 0; r < NR; ++r) { + float qg[UNIT], qu[UNIT]; + float4 cg, cu; + Dec::run(g0 + r * a.gu_row, u, qg, cg, tab); + Dec::run(u0 + r * a.gu_row, u, qu, cu, tab); + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + ag[r][t] += unit_combine(qg, cg, xs, xsum); + au[r][t] += unit_combine(qu, cu, xs, xsum); + } + } + } else if (NR == 1 && NE == 4 && (F == T_IQ2_S || F == T_IQ2_XXS || F == T_IQ2_XS || F == T_IQ3_S || F == T_IQ3_XXS)) { + const int boff = (u >> 4) * fmt::BB, ui = (LPR % 16 == 0) ? (li & 15) : (u & 15); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + uint4 gg, gu; + uint sg, su; + const float dg = deq_iq_packed(g0 + r * a.gu_row + boff, ui, gg, sg, tab); + const float du = deq_iq_packed(u0 + r * a.gu_row + boff, ui, gu, su, tab); + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + ag[r][t] += iq_combine(gg, sg, dg, xs); + au[r][t] += iq_combine(gu, su, du, xs); + } + } + } else { + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 c; + Dec::run(g0 + r * a.gu_row, u, q, c, tab); + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + ag[r][t] += unit_combine(q, c, xs, xsum); + } + Dec::run(u0 + r * a.gu_row, u, q, c, tab); + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + au[r][t] += unit_combine(q, c, xs, xsum); + } + } + } + } + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const float gv = lane_sum(ag[r][t]), uv = lane_sum(au[r][t]); + if (li == 0) h[(c + t) * int(a.n_ff) + row0 + r] = gv / (1.0f + exp(-gv)) * uv; + } +} + +// function constant 27, with 26's plan: 1 = only the groups resident at routing time (indirect grid P_IGU_R / P_IDN_R), +// 2 = only the missing ones (P_IGU_M / P_IDN_M); 0 or unset = all (P_IGU / P_IDN) +constant int kPhaseFc [[function_constant(27)]]; +constant int kPhase = is_function_constant_defined(kPhaseFc) ? kPhaseFc : 0; + +// 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; + +template +kernel void moe_gu(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]], + device float* h [[buffer(2)]], device const int* plan [[buffer(3)]], + constant MoeArgs& a [[buffer(4)]], device const float* xs16 [[buffer(5)]], + device const uint* smap [[buffer(6)]], device const uchar* W1 [[buffer(7)]], + device const uchar* W2 [[buffer(8)]], device const uchar* W3 [[buffer(9)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1 && moe_probe != 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int g = int(tg.y) + (kPhase == 2 ? plan[P_SPLIT] : 0); + if (g >= (kPhase == 1 ? plan[P_SPLIT] : plan[P_NG])) return; + const int li = lane % LPR; + const int row0 = ((int(tg.x) * SG + sg) * RG + lane / LPR) * NR; + const ulong ex = (ulong)plan[P_GEXP + g]; + const int c = plan[P_GSTART + g], ne = plan[P_GSTART + g + 1] - c; + device const int* etok = plan + P_ETOK + c; + const uint scode = (kArena && kPhase == 1) ? uint(plan[P_GSLOT + g]) : (kArena ? smap[ex] : 0u); + device const uchar* E = kArena ? arena_slot(scode, W, W1, W2, W3) : W + ex * a.gu_stride; + 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) { + case 1: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + case 2: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + case 3: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + case 4: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + default: + for (int k = 0; k < ne; k += 4) { + const int m = min(4, ne - k); + if (m == 4) gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + else if (m == 3) gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + else if (m == 2) gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + else gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + } + } +} + +// ---- down --------------------------------------------------------------------------------------------------------- +template +static inline void dn_body(device const uchar* d0, device const float* h, device const float* hs16, device float* ye, + device const int* plan, int c, int row0, int li, constant MoeArgs& a, + threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT; + float ac[NR][NE]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + #pragma clang loop unroll(full) + 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) { + if (F == T_Q2_0 || F == T_Q2_0R) { + device const uchar* bu = d0 + (u >> 2) * 18; + const int qoff = 2 + 4 * (u & 3); + ushort2 pk[NR]; + float d[NR]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + device const uchar* b = bu + r * a.d_row; + device const ushort* p = (device const ushort*)(b + qoff); + pk[r] = ushort2(p[0], p[1]); + d[r] = rdh(b); + } + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + float xs[UNIT], xsum[2]; + X::load(h + (c + t) * int(a.n_ff), hs16 + (c + t) * int(a.n_ff / 16), u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + ac[r][t] += q2_combine(pk[r], d[r], xs, xsum[0]); + } + } else if (NE == 1) { + float xs[UNIT], xsum[2]; + X::load(h + c * int(a.n_ff), hs16 + c * int(a.n_ff / 16), u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 sc; + Dec::run(d0 + r * a.d_row, u, q, sc, tab); + ac[r][0] += unit_combine(q, sc, xs, xsum); + } + } else { + float q[NR][UNIT]; + float4 sc[NR]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(d0 + r * a.d_row, u, q[r], sc[r], tab); + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + float xs[UNIT], xsum[2]; + X::load(h + (c + t) * int(a.n_ff), hs16 + (c + t) * int(a.n_ff / 16), u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) ac[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + } + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const float v = lane_sum(ac[r][t]); + if (li == 0) ye[(c + t) * int(a.n_embd) + row0 + r] = as_type(plan[P_EW + c + t]) * v; + } +} + +template +kernel void moe_down(device const uchar* W [[buffer(0)]], device const float* h [[buffer(1)]], + device float* ye [[buffer(2)]], device const int* plan [[buffer(3)]], + constant MoeArgs& a [[buffer(4)]], device const float* hs16 [[buffer(5)]], + device const uint* smap [[buffer(6)]], device const uchar* W1 [[buffer(7)]], + device const uchar* W2 [[buffer(8)]], device const uchar* W3 [[buffer(9)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int g = int(tg.y) + (kPhase == 2 ? plan[P_SPLIT] : 0); + if (g >= (kPhase == 1 ? plan[P_SPLIT] : plan[P_NG])) return; + const int li = lane % LPR; + const int row0 = ((int(tg.x) * SG + sg) * RG + lane / LPR) * NR; + const ulong ex = (ulong)plan[P_GEXP + g]; + const int c = plan[P_GSTART + g], ne = plan[P_GSTART + g + 1] - c; + const uint scode = (kArena && kPhase == 1) ? uint(plan[P_GSLOT + g]) : (kArena ? smap[ex] : 0u); + device const uchar* d0 = (kArena ? arena_slot(scode, W, W1, W2, W3) : W + ex * a.d_stride) + a.down_off + (ulong)row0 * a.d_row; + switch (ne) { + case 1: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + case 2: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + case 3: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + case 4: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + default: + for (int k = 0; k < ne; k += 4) { + const int m = min(4, ne - k); + if (m == 4) dn_body(d0, h, hs16, ye, plan, c + k, row0, li, a, tab); + else if (m == 3) dn_body(d0, h, hs16, ye, plan, c + k, row0, li, a, tab); + else if (m == 2) dn_body(d0, h, hs16, ye, plan, c + k, row0, li, a, tab); + else dn_body(d0, h, hs16, ye, plan, c + k, row0, li, a, tab); + } + } +} + +kernel void moe_combine(device const float* ye [[buffer(0)]], device float* y [[buffer(1)]], + device const int* plan [[buffer(2)]], constant MoeArgs& a [[buffer(3)]], + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.n_embd, k = gid % a.n_embd; + if (t >= a.n_tok) return; + device const int* te = plan + P_TOKENT + t * a.topk; + float s = 0.0f; + if (a.topk == 10) { +#pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += ye[te[j] * int(a.n_embd) + k]; + } else { + for (uint j = 0; j < a.topk; ++j) s += ye[te[j] * int(a.n_embd) + k]; + } + y[gid] = s; +} + +// ---- instantiations: moe_gu____, moe_down____ ------------------------------- +#define GU(F, L, R, S) template [[host_name("moe_gu_" #F "_" #L "_" #R "_" #S)]] kernel decltype(moe_gu) moe_gu; +#define DN(F, L, R, S) template [[host_name("moe_down_" #F "_" #L "_" #R "_" #S)]] kernel decltype(moe_down) moe_down; +#define GU_SET(F) GU(F, 32, 1, 4) GU(F, 32, 2, 4) GU(F, 32, 4, 4) GU(F, 32, 2, 2) GU(F, 32, 2, 8) GU(F, 32, 1, 8) \ + GU(F, 16, 1, 4) GU(F, 16, 2, 4) GU(F, 8, 1, 4) GU(F, 8, 2, 4) GU(F, 8, 1, 2) +#define DN_SET(F) DN(F, 8, 1, 4) DN(F, 8, 2, 4) DN(F, 8, 4, 4) DN(F, 8, 1, 8) DN(F, 8, 2, 8) DN(F, 4, 1, 4) DN(F, 4, 2, 4) \ + DN(F, 4, 4, 4) DN(F, 16, 1, 4) DN(F, 16, 2, 4) DN(F, 32, 1, 4) DN(F, 2, 2, 4) +GU_SET(T_Q2_0) +GU_SET(T_IQ3_XXS) +GU_SET(T_IQ2_XXS) +GU_SET(T_IQ2_XS) +GU_SET(T_IQ2_S) +GU_SET(T_IQ3_S) +GU_SET(T_IQ4_XS) // the IQ3_S pack's gate/up in one layer +GU_SET(T_Q4_K) GU_SET(T_Q5_K) GU_SET(T_Q6_K) GU_SET(T_Q8_0) // the UD-Q4/5/6_K_XL packs +DN_SET(T_Q5_1) DN_SET(T_Q8_0) +DN_SET(T_Q2_0) +DN_SET(T_Q2_0R) +DN_SET(T_IQ4_NL) diff --git a/results/round39/astra_r39_qsa/baseline-kernels/moe_mm.metal b/results/round39/astra_r39_qsa/baseline-kernels/moe_mm.metal new file mode 100644 index 0000000..03ac90a --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/moe_mm.metal @@ -0,0 +1,1312 @@ +// metal/kernels/moe_mm.metal - routed experts for a PROMPT chunk (after quant/units/mma.metal): a grouped GEMM on +// the matrix units through Metal 4's MetalPerformancePrimitives matmul2d (on the M5 its neural accelerators). +// +// Plan (deterministic, three passes): the chunk's T x topk (token, expert) entries are sorted by expert, then by +// entry index; each expert's run is padded to a multiple of TM rows; tile i = (expert, first padded row). +// mm_hist per block of 1024 entries: a histogram of experts -> bcnt[block][expert] +// mm_scan one threadgroup (thread = expert): block offsets, the padded expert offsets, the tile list +// mm_place per block: an entry's padded row = expert offset + earlier blocks' count + rank in its block; +// srt_tok[row] = token (-1 on padding), srt_w[row] = weight, pos[entry] = row +// GEMMs: a threadgroup takes a tile (TM rows of one expert) and TN output rows; per K chunk of TK it converts the +// tile's activations to fp16 and dequantizes TN x TK weights (ref_quant's arithmetic, then rounded to fp16) into +// threadgroup memory, and matmul2d multiplies them, fp32 accumulation. +// mm_gu H[row][f] = silu(x . gate_f) * (x . up_f), fp32, for the n_ff rows +// mm_down YE[row][d] = h . down_d, fp32 +// mm_combine y[t][d] = sum_k w(t, k) * YE[pos(t, k)][d], k in order +// +// The idea of feeding matmul2d tiles of dequantized weights from threadgroup memory follows yibie/strata-mlx's +// tile_matmul (MIT); the code is our own. +#include +using namespace mpp::tensor_ops; + +struct MmPlanArgs { + uint n_ent; // T * topk + uint topk; + uint n_expert; + uint tm; // the tile height (rows of one expert) + uint n_blocks; // ceil(n_ent / 1024) +}; + +// counts per block: bcnt[block * n_expert + e] +kernel void mm_hist(device const int* ids [[buffer(0)]], device uint* bcnt [[buffer(1)]], + constant MmPlanArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup atomic_uint h[1024]; + if (tid < a.n_expert) atomic_store_explicit(&h[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint i = tg.x * 1024 + tid; + if (i < a.n_ent && ids[i] >= 0) atomic_fetch_add_explicit(&h[ids[i]], 1u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < a.n_expert) + bcnt[tg.x * a.n_expert + tid] = atomic_load_explicit(&h[tid], memory_order_relaxed); +} + +// one threadgroup of n_expert (<= 1024) threads. boff[block][e] = expert e's rows before this block (within e); +// eoff[e] = expert e's first padded row; tiles[i] = (expert, row0); meta[0] = n_tiles, meta[1] = padded rows; +// indirect grids: meta[2..4] = (gx_gu, n_tiles, 1) for the gate/up GEMM, meta[5..7] = (gx_down, n_tiles, 1) for the +// down GEMM (gx_gu in meta[2], gx_down in meta[5] written by the host once), meta[8..10] = (rows, 1, 1) for mm_gather. +// Padding rows [r0 + tot, r0 + rows) of srt_tok (when bound at buffer 6) are cleared to -1 here so mm_clear is not needed. +kernel void mm_scan(device const uint* bcnt [[buffer(0)]], device uint* boff [[buffer(1)]], + device uint* eoff [[buffer(2)]], device int2* tiles [[buffer(3)]], device uint* meta [[buffer(4)]], + constant MmPlanArgs& a [[buffer(5)]], device int* srt_tok [[buffer(6)]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort nsg [[simdgroups_per_threadgroup]]) { + threadgroup uint s_rows[32], s_tiles[32]; + const uint e = tid; + uint tot = 0; + if (e < a.n_expert) + for (uint b = 0; b < a.n_blocks; ++b) { + boff[b * a.n_expert + e] = tot; + tot += bcnt[b * a.n_expert + e]; + } + const uint nt = (tot + a.tm - 1) / a.tm, rows = nt * a.tm; + const uint pr = simd_prefix_exclusive_sum(rows), pt = simd_prefix_exclusive_sum(nt); + if (lane == 31) { s_rows[sg] = pr + rows; s_tiles[sg] = pt + nt; } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint orow = 0, otile = 0, all_rows = 0, all_tiles = 0; + for (uint i = 0; i < nsg; ++i) { + if (i < sg) { orow += s_rows[i]; otile += s_tiles[i]; } + all_rows += s_rows[i]; + all_tiles += s_tiles[i]; + } + if (e < a.n_expert) { + const uint r0 = orow + pr, t0 = otile + pt; + eoff[e] = r0; + for (uint k = 0; k < nt; ++k) tiles[t0 + k] = int2(int(e), int(r0 + k * a.tm)); + if (srt_tok) + for (uint r = tot; r < rows; ++r) srt_tok[r0 + r] = -1; + } + if (tid == 0) { + meta[0] = all_tiles; + meta[1] = all_rows; + meta[3] = all_tiles; // indirect grid of the gate/up GEMM: (meta[2] given, tiles, 1) + meta[4] = 1; + meta[6] = all_tiles; // of the down GEMM: (meta[5] given, tiles, 1) + meta[7] = 1; + meta[8] = all_rows; // of the gather: (rows, 1, 1) + meta[9] = 1; + meta[10] = 1; + } +} + +static inline uint mm_sum_nibbles(uint c) { + c = (c & 0x0F0F0F0Fu) + ((c >> 4) & 0x0F0F0F0Fu); + c = (c + (c >> 8)) & 0x00FF00FFu; + return (c + (c >> 16)) & 0xFFFFu; +} + +// grid (n_blocks), 1024 threads. Two-level deterministic intra-block rank: 32-lane sub-block nibble counts in p_cnt +// plus simd_broadcast scan within the caller's 32-lane simdgroup (no s_id[1024] threadgroup array). +kernel void mm_place(device const int* ids [[buffer(0)]], device const float* w [[buffer(1)]], + device const uint* boff [[buffer(2)]], device const uint* eoff [[buffer(3)]], + device int* srt_tok [[buffer(4)]], device float* srt_w [[buffer(5)]], device int* pos [[buffer(6)]], + constant MmPlanArgs& a [[buffer(7)]], + uint tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint p_cnt[4096]; + const uint n4e = 4 * a.n_expert; + if (a.n_expert == 512) { + atomic_store_explicit(&p_cnt[tid], 0u, memory_order_relaxed); + atomic_store_explicit(&p_cnt[tid + 1024], 0u, memory_order_relaxed); + } else { + for (uint k = tid; k < n4e; k += 1024) atomic_store_explicit(&p_cnt[k], 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint i = tg * 1024 + tid; + const int e = i < a.n_ent ? ids[i] : -1; + const uint sub = tid >> 5; + if (e >= 0 && sub < 31) { + const uint idx = (sub >> 3) * a.n_expert + uint(e); + const uint sh = (sub & 7u) * 4u; + atomic_fetch_add_explicit(&p_cnt[idx], 1u << sh, memory_order_relaxed); + } + uint rank = 0; + #pragma clang loop unroll(full) + for (ushort l = 0; l < 32; ++l) { + const int el = simd_broadcast(e, l); + if (l < lane && el == e) ++rank; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (i >= a.n_ent) return; + if (e < 0) { pos[i] = -1; return; } + for (uint g = 0; g < (sub >> 3); ++g) { + const uint c = atomic_load_explicit(&p_cnt[g * a.n_expert + uint(e)], memory_order_relaxed); + rank += mm_sum_nibbles(c); + } + if (sub & 7u) { + const uint m = (1u << ((sub & 7u) * 4u)) - 1u; + const uint c = atomic_load_explicit(&p_cnt[(sub >> 3) * a.n_expert + uint(e)], memory_order_relaxed) & m; + rank += mm_sum_nibbles(c); + } + const uint row = eoff[e] + boff[tg * a.n_expert + e] + rank; + srt_tok[row] = int(i / a.topk); + srt_w[row] = w[i]; + pos[i] = int(row); +} + +kernel void mm_clear(device int* srt_tok [[buffer(0)]], constant uint& n [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + if (gid < n) srt_tok[gid] = -1; +} + +#ifndef MM_GATHER_ONLY +// ---- the GEMMs ----------------------------------------------------------------------------------------------------- +// weights of one expert: gate at W + gate_off + e * gu_stride (+ up_off for up), down at W + down_off + e * d_stride +struct MmArgs { + ulong gate_off; + ulong up_off; + ulong down_off; + ulong gu_stride; + ulong d_stride; + uint gu_row; // bytes per gate/up row (n_embd values) + uint d_row; // bytes per down row (n_ff values) + uint n_embd; + uint n_ff; +}; + +// TN x TK weights of `nrow`-strided rows into threadgroup half: thread-strided units of the format +template +static inline void stage_w(device const uchar* w0, uint row_bytes, int k0, threadgroup half* dst, + threadgroup const uint* tab, ushort tid, ushort nth) { + constexpr int UNIT = fmt::UNIT, UPC = TK / UNIT, H = UNIT / 2; + static_assert(TK % UNIT == 0, "the K chunk must hold whole units of the format"); + for (int it = tid; it < TN * UPC; it += nth) { + const int r = it / UPC, uu = it - r * UPC; + float wv[UNIT]; + int p0, p1; + DW::run(w0 + ulong(r) * row_bytes, k0 / UNIT + uu, k0, wv, p0, p1, tab); +#pragma clang loop unroll(full) + for (int j = 0; j < H; ++j) { + dst[r * TK + p0 + j] = half(wv[j]); + dst[r * TK + p1 + j] = half(wv[H + j]); + } + } +} + +#ifndef MM_V3_ONLY +template +kernel void mm_gu(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]], + device float* H [[buffer(2)]], device const int* srt_tok [[buffer(3)]], + device const int2* tiles [[buffer(4)]], constant MmArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int NTH = SG * 32; + threadgroup half xa[TM * TK]; + threadgroup half wg[TN * TK]; + threadgroup half wu[TN * TK]; + threadgroup uint tab[fmt::TAB]; + threadgroup int tok[TM]; + load_tab(tab, tid, NTH); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, f0 = int(tg.x) * TN; + for (int m = tid; m < TM; m += NTH) tok[m] = srt_tok[row0 + m]; + device const uchar* g0 = W + a.gate_off + ulong(e) * a.gu_stride + ulong(f0) * a.gu_row; + device const uchar* u0 = W + a.up_off + ulong(e) * a.gu_stride + ulong(f0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false); + matmul2d> op; + tensor, tensor_inline> tA(xa, dextents(TK, TM)); + tensor, tensor_inline> tG(wg, dextents(TK, TN)); + tensor, tensor_inline> tU(wu, dextents(TK, TN)); + auto accg = op.template get_destination_cooperative_tensor(); + auto accu = op.template get_destination_cooperative_tensor(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + for (int it = tid; it < TM * TK; it += NTH) { + const int m = it / TK, k = it - m * TK; + xa[it] = tok[m] >= 0 ? half(x[ulong(tok[m]) * a.n_embd + k0 + k]) : half(0.0f); + } + stage_w(g0, a.gu_row, k0, wg, tab, tid, NTH); + stage_w(u0, a.gu_row, k0, wu, tab, tid, NTH); + threadgroup_barrier(mem_flags::mem_threadgroup); + auto pg = op.template get_destination_cooperative_tensor(); + auto pu = op.template get_destination_cooperative_tensor(); + op.run(tA, tG, pg); + op.run(tA, tU, pu); + for (uint16_t i = 0; i < pg.get_capacity(); ++i) + if (pg.is_valid_element(i)) { accg[i] += pg[i]; accu[i] += pu[i]; } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = accg[i]; + accg[i] = g / (1.0f + exp(-g)) * accu[i]; + } + tensor, tensor_inline> tH(H, dextents(a.n_ff, 1 << 30)); + accg.store(tH.slice(f0, row0)); +} + +template +kernel void mm_down(device const uchar* W [[buffer(0)]], device const float* H [[buffer(1)]], + device float* YE [[buffer(2)]], device const int2* tiles [[buffer(3)]], + constant MmArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int NTH = SG * 32; + threadgroup half ha[TM * TK]; + threadgroup half wd[TN * TK]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, NTH); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, d0 = int(tg.x) * TN; + device const uchar* w0 = W + a.down_off + ulong(e) * a.d_stride + ulong(d0) * a.d_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false); + matmul2d> op; + tensor, tensor_inline> tA(ha, dextents(TK, TM)); + tensor, tensor_inline> tW(wd, dextents(TK, TN)); + auto acc = op.template get_destination_cooperative_tensor(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) { + for (int it = tid; it < TM * TK; it += NTH) { + const int m = it / TK, k = it - m * TK; + ha[it] = half(H[ulong(row0 + m) * a.n_ff + k0 + k]); + } + stage_w(w0, a.d_row, k0, wd, tab, tid, NTH); + threadgroup_barrier(mem_flags::mem_threadgroup); + auto p = op.template get_destination_cooperative_tensor(); + op.run(tA, tW, p); + for (uint16_t i = 0; i < p.get_capacity(); ++i) + if (p.is_valid_element(i)) acc[i] += p[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + tensor, tensor_inline> tY(YE, dextents(a.n_embd, 1 << 30)); + acc.store(tY.slice(d0, row0)); +} + +// grid (T * n_embd / 256), 256 threads +#endif +kernel void mm_combine(device const float* YE [[buffer(0)]], device const int* pos [[buffer(1)]], + device const float* srt_w [[buffer(2)]], device float* y [[buffer(3)]], + constant uint4& a [[buffer(4)]], // (T, n_embd, topk, -) + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.y, d = gid % a.y; + if (t >= a.x) return; + float s = 0.0f; + for (uint k = 0; k < a.z; ++k) { + const int r = pos[t * a.z + k]; + if (r >= 0) s = fma(srt_w[r], YE[ulong(r) * a.y + d], s); + } + y[gid] = s; +} + +#ifndef MM_V3_ONLY +#define MMGU_(F, M, N, K, S) template [[host_name("mm_gu_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu) mm_gu; +#define MMDN_(F, M, N, K, S) template [[host_name("mm_down_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down) mm_down; +#define MMGU(F, M, N, K, S) MMGU_(F, M, N, K, S) +#define MMDN(F, M, N, K, S) MMDN_(F, M, N, K, S) +#define MM_GU_SET(F) MMGU(F, 32, 64, 64, 4) MMGU(F, 64, 64, 64, 4) MMGU(F, 16, 64, 64, 2) MMGU(F, 32, 32, 64, 2) +#define MM_DN_SET(F) MMDN(F, 32, 64, 64, 4) MMDN(F, 64, 64, 64, 4) MMDN(F, 16, 64, 64, 2) MMDN(F, 32, 32, 64, 2) +MM_GU_SET(T_Q2_0) MM_GU_SET(T_IQ3_XXS) MM_GU_SET(T_IQ2_XXS) MM_GU_SET(T_IQ2_XS) MM_GU_SET(T_IQ2_S) MM_GU_SET(T_IQ3_S) +MM_DN_SET(T_Q2_0) MM_DN_SET(T_IQ4_NL) +#endif + +#endif +// ---- v2: activations gathered once per layer as fp16 rows, read by matmul2d from device memory --------------------- +// mm_gather: XS[row][k] = half(x[srt_tok[row]][k]) (0 on padding rows). grid (rows), 256 threads. +kernel void mm_gather(device const float* x [[buffer(0)]], device const int* srt_tok [[buffer(1)]], + device half* XS [[buffer(2)]], constant uint& n [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const int t = srt_tok[tg.x]; + device half* o = XS + ulong(tg.x) * n; + if (n == 2560) { + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const uint k = tid + i * 256; + o[k] = t >= 0 ? half(x[ulong(t) * n + k]) : half(0.0f); + } + } else { + for (uint k = tid; k < n; k += 256) o[k] = t >= 0 ? half(x[ulong(t) * n + k]) : half(0.0f); + } +} + +// mm_gather_h: gathers straight from fp16 rows (mixh) with 128-bit uint4 copies across nth threads per row. +kernel void mm_gather_h(device const half* xh [[buffer(0)]], device const int* srt_tok [[buffer(1)]], + device half* XS [[buffer(2)]], constant uint& n [[buffer(3)]], + uint tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort nth [[threads_per_threadgroup]]) { + const int t = srt_tok[tg]; + device uint4* o8 = (device uint4*)(XS + ulong(tg) * n); + const uint n8 = n >> 3; + if (n == 2560) { + device const uint4* x8 = (device const uint4*)(xh + ulong(max(t, 0)) * n); + if (nth == 32) { + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) o8[tid + i * 32] = t >= 0 ? x8[tid + i * 32] : uint4(0u); + return; + } + if (nth == 64) { + #pragma clang loop unroll(full) + for (uint i = 0; i < 5; ++i) o8[tid + i * 64] = t >= 0 ? x8[tid + i * 64] : uint4(0u); + return; + } + if (nth == 160) { + o8[tid] = t >= 0 ? x8[tid] : uint4(0u); + o8[tid + 160] = t >= 0 ? x8[tid + 160] : uint4(0u); + return; + } + if (nth == 256) { + o8[tid] = t >= 0 ? x8[tid] : uint4(0u); + if (tid < 64) o8[tid + 256] = t >= 0 ? x8[tid + 256] : uint4(0u); + return; + } + } + if (t >= 0) { + device const uint4* x8 = (device const uint4*)(xh + ulong(t) * n); + for (uint k = tid; k < n8; k += nth) o8[k] = x8[k]; + } else { + for (uint k = tid; k < n8; k += nth) o8[k] = uint4(0u); + } +} + +#ifndef MM_GATHER_ONLY +// The SSD-streaming arena (as in moe.metal): with function constant 20, expert e's slot is smap[e] in the chunks W, +// W1..W3 (buffers 5..8 of mm_gu2 / mm_down2 / mm_gu3 / mm_down3); gate_off / up_off / down_off are offsets in the slot. +#ifndef STRATA_ARENA_METAL +#define STRATA_ARENA_METAL +constant bool kArenaFc [[function_constant(20)]]; +constant bool kArena = is_function_constant_defined(kArenaFc) && kArenaFc; +static inline device const uchar* arena_slot(uint v, device const uchar* W, device const uchar* W1, device const uchar* W2, + device const uchar* W3) { + const uint c = v >> 28; + device const uchar* b = c == 0 ? W : (c == 1 ? W1 : (c == 2 ? W2 : W3)); + return b + (ulong)(v & 0x0FFFFFFFu) * 16384ul; +} +#endif + +// a probe for the bench (function constant 5): 1 = the weights not decoded (a cheap fill), 2 = no matrix multiply +constant int MM_PROBE [[function_constant(5)]]; +constant int mm_probe = is_function_constant_defined(MM_PROBE) ? MM_PROBE : 0; + +#ifndef MM_V3_ONLY +template +kernel void mm_gu2(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], + device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int NTH = SG * 32; + threadgroup half wg[TN * TK]; + threadgroup half wu[TN * TK]; + threadgroup float ot[TM * TN]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, NTH); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, f0 = int(tg.x) * TN; + device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride; + device const uchar* g0 = E + a.gate_off + ulong(f0) * a.gu_row; + device const uchar* u0 = E + a.up_off + ulong(f0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d> op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + tensor, tensor_inline> tG(wg, dextents(TK, TN)); + tensor, tensor_inline> tU(wu, dextents(TK, TN)); + auto tA = tX.template slice(0, row0); + auto accg = op.template get_destination_cooperative_tensor(); + auto accu = op.template get_destination_cooperative_tensor(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + if (mm_probe == 1) { + for (int it = tid; it < TN * TK; it += NTH) { wg[it] = half(k0 & 7); wu[it] = half(it & 7); } + } else { + stage_w(g0, a.gu_row, k0, wg, tab, tid, NTH); + stage_w(u0, a.gu_row, k0, wu, tab, tid, NTH); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + tA = tX.template slice(k0, row0); + if (mm_probe != 2) { + op.run(tA, tG, accg); + op.run(tA, tU, accu); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = accg[i]; + accg[i] = g / (1.0f + exp(-g)) * accu[i]; + } + tensor, tensor_inline> tO(ot, dextents(TN, TM)); + accg.store(tO); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int it = tid; it < TM * TN; it += NTH) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + m) * a.n_ff + f0 + f] = half(ot[it]); + } +} + +template +kernel void mm_down2(device const uchar* W [[buffer(0)]], device const half* HS [[buffer(1)]], + device float* YE [[buffer(2)]], device const int2* tiles [[buffer(3)]], + constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], + device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int NTH = SG * 32; + threadgroup half wd[TN * TK]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, NTH); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, d0 = int(tg.x) * TN; + device const uchar* w0 = (kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.d_stride) + a.down_off + ulong(d0) * a.d_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d> op; + tensor, tensor_inline> tH((device half*)HS, dextents(a.n_ff, 1 << 30)); + tensor, tensor_inline> tW(wd, dextents(TK, TN)); + auto tA = tH.template slice(0, row0); + auto acc = op.template get_destination_cooperative_tensor(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) { + stage_w(w0, a.d_row, k0, wd, tab, tid, NTH); + threadgroup_barrier(mem_flags::mem_threadgroup); + tA = tH.template slice(k0, row0); + op.run(tA, tW, acc); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + tensor, tensor_inline> tY(YE, dextents(a.n_embd, 1 << 30)); + acc.store(tY.slice(d0, row0)); +} + +#define MMGU2_(F, M, N, K, S) template [[host_name("mm_gu2_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu2) mm_gu2; +#define MMDN2_(F, M, N, K, S) template [[host_name("mm_down2_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down2) mm_down2; +#define MMGU2(F, M, N, K, S) MMGU2_(F, M, N, K, S) +#define MMDN2(F, M, N, K, S) MMDN2_(F, M, N, K, S) +#define MM_GU2_SET(F) MMGU2(F, 32, 64, 32, 2) MMGU2(F, 32, 32, 32, 2) MMGU2(F, 32, 64, 16, 2) MMGU2(F, 64, 32, 16, 2) MMGU2(F, 16, 32, 32, 2) MMGU2(F, 16, 64, 16, 2) MMGU2(F, 16, 128, 32, 4) MMGU2(F, 32, 64, 16, 4) MMGU2(F, 16, 32, 32, 1) MMGU2(F, 16, 64, 32, 4) MMGU2(F, 16, 64, 64, 2) MMGU2(F, 32, 64, 64, 4) MMGU2(F, 64, 64, 64, 4) MMGU2(F, 16, 32, 64, 2) MMGU2(F, 32, 64, 32, 4) MMGU2(F, 64, 64, 32, 4) MMGU2(F, 16, 64, 32, 2) +#define MM_DN2_SET(F) MMDN2(F, 32, 64, 32, 2) MMDN2(F, 32, 32, 32, 2) MMDN2(F, 32, 64, 16, 2) MMDN2(F, 64, 32, 16, 2) MMDN2(F, 16, 32, 32, 2) MMDN2(F, 16, 64, 16, 2) MMDN2(F, 16, 128, 32, 4) MMDN2(F, 32, 64, 16, 4) MMDN2(F, 16, 32, 32, 1) MMDN2(F, 16, 64, 32, 4) MMDN2(F, 16, 64, 64, 2) MMDN2(F, 32, 64, 64, 4) MMDN2(F, 64, 64, 64, 4) MMDN2(F, 16, 32, 64, 2) MMDN2(F, 32, 64, 32, 4) MMDN2(F, 64, 64, 32, 4) MMDN2(F, 16, 64, 32, 2) +#define MM_GU2_U32(F) MMGU2(F, 16, 64, 32, 2) MMGU2(F, 16, 64, 64, 2) MMGU2(F, 32, 64, 32, 2) MMGU2(F, 16, 32, 32, 2) \ + MMGU2(F, 64, 64, 32, 4) +#define MM_DN2_U32(F) MMDN2(F, 16, 64, 32, 2) MMDN2(F, 16, 64, 64, 2) MMDN2(F, 16, 32, 32, 2) MMDN2(F, 32, 32, 32, 2) \ + MMDN2(F, 16, 128, 32, 4) MMDN2(F, 32, 64, 32, 2) MMDN2(F, 64, 64, 32, 4) +// K-quants: the staged K chunk must hold whole units as DW places them (128 wide): TK 128 +#define MM_GU2_KQ(F) MMGU2(F, 16, 32, 128, 2) MMGU2(F, 16, 16, 128, 2) MMGU2(F, 32, 32, 128, 2) MMGU2(F, 16, 32, 128, 4) +MM_GU2_KQ(T_Q4_K) MM_GU2_KQ(T_Q5_K) MM_GU2_KQ(T_Q6_K) MM_GU2_U32(T_Q8_0) MM_DN2_U32(T_Q5_1) MM_DN2_U32(T_Q8_0) +// IQ4_XS (UNIT 32): the K chunk must be >= 32 +MMGU2(T_IQ4_XS, 16, 64, 32, 2) MMGU2(T_IQ4_XS, 16, 64, 64, 2) MMGU2(T_IQ4_XS, 32, 64, 32, 2) MMGU2(T_IQ4_XS, 16, 32, 32, 2) +MMGU2(T_IQ4_XS, 64, 64, 32, 4) +MM_GU2_SET(T_Q2_0) MM_GU2_SET(T_IQ3_XXS) MM_GU2_SET(T_IQ2_XXS) MM_GU2_SET(T_IQ2_XS) MM_GU2_SET(T_IQ2_S) MM_GU2_SET(T_IQ3_S) +MM_DN2_SET(T_Q2_0) +MMDN2(T_IQ4_NL, 16, 64, 32, 2) MMDN2(T_IQ4_NL, 16, 64, 64, 2) MMDN2(T_IQ4_NL, 16, 32, 32, 2) MMDN2(T_IQ4_NL, 32, 32, 32, 2) MMDN2(T_IQ4_NL, 16, 128, 32, 4) + +#endif + +// ---- v3: the weights decoded straight into matmul2d's right-input cooperative tensor (single-simdgroup scope), +// no threadgroup staging and no barriers in the K loop. Each simdgroup owns TN output rows; a threadgroup SG * TN. +// The right input's element layout (tests/coop_layout): runs of 4 consecutive k at one n, so a lane decodes 4 values +// at a time (W4). Activations: the gathered fp16 rows from device memory. +#ifndef STRATA_COOP_K4 +#define STRATA_COOP_K4 +// function constant 7: the right-input cooperative tensor holds runs of 4 consecutive k at one n per lane (the M5's +// layout; tests/coop_layout). false (the M4: runs of 2 n at one k): every element decoded on its own. +constant bool COOP_K4 [[function_constant(7)]]; +constant bool coop_k4 = is_function_constant_defined(COOP_K4) ? COOP_K4 : true; +#endif +template struct W4; +// the four weights k .. k + 3 (k % 4 == 0) of a row, in quant.metal's deq arithmetic (code * scale in fp32), as fp16 +// the threadgroup forms of quant.metal's field readers (blocks staged in threadgroup memory) +static inline uint rd32(threadgroup const uchar* p) { + threadgroup const ushort* s = (threadgroup const ushort*)p; + return uint(s[0]) | (uint(s[1]) << 16); +} +static inline float rdh(threadgroup const uchar* p) { return float(*(threadgroup const half*)p); } + +static inline half4 w4_signed(uint g, uint s, float sc) { + const uchar4 v = as_type(g); + const float4 q = float4((s & 1) ? -float(v[0]) : float(v[0]), (s & 2) ? -float(v[1]) : float(v[1]), + (s & 4) ? -float(v[2]) : float(v[2]), (s & 8) ? -float(v[3]) : float(v[3])); + return half4(q * sc); +} + +template <> struct W4 { // 64 values a block: fp16 d, 16 bytes of 2-bit codes (value v at byte v / 4) + // the codes c land in fp16 mantissas (1024 + c), so c - 1 and its product with d are exact in fp16: bit for bit the + // fp32 decode for every d and code byte (checked exhaustively), and 8-11% off mm_down3 at T 1024-2048 + static inline half4 dec(uint q, half d) { + const uint2 c = (uint2(q | (q << 14), (q >> 4) | (q << 10)) & 0x00030003u) | 0x64006400u; + return (as_type(c) - 1025.0h) * d; + } + template + static inline half4 run_q(device const uchar* b, uint kl4) { + return dec((b + kl4)[OFF], *(device const half*)b); + } + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + return run_q<2>(row + (k >> 6) * 18, uint(k & 63) >> 2); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 98; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint aux = rd32(b + 2 + 64 + 4 * ib); + const float sc = rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f); + return w4_signed(g_iq3xxs(tab, b[2 + ((k & 255) >> 2)]), sign4_iq2xs((aux >> (7 * l)) & 127, hh), sc); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 66; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint a0 = rd32(b + 2 + 8 * ib), a1 = rd32(b + 2 + 8 * ib + 4); + const uint idx = (a0 >> (8 * l)) & 255; + const float sc = rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f); + return w4_signed(g_iq2xxs(tab, idx, hh), sign4_iq2xs((a1 >> (7 * l)) & 127, hh), sc); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 74; + const int ib = (k & 255) >> 5, hh = (k & 7) >> 2; + const uint v = ((device const ushort*)(b + 2))[(k & 255) >> 3]; + const uint scb = b[66 + ib]; + const float sc = rdh(b) * fma(float((scb >> ((k & 16) >> 2)) & 15), 0.25f, 0.125f); + return w4_signed(g_iq2xs(tab, v & 511, hh), sign4_iq2xs(v >> 9, hh), sc); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 82; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint qh = b[66 + ib], o = (k & 255) >> 3; + const uint idx = uint(b[2 + o]) | ((qh << (8 - 2 * l)) & 0x300); + const uint s = b[34 + o]; + const uint scb = b[74 + ib]; + const float sc = rdh(b) * fma(float((scb >> ((k & 16) >> 2)) & 15), 0.25f, 0.125f); + return w4_signed(g_iq2s(tab, idx, hh), hh ? s >> 4 : s, sc); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 110; + const int ib = (k & 255) >> 5, hh = (k & 7) >> 2; + const uint qh = b[66 + ib]; + const uint i = uint(b[2 + ((k & 255) >> 2)]) | ((qh << (8 - ((k & 31) >> 2))) & 256u); + const uint s = b[74 + ((k & 255) >> 3)]; + const uint scb = b[106 + (ib >> 1)]; + const float sc = rdh(b) * float(1 + 2 * ((scb >> (4 * (ib & 1))) & 15)); + return w4_signed(g_iq3s(tab, i), hh ? s >> 4 : s, sc); + } +}; +template <> struct W4 { // 256 a block: d, scales_h, scales_l[4], qs[128]; 32-groups as IQ4_NL's + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 136; + const int j = k & 255, ib = j >> 5, o = j & 15, hi = (j >> 4) & 1; + const uint sh = *(device const ushort*)(b + 2); + const int ls = int((b[4 + (ib >> 1)] >> (4 * (ib & 1))) & 15) | int(((sh >> (2 * ib)) & 3) << 4); + const float s = rdh(b) * float(ls - 32); + device const uchar* qs = b + 8 + 16 * ib; + float4 q; + for (int c = 0; c < 4; ++c) { + const half2 w = as_type(tab[qs[o + c]]); + q[c] = float(hi ? w.y : w.x); + } + return half4(q * s); + } +}; +template <> struct W4 { // 32 a block: d, 16 bytes; low nibbles are elements 0..15, high 16..31 + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 5) * 18; + const int j = k & 31, hi = j >> 4, o = j & 15; + const float d = rdh(b); + float4 q; + for (int c = 0; c < 4; ++c) { + const half2 w = as_type(tab[b[2 + o + c]]); + q[c] = float(hi ? w.y : w.x); + } + return half4(q * d); + } +}; + +// W8: the runs k .. k + 3 and k + 16 .. k + 19 of a row (k % 32 < 16): one 32-group, so the per-group fields (scale, +// sign word) are read once. The generic one is two W4. +// ---- the K-quant / legacy-quant W4 decoders (the UD-Q4/5/6_K_XL packs' experts): values k .. k + 3 (k % 4 == 0) +// of a row in ref_quant's arithmetic (float), then fp16 +template <> struct W4 { // 32 a block: d, int8 qs[32] + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 5) * 34; + const float d = rdh(b); + device const char* q = (device const char*)(b + 2 + (k & 31)); + return half4(float4(float(q[0]), float(q[1]), float(q[2]), float(q[3])) * d); + } +}; +template <> struct W4 { // 32 a block: d, m, qh (32 bits), qs[16] (low nibbles 0..15, high 16..31) + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 5) * 24; + const float d = rdh(b), m = rdh(b + 2); + const uint hb = rd32(b + 4); + const int j = k & 31, o = j & 15, hi = j >> 4; + float4 w; + for (int c = 0; c < 4; ++c) { + const uint qv = b[8 + o + c]; + const int x = int(hi ? (qv >> 4) : (qv & 15)) | int(((hb >> (j + c)) & 1) << 4); + w[c] = float(x) * d + m; + } + return half4(w); + } +}; +template <> struct W4 { // 256 a block: d, dmin, scales[12], qs[128] (64-chunks: low nibbles then high) + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 8) * 144; + const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31; + float s, mm; + scale_min_k4(2 * c + hi, b + 4, s, mm); + const float d1 = rdh(b) * s, m1 = rdh(b + 2) * mm; + device const uchar* q = b + 16 + 32 * c + l; + float4 w; + for (int i = 0; i < 4; ++i) w[i] = d1 * float(hi ? (q[i] >> 4) : (q[i] & 15)) - m1; + return half4(w); + } +}; +template <> struct W4 { // 256 a block: d, dmin, scales[12], qh[32], qs[128] + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 8) * 176; + const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31; + float s, mm; + scale_min_k4(2 * c + hi, b + 4, s, mm); + const float d1 = rdh(b) * s, m1 = rdh(b + 2) * mm; + device const uchar* qh = b + 16 + l; + device const uchar* ql = b + 48 + 32 * c + l; + const uint bit = 1u << (2 * c + hi); + float4 w; + for (int i = 0; i < 4; ++i) + w[i] = d1 * float(int(hi ? (ql[i] >> 4) : (ql[i] & 15)) + ((qh[i] & bit) ? 16 : 0)) - m1; + return half4(w); + } +}; +template <> struct W4 { // 256 a block: ql[128], qh[64], int8 scales[16], d; halves of 128 + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 8) * 210; + const uint j = uint(k & 255), j4 = j >> 4; + const uint la = rd32(b + (((j >> 1) & 64u) | (j & 63u))); + const uint ha = rd32(b + 128u + (((j >> 2) & 32u) | (j & 31u))); + const float ds = rdh(b + 208) * float(((device const char*)(b + 192))[j4]); + const uint qa = ((la >> (j4 & 4u)) & 0x0F0F0F0Fu) | (((ha >> (j4 & 6u)) & 0x03030303u) << 4); + return half4(ds * (float4(as_type(qa)) - 32.0f)); + } +}; + +template struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + a = W4::run(row, k, tab); + b = W4::run(row, k + 16, tab); + } +}; +// (a W8 sharing the block scale was measured slower than the two W4 of the generic one: 5.5 vs 3.5-4.2 ms) +template <> struct W8 { + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + const P bl = row + (k >> 8) * 98; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint aux = rd32(bl + 2 + 64 + 4 * ib); + const float sc = rdh(bl) * fma(float(aux >> 28), 0.5f, 0.25f); + const P qs = bl + 2 + ((k & 255) >> 2); + if (mm_probe == 3) { // bench probe: no threadgroup-table lookups (wrong values) + a = w4_signed(qs[0] * 0x01010101u, aux, sc); + b = w4_signed(qs[4] * 0x01010101u, aux >> 7, sc); + return; + } + const uint j0 = (aux >> (7 * l)) & 127, j1 = (aux >> (7 * (l + 2))) & 127; + a = w4_signed(g_iq3xxs(tab, qs[0]), sign4_iq2xs(j0, hh), sc); + b = w4_signed(g_iq3xxs(tab, qs[4]), sign4_iq2xs(j1, hh), sc); + } +}; +template <> struct W8 { + template + static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { + const int ib = (k & 255) >> 5, l = (k & 31) >> 3; + const P p = bl + 2 + 8 * ib; + const uint a1 = rd32(p + 4); + const float sc = d * fma(float(a1 >> 28), 0.25f, 0.125f); + const uint i0 = p[l], i1 = p[l + 2]; + const uint j0 = (a1 >> (7 * l)) & 127, j1 = (a1 >> (7 * (l + 2))) & 127; + threadgroup const uint* th = tab + hh; + a = w4_signed(ctab ? uint(iq2xxs_grid[i0] >> (32 * hh)) : th[2 * i0], sign4_iq2xs(j0, hh), sc); + b = w4_signed(ctab ? uint(iq2xxs_grid[i1] >> (32 * hh)) : th[2 * i1], sign4_iq2xs(j1, hh), sc); + } + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + const P bl = row + (k >> 8) * 66; + block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b); + } +}; +template <> struct W8 { + static inline void block(device const uchar* bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { + const int ib = (k & 255) >> 5; + device const ushort* q = (device const ushort*)(bl + 2) + ((k & 255) >> 3); + const uint v0 = q[0], v1 = q[2]; + const uint scb = bl[66 + ib]; + const float sc0 = d * fma(float(scb & 15), 0.25f, 0.125f); // l < 2 (k % 32 < 16): runs l, l + 2 take nibbles 0, 1 + const float sc1 = d * fma(float(scb >> 4), 0.25f, 0.125f); + threadgroup const uint* th = tab + hh; + a = w4_signed(ctab ? uint(iq2xs_grid[v0 & 511] >> (32 * hh)) : th[2 * (v0 & 511)], sign4_iq2xs(v0 >> 9, hh), sc0); + b = w4_signed(ctab ? uint(iq2xs_grid[v1 & 511] >> (32 * hh)) : th[2 * (v1 & 511)], sign4_iq2xs(v1 >> 9, hh), sc1); + } + static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 74; + block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b); + } +}; +template <> struct W8 { + // grid words read through tab + hh and the sign byte as s >> 4 * hh (not a select): 8% off mm_gu3 at T 2048; + // the same for IQ2_XXS / IQ2_XS, whose sign4_iq2xs select stays (a shift form there measured 10% slower) + template + static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { + const int ib = (k & 255) >> 5, l = (k & 31) >> 3; + threadgroup const uint* th = tab + hh; + const uint sh = uint(hh) << 2; + const uint qh = bl[66 + ib], scb = bl[74 + ib]; + const P p = bl + 2 + ((k & 255) >> 3); + const uint i0 = uint(p[0]) | ((qh << (8 - 2 * l)) & 0x300); + const uint i1 = uint(p[2]) | ((qh << (4 - 2 * l)) & 0x300); + const uint s0 = p[32], s1 = p[34]; + const float sc0 = d * fma(float(scb & 15), 0.25f, 0.125f); + const float sc1 = d * fma(float(scb >> 4), 0.25f, 0.125f); + a = w4_signed(ctab ? uint(iq2s_grid[i0] >> (32 * hh)) : th[2 * i0], s0 >> sh, sc0); + b = w4_signed(ctab ? uint(iq2s_grid[i1] >> (32 * hh)) : th[2 * i1], s1 >> sh, sc1); + } + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + const P bl = row + (k >> 8) * 82; + block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b); + } +}; +template <> struct W8 { + template + static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { + const int ib = k >> 5; + const uint qh = bl[66 + ib], scb = bl[106 + (ib >> 1)]; + const float sc = d * fma(float((scb >> (4 * (ib & 1))) & 15), 2.0f, 1.0f); + const P qs = bl + 2 + (k >> 2); + const int sh0 = 8 - ((k & 31) >> 2); + const uint i0 = uint(qs[0]) | ((qh << sh0) & 256u); + const uint i1 = uint(qs[4]) | ((qh << (sh0 - 4)) & 256u); + const P sp = bl + 74 + (k >> 3); + const uint s0 = sp[0], s1 = sp[2]; + if (mm_probe == 5) { // bench probe: the loads and the scale, no grid lookup / sign math (wrong values) + a = half4(half(float(i0 ^ s0) * sc)); + b = half4(half(float(i1 ^ s1) * sc)); + return; + } + a = w4_signed(g_iq3s(tab, i0), s0 >> (4 * hh), sc); + b = w4_signed(g_iq3s(tab, i1), s1 >> (4 * hh), sc); + } + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + if (mm_probe == 7) { + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint4 r = *(device const uint4*)((ulong(row) + ulong(k >> 5) * 16) & ~ulong(15)); + const uint qhv = r.z & 0xFF, scv = (r.z >> 8) & 0xFF; + const float scr = float(as_type(ushort(r.w & 0xFFFF))) * float(1 + 2 * ((scv >> (4 * (ib & 1))) & 15)); + const uint j0 = ((r.x >> (8 * (l & 1))) & 0xFF) | ((qhv << (8 - 2 * l)) & 256); + const uint j1 = ((r.x >> (16 + 8 * (l & 1))) & 0xFF) | ((qhv << (7 - 2 * l)) & 256); + a = w4_signed(tab[j0], (r.y >> (8 * hh)) & 0xFF, scr); + b = w4_signed(tab[j1], (r.y >> (16 + 8 * hh)) & 0xFF, scr); + return; + } + const P bl = row + (k >> 8) * 110; + block(bl, k & 255, rdh(bl), tab, (k & 7) >> 2, a, b); + } +}; + +// W8 for the K-quants and Q5_1: the lane's runs k and k + 16 (k % 32 < 16) share the block's scales / mins / high +// bits; the 4-byte code runs read as one uint where the block layout keeps them 4-aligned (Q4_K 144-, Q5_K 176-, +// Q5_1 24-byte blocks: rows of whole blocks start 4-aligned). The same per-value arithmetic as W4. +static inline float4 nib4(uint v, int hi) { const uint s = hi ? (v >> 4) : v; return float4(as_type(s & 0x0F0F0F0Fu)); } +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 144; + const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31; + float s, mm; + scale_min_k4(2 * c + hi, bl + 4, s, mm); + const float d1 = rdh(bl) * s, m1 = rdh(bl + 2) * mm; + device const uint* q = (device const uint*)(bl + 16 + 32 * c + l); + const float4 qa = nib4(q[0], hi), qb = nib4(q[4], hi); + a = half4(d1 * qa - m1); + b = half4(d1 * qb - m1); + } +}; +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 176; + const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31; + float s, mm; + scale_min_k4(2 * c + hi, bl + 4, s, mm); + const float d1 = rdh(bl) * s, m1 = rdh(bl + 2) * mm; + device const uint* qh = (device const uint*)(bl + 16 + l); + device const uint* ql = (device const uint*)(bl + 48 + 32 * c + l); + const uint sh = uint(2 * c + hi); + const float4 ha = float4(as_type((qh[0] >> sh) & 0x01010101u)) * 16.0f; + const float4 hb = float4(as_type((qh[4] >> sh) & 0x01010101u)) * 16.0f; + a = half4(d1 * (nib4(ql[0], hi) + ha) - m1); + b = half4(d1 * (nib4(ql[4], hi) + hb) - m1); + } +}; +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 5) * 24; + const float d = rdh(bl), m = rdh(bl + 2); + const uint hb = rd32(bl + 4); + const int j = k & 15; // runs at j and j + 16: the low nibbles of qs[j..j+3], then the high ones + const uint v = *(device const uint*)(bl + 8 + j); + float4 wa, wb; + for (int c = 0; c < 4; ++c) { + const uint qv = (v >> (8 * c)) & 0xFFu; + wa[c] = float(int(qv & 15) | int(((hb >> (j + c)) & 1) << 4)) * d + m; + wb[c] = float(int(qv >> 4) | int(((hb >> (j + 16 + c)) & 1) << 4)) * d + m; + } + a = half4(wa); + b = half4(wb); + } +}; +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 210; + const uint j = uint(k & 255), j4 = j >> 4; // l = j & 31 < 16; the pair at l + 16 + device const uchar* ql = bl + (((j >> 1) & 64u) | (j & 63u)); + device const uchar* qh = bl + 128u + (((j >> 2) & 32u) | (j & 31u)); + const float d = rdh(bl + 208); + device const char* sc = (device const char*)(bl + 192u + j4); + const float dsa = d * float(sc[0]), dsb = d * float(sc[1]); + const uint la = rd32(ql), lb = rd32(ql + 16); + const uint ha = rd32(qh), hb = rd32(qh + 16); + const uint sl = j4 & 4u, shh = j4 & 6u; + const uint qa = ((la >> sl) & 0x0F0F0F0Fu) | (((ha >> shh) & 0x03030303u) << 4); + const uint qb = ((lb >> sl) & 0x0F0F0F0Fu) | (((hb >> shh) & 0x03030303u) << 4); + a = half4(dsa * (float4(as_type(qa)) - 32.0f)); + b = half4(dsb * (float4(as_type(qb)) - 32.0f)); + } +}; + +template +kernel void mm_gu3(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], + device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float ot[SG][TM * TN]; + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN; + device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride; + device const uchar* g0 = E + a.gate_off + ulong(n0) * a.gu_row; + device const uchar* u0 = E + a.up_off + ulong(n0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + auto rw = op.template get_right_input_cooperative_tensor(); + auto tA = tX.template slice(0, row0); + auto accg = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + auto accu = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + if (coop_k4 && TK == 32 && mm_probe != 1) { + const uint16_t half_cap = rw.get_capacity() / 2; + const auto ix0 = rw.get_multidimensional_index(uint16_t(0)); + const int kl = ix0[0], hh = (kl >> 2) & 1; + const ulong roff0 = ulong(ix0[1]) * a.gu_row; + device const uchar* g0_l = g0 + roff0; + device const uchar* u0_l = u0 + roff0; + const uint step8 = 8u * a.gu_row; + // IQ2 formats: a 256 super-block at a time, its d read once; the lane's hh is fixed for the whole K loop + constexpr bool sb = F == T_IQ2_S || F == T_IQ2_XXS || F == T_IQ2_XS; + for (int kb = 0; kb < int(a.n_embd); kb += (sb ? 256 : TK)) { + const uint sboff = uint(kb >> 8) * fmt::BB; + float dg[TN * TK / 256], du[TN * TK / 256]; + if (sb) { +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap / 4; ++i) { + dg[i] = rdh(g0_l + i * step8 + sboff); + du[i] = rdh(u0_l + i * step8 + sboff); + } + } + for (int ki = 0; ki < (sb ? 256 : TK); ki += TK) { + const int k0 = kb + ki; + tA = tX.template slice(k0, row0); + const uint boff = F == T_IQ2_XS ? sboff : (fmt::QK == 256 ? uint(k0 >> 8) * fmt::BB : (fmt::QK == 64 ? uint(k0 >> 6) * fmt::BB : 0u)); + const int j_in = fmt::QK == 256 ? ((k0 + kl) & 255) : (fmt::QK == 64 ? ((k0 + kl) & 63) : (k0 + kl)); + device const uchar* rgw = fmt::QK >= 64 ? g0_l + boff : g0_l; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, rgw += step8) { + half4 wg, wg2; + if (F == T_IQ2_S) W8::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2); + else if (F == T_IQ2_XXS) W8::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2); + else if (F == T_IQ2_XS) W8::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2); + else if (F == T_IQ3_S) W8::block(rgw, j_in, rdh(rgw), tab, hh, wg, wg2); + else W8::run(rgw, j_in, tab, wg, wg2); + rw[i] = wg.x; rw[i + 1] = wg.y; rw[i + 2] = wg.z; rw[i + 3] = wg.w; + const uint16_t j = i + half_cap; + rw[j] = wg2.x; rw[j + 1] = wg2.y; rw[j + 2] = wg2.z; rw[j + 3] = wg2.w; + } + op.run(tA, rw, accg); + device const uchar* ruw = fmt::QK >= 64 ? u0_l + boff : u0_l; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, ruw += step8) { + half4 wu, wu2; + if (F == T_IQ2_S) W8::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2); + else if (F == T_IQ2_XXS) W8::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2); + else if (F == T_IQ2_XS) W8::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2); + else if (F == T_IQ3_S) W8::block(ruw, j_in, rdh(ruw), tab, hh, wu, wu2); + else W8::run(ruw, j_in, tab, wu, wu2); + rw[i] = wu.x; rw[i + 1] = wu.y; rw[i + 2] = wu.z; rw[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rw[j] = wu2.x; rw[j + 1] = wu2.y; rw[j + 2] = wu2.z; rw[j + 3] = wu2.w; + } + op.run(tA, rw, accu); + } + } + } else + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + const uint16_t cap = rw.get_capacity(), half_cap = cap / 2; + tA = tX.template slice(k0, row0); + if (!coop_k4) { + for (uint16_t i = 0; i < cap; ++i) { + const auto ix = rw.get_multidimensional_index(i); + const int k = k0 + ix[0], n = ix[1]; + rw[i] = W4::run(g0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3]; + } + op.run(tA, rw, accg); + for (uint16_t i = 0; i < cap; ++i) { + const auto ix = rw.get_multidimensional_index(i); + const int k = k0 + ix[0], n = ix[1]; + rw[i] = W4::run(u0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3]; + } + op.run(tA, rw, accu); + } else { +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rw.get_multidimensional_index(i); + const auto iy = rw.get_multidimensional_index(i + half_cap); + const int k = k0 + ix[0], n = ix[1]; + device const uchar* rgw = g0 + ulong(n) * a.gu_row; + half4 wg, wg2; + if (mm_probe == 1) { + wg = half4(half(n & 7)); wg2 = half4(half(k & 7)); + } else if (TK == 32) { + W8::run(rgw, k, tab, wg, wg2); + } else { + const int k2 = k0 + iy[0]; + wg = W4::run(rgw, k, tab); + wg2 = W4::run(g0 + ulong(iy[1]) * a.gu_row, k2, tab); + } + rw[i] = wg.x; rw[i + 1] = wg.y; rw[i + 2] = wg.z; rw[i + 3] = wg.w; + const uint16_t j = i + half_cap; + rw[j] = wg2.x; rw[j + 1] = wg2.y; rw[j + 2] = wg2.z; rw[j + 3] = wg2.w; + } + op.run(tA, rw, accg); +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rw.get_multidimensional_index(i); + const auto iy = rw.get_multidimensional_index(i + half_cap); + const int k = k0 + ix[0], n = ix[1]; + device const uchar* ruw = u0 + ulong(n) * a.gu_row; + half4 wu, wu2; + if (mm_probe == 1) { + wu = half4(half(n & 7)); wu2 = half4(half(k & 7)); + } else if (TK == 32) { + W8::run(ruw, k, tab, wu, wu2); + } else { + const int k2 = k0 + iy[0]; + wu = W4::run(ruw, k, tab); + wu2 = W4::run(u0 + ulong(iy[1]) * a.gu_row, k2, tab); + } + rw[i] = wu.x; rw[i + 1] = wu.y; rw[i + 2] = wu.z; rw[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rw[j] = wu2.x; rw[j + 1] = wu2.y; rw[j + 2] = wu2.z; rw[j + 3] = wu2.w; + } + op.run(tA, rw, accu); + } + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = accg[i]; + accg[i] = g / (1.0f + exp(-g)) * accu[i]; + } + tensor, tensor_inline> tO(ot[sg], dextents(TN, TM)); + accg.store(tO); + simdgroup_barrier(mem_flags::mem_threadgroup); + if ((F == T_IQ2_XS || F == T_IQ3_XXS || F == T_IQ2_XXS || F == T_Q2_0 || F == T_IQ2_S || F == T_IQ3_S) && 128 % TN == 0 && (TM * TN) % 128 == 0) { + // a lane stores 4 adjacent outputs a pass, 128 a pass over the simdgroup + constexpr int RP = 128 / TN; // rows a pass + const int mr = lane / (TN / 4), f = (lane % (TN / 4)) * 4; + threadgroup const float4* ot_p = (threadgroup const float4*)(&ot[sg][mr * TN + f]); + device half* hs_p = HS + ulong(row0 + mr) * a.n_ff + n0 + f; + const uint hs_step = uint(RP) * a.n_ff; +#pragma clang loop unroll(full) + for (int m = 0; m < TM / RP; ++m, hs_p += hs_step) + *(device half4*)hs_p = half4(ot_p[m * (RP * TN / 4)]); + } else + for (int it = lane; it < TM * TN; it += 32) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + m) * a.n_ff + n0 + f] = half(ot[sg][it]); + } +} + +template +kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS [[buffer(1)]], + device float* YE [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], + device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN; + device const uchar* w0 = (kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.d_stride) + a.down_off + ulong(n0) * a.d_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tH((device half*)HS, dextents(a.n_ff, 1 << 30)); + auto rw = op.template get_right_input_cooperative_tensor(); + auto tA = tH.template slice(0, row0); + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + const auto ix0 = rw.get_multidimensional_index(uint16_t(0)); + const int kl = ix0[0]; + device const uchar* w0_l = w0 + ulong(ix0[1]) * a.d_row; + const uint step8 = 8u * a.d_row; + if (coop_k4 && F == T_Q2_0 && TK == 32) { + const uint16_t half_cap = rw.get_capacity() / 2; + const uint kl4 = uint(kl) >> 2; + for (int k0 = 0; k0 < int(a.n_ff); k0 += 64, w0_l += 18) { + device const uchar* rdw = w0_l; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) { + const half4 w = W4::run_q<2>(rdw, kl4); + const half4 w2 = W4::run_q<6>(rdw, kl4); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + const uint16_t j = i + half_cap; + rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; + } + tA = tH.template slice(k0, row0); + op.run(tA, rw, acc); + rdw = w0_l; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) { + const half4 w = W4::run_q<10>(rdw, kl4); + const half4 w2 = W4::run_q<14>(rdw, kl4); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + const uint16_t j = i + half_cap; + rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; + } + tA = tH.template slice(k0 + 32, row0); + op.run(tA, rw, acc); + } + } else + for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) { + if (!coop_k4) { + for (uint16_t i = 0; i < rw.get_capacity(); ++i) { + const auto ix = rw.get_multidimensional_index(i); + const int k = k0 + ix[0]; + rw[i] = W4::run(w0 + ulong(ix[1]) * a.d_row, k & ~3, tab)[k & 3]; + } + } else if (TK == 32) { + // the paired runs (element i + cap / 2 is (k + 16, n), the COOP_K4 layout): one W8 a pair + const uint16_t half_cap = rw.get_capacity() / 2; + const uint boff = fmt::QK == 64 ? uint(k0 >> 6) * fmt::BB : (fmt::QK == 32 ? uint(k0 >> 5) * fmt::BB : 0u); + const int j_in = fmt::QK == 64 ? ((k0 + kl) & 63) : (fmt::QK == 32 ? kl : (k0 + kl)); + device const uchar* rdw = w0_l + boff; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) { + half4 w, w2; + W8::run(rdw, j_in, tab, w, w2); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + const uint16_t j = i + half_cap; + rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rw.get_capacity(); i += 4) { + const auto ix = rw.get_multidimensional_index(i); + const half4 w = W4::run(w0 + ulong(ix[1]) * a.d_row, k0 + ix[0], tab); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + } + tA = tH.template slice(k0, row0); + op.run(tA, rw, acc); + } + tensor, tensor_inline> tY(YE, dextents(a.n_embd, 1 << 30)); + acc.store(tY.slice(n0, row0)); +} + +#define MMGU3_(F, M, N, K, S) template [[host_name("mm_gu3_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu3) mm_gu3; +#define MMDN3_(F, M, N, K, S) template [[host_name("mm_down3_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down3) mm_down3; +#define MMGU3(F, M, N, K, S) MMGU3_(F, M, N, K, S) +#define MMDN3(F, M, N, K, S) MMDN3_(F, M, N, K, S) +#define MM3_ONE(F, M, N, K, S) MMGU3(F, M, N, K, S) MMDN3(F, M, N, K, S) +#define MM3_GU(F) MMGU3(F, 16, 64, 32, 2) MMGU3(F, 16, 32, 32, 4) MMGU3(F, 32, 64, 32, 2) MMGU3(F, 16, 64, 64, 2) \ + MMGU3(F, 16, 32, 16, 2) MMGU3(F, 32, 32, 32, 2) MMGU3(F, 16, 64, 32, 1) MMGU3(F, 16, 64, 32, 4) MMGU3(F, 16, 32, 16, 4) MMGU3(F, 16, 32, 32, 5) MMGU3(F, 16, 32, 16, 5) MMGU3(F, 16, 16, 32, 4) MMGU3(F, 16, 16, 32, 8) MMGU3(F, 16, 16, 16, 8) +#define MM3_DN(F) MMDN3(F, 16, 64, 32, 2) MMDN3(F, 16, 32, 32, 4) MMDN3(F, 32, 64, 32, 2) MMDN3(F, 16, 64, 64, 2) \ + MMDN3(F, 16, 32, 16, 2) MMDN3(F, 32, 32, 32, 2) MMDN3(F, 16, 64, 32, 1) MMDN3(F, 16, 64, 32, 4) MMDN3(F, 16, 64, 16, 2) MMDN3(F, 16, 32, 16, 4) MMDN3(F, 16, 64, 32, 5) MMDN3(F, 16, 64, 16, 4) +MM3_GU(T_Q4_K) MM3_GU(T_Q5_K) MM3_GU(T_Q6_K) MM3_GU(T_Q8_0) MM3_GU(T_IQ4_XS) MM3_GU(T_Q2_0) MM3_GU(T_IQ3_XXS) MM3_GU(T_IQ2_XXS) MM3_GU(T_IQ2_XS) MM3_GU(T_IQ2_S) MM3_GU(T_IQ3_S) +MM3_DN(T_Q2_0) MM3_DN(T_IQ4_NL) MM3_DN(T_Q5_1) MM3_DN(T_Q8_0) + +#ifndef MM_V3_ONLY +// mm_gu4: mm_gu3 with each decoded right operand (TN x TK weights of gate and of up) applied to MS 16-row token +// tiles (a tile of 16 * MS rows from the plan, tm = 16 * MS): the decode, which bounds mm_gu3, is shared by MS times +// as many token rows (more padding rows per expert, but the matrix ops are cheap next to the decode). The same +// per-row arithmetic as mm_gu3 (each 16-row op is mm_gu3's): bitwise for the rows both compute. COOP_K4 layout only. +template +kernel void mm_gu4(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int TM = 16; + threadgroup float ot[SG][TM * TN]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN; + device const uchar* g0 = W + a.gate_off + ulong(e) * a.gu_stride + ulong(n0) * a.gu_row; + device const uchar* u0 = W + a.up_off + ulong(e) * a.gu_stride + ulong(n0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + auto rg = op.template get_right_input_cooperative_tensor(); + auto ru = op.template get_right_input_cooperative_tensor(); + auto tA0 = tX.template slice(0, row0); + auto tA1 = tX.template slice(0, row0 + TM); +#define MM4_DST op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>() + auto accg0 = MM4_DST; + auto accu0 = MM4_DST; + auto accg1 = MM4_DST; + auto accu1 = MM4_DST; +#undef MM4_DST + using DT = decltype(accg0); + for (uint16_t i = 0; i < accg0.get_capacity(); ++i) + if (accg0.is_valid_element(i)) { accg0[i] = 0.0f; accu0[i] = 0.0f; accg1[i] = 0.0f; accu1[i] = 0.0f; } + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + const uint16_t cap = rg.get_capacity(), half_cap = cap / 2; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rg.get_multidimensional_index(i); + const int k = k0 + ix[0], n = ix[1]; + device const uchar* rgw = g0 + ulong(n) * a.gu_row; + device const uchar* ruw = u0 + ulong(n) * a.gu_row; + half4 wg, wu, wg2, wu2; + W8::run(rgw, k, tab, wg, wg2); + W8::run(ruw, k, tab, wu, wu2); + rg[i] = wg.x; rg[i + 1] = wg.y; rg[i + 2] = wg.z; rg[i + 3] = wg.w; + ru[i] = wu.x; ru[i + 1] = wu.y; ru[i + 2] = wu.z; ru[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rg[j] = wg2.x; rg[j + 1] = wg2.y; rg[j + 2] = wg2.z; rg[j + 3] = wg2.w; + ru[j] = wu2.x; ru[j + 1] = wu2.y; ru[j + 2] = wu2.z; ru[j + 3] = wu2.w; + } + tA0 = tX.template slice(k0, row0); + op.run(tA0, rg, accg0); + op.run(tA0, ru, accu0); + if (MS > 1) { + tA1 = tX.template slice(k0, row0 + TM); + op.run(tA1, rg, accg1); + op.run(tA1, ru, accu1); + } + } + for (int s = 0; s < MS; ++s) { + thread DT& ag = s == 0 ? accg0 : accg1; + thread DT& au = s == 0 ? accu0 : accu1; + for (uint16_t i = 0; i < ag.get_capacity(); ++i) + if (ag.is_valid_element(i)) { + const float g = ag[i]; + ag[i] = g / (1.0f + exp(-g)) * au[i]; + } + tensor, tensor_inline> tO(ot[sg], dextents(TN, TM)); + simdgroup_barrier(mem_flags::mem_threadgroup); + ag.store(tO); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int it = lane; it < TM * TN; it += 32) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + s * TM + m) * a.n_ff + n0 + f] = half(ot[sg][it]); + } + } +} +#define MMGU4_(F, M, N, K, S) template [[host_name("mm_gu4_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu4) mm_gu4; +#define MMGU4(F, M, N, K, S) MMGU4_(F, M, N, K, S) +#define MM4_GU(F) MMGU4(F, 2, 32, 32, 4) MMGU4(F, 2, 32, 32, 2) MMGU4(F, 2, 16, 32, 4) MMGU4(F, 1, 32, 32, 4) +MM4_GU(T_Q2_0) MM4_GU(T_IQ3_XXS) MM4_GU(T_IQ2_XXS) MM4_GU(T_IQ2_XS) MM4_GU(T_IQ2_S) MM4_GU(T_IQ3_S) MM4_GU(T_IQ4_XS) + +// mm_gu5: mm_gu3 (COOP_K4 layout, TK 32) with the threadgroup's weight rows staged block by block (QK = 256 values a +// row, gate and up) into threadgroup memory by coalesced ushort loads, the lanes decoding from there: the decode's +// many small scattered device loads (2-6 a run) become threadgroup reads. Same decoded values and the same matrix +// ops as mm_gu3: bitwise. Formats with 256-value blocks and a templated W8 (IQ3_S, IQ3_XXS, IQ2_S). +template +kernel void mm_gu5(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], + device const uint* smap [[buffer(5)]], device const uchar* W1 [[buffer(6)]], + device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int TM = 16, TK = 32, BB = fmt::BB, ROWS = SG * TN, RU = BB / 2; + threadgroup float ot[SG][TM * TN]; + threadgroup uint tab[fmt::TAB]; + threadgroup ushort stg[2][ROWS * RU]; + load_tab(tab, tid, SG * 32); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, nt0 = int(tg.x) * ROWS, n0 = nt0 + sg * TN; + device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride; + device const uchar* gt = E + a.gate_off + ulong(nt0) * a.gu_row; + device const uchar* ut = E + a.up_off + ulong(nt0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + auto rg = op.template get_right_input_cooperative_tensor(); + auto ru = op.template get_right_input_cooperative_tensor(); + auto tA = tX.template slice(0, row0); + auto accg = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + auto accu = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + const int NB = int(a.n_embd) / 256; + for (int kb = 0; kb < NB; ++kb) { + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < ROWS * RU; idx += SG * 32) { + const int r = idx / RU, w = idx - r * RU; + stg[0][idx] = ((device const ushort*)(gt + ulong(r) * a.gu_row + kb * BB))[w]; + stg[1][idx] = ((device const ushort*)(ut + ulong(r) * a.gu_row + kb * BB))[w]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int ks = 0; ks < 8; ++ks) { + const int k0 = kb * 256 + ks * TK; + const uint16_t half_cap = rg.get_capacity() / 2; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rg.get_multidimensional_index(i); + const int kl = ks * TK + ix[0], n = sg * TN + ix[1]; + // the block's row in the stage; W8 indexes block (k >> 8) = 0 of it with k = kl + threadgroup const uchar* rgw = (threadgroup const uchar*)(stg[0] + n * RU); + threadgroup const uchar* ruw = (threadgroup const uchar*)(stg[1] + n * RU); + half4 wg, wu, wg2, wu2; + W8::run(rgw, kl, tab, wg, wg2); + W8::run(ruw, kl, tab, wu, wu2); + rg[i] = wg.x; rg[i + 1] = wg.y; rg[i + 2] = wg.z; rg[i + 3] = wg.w; + ru[i] = wu.x; ru[i + 1] = wu.y; ru[i + 2] = wu.z; ru[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rg[j] = wg2.x; rg[j + 1] = wg2.y; rg[j + 2] = wg2.z; rg[j + 3] = wg2.w; + ru[j] = wu2.x; ru[j + 1] = wu2.y; ru[j + 2] = wu2.z; ru[j + 3] = wu2.w; + } + tA = tX.template slice(k0, row0); + op.run(tA, rg, accg); + op.run(tA, ru, accu); + } + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = accg[i]; + accg[i] = g / (1.0f + exp(-g)) * accu[i]; + } + tensor, tensor_inline> tO(ot[sg], dextents(TN, TM)); + accg.store(tO); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int it = lane; it < TM * TN; it += 32) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + m) * a.n_ff + n0 + f] = half(ot[sg][it]); + } +} +#define MMGU5_(F, N, S) template [[host_name("mm_gu5_" #F "_" #N "_" #S)]] kernel decltype(mm_gu5) mm_gu5; +#define MMGU5(F, N, S) MMGU5_(F, N, S) +MMGU5(T_IQ3_S, 32, 2) MMGU5(T_IQ3_S, 32, 1) MMGU5(T_IQ3_S, 16, 4) MMGU5(T_IQ3_XXS, 32, 2) MMGU5(T_IQ3_XXS, 32, 1) +MMGU5(T_IQ3_XXS, 16, 4) MMGU5(T_IQ2_S, 32, 2) MMGU5(T_IQ2_S, 32, 1) MMGU5(T_IQ2_S, 16, 4) MMGU5(T_IQ2_XXS, 32, 2) MMGU5(T_IQ2_XXS, 16, 4) +#endif +#endif diff --git a/results/round39/astra_r39_qsa/baseline-kernels/prompt.metal b/results/round39/astra_r39_qsa/baseline-kernels/prompt.metal new file mode 100644 index 0000000..a3fd96b --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/prompt.metal @@ -0,0 +1,690 @@ +// metal/kernels/prompt.metal - the elementwise steps of a prompt chunk around the GEMMs (after quant/units.metal): +// +// route_topk per token: softmax over the 512 logits, top 10 (ties: the lower expert), weights renormalised +// over the selected with a 2^-14 floor - moe_route's rule; one simdgroup a token +// gr_epi_down after the HC down GEMM (V = xn . [w_down; w_inject], T x (hc_lr + hc)): lo_h = half(silu(v / hc)), +// inject = v +// gr_epi_up after the HC up GEMM (G = lo . w_up, T x hc*n): mixed = mean_c xn[c] * sigmoid(G[c]) -> float and +// half (the next GEMMs' operand) +// swiglu_h h = half(silu(g) * u) +// gemm_hc_down, gemm_swiglu_up (after gemm_hc_up) the HC down and shared-expert up GEMMs with gr_epi_down / +// swiglu_h as register epilogues +// moe_out bo = y_routed + sigmoid(sg) * y_shared (prompt_chain.inc; the prompt chunk folds it into the next +// read: gr_norm_h_moe, gr_apply_moe, ple_moe_gate_norm; with mm_combine folded in too: gr_norm_h_moec, +// gr_apply_moec) + +kernel void route_topk(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]], + device float* w [[buffer(2)]], constant uint& n_tok_packed [[buffer(3)]], + device const uint* smap [[buffer(4)]], constant float& prompt_moce [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + const uint n_tok = n_tok_packed & 0xFFFFu; + const int n_exp = (n_tok_packed >> 16) ? int(n_tok_packed >> 16) : 512; + const uint t = tg.x * 8 + sg; + if (t >= n_tok) return; + device const float* l = logits + ulong(t) * ulong(n_exp); + float v[16]; + float mx = -INFINITY; + float sum = 0.0f; + if (n_exp == 512) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] = l[lane + 32 * i]; + mx = max(mx, v[i]); + } + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] = exp(v[i] - mx); + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + } else { + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + v[i] = idx < n_exp ? l[idx] : -INFINITY; + } + for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + for (int i = 0; i < 16; ++i) { + v[i] = (lane + 32 * i < n_exp) ? exp(v[i] - mx) : 0.0f; + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + for (int i = 0; i < 16; ++i) { + if (lane + 32 * i >= n_exp) { + v[i] = -INFINITY; + } else { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + } + } + const int tail_exact = int(prompt_moce) / 10000; + const float rem_moce = prompt_moce - float(tail_exact * 10000); + const int body_prot = int(rem_moce) / 100; + const float body_bonus = rem_moce - float(body_prot * 100); + const bool in_tail = tail_exact > 0 && (t + uint(tail_exact) >= n_tok); + const int p_prot = in_tail ? 10 : body_prot; + const float p_bonus = in_tail ? 0.0f : body_bonus; + const float moce_scale = p_bonus > 0.0f ? exp(min(p_bonus, 40.0f)) : 1.0f; + const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f; + ushort scaled_mask = 0; + float selected = 0.0f, ssum = 0.0f; + for (int rank = 0; rank < 10; ++rank) { + if (rank == p_prot && moce_scale > 1.0f) { + if (n_exp == 512) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if (v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) { + v[i] *= moce_scale; + scaled_mask |= ushort(1u << i); + } + } + } else { + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if (idx < n_exp && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) { + v[i] *= moce_scale; + scaled_mask |= ushort(1u << i); + } + } + } + } + float best = v[0]; + int ex = lane; + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + float true_prob = best; + if (rank >= p_prot && moce_scale > 1.0f && best > 0.0f) { + const ushort win_lane = ushort(ex & 31); + 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 * rcp_moce; + } + if ((ex & 31) == lane) { + // v stays in registers: no runtime index into it + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) v[i] = i == (ex >> 5) ? -INFINITY : v[i]; + ids[t * 10 + rank] = ex; + ssum += true_prob; + } + if (rank == lane) selected = true_prob; + } + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + if (prompt_moce < 0.0f) { + const int skip_tail = int(-prompt_moce); + const float tau = (-prompt_moce) - float(skip_tail); + const bool skip_active = tau > 0.0f && !(skip_tail > 0 && (t + uint(skip_tail) >= n_tok)); + if (skip_active) { + float kept_prob = (lane < 10) ? selected : 0.0f; + if (lane >= 3 && lane < 10) { + const int ex_r = ids[t * 10 + lane]; + if (smap[ex_r] == 0xFFFFFFFFu && selected < tau * ssum) { + ids[t * 10 + lane] = -1; + kept_prob = 0.0f; + } + } + selected = kept_prob; + ssum = kept_prob; + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + } + } + ssum = max(ssum, 6.103515625e-5f); + if (lane < 10) w[t * 10 + lane] = selected * (1.0f / ssum); +} + +struct EpiArgs { + uint n_tok; + uint n_embd; // 2560 + uint hc; // 4 + uint hc_lr; // 320 +}; + +// grid (T * (hc_lr + hc) / 256) +kernel void gr_epi_down(device const float* V [[buffer(0)]], device half* lo_h [[buffer(1)]], + device float* inject [[buffer(2)]], constant EpiArgs& a [[buffer(3)]], + uint gid [[thread_position_in_grid]]) { + const uint nr = a.hc_lr + a.hc, t = gid / nr, r = gid % nr; + if (t >= a.n_tok) return; + const float v = V[gid]; + if (r < a.hc_lr) { + const float z = v / float(a.hc); + lo_h[t * a.hc_lr + r] = half(z / (1.0f + exp(-z))); + } else { + inject[t * a.hc + r - a.hc_lr] = v; + } +} + +// grid (T * n_embd / 256) +kernel void gr_epi_up(device const float* G [[buffer(0)]], device const float* xn [[buffer(1)]], + device float* mixed [[buffer(2)]], device half* mixed_h [[buffer(3)]], + constant EpiArgs& a [[buffer(4)]], uint gid [[thread_position_in_grid]]) { + const uint n = a.n_embd, t = gid / n, d = gid % n; + if (t >= a.n_tok) return; + float m = 0.0f; + for (uint c = 0; c < a.hc; ++c) { + const float g = G[ulong(t) * a.hc * n + c * n + d]; + m = fma(xn[ulong(t) * a.hc * n + c * n + d], 1.0f / (1.0f + exp(-g)), m); + } + m = m / float(a.hc); + mixed[gid] = m; + mixed_h[gid] = half(m); +} + +kernel void swiglu_h(device const float* g [[buffer(0)]], device const float* u [[buffer(1)]], device half* h [[buffer(2)]], + constant uint& n [[buffer(3)]], uint gid [[thread_position_in_grid]]) { + if (gid >= n) return; + const float gv = g[gid]; + h[gid] = half(gv / (1.0f + exp(-gv)) * u[gid]); +} + +// grid (T * n_embd / 256); sg: T values +kernel void moe_out(device const float* yr [[buffer(0)]], device const float* ys [[buffer(1)]], + device const float* sg [[buffer(2)]], device float* bo [[buffer(3)]], constant uint2& a [[buffer(4)]], + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.y; + if (t >= a.x) return; + bo[gid] = fma(1.0f / (1.0f + exp(-sg[t])), ys[gid], yr[gid]); +} + +// ---- the HC read with fewer passes (after dense_mm.metal: GemmArgs, MetalPerformancePrimitives) --------------------- +// gr_norm_h: gr_norm with xn written as fp16 only (the GEMM operand; gemm_hc_up's epilogue reads it too). +// one threadgroup per (stream c, token t), 256 threads; args as gr_norm's (GrArgs, gr.metal) +kernel void gr_norm_h(device float* R [[buffer(0)]], device const float* bo [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device half* xnh [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint c = tg.x, t = tg.y, n = a.n_embd; + device float* r = R + (t * a.hc + c) * n; + const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; + float rv[10], ss = 0.0f; + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + float v = r[d]; + if (a.apply) { + v = fma(bo[t * n + d], g, v); + r[d] = v; + } + rv[i] = v; + ss = fma(v, v, ss); + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; +#pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device half* o = xnh + (t * a.hc + c) * n; + device const float* w = w_norm + c * n; + for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[i] * w[d] * rs); +} + +// gr_norm_h's tail for one token's hc streams: the sums of squares ss[c] (per thread) to the norm, xnh = R' * w * rs +inline void gr_norm_h_out(thread const float (*rv)[10], device const float* w_norm, device half* xnh, uint t, uint hc, + uint n, float eps, thread const float* ss, threadgroup float (*part)[8], ushort tid, + ushort sg, ushort lane) { + for (uint c = 0; c < hc; ++c) { + const float s = simd_sum(ss[c]); + if (lane == 0) part[c][sg] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint c = 0; c < hc; ++c) { + float tot = 0.0f; +#pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[c][i]; + const float rs = rsqrt(tot / float(n) + eps); + device half* o = xnh + (t * hc + c) * n; + device const float* w = w_norm + c * n; + for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[c][i] * w[d] * rs); + } +} + +// stream-wise reduction keeps one residual stream plus bo live, rather than all four residual streams. +inline void gr_norm_h_stream_out(thread const float* rv, device const float* w, device half* o, uint n, float eps, + float ss, threadgroup float* part, ushort tid, ushort sg, ushort lane) { + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; +#pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + eps); + if (n == 2560) { +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const uint d = tid + i * 256; + o[d] = half(rv[i] * w[d] * rs); + } + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[i] * w[d] * rs); + } +} + +// gr_norm_h4: gr_norm_h with one threadgroup per token over all hc streams, so bo is read once a token rather than +// hc times; R', xnh bitwise gr_norm_h's. grid (1, T), 256 threads; hc <= 4 +kernel void gr_norm_h4(device float* R [[buffer(0)]], device const float* bo [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device half* xnh [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + const uint t = tg.y, n = a.n_embd, hc = a.hc; + device float* r = R + t * hc * n; + if (!a.apply) { + for (uint c = 0; c < hc; ++c) { + float rv[10], ss = 0.0f; + if (n == 2560) { + device const float* rc = r + c * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const float v = rc[i * 256]; + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + const float v = r[c * n + d]; + rv[i] = v; + ss = fma(v, v, ss); + } + } + gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); + } + return; + } + float bv[10]; + if (n == 2560) { + device const float* brow = bo + t * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) bv[i] = brow[i * 256]; + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) bv[i] = bo[t * n + d]; + } + for (uint c = 0; c < hc; ++c) { + const float g = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); + float rv[10], ss = 0.0f; + if (n == 2560) { + device float* rc = r + c * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const float v = fma(bv[i], g, rc[i * 256]); + rc[i * 256] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + const float v = fma(bv[i], g, r[c * n + d]); + r[c * n + d] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } + gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); + } +} + +// gr_norm_h_moe: gr_norm_h with the sublayer output bo formed on the fly as moe_out - bo = fma(sigmoid(sg[t]), ys, yr) +// - so the MoE's write never round-trips through bo; R', xnh bitwise moe_out + gr_norm_h. Always applies. +// One threadgroup per token t (all hc streams: bo is formed once a token), 256 threads; hc <= 4. +kernel void gr_norm_h_moe(device float* R [[buffer(0)]], device const float* yr [[buffer(1)]], + device const float* ys [[buffer(2)]], device const float* sgv [[buffer(3)]], + device const float* inj [[buffer(4)]], device const float* w_norm [[buffer(5)]], + device half* xnh [[buffer(6)]], constant GrArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + const uint t = tg.y, n = a.n_embd, hc = a.hc; + device float* r = R + t * hc * n; + float rv[4][10], g[4], ss[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (uint c = 0; c < hc; ++c) g[c] = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + const float bo = fma(gs, ys[t * n + d], yr[t * n + d]); + for (uint c = 0; c < hc; ++c) { + const float v = fma(bo, g[c], r[c * n + d]); + r[c * n + d] = v; + rv[c][i] = v; + ss[c] = fma(v, v, ss[c]); + } + } + gr_norm_h_out(rv, w_norm, xnh, t, hc, n, a.eps, ss, part, tid, sg, lane); +} + +// token t's routed slots for moe_comb: compacted valid YE row offsets (uint(r) * n) and their weights in threadgroup memory +inline uint moe_comb_slots(device const int* pos, device const float* srt_w, uint t, uint topk, uint n, + threadgroup uint* roff, threadgroup float* wk, ushort tid) { + int r = -1; + float w = 0.0f; + if (tid < topk) { + r = pos[t * topk + tid]; + if (r >= 0) w = srt_w[r]; + } + const bool valid = r >= 0; + const ushort prefix = simd_prefix_exclusive_sum(ushort(valid)); + const ushort cnt = simd_sum(ushort(valid)); + if (valid) { + roff[prefix] = uint(r) * n; + wk[prefix] = w; + } + if (tid == 0) roff[16] = cnt; + threadgroup_barrier(mem_flags::mem_threadgroup); + return roff[16]; +} + +// the routed experts' write at (t, d) as mm_combine (moe_mm.metal) forms it: sum_k w(t, k) * YE[pos(t, k)][d], k in order +inline float moe_comb(device const float* YE, threadgroup const uint* roff, threadgroup const float* wk, uint cnt, + uint d) { + float s = 0.0f; + for (uint k = 0; k < cnt; ++k) s = fma(wk[k], YE[roff[k] + d], s); + return s; +} + +// gr_norm_h_moec: gr_norm_h_moe with mm_combine folded in - yr formed from the expert rows YE (pos, srt_w: mm_place's) +// as it is read, so it is never written; R', xnh bitwise mm_combine + gr_norm_h_moe. grid (1, T), 256 threads; +// hc <= 4, topk <= 16 +kernel void gr_norm_h_moec(device float* R [[buffer(0)]], device const float* YE [[buffer(1)]], + device const int* pos [[buffer(2)]], device const float* srt_w [[buffer(3)]], + device const float* ys [[buffer(4)]], device const float* sgv [[buffer(5)]], + device const float* inj [[buffer(6)]], device const float* w_norm [[buffer(7)]], + device half* xnh [[buffer(8)]], constant GrArgs& a [[buffer(9)]], constant uint& topk [[buffer(10)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + threadgroup uint roff[17]; + threadgroup float wk[16]; + const uint t = tg.y, n = a.n_embd, hc = a.hc; + const uint cnt = moe_comb_slots(pos, srt_w, t, topk, n, roff, wk, tid); + device float* r = R + t * hc * n; + float bv[10]; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + if (n == 2560) { + // k outer: each slot's weight and row read once, its ten loads in flight together; s[i]'s fma order is moe_comb's + float s[10] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + for (uint k = 0; k < cnt; ++k) { + const float w = wk[k]; + device const float* yk = YE + roff[k] + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) s[i] = fma(w, yk[i * 256], s[i]); + } + device const float* yrow = ys + t * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) bv[i] = fma(gs, yrow[i * 256], s[i]); + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) bv[i] = fma(gs, ys[t * n + d], moe_comb(YE, roff, wk, cnt, d)); + } + for (uint c = 0; c < hc; ++c) { + const float g = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); + float rv[10], ss = 0.0f; + if (n == 2560) { + device float* rc = r + c * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const float v = fma(bv[i], g, rc[i * 256]); + rc[i * 256] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + const float v = fma(bv[i], g, r[c * n + d]); + r[c * n + d] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } + gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); + } +} + +// gr_apply_moe: the last layer's MoE write applied to R alone (R' of moe_out + gr_norm; its xn has no reader). +// grid (hc * n / 256, T), 256 threads +kernel void gr_apply_moe(device float* R [[buffer(0)]], device const float* yr [[buffer(1)]], + device const float* ys [[buffer(2)]], device const float* sgv [[buffer(3)]], + device const float* inj [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const uint ch = tg.x * 256 + tid, t = tg.y, n = a.n_embd; + if (ch >= a.hc * n) return; + const uint c = ch / n, d = ch - c * n; + const float g = 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))); + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + device float* r = R + t * a.hc * n + ch; + *r = fma(fma(gs, ys[t * n + d], yr[t * n + d]), g, *r); +} + +// gr_apply_moec: gr_apply_moe with mm_combine folded in (as gr_norm_h_moec); R' bitwise mm_combine + gr_apply_moe. +// grid (ceil(n / 256), T), 256 threads: a thread's d over all hc streams (yr formed once); topk <= 16 +kernel void gr_apply_moec(device float* R [[buffer(0)]], device const float* YE [[buffer(1)]], + device const int* pos [[buffer(2)]], device const float* srt_w [[buffer(3)]], + device const float* ys [[buffer(4)]], device const float* sgv [[buffer(5)]], + device const float* inj [[buffer(6)]], constant GrArgs& a [[buffer(7)]], + constant uint& topk [[buffer(8)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint roff[17]; + threadgroup float wk[16]; + const uint d = tg.x * 256 + tid, t = tg.y, n = a.n_embd, hc = a.hc; + const uint cnt = moe_comb_slots(pos, srt_w, t, topk, n, roff, wk, tid); + if (d >= n) return; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + const float bo = fma(gs, ys[t * n + d], moe_comb(YE, roff, wk, cnt, d)); + device float* r = R + t * hc * n + d; + for (uint c = 0; c < hc; ++c) r[c * n] = fma(bo, 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))), r[c * n]); +} + +// gemm_hc_up: G = lo . w_up^T with w_up's rows interleaved by dq_f16 (perm = hc: row d * hc + c), so a TN-column +// tile holds TN / hc outputs d with all their streams; the epilogue mixed[t][d] = mean_c xn[t][c][d] * +// sigmoid(G[t][d * hc + c]) -> fp16 and, if mixed is bound, fp32. G is never written. +// grid (hc * n / TN, ceil(T / TM)), 32 * SG threads. +template +inline void hc_up_tile(device half* A, device half* B, device const half* xnh, device float* mixed, + device half* mixh, constant GemmArgs& g, constant EpiArgs& e, + int t0, int c0, ushort tid, threadgroup float* ot) { + tensor, tensor_inline> tA(A, dextents(g.K, g.M)); + tensor, tensor_inline> tB(B, dextents(g.K, g.N)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, t0); + auto mB = tB.slice(0, c0); + if constexpr (REG) { + auto acc = op.template get_destination_cooperative_tensor(); + op.run(mA, mB, acc); + const int n = int(e.n_embd), d0 = c0 / 4; + // On M5, each lane owns four consecutive output columns: one output's HC streams. + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { + const auto ix = acc.get_multidimensional_index(i); + const int t = t0 + ix[1], d = d0 + ix[0] / 4; + if (t >= int(e.n_tok)) continue; + device const half* xnh_t = xnh + ulong(t) * 4 * n + d; + float s = 0.0f; + #pragma clang loop unroll(full) + for (int c = 0; c < 4; ++c) + s = fma(float(xnh_t[c * n]), 1.0f / (1.0f + exp(-acc[i + c])), s); + s *= 0.25f; + const ulong out_idx = ulong(t) * n + d; + if (mixed) mixed[out_idx] = s; + mixh[out_idx] = half(s); + } + } else { + tensor, tensor_inline> tO(ot, dextents(TN, TM)); + op.run(mA, mB, tO); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int hc = int(e.hc), n = int(e.n_embd), DPT = TN / hc; + const int d0 = c0 / hc; + for (int idx = tid; idx < TM * DPT; idx += SG * 32) { + const int m = idx / DPT, dd = idx - m * DPT, t = t0 + m; + if (t >= int(e.n_tok)) continue; + float s = 0.0f; + for (int c = 0; c < hc; ++c) + s = fma(float(xnh[ulong(t) * hc * n + c * n + d0 + dd]), 1.0f / (1.0f + exp(-ot[m * TN + dd * hc + c])), s); + s = s / float(hc); + if (mixed) mixed[ulong(t) * n + d0 + dd] = s; + mixh[ulong(t) * n + d0 + dd] = half(s); + } + } +} +template +kernel void gemm_hc_up(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device const half* xnh [[buffer(2)]], + device float* mixed [[buffer(3)]], device half* mixh [[buffer(4)]], constant GemmArgs& g [[buffer(5)]], + constant EpiArgs& e [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float ot[REG ? 1 : TM * TN]; + const int t0 = int(tg.y) * TM, c0 = int(tg.x) * TN; + if constexpr (REG && TM > 16) { + if (t0 + 16 >= int(e.n_tok)) { + hc_up_tile<16, TN, SG, REG>(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot); + return; + } + } + if constexpr (REG && TM > 32) { + if (t0 + 32 >= int(e.n_tok)) { + hc_up_tile<32, TN, SG, REG>(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot); + return; + } + } + hc_up_tile(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot); +} +template [[host_name("gemm_hc_up_64_64_4")]] kernel decltype(gemm_hc_up<64, 64, 4>) gemm_hc_up<64, 64, 4>; +template [[host_name("gemm_hc_up_64_32_4")]] kernel decltype(gemm_hc_up<64, 32, 4>) gemm_hc_up<64, 32, 4>; +template [[host_name("gemm_hc_up_32_64_4")]] kernel decltype(gemm_hc_up<32, 64, 4>) gemm_hc_up<32, 64, 4>; +template [[host_name("gemm_hc_up_128_64_4")]] kernel decltype(gemm_hc_up<128, 64, 4>) gemm_hc_up<128, 64, 4>; +template [[host_name("gemm_hc_up_64_128_4")]] kernel decltype(gemm_hc_up<64, 128, 4>) gemm_hc_up<64, 128, 4>; +template [[host_name("gemm_hc_up_reg_64_64_4")]] kernel decltype(gemm_hc_up<64, 64, 4, true>) gemm_hc_up<64, 64, 4, true>; +template [[host_name("gemm_hc_up_reg_128_64_4")]] kernel decltype(gemm_hc_up<128, 64, 4, true>) gemm_hc_up<128, 64, 4, true>; +template [[host_name("gemm_hc_up_reg_64_128_4")]] kernel decltype(gemm_hc_up<64, 128, 4, true>) gemm_hc_up<64, 128, 4, true>; +template [[host_name("gemm_hc_up_reg_64_32_4")]] kernel decltype(gemm_hc_up<64, 32, 4, true>) gemm_hc_up<64, 32, 4, true>; +template [[host_name("gemm_hc_up_reg_32_128_4")]] kernel decltype(gemm_hc_up<32, 128, 4, true>) gemm_hc_up<32, 128, 4, true>; + +// gemm_hc_down: the HC down GEMM V = xn . [w_down; w_inject]^T (T x (hc_lr + hc), K = hc * n_embd) with gr_epi_down +// as its register epilogue: lo_h = half(silu(v / hc)), inject = v. V is never written; lo_h, inject bitwise gemm_f16 + +// gr_epi_down. M5 layout only (coop_k4: a lane's elements 4j .. 4j + 3 are four consecutive columns of one row); +// hc_lr % 4 == 0, hc == 4. grid (ceil((hc_lr + hc) / TN), ceil(T / TM)), 32 * SG threads +template +inline void hc_down_tile(device half* A, device half* B, device half* lo_h, device float* inject, + constant EpiArgs& e, int t0, int c0) { + const int K = int(e.hc * e.n_embd), nr = int(e.hc_lr + e.hc); + tensor, tensor_inline> tA(A, dextents(K, int(e.n_tok))); + tensor, tensor_inline> tB(B, dextents(K, nr)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, t0); + auto mB = tB.slice(0, c0); + auto acc = op.template get_destination_cooperative_tensor(); + op.run(mA, mB, acc); + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { + const auto ix = acc.get_multidimensional_index(i); + const int r = c0 + ix[0], t = t0 + ix[1]; + if (t >= int(e.n_tok) || r >= nr) continue; + if (r < int(e.hc_lr)) { + half4 o; + #pragma clang loop unroll(full) + for (int c = 0; c < 4; ++c) { + const float z = acc[i + c] * 0.25f; // / hc, exact: hc == 4 + o[c] = half(z / (1.0f + exp(-z))); + } + *(device half4*)(lo_h + ulong(t) * e.hc_lr + r) = o; + } else { + *(device float4*)(inject + ulong(t) * e.hc) = float4(acc[i], acc[i + 1], acc[i + 2], acc[i + 3]); + } + } +} +template +inline void hc_down_rows(device half* A, device half* B, device half* lo_h, device float* inject, + constant EpiArgs& e, int t0, int c0) { + // Trimming 32-row parents regresses full tiles on M5. + if constexpr (TM > 32) { + if (t0 + 16 >= int(e.n_tok)) { + hc_down_tile<16, TN, SG>(A, B, lo_h, inject, e, t0, c0); + return; + } + if (t0 + 32 >= int(e.n_tok)) { + hc_down_tile<32, TN, SG>(A, B, lo_h, inject, e, t0, c0); + return; + } + } + hc_down_tile(A, B, lo_h, inject, e, t0, c0); +} +template +kernel void gemm_hc_down(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device half* lo_h [[buffer(2)]], + device float* inject [[buffer(3)]], constant EpiArgs& e [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int c0 = int(tg.x) * TN, t0 = int(tg.y) * TM; + if constexpr (TN == 128) { + // A single 96-column tail beats serial 64+16 on M5, including SG=2. + if (e.hc == 4 && c0 + 64 == int(e.hc_lr)) { + hc_down_rows(A, B, lo_h, inject, e, t0, c0); + return; + } + } + hc_down_rows(A, B, lo_h, inject, e, t0, c0); +} +template [[host_name("gemm_hc_down_16_64_2")]] kernel decltype(gemm_hc_down<16, 64, 2>) gemm_hc_down<16, 64, 2>; +template [[host_name("gemm_hc_down_32_64_2")]] kernel decltype(gemm_hc_down<32, 64, 2>) gemm_hc_down<32, 64, 2>; +template [[host_name("gemm_hc_down_32_128_4")]] kernel decltype(gemm_hc_down<32, 128, 4>) gemm_hc_down<32, 128, 4>; +template [[host_name("gemm_hc_down_64_128_4")]] kernel decltype(gemm_hc_down<64, 128, 4>) gemm_hc_down<64, 128, 4>; +template [[host_name("gemm_hc_down_32_128_2")]] kernel decltype(gemm_hc_down<32, 128, 2>) gemm_hc_down<32, 128, 2>; + +// gemm_swiglu_up: the shared expert's up GEMM u = x . w_up^T (M x N, K) with swiglu_h as its register epilogue against +// the gate GEMM's g: h = half(silu(g) * u). u is never written; h bitwise gemm_f16 + swiglu_h. One accumulator: the +// gate stays a plain gemm_f16 before it. M5 layout only (as gemm_hc_down); N % 4 == 0. +// grid (ceil(N / TN), ceil(M / TM)), 32 * SG threads +template +inline void swiglu_up_tile(device half* A, device half* B, device const float* g, device half* h, + constant GemmArgs& a, int t0, int c0) { + tensor, tensor_inline> tA(A, dextents(a.K, a.M)); + tensor, tensor_inline> tB(B, dextents(a.K, a.N)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, t0); + auto mB = tB.slice(0, c0); + auto acc = op.template get_destination_cooperative_tensor(); + op.run(mA, mB, acc); + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { + const auto ix = acc.get_multidimensional_index(i); + const int c = c0 + ix[0], t = t0 + ix[1]; + if (t >= a.M || c >= a.N) continue; + const ulong idx = ulong(t) * uint(a.N) + uint(c); + const float4 gv = *(device const float4*)(g + idx); + half4 o; + #pragma clang loop unroll(full) + for (int k = 0; k < 4; ++k) o[k] = half(gv[k] / (1.0f + exp(-gv[k])) * acc[i + k]); + *(device half4*)(h + idx) = o; + } +} +template +kernel void gemm_swiglu_up(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device const float* g [[buffer(2)]], + device half* h [[buffer(3)]], constant GemmArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int t0 = int(tg.y) * TM, c0 = int(tg.x) * TN; + if constexpr (TM > 32) { + if (t0 + 16 >= a.M) { + swiglu_up_tile<16, TN, SG>(A, B, g, h, a, t0, c0); + return; + } + if (t0 + 32 >= a.M) { + swiglu_up_tile<32, TN, SG>(A, B, g, h, a, t0, c0); + return; + } + } + swiglu_up_tile(A, B, g, h, a, t0, c0); +} +template [[host_name("gemm_swiglu_up_32_64_2")]] kernel decltype(gemm_swiglu_up<32, 64, 2>) gemm_swiglu_up<32, 64, 2>; +template [[host_name("gemm_swiglu_up_64_128_4")]] kernel decltype(gemm_swiglu_up<64, 128, 4>) gemm_swiglu_up<64, 128, 4>; diff --git a/results/round39/astra_r39_qsa/baseline-kernels/qsa.metal b/results/round39/astra_r39_qsa/baseline-kernels/qsa.metal new file mode 100644 index 0000000..63259b4 --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/qsa.metal @@ -0,0 +1,492 @@ +// metal/kernels/qsa.metal - QSA (sparse full attention) decode for a window of tokens, between the projections. +// Contract of include/strata/kernels/qsa.hpp / qsa_select.hpp / qsa_decode_attn.hpp / kv_q8.hpp (the engine's +// native path), in a Metal layout (no paging: unified memory): +// +// qsa_prep per token t (cell pos0 + t): q heads rms_norm * q_norm, NEOX RoPE on 64 dims -> qcur; +// k rms_norm * k_norm, RoPE -> int8 KV append (fp16 scale per 64 = fp16(amax/127), rint, clamp); +// v -> int8 append; indexer q rms_norm * iq_norm, RoPE -> qidx. The indexer KEYS (last +// threadgroup, all window tokens in order): raw tail [cell % 4]; on a block completion +// pooled[b] = rope(rms_norm(mean of the block's 4 raw keys) * ik_norm, pos of the block's first cell); +// cell 0 also sets dead = rms_norm(raw) * ik_norm (the tail block's key). +// qsa_score per token, block b <= n_bid: sum over the 4 indexer heads of relu(qidx . key_b); key_b = pooled[b], +// the tail block n_bid scores `dead` and gets +1e9 when it has cells. +// qsa_topk per token: width = min(n_kv, 2051) cells; n_kv <= 2051: all; else the width cells of largest +// score (a cell inherits its block's), ties to the lowest cell, emitted ascending (radix select). +// qsa_attn split-K: (chunk of CH cells, kv head, token): 12 query heads; s = q.k / 16, chunk max m, sum l, +// acc = p . v -> partials +// qsa_merge per (head, token): log-sum-exp merge, * sigmoid(gate), written as prepared activations (xprep +// planes) for attn_output. +// +// RoPE: cos/sin from a host table [pos][32] (double precision, as build_rope_table). + +#define QH 24 +#define QKV 2 +#define QD 256 +#define QROT 64 +#define IH 4 +#define ID 128 +#define QBUDGET 2051 +#define QCH 64 + +constant int QSA_NTOK [[function_constant(3)]]; +// QSA_MROPE (function constant 24): interleaved M-RoPE (the vision path; strata/kernels/mrope.hpp, ggml IMROPE, +// sections 11/11/10): buffer 20 of qsa_prep is a table int32 [cells][3] (t, h, w), and rotary pair i (0..31) of the row +// at (pos_base + cell) rotates by the cos/sin table's row tab[(pos_base + cell) * 3 + i % 3] instead of its own - +// the pooled indexer keys too (position pos_base + 4 b of a block). Off: exactly the plain RoPE. +constant bool QSA_MROPE [[function_constant(24)]]; +constant bool qsa_mrope = is_function_constant_defined(QSA_MROPE) && QSA_MROPE; +// QSA_Q_F16_FC (function constant 25): qsa_prep writes qcur directly as fp16 (device half*) for qsa_flash, skipping qsa_q_f16. +constant bool QSA_Q_F16_FC [[function_constant(25)]]; +constant bool qsa_qf16 = is_function_constant_defined(QSA_Q_F16_FC) ? QSA_Q_F16_FC : false; +// QSA_ROW0 (function constant 28): qsa_prep's threadgroup x is row QSA_ROW0 + x (24: the K/V rows only, for rows whose +// queries nobody reads). Undefined: 0. +constant int QSA_ROW0 [[function_constant(28)]]; +constant int qsa_row0 = is_function_constant_defined(QSA_ROW0) ? QSA_ROW0 : 0; + +struct QsaArgs { + uint pos0; // the window's first cell (relative to the sequence) + uint pos_base; // absolute position of cell 0 (RoPE) + uint max_cells; + uint n_tok; + float eps; +}; + +// cache layout per layer +// kq/vq [cell][kv head][256] int8, ks/vs [cell][kv head][4] half +// tail [4][128] f32 (slot = cell % 4), dead [128], pooled [max_cells / 4 + 1][128] + +static inline float tg_sum256(float v, threadgroup float* part, ushort sg, ushort lane, ushort nsg) { + v = simd_sum(v); + if (lane == 0) part[sg] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + float t = 0.0f; + for (int i = 0; i < nsg; ++i) t += part[i]; + return t; +} + +// one row (width w <= 256 = threads) : out = rope((rs * x) * g) (RoPE on the first 64 dims, NEOX pairs i, i+32) +static inline float norm_rope(float x, float xhi_unused, device const float* g, threadgroup float* row, + threadgroup float* part, ushort tid, ushort sg, ushort lane, int w, float eps, + device const float2* cs, uint absp, bool rope, device const int* mtab, + bool reuse = false) { + const float ss = tg_sum256(tid < w ? x * x : 0.0f, part, sg, lane, 8); + const float rs = rsqrt(ss / float(w) + eps); + const float y = tid < w ? (rs * x) * g[tid] : 0.0f; + float o = y; + if (rope) { + if (tid < QROT) row[tid] = y; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < QROT) { + const int p = tid & 31; + const uint pp = qsa_mrope ? uint(mtab[absp * 3 + uint(p % 3)]) : absp; + const float2 c = cs[pp * 32 + p]; + const float a = row[p], b = row[p + 32]; + o = tid < 32 ? a * c.x - b * c.y : a * c.y + b * c.x; + } + } + if (reuse) threadgroup_barrier(mem_flags::mem_threadgroup); + (void)xhi_unused; + return o; +} + +// QSA_KV_STRIDE (function constant 26): kcur, vcur, iqraw and ikraw are column slices of one [n_tok][QSA_KV_STRIDE] +// projection (prompt_chunk's fused K | V | indexer K | indexer Q GEMM, 1152 or 1664 columns). Undefined: their own +// [n_tok][rows] buffers. +constant int QSA_KV_STRIDE [[function_constant(26)]]; +constant int qsa_kv_s = is_function_constant_defined(QSA_KV_STRIDE) ? QSA_KV_STRIDE : 0; + +// grid (32 rows + 1, n_tok): rows 0..23 q heads, 24..25 k heads, 26..27 v heads, 28..31 indexer q heads; +// threadgroup 32 (x = 32, y = 0) is the indexer-key threadgroup. 256 threads. A grid of 28 (or of 4 from QSA_ROW0 = +// 24) writes no indexer state: the MTP layer's attention is a dense window, never scored. +kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float* kcur [[buffer(1)]], + device const float* vcur [[buffer(2)]], device const float* iqraw [[buffer(3)]], + device const float* ikraw [[buffer(4)]], device const float* qn [[buffer(5)]], + device const float* kn [[buffer(6)]], device const float* iqn [[buffer(7)]], + device const float* ikn [[buffer(8)]], device const float2* cs [[buffer(9)]], + device float* qcur [[buffer(10)]], device float* qidx [[buffer(11)]], + device char* kq [[buffer(12)]], device char* vq [[buffer(13)]], device half* ksc [[buffer(14)]], + device half* vsc [[buffer(15)]], device float* tail [[buffer(16)]], device float* dead [[buffer(17)]], + device float* pooled [[buffer(18)]], constant QsaArgs& a [[buffer(19)]], + device const int* mtab [[buffer(20)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + threadgroup float row[QROT]; + threadgroup float amax[8]; + const int r = int(tg.x) + qsa_row0; + // the row strides of kcur / vcur, ikraw, iqraw + const int s_kv = qsa_kv_s ? qsa_kv_s : QKV * QD, s_ik = qsa_kv_s ? qsa_kv_s : ID; + const int s_iq = qsa_kv_s ? qsa_kv_s : IH * ID; + if (r == 32) { + // the indexer keys: threadgroup (32, 0) handles t < 3 (boundary with prior tail + cell 0 dead) and final tail; + // every threadgroup (32, t >= 3) with (cell & 3) == 3 pools its 4 in-window keys in parallel. + const int t_self = int(tg.y); + if (t_self >= 3) { + const uint cell = a.pos0 + uint(t_self); + if ((cell & 3) == 3) { + float m = 0.0f; + if (tid < ID) { + m = (((ikraw[(t_self - 3) * s_ik + tid] + ikraw[(t_self - 2) * s_ik + tid]) + + ikraw[(t_self - 1) * s_ik + tid]) + ikraw[t_self * s_ik + tid]) * 0.25f; + } + const uint b = cell >> 2; + const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab); + if (tid < ID) pooled[b * ID + tid] = o; + } + return; + } + if (t_self != 0) return; + const int n_pre = min(int(a.n_tok), 3); + for (int t = 0; t < n_pre; ++t) { + const uint cell = a.pos0 + uint(t); + const float x = tid < ID ? ikraw[t * s_ik + tid] : 0.0f; + if (tid < ID) tail[(cell & 3) * ID + tid] = x; + threadgroup_barrier(mem_flags::mem_device); + if (cell == 0) { + const float o = norm_rope(x, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, 0, false, mtab, true); + if (tid < ID) dead[tid] = o; + } + if ((cell & 3) == 3) { + float m = 0.0f; + if (tid < ID) m = (((tail[tid] + tail[ID + tid]) + tail[2 * ID + tid]) + tail[3 * ID + tid]) * 0.25f; + const uint b = cell >> 2; + const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab, true); + if (tid < ID) pooled[b * ID + tid] = o; + } + threadgroup_barrier(mem_flags::mem_device); + } + const int t_tail0 = max(n_pre, int(a.n_tok) - 4); + for (int t = t_tail0; t < int(a.n_tok); ++t) { + const uint cell = a.pos0 + uint(t); + if (tid < ID) tail[(cell & 3) * ID + tid] = ikraw[t * s_ik + tid]; + } + return; + } + const int t = int(tg.y); + const uint cell = a.pos0 + uint(t), absp = a.pos_base + cell; + if (r < QH) { + const float x = qfull[(t * QH + r) * 2 * QD + tid]; + const float o = norm_rope(x, 0.0f, qn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); + if (qsa_qf16) ((device half*)qcur)[(t * QH + r) * QD + tid] = half(o); + else qcur[(t * QH + r) * QD + tid] = o; + } else if (r < QH + 2 * QKV) { + const bool isv = r >= QH + QKV; + const int h = r - QH - (isv ? QKV : 0); + float x = (isv ? vcur : kcur)[t * s_kv + h * QD + tid]; + if (!isv) x = norm_rope(x, 0.0f, kn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); + // int8 with an fp16 scale per 64 values: groups = simdgroup pairs (64 threads) + float m = simd_max(fabs(x)); + if (lane == 0) amax[sg] = m; + threadgroup_barrier(mem_flags::mem_threadgroup); + const int g = tid >> 6; + m = max(amax[2 * g], amax[2 * g + 1]); + const half sh = half(m / 127.0f); + const float sf = float(sh); + int q = 0; + if (sf > 0.0f) q = clamp(int(rint(x / sf)), -127, 127); + const uint rowi = cell * QKV + uint(h); + (isv ? vq : kq)[rowi * QD + tid] = char(q); + if ((tid & 63) == 0) (isv ? vsc : ksc)[rowi * 4 + g] = sh; + } else { + if (a.pos0 + a.n_tok <= uint(QBUDGET)) return; + const int h = r - QH - 2 * QKV; + const float x = tid < ID ? iqraw[t * s_iq + 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); + if (tid < ID) qidx[(t * IH + h) * ID + tid] = o; + } +} + +// scores [n_tok][max_blocks]; grid (ceil(max_blocks / 8), n_tok), 256 threads: a simdgroup per block. +kernel void qsa_score(device const float* qidx [[buffer(0)]], device const float* pooled [[buffer(1)]], + device const float* dead [[buffer(2)]], device float* scores [[buffer(3)]], + constant QsaArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + const int t = int(tg.y); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + const uint b = tg.x * 8 + sg; + if (b > n_bid) return; + device const float* key = b == n_bid ? dead : pooled + b * ID; + device const float* qp = qidx + (t * IH) * ID + lane; + float k[4], q[IH][4]; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) k[i] = key[lane + 32 * i]; + #pragma clang loop unroll(full) + for (int h = 0; h < IH; ++h) { + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) q[h][i] = qp[h * ID + 32 * i]; + } + float s = 0.0f; + #pragma clang loop unroll(full) + for (int h = 0; h < IH; ++h) { + float d = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) d = fma(q[h][i], k[i], d); + d = simd_sum(d); + s += max(d, 0.0f); + } + if (b == n_bid && (n_kv & 3) != 0) s += 1e9f; + if (lane == 0) scores[uint(t) * (a.max_cells / 4 + 2) + b] = s; +} + +// ids [n_tok][QBUDGET], ascending. One threadgroup of 1024 threads per token. Order key of a float: larger score = +// larger key; a radix select (4 passes of 8 bits) over the blocks with weights = their cell counts finds the +// threshold key T and how many cells with key == T to take (the lowest cells first). +static inline uint okey(float f) { + const uint u = as_type(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]], + constant QsaArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint hist[256]; + threadgroup uint sel[4]; + threadgroup uint2 pref[32]; + threadgroup uint2 totals; + const int t = int(tg.x); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2, nb = n_bid + 1; + device int* out = ids + t * QBUDGET; + device const float* sc = scores + uint(t) * (a.max_cells / 4 + 2); + if (n_kv <= QBUDGET) { + for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); + return; + } + const bool cached = nb <= 4096; + uint keys[4]; + if (cached) + for (uint b = tid, i = 0; b < nb; b += 1024, ++i) keys[i] = okey(sc[b]); + // radix select of the `need`-th largest cell key + uint prefix = 0, pmask = 0, need = QBUDGET; + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint b = tid; b < nb; b += 1024) { + const uint k = cached ? keys[b / 1024] : okey(sc[b]); + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + uint counts[8], sum = 0; + for (uint j = 0; j < 8; ++j) { + counts[j] = atomic_load_explicit(&hist[255 - (lane * 8 + j)], memory_order_relaxed); + sum += counts[j]; + } + uint acc = simd_prefix_exclusive_sum(sum); + if (acc < need && acc + sum >= need) { + for (uint j = 0; j < 8; ++j) { + if (acc + counts[j] >= need) { + sel[0] = 255 - (lane * 8 + j); + sel[1] = need - acc; + break; + } + acc += counts[j]; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; + } + // T = prefix; take every cell with key > T and the first `need` cells (lowest index) with key == T, ascending. + // Pass over blocks in order of index with a prefix scan of each block's selected cell count. + const uint T = prefix; + uint base = 0, eq_taken = 0; + for (uint b0 = 0; b0 < nb; b0 += 1024) { + const uint b = b0 + tid; + uint w = 0, weq = 0; + if (b < nb) { + const uint k = cached ? keys[b / 1024] : okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + if (k > T) w = cells; + else if (k == T) weq = cells; + } + const uint eq_before = simd_prefix_exclusive_sum(weq); + const uint gt_before = simd_prefix_exclusive_sum(w); + if (lane == 31) pref[sg] = uint2(eq_before + weq, gt_before + w); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + const uint2 count = pref[lane]; + const uint2 before = uint2(simd_prefix_exclusive_sum(count.x), + simd_prefix_exclusive_sum(count.y)); + pref[lane] = before; + if (lane == 31) totals = before + count; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint2 group_before = pref[sg], count = totals; + simdgroup_barrier(mem_flags::mem_threadgroup); + const uint eq0 = eq_taken + group_before.x + eq_before; + const uint take_eq = eq0 < need ? min(weq, need - eq0) : 0u; + const uint mine = w + take_eq; + // Clipping the equal-key prefix preserves the lowest-cell tie break. + const uint remaining = need - min(need, eq_taken); + const uint off = base + group_before.y + gt_before + + min(remaining, group_before.x + eq_before); + for (uint c = 0; c < mine; ++c) out[off + c] = int(4 * b + c); + base += count.y + min(remaining, count.x); + eq_taken += count.x; + } +} + +// partials: [n_tok][chunks][QH][QD + 2] (acc, then m, l). grid (chunks, QKV, n_tok), 256 threads. +kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + 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 HPG = QH / QKV; // 12 query heads per kv head + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH]; + threadgroup int cid[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; + const int nchunks = (QBUDGET + QCH - 1) / QCH; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + if (uint(c0) >= width) { + // an empty chunk: m = -inf, l = 0 (the merge skips it) + for (int h = 0; h < HPG; ++h) + if (tid == 0) { pp[(g * HPG + h) * (QD + 2) + QD] = -INFINITY; pp[(g * HPG + h) * (QD + 2) + QD + 1] = 0.0f; } + return; + } + const int nc = min(QCH, int(width) - c0); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + g * HPG + h) * QD + tid]; + if (tid < QCH) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + // scores: 12 x 64 dots of 256; simdgroup sg takes cells sg, sg + 8, ...; lanes split the 256 dims (8 each) + threadgroup const float* q_lane = &qs[0][lane * 8]; + for (int c = sg; c < nc; c += 8) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char* kr = kq + rowi * QD + lane * 8; + const char4 k0 = *(device const char4*)(kr); + const char4 k1 = *(device const char4*)(kr + 4); + const float scl = float(ksc[rowi * 4 + (lane >> 3)]); + const float kv[8] = { + float(k0.x) * scl, float(k0.y) * scl, float(k0.z) * scl, float(k0.w) * scl, + float(k1.x) * scl, float(k1.y) * scl, float(k1.z) * scl, float(k1.w) * scl + }; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + threadgroup const float* qh = q_lane + h * QD; + float d = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) d = fma(qh[i], kv[i], d); + d = simd_sum(d); + if (lane == 0) s[h][c] = d * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // per head: max, exp, sum (simdgroup h % 8 handles heads h, h + 8) + threadgroup float mh[HPG], lh[HPG]; + // nc <= QCH = 64: a lane holds cells lane and lane + 32 in registers (the same max and sum as a strided loop) + const bool in0 = int(lane) < nc, in1 = int(lane) + 32 < nc; + for (int h = sg; h < HPG; h += 8) { + const float s0 = in0 ? s[h][lane] : -INFINITY, s1 = in1 ? s[h][lane + 32] : -INFINITY; + const float m = simd_max(max(s0, s1)); + const float p0 = in0 ? exp(s0 - m) : 0.0f, p1 = in1 ? exp(s1 - m) : 0.0f; + if (in0) s[h][lane] = p0; + if (in1) s[h][lane + 32] = p1; + const float l = simd_sum(p0 + p1); + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // acc[h][j] = sum_c p[h][c] v[c][j]: thread = dim j + float acc[HPG]; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + // four cells' V loads in flight at once; the fma order per head stays c, c + 1, ... + auto vcell = [&](int c) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + return float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); + }; + int c = 0; + for (; c + 4 <= nc; c += 4) { + const float v0 = vcell(c), v1 = vcell(c + 1), v2 = vcell(c + 2), v3 = vcell(c + 3); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + float x = fma(s[h][c], v0, acc[h]); + x = fma(s[h][c + 1], v1, x); + x = fma(s[h][c + 2], v2, x); + acc[h] = fma(s[h][c + 3], v3, x); + } + } + for (; c < nc; ++c) { + const float v = vcell(c); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][c], v, acc[h]); + } + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (g * HPG + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} + +constant float qsa_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; +constant float qsa_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + +static inline void qsa_write_prepared(float y, device float* xp, constant QsaArgs& a, + int t, int h, ushort tid) { + // prepared planes (xprep layout) of the n = QH * QD vector + const uint n = QH * QD, plane = a.n_tok * n, d = uint(h) * QD + tid; + const int j = tid & 15; + device float* raw = xp + uint(t) * n + d; + raw[0] = y; + raw[plane] = y * qsa_q2s[j & 7]; + raw[2 * plane] = y * qsa_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? y * 0.00390625f : y; + float s0 = 0.0f, s1 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(y, ushort(k)); + s1 += simd_shuffle_down(y, ushort(k + 1)); + } + if (j == 0) xp[4 * plane + uint(t) * (n / 16) + d / 16] = s0 + s1; +} + +// grid (QH, n_tok), 256 threads: out = merge * sigmoid(gate), as prepared activations of n = QH * QD per token +kernel void qsa_merge(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], + device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const int h = int(tg.x), t = int(tg.y); + const int nchunks = (QBUDGET + QCH - 1) / QCH; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nc = int((width + QCH - 1) / QCH); + constexpr uint STRIDE = QH * (QD + 2); + device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2); + // the chunks' (m, l) and weights exp(m - M) once per threadgroup, not once per thread (nc <= 33 < 64) + threadgroup float mk[64], lk[64], ek[64], L_sh; + if (tid < uint(nc)) { + const float2 ml = *(device const float2*)(pp + tid * STRIDE + QD); + mk[tid] = ml.x; + lk[tid] = ml.y; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float M = -INFINITY; + for (int k = tid & 31; k < nc; k += 32) M = max(M, mk[k]); + M = simd_max(M); + if (tid < uint(nc)) ek[tid] = exp(mk[tid] - M); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + float L = 0.0f; + for (int k = 0; k < nc; ++k) L = fma(lk[k], ek[k], L); + L_sh = L; + } + float o = 0.0f; + device const float* q = pp + tid; + for (int k = 0; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); + threadgroup_barrier(mem_flags::mem_threadgroup); + o = o / L_sh; + const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; + const float y = o * (1.0f / (1.0f + exp(-gate))); + qsa_write_prepared(y, xp, a, t, h, tid); +} + diff --git a/results/round39/astra_r39_qsa/baseline-kernels/qsa2.metal b/results/round39/astra_r39_qsa/baseline-kernels/qsa2.metal new file mode 100644 index 0000000..d5b594c --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/qsa2.metal @@ -0,0 +1,338 @@ +// metal/kernels/qsa2.metal - QSA decode v2 (after qsa.metal): the same selection and attention math, more parallel. +// +// qsa_attn2 chunks of 32 cells, the 12 query heads of a kv head in 2 groups of 6: 4x the threadgroups of +// qsa_attn (at a 1K context 128 instead of 32); V loads 8 cells at a time. +// qsa_merge2 qsa_merge for that chunking. +// qsa_topk2 the digit of each radix pass found by a parallel suffix scan of the 256 bins (simdgroup prefix +// sums) instead of one thread; cells emitted from contiguous block ranges per thread after one scan. + +#define QCH2 32 +#define QHG 2 + +// partials: [n_tok][chunks2][QH][QD + 2]. grid (chunks2, QKV * QHG, n_tok), 256 threads. +kernel void qsa_attn2(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + 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 HPG = QH / QKV / QHG; // 6 + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH2]; + threadgroup float mh[HPG], lh[HPG]; + threadgroup int cid[QCH2]; + const int chunk = int(tg.x), g = int(tg.y) / QHG, hg = int(tg.y) % QHG, t = int(tg.z); + const int h0 = g * (QH / QKV) + hg * HPG; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH2; + const int nchunks = (QBUDGET + QCH2 - 1) / QCH2; + if (uint(c0) >= width) return; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + const int nc = min(QCH2, int(width) - c0); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + h0 + h) * QD + tid]; + if (tid < QCH2) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup const float* q_lane = &qs[0][lane * 8]; + for (int c = sg; c < nc; c += 8) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char4* kr = (device const char4*)(kq + rowi * QD + lane * 8); + const char4 k0 = kr[0], k1 = kr[1]; + const float scl = float(ksc[rowi * 4 + (lane >> 3)]); + const float kv[8] = {float(k0.x) * scl, float(k0.y) * scl, float(k0.z) * scl, float(k0.w) * scl, + float(k1.x) * scl, float(k1.y) * scl, float(k1.z) * scl, float(k1.w) * scl}; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + float d = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) d = fma(q_lane[h * QD + i], kv[i], d); + d = simd_sum(d); + if (lane == 0) s[h][c] = d * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg < HPG) { + const int h = sg; + const float sv = lane < nc ? s[h][lane] : -INFINITY; + const float m = simd_max(sv); + const float p = lane < nc ? exp(sv - m) : 0.0f; + const float l = simd_sum(p); + if (lane < nc) s[h][lane] = p; + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float acc[HPG]; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + const uint v_scale_off = tid >> 6; + int cb = 0; + for (; cb + 8 <= nc; cb += 8) { + float v[8]; + #pragma clang loop unroll(full) + for (int c = 0; c < 8; ++c) { + const uint rowi = uint(cid[cb + c]) * QKV + uint(g); + v[c] = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + v_scale_off]); + } + #pragma clang loop unroll(full) + for (int c = 0; c < 8; ++c) + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb + c], v[c], acc[h]); + } + for (; cb < nc; ++cb) { + const uint rowi = uint(cid[cb]) * QKV + uint(g); + const float v = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + v_scale_off]); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb], v, acc[h]); + } + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (h0 + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} + +kernel void qsa_merge2(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], + device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float mk[128], lk[128], ek[128], L_sh; + const int h = int(tg.x), t = int(tg.y); + const int nchunks = (QBUDGET + QCH2 - 1) / QCH2; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nc = int((width + QCH2 - 1) / QCH2); + constexpr uint STRIDE = QH * (QD + 2); + device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2); + if (tid < uint(nc)) { + const float2 ml = *(device const float2*)(pp + tid * STRIDE + QD); + mk[tid] = ml.x; + lk[tid] = ml.y; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float M = -INFINITY; + for (int k = tid & 31; k < nc; k += 32) M = max(M, mk[k]); + M = simd_max(M); + if (tid < uint(nc)) ek[tid] = exp(mk[tid] - M); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + float L = 0.0f; + for (int k = 0; k < nc; ++k) L = fma(lk[k], ek[k], L); + L_sh = L; + } + float o = 0.0f; + device const float* q = pp + tid; + for (int k = 0; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); + threadgroup_barrier(mem_flags::mem_threadgroup); + o = o / L_sh; + const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; + const float y = o * (1.0f / (1.0f + exp(-gate))); + qsa_write_prepared(y, xp, a, t, h, tid); +} + +kernel void qsa_topk2(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]], + constant QsaArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint hist[256]; + threadgroup uint sgs[32], sgs2[32]; + threadgroup uint sel[2]; + const int t = int(tg.x); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2, nb = n_bid + 1; + device int* out = ids + t * QBUDGET; + device const float* sc = scores + uint(t) * (a.max_cells / 4 + 2); + if (n_kv <= QBUDGET) { + for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); + return; + } + if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + uint prefix = 0, pmask = 0, need = QBUDGET; + if (nb <= 4096) { + uint keys[4]; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) { + const uint b = tid + 1024u * uint(i); + keys[i] = b < nb ? okey(sc[b]) : 0u; + } + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) { + const uint b = tid + 1024u * uint(i); + if (b >= nb) break; + const uint k = keys[i]; + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint cnt = 0, ex = 0; + if (tid < 256) { + cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); + if (pass < 3) atomic_store_explicit(&hist[255 - tid], 0u, memory_order_relaxed); + ex = simd_prefix_exclusive_sum(cnt); + if (lane == 31) sgs[sg] = ex + cnt; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 256) { + uint above = ex; + for (int i = 0; i < int(sg); ++i) above += sgs[i]; + if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; + } + } else { + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + for (uint b = tid; b < nb; b += 1024) { + const uint k = okey(sc[b]); + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint cnt = 0, ex = 0; + if (tid < 256) { + cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); + if (pass < 3) atomic_store_explicit(&hist[255 - tid], 0u, memory_order_relaxed); + ex = simd_prefix_exclusive_sum(cnt); + if (lane == 31) sgs[sg] = ex + cnt; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 256) { + uint above = ex; + for (int i = 0; i < int(sg); ++i) above += sgs[i]; + if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; + } + } + const uint T = prefix; + const uint per = (nb + 1023) / 1024, b0 = min(nb, uint(tid) * per), b1 = min(nb, b0 + per); + uint gt = 0, eq = 0; + for (uint b = b0; b < b1; ++b) { + const uint k = okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + if (k > T) gt += cells; + else if (k == T) eq += cells; + } + const uint eqb = simd_prefix_exclusive_sum(eq); + if (lane == 31) sgs2[sg] = eqb + eq; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint eq_before = eqb; + for (int i = 0; i < int(sg); ++i) eq_before += sgs2[i]; + const uint eq_take = eq_before >= need ? 0u : min(eq, need - eq_before); + const uint mine = gt + eq_take; + const uint mb = simd_prefix_exclusive_sum(mine); + if (lane == 31) sgs[sg] = mb + mine; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint off = mb; + for (int i = 0; i < int(sg); ++i) off += sgs[i]; + uint left = eq_take; + for (uint b = b0; b < b1; ++b) { + const uint k = okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + uint take = 0; + if (k > T) take = cells; + else if (k == T) { take = min(cells, left); left -= take; } + for (uint c = 0; c < take; ++c) out[off + c] = int(4 * b + c); + off += take; + } +} + +// qsa_attn3: qsa_attn's chunking and partial layout (qsa_merge reads it), without a reduction per (cell, head): +// for q.k a lane owns a cell (8 simdgroups = 2 halves of the 64 cells x 4 groups of 3 heads; q from threadgroup +// memory, read by every lane alike), for p.v a thread owns a dimension. grid (33, QKV, n_tok), 256 threads. +kernel void qsa_attn3(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + 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 HPG = QH / QKV; // 12 + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH]; + threadgroup float mh[HPG], lh[HPG]; + threadgroup int cid[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; + const int nchunks = (QBUDGET + QCH - 1) / QCH; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + if (uint(c0) >= width) return; + const int nc = min(QCH, int(width) - c0); + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + g * HPG + h) * QD + tid]; + if (tid < QCH) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + { + const int half_ = sg & 1, hgrp = sg >> 1; // cells 32 * half_ + lane, heads 3 * hgrp .. + 3 + const int c = 32 * half_ + lane; + if (c < nc) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char4* kr = (device const char4*)(kq + rowi * QD); + float d0 = 0.0f, d1 = 0.0f, d2 = 0.0f; + for (int gq = 0; gq < 4; ++gq) { + const float scl = float(ksc[rowi * 4 + gq]); + float e0 = 0.0f, e1 = 0.0f, e2 = 0.0f; + for (int i = 0; i < 16; ++i) { + const char4 k4 = kr[gq * 16 + i]; + const int base = gq * 64 + 4 * i; + const float kv[4] = {float(k4.x), float(k4.y), float(k4.z), float(k4.w)}; + for (int e = 0; e < 4; ++e) { + e0 = fma(qs[3 * hgrp][base + e], kv[e], e0); + e1 = fma(qs[3 * hgrp + 1][base + e], kv[e], e1); + e2 = fma(qs[3 * hgrp + 2][base + e], kv[e], e2); + } + } + d0 = fma(e0, scl, d0); + d1 = fma(e1, scl, d1); + d2 = fma(e2, scl, d2); + } + s[3 * hgrp][c] = d0 * 0.0625f; + s[3 * hgrp + 1][c] = d1 * 0.0625f; + s[3 * hgrp + 2][c] = d2 * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int h = sg; h < HPG; h += 8) { + float m = -INFINITY; + for (int c = lane; c < nc; c += 32) m = max(m, s[h][c]); + m = simd_max(m); + float l = 0.0f; + for (int c = lane; c < nc; c += 32) { + const float p = exp(s[h][c] - m); + s[h][c] = p; + l += p; + } + l = simd_sum(l); + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float acc[HPG]; + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + for (int cb = 0; cb < nc; cb += 8) { + float v[8]; + const int m = min(8, nc - cb); + for (int c = 0; c < 8; ++c) { + v[c] = 0.0f; + if (c < m) { + const uint rowi = uint(cid[cb + c]) * QKV + uint(g); + v[c] = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); + } + } + for (int c = 0; c < m; ++c) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb + c], v[c], acc[h]); + } + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (g * HPG + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} diff --git a/results/round39/astra_r39_qsa/baseline-kernels/qsa_prompt.metal b/results/round39/astra_r39_qsa/baseline-kernels/qsa_prompt.metal new file mode 100644 index 0000000..444e0c4 --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/qsa_prompt.metal @@ -0,0 +1,453 @@ +// metal/kernels/qsa_prompt.metal - QSA attention for a PROMPT chunk on the matrix units (after qsa.metal: QH, QKV, QD, +// QBUDGET, QCH, QsaArgs). The same split-K partials as qsa_attn (part [t][chunk][QH][QD + 2]: acc, then m, l), so +// qsa_merge is unchanged; per (chunk of QCH = 64 selected cells, kv head g, token t) one simdgroup: +// S[12 x 64] = Qh[12 x 256] . K^T / 16 single-simdgroup matmul2d, K (int8 * fp16 scale per 64) decoded into +// the right-input cooperative tensor, 32 dims a step +// per head: m = max, p = exp(s - m), l = sum p (fp32, through threadgroup memory), P -> fp16 +// acc[12 x 256] = P[12 x 64] . V[64 x 256] four 64-dim matmuls, V decoded into the right input +// The operands are fp16 (q rounded from fp32; k, v = int8 * scale rounded): to rounding of qsa_attn's fp32. +#include +using namespace mpp::tensor_ops; +#ifndef STRATA_COOP_K4 +#define STRATA_COOP_K4 +// function constant 7: the right-input cooperative tensor holds runs of 4 consecutive k at one n per lane (the M5's +// layout; tests/coop_layout). false (the M4: runs of 2 n at one k): every element decoded on its own. +constant bool COOP_K4 [[function_constant(7)]]; +constant bool coop_k4 = is_function_constant_defined(COOP_K4) ? COOP_K4 : true; +#endif + +// a bisection probe (function constant 8): stop at stage n and write m = n (1: after the ids, 2: after S, 3: after +// the softmax); 0 = normal +constant int QSA_STOP [[function_constant(8)]]; +constant int qsa_stop = is_function_constant_defined(QSA_STOP) ? QSA_STOP : 0; + +// grid (chunks, QKV, n_tok), 32 threads. qh: the q heads as fp16 [t][QH][QD] (qcur rounded). +kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int HPG = QH / QKV, NCH = (QBUDGET + QCH - 1) / QCH; + threadgroup float s[16 * QCH]; + threadgroup half p[16 * QCH]; + threadgroup int cid[QCH]; // coop_k4: byte offsets; otherwise cell ids + threadgroup half4 ks[QCH], vs[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; + device float* pp = part + ((uint(t) * NCH + chunk) * QH + g * HPG) * (QD + 2); + if (uint(c0) >= width) { + if (lane < HPG) { pp[lane * (QD + 2) + QD] = -INFINITY; pp[lane * (QD + 2) + QD + 1] = 0.0f; } + return; + } + const int nc = min(QCH, int(width) - c0); + for (int c = lane; c < QCH; c += 32) { + const int id = c < nc ? (n_kv <= QBUDGET ? c0 + c : ids[t * QBUDGET + c0 + c]) : -1; + cid[c] = coop_k4 ? int((uint(max(id, 0)) * QKV + g) * QD) : id; + if (id >= 0) { + const uint row4 = (uint(id) * QKV + g) * 4; + if (coop_k4) ks[c] = *(device const half4*)(ksc + row4); + vs[c] = *(device const half4*)(vsc + row4); + } else { + ks[c] = half4(0.0h); + vs[c] = half4(0.0h); + } + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (qsa_stop == 1) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 1.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } + // ---- S = Q K^T + { + constexpr auto desc = matmul2d_descriptor(16, QCH, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tQ((device half*)qh + (uint(t) * QH + g * HPG) * QD, + dextents(QD, HPG)); + auto rk = op.template get_right_input_cooperative_tensor(); + auto tA = tQ.template slice<32, 16>(0, 0); + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + // the second half repeats these eight cells, 16 k later; empty cells read row 0 at scale 0. + uint ko[8]; + float ks8[8][4]; + if (coop_k4) +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = rk.get_multidimensional_index(uint16_t(4 * j))[1]; + ko[j] = uint(cid[c]); +#pragma clang loop unroll(full) + for (int h = 0; h < 4; ++h) ks8[j][h] = float(ks[c][h]); + } + for (int k0 = 0; k0 < QD; k0 += 32) { + if (!coop_k4) { + for (uint16_t i = 0; i < rk.get_capacity(); ++i) { + const auto ix = rk.get_multidimensional_index(i); + const int k = k0 + ix[0], id = cid[ix[1]]; + const uint row = uint(id) * QKV + g; + rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)])) : half(0.0f); + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rk.get_capacity(); i += 4) { + const int k = k0 + rk.get_multidimensional_index(i)[0], j = (i >> 2) & 7; + const half4 w = half4(float4(*(device const char4*)(kq + (ko[j] + uint(k)))) * ks8[j][k >> 6]); + rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w; + } + tA = tQ.template slice<32, 16>(k0, 0); + op.run(tA, rk, acc); + } + tensor, tensor_inline> tS(s, dextents(QCH, 16)); + acc.store(tS); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (qsa_stop == 2) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 2.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } + // ---- softmax pieces per head (lane = head for the reductions; 12 heads) + float mh = -INFINITY, lh = 0.0f; + if (lane < HPG) { +#pragma clang loop unroll_count(4) + for (int c = 0; c < nc; ++c) mh = max(mh, s[lane * QCH + c] * 0.0625f); +#pragma clang loop unroll_count(4) + for (int c = 0; c < QCH; ++c) { + const float e = c < nc ? exp(s[lane * QCH + c] * 0.0625f - mh) : 0.0f; + lh += e; + p[lane * QCH + c] = half(e); + } + } else if (lane < 16) { +#pragma clang loop unroll_count(4) + for (int c = 0; c < QCH; ++c) p[lane * QCH + c] = half(0.0f); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (qsa_stop == 3) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 3.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } + // ---- acc = P V, 64 dims at a time + { + // K = 32 cells a run (two runs a 64-cell chunk): K = 64 runs compute garbage on the M4 + constexpr auto desc = matmul2d_descriptor(16, 64, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tP(p, dextents(QCH, 16)); + tensor, tensor_inline> tO(pp, dextents(QD + 2, HPG)); + for (int d0 = 0; d0 < QD; d0 += 64) { + auto rv = op.template get_right_input_cooperative_tensor(); + auto tPh = tP.template slice<32, 16>(0, 0); + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + for (int c0 = 0; c0 < nc; c0 += 32) { + if (!coop_k4) { +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); ++i) { + const auto ix = rv.get_multidimensional_index(i); // (cell, dim) + const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; + rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vs[c][d >> 6])) : half(0.0f); + } + } else { + uint vo[8]; + float vs8[8]; + const int cb = c0 + rv.get_multidimensional_index(uint16_t(0))[0]; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = cb + (j & 3) + 16 * (j >> 2); + vo[j] = uint(cid[c]) + uint(d0); + vs8[j] = float(vs[c][d0 >> 6]); + } +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); i += 4) { + const int d = rv.get_multidimensional_index(i)[1], r = i < rv.get_capacity() / 2 ? 0 : 4; +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) rv[i + j] = half(float(vq[vo[r + j] + uint(d)]) * vs8[r + j]); + } + } + tPh = tP.template slice<32, 16>(c0, 0); + op.run(tPh, rv, acc); + } + if (qsa_stop == 4) { // the bisection: the old direct store into the 258-float rows + auto mO = tO.slice(d0, 0); + acc.store(mO); + continue; + } + // through threadgroup memory (s is free now): a matmul2d store into rows of QD + 2 floats loses the + // whole threadgroup on the M4 + tensor, tensor_inline> tT(s, dextents(64, 16)); + simdgroup_barrier(mem_flags::mem_threadgroup); + acc.store(tT); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int it = lane; it < HPG * 16; it += 32) { + const int h = it >> 4, d = (it & 15) * 4; + const float4 v = *(threadgroup const float4*)(s + h * 64 + d); + // odd heads start on an 8-byte boundary, not a float4 boundary. + device float* dst = pp + h * (QD + 2) + d0 + d; + dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w; + } + } + } + if (lane < HPG) { pp[lane * (QD + 2) + QD] = mh; pp[lane * (QD + 2) + QD + 1] = lh; } +} + +// qcur (fp32) -> fp16 for qsa_attn_mm +kernel void qsa_q_f16(device const float* q [[buffer(0)]], device half* qh [[buffer(1)]], constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + if (gid < n) qh[gid] = half(q[gid]); +} + +// qsa_score as a GEMM (fp32 operands, so the selection sees qsa_score's values to fp32 rounding): S[(t, h)][b] = +// qidx[t][h] . pooled[b] for 16 tokens x IH heads (64 rows) by 64 blocks a threadgroup; the epilogue sums relu over +// the heads. The tail block n_bid (the `dead` key, + 1e9 when it has cells) and blocks past it are left to the +// scalar path: the tail is scored here by one simdgroup a token; blocks > n_bid are not written (qsa_topk reads +// b <= n_bid). grid (ceil(max blocks / 64), ceil(T / 16)), 128 threads. +kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const float* pooled [[buffer(1)]], + device const float* dead [[buffer(2)]], device float* scores [[buffer(3)]], + constant QsaArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float ot[64 * 64]; + const int b0 = int(tg.x) * 64, t0 = int(tg.y) * 16; + const uint max_t = min(uint(t0 + 15), a.n_tok - 1u); + const uint max_bid = (a.pos0 + max_t + 1u) >> 2; + // column zero still owns every token's tail score. + if (tg.x != 0 && uint(b0) >= max_bid) return; + const uint nb_all = (a.pos0 + a.n_tok) / 4 + 1; // blocks any token of the chunk can score + tensor, tensor_inline> tQ((device float*)qidx, dextents(ID, int(a.n_tok) * IH)); + tensor, tensor_inline> tK((device float*)pooled, dextents(ID, int(nb_all))); + tensor, tensor_inline> tO(ot, dextents(64, 64)); + constexpr auto desc = matmul2d_descriptor(64, 64, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tQ.slice(0, t0 * IH); + auto mB = tK.slice(0, b0); + op.run(mA, mB, tO); + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint stride = a.max_cells / 4 + 2; + const int j = tid & 63, b = b0 + j; + for (int tt = tid >> 6; tt < 16; tt += 2) { + const int t = t0 + tt; + if (t >= int(a.n_tok)) break; + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + if (uint(b) >= n_bid) continue; + float s = 0.0f; + for (int h = 0; h < IH; ++h) s += max(ot[(tt * IH + h) * 64 + j], 0.0f); + scores[uint(t) * stride + b] = s; + } + // the tail block of each token in this threadgroup's block range (the first column threadgroup does it) + if (tg.x == 0) { + float k[4]; + for (int i = 0; i < 4; ++i) k[i] = dead[lane + 32 * i]; + for (int tt = sg; tt < 16; tt += 4) { + const int t = t0 + tt; + if (t >= int(a.n_tok)) break; + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + float s = 0.0f; + for (int h = 0; h < IH; ++h) { + float d = 0.0f; + for (int i = 0; i < 4; ++i) d = fma(qidx[(t * IH + h) * ID + lane + 32 * i], k[i], d); + d = simd_sum(d); + s += max(d, 0.0f); + } + if ((n_kv & 3) != 0) s += 1e9f; + if (lane == 0) scores[uint(t) * stride + n_bid] = s; + } + } +} + +// qsa_flash: the whole attention of a (kv head g, token t) in one threadgroup (no split-K partials, no merge): the +// width selected cells in chunks of 64 with an online softmax; 4 simdgroups, simdgroup s owns dims 64 s .. 64 s + 63 +// (its slice of S's K sum and its 64 output dims). Out: xq = (acc / l) * sigmoid(gate) as fp16 (attn_output's GEMM +// operand). grid (QKV, n_tok), 128 threads. +kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device const float* qfull [[buffer(6)]], device half* xqh [[buffer(7)]], constant QsaArgs& a [[buffer(8)]], + 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 HPG = QH / QKV; + threadgroup float sp[4][HPG * QCH]; + threadgroup half* p = (threadgroup half*)sp[1]; // P [16][QCH] in fp16 over a dead S partial + threadgroup int cid[QCH]; // coop_k4: byte offsets; otherwise cell ids + threadgroup half4 ks_tg[QCH], vs_tg[QCH]; + threadgroup float fac[16], lsum[16], mrun[16]; + const int g = int(tg.x), t = int(tg.y); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nch = int((width + QCH - 1) / QCH); + if (tid < 16) { mrun[tid] = -INFINITY; lsum[tid] = 0.0f; } + constexpr auto dS = matmul2d_descriptor(16, QCH, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + constexpr auto dO = matmul2d_descriptor(16, 64, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d opS; + matmul2d opO; + tensor, tensor_inline> tQ((device half*)qh + (uint(t) * QH + g * HPG) * QD, + dextents(QD, HPG)); + tensor, tensor_inline> tP(p, dextents(QCH, 16)); + auto rv = opO.template get_right_input_cooperative_tensor(); + auto tPh = tP.template slice<32, 16>(0, 0); + auto acc = opO.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); + // (a second destination tensor, not a copy of acc: the MetalPerformancePrimitives headers of macOS 26.3 cannot copy a + // cooperative tensor - "cannot initialize a parameter of type 'void *' with an rvalue of type 'const void *'", #2323; + // tot is zeroed right below, so the copy was only there for its type) + auto tot = opO.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); // the running output; acc takes one chunk (fresh each chunk: on the M4 a run does not accumulate + // into elements rescaled in place) + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) tot[i] = 0.0f; + const int d0 = sg * 64; + if (tid < 16 - HPG) fac[HPG + tid] = 0.0f; // the padding rows' factor (P's rows are zeroed each chunk) + for (int ch = 0; ch < nch; ++ch) { + const int c0 = ch * QCH, nc = min(QCH, int(width) - c0); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < QCH) { + const int id = tid < nc ? (n_kv <= QBUDGET ? c0 + tid : ids[t * QBUDGET + c0 + tid]) : -1; + cid[tid] = coop_k4 ? int((uint(max(id, 0)) * QKV + g) * QD) : id; + if (id >= 0) { + const uint row4 = (uint(id) * QKV + g) * 4; + ks_tg[tid] = *(device const half4*)(ksc + row4); + vs_tg[tid] = *(device const half4*)(vsc + row4); + } else { + ks_tg[tid] = half4(0.0h); + vs_tg[tid] = half4(0.0h); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // Eight lanes per head; keep the denominator's original sequential sum. + const int h = sg * 4 + (lane >> 3), hlane = lane & 7; + float sc[8]; + // S partial over this simdgroup's 64 dims + { + auto rk = opS.template get_right_input_cooperative_tensor(); + auto tA = tQ.template slice<32, 16>(d0, 0); + auto s = opS.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < s.get_capacity(); ++i) + if (s.is_valid_element(i)) s[i] = 0.0f; + // the lane's eight cells (one a run of four k; the second half repeats them 16 k on): row pointers and + // scales once a chunk, unconditional loads (an empty cell reads row 0 at scale 0, and a +-0 operand sums to + // the same value as the skipped one; the branchy loads cost half the kernel) + uint ko[8]; + float ks8[8]; + if (coop_k4) { +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = rk.get_multidimensional_index(uint16_t(4 * j))[1]; + ko[j] = uint(cid[c]) + uint(d0 + rk.get_multidimensional_index(uint16_t(0))[0]); + ks8[j] = float(ks_tg[c][sg]); + } + } +#pragma clang loop unroll(full) + for (int k0 = d0; k0 < d0 + 64; k0 += 32) { + if (!coop_k4) { + for (uint16_t i = 0; i < rk.get_capacity(); ++i) { + const auto ix = rk.get_multidimensional_index(i); + const int k = k0 + ix[0], id = cid[ix[1]]; + const uint row = uint(id) * QKV + g; + rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ks_tg[ix[1]][k >> 6])) : half(0.0f); + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rk.get_capacity(); i += 4) { + const int k = k0 - d0 + (i >= rk.get_capacity() / 2 ? 16 : 0), j = (i >> 2) & 7; + const half4 w = half4(float4(*(device const char4*)(kq + ko[j] + uint(k))) * ks8[j]); + rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w; + } + tA = tQ.template slice<32, 16>(k0, 0); + opS.run(tA, rk, s); + } + // Preserve ((s0 + s1) + (s2 + s3)) while storing all partials together. + tensor, tensor_inline> tS(sp[sg], dextents(QCH, HPG)); + s.store(tS); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (h < HPG) +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) { + const int q = h * QCH + hlane + 8 * k; + sc[k] = ((sp[0][q] + sp[1][q]) + (sp[2][q] + sp[3][q])) * 0.0625f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); // P overwrites sp[1] + if (h < HPG) { + float m = mrun[h]; +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) + if (hlane + 8 * k < nc) m = max(m, sc[k]); + m = max(m, simd_shuffle_xor(m, 4)); + m = max(m, simd_shuffle_xor(m, 2)); + m = max(m, simd_shuffle_xor(m, 1)); +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) { + const int c = hlane + 8 * k, q = h * QCH + c; + const float e = c < nc ? exp(sc[k] - m) : 0.0f; + p[q] = half(e); + sp[0][q] = e; + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (hlane == 0) { + const float f = exp(mrun[h] - m); + float l = lsum[h] * f; +#pragma clang loop unroll(full) + for (int c = 0; c < QCH; c += 4) { // in order, four loads at a time + const float4 v = *(threadgroup const float4*)(sp[0] + h * QCH + c); + l += v.x; l += v.y; l += v.z; l += v.w; + } + mrun[h] = m; + lsum[h] = l; + fac[h] = f; + } + } else { + for (int c = hlane; c < QCH; c += 8) p[h * QCH + c] = half(0.0f); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // rescale this simdgroup's accumulator rows, then acc += P V over its 64 dims + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + for (int c0 = 0; c0 < nc; c0 += 32) { // K = 32 cells a run (K = 64 is wrong on the M4) + if (!coop_k4) { +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); ++i) { + const auto ix = rv.get_multidimensional_index(i); // (cell, dim) + const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; + rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vs_tg[c][sg])) : half(0.0f); + } + } else { + // a lane holds the same eight cells (two runs of four, 16 apart) at each of its dims: as for rk, row + // pointers and scales once a run, unconditional loads + uint vo[8]; + float vs8[8]; + const int cb = c0 + rv.get_multidimensional_index(uint16_t(0))[0]; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = cb + (j & 3) + 16 * (j >> 2); + vo[j] = uint(cid[c]) + uint(d0); + vs8[j] = float(vs_tg[c][sg]); + } +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); i += 4) { + const int d = rv.get_multidimensional_index(i)[1], r = i < rv.get_capacity() / 2 ? 0 : 4; +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) rv[i + j] = half(float(vq[vo[r + j] + uint(d)]) * vs8[r + j]); + } + } + tPh = tP.template slice<32, 16>(c0, 0); + opO.run(tPh, rv, acc); + } + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) tot[i] = fma(tot[i], fac[acc.get_multidimensional_index(i)[1]], acc[i]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (!coop_k4) { + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) { + const auto ix = acc.get_multidimensional_index(i); // (dim, row) + const int h = ix[1], d = d0 + ix[0]; + if (h < HPG) { + const int hh = g * HPG + h; + const float gate = qfull[(uint(t) * QH + hh) * 2 * QD + QD + d]; + xqh[(uint(t) * QH + hh) * QD + d] = half(tot[i] / lsum[h] * (1.0f / (1.0f + exp(-gate)))); + } + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { // runs of four dims of one row + const auto ix = acc.get_multidimensional_index(i); + const int h = ix[1], d = d0 + ix[0]; + if (h < HPG) { + const int hh = g * HPG + h; + const float l = lsum[h]; + const float4 gate = *(device const float4*)(qfull + (uint(t) * QH + hh) * 2 * QD + QD + d); + const float4 v = float4(tot[i], tot[i + 1], tot[i + 2], tot[i + 3]) / l * (1.0f / (1.0f + exp(-gate))); + *(device half4*)(xqh + (uint(t) * QH + hh) * QD + d) = half4(v); + } + } +} diff --git a/results/round39/astra_r39_qsa/baseline-kernels/quant.metal b/results/round39/astra_r39_qsa/baseline-kernels/quant.metal new file mode 100644 index 0000000..fe667be --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/quant.metal @@ -0,0 +1,316 @@ +// metal/kernels/quant.metal - the weight formats as "units": a unit is UNIT consecutive weights of one row that one +// SIMD lane decodes at a time into raw values q[0..UNIT) and one scale s, so that the unit's dot with x is +// s * sum(q[j] * x[j]). Every GEMV kernel (moe.metal, gemv.metal) is written once over these. +// +// Decode rules transcribed from ggml-quants.c dequantize_row_* (3cf03257) and Strata's Q2_0 contract +// (docs/q2_0-contract.md); the grids and sign tables are ggml-common.h's (prepended by the runtime). The tables a +// format needs are copied to threadgroup memory once per threadgroup (constant-memory lookups with divergent indices +// serialise on Apple GPUs). +// +// Blocks start at 2-byte aligned addresses (18, 66, 74, 82, 98, 110 bytes), so every multi-byte field is read as +// ushorts. + +#define T_Q4_0 2 +#define T_Q3_K 11 +#define T_BF16 30 +#define T_F32 0 +#define T_Q5_0 6 +#define T_Q5_1 7 +#define T_Q8_0 8 +#define T_Q4_K 12 +#define T_Q5_K 13 +#define T_Q6_K 14 +#define T_IQ2_XXS 16 +#define T_IQ2_XS 17 +#define T_IQ3_XXS 18 +#define T_IQ4_NL 20 +#define T_IQ3_S 21 +#define T_IQ2_S 22 +#define T_IQ4_XS 23 +#define T_Q2_0 42 +#define T_Q2_0R 142 // Q2_0 read with raw activations (codes - 1 in the decoder; no prepared x needed) + +template struct fmt; +// QK block elements, BB block bytes, UNIT weights per unit, TAB words of threadgroup table +template <> struct fmt { enum { QK = 64, BB = 18, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 64, BB = 18, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 98, UNIT = 16, TAB = 256 }; }; +template <> struct fmt { enum { QK = 256, BB = 66, UNIT = 16, TAB = 512 }; }; +template <> struct fmt { enum { QK = 256, BB = 74, UNIT = 16, TAB = 1024 }; }; +template <> struct fmt { enum { QK = 256, BB = 82, UNIT = 16, TAB = 2048 }; }; +template <> struct fmt { enum { QK = 256, BB = 110, UNIT = 16, TAB = 512 }; }; +template <> struct fmt { enum { QK = 32, BB = 18, UNIT = 32, TAB = 256 }; }; +template <> struct fmt { enum { QK = 256, BB = 136, UNIT = 32, TAB = 256 }; }; +template <> struct fmt { enum { QK = 256, BB = 110, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 144, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 176, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 210, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 16, BB = 32, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 16, BB = 64, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 32, BB = 34, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 32, BB = 18, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 32, BB = 22, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 32, BB = 24, UNIT = 32, TAB = 1 }; }; + +static inline uint rd32(device const uchar* p) { + device const ushort* s = (device const ushort*)p; + return uint(s[0]) | (uint(s[1]) << 16); +} +static inline float rdh(device const uchar* p) { return float(*(device const half*)p); } + +// ---- threadgroup tables ------------------------------------------------------------------------------------------- +// layout in `tab` (uint words): [0, G) the grid; IQ4_NL/XS: 256 half2. +template static inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth); + +// QUANT_CTAB (function constant 23): the i-quant grids read from ggml-common's constant arrays +// instead of a threadgroup copy (no table load and no barrier wait at the threadgroup's start); the same values. +constant bool QUANT_CTAB [[function_constant(23)]]; +constant bool ctab = is_function_constant_defined(QUANT_CTAB) && QUANT_CTAB; +// ksigns_iq2xs[j] = j | parity(j) << 7 (an even count of minus signs): the low 4 bits of run hh (0 or 1). +static inline uint sign4_iq2xs(uint j, int hh) { + return hh ? ((j >> 4) | ((popcount(j) & 1u) << 3)) : j; +} +static inline uint g_iq3xxs(threadgroup const uint* tab, uint i) { return ctab ? iq3xxs_grid[i] : tab[i]; } +static inline uint g_iq3s(threadgroup const uint* tab, uint i) { return ctab ? iq3s_grid[i] : tab[i]; } +static inline uint g_iq2xxs(threadgroup const uint* tab, uint i, int w) { return ctab ? uint(iq2xxs_grid[i] >> (32 * w)) : tab[2 * i + w]; } +static inline uint g_iq2xs(threadgroup const uint* tab, uint i, int w) { return ctab ? uint(iq2xs_grid[i] >> (32 * w)) : tab[2 * i + w]; } +static inline uint g_iq2s(threadgroup const uint* tab, uint i, int w) { return ctab ? uint(iq2s_grid[i] >> (32 * w)) : tab[2 * i + w]; } + +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} + +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 256; i += nth) tab[i] = iq3xxs_grid[i]; +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 512; i += nth) tab[i] = iq3s_grid[i]; +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 256; i += nth) { tab[2 * i] = uint(iq2xxs_grid[i]); tab[2 * i + 1] = uint(iq2xxs_grid[i] >> 32); } +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 512; i += nth) { tab[2 * i] = uint(iq2xs_grid[i]); tab[2 * i + 1] = uint(iq2xs_grid[i] >> 32); } +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 1024; i += nth) { tab[2 * i] = uint(iq2s_grid[i]); tab[2 * i + 1] = uint(iq2s_grid[i] >> 32); } +} +// IQ4: a byte -> (codebook[low nibble], codebook[high nibble]) as half2, so one lookup gives two weights +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + for (int i = tid; i < 256; i += nth) + tab[i] = as_type(half2(half(kvalues_iq4nl[i & 15]), half(kvalues_iq4nl[i >> 4]))); +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + load_tab(tab, tid, nth); +} + +// ---- unit decoders: return s, fill q[UNIT] ------------------------------------------------------------------------ +// ui = unit index inside the block (0 .. QK/UNIT - 1). +template static inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab); + +// Q2_0: 64 = 4 units of 16; code c (2 bits, LSB first) -> c - 1; scale d. +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint*) { + device const ushort* p = (device const ushort*)(b + 2 + 4 * ui); + const ushort a = p[0], c = p[1]; + for (int j = 0; j < 8; ++j) { + q[j] = float((a >> (2 * j)) & 3) - 1.0f; + q[8 + j] = float((c >> (2 * j)) & 3) - 1.0f; + } + return rdh(b); +} + +// 8 signed grid bytes from a 4-byte grid word pair helper +static inline void put_signed4(thread float* q, uint g, uint s) { + const uchar4 v = as_type(g); + q[0] = (s & 1) ? -float(v[0]) : float(v[0]); + q[1] = (s & 2) ? -float(v[1]) : float(v[1]); + q[2] = (s & 4) ? -float(v[2]) : float(v[2]); + q[3] = (s & 8) ? -float(v[3]) : float(v[3]); +} + +// IQ3_XXS: 256 = 8 x 32; unit = half of a 32 (two grid pairs). +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + const uint aux = rd32(b + 2 + 64 + 4 * ib), s0 = aux >> (14 * h); + device const uchar* qs = b + 2 + 4 * ui; + #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { + const uint s = (s0 >> (7 * l2)) & 127; + put_signed4(q + 8 * l2, g_iq3xxs(tab, qs[2 * l2]), s); + put_signed4(q + 8 * l2 + 4, g_iq3xxs(tab, qs[2 * l2 + 1]), sign4_iq2xs(s, 1)); + } + return rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f); +} + +// IQ2_XXS: per 32: two u32 (4 grid indices, then 4 x 7 sign bits + 4-bit scale). +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + device const uchar* qs = b + 2 + 4 * ib + 2 * ui; + const uint a1 = rd32(b + 2 + 8 * ib + 4), s0 = a1 >> (14 * h); + #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { + const uint idx = qs[l2], s = (s0 >> (7 * l2)) & 127; + const uint2 g = ctab ? as_type(iq2xxs_grid[idx]) : ((threadgroup const uint2*)tab)[idx]; + put_signed4(q + 8 * l2, g.x, s); + put_signed4(q + 8 * l2 + 4, g.y, sign4_iq2xs(s, 1)); + } + return rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f); +} + +// IQ2_XS: uint16 qs[32] (9-bit grid index + 7-bit sign index), scales[8] nibbles. +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + device const ushort* qs = (device const ushort*)(b + 2) + 2 * ui; + #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { + const uint v = qs[l2]; + const uint idx = v & 511, s = v >> 9; + const uint2 g = ctab ? as_type(iq2xs_grid[idx]) : ((threadgroup const uint2*)tab)[idx]; + put_signed4(q + 8 * l2, g.x, s); + put_signed4(q + 8 * l2 + 4, g.y, sign4_iq2xs(s, 1)); + } + const uint sc = b[66 + ib]; + return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); +} + +// IQ2_S: qs[32] low 8 bits of the 10-bit index, signs[32], qh[8] (2 high bits per group of 8), scales[8]. +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + const uint q2 = *(device const ushort*)(b + 2 + 2 * ui); + const uint s2 = *(device const ushort*)(b + 34 + 2 * ui); + const uint qh = uint(b[66 + ib]) >> (4 * h); + #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { + const uint idx = ((q2 >> (8 * l2)) & 255) | ((qh << (8 - 2 * l2)) & 0x300); + const uint s = (s2 >> (8 * l2)) & 255; + const uint2 g = ctab ? as_type(iq2s_grid[idx]) : ((threadgroup const uint2*)tab)[idx]; + put_signed4(q + 8 * l2, g.x, s); + put_signed4(q + 8 * l2 + 4, g.y, s >> 4); + } + const uint sc = b[74 + ib]; + return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); +} + +// IQ3_S: qs[64] low 8 bits of 9-bit indices, qh[8], signs[32], scales[4] (nibble per 32). +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + const uint s2 = *(device const ushort*)(b + 74 + 2 * ui); + const uint qh = b[66 + ib]; + #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { + const int l = 2 * h + l2; + const uint i1 = uint(b[2 + 8 * ib + 2 * l]) | ((qh << (8 - 2 * l)) & 256); + const uint i2 = uint(b[2 + 8 * ib + 2 * l + 1]) | ((qh << (7 - 2 * l)) & 256); + const uint s = (s2 >> (8 * l2)) & 255; + put_signed4(q + 8 * l2, g_iq3s(tab, i1), s); + put_signed4(q + 8 * l2 + 4, g_iq3s(tab, i2), s >> 4); + } + const uint sc = b[106 + (ib >> 1)]; + return rdh(b) * fma(float((sc >> (4 * (ib & 1))) & 15), 2.0f, 1.0f); +} + +// IQ4_NL: 32 = 1 unit; low nibbles are elements 0..15, high 16..31 (one half2 lookup per byte). +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint* tab) { + device const ushort* p = (device const ushort*)(b + 2); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); + q[2 * j] = float(w0.x); + q[2 * j + 16] = float(w0.y); + q[2 * j + 1] = float(w1.x); + q[2 * j + 17] = float(w1.y); + } + return rdh(b); +} + +// IQ4_XS: 256 = 8 units of 32; 6-bit scale from scales_l nibble + 2 bits of scales_h, minus 32. +static inline float iq4xs_scale(device const uchar* b, int ui) { + const uint sh = *(device const ushort*)(b + 2); + const int ls = int((b[4 + (ui >> 1)] >> (4 * (ui & 1))) & 15) | int(((sh >> (2 * ui)) & 3) << 4); + return rdh(b) * float(ls - 32); +} +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + device const ushort* p = (device const ushort*)(b + 8 + 16 * ui); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); + q[2 * j] = float(w0.x); + q[2 * j + 16] = float(w0.y); + q[2 * j + 1] = float(w1.x); + q[2 * j + 17] = float(w1.y); + } + return iq4xs_scale(b, ui); +} + +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { + device const ushort* p = (device const ushort*)(b + 2); + for (int j = 0; j < 16; ++j) { + const ushort v = p[j]; + q[2 * j] = float(as_type(uchar(v & 255))); + q[2 * j + 1] = float(as_type(uchar(v >> 8))); + } + return rdh(b); +} + +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { + device const ushort* p = (device const ushort*)(b + 2); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + q[2 * j] = float(v & 15) - 8.0f; + q[2 * j + 16] = float((v >> 4) & 15) - 8.0f; + q[2 * j + 1] = float((v >> 8) & 15) - 8.0f; + q[2 * j + 17] = float(v >> 12) - 8.0f; + } + return rdh(b); +} + +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { + const uint hb = rd32(b + 2); + device const ushort* p = (device const ushort*)(b + 6); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const int j0 = 2 * j, j1 = 2 * j + 1; + q[j0] = float(int(v & 15) | int(((hb >> j0) & 1) << 4)) - 16.0f; + q[j0 + 16] = float(int((v >> 4) & 15) | int(((hb >> (j0 + 16)) & 1) << 4)) - 16.0f; + q[j1] = float(int((v >> 8) & 15) | int(((hb >> j1) & 1) << 4)) - 16.0f; + q[j1 + 16] = float(int(v >> 12) | int(((hb >> (j1 + 16)) & 1) << 4)) - 16.0f; + } + return rdh(b); +} + +// Q5_1: 32 = 1 unit; q = the 5-bit code (no offset; the min m is applied through the x sums), returns d +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { + const uint hb = rd32(b + 4); + device const ushort* p = (device const ushort*)(b + 8); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const int j0 = 2 * j, j1 = 2 * j + 1; + q[j0] = float(int(v & 15) | int(((hb >> j0) & 1) << 4)); + q[j0 + 16] = float(int((v >> 4) & 15) | int(((hb >> (j0 + 16)) & 1) << 4)); + q[j1] = float(int((v >> 8) & 15) | int(((hb >> j1) & 1) << 4)); + q[j1 + 16] = float(int(v >> 12) | int(((hb >> (j1 + 16)) & 1) << 4)); + } + return rdh(b); +} + +// sum over groups of LPR consecutive lanes (LPR = 1, 2, 4, 8, 16 or 32); the result is valid in every lane +template +static inline float lane_sum(float v) { + if (LPR == 32) return simd_sum(v); + for (ushort o = LPR / 2; o > 0; o >>= 1) v += simd_shuffle_xor(v, o); + return v; +} diff --git a/results/round39/astra_r39_qsa/baseline-kernels/router.metal b/results/round39/astra_r39_qsa/baseline-kernels/router.metal new file mode 100644 index 0000000..5c6bf59 --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/router.metal @@ -0,0 +1,173 @@ +// metal/kernels/router.metal - the router GEMV and moe_route in ONE dispatch (after quant/units/moe/gemv.metal): the +// logits rows as gemv's (LPR = 32 lanes a row, KS = 1: bitwise gemv__32_1_*_1's), 16 rows a threadgroup of 512 +// threads; every threadgroup bumps a device counter after its rows (device-scope fences around it), and the LAST one +// runs moe_route's body (the softmax / top-10 / weights, then the plan) with its 512 threads (= n_expert) and resets +// the counter. No threadgroup waits for another. ids / w / plan are bitwise moe_route's. +// Buffers: W, xp, s16, logits, args (GemvArgs: n_out = n_expert), ids, w, plan, MoeArgs, counter (one uint, zero). + +#ifndef STRATA_MOCE_FC_METAL +#define STRATA_MOCE_FC_METAL +constant int kMoceProtectFc [[function_constant(24)]]; +constant int kMoceProtect = is_function_constant_defined(kMoceProtectFc) ? kMoceProtectFc : 3; +constant uint kMoceScaleR1BitsFc [[function_constant(25)]]; +constant float kMoceScaleR1 = is_function_constant_defined(kMoceScaleR1BitsFc) ? as_type(kMoceScaleR1BitsFc) : 1.0f; +#endif +static inline void route_body(device const float* logits, device int* ids, device float* w, device int* plan, + constant MoeArgs& a, device const uint* smap, ushort tid, ushort nth, ushort sg, ushort lane, + threadgroup int* s_grp, threadgroup int* s_ent) { + if (sg < a.n_tok) { + device const float* l = logits + sg * 512; + float v[16]; + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + v[i] = idx < int(a.n_expert) ? l[idx] : -INFINITY; + } + float mx = -INFINITY; + for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + float sum = 0.0f; + for (int i = 0; i < 16; ++i) { + v[i] = (lane + 32 * i < int(a.n_expert)) ? exp(v[i] - mx) : 0.0f; + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + for (int i = 0; i < 16; ++i) { + if (lane + 32 * i >= int(a.n_expert)) { + v[i] = -INFINITY; + } else { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + } + const bool draft_moce = (a.pad & 0x80000000u) != 0u && sg > 0; + const float base_moce = (a.pad & 0x7FFFFFFFu) != 0u ? as_type(a.pad & 0x7FFFFFFFu) : 1.0f; + const float moce_scale = draft_moce ? (base_moce * base_moce) : base_moce; + const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f; + const int r_prot = draft_moce ? 1 : kMoceProtect; + const float r1_scale = kMoceScaleR1; + const float rcp_r1 = r1_scale > 1.0f ? (1.0f / r1_scale) : 1.0f; + const float r_step_scale = (r1_scale > 1.0f && moce_scale > 1.0f) ? (moce_scale * rcp_r1) : moce_scale; + ushort scaled_mask = 0; + float selected = 0.0f, ssum = 0.0f; + for (int rank = 0; rank < 10; ++rank) { + if (r1_scale > 1.0f && r_prot > 1) { + if (rank == 1 && smap != nullptr) { + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) { + v[i] *= r1_scale; + scaled_mask |= ushort(1u << i); + } + } + } else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) { + for (int i = 0; i < 16; ++i) { + if ((scaled_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) { + v[i] *= r_step_scale; + } + } + } + } else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) { + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) { + v[i] *= moce_scale; + scaled_mask |= ushort(1u << i); + } + } + } + float best = v[0]; + int ex = lane; + #pragma clang loop unroll(full) + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + float true_prob = best; + const int min_scaled_rank = (r1_scale > 1.0f && r_prot > 1) ? 1 : r_prot; + if (rank >= min_scaled_rank && (moce_scale > 1.0f || r1_scale > 1.0f) && best > 0.0f) { + const ushort win_lane = ushort(ex & 31); + 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) { + // 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); + } + } + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) v[i] = (ex == lane + 32 * i) ? -INFINITY : v[i]; + if ((ex & 31) == lane) { + ids[sg * 10 + rank] = ex; + ssum += true_prob; + } + if (rank == lane) selected = true_prob; + } + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + ssum = max(ssum, 6.103515625e-5f); + const float inv = 1.0f / ssum; + if (lane < 10) w[sg * 10 + lane] = selected * inv; + } + threadgroup_barrier(mem_flags::mem_device); + plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent, smap); +} + +template +kernel void router_route(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + 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)]], + device atomic_uint* counter [[buffer(9)]], device const uint* smap [[buffer(10)]], + uint tgx [[threadgroup_position_in_grid]], uint ngx [[threadgroups_per_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort nth [[threads_per_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int UNIT = fmt::UNIT; + threadgroup uint tab[fmt::TAB]; + threadgroup int s_grp[64], s_ent[64]; + threadgroup uint last; + load_tab(tab, tid, nth); + threadgroup_barrier(mem_flags::mem_threadgroup); + // the GEMV: simdgroup sg -> row tg.x * 16 + sg, 32 lanes over its units + const int row = int(tgx) * 16 + sg; + const int rr = min(row, int(ga.n_out) - 1); + device const uchar* rp = W + (ulong)rr * ga.row_bytes; + const int U = int(ga.n_in) / UNIT, s16n = int(ga.n_in) / 16; + float acc[8]; +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[t] = 0.0f; + for (int u = lane; u < U; u += 32) { + float q[UNIT]; + float4 sc; + Dec::run(rp, u, q, sc, tab); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(ga.x_stride), s16 + t * s16n, u, xs, xsum); + acc[t] += unit_combine(q, sc, xs, xsum); + } + } +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum<32>(acc[t]); + if (lane == 0 && row < int(ga.n_out)) logits[t * int(ga.y_stride) + row] = v; + } + // the last threadgroup to finish routes + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + threadgroup_barrier(mem_flags::mem_device); + if (tid == 0) last = atomic_fetch_add_explicit(counter, 1u, memory_order_relaxed) == ngx - 1 ? 1u : 0u; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (!last) return; + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + route_body(logits, ids, w, plan, a, smap, tid, nth, sg, lane, s_grp, s_ent); + if (tid == 0) atomic_store_explicit(counter, 0u, memory_order_relaxed); +} + +template [[host_name("router_route_30")]] kernel decltype(router_route) router_route; +template [[host_name("router_route_0")]] kernel decltype(router_route) router_route; + +template [[host_name("router_route_8")]] kernel decltype(router_route) router_route; diff --git a/results/round39/astra_r39_qsa/baseline-kernels/rt_bench.metal b/results/round39/astra_r39_qsa/baseline-kernels/rt_bench.metal new file mode 100644 index 0000000..bdce780 --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/rt_bench.metal @@ -0,0 +1,67 @@ +// metal/kernels/rt_bench.metal - bandwidth and dispatch-cost probes (rt_bench.cpp). + +// Grid-stride streaming read: thread i reads src[i], src[i + nt], ... (coalesced across the grid). +kernel void k_read(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]], uint nt [[threads_per_grid]]) { + uint4 a = 0; + for (uint i = gid; i < n; i += nt) a += src[i]; + out[gid] = a.x ^ a.y ^ a.z ^ a.w; +} + +// The same, 4 independent loads in flight per iteration. +kernel void k_read_u4(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]], uint nt [[threads_per_grid]]) { + uint4 a = 0, b = 0, c = 0, d = 0; + uint i = gid; + for (; i + 3 * nt < n; i += 4 * nt) { + a += src[i]; + b += src[i + nt]; + c += src[i + 2 * nt]; + d += src[i + 3 * nt]; + } + for (; i < n; i += nt) a += src[i]; + a += b + c + d; + out[gid] = a.x ^ a.y ^ a.z ^ a.w; +} + +// Contiguous chunk per threadgroup (how a GEMV row block reads): threadgroup g reads [g*chunk, (g+1)*chunk). +kernel void k_read_chunk(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + constant uint& n [[buffer(2)]], constant uint& chunk [[buffer(3)]], + uint tid [[thread_position_in_threadgroup]], uint tgs [[threads_per_threadgroup]], + uint g [[threadgroup_position_in_grid]], uint gid [[thread_position_in_grid]]) { + uint4 a = 0, b = 0; + const uint base = g * chunk, end = min(base + chunk, n); + uint i = base + tid; + for (; i + tgs < end; i += 2 * tgs) { + a += src[i]; + b += src[i + tgs]; + } + for (; i < end; i += tgs) a += src[i]; + a += b; + out[gid] = a.x ^ a.y ^ a.z ^ a.w; +} + +// Read + write (a state update): dst[i] = src[i] * 1.0001 for float4. +kernel void k_copy(device const float4* src [[buffer(0)]], device float4* dst [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]], uint nt [[threads_per_grid]]) { + for (uint i = gid; i < n; i += nt) dst[i] = src[i] * 1.0001f; +} + +// Dispatch-cost probes. +kernel void k_empty(device uint* out [[buffer(0)]], uint gid [[thread_position_in_grid]]) { + if (gid == 0xFFFFFFFFu) out[0] = 1; +} + +// A dependent chain: each dispatch reads what the previous one wrote. +kernel void k_chain(device uint* buf [[buffer(0)]], uint gid [[thread_position_in_grid]]) { + if (gid == 0) buf[0] = buf[0] + 1; +} + +// A small but real kernel: 2560 floats read, 2560 written (an elementwise op on one token). +kernel void k_small(device const float* x [[buffer(0)]], device float* y [[buffer(1)]], + uint gid [[thread_position_in_grid]]) { + if (gid < 2560) y[gid] = x[gid] * 0.5f + y[gid]; +} diff --git a/results/round39/astra_r39_qsa/baseline-kernels/sample.metal b/results/round39/astra_r39_qsa/baseline-kernels/sample.metal new file mode 100644 index 0000000..4531ca2 --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/sample.metal @@ -0,0 +1,61 @@ +// metal/kernels/sample.metal - greedy sampling over the LM head's logits: the largest logit, ties to the lowest +// token id (NaN never wins). Two stages: argmax_part (grid (parts, n_tok), 256 threads) writes one (value, id) +// pair per threadgroup; argmax_final (grid (n_tok), 256 threads) reduces them in order. + +struct ArgArgs { + uint n_vocab; + uint parts; +}; + +static inline void better(thread float& bv, thread int& bi, float v, int i) { + if (v > bv || (v == bv && i < bi)) { bv = v; bi = i; } +} + +static inline void tg_argmax(thread float& bv, thread int& bi, threadgroup float* sv, threadgroup int* si, ushort sg, + ushort lane) { +#pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) { + const float ov = simd_shuffle_xor(bv, m); + const int oi = simd_shuffle_xor(bi, m); + better(bv, bi, ov, oi); + } + if (lane == 0) { sv[sg] = bv; si[sg] = bi; } + threadgroup_barrier(mem_flags::mem_threadgroup); + bv = -INFINITY; + bi = 0x7fffffff; +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) better(bv, bi, sv[k], si[k]); +} + +kernel void argmax_part(device const float* logits [[buffer(0)]], device float* pv [[buffer(1)]], + device int* pi [[buffer(2)]], constant ArgArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float sv[8]; + threadgroup int si[8]; + const uint t = tg.y, per = (a.n_vocab + a.parts - 1) / a.parts; + const uint b0 = tg.x * per, b1 = min(b0 + per, a.n_vocab); + device const float* l = logits + t * a.n_vocab; + float bv = -INFINITY; + int bi = 0x7fffffff; + for (uint i = b0 + tid; i < b1; i += 256) { + const float v = l[i]; + if (!isnan(v)) better(bv, bi, v, int(i)); + } + tg_argmax(bv, bi, sv, si, sg, lane); + if (tid == 0) { pv[t * a.parts + tg.x] = bv; pi[t * a.parts + tg.x] = bi; } +} + +kernel void argmax_final(device const float* pv [[buffer(0)]], device const int* pi [[buffer(1)]], + device int* out [[buffer(2)]], constant ArgArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float sv[8]; + threadgroup int si[8]; + const uint t = tg.x; + float bv = -INFINITY; + int bi = 0x7fffffff; + for (uint k = tid; k < a.parts; k += 256) better(bv, bi, pv[t * a.parts + k], pi[t * a.parts + k]); + tg_argmax(bv, bi, sv, si, sg, lane); + if (tid == 0) out[t] = bi; +} diff --git a/results/round39/astra_r39_qsa/baseline-kernels/units.metal b/results/round39/astra_r39_qsa/baseline-kernels/units.metal new file mode 100644 index 0000000..25d0761 --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-kernels/units.metal @@ -0,0 +1,328 @@ +// metal/kernels/units.metal - the dot of one unit with the window's tokens (after quant.metal). +// +// X::load(xrow, u, xs, xsum) prepares unit u's x slice of one token (xrow = the token's row start), once per unit +// and shared by every weight row the lane reads: the UNIT x values in the order the decoder produces weights, and +// up to two partial sums (for the formats with an offset or a min). Where the codes are read with the mask trick +// (an AND keeps a field in place, so the code arrives multiplied by 2^k), the x values are pre-scaled by 2^-k here: +// exact in float. + +// the threadgroup-memory form (x staged once per threadgroup; the same values) +static inline void load16(threadgroup const float* p, thread float* xs) { + threadgroup const float4* v = (threadgroup const float4*)p; + for (int j = 0; j < 4; ++j) { + const float4 a = v[j]; + xs[4 * j] = a.x; + xs[4 * j + 1] = a.y; + xs[4 * j + 2] = a.z; + xs[4 * j + 3] = a.w; + } +} +static inline void load16(device const float* p, thread float* xs) { + device const float4* v = (device const float4*)p; + for (int j = 0; j < 4; ++j) { + const float4 a = v[j]; + xs[4 * j] = a.x; + xs[4 * j + 1] = a.y; + xs[4 * j + 2] = a.z; + xs[4 * j + 3] = a.w; + } +} + +// ---- prepared x (xprep): XP holds, per token row of n floats, the kinds +// XK_RAW x (every format without the mask trick) +// XK_Q2 x * 4^-(i % 8) (Q2_0: 2-bit fields in place in a 16-bit word) +// XK_P4 x * 2^-(8 (i % 4)) (Q4_K, Q5_K: byte fields in place in a 32-bit word) +// XK_P2 x * 2^-(8 (i % 2)) (Q6_K, Q3_K: byte fields in place in a 16-bit word) +// and S16 = the sum of every 16 consecutive raw x (the formats with an offset or a min). +#define XK_RAW 0 +#define XK_Q2 1 +#define XK_P4 2 +#define XK_P2 3 +template struct XKind { enum { K = XK_RAW }; }; +template <> struct XKind { enum { K = XK_Q2 }; }; +template <> struct XKind { enum { K = XK_P4 }; }; +template <> struct XKind { enum { K = XK_P4 }; }; +template <> struct XKind { enum { K = XK_P2 }; }; +template <> struct XKind { enum { K = XK_P2 }; }; + +// X::load(xp, s16, u, xs, xsum): unit u's prepared x (xp = the token's row of the format's kind, s16 = its sums) +template struct X { + template static inline void load(P xp, P, int u, thread float* xs, thread float* xsum) { + constexpr int UNIT = fmt::UNIT; + for (int k = 0; k < UNIT; k += 16) load16(xp + u * UNIT + k, xs + k); + xsum[0] = xsum[1] = 0.0f; + } +}; +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + load16(xp + 16 * u, xs); + xsum[0] = s16[u]; + xsum[1] = 0.0f; + } +}; +// Q4_K / Q5_K: unit (block b, chunk c, half m): run0 at 256b + 64c + 16m, run1 = run0 + 32 +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + const int p0 = 256 * (u >> 3) + 64 * ((u >> 1) & 3) + 16 * (u & 1); + load16(xp + p0, xs); + load16(xp + p0 + 32, xs + 16); + xsum[0] = s16[p0 >> 4]; + xsum[1] = s16[(p0 >> 4) + 2]; + } +}; +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + X::load(xp, s16, u, xs, xsum); + } +}; +// Q6_K: unit (block b, half n, m, p): run0 at 256b + 128n + 32p + 16m, run1 = run0 + 64 +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + const int p0 = 256 * (u >> 3) + 128 * ((u >> 2) & 1) + 32 * (u & 1) + 16 * ((u >> 1) & 1); + load16(xp + p0, xs); + load16(xp + p0 + 64, xs + 16); + xsum[0] = s16[p0 >> 4]; + xsum[1] = s16[(p0 >> 4) + 4]; + } +}; +// Q5_1: 32 contiguous raw x, the two per-16 sums formed from them (the GEMVs that read Q5_1 - the routed down - get +// h without prepared sums) +template <> struct X { + template static inline void load(P xp, P, int u, thread float* xs, thread float* xsum) { + load16(xp + 32 * u, xs); + load16(xp + 32 * u + 16, xs + 16); + float a = 0.0f, b = 0.0f; + for (int j = 0; j < 16; ++j) { a += xs[j]; b += xs[16 + j]; } + xsum[0] = a; + xsum[1] = b; + } +}; +// Q3_K: 32 contiguous, sums per 16 +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + load16(xp + 32 * u, xs); + load16(xp + 32 * u + 16, xs + 16); + xsum[0] = s16[2 * u]; + xsum[1] = s16[2 * u + 1]; + } +}; + +// xprep: one thread per 16 consecutive values of one token row; XP = [4 kinds][n_tok][n] then S16 [n_tok][n / 16]. +// raw and plane are 16-float aligned: the four planes are written as float4s (mtp.metal's mtp_prepare, bit for bit) +kernel void xprep(device const float* x [[buffer(0)]], device float* xp [[buffer(1)]], constant uint& n [[buffer(2)]], + constant uint& n_tok [[buffer(3)]], uint gid [[thread_position_in_grid]]) { + const uint per = n / 16; + if (gid >= per * n_tok) return; + const uint t = gid / per, i = gid % per; + device const float4* src = (device const float4*)(x + t * n + 16 * i); + const uint plane = n_tok * n; + device float* raw = xp + t * n + 16 * i; + device float4* r0 = (device float4*)raw; + device float4* r1 = (device float4*)(raw + plane); + device float4* r2 = (device float4*)(raw + 2 * plane); + device float4* r3 = (device float4*)(raw + 3 * plane); + const float4 q2lo = float4(1.0f, 0.25f, 0.0625f, 0.015625f), + q2hi = float4(0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f), + p4 = float4(1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f); + float s0 = 0, s1 = 0; +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) { + const float4 v = src[j]; + s0 += v.x; s0 += v.z; + s1 += v.y; s1 += v.w; + r0[j] = v; + r1[j] = v * ((j & 1) ? q2hi : q2lo); + r2[j] = v * p4; + r3[j] = float4(v.x, v.y * 0.00390625f, v.z, v.w * 0.00390625f); + } + xp[4 * plane + t * per + i] = s0 + s1; +} + +// ---- decoders: Dec::run(row, u, q, sc) decodes unit u of a row into q[UNIT] and sc = (d1, d2, m1, m2), so that +// for any token: dot = d1 * sum_{j < UNIT/2} q[j] xs[j] + d2 * sum_{j >= UNIT/2} q[j] xs[j] - (m1 xsum0 + m2 xsum1) +// with xs, xsum from X::load. A unit's weights are decoded once per row and applied to every token. +template struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT, UPB = fmt::QK / UNIT; + const float s = deq(row + (u / UPB) * fmt::BB, u % UPB, q, tab); + sc = float4(s, s, 0.0f, 0.0f); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 2) * 18; + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + const ushort a = p[0], c = p[1]; + for (int j = 0; j < 8; ++j) { + q[j] = float(a & ushort(3u << (2 * j))); + q[8 + j] = float(c & ushort(3u << (2 * j))); + } + const float d = rdh(b); + sc = float4(d, d, d, 0.0f); + } +}; + +// Q2_0 with raw x: q = c - 1 (shift, and, convert, subtract), one scale +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 2) * 18; + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + const ushort a = p[0], c = p[1]; + for (int j = 0; j < 8; ++j) { + q[j] = float((a >> (2 * j)) & 3) - 1.0f; + q[8 + j] = float((c >> (2 * j)) & 3) - 1.0f; + } + const float d = rdh(b); + sc = float4(d, d, 0.0f, 0.0f); + } +}; + +static inline void scale_min_k4(int j, device const uchar* q, thread float& d, thread float& m) { + if (j < 4) { + d = float(q[j] & 63); + m = float(q[j + 4] & 63); + } else { + d = float((q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4)); + m = float((q[j + 4] >> 4) | ((q[j] >> 6) << 4)); + } +} +static inline void bytes4(uint w, thread float* q) { + q[0] = float(w & 0xFFu); + q[1] = float(w & 0xFF00u); + q[2] = float(w & 0xFF0000u); + q[3] = float(w & 0xFF000000u); +} +static inline void bytes2(ushort w, thread float* q) { + q[0] = float(w & 0xFF); + q[1] = float(w & 0xFF00); +} + +template <> struct Dec { // w = d q + m: sum = d (q . x) + m (sum x) + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint* tab) { + device const uchar* b = row + u * 24; + const float d = deq(b, 0, q, tab), m = rdh(b + 2); + sc = float4(d, d, -m, -m); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 144; + const int c = (u >> 1) & 3, m = u & 1; + const float d = rdh(b), dm = rdh(b + 2); + float s1, m1, s2, m2; + scale_min_k4(2 * c, b + 4, s1, m1); + scale_min_k4(2 * c + 1, b + 4, s2, m2); + device const uint* qp = (device const uint*)(b + 16 + 32 * c + 16 * m); + for (int w = 0; w < 4; ++w) { + const uint v = qp[w]; + bytes4(v & 0x0F0F0F0Fu, q + 4 * w); + bytes4((v >> 4) & 0x0F0F0F0Fu, q + 16 + 4 * w); + } + sc = float4(d * s1, d * s2, dm * m1, dm * m2); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 176; + const int c = (u >> 1) & 3, m = u & 1; + const float d = rdh(b), dm = rdh(b + 2); + float s1, m1, s2, m2; + scale_min_k4(2 * c, b + 4, s1, m1); + scale_min_k4(2 * c + 1, b + 4, s2, m2); + device const uint* qp = (device const uint*)(b + 48 + 32 * c + 16 * m); + device const uint* hp = (device const uint*)(b + 16 + 16 * m); + for (int w = 0; w < 4; ++w) { + const uint v = qp[w], h = hp[w]; + bytes4((v & 0x0F0F0F0Fu) | (((h >> (2 * c)) & 0x01010101u) << 4), q + 4 * w); + bytes4(((v >> 4) & 0x0F0F0F0Fu) | (((h >> (2 * c + 1)) & 0x01010101u) << 4), q + 16 + 4 * w); + } + sc = float4(d * s1, d * s2, dm * m1, dm * m2); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 210; + const int n = (u >> 2) & 1, m = (u >> 1) & 1, p = u & 1; + device const ushort* ql = (device const ushort*)(b + 64 * n + 32 * p + 16 * m); + device const ushort* qh = (device const ushort*)(b + 128 + 32 * n + 16 * m); + device const char* scs = (device const char*)(b + 192); + const float d = rdh(b + 208); + for (int w = 0; w < 8; ++w) { + const ushort l = ql[w], h = qh[w]; + bytes2((l & 0x0F0F) | (((h >> (2 * p)) & 0x0303) << 4), q + 2 * w); + bytes2(((l >> 4) & 0x0F0F) | (((h >> (2 * p + 4)) & 0x0303) << 4), q + 16 + 2 * w); + } + const float d1 = d * float(scs[8 * n + m + 2 * p]), d2 = d * float(scs[8 * n + m + 2 * p + 4]); + sc = float4(d1, d2, 32.0f * d1, 32.0f * d2); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 110; + const int n = (u >> 2) & 1, j = u & 3; + // qs and hmask as packed (2-byte aligned) vectors: blocks are 110 bytes, so odd blocks are only 2-aligned + device const packed_ushort4* qv = (device const packed_ushort4*)(b + 32 + 32 * n); + device const packed_ushort4* hv = (device const packed_ushort4*)b; + ushort qq[16], hm[16]; + for (int v = 0; v < 4; ++v) { + const ushort4 a4 = ushort4(qv[v]), h4 = ushort4(hv[v]); + qq[4 * v] = a4.x; qq[4 * v + 1] = a4.y; qq[4 * v + 2] = a4.z; qq[4 * v + 3] = a4.w; + hm[4 * v] = h4.x; hm[4 * v + 1] = h4.y; hm[4 * v + 2] = h4.z; hm[4 * v + 3] = h4.w; + } + device const uchar* s = b + 96; + const int i1 = 8 * n + 2 * j, i2 = i1 + 1; + const int l1 = (i1 < 8 ? s[i1] & 0xF : s[i1 - 8] >> 4) | (((s[8 + (i1 & 3)] >> (2 * (i1 >> 2))) & 3) << 4); + const int l2 = (i2 < 8 ? s[i2] & 0xF : s[i2 - 8] >> 4) | (((s[8 + (i2 & 3)] >> (2 * (i2 >> 2))) & 3) << 4); + const float d = rdh(b + 108); + const int hs = 4 * n + j; + for (int w = 0; w < 8; ++w) { + bytes2(((qq[w] >> (2 * j)) & 0x0303) | (((hm[w] >> hs) & 0x0101) << 2), q + 2 * w); + bytes2(((qq[w + 8] >> (2 * j)) & 0x0303) | (((hm[w + 8] >> hs) & 0x0101) << 2), q + 16 + 2 * w); + } + const float d1 = d * float(l1 - 32), d2 = d * float(l2 - 32); + sc = float4(d1, d2, 4.0f * d1, 4.0f * d2); + } +}; + +template <> struct Dec { // 16 floats a unit (the Unsloth files' router / alpha / beta / inject rows) + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const float4* p = (device const float4*)(row + 64 * u); + for (int k = 0; k < 4; ++k) { + const float4 v = p[k]; + q[4 * k] = v.x; q[4 * k + 1] = v.y; q[4 * k + 2] = v.z; q[4 * k + 3] = v.w; + } + sc = float4(1.0f, 1.0f, 0.0f, 0.0f); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uint4* p = (device const uint4*)(row + 32 * u); + const uint4 a = p[0], c = p[1]; + const uint v[8] = {a.x, a.y, a.z, a.w, c.x, c.y, c.z, c.w}; + for (int k = 0; k < 8; ++k) { + q[2 * k] = as_type(v[k] << 16); + q[2 * k + 1] = as_type(v[k] & 0xFFFF0000u); + } + sc = float4(1.0f, 1.0f, 0.0f, 0.0f); + } +}; + +// acc += unit dot of decoded q / sc with one token's prepared slice +template +static inline float unit_combine(thread const float* q, float4 sc, thread const float* xs, thread const float* xsum) { + constexpr int H = UNIT / 2; + float s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int j = 0; j < H; j += 2) { + s0 = fma(q[j], xs[j], s0); + s1 = fma(q[j + 1], xs[j + 1], s1); + s2 = fma(q[H + j], xs[H + j], s2); + s3 = fma(q[H + j + 1], xs[H + j + 1], s3); + } + return sc.x * (s0 + s1) + sc.y * (s2 + s3) - (sc.z * xsum[0] + sc.w * xsum[1]); +} diff --git a/results/round39/astra_r39_qsa/baseline-moe_mm_parity.cpp b/results/round39/astra_r39_qsa/baseline-moe_mm_parity.cpp new file mode 100644 index 0000000..8e06a77 --- /dev/null +++ b/results/round39/astra_r39_qsa/baseline-moe_mm_parity.cpp @@ -0,0 +1,506 @@ +// metal/tests/moe_mm_parity.cpp - the prompt-chunk routed experts (moe_mm.metal): plan, gate/up GEMM, down GEMM, +// combine, against a CPU reference that rounds the activations and weights to fp16 as the kernels do (the matrix +// units' operands) - and against the exact float reference for the size of that rounding; then the speed. +// +// moe_mm_parity [--gu Q2_0] [--dn Q2_0] [--T 256] [--tile 32,64,64,4] parity (64 experts) +// moe_mm_parity --gather-bench [T=2048] [experts=512] [width=2560] short plan/copy sweep +// MM_BASELINE_KERNEL=/path/to/moe_mm.metal compares every plan buffer byte-for-byte +// moe_mm_parity --bench [--gu ..] [--dn ..] [--T 2048] [--tile ..] per-layer time, 512 experts +#include "coop_probe.h" +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { +constexpr int N_EMBD = 2560, N_FF = 640, TOPK = 10; +struct MmPlanArgs { uint32_t n_ent, topk, n_expert, tm, n_blocks; }; +struct MmArgs { uint64_t gate_off, up_off, down_off, gu_stride, d_stride; uint32_t gu_row, d_row, n_embd, n_ff; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +int type_of(const std::string& s) { + for (int t : {ref::Q2_0, ref::IQ3_XXS, ref::IQ2_XXS, ref::IQ2_XS, ref::IQ2_S, ref::IQ3_S, ref::IQ4_NL, ref::IQ4_XS, ref::Q4_K, ref::Q5_K, ref::Q6_K, ref::Q8_0, ref::Q5_1}) + if (s == ref::type_name(t)) return t; + std::exit(2); +} +float h16(float v) { return ref::f16_to_f32(ref::f32_to_f16(v)); } + +int gather_bench(int T, int NE, int n) { + if (T <= 0 || NE < TOPK || NE > 1024 || NE % 32 || n <= 0 || n % 8) { + std::fprintf(stderr, "gather-bench needs T > 0, experts in [32, 1024] divisible by 32, and width divisible by 8\n"); + return 2; + } + mtl::Context cx; + const int lib = cx.compile({"moe_mm.metal"}, "#define NO_GGML_COMMON 1\n#define MM_GATHER_ONLY 1\n"); + mtl::Pipeline phist = cx.pipeline(lib, "mm_hist"), pscan = cx.pipeline(lib, "mm_scan"), pplace = cx.pipeline(lib, "mm_place"), + pgat = cx.pipeline(lib, "mm_gather_h"), pgat_f = cx.pipeline(lib, "mm_gather"); + const char* baseline = std::getenv("MM_BASELINE_KERNEL"); + const int bl = baseline ? cx.compile({baseline}, "#define NO_GGML_COMMON 1\n#define MM_GATHER_ONLY 1\n") : lib; + mtl::Pipeline bh = cx.pipeline(bl, "mm_hist"), bs = cx.pipeline(bl, "mm_scan"), bp = cx.pipeline(bl, "mm_place"), + bg = cx.pipeline(bl, "mm_gather_h"), bf = cx.pipeline(bl, "mm_gather"); + std::printf("device %s; mm_plan + mm_gather_h bitwise parity & benchmark\n", cx.device_name().c_str()); + const int TM = 16; + const int nent = T * TOPK, nblk = (nent + 1023) / 1024, maxrows = nent + NE * TM; + mtl::Buffer x = cx.alloc(4 * (size_t)T * n), xh = cx.alloc(2 * (size_t)T * n), + ids = cx.alloc(4 * nent), w = cx.alloc(4 * nent), + bcnt = cx.alloc(4 * (size_t)nblk * NE), boff = cx.alloc(4 * (size_t)nblk * NE), eoff = cx.alloc(4 * NE), + tiles = cx.alloc(8 * (size_t)(maxrows / TM + 1)), meta = cx.alloc(64), stok = cx.alloc(4 * (size_t)maxrows), + sw = cx.alloc(4 * (size_t)maxrows), pos = cx.alloc(4 * nent), + XS = cx.alloc(2 * (size_t)maxrows * n), XS_ref = cx.alloc(2 * (size_t)maxrows * n); + ref::Rng r(7); + for (size_t i = 0; i < (size_t)T * n; ++i) { + const float v = r.sym(); + x.as()[i] = v; + xh.as()[i] = ref::f32_to_f16(v); + } + for (int t = 0; t < T; ++t) { + std::set used; + float s = 0; + for (int k = 0; k < TOPK; ++k) { + int e; + do e = (int)(r.next() % NE); while (used.count(e)); + used.insert(e); + ids.as()[t * TOPK + k] = e; + s += (w.as()[t * TOPK + k] = 0.2f + r.uni()); + } + for (int k = 0; k < TOPK; ++k) w.as()[t * TOPK + k] /= s; + } + MmPlanArgs pa{(uint32_t)nent, TOPK, (uint32_t)NE, (uint32_t)TM, (uint32_t)nblk}; + std::vector ds_plan; + mtl::Dispatch h; h.p = &phist; h.args = {mtl::Arg::B(ids), mtl::Arg::B(bcnt), mtl::Arg::V(pa)}; h.grid[0] = nblk; h.tg[0] = 1024; ds_plan.push_back(h); + mtl::Dispatch sc; sc.p = &pscan; sc.args = {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta), mtl::Arg::V(pa), mtl::Arg::B(stok)}; sc.tg[0] = NE; ds_plan.push_back(sc); + mtl::Dispatch pl; pl.p = &pplace; pl.args = {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok), mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}; pl.grid[0] = nblk; pl.tg[0] = 1024; ds_plan.push_back(pl); + bool ok_bits = true; + if (baseline) { + auto bd = ds_plan; + bd[0].p = &bh; bd[1].p = &bs; bd[2].p = &bp; + cx.run(bd); + std::vector outputs{&bcnt, &boff, &eoff, &tiles, &meta, &stok, &sw, &pos}; + std::vector sizes{4u * nblk * NE, 4u * nblk * NE, 4u * NE, + 8u * (maxrows / TM + 1), 64, 4u * maxrows, 4u * maxrows, 4u * nent}; + std::vector> expected; + for (size_t i = 0; i < outputs.size(); ++i) { + const auto* bytes = outputs[i]->as(); + expected.emplace_back(bytes, bytes + sizes[i]); + std::memset(outputs[i]->ptr, 0, sizes[i]); + } + cx.run(ds_plan); + const char* names[]{"bcnt", "boff", "eoff", "tiles", "meta", "srt_tok", "srt_w", "pos"}; + for (size_t i = 0; i < outputs.size(); ++i) { + if (std::memcmp(outputs[i]->ptr, expected[i].data(), sizes[i]) != 0) { + std::fprintf(stderr, "%s memcmp failed\n", names[i]); + ok_bits = false; + } + } + cx.warm(bd, 0.02); + cx.warm(ds_plan, 0.02); + std::vector bt, ct; + for (int i = 0; i < 7; ++i) { + if (i & 1) { + ct.push_back(cx.run(ds_plan, 8, false, true)); + bt.push_back(cx.run(bd, 8, false, true)); + } else { + bt.push_back(cx.run(bd, 8, false, true)); + ct.push_back(cx.run(ds_plan, 8, false, true)); + } + } + std::printf(" plan baseline %.3f candidate %.3f us/L; all 8 buffers memcmp=%s\n", + median(bt) * 1e6 / 8, median(ct) * 1e6 / 8, ok_bits ? "PASS" : "FAIL"); + } else cx.run(ds_plan); + const uint32_t tot_rows = meta.as()[1], ne = n; + const size_t xs_bytes = 2 * (size_t)tot_rows * n; + bool ok_plan = true; + std::vector seen(tot_rows, 0); + for (int i = 0; i < nent; ++i) { + const int row = pos.as()[i]; + if (row < 0 || row >= (int)tot_rows || stok.as()[row] != i / TOPK) ok_plan = false; + else seen[row]++; + } + for (uint32_t r0 = 0; r0 < tot_rows; ++r0) + if (seen[r0] > 1 || (seen[r0] == 0 && stok.as()[r0] != -1)) ok_plan = false; + mtl::Dispatch d_ref; d_ref.p = &pgat_f; d_ref.args = {mtl::Arg::B(x), mtl::Arg::B(stok), mtl::Arg::B(XS_ref), mtl::Arg::V(ne)}; d_ref.grid[0] = tot_rows; d_ref.tg[0] = 256; + cx.run({d_ref}); + auto b_ref = d_ref; b_ref.p = &bf; b_ref.args[2] = mtl::Arg::B(XS); + cx.run({b_ref}); + bool ok_xs = std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) == 0; + for (uint32_t tg_sz : {32u, 64u, 80u, 128u, 160u, 256u}) { + std::memset(XS.ptr, 0x55, xs_bytes); + mtl::Dispatch d; d.p = &pgat; d.args = {mtl::Arg::B(xh), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; d.grid[0] = tot_rows; d.tg[0] = tg_sz; + cx.run({d}); + if (std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) != 0) ok_xs = false; + d.grid[0] = 1; d.indirect = &meta; d.ind_off = 8 * 4; + cx.run({d}); + if (std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) != 0) ok_xs = false; + auto b_gat = d; b_gat.p = &bg; + cx.warm({b_gat, d}, 0.02); + std::vector bt, ct; + for (int i = 0; i < 7; ++i) { + if (i & 1) { + ct.push_back(cx.run({d}, 8, false, true)); + bt.push_back(cx.run({b_gat}, 8, false, true)); + } else { + bt.push_back(cx.run({b_gat}, 8, false, true)); + ct.push_back(cx.run({d}, 8, false, true)); + } + } + std::printf(" T=%d NE=%d rows=%u tg=%u baseline %.3f candidate %.3f us/L XS_memcmp=%s\n", + T, NE, tot_rows, tg_sz, median(bt) * 1e6 / 8, median(ct) * 1e6 / 8, ok_xs ? "PASS" : "FAIL"); + } + auto combined = ds_plan; + mtl::Dispatch gather = d_ref; + gather.p = &pgat; + gather.args[0] = mtl::Arg::B(xh); + gather.args[2] = mtl::Arg::B(XS); + gather.indirect = &meta; gather.ind_off = 8 * 4; + combined.push_back(gather); + std::vector times; + for (int i = 0; i < 3; ++i) times.push_back(cx.run(combined, 8, false, true)); + ok_xs &= std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) == 0; + std::printf(" combined plan+gather tg256 %.3f us/L plan=%s XS_memcmp=%s\n", + median(times) * 1e6 / 8, ok_plan ? "PASS" : "FAIL", ok_xs ? "PASS" : "FAIL"); + if (!ok_plan || !ok_xs || !ok_bits) return 1; + return 0; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc > 1 && std::string(argv[1]) == "--gather-bench") { + return gather_bench(argc > 2 ? std::atoi(argv[2]) : 2048, + argc > 3 ? std::atoi(argv[3]) : 512, + argc > 4 ? std::atoi(argv[4]) : N_EMBD); + } + int gu = ref::Q2_0, dn = ref::Q2_0, T = 256, TM = 16, TN = 64, TK = 16, SG = 2, DTN = -1, DTK = -1, DSG = -1; + bool bench = false, v2 = false, v3 = false, v4 = false, v5 = false; + int MS = 1; + bool gu_only = false, dn_only = false; + const char* dump = nullptr, *ab_dir = nullptr; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + if (a == "--gu") gu = type_of(argv[++i]); + else if (a == "--dn") dn = type_of(argv[++i]); + else if (a == "--T" || a == "-T") T = std::atoi(argv[++i]); + else if (a == "--tile") std::sscanf(argv[++i], "%d,%d,%d,%d", &TM, &TN, &TK, &SG); + else if (a == "--dtile") { + int m, n, k, s; + const char* value = argv[++i]; + if (std::sscanf(value, "%d,%d,%d,%d", &m, &n, &k, &s) == 4) { + if (m != TM) { + std::fprintf(stderr, "down tile TM must match gate/up TM\n"); + return 2; + } + DTN = n; DTK = k; DSG = s; + } else std::sscanf(value, "%d,%d,%d", &DTN, &DTK, &DSG); + } + else if (a == "--bench") bench = true; + else if (a == "--gu-bench") bench = gu_only = true; // --bench timing the gate/up GEMM alone + else if (a == "--dn-bench") bench = dn_only = true; + else if (a == "--v2") v2 = true; + else if (a == "--v3") v2 = v3 = true; // v2's gather + the register-operand GEMMs (TN per simdgroup) + else if (a == "--v4") { v2 = v3 = v4 = true; MS = std::atoi(argv[++i]); } // mm_gu4 with MS 16-row subtiles + else if (a == "--v5") { v2 = v3 = true; v5 = true; } // mm_gu5 (weight blocks staged in threadgroup memory) + else if (a == "--dump") dump = argv[++i]; // HS and y to a file: bitwise comparison of two kernel versions + else if (a == "--ab") ab_dir = argv[++i]; + } + const bool isolated_bench = gu_only || dn_only; + if (bench && T == 256) T = 2048; + if (v4) TM = 16 * MS; // the plan's tile height; mm_gu4's name carries MS instead + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "moe_mm.metal"}, (v3 && !v4 && !v5) ? "#define MM_V3_ONLY 1\n" : ""); + const int NE = bench ? 512 : 64, LAYERS = bench && !isolated_bench ? 4 : 1; + const size_t gu_row = ref::row_bytes(gu, N_EMBD), d_row = ref::row_bytes(dn, N_FF); + const size_t up_off = (size_t)N_FF * gu_row, down_off = 2 * up_off, blob = down_off + (size_t)N_EMBD * d_row; + std::printf("device %s; %s/%s, T %d, %d experts, tile TM %d TN %d TK %d SG %d\n", cx.device_name().c_str(), + ref::type_name(gu), ref::type_name(dn), T, NE, TM, TN, TK, SG); + mtl::Buffer W = cx.alloc(blob * NE * LAYERS); + { + std::vector tmp(blob * 16); + for (int e = 0; e < 16; ++e) { + ref::random_blocks(gu, tmp.data() + e * blob, 2 * N_FF, N_EMBD, 100 + e, 0.02f); + ref::random_blocks(dn, tmp.data() + e * blob + down_off, N_EMBD, N_FF, 200 + e, 0.02f); + } + for (size_t e = 0; e < (size_t)NE * LAYERS; ++e) std::memcpy(W.as() + e * blob, tmp.data() + (e % 16) * blob, blob); + } + const int nent = T * TOPK, nblk = (nent + 1023) / 1024; + const int maxrows = nent + NE * TM; + mtl::Buffer x = cx.alloc(4 * (size_t)T * N_EMBD), xh = cx.alloc(2 * (size_t)T * N_EMBD), ids = cx.alloc(4 * nent), w = cx.alloc(4 * nent), + bcnt = cx.alloc(4 * (size_t)nblk * NE), boff = cx.alloc(4 * (size_t)nblk * NE), eoff = cx.alloc(4 * NE), + tiles = cx.alloc(8 * (size_t)(maxrows / TM + 1)), meta = cx.alloc(64), stok = cx.alloc(4 * (size_t)maxrows), + sw = cx.alloc(4 * (size_t)maxrows), pos = cx.alloc(4 * nent), H = cx.alloc(4 * (size_t)maxrows * N_FF), + YE = cx.alloc(4 * (size_t)maxrows * N_EMBD), y = cx.alloc(4 * (size_t)T * N_EMBD), + XS = cx.alloc(2 * (size_t)maxrows * N_EMBD), HS = cx.alloc(2 * (size_t)maxrows * N_FF); + ref::Rng r(7); + for (size_t i = 0; i < (size_t)T * N_EMBD; ++i) { + const float v = r.sym(); + x.as()[i] = v; + xh.as()[i] = ref::f32_to_f16(v); + } + for (int t = 0; t < T; ++t) { + std::set used; + float s = 0; + for (int k = 0; k < TOPK; ++k) { + int e; + do e = (int)(r.next() % NE); while (used.count(e)); + used.insert(e); + ids.as()[t * TOPK + k] = e; + s += (w.as()[t * TOPK + k] = 0.2f + r.uni()); + } + for (int k = 0; k < TOPK; ++k) w.as()[t * TOPK + k] /= s; + } + mtl::Pipeline phist = cx.pipeline(lib, "mm_hist"), pscan = cx.pipeline(lib, "mm_scan"), pplace = cx.pipeline(lib, "mm_place"), + pcomb = cx.pipeline(lib, "mm_combine"); + if (DTN < 0) { DTN = TN; DTK = TK; DSG = SG; } + if (dn == ref::IQ4_NL && DTK < 32) DTK = 32; + const std::string tl = std::to_string(TM) + "_" + std::to_string(TN) + "_" + std::to_string(TK) + "_" + std::to_string(SG); + const std::string dtl = std::to_string(TM) + "_" + std::to_string(DTN) + "_" + std::to_string(DTK) + "_" + std::to_string(DSG); + mtl::Constants probe; + probe.i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0); + const bool k4 = coop_k4(cx); + if (v3) std::printf("right-input layout: %s\n", k4 ? "runs of 4 k" : "other (per-element decode)"); + mtl::Pipeline pgu = v5 ? cx.pipeline(lib, "mm_gu5_" + std::to_string(gu) + "_" + std::to_string(TN) + "_" + std::to_string(SG), + mtl::Constants().b(7, k4)) + : v4 ? cx.pipeline(lib, "mm_gu4_" + std::to_string(gu) + "_" + std::to_string(MS) + tl.substr(tl.find('_')), + mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0)) + : v3 ? cx.pipeline(lib, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)) + : cx.pipeline(lib, (v2 ? "mm_gu2_" : "mm_gu_") + std::to_string(gu) + "_" + tl, probe), + pdn = v3 ? cx.pipeline(lib, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4)) + : cx.pipeline(lib, (v2 ? "mm_down2_" : "mm_down_") + std::to_string(dn) + "_" + (v2 ? dtl : tl)), + pgat = cx.pipeline(lib, "mm_gather_h"); + MmPlanArgs pa{(uint32_t)nent, TOPK, (uint32_t)NE, (uint32_t)TM, (uint32_t)nblk}; + auto encode = [&](int layer, std::vector& ds, bool plan, bool gemm) { + if (plan) { + mtl::Dispatch h; h.p = &phist; h.args = {mtl::Arg::B(ids), mtl::Arg::B(bcnt), mtl::Arg::V(pa)}; h.grid[0] = nblk; h.tg[0] = 1024; ds.push_back(h); + meta.as()[2] = N_FF / TN; + mtl::Dispatch sc; sc.p = &pscan; sc.args = {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta), mtl::Arg::V(pa), mtl::Arg::B(stok)}; sc.tg[0] = NE; ds.push_back(sc); + mtl::Dispatch pl; pl.p = &pplace; pl.args = {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok), mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}; pl.grid[0] = nblk; pl.tg[0] = 1024; ds.push_back(pl); + } + if (gemm) { + MmArgs ma{(uint64_t)layer * NE * blob, (uint64_t)layer * NE * blob + up_off, (uint64_t)layer * NE * blob + down_off, + blob, blob, (uint32_t)gu_row, (uint32_t)d_row, N_EMBD, N_FF}; + if (v2) { + const uint32_t ne = N_EMBD; + mtl::Dispatch ga; ga.p = &pgat; ga.args = {mtl::Arg::B(xh), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; + ga.grid[0] = meta.as()[1]; ga.tg[0] = 256; ds.push_back(ga); + } + mtl::Dispatch g; g.p = &pgu; + if (v2) g.args = {mtl::Arg::B(W), mtl::Arg::B(XS), mtl::Arg::B(HS), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + else g.args = {mtl::Arg::B(W), mtl::Arg::B(x), mtl::Arg::B(H), mtl::Arg::B(stok), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + g.grid[0] = N_FF / (v3 ? TN * SG : TN); g.grid[1] = meta.as()[0]; g.tg[0] = 32 * SG; ds.push_back(g); + mtl::Dispatch d; d.p = &pdn; d.args = {mtl::Arg::B(W), mtl::Arg::B(v2 ? HS : H), mtl::Arg::B(YE), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + d.grid[0] = N_EMBD / (v3 ? DTN * DSG : v2 ? DTN : TN); d.grid[1] = meta.as()[0]; d.tg[0] = 32 * (v2 ? DSG : SG); ds.push_back(d); + const uint32_t ca[4] = {(uint32_t)T, N_EMBD, TOPK, 0}; + mtl::Dispatch cb; cb.p = &pcomb; cb.args = {mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(y), mtl::Arg::V(ca)}; + cb.grid[0] = (T * N_EMBD + 255) / 256; cb.tg[0] = 256; ds.push_back(cb); + } + }; + { + std::vector ds; + encode(0, ds, true, false); + cx.run(ds); // the plan first: the GEMMs' grid is the tile count + } + const uint32_t ntiles = meta.as()[0]; + std::printf("plan: %u tiles of %d, %u padded rows for %d entries\n", ntiles, TM, meta.as()[1], nent); + auto compare_ab = [&](const std::vector& ds, mtl::Pipeline& pgu_b, mtl::Pipeline& pdn_b, size_t pr) { + std::vector hs_a(HS.as(), HS.as() + 2 * pr * N_FF); + std::vector ye_a(YE.as(), YE.as() + 4 * pr * N_EMBD); + std::vector y_a(y.as(), y.as() + 4 * (size_t)T * N_EMBD); + std::memset(HS.ptr, 0, hs_a.size()); + std::memset(YE.ptr, 0, ye_a.size()); + std::memset(y.ptr, 0, y_a.size()); + auto ds_b = ds; + ds_b[v2 ? 1 : 0].p = &pgu_b; + ds_b[v2 ? 2 : 1].p = &pdn_b; + cx.run(ds_b); + const bool ok_hs = std::memcmp(HS.ptr, hs_a.data(), hs_a.size()) == 0; + const bool ok_ye = std::memcmp(YE.ptr, ye_a.data(), ye_a.size()) == 0; + const bool ok_y = std::memcmp(y.ptr, y_a.data(), y_a.size()) == 0; + std::printf("ab bitwise (T=%d): HS_memcmp=%s YE_memcmp=%s Y_memcmp=%s\n", + T, ok_hs ? "PASS" : "FAIL", ok_ye ? "PASS" : "FAIL", ok_y ? "PASS" : "FAIL"); + return ok_hs && ok_ye && ok_y; + }; + if (!bench) { + // the plan: every entry placed once, rows of an expert contiguous and in entry order + bool ok_plan = true; + std::vector seen(maxrows, 0); + for (int i = 0; i < nent; ++i) { + const int row = pos.as()[i]; + if (row < 0 || row >= (int)meta.as()[1] || stok.as()[row] != i / TOPK) ok_plan = false; + else seen[row]++; + } + for (int i = 0; i < maxrows; ++i) if (seen[i] > 1) ok_plan = false; + std::vector ds; + encode(0, ds, false, true); + cx.run(ds); + const size_t pr = meta.as()[1]; + if (dump) { + FILE* f = std::fopen(dump, "wb"); + if (v2) std::fwrite(HS.as(), 2, pr * N_FF, f); + std::fwrite(YE.as(), 4, pr * N_EMBD, f); + std::fwrite(y.as(), 4, (size_t)T * N_EMBD, f); + std::fclose(f); + return 0; + } + if (ab_dir) { + const std::string kd(ab_dir); + const int lib_b = cx.compile({kd + "/quant.metal", kd + "/units.metal", kd + "/mma.metal", kd + "/moe_mm.metal"}, (v3 && !v4 && !v5) ? "#define MM_V3_ONLY 1\n" : ""); + mtl::Pipeline pgu_b = cx.pipeline(lib_b, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)); + mtl::Pipeline pdn_b = cx.pipeline(lib_b, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4)); + return (ok_plan && compare_ab(ds, pgu_b, pdn_b, pr)) ? 0 : 1; + } + // reference: fp16 operands (x, weights, h), fp32+ accumulation; and the exact one + double se16 = 0, se = 0, sr = 0; + std::vector wg((size_t)N_FF * N_EMBD), wu(wg.size()), wd((size_t)N_EMBD * N_FF); + std::vector yr16((size_t)T * N_EMBD, 0.0), yr((size_t)T * N_EMBD, 0.0); + for (int e = 0; e < NE; ++e) { + const uint8_t* b = W.as() + e * blob; + bool any = false; + for (int i = 0; i < nent && !any; ++i) any = ids.as()[i] == e; + if (!any) continue; + ref::dequant(gu, b, wg.data(), wg.size()); + ref::dequant(gu, b + up_off, wu.data(), wu.size()); + ref::dequant(dn, b + down_off, wd.data(), wd.size()); + for (int i = 0; i < nent; ++i) { + if (ids.as()[i] != e) continue; + const int t = i / TOPK; + const float* xt = x.as() + (size_t)t * N_EMBD; + std::vector h16v(N_FF), hv(N_FF); + for (int f = 0; f < N_FF; ++f) { + double g = 0, u = 0, g16 = 0, u16 = 0; + for (int k = 0; k < N_EMBD; ++k) { + const float xk = xt[k], xh = h16(xk); + g += (double)wg[(size_t)f * N_EMBD + k] * xk; + u += (double)wu[(size_t)f * N_EMBD + k] * xk; + g16 += (double)h16(wg[(size_t)f * N_EMBD + k]) * xh; + u16 += (double)h16(wu[(size_t)f * N_EMBD + k]) * xh; + } + hv[f] = g / (1 + std::exp(-g)) * u; + h16v[f] = h16((float)(g16 / (1 + std::exp(-g16)) * u16)); + } + const double wt = w.as()[i]; + for (int d = 0; d < N_EMBD; ++d) { + double s = 0, s16 = 0; + for (int f = 0; f < N_FF; ++f) { + s += (double)wd[(size_t)d * N_FF + f] * hv[f]; + s16 += (double)h16(wd[(size_t)d * N_FF + f]) * h16v[f]; + } + yr[(size_t)t * N_EMBD + d] += wt * s; + yr16[(size_t)t * N_EMBD + d] += wt * s16; + } + } + } + for (size_t i = 0; i < yr.size(); ++i) { + const double g = y.as()[i]; + se16 += (g - yr16[i]) * (g - yr16[i]); + se += (g - yr[i]) * (g - yr[i]); + sr += yr[i] * yr[i]; + } + const double rel16 = std::sqrt(se16 / sr), rel = std::sqrt(se / sr); + const bool ok = ok_plan && rel16 < 2e-4; + std::printf("parity: plan %s; y vs fp16-operand reference rel L2 %.2e, vs exact float reference %.2e %s\n", + ok_plan ? "ok" : "WRONG", rel16, rel, ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; + } + // bench: the plan, the two GEMMs and the combine, LAYERS distinct weight sets rotated + double flops = 0; + { + const double rows = meta.as()[1]; + flops = 2.0 * nent * (2.0 * N_FF * N_EMBD + (double)N_EMBD * N_FF); + std::printf("useful GFLOP per layer %.1f (padded rows %.0f = %.2fx entries)\n", flops / 1e9, rows, rows / nent); + } + auto timeit = [&](bool plan, bool gemm, int part) { + std::vector ds; + const int repeats = isolated_bench ? 4 : 8; + for (int l = 0; l < repeats; ++l) { + std::vector one; + encode(l % LAYERS, one, plan, gemm); + if (part >= 0 && gemm) { ds.push_back(one[part]); continue; } + for (auto& d : one) ds.push_back(d); + } + cx.warm(ds, isolated_bench ? 0.1 : 0.3); + std::vector t; + for (int i = 0; i < (isolated_bench ? 3 : 5); ++i) t.push_back(cx.run(ds) / repeats); + return median(t); + }; + const int o = v2 ? 1 : 0; + if (isolated_bench) { + std::vector ds; + encode(0, ds, false, true); + cx.run(ds); + } + if (ab_dir && isolated_bench) { + const std::string kd(ab_dir); + const int lib_b = cx.compile({kd + "/quant.metal", kd + "/units.metal", kd + "/mma.metal", kd + "/moe_mm.metal"}, (v3 && !v4 && !v5) ? "#define MM_V3_ONLY 1\n" : ""); + mtl::Pipeline pgu_b = cx.pipeline(lib_b, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)); + mtl::Pipeline pdn_b = cx.pipeline(lib_b, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4)); + std::vector parity; + encode(0, parity, false, true); + const size_t pr = meta.as()[1]; + std::memset(HS.ptr, 0, 2 * pr * N_FF); + std::memset(YE.ptr, 0, 4 * pr * N_EMBD); + cx.run(parity); + if (!compare_ab(parity, pgu_b, pdn_b, pr)) return 1; + const int part = dn_only ? (o + 1) : o; + std::vector one; + encode(0, one, false, true); + std::vector ds_a{one[part]}; + mtl::Dispatch db = one[part]; + db.p = dn_only ? &pdn_b : &pgu_b; + std::vector ds_b{db}; + cx.warm(ds_a, 0.05); + cx.warm(ds_b, 0.05); + std::vector ta, tb; + double best_pair = 1e30; + for (int i = 0; i < 19; ++i) { + double va, vb; + if (i & 1) { + vb = cx.run(ds_b, 2) / 2; + va = cx.run(ds_a, 2) / 2; + } else { + va = cx.run(ds_a, 2) / 2; + vb = cx.run(ds_b, 2) / 2; + } + best_pair = std::min(best_pair, std::min(va, vb)); + ta.push_back(va); + tb.push_back(vb); + } + std::vector fa, fb; + for (size_t i = 0; i < ta.size(); ++i) { + if (std::max(ta[i], tb[i]) <= 1.10 * best_pair) { + fa.push_back(ta[i]); + fb.push_back(tb[i]); + } + } + if (fa.size() >= 3) { ta = fa; tb = fb; } + const double ma = median(ta) * 1e6, mb = median(tb) * 1e6; + const double mina = *std::min_element(ta.begin(), ta.end()) * 1e6; + const double minb = *std::min_element(tb.begin(), tb.end()) * 1e6; + std::printf("%s A/B (T=%d): A med=%.1f (min=%.1f) | B med=%.1f (min=%.1f) | diff med=%+.2f%% min=%+.2f%%\n", + dn_only ? "down" : "gate/up", T, ma, mina, mb, minb, (mb - ma) / ma * 100.0, (minb - mina) / mina * 100.0); + return 0; + } + if (dn_only) { + const double td = timeit(false, true, o + 1); + std::printf("down %.1f us (%.2f TFLOP/s useful)\n", td * 1e6, 2.0 * nent * N_FF * N_EMBD / td / 1e12); + return 0; + } + if (gu_only) { + const double tg = timeit(false, true, o); + std::printf("gate/up %.1f us (%.2f TFLOP/s useful)\n", tg * 1e6, 2.0 * nent * 2.0 * N_FF * N_EMBD / tg / 1e12); + return 0; + } + const double tp = timeit(true, false, -1), tg = timeit(false, true, o), td = timeit(false, true, o + 1), + tc = timeit(false, true, o + 2), ta = timeit(false, true, -1); + if (v2) std::printf("gather %.1f us\n", timeit(false, true, 0) * 1e6); + const double fgu = 2.0 * nent * 2.0 * N_FF * N_EMBD, fdn = 2.0 * nent * N_EMBD * N_FF; + std::printf("per layer: plan %.1f us, gate/up %.1f us (%.1f TFLOP/s useful), down %.1f us (%.1f TFLOP/s), combine %.1f us;" + " GEMMs + combine %.1f us = %.1f TFLOP/s; %.1f us per token per layer -> 48 layers: %.0f tok/s (experts only)\n", + tp * 1e6, tg * 1e6, fgu / tg / 1e12, td * 1e6, fdn / td / 1e12, tc * 1e6, ta * 1e6, flops / ta / 1e12, + (ta + tp) / T * 1e6, T / (48 * (ta + tp))); + return 0; +} diff --git a/results/round39/astra_r39_qsa/evaluator/metal/src/ref_quant.cpp b/results/round39/astra_r39_qsa/evaluator/metal/src/ref_quant.cpp new file mode 100644 index 0000000..ede5e5b --- /dev/null +++ b/results/round39/astra_r39_qsa/evaluator/metal/src/ref_quant.cpp @@ -0,0 +1,445 @@ +// metal/src/ref_quant.cpp - see ref_quant.h. The i-quant loops are transcribed from ggml-quants.c +// (dequantize_row_iq2_xxs / iq2_xs / iq2_s / iq3_xxs / iq3_s / iq4_nl / iq4_xs) at 3cf03257; the K-quants and the +// 32-blocks from Strata's include/strata/artifact/dequant.hpp (themselves transcribed from the same commit). +#include "ref_quant.h" + +#include +#include +#include +#include +#include + +#define GGML_COMMON_DECL_CPP +#define GGML_COMMON_IMPL_CPP +#include "ggml-common.h" + +namespace ref { + +const char* type_name(int t) { + switch (t) { + case F32: return "F32"; case F16: return "F16"; case BF16: return "BF16"; + case Q4_0: return "Q4_0"; case Q5_0: return "Q5_0"; case Q5_1: return "Q5_1"; case Q8_0: return "Q8_0"; + case Q3_K: return "Q3_K"; case Q4_K: return "Q4_K"; case Q5_K: return "Q5_K"; case Q6_K: return "Q6_K"; + case IQ2_XXS: return "IQ2_XXS"; case IQ2_XS: return "IQ2_XS"; case IQ2_S: return "IQ2_S"; + case IQ3_XXS: return "IQ3_XXS"; case IQ3_S: return "IQ3_S"; case IQ4_NL: return "IQ4_NL"; + case IQ4_XS: return "IQ4_XS"; case Q2_0: return "Q2_0"; + } + return "?"; +} + +int block_elems(int t) { + switch (t) { + case F32: case F16: case BF16: return 1; + case Q4_0: case Q5_0: case Q5_1: case Q8_0: case IQ4_NL: return 32; + case Q2_0: return 64; + default: return 256; + } +} + +int block_bytes(int t) { + switch (t) { + case F32: return 4; case F16: case BF16: return 2; + case Q4_0: return 18; case Q5_0: return 22; case Q5_1: return 24; case Q8_0: return 34; case IQ4_NL: return 18; case Q2_0: return 18; + case Q3_K: return 110; case Q4_K: return 144; case Q5_K: return 176; case Q6_K: return 210; + case IQ2_XXS: return 66; case IQ2_XS: return 74; case IQ2_S: return 82; case IQ3_XXS: return 98; + case IQ3_S: return 110; case IQ4_XS: return 136; + } + std::fprintf(stderr, "ref: unknown type %d\n", t); + std::exit(2); +} + +float f16_to_f32(uint16_t h) { + const uint32_t sign = (uint32_t)(h >> 15) & 1u; + uint32_t exp = (h >> 10) & 0x1Fu, man = h & 0x3FFu, f; + if (exp == 0) { + if (man == 0) + f = sign << 31; + else { + exp = 127 - 15 + 1; + while (!(man & 0x400u)) { man <<= 1; --exp; } + man &= 0x3FFu; + f = (sign << 31) | (exp << 23) | (man << 13); + } + } else if (exp == 31) + f = (sign << 31) | 0x7F800000u | (man << 13); + else + f = (sign << 31) | ((exp - 15 + 127) << 23) | (man << 13); + float out; + std::memcpy(&out, &f, 4); + return out; +} + +uint16_t f32_to_f16(float f) { + _Float16 h = (_Float16)f; + uint16_t r; + std::memcpy(&r, &h, 2); + return r; +} + +float bf16_to_f32(uint16_t h) { + const uint32_t f = (uint32_t)h << 16; + float out; + std::memcpy(&out, &f, 4); + return out; +} + +uint16_t f32_to_bf16(float f) { + uint32_t u; + std::memcpy(&u, &f, 4); + u += 0x7FFFu + ((u >> 16) & 1u); + return (uint16_t)(u >> 16); +} + +static inline uint16_t rd16(const uint8_t* p) { return (uint16_t)(p[0] | (p[1] << 8)); } + +static void get_scale_min_k4(int j, const uint8_t* q, uint8_t& d, uint8_t& m) { + if (j < 4) { + d = q[j] & 63; + m = q[j + 4] & 63; + } else { + d = (q[j + 4] & 0x0F) | ((q[j - 4] >> 6) << 4); + m = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4); + } +} + +static void deq_block(int t, const uint8_t* b, float* y) { + switch (t) { + case Q2_0: { + const float d = f16_to_f32(rd16(b)); + for (int j = 0; j < 64; ++j) y[j] = (float)(((b[2 + j / 4] >> ((j % 4) * 2)) & 3) - 1) * d; + return; + } + case Q8_0: { + const float d = f16_to_f32(rd16(b)); + for (int j = 0; j < 32; ++j) y[j] = (float)(int8_t)b[2 + j] * d; + return; + } + case Q4_0: { + const float d = f16_to_f32(rd16(b)); + for (int j = 0; j < 16; ++j) { + y[j] = (float)((b[2 + j] & 15) - 8) * d; + y[j + 16] = (float)((b[2 + j] >> 4) - 8) * d; + } + return; + } + case Q5_0: { + const float d = f16_to_f32(rd16(b)); + uint32_t h; + std::memcpy(&h, b + 2, 4); + const uint8_t* qs = b + 6; + for (int j = 0; j < 16; ++j) { + const int x0 = (qs[j] & 15) | (((h >> j) & 1) << 4), x1 = (qs[j] >> 4) | (((h >> (j + 16)) & 1) << 4); + y[j] = (float)(x0 - 16) * d; + y[j + 16] = (float)(x1 - 16) * d; + } + return; + } + case Q5_1: { // ggml dequantize_row_q5_1: y = x * d + m + const float d = f16_to_f32(rd16(b)), m = f16_to_f32(rd16(b + 2)); + uint32_t h; + std::memcpy(&h, b + 4, 4); + const uint8_t* qs = b + 8; + for (int j = 0; j < 16; ++j) { + const int x0 = (qs[j] & 15) | (((h >> j) & 1) << 4), x1 = (qs[j] >> 4) | (((h >> (j + 16)) & 1) << 4); + y[j] = (float)x0 * d + m; + y[j + 16] = (float)x1 * d + m; + } + return; + } + case IQ4_NL: { + const float d = f16_to_f32(rd16(b)); + for (int j = 0; j < 16; ++j) { + y[j] = d * kvalues_iq4nl[b[2 + j] & 15]; + y[j + 16] = d * kvalues_iq4nl[b[2 + j] >> 4]; + } + return; + } + case IQ4_XS: { + const float d = f16_to_f32(rd16(b)); + const uint16_t sh = rd16(b + 2); + const uint8_t* sl = b + 4; + const uint8_t* qs = b + 8; + for (int ib = 0; ib < 8; ++ib) { + const int ls = ((sl[ib / 2] >> (4 * (ib % 2))) & 15) | (((sh >> (2 * ib)) & 3) << 4); + const float dl = d * (ls - 32); + for (int j = 0; j < 16; ++j) { + y[j] = dl * kvalues_iq4nl[qs[j] & 15]; + y[j + 16] = dl * kvalues_iq4nl[qs[j] >> 4]; + } + y += 32; + qs += 16; + } + return; + } + case Q6_K: { + const uint8_t* ql = b; + const uint8_t* qh = b + 128; + const int8_t* sc = (const int8_t*)(b + 192); + const float d = f16_to_f32(rd16(b + 208)); + for (int n = 0; n < 256; n += 128) { + for (int l = 0; l < 32; ++l) { + const int is = l / 16; + const int q1 = (int)((ql[l] & 15) | (((qh[l] >> 0) & 3) << 4)) - 32; + const int q2 = (int)((ql[l + 32] & 15) | (((qh[l] >> 2) & 3) << 4)) - 32; + const int q3 = (int)((ql[l] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + const int q4 = (int)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + y[l] = d * sc[is] * q1; + y[l + 32] = d * sc[is + 2] * q2; + y[l + 64] = d * sc[is + 4] * q3; + y[l + 96] = d * sc[is + 6] * q4; + } + y += 128; + ql += 64; + qh += 32; + sc += 8; + } + return; + } + case Q4_K: { + const float d = f16_to_f32(rd16(b)), mn = f16_to_f32(rd16(b + 2)); + const uint8_t* sc = b + 4; + const uint8_t* q = b + 16; + int is = 0; + for (int j = 0; j < 256; j += 64) { + uint8_t s, m; + get_scale_min_k4(is, sc, s, m); + const float d1 = d * s, m1 = mn * m; + get_scale_min_k4(is + 1, sc, s, m); + const float d2 = d * s, m2 = mn * m; + for (int l = 0; l < 32; ++l) *y++ = d1 * (q[l] & 15) - m1; + for (int l = 0; l < 32; ++l) *y++ = d2 * (q[l] >> 4) - m2; + q += 32; + is += 2; + } + return; + } + case Q5_K: { + const float d = f16_to_f32(rd16(b)), mn = f16_to_f32(rd16(b + 2)); + const uint8_t* sc = b + 4; + const uint8_t* qh = b + 16; + const uint8_t* ql = b + 48; + int is = 0; + uint8_t u1 = 1, u2 = 2; + for (int j = 0; j < 256; j += 64) { + uint8_t s, m; + get_scale_min_k4(is, sc, s, m); + const float d1 = d * s, m1 = mn * m; + get_scale_min_k4(is + 1, sc, s, m); + const float d2 = d * s, m2 = mn * m; + for (int l = 0; l < 32; ++l) *y++ = d1 * ((ql[l] & 15) + ((qh[l] & u1) ? 16 : 0)) - m1; + for (int l = 0; l < 32; ++l) *y++ = d2 * ((ql[l] >> 4) + ((qh[l] & u2) ? 16 : 0)) - m2; + ql += 32; + is += 2; + u1 <<= 2; + u2 <<= 2; + } + return; + } + case Q3_K: { + const uint8_t* hm = b; + const uint8_t* q = b + 32; + const float d_all = f16_to_f32(rd16(b + 108)); + const uint32_t kmask1 = 0x03030303u, kmask2 = 0x0f0f0f0fu; + uint32_t aux[4] = {0, 0, 0, 0}; + std::memcpy(aux, b + 96, 12); + const uint32_t tmp = aux[2]; + aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4); + aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4); + aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4); + aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4); + const int8_t* scales = (const int8_t*)aux; + int is = 0; + uint8_t m = 1; + for (int n = 0; n < 256; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + float dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) *y++ = dl * ((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)); + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) + *y++ = dl * ((int8_t)((q[l + 16] >> shift) & 3) - ((hm[l + 16] & m) ? 0 : 4)); + shift += 2; + m <<= 1; + } + q += 32; + } + return; + } + case IQ2_XXS: { + const float d = f16_to_f32(rd16(b)); + const uint8_t* qs = b + 2; + for (int ib32 = 0; ib32 < 8; ++ib32) { + uint32_t aux32[2]; + std::memcpy(aux32, qs + 8 * ib32, 8); + const uint8_t* aux8 = (const uint8_t*)aux32; + const float db = d * (0.5f + (aux32[1] >> 28)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const uint8_t* grid = (const uint8_t*)(iq2xxs_grid + aux8[l]); + const uint8_t signs = ksigns_iq2xs[(aux32[1] >> 7 * l) & 127]; + for (int j = 0; j < 8; ++j) y[j] = db * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); + y += 8; + } + } + return; + } + case IQ2_XS: { + const float d = f16_to_f32(rd16(b)); + const uint8_t* q8 = b + 2; + const uint8_t* scales = b + 66; + for (int ib32 = 0; ib32 < 8; ++ib32) { + const float db0 = d * (0.5f + (scales[ib32] & 15)) * 0.25f; + const float db1 = d * (0.5f + (scales[ib32] >> 4)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const uint16_t q = rd16(q8 + 2 * (4 * ib32 + l)); + const uint8_t* grid = (const uint8_t*)(iq2xs_grid + (q & 511)); + const uint8_t signs = ksigns_iq2xs[q >> 9]; + const float db = l < 2 ? db0 : db1; + for (int j = 0; j < 8; ++j) y[j] = db * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); + y += 8; + } + } + return; + } + case IQ2_S: { + const float d = f16_to_f32(rd16(b)); + const uint8_t* qs = b + 2; + const uint8_t* qh = b + 66; + const uint8_t* scales = b + 74; + const uint8_t* signs = qs + 32; + for (int ib32 = 0; ib32 < 8; ++ib32) { + const float db0 = d * (0.5f + (scales[ib32] & 15)) * 0.25f; + const float db1 = d * (0.5f + (scales[ib32] >> 4)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const float dl = l < 2 ? db0 : db1; + const uint8_t* grid = (const uint8_t*)(iq2s_grid + (qs[l] | ((qh[ib32] << (8 - 2 * l)) & 0x300))); + for (int j = 0; j < 8; ++j) y[j] = dl * grid[j] * (signs[l] & kmask_iq2xs[j] ? -1.f : 1.f); + y += 8; + } + qs += 4; + signs += 4; + } + return; + } + case IQ3_XXS: { + const float d = f16_to_f32(rd16(b)); + const uint8_t* qs = b + 2; + const uint8_t* sas = b + 2 + 64; + for (int ib32 = 0; ib32 < 8; ++ib32) { + uint32_t aux32; + std::memcpy(&aux32, sas + 4 * ib32, 4); + const float db = d * (0.5f + (aux32 >> 28)) * 0.5f; + for (int l = 0; l < 4; ++l) { + const uint8_t signs = ksigns_iq2xs[(aux32 >> 7 * l) & 127]; + const uint8_t* g1 = (const uint8_t*)(iq3xxs_grid + qs[2 * l + 0]); + const uint8_t* g2 = (const uint8_t*)(iq3xxs_grid + qs[2 * l + 1]); + for (int j = 0; j < 4; ++j) { + y[j + 0] = db * g1[j] * (signs & kmask_iq2xs[j + 0] ? -1.f : 1.f); + y[j + 4] = db * g2[j] * (signs & kmask_iq2xs[j + 4] ? -1.f : 1.f); + } + y += 8; + } + qs += 8; + } + return; + } + case IQ3_S: { + const float d = f16_to_f32(rd16(b)); + const uint8_t* qs = b + 2; + const uint8_t* qh = b + 66; + const uint8_t* signs = b + 74; + const uint8_t* scales = b + 106; + for (int ib32 = 0; ib32 < 8; ib32 += 2) { + const float db1 = d * (1 + 2 * (scales[ib32 / 2] & 15)); + const float db2 = d * (1 + 2 * (scales[ib32 / 2] >> 4)); + for (int l = 0; l < 4; ++l) { + const uint8_t* g1 = (const uint8_t*)(iq3s_grid + (qs[2 * l + 0] | ((qh[0] << (8 - 2 * l)) & 256))); + const uint8_t* g2 = (const uint8_t*)(iq3s_grid + (qs[2 * l + 1] | ((qh[0] << (7 - 2 * l)) & 256))); + for (int j = 0; j < 4; ++j) { + y[j + 0] = db1 * g1[j] * (signs[l] & kmask_iq2xs[j + 0] ? -1.f : 1.f); + y[j + 4] = db1 * g2[j] * (signs[l] & kmask_iq2xs[j + 4] ? -1.f : 1.f); + } + y += 8; + } + qs += 8; + signs += 4; + for (int l = 0; l < 4; ++l) { + const uint8_t* g1 = (const uint8_t*)(iq3s_grid + (qs[2 * l + 0] | ((qh[1] << (8 - 2 * l)) & 256))); + const uint8_t* g2 = (const uint8_t*)(iq3s_grid + (qs[2 * l + 1] | ((qh[1] << (7 - 2 * l)) & 256))); + for (int j = 0; j < 4; ++j) { + y[j + 0] = db2 * g1[j] * (signs[l] & kmask_iq2xs[j + 0] ? -1.f : 1.f); + y[j + 4] = db2 * g2[j] * (signs[l] & kmask_iq2xs[j + 4] ? -1.f : 1.f); + } + y += 8; + } + qh += 2; + qs += 8; + signs += 4; + } + return; + } + } + std::fprintf(stderr, "ref: dequant of %s not implemented\n", type_name(t)); + std::exit(2); +} + +void dequant(int t, const void* src, float* dst, size_t n) { + const uint8_t* s = (const uint8_t*)src; + if (t == F32) { std::memcpy(dst, s, n * 4); return; } + if (t == F16) { for (size_t i = 0; i < n; ++i) dst[i] = f16_to_f32(rd16(s + 2 * i)); return; } + if (t == BF16) { for (size_t i = 0; i < n; ++i) dst[i] = bf16_to_f32(rd16(s + 2 * i)); return; } + const int be = block_elems(t), bb = block_bytes(t); + for (size_t i = 0; i < n / be; ++i) deq_block(t, s + i * bb, dst + i * be); +} + +void quantize_q8_0(const float* x, size_t n, uint8_t* out) { + for (size_t b = 0; b < n / 32; ++b) { + float amax = 0.0f; + for (int j = 0; j < 32; ++j) amax = std::max(amax, std::fabs(x[b * 32 + j])); + const float d = amax / 127.0f, id = d != 0.0f ? 1.0f / d : 0.0f; + const uint16_t h = f32_to_f16(d); + uint8_t* o = out + b * 34; + o[0] = (uint8_t)h; + o[1] = (uint8_t)(h >> 8); + for (int j = 0; j < 32; ++j) o[2 + j] = (uint8_t)(int8_t)std::roundf(x[b * 32 + j] * id); + } +} + +static void put16(uint8_t* p, uint16_t v) { p[0] = (uint8_t)v; p[1] = (uint8_t)(v >> 8); } + +void random_blocks(int t, uint8_t* dst, int rows, int cols, uint64_t seed, float scale) { + Rng r(seed); + const size_t n = (size_t)rows * cols; + if (t == F32 || t == F16 || t == BF16) { + for (size_t i = 0; i < n; ++i) { + const float v = r.sym() * scale * 2.0f; + if (t == F32) std::memcpy(dst + 4 * i, &v, 4); + else put16(dst + 2 * i, t == F16 ? f32_to_f16(v) : f32_to_bf16(v)); + } + return; + } + const int be = block_elems(t), bb = block_bytes(t); + const size_t nb = n / be; + for (size_t i = 0; i < nb; ++i) { + uint8_t* b = dst + i * bb; + for (int k = 0; k < bb; ++k) b[k] = (uint8_t)r.next(); + // the scale fields: magnitude so that |w| ~ scale. Typical code magnitudes differ per type. + const float u = 0.5f + r.uni(); + switch (t) { + case Q2_0: put16(b, f32_to_f16(scale * u)); break; + case Q8_0: put16(b, f32_to_f16(scale * u / 64)); break; + case Q4_0: case IQ4_NL: put16(b, f32_to_f16(scale * u / (t == Q4_0 ? 4 : 64))); break; + case Q5_0: put16(b, f32_to_f16(scale * u / 8)); break; + case Q5_1: put16(b, f32_to_f16(scale * u / 16)); put16(b + 2, f32_to_f16(-scale * u)); break; + case IQ4_XS: put16(b, f32_to_f16(scale * u / 64 / 16)); break; + case Q4_K: case Q5_K: put16(b, f32_to_f16(scale * u / 32 / 8)); put16(b + 2, f32_to_f16(scale * u / 32 / 2)); break; + case Q6_K: put16(b + 208, f32_to_f16(scale * u / 32 / 64)); break; + case Q3_K: put16(b + 108, f32_to_f16(scale * u / 32 / 2)); break; + case IQ2_XXS: case IQ2_XS: case IQ2_S: put16(b, f32_to_f16(scale * u / 4)); break; + case IQ3_XXS: put16(b, f32_to_f16(scale * u / 8)); break; + case IQ3_S: put16(b, f32_to_f16(scale * u / 64)); break; + default: break; + } + } +} + +} // namespace ref diff --git a/results/round39/astra_r39_qsa/evaluator/metal/tests/qsa_parity.cpp b/results/round39/astra_r39_qsa/evaluator/metal/tests/qsa_parity.cpp new file mode 100644 index 0000000..2897b5c --- /dev/null +++ b/results/round39/astra_r39_qsa/evaluator/metal/tests/qsa_parity.cpp @@ -0,0 +1,489 @@ +// metal/tests/qsa_parity.cpp - QSA decode (qsa.metal) stage by stage against a CPU reference, and its speed. +// +// qsa_parity [P] a context of P cells built by qsa_prep (default 3000: the selection is active above +// 2051), then a window of 3 tokens: prep, score, top-k, attention, merge. Each stage is +// checked against the CPU fed with the GPU's inputs to that stage; the ids exactly. +// qsa_parity --bench [P] GPU time of the window's kernels at context P, 1 and 3 tokens +// qsa_parity --bitwise REF.metal +// qsa_attn + qsa_merge against the same kernels compiled from REF.metal (e.g. the file +// before an edit): part and xp memcmp-equal over pos0 0..30000 and 1, 2, 4, 8 tokens on +// random KV codes and ids, then paired GPU times of both +// QSA_ATTN=2 tests qsa_attn2 + qsa_merge2; REF is then qsa2.metal +// qsa_parity --paired REF_DIR [P] [--tokens N] [--prep-only] +// full-buffer memcmp and paired per-stage timings; prep-only avoids attention allocation +// qsa_parity --ab REF_DIR [P] +// paired A/B across decode T = 1, 2, 3, 4, 8 and prompt T = 1984, 2048 (with warmup + 12x reps) +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { +constexpr int QH = 24, QKV = 2, QD = 256, IH = 4, ID = 128, BUDGET = 2051, CH = 64, NCH = (BUDGET + CH - 1) / CH; +struct QsaArgs { uint32_t pos0, pos_base, max_cells, n_tok; float eps; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +struct Err { + double se = 0, sr = 0; + void add(double g, double r) { se += (g - r) * (g - r); sr += r * r; } + double l2() const { return sr > 0 ? std::sqrt(se / sr) : std::sqrt(se); } +}; +float h2f(uint16_t h) { return ref::f16_to_f32(h); } +// --bitwise: the attention and merge of this tree against REF's, bit for bit on the same inputs, and their times +int bitwise(const char* ref_path) { + constexpr int MAXT = 8, CELLS = 32768; + mtl::Context cx; + const bool v2 = std::getenv("QSA_ATTN") && std::atoi(std::getenv("QSA_ATTN")) == 2; + const int nch = v2 ? (BUDGET + 31) / 32 : NCH; + std::vector src_new{"qsa.metal"}, src_ref; + if (v2) { src_new.push_back("qsa2.metal"); src_ref.push_back("qsa.metal"); } + src_ref.push_back(ref_path); + const int lnew = cx.compile(src_new, "#define NO_GGML_COMMON 1\n"); + const int lref = cx.compile(src_ref, "#define NO_GGML_COMMON 1\n"); + const char* attn = v2 ? "qsa_attn2" : "qsa_attn"; + const char* merge = v2 ? "qsa_merge2" : "qsa_merge"; + mtl::Pipeline an = cx.pipeline(lnew, attn), mn = cx.pipeline(lnew, merge), + ar = cx.pipeline(lref, attn), mr = cx.pipeline(lref, merge); + const size_t part_bytes = 4 * (size_t)MAXT * nch * QH * (QD + 2), xp_bytes = 4 * (4 * MAXT * QH * QD + MAXT * QH * QD / 16); + mtl::Buffer qcur = cx.alloc(4 * MAXT * QH * QD), qfull = cx.alloc(4 * MAXT * QH * 2 * QD), + kq = cx.alloc((size_t)CELLS * QKV * QD), vq = cx.alloc((size_t)CELLS * QKV * QD), + ksc = cx.alloc(2 * (size_t)CELLS * QKV * 4), vsc = cx.alloc(2 * (size_t)CELLS * QKV * 4), + ids = cx.alloc(4 * MAXT * BUDGET), part[2] = {cx.alloc(part_bytes), cx.alloc(part_bytes)}, + xp[2] = {cx.alloc(xp_bytes), cx.alloc(xp_bytes)}; + ref::Rng r(41); + for (size_t i = 0; i < qcur.bytes / 4; ++i) qcur.as()[i] = r.sym() * 1.5f; + for (size_t i = 0; i < qfull.bytes / 4; ++i) qfull.as()[i] = r.sym() * 2.0f; + for (size_t i = 0; i < kq.bytes; ++i) { kq.as()[i] = (int8_t)(r.next() % 255 - 127); vq.as()[i] = (int8_t)(r.next() % 255 - 127); } + for (size_t i = 0; i < ksc.bytes / 2; ++i) { + ksc.as()[i] = ref::f32_to_f16(0.004f + 0.012f * r.uni()); + vsc.as()[i] = ref::f32_to_f16(0.004f + 0.012f * r.uni()); + } + auto fill_ids = [&](uint32_t pos0, int nt) { + for (int t = 0; t < nt; ++t) { + const int n_kv = (int)pos0 + t + 1, width = std::min(n_kv, BUDGET); + int* o = ids.as() + t * BUDGET; + if (n_kv <= BUDGET) { for (int c = 0; c < width; ++c) o[c] = c; continue; } + // width distinct cells of [0, n_kv) ascending, as the top-k emits them (folded into the KV buffer) + std::vector pick; + for (int c = 0, left = width; c < n_kv && left; ++c) + if (r.next() % (uint64_t)(n_kv - c) < (uint64_t)left) { pick.push_back(c % CELLS); --left; } + std::copy(pick.begin(), pick.end(), o); + } + }; + auto dispatches = [&](int v, uint32_t pos0, int nt) { + QsaArgs a{pos0, 0, (uint32_t)CELLS, (uint32_t)nt, 1e-6f}; + mtl::Dispatch at, m; + at.p = v ? &an : &ar; + at.args = {mtl::Arg::B(qcur), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), + mtl::Arg::B(ids), mtl::Arg::B(part[v]), mtl::Arg::V(a)}; + at.grid[0] = nch; at.grid[1] = v2 ? QKV * 2 : QKV; at.grid[2] = nt; at.tg[0] = 256; + m.p = v ? &mn : &mr; + m.args = {mtl::Arg::B(part[v]), mtl::Arg::B(qfull), mtl::Arg::B(xp[v]), mtl::Arg::V(a)}; + m.grid[0] = QH; m.grid[1] = nt; m.tg[0] = 256; + return std::vector{at, m}; + }; + int fails = 0; + for (uint32_t pos0 : {0u, 1u, 63u, 64u, 512u, 2047u, 2048u, 2050u, 8192u, 30000u}) + for (int nt : {1, 2, 4, 8}) { + fill_ids(pos0, nt); + for (int v = 0; v < 2; ++v) { + std::memset(part[v].ptr, 0, part_bytes); + std::memset(xp[v].ptr, 0, xp_bytes); + cx.run(dispatches(v, pos0, nt)); + } + const bool part_ok = !std::memcmp(part[0].ptr, part[1].ptr, part_bytes); + const bool xp_ok = !std::memcmp(xp[0].ptr, xp[1].ptr, xp_bytes); + const bool ok = part_ok && xp_ok; + fails += !ok; + if (!ok) std::printf("pos0 %5u ntok %d: part %s xp %s FAIL\n", pos0, nt, + part_ok ? "PASS" : "FAIL", xp_ok ? "PASS" : "FAIL"); + } + std::printf("bitwise: %d of 40 (pos0, ntok) cases differ %s\n", fails, fails ? "FAIL" : "PASS"); + std::printf("device %s: GPU us per dispatch, median of 9 paired rounds of 32 (ref / this)\n", cx.device_name().c_str()); + for (uint32_t pos0 : {512u, 2048u, 8192u}) + for (int nt : {1, 4}) { + fill_ids(pos0, nt); + std::vector d[2] = {dispatches(0, pos0, nt), dispatches(1, pos0, nt)}; + double us[2][2]; + for (int i = 0; i < 2; ++i) { + std::vector one[2] = {std::vector(32, d[0][i]), std::vector(32, d[1][i])}; + cx.warm(one[0], 0.05); + std::vector t[2]; + for (int k = 0; k < 9; ++k) + for (int turn = 0; turn < 2; ++turn) { + const int v = (k + turn) & 1; + t[v].push_back(cx.run(one[v]) / 32); + } + for (int v = 0; v < 2; ++v) us[i][v] = median(t[v]) * 1e6; + } + std::printf(" pos0 %5u ntok %d attn %6.2f / %6.2f merge %5.2f / %5.2f\n", pos0, nt, us[0][0], us[0][1], + us[1][0], us[1][1]); + } + return fails ? 1 : 0; +} +} // namespace + +int main(int argc, char** argv) { + bool bench = false; + bool ab_suite = false; + const char* paired_ref = nullptr; + bool prep_only = false; + int paired_nt = 3; + int P = 3000; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + if (a == "--prep-only") { prep_only = true; continue; } + if (a == "--ab" && i + 1 < argc) { paired_ref = argv[++i]; ab_suite = true; continue; } + if (a == "--paired" && i + 1 < argc) { paired_ref = argv[++i]; continue; } + if (a == "--tokens" && i + 1 < argc) { paired_nt = std::atoi(argv[++i]); continue; } + if (a == "--bitwise" && i + 1 < argc) return bitwise(argv[i + 1]); + if (a == "--bench") bench = true; + else P = std::atoi(argv[i]); + } + if ((prep_only && !paired_ref) || (paired_nt > 128 && !prep_only)) return 2; + if (paired_ref && (paired_nt < 1 || paired_nt > 2048 || P < 0)) return 2; + const int MAXT = ab_suite ? 2048 : std::max(8, paired_nt), MAXC = std::max(4096, (P + MAXT + 64) / 4 * 4 + 64); + mtl::Context cx; + const int lib = cx.compile({"qsa.metal", "qsa2.metal"}, "#define NO_GGML_COMMON 1\n"); + // QSA_TOPK=1: the v1 top-k; QSA_ATTN=1|2|3 (default 3): which attention (2 has its own chunking and merge) + const int av = std::getenv("QSA_ATTN") ? std::atoi(std::getenv("QSA_ATTN")) : 1; + const bool v2 = av == 2; + mtl::Pipeline pprep = cx.pipeline(lib, "qsa_prep", mtl::Constants().b(24, std::getenv("QSA_MROPE") != nullptr)), pscore = cx.pipeline(lib, "qsa_score"), + ptopk = cx.pipeline(lib, std::getenv("QSA_TOPK") ? "qsa_topk" : "qsa_topk2"), + pattn = cx.pipeline(lib, av == 1 ? "qsa_attn" : av == 2 ? "qsa_attn2" : "qsa_attn3"), + pmerge = cx.pipeline(lib, v2 ? "qsa_merge2" : "qsa_merge"); + const int MB = MAXC / 4 + 2; + const int AT = prep_only ? 1 : MAXT; + // per-cell inputs for every cell (context + window) + const int NC = P + MAXT; + mtl::Buffer qfull = cx.alloc(4 * (size_t)NC * QH * 2 * QD), kc = cx.alloc(4 * (size_t)NC * QKV * QD), + vc = cx.alloc(4 * (size_t)NC * QKV * QD), iq = cx.alloc(4 * (size_t)NC * IH * ID), + ik = cx.alloc(4 * (size_t)NC * ID), qn = cx.alloc(4 * QD), kn = cx.alloc(4 * QD), iqn = cx.alloc(4 * ID), + ikn = cx.alloc(4 * ID), cs = cx.alloc(8 * (size_t)(MAXC + 16) * 32), qcur = cx.alloc(4 * MAXT * QH * QD), + qidx = cx.alloc(4 * MAXT * IH * ID), kq = cx.alloc((size_t)MAXC * QKV * QD), vq = cx.alloc((size_t)MAXC * QKV * QD), + ksc = cx.alloc(2 * (size_t)MAXC * QKV * 4), vsc = cx.alloc(2 * (size_t)MAXC * QKV * 4), tail = cx.alloc(4 * 4 * ID), + dead = cx.alloc(4 * ID), pooled = cx.alloc(4 * (size_t)MB * ID), scores = cx.alloc(4 * (size_t)MAXT * MB), + ids = cx.alloc(4 * AT * BUDGET), part = cx.alloc(4 * (size_t)AT * 2 * NCH * QH * (QD + 2)), + xp = cx.alloc(4 * (4 * AT * QH * QD + AT * QH * QD / 16)); + ref::Rng r(23); + for (size_t i = 0; i < qfull.bytes / 4; ++i) qfull.as()[i] = r.sym() * 2.0f; + for (size_t i = 0; i < kc.bytes / 4; ++i) kc.as()[i] = r.sym() * 2.0f; + for (size_t i = 0; i < vc.bytes / 4; ++i) vc.as()[i] = r.sym(); + for (size_t i = 0; i < iq.bytes / 4; ++i) iq.as()[i] = r.sym(); + for (size_t i = 0; i < ik.bytes / 4; ++i) ik.as()[i] = r.sym(); + for (int i = 0; i < QD; ++i) { qn.as()[i] = 1.0f + 0.2f * r.sym(); kn.as()[i] = 1.0f + 0.2f * r.sym(); } + for (int i = 0; i < ID; ++i) { iqn.as()[i] = 1.0f + 0.2f * r.sym(); ikn.as()[i] = 1.0f + 0.2f * r.sym(); } + const uint32_t pos_base = 0; + // QSA_MROPE=1: qsa_prep with the M-RoPE table (function constant 24): random (t, h, w) positions per cell, the CPU + // reference rotating pair i by row tab[cell * 3 + i % 3]; QSA_MROPE=id: the identity table, checked bitwise + // against the plain kernel + const char* mre = std::getenv("QSA_MROPE"); + const bool mrope = mre != nullptr, mid = mrope && std::string(mre) == "id"; + mtl::Buffer mtab = cx.alloc(4 * 3 * (size_t)(MAXC + 16)); + { + ref::Rng rm(77); + for (int c = 0; c < MAXC + 16; ++c) + for (int k = 0; k < 3; ++k) mtab.as()[c * 3 + k] = mid ? c : (int)(rm.next() % (uint64_t)(MAXC + 16)); + } + for (int p = 0; p < MAXC + 16; ++p) + for (int i = 0; i < 32; ++i) { + const double inv = std::pow(1e7, -2.0 * i / 64.0), th = (double)p * inv; + cs.as()[(p * 32 + i) * 2] = (float)std::cos(th); + cs.as()[(p * 32 + i) * 2 + 1] = (float)std::sin(th); + } + auto prep = [&](uint32_t pos0, int nt) { + QsaArgs a{pos0, pos_base, (uint32_t)MAXC, (uint32_t)nt, 1e-6f}; + mtl::Dispatch d; + d.p = &pprep; + d.args = {mtl::Arg::B(qfull, 4 * (size_t)pos0 * QH * 2 * QD), mtl::Arg::B(kc, 4 * (size_t)pos0 * QKV * QD), + mtl::Arg::B(vc, 4 * (size_t)pos0 * QKV * QD), mtl::Arg::B(iq, 4 * (size_t)pos0 * IH * ID), + mtl::Arg::B(ik, 4 * (size_t)pos0 * ID), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(iqn), + mtl::Arg::B(ikn), mtl::Arg::B(cs), mtl::Arg::B(qcur), mtl::Arg::B(qidx), mtl::Arg::B(kq), + mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), mtl::Arg::B(tail), mtl::Arg::B(dead), + mtl::Arg::B(pooled), mtl::Arg::V(a), mtl::Arg::B(mtab)}; + d.grid[0] = 33; + d.grid[1] = nt; + d.tg[0] = 256; + return d; + }; + auto window = [&](uint32_t pos0, int nt, std::vector& ds) { + QsaArgs a{pos0, pos_base, (uint32_t)MAXC, (uint32_t)nt, 1e-6f}; + ds.push_back(prep(pos0, nt)); + mtl::Dispatch s; + s.p = &pscore; + s.args = {mtl::Arg::B(qidx), mtl::Arg::B(pooled), mtl::Arg::B(dead), mtl::Arg::B(scores), mtl::Arg::V(a)}; + s.grid[0] = ((pos0 + nt) / 4 + 1 + 7) / 8; + s.grid[1] = nt; + s.tg[0] = 256; + ds.push_back(s); + mtl::Dispatch k; + k.p = &ptopk; + k.args = {mtl::Arg::B(scores), mtl::Arg::B(ids), mtl::Arg::V(a)}; + k.grid[0] = nt; + k.tg[0] = 1024; + ds.push_back(k); + mtl::Dispatch at; + at.p = &pattn; + at.args = {mtl::Arg::B(qcur), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), + mtl::Arg::B(ids), mtl::Arg::B(part), mtl::Arg::V(a)}; + at.grid[0] = v2 ? (BUDGET + 31) / 32 : NCH; + at.grid[1] = v2 ? QKV * 2 : QKV; + at.grid[2] = nt; + at.tg[0] = 256; + ds.push_back(at); + mtl::Dispatch m; + m.p = &pmerge; + m.args = {mtl::Arg::B(part), mtl::Arg::B(qfull, 4 * (size_t)pos0 * QH * 2 * QD), mtl::Arg::B(xp), mtl::Arg::V(a)}; + m.grid[0] = QH; + m.grid[1] = nt; + m.tg[0] = 256; + ds.push_back(m); + }; + if (paired_ref) { + const std::string dir = paired_ref; + const int lr = cx.compile({dir + "/qsa.metal", dir + "/qsa2.metal"}, "#define NO_GGML_COMMON 1\n"); + mtl::Pipeline pr[] = { + cx.pipeline(lr, "qsa_prep", mtl::Constants().b(24, mrope)), cx.pipeline(lr, "qsa_score"), + cx.pipeline(lr, std::getenv("QSA_TOPK") ? "qsa_topk" : "qsa_topk2"), + cx.pipeline(lr, av == 1 ? "qsa_attn" : av == 2 ? "qsa_attn2" : "qsa_attn3"), + cx.pipeline(lr, v2 ? "qsa_merge2" : "qsa_merge")}; + std::vector outputs{&qcur, &qidx, &kq, &vq, &ksc, &vsc, &tail, &dead, + &pooled, &scores, &ids, &part, &xp}; + const char* names[] = {"qcur", "qidx", "kq", "vq", "ksc", "vsc", "tail", "dead", + "pooled", "scores", "ids", "part", "xp"}; + auto run_case = [&](int nt, bool p_only) -> int { + std::vector ds[2]; + window(P, nt, ds[1]); + if (p_only) ds[1].resize(1); + ds[0] = ds[1]; + for (size_t i = 0; i < ds[0].size(); ++i) ds[0][i].p = &pr[i]; + std::vector> expected; + int fails = 0; + for (int v = 0; v < 2; ++v) { + for (auto* b : outputs) std::memset(b->ptr, 0, b->bytes); + ref::Rng history(103); + for (auto* b : {&tail, &dead, &pooled}) + for (size_t i = 0; i < b->bytes / 4; ++i) b->as()[i] = history.sym(); + cx.run(ds[v]); + for (size_t i = 0; i < outputs.size(); ++i) { + auto* b = outputs[i]; + const auto* bytes = static_cast(b->ptr); + if (!v) expected.emplace_back(bytes, bytes + b->bytes); + else { + const bool ok = !std::memcmp(bytes, expected[i].data(), b->bytes); + if (!ab_suite || !ok) std::printf("%s %s\n", names[i], ok ? "PASS" : "FAIL"); + fails += !ok; + } + } + } + if (fails) return 1; + std::printf("paired P=%d nt=%d mrope=%d attn=%d%s, ref / candidate GPU us\n", P, nt, mrope, av, + ab_suite ? " (memcmp PASS)" : ""); + const char* stages[] = {"prep", "score", "topk", "attn", "merge"}; + for (size_t i = 0; i < ds[0].size(); ++i) { + const int repeats = nt > 128 ? 4 : 16, pairs = ab_suite ? 5 : 9; + cx.warm({ds[0][i], ds[1][i]}, 0.01); + cx.run({ds[0][i]}, repeats); + cx.run({ds[1][i]}, repeats); + std::vector times[2], ratios; + for (int k = 0; k < pairs; ++k) { + double t0 = 0, t1 = 0; + for (int turn = 0; turn < 2; ++turn) { + const int v = (k + turn) & 1; + const double dt = cx.run({ds[v][i]}, repeats) * 1e6 / repeats; + if (v == 0) t0 = dt; else t1 = dt; + times[v].push_back(dt); + } + ratios.push_back(t1 / t0); + } + std::printf("%s %.3f / %.3f (ratio %.4f)\n", stages[i], median(times[0]), median(times[1]), median(ratios)); + } + return 0; + }; + if (ab_suite) { + for (int nt : {1, 2, 3, 4, 8, 1984, 2048}) + if (run_case(nt, prep_only || nt > 8)) return 1; + return 0; + } + return run_case(paired_nt, prep_only); + } + // the context, one cell at a time + { + std::vector ds; + for (int p = 0; p < P; ++p) ds.push_back(prep(p, 1)); + cx.run(ds); + } + if (bench) { + std::printf("device %s, context %d cells\n", cx.device_name().c_str(), P); + for (int nt : {1, 3}) { + std::vector ds; + window(P, nt, ds); + const char* names[] = {"prep", "score", "topk", "attn", "merge"}; + double total = 0; + for (int i = 0; i < 5; ++i) { + std::vector one(32, ds[i]); + cx.warm(one, 0.05); + std::vector t; + for (int k = 0; k < 7; ++k) t.push_back(cx.run(one) / 32); + std::printf(" ntok %d %-6s %7.2f us\n", nt, names[i], median(t) * 1e6); + total += median(t); + } + std::vector t; + for (int k = 0; k < 7; ++k) t.push_back(cx.run(ds)); + std::printf(" ntok %d window (5 dependent dispatches): %.2f us (sum of parts %.2f)\n", nt, median(t) * 1e6, + total * 1e6); + } + return 0; + } + const int NT = 3; + std::vector ds; + window(P, NT, ds); + // stage by stage, re-reading GPU state in between + cx.run({ds[0]}); + int fails = 0; + // ---- prep: qcur, qidx, the window cells' KV codes and the pooled keys vs CPU + { + Err eq, ei, ekv, ep; + auto normrope = [&](const float* x, const float* g, int w, int absp, bool rope, std::vector& out) { + double ss = 0; + for (int i = 0; i < w; ++i) ss += (double)x[i] * x[i]; + const double rs = 1.0 / std::sqrt(ss / w + 1e-6); + out.assign(w, 0.0); + for (int i = 0; i < w; ++i) out[i] = rs * x[i] * g[i]; + if (rope) + for (int i = 0; i < 32; ++i) { + const int pp = mrope ? mtab.as()[absp * 3 + i % 3] : absp; + const double c = cs.as()[(pp * 32 + i) * 2], s = cs.as()[(pp * 32 + i) * 2 + 1]; + const double a = out[i], b = out[i + 32]; + out[i] = a * c - b * s; + out[i + 32] = a * s + b * c; + } + }; + std::vector o; + for (int t = 0; t < NT; ++t) { + const int cell = P + t; + for (int h = 0; h < QH; ++h) { + normrope(qfull.as() + ((size_t)cell * QH + h) * 2 * QD, qn.as(), QD, cell, true, o); + for (int i = 0; i < QD; ++i) eq.add(qcur.as()[(t * QH + h) * QD + i], o[i]); + } + for (int h = 0; h < IH; ++h) { + normrope(iq.as() + ((size_t)cell * IH + h) * ID, iqn.as(), ID, cell, true, o); + for (int i = 0; i < ID; ++i) ei.add(qidx.as()[(t * IH + h) * ID + i], o[i]); + } + for (int h = 0; h < QKV; ++h) { + normrope(kc.as() + ((size_t)cell * QKV + h) * QD, kn.as(), QD, cell, true, o); + for (int i = 0; i < QD; ++i) { + const size_t row = (size_t)cell * QKV + h; + const double deq = (double)(int8_t)kq.as()[row * QD + i] * h2f(ksc.as()[row * 4 + i / 64]); + ekv.add(deq, o[i]); + } + } + } + // pooled keys of the last few completed blocks + const int nb = (P + NT) / 4; + for (int b = std::max(0, nb - 8); b < nb; ++b) { + std::vector m(ID); + for (int i = 0; i < ID; ++i) + m[i] = (((ik.as()[(4 * b) * ID + i] + ik.as()[(4 * b + 1) * ID + i]) + + ik.as()[(4 * b + 2) * ID + i]) + ik.as()[(4 * b + 3) * ID + i]) * 0.25f; + normrope(m.data(), ikn.as(), ID, 4 * b, true, o); + for (int i = 0; i < ID; ++i) ep.add(pooled.as()[b * ID + i], o[i]); + } + // int8 quantization error is ~1/254 of max|x| per group: the KV check is a quantization tolerance + const bool ok = eq.l2() < 1e-5 && ei.l2() < 1e-5 && ep.l2() < 1e-5 && ekv.l2() < 6e-3; + fails += !ok; + std::printf("prep: q rel L2 %.2e, indexer q %.2e, pooled keys %.2e, K int8 (vs float) %.2e %s\n", eq.l2(), + ei.l2(), ep.l2(), ekv.l2(), ok ? "PASS" : "FAIL"); + } + // ---- score + cx.run({ds[1]}); + { + Err es; + for (int t = 0; t < NT; ++t) { + const int n_kv = P + t + 1, n_bid = n_kv / 4; + for (int b = 0; b <= n_bid; ++b) { + const float* key = b == n_bid ? dead.as() : pooled.as() + b * ID; + double s = 0; + for (int h = 0; h < IH; ++h) { + double d = 0; + for (int i = 0; i < ID; ++i) d += (double)qidx.as()[(t * IH + h) * ID + i] * key[i]; + s += std::max(d, 0.0); + } + if (b == n_bid && (n_kv & 3)) { + // float(s + 1e9) keeps no bits of s: only check that the bias is there + if (!(scores.as()[t * MB + b] >= 1e9f)) es.add(1.0, 0.0); + } else { + es.add(scores.as()[t * MB + b], s); + } + } + } + const bool ok = es.l2() < 1e-5; + fails += !ok; + std::printf("score: rel L2 %.2e (tail +1e9 checked separately) %s\n", es.l2(), ok ? "PASS" : "FAIL"); + } + // ---- top-k (exact, from the GPU's scores) + cx.run({ds[2]}); + { + int bad = 0; + for (int t = 0; t < NT; ++t) { + const int n_kv = P + t + 1, n_bid = n_kv / 4, width = std::min(n_kv, BUDGET); + std::vector cells(n_kv); + for (int c = 0; c < n_kv; ++c) cells[c] = c; + if (n_kv > BUDGET) { + const float* sc = scores.as() + t * MB; + std::stable_sort(cells.begin(), cells.end(), [&](int x, int y) { return sc[x / 4] > sc[y / 4]; }); + cells.resize(width); + std::sort(cells.begin(), cells.end()); + } + (void)n_bid; + for (int i = 0; i < width; ++i) bad += ids.as()[t * BUDGET + i] != cells[i]; + } + fails += bad != 0; + std::printf("top-k: %d of %d ids differ (context %d: selection %s) %s\n", bad, NT * std::min(P + 1, BUDGET), P, + P + 1 > BUDGET ? "active" : "identity", bad ? "FAIL" : "PASS"); + } + // ---- attention + merge (from the GPU's qcur, ids and KV codes) + cx.run({ds[3], ds[4]}); + { + Err eo; + for (int t = 0; t < NT; ++t) { + const int n_kv = P + t + 1, width = std::min(n_kv, BUDGET); + for (int h = 0; h < QH; ++h) { + const int g = h / (QH / QKV); + std::vector s(width); + double mx = -1e300; + for (int c = 0; c < width; ++c) { + const size_t row = (size_t)ids.as()[t * BUDGET + c] * QKV + g; + double d = 0; + for (int i = 0; i < QD; ++i) + d += (double)qcur.as()[(t * QH + h) * QD + i] * (int8_t)kq.as()[row * QD + i] * + h2f(ksc.as()[row * 4 + i / 64]); + s[c] = d / 16.0; + mx = std::max(mx, s[c]); + } + double l = 0; + for (int c = 0; c < width; ++c) l += (s[c] = std::exp(s[c] - mx)); + for (int i = 0; i < QD; ++i) { + double o = 0; + for (int c = 0; c < width; ++c) { + const size_t row = (size_t)ids.as()[t * BUDGET + c] * QKV + g; + o += s[c] * (int8_t)vq.as()[row * QD + i] * h2f(vsc.as()[row * 4 + i / 64]); + } + o /= l; + const double gate = qfull.as()[((size_t)(P + t) * QH + h) * 2 * QD + QD + i]; + eo.add(xp.as()[t * QH * QD + h * QD + i], o / (1.0 + std::exp(-gate))); + } + } + } + const bool ok = eo.l2() < 1e-5; + fails += !ok; + std::printf("attention + gate: rel L2 %.2e %s\n", eo.l2(), ok ? "PASS" : "FAIL"); + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; +} diff --git a/results/round39/astra_r39_qsa/evaluator/metal/tests/qsa_prompt_parity.cpp b/results/round39/astra_r39_qsa/evaluator/metal/tests/qsa_prompt_parity.cpp new file mode 100644 index 0000000..02dde83 --- /dev/null +++ b/results/round39/astra_r39_qsa/evaluator/metal/tests/qsa_prompt_parity.cpp @@ -0,0 +1,620 @@ +// metal/tests/qsa_prompt_parity.cpp - qsa_attn_mm (qsa_prompt.metal, the matrix units, fp16 operands) against qsa_attn +// (fp32, itself checked by qsa_parity), both followed by qsa_merge: the attention output of T query tokens on random +// int8 caches, within the budget (all cells, causal) and past it (2051 selected cells per token); then the speed. +// qsa_prompt_parity [--bench | --flash-compare T pos0 baseline_qsa_prompt.metal [layers pairs]] +// qsa_prompt_parity --score-compare T pos0 baseline_qsa_prompt.metal baseline_qsa.metal pairs +// qsa_prompt_parity --topk-compare T pos0 baseline_qsa.metal pairs +// QSA_TOPK_EDGE=1: ties/signed zero; QSA_FLASH_EDGE=1|2: signed zero/tiny cancellation +// qsa_prompt_parity --mm-compare attn|score T pos0 baseline_qsa_prompt.metal pairs +// qsa_prompt_parity --ab T pos0 baseline_qsa_prompt.metal baseline_qsa.metal pairs +#include "coop_probe.h" +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { +constexpr int QH = 24, QKV = 2, QD = 256, QB = 2051, NCH = 33; +struct QsaArgs { uint32_t pos0, pos_base, max_cells, n_tok; float eps; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } + +int topk_compare(int T, uint32_t pos0, const std::string& baseline, int pairs) { + if (T < 1 || T > 1984 || pos0 > 30000 || pairs < 1 || pairs > 15) return 2; + mtl::Context cx; + const int lib = cx.compile({"qsa.metal"}, "#define NO_GGML_COMMON 1\n"); + const int old = cx.compile({baseline}, "#define NO_GGML_COMMON 1\n"); + const uint32_t cells = pos0 + T, stride = cells / 4 + 2; + auto scores = cx.alloc(4 * size_t(T) * stride), out = cx.alloc(4 * size_t(T) * QB), + ref = cx.alloc(out.bytes); + ref::Rng r(5); + const bool edge = std::getenv("QSA_TOPK_EDGE") != nullptr; + for (int t = 0; t < T; ++t) + for (uint32_t b = 0; b < stride; ++b) { + float v = std::fabs(r.sym()) * 32.0f; + if (edge) { + if (t % 4 == 0) v = 1.0f; + if (t % 4 == 1) v = float(int(r.next() % 7) - 3); + if (t % 4 == 2) v = b & 1 ? -0.0f : 0.0f; + } + if (!edge && b == (pos0 + t + 1) / 4 && ((pos0 + t + 1) & 3)) v += 1e9f; + scores.as()[size_t(t) * stride + b] = v; + } + std::memset(out.ptr, 0xa5, out.bytes); std::memset(ref.ptr, 0xa5, ref.bytes); + auto p = cx.pipeline(lib, "qsa_topk"), pold = cx.pipeline(old, "qsa_topk"); + QsaArgs a{pos0, 0, cells, uint32_t(T), 1e-6f}; + mtl::Dispatch d; + d.p = &pold; + d.args = {mtl::Arg::B(scores), mtl::Arg::B(ref), mtl::Arg::V(a)}; + d.grid[0] = T; d.tg[0] = 1024; + auto candidate = d; + candidate.p = &p; candidate.args[1] = mtl::Arg::B(out); + cx.run({d}); cx.run({candidate}); + const bool equal = std::memcmp(out.ptr, ref.ptr, out.bytes) == 0; + std::vector before, after; + for (int i = 0; i < pairs; ++i) { + if (i & 1) { after.push_back(cx.run({candidate})); before.push_back(cx.run({d})); } + else { before.push_back(cx.run({d})); after.push_back(cx.run({candidate})); } + } + std::printf("device %s; topk T %d pos0 %u memcmp %s before %.6f after %.6f ms\n", + cx.device_name().c_str(), T, pos0, equal ? "PASS" : "FAIL", median(before) * 1e3, median(after) * 1e3); + for (int i = 0; i < pairs; ++i) + std::printf("pair %d before %.6f after %.6f ms\n", i, before[i] * 1e3, after[i] * 1e3); + return equal ? 0 : 1; +} + +int score_compare(int T, uint32_t pos0, const std::string& baseline_prompt, + const std::string& baseline_qsa, int pairs) { + if (T < 1 || T > 1984 || pos0 > 30000 || pairs < 1 || pairs > 15) return 2; + mtl::Context cx; + const int lib = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + const int old = cx.compile({baseline_qsa, baseline_prompt}, "#define NO_GGML_COMMON 1\n"); + const size_t cells = pos0 + T, nb = cells / 4 + 1, stride = cells / 4 + 2; + auto q = cx.alloc(4 * size_t(T) * 4 * 128), k = cx.alloc(4 * nb * 128), dead = cx.alloc(4 * 128); + auto scores = cx.alloc(4 * size_t(T) * stride), ref_scores = cx.alloc(scores.bytes), + ids = cx.alloc(4 * size_t(T) * QB), ref_ids = cx.alloc(ids.bytes); + ref::Rng r(5); + for (auto* b : {&q, &k, &dead}) + for (size_t i = 0; i < b->bytes / 4; ++i) b->as()[i] = r.sym(); + for (auto* b : {&scores, &ref_scores, &ids, &ref_ids}) std::memset(b->ptr, 0xa5, b->bytes); + const auto constants = mtl::Constants().b(7, coop_k4(cx)); + auto score = cx.pipeline(lib, "qsa_score_mm", constants), + old_score = cx.pipeline(old, "qsa_score_mm", constants), + topk = cx.pipeline(lib, "qsa_topk"), old_topk = cx.pipeline(old, "qsa_topk"); + QsaArgs a{pos0, 0, uint32_t(cells), uint32_t(T), 1e-6f}; + mtl::Dispatch ds, dt; + ds.p = &old_score; + ds.args = {mtl::Arg::B(q), mtl::Arg::B(k), mtl::Arg::B(dead), mtl::Arg::B(ref_scores), mtl::Arg::V(a)}; + ds.grid[0] = (nb + 63) / 64; ds.grid[1] = (T + 15) / 16; ds.tg[0] = 128; + dt.p = &old_topk; + dt.args = {mtl::Arg::B(ref_scores), mtl::Arg::B(ref_ids), mtl::Arg::V(a)}; + dt.grid[0] = T; dt.tg[0] = 1024; + auto cs = ds, ct = dt; + cs.p = &score; cs.args[3] = mtl::Arg::B(scores); + ct.p = &topk; ct.args[0] = mtl::Arg::B(scores); ct.args[1] = mtl::Arg::B(ids); + cx.run({ds, dt}, 1, false, true); cx.run({cs, ct}, 1, false, true); + const bool same_scores = std::memcmp(scores.ptr, ref_scores.ptr, scores.bytes) == 0; + const bool same_ids = std::memcmp(ids.ptr, ref_ids.ptr, ids.bytes) == 0; + std::printf("device %s; score+topk T %d pos0 %u scores %s qids %s tg %d -> %d bytes\n", + cx.device_name().c_str(), T, pos0, same_scores ? "PASS" : "FAIL", same_ids ? "PASS" : "FAIL", + old_score.static_tg_mem, score.static_tg_mem); + const std::vector before_dispatches[] = {{ds}, {dt}, {ds, dt}}; + const std::vector after_dispatches[] = {{cs}, {ct}, {cs, ct}}; + for (int stage = 0; stage < 3; ++stage) { + const auto& before_dispatch = before_dispatches[stage]; + const auto& after_dispatch = after_dispatches[stage]; + cx.run(before_dispatch, 1, false, true); cx.run(after_dispatch, 1, false, true); + std::vector before, after; + for (int i = 0; i < pairs; ++i) { + if (i & 1) { + after.push_back(cx.run(after_dispatch, 1, false, true)); + before.push_back(cx.run(before_dispatch, 1, false, true)); + } else { + before.push_back(cx.run(before_dispatch, 1, false, true)); + after.push_back(cx.run(after_dispatch, 1, false, true)); + } + } + std::printf("%s before %.6f after %.6f ms\n", stage == 0 ? "score" : stage == 1 ? "topk" : "combined", + median(before) * 1e3, median(after) * 1e3); + for (int i = 0; i < pairs; ++i) + std::printf("pair %d before %.6f after %.6f ms\n", i, before[i] * 1e3, after[i] * 1e3); + } + const bool stable = std::memcmp(scores.ptr, ref_scores.ptr, scores.bytes) == 0 && + std::memcmp(ids.ptr, ref_ids.ptr, ids.bytes) == 0; + std::printf("post-timing memcmp %s\n", stable ? "PASS" : "FAIL"); + return same_scores && same_ids && stable ? 0 : 1; +} + +int flash_compare(int T, uint32_t pos0, const std::string& baseline, int layers, int pairs) { + if (T < 1 || T > 1984 || pos0 > 30000 || layers < 1 || pairs < 1) return 2; + mtl::Context cx; + const int lib = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + const int old = cx.compile({"qsa.metal", baseline}, "#define NO_GGML_COMMON 1\n"); + const size_t cells = pos0 + T, nq = (size_t)T * QH * QD; + mtl::Buffer kq = cx.alloc(cells * QKV * QD), vq = cx.alloc(kq.bytes), + ks = cx.alloc(2 * cells * QKV * 4), vs = cx.alloc(ks.bytes), + qh = cx.alloc(2 * nq), gate = cx.alloc(8 * nq), ids = cx.alloc(4 * (size_t)T * QB), + out = cx.alloc(2 * nq), ref = cx.alloc(out.bytes); + ref::Rng r(5); + for (size_t i = 0; i < kq.bytes; ++i) { + kq.as()[i] = int8_t(int(r.next() % 255) - 127); + vq.as()[i] = int8_t(int(r.next() % 255) - 127); + } + for (size_t i = 0; i < ks.bytes / 2; ++i) { + ks.as()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); + vs.as()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); + } + for (size_t i = 0; i < nq; ++i) qh.as()[i] = ref::f32_to_f16(r.sym() * 3.0f); + for (size_t i = 0; i < 2 * nq; ++i) gate.as()[i] = r.sym(); + for (int t = 0; t < T; ++t) { + const uint32_t n = pos0 + t + 1; + int* o = ids.as() + (size_t)t * QB; + if (n <= QB) { + for (uint32_t c = 0; c < n; ++c) o[c] = int(c); + } else { + std::set selected; + while (selected.size() < QB) selected.insert(int(r.next() % n)); + std::copy(selected.begin(), selected.end(), o); + } + } + if (const char* edge = std::getenv("QSA_FLASH_EDGE")) { + const int mode = std::atoi(edge); + for (size_t i = 0; i < kq.bytes; ++i) { + kq.as()[i] = 0; + vq.as()[i] = mode == 2 ? int8_t((i / (QKV * QD)) & 1 ? -1 : 1) : 0; + } + for (size_t i = 0; i < ks.bytes / 2; ++i) { + ks.as()[i] = i & 1 ? 0x8000 : 0; + vs.as()[i] = mode == 2 ? 0x0400 : (i & 1 ? 0x8000 : 0); + } + for (size_t i = 0; i < nq; ++i) qh.as()[i] = i & 1 ? 0x8000 : 0; + } + const auto constants = mtl::Constants().b(7, coop_k4(cx)); + auto p = cx.pipeline(lib, "qsa_flash", constants), pold = cx.pipeline(old, "qsa_flash", constants); + QsaArgs a{pos0, 0, uint32_t(cells), uint32_t(T), 1e-6f}; + mtl::Dispatch d; + d.p = &pold; + d.args = {mtl::Arg::B(qh), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ks), mtl::Arg::B(vs), + mtl::Arg::B(ids), mtl::Arg::B(gate), mtl::Arg::B(ref), mtl::Arg::V(a)}; + d.grid[0] = QKV; d.grid[1] = T; d.tg[0] = 128; + cx.run({d}); + // Poison only implicit rows, including the early rows of a chunk crossing the budget. + for (int t = 0; t < T && pos0 + t + 1 <= QB; ++t) + std::fill_n(ids.as() + (size_t)t * QB, QB, -1); + auto candidate = d; + candidate.p = &p; candidate.args[7] = mtl::Arg::B(out); + cx.run({candidate}); + const bool equal = std::memcmp(out.ptr, ref.ptr, out.bytes) == 0; + // Baseline needs explicit ids for the paired timings below. + for (int t = 0; t < T && pos0 + t + 1 <= QB; ++t) + for (uint32_t c = 0; c < pos0 + t + 1; ++c) ids.as()[(size_t)t * QB + c] = int(c); + // paired, alternating order: each sample is one command buffer of `layers` back-to-back dispatches (the 12 QSA + // layers of a prompt chunk) + std::vector before, after; + for (int i = 0; i < pairs; ++i) { + if (i & 1) { after.push_back(cx.run({candidate}, layers)); before.push_back(cx.run({d}, layers)); } + else { before.push_back(cx.run({d}, layers)); after.push_back(cx.run({candidate}, layers)); } + } + std::printf("device %s; T %d pos0 %u: memcmp %s; flash x%d before median %.3f min %.3f ms, after median %.3f min %.3f " + "ms; tg %d -> %d bytes\n", cx.device_name().c_str(), T, pos0, equal ? "PASS" : "FAIL", layers, + median(before) * 1e3, *std::min_element(before.begin(), before.end()) * 1e3, median(after) * 1e3, + *std::min_element(after.begin(), after.end()) * 1e3, pold.static_tg_mem, p.static_tg_mem); + for (int i = 0; i < pairs; ++i) + std::printf("pair %d before %.6f after %.6f ms\n", i, before[i] * 1e3, after[i] * 1e3); + return equal ? 0 : 1; +} + +int mm_compare(bool score, int T, uint32_t pos0, const std::string& baseline, int pairs) { + if (T < 1 || T > 1984 || pos0 > 30000 || pairs < 1 || pairs > 15) return 2; + mtl::Context cx; + const int lib = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + const int old = cx.compile({"qsa.metal", baseline}, "#define NO_GGML_COMMON 1\n"); + const size_t cells = pos0 + T; + QsaArgs a{pos0, 0, uint32_t(cells), uint32_t(T), 1e-6f}; + const bool k4 = coop_k4(cx); + const auto constants = mtl::Constants().b(7, k4); + std::printf("device %s coop_k4 %d\n", cx.device_name().c_str(), int(k4)); + const std::string name = score ? "qsa_score_mm" : "qsa_attn_mm"; + auto p = cx.pipeline(lib, name.c_str(), constants), pold = cx.pipeline(old, name.c_str(), constants); + mtl::Dispatch d; + d.p = &pold; + mtl::Buffer out, ref; + auto compare = [&](const mtl::Dispatch& before_dispatch, const mtl::Dispatch& after_dispatch) { + cx.run({before_dispatch}); cx.run({after_dispatch}); + const bool equal = std::memcmp(out.ptr, ref.ptr, out.bytes) == 0; + cx.warm({before_dispatch, after_dispatch}, 0.05); + std::vector before, after; + for (int i = 0; i < pairs; ++i) { + if (i & 1) { + after.push_back(cx.run({after_dispatch})); before.push_back(cx.run({before_dispatch})); + } else { + before.push_back(cx.run({before_dispatch})); after.push_back(cx.run({after_dispatch})); + } + } + std::printf("%s T %d pos0 %u memcmp %s before %.6f after %.6f ms tg %d -> %d\n", name.c_str(), T, pos0, + equal ? "PASS" : "FAIL", median(before) * 1e3, median(after) * 1e3, pold.static_tg_mem, p.static_tg_mem); + for (int i = 0; i < pairs; ++i) + std::printf("pair %d before %.6f after %.6f ms\n", i, before[i] * 1e3, after[i] * 1e3); + const bool stable = std::memcmp(out.ptr, ref.ptr, out.bytes) == 0; + std::printf("post-timing memcmp %s\n", stable ? "PASS" : "FAIL"); + return equal && stable ? 0 : 1; + }; + ref::Rng r(5); + if (score) { + const size_t nb = cells / 4 + 1, stride = cells / 4 + 2; + auto q = cx.alloc(4 * size_t(T) * 4 * 128), k = cx.alloc(4 * nb * 128), dead = cx.alloc(4 * 128); + for (auto* b : {&q, &k, &dead}) + for (size_t i = 0; i < b->bytes / 4; ++i) b->as()[i] = r.sym(); + out = cx.alloc(4 * size_t(T) * stride); ref = cx.alloc(out.bytes); + std::memset(out.ptr, 0xa5, out.bytes); std::memset(ref.ptr, 0xa5, ref.bytes); + d.args = {mtl::Arg::B(q), mtl::Arg::B(k), mtl::Arg::B(dead), mtl::Arg::B(ref), mtl::Arg::V(a)}; + d.grid[0] = (nb + 63) / 64; d.grid[1] = (T + 15) / 16; d.tg[0] = 128; + auto candidate = d; + candidate.p = &p; candidate.args[3] = mtl::Arg::B(out); + return compare(d, candidate); + } + auto k = cx.alloc(cells * QKV * QD), v = cx.alloc(k.bytes), ks = cx.alloc(2 * cells * QKV * 4), vs = cx.alloc(ks.bytes), + q = cx.alloc(2 * size_t(T) * QH * QD), ids = cx.alloc(4 * size_t(T) * QB); + for (size_t i = 0; i < k.bytes; ++i) { + k.as()[i] = int8_t(int(r.next() % 255) - 127); v.as()[i] = int8_t(int(r.next() % 255) - 127); + } + for (auto* b : {&ks, &vs}) + for (size_t i = 0; i < b->bytes / 2; ++i) b->as()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); + for (size_t i = 0; i < q.bytes / 2; ++i) q.as()[i] = ref::f32_to_f16(r.sym() * 3.0f); + for (int t = 0; t < T; ++t) { + const uint32_t n = pos0 + t + 1; + for (int c = 0; c < QB; ++c) ids.as()[size_t(t) * QB + c] = n <= QB ? -1 : int(uint64_t(c) * n / QB); + } + out = cx.alloc(4 * size_t(T) * NCH * QH * (QD + 2)); ref = cx.alloc(out.bytes); + std::memset(out.ptr, 0xa5, out.bytes); std::memset(ref.ptr, 0xa5, ref.bytes); + d.args = {mtl::Arg::B(q), mtl::Arg::B(k), mtl::Arg::B(v), mtl::Arg::B(ks), mtl::Arg::B(vs), + mtl::Arg::B(ids), mtl::Arg::B(ref), mtl::Arg::V(a)}; + d.grid[0] = NCH; d.grid[1] = QKV; d.grid[2] = T; d.tg[0] = 32; + auto candidate = d; + candidate.p = &p; candidate.args[6] = mtl::Arg::B(out); + const int result = compare(d, candidate); + const size_t nq = size_t(T) * QH * QD; + auto gate = cx.alloc(8 * nq), merged = cx.alloc(4 * (4 * nq + nq / 16)), merged_ref = cx.alloc(merged.bytes); + for (size_t i = 0; i < gate.bytes / 4; ++i) gate.as()[i] = r.sym(); + auto merge = cx.pipeline(lib, "qsa_merge"); + mtl::Dispatch dm; + dm.p = &merge; + dm.args = {mtl::Arg::B(ref), mtl::Arg::B(gate), mtl::Arg::B(merged_ref), mtl::Arg::V(a)}; + dm.grid[0] = QH; dm.grid[1] = T; dm.tg[0] = 256; + auto cm = dm; + cm.args[0] = mtl::Arg::B(out); cm.args[2] = mtl::Arg::B(merged); + cx.run({dm}); cx.run({cm}); + const bool same_merge = std::memcmp(merged.ptr, merged_ref.ptr, merged.bytes) == 0; + std::printf("merge+prepared memcmp %s\n", same_merge ? "PASS" : "FAIL"); + return result || !same_merge; +} +} // namespace + +int ab_compare(int T, uint32_t pos0, const std::string& baseline_prompt, const std::string& baseline_qsa, int pairs) { + if (T < 1 || T > 2048 || pos0 > 30000 || pairs < 1 || pairs > 15) return 2; + mtl::Context cx; + const int lib = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + const int old = cx.compile({baseline_qsa, baseline_prompt}, "#define NO_GGML_COMMON 1\n"); + const size_t cells = pos0 + T, nb = cells / 4 + 1, stride = cells / 4 + 2; + QsaArgs a{pos0, 0, uint32_t(cells), uint32_t(T), 1e-6f}; + const bool mrope = std::getenv("QSA_MROPE") != nullptr; + const bool qf16 = std::getenv("QSA_QF16") != nullptr; + const auto prep_c = mtl::Constants().b(24, mrope).b(25, qf16); + const auto score_c = mtl::Constants().b(7, coop_k4(cx)); + auto prep = cx.pipeline(lib, "qsa_prep", prep_c), old_prep = cx.pipeline(old, "qsa_prep", prep_c); + auto score = cx.pipeline(lib, "qsa_score_mm", score_c), old_score = cx.pipeline(old, "qsa_score_mm", score_c); + mtl::Buffer qfull = cx.alloc(4 * size_t(T) * QH * 2 * QD), kc = cx.alloc(4 * size_t(T) * QKV * QD), + vc = cx.alloc(4 * size_t(T) * QKV * QD), iq = cx.alloc(4 * size_t(T) * 4 * 128), + ik = cx.alloc(4 * size_t(T) * 128), qn = cx.alloc(4 * QD), kn = cx.alloc(4 * QD), + iqn = cx.alloc(4 * 128), ikn = cx.alloc(4 * 128), cs = cx.alloc(8 * (cells + 16) * 32), + mtab = cx.alloc(4 * 3 * (cells + 16)); + mtl::Buffer qcur[2] = {cx.alloc((qf16 ? 2 : 4) * size_t(T) * QH * QD), cx.alloc((qf16 ? 2 : 4) * size_t(T) * QH * QD)}, + qidx[2] = {cx.alloc(4 * size_t(T) * 4 * 128), cx.alloc(4 * size_t(T) * 4 * 128)}, + kq[2] = {cx.alloc(cells * QKV * QD), cx.alloc(cells * QKV * QD)}, + vq[2] = {cx.alloc(cells * QKV * QD), cx.alloc(cells * QKV * QD)}, + ksc[2] = {cx.alloc(2 * cells * QKV * 4), cx.alloc(2 * cells * QKV * 4)}, + vsc[2] = {cx.alloc(2 * cells * QKV * 4), cx.alloc(2 * cells * QKV * 4)}, + tail[2] = {cx.alloc(4 * 4 * 128), cx.alloc(4 * 4 * 128)}, + dead[2] = {cx.alloc(4 * 128), cx.alloc(4 * 128)}, + pooled[2] = {cx.alloc(4 * stride * 128), cx.alloc(4 * stride * 128)}, + scores[2] = {cx.alloc(4 * size_t(T) * stride), cx.alloc(4 * size_t(T) * stride)}; + ref::Rng r(5); + for (auto* b : {&qfull, &kc, &vc, &iq, &ik}) + for (size_t i = 0; i < b->bytes / 4; ++i) b->as()[i] = r.sym(); + for (int i = 0; i < QD; ++i) { qn.as()[i] = 1.0f + 0.2f * r.sym(); kn.as()[i] = 1.0f + 0.2f * r.sym(); } + for (int i = 0; i < 128; ++i) { iqn.as()[i] = 1.0f + 0.2f * r.sym(); ikn.as()[i] = 1.0f + 0.2f * r.sym(); } + for (size_t p = 0; p < cells + 16; ++p) { + for (int k = 0; k < 3; ++k) mtab.as()[p * 3 + k] = int(r.next() % (cells + 16)); + for (int i = 0; i < 32; ++i) { + const double inv = std::pow(1e7, -2.0 * i / 64.0), th = double(p) * inv; + cs.as()[(p * 32 + i) * 2] = float(std::cos(th)); + cs.as()[(p * 32 + i) * 2 + 1] = float(std::sin(th)); + } + } + auto reset_state = [&](int v) { + ref::Rng h(103); + for (auto* b : {&qcur[v], &qidx[v], &kq[v], &vq[v], &ksc[v], &vsc[v], &scores[v]}) + std::memset(b->ptr, 0xa5, b->bytes); + for (auto* b : {&tail[v], &dead[v], &pooled[v]}) + for (size_t i = 0; i < b->bytes / 4; ++i) b->as()[i] = h.sym(); + }; + mtl::Dispatch dp[2], ds[2]; + for (int v = 0; v < 2; ++v) { + dp[v].p = v ? &prep : &old_prep; + dp[v].args = {mtl::Arg::B(qfull), mtl::Arg::B(kc), mtl::Arg::B(vc), mtl::Arg::B(iq), mtl::Arg::B(ik), + mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(iqn), mtl::Arg::B(ikn), mtl::Arg::B(cs), + mtl::Arg::B(qcur[v]), mtl::Arg::B(qidx[v]), mtl::Arg::B(kq[v]), mtl::Arg::B(vq[v]), + mtl::Arg::B(ksc[v]), mtl::Arg::B(vsc[v]), mtl::Arg::B(tail[v]), mtl::Arg::B(dead[v]), + mtl::Arg::B(pooled[v]), mtl::Arg::V(a), mtl::Arg::B(mtab)}; + dp[v].grid[0] = 33; dp[v].grid[1] = T; dp[v].tg[0] = 256; + ds[v].p = v ? &score : &old_score; + ds[v].args = {mtl::Arg::B(qidx[v]), mtl::Arg::B(pooled[v]), mtl::Arg::B(dead[v]), mtl::Arg::B(scores[v]), mtl::Arg::V(a)}; + ds[v].grid[0] = (nb + 63) / 64; ds[v].grid[1] = (T + 15) / 16; ds[v].tg[0] = 128; + reset_state(v); + cx.run({dp[v], ds[v]}, 1, false, true); + } + const mtl::Buffer* b0[] = {&qcur[0], &qidx[0], &kq[0], &vq[0], &ksc[0], &vsc[0], &tail[0], &dead[0], &pooled[0], &scores[0]}; + const mtl::Buffer* b1[] = {&qcur[1], &qidx[1], &kq[1], &vq[1], &ksc[1], &vsc[1], &tail[1], &dead[1], &pooled[1], &scores[1]}; + bool ok = true; + for (int i = 0; i < 10; ++i) ok = ok && (std::memcmp(b0[i]->ptr, b1[i]->ptr, b0[i]->bytes) == 0); + std::printf("device %s; ab T %d pos0 %u mrope %d qf16 %d memcmp %s\n", + cx.device_name().c_str(), T, pos0, int(mrope), int(qf16), ok ? "PASS" : "FAIL"); + if (!ok) return 1; + const char* sname[] = {"qsa_prep", "qsa_score_mm", "combined"}; + const std::vector d0[] = {{dp[0]}, {ds[0]}, {dp[0], ds[0]}}; + const std::vector d1[] = {{dp[1]}, {ds[1]}, {dp[1], ds[1]}}; + for (int s = 0; s < 3; ++s) { + cx.warm(d0[s], 0.05); + cx.run(d0[s], 12, false, true); + cx.run(d1[s], 12, false, true); + std::vector before, after, ratios; + for (int i = 0; i < pairs; ++i) { + double t0 = 0, t1 = 0; + if (i & 1) { + t1 = cx.run(d1[s], 12, false, true) * 1e3 / 12.0; + t0 = cx.run(d0[s], 12, false, true) * 1e3 / 12.0; + } else { + t0 = cx.run(d0[s], 12, false, true) * 1e3 / 12.0; + t1 = cx.run(d1[s], 12, false, true) * 1e3 / 12.0; + } + before.push_back(t0); after.push_back(t1); ratios.push_back(t1 / t0); + } + std::printf(" %-12s before %.4f after %.4f ms (ratio %.4f)\n", sname[s], median(before), median(after), median(ratios)); + } + return 0; +} + +int main(int argc, char** argv) { + if (argc == 7 && std::string(argv[1]) == "--ab") + return ab_compare(std::atoi(argv[2]), uint32_t(std::strtoul(argv[3], nullptr, 10)), + argv[4], argv[5], std::atoi(argv[6])); + if (argc == 7 && std::string(argv[1]) == "--score-compare") + return score_compare(std::atoi(argv[2]), uint32_t(std::strtoul(argv[3], nullptr, 10)), + argv[4], argv[5], std::atoi(argv[6])); + if (argc == 6 && std::string(argv[1]) == "--topk-compare") + return topk_compare(std::atoi(argv[2]), uint32_t(std::strtoul(argv[3], nullptr, 10)), + argv[4], std::atoi(argv[5])); + if (argc == 7 && std::string(argv[1]) == "--mm-compare") { + const std::string target = argv[2]; + if (target != "attn" && target != "score") return 2; + return mm_compare(target == "score", std::atoi(argv[3]), + uint32_t(std::strtoul(argv[4], nullptr, 10)), argv[5], std::atoi(argv[6])); + } + if ((argc == 5 || argc == 7) && std::string(argv[1]) == "--flash-compare") + return flash_compare(std::atoi(argv[2]), uint32_t(std::strtoul(argv[3], nullptr, 10)), argv[4], + argc == 7 ? std::atoi(argv[5]) : 1, argc == 7 ? std::atoi(argv[6]) : 3); + const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; + std::setvbuf(stdout, nullptr, _IOLBF, 0); + const bool noflash = std::getenv("QSA_NOFLASH") != nullptr; + mtl::Context cx; + const int lib = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + std::printf("device %s\n", cx.device_name().c_str()); + const int T = bench ? 2048 : 96, MAXC = 40000; + mtl::Buffer kq = cx.alloc((size_t)MAXC * QKV * QD), vq = cx.alloc((size_t)MAXC * QKV * QD), ksc = cx.alloc(2 * (size_t)MAXC * QKV * 4), + vsc = cx.alloc(2 * (size_t)MAXC * QKV * 4), qcur = cx.alloc(4 * (size_t)T * QH * QD), qh = cx.alloc(2 * (size_t)T * QH * QD), + qfull = cx.alloc(4 * (size_t)T * QH * 2 * QD), ids = cx.alloc(4 * (size_t)T * QB), + part = cx.alloc(4 * (size_t)T * NCH * QH * (QD + 2)), xa = cx.alloc(4 * ((size_t)4 * T * QH * QD + (size_t)T * QH * QD / 16)), + xb = cx.alloc(4 * ((size_t)4 * T * QH * QD + (size_t)T * QH * QD / 16)); + ref::Rng r(5); + for (size_t i = 0; i < kq.bytes; ++i) { kq.as()[i] = (int8_t)((int)(r.next() % 255) - 127); vq.as()[i] = (int8_t)((int)(r.next() % 255) - 127); } + for (size_t i = 0; i < ksc.bytes / 2; ++i) { ksc.as()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); vsc.as()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); } + for (size_t i = 0; i < (size_t)T * QH * QD; ++i) qcur.as()[i] = r.sym() * 3.0f; + for (size_t i = 0; i < (size_t)T * QH * 2 * QD; ++i) qfull.as()[i] = r.sym(); + mtl::Pipeline pa = cx.pipeline(lib, "qsa_attn"), pm = cx.pipeline(lib, "qsa_merge"), pmm = cx.pipeline(lib, "qsa_attn_mm", mtl::Constants().b(7, coop_k4(cx)).i(8, std::getenv("QSA_STOP") ? std::atoi(std::getenv("QSA_STOP")) : 0)), + pq = cx.pipeline(lib, "qsa_q_f16"), pf = cx.pipeline(lib, "qsa_flash", mtl::Constants().b(7, coop_k4(cx))); + mtl::Buffer xf = cx.alloc(2 * (size_t)T * QH * QD); + std::printf("pipelines: qsa_attn_mm max threads %d, tg mem %d; qsa_flash max threads %d, tg mem %d\n", pmm.max_threads, + pmm.static_tg_mem, pf.max_threads, pf.static_tg_mem); + int fails = 0; + for (uint32_t pos0 : {0u, 1984u, 2050u, 30000u}) { + // ids: within the budget the identity; past it 2051 distinct ascending cells < n_kv + for (int t = 0; t < T; ++t) { + const uint32_t n_kv = pos0 + t + 1; + int* o = ids.as() + (size_t)t * QB; + if (n_kv <= (uint32_t)QB) { for (uint32_t c = 0; c < n_kv; ++c) o[c] = (int)c; continue; } + std::set s; + while ((int)s.size() < QB) s.insert((int)(r.next() % n_kv)); + int j = 0; + for (int c : s) o[j++] = c; + } + QsaArgs a{pos0, 0, (uint32_t)MAXC, (uint32_t)T, 1e-6f}; + auto D = [&](const mtl::Pipeline& p, std::vector args, uint32_t gx, uint32_t gy, uint32_t gz, uint32_t tg) { + mtl::Dispatch d; d.p = &p; d.args = std::move(args); d.grid[0] = gx; d.grid[1] = gy; d.grid[2] = gz; d.tg[0] = tg; return d; + }; + const uint32_t nq = T * QH * QD; + std::vector ref_ds = { + D(pa, {mtl::Arg::B(qcur), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), mtl::Arg::B(ids), mtl::Arg::B(part), mtl::Arg::V(a)}, NCH, QKV, T, 256), + D(pm, {mtl::Arg::B(part), mtl::Arg::B(qfull), mtl::Arg::B(xa), mtl::Arg::V(a)}, QH, T, 1, 256)}; + std::vector mm_ds = { + D(pq, {mtl::Arg::B(qcur), mtl::Arg::B(qh), mtl::Arg::V(nq)}, (nq + 255) / 256, 1, 1, 256), + D(pmm, {mtl::Arg::B(qh), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), mtl::Arg::B(ids), mtl::Arg::B(part), mtl::Arg::V(a)}, NCH, QKV, T, 32), + D(pm, {mtl::Arg::B(part), mtl::Arg::B(qfull), mtl::Arg::B(xb), mtl::Arg::V(a)}, QH, T, 1, 256)}; + const mtl::Dispatch fl = D(pf, {mtl::Arg::B(qh), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), + mtl::Arg::B(ids), mtl::Arg::B(qfull), mtl::Arg::B(xf), mtl::Arg::V(a)}, QKV, T, 1, 128); + if (bench) { + auto timeit = [&](const std::vector& ds) { + cx.warm(ds, 0.3); + std::vector t; + for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds)); + return median(t); + }; + std::printf("pos0 %5u, T %d: qsa_attn %.2f ms, qsa_attn_mm %.2f ms, merge %.2f ms; qsa_flash (no merge) %.2f ms\n", pos0, + T, timeit({ref_ds[0]}) * 1e3, timeit({mm_ds[1]}) * 1e3, timeit({ref_ds[1]}) * 1e3, timeit({fl}) * 1e3); + continue; + } + cx.run(ref_ds); + std::fill(part.as(), part.as() + part.bytes / 4, std::nanf("")); // a kernel that writes nothing shows + for (int t = 0; t < T && pos0 + t + 1 <= QB; ++t) + std::fill_n(ids.as() + (size_t)t * QB, QB, -1); + cx.run(mm_ds); + if (std::getenv("QSA_DEBUG")) { + // how much of the partials the matrix-unit kernel wrote, and token 5's first chunk (m, l, acc[0..3]) + size_t written = 0, nonfinite = 0; + const float* pp = part.as(); + for (size_t i = 0; i < (size_t)T * NCH * QH * (QD + 2); ++i) { + if (!std::isnan(pp[i])) ++written; + if (!std::isfinite(pp[i]) && !std::isnan(pp[i])) ++nonfinite; + } + const float* q = pp + ((size_t)5 * NCH * QH) * (QD + 2); + std::printf("debug pos0 %u: %zu of %zu partial floats written (%zu inf); t5 h0: m %g l %g acc %g %g %g %g\n", pos0, + written, (size_t)T * NCH * QH * (QD + 2), nonfinite, q[QD], q[QD + 1], q[0], q[1], q[2], q[3]); + } + double se = 0, sr = 0; + for (size_t i = 0; i < (size_t)T * QH * QD; ++i) { + const double e = xb.as()[i] - xa.as()[i]; + se += e * e; + sr += (double)xa.as()[i] * xa.as()[i]; + } + const double rel = std::sqrt(se / sr); + const bool ok = rel < 3e-3; + fails += !ok; + std::printf("pos0 %5u, T %d: qsa_attn_mm vs qsa_attn (after merge) rel L2 %.2e %s\n", pos0, T, rel, ok ? "PASS" : "FAIL"); + if (noflash) continue; + cx.run({fl}); + double sf = 0; + for (size_t i = 0; i < (size_t)T * QH * QD; ++i) { + const double e = ref::f16_to_f32(xf.as()[i]) - xa.as()[i]; + sf += e * e; + } + const double relf = std::sqrt(sf / sr); + const bool okf = relf < 3e-3; + fails += !okf; + std::printf("pos0 %5u, T %d: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 %.2e %s\n", pos0, T, relf, okf ? "PASS" : "FAIL"); + } + // qsa_score_mm vs qsa_score at pos0 30000: the scores of every block <= n_bid + { + const uint32_t pos0 = 30000, MB = MAXC / 4 + 2, nb = (pos0 + T) / 4 + 1; + mtl::Buffer qidx = cx.alloc(4 * (size_t)T * 4 * 128), pooled = cx.alloc(4 * (size_t)MB * 128), dead = cx.alloc(4 * 128), + sa = cx.alloc(4 * (size_t)T * MB), sb = cx.alloc(4 * (size_t)T * MB); + for (size_t i = 0; i < (size_t)T * 4 * 128; ++i) qidx.as()[i] = r.sym(); + for (size_t i = 0; i < (size_t)MB * 128; ++i) pooled.as()[i] = r.sym(); + for (int i = 0; i < 128; ++i) dead.as()[i] = r.sym(); + QsaArgs a{pos0, 0, (uint32_t)MAXC, (uint32_t)T, 1e-6f}; + mtl::Pipeline ps = cx.pipeline(lib, "qsa_score"), psm = cx.pipeline(lib, "qsa_score_mm"); + mtl::Dispatch d1; d1.p = &ps; d1.args = {mtl::Arg::B(qidx), mtl::Arg::B(pooled), mtl::Arg::B(dead), mtl::Arg::B(sa), mtl::Arg::V(a)}; + d1.grid[0] = (nb + 7) / 8; d1.grid[1] = T; d1.tg[0] = 256; + mtl::Dispatch d2; d2.p = &psm; d2.args = {mtl::Arg::B(qidx), mtl::Arg::B(pooled), mtl::Arg::B(dead), mtl::Arg::B(sb), mtl::Arg::V(a)}; + d2.grid[0] = (nb + 63) / 64; d2.grid[1] = (T + 15) / 16; d2.tg[0] = 128; + if (bench) { + auto timeit = [&](const mtl::Dispatch& d) { + cx.warm({d}, 0.3); + std::vector t; + for (int i = 0; i < 5; ++i) t.push_back(cx.run({d})); + return median(t); + }; + std::printf("score pos0 %u, T %d: qsa_score %.2f ms, qsa_score_mm %.2f ms\n", pos0, T, timeit(d1) * 1e3, timeit(d2) * 1e3); + } else { + cx.run({d1}); + cx.run({d2}); + double worst = 0; + for (int t = 0; t < T; ++t) { + const uint32_t nbid = (pos0 + t + 1) / 4; + for (uint32_t b = 0; b <= nbid; ++b) { + const double x = sa.as()[(size_t)t * MB + b], y = sb.as()[(size_t)t * MB + b]; + worst = std::max(worst, std::fabs(x - y) / (std::fabs(x) + 1.0)); // scores are O(10) + } + } + const bool ok = worst < 1e-5; + fails += !ok; + std::printf("score pos0 %u, T %d: qsa_score_mm vs qsa_score worst rel %.2e %s\n", pos0, T, worst, ok ? "PASS" : "FAIL"); + } + } + // qsa_prep with QSA_Q_F16_FC=true vs qsa_prep (fp32) + qsa_q_f16: 100% bitwise parity on qcurh and qsa_flash output + { + const uint32_t pos0 = 0, MB = MAXC / 4 + 2; + mtl::Buffer kc = cx.alloc(4 * (size_t)T * QKV * QD), vc = cx.alloc(4 * (size_t)T * QKV * QD), + iq = cx.alloc(4 * (size_t)T * 4 * 128), ik = cx.alloc(4 * (size_t)T * 128), + qn = cx.alloc(4 * QD), kn = cx.alloc(4 * QD), iqn = cx.alloc(4 * 128), ikn = cx.alloc(4 * 128), + cs = cx.alloc(8 * (size_t)(T + 16) * 32), qidx = cx.alloc(4 * (size_t)T * 4 * 128), + tail = cx.alloc(4 * 4 * 128), dead = cx.alloc(4 * 128), pooled = cx.alloc(4 * (size_t)MB * 128), + mtab = cx.alloc(4 * 3 * (size_t)(T + 16)), + qh_ref = cx.alloc(2 * (size_t)T * QH * QD), qh_fused = cx.alloc(2 * (size_t)T * QH * QD), + xf_ref = cx.alloc(2 * (size_t)T * QH * QD), xf_fused = cx.alloc(2 * (size_t)T * QH * QD); + for (size_t i = 0; i < kc.bytes / 4; ++i) { kc.as()[i] = r.sym(); vc.as()[i] = r.sym(); } + for (size_t i = 0; i < iq.bytes / 4; ++i) iq.as()[i] = r.sym(); + for (size_t i = 0; i < ik.bytes / 4; ++i) ik.as()[i] = r.sym(); + for (int i = 0; i < QD; ++i) { qn.as()[i] = 1.0f + 0.2f * r.sym(); kn.as()[i] = 1.0f + 0.2f * r.sym(); } + for (int i = 0; i < 128; ++i) { iqn.as()[i] = 1.0f + 0.2f * r.sym(); ikn.as()[i] = 1.0f + 0.2f * r.sym(); } + for (int p = 0; p < T + 16; ++p) + for (int i = 0; i < 32; ++i) { + const double inv = std::pow(1e7, -2.0 * i / 64.0), th = (double)p * inv; + cs.as()[(p * 32 + i) * 2] = (float)std::cos(th); + cs.as()[(p * 32 + i) * 2 + 1] = (float)std::sin(th); + } + QsaArgs a{pos0, 0, (uint32_t)MAXC, (uint32_t)T, 1e-6f}; + const uint32_t nq = T * QH * QD; + mtl::Pipeline pprep_f32 = cx.pipeline(lib, "qsa_prep"), + pprep_f16 = cx.pipeline(lib, "qsa_prep", mtl::Constants().b(25, true)); + auto make_prep = [&](const mtl::Pipeline& pp, const mtl::Buffer& qout) { + mtl::Dispatch d; + d.p = &pp; + d.args = {mtl::Arg::B(qfull), mtl::Arg::B(kc), mtl::Arg::B(vc), mtl::Arg::B(iq), mtl::Arg::B(ik), + mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(iqn), mtl::Arg::B(ikn), mtl::Arg::B(cs), + mtl::Arg::B(qout), mtl::Arg::B(qidx), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), + mtl::Arg::B(vsc), mtl::Arg::B(tail), mtl::Arg::B(dead), mtl::Arg::B(pooled), mtl::Arg::V(a), + mtl::Arg::B(mtab)}; + d.grid[0] = 33; d.grid[1] = T; d.tg[0] = 256; + return d; + }; + mtl::Dispatch d_prep_f32 = make_prep(pprep_f32, qcur); + mtl::Dispatch d_qf16; d_qf16.p = &pq; d_qf16.args = {mtl::Arg::B(qcur), mtl::Arg::B(qh_ref), mtl::Arg::V(nq)}; + d_qf16.grid[0] = (nq + 255) / 256; d_qf16.tg[0] = 256; + mtl::Dispatch d_prep_f16 = make_prep(pprep_f16, qh_fused); + for (int t = 0; t < T; ++t) { + int* o = ids.as() + (size_t)t * QB; + for (int c = 0; c <= t && c < QB; ++c) o[c] = c; + } + mtl::Dispatch d_fl_ref; d_fl_ref.p = &pf; + d_fl_ref.args = {mtl::Arg::B(qh_ref), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), + mtl::Arg::B(ids), mtl::Arg::B(qfull), mtl::Arg::B(xf_ref), mtl::Arg::V(a)}; + d_fl_ref.grid[0] = QKV; d_fl_ref.grid[1] = T; d_fl_ref.tg[0] = 128; + mtl::Dispatch d_fl_fused = d_fl_ref; + d_fl_fused.args[0] = mtl::Arg::B(qh_fused); + d_fl_fused.args[7] = mtl::Arg::B(xf_fused); + if (bench) { + auto timeit = [&](const std::vector& ds) { + cx.warm(ds, 0.3); + std::vector t; + for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds, 1, false, true)); + return median(t); + }; + std::printf("prep T %d: qsa_prep(f32)+qsa_q_f16 %.3f ms (prep %.3f ms, q_f16 %.3f ms), qsa_prep(qf16) %.3f ms\n", + T, timeit({d_prep_f32, d_qf16}) * 1e3, timeit({d_prep_f32}) * 1e3, timeit({d_qf16}) * 1e3, timeit({d_prep_f16}) * 1e3); + } else { + cx.run({d_prep_f32, d_qf16, d_fl_ref}, 1, false, true); + cx.run({d_prep_f16, d_fl_fused}, 1, false, true); + const bool ok_qh = std::memcmp(qh_ref.ptr, qh_fused.ptr, qh_ref.bytes) == 0; + const bool ok_xf = std::memcmp(xf_ref.ptr, xf_fused.ptr, xf_ref.bytes) == 0; + fails += (!ok_qh || !ok_xf); + std::printf("qsa_prep(QSA_Q_F16_FC) vs qsa_prep+qsa_q_f16: qcurh memcmp %s, qsa_flash memcmp %s\n", + ok_qh ? "PASS" : "FAIL", ok_xf ? "PASS" : "FAIL"); + } + } + if (!bench) std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; +} diff --git a/results/round39/astra_r39_qsa/final/metal/kernels/qsa.metal b/results/round39/astra_r39_qsa/final/metal/kernels/qsa.metal new file mode 100644 index 0000000..63259b4 --- /dev/null +++ b/results/round39/astra_r39_qsa/final/metal/kernels/qsa.metal @@ -0,0 +1,492 @@ +// metal/kernels/qsa.metal - QSA (sparse full attention) decode for a window of tokens, between the projections. +// Contract of include/strata/kernels/qsa.hpp / qsa_select.hpp / qsa_decode_attn.hpp / kv_q8.hpp (the engine's +// native path), in a Metal layout (no paging: unified memory): +// +// qsa_prep per token t (cell pos0 + t): q heads rms_norm * q_norm, NEOX RoPE on 64 dims -> qcur; +// k rms_norm * k_norm, RoPE -> int8 KV append (fp16 scale per 64 = fp16(amax/127), rint, clamp); +// v -> int8 append; indexer q rms_norm * iq_norm, RoPE -> qidx. The indexer KEYS (last +// threadgroup, all window tokens in order): raw tail [cell % 4]; on a block completion +// pooled[b] = rope(rms_norm(mean of the block's 4 raw keys) * ik_norm, pos of the block's first cell); +// cell 0 also sets dead = rms_norm(raw) * ik_norm (the tail block's key). +// qsa_score per token, block b <= n_bid: sum over the 4 indexer heads of relu(qidx . key_b); key_b = pooled[b], +// the tail block n_bid scores `dead` and gets +1e9 when it has cells. +// qsa_topk per token: width = min(n_kv, 2051) cells; n_kv <= 2051: all; else the width cells of largest +// score (a cell inherits its block's), ties to the lowest cell, emitted ascending (radix select). +// qsa_attn split-K: (chunk of CH cells, kv head, token): 12 query heads; s = q.k / 16, chunk max m, sum l, +// acc = p . v -> partials +// qsa_merge per (head, token): log-sum-exp merge, * sigmoid(gate), written as prepared activations (xprep +// planes) for attn_output. +// +// RoPE: cos/sin from a host table [pos][32] (double precision, as build_rope_table). + +#define QH 24 +#define QKV 2 +#define QD 256 +#define QROT 64 +#define IH 4 +#define ID 128 +#define QBUDGET 2051 +#define QCH 64 + +constant int QSA_NTOK [[function_constant(3)]]; +// QSA_MROPE (function constant 24): interleaved M-RoPE (the vision path; strata/kernels/mrope.hpp, ggml IMROPE, +// sections 11/11/10): buffer 20 of qsa_prep is a table int32 [cells][3] (t, h, w), and rotary pair i (0..31) of the row +// at (pos_base + cell) rotates by the cos/sin table's row tab[(pos_base + cell) * 3 + i % 3] instead of its own - +// the pooled indexer keys too (position pos_base + 4 b of a block). Off: exactly the plain RoPE. +constant bool QSA_MROPE [[function_constant(24)]]; +constant bool qsa_mrope = is_function_constant_defined(QSA_MROPE) && QSA_MROPE; +// QSA_Q_F16_FC (function constant 25): qsa_prep writes qcur directly as fp16 (device half*) for qsa_flash, skipping qsa_q_f16. +constant bool QSA_Q_F16_FC [[function_constant(25)]]; +constant bool qsa_qf16 = is_function_constant_defined(QSA_Q_F16_FC) ? QSA_Q_F16_FC : false; +// QSA_ROW0 (function constant 28): qsa_prep's threadgroup x is row QSA_ROW0 + x (24: the K/V rows only, for rows whose +// queries nobody reads). Undefined: 0. +constant int QSA_ROW0 [[function_constant(28)]]; +constant int qsa_row0 = is_function_constant_defined(QSA_ROW0) ? QSA_ROW0 : 0; + +struct QsaArgs { + uint pos0; // the window's first cell (relative to the sequence) + uint pos_base; // absolute position of cell 0 (RoPE) + uint max_cells; + uint n_tok; + float eps; +}; + +// cache layout per layer +// kq/vq [cell][kv head][256] int8, ks/vs [cell][kv head][4] half +// tail [4][128] f32 (slot = cell % 4), dead [128], pooled [max_cells / 4 + 1][128] + +static inline float tg_sum256(float v, threadgroup float* part, ushort sg, ushort lane, ushort nsg) { + v = simd_sum(v); + if (lane == 0) part[sg] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + float t = 0.0f; + for (int i = 0; i < nsg; ++i) t += part[i]; + return t; +} + +// one row (width w <= 256 = threads) : out = rope((rs * x) * g) (RoPE on the first 64 dims, NEOX pairs i, i+32) +static inline float norm_rope(float x, float xhi_unused, device const float* g, threadgroup float* row, + threadgroup float* part, ushort tid, ushort sg, ushort lane, int w, float eps, + device const float2* cs, uint absp, bool rope, device const int* mtab, + bool reuse = false) { + const float ss = tg_sum256(tid < w ? x * x : 0.0f, part, sg, lane, 8); + const float rs = rsqrt(ss / float(w) + eps); + const float y = tid < w ? (rs * x) * g[tid] : 0.0f; + float o = y; + if (rope) { + if (tid < QROT) row[tid] = y; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < QROT) { + const int p = tid & 31; + const uint pp = qsa_mrope ? uint(mtab[absp * 3 + uint(p % 3)]) : absp; + const float2 c = cs[pp * 32 + p]; + const float a = row[p], b = row[p + 32]; + o = tid < 32 ? a * c.x - b * c.y : a * c.y + b * c.x; + } + } + if (reuse) threadgroup_barrier(mem_flags::mem_threadgroup); + (void)xhi_unused; + return o; +} + +// QSA_KV_STRIDE (function constant 26): kcur, vcur, iqraw and ikraw are column slices of one [n_tok][QSA_KV_STRIDE] +// projection (prompt_chunk's fused K | V | indexer K | indexer Q GEMM, 1152 or 1664 columns). Undefined: their own +// [n_tok][rows] buffers. +constant int QSA_KV_STRIDE [[function_constant(26)]]; +constant int qsa_kv_s = is_function_constant_defined(QSA_KV_STRIDE) ? QSA_KV_STRIDE : 0; + +// grid (32 rows + 1, n_tok): rows 0..23 q heads, 24..25 k heads, 26..27 v heads, 28..31 indexer q heads; +// threadgroup 32 (x = 32, y = 0) is the indexer-key threadgroup. 256 threads. A grid of 28 (or of 4 from QSA_ROW0 = +// 24) writes no indexer state: the MTP layer's attention is a dense window, never scored. +kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float* kcur [[buffer(1)]], + device const float* vcur [[buffer(2)]], device const float* iqraw [[buffer(3)]], + device const float* ikraw [[buffer(4)]], device const float* qn [[buffer(5)]], + device const float* kn [[buffer(6)]], device const float* iqn [[buffer(7)]], + device const float* ikn [[buffer(8)]], device const float2* cs [[buffer(9)]], + device float* qcur [[buffer(10)]], device float* qidx [[buffer(11)]], + device char* kq [[buffer(12)]], device char* vq [[buffer(13)]], device half* ksc [[buffer(14)]], + device half* vsc [[buffer(15)]], device float* tail [[buffer(16)]], device float* dead [[buffer(17)]], + device float* pooled [[buffer(18)]], constant QsaArgs& a [[buffer(19)]], + device const int* mtab [[buffer(20)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + threadgroup float row[QROT]; + threadgroup float amax[8]; + const int r = int(tg.x) + qsa_row0; + // the row strides of kcur / vcur, ikraw, iqraw + const int s_kv = qsa_kv_s ? qsa_kv_s : QKV * QD, s_ik = qsa_kv_s ? qsa_kv_s : ID; + const int s_iq = qsa_kv_s ? qsa_kv_s : IH * ID; + if (r == 32) { + // the indexer keys: threadgroup (32, 0) handles t < 3 (boundary with prior tail + cell 0 dead) and final tail; + // every threadgroup (32, t >= 3) with (cell & 3) == 3 pools its 4 in-window keys in parallel. + const int t_self = int(tg.y); + if (t_self >= 3) { + const uint cell = a.pos0 + uint(t_self); + if ((cell & 3) == 3) { + float m = 0.0f; + if (tid < ID) { + m = (((ikraw[(t_self - 3) * s_ik + tid] + ikraw[(t_self - 2) * s_ik + tid]) + + ikraw[(t_self - 1) * s_ik + tid]) + ikraw[t_self * s_ik + tid]) * 0.25f; + } + const uint b = cell >> 2; + const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab); + if (tid < ID) pooled[b * ID + tid] = o; + } + return; + } + if (t_self != 0) return; + const int n_pre = min(int(a.n_tok), 3); + for (int t = 0; t < n_pre; ++t) { + const uint cell = a.pos0 + uint(t); + const float x = tid < ID ? ikraw[t * s_ik + tid] : 0.0f; + if (tid < ID) tail[(cell & 3) * ID + tid] = x; + threadgroup_barrier(mem_flags::mem_device); + if (cell == 0) { + const float o = norm_rope(x, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, 0, false, mtab, true); + if (tid < ID) dead[tid] = o; + } + if ((cell & 3) == 3) { + float m = 0.0f; + if (tid < ID) m = (((tail[tid] + tail[ID + tid]) + tail[2 * ID + tid]) + tail[3 * ID + tid]) * 0.25f; + const uint b = cell >> 2; + const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab, true); + if (tid < ID) pooled[b * ID + tid] = o; + } + threadgroup_barrier(mem_flags::mem_device); + } + const int t_tail0 = max(n_pre, int(a.n_tok) - 4); + for (int t = t_tail0; t < int(a.n_tok); ++t) { + const uint cell = a.pos0 + uint(t); + if (tid < ID) tail[(cell & 3) * ID + tid] = ikraw[t * s_ik + tid]; + } + return; + } + const int t = int(tg.y); + const uint cell = a.pos0 + uint(t), absp = a.pos_base + cell; + if (r < QH) { + const float x = qfull[(t * QH + r) * 2 * QD + tid]; + const float o = norm_rope(x, 0.0f, qn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); + if (qsa_qf16) ((device half*)qcur)[(t * QH + r) * QD + tid] = half(o); + else qcur[(t * QH + r) * QD + tid] = o; + } else if (r < QH + 2 * QKV) { + const bool isv = r >= QH + QKV; + const int h = r - QH - (isv ? QKV : 0); + float x = (isv ? vcur : kcur)[t * s_kv + h * QD + tid]; + if (!isv) x = norm_rope(x, 0.0f, kn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); + // int8 with an fp16 scale per 64 values: groups = simdgroup pairs (64 threads) + float m = simd_max(fabs(x)); + if (lane == 0) amax[sg] = m; + threadgroup_barrier(mem_flags::mem_threadgroup); + const int g = tid >> 6; + m = max(amax[2 * g], amax[2 * g + 1]); + const half sh = half(m / 127.0f); + const float sf = float(sh); + int q = 0; + if (sf > 0.0f) q = clamp(int(rint(x / sf)), -127, 127); + const uint rowi = cell * QKV + uint(h); + (isv ? vq : kq)[rowi * QD + tid] = char(q); + if ((tid & 63) == 0) (isv ? vsc : ksc)[rowi * 4 + g] = sh; + } else { + if (a.pos0 + a.n_tok <= uint(QBUDGET)) return; + const int h = r - QH - 2 * QKV; + const float x = tid < ID ? iqraw[t * s_iq + 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); + if (tid < ID) qidx[(t * IH + h) * ID + tid] = o; + } +} + +// scores [n_tok][max_blocks]; grid (ceil(max_blocks / 8), n_tok), 256 threads: a simdgroup per block. +kernel void qsa_score(device const float* qidx [[buffer(0)]], device const float* pooled [[buffer(1)]], + device const float* dead [[buffer(2)]], device float* scores [[buffer(3)]], + constant QsaArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + const int t = int(tg.y); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + const uint b = tg.x * 8 + sg; + if (b > n_bid) return; + device const float* key = b == n_bid ? dead : pooled + b * ID; + device const float* qp = qidx + (t * IH) * ID + lane; + float k[4], q[IH][4]; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) k[i] = key[lane + 32 * i]; + #pragma clang loop unroll(full) + for (int h = 0; h < IH; ++h) { + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) q[h][i] = qp[h * ID + 32 * i]; + } + float s = 0.0f; + #pragma clang loop unroll(full) + for (int h = 0; h < IH; ++h) { + float d = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) d = fma(q[h][i], k[i], d); + d = simd_sum(d); + s += max(d, 0.0f); + } + if (b == n_bid && (n_kv & 3) != 0) s += 1e9f; + if (lane == 0) scores[uint(t) * (a.max_cells / 4 + 2) + b] = s; +} + +// ids [n_tok][QBUDGET], ascending. One threadgroup of 1024 threads per token. Order key of a float: larger score = +// larger key; a radix select (4 passes of 8 bits) over the blocks with weights = their cell counts finds the +// threshold key T and how many cells with key == T to take (the lowest cells first). +static inline uint okey(float f) { + const uint u = as_type(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]], + constant QsaArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint hist[256]; + threadgroup uint sel[4]; + threadgroup uint2 pref[32]; + threadgroup uint2 totals; + const int t = int(tg.x); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2, nb = n_bid + 1; + device int* out = ids + t * QBUDGET; + device const float* sc = scores + uint(t) * (a.max_cells / 4 + 2); + if (n_kv <= QBUDGET) { + for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); + return; + } + const bool cached = nb <= 4096; + uint keys[4]; + if (cached) + for (uint b = tid, i = 0; b < nb; b += 1024, ++i) keys[i] = okey(sc[b]); + // radix select of the `need`-th largest cell key + uint prefix = 0, pmask = 0, need = QBUDGET; + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint b = tid; b < nb; b += 1024) { + const uint k = cached ? keys[b / 1024] : okey(sc[b]); + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + uint counts[8], sum = 0; + for (uint j = 0; j < 8; ++j) { + counts[j] = atomic_load_explicit(&hist[255 - (lane * 8 + j)], memory_order_relaxed); + sum += counts[j]; + } + uint acc = simd_prefix_exclusive_sum(sum); + if (acc < need && acc + sum >= need) { + for (uint j = 0; j < 8; ++j) { + if (acc + counts[j] >= need) { + sel[0] = 255 - (lane * 8 + j); + sel[1] = need - acc; + break; + } + acc += counts[j]; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; + } + // T = prefix; take every cell with key > T and the first `need` cells (lowest index) with key == T, ascending. + // Pass over blocks in order of index with a prefix scan of each block's selected cell count. + const uint T = prefix; + uint base = 0, eq_taken = 0; + for (uint b0 = 0; b0 < nb; b0 += 1024) { + const uint b = b0 + tid; + uint w = 0, weq = 0; + if (b < nb) { + const uint k = cached ? keys[b / 1024] : okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + if (k > T) w = cells; + else if (k == T) weq = cells; + } + const uint eq_before = simd_prefix_exclusive_sum(weq); + const uint gt_before = simd_prefix_exclusive_sum(w); + if (lane == 31) pref[sg] = uint2(eq_before + weq, gt_before + w); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + const uint2 count = pref[lane]; + const uint2 before = uint2(simd_prefix_exclusive_sum(count.x), + simd_prefix_exclusive_sum(count.y)); + pref[lane] = before; + if (lane == 31) totals = before + count; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint2 group_before = pref[sg], count = totals; + simdgroup_barrier(mem_flags::mem_threadgroup); + const uint eq0 = eq_taken + group_before.x + eq_before; + const uint take_eq = eq0 < need ? min(weq, need - eq0) : 0u; + const uint mine = w + take_eq; + // Clipping the equal-key prefix preserves the lowest-cell tie break. + const uint remaining = need - min(need, eq_taken); + const uint off = base + group_before.y + gt_before + + min(remaining, group_before.x + eq_before); + for (uint c = 0; c < mine; ++c) out[off + c] = int(4 * b + c); + base += count.y + min(remaining, count.x); + eq_taken += count.x; + } +} + +// partials: [n_tok][chunks][QH][QD + 2] (acc, then m, l). grid (chunks, QKV, n_tok), 256 threads. +kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + 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 HPG = QH / QKV; // 12 query heads per kv head + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH]; + threadgroup int cid[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; + const int nchunks = (QBUDGET + QCH - 1) / QCH; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + if (uint(c0) >= width) { + // an empty chunk: m = -inf, l = 0 (the merge skips it) + for (int h = 0; h < HPG; ++h) + if (tid == 0) { pp[(g * HPG + h) * (QD + 2) + QD] = -INFINITY; pp[(g * HPG + h) * (QD + 2) + QD + 1] = 0.0f; } + return; + } + const int nc = min(QCH, int(width) - c0); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + g * HPG + h) * QD + tid]; + if (tid < QCH) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + // scores: 12 x 64 dots of 256; simdgroup sg takes cells sg, sg + 8, ...; lanes split the 256 dims (8 each) + threadgroup const float* q_lane = &qs[0][lane * 8]; + for (int c = sg; c < nc; c += 8) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char* kr = kq + rowi * QD + lane * 8; + const char4 k0 = *(device const char4*)(kr); + const char4 k1 = *(device const char4*)(kr + 4); + const float scl = float(ksc[rowi * 4 + (lane >> 3)]); + const float kv[8] = { + float(k0.x) * scl, float(k0.y) * scl, float(k0.z) * scl, float(k0.w) * scl, + float(k1.x) * scl, float(k1.y) * scl, float(k1.z) * scl, float(k1.w) * scl + }; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + threadgroup const float* qh = q_lane + h * QD; + float d = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) d = fma(qh[i], kv[i], d); + d = simd_sum(d); + if (lane == 0) s[h][c] = d * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // per head: max, exp, sum (simdgroup h % 8 handles heads h, h + 8) + threadgroup float mh[HPG], lh[HPG]; + // nc <= QCH = 64: a lane holds cells lane and lane + 32 in registers (the same max and sum as a strided loop) + const bool in0 = int(lane) < nc, in1 = int(lane) + 32 < nc; + for (int h = sg; h < HPG; h += 8) { + const float s0 = in0 ? s[h][lane] : -INFINITY, s1 = in1 ? s[h][lane + 32] : -INFINITY; + const float m = simd_max(max(s0, s1)); + const float p0 = in0 ? exp(s0 - m) : 0.0f, p1 = in1 ? exp(s1 - m) : 0.0f; + if (in0) s[h][lane] = p0; + if (in1) s[h][lane + 32] = p1; + const float l = simd_sum(p0 + p1); + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // acc[h][j] = sum_c p[h][c] v[c][j]: thread = dim j + float acc[HPG]; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + // four cells' V loads in flight at once; the fma order per head stays c, c + 1, ... + auto vcell = [&](int c) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + return float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); + }; + int c = 0; + for (; c + 4 <= nc; c += 4) { + const float v0 = vcell(c), v1 = vcell(c + 1), v2 = vcell(c + 2), v3 = vcell(c + 3); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + float x = fma(s[h][c], v0, acc[h]); + x = fma(s[h][c + 1], v1, x); + x = fma(s[h][c + 2], v2, x); + acc[h] = fma(s[h][c + 3], v3, x); + } + } + for (; c < nc; ++c) { + const float v = vcell(c); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][c], v, acc[h]); + } + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (g * HPG + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} + +constant float qsa_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; +constant float qsa_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + +static inline void qsa_write_prepared(float y, device float* xp, constant QsaArgs& a, + int t, int h, ushort tid) { + // prepared planes (xprep layout) of the n = QH * QD vector + const uint n = QH * QD, plane = a.n_tok * n, d = uint(h) * QD + tid; + const int j = tid & 15; + device float* raw = xp + uint(t) * n + d; + raw[0] = y; + raw[plane] = y * qsa_q2s[j & 7]; + raw[2 * plane] = y * qsa_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? y * 0.00390625f : y; + float s0 = 0.0f, s1 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(y, ushort(k)); + s1 += simd_shuffle_down(y, ushort(k + 1)); + } + if (j == 0) xp[4 * plane + uint(t) * (n / 16) + d / 16] = s0 + s1; +} + +// grid (QH, n_tok), 256 threads: out = merge * sigmoid(gate), as prepared activations of n = QH * QD per token +kernel void qsa_merge(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], + device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const int h = int(tg.x), t = int(tg.y); + const int nchunks = (QBUDGET + QCH - 1) / QCH; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nc = int((width + QCH - 1) / QCH); + constexpr uint STRIDE = QH * (QD + 2); + device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2); + // the chunks' (m, l) and weights exp(m - M) once per threadgroup, not once per thread (nc <= 33 < 64) + threadgroup float mk[64], lk[64], ek[64], L_sh; + if (tid < uint(nc)) { + const float2 ml = *(device const float2*)(pp + tid * STRIDE + QD); + mk[tid] = ml.x; + lk[tid] = ml.y; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float M = -INFINITY; + for (int k = tid & 31; k < nc; k += 32) M = max(M, mk[k]); + M = simd_max(M); + if (tid < uint(nc)) ek[tid] = exp(mk[tid] - M); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + float L = 0.0f; + for (int k = 0; k < nc; ++k) L = fma(lk[k], ek[k], L); + L_sh = L; + } + float o = 0.0f; + device const float* q = pp + tid; + for (int k = 0; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); + threadgroup_barrier(mem_flags::mem_threadgroup); + o = o / L_sh; + const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; + const float y = o * (1.0f / (1.0f + exp(-gate))); + qsa_write_prepared(y, xp, a, t, h, tid); +} + diff --git a/results/round39/astra_r39_qsa/final/metal/kernels/qsa2.metal b/results/round39/astra_r39_qsa/final/metal/kernels/qsa2.metal new file mode 100644 index 0000000..30439a1 --- /dev/null +++ b/results/round39/astra_r39_qsa/final/metal/kernels/qsa2.metal @@ -0,0 +1,346 @@ +// metal/kernels/qsa2.metal - QSA decode v2 (after qsa.metal): the same selection and attention math, more parallel. +// +// qsa_attn2 chunks of 32 cells, the 12 query heads of a kv head in 2 groups of 6: 4x the threadgroups of +// qsa_attn (at a 1K context 128 instead of 32); V loads 8 cells at a time. +// qsa_merge2 qsa_merge for that chunking. +// qsa_topk2 the digit of each radix pass found by a parallel suffix scan of the 256 bins (simdgroup prefix +// sums) instead of one thread; cells emitted from contiguous block ranges per thread after one scan. + +#define QCH2 32 +#define QHG 2 + +// partials: [n_tok][chunks2][QH][QD + 2]. grid (chunks2, QKV * QHG, n_tok), 256 threads. +kernel void qsa_attn2(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + 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 HPG = QH / QKV / QHG; // 6 + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH2]; + threadgroup float mh[HPG], lh[HPG]; + threadgroup int cid[QCH2]; + const int chunk = int(tg.x), g = int(tg.y) / QHG, hg = int(tg.y) % QHG, t = int(tg.z); + const int h0 = g * (QH / QKV) + hg * HPG; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH2; + const int nchunks = (QBUDGET + QCH2 - 1) / QCH2; + if (uint(c0) >= width) return; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + const int nc = min(QCH2, int(width) - c0); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + h0 + h) * QD + tid]; + if (tid < QCH2) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup const float* q_lane = &qs[0][lane * 8]; + for (int c = sg; c < nc; c += 8) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char4* kr = (device const char4*)(kq + rowi * QD + lane * 8); + const char4 k0 = kr[0], k1 = kr[1]; + const float scl = float(ksc[rowi * 4 + (lane >> 3)]); + const float kv[8] = {float(k0.x) * scl, float(k0.y) * scl, float(k0.z) * scl, float(k0.w) * scl, + float(k1.x) * scl, float(k1.y) * scl, float(k1.z) * scl, float(k1.w) * scl}; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + float d = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) d = fma(q_lane[h * QD + i], kv[i], d); + d = simd_sum(d); + if (lane == 0) s[h][c] = d * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg < HPG) { + const int h = sg; + const float sv = lane < nc ? s[h][lane] : -INFINITY; + const float m = simd_max(sv); + const float p = lane < nc ? exp(sv - m) : 0.0f; + const float l = simd_sum(p); + if (lane < nc) s[h][lane] = p; + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float acc[HPG]; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + const uint v_scale_off = tid >> 6; + int cb = 0; + for (; cb + 8 <= nc; cb += 8) { + float v[8]; + #pragma clang loop unroll(full) + for (int c = 0; c < 8; ++c) { + const uint rowi = uint(cid[cb + c]) * QKV + uint(g); + v[c] = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + v_scale_off]); + } + #pragma clang loop unroll(full) + for (int c = 0; c < 8; ++c) + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb + c], v[c], acc[h]); + } + for (; cb < nc; ++cb) { + const uint rowi = uint(cid[cb]) * QKV + uint(g); + const float v = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + v_scale_off]); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb], v, acc[h]); + } + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (h0 + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} + +kernel void qsa_merge2(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], + device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float mk[128], lk[128], ek[128], L_sh; + const int h = int(tg.x), t = int(tg.y); + const int nchunks = (QBUDGET + QCH2 - 1) / QCH2; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nc = int((width + QCH2 - 1) / QCH2); + constexpr uint STRIDE = QH * (QD + 2); + device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2); + if (tid < uint(nc)) { + const float2 ml = *(device const float2*)(pp + tid * STRIDE + QD); + mk[tid] = ml.x; + lk[tid] = ml.y; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float M = -INFINITY; + for (int k = tid & 31; k < nc; k += 32) M = max(M, mk[k]); + M = simd_max(M); + if (tid < uint(nc)) ek[tid] = exp(mk[tid] - M); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + float L = 0.0f; + for (int k = 0; k < nc; ++k) L = fma(lk[k], ek[k], L); + L_sh = L; + } + float o = 0.0f; + device const float* q = pp + tid; + int k = 0; + for (; k + 3 < nc; k += 4, q += 4 * STRIDE) { + const float p0 = q[0], p1 = q[STRIDE], p2 = q[2 * STRIDE], p3 = q[3 * STRIDE]; + o = fma(p0, ek[k], o); + o = fma(p1, ek[k + 1], o); + o = fma(p2, ek[k + 2], o); + o = fma(p3, ek[k + 3], o); + } + for (; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); + threadgroup_barrier(mem_flags::mem_threadgroup); + o = o / L_sh; + const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; + const float y = o * (1.0f / (1.0f + exp(-gate))); + qsa_write_prepared(y, xp, a, t, h, tid); +} + +kernel void qsa_topk2(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]], + constant QsaArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint hist[256]; + threadgroup uint sgs[32], sgs2[32]; + threadgroup uint sel[2]; + const int t = int(tg.x); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2, nb = n_bid + 1; + device int* out = ids + t * QBUDGET; + device const float* sc = scores + uint(t) * (a.max_cells / 4 + 2); + if (n_kv <= QBUDGET) { + for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); + return; + } + if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + uint prefix = 0, pmask = 0, need = QBUDGET; + if (nb <= 4096) { + uint keys[4]; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) { + const uint b = tid + 1024u * uint(i); + keys[i] = b < nb ? okey(sc[b]) : 0u; + } + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) { + const uint b = tid + 1024u * uint(i); + if (b >= nb) break; + const uint k = keys[i]; + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint cnt = 0, ex = 0; + if (tid < 256) { + cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); + if (pass < 3) atomic_store_explicit(&hist[255 - tid], 0u, memory_order_relaxed); + ex = simd_prefix_exclusive_sum(cnt); + if (lane == 31) sgs[sg] = ex + cnt; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 256) { + uint above = ex; + for (int i = 0; i < int(sg); ++i) above += sgs[i]; + if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; + } + } else { + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + for (uint b = tid; b < nb; b += 1024) { + const uint k = okey(sc[b]); + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint cnt = 0, ex = 0; + if (tid < 256) { + cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); + if (pass < 3) atomic_store_explicit(&hist[255 - tid], 0u, memory_order_relaxed); + ex = simd_prefix_exclusive_sum(cnt); + if (lane == 31) sgs[sg] = ex + cnt; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 256) { + uint above = ex; + for (int i = 0; i < int(sg); ++i) above += sgs[i]; + if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; + } + } + const uint T = prefix; + const uint per = (nb + 1023) / 1024, b0 = min(nb, uint(tid) * per), b1 = min(nb, b0 + per); + uint gt = 0, eq = 0; + for (uint b = b0; b < b1; ++b) { + const uint k = okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + if (k > T) gt += cells; + else if (k == T) eq += cells; + } + const uint eqb = simd_prefix_exclusive_sum(eq); + if (lane == 31) sgs2[sg] = eqb + eq; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint eq_before = eqb; + for (int i = 0; i < int(sg); ++i) eq_before += sgs2[i]; + const uint eq_take = eq_before >= need ? 0u : min(eq, need - eq_before); + const uint mine = gt + eq_take; + const uint mb = simd_prefix_exclusive_sum(mine); + if (lane == 31) sgs[sg] = mb + mine; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint off = mb; + for (int i = 0; i < int(sg); ++i) off += sgs[i]; + uint left = eq_take; + for (uint b = b0; b < b1; ++b) { + const uint k = okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + uint take = 0; + if (k > T) take = cells; + else if (k == T) { take = min(cells, left); left -= take; } + for (uint c = 0; c < take; ++c) out[off + c] = int(4 * b + c); + off += take; + } +} + +// qsa_attn3: qsa_attn's chunking and partial layout (qsa_merge reads it), without a reduction per (cell, head): +// for q.k a lane owns a cell (8 simdgroups = 2 halves of the 64 cells x 4 groups of 3 heads; q from threadgroup +// memory, read by every lane alike), for p.v a thread owns a dimension. grid (33, QKV, n_tok), 256 threads. +kernel void qsa_attn3(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + 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 HPG = QH / QKV; // 12 + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH]; + threadgroup float mh[HPG], lh[HPG]; + threadgroup int cid[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; + const int nchunks = (QBUDGET + QCH - 1) / QCH; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + if (uint(c0) >= width) return; + const int nc = min(QCH, int(width) - c0); + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + g * HPG + h) * QD + tid]; + if (tid < QCH) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + { + const int half_ = sg & 1, hgrp = sg >> 1; // cells 32 * half_ + lane, heads 3 * hgrp .. + 3 + const int c = 32 * half_ + lane; + if (c < nc) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char4* kr = (device const char4*)(kq + rowi * QD); + float d0 = 0.0f, d1 = 0.0f, d2 = 0.0f; + for (int gq = 0; gq < 4; ++gq) { + const float scl = float(ksc[rowi * 4 + gq]); + float e0 = 0.0f, e1 = 0.0f, e2 = 0.0f; + for (int i = 0; i < 16; ++i) { + const char4 k4 = kr[gq * 16 + i]; + const int base = gq * 64 + 4 * i; + const float kv[4] = {float(k4.x), float(k4.y), float(k4.z), float(k4.w)}; + for (int e = 0; e < 4; ++e) { + e0 = fma(qs[3 * hgrp][base + e], kv[e], e0); + e1 = fma(qs[3 * hgrp + 1][base + e], kv[e], e1); + e2 = fma(qs[3 * hgrp + 2][base + e], kv[e], e2); + } + } + d0 = fma(e0, scl, d0); + d1 = fma(e1, scl, d1); + d2 = fma(e2, scl, d2); + } + s[3 * hgrp][c] = d0 * 0.0625f; + s[3 * hgrp + 1][c] = d1 * 0.0625f; + s[3 * hgrp + 2][c] = d2 * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int h = sg; h < HPG; h += 8) { + float m = -INFINITY; + for (int c = lane; c < nc; c += 32) m = max(m, s[h][c]); + m = simd_max(m); + float l = 0.0f; + for (int c = lane; c < nc; c += 32) { + const float p = exp(s[h][c] - m); + s[h][c] = p; + l += p; + } + l = simd_sum(l); + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float acc[HPG]; + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + for (int cb = 0; cb < nc; cb += 8) { + float v[8]; + const int m = min(8, nc - cb); + for (int c = 0; c < 8; ++c) { + v[c] = 0.0f; + if (c < m) { + const uint rowi = uint(cid[cb + c]) * QKV + uint(g); + v[c] = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); + } + } + for (int c = 0; c < m; ++c) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb + c], v[c], acc[h]); + } + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (g * HPG + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} diff --git a/results/round39/astra_r39_qsa/final/metal/kernels/qsa_prompt.metal b/results/round39/astra_r39_qsa/final/metal/kernels/qsa_prompt.metal new file mode 100644 index 0000000..2ee4a4b --- /dev/null +++ b/results/round39/astra_r39_qsa/final/metal/kernels/qsa_prompt.metal @@ -0,0 +1,461 @@ +// metal/kernels/qsa_prompt.metal - QSA attention for a PROMPT chunk on the matrix units (after qsa.metal: QH, QKV, QD, +// QBUDGET, QCH, QsaArgs). The same split-K partials as qsa_attn (part [t][chunk][QH][QD + 2]: acc, then m, l), so +// qsa_merge is unchanged; per (chunk of QCH = 64 selected cells, kv head g, token t) one simdgroup: +// S[12 x 64] = Qh[12 x 256] . K^T / 16 single-simdgroup matmul2d, K (int8 * fp16 scale per 64) decoded into +// the right-input cooperative tensor, 32 dims a step +// per head: m = max, p = exp(s - m), l = sum p (fp32, through threadgroup memory), P -> fp16 +// acc[12 x 256] = P[12 x 64] . V[64 x 256] four 64-dim matmuls, V decoded into the right input +// The operands are fp16 (q rounded from fp32; k, v = int8 * scale rounded): to rounding of qsa_attn's fp32. +#include +using namespace mpp::tensor_ops; +#ifndef STRATA_COOP_K4 +#define STRATA_COOP_K4 +// function constant 7: the right-input cooperative tensor holds runs of 4 consecutive k at one n per lane (the M5's +// layout; tests/coop_layout). false (the M4: runs of 2 n at one k): every element decoded on its own. +constant bool COOP_K4 [[function_constant(7)]]; +constant bool coop_k4 = is_function_constant_defined(COOP_K4) ? COOP_K4 : true; +#endif + +// a bisection probe (function constant 8): stop at stage n and write m = n (1: after the ids, 2: after S, 3: after +// the softmax); 0 = normal +constant int QSA_STOP [[function_constant(8)]]; +constant int qsa_stop = is_function_constant_defined(QSA_STOP) ? QSA_STOP : 0; + +// grid (chunks, QKV, n_tok), 32 threads. qh: the q heads as fp16 [t][QH][QD] (qcur rounded). +kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int HPG = QH / QKV, NCH = (QBUDGET + QCH - 1) / QCH; + threadgroup float s[16 * QCH]; + threadgroup half p[16 * QCH]; + threadgroup int cid[QCH]; // coop_k4: byte offsets; otherwise cell ids + threadgroup half4 ks[QCH], vs[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; + device float* pp = part + ((uint(t) * NCH + chunk) * QH + g * HPG) * (QD + 2); + if (uint(c0) >= width) { + if (lane < HPG) { pp[lane * (QD + 2) + QD] = -INFINITY; pp[lane * (QD + 2) + QD + 1] = 0.0f; } + return; + } + const int nc = min(QCH, int(width) - c0); + for (int c = lane; c < QCH; c += 32) { + const int id = c < nc ? (n_kv <= QBUDGET ? c0 + c : ids[t * QBUDGET + c0 + c]) : -1; + cid[c] = coop_k4 ? int((uint(max(id, 0)) * QKV + g) * QD) : id; + if (id >= 0) { + const uint row4 = (uint(id) * QKV + g) * 4; + if (coop_k4) ks[c] = *(device const half4*)(ksc + row4); + vs[c] = *(device const half4*)(vsc + row4); + } else { + ks[c] = half4(0.0h); + vs[c] = half4(0.0h); + } + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (qsa_stop == 1) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 1.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } + // ---- S = Q K^T + { + constexpr auto desc = matmul2d_descriptor(16, QCH, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tQ((device half*)qh + (uint(t) * QH + g * HPG) * QD, + dextents(QD, HPG)); + auto rk = op.template get_right_input_cooperative_tensor(); + auto tA = tQ.template slice<32, 16>(0, 0); + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + // the second half repeats these eight cells, 16 k later; empty cells read row 0 at scale 0. + uint ko[8]; + float ks8[8][4]; + if (coop_k4) +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = rk.get_multidimensional_index(uint16_t(4 * j))[1]; + ko[j] = uint(cid[c]); +#pragma clang loop unroll(full) + for (int h = 0; h < 4; ++h) ks8[j][h] = float(ks[c][h]); + } + for (int k0 = 0; k0 < QD; k0 += 32) { + if (!coop_k4) { + for (uint16_t i = 0; i < rk.get_capacity(); ++i) { + const auto ix = rk.get_multidimensional_index(i); + const int k = k0 + ix[0], id = cid[ix[1]]; + // the row is clamped and the select is on the value: a select with the device load in the arm + // loses the threadgroup on the M1 (macOS 26, measured); the discarded value cannot matter + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)]); + rk[i] = id >= 0 ? half(v) : half(0.0f); + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rk.get_capacity(); i += 4) { + const int k = k0 + rk.get_multidimensional_index(i)[0], j = (i >> 2) & 7; + const half4 w = half4(float4(*(device const char4*)(kq + (ko[j] + uint(k)))) * ks8[j][k >> 6]); + rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w; + } + tA = tQ.template slice<32, 16>(k0, 0); + op.run(tA, rk, acc); + } + tensor, tensor_inline> tS(s, dextents(QCH, 16)); + acc.store(tS); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (qsa_stop == 2) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 2.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } + // ---- softmax pieces per head (lane = head for the reductions; 12 heads) + float mh = -INFINITY, lh = 0.0f; + if (lane < HPG) { +#pragma clang loop unroll_count(4) + for (int c = 0; c < nc; ++c) mh = max(mh, s[lane * QCH + c] * 0.0625f); +#pragma clang loop unroll_count(4) + for (int c = 0; c < QCH; ++c) { + const float e = c < nc ? exp(s[lane * QCH + c] * 0.0625f - mh) : 0.0f; + lh += e; + p[lane * QCH + c] = half(e); + } + } else if (lane < 16) { +#pragma clang loop unroll_count(4) + for (int c = 0; c < QCH; ++c) p[lane * QCH + c] = half(0.0f); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (qsa_stop == 3) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 3.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } + // ---- acc = P V, 64 dims at a time + { + // K = 32 cells a run (two runs a 64-cell chunk): K = 64 runs compute garbage on the M4 + constexpr auto desc = matmul2d_descriptor(16, 64, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tP(p, dextents(QCH, 16)); + tensor, tensor_inline> tO(pp, dextents(QD + 2, HPG)); + for (int d0 = 0; d0 < QD; d0 += 64) { + auto rv = op.template get_right_input_cooperative_tensor(); + auto tPh = tP.template slice<32, 16>(0, 0); + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + for (int c0 = 0; c0 < nc; c0 += 32) { + if (!coop_k4) { +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); ++i) { + const auto ix = rv.get_multidimensional_index(i); // (cell, dim) + const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(vq[row * QD + d]) * float(vs[c][d >> 6]); + rv[i] = id >= 0 ? half(v) : half(0.0f); + } + } else { + uint vo[8]; + float vs8[8]; + const int cb = c0 + rv.get_multidimensional_index(uint16_t(0))[0]; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = cb + (j & 3) + 16 * (j >> 2); + vo[j] = uint(cid[c]) + uint(d0); + vs8[j] = float(vs[c][d0 >> 6]); + } +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); i += 4) { + const int d = rv.get_multidimensional_index(i)[1], r = i < rv.get_capacity() / 2 ? 0 : 4; +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) rv[i + j] = half(float(vq[vo[r + j] + uint(d)]) * vs8[r + j]); + } + } + tPh = tP.template slice<32, 16>(c0, 0); + op.run(tPh, rv, acc); + } + if (qsa_stop == 4) { // the bisection: the old direct store into the 258-float rows + auto mO = tO.slice(d0, 0); + acc.store(mO); + continue; + } + // through threadgroup memory (s is free now): a matmul2d store into rows of QD + 2 floats loses the + // whole threadgroup on the M4 + tensor, tensor_inline> tT(s, dextents(64, 16)); + simdgroup_barrier(mem_flags::mem_threadgroup); + acc.store(tT); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int it = lane; it < HPG * 16; it += 32) { + const int h = it >> 4, d = (it & 15) * 4; + const float4 v = *(threadgroup const float4*)(s + h * 64 + d); + // odd heads start on an 8-byte boundary, not a float4 boundary. + device float* dst = pp + h * (QD + 2) + d0 + d; + dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w; + } + } + } + if (lane < HPG) { pp[lane * (QD + 2) + QD] = mh; pp[lane * (QD + 2) + QD + 1] = lh; } +} + +// qcur (fp32) -> fp16 for qsa_attn_mm +kernel void qsa_q_f16(device const float* q [[buffer(0)]], device half* qh [[buffer(1)]], constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + if (gid < n) qh[gid] = half(q[gid]); +} + +// qsa_score as a GEMM (fp32 operands, so the selection sees qsa_score's values to fp32 rounding): S[(t, h)][b] = +// qidx[t][h] . pooled[b] for 16 tokens x IH heads (64 rows) by 64 blocks a threadgroup; the epilogue sums relu over +// the heads. The tail block n_bid (the `dead` key, + 1e9 when it has cells) and blocks past it are left to the +// scalar path: the tail is scored here by one simdgroup a token; blocks > n_bid are not written (qsa_topk reads +// b <= n_bid). grid (ceil(max blocks / 64), ceil(T / 16)), 128 threads. +kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const float* pooled [[buffer(1)]], + device const float* dead [[buffer(2)]], device float* scores [[buffer(3)]], + constant QsaArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float ot[64 * 64]; + const int b0 = int(tg.x) * 64, t0 = int(tg.y) * 16; + const uint max_t = min(uint(t0 + 15), a.n_tok - 1u); + const uint max_bid = (a.pos0 + max_t + 1u) >> 2; + // column zero still owns every token's tail score. + if (tg.x != 0 && uint(b0) >= max_bid) return; + const uint nb_all = (a.pos0 + a.n_tok) / 4 + 1; // blocks any token of the chunk can score + tensor, tensor_inline> tQ((device float*)qidx, dextents(ID, int(a.n_tok) * IH)); + tensor, tensor_inline> tK((device float*)pooled, dextents(ID, int(nb_all))); + tensor, tensor_inline> tO(ot, dextents(64, 64)); + constexpr auto desc = matmul2d_descriptor(64, 64, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tQ.slice(0, t0 * IH); + auto mB = tK.slice(0, b0); + op.run(mA, mB, tO); + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint stride = a.max_cells / 4 + 2; + const int j = tid & 63, b = b0 + j; + for (int tt = tid >> 6; tt < 16; tt += 2) { + const int t = t0 + tt; + if (t >= int(a.n_tok)) break; + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + if (uint(b) >= n_bid) continue; + float s = 0.0f; + for (int h = 0; h < IH; ++h) s += max(ot[(tt * IH + h) * 64 + j], 0.0f); + scores[uint(t) * stride + b] = s; + } + // the tail block of each token in this threadgroup's block range (the first column threadgroup does it) + if (tg.x == 0) { + float k[4]; + for (int i = 0; i < 4; ++i) k[i] = dead[lane + 32 * i]; + for (int tt = sg; tt < 16; tt += 4) { + const int t = t0 + tt; + if (t >= int(a.n_tok)) break; + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + float s = 0.0f; + for (int h = 0; h < IH; ++h) { + float d = 0.0f; + for (int i = 0; i < 4; ++i) d = fma(qidx[(t * IH + h) * ID + lane + 32 * i], k[i], d); + d = simd_sum(d); + s += max(d, 0.0f); + } + if ((n_kv & 3) != 0) s += 1e9f; + if (lane == 0) scores[uint(t) * stride + n_bid] = s; + } + } +} + +// qsa_flash: the whole attention of a (kv head g, token t) in one threadgroup (no split-K partials, no merge): the +// width selected cells in chunks of 64 with an online softmax; 4 simdgroups, simdgroup s owns dims 64 s .. 64 s + 63 +// (its slice of S's K sum and its 64 output dims). Out: xq = (acc / l) * sigmoid(gate) as fp16 (attn_output's GEMM +// operand). grid (QKV, n_tok), 128 threads. +kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device const float* qfull [[buffer(6)]], device half* xqh [[buffer(7)]], constant QsaArgs& a [[buffer(8)]], + 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 HPG = QH / QKV; + threadgroup float sp[4][HPG * QCH]; + threadgroup half* p = (threadgroup half*)sp[1]; // P [16][QCH] in fp16 over a dead S partial + threadgroup int cid[QCH]; // coop_k4: byte offsets; otherwise cell ids + threadgroup half4 ks_tg[QCH], vs_tg[QCH]; + threadgroup float fac[16], lsum[16], mrun[16]; + const int g = int(tg.x), t = int(tg.y); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nch = int((width + QCH - 1) / QCH); + if (tid < 16) { mrun[tid] = -INFINITY; lsum[tid] = 0.0f; } + constexpr auto dS = matmul2d_descriptor(16, QCH, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + constexpr auto dO = matmul2d_descriptor(16, 64, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d opS; + matmul2d opO; + tensor, tensor_inline> tQ((device half*)qh + (uint(t) * QH + g * HPG) * QD, + dextents(QD, HPG)); + tensor, tensor_inline> tP(p, dextents(QCH, 16)); + auto rv = opO.template get_right_input_cooperative_tensor(); + auto tPh = tP.template slice<32, 16>(0, 0); + auto acc = opO.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); + // (a second destination tensor, not a copy of acc: the MetalPerformancePrimitives headers of macOS 26.3 cannot copy a + // cooperative tensor - "cannot initialize a parameter of type 'void *' with an rvalue of type 'const void *'", #2323; + // tot is zeroed right below, so the copy was only there for its type) + auto tot = opO.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); // the running output; acc takes one chunk (fresh each chunk: on the M4 a run does not accumulate + // into elements rescaled in place) + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) tot[i] = 0.0f; + const int d0 = sg * 64; + if (tid < 16 - HPG) fac[HPG + tid] = 0.0f; // the padding rows' factor (P's rows are zeroed each chunk) + for (int ch = 0; ch < nch; ++ch) { + const int c0 = ch * QCH, nc = min(QCH, int(width) - c0); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < QCH) { + const int id = tid < nc ? (n_kv <= QBUDGET ? c0 + tid : ids[t * QBUDGET + c0 + tid]) : -1; + cid[tid] = coop_k4 ? int((uint(max(id, 0)) * QKV + g) * QD) : id; + if (id >= 0) { + const uint row4 = (uint(id) * QKV + g) * 4; + ks_tg[tid] = *(device const half4*)(ksc + row4); + vs_tg[tid] = *(device const half4*)(vsc + row4); + } else { + ks_tg[tid] = half4(0.0h); + vs_tg[tid] = half4(0.0h); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // Eight lanes per head; keep the denominator's original sequential sum. + const int h = sg * 4 + (lane >> 3), hlane = lane & 7; + float sc[8]; + // S partial over this simdgroup's 64 dims + { + auto rk = opS.template get_right_input_cooperative_tensor(); + auto tA = tQ.template slice<32, 16>(d0, 0); + auto s = opS.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < s.get_capacity(); ++i) + if (s.is_valid_element(i)) s[i] = 0.0f; + // the lane's eight cells (one a run of four k; the second half repeats them 16 k on): row pointers and + // scales once a chunk, unconditional loads (an empty cell reads row 0 at scale 0, and a +-0 operand sums to + // the same value as the skipped one; the branchy loads cost half the kernel) + uint ko[8]; + float ks8[8]; + if (coop_k4) { +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = rk.get_multidimensional_index(uint16_t(4 * j))[1]; + ko[j] = uint(cid[c]) + uint(d0 + rk.get_multidimensional_index(uint16_t(0))[0]); + ks8[j] = float(ks_tg[c][sg]); + } + } +#pragma clang loop unroll(full) + for (int k0 = d0; k0 < d0 + 64; k0 += 32) { + if (!coop_k4) { + for (uint16_t i = 0; i < rk.get_capacity(); ++i) { + const auto ix = rk.get_multidimensional_index(i); + const int k = k0 + ix[0], id = cid[ix[1]]; + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(kq[row * QD + k]) * float(ks_tg[ix[1]][k >> 6]); + rk[i] = id >= 0 ? half(v) : half(0.0f); + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rk.get_capacity(); i += 4) { + const int k = k0 - d0 + (i >= rk.get_capacity() / 2 ? 16 : 0), j = (i >> 2) & 7; + const half4 w = half4(float4(*(device const char4*)(kq + ko[j] + uint(k))) * ks8[j]); + rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w; + } + tA = tQ.template slice<32, 16>(k0, 0); + opS.run(tA, rk, s); + } + // Preserve ((s0 + s1) + (s2 + s3)) while storing all partials together. + tensor, tensor_inline> tS(sp[sg], dextents(QCH, HPG)); + s.store(tS); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (h < HPG) +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) { + const int q = h * QCH + hlane + 8 * k; + sc[k] = ((sp[0][q] + sp[1][q]) + (sp[2][q] + sp[3][q])) * 0.0625f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); // P overwrites sp[1] + if (h < HPG) { + float m = mrun[h]; +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) + if (hlane + 8 * k < nc) m = max(m, sc[k]); + m = max(m, simd_shuffle_xor(m, 4)); + m = max(m, simd_shuffle_xor(m, 2)); + m = max(m, simd_shuffle_xor(m, 1)); +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) { + const int c = hlane + 8 * k, q = h * QCH + c; + const float e = c < nc ? exp(sc[k] - m) : 0.0f; + p[q] = half(e); + sp[0][q] = e; + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (hlane == 0) { + const float f = exp(mrun[h] - m); + float l = lsum[h] * f; +#pragma clang loop unroll(full) + for (int c = 0; c < QCH; c += 4) { // in order, four loads at a time + const float4 v = *(threadgroup const float4*)(sp[0] + h * QCH + c); + l += v.x; l += v.y; l += v.z; l += v.w; + } + mrun[h] = m; + lsum[h] = l; + fac[h] = f; + } + } else { + for (int c = hlane; c < QCH; c += 8) p[h * QCH + c] = half(0.0f); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // rescale this simdgroup's accumulator rows, then acc += P V over its 64 dims + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + for (int c0 = 0; c0 < nc; c0 += 32) { // K = 32 cells a run (K = 64 is wrong on the M4) + if (!coop_k4) { +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); ++i) { + const auto ix = rv.get_multidimensional_index(i); // (cell, dim) + const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(vq[row * QD + d]) * float(vs_tg[c][sg]); + rv[i] = id >= 0 ? half(v) : half(0.0f); + } + } else { + // a lane holds the same eight cells (two runs of four, 16 apart) at each of its dims: as for rk, row + // pointers and scales once a run, unconditional loads + uint vo[8]; + float vs8[8]; + const int cb = c0 + rv.get_multidimensional_index(uint16_t(0))[0]; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = cb + (j & 3) + 16 * (j >> 2); + vo[j] = uint(cid[c]) + uint(d0); + vs8[j] = float(vs_tg[c][sg]); + } +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); i += 4) { + const int d = rv.get_multidimensional_index(i)[1], r = i < rv.get_capacity() / 2 ? 0 : 4; +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) rv[i + j] = half(float(vq[vo[r + j] + uint(d)]) * vs8[r + j]); + } + } + tPh = tP.template slice<32, 16>(c0, 0); + opO.run(tPh, rv, acc); + } + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) tot[i] = fma(tot[i], fac[acc.get_multidimensional_index(i)[1]], acc[i]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (!coop_k4) { + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) { + const auto ix = acc.get_multidimensional_index(i); // (dim, row) + const int h = ix[1], d = d0 + ix[0]; + if (h < HPG) { + const int hh = g * HPG + h; + const float gate = qfull[(uint(t) * QH + hh) * 2 * QD + QD + d]; + xqh[(uint(t) * QH + hh) * QD + d] = half(tot[i] / lsum[h] * (1.0f / (1.0f + exp(-gate)))); + } + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { // runs of four dims of one row + const auto ix = acc.get_multidimensional_index(i); + const int h = ix[1], d = d0 + ix[0]; + if (h < HPG) { + const int hh = g * HPG + h; + const float l = lsum[h]; + const float4 gate = *(device const float4*)(qfull + (uint(t) * QH + hh) * 2 * QD + QD + d); + const float4 v = float4(tot[i], tot[i + 1], tot[i + 2], tot[i + 3]) / l * (1.0f / (1.0f + exp(-gate))); + *(device half4*)(xqh + (uint(t) * QH + hh) * QD + d) = half4(v); + } + } +} diff --git a/results/round39/astra_r39_qsa/final/metal/tests/moe_mm_parity.cpp b/results/round39/astra_r39_qsa/final/metal/tests/moe_mm_parity.cpp new file mode 100644 index 0000000..a44397c --- /dev/null +++ b/results/round39/astra_r39_qsa/final/metal/tests/moe_mm_parity.cpp @@ -0,0 +1,506 @@ +// metal/tests/moe_mm_parity.cpp - the prompt-chunk routed experts (moe_mm.metal): plan, gate/up GEMM, down GEMM, +// combine, against a CPU reference that rounds the activations and weights to fp16 as the kernels do (the matrix +// units' operands) - and against the exact float reference for the size of that rounding; then the speed. +// +// moe_mm_parity [--gu Q2_0] [--dn Q2_0] [--T 256] [--tile 32,64,64,4] parity (64 experts) +// moe_mm_parity --gather-bench [T=2048] [experts=512] [width=2560] short plan/copy sweep +// MM_BASELINE_KERNEL=/path/to/moe_mm.metal compares every plan buffer byte-for-byte +// moe_mm_parity --bench [--gu ..] [--dn ..] [--T 2048] [--tile ..] per-layer time, 512 experts +#include "coop_probe.h" +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { +constexpr int N_EMBD = 2560, N_FF = 640, TOPK = 10; +struct MmPlanArgs { uint32_t n_ent, topk, n_expert, tm, n_blocks; }; +struct MmArgs { uint64_t gate_off, up_off, down_off, gu_stride, d_stride; uint32_t gu_row, d_row, n_embd, n_ff; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +int type_of(const std::string& s) { + for (int t : {ref::Q2_0, ref::IQ3_XXS, ref::IQ2_XXS, ref::IQ2_XS, ref::IQ2_S, ref::IQ3_S, ref::IQ4_NL, ref::IQ4_XS, ref::Q4_K, ref::Q5_K, ref::Q6_K, ref::Q8_0, ref::Q5_1}) + if (s == ref::type_name(t)) return t; + std::exit(2); +} +float h16(float v) { return ref::f16_to_f32(ref::f32_to_f16(v)); } + +int gather_bench(int T, int NE, int n) { + if (T <= 0 || NE < TOPK || NE > 1024 || NE % 32 || n <= 0 || n % 8) { + std::fprintf(stderr, "gather-bench needs T > 0, experts in [32, 1024] divisible by 32, and width divisible by 8\n"); + return 2; + } + mtl::Context cx; + const int lib = cx.compile({"moe_mm.metal"}, "#define NO_GGML_COMMON 1\n#define MM_GATHER_ONLY 1\n"); + mtl::Pipeline phist = cx.pipeline(lib, "mm_hist"), pscan = cx.pipeline(lib, "mm_scan"), pplace = cx.pipeline(lib, "mm_place"), + pgat = cx.pipeline(lib, "mm_gather_h"), pgat_f = cx.pipeline(lib, "mm_gather"); + const char* baseline = std::getenv("MM_BASELINE_KERNEL"); + const int bl = baseline ? cx.compile({baseline}, "#define NO_GGML_COMMON 1\n#define MM_GATHER_ONLY 1\n") : lib; + mtl::Pipeline bh = cx.pipeline(bl, "mm_hist"), bs = cx.pipeline(bl, "mm_scan"), bp = cx.pipeline(bl, "mm_place"), + bg = cx.pipeline(bl, "mm_gather_h"), bf = cx.pipeline(bl, "mm_gather"); + std::printf("device %s; mm_plan + mm_gather_h bitwise parity & benchmark\n", cx.device_name().c_str()); + const int TM = 16; + const int nent = T * TOPK, nblk = (nent + 1023) / 1024, maxrows = nent + NE * TM; + mtl::Buffer x = cx.alloc(4 * (size_t)T * n), xh = cx.alloc(2 * (size_t)T * n), + ids = cx.alloc(4 * nent), w = cx.alloc(4 * nent), + bcnt = cx.alloc(4 * (size_t)nblk * NE), boff = cx.alloc(4 * (size_t)nblk * NE), eoff = cx.alloc(4 * NE), + tiles = cx.alloc(8 * (size_t)(maxrows / TM + 1)), meta = cx.alloc(64), stok = cx.alloc(4 * (size_t)maxrows), + sw = cx.alloc(4 * (size_t)maxrows), pos = cx.alloc(4 * nent), + XS = cx.alloc(2 * (size_t)maxrows * n), XS_ref = cx.alloc(2 * (size_t)maxrows * n); + ref::Rng r(7); + for (size_t i = 0; i < (size_t)T * n; ++i) { + const float v = r.sym(); + x.as()[i] = v; + xh.as()[i] = ref::f32_to_f16(v); + } + for (int t = 0; t < T; ++t) { + std::set used; + float s = 0; + for (int k = 0; k < TOPK; ++k) { + int e; + do e = (int)(r.next() % NE); while (used.count(e)); + used.insert(e); + ids.as()[t * TOPK + k] = e; + s += (w.as()[t * TOPK + k] = 0.2f + r.uni()); + } + for (int k = 0; k < TOPK; ++k) w.as()[t * TOPK + k] /= s; + } + MmPlanArgs pa{(uint32_t)nent, TOPK, (uint32_t)NE, (uint32_t)TM, (uint32_t)nblk}; + std::vector ds_plan; + mtl::Dispatch h; h.p = &phist; h.args = {mtl::Arg::B(ids), mtl::Arg::B(bcnt), mtl::Arg::V(pa)}; h.grid[0] = nblk; h.tg[0] = 1024; ds_plan.push_back(h); + mtl::Dispatch sc; sc.p = &pscan; sc.args = {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta), mtl::Arg::V(pa), mtl::Arg::B(stok)}; sc.tg[0] = NE; ds_plan.push_back(sc); + mtl::Dispatch pl; pl.p = &pplace; pl.args = {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok), mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}; pl.grid[0] = nblk; pl.tg[0] = 1024; ds_plan.push_back(pl); + bool ok_bits = true; + if (baseline) { + auto bd = ds_plan; + bd[0].p = &bh; bd[1].p = &bs; bd[2].p = &bp; + cx.run(bd); + std::vector outputs{&bcnt, &boff, &eoff, &tiles, &meta, &stok, &sw, &pos}; + std::vector sizes{4u * nblk * NE, 4u * nblk * NE, 4u * NE, + 8u * (maxrows / TM + 1), 64, 4u * maxrows, 4u * maxrows, 4u * nent}; + std::vector> expected; + for (size_t i = 0; i < outputs.size(); ++i) { + const auto* bytes = outputs[i]->as(); + expected.emplace_back(bytes, bytes + sizes[i]); + std::memset(outputs[i]->ptr, 0, sizes[i]); + } + cx.run(ds_plan); + const char* names[]{"bcnt", "boff", "eoff", "tiles", "meta", "srt_tok", "srt_w", "pos"}; + for (size_t i = 0; i < outputs.size(); ++i) { + if (std::memcmp(outputs[i]->ptr, expected[i].data(), sizes[i]) != 0) { + std::fprintf(stderr, "%s memcmp failed\n", names[i]); + ok_bits = false; + } + } + cx.warm(bd, 0.02); + cx.warm(ds_plan, 0.02); + std::vector bt, ct; + for (int i = 0; i < 7; ++i) { + if (i & 1) { + ct.push_back(cx.run(ds_plan, 8, false, true)); + bt.push_back(cx.run(bd, 8, false, true)); + } else { + bt.push_back(cx.run(bd, 8, false, true)); + ct.push_back(cx.run(ds_plan, 8, false, true)); + } + } + std::printf(" plan baseline %.3f candidate %.3f us/L; all 8 buffers memcmp=%s\n", + median(bt) * 1e6 / 8, median(ct) * 1e6 / 8, ok_bits ? "PASS" : "FAIL"); + } else cx.run(ds_plan); + const uint32_t tot_rows = meta.as()[1], ne = n; + const size_t xs_bytes = 2 * (size_t)tot_rows * n; + bool ok_plan = true; + std::vector seen(tot_rows, 0); + for (int i = 0; i < nent; ++i) { + const int row = pos.as()[i]; + if (row < 0 || row >= (int)tot_rows || stok.as()[row] != i / TOPK) ok_plan = false; + else seen[row]++; + } + for (uint32_t r0 = 0; r0 < tot_rows; ++r0) + if (seen[r0] > 1 || (seen[r0] == 0 && stok.as()[r0] != -1)) ok_plan = false; + mtl::Dispatch d_ref; d_ref.p = &pgat_f; d_ref.args = {mtl::Arg::B(x), mtl::Arg::B(stok), mtl::Arg::B(XS_ref), mtl::Arg::V(ne)}; d_ref.grid[0] = tot_rows; d_ref.tg[0] = 256; + cx.run({d_ref}); + auto b_ref = d_ref; b_ref.p = &bf; b_ref.args[2] = mtl::Arg::B(XS); + cx.run({b_ref}); + bool ok_xs = std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) == 0; + for (uint32_t tg_sz : {32u, 64u, 80u, 128u, 160u, 256u}) { + std::memset(XS.ptr, 0x55, xs_bytes); + mtl::Dispatch d; d.p = &pgat; d.args = {mtl::Arg::B(xh), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; d.grid[0] = tot_rows; d.tg[0] = tg_sz; + cx.run({d}); + if (std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) != 0) ok_xs = false; + d.grid[0] = 1; d.indirect = &meta; d.ind_off = 8 * 4; + cx.run({d}); + if (std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) != 0) ok_xs = false; + auto b_gat = d; b_gat.p = &bg; + cx.warm({b_gat, d}, 0.02); + std::vector bt, ct; + for (int i = 0; i < 7; ++i) { + if (i & 1) { + ct.push_back(cx.run({d}, 8, false, true)); + bt.push_back(cx.run({b_gat}, 8, false, true)); + } else { + bt.push_back(cx.run({b_gat}, 8, false, true)); + ct.push_back(cx.run({d}, 8, false, true)); + } + } + std::printf(" T=%d NE=%d rows=%u tg=%u baseline %.3f candidate %.3f us/L XS_memcmp=%s\n", + T, NE, tot_rows, tg_sz, median(bt) * 1e6 / 8, median(ct) * 1e6 / 8, ok_xs ? "PASS" : "FAIL"); + } + auto combined = ds_plan; + mtl::Dispatch gather = d_ref; + gather.p = &pgat; + gather.args[0] = mtl::Arg::B(xh); + gather.args[2] = mtl::Arg::B(XS); + gather.indirect = &meta; gather.ind_off = 8 * 4; + combined.push_back(gather); + std::vector times; + for (int i = 0; i < 3; ++i) times.push_back(cx.run(combined, 8, false, true)); + ok_xs &= std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) == 0; + std::printf(" combined plan+gather tg256 %.3f us/L plan=%s XS_memcmp=%s\n", + median(times) * 1e6 / 8, ok_plan ? "PASS" : "FAIL", ok_xs ? "PASS" : "FAIL"); + if (!ok_plan || !ok_xs || !ok_bits) return 1; + return 0; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc > 1 && std::string(argv[1]) == "--gather-bench") { + return gather_bench(argc > 2 ? std::atoi(argv[2]) : 2048, + argc > 3 ? std::atoi(argv[3]) : 512, + argc > 4 ? std::atoi(argv[4]) : N_EMBD); + } + int gu = ref::Q2_0, dn = ref::Q2_0, T = 256, TM = 16, TN = 64, TK = 64, SG = 2, DTN = -1, DTK = -1, DSG = -1; + bool bench = false, v2 = false, v3 = false, v4 = false, v5 = false; + int MS = 1; + bool gu_only = false, dn_only = false; + const char* dump = nullptr, *ab_dir = nullptr; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + if (a == "--gu") gu = type_of(argv[++i]); + else if (a == "--dn") dn = type_of(argv[++i]); + else if (a == "--T" || a == "-T") T = std::atoi(argv[++i]); + else if (a == "--tile") std::sscanf(argv[++i], "%d,%d,%d,%d", &TM, &TN, &TK, &SG); + else if (a == "--dtile") { + int m, n, k, s; + const char* value = argv[++i]; + if (std::sscanf(value, "%d,%d,%d,%d", &m, &n, &k, &s) == 4) { + if (m != TM) { + std::fprintf(stderr, "down tile TM must match gate/up TM\n"); + return 2; + } + DTN = n; DTK = k; DSG = s; + } else std::sscanf(value, "%d,%d,%d", &DTN, &DTK, &DSG); + } + else if (a == "--bench") bench = true; + else if (a == "--gu-bench") bench = gu_only = true; // --bench timing the gate/up GEMM alone + else if (a == "--dn-bench") bench = dn_only = true; + else if (a == "--v2") v2 = true; + else if (a == "--v3") v2 = v3 = true; // v2's gather + the register-operand GEMMs (TN per simdgroup) + else if (a == "--v4") { v2 = v3 = v4 = true; MS = std::atoi(argv[++i]); } // mm_gu4 with MS 16-row subtiles + else if (a == "--v5") { v2 = v3 = true; v5 = true; } // mm_gu5 (weight blocks staged in threadgroup memory) + else if (a == "--dump") dump = argv[++i]; // HS and y to a file: bitwise comparison of two kernel versions + else if (a == "--ab") ab_dir = argv[++i]; + } + const bool isolated_bench = gu_only || dn_only; + if (bench && T == 256) T = 2048; + if (v4) TM = 16 * MS; // the plan's tile height; mm_gu4's name carries MS instead + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "moe_mm.metal"}, (v3 && !v4 && !v5) ? "#define MM_V3_ONLY 1\n" : ""); + const int NE = bench ? 512 : 64, LAYERS = bench && !isolated_bench ? 4 : 1; + const size_t gu_row = ref::row_bytes(gu, N_EMBD), d_row = ref::row_bytes(dn, N_FF); + const size_t up_off = (size_t)N_FF * gu_row, down_off = 2 * up_off, blob = down_off + (size_t)N_EMBD * d_row; + std::printf("device %s; %s/%s, T %d, %d experts, tile TM %d TN %d TK %d SG %d\n", cx.device_name().c_str(), + ref::type_name(gu), ref::type_name(dn), T, NE, TM, TN, TK, SG); + mtl::Buffer W = cx.alloc(blob * NE * LAYERS); + { + std::vector tmp(blob * 16); + for (int e = 0; e < 16; ++e) { + ref::random_blocks(gu, tmp.data() + e * blob, 2 * N_FF, N_EMBD, 100 + e, 0.02f); + ref::random_blocks(dn, tmp.data() + e * blob + down_off, N_EMBD, N_FF, 200 + e, 0.02f); + } + for (size_t e = 0; e < (size_t)NE * LAYERS; ++e) std::memcpy(W.as() + e * blob, tmp.data() + (e % 16) * blob, blob); + } + const int nent = T * TOPK, nblk = (nent + 1023) / 1024; + const int maxrows = nent + NE * TM; + mtl::Buffer x = cx.alloc(4 * (size_t)T * N_EMBD), xh = cx.alloc(2 * (size_t)T * N_EMBD), ids = cx.alloc(4 * nent), w = cx.alloc(4 * nent), + bcnt = cx.alloc(4 * (size_t)nblk * NE), boff = cx.alloc(4 * (size_t)nblk * NE), eoff = cx.alloc(4 * NE), + tiles = cx.alloc(8 * (size_t)(maxrows / TM + 1)), meta = cx.alloc(64), stok = cx.alloc(4 * (size_t)maxrows), + sw = cx.alloc(4 * (size_t)maxrows), pos = cx.alloc(4 * nent), H = cx.alloc(4 * (size_t)maxrows * N_FF), + YE = cx.alloc(4 * (size_t)maxrows * N_EMBD), y = cx.alloc(4 * (size_t)T * N_EMBD), + XS = cx.alloc(2 * (size_t)maxrows * N_EMBD), HS = cx.alloc(2 * (size_t)maxrows * N_FF); + ref::Rng r(7); + for (size_t i = 0; i < (size_t)T * N_EMBD; ++i) { + const float v = r.sym(); + x.as()[i] = v; + xh.as()[i] = ref::f32_to_f16(v); + } + for (int t = 0; t < T; ++t) { + std::set used; + float s = 0; + for (int k = 0; k < TOPK; ++k) { + int e; + do e = (int)(r.next() % NE); while (used.count(e)); + used.insert(e); + ids.as()[t * TOPK + k] = e; + s += (w.as()[t * TOPK + k] = 0.2f + r.uni()); + } + for (int k = 0; k < TOPK; ++k) w.as()[t * TOPK + k] /= s; + } + mtl::Pipeline phist = cx.pipeline(lib, "mm_hist"), pscan = cx.pipeline(lib, "mm_scan"), pplace = cx.pipeline(lib, "mm_place"), + pcomb = cx.pipeline(lib, "mm_combine"); + if (DTN < 0) { DTN = TN; DTK = TK; DSG = SG; } + if (dn == ref::IQ4_NL && DTK < 32) DTK = 32; + const std::string tl = std::to_string(TM) + "_" + std::to_string(TN) + "_" + std::to_string(TK) + "_" + std::to_string(SG); + const std::string dtl = std::to_string(TM) + "_" + std::to_string(DTN) + "_" + std::to_string(DTK) + "_" + std::to_string(DSG); + mtl::Constants probe; + probe.i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0); + const bool k4 = coop_k4(cx); + if (v3) std::printf("right-input layout: %s\n", k4 ? "runs of 4 k" : "other (per-element decode)"); + mtl::Pipeline pgu = v5 ? cx.pipeline(lib, "mm_gu5_" + std::to_string(gu) + "_" + std::to_string(TN) + "_" + std::to_string(SG), + mtl::Constants().b(7, k4)) + : v4 ? cx.pipeline(lib, "mm_gu4_" + std::to_string(gu) + "_" + std::to_string(MS) + tl.substr(tl.find('_')), + mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0)) + : v3 ? cx.pipeline(lib, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)) + : cx.pipeline(lib, (v2 ? "mm_gu2_" : "mm_gu_") + std::to_string(gu) + "_" + tl, probe), + pdn = v3 ? cx.pipeline(lib, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4)) + : cx.pipeline(lib, (v2 ? "mm_down2_" : "mm_down_") + std::to_string(dn) + "_" + (v2 ? dtl : tl)), + pgat = cx.pipeline(lib, "mm_gather_h"); + MmPlanArgs pa{(uint32_t)nent, TOPK, (uint32_t)NE, (uint32_t)TM, (uint32_t)nblk}; + auto encode = [&](int layer, std::vector& ds, bool plan, bool gemm) { + if (plan) { + mtl::Dispatch h; h.p = &phist; h.args = {mtl::Arg::B(ids), mtl::Arg::B(bcnt), mtl::Arg::V(pa)}; h.grid[0] = nblk; h.tg[0] = 1024; ds.push_back(h); + meta.as()[2] = N_FF / TN; + mtl::Dispatch sc; sc.p = &pscan; sc.args = {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta), mtl::Arg::V(pa), mtl::Arg::B(stok)}; sc.tg[0] = NE; ds.push_back(sc); + mtl::Dispatch pl; pl.p = &pplace; pl.args = {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok), mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}; pl.grid[0] = nblk; pl.tg[0] = 1024; ds.push_back(pl); + } + if (gemm) { + MmArgs ma{(uint64_t)layer * NE * blob, (uint64_t)layer * NE * blob + up_off, (uint64_t)layer * NE * blob + down_off, + blob, blob, (uint32_t)gu_row, (uint32_t)d_row, N_EMBD, N_FF}; + if (v2 || v3) { + const uint32_t ne = N_EMBD; + mtl::Dispatch ga; ga.p = &pgat; ga.args = {mtl::Arg::B(xh), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; + ga.grid[0] = meta.as()[1]; ga.tg[0] = 256; ds.push_back(ga); + } + mtl::Dispatch g; g.p = &pgu; + if (v2 || v3) g.args = {mtl::Arg::B(W), mtl::Arg::B(XS), mtl::Arg::B(HS), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + else g.args = {mtl::Arg::B(W), mtl::Arg::B(x), mtl::Arg::B(H), mtl::Arg::B(stok), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + g.grid[0] = N_FF / (v3 ? TN * SG : TN); g.grid[1] = meta.as()[0]; g.tg[0] = 32 * SG; ds.push_back(g); + mtl::Dispatch d; d.p = &pdn; d.args = {mtl::Arg::B(W), mtl::Arg::B(v2 || v3 ? HS : H), mtl::Arg::B(YE), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + d.grid[0] = N_EMBD / (v3 ? DTN * DSG : v2 ? DTN : TN); d.grid[1] = meta.as()[0]; d.tg[0] = 32 * ((v2 || v3) ? DSG : SG); ds.push_back(d); + const uint32_t ca[4] = {(uint32_t)T, N_EMBD, TOPK, 0}; + mtl::Dispatch cb; cb.p = &pcomb; cb.args = {mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(y), mtl::Arg::V(ca)}; + cb.grid[0] = (T * N_EMBD + 255) / 256; cb.tg[0] = 256; ds.push_back(cb); + } + }; + { + std::vector ds; + encode(0, ds, true, false); + cx.run(ds); // the plan first: the GEMMs' grid is the tile count + } + const uint32_t ntiles = meta.as()[0]; + std::printf("plan: %u tiles of %d, %u padded rows for %d entries\n", ntiles, TM, meta.as()[1], nent); + auto compare_ab = [&](const std::vector& ds, mtl::Pipeline& pgu_b, mtl::Pipeline& pdn_b, size_t pr) { + std::vector hs_a(HS.as(), HS.as() + 2 * pr * N_FF); + std::vector ye_a(YE.as(), YE.as() + 4 * pr * N_EMBD); + std::vector y_a(y.as(), y.as() + 4 * (size_t)T * N_EMBD); + std::memset(HS.ptr, 0, hs_a.size()); + std::memset(YE.ptr, 0, ye_a.size()); + std::memset(y.ptr, 0, y_a.size()); + auto ds_b = ds; + ds_b[v2 ? 1 : 0].p = &pgu_b; + ds_b[v2 ? 2 : 1].p = &pdn_b; + cx.run(ds_b); + const bool ok_hs = std::memcmp(HS.ptr, hs_a.data(), hs_a.size()) == 0; + const bool ok_ye = std::memcmp(YE.ptr, ye_a.data(), ye_a.size()) == 0; + const bool ok_y = std::memcmp(y.ptr, y_a.data(), y_a.size()) == 0; + std::printf("ab bitwise (T=%d): HS_memcmp=%s YE_memcmp=%s Y_memcmp=%s\n", + T, ok_hs ? "PASS" : "FAIL", ok_ye ? "PASS" : "FAIL", ok_y ? "PASS" : "FAIL"); + return ok_hs && ok_ye && ok_y; + }; + if (!bench) { + // the plan: every entry placed once, rows of an expert contiguous and in entry order + bool ok_plan = true; + std::vector seen(maxrows, 0); + for (int i = 0; i < nent; ++i) { + const int row = pos.as()[i]; + if (row < 0 || row >= (int)meta.as()[1] || stok.as()[row] != i / TOPK) ok_plan = false; + else seen[row]++; + } + for (int i = 0; i < maxrows; ++i) if (seen[i] > 1) ok_plan = false; + std::vector ds; + encode(0, ds, false, true); + cx.run(ds); + const size_t pr = meta.as()[1]; + if (dump) { + FILE* f = std::fopen(dump, "wb"); + if (v2) std::fwrite(HS.as(), 2, pr * N_FF, f); + std::fwrite(YE.as(), 4, pr * N_EMBD, f); + std::fwrite(y.as(), 4, (size_t)T * N_EMBD, f); + std::fclose(f); + return 0; + } + if (ab_dir) { + const std::string kd(ab_dir); + const int lib_b = cx.compile({kd + "/quant.metal", kd + "/units.metal", kd + "/mma.metal", kd + "/moe_mm.metal"}, (v3 && !v4 && !v5) ? "#define MM_V3_ONLY 1\n" : ""); + mtl::Pipeline pgu_b = cx.pipeline(lib_b, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)); + mtl::Pipeline pdn_b = cx.pipeline(lib_b, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4)); + return (ok_plan && compare_ab(ds, pgu_b, pdn_b, pr)) ? 0 : 1; + } + // reference: fp16 operands (x, weights, h), fp32+ accumulation; and the exact one + double se16 = 0, se = 0, sr = 0; + std::vector wg((size_t)N_FF * N_EMBD), wu(wg.size()), wd((size_t)N_EMBD * N_FF); + std::vector yr16((size_t)T * N_EMBD, 0.0), yr((size_t)T * N_EMBD, 0.0); + for (int e = 0; e < NE; ++e) { + const uint8_t* b = W.as() + e * blob; + bool any = false; + for (int i = 0; i < nent && !any; ++i) any = ids.as()[i] == e; + if (!any) continue; + ref::dequant(gu, b, wg.data(), wg.size()); + ref::dequant(gu, b + up_off, wu.data(), wu.size()); + ref::dequant(dn, b + down_off, wd.data(), wd.size()); + for (int i = 0; i < nent; ++i) { + if (ids.as()[i] != e) continue; + const int t = i / TOPK; + const float* xt = x.as() + (size_t)t * N_EMBD; + std::vector h16v(N_FF), hv(N_FF); + for (int f = 0; f < N_FF; ++f) { + double g = 0, u = 0, g16 = 0, u16 = 0; + for (int k = 0; k < N_EMBD; ++k) { + const float xk = xt[k], xh = h16(xk); + g += (double)wg[(size_t)f * N_EMBD + k] * xk; + u += (double)wu[(size_t)f * N_EMBD + k] * xk; + g16 += (double)h16(wg[(size_t)f * N_EMBD + k]) * xh; + u16 += (double)h16(wu[(size_t)f * N_EMBD + k]) * xh; + } + hv[f] = g / (1 + std::exp(-g)) * u; + h16v[f] = h16((float)(g16 / (1 + std::exp(-g16)) * u16)); + } + const double wt = w.as()[i]; + for (int d = 0; d < N_EMBD; ++d) { + double s = 0, s16 = 0; + for (int f = 0; f < N_FF; ++f) { + s += (double)wd[(size_t)d * N_FF + f] * hv[f]; + s16 += (double)h16(wd[(size_t)d * N_FF + f]) * h16v[f]; + } + yr[(size_t)t * N_EMBD + d] += wt * s; + yr16[(size_t)t * N_EMBD + d] += wt * s16; + } + } + } + for (size_t i = 0; i < yr.size(); ++i) { + const double g = y.as()[i]; + se16 += (g - yr16[i]) * (g - yr16[i]); + se += (g - yr[i]) * (g - yr[i]); + sr += yr[i] * yr[i]; + } + const double rel16 = std::sqrt(se16 / sr), rel = std::sqrt(se / sr); + const bool ok = ok_plan && rel16 < 2e-4; + std::printf("parity: plan %s; y vs fp16-operand reference rel L2 %.2e, vs exact float reference %.2e %s\n", + ok_plan ? "ok" : "WRONG", rel16, rel, ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; + } + // bench: the plan, the two GEMMs and the combine, LAYERS distinct weight sets rotated + double flops = 0; + { + const double rows = meta.as()[1]; + flops = 2.0 * nent * (2.0 * N_FF * N_EMBD + (double)N_EMBD * N_FF); + std::printf("useful GFLOP per layer %.1f (padded rows %.0f = %.2fx entries)\n", flops / 1e9, rows, rows / nent); + } + auto timeit = [&](bool plan, bool gemm, int part) { + std::vector ds; + const int repeats = isolated_bench ? 4 : 8; + for (int l = 0; l < repeats; ++l) { + std::vector one; + encode(l % LAYERS, one, plan, gemm); + if (part >= 0 && gemm) { ds.push_back(one[part]); continue; } + for (auto& d : one) ds.push_back(d); + } + cx.warm(ds, isolated_bench ? 0.1 : 0.3); + std::vector t; + for (int i = 0; i < (isolated_bench ? 3 : 5); ++i) t.push_back(cx.run(ds) / repeats); + return median(t); + }; + const int o = v2 ? 1 : 0; + if (isolated_bench) { + std::vector ds; + encode(0, ds, false, true); + cx.run(ds); + } + if (ab_dir && isolated_bench) { + const std::string kd(ab_dir); + const int lib_b = cx.compile({kd + "/quant.metal", kd + "/units.metal", kd + "/mma.metal", kd + "/moe_mm.metal"}, (v3 && !v4 && !v5) ? "#define MM_V3_ONLY 1\n" : ""); + mtl::Pipeline pgu_b = cx.pipeline(lib_b, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)); + mtl::Pipeline pdn_b = cx.pipeline(lib_b, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4)); + std::vector parity; + encode(0, parity, false, true); + const size_t pr = meta.as()[1]; + std::memset(HS.ptr, 0, 2 * pr * N_FF); + std::memset(YE.ptr, 0, 4 * pr * N_EMBD); + cx.run(parity); + if (!compare_ab(parity, pgu_b, pdn_b, pr)) return 1; + const int part = dn_only ? (o + 1) : o; + std::vector one; + encode(0, one, false, true); + std::vector ds_a{one[part]}; + mtl::Dispatch db = one[part]; + db.p = dn_only ? &pdn_b : &pgu_b; + std::vector ds_b{db}; + cx.warm(ds_a, 0.05); + cx.warm(ds_b, 0.05); + std::vector ta, tb; + double best_pair = 1e30; + for (int i = 0; i < 19; ++i) { + double va, vb; + if (i & 1) { + vb = cx.run(ds_b, 2) / 2; + va = cx.run(ds_a, 2) / 2; + } else { + va = cx.run(ds_a, 2) / 2; + vb = cx.run(ds_b, 2) / 2; + } + best_pair = std::min(best_pair, std::min(va, vb)); + ta.push_back(va); + tb.push_back(vb); + } + std::vector fa, fb; + for (size_t i = 0; i < ta.size(); ++i) { + if (std::max(ta[i], tb[i]) <= 1.10 * best_pair) { + fa.push_back(ta[i]); + fb.push_back(tb[i]); + } + } + if (fa.size() >= 3) { ta = fa; tb = fb; } + const double ma = median(ta) * 1e6, mb = median(tb) * 1e6; + const double mina = *std::min_element(ta.begin(), ta.end()) * 1e6; + const double minb = *std::min_element(tb.begin(), tb.end()) * 1e6; + std::printf("%s A/B (T=%d): A med=%.1f (min=%.1f) | B med=%.1f (min=%.1f) | diff med=%+.2f%% min=%+.2f%%\n", + dn_only ? "down" : "gate/up", T, ma, mina, mb, minb, (mb - ma) / ma * 100.0, (minb - mina) / mina * 100.0); + return 0; + } + if (dn_only) { + const double td = timeit(false, true, o + 1); + std::printf("down %.1f us (%.2f TFLOP/s useful)\n", td * 1e6, 2.0 * nent * N_FF * N_EMBD / td / 1e12); + return 0; + } + if (gu_only) { + const double tg = timeit(false, true, o); + std::printf("gate/up %.1f us (%.2f TFLOP/s useful)\n", tg * 1e6, 2.0 * nent * 2.0 * N_FF * N_EMBD / tg / 1e12); + return 0; + } + const double tp = timeit(true, false, -1), tg = timeit(false, true, o), td = timeit(false, true, o + 1), + tc = timeit(false, true, o + 2), ta = timeit(false, true, -1); + if (v2) std::printf("gather %.1f us\n", timeit(false, true, 0) * 1e6); + const double fgu = 2.0 * nent * 2.0 * N_FF * N_EMBD, fdn = 2.0 * nent * N_EMBD * N_FF; + std::printf("per layer: plan %.1f us, gate/up %.1f us (%.1f TFLOP/s useful), down %.1f us (%.1f TFLOP/s), combine %.1f us;" + " GEMMs + combine %.1f us = %.1f TFLOP/s; %.1f us per token per layer -> 48 layers: %.0f tok/s (experts only)\n", + tp * 1e6, tg * 1e6, fgu / tg / 1e12, td * 1e6, fdn / td / 1e12, tc * 1e6, ta * 1e6, flops / ta / 1e12, + (ta + tp) / T * 1e6, T / (48 * (ta + tp))); + return 0; +} diff --git a/results/round39/astra_r39_qsa/logs/baseline-ab.log b/results/round39/astra_r39_qsa/logs/baseline-ab.log new file mode 100644 index 0000000..2e03c87 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/baseline-ab.log @@ -0,0 +1,34 @@ +paired P=3000 nt=1 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 7.367 / 7.112 (ratio 0.9895) +score 10.122 / 10.008 (ratio 1.0039) +topk 32.169 / 31.419 (ratio 0.9823) +attn 108.711 / 108.607 (ratio 1.0004) +merge 9.823 / 9.815 (ratio 0.9997) +paired P=3000 nt=2 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 9.133 / 9.141 (ratio 1.0000) +score 16.008 / 16.260 (ratio 1.0158) +topk 31.391 / 32.273 (ratio 1.0208) +attn 194.680 / 194.924 (ratio 1.0005) +merge 8.359 / 8.383 (ratio 1.0037) +paired P=3000 nt=3 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 11.125 / 10.984 (ratio 0.9874) +score 21.109 / 21.292 (ratio 1.0041) +topk 32.008 / 31.974 (ratio 0.9883) +attn 124.969 / 124.594 (ratio 0.9971) +merge 16.414 / 16.414 (ratio 1.0000) +paired P=3000 nt=4 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 13.336 / 13.477 (ratio 1.0152) +score 25.703 / 25.924 (ratio 1.0080) +topk 31.875 / 32.083 (ratio 1.0086) +attn 160.729 / 160.659 (ratio 0.9989) +merge 23.516 / 23.555 (ratio 1.0028) +paired P=3000 nt=8 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 24.187 / 23.508 (ratio 0.9850) +score 46.758 / 46.628 (ratio 0.9972) +topk 31.792 / 32.055 (ratio 1.0088) +attn 316.414 / 316.109 (ratio 0.9991) +merge 46.141 / 45.961 (ratio 0.9961) +paired P=3000 nt=1984 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 985.771 / 981.042 (ratio 1.0134) +paired P=3000 nt=2048 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 1004.865 / 1006.500 (ratio 1.0016) diff --git a/results/round39/astra_r39_qsa/logs/baseline-moe.log b/results/round39/astra_r39_qsa/logs/baseline-moe.log new file mode 100644 index 0000000..f93da6b --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/baseline-moe.log @@ -0,0 +1,2 @@ +mtl: no function mm_gu_42_16_64_16_2 Function mm_gu_42_16_64_16_2 was not found in the library +device Apple M5; Q2_0/Q2_0, T 256, 64 experts, tile TM 16 TN 64 TK 16 SG 2 diff --git a/results/round39/astra_r39_qsa/logs/confirm-merge1-ab-10180.log b/results/round39/astra_r39_qsa/logs/confirm-merge1-ab-10180.log new file mode 100644 index 0000000..3d8b8bd --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/confirm-merge1-ab-10180.log @@ -0,0 +1,34 @@ +paired P=10180 nt=1 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 7.260 / 7.305 (ratio 0.9993) +score 5.266 / 5.281 (ratio 1.0030) +topk 8.297 / 7.878 (ratio 0.9482) +attn 46.211 / 46.466 (ratio 1.0017) +merge 5.760 / 4.130 (ratio 0.7170) +paired P=10180 nt=2 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 8.953 / 8.914 (ratio 1.0026) +score 39.716 / 38.818 (ratio 0.9977) +topk 8.440 / 8.391 (ratio 0.9883) +attn 83.565 / 83.771 (ratio 1.0018) +merge 6.781 / 5.469 (ratio 0.8099) +paired P=10180 nt=3 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 10.867 / 10.898 (ratio 1.0029) +score 57.008 / 57.117 (ratio 1.0002) +topk 39.073 / 39.734 (ratio 1.0099) +attn 125.500 / 125.396 (ratio 0.9989) +merge 16.398 / 17.466 (ratio 1.0660) +paired P=10180 nt=4 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 12.206 / 12.245 (ratio 1.0071) +score 16.836 / 16.833 (ratio 1.0014) +topk 8.242 / 8.302 (ratio 1.0025) +attn 161.979 / 161.953 (ratio 1.0000) +merge 23.547 / 22.190 (ratio 0.9424) +paired P=10180 nt=8 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 7.641 / 7.750 (ratio 1.0087) +score 32.688 / 32.781 (ratio 1.0019) +topk 8.260 / 8.143 (ratio 0.9987) +attn 325.305 / 325.188 (ratio 0.9995) +merge 46.406 / 44.102 (ratio 0.9527) +paired P=10180 nt=1984 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 976.844 / 988.021 (ratio 1.0023) +paired P=10180 nt=2048 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 1005.812 / 1016.844 (ratio 1.0068) diff --git a/results/round39/astra_r39_qsa/logs/confirm-merge1-ab-3000.log b/results/round39/astra_r39_qsa/logs/confirm-merge1-ab-3000.log new file mode 100644 index 0000000..81d9baa --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/confirm-merge1-ab-3000.log @@ -0,0 +1,34 @@ +paired P=3000 nt=1 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 7.315 / 7.289 (ratio 1.0028) +score 6.859 / 6.810 (ratio 0.9924) +topk 15.482 / 15.500 (ratio 1.0005) +attn 86.906 / 86.859 (ratio 0.9996) +merge 8.078 / 5.469 (ratio 0.6769) +paired P=3000 nt=2 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 8.906 / 8.742 (ratio 0.9894) +score 3.458 / 3.456 (ratio 0.9985) +topk 6.732 / 6.750 (ratio 1.0015) +attn 82.930 / 83.094 (ratio 1.0013) +merge 6.661 / 5.469 (ratio 0.8210) +paired P=3000 nt=3 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 10.706 / 10.664 (ratio 1.0144) +score 13.659 / 13.661 (ratio 0.9979) +topk 6.612 / 6.555 (ratio 0.9831) +attn 124.562 / 124.687 (ratio 1.0006) +merge 16.391 / 17.211 (ratio 1.0551) +paired P=3000 nt=4 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 13.273 / 13.461 (ratio 1.0035) +score 10.508 / 10.500 (ratio 0.9993) +topk 6.544 / 6.638 (ratio 1.0048) +attn 160.992 / 160.945 (ratio 0.9996) +merge 23.477 / 22.227 (ratio 0.9452) +paired P=3000 nt=8 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 24.237 / 24.201 (ratio 1.0000) +score 12.984 / 12.990 (ratio 0.9994) +topk 6.706 / 6.724 (ratio 1.0194) +attn 316.211 / 316.211 (ratio 1.0001) +merge 45.961 / 44.016 (ratio 0.9569) +paired P=3000 nt=1984 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 978.250 / 996.896 (ratio 1.0102) +paired P=3000 nt=2048 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 1008.271 / 1018.385 (ratio 1.0100) diff --git a/results/round39/astra_r39_qsa/logs/confirm-merge1-ab-30000.log b/results/round39/astra_r39_qsa/logs/confirm-merge1-ab-30000.log new file mode 100644 index 0000000..35ca3d3 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/confirm-merge1-ab-30000.log @@ -0,0 +1,34 @@ +paired P=30000 nt=1 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 7.141 / 7.320 (ratio 1.0089) +score 58.344 / 58.073 (ratio 1.0054) +topk 14.534 / 14.539 (ratio 1.0020) +attn 46.711 / 46.901 (ratio 1.0025) +merge 5.966 / 4.052 (ratio 0.6784) +paired P=30000 nt=2 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 9.141 / 8.896 (ratio 0.9881) +score 42.872 / 27.826 (ratio 1.0010) +topk 14.305 / 14.081 (ratio 0.9878) +attn 84.143 / 84.201 (ratio 1.0000) +merge 6.687 / 5.372 (ratio 0.8106) +paired P=30000 nt=3 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 10.883 / 10.865 (ratio 0.9927) +score 53.570 / 51.617 (ratio 0.9994) +topk 14.102 / 14.656 (ratio 1.0306) +attn 126.227 / 126.073 (ratio 0.9985) +merge 16.419 / 17.349 (ratio 1.0563) +paired P=30000 nt=4 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 12.599 / 12.344 (ratio 0.9806) +score 62.320 / 47.906 (ratio 0.9993) +topk 14.078 / 14.492 (ratio 1.0086) +attn 165.086 / 165.195 (ratio 1.0005) +merge 23.615 / 22.219 (ratio 0.9409) +paired P=30000 nt=8 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 22.362 / 23.547 (ratio 1.0209) +score 108.987 / 101.703 (ratio 0.9996) +topk 14.385 / 14.544 (ratio 0.9903) +attn 332.862 / 332.672 (ratio 0.9997) +merge 45.729 / 43.919 (ratio 0.9579) +paired P=30000 nt=1984 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 1003.563 / 991.687 (ratio 1.0010) +paired P=30000 nt=2048 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 1042.229 / 1036.750 (ratio 0.9982) diff --git a/results/round39/astra_r39_qsa/logs/confirm-merge2-ab-10180.log b/results/round39/astra_r39_qsa/logs/confirm-merge2-ab-10180.log new file mode 100644 index 0000000..c002f4a --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/confirm-merge2-ab-10180.log @@ -0,0 +1,34 @@ +paired P=10180 nt=1 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 7.273 / 7.318 (ratio 0.9972) +score 5.289 / 5.273 (ratio 0.9970) +topk 8.365 / 8.406 (ratio 0.9896) +attn 50.875 / 50.437 (ratio 0.9921) +merge 10.089 / 6.523 (ratio 0.6479) +paired P=10180 nt=2 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 9.083 / 9.008 (ratio 1.0087) +score 8.987 / 8.979 (ratio 0.9977) +topk 8.219 / 8.357 (ratio 1.0285) +attn 99.315 / 98.672 (ratio 0.9917) +merge 21.937 / 21.549 (ratio 0.9823) +paired P=10180 nt=3 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 11.016 / 10.904 (ratio 0.9922) +score 21.297 / 21.302 (ratio 1.0000) +topk 8.271 / 8.271 (ratio 0.9819) +attn 147.648 / 146.031 (ratio 0.9890) +merge 33.617 / 32.982 (ratio 0.9843) +paired P=10180 nt=4 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 6.654 / 6.542 (ratio 0.9968) +score 16.953 / 16.852 (ratio 0.9952) +topk 8.477 / 7.911 (ratio 0.9396) +attn 195.630 / 193.781 (ratio 0.9906) +merge 45.521 / 43.383 (ratio 0.9528) +paired P=10180 nt=8 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 24.161 / 24.185 (ratio 0.9989) +score 32.828 / 32.745 (ratio 0.9975) +topk 8.466 / 8.445 (ratio 0.9890) +attn 385.242 / 382.505 (ratio 0.9929) +merge 90.805 / 86.240 (ratio 0.9495) +paired P=10180 nt=1984 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 985.000 / 992.104 (ratio 1.0021) +paired P=10180 nt=2048 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 1028.552 / 1048.875 (ratio 1.0140) diff --git a/results/round39/astra_r39_qsa/logs/confirm-merge2-ab-3000.log b/results/round39/astra_r39_qsa/logs/confirm-merge2-ab-3000.log new file mode 100644 index 0000000..7e3ff54 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/confirm-merge2-ab-3000.log @@ -0,0 +1,34 @@ +paired P=3000 nt=1 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 2.242 / 2.242 (ratio 1.0000) +score 2.922 / 2.914 (ratio 0.9973) +topk 6.523 / 6.737 (ratio 1.0164) +attn 50.818 / 50.352 (ratio 0.9909) +merge 10.195 / 6.294 (ratio 0.6273) +paired P=3000 nt=2 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 9.177 / 9.182 (ratio 1.0006) +score 3.453 / 3.466 (ratio 1.0030) +topk 6.766 / 6.531 (ratio 0.9654) +attn 98.906 / 98.271 (ratio 0.9951) +merge 21.927 / 21.594 (ratio 0.9794) +paired P=3000 nt=3 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 10.773 / 10.659 (ratio 0.9894) +score 4.586 / 4.589 (ratio 1.0226) +topk 6.776 / 6.745 (ratio 0.9942) +attn 146.336 / 144.961 (ratio 0.9906) +merge 33.615 / 33.089 (ratio 0.9845) +paired P=3000 nt=4 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 12.813 / 13.151 (ratio 0.9820) +score 25.451 / 25.201 (ratio 0.9970) +topk 6.607 / 6.703 (ratio 0.9828) +attn 193.346 / 191.594 (ratio 0.9908) +merge 45.727 / 43.313 (ratio 0.9480) +paired P=3000 nt=8 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 23.443 / 23.995 (ratio 1.0236) +score 10.497 / 10.523 (ratio 0.9985) +topk 6.664 / 6.724 (ratio 1.0246) +attn 380.055 / 377.034 (ratio 0.9922) +merge 90.633 / 86.607 (ratio 0.9549) +paired P=3000 nt=1984 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 980.542 / 983.604 (ratio 1.0031) +paired P=3000 nt=2048 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 1023.500 / 1023.094 (ratio 0.9929) diff --git a/results/round39/astra_r39_qsa/logs/confirm-merge2-ab-30000.log b/results/round39/astra_r39_qsa/logs/confirm-merge2-ab-30000.log new file mode 100644 index 0000000..e8c48a5 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/confirm-merge2-ab-30000.log @@ -0,0 +1,34 @@ +paired P=30000 nt=1 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 2.688 / 2.695 (ratio 1.0039) +score 26.898 / 25.477 (ratio 0.9985) +topk 14.148 / 13.997 (ratio 1.0031) +attn 50.883 / 50.461 (ratio 0.9916) +merge 9.945 / 6.445 (ratio 0.6489) +paired P=30000 nt=2 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 9.062 / 9.062 (ratio 1.0000) +score 25.323 / 25.271 (ratio 0.9978) +topk 14.289 / 14.023 (ratio 0.9929) +attn 99.570 / 98.607 (ratio 0.9899) +merge 21.958 / 21.448 (ratio 0.9747) +paired P=30000 nt=3 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 10.880 / 10.857 (ratio 0.9986) +score 36.188 / 36.133 (ratio 0.9982) +topk 14.065 / 14.484 (ratio 1.0298) +attn 147.904 / 146.596 (ratio 0.9910) +merge 33.594 / 33.279 (ratio 0.9902) +paired P=30000 nt=4 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 13.492 / 12.823 (ratio 0.9401) +score 47.870 / 47.852 (ratio 0.9996) +topk 14.307 / 14.312 (ratio 0.9813) +attn 196.112 / 194.602 (ratio 0.9931) +merge 45.367 / 43.659 (ratio 0.9620) +paired P=30000 nt=8 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 22.393 / 23.977 (ratio 1.0531) +score 94.857 / 94.878 (ratio 0.9999) +topk 14.279 / 14.273 (ratio 1.0162) +attn 388.385 / 386.602 (ratio 0.9956) +merge 91.096 / 87.198 (ratio 0.9549) +paired P=30000 nt=1984 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 1033.010 / 994.500 (ratio 0.9841) +paired P=30000 nt=2048 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 1038.406 / 1010.396 (ratio 0.9621) diff --git a/results/round39/astra_r39_qsa/logs/final-merge2-ab-10180.log b/results/round39/astra_r39_qsa/logs/final-merge2-ab-10180.log new file mode 100644 index 0000000..c082a2c --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/final-merge2-ab-10180.log @@ -0,0 +1,34 @@ +paired P=10180 nt=1 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 7.026 / 7.073 (ratio 1.0056) +score 5.286 / 5.292 (ratio 1.0015) +topk 8.042 / 8.271 (ratio 1.0285) +attn 51.812 / 51.010 (ratio 0.9898) +merge 10.961 / 6.844 (ratio 0.6240) +paired P=10180 nt=2 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 9.039 / 9.115 (ratio 1.0121) +score 12.099 / 12.117 (ratio 1.0017) +topk 9.729 / 9.695 (ratio 0.9968) +attn 100.708 / 99.680 (ratio 0.9901) +merge 21.977 / 21.589 (ratio 0.9834) +paired P=10180 nt=3 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 10.255 / 10.273 (ratio 1.0025) +score 55.672 / 55.685 (ratio 1.0004) +topk 36.477 / 36.427 (ratio 0.9990) +attn 148.958 / 147.414 (ratio 0.9893) +merge 33.732 / 32.961 (ratio 0.9772) +paired P=10180 nt=4 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 13.375 / 13.422 (ratio 0.9970) +score 16.859 / 16.852 (ratio 0.9995) +topk 8.219 / 8.195 (ratio 0.9810) +attn 197.557 / 195.672 (ratio 0.9905) +merge 45.344 / 43.313 (ratio 0.9546) +paired P=10180 nt=8 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 19.945 / 19.906 (ratio 0.9957) +score 142.211 / 142.211 (ratio 1.0002) +topk 36.568 / 36.487 (ratio 0.9983) +attn 388.375 / 385.505 (ratio 0.9927) +merge 91.393 / 86.607 (ratio 0.9493) +paired P=10180 nt=1984 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 985.844 / 977.042 (ratio 1.0114) +paired P=10180 nt=2048 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 1032.750 / 1019.031 (ratio 0.9805) diff --git a/results/round39/astra_r39_qsa/logs/final-merge2-ab-3000.log b/results/round39/astra_r39_qsa/logs/final-merge2-ab-3000.log new file mode 100644 index 0000000..2d3001d --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/final-merge2-ab-3000.log @@ -0,0 +1,34 @@ +paired P=3000 nt=1 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 6.828 / 6.846 (ratio 1.0027) +score 9.732 / 9.740 (ratio 0.9997) +topk 28.977 / 28.914 (ratio 0.9978) +attn 118.224 / 75.477 (ratio 0.9881) +merge 15.143 / 9.273 (ratio 0.6122) +paired P=3000 nt=2 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 8.867 / 8.953 (ratio 1.0025) +score 15.464 / 15.008 (ratio 0.9924) +topk 6.701 / 6.417 (ratio 0.9648) +attn 100.195 / 99.367 (ratio 0.9911) +merge 21.924 / 21.516 (ratio 0.9812) +paired P=3000 nt=3 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 10.935 / 10.857 (ratio 0.9986) +score 6.193 / 6.195 (ratio 1.0017) +topk 7.247 / 7.221 (ratio 0.9968) +attn 147.727 / 146.453 (ratio 0.9917) +merge 33.688 / 32.969 (ratio 0.9791) +paired P=3000 nt=4 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 12.112 / 12.039 (ratio 0.9939) +score 24.513 / 24.578 (ratio 1.0027) +topk 29.099 / 29.016 (ratio 0.9980) +attn 195.198 / 193.445 (ratio 0.9916) +merge 45.458 / 43.227 (ratio 0.9522) +paired P=3000 nt=8 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 19.992 / 19.859 (ratio 0.9965) +score 45.008 / 45.003 (ratio 0.9995) +topk 29.115 / 29.055 (ratio 0.9985) +attn 384.055 / 380.742 (ratio 0.9914) +merge 91.102 / 86.547 (ratio 0.9501) +paired P=3000 nt=1984 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 997.938 / 995.854 (ratio 1.0029) +paired P=3000 nt=2048 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 1015.865 / 1026.094 (ratio 1.0062) diff --git a/results/round39/astra_r39_qsa/logs/final-merge2-ab-30000.log b/results/round39/astra_r39_qsa/logs/final-merge2-ab-30000.log new file mode 100644 index 0000000..0971eb3 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/final-merge2-ab-30000.log @@ -0,0 +1,34 @@ +paired P=30000 nt=1 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 7.078 / 7.107 (ratio 1.0029) +score 32.344 / 32.388 (ratio 0.9983) +topk 43.786 / 43.844 (ratio 1.0013) +attn 132.607 / 118.047 (ratio 0.9911) +merge 17.180 / 10.807 (ratio 0.6295) +paired P=30000 nt=2 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 8.456 / 8.471 (ratio 1.0019) +score 105.219 / 105.216 (ratio 1.0000) +topk 63.245 / 63.401 (ratio 1.0033) +attn 101.031 / 100.091 (ratio 0.9910) +merge 22.125 / 21.547 (ratio 0.9739) +paired P=30000 nt=3 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 10.674 / 10.648 (ratio 0.9978) +score 108.435 / 83.169 (ratio 0.9991) +topk 26.979 / 27.034 (ratio 1.0026) +attn 172.091 / 170.409 (ratio 0.9899) +merge 33.500 / 33.471 (ratio 0.9955) +paired P=30000 nt=4 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 12.654 / 13.430 (ratio 1.0502) +score 207.747 / 208.276 (ratio 1.0011) +topk 33.771 / 33.904 (ratio 1.0034) +attn 198.734 / 196.766 (ratio 0.9902) +merge 45.469 / 43.544 (ratio 0.9577) +paired P=30000 nt=8 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 23.466 / 23.810 (ratio 1.0229) +score 94.807 / 94.805 (ratio 0.9999) +topk 14.232 / 14.391 (ratio 1.0062) +attn 390.703 / 389.167 (ratio 0.9963) +merge 91.594 / 88.021 (ratio 0.9556) +paired P=30000 nt=1984 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 997.521 / 1009.469 (ratio 0.9897) +paired P=30000 nt=2048 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 1069.354 / 1133.031 (ratio 0.9944) diff --git a/results/round39/astra_r39_qsa/logs/final-prompt-k4off.log b/results/round39/astra_r39_qsa/logs/final-prompt-k4off.log new file mode 100644 index 0000000..38a5495 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/final-prompt-k4off.log @@ -0,0 +1,13 @@ +device Apple M5 +pipelines: qsa_attn_mm max threads 1024, tg mem 6912; qsa_flash max threads 1024, tg mem 13760 +pos0 0, T 96: qsa_attn_mm vs qsa_attn (after merge) rel L2 3.09e-04 PASS +pos0 0, T 96: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 3.72e-04 PASS +pos0 1984, T 96: qsa_attn_mm vs qsa_attn (after merge) rel L2 3.61e-04 PASS +pos0 1984, T 96: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 4.27e-04 PASS +pos0 2050, T 96: qsa_attn_mm vs qsa_attn (after merge) rel L2 3.62e-04 PASS +pos0 2050, T 96: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 4.28e-04 PASS +pos0 30000, T 96: qsa_attn_mm vs qsa_attn (after merge) rel L2 3.57e-04 PASS +pos0 30000, T 96: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 4.23e-04 PASS +score pos0 30000, T 96: qsa_score_mm vs qsa_score worst rel 2.76e-06 PASS +qsa_prep(QSA_Q_F16_FC) vs qsa_prep+qsa_q_f16: qcurh memcmp PASS, qsa_flash memcmp PASS +ALL PASS diff --git a/results/round39/astra_r39_qsa/logs/fixed-moe-default-full.log b/results/round39/astra_r39_qsa/logs/fixed-moe-default-full.log new file mode 100644 index 0000000..66eac20 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/fixed-moe-default-full.log @@ -0,0 +1,3 @@ +device Apple M5; Q2_0/Q2_0, T 256, 64 experts, tile TM 16 TN 64 TK 64 SG 2 +plan: 190 tiles of 16, 3040 padded rows for 2560 entries +parity: plan ok; y vs fp16-operand reference rel L2 2.03e-06, vs exact float reference 2.37e-04 PASS diff --git a/results/round39/astra_r39_qsa/logs/fixed-moe-default.log b/results/round39/astra_r39_qsa/logs/fixed-moe-default.log new file mode 100644 index 0000000..9ca2225 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/fixed-moe-default.log @@ -0,0 +1,2 @@ + +[deadline exceeded after 15s; process group 90854 terminated after SIGTERM] diff --git a/results/round39/astra_r39_qsa/logs/fixed-moe-v3-full.log b/results/round39/astra_r39_qsa/logs/fixed-moe-v3-full.log new file mode 100644 index 0000000..f6f610d --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/fixed-moe-v3-full.log @@ -0,0 +1,4 @@ +device Apple M5; IQ2_S/Q2_0, T 256, 64 experts, tile TM 16 TN 32 TK 32 SG 4 +right-input layout: runs of 4 k +plan: 190 tiles of 16, 3040 padded rows for 2560 entries +parity: plan ok; y vs fp16-operand reference rel L2 1.60e-05, vs exact float reference 4.38e-04 PASS diff --git a/results/round39/astra_r39_qsa/logs/fixed-moe-v3.log b/results/round39/astra_r39_qsa/logs/fixed-moe-v3.log new file mode 100644 index 0000000..2f1123c --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/fixed-moe-v3.log @@ -0,0 +1,2 @@ + +[deadline exceeded after 20s; process group 90861 terminated after SIGTERM] diff --git a/results/round39/astra_r39_qsa/logs/fixed-prompt.log b/results/round39/astra_r39_qsa/logs/fixed-prompt.log new file mode 100644 index 0000000..32edaa0 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/fixed-prompt.log @@ -0,0 +1,13 @@ +device Apple M5 +pipelines: qsa_attn_mm max threads 1024, tg mem 7424; qsa_flash max threads 1024, tg mem 13760 +pos0 0, T 96: qsa_attn_mm vs qsa_attn (after merge) rel L2 3.09e-04 PASS +pos0 0, T 96: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 3.72e-04 PASS +pos0 1984, T 96: qsa_attn_mm vs qsa_attn (after merge) rel L2 3.61e-04 PASS +pos0 1984, T 96: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 4.27e-04 PASS +pos0 2050, T 96: qsa_attn_mm vs qsa_attn (after merge) rel L2 3.62e-04 PASS +pos0 2050, T 96: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 4.28e-04 PASS +pos0 30000, T 96: qsa_attn_mm vs qsa_attn (after merge) rel L2 3.57e-04 PASS +pos0 30000, T 96: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 4.23e-04 PASS +score pos0 30000, T 96: qsa_score_mm vs qsa_score worst rel 2.76e-06 PASS +qsa_prep(QSA_Q_F16_FC) vs qsa_prep+qsa_q_f16: qcurh memcmp PASS, qsa_flash memcmp PASS +ALL PASS diff --git a/results/round39/astra_r39_qsa/logs/merge1-ab-10180.log b/results/round39/astra_r39_qsa/logs/merge1-ab-10180.log new file mode 100644 index 0000000..c609011 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/merge1-ab-10180.log @@ -0,0 +1,34 @@ +paired P=10180 nt=1 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 7.354 / 7.141 (ratio 0.9979) +score 26.320 / 26.518 (ratio 1.0064) +topk 39.250 / 38.987 (ratio 0.9894) +attn 107.724 / 94.398 (ratio 0.9966) +merge 8.102 / 5.492 (ratio 0.6776) +paired P=10180 nt=2 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 8.914 / 9.078 (ratio 1.0011) +score 8.974 / 8.977 (ratio 1.0003) +topk 8.250 / 8.195 (ratio 1.0195) +attn 83.227 / 83.190 (ratio 1.0016) +merge 6.898 / 5.492 (ratio 0.8041) +paired P=10180 nt=3 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 10.776 / 10.841 (ratio 1.0029) +score 12.898 / 12.935 (ratio 1.0089) +topk 8.289 / 8.297 (ratio 1.0155) +attn 125.417 / 125.102 (ratio 0.9965) +merge 16.367 / 17.435 (ratio 1.0652) +paired P=10180 nt=4 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 13.078 / 12.836 (ratio 0.9862) +score 16.836 / 16.828 (ratio 1.0015) +topk 7.891 / 8.435 (ratio 1.0693) +attn 162.078 / 161.914 (ratio 0.9995) +merge 23.547 / 22.120 (ratio 0.9428) +paired P=10180 nt=8 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 22.227 / 23.302 (ratio 0.9868) +score 48.805 / 48.820 (ratio 1.0008) +topk 10.539 / 10.547 (ratio 1.0007) +attn 324.810 / 324.659 (ratio 0.9989) +merge 46.128 / 43.958 (ratio 0.9549) +paired P=10180 nt=1984 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 989.500 / 995.594 (ratio 1.0028) +paired P=10180 nt=2048 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 1029.563 / 1032.083 (ratio 0.9913) diff --git a/results/round39/astra_r39_qsa/logs/merge1-ab-3000.log b/results/round39/astra_r39_qsa/logs/merge1-ab-3000.log new file mode 100644 index 0000000..0db910c --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/merge1-ab-3000.log @@ -0,0 +1,34 @@ +paired P=3000 nt=1 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 7.143 / 7.289 (ratio 1.0075) +score 6.844 / 6.831 (ratio 1.0000) +topk 15.458 / 15.424 (ratio 0.9971) +attn 77.891 / 77.836 (ratio 0.9994) +merge 8.078 / 5.477 (ratio 0.6779) +paired P=3000 nt=2 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 9.130 / 8.690 (ratio 0.9806) +score 3.432 / 3.430 (ratio 0.9977) +topk 6.430 / 6.664 (ratio 1.0320) +attn 83.362 / 83.508 (ratio 1.0019) +merge 6.698 / 5.456 (ratio 0.8110) +paired P=3000 nt=3 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 10.898 / 10.773 (ratio 0.9964) +score 4.508 / 4.565 (ratio 1.0173) +topk 6.599 / 6.578 (ratio 1.0145) +attn 124.635 / 124.799 (ratio 1.0022) +merge 16.437 / 17.208 (ratio 1.0483) +paired P=3000 nt=4 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 12.297 / 12.938 (ratio 1.0136) +score 5.734 / 5.758 (ratio 1.0054) +topk 6.440 / 6.669 (ratio 1.0348) +attn 160.836 / 160.805 (ratio 1.0007) +merge 23.490 / 22.161 (ratio 0.9434) +paired P=3000 nt=8 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 23.086 / 22.891 (ratio 1.0161) +score 10.453 / 10.437 (ratio 0.9983) +topk 6.581 / 6.625 (ratio 0.9996) +attn 315.937 / 316.089 (ratio 1.0007) +merge 46.117 / 43.852 (ratio 0.9509) +paired P=3000 nt=1984 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 971.292 / 984.625 (ratio 1.0209) +paired P=3000 nt=2048 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 1007.854 / 1018.385 (ratio 1.0104) diff --git a/results/round39/astra_r39_qsa/logs/merge1-ab-30000.log b/results/round39/astra_r39_qsa/logs/merge1-ab-30000.log new file mode 100644 index 0000000..eda6104 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/merge1-ab-30000.log @@ -0,0 +1,34 @@ +paired P=30000 nt=1 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 7.083 / 7.135 (ratio 1.0074) +score 59.336 / 60.596 (ratio 1.0212) +topk 65.310 / 65.651 (ratio 1.0048) +attn 109.383 / 109.302 (ratio 0.9993) +merge 11.102 / 7.503 (ratio 0.6758) +paired P=30000 nt=2 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 8.917 / 8.875 (ratio 1.0035) +score 109.680 / 110.154 (ratio 1.0061) +topk 65.401 / 65.898 (ratio 1.0070) +attn 196.049 / 195.747 (ratio 0.9981) +merge 6.961 / 5.258 (ratio 0.7553) +paired P=30000 nt=3 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 11.237 / 11.180 (ratio 0.9956) +score 83.232 / 83.201 (ratio 0.9997) +topk 33.846 / 33.930 (ratio 1.0018) +attn 126.208 / 126.227 (ratio 1.0002) +merge 16.417 / 17.440 (ratio 1.0588) +paired P=30000 nt=4 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 13.438 / 13.594 (ratio 1.0116) +score 211.352 / 212.031 (ratio 1.0040) +topk 64.773 / 44.307 (ratio 0.9977) +attn 165.656 / 165.573 (ratio 0.9917) +merge 23.529 / 22.190 (ratio 0.9434) +paired P=30000 nt=8 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 24.268 / 24.477 (ratio 0.9965) +score 362.250 / 320.018 (ratio 0.9490) +topk 27.164 / 27.169 (ratio 0.9988) +attn 336.633 / 334.083 (ratio 0.9992) +merge 45.935 / 44.151 (ratio 0.9621) +paired P=30000 nt=1984 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 1001.917 / 991.240 (ratio 1.0044) +paired P=30000 nt=2048 mrope=0 attn=1 (memcmp PASS), ref / candidate GPU us +prep 1038.167 / 1015.844 (ratio 1.0015) diff --git a/results/round39/astra_r39_qsa/logs/merge2-ab-10180.log b/results/round39/astra_r39_qsa/logs/merge2-ab-10180.log new file mode 100644 index 0000000..51edafd --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/merge2-ab-10180.log @@ -0,0 +1,34 @@ +paired P=10180 nt=1 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 7.039 / 7.109 (ratio 1.0100) +score 5.286 / 5.286 (ratio 0.9995) +topk 8.500 / 8.375 (ratio 0.9920) +attn 50.893 / 50.385 (ratio 0.9903) +merge 10.271 / 6.148 (ratio 0.5984) +paired P=10180 nt=2 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 8.870 / 9.016 (ratio 1.0068) +score 12.107 / 12.107 (ratio 0.9996) +topk 8.070 / 8.185 (ratio 1.0142) +attn 99.880 / 98.500 (ratio 0.9862) +merge 21.911 / 21.523 (ratio 0.9864) +paired P=10180 nt=3 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 10.698 / 10.724 (ratio 1.0046) +score 12.948 / 12.883 (ratio 0.9930) +topk 8.219 / 8.448 (ratio 1.0279) +attn 147.336 / 146.255 (ratio 0.9927) +merge 33.529 / 33.125 (ratio 0.9884) +paired P=10180 nt=4 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 13.107 / 12.729 (ratio 0.9971) +score 16.846 / 16.846 (ratio 1.0000) +topk 8.414 / 8.435 (ratio 1.0096) +attn 195.586 / 193.758 (ratio 0.9904) +merge 45.708 / 43.352 (ratio 0.9484) +paired P=10180 nt=8 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 22.922 / 23.500 (ratio 0.9804) +score 32.771 / 32.742 (ratio 1.0003) +topk 8.068 / 8.286 (ratio 1.0271) +attn 384.734 / 382.318 (ratio 0.9937) +merge 90.992 / 86.070 (ratio 0.9488) +paired P=10180 nt=1984 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 971.354 / 986.990 (ratio 1.0108) +paired P=10180 nt=2048 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 1026.281 / 1034.135 (ratio 1.0249) diff --git a/results/round39/astra_r39_qsa/logs/merge2-ab-3000.log b/results/round39/astra_r39_qsa/logs/merge2-ab-3000.log new file mode 100644 index 0000000..e30c137 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/merge2-ab-3000.log @@ -0,0 +1,34 @@ +paired P=3000 nt=1 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 7.125 / 7.190 (ratio 1.0044) +score 5.227 / 5.234 (ratio 1.0005) +topk 6.471 / 6.719 (ratio 1.0241) +attn 50.724 / 50.276 (ratio 0.9912) +merge 10.266 / 6.078 (ratio 0.5912) +paired P=3000 nt=2 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 9.253 / 8.867 (ratio 0.9550) +score 10.146 / 10.201 (ratio 0.9970) +topk 6.753 / 6.516 (ratio 0.9649) +attn 99.292 / 98.451 (ratio 0.9898) +merge 21.969 / 21.492 (ratio 0.9780) +paired P=3000 nt=3 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 10.687 / 10.841 (ratio 0.9985) +score 4.591 / 4.591 (ratio 0.9965) +topk 6.570 / 6.758 (ratio 1.0121) +attn 146.141 / 145.052 (ratio 0.9927) +merge 33.766 / 33.000 (ratio 0.9793) +paired P=3000 nt=4 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 12.849 / 13.031 (ratio 1.0246) +score 5.734 / 5.747 (ratio 1.0009) +topk 6.781 / 6.516 (ratio 0.9590) +attn 193.172 / 191.469 (ratio 0.9913) +merge 45.599 / 43.302 (ratio 0.9488) +paired P=3000 nt=8 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 23.305 / 22.266 (ratio 0.9843) +score 10.445 / 10.437 (ratio 1.0000) +topk 6.633 / 6.763 (ratio 1.0342) +attn 379.750 / 376.664 (ratio 0.9921) +merge 90.648 / 86.599 (ratio 0.9554) +paired P=3000 nt=1984 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 973.813 / 1011.719 (ratio 1.0276) +paired P=3000 nt=2048 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 1010.000 / 1014.750 (ratio 1.0047) diff --git a/results/round39/astra_r39_qsa/logs/merge2-ab-30000.log b/results/round39/astra_r39_qsa/logs/merge2-ab-30000.log new file mode 100644 index 0000000..bffb3bf --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/merge2-ab-30000.log @@ -0,0 +1,34 @@ +paired P=30000 nt=1 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 7.398 / 7.289 (ratio 0.9842) +score 25.516 / 25.477 (ratio 0.9996) +topk 14.164 / 14.010 (ratio 0.9775) +attn 50.865 / 50.469 (ratio 0.9917) +merge 10.125 / 6.055 (ratio 0.5957) +paired P=30000 nt=2 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 9.039 / 9.151 (ratio 1.0054) +score 41.872 / 27.826 (ratio 0.9976) +topk 14.268 / 14.091 (ratio 0.9838) +attn 99.729 / 98.557 (ratio 0.9885) +merge 22.018 / 21.464 (ratio 0.9762) +paired P=30000 nt=3 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 10.911 / 10.750 (ratio 0.9913) +score 36.117 / 36.112 (ratio 0.9999) +topk 14.060 / 14.445 (ratio 1.0262) +attn 148.052 / 146.667 (ratio 0.9899) +merge 33.581 / 33.560 (ratio 1.0008) +paired P=30000 nt=4 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 13.594 / 13.641 (ratio 1.0200) +score 208.391 / 208.427 (ratio 0.9988) +topk 33.667 / 33.841 (ratio 1.0042) +attn 196.599 / 194.807 (ratio 0.9905) +merge 45.617 / 43.651 (ratio 0.9558) +paired P=30000 nt=8 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 24.669 / 24.594 (ratio 1.0049) +score 117.820 / 112.247 (ratio 0.9995) +topk 14.414 / 14.281 (ratio 0.9880) +attn 388.141 / 386.224 (ratio 0.9954) +merge 90.964 / 86.977 (ratio 0.9556) +paired P=30000 nt=1984 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 995.479 / 992.313 (ratio 0.9886) +paired P=30000 nt=2048 mrope=0 attn=2 (memcmp PASS), ref / candidate GPU us +prep 1027.844 / 1014.938 (ratio 0.9874) diff --git a/results/round39/astra_r39_qsa/logs/merge2-bitwise.log b/results/round39/astra_r39_qsa/logs/merge2-bitwise.log new file mode 100644 index 0000000..be22578 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/merge2-bitwise.log @@ -0,0 +1,8 @@ +bitwise: 0 of 40 (pos0, ntok) cases differ PASS +device Apple M5: GPU us per dispatch, median of 9 paired rounds of 32 (ref / this) + pos0 512 ntok 1 attn 14.79 / 14.79 merge 3.71 / 2.83 + pos0 512 ntok 4 attn 50.16 / 50.18 merge 8.29 / 6.55 + pos0 2048 ntok 1 attn 50.46 / 50.53 merge 10.13 / 6.23 + pos0 2048 ntok 4 attn 190.94 / 190.92 merge 45.72 / 43.38 + pos0 8192 ntok 1 attn 50.71 / 51.02 merge 9.94 / 6.06 + pos0 8192 ntok 4 attn 193.76 / 193.71 merge 45.88 / 43.66 diff --git a/results/round39/astra_r39_qsa/logs/moe-default-T1.log b/results/round39/astra_r39_qsa/logs/moe-default-T1.log new file mode 100644 index 0000000..c483587 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/moe-default-T1.log @@ -0,0 +1,3 @@ +device Apple M5; Q2_0/Q2_0, T 1, 64 experts, tile TM 16 TN 64 TK 64 SG 2 +plan: 10 tiles of 16, 160 padded rows for 10 entries +parity: plan ok; y vs fp16-operand reference rel L2 1.75e-06, vs exact float reference 4.47e-04 PASS diff --git a/results/round39/astra_r39_qsa/logs/moe-dn-bench.log b/results/round39/astra_r39_qsa/logs/moe-dn-bench.log new file mode 100644 index 0000000..3a87592 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/moe-dn-bench.log @@ -0,0 +1,6 @@ +device Apple M5; IQ2_S/Q2_0, T 1, 512 experts, tile TM 16 TN 32 TK 32 SG 4 +right-input layout: runs of 4 k +plan: 10 tiles of 16, 160 padded rows for 10 entries +useful GFLOP per layer 0.1 (padded rows 160 = 16.00x entries) +ab bitwise (T=1): HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS +down A/B (T=1): A med=61.4 (min=58.0) | B med=61.7 (min=58.7) | diff med=+0.51% min=+1.15% diff --git a/results/round39/astra_r39_qsa/logs/moe-gu-bench.log b/results/round39/astra_r39_qsa/logs/moe-gu-bench.log new file mode 100644 index 0000000..b9c02a6 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/moe-gu-bench.log @@ -0,0 +1,6 @@ +device Apple M5; IQ2_S/Q2_0, T 1, 512 experts, tile TM 16 TN 32 TK 32 SG 4 +right-input layout: runs of 4 k +plan: 10 tiles of 16, 160 padded rows for 10 entries +useful GFLOP per layer 0.1 (padded rows 160 = 16.00x entries) +ab bitwise (T=1): HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS +gate/up A/B (T=1): A med=584.6 (min=550.0) | B med=548.5 (min=545.7) | diff med=-6.17% min=-0.78% diff --git a/results/round39/astra_r39_qsa/logs/moe-v3-T1.log b/results/round39/astra_r39_qsa/logs/moe-v3-T1.log new file mode 100644 index 0000000..0f5dd63 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/moe-v3-T1.log @@ -0,0 +1,4 @@ +device Apple M5; IQ2_S/Q2_0, T 1, 64 experts, tile TM 16 TN 32 TK 32 SG 4 +right-input layout: runs of 4 k +plan: 10 tiles of 16, 160 padded rows for 10 entries +parity: plan ok; y vs fp16-operand reference rel L2 1.62e-05, vs exact float reference 4.61e-04 PASS diff --git a/results/round39/astra_r39_qsa/logs/moe-v3-ab.log b/results/round39/astra_r39_qsa/logs/moe-v3-ab.log new file mode 100644 index 0000000..f67dd9d --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/moe-v3-ab.log @@ -0,0 +1,4 @@ +device Apple M5; IQ2_S/Q2_0, T 256, 64 experts, tile TM 16 TN 32 TK 32 SG 4 +right-input layout: runs of 4 k +plan: 190 tiles of 16, 3040 padded rows for 2560 entries +ab bitwise (T=256): HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS diff --git a/results/round39/astra_r39_qsa/logs/prompt-flash-0.log b/results/round39/astra_r39_qsa/logs/prompt-flash-0.log new file mode 100644 index 0000000..9ef9979 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/prompt-flash-0.log @@ -0,0 +1,4 @@ +device Apple M5; T 96 pos0 0: memcmp PASS; flash x1 before median 0.115 min 0.093 ms, after median 0.115 min 0.094 ms; tg 13760 -> 13760 bytes +pair 0 before 0.093292 after 0.093583 ms +pair 1 before 0.115000 after 0.117458 ms +pair 2 before 0.115125 after 0.114875 ms diff --git a/results/round39/astra_r39_qsa/logs/prompt-flash-30000.log b/results/round39/astra_r39_qsa/logs/prompt-flash-30000.log new file mode 100644 index 0000000..6a45d52 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/prompt-flash-30000.log @@ -0,0 +1,4 @@ +device Apple M5; T 96 pos0 30000: memcmp PASS; flash x1 before median 2.836 min 2.637 ms, after median 2.836 min 2.386 ms; tg 13760 -> 13760 bytes +pair 0 before 2.864583 after 2.385708 ms +pair 1 before 2.637458 after 2.835625 ms +pair 2 before 2.836375 after 2.846625 ms diff --git a/results/round39/astra_r39_qsa/logs/prompt-flash-k4off-0.log b/results/round39/astra_r39_qsa/logs/prompt-flash-k4off-0.log new file mode 100644 index 0000000..ab91631 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/prompt-flash-k4off-0.log @@ -0,0 +1,4 @@ +device Apple M5; T 96 pos0 0: memcmp PASS; flash x1 before median 0.805 min 0.803 ms, after median 0.379 min 0.378 ms; tg 13760 -> 13760 bytes +pair 0 before 0.803083 after 0.377625 ms +pair 1 before 0.805250 after 0.382875 ms +pair 2 before 0.804875 after 0.379292 ms diff --git a/results/round39/astra_r39_qsa/logs/prompt-flash-k4off-30000.log b/results/round39/astra_r39_qsa/logs/prompt-flash-k4off-30000.log new file mode 100644 index 0000000..9a50e92 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/prompt-flash-k4off-30000.log @@ -0,0 +1,4 @@ +device Apple M5; T 96 pos0 30000: memcmp PASS; flash x1 before median 5.685 min 5.442 ms, after median 1.828 min 1.808 ms; tg 13760 -> 13760 bytes +pair 0 before 5.731000 after 1.807667 ms +pair 1 before 5.684625 after 1.835375 ms +pair 2 before 5.442500 after 1.827583 ms diff --git a/results/round39/astra_r39_qsa/logs/prompt-mm-0.log b/results/round39/astra_r39_qsa/logs/prompt-mm-0.log new file mode 100644 index 0000000..8996140 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/prompt-mm-0.log @@ -0,0 +1,7 @@ +device Apple M5 coop_k4 1 +qsa_attn_mm T 96 pos0 0 memcmp PASS before 0.301000 after 0.295875 ms tg 7424 -> 7424 +pair 0 before 0.303667 after 0.295875 ms +pair 1 before 0.286625 after 0.291125 ms +pair 2 before 0.301000 after 0.296125 ms +post-timing memcmp PASS +merge+prepared memcmp PASS diff --git a/results/round39/astra_r39_qsa/logs/prompt-mm-30000.log b/results/round39/astra_r39_qsa/logs/prompt-mm-30000.log new file mode 100644 index 0000000..c3accd6 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/prompt-mm-30000.log @@ -0,0 +1,7 @@ +device Apple M5 coop_k4 1 +qsa_attn_mm T 96 pos0 30000 memcmp PASS before 1.421292 after 1.415875 ms tg 7424 -> 7424 +pair 0 before 1.421292 after 1.426250 ms +pair 1 before 1.422000 after 1.415625 ms +pair 2 before 1.411750 after 1.415875 ms +post-timing memcmp PASS +merge+prepared memcmp PASS diff --git a/results/round39/astra_r39_qsa/logs/prompt-mm-k4off-0.log b/results/round39/astra_r39_qsa/logs/prompt-mm-k4off-0.log new file mode 100644 index 0000000..4c5cda0 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/prompt-mm-k4off-0.log @@ -0,0 +1,7 @@ +device Apple M5 coop_k4 0 +qsa_attn_mm T 96 pos0 0 memcmp PASS before 0.866917 after 0.350000 ms tg 6912 -> 6912 +pair 0 before 0.871917 after 0.350000 ms +pair 1 before 0.855875 after 0.350750 ms +pair 2 before 0.866917 after 0.321417 ms +post-timing memcmp PASS +merge+prepared memcmp PASS diff --git a/results/round39/astra_r39_qsa/logs/prompt-mm-k4off-30000.log b/results/round39/astra_r39_qsa/logs/prompt-mm-k4off-30000.log new file mode 100644 index 0000000..6c56df3 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/prompt-mm-k4off-30000.log @@ -0,0 +1,7 @@ +device Apple M5 coop_k4 0 +qsa_attn_mm T 96 pos0 30000 memcmp PASS before 4.954000 after 1.536500 ms tg 6912 -> 6912 +pair 0 before 4.954000 after 1.535708 ms +pair 1 before 4.955250 after 1.536500 ms +pair 2 before 4.953292 after 1.542250 ms +post-timing memcmp PASS +merge+prepared memcmp PASS diff --git a/results/round39/astra_r39_qsa/logs/score-ab-3000.log b/results/round39/astra_r39_qsa/logs/score-ab-3000.log new file mode 100644 index 0000000..d3e8917 --- /dev/null +++ b/results/round39/astra_r39_qsa/logs/score-ab-3000.log @@ -0,0 +1 @@ +scores FAIL diff --git a/results/round39/astra_r39_qsa/metadata.json b/results/round39/astra_r39_qsa/metadata.json new file mode 100644 index 0000000..49309d3 --- /dev/null +++ b/results/round39/astra_r39_qsa/metadata.json @@ -0,0 +1,488 @@ +{ + "baseline_git_head": "296b59d2f73c8f04a805df29c2fa97655adace78", + "baseline_working_files_sha256": { + "metal/kernels/qsa.metal": "d2d3f311030eae507ab83148afe267b7746a15ac2231aa50aa0a22cc4403db4c", + "metal/kernels/qsa2.metal": "1e1937f95e2224ce09944c8cb549af82edc02a0ab38a1cceeb203f4f7b58d59b", + "metal/kernels/qsa_prompt.metal": "9c141a1b65d5c2fa8bbff84027cd8757ed85cc1d72e80143cc94fe7255414e26", + "metal/tests/moe_mm_parity.cpp": "692d1eea87599cbe5c0659a8052bb3416fd10dd6b695fb4f4ba07332e76bf3a0" + }, + "system": "ProductName:\t\tmacOS\nProductVersion:\t\t27.0\nBuildVersion:\t\t26A428\n", + "device": "Apple M5", + "test_jobs": [ + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "3000" + ], + "exit_code": 0, + "duration": 14.704230582981836, + "timed_out": false + }, + { + "argv": [ + "./build-metal/moe_mm_parity" + ], + "exit_code": 3, + "duration": 4.560703166993335, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "3000" + ], + "exit_code": 1, + "duration": 1.7547417919849977, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "3000" + ], + "exit_code": 0, + "duration": 5.871507915988332, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "10180" + ], + "exit_code": 0, + "duration": 6.856663957994897, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "30000" + ], + "exit_code": 0, + "duration": 11.867526749992976, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "3000" + ], + "exit_code": 0, + "duration": 3.023166750004748, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "10180" + ], + "exit_code": 0, + "duration": 5.168871124973521, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "30000" + ], + "exit_code": 0, + "duration": 8.313932707998902, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "3000" + ], + "exit_code": 0, + "duration": 3.2664715830178466, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "10180" + ], + "exit_code": 0, + "duration": 4.757935791014461, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "30000" + ], + "exit_code": 0, + "duration": 7.467346207995433, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "3000" + ], + "exit_code": 0, + "duration": 2.6887030830257572, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "10180" + ], + "exit_code": 0, + "duration": 4.4522730839962605, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "30000" + ], + "exit_code": 0, + "duration": 6.560153625003295, + "timed_out": false + }, + { + "argv": [ + "./build-metal/moe_mm_parity" + ], + "exit_code": -15, + "duration": 15.026818915997865, + "timed_out": true + }, + { + "argv": [ + "./build-metal/moe_mm_parity", + "--v3", + "--gu", + "IQ2_S", + "--dn", + "Q2_0", + "--tile", + "16,32,32,4", + "--dtile", + "64,32,2" + ], + "exit_code": -15, + "duration": 20.02672729198821, + "timed_out": true + }, + { + "argv": [ + "./build-metal/qsa_prompt_parity" + ], + "exit_code": 0, + "duration": 0.9771465419908054, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_prompt_parity", + "--mm-compare", + "attn", + "96", + "0", + "/tmp/astra_r39_qsa/baseline/qsa_prompt.metal", + "3" + ], + "exit_code": 0, + "duration": 0.4449748329934664, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_prompt_parity", + "--flash-compare", + "96", + "0", + "/tmp/astra_r39_qsa/baseline/qsa_prompt.metal", + "1", + "3" + ], + "exit_code": 0, + "duration": 0.035576458991272375, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_prompt_parity", + "--mm-compare", + "attn", + "96", + "30000", + "/tmp/astra_r39_qsa/baseline/qsa_prompt.metal", + "3" + ], + "exit_code": 0, + "duration": 0.1662420420034323, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_prompt_parity", + "--flash-compare", + "96", + "30000", + "/tmp/astra_r39_qsa/baseline/qsa_prompt.metal", + "1", + "3" + ], + "exit_code": 0, + "duration": 0.11264241699245758, + "timed_out": false + }, + { + "argv": [ + "./build-metal/moe_mm_parity", + "--v3", + "--gu", + "IQ2_S", + "--dn", + "Q2_0", + "--tile", + "16,32,32,4", + "--dtile", + "64,32,2", + "--ab", + "/tmp/astra_r39_qsa/baseline" + ], + "exit_code": 0, + "duration": 2.633624290989246, + "timed_out": false + }, + { + "argv": [ + "./build-metal/moe_mm_parity", + "--T", + "1" + ], + "exit_code": 0, + "duration": 4.224036582978442, + "timed_out": false + }, + { + "argv": [ + "./build-metal/moe_mm_parity", + "--v3", + "--gu", + "IQ2_S", + "--dn", + "Q2_0", + "--tile", + "16,32,32,4", + "--dtile", + "64,32,2", + "--T", + "1" + ], + "exit_code": 0, + "duration": 0.1989670839975588, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--bitwise", + "/tmp/astra_r39_qsa/baseline/qsa2.metal" + ], + "exit_code": 0, + "duration": 1.3079375409870408, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_prompt_parity", + "--mm-compare", + "attn", + "96", + "0", + "/tmp/astra_r39_qsa/baseline/qsa_prompt.metal", + "3" + ], + "exit_code": 0, + "duration": 0.25848625000799075, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_prompt_parity", + "--flash-compare", + "96", + "0", + "/tmp/astra_r39_qsa/baseline/qsa_prompt.metal", + "1", + "3" + ], + "exit_code": 0, + "duration": 0.31251912500010803, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_prompt_parity", + "--mm-compare", + "attn", + "96", + "30000", + "/tmp/astra_r39_qsa/baseline/qsa_prompt.metal", + "3" + ], + "exit_code": 0, + "duration": 0.18658049998339266, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_prompt_parity", + "--flash-compare", + "96", + "30000", + "/tmp/astra_r39_qsa/baseline/qsa_prompt.metal", + "1", + "3" + ], + "exit_code": 0, + "duration": 0.131041875021765, + "timed_out": false + }, + { + "argv": [ + "./build-metal/moe_mm_parity", + "--v3", + "--gu", + "IQ2_S", + "--dn", + "Q2_0", + "--tile", + "16,32,32,4", + "--dtile", + "64,32,2", + "--T", + "1", + "--gu-bench", + "--ab", + "/tmp/astra_r39_qsa/baseline" + ], + "exit_code": 0, + "duration": 0.34804070901009254, + "timed_out": false + }, + { + "argv": [ + "./build-metal/moe_mm_parity", + "--v3", + "--gu", + "IQ2_S", + "--dn", + "Q2_0", + "--tile", + "16,32,32,4", + "--dtile", + "64,32,2", + "--T", + "1", + "--dn-bench", + "--ab", + "/tmp/astra_r39_qsa/baseline" + ], + "exit_code": 0, + "duration": 0.26160704100038856, + "timed_out": false + }, + { + "argv": [ + "./build-metal/moe_mm_parity" + ], + "exit_code": 0, + "duration": 23.961265375022776, + "timed_out": false + }, + { + "argv": [ + "./build-metal/moe_mm_parity", + "--v3", + "--gu", + "IQ2_S", + "--dn", + "Q2_0", + "--tile", + "16,32,32,4", + "--dtile", + "64,32,2" + ], + "exit_code": 0, + "duration": 20.010784165991936, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "3000" + ], + "exit_code": 0, + "duration": 8.637194999988424, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "10180" + ], + "exit_code": 0, + "duration": 7.093214749998879, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "30000" + ], + "exit_code": 0, + "duration": 9.387072458019247, + "timed_out": false + }, + { + "argv": [ + "./build-metal/qsa_prompt_parity" + ], + "exit_code": 0, + "duration": 0.36627699999371544, + "timed_out": false + } + ] +} \ No newline at end of file diff --git a/results/round39/astra_r39_qsa/plan.md b/results/round39/astra_r39_qsa/plan.md new file mode 100644 index 0000000..2a79119 --- /dev/null +++ b/results/round39/astra_r39_qsa/plan.md @@ -0,0 +1 @@ +owner: metal qsa kernels and moe parity harness only. preserve layouts, dispatch contracts, flags, fma order. baseline: copied working tree. test paired memcmp and 12-round median A/B at P=3000,10180,30000; reject any mismatch or repeatable slowdown. GPU tests serialized, each <=15s (cold compile <=20s). no server/full-model access. two confirmation sweeps for neutral/faster candidates. prompt/moe references are tolerance checks, not bitwise checks. diff --git a/results/round39/astra_r39_qsa/rejected-merge1.metal b/results/round39/astra_r39_qsa/rejected-merge1.metal new file mode 100644 index 0000000..a843413 --- /dev/null +++ b/results/round39/astra_r39_qsa/rejected-merge1.metal @@ -0,0 +1,500 @@ +// metal/kernels/qsa.metal - QSA (sparse full attention) decode for a window of tokens, between the projections. +// Contract of include/strata/kernels/qsa.hpp / qsa_select.hpp / qsa_decode_attn.hpp / kv_q8.hpp (the engine's +// native path), in a Metal layout (no paging: unified memory): +// +// qsa_prep per token t (cell pos0 + t): q heads rms_norm * q_norm, NEOX RoPE on 64 dims -> qcur; +// k rms_norm * k_norm, RoPE -> int8 KV append (fp16 scale per 64 = fp16(amax/127), rint, clamp); +// v -> int8 append; indexer q rms_norm * iq_norm, RoPE -> qidx. The indexer KEYS (last +// threadgroup, all window tokens in order): raw tail [cell % 4]; on a block completion +// pooled[b] = rope(rms_norm(mean of the block's 4 raw keys) * ik_norm, pos of the block's first cell); +// cell 0 also sets dead = rms_norm(raw) * ik_norm (the tail block's key). +// qsa_score per token, block b <= n_bid: sum over the 4 indexer heads of relu(qidx . key_b); key_b = pooled[b], +// the tail block n_bid scores `dead` and gets +1e9 when it has cells. +// qsa_topk per token: width = min(n_kv, 2051) cells; n_kv <= 2051: all; else the width cells of largest +// score (a cell inherits its block's), ties to the lowest cell, emitted ascending (radix select). +// qsa_attn split-K: (chunk of CH cells, kv head, token): 12 query heads; s = q.k / 16, chunk max m, sum l, +// acc = p . v -> partials +// qsa_merge per (head, token): log-sum-exp merge, * sigmoid(gate), written as prepared activations (xprep +// planes) for attn_output. +// +// RoPE: cos/sin from a host table [pos][32] (double precision, as build_rope_table). + +#define QH 24 +#define QKV 2 +#define QD 256 +#define QROT 64 +#define IH 4 +#define ID 128 +#define QBUDGET 2051 +#define QCH 64 + +constant int QSA_NTOK [[function_constant(3)]]; +// QSA_MROPE (function constant 24): interleaved M-RoPE (the vision path; strata/kernels/mrope.hpp, ggml IMROPE, +// sections 11/11/10): buffer 20 of qsa_prep is a table int32 [cells][3] (t, h, w), and rotary pair i (0..31) of the row +// at (pos_base + cell) rotates by the cos/sin table's row tab[(pos_base + cell) * 3 + i % 3] instead of its own - +// the pooled indexer keys too (position pos_base + 4 b of a block). Off: exactly the plain RoPE. +constant bool QSA_MROPE [[function_constant(24)]]; +constant bool qsa_mrope = is_function_constant_defined(QSA_MROPE) && QSA_MROPE; +// QSA_Q_F16_FC (function constant 25): qsa_prep writes qcur directly as fp16 (device half*) for qsa_flash, skipping qsa_q_f16. +constant bool QSA_Q_F16_FC [[function_constant(25)]]; +constant bool qsa_qf16 = is_function_constant_defined(QSA_Q_F16_FC) ? QSA_Q_F16_FC : false; +// QSA_ROW0 (function constant 28): qsa_prep's threadgroup x is row QSA_ROW0 + x (24: the K/V rows only, for rows whose +// queries nobody reads). Undefined: 0. +constant int QSA_ROW0 [[function_constant(28)]]; +constant int qsa_row0 = is_function_constant_defined(QSA_ROW0) ? QSA_ROW0 : 0; + +struct QsaArgs { + uint pos0; // the window's first cell (relative to the sequence) + uint pos_base; // absolute position of cell 0 (RoPE) + uint max_cells; + uint n_tok; + float eps; +}; + +// cache layout per layer +// kq/vq [cell][kv head][256] int8, ks/vs [cell][kv head][4] half +// tail [4][128] f32 (slot = cell % 4), dead [128], pooled [max_cells / 4 + 1][128] + +static inline float tg_sum256(float v, threadgroup float* part, ushort sg, ushort lane, ushort nsg) { + v = simd_sum(v); + if (lane == 0) part[sg] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + float t = 0.0f; + for (int i = 0; i < nsg; ++i) t += part[i]; + return t; +} + +// one row (width w <= 256 = threads) : out = rope((rs * x) * g) (RoPE on the first 64 dims, NEOX pairs i, i+32) +static inline float norm_rope(float x, float xhi_unused, device const float* g, threadgroup float* row, + threadgroup float* part, ushort tid, ushort sg, ushort lane, int w, float eps, + device const float2* cs, uint absp, bool rope, device const int* mtab, + bool reuse = false) { + const float ss = tg_sum256(tid < w ? x * x : 0.0f, part, sg, lane, 8); + const float rs = rsqrt(ss / float(w) + eps); + const float y = tid < w ? (rs * x) * g[tid] : 0.0f; + float o = y; + if (rope) { + if (tid < QROT) row[tid] = y; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < QROT) { + const int p = tid & 31; + const uint pp = qsa_mrope ? uint(mtab[absp * 3 + uint(p % 3)]) : absp; + const float2 c = cs[pp * 32 + p]; + const float a = row[p], b = row[p + 32]; + o = tid < 32 ? a * c.x - b * c.y : a * c.y + b * c.x; + } + } + if (reuse) threadgroup_barrier(mem_flags::mem_threadgroup); + (void)xhi_unused; + return o; +} + +// QSA_KV_STRIDE (function constant 26): kcur, vcur, iqraw and ikraw are column slices of one [n_tok][QSA_KV_STRIDE] +// projection (prompt_chunk's fused K | V | indexer K | indexer Q GEMM, 1152 or 1664 columns). Undefined: their own +// [n_tok][rows] buffers. +constant int QSA_KV_STRIDE [[function_constant(26)]]; +constant int qsa_kv_s = is_function_constant_defined(QSA_KV_STRIDE) ? QSA_KV_STRIDE : 0; + +// grid (32 rows + 1, n_tok): rows 0..23 q heads, 24..25 k heads, 26..27 v heads, 28..31 indexer q heads; +// threadgroup 32 (x = 32, y = 0) is the indexer-key threadgroup. 256 threads. A grid of 28 (or of 4 from QSA_ROW0 = +// 24) writes no indexer state: the MTP layer's attention is a dense window, never scored. +kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float* kcur [[buffer(1)]], + device const float* vcur [[buffer(2)]], device const float* iqraw [[buffer(3)]], + device const float* ikraw [[buffer(4)]], device const float* qn [[buffer(5)]], + device const float* kn [[buffer(6)]], device const float* iqn [[buffer(7)]], + device const float* ikn [[buffer(8)]], device const float2* cs [[buffer(9)]], + device float* qcur [[buffer(10)]], device float* qidx [[buffer(11)]], + device char* kq [[buffer(12)]], device char* vq [[buffer(13)]], device half* ksc [[buffer(14)]], + device half* vsc [[buffer(15)]], device float* tail [[buffer(16)]], device float* dead [[buffer(17)]], + device float* pooled [[buffer(18)]], constant QsaArgs& a [[buffer(19)]], + device const int* mtab [[buffer(20)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + threadgroup float row[QROT]; + threadgroup float amax[8]; + const int r = int(tg.x) + qsa_row0; + // the row strides of kcur / vcur, ikraw, iqraw + const int s_kv = qsa_kv_s ? qsa_kv_s : QKV * QD, s_ik = qsa_kv_s ? qsa_kv_s : ID; + const int s_iq = qsa_kv_s ? qsa_kv_s : IH * ID; + if (r == 32) { + // the indexer keys: threadgroup (32, 0) handles t < 3 (boundary with prior tail + cell 0 dead) and final tail; + // every threadgroup (32, t >= 3) with (cell & 3) == 3 pools its 4 in-window keys in parallel. + const int t_self = int(tg.y); + if (t_self >= 3) { + const uint cell = a.pos0 + uint(t_self); + if ((cell & 3) == 3) { + float m = 0.0f; + if (tid < ID) { + m = (((ikraw[(t_self - 3) * s_ik + tid] + ikraw[(t_self - 2) * s_ik + tid]) + + ikraw[(t_self - 1) * s_ik + tid]) + ikraw[t_self * s_ik + tid]) * 0.25f; + } + const uint b = cell >> 2; + const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab); + if (tid < ID) pooled[b * ID + tid] = o; + } + return; + } + if (t_self != 0) return; + const int n_pre = min(int(a.n_tok), 3); + for (int t = 0; t < n_pre; ++t) { + const uint cell = a.pos0 + uint(t); + const float x = tid < ID ? ikraw[t * s_ik + tid] : 0.0f; + if (tid < ID) tail[(cell & 3) * ID + tid] = x; + threadgroup_barrier(mem_flags::mem_device); + if (cell == 0) { + const float o = norm_rope(x, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, 0, false, mtab, true); + if (tid < ID) dead[tid] = o; + } + if ((cell & 3) == 3) { + float m = 0.0f; + if (tid < ID) m = (((tail[tid] + tail[ID + tid]) + tail[2 * ID + tid]) + tail[3 * ID + tid]) * 0.25f; + const uint b = cell >> 2; + const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab, true); + if (tid < ID) pooled[b * ID + tid] = o; + } + threadgroup_barrier(mem_flags::mem_device); + } + const int t_tail0 = max(n_pre, int(a.n_tok) - 4); + for (int t = t_tail0; t < int(a.n_tok); ++t) { + const uint cell = a.pos0 + uint(t); + if (tid < ID) tail[(cell & 3) * ID + tid] = ikraw[t * s_ik + tid]; + } + return; + } + const int t = int(tg.y); + const uint cell = a.pos0 + uint(t), absp = a.pos_base + cell; + if (r < QH) { + const float x = qfull[(t * QH + r) * 2 * QD + tid]; + const float o = norm_rope(x, 0.0f, qn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); + if (qsa_qf16) ((device half*)qcur)[(t * QH + r) * QD + tid] = half(o); + else qcur[(t * QH + r) * QD + tid] = o; + } else if (r < QH + 2 * QKV) { + const bool isv = r >= QH + QKV; + const int h = r - QH - (isv ? QKV : 0); + float x = (isv ? vcur : kcur)[t * s_kv + h * QD + tid]; + if (!isv) x = norm_rope(x, 0.0f, kn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); + // int8 with an fp16 scale per 64 values: groups = simdgroup pairs (64 threads) + float m = simd_max(fabs(x)); + if (lane == 0) amax[sg] = m; + threadgroup_barrier(mem_flags::mem_threadgroup); + const int g = tid >> 6; + m = max(amax[2 * g], amax[2 * g + 1]); + const half sh = half(m / 127.0f); + const float sf = float(sh); + int q = 0; + if (sf > 0.0f) q = clamp(int(rint(x / sf)), -127, 127); + const uint rowi = cell * QKV + uint(h); + (isv ? vq : kq)[rowi * QD + tid] = char(q); + if ((tid & 63) == 0) (isv ? vsc : ksc)[rowi * 4 + g] = sh; + } else { + if (a.pos0 + a.n_tok <= uint(QBUDGET)) return; + const int h = r - QH - 2 * QKV; + const float x = tid < ID ? iqraw[t * s_iq + 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); + if (tid < ID) qidx[(t * IH + h) * ID + tid] = o; + } +} + +// scores [n_tok][max_blocks]; grid (ceil(max_blocks / 8), n_tok), 256 threads: a simdgroup per block. +kernel void qsa_score(device const float* qidx [[buffer(0)]], device const float* pooled [[buffer(1)]], + device const float* dead [[buffer(2)]], device float* scores [[buffer(3)]], + constant QsaArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + const int t = int(tg.y); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + const uint b = tg.x * 8 + sg; + if (b > n_bid) return; + device const float* key = b == n_bid ? dead : pooled + b * ID; + device const float* qp = qidx + (t * IH) * ID + lane; + float k[4], q[IH][4]; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) k[i] = key[lane + 32 * i]; + #pragma clang loop unroll(full) + for (int h = 0; h < IH; ++h) { + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) q[h][i] = qp[h * ID + 32 * i]; + } + float s = 0.0f; + #pragma clang loop unroll(full) + for (int h = 0; h < IH; ++h) { + float d = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) d = fma(q[h][i], k[i], d); + d = simd_sum(d); + s += max(d, 0.0f); + } + if (b == n_bid && (n_kv & 3) != 0) s += 1e9f; + if (lane == 0) scores[uint(t) * (a.max_cells / 4 + 2) + b] = s; +} + +// ids [n_tok][QBUDGET], ascending. One threadgroup of 1024 threads per token. Order key of a float: larger score = +// larger key; a radix select (4 passes of 8 bits) over the blocks with weights = their cell counts finds the +// threshold key T and how many cells with key == T to take (the lowest cells first). +static inline uint okey(float f) { + const uint u = as_type(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]], + constant QsaArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint hist[256]; + threadgroup uint sel[4]; + threadgroup uint2 pref[32]; + threadgroup uint2 totals; + const int t = int(tg.x); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2, nb = n_bid + 1; + device int* out = ids + t * QBUDGET; + device const float* sc = scores + uint(t) * (a.max_cells / 4 + 2); + if (n_kv <= QBUDGET) { + for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); + return; + } + const bool cached = nb <= 4096; + uint keys[4]; + if (cached) + for (uint b = tid, i = 0; b < nb; b += 1024, ++i) keys[i] = okey(sc[b]); + // radix select of the `need`-th largest cell key + uint prefix = 0, pmask = 0, need = QBUDGET; + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint b = tid; b < nb; b += 1024) { + const uint k = cached ? keys[b / 1024] : okey(sc[b]); + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + uint counts[8], sum = 0; + for (uint j = 0; j < 8; ++j) { + counts[j] = atomic_load_explicit(&hist[255 - (lane * 8 + j)], memory_order_relaxed); + sum += counts[j]; + } + uint acc = simd_prefix_exclusive_sum(sum); + if (acc < need && acc + sum >= need) { + for (uint j = 0; j < 8; ++j) { + if (acc + counts[j] >= need) { + sel[0] = 255 - (lane * 8 + j); + sel[1] = need - acc; + break; + } + acc += counts[j]; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; + } + // T = prefix; take every cell with key > T and the first `need` cells (lowest index) with key == T, ascending. + // Pass over blocks in order of index with a prefix scan of each block's selected cell count. + const uint T = prefix; + uint base = 0, eq_taken = 0; + for (uint b0 = 0; b0 < nb; b0 += 1024) { + const uint b = b0 + tid; + uint w = 0, weq = 0; + if (b < nb) { + const uint k = cached ? keys[b / 1024] : okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + if (k > T) w = cells; + else if (k == T) weq = cells; + } + const uint eq_before = simd_prefix_exclusive_sum(weq); + const uint gt_before = simd_prefix_exclusive_sum(w); + if (lane == 31) pref[sg] = uint2(eq_before + weq, gt_before + w); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + const uint2 count = pref[lane]; + const uint2 before = uint2(simd_prefix_exclusive_sum(count.x), + simd_prefix_exclusive_sum(count.y)); + pref[lane] = before; + if (lane == 31) totals = before + count; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint2 group_before = pref[sg], count = totals; + simdgroup_barrier(mem_flags::mem_threadgroup); + const uint eq0 = eq_taken + group_before.x + eq_before; + const uint take_eq = eq0 < need ? min(weq, need - eq0) : 0u; + const uint mine = w + take_eq; + // Clipping the equal-key prefix preserves the lowest-cell tie break. + const uint remaining = need - min(need, eq_taken); + const uint off = base + group_before.y + gt_before + + min(remaining, group_before.x + eq_before); + for (uint c = 0; c < mine; ++c) out[off + c] = int(4 * b + c); + base += count.y + min(remaining, count.x); + eq_taken += count.x; + } +} + +// partials: [n_tok][chunks][QH][QD + 2] (acc, then m, l). grid (chunks, QKV, n_tok), 256 threads. +kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + 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 HPG = QH / QKV; // 12 query heads per kv head + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH]; + threadgroup int cid[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; + const int nchunks = (QBUDGET + QCH - 1) / QCH; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + if (uint(c0) >= width) { + // an empty chunk: m = -inf, l = 0 (the merge skips it) + for (int h = 0; h < HPG; ++h) + if (tid == 0) { pp[(g * HPG + h) * (QD + 2) + QD] = -INFINITY; pp[(g * HPG + h) * (QD + 2) + QD + 1] = 0.0f; } + return; + } + const int nc = min(QCH, int(width) - c0); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + g * HPG + h) * QD + tid]; + if (tid < QCH) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + // scores: 12 x 64 dots of 256; simdgroup sg takes cells sg, sg + 8, ...; lanes split the 256 dims (8 each) + threadgroup const float* q_lane = &qs[0][lane * 8]; + for (int c = sg; c < nc; c += 8) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char* kr = kq + rowi * QD + lane * 8; + const char4 k0 = *(device const char4*)(kr); + const char4 k1 = *(device const char4*)(kr + 4); + const float scl = float(ksc[rowi * 4 + (lane >> 3)]); + const float kv[8] = { + float(k0.x) * scl, float(k0.y) * scl, float(k0.z) * scl, float(k0.w) * scl, + float(k1.x) * scl, float(k1.y) * scl, float(k1.z) * scl, float(k1.w) * scl + }; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + threadgroup const float* qh = q_lane + h * QD; + float d = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) d = fma(qh[i], kv[i], d); + d = simd_sum(d); + if (lane == 0) s[h][c] = d * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // per head: max, exp, sum (simdgroup h % 8 handles heads h, h + 8) + threadgroup float mh[HPG], lh[HPG]; + // nc <= QCH = 64: a lane holds cells lane and lane + 32 in registers (the same max and sum as a strided loop) + const bool in0 = int(lane) < nc, in1 = int(lane) + 32 < nc; + for (int h = sg; h < HPG; h += 8) { + const float s0 = in0 ? s[h][lane] : -INFINITY, s1 = in1 ? s[h][lane + 32] : -INFINITY; + const float m = simd_max(max(s0, s1)); + const float p0 = in0 ? exp(s0 - m) : 0.0f, p1 = in1 ? exp(s1 - m) : 0.0f; + if (in0) s[h][lane] = p0; + if (in1) s[h][lane + 32] = p1; + const float l = simd_sum(p0 + p1); + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // acc[h][j] = sum_c p[h][c] v[c][j]: thread = dim j + float acc[HPG]; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + // four cells' V loads in flight at once; the fma order per head stays c, c + 1, ... + auto vcell = [&](int c) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + return float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); + }; + int c = 0; + for (; c + 4 <= nc; c += 4) { + const float v0 = vcell(c), v1 = vcell(c + 1), v2 = vcell(c + 2), v3 = vcell(c + 3); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + float x = fma(s[h][c], v0, acc[h]); + x = fma(s[h][c + 1], v1, x); + x = fma(s[h][c + 2], v2, x); + acc[h] = fma(s[h][c + 3], v3, x); + } + } + for (; c < nc; ++c) { + const float v = vcell(c); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][c], v, acc[h]); + } + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (g * HPG + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} + +constant float qsa_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; +constant float qsa_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + +static inline void qsa_write_prepared(float y, device float* xp, constant QsaArgs& a, + int t, int h, ushort tid) { + // prepared planes (xprep layout) of the n = QH * QD vector + const uint n = QH * QD, plane = a.n_tok * n, d = uint(h) * QD + tid; + const int j = tid & 15; + device float* raw = xp + uint(t) * n + d; + raw[0] = y; + raw[plane] = y * qsa_q2s[j & 7]; + raw[2 * plane] = y * qsa_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? y * 0.00390625f : y; + float s0 = 0.0f, s1 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(y, ushort(k)); + s1 += simd_shuffle_down(y, ushort(k + 1)); + } + if (j == 0) xp[4 * plane + uint(t) * (n / 16) + d / 16] = s0 + s1; +} + +// grid (QH, n_tok), 256 threads: out = merge * sigmoid(gate), as prepared activations of n = QH * QD per token +kernel void qsa_merge(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], + device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const int h = int(tg.x), t = int(tg.y); + const int nchunks = (QBUDGET + QCH - 1) / QCH; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nc = int((width + QCH - 1) / QCH); + constexpr uint STRIDE = QH * (QD + 2); + device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2); + // the chunks' (m, l) and weights exp(m - M) once per threadgroup, not once per thread (nc <= 33 < 64) + threadgroup float mk[64], lk[64], ek[64], L_sh; + if (tid < uint(nc)) { + const float2 ml = *(device const float2*)(pp + tid * STRIDE + QD); + mk[tid] = ml.x; + lk[tid] = ml.y; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float M = -INFINITY; + for (int k = tid & 31; k < nc; k += 32) M = max(M, mk[k]); + M = simd_max(M); + if (tid < uint(nc)) ek[tid] = exp(mk[tid] - M); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + float L = 0.0f; + for (int k = 0; k < nc; ++k) L = fma(lk[k], ek[k], L); + L_sh = L; + } + float o = 0.0f; + device const float* q = pp + tid; + int k = 0; + for (; k + 3 < nc; k += 4, q += 4 * STRIDE) { + const float p0 = q[0], p1 = q[STRIDE], p2 = q[2 * STRIDE], p3 = q[3 * STRIDE]; + o = fma(p0, ek[k], o); + o = fma(p1, ek[k + 1], o); + o = fma(p2, ek[k + 2], o); + o = fma(p3, ek[k + 3], o); + } + for (; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); + threadgroup_barrier(mem_flags::mem_threadgroup); + o = o / L_sh; + const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; + const float y = o * (1.0f / (1.0f + exp(-gate))); + qsa_write_prepared(y, xp, a, t, h, tid); +} + diff --git a/results/round39/astra_r39_qsa/rejected-score.metal b/results/round39/astra_r39_qsa/rejected-score.metal new file mode 100644 index 0000000..523e823 --- /dev/null +++ b/results/round39/astra_r39_qsa/rejected-score.metal @@ -0,0 +1,502 @@ +// metal/kernels/qsa.metal - QSA (sparse full attention) decode for a window of tokens, between the projections. +// Contract of include/strata/kernels/qsa.hpp / qsa_select.hpp / qsa_decode_attn.hpp / kv_q8.hpp (the engine's +// native path), in a Metal layout (no paging: unified memory): +// +// qsa_prep per token t (cell pos0 + t): q heads rms_norm * q_norm, NEOX RoPE on 64 dims -> qcur; +// k rms_norm * k_norm, RoPE -> int8 KV append (fp16 scale per 64 = fp16(amax/127), rint, clamp); +// v -> int8 append; indexer q rms_norm * iq_norm, RoPE -> qidx. The indexer KEYS (last +// threadgroup, all window tokens in order): raw tail [cell % 4]; on a block completion +// pooled[b] = rope(rms_norm(mean of the block's 4 raw keys) * ik_norm, pos of the block's first cell); +// cell 0 also sets dead = rms_norm(raw) * ik_norm (the tail block's key). +// qsa_score per token, block b <= n_bid: sum over the 4 indexer heads of relu(qidx . key_b); key_b = pooled[b], +// the tail block n_bid scores `dead` and gets +1e9 when it has cells. +// qsa_topk per token: width = min(n_kv, 2051) cells; n_kv <= 2051: all; else the width cells of largest +// score (a cell inherits its block's), ties to the lowest cell, emitted ascending (radix select). +// qsa_attn split-K: (chunk of CH cells, kv head, token): 12 query heads; s = q.k / 16, chunk max m, sum l, +// acc = p . v -> partials +// qsa_merge per (head, token): log-sum-exp merge, * sigmoid(gate), written as prepared activations (xprep +// planes) for attn_output. +// +// RoPE: cos/sin from a host table [pos][32] (double precision, as build_rope_table). + +#define QH 24 +#define QKV 2 +#define QD 256 +#define QROT 64 +#define IH 4 +#define ID 128 +#define QBUDGET 2051 +#define QCH 64 + +constant int QSA_NTOK [[function_constant(3)]]; +// QSA_MROPE (function constant 24): interleaved M-RoPE (the vision path; strata/kernels/mrope.hpp, ggml IMROPE, +// sections 11/11/10): buffer 20 of qsa_prep is a table int32 [cells][3] (t, h, w), and rotary pair i (0..31) of the row +// at (pos_base + cell) rotates by the cos/sin table's row tab[(pos_base + cell) * 3 + i % 3] instead of its own - +// the pooled indexer keys too (position pos_base + 4 b of a block). Off: exactly the plain RoPE. +constant bool QSA_MROPE [[function_constant(24)]]; +constant bool qsa_mrope = is_function_constant_defined(QSA_MROPE) && QSA_MROPE; +// QSA_Q_F16_FC (function constant 25): qsa_prep writes qcur directly as fp16 (device half*) for qsa_flash, skipping qsa_q_f16. +constant bool QSA_Q_F16_FC [[function_constant(25)]]; +constant bool qsa_qf16 = is_function_constant_defined(QSA_Q_F16_FC) ? QSA_Q_F16_FC : false; +// QSA_ROW0 (function constant 28): qsa_prep's threadgroup x is row QSA_ROW0 + x (24: the K/V rows only, for rows whose +// queries nobody reads). Undefined: 0. +constant int QSA_ROW0 [[function_constant(28)]]; +constant int qsa_row0 = is_function_constant_defined(QSA_ROW0) ? QSA_ROW0 : 0; + +struct QsaArgs { + uint pos0; // the window's first cell (relative to the sequence) + uint pos_base; // absolute position of cell 0 (RoPE) + uint max_cells; + uint n_tok; + float eps; +}; + +// cache layout per layer +// kq/vq [cell][kv head][256] int8, ks/vs [cell][kv head][4] half +// tail [4][128] f32 (slot = cell % 4), dead [128], pooled [max_cells / 4 + 1][128] + +static inline float tg_sum256(float v, threadgroup float* part, ushort sg, ushort lane, ushort nsg) { + v = simd_sum(v); + if (lane == 0) part[sg] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + float t = 0.0f; + for (int i = 0; i < nsg; ++i) t += part[i]; + return t; +} + +// one row (width w <= 256 = threads) : out = rope((rs * x) * g) (RoPE on the first 64 dims, NEOX pairs i, i+32) +static inline float norm_rope(float x, float xhi_unused, device const float* g, threadgroup float* row, + threadgroup float* part, ushort tid, ushort sg, ushort lane, int w, float eps, + device const float2* cs, uint absp, bool rope, device const int* mtab, + bool reuse = false) { + const float ss = tg_sum256(tid < w ? x * x : 0.0f, part, sg, lane, 8); + const float rs = rsqrt(ss / float(w) + eps); + const float y = tid < w ? (rs * x) * g[tid] : 0.0f; + float o = y; + if (rope) { + if (tid < QROT) row[tid] = y; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < QROT) { + const int p = tid & 31; + const uint pp = qsa_mrope ? uint(mtab[absp * 3 + uint(p % 3)]) : absp; + const float2 c = cs[pp * 32 + p]; + const float a = row[p], b = row[p + 32]; + o = tid < 32 ? a * c.x - b * c.y : a * c.y + b * c.x; + } + } + if (reuse) threadgroup_barrier(mem_flags::mem_threadgroup); + (void)xhi_unused; + return o; +} + +// QSA_KV_STRIDE (function constant 26): kcur, vcur, iqraw and ikraw are column slices of one [n_tok][QSA_KV_STRIDE] +// projection (prompt_chunk's fused K | V | indexer K | indexer Q GEMM, 1152 or 1664 columns). Undefined: their own +// [n_tok][rows] buffers. +constant int QSA_KV_STRIDE [[function_constant(26)]]; +constant int qsa_kv_s = is_function_constant_defined(QSA_KV_STRIDE) ? QSA_KV_STRIDE : 0; + +// grid (32 rows + 1, n_tok): rows 0..23 q heads, 24..25 k heads, 26..27 v heads, 28..31 indexer q heads; +// threadgroup 32 (x = 32, y = 0) is the indexer-key threadgroup. 256 threads. A grid of 28 (or of 4 from QSA_ROW0 = +// 24) writes no indexer state: the MTP layer's attention is a dense window, never scored. +kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float* kcur [[buffer(1)]], + device const float* vcur [[buffer(2)]], device const float* iqraw [[buffer(3)]], + device const float* ikraw [[buffer(4)]], device const float* qn [[buffer(5)]], + device const float* kn [[buffer(6)]], device const float* iqn [[buffer(7)]], + device const float* ikn [[buffer(8)]], device const float2* cs [[buffer(9)]], + device float* qcur [[buffer(10)]], device float* qidx [[buffer(11)]], + device char* kq [[buffer(12)]], device char* vq [[buffer(13)]], device half* ksc [[buffer(14)]], + device half* vsc [[buffer(15)]], device float* tail [[buffer(16)]], device float* dead [[buffer(17)]], + device float* pooled [[buffer(18)]], constant QsaArgs& a [[buffer(19)]], + device const int* mtab [[buffer(20)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + threadgroup float row[QROT]; + threadgroup float amax[8]; + const int r = int(tg.x) + qsa_row0; + // the row strides of kcur / vcur, ikraw, iqraw + const int s_kv = qsa_kv_s ? qsa_kv_s : QKV * QD, s_ik = qsa_kv_s ? qsa_kv_s : ID; + const int s_iq = qsa_kv_s ? qsa_kv_s : IH * ID; + if (r == 32) { + // the indexer keys: threadgroup (32, 0) handles t < 3 (boundary with prior tail + cell 0 dead) and final tail; + // every threadgroup (32, t >= 3) with (cell & 3) == 3 pools its 4 in-window keys in parallel. + const int t_self = int(tg.y); + if (t_self >= 3) { + const uint cell = a.pos0 + uint(t_self); + if ((cell & 3) == 3) { + float m = 0.0f; + if (tid < ID) { + m = (((ikraw[(t_self - 3) * s_ik + tid] + ikraw[(t_self - 2) * s_ik + tid]) + + ikraw[(t_self - 1) * s_ik + tid]) + ikraw[t_self * s_ik + tid]) * 0.25f; + } + const uint b = cell >> 2; + const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab); + if (tid < ID) pooled[b * ID + tid] = o; + } + return; + } + if (t_self != 0) return; + const int n_pre = min(int(a.n_tok), 3); + for (int t = 0; t < n_pre; ++t) { + const uint cell = a.pos0 + uint(t); + const float x = tid < ID ? ikraw[t * s_ik + tid] : 0.0f; + if (tid < ID) tail[(cell & 3) * ID + tid] = x; + threadgroup_barrier(mem_flags::mem_device); + if (cell == 0) { + const float o = norm_rope(x, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, 0, false, mtab, true); + if (tid < ID) dead[tid] = o; + } + if ((cell & 3) == 3) { + float m = 0.0f; + if (tid < ID) m = (((tail[tid] + tail[ID + tid]) + tail[2 * ID + tid]) + tail[3 * ID + tid]) * 0.25f; + const uint b = cell >> 2; + const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab, true); + if (tid < ID) pooled[b * ID + tid] = o; + } + threadgroup_barrier(mem_flags::mem_device); + } + const int t_tail0 = max(n_pre, int(a.n_tok) - 4); + for (int t = t_tail0; t < int(a.n_tok); ++t) { + const uint cell = a.pos0 + uint(t); + if (tid < ID) tail[(cell & 3) * ID + tid] = ikraw[t * s_ik + tid]; + } + return; + } + const int t = int(tg.y); + const uint cell = a.pos0 + uint(t), absp = a.pos_base + cell; + if (r < QH) { + const float x = qfull[(t * QH + r) * 2 * QD + tid]; + const float o = norm_rope(x, 0.0f, qn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); + if (qsa_qf16) ((device half*)qcur)[(t * QH + r) * QD + tid] = half(o); + else qcur[(t * QH + r) * QD + tid] = o; + } else if (r < QH + 2 * QKV) { + const bool isv = r >= QH + QKV; + const int h = r - QH - (isv ? QKV : 0); + float x = (isv ? vcur : kcur)[t * s_kv + h * QD + tid]; + if (!isv) x = norm_rope(x, 0.0f, kn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); + // int8 with an fp16 scale per 64 values: groups = simdgroup pairs (64 threads) + float m = simd_max(fabs(x)); + if (lane == 0) amax[sg] = m; + threadgroup_barrier(mem_flags::mem_threadgroup); + const int g = tid >> 6; + m = max(amax[2 * g], amax[2 * g + 1]); + const half sh = half(m / 127.0f); + const float sf = float(sh); + int q = 0; + if (sf > 0.0f) q = clamp(int(rint(x / sf)), -127, 127); + const uint rowi = cell * QKV + uint(h); + (isv ? vq : kq)[rowi * QD + tid] = char(q); + if ((tid & 63) == 0) (isv ? vsc : ksc)[rowi * 4 + g] = sh; + } else { + if (a.pos0 + a.n_tok <= uint(QBUDGET)) return; + const int h = r - QH - 2 * QKV; + const float x = tid < ID ? iqraw[t * s_iq + 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); + if (tid < ID) qidx[(t * IH + h) * ID + tid] = o; + } +} + +// scores [n_tok][max_blocks]; grid (ceil(max_blocks / 8), n_tok), 256 threads: a simdgroup per block. +kernel void qsa_score(device const float* qidx [[buffer(0)]], device const float* pooled [[buffer(1)]], + device const float* dead [[buffer(2)]], device float* scores [[buffer(3)]], + constant QsaArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + const int t = int(tg.y); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + const uint b = tg.x * 8 + sg; + if (b > n_bid) return; + device const float* key = b == n_bid ? dead : pooled + b * ID; + device const float* qp = qidx + (t * IH) * ID + lane; + float k[4], q[IH][4]; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) k[i] = key[lane + 32 * i]; + #pragma clang loop unroll(full) + for (int h = 0; h < IH; ++h) { + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) q[h][i] = qp[h * ID + 32 * i]; + } + float d[IH]; + #pragma clang loop unroll(full) + for (int h = 0; h < IH; ++h) { + float dh = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) dh = fma(q[h][i], k[i], dh); + d[h] = dh; + } + const bool hi16 = (lane & 16) != 0, hi8 = (lane & 8) != 0; + float r2_0 = (hi16 ? d[2] : d[0]) + simd_shuffle_xor(hi16 ? d[0] : d[2], ushort(16)); + float r2_1 = (hi16 ? d[3] : d[1]) + simd_shuffle_xor(hi16 ? d[1] : d[3], ushort(16)); + float r = (hi8 ? r2_1 : r2_0) + simd_shuffle_xor(hi8 ? r2_0 : r2_1, ushort(8)); + r += simd_shuffle_xor(r, ushort(4)); + r += simd_shuffle_xor(r, ushort(2)); + r += simd_shuffle_xor(r, ushort(1)); + const float s1 = simd_broadcast(r, ushort(8)); + const float s2 = simd_broadcast(r, ushort(16)); + const float s3 = simd_broadcast(r, ushort(24)); + float s = max(r, 0.0f) + max(s1, 0.0f) + max(s2, 0.0f) + max(s3, 0.0f); + if (b == n_bid && (n_kv & 3) != 0) s += 1e9f; + if (lane == 0) scores[uint(t) * (a.max_cells / 4 + 2) + b] = s; +} + +// ids [n_tok][QBUDGET], ascending. One threadgroup of 1024 threads per token. Order key of a float: larger score = +// larger key; a radix select (4 passes of 8 bits) over the blocks with weights = their cell counts finds the +// threshold key T and how many cells with key == T to take (the lowest cells first). +static inline uint okey(float f) { + const uint u = as_type(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]], + constant QsaArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint hist[256]; + threadgroup uint sel[4]; + threadgroup uint2 pref[32]; + threadgroup uint2 totals; + const int t = int(tg.x); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2, nb = n_bid + 1; + device int* out = ids + t * QBUDGET; + device const float* sc = scores + uint(t) * (a.max_cells / 4 + 2); + if (n_kv <= QBUDGET) { + for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); + return; + } + const bool cached = nb <= 4096; + uint keys[4]; + if (cached) + for (uint b = tid, i = 0; b < nb; b += 1024, ++i) keys[i] = okey(sc[b]); + // radix select of the `need`-th largest cell key + uint prefix = 0, pmask = 0, need = QBUDGET; + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint b = tid; b < nb; b += 1024) { + const uint k = cached ? keys[b / 1024] : okey(sc[b]); + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + uint counts[8], sum = 0; + for (uint j = 0; j < 8; ++j) { + counts[j] = atomic_load_explicit(&hist[255 - (lane * 8 + j)], memory_order_relaxed); + sum += counts[j]; + } + uint acc = simd_prefix_exclusive_sum(sum); + if (acc < need && acc + sum >= need) { + for (uint j = 0; j < 8; ++j) { + if (acc + counts[j] >= need) { + sel[0] = 255 - (lane * 8 + j); + sel[1] = need - acc; + break; + } + acc += counts[j]; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; + } + // T = prefix; take every cell with key > T and the first `need` cells (lowest index) with key == T, ascending. + // Pass over blocks in order of index with a prefix scan of each block's selected cell count. + const uint T = prefix; + uint base = 0, eq_taken = 0; + for (uint b0 = 0; b0 < nb; b0 += 1024) { + const uint b = b0 + tid; + uint w = 0, weq = 0; + if (b < nb) { + const uint k = cached ? keys[b / 1024] : okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + if (k > T) w = cells; + else if (k == T) weq = cells; + } + const uint eq_before = simd_prefix_exclusive_sum(weq); + const uint gt_before = simd_prefix_exclusive_sum(w); + if (lane == 31) pref[sg] = uint2(eq_before + weq, gt_before + w); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + const uint2 count = pref[lane]; + const uint2 before = uint2(simd_prefix_exclusive_sum(count.x), + simd_prefix_exclusive_sum(count.y)); + pref[lane] = before; + if (lane == 31) totals = before + count; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint2 group_before = pref[sg], count = totals; + simdgroup_barrier(mem_flags::mem_threadgroup); + const uint eq0 = eq_taken + group_before.x + eq_before; + const uint take_eq = eq0 < need ? min(weq, need - eq0) : 0u; + const uint mine = w + take_eq; + // Clipping the equal-key prefix preserves the lowest-cell tie break. + const uint remaining = need - min(need, eq_taken); + const uint off = base + group_before.y + gt_before + + min(remaining, group_before.x + eq_before); + for (uint c = 0; c < mine; ++c) out[off + c] = int(4 * b + c); + base += count.y + min(remaining, count.x); + eq_taken += count.x; + } +} + +// partials: [n_tok][chunks][QH][QD + 2] (acc, then m, l). grid (chunks, QKV, n_tok), 256 threads. +kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + 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 HPG = QH / QKV; // 12 query heads per kv head + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH]; + threadgroup int cid[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; + const int nchunks = (QBUDGET + QCH - 1) / QCH; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + if (uint(c0) >= width) { + // an empty chunk: m = -inf, l = 0 (the merge skips it) + for (int h = 0; h < HPG; ++h) + if (tid == 0) { pp[(g * HPG + h) * (QD + 2) + QD] = -INFINITY; pp[(g * HPG + h) * (QD + 2) + QD + 1] = 0.0f; } + return; + } + const int nc = min(QCH, int(width) - c0); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + g * HPG + h) * QD + tid]; + if (tid < QCH) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + // scores: 12 x 64 dots of 256; simdgroup sg takes cells sg, sg + 8, ...; lanes split the 256 dims (8 each) + threadgroup const float* q_lane = &qs[0][lane * 8]; + for (int c = sg; c < nc; c += 8) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char* kr = kq + rowi * QD + lane * 8; + const char4 k0 = *(device const char4*)(kr); + const char4 k1 = *(device const char4*)(kr + 4); + const float scl = float(ksc[rowi * 4 + (lane >> 3)]); + const float kv[8] = { + float(k0.x) * scl, float(k0.y) * scl, float(k0.z) * scl, float(k0.w) * scl, + float(k1.x) * scl, float(k1.y) * scl, float(k1.z) * scl, float(k1.w) * scl + }; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + threadgroup const float* qh = q_lane + h * QD; + float d = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) d = fma(qh[i], kv[i], d); + d = simd_sum(d); + if (lane == 0) s[h][c] = d * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // per head: max, exp, sum (simdgroup h % 8 handles heads h, h + 8) + threadgroup float mh[HPG], lh[HPG]; + // nc <= QCH = 64: a lane holds cells lane and lane + 32 in registers (the same max and sum as a strided loop) + const bool in0 = int(lane) < nc, in1 = int(lane) + 32 < nc; + for (int h = sg; h < HPG; h += 8) { + const float s0 = in0 ? s[h][lane] : -INFINITY, s1 = in1 ? s[h][lane + 32] : -INFINITY; + const float m = simd_max(max(s0, s1)); + const float p0 = in0 ? exp(s0 - m) : 0.0f, p1 = in1 ? exp(s1 - m) : 0.0f; + if (in0) s[h][lane] = p0; + if (in1) s[h][lane + 32] = p1; + const float l = simd_sum(p0 + p1); + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // acc[h][j] = sum_c p[h][c] v[c][j]: thread = dim j + float acc[HPG]; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + // four cells' V loads in flight at once; the fma order per head stays c, c + 1, ... + auto vcell = [&](int c) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + return float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); + }; + int c = 0; + for (; c + 4 <= nc; c += 4) { + const float v0 = vcell(c), v1 = vcell(c + 1), v2 = vcell(c + 2), v3 = vcell(c + 3); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + float x = fma(s[h][c], v0, acc[h]); + x = fma(s[h][c + 1], v1, x); + x = fma(s[h][c + 2], v2, x); + acc[h] = fma(s[h][c + 3], v3, x); + } + } + for (; c < nc; ++c) { + const float v = vcell(c); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][c], v, acc[h]); + } + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (g * HPG + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} + +constant float qsa_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; +constant float qsa_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + +static inline void qsa_write_prepared(float y, device float* xp, constant QsaArgs& a, + int t, int h, ushort tid) { + // prepared planes (xprep layout) of the n = QH * QD vector + const uint n = QH * QD, plane = a.n_tok * n, d = uint(h) * QD + tid; + const int j = tid & 15; + device float* raw = xp + uint(t) * n + d; + raw[0] = y; + raw[plane] = y * qsa_q2s[j & 7]; + raw[2 * plane] = y * qsa_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? y * 0.00390625f : y; + float s0 = 0.0f, s1 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(y, ushort(k)); + s1 += simd_shuffle_down(y, ushort(k + 1)); + } + if (j == 0) xp[4 * plane + uint(t) * (n / 16) + d / 16] = s0 + s1; +} + +// grid (QH, n_tok), 256 threads: out = merge * sigmoid(gate), as prepared activations of n = QH * QD per token +kernel void qsa_merge(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], + device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const int h = int(tg.x), t = int(tg.y); + const int nchunks = (QBUDGET + QCH - 1) / QCH; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nc = int((width + QCH - 1) / QCH); + constexpr uint STRIDE = QH * (QD + 2); + device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2); + // the chunks' (m, l) and weights exp(m - M) once per threadgroup, not once per thread (nc <= 33 < 64) + threadgroup float mk[64], lk[64], ek[64], L_sh; + if (tid < uint(nc)) { + const float2 ml = *(device const float2*)(pp + tid * STRIDE + QD); + mk[tid] = ml.x; + lk[tid] = ml.y; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float M = -INFINITY; + for (int k = tid & 31; k < nc; k += 32) M = max(M, mk[k]); + M = simd_max(M); + if (tid < uint(nc)) ek[tid] = exp(mk[tid] - M); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + float L = 0.0f; + for (int k = 0; k < nc; ++k) L = fma(lk[k], ek[k], L); + L_sh = L; + } + float o = 0.0f; + device const float* q = pp + tid; + for (int k = 0; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); + threadgroup_barrier(mem_flags::mem_threadgroup); + o = o / L_sh; + const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; + const float y = o * (1.0f / (1.0f + exp(-gate))); + qsa_write_prepared(y, xp, a, t, h, tid); +} + diff --git a/results/round39/astra_r39_qsa/report.md b/results/round39/astra_r39_qsa/report.md new file mode 100644 index 0000000..56ba472 --- /dev/null +++ b/results/round39/astra_r39_qsa/report.md @@ -0,0 +1,71 @@ +# round 39 track a — qsa upstream fixes and local merge experiment + +## decision + +- keep `831ba47a`'s clamped-row/unconditional-load/value-select fix in the four non-cooperative prompt fills. preserve our staged scales and cooperative paths. +- keep `ad734531`'s `TK = 64` default and explicit `v2 || v3` gather/HS/binding conditions. this tree already sets `v2 = v3 = true` for `--v3`, so the binding change is documentary/defensive here, not a newly repaired v3 failure. existing half-input gather and `--gu-bench`, `--dn-bench`, `--ab`, `--dtile` remain intact. +- keep the **local** four-partial-load batching experiment in `qsa_merge2`. this is inspired by `731528be`, not a literal upstream transplant. active-chunk bounds already existed and were not changed. +- reject the exact requested six-shuffle `qsa_score` butterfly: `scores FAIL`, exit 1, on its first `--ab` case at p=3000. no speed claim for a numerically invalid candidate. cuda's reduction-order argument does not establish equivalence to metal `simd_sum`; the mismatch is measured, its compiler-level cause is not proven. +- reject four-partial batching in `qsa_merge`: bitwise correct, but nt=3 regressed repeatedly by 4.8–6.6% across all three contexts. `qsa.metal` is entirely unchanged in the delivered patch. + +## artifacts and reproducibility + +- patch: `/tmp/astra_r39_qsa/astra_r39_qsa.patch` +- durable copy, source snapshots, evaluators, metadata, commands, exit statuses and full logs: `results/round39/astra_r39_qsa/` +- exact baseline: the working-tree snapshot, not merely git head; `metadata.json` records its head and four owned-file sha256 values. `baseline-kernels/` and `baseline-moe_mm_parity.cpp` preserve inputs. `rejected-score.metal` and `rejected-merge1.metal` preserve rejected experiments; `final/` preserves the delivered sources. +- all editing/building ran in `/tmp/astra_r39_qsa/strata`; no tracked vendor files, live service or full-model tests were touched. +- device: apple m5; system version recorded in `metadata.json`. apple clang 21.0.0, release `-O2`. build used `cmake -S metal -B build-metal -G Ninja` followed by `cmake --build build-metal --target qsa_parity qsa_prompt_parity moe_mm_parity -j 3`. +- tests were serialized on the gpu. initial moe runs exceeded the original 15s/20s limits and were terminated; after parent explicitly extended CPU-reference allowance to 180s, both full tests completed. all unsuccessful attempts remain in the logs. +- `git -C vendor/strata apply --check /tmp/astra_r39_qsa/astra_r39_qsa.patch`: exit 0. scratch `git diff --check`: exit 0. final visual diff/context review completed. + +## correctness and test results + +all commands below run from the scratch tree; `BASE=/tmp/astra_r39_qsa/baseline`. + +| test | result | +|---|---| +| baseline `qsa_parity --ab $BASE 3000` | exit 0, all memcmps pass | +| baseline `moe_mm_parity` | exit 3, missing `mm_gu_42_16_64_16_2` | +| final `moe_mm_parity` (no args) | exit 0, 23.961s, `PASS`; fp16-reference rel l2 2.03e-06, exact-reference 2.37e-04 | +| final `moe_mm_parity --v3 --gu IQ2_S --dn Q2_0 --tile 16,32,32,4 --dtile 64,32,2` | exit 0, 20.011s, `PASS`; fp16-reference rel l2 1.60e-05, exact-reference 4.38e-04 | +| same v3 command plus `--ab $BASE` | exit 0, HS/YE/Y memcmp pass, t=256 | +| default and specified v3 plus `--T 1` | exit 0, CPU-reference parity pass | +| v3 plus `--T 1 --gu-bench --ab $BASE`, and `--dn-bench` | both exit 0; HS/YE/Y memcmp pass; flags preserved | +| `qsa_prompt_parity` | exit 0, `ALL PASS`, default cooperative layout | +| `COOP_K4=0 qsa_prompt_parity` | exit 0, `ALL PASS`, edited non-cooperative layout | +| `qsa_prompt_parity --mm-compare attn 96 P $BASE/qsa_prompt.metal 3` | p=0 and 30000, both COOP_K4=0 and default: exit 0; output, post-timing and merge+prepared memcmp pass | +| `qsa_prompt_parity --flash-compare 96 P $BASE/qsa_prompt.metal 1 3` | same four combinations: exit 0, memcmp pass | +| `QSA_ATTN=2 qsa_parity --ab $BASE P` | p=3000,10180,30000: all full-buffer memcmps pass; nt=1,2,3,4,8 plus prompt prep nt=1984,2048 | +| `QSA_ATTN=2 qsa_parity --bitwise $BASE/qsa2.metal` | exit 0; zero of 40 cases differ, including short contexts/remainders | + +important distinction: moe's no-argument and specified-v3 CPU checks are tolerance checks and print `PASS`, not `ALL PASS` or bitwise equality to CPU. the separate v3 A/B command provides bitwise equality to the baseline kernels. prompt's default suite likewise includes tolerance checks; the comparison modes provide bitwise evidence for both edited kernels. no unsupported claim of IEEE equality to a differently accumulated CPU reference is made. + +## merge2 paired timings + +three sweeps: initial, confirmation, then final with rejected merge1 restored. all pass. the final sweep below uses only retained changes. measurements are median GPU µs; paired ratio is median(candidate/baseline), not necessarily the quotient of medians. inspected evaluator code uses **five paired rounds**, 16 dispatch repetitions for decode (the file header's 12-repetition description is stale). each stage is warmed before interleaved A/B timing. + +| p | nt | baseline µs | candidate µs | paired candidate/baseline | +|---:|---:|---:|---:|---:| +| 3000 | 1 | 15.143 | 9.273 | 0.6122 | +| 3000 | 2 | 21.924 | 21.516 | 0.9812 | +| 3000 | 3 | 33.688 | 32.969 | 0.9791 | +| 3000 | 4 | 45.458 | 43.227 | 0.9522 | +| 3000 | 8 | 91.102 | 86.547 | 0.9501 | +| 10180 | 1 | 10.961 | 6.844 | 0.6240 | +| 10180 | 2 | 21.977 | 21.589 | 0.9834 | +| 10180 | 3 | 33.732 | 32.961 | 0.9772 | +| 10180 | 4 | 45.344 | 43.313 | 0.9546 | +| 10180 | 8 | 91.393 | 86.607 | 0.9493 | +| 30000 | 1 | 17.180 | 10.807 | 0.6295 | +| 30000 | 2 | 22.125 | 21.547 | 0.9739 | +| 30000 | 3 | 33.500 | 33.471 | 0.9955 | +| 30000 | 4 | 45.469 | 43.544 | 0.9577 | +| 30000 | 8 | 91.594 | 88.021 | 0.9556 | + +nt=1 improves 35–41% across the earlier sweeps and 37–39% in the final sweep; other cases are neutral to about 5% faster. one initial p=30000, nt=3 ratio was 1.0008 (0.08% noise); subsequent ratios were 0.9902 and 0.9955. no repeatable regression. absolute clocks/times drift (even unchanged stages), so use paired ratios, not cross-run raw times. this is a standalone-kernel result, not evidence of whole-model throughput improvement. + +for the rejected merge1, confirmation nt=3 ratios were 1.0551, 1.0660, 1.0563 at p=3000,10180,30000. despite wins elsewhere, these violate the requested neutral-to-faster-across-suite rule. + +## prompt safety fix timing observation + +with `COOP_K4=0`, p=30000/t=96 non-cooperative attn median fell from 4.954 to 1.537ms, flash from 5.685 to 1.828ms, with memcmp pass. p=0 also improved. these three-pair observations are secondary to the safety fix, not broad performance claims. cooperative mode remains bitwise identical. actual m1 hardware was not available; the m1 fault rationale comes from upstream `831ba47a`, while our runtime evidence is m5 with both constant settings. diff --git a/results/round39/astra_r39_qsa/runs.json b/results/round39/astra_r39_qsa/runs.json new file mode 100644 index 0000000..1405311 --- /dev/null +++ b/results/round39/astra_r39_qsa/runs.json @@ -0,0 +1,671 @@ +[ + { + "name": "baseline-ab", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "3000" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": 0, + "duration_s": 14.704230582981836, + "timed_out": false, + "log": "logs/baseline-ab.log" + }, + { + "name": "baseline-moe", + "argv": [ + "./build-metal/moe_mm_parity" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": 3, + "duration_s": 4.560703166993335, + "timed_out": false, + "log": "logs/baseline-moe.log" + }, + { + "name": "confirm-merge1-ab-10180", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "10180" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "QSA_ATTN": "1" + }, + "exit_code": 0, + "duration_s": 4.757935791014461, + "timed_out": false, + "log": "logs/confirm-merge1-ab-10180.log" + }, + { + "name": "confirm-merge1-ab-3000", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "3000" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "QSA_ATTN": "1" + }, + "exit_code": 0, + "duration_s": 3.2664715830178466, + "timed_out": false, + "log": "logs/confirm-merge1-ab-3000.log" + }, + { + "name": "confirm-merge1-ab-30000", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "30000" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "QSA_ATTN": "1" + }, + "exit_code": 0, + "duration_s": 7.467346207995433, + "timed_out": false, + "log": "logs/confirm-merge1-ab-30000.log" + }, + { + "name": "confirm-merge2-ab-10180", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "10180" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "QSA_ATTN": "2" + }, + "exit_code": 0, + "duration_s": 4.4522730839962605, + "timed_out": false, + "log": "logs/confirm-merge2-ab-10180.log" + }, + { + "name": "confirm-merge2-ab-3000", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "3000" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "QSA_ATTN": "2" + }, + "exit_code": 0, + "duration_s": 2.6887030830257572, + "timed_out": false, + "log": "logs/confirm-merge2-ab-3000.log" + }, + { + "name": "confirm-merge2-ab-30000", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "30000" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "QSA_ATTN": "2" + }, + "exit_code": 0, + "duration_s": 6.560153625003295, + "timed_out": false, + "log": "logs/confirm-merge2-ab-30000.log" + }, + { + "name": "final-merge2-ab-10180", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "10180" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "QSA_ATTN": "2" + }, + "exit_code": 0, + "duration_s": 7.093214749998879, + "timed_out": false, + "log": "logs/final-merge2-ab-10180.log" + }, + { + "name": "final-merge2-ab-3000", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "3000" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "QSA_ATTN": "2" + }, + "exit_code": 0, + "duration_s": 8.637194999988424, + "timed_out": false, + "log": "logs/final-merge2-ab-3000.log" + }, + { + "name": "final-merge2-ab-30000", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "30000" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "QSA_ATTN": "2" + }, + "exit_code": 0, + "duration_s": 9.387072458019247, + "timed_out": false, + "log": "logs/final-merge2-ab-30000.log" + }, + { + "name": "final-prompt-k4off", + "argv": [ + "./build-metal/qsa_prompt_parity" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "COOP_K4": "0" + }, + "exit_code": 0, + "duration_s": 0.36627699999371544, + "timed_out": false, + "log": "logs/final-prompt-k4off.log" + }, + { + "name": "fixed-moe-default-full", + "argv": [ + "./build-metal/moe_mm_parity" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": 0, + "duration_s": 23.961265375022776, + "timed_out": false, + "log": "logs/fixed-moe-default-full.log" + }, + { + "name": "fixed-moe-default", + "argv": [ + "./build-metal/moe_mm_parity" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": -15, + "duration_s": 15.026818915997865, + "timed_out": true, + "log": "logs/fixed-moe-default.log" + }, + { + "name": "fixed-moe-v3-full", + "argv": [ + "./build-metal/moe_mm_parity", + "--v3", + "--gu", + "IQ2_S", + "--dn", + "Q2_0", + "--tile", + "16,32,32,4", + "--dtile", + "64,32,2" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": 0, + "duration_s": 20.010784165991936, + "timed_out": false, + "log": "logs/fixed-moe-v3-full.log" + }, + { + "name": "fixed-moe-v3", + "argv": [ + "./build-metal/moe_mm_parity", + "--v3", + "--gu", + "IQ2_S", + "--dn", + "Q2_0", + "--tile", + "16,32,32,4", + "--dtile", + "64,32,2" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": -15, + "duration_s": 20.02672729198821, + "timed_out": true, + "log": "logs/fixed-moe-v3.log" + }, + { + "name": "fixed-prompt", + "argv": [ + "./build-metal/qsa_prompt_parity" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": 0, + "duration_s": 0.9771465419908054, + "timed_out": false, + "log": "logs/fixed-prompt.log" + }, + { + "name": "merge1-ab-10180", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "10180" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "QSA_ATTN": "1" + }, + "exit_code": 0, + "duration_s": 6.856663957994897, + "timed_out": false, + "log": "logs/merge1-ab-10180.log" + }, + { + "name": "merge1-ab-3000", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "3000" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "QSA_ATTN": "1" + }, + "exit_code": 0, + "duration_s": 5.871507915988332, + "timed_out": false, + "log": "logs/merge1-ab-3000.log" + }, + { + "name": "merge1-ab-30000", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "30000" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "QSA_ATTN": "1" + }, + "exit_code": 0, + "duration_s": 11.867526749992976, + "timed_out": false, + "log": "logs/merge1-ab-30000.log" + }, + { + "name": "merge2-ab-10180", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "10180" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "QSA_ATTN": "2" + }, + "exit_code": 0, + "duration_s": 5.168871124973521, + "timed_out": false, + "log": "logs/merge2-ab-10180.log" + }, + { + "name": "merge2-ab-3000", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "3000" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "QSA_ATTN": "2" + }, + "exit_code": 0, + "duration_s": 3.023166750004748, + "timed_out": false, + "log": "logs/merge2-ab-3000.log" + }, + { + "name": "merge2-ab-30000", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "30000" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "QSA_ATTN": "2" + }, + "exit_code": 0, + "duration_s": 8.313932707998902, + "timed_out": false, + "log": "logs/merge2-ab-30000.log" + }, + { + "name": "merge2-bitwise", + "argv": [ + "./build-metal/qsa_parity", + "--bitwise", + "/tmp/astra_r39_qsa/baseline/qsa2.metal" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "QSA_ATTN": "2" + }, + "exit_code": 0, + "duration_s": 1.3079375409870408, + "timed_out": false, + "log": "logs/merge2-bitwise.log" + }, + { + "name": "moe-default-T1", + "argv": [ + "./build-metal/moe_mm_parity", + "--T", + "1" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": 0, + "duration_s": 4.224036582978442, + "timed_out": false, + "log": "logs/moe-default-T1.log" + }, + { + "name": "moe-dn-bench", + "argv": [ + "./build-metal/moe_mm_parity", + "--v3", + "--gu", + "IQ2_S", + "--dn", + "Q2_0", + "--tile", + "16,32,32,4", + "--dtile", + "64,32,2", + "--T", + "1", + "--dn-bench", + "--ab", + "/tmp/astra_r39_qsa/baseline" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": 0, + "duration_s": 0.26160704100038856, + "timed_out": false, + "log": "logs/moe-dn-bench.log" + }, + { + "name": "moe-gu-bench", + "argv": [ + "./build-metal/moe_mm_parity", + "--v3", + "--gu", + "IQ2_S", + "--dn", + "Q2_0", + "--tile", + "16,32,32,4", + "--dtile", + "64,32,2", + "--T", + "1", + "--gu-bench", + "--ab", + "/tmp/astra_r39_qsa/baseline" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": 0, + "duration_s": 0.34804070901009254, + "timed_out": false, + "log": "logs/moe-gu-bench.log" + }, + { + "name": "moe-v3-T1", + "argv": [ + "./build-metal/moe_mm_parity", + "--v3", + "--gu", + "IQ2_S", + "--dn", + "Q2_0", + "--tile", + "16,32,32,4", + "--dtile", + "64,32,2", + "--T", + "1" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": 0, + "duration_s": 0.1989670839975588, + "timed_out": false, + "log": "logs/moe-v3-T1.log" + }, + { + "name": "moe-v3-ab", + "argv": [ + "./build-metal/moe_mm_parity", + "--v3", + "--gu", + "IQ2_S", + "--dn", + "Q2_0", + "--tile", + "16,32,32,4", + "--dtile", + "64,32,2", + "--ab", + "/tmp/astra_r39_qsa/baseline" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": 0, + "duration_s": 2.633624290989246, + "timed_out": false, + "log": "logs/moe-v3-ab.log" + }, + { + "name": "prompt-flash-0", + "argv": [ + "./build-metal/qsa_prompt_parity", + "--flash-compare", + "96", + "0", + "/tmp/astra_r39_qsa/baseline/qsa_prompt.metal", + "1", + "3" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": 0, + "duration_s": 0.035576458991272375, + "timed_out": false, + "log": "logs/prompt-flash-0.log" + }, + { + "name": "prompt-flash-30000", + "argv": [ + "./build-metal/qsa_prompt_parity", + "--flash-compare", + "96", + "30000", + "/tmp/astra_r39_qsa/baseline/qsa_prompt.metal", + "1", + "3" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": 0, + "duration_s": 0.11264241699245758, + "timed_out": false, + "log": "logs/prompt-flash-30000.log" + }, + { + "name": "prompt-flash-k4off-0", + "argv": [ + "./build-metal/qsa_prompt_parity", + "--flash-compare", + "96", + "0", + "/tmp/astra_r39_qsa/baseline/qsa_prompt.metal", + "1", + "3" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "COOP_K4": "0" + }, + "exit_code": 0, + "duration_s": 0.31251912500010803, + "timed_out": false, + "log": "logs/prompt-flash-k4off-0.log" + }, + { + "name": "prompt-flash-k4off-30000", + "argv": [ + "./build-metal/qsa_prompt_parity", + "--flash-compare", + "96", + "30000", + "/tmp/astra_r39_qsa/baseline/qsa_prompt.metal", + "1", + "3" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "COOP_K4": "0" + }, + "exit_code": 0, + "duration_s": 0.131041875021765, + "timed_out": false, + "log": "logs/prompt-flash-k4off-30000.log" + }, + { + "name": "prompt-mm-0", + "argv": [ + "./build-metal/qsa_prompt_parity", + "--mm-compare", + "attn", + "96", + "0", + "/tmp/astra_r39_qsa/baseline/qsa_prompt.metal", + "3" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": 0, + "duration_s": 0.4449748329934664, + "timed_out": false, + "log": "logs/prompt-mm-0.log" + }, + { + "name": "prompt-mm-30000", + "argv": [ + "./build-metal/qsa_prompt_parity", + "--mm-compare", + "attn", + "96", + "30000", + "/tmp/astra_r39_qsa/baseline/qsa_prompt.metal", + "3" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": 0, + "duration_s": 0.1662420420034323, + "timed_out": false, + "log": "logs/prompt-mm-30000.log" + }, + { + "name": "prompt-mm-k4off-0", + "argv": [ + "./build-metal/qsa_prompt_parity", + "--mm-compare", + "attn", + "96", + "0", + "/tmp/astra_r39_qsa/baseline/qsa_prompt.metal", + "3" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "COOP_K4": "0" + }, + "exit_code": 0, + "duration_s": 0.25848625000799075, + "timed_out": false, + "log": "logs/prompt-mm-k4off-0.log" + }, + { + "name": "prompt-mm-k4off-30000", + "argv": [ + "./build-metal/qsa_prompt_parity", + "--mm-compare", + "attn", + "96", + "30000", + "/tmp/astra_r39_qsa/baseline/qsa_prompt.metal", + "3" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": { + "COOP_K4": "0" + }, + "exit_code": 0, + "duration_s": 0.18658049998339266, + "timed_out": false, + "log": "logs/prompt-mm-k4off-30000.log" + }, + { + "name": "score-ab-3000", + "argv": [ + "./build-metal/qsa_parity", + "--ab", + "/tmp/astra_r39_qsa/baseline", + "3000" + ], + "cwd": "/tmp/astra_r39_qsa/strata", + "env_added": {}, + "exit_code": 1, + "duration_s": 1.7547417919849977, + "timed_out": false, + "log": "logs/score-ab-3000.log" + } +] \ No newline at end of file diff --git a/results/round39/complete-patch-check.log b/results/round39/complete-patch-check.log new file mode 100644 index 0000000..e69de29 diff --git a/results/round39/complete-patch-replay.json b/results/round39/complete-patch-replay.json new file mode 100644 index 0000000..2fe0e31 --- /dev/null +++ b/results/round39/complete-patch-replay.json @@ -0,0 +1,33 @@ +{ + "apply_exit": 0, + "compared": [ + "serve/frontend.py", + "serve/logit_bias.py", + "serve/runconfig.py", + "serve/server.py", + "serve/telemetry.py", + "serve/test_admit_order.py", + "serve/test_effort_paths.py", + "serve/test_logit_bias.py", + "serve/test_model_card_sampling.py", + "serve/test_parallel.py", + "serve/test_parser_call_scan.py", + "serve/test_parser_hold.py", + "serve/test_prompt_encoder.py", + "serve/test_q36_config.py", + "serve/test_reasoning_loop_recovery.py", + "serve/test_server.py", + "serve/test_slot_alloc.py", + "serve/test_web_sampling.py", + "serve/web/app.js", + "tools/strata_tokenizer.py", + "tools/test_strata_tokenizer.py", + "serve/test_image_omission.py", + "metal/kernels/qsa2.metal", + "metal/kernels/qsa_prompt.metal", + "metal/tests/moe_mm_parity.cpp", + "src/kernels/ngram.cpp", + "src/ngram/ple_reader_test.cpp" + ], + "mismatches": [] +} diff --git a/results/round39/integration-build.log b/results/round39/integration-build.log new file mode 100644 index 0000000..41c048d --- /dev/null +++ b/results/round39/integration-build.log @@ -0,0 +1,7 @@ +ninja: Entering directory `vendor/strata/build-metal' +[0/2] Re-checking globbed directories... +[1/6] Building CXX object CMakeFiles/strata_metal_host.dir/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/src/kernels/ngram.cpp.o +[2/6] Linking CXX static library libstrata_metal_host.a +[3/6] Linking CXX executable strata_window_cli +[4/6] Building CXX object CMakeFiles/moe_mm_parity.dir/tests/moe_mm_parity.cpp.o +[5/6] Linking CXX executable moe_mm_parity diff --git a/results/round39/integration-source-sha256.json b/results/round39/integration-source-sha256.json new file mode 100644 index 0000000..989600f --- /dev/null +++ b/results/round39/integration-source-sha256.json @@ -0,0 +1,29 @@ +{ + "serve/frontend.py": "1f40ac6fc9f4ad021ac0c202f35f184f5784009355dd7485ee0e44746bfd75e1", + "serve/logit_bias.py": "b6564c97e91c987e5787349b3df209fa0cf00fc8e80357ad93d40fd1e594da09", + "serve/runconfig.py": "58aa118367a9d6dbb534cae5f35d6f418112f96985db85b048fa23eb042e403a", + "serve/server.py": "7ee5536c9293583244abb3774851fb16aa2ac7e7e55e200fef13bb4be4aac19b", + "serve/telemetry.py": "973772d0209167bf0dbaa8fd98fab8e64801f12dd241c72d5e06bbaaf96b646e", + "serve/test_admit_order.py": "bbfe80b214f0dfcf2339bb4aab08e1549330951bcf5b7a607acea4f436cae47b", + "serve/test_effort_paths.py": "9501d8787e396b95f5143eb298334973a8553b298feff0e93b8708dfcc4f6db9", + "serve/test_logit_bias.py": "d4afa9e4f50bfa3595db67bc9e2487adabcaf7e902be2c3dca3bbccec434e92a", + "serve/test_model_card_sampling.py": "38a90fe3d2339d1bca61b073a5e0bafae936c36ab1b1953ffdc6f756a06bf5f3", + "serve/test_parallel.py": "0f6147f1e74badb61f974d0fd71d2b39e4c141eaeee6dc3cfbf5f5bbb585eeec", + "serve/test_parser_call_scan.py": "71aab3ee7ea39bf8810099e3d59d1c8580836923b7aa3446a72c281b6398232d", + "serve/test_parser_hold.py": "d50e799c0bbf003318b51938f09e0b40fb11b8d572a7377ddd4617b9ddf9d7cf", + "serve/test_prompt_encoder.py": "fad128ea329949fd38ce034a026b3f86d060c484fcfe1c0fea2a13f94524a73c", + "serve/test_q36_config.py": "278dae2f526ad8cda6fe8a719e075456c71824bbd8b2362b87fe62139b6b0ea8", + "serve/test_reasoning_loop_recovery.py": "c448df0757ae1bffb0a4fe8b48f45ee57ae0f06e7056d6bc96885dc8b024bf7b", + "serve/test_server.py": "576f7e5eb3ea07cb8dc0939714d282b59e9f51d0d4c6e00549394f4cd9753d17", + "serve/test_slot_alloc.py": "7ea0cb07ed01f95d720fb0ae1d945cee705a31a34213bd5ada22dfa35c68d40c", + "serve/test_web_sampling.py": "8e4e717aaab9049b99a1a44beff2bae1d48e6f60d08de98a28e2d875d2eba18d", + "serve/web/app.js": "7d15c4efeccbc4b1f4944494dc76987e18e7a6cf40d12279978dfe93951bb8fd", + "tools/strata_tokenizer.py": "e0cd8d6ae38dda0a4ad7c14f812eb64d5fd6f4fed9c920d4427a14bf585dc9b3", + "tools/test_strata_tokenizer.py": "6a3cadf78c1972e83acd0ef90a8f5119868c7c3a4613fcfa000faae1195017b5", + "serve/test_image_omission.py": "88c88b3a2f15d41958a2b60760734897c4656d275353b906b6c49d54591d4b39", + "metal/kernels/qsa2.metal": "ceb7fb66e2acd2e05e26ff0ce27e6533831c01c1d0adbb6d1917ec5185ccc172", + "metal/kernels/qsa_prompt.metal": "6af1d7979169bcd09de6e4ac7fe6ed5a6aa4e080917577c7e4731b8b6f466c94", + "metal/tests/moe_mm_parity.cpp": "e8fec719467a5d236857ad141e6a7318fc25f6abcad79efb41f33417d4538217", + "src/kernels/ngram.cpp": "f76bb2b9d8a082ad9489252ff71a7dd9e13cedfc01c92c27fffbf2c95a29bfd3", + "src/ngram/ple_reader_test.cpp": "6eede85aea6a594d091b542b0be5d286681c8c083eff8c8cea86cdbd12d0e84b" +} diff --git a/results/round39/keepawake-mechanism-baseline.log b/results/round39/keepawake-mechanism-baseline.log new file mode 100644 index 0000000..fdf7c24 --- /dev/null +++ b/results/round39/keepawake-mechanism-baseline.log @@ -0,0 +1,3 @@ +[strata] done: 3 tokens in 0 s (125829.1 tok/s) (stop, cancel=False) +request completed before worker condition lock: count= 0 calls= [] +after close: [2147483648] diff --git a/results/round39/keepawake-mechanism-candidate.log b/results/round39/keepawake-mechanism-candidate.log new file mode 100644 index 0000000..6266527 --- /dev/null +++ b/results/round39/keepawake-mechanism-candidate.log @@ -0,0 +1,3 @@ +[strata] done: 3 tokens in 0 s (106634.8 tok/s) (stop, cancel=False) +request completed before worker condition lock: count= 0 calls= [] +after close: [2147483648] diff --git a/results/round39/live-chat-first.json b/results/round39/live-chat-first.json new file mode 100644 index 0000000..1705c44 --- /dev/null +++ b/results/round39/live-chat-first.json @@ -0,0 +1 @@ +{"id": "chatcmpl-15093e04ea61492a8a4a78db", "object": "chat.completion", "created": 1792314432, "model": "swift-1.5-k512-iq2xs", "choices": [{"index": 0, "message": {"role": "assistant", "content": "ready"}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 17, "completion_tokens": 2, "total_tokens": 19, "prompt_tokens_details": {"cached_tokens": 0}}, "timings": {"cache_n": 0, "prompt_n": 17, "prompt_ms": 1226.0, "prompt_per_token_ms": 72.118, "prompt_per_second": 13.9, "predicted_n": 2, "predicted_ms": 201.9, "predicted_per_token_ms": 100.95, "predicted_per_second": 9.9, "draft_n": 0, "draft_n_accepted": 0}} \ No newline at end of file diff --git a/results/round39/live-chat-second.json b/results/round39/live-chat-second.json new file mode 100644 index 0000000..7038234 --- /dev/null +++ b/results/round39/live-chat-second.json @@ -0,0 +1 @@ +{"id": "chatcmpl-ee137d5e12f349f998471b11", "object": "chat.completion", "created": 1792314434, "model": "swift-1.5-k512-iq2xs", "choices": [{"index": 0, "message": {"role": "assistant", "content": "still ready"}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 39, "completion_tokens": 3, "total_tokens": 42, "prompt_tokens_details": {"cached_tokens": 18}}, "timings": {"cache_n": 18, "prompt_n": 21, "prompt_ms": 278.0, "prompt_per_token_ms": 13.238, "prompt_per_second": 75.5, "predicted_n": 3, "predicted_ms": 139.7, "predicted_per_token_ms": 46.567, "predicted_per_second": 21.5, "draft_n": 0, "draft_n_accepted": 0}} \ No newline at end of file diff --git a/results/round39/model-ab/10k-0-baseline.log b/results/round39/model-ab/10k-0-baseline.log new file mode 100644 index 0000000..5a626ef --- /dev/null +++ b/results/round39/model-ab/10k-0-baseline.log @@ -0,0 +1,46 @@ +strata metal stream: experts read from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf.experts (one request per expert) +strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): 290 matrices, 680 MB, made in 138 ms +strata metal stream: locked 2.27 GB of active GPU dense tensors (unpinned 0.34 GB CPU token_embd + 1.28 GB skipped BF16 HC in 3.90 GB file map) +strata metal stream: 1079 dense tensors wrapped zero-copy from file (3.90 GB clean file-backed map, 123.6 ms) +strata metal: decode window on Apple M5: /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf streamed from the SSD (experts) (3.9 GB), 32896 cells, PLE on, ready in 538 ms +strata metal: MTP map: 0.71 GB wired (mlock; 0.00 GB failed), 0.18 GB unused released +strata metal: MTP draft round in one command buffer (embedding type 23 on the GPU: max |GPU - host| 0 over 8 rows) +strata metal: MTP draft layer from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf (0.89 GB, 9 norms folded), draft head over 40525 tokens +strata metal stream: 0.00 GB of the host code's device weight copies given back (unused by the Metal window) +strata metal stream: 2.60 GB allocated before the arena (the process holds 2.24 GB, 5.40 GB resident); RAM 25.8 GB (12.0 GB available, 7.0 GB kept for macOS), GPU working set 19.1 GB +strata metal stream: planned resident footprint 9.6 GB decode (10.7 GB virtual peak during prefill with 1.12 GB transient prompt scratch) = arena 6.5 + 2.6 held before it + 0.5 reserved for the n-gram table; cap 18.8 GB = the lesser of the GPU working set 19.1 GB and the RAM 25.8 GB less 7.0 GB left to macOS and its file cache +strata metal stream: expert arena 6.50 GB, 4129 slots in 2 size classes, 4129 experts read from the profile's top (24576 ranked) in 1.5 s (4.14 GB/s); 8 I/O threads, read-ahead top 6; 1 buffers +strata metal stream: 6.5 GB of the arena wired (mlock) in 575 ms; macOS cannot compress it while the engine idles (STRATA_METAL_WIRE_ARENA=0 turns it off) +strata metal stream: resident footprint after the fill: the process holds 8.7 GB (10.4 GB resident), planned 9.6 GB, cap 18.8 GB +strata metal stream: prompt scratch 1.12 GB mapped in separate anonymous mmap +strata metal: prompts on the matrix units in chunks of 2048 (experts v3, QSA flash, matmul2d layout K4; 1.12 GB of buffers, in a separate transient mmap while a prompt runs) +strata metal: prompt chunk pos 0, 2048 rows: 4592 ms (embd 26 ms, ple 3 ms, graph build 61 ms) +strata metal: prompt chunk pos 2048, 2048 rows: 4455 ms (embd 7 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 4096, 2048 rows: 4593 ms (embd 4 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 6144, 2048 rows: 4752 ms (embd 3 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 8192, 1987 rows: 4686 ms (embd 4 ms, ple 0 ms, graph build 1 ms) +strata metal stream: prompt scratch 1.12 GB unmapped after prefill (phys_footprint 8.89 -> 8.79 GB, resident RSS 9.61 -> 9.52 GB) +strata metal: window pos 10179 token 198 -> 760 (top 760 23.8239, 2nd 35 23.7743, margin 0.0496) 142.06 ms (GPU 102.97 ms, 1233 dispatches) +expert stream prompts 5 chunks: 82802 distinct lookups (345.0 per layer), in the arena 0.1981, streamed ahead 0.6932, missed 0.1088; 66194 streamed in (22010.2 MB/chunk, 75051 reads); host 465.2 ms per chunk, of it waiting 434.6 ms +expert stream 1 windows: 480 distinct lookups (480.0 per window), hits 0.5417, read-ahead hits 0.2875 (2 of them late), misses 0.1708 (82.0 per window); read-aheads 176 (176.0 per window, 148 dropped); 375.2 MB read per window (258 reads); host 24.50 ms per window, of it waiting 23.74 ms (+ 1.13 ms after the GPU went on) +stream windows T=1 1 (decode): 141.80 ms each (GPU incl. the host's turns 102.97 ms) +strata metal: window pos 10180 token 760 -> 1156 (top 1156 29.0186, 2nd 6353 15.3183, margin 13.7003) 85.37 ms (GPU 82.30 ms, 1233 dispatches) +strata metal: window pos 10183 token 33362 -> 11 (top 11 27.0843, 2nd 13 19.1119, margin 7.9724) 85.95 ms (GPU 83.36 ms, 1269 dispatches) +strata metal: window pos 10186 token 8523 -> 728 (top 728 26.5949, 2nd 279 17.2821, margin 9.3128) 77.07 ms (GPU 74.19 ms, 1269 dispatches) +strata metal: window pos 10188 token 13 -> 2838 (top 2838 22.9993, 2nd 2302 20.3597, margin 2.6395) 51.83 ms (GPU 49.43 ms, 1269 dispatches) +strata metal: window pos 10189 token 2838 -> 579 (top 579 24.7598, 2nd 42315 19.9055, margin 4.8543) 60.04 ms (GPU 58.52 ms, 1233 dispatches) +strata metal: window pos 10191 token 5315 -> 1292 (top 1292 22.6356, 2nd 15131 21.4028, margin 1.2328) 61.29 ms (GPU 58.39 ms, 1269 dispatches) +strata metal: window pos 10193 token 37444 -> 685 (top 685 25.6466, 2nd 321 15.4891, margin 10.1575) 68.79 ms (GPU 64.26 ms, 1269 dispatches) +strata metal: window pos 10196 token 9859 -> 1092 (top 1092 23.8715, 2nd 62431 16.8946, margin 6.9769) 59.56 ms (GPU 58.04 ms, 1269 dispatches) +strata metal: window pos 10199 token 2908 -> 978 (top 978 27.4026, 2nd 2725 16.9151, margin 10.4874) 59.12 ms (GPU 57.60 ms, 1269 dispatches) +strata metal: window pos 10202 token 13 -> 1061 (top 1061 21.7433, 2nd 2532 20.4921, margin 1.2513) 63.91 ms (GPU 62.29 ms, 1269 dispatches) +strata metal: window pos 10205 token 264 -> 15855 (top 15855 26.3293, 2nd 3471 19.5190, margin 6.8104) 74.30 ms (GPU 72.17 ms, 1269 dispatches) +strata metal: window pos 10207 token 40719 -> 13 (top 13 21.3241, 2nd 11 20.2257, margin 1.0984) 57.98 ms (GPU 55.96 ms, 1269 dispatches) +strata metal: window pos 10208 token 13 -> 353 (top 353 23.4751, 2nd 6558 22.0544, margin 1.4206) 63.10 ms (GPU 60.21 ms, 1269 dispatches) +STRATA_RESULT init_ms=3103.75 prompt_tokens=10180 prefill_ms=23415.08 prefill_tok_s=434.76 decode_tokens=31 decode_ms=908.29 decode_tok_s=34.13 decode_ms_per_tok=29.30 +STRATA_SPEC rounds=13 drafted=23 accepted=18 accept_pct=78.3 tok_per_round=2.38 mtp_ms=113.46 +GEN_IDS: 760 1156 369 33362 11 852 8523 728 13 2838 579 5315 1292 37444 685 321 9859 1092 353 2908 978 3604 13 1061 369 264 15855 40719 13 353 1220 5707 +expert stream 13 windows: 11692 distinct lookups (899.4 per window), hits 0.8885, read-ahead hits 0.0457 (69 of them late), misses 0.0659 (59.2 per window); read-aheads 1303 (100.2 per window, 0 dropped); 232.8 MB read per window (2073 reads); host 17.68 ms per window, of it waiting 17.08 ms (+ 0.92 ms after the GPU went on) +stream windows T=1 1 (decode): 51.60 ms each (GPU incl. the host's turns 49.43 ms) +stream windows T=2 1 (decode): 57.75 ms each (GPU incl. the host's turns 55.96 ms) +stream windows T=3 11 (decode): 68.75 ms each (GPU incl. the host's turns 66.49 ms) diff --git a/results/round39/model-ab/10k-1-candidate.log b/results/round39/model-ab/10k-1-candidate.log new file mode 100644 index 0000000..f899b66 --- /dev/null +++ b/results/round39/model-ab/10k-1-candidate.log @@ -0,0 +1,46 @@ +strata metal stream: experts read from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf.experts (one request per expert) +strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): 290 matrices, 680 MB, made in 153 ms +strata metal stream: locked 2.27 GB of active GPU dense tensors (unpinned 0.34 GB CPU token_embd + 1.28 GB skipped BF16 HC in 3.90 GB file map) +strata metal stream: 1079 dense tensors wrapped zero-copy from file (3.90 GB clean file-backed map, 172.7 ms) +strata metal: decode window on Apple M5: /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf streamed from the SSD (experts) (3.9 GB), 32896 cells, PLE on, ready in 674 ms +strata metal: MTP map: 0.71 GB wired (mlock; 0.00 GB failed), 0.18 GB unused released +strata metal: MTP draft round in one command buffer (embedding type 23 on the GPU: max |GPU - host| 0 over 8 rows) +strata metal: MTP draft layer from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf (0.89 GB, 9 norms folded), draft head over 40525 tokens +strata metal stream: 0.00 GB of the host code's device weight copies given back (unused by the Metal window) +strata metal stream: 2.60 GB allocated before the arena (the process holds 2.24 GB, 5.40 GB resident); RAM 25.8 GB (12.1 GB available, 7.0 GB kept for macOS), GPU working set 19.1 GB +strata metal stream: planned resident footprint 9.6 GB decode (10.7 GB virtual peak during prefill with 1.12 GB transient prompt scratch) = arena 6.5 + 2.6 held before it + 0.5 reserved for the n-gram table; cap 18.8 GB = the lesser of the GPU working set 19.1 GB and the RAM 25.8 GB less 7.0 GB left to macOS and its file cache +strata metal stream: expert arena 6.50 GB, 4129 slots in 2 size classes, 4129 experts read from the profile's top (24576 ranked) in 1.9 s (3.25 GB/s); 8 I/O threads, read-ahead top 6; 1 buffers +strata metal stream: 6.5 GB of the arena wired (mlock) in 888 ms; macOS cannot compress it while the engine idles (STRATA_METAL_WIRE_ARENA=0 turns it off) +strata metal stream: resident footprint after the fill: the process holds 8.7 GB (9.8 GB resident), planned 9.6 GB, cap 18.8 GB +strata metal stream: prompt scratch 1.12 GB mapped in separate anonymous mmap +strata metal: prompts on the matrix units in chunks of 2048 (experts v3, QSA flash, matmul2d layout K4; 1.12 GB of buffers, in a separate transient mmap while a prompt runs) +strata metal: prompt chunk pos 0, 2048 rows: 4799 ms (embd 20 ms, ple 0 ms, graph build 54 ms) +strata metal: prompt chunk pos 2048, 2048 rows: 4564 ms (embd 7 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 4096, 2048 rows: 4646 ms (embd 6 ms, ple 0 ms, graph build 2 ms) +strata metal: prompt chunk pos 6144, 2048 rows: 4838 ms (embd 4 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 8192, 1987 rows: 4882 ms (embd 4 ms, ple 0 ms, graph build 1 ms) +strata metal stream: prompt scratch 1.12 GB unmapped after prefill (phys_footprint 8.89 -> 8.79 GB, resident RSS 9.61 -> 9.53 GB) +strata metal: window pos 10179 token 198 -> 760 (top 760 23.8239, 2nd 35 23.7743, margin 0.0496) 160.61 ms (GPU 118.47 ms, 1233 dispatches) +expert stream prompts 5 chunks: 82802 distinct lookups (345.0 per layer), in the arena 0.1981, streamed ahead 0.6932, missed 0.1088; 66194 streamed in (21984.8 MB/chunk, 74964 reads); host 480.0 ms per chunk, of it waiting 445.9 ms +expert stream 1 windows: 480 distinct lookups (480.0 per window), hits 0.5417, read-ahead hits 0.2875 (4 of them late), misses 0.1708 (82.0 per window); read-aheads 176 (176.0 per window, 235 dropped); 375.2 MB read per window (258 reads); host 26.58 ms per window, of it waiting 25.73 ms (+ 1.32 ms after the GPU went on) +stream windows T=1 1 (decode): 160.34 ms each (GPU incl. the host's turns 118.47 ms) +strata metal: window pos 10180 token 760 -> 1156 (top 1156 29.0186, 2nd 6353 15.3183, margin 13.7003) 92.91 ms (GPU 89.53 ms, 1233 dispatches) +strata metal: window pos 10183 token 33362 -> 11 (top 11 27.0843, 2nd 13 19.1119, margin 7.9724) 88.00 ms (GPU 85.02 ms, 1269 dispatches) +strata metal: window pos 10186 token 8523 -> 728 (top 728 26.5949, 2nd 279 17.2821, margin 9.3128) 79.03 ms (GPU 76.25 ms, 1269 dispatches) +strata metal: window pos 10188 token 13 -> 2838 (top 2838 22.9993, 2nd 2302 20.3597, margin 2.6395) 53.14 ms (GPU 50.53 ms, 1269 dispatches) +strata metal: window pos 10189 token 2838 -> 579 (top 579 24.7598, 2nd 42315 19.9055, margin 4.8543) 62.31 ms (GPU 60.03 ms, 1233 dispatches) +strata metal: window pos 10191 token 5315 -> 1292 (top 1292 22.6356, 2nd 15131 21.4028, margin 1.2328) 62.35 ms (GPU 59.18 ms, 1269 dispatches) +strata metal: window pos 10193 token 37444 -> 685 (top 685 25.6466, 2nd 321 15.4891, margin 10.1575) 67.21 ms (GPU 64.21 ms, 1269 dispatches) +strata metal: window pos 10196 token 9859 -> 1092 (top 1092 23.8715, 2nd 62431 16.8946, margin 6.9769) 60.54 ms (GPU 58.23 ms, 1269 dispatches) +strata metal: window pos 10199 token 2908 -> 978 (top 978 27.4026, 2nd 2725 16.9151, margin 10.4874) 59.16 ms (GPU 56.94 ms, 1269 dispatches) +strata metal: window pos 10202 token 13 -> 1061 (top 1061 21.7433, 2nd 2532 20.4921, margin 1.2513) 61.83 ms (GPU 60.27 ms, 1269 dispatches) +strata metal: window pos 10205 token 264 -> 15855 (top 15855 26.3293, 2nd 3471 19.5190, margin 6.8104) 73.30 ms (GPU 71.74 ms, 1269 dispatches) +strata metal: window pos 10207 token 40719 -> 13 (top 13 21.3241, 2nd 11 20.2257, margin 1.0984) 59.08 ms (GPU 55.78 ms, 1269 dispatches) +strata metal: window pos 10208 token 13 -> 353 (top 353 23.4751, 2nd 6558 22.0544, margin 1.4206) 62.02 ms (GPU 59.79 ms, 1269 dispatches) +STRATA_RESULT init_ms=4106.98 prompt_tokens=10180 prefill_ms=24080.38 prefill_tok_s=422.75 decode_tokens=31 decode_ms=920.70 decode_tok_s=33.67 decode_ms_per_tok=29.70 +STRATA_SPEC rounds=13 drafted=23 accepted=18 accept_pct=78.3 tok_per_round=2.38 mtp_ms=115.79 +GEN_IDS: 760 1156 369 33362 11 852 8523 728 13 2838 579 5315 1292 37444 685 321 9859 1092 353 2908 978 3604 13 1061 369 264 15855 40719 13 353 1220 5707 +expert stream 13 windows: 11692 distinct lookups (899.4 per window), hits 0.8885, read-ahead hits 0.0457 (78 of them late), misses 0.0659 (59.2 per window); read-aheads 1303 (100.2 per window, 0 dropped); 232.8 MB read per window (2073 reads); host 18.28 ms per window, of it waiting 17.66 ms (+ 0.95 ms after the GPU went on) +stream windows T=1 1 (decode): 52.89 ms each (GPU incl. the host's turns 50.53 ms) +stream windows T=2 1 (decode): 58.85 ms each (GPU incl. the host's turns 55.78 ms) +stream windows T=3 11 (decode): 69.65 ms each (GPU incl. the host's turns 67.38 ms) diff --git a/results/round39/model-ab/10k-2-candidate.log b/results/round39/model-ab/10k-2-candidate.log new file mode 100644 index 0000000..21d10ea --- /dev/null +++ b/results/round39/model-ab/10k-2-candidate.log @@ -0,0 +1,46 @@ +strata metal stream: experts read from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf.experts (one request per expert) +strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): 290 matrices, 680 MB, made in 155 ms +strata metal stream: locked 2.27 GB of active GPU dense tensors (unpinned 0.34 GB CPU token_embd + 1.28 GB skipped BF16 HC in 3.90 GB file map) +strata metal stream: 1079 dense tensors wrapped zero-copy from file (3.90 GB clean file-backed map, 178.6 ms) +strata metal: decode window on Apple M5: /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf streamed from the SSD (experts) (3.9 GB), 32896 cells, PLE on, ready in 512 ms +strata metal: MTP map: 0.71 GB wired (mlock; 0.00 GB failed), 0.18 GB unused released +strata metal: MTP draft round in one command buffer (embedding type 23 on the GPU: max |GPU - host| 0 over 8 rows) +strata metal: MTP draft layer from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf (0.89 GB, 9 norms folded), draft head over 40525 tokens +strata metal stream: 0.00 GB of the host code's device weight copies given back (unused by the Metal window) +strata metal stream: 2.60 GB allocated before the arena (the process holds 2.24 GB, 5.40 GB resident); RAM 25.8 GB (12.3 GB available, 7.0 GB kept for macOS), GPU working set 19.1 GB +strata metal stream: planned resident footprint 9.6 GB decode (10.7 GB virtual peak during prefill with 1.12 GB transient prompt scratch) = arena 6.5 + 2.6 held before it + 0.5 reserved for the n-gram table; cap 18.8 GB = the lesser of the GPU working set 19.1 GB and the RAM 25.8 GB less 7.0 GB left to macOS and its file cache +strata metal stream: expert arena 6.50 GB, 4129 slots in 2 size classes, 4129 experts read from the profile's top (24576 ranked) in 1.5 s (4.30 GB/s); 8 I/O threads, read-ahead top 6; 1 buffers +strata metal stream: 6.5 GB of the arena wired (mlock) in 796 ms; macOS cannot compress it while the engine idles (STRATA_METAL_WIRE_ARENA=0 turns it off) +strata metal stream: resident footprint after the fill: the process holds 8.7 GB (10.3 GB resident), planned 9.6 GB, cap 18.8 GB +strata metal stream: prompt scratch 1.12 GB mapped in separate anonymous mmap +strata metal: prompts on the matrix units in chunks of 2048 (experts v3, QSA flash, matmul2d layout K4; 1.12 GB of buffers, in a separate transient mmap while a prompt runs) +strata metal: prompt chunk pos 0, 2048 rows: 4804 ms (embd 24 ms, ple 0 ms, graph build 32 ms) +strata metal: prompt chunk pos 2048, 2048 rows: 4577 ms (embd 10 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 4096, 2048 rows: 4704 ms (embd 4 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 6144, 2048 rows: 4938 ms (embd 4 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 8192, 1987 rows: 4905 ms (embd 4 ms, ple 0 ms, graph build 1 ms) +strata metal stream: prompt scratch 1.12 GB unmapped after prefill (phys_footprint 8.89 -> 8.79 GB, resident RSS 9.64 -> 9.56 GB) +strata metal: window pos 10179 token 198 -> 760 (top 760 23.8239, 2nd 35 23.7743, margin 0.0496) 142.09 ms (GPU 104.11 ms, 1233 dispatches) +expert stream prompts 5 chunks: 82802 distinct lookups (345.0 per layer), in the arena 0.1981, streamed ahead 0.6932, missed 0.1088; 66194 streamed in (22007.9 MB/chunk, 75042 reads); host 468.8 ms per chunk, of it waiting 437.6 ms +expert stream 1 windows: 480 distinct lookups (480.0 per window), hits 0.5417, read-ahead hits 0.2875 (11 of them late), misses 0.1708 (82.0 per window); read-aheads 176 (176.0 per window, 157 dropped); 375.2 MB read per window (258 reads); host 24.99 ms per window, of it waiting 24.29 ms (+ 1.31 ms after the GPU went on) +stream windows T=1 1 (decode): 141.86 ms each (GPU incl. the host's turns 104.11 ms) +strata metal: window pos 10180 token 760 -> 1156 (top 1156 29.0186, 2nd 6353 15.3183, margin 13.7003) 92.75 ms (GPU 89.78 ms, 1233 dispatches) +strata metal: window pos 10183 token 33362 -> 11 (top 11 27.0843, 2nd 13 19.1119, margin 7.9724) 86.09 ms (GPU 83.86 ms, 1269 dispatches) +strata metal: window pos 10186 token 8523 -> 728 (top 728 26.5949, 2nd 279 17.2821, margin 9.3128) 78.57 ms (GPU 74.37 ms, 1269 dispatches) +strata metal: window pos 10188 token 13 -> 2838 (top 2838 22.9993, 2nd 2302 20.3597, margin 2.6395) 52.04 ms (GPU 49.68 ms, 1269 dispatches) +strata metal: window pos 10189 token 2838 -> 579 (top 579 24.7598, 2nd 42315 19.9055, margin 4.8543) 61.74 ms (GPU 59.61 ms, 1233 dispatches) +strata metal: window pos 10191 token 5315 -> 1292 (top 1292 22.6356, 2nd 15131 21.4028, margin 1.2328) 62.43 ms (GPU 58.29 ms, 1269 dispatches) +strata metal: window pos 10193 token 37444 -> 685 (top 685 25.6466, 2nd 321 15.4891, margin 10.1575) 66.66 ms (GPU 63.97 ms, 1269 dispatches) +strata metal: window pos 10196 token 9859 -> 1092 (top 1092 23.8715, 2nd 62431 16.8946, margin 6.9769) 59.43 ms (GPU 57.90 ms, 1269 dispatches) +strata metal: window pos 10199 token 2908 -> 978 (top 978 27.4026, 2nd 2725 16.9151, margin 10.4874) 58.85 ms (GPU 55.51 ms, 1269 dispatches) +strata metal: window pos 10202 token 13 -> 1061 (top 1061 21.7433, 2nd 2532 20.4921, margin 1.2513) 62.63 ms (GPU 61.09 ms, 1269 dispatches) +strata metal: window pos 10205 token 264 -> 15855 (top 15855 26.3293, 2nd 3471 19.5190, margin 6.8104) 73.56 ms (GPU 71.70 ms, 1269 dispatches) +strata metal: window pos 10207 token 40719 -> 13 (top 13 21.3241, 2nd 11 20.2257, margin 1.0984) 58.30 ms (GPU 55.54 ms, 1269 dispatches) +strata metal: window pos 10208 token 13 -> 353 (top 353 23.4751, 2nd 6558 22.0544, margin 1.4206) 62.86 ms (GPU 60.14 ms, 1269 dispatches) +STRATA_RESULT init_ms=3115.54 prompt_tokens=10180 prefill_ms=24248.45 prefill_tok_s=419.82 decode_tokens=31 decode_ms=915.71 decode_tok_s=33.85 decode_ms_per_tok=29.54 +STRATA_SPEC rounds=13 drafted=23 accepted=18 accept_pct=78.3 tok_per_round=2.38 mtp_ms=113.97 +GEN_IDS: 760 1156 369 33362 11 852 8523 728 13 2838 579 5315 1292 37444 685 321 9859 1092 353 2908 978 3604 13 1061 369 264 15855 40719 13 353 1220 5707 +expert stream 13 windows: 11692 distinct lookups (899.4 per window), hits 0.8885, read-ahead hits 0.0457 (66 of them late), misses 0.0659 (59.2 per window); read-aheads 1303 (100.2 per window, 0 dropped); 232.8 MB read per window (2073 reads); host 17.81 ms per window, of it waiting 17.22 ms (+ 0.89 ms after the GPU went on) +stream windows T=1 1 (decode): 51.81 ms each (GPU incl. the host's turns 49.68 ms) +stream windows T=2 1 (decode): 58.07 ms each (GPU incl. the host's turns 55.54 ms) +stream windows T=3 11 (decode): 69.40 ms each (GPU incl. the host's turns 66.93 ms) diff --git a/results/round39/model-ab/10k-3-baseline.log b/results/round39/model-ab/10k-3-baseline.log new file mode 100644 index 0000000..f33bc48 --- /dev/null +++ b/results/round39/model-ab/10k-3-baseline.log @@ -0,0 +1,46 @@ +strata metal stream: experts read from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf.experts (one request per expert) +strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): 290 matrices, 680 MB, made in 146 ms +strata metal stream: locked 2.27 GB of active GPU dense tensors (unpinned 0.34 GB CPU token_embd + 1.28 GB skipped BF16 HC in 3.90 GB file map) +strata metal stream: 1079 dense tensors wrapped zero-copy from file (3.90 GB clean file-backed map, 177.6 ms) +strata metal: decode window on Apple M5: /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf streamed from the SSD (experts) (3.9 GB), 32896 cells, PLE on, ready in 657 ms +strata metal: MTP map: 0.71 GB wired (mlock; 0.00 GB failed), 0.18 GB unused released +strata metal: MTP draft round in one command buffer (embedding type 23 on the GPU: max |GPU - host| 0 over 8 rows) +strata metal: MTP draft layer from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf (0.89 GB, 9 norms folded), draft head over 40525 tokens +strata metal stream: 0.00 GB of the host code's device weight copies given back (unused by the Metal window) +strata metal stream: 2.60 GB allocated before the arena (the process holds 2.24 GB, 5.40 GB resident); RAM 25.8 GB (12.1 GB available, 7.0 GB kept for macOS), GPU working set 19.1 GB +strata metal stream: planned resident footprint 9.6 GB decode (10.7 GB virtual peak during prefill with 1.12 GB transient prompt scratch) = arena 6.5 + 2.6 held before it + 0.5 reserved for the n-gram table; cap 18.8 GB = the lesser of the GPU working set 19.1 GB and the RAM 25.8 GB less 7.0 GB left to macOS and its file cache +strata metal stream: expert arena 6.50 GB, 4129 slots in 2 size classes, 4129 experts read from the profile's top (24576 ranked) in 2.0 s (3.06 GB/s); 8 I/O threads, read-ahead top 6; 1 buffers +strata metal stream: 6.5 GB of the arena wired (mlock) in 874 ms; macOS cannot compress it while the engine idles (STRATA_METAL_WIRE_ARENA=0 turns it off) +strata metal stream: resident footprint after the fill: the process holds 8.7 GB (10.0 GB resident), planned 9.6 GB, cap 18.8 GB +strata metal stream: prompt scratch 1.12 GB mapped in separate anonymous mmap +strata metal: prompts on the matrix units in chunks of 2048 (experts v3, QSA flash, matmul2d layout K4; 1.12 GB of buffers, in a separate transient mmap while a prompt runs) +strata metal: prompt chunk pos 0, 2048 rows: 5274 ms (embd 19 ms, ple 4 ms, graph build 62 ms) +strata metal: prompt chunk pos 2048, 2048 rows: 5037 ms (embd 9 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 4096, 2048 rows: 5278 ms (embd 4 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 6144, 2048 rows: 5340 ms (embd 4 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 8192, 1987 rows: 5341 ms (embd 5 ms, ple 0 ms, graph build 1 ms) +strata metal stream: prompt scratch 1.12 GB unmapped after prefill (phys_footprint 8.89 -> 8.79 GB, resident RSS 9.61 -> 9.53 GB) +strata metal: window pos 10179 token 198 -> 760 (top 760 23.8239, 2nd 35 23.7743, margin 0.0496) 153.32 ms (GPU 116.82 ms, 1233 dispatches) +expert stream prompts 5 chunks: 82802 distinct lookups (345.0 per layer), in the arena 0.1981, streamed ahead 0.6932, missed 0.1088; 66194 streamed in (22038.1 MB/chunk, 75142 reads); host 469.2 ms per chunk, of it waiting 440.0 ms +expert stream 1 windows: 480 distinct lookups (480.0 per window), hits 0.5417, read-ahead hits 0.2875 (4 of them late), misses 0.1708 (82.0 per window); read-aheads 176 (176.0 per window, 57 dropped); 375.2 MB read per window (258 reads); host 26.50 ms per window, of it waiting 25.70 ms (+ 1.29 ms after the GPU went on) +stream windows T=1 1 (decode): 153.06 ms each (GPU incl. the host's turns 116.82 ms) +strata metal: window pos 10180 token 760 -> 1156 (top 1156 29.0186, 2nd 6353 15.3183, margin 13.7003) 90.18 ms (GPU 87.55 ms, 1233 dispatches) +strata metal: window pos 10183 token 33362 -> 11 (top 11 27.0843, 2nd 13 19.1119, margin 7.9724) 85.62 ms (GPU 83.06 ms, 1269 dispatches) +strata metal: window pos 10186 token 8523 -> 728 (top 728 26.5949, 2nd 279 17.2821, margin 9.3128) 76.34 ms (GPU 73.74 ms, 1269 dispatches) +strata metal: window pos 10188 token 13 -> 2838 (top 2838 22.9993, 2nd 2302 20.3597, margin 2.6395) 51.51 ms (GPU 49.35 ms, 1269 dispatches) +strata metal: window pos 10189 token 2838 -> 579 (top 579 24.7598, 2nd 42315 19.9055, margin 4.8543) 60.41 ms (GPU 58.53 ms, 1233 dispatches) +strata metal: window pos 10191 token 5315 -> 1292 (top 1292 22.6356, 2nd 15131 21.4028, margin 1.2328) 60.86 ms (GPU 58.11 ms, 1269 dispatches) +strata metal: window pos 10193 token 37444 -> 685 (top 685 25.6466, 2nd 321 15.4891, margin 10.1575) 67.87 ms (GPU 65.19 ms, 1269 dispatches) +strata metal: window pos 10196 token 9859 -> 1092 (top 1092 23.8715, 2nd 62431 16.8946, margin 6.9769) 63.95 ms (GPU 61.82 ms, 1269 dispatches) +strata metal: window pos 10199 token 2908 -> 978 (top 978 27.4026, 2nd 2725 16.9151, margin 10.4874) 62.93 ms (GPU 60.77 ms, 1269 dispatches) +strata metal: window pos 10202 token 13 -> 1061 (top 1061 21.7433, 2nd 2532 20.4921, margin 1.2513) 71.99 ms (GPU 69.75 ms, 1269 dispatches) +strata metal: window pos 10205 token 264 -> 15855 (top 15855 26.3293, 2nd 3471 19.5190, margin 6.8104) 80.68 ms (GPU 78.47 ms, 1269 dispatches) +strata metal: window pos 10207 token 40719 -> 13 (top 13 21.3241, 2nd 11 20.2257, margin 1.0984) 62.65 ms (GPU 59.87 ms, 1269 dispatches) +strata metal: window pos 10208 token 13 -> 353 (top 353 23.4751, 2nd 6558 22.0544, margin 1.4206) 72.28 ms (GPU 69.26 ms, 1269 dispatches) +STRATA_RESULT init_ms=4189.01 prompt_tokens=10180 prefill_ms=26637.98 prefill_tok_s=382.16 decode_tokens=31 decode_ms=949.61 decode_tok_s=32.64 decode_ms_per_tok=30.63 +STRATA_SPEC rounds=13 drafted=23 accepted=18 accept_pct=78.3 tok_per_round=2.38 mtp_ms=118.45 +GEN_IDS: 760 1156 369 33362 11 852 8523 728 13 2838 579 5315 1292 37444 685 321 9859 1092 353 2908 978 3604 13 1061 369 264 15855 40719 13 353 1220 5707 +expert stream 13 windows: 11692 distinct lookups (899.4 per window), hits 0.8885, read-ahead hits 0.0457 (64 of them late), misses 0.0659 (59.2 per window); read-aheads 1303 (100.2 per window, 0 dropped); 232.8 MB read per window (2073 reads); host 17.86 ms per window, of it waiting 17.32 ms (+ 0.88 ms after the GPU went on) +stream windows T=1 1 (decode): 51.30 ms each (GPU incl. the host's turns 49.35 ms) +stream windows T=2 1 (decode): 62.41 ms each (GPU incl. the host's turns 59.87 ms) +stream windows T=3 11 (decode): 71.89 ms each (GPU incl. the host's turns 69.66 ms) diff --git a/results/round39/model-ab/10k-4-baseline.log b/results/round39/model-ab/10k-4-baseline.log new file mode 100644 index 0000000..183ecd9 --- /dev/null +++ b/results/round39/model-ab/10k-4-baseline.log @@ -0,0 +1,46 @@ +strata metal stream: experts read from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf.experts (one request per expert) +strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): 290 matrices, 680 MB, made in 136 ms +strata metal stream: locked 2.27 GB of active GPU dense tensors (unpinned 0.34 GB CPU token_embd + 1.28 GB skipped BF16 HC in 3.90 GB file map) +strata metal stream: 1079 dense tensors wrapped zero-copy from file (3.90 GB clean file-backed map, 130.3 ms) +strata metal: decode window on Apple M5: /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf streamed from the SSD (experts) (3.9 GB), 32896 cells, PLE on, ready in 407 ms +strata metal: MTP map: 0.71 GB wired (mlock; 0.00 GB failed), 0.18 GB unused released +strata metal: MTP draft round in one command buffer (embedding type 23 on the GPU: max |GPU - host| 0 over 8 rows) +strata metal: MTP draft layer from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf (0.89 GB, 9 norms folded), draft head over 40525 tokens +strata metal stream: 0.00 GB of the host code's device weight copies given back (unused by the Metal window) +strata metal stream: 2.60 GB allocated before the arena (the process holds 2.24 GB, 5.40 GB resident); RAM 25.8 GB (12.1 GB available, 7.0 GB kept for macOS), GPU working set 19.1 GB +strata metal stream: planned resident footprint 9.6 GB decode (10.7 GB virtual peak during prefill with 1.12 GB transient prompt scratch) = arena 6.5 + 2.6 held before it + 0.5 reserved for the n-gram table; cap 18.8 GB = the lesser of the GPU working set 19.1 GB and the RAM 25.8 GB less 7.0 GB left to macOS and its file cache +strata metal stream: expert arena 6.50 GB, 4129 slots in 2 size classes, 4129 experts read from the profile's top (24576 ranked) in 1.5 s (4.21 GB/s); 8 I/O threads, read-ahead top 6; 1 buffers +strata metal stream: 6.5 GB of the arena wired (mlock) in 741 ms; macOS cannot compress it while the engine idles (STRATA_METAL_WIRE_ARENA=0 turns it off) +strata metal stream: resident footprint after the fill: the process holds 8.7 GB (10.3 GB resident), planned 9.6 GB, cap 18.8 GB +strata metal stream: prompt scratch 1.12 GB mapped in separate anonymous mmap +strata metal: prompts on the matrix units in chunks of 2048 (experts v3, QSA flash, matmul2d layout K4; 1.12 GB of buffers, in a separate transient mmap while a prompt runs) +strata metal: prompt chunk pos 0, 2048 rows: 5235 ms (embd 16 ms, ple 0 ms, graph build 36 ms) +strata metal: prompt chunk pos 2048, 2048 rows: 5593 ms (embd 15 ms, ple 0 ms, graph build 2 ms) +strata metal: prompt chunk pos 4096, 2048 rows: 5574 ms (embd 5 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 6144, 2048 rows: 5884 ms (embd 4 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 8192, 1987 rows: 5821 ms (embd 7 ms, ple 0 ms, graph build 2 ms) +strata metal stream: prompt scratch 1.12 GB unmapped after prefill (phys_footprint 8.89 -> 8.79 GB, resident RSS 9.63 -> 9.55 GB) +strata metal: window pos 10179 token 198 -> 760 (top 760 23.8239, 2nd 35 23.7743, margin 0.0496) 168.29 ms (GPU 130.09 ms, 1233 dispatches) +expert stream prompts 5 chunks: 82802 distinct lookups (345.0 per layer), in the arena 0.1981, streamed ahead 0.6932, missed 0.1088; 66194 streamed in (22014.1 MB/chunk, 75063 reads); host 519.4 ms per chunk, of it waiting 486.1 ms +expert stream 1 windows: 480 distinct lookups (480.0 per window), hits 0.5417, read-ahead hits 0.2875 (2 of them late), misses 0.1708 (82.0 per window); read-aheads 176 (176.0 per window, 136 dropped); 375.2 MB read per window (258 reads); host 28.04 ms per window, of it waiting 27.14 ms (+ 1.41 ms after the GPU went on) +stream windows T=1 1 (decode): 167.95 ms each (GPU incl. the host's turns 130.09 ms) +strata metal: window pos 10180 token 760 -> 1156 (top 1156 29.0186, 2nd 6353 15.3183, margin 13.7003) 96.12 ms (GPU 93.38 ms, 1233 dispatches) +strata metal: window pos 10183 token 33362 -> 11 (top 11 27.0843, 2nd 13 19.1119, margin 7.9724) 88.32 ms (GPU 85.18 ms, 1269 dispatches) +strata metal: window pos 10186 token 8523 -> 728 (top 728 26.5949, 2nd 279 17.2821, margin 9.3128) 78.47 ms (GPU 75.45 ms, 1269 dispatches) +strata metal: window pos 10188 token 13 -> 2838 (top 2838 22.9993, 2nd 2302 20.3597, margin 2.6395) 53.79 ms (GPU 51.22 ms, 1269 dispatches) +strata metal: window pos 10189 token 2838 -> 579 (top 579 24.7598, 2nd 42315 19.9055, margin 4.8543) 62.41 ms (GPU 60.12 ms, 1233 dispatches) +strata metal: window pos 10191 token 5315 -> 1292 (top 1292 22.6356, 2nd 15131 21.4028, margin 1.2328) 64.74 ms (GPU 61.62 ms, 1269 dispatches) +strata metal: window pos 10193 token 37444 -> 685 (top 685 25.6466, 2nd 321 15.4891, margin 10.1575) 72.78 ms (GPU 69.52 ms, 1269 dispatches) +strata metal: window pos 10196 token 9859 -> 1092 (top 1092 23.8715, 2nd 62431 16.8946, margin 6.9769) 70.04 ms (GPU 67.63 ms, 1269 dispatches) +strata metal: window pos 10199 token 2908 -> 978 (top 978 27.4026, 2nd 2725 16.9151, margin 10.4874) 67.69 ms (GPU 65.37 ms, 1269 dispatches) +strata metal: window pos 10202 token 13 -> 1061 (top 1061 21.7433, 2nd 2532 20.4921, margin 1.2513) 73.24 ms (GPU 70.84 ms, 1269 dispatches) +strata metal: window pos 10205 token 264 -> 15855 (top 15855 26.3293, 2nd 3471 19.5190, margin 6.8104) 84.76 ms (GPU 82.33 ms, 1269 dispatches) +strata metal: window pos 10207 token 40719 -> 13 (top 13 21.3241, 2nd 11 20.2257, margin 1.0984) 65.95 ms (GPU 62.92 ms, 1269 dispatches) +strata metal: window pos 10208 token 13 -> 353 (top 353 23.4751, 2nd 6558 22.0544, margin 1.4206) 74.43 ms (GPU 71.31 ms, 1269 dispatches) +STRATA_RESULT init_ms=2935.91 prompt_tokens=10180 prefill_ms=28519.16 prefill_tok_s=356.95 decode_tokens=31 decode_ms=996.64 decode_tok_s=31.10 decode_ms_per_tok=32.15 +STRATA_SPEC rounds=13 drafted=23 accepted=18 accept_pct=78.3 tok_per_round=2.38 mtp_ms=125.06 +GEN_IDS: 760 1156 369 33362 11 852 8523 728 13 2838 579 5315 1292 37444 685 321 9859 1092 353 2908 978 3604 13 1061 369 264 15855 40719 13 353 1220 5707 +expert stream 13 windows: 11692 distinct lookups (899.4 per window), hits 0.8885, read-ahead hits 0.0457 (62 of them late), misses 0.0659 (59.2 per window); read-aheads 1303 (100.2 per window, 0 dropped); 232.8 MB read per window (2073 reads); host 18.55 ms per window, of it waiting 17.91 ms (+ 0.99 ms after the GPU went on) +stream windows T=1 1 (decode): 53.53 ms each (GPU incl. the host's turns 51.22 ms) +stream windows T=2 1 (decode): 65.69 ms each (GPU incl. the host's turns 62.92 ms) +stream windows T=3 11 (decode): 75.48 ms each (GPU incl. the host's turns 72.98 ms) diff --git a/results/round39/model-ab/10k-5-candidate.log b/results/round39/model-ab/10k-5-candidate.log new file mode 100644 index 0000000..4bbe146 --- /dev/null +++ b/results/round39/model-ab/10k-5-candidate.log @@ -0,0 +1,46 @@ +strata metal stream: experts read from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf.experts (one request per expert) +strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): 290 matrices, 680 MB, made in 136 ms +strata metal stream: locked 2.27 GB of active GPU dense tensors (unpinned 0.34 GB CPU token_embd + 1.28 GB skipped BF16 HC in 3.90 GB file map) +strata metal stream: 1079 dense tensors wrapped zero-copy from file (3.90 GB clean file-backed map, 123.2 ms) +strata metal: decode window on Apple M5: /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf streamed from the SSD (experts) (3.9 GB), 32896 cells, PLE on, ready in 533 ms +strata metal: MTP map: 0.71 GB wired (mlock; 0.00 GB failed), 0.18 GB unused released +strata metal: MTP draft round in one command buffer (embedding type 23 on the GPU: max |GPU - host| 0 over 8 rows) +strata metal: MTP draft layer from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf (0.89 GB, 9 norms folded), draft head over 40525 tokens +strata metal stream: 0.00 GB of the host code's device weight copies given back (unused by the Metal window) +strata metal stream: 2.60 GB allocated before the arena (the process holds 2.24 GB, 5.40 GB resident); RAM 25.8 GB (12.1 GB available, 7.0 GB kept for macOS), GPU working set 19.1 GB +strata metal stream: planned resident footprint 9.6 GB decode (10.7 GB virtual peak during prefill with 1.12 GB transient prompt scratch) = arena 6.5 + 2.6 held before it + 0.5 reserved for the n-gram table; cap 18.8 GB = the lesser of the GPU working set 19.1 GB and the RAM 25.8 GB less 7.0 GB left to macOS and its file cache +strata metal stream: expert arena 6.50 GB, 4129 slots in 2 size classes, 4129 experts read from the profile's top (24576 ranked) in 1.7 s (3.64 GB/s); 8 I/O threads, read-ahead top 6; 1 buffers +strata metal stream: 6.5 GB of the arena wired (mlock) in 961 ms; macOS cannot compress it while the engine idles (STRATA_METAL_WIRE_ARENA=0 turns it off) +strata metal stream: resident footprint after the fill: the process holds 8.7 GB (10.2 GB resident), planned 9.6 GB, cap 18.8 GB +strata metal stream: prompt scratch 1.12 GB mapped in separate anonymous mmap +strata metal: prompts on the matrix units in chunks of 2048 (experts v3, QSA flash, matmul2d layout K4; 1.12 GB of buffers, in a separate transient mmap while a prompt runs) +strata metal: prompt chunk pos 0, 2048 rows: 5804 ms (embd 24 ms, ple 0 ms, graph build 64 ms) +strata metal: prompt chunk pos 2048, 2048 rows: 5869 ms (embd 15 ms, ple 1 ms, graph build 1 ms) +strata metal: prompt chunk pos 4096, 2048 rows: 6004 ms (embd 7 ms, ple 0 ms, graph build 2 ms) +strata metal: prompt chunk pos 6144, 2048 rows: 6145 ms (embd 7 ms, ple 0 ms, graph build 2 ms) +strata metal: prompt chunk pos 8192, 1987 rows: 6153 ms (embd 6 ms, ple 0 ms, graph build 2 ms) +strata metal stream: prompt scratch 1.12 GB unmapped after prefill (phys_footprint 8.89 -> 8.79 GB, resident RSS 9.61 -> 9.53 GB) +strata metal: window pos 10179 token 198 -> 760 (top 760 23.8239, 2nd 35 23.7743, margin 0.0496) 172.42 ms (GPU 130.96 ms, 1233 dispatches) +expert stream prompts 5 chunks: 82802 distinct lookups (345.0 per layer), in the arena 0.1981, streamed ahead 0.6932, missed 0.1088; 66194 streamed in (22039.9 MB/chunk, 75149 reads); host 486.5 ms per chunk, of it waiting 449.2 ms +expert stream 1 windows: 480 distinct lookups (480.0 per window), hits 0.5417, read-ahead hits 0.2875 (1 of them late), misses 0.1708 (82.0 per window); read-aheads 176 (176.0 per window, 50 dropped); 375.2 MB read per window (258 reads); host 26.93 ms per window, of it waiting 25.98 ms (+ 1.51 ms after the GPU went on) +stream windows T=1 1 (decode): 172.13 ms each (GPU incl. the host's turns 130.96 ms) +strata metal: window pos 10180 token 760 -> 1156 (top 1156 29.0186, 2nd 6353 15.3183, margin 13.7003) 99.86 ms (GPU 96.92 ms, 1233 dispatches) +strata metal: window pos 10183 token 33362 -> 11 (top 11 27.0843, 2nd 13 19.1119, margin 7.9724) 88.55 ms (GPU 85.31 ms, 1269 dispatches) +strata metal: window pos 10186 token 8523 -> 728 (top 728 26.5949, 2nd 279 17.2821, margin 9.3128) 79.53 ms (GPU 76.39 ms, 1269 dispatches) +strata metal: window pos 10188 token 13 -> 2838 (top 2838 22.9993, 2nd 2302 20.3597, margin 2.6395) 53.83 ms (GPU 51.07 ms, 1269 dispatches) +strata metal: window pos 10189 token 2838 -> 579 (top 579 24.7598, 2nd 42315 19.9055, margin 4.8543) 64.23 ms (GPU 61.80 ms, 1233 dispatches) +strata metal: window pos 10191 token 5315 -> 1292 (top 1292 22.6356, 2nd 15131 21.4028, margin 1.2328) 68.08 ms (GPU 64.84 ms, 1269 dispatches) +strata metal: window pos 10193 token 37444 -> 685 (top 685 25.6466, 2nd 321 15.4891, margin 10.1575) 74.50 ms (GPU 71.45 ms, 1269 dispatches) +strata metal: window pos 10196 token 9859 -> 1092 (top 1092 23.8715, 2nd 62431 16.8946, margin 6.9769) 70.08 ms (GPU 67.67 ms, 1269 dispatches) +strata metal: window pos 10199 token 2908 -> 978 (top 978 27.4026, 2nd 2725 16.9151, margin 10.4874) 69.01 ms (GPU 66.63 ms, 1269 dispatches) +strata metal: window pos 10202 token 13 -> 1061 (top 1061 21.7433, 2nd 2532 20.4921, margin 1.2513) 75.82 ms (GPU 73.16 ms, 1269 dispatches) +strata metal: window pos 10205 token 264 -> 15855 (top 15855 26.3293, 2nd 3471 19.5190, margin 6.8104) 88.19 ms (GPU 85.33 ms, 1269 dispatches) +strata metal: window pos 10207 token 40719 -> 13 (top 13 21.3241, 2nd 11 20.2257, margin 1.0984) 69.23 ms (GPU 65.43 ms, 1269 dispatches) +strata metal: window pos 10208 token 13 -> 353 (top 353 23.4751, 2nd 6558 22.0544, margin 1.4206) 78.49 ms (GPU 74.75 ms, 1269 dispatches) +STRATA_RESULT init_ms=3740.62 prompt_tokens=10180 prefill_ms=30410.22 prefill_tok_s=334.76 decode_tokens=31 decode_ms=1024.55 decode_tok_s=30.26 decode_ms_per_tok=33.05 +STRATA_SPEC rounds=13 drafted=23 accepted=18 accept_pct=78.3 tok_per_round=2.38 mtp_ms=123.96 +GEN_IDS: 760 1156 369 33362 11 852 8523 728 13 2838 579 5315 1292 37444 685 321 9859 1092 353 2908 978 3604 13 1061 369 264 15855 40719 13 353 1220 5707 +expert stream 13 windows: 11692 distinct lookups (899.4 per window), hits 0.8885, read-ahead hits 0.0457 (68 of them late), misses 0.0659 (59.2 per window); read-aheads 1303 (100.2 per window, 0 dropped); 232.8 MB read per window (2073 reads); host 18.60 ms per window, of it waiting 17.93 ms (+ 1.08 ms after the GPU went on) +stream windows T=1 1 (decode): 53.57 ms each (GPU incl. the host's turns 51.07 ms) +stream windows T=2 1 (decode): 68.94 ms each (GPU incl. the host's turns 65.43 ms) +stream windows T=3 11 (decode): 77.59 ms each (GPU incl. the host's turns 74.93 ms) diff --git a/results/round39/model-ab/10k-warm-0-baseline.log b/results/round39/model-ab/10k-warm-0-baseline.log new file mode 100644 index 0000000..1706018 --- /dev/null +++ b/results/round39/model-ab/10k-warm-0-baseline.log @@ -0,0 +1,46 @@ +strata metal stream: experts read from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf.experts (one request per expert) +strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): 290 matrices, 680 MB, made in 130 ms +strata metal stream: locked 2.27 GB of active GPU dense tensors (unpinned 0.34 GB CPU token_embd + 1.28 GB skipped BF16 HC in 3.90 GB file map) +strata metal stream: 1079 dense tensors wrapped zero-copy from file (3.90 GB clean file-backed map, 111.3 ms) +strata metal: decode window on Apple M5: /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf streamed from the SSD (experts) (3.9 GB), 32896 cells, PLE on, ready in 5323 ms +strata metal: MTP map: 0.71 GB wired (mlock; 0.00 GB failed), 0.18 GB unused released +strata metal: MTP draft round in one command buffer (embedding type 23 on the GPU: max |GPU - host| 0 over 8 rows) +strata metal: MTP draft layer from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf (0.89 GB, 9 norms folded), draft head over 40525 tokens +strata metal stream: 0.00 GB of the host code's device weight copies given back (unused by the Metal window) +strata metal stream: 2.60 GB allocated before the arena (the process holds 2.24 GB, 5.40 GB resident); RAM 25.8 GB (12.0 GB available, 7.0 GB kept for macOS), GPU working set 19.1 GB +strata metal stream: planned resident footprint 9.6 GB decode (10.7 GB virtual peak during prefill with 1.12 GB transient prompt scratch) = arena 6.5 + 2.6 held before it + 0.5 reserved for the n-gram table; cap 18.8 GB = the lesser of the GPU working set 19.1 GB and the RAM 25.8 GB less 7.0 GB left to macOS and its file cache +strata metal stream: expert arena 6.50 GB, 4129 slots in 2 size classes, 4129 experts read from the profile's top (24576 ranked) in 0.7 s (8.36 GB/s); 8 I/O threads, read-ahead top 6; 1 buffers +strata metal stream: 6.5 GB of the arena wired (mlock) in 149 ms; macOS cannot compress it while the engine idles (STRATA_METAL_WIRE_ARENA=0 turns it off) +strata metal stream: resident footprint after the fill: the process holds 8.7 GB (11.9 GB resident), planned 9.6 GB, cap 18.8 GB +strata metal stream: prompt scratch 1.12 GB mapped in separate anonymous mmap +strata metal: prompts on the matrix units in chunks of 2048 (experts v3, QSA flash, matmul2d layout K4; 1.12 GB of buffers, in a separate transient mmap while a prompt runs) +strata metal: prompt chunk pos 0, 2048 rows: 4747 ms (embd 19 ms, ple 0 ms, graph build 87 ms) +strata metal: prompt chunk pos 2048, 2048 rows: 4032 ms (embd 8 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 4096, 2048 rows: 4077 ms (embd 3 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 6144, 2048 rows: 4202 ms (embd 3 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 8192, 1987 rows: 4235 ms (embd 3 ms, ple 0 ms, graph build 1 ms) +strata metal stream: prompt scratch 1.12 GB unmapped after prefill (phys_footprint 8.89 -> 8.79 GB, resident RSS 10.30 -> 10.22 GB) +strata metal: window pos 10179 token 198 -> 760 (top 760 23.8239, 2nd 35 23.7743, margin 0.0496) 121.14 ms (GPU 91.32 ms, 1233 dispatches) +expert stream prompts 5 chunks: 82802 distinct lookups (345.0 per layer), in the arena 0.1981, streamed ahead 0.6932, missed 0.1088; 66194 streamed in (21914.0 MB/chunk, 74722 reads); host 527.5 ms per chunk, of it waiting 500.7 ms +expert stream 1 windows: 480 distinct lookups (480.0 per window), hits 0.5417, read-ahead hits 0.2875 (15 of them late), misses 0.1708 (82.0 per window); read-aheads 176 (176.0 per window, 477 dropped); 375.2 MB read per window (258 reads); host 23.40 ms per window, of it waiting 22.85 ms (+ 0.95 ms after the GPU went on) +stream windows T=1 1 (decode): 120.92 ms each (GPU incl. the host's turns 91.32 ms) +strata metal: window pos 10180 token 760 -> 1156 (top 1156 29.0186, 2nd 6353 15.3183, margin 13.7003) 85.55 ms (GPU 83.10 ms, 1233 dispatches) +strata metal: window pos 10183 token 33362 -> 11 (top 11 27.0843, 2nd 13 19.1119, margin 7.9724) 84.11 ms (GPU 81.84 ms, 1269 dispatches) +strata metal: window pos 10186 token 8523 -> 728 (top 728 26.5949, 2nd 279 17.2821, margin 9.3128) 75.88 ms (GPU 73.65 ms, 1269 dispatches) +strata metal: window pos 10188 token 13 -> 2838 (top 2838 22.9993, 2nd 2302 20.3597, margin 2.6395) 52.19 ms (GPU 50.55 ms, 1269 dispatches) +strata metal: window pos 10189 token 2838 -> 579 (top 579 24.7598, 2nd 42315 19.9055, margin 4.8543) 59.31 ms (GPU 57.80 ms, 1233 dispatches) +strata metal: window pos 10191 token 5315 -> 1292 (top 1292 22.6356, 2nd 15131 21.4028, margin 1.2328) 59.60 ms (GPU 57.39 ms, 1269 dispatches) +strata metal: window pos 10193 token 37444 -> 685 (top 685 25.6466, 2nd 321 15.4891, margin 10.1575) 65.09 ms (GPU 62.87 ms, 1269 dispatches) +strata metal: window pos 10196 token 9859 -> 1092 (top 1092 23.8715, 2nd 62431 16.8946, margin 6.9769) 60.05 ms (GPU 58.51 ms, 1269 dispatches) +strata metal: window pos 10199 token 2908 -> 978 (top 978 27.4026, 2nd 2725 16.9151, margin 10.4874) 55.33 ms (GPU 53.81 ms, 1269 dispatches) +strata metal: window pos 10202 token 13 -> 1061 (top 1061 21.7433, 2nd 2532 20.4921, margin 1.2513) 62.14 ms (GPU 60.62 ms, 1269 dispatches) +strata metal: window pos 10205 token 264 -> 15855 (top 15855 26.3293, 2nd 3471 19.5190, margin 6.8104) 71.95 ms (GPU 70.41 ms, 1269 dispatches) +strata metal: window pos 10207 token 40719 -> 13 (top 13 21.3241, 2nd 11 20.2257, margin 1.0984) 56.72 ms (GPU 54.62 ms, 1269 dispatches) +strata metal: window pos 10208 token 13 -> 353 (top 353 23.4751, 2nd 6558 22.0544, margin 1.4206) 61.39 ms (GPU 59.10 ms, 1269 dispatches) +STRATA_RESULT init_ms=14139.98 prompt_tokens=10180 prefill_ms=21641.96 prefill_tok_s=470.38 decode_tokens=31 decode_ms=888.16 decode_tok_s=34.90 decode_ms_per_tok=28.65 +STRATA_SPEC rounds=13 drafted=23 accepted=18 accept_pct=78.3 tok_per_round=2.38 mtp_ms=109.42 +GEN_IDS: 760 1156 369 33362 11 852 8523 728 13 2838 579 5315 1292 37444 685 321 9859 1092 353 2908 978 3604 13 1061 369 264 15855 40719 13 353 1220 5707 +expert stream 13 windows: 11692 distinct lookups (899.4 per window), hits 0.8885, read-ahead hits 0.0457 (62 of them late), misses 0.0659 (59.2 per window); read-aheads 1303 (100.2 per window, 0 dropped); 232.8 MB read per window (2073 reads); host 17.36 ms per window, of it waiting 16.95 ms (+ 0.83 ms after the GPU went on) +stream windows T=1 1 (decode): 52.05 ms each (GPU incl. the host's turns 50.55 ms) +stream windows T=2 1 (decode): 56.55 ms each (GPU incl. the host's turns 54.62 ms) +stream windows T=3 11 (decode): 67.13 ms each (GPU incl. the host's turns 65.37 ms) diff --git a/results/round39/model-ab/10k-warm-1-candidate.log b/results/round39/model-ab/10k-warm-1-candidate.log new file mode 100644 index 0000000..856b113 --- /dev/null +++ b/results/round39/model-ab/10k-warm-1-candidate.log @@ -0,0 +1,46 @@ +strata metal stream: experts read from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf.experts (one request per expert) +strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): 290 matrices, 680 MB, made in 135 ms +strata metal stream: locked 2.27 GB of active GPU dense tensors (unpinned 0.34 GB CPU token_embd + 1.28 GB skipped BF16 HC in 3.90 GB file map) +strata metal stream: 1079 dense tensors wrapped zero-copy from file (3.90 GB clean file-backed map, 117.9 ms) +strata metal: decode window on Apple M5: /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf streamed from the SSD (experts) (3.9 GB), 32896 cells, PLE on, ready in 520 ms +strata metal: MTP map: 0.71 GB wired (mlock; 0.00 GB failed), 0.18 GB unused released +strata metal: MTP draft round in one command buffer (embedding type 23 on the GPU: max |GPU - host| 0 over 8 rows) +strata metal: MTP draft layer from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf (0.89 GB, 9 norms folded), draft head over 40525 tokens +strata metal stream: 0.00 GB of the host code's device weight copies given back (unused by the Metal window) +strata metal stream: 2.60 GB allocated before the arena (the process holds 2.24 GB, 5.40 GB resident); RAM 25.8 GB (12.2 GB available, 7.0 GB kept for macOS), GPU working set 19.1 GB +strata metal stream: planned resident footprint 9.6 GB decode (10.7 GB virtual peak during prefill with 1.12 GB transient prompt scratch) = arena 6.5 + 2.6 held before it + 0.5 reserved for the n-gram table; cap 18.8 GB = the lesser of the GPU working set 19.1 GB and the RAM 25.8 GB less 7.0 GB left to macOS and its file cache +strata metal stream: expert arena 6.50 GB, 4129 slots in 2 size classes, 4129 experts read from the profile's top (24576 ranked) in 1.3 s (4.64 GB/s); 8 I/O threads, read-ahead top 6; 1 buffers +strata metal stream: 6.5 GB of the arena wired (mlock) in 677 ms; macOS cannot compress it while the engine idles (STRATA_METAL_WIRE_ARENA=0 turns it off) +strata metal stream: resident footprint after the fill: the process holds 8.7 GB (10.3 GB resident), planned 9.6 GB, cap 18.8 GB +strata metal stream: prompt scratch 1.12 GB mapped in separate anonymous mmap +strata metal: prompts on the matrix units in chunks of 2048 (experts v3, QSA flash, matmul2d layout K4; 1.12 GB of buffers, in a separate transient mmap while a prompt runs) +strata metal: prompt chunk pos 0, 2048 rows: 4378 ms (embd 26 ms, ple 0 ms, graph build 62 ms) +strata metal: prompt chunk pos 2048, 2048 rows: 4181 ms (embd 6 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 4096, 2048 rows: 4324 ms (embd 4 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 6144, 2048 rows: 4501 ms (embd 4 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 8192, 1987 rows: 4508 ms (embd 4 ms, ple 0 ms, graph build 1 ms) +strata metal stream: prompt scratch 1.12 GB unmapped after prefill (phys_footprint 8.89 -> 8.79 GB, resident RSS 9.61 -> 9.53 GB) +strata metal: window pos 10179 token 198 -> 760 (top 760 23.8239, 2nd 35 23.7743, margin 0.0496) 124.04 ms (GPU 94.11 ms, 1233 dispatches) +expert stream prompts 5 chunks: 82802 distinct lookups (345.0 per layer), in the arena 0.1981, streamed ahead 0.6932, missed 0.1088; 66194 streamed in (21993.9 MB/chunk, 74996 reads); host 466.7 ms per chunk, of it waiting 438.1 ms +expert stream 1 windows: 480 distinct lookups (480.0 per window), hits 0.5417, read-ahead hits 0.2875 (6 of them late), misses 0.1708 (82.0 per window); read-aheads 176 (176.0 per window, 203 dropped); 375.2 MB read per window (258 reads); host 23.88 ms per window, of it waiting 23.30 ms (+ 0.89 ms after the GPU went on) +stream windows T=1 1 (decode): 123.82 ms each (GPU incl. the host's turns 94.11 ms) +strata metal: window pos 10180 token 760 -> 1156 (top 1156 29.0186, 2nd 6353 15.3183, margin 13.7003) 86.99 ms (GPU 84.53 ms, 1233 dispatches) +strata metal: window pos 10183 token 33362 -> 11 (top 11 27.0843, 2nd 13 19.1119, margin 7.9724) 84.19 ms (GPU 81.93 ms, 1269 dispatches) +strata metal: window pos 10186 token 8523 -> 728 (top 728 26.5949, 2nd 279 17.2821, margin 9.3128) 75.33 ms (GPU 72.91 ms, 1269 dispatches) +strata metal: window pos 10188 token 13 -> 2838 (top 2838 22.9993, 2nd 2302 20.3597, margin 2.6395) 50.31 ms (GPU 48.60 ms, 1269 dispatches) +strata metal: window pos 10189 token 2838 -> 579 (top 579 24.7598, 2nd 42315 19.9055, margin 4.8543) 58.94 ms (GPU 57.41 ms, 1233 dispatches) +strata metal: window pos 10191 token 5315 -> 1292 (top 1292 22.6356, 2nd 15131 21.4028, margin 1.2328) 59.67 ms (GPU 57.39 ms, 1269 dispatches) +strata metal: window pos 10193 token 37444 -> 685 (top 685 25.6466, 2nd 321 15.4891, margin 10.1575) 65.38 ms (GPU 63.14 ms, 1269 dispatches) +strata metal: window pos 10196 token 9859 -> 1092 (top 1092 23.8715, 2nd 62431 16.8946, margin 6.9769) 59.56 ms (GPU 58.09 ms, 1269 dispatches) +strata metal: window pos 10199 token 2908 -> 978 (top 978 27.4026, 2nd 2725 16.9151, margin 10.4874) 56.78 ms (GPU 55.28 ms, 1269 dispatches) +strata metal: window pos 10202 token 13 -> 1061 (top 1061 21.7433, 2nd 2532 20.4921, margin 1.2513) 63.83 ms (GPU 62.30 ms, 1269 dispatches) +strata metal: window pos 10205 token 264 -> 15855 (top 15855 26.3293, 2nd 3471 19.5190, margin 6.8104) 71.70 ms (GPU 70.10 ms, 1269 dispatches) +strata metal: window pos 10207 token 40719 -> 13 (top 13 21.3241, 2nd 11 20.2257, margin 1.0984) 58.12 ms (GPU 56.03 ms, 1269 dispatches) +strata metal: window pos 10208 token 13 -> 353 (top 353 23.4751, 2nd 6558 22.0544, margin 1.4206) 61.57 ms (GPU 59.27 ms, 1269 dispatches) +STRATA_RESULT init_ms=3039.06 prompt_tokens=10180 prefill_ms=22208.75 prefill_tok_s=458.38 decode_tokens=31 decode_ms=891.67 decode_tok_s=34.77 decode_ms_per_tok=28.76 +STRATA_SPEC rounds=13 drafted=23 accepted=18 accept_pct=78.3 tok_per_round=2.38 mtp_ms=99.78 +GEN_IDS: 760 1156 369 33362 11 852 8523 728 13 2838 579 5315 1292 37444 685 321 9859 1092 353 2908 978 3604 13 1061 369 264 15855 40719 13 353 1220 5707 +expert stream 13 windows: 11692 distinct lookups (899.4 per window), hits 0.8885, read-ahead hits 0.0457 (56 of them late), misses 0.0659 (59.2 per window); read-aheads 1303 (100.2 per window, 0 dropped); 232.8 MB read per window (2073 reads); host 16.88 ms per window, of it waiting 16.43 ms (+ 0.87 ms after the GPU went on) +stream windows T=1 1 (decode): 50.13 ms each (GPU incl. the host's turns 48.60 ms) +stream windows T=2 1 (decode): 57.94 ms each (GPU incl. the host's turns 56.03 ms) +stream windows T=3 11 (decode): 67.46 ms each (GPU incl. the host's turns 65.67 ms) diff --git a/results/round39/model-ab/2k-0-baseline.log b/results/round39/model-ab/2k-0-baseline.log new file mode 100644 index 0000000..062aecb --- /dev/null +++ b/results/round39/model-ab/2k-0-baseline.log @@ -0,0 +1,68 @@ +strata metal stream: experts read from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf.experts (one request per expert) +strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): 290 matrices, 680 MB, made in 134 ms +strata metal stream: locked 2.27 GB of active GPU dense tensors (unpinned 0.34 GB CPU token_embd + 1.28 GB skipped BF16 HC in 3.90 GB file map) +strata metal stream: 1079 dense tensors wrapped zero-copy from file (3.90 GB clean file-backed map, 133.1 ms) +strata metal: decode window on Apple M5: /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf streamed from the SSD (experts) (3.9 GB), 32896 cells, PLE on, ready in 553 ms +strata metal: MTP map: 0.71 GB wired (mlock; 0.00 GB failed), 0.18 GB unused released +strata metal: MTP draft round in one command buffer (embedding type 23 on the GPU: max |GPU - host| 0 over 8 rows) +strata metal: MTP draft layer from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf (0.89 GB, 9 norms folded), draft head over 40525 tokens +strata metal stream: 0.00 GB of the host code's device weight copies given back (unused by the Metal window) +strata metal stream: 2.60 GB allocated before the arena (the process holds 2.24 GB, 5.40 GB resident); RAM 25.8 GB (12.1 GB available, 7.0 GB kept for macOS), GPU working set 19.1 GB +strata metal stream: planned resident footprint 9.6 GB decode (10.7 GB virtual peak during prefill with 1.12 GB transient prompt scratch) = arena 6.5 + 2.6 held before it + 0.5 reserved for the n-gram table; cap 18.8 GB = the lesser of the GPU working set 19.1 GB and the RAM 25.8 GB less 7.0 GB left to macOS and its file cache +strata metal stream: expert arena 6.50 GB, 4129 slots in 2 size classes, 4129 experts read from the profile's top (24576 ranked) in 1.5 s (4.13 GB/s); 8 I/O threads, read-ahead top 6; 1 buffers +strata metal stream: 6.5 GB of the arena wired (mlock) in 897 ms; macOS cannot compress it while the engine idles (STRATA_METAL_WIRE_ARENA=0 turns it off) +strata metal stream: resident footprint after the fill: the process holds 8.7 GB (10.2 GB resident), planned 9.6 GB, cap 18.8 GB +strata metal stream: prompt scratch 1.12 GB mapped in separate anonymous mmap +strata metal: prompts on the matrix units in chunks of 2048 (experts v3, QSA flash, matmul2d layout K4; 1.12 GB of buffers, in a separate transient mmap while a prompt runs) +strata metal: prompt chunk pos 0, 1983 rows: 5882 ms (embd 21 ms, ple 1 ms, graph build 63 ms) +strata metal stream: prompt scratch 1.12 GB unmapped after prefill (phys_footprint 8.89 -> 8.78 GB, resident RSS 9.52 -> 9.52 GB) +strata metal: window pos 1983 token 28 -> 16 (top 16 19.7624, 2nd 248046 19.6264, margin 0.1360) 167.18 ms (GPU 123.11 ms, 1197 dispatches) +expert stream prompts 1 chunks: 21669 distinct lookups (451.4 per layer), in the arena 0.1383, streamed ahead 0.7971, missed 0.0647; 18410 streamed in (29000.1 MB/chunk, 19751 reads); host 461.4 ms per chunk, of it waiting 398.9 ms +expert stream 1 windows: 480 distinct lookups (480.0 per window), hits 0.5146, read-ahead hits 0.3125 (5 of them late), misses 0.1729 (83.0 per window); read-aheads 228 (228.0 per window, 60 dropped); 455.8 MB read per window (311 reads); host 27.92 ms per window, of it waiting 26.75 ms (+ 1.70 ms after the GPU went on) +stream windows T=1 1 (decode): 166.84 ms each (GPU incl. the host's turns 123.11 ms) +strata metal: window pos 1984 token 16 -> 248046 (top 248046 20.9569, 2nd 15 19.4699, margin 1.4870) 69.50 ms (GPU 67.22 ms, 1197 dispatches) +strata metal: window pos 1985 token 248046 -> 198 (top 198 30.4136, 2nd 248044 13.3331, margin 17.0805) 97.21 ms (GPU 94.23 ms, 1197 dispatches) +strata metal: window pos 1988 token 248046 -> 198 (top 198 34.6470, 2nd 271 17.3126, margin 17.3344) 70.43 ms (GPU 67.19 ms, 1233 dispatches) +strata metal: window pos 1991 token 74455 -> 198 (top 198 28.0621, 2nd 248046 18.3070, margin 9.7551) 77.54 ms (GPU 73.95 ms, 1233 dispatches) +strata metal: window pos 1994 token 271 -> 248069 (top 248069 34.4544, 2nd 510 14.2772, margin 20.1772) 59.86 ms (GPU 56.30 ms, 1233 dispatches) +strata metal: window pos 1997 token 71093 -> 2164 (top 2164 23.4538, 2nd 198 16.9753, margin 6.4785) 74.07 ms (GPU 71.50 ms, 1233 dispatches) +strata metal: window pos 2000 token 90 -> 198 (top 198 21.8618, 2nd 248046 16.6073, margin 5.2545) 80.28 ms (GPU 77.56 ms, 1233 dispatches) +strata metal: window pos 2003 token 328 -> 10083 (top 10083 21.7508, 2nd 7945 21.2813, margin 0.4695) 87.67 ms (GPU 85.14 ms, 1233 dispatches) +strata metal: window pos 2004 token 10083 -> 763 (top 763 23.1994, 2nd 3147 13.6225, margin 9.5770) 104.81 ms (GPU 101.76 ms, 1233 dispatches) +strata metal: window pos 2008 token 799 -> 1534 (top 1534 23.7638, 2nd 1483 16.7971, margin 6.9667) 115.65 ms (GPU 112.35 ms, 1233 dispatches) +strata metal: window pos 2012 token 318 -> 89478 (top 89478 23.2369, 2nd 13259 17.0454, margin 6.1915) 122.15 ms (GPU 118.72 ms, 1233 dispatches) +strata metal: window pos 2016 token 11 -> 4713 (top 4713 23.7222, 2nd 14408 19.8758, margin 3.8464) 94.24 ms (GPU 91.72 ms, 1233 dispatches) +strata metal: window pos 2020 token 11 -> 14408 (top 14408 26.5939, 2nd 321 17.7990, margin 8.7950) 98.10 ms (GPU 95.57 ms, 1233 dispatches) +strata metal: window pos 2024 token 11836 -> 681 (top 681 22.5685, 2nd 11 19.3141, margin 3.2543) 105.32 ms (GPU 102.79 ms, 1233 dispatches) +strata metal: window pos 2028 token 279 -> 3274 (top 3274 24.5338, 2nd 4434 15.0477, margin 9.4860) 111.79 ms (GPU 109.05 ms, 1233 dispatches) +strata metal: window pos 2032 token 198 -> 220 (top 220 23.5962, 2nd 1 16.3865, margin 7.2098) 96.08 ms (GPU 93.34 ms, 1233 dispatches) +strata metal: window pos 2036 token 763 -> 498 (top 498 23.6984, 2nd 9774 18.6071, margin 5.0914) 94.11 ms (GPU 91.43 ms, 1233 dispatches) +strata metal: window pos 2040 token 313 -> 198 (top 198 24.4607, 2nd 328 18.3975, margin 6.0631) 102.97 ms (GPU 100.34 ms, 1233 dispatches) +strata metal: window pos 2044 token 432 -> 582 (top 582 23.2588, 2nd 39903 12.5246, margin 10.7341) 112.30 ms (GPU 108.41 ms, 1233 dispatches) +strata metal: window pos 2045 token 582 -> 53188 (top 53188 28.0282, 2nd 7206 19.4190, margin 8.6091) 98.82 ms (GPU 94.77 ms, 1233 dispatches) +strata metal: window pos 2046 token 53188 -> 763 (top 763 24.6892, 2nd 3147 15.0103, margin 9.6789) 104.58 ms (GPU 101.28 ms, 1233 dispatches) +strata metal: window pos 2050 token 359 -> 8442 (top 8442 23.6402, 2nd 95429 16.1098, margin 7.5303) 115.30 ms (GPU 110.82 ms, 1269 dispatches) +strata metal: window pos 2054 token 31361 -> 3381 (top 3381 22.9197, 2nd 1981 14.6839, margin 8.2359) 115.60 ms (GPU 111.38 ms, 1269 dispatches) +strata metal: window pos 2058 token 2502 -> 1228 (top 1228 25.2811, 2nd 18997 18.3873, margin 6.8938) 117.91 ms (GPU 114.79 ms, 1269 dispatches) +strata metal: window pos 2062 token 593 -> 2911 (top 2911 25.1616, 2nd 41490 19.8592, margin 5.3024) 121.03 ms (GPU 118.02 ms, 1269 dispatches) +strata metal: window pos 2066 token 3778 -> 976 (top 976 23.8748, 2nd 248046 17.3181, margin 6.5567) 113.67 ms (GPU 110.57 ms, 1269 dispatches) +strata metal: window pos 2070 token 5454 -> 53294 (top 53294 23.8362, 2nd 248046 16.4544, margin 7.3819) 113.94 ms (GPU 110.72 ms, 1269 dispatches) +strata metal: window pos 2074 token 33684 -> 6 (top 6 24.3685, 2nd 318 16.9675, margin 7.4010) 107.39 ms (GPU 104.40 ms, 1269 dispatches) +strata metal: window pos 2078 token 7018 -> 1152 (top 1152 24.8221, 2nd 13259 17.0250, margin 7.7970) 94.28 ms (GPU 89.86 ms, 1269 dispatches) +strata metal: window pos 2080 token 1331 -> 25 (top 25 24.1473, 2nd 314 14.1257, margin 10.0216) 107.73 ms (GPU 103.51 ms, 1269 dispatches) +strata metal: window pos 2084 token 471 -> 66 (top 66 26.1635, 2nd 248046 16.1265, margin 10.0370) 104.44 ms (GPU 101.47 ms, 1269 dispatches) +strata metal: window pos 2088 token 2911 -> 80334 (top 80334 26.8637, 2nd 14 16.2198, margin 10.6438) 91.60 ms (GPU 87.80 ms, 1269 dispatches) +strata metal: window pos 2090 token 53294 -> 3778 (top 3778 24.5964, 2nd 13 13.7553, margin 10.8411) 93.13 ms (GPU 87.37 ms, 1269 dispatches) +strata metal: window pos 2092 token 8 -> 5639 (top 5639 25.4322, 2nd 248046 17.5445, margin 7.8877) 103.67 ms (GPU 100.65 ms, 1269 dispatches) +strata metal: window pos 2096 token 1627 -> 359 (top 359 23.6442, 2nd 7018 23.0138, margin 0.6305) 81.09 ms (GPU 77.83 ms, 1269 dispatches) +strata metal: window pos 2098 token 4277 -> 593 (top 593 22.6791, 2nd 7658 18.9420, margin 3.7372) 96.27 ms (GPU 93.49 ms, 1269 dispatches) +strata metal: window pos 2102 token 11836 -> 12091 (top 12091 23.7703, 2nd 318 17.5935, margin 6.1768) 98.68 ms (GPU 96.06 ms, 1269 dispatches) +strata metal: window pos 2106 token 18997 -> 471 (top 471 24.9212, 2nd 248046 14.8304, margin 10.0908) 83.35 ms (GPU 80.27 ms, 1269 dispatches) +strata metal: window pos 2108 token 16 -> 593 (top 593 24.9663, 2nd 49 15.3718, margin 9.5945) 84.12 ms (GPU 81.30 ms, 1269 dispatches) +STRATA_RESULT init_ms=3523.84 prompt_tokens=1984 prefill_ms=6219.43 prefill_tok_s=319.00 decode_tokens=127 decode_ms=3997.56 decode_tok_s=31.77 decode_ms_per_tok=31.48 +STRATA_SPEC rounds=39 drafted=99 accepted=88 accept_pct=88.9 tok_per_round=3.26 mtp_ms=253.96 +GEN_IDS: 16 248046 198 248045 248046 198 248045 74455 198 248068 271 248069 271 71093 2164 198 90 198 220 328 10083 763 328 6516 799 1534 57287 18944 318 89478 11 1331 11 4713 2552 52186 11 14408 314 7658 11836 681 1179 1804 279 3274 4434 10152 198 220 328 23508 763 498 198 262 313 198 406 328 432 582 53188 763 328 2949 359 8442 64697 43681 31361 3381 1981 6 2502 1228 471 69 593 2911 80334 53294 3778 976 1627 359 5454 53294 3778 25 33684 6 2502 1627 7018 1152 1331 25 4771 23277 471 66 361 593 2911 80334 53294 3778 8 5639 2037 2502 1627 359 4277 593 2911 76644 11836 12091 6 2502 18997 471 16 593 2911 76644 +expert stream 39 windows: 42631 distinct lookups (1093.1 per window), hits 0.8959, read-ahead hits 0.0405 (181 of them late), misses 0.0635 (69.5 per window); read-aheads 5239 (134.3 per window, 0 dropped); 297.4 MB read per window (7948 reads); host 25.40 ms per window, of it waiting 24.44 ms (+ 1.54 ms after the GPU went on) +stream windows T=1 1 (decode): 69.27 ms each (GPU incl. the host's turns 67.22 ms) +stream windows T=3 15 (decode): 85.32 ms each (GPU incl. the host's turns 82.23 ms) +stream windows T=4 23 (decode): 106.85 ms each (GPU incl. the host's turns 104.10 ms) diff --git a/results/round39/model-ab/2k-1-candidate.log b/results/round39/model-ab/2k-1-candidate.log new file mode 100644 index 0000000..63d6657 --- /dev/null +++ b/results/round39/model-ab/2k-1-candidate.log @@ -0,0 +1,68 @@ +strata metal stream: experts read from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf.experts (one request per expert) +strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): 290 matrices, 680 MB, made in 133 ms +strata metal stream: locked 2.27 GB of active GPU dense tensors (unpinned 0.34 GB CPU token_embd + 1.28 GB skipped BF16 HC in 3.90 GB file map) +strata metal stream: 1079 dense tensors wrapped zero-copy from file (3.90 GB clean file-backed map, 117.7 ms) +strata metal: decode window on Apple M5: /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf streamed from the SSD (experts) (3.9 GB), 32896 cells, PLE on, ready in 523 ms +strata metal: MTP map: 0.71 GB wired (mlock; 0.00 GB failed), 0.18 GB unused released +strata metal: MTP draft round in one command buffer (embedding type 23 on the GPU: max |GPU - host| 0 over 8 rows) +strata metal: MTP draft layer from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf (0.89 GB, 9 norms folded), draft head over 40525 tokens +strata metal stream: 0.00 GB of the host code's device weight copies given back (unused by the Metal window) +strata metal stream: 2.60 GB allocated before the arena (the process holds 2.24 GB, 5.40 GB resident); RAM 25.8 GB (11.9 GB available, 7.0 GB kept for macOS), GPU working set 19.1 GB +strata metal stream: planned resident footprint 9.6 GB decode (10.7 GB virtual peak during prefill with 1.12 GB transient prompt scratch) = arena 6.5 + 2.6 held before it + 0.5 reserved for the n-gram table; cap 18.8 GB = the lesser of the GPU working set 19.1 GB and the RAM 25.8 GB less 7.0 GB left to macOS and its file cache +strata metal stream: expert arena 6.50 GB, 4129 slots in 2 size classes, 4129 experts read from the profile's top (24576 ranked) in 1.8 s (3.42 GB/s); 8 I/O threads, read-ahead top 6; 1 buffers +strata metal stream: 6.5 GB of the arena wired (mlock) in 755 ms; macOS cannot compress it while the engine idles (STRATA_METAL_WIRE_ARENA=0 turns it off) +strata metal stream: resident footprint after the fill: the process holds 8.7 GB (10.0 GB resident), planned 9.6 GB, cap 18.8 GB +strata metal stream: prompt scratch 1.12 GB mapped in separate anonymous mmap +strata metal: prompts on the matrix units in chunks of 2048 (experts v3, QSA flash, matmul2d layout K4; 1.12 GB of buffers, in a separate transient mmap while a prompt runs) +strata metal: prompt chunk pos 0, 1983 rows: 6502 ms (embd 35 ms, ple 1 ms, graph build 66 ms) +strata metal stream: prompt scratch 1.12 GB unmapped after prefill (phys_footprint 8.89 -> 8.78 GB, resident RSS 9.53 -> 9.53 GB) +strata metal: window pos 1983 token 28 -> 16 (top 16 19.7624, 2nd 248046 19.6264, margin 0.1360) 169.32 ms (GPU 125.83 ms, 1197 dispatches) +expert stream prompts 1 chunks: 21669 distinct lookups (451.4 per layer), in the arena 0.1383, streamed ahead 0.7971, missed 0.0647; 18410 streamed in (28894.6 MB/chunk, 19678 reads); host 540.3 ms per chunk, of it waiting 476.4 ms +expert stream 1 windows: 480 distinct lookups (480.0 per window), hits 0.5146, read-ahead hits 0.3125 (8 of them late), misses 0.1729 (83.0 per window); read-aheads 228 (228.0 per window, 133 dropped); 455.8 MB read per window (311 reads); host 29.99 ms per window, of it waiting 28.88 ms (+ 1.52 ms after the GPU went on) +stream windows T=1 1 (decode): 168.92 ms each (GPU incl. the host's turns 125.83 ms) +strata metal: window pos 1984 token 16 -> 248046 (top 248046 20.9569, 2nd 15 19.4699, margin 1.4870) 71.19 ms (GPU 68.05 ms, 1197 dispatches) +strata metal: window pos 1985 token 248046 -> 198 (top 198 30.4136, 2nd 248044 13.3331, margin 17.0805) 104.79 ms (GPU 101.51 ms, 1197 dispatches) +strata metal: window pos 1988 token 248046 -> 198 (top 198 34.6470, 2nd 271 17.3126, margin 17.3344) 71.32 ms (GPU 67.79 ms, 1233 dispatches) +strata metal: window pos 1991 token 74455 -> 198 (top 198 28.0621, 2nd 248046 18.3070, margin 9.7551) 80.62 ms (GPU 76.95 ms, 1233 dispatches) +strata metal: window pos 1994 token 271 -> 248069 (top 248069 34.4544, 2nd 510 14.2772, margin 20.1772) 60.05 ms (GPU 56.32 ms, 1233 dispatches) +strata metal: window pos 1997 token 71093 -> 2164 (top 2164 23.4538, 2nd 198 16.9753, margin 6.4785) 72.45 ms (GPU 69.74 ms, 1233 dispatches) +strata metal: window pos 2000 token 90 -> 198 (top 198 21.8618, 2nd 248046 16.6073, margin 5.2545) 81.52 ms (GPU 78.80 ms, 1233 dispatches) +strata metal: window pos 2003 token 328 -> 10083 (top 10083 21.7508, 2nd 7945 21.2813, margin 0.4695) 89.14 ms (GPU 86.50 ms, 1233 dispatches) +strata metal: window pos 2004 token 10083 -> 763 (top 763 23.1994, 2nd 3147 13.6225, margin 9.5770) 103.66 ms (GPU 100.34 ms, 1233 dispatches) +strata metal: window pos 2008 token 799 -> 1534 (top 1534 23.7638, 2nd 1483 16.7971, margin 6.9667) 117.06 ms (GPU 114.03 ms, 1233 dispatches) +strata metal: window pos 2012 token 318 -> 89478 (top 89478 23.2369, 2nd 13259 17.0454, margin 6.1915) 122.89 ms (GPU 119.71 ms, 1233 dispatches) +strata metal: window pos 2016 token 11 -> 4713 (top 4713 23.7222, 2nd 14408 19.8758, margin 3.8464) 98.22 ms (GPU 95.72 ms, 1233 dispatches) +strata metal: window pos 2020 token 11 -> 14408 (top 14408 26.5939, 2nd 321 17.7990, margin 8.7950) 97.36 ms (GPU 94.75 ms, 1233 dispatches) +strata metal: window pos 2024 token 11836 -> 681 (top 681 22.5685, 2nd 11 19.3141, margin 3.2543) 107.65 ms (GPU 104.98 ms, 1233 dispatches) +strata metal: window pos 2028 token 279 -> 3274 (top 3274 24.5338, 2nd 4434 15.0477, margin 9.4860) 105.33 ms (GPU 102.76 ms, 1233 dispatches) +strata metal: window pos 2032 token 198 -> 220 (top 220 23.5962, 2nd 1 16.3865, margin 7.2098) 95.49 ms (GPU 92.77 ms, 1233 dispatches) +strata metal: window pos 2036 token 763 -> 498 (top 498 23.6984, 2nd 9774 18.6071, margin 5.0914) 92.36 ms (GPU 89.61 ms, 1233 dispatches) +strata metal: window pos 2040 token 313 -> 198 (top 198 24.4607, 2nd 328 18.3975, margin 6.0631) 102.89 ms (GPU 100.30 ms, 1233 dispatches) +strata metal: window pos 2044 token 432 -> 582 (top 582 23.2588, 2nd 39903 12.5246, margin 10.7341) 115.11 ms (GPU 111.04 ms, 1233 dispatches) +strata metal: window pos 2045 token 582 -> 53188 (top 53188 28.0282, 2nd 7206 19.4190, margin 8.6091) 104.69 ms (GPU 99.94 ms, 1233 dispatches) +strata metal: window pos 2046 token 53188 -> 763 (top 763 24.6892, 2nd 3147 15.0103, margin 9.6789) 105.89 ms (GPU 102.63 ms, 1233 dispatches) +strata metal: window pos 2050 token 359 -> 8442 (top 8442 23.6402, 2nd 95429 16.1098, margin 7.5303) 114.34 ms (GPU 110.00 ms, 1269 dispatches) +strata metal: window pos 2054 token 31361 -> 3381 (top 3381 22.9197, 2nd 1981 14.6839, margin 8.2359) 124.49 ms (GPU 120.13 ms, 1269 dispatches) +strata metal: window pos 2058 token 2502 -> 1228 (top 1228 25.2811, 2nd 18997 18.3873, margin 6.8938) 118.90 ms (GPU 115.38 ms, 1269 dispatches) +strata metal: window pos 2062 token 593 -> 2911 (top 2911 25.1616, 2nd 41490 19.8592, margin 5.3024) 119.90 ms (GPU 116.74 ms, 1269 dispatches) +strata metal: window pos 2066 token 3778 -> 976 (top 976 23.8748, 2nd 248046 17.3181, margin 6.5567) 116.45 ms (GPU 113.17 ms, 1269 dispatches) +strata metal: window pos 2070 token 5454 -> 53294 (top 53294 23.8362, 2nd 248046 16.4544, margin 7.3819) 119.89 ms (GPU 116.86 ms, 1269 dispatches) +strata metal: window pos 2074 token 33684 -> 6 (top 6 24.3685, 2nd 318 16.9675, margin 7.4010) 116.32 ms (GPU 113.01 ms, 1269 dispatches) +strata metal: window pos 2078 token 7018 -> 1152 (top 1152 24.8221, 2nd 13259 17.0250, margin 7.7970) 100.57 ms (GPU 95.92 ms, 1269 dispatches) +strata metal: window pos 2080 token 1331 -> 25 (top 25 24.1473, 2nd 314 14.1257, margin 10.0216) 110.76 ms (GPU 106.35 ms, 1269 dispatches) +strata metal: window pos 2084 token 471 -> 66 (top 66 26.1635, 2nd 248046 16.1265, margin 10.0370) 106.92 ms (GPU 103.55 ms, 1269 dispatches) +strata metal: window pos 2088 token 2911 -> 80334 (top 80334 26.8637, 2nd 14 16.2198, margin 10.6438) 91.00 ms (GPU 86.23 ms, 1269 dispatches) +strata metal: window pos 2090 token 53294 -> 3778 (top 3778 24.5964, 2nd 13 13.7553, margin 10.8411) 93.12 ms (GPU 88.43 ms, 1269 dispatches) +strata metal: window pos 2092 token 8 -> 5639 (top 5639 25.4322, 2nd 248046 17.5445, margin 7.8877) 103.49 ms (GPU 100.07 ms, 1269 dispatches) +strata metal: window pos 2096 token 1627 -> 359 (top 359 23.6442, 2nd 7018 23.0138, margin 0.6305) 81.94 ms (GPU 78.42 ms, 1269 dispatches) +strata metal: window pos 2098 token 4277 -> 593 (top 593 22.6791, 2nd 7658 18.9420, margin 3.7372) 96.84 ms (GPU 93.84 ms, 1269 dispatches) +strata metal: window pos 2102 token 11836 -> 12091 (top 12091 23.7703, 2nd 318 17.5935, margin 6.1768) 96.78 ms (GPU 93.39 ms, 1269 dispatches) +strata metal: window pos 2106 token 18997 -> 471 (top 471 24.9212, 2nd 248046 14.8304, margin 10.0908) 82.07 ms (GPU 78.46 ms, 1269 dispatches) +strata metal: window pos 2108 token 16 -> 593 (top 593 24.9663, 2nd 49 15.3718, margin 9.5945) 85.57 ms (GPU 82.46 ms, 1269 dispatches) +STRATA_RESULT init_ms=3630.76 prompt_tokens=1984 prefill_ms=6868.23 prefill_tok_s=288.87 decode_tokens=127 decode_ms=4056.08 decode_tok_s=31.31 decode_ms_per_tok=31.94 +STRATA_SPEC rounds=39 drafted=99 accepted=88 accept_pct=88.9 tok_per_round=3.26 mtp_ms=257.21 +GEN_IDS: 16 248046 198 248045 248046 198 248045 74455 198 248068 271 248069 271 71093 2164 198 90 198 220 328 10083 763 328 6516 799 1534 57287 18944 318 89478 11 1331 11 4713 2552 52186 11 14408 314 7658 11836 681 1179 1804 279 3274 4434 10152 198 220 328 23508 763 498 198 262 313 198 406 328 432 582 53188 763 328 2949 359 8442 64697 43681 31361 3381 1981 6 2502 1228 471 69 593 2911 80334 53294 3778 976 1627 359 5454 53294 3778 25 33684 6 2502 1627 7018 1152 1331 25 4771 23277 471 66 361 593 2911 80334 53294 3778 8 5639 2037 2502 1627 359 4277 593 2911 76644 11836 12091 6 2502 18997 471 16 593 2911 76644 +expert stream 39 windows: 42631 distinct lookups (1093.1 per window), hits 0.8959, read-ahead hits 0.0405 (194 of them late), misses 0.0635 (69.5 per window); read-aheads 5239 (134.3 per window, 4 dropped); 297.2 MB read per window (7944 reads); host 26.52 ms per window, of it waiting 25.54 ms (+ 1.66 ms after the GPU went on) +stream windows T=1 1 (decode): 70.96 ms each (GPU incl. the host's turns 68.05 ms) +stream windows T=3 15 (decode): 87.20 ms each (GPU incl. the host's turns 83.90 ms) +stream windows T=4 23 (decode): 108.07 ms each (GPU incl. the host's turns 105.22 ms) diff --git a/results/round39/model-ab/2k-2-candidate.log b/results/round39/model-ab/2k-2-candidate.log new file mode 100644 index 0000000..7f5b7f5 --- /dev/null +++ b/results/round39/model-ab/2k-2-candidate.log @@ -0,0 +1,68 @@ +strata metal stream: experts read from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf.experts (one request per expert) +strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): 290 matrices, 680 MB, made in 141 ms +strata metal stream: locked 2.27 GB of active GPU dense tensors (unpinned 0.34 GB CPU token_embd + 1.28 GB skipped BF16 HC in 3.90 GB file map) +strata metal stream: 1079 dense tensors wrapped zero-copy from file (3.90 GB clean file-backed map, 128.9 ms) +strata metal: decode window on Apple M5: /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf streamed from the SSD (experts) (3.9 GB), 32896 cells, PLE on, ready in 409 ms +strata metal: MTP map: 0.71 GB wired (mlock; 0.00 GB failed), 0.18 GB unused released +strata metal: MTP draft round in one command buffer (embedding type 23 on the GPU: max |GPU - host| 0 over 8 rows) +strata metal: MTP draft layer from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf (0.89 GB, 9 norms folded), draft head over 40525 tokens +strata metal stream: 0.00 GB of the host code's device weight copies given back (unused by the Metal window) +strata metal stream: 2.60 GB allocated before the arena (the process holds 2.24 GB, 5.40 GB resident); RAM 25.8 GB (12.0 GB available, 7.0 GB kept for macOS), GPU working set 19.1 GB +strata metal stream: planned resident footprint 9.6 GB decode (10.7 GB virtual peak during prefill with 1.12 GB transient prompt scratch) = arena 6.5 + 2.6 held before it + 0.5 reserved for the n-gram table; cap 18.8 GB = the lesser of the GPU working set 19.1 GB and the RAM 25.8 GB less 7.0 GB left to macOS and its file cache +strata metal stream: expert arena 6.50 GB, 4129 slots in 2 size classes, 4129 experts read from the profile's top (24576 ranked) in 2.0 s (3.10 GB/s); 8 I/O threads, read-ahead top 6; 1 buffers +strata metal stream: 6.5 GB of the arena wired (mlock) in 649 ms; macOS cannot compress it while the engine idles (STRATA_METAL_WIRE_ARENA=0 turns it off) +strata metal stream: resident footprint after the fill: the process holds 8.7 GB (10.0 GB resident), planned 9.6 GB, cap 18.8 GB +strata metal stream: prompt scratch 1.12 GB mapped in separate anonymous mmap +strata metal: prompts on the matrix units in chunks of 2048 (experts v3, QSA flash, matmul2d layout K4; 1.12 GB of buffers, in a separate transient mmap while a prompt runs) +strata metal: prompt chunk pos 0, 1983 rows: 6218 ms (embd 20 ms, ple 0 ms, graph build 33 ms) +strata metal stream: prompt scratch 1.12 GB unmapped after prefill (phys_footprint 8.89 -> 8.78 GB, resident RSS 9.54 -> 9.54 GB) +strata metal: window pos 1983 token 28 -> 16 (top 16 19.7624, 2nd 248046 19.6264, margin 0.1360) 175.35 ms (GPU 131.70 ms, 1197 dispatches) +expert stream prompts 1 chunks: 21669 distinct lookups (451.4 per layer), in the arena 0.1383, streamed ahead 0.7971, missed 0.0647; 18410 streamed in (28911.8 MB/chunk, 19689 reads); host 566.3 ms per chunk, of it waiting 489.8 ms +expert stream 1 windows: 480 distinct lookups (480.0 per window), hits 0.5146, read-ahead hits 0.3125 (5 of them late), misses 0.1729 (83.0 per window); read-aheads 228 (228.0 per window, 122 dropped); 455.8 MB read per window (311 reads); host 31.55 ms per window, of it waiting 30.42 ms (+ 1.50 ms after the GPU went on) +stream windows T=1 1 (decode): 174.95 ms each (GPU incl. the host's turns 131.70 ms) +strata metal: window pos 1984 token 16 -> 248046 (top 248046 20.9569, 2nd 15 19.4699, margin 1.4870) 71.78 ms (GPU 69.40 ms, 1197 dispatches) +strata metal: window pos 1985 token 248046 -> 198 (top 198 30.4136, 2nd 248044 13.3331, margin 17.0805) 105.05 ms (GPU 101.66 ms, 1197 dispatches) +strata metal: window pos 1988 token 248046 -> 198 (top 198 34.6470, 2nd 271 17.3126, margin 17.3344) 72.37 ms (GPU 69.05 ms, 1233 dispatches) +strata metal: window pos 1991 token 74455 -> 198 (top 198 28.0621, 2nd 248046 18.3070, margin 9.7551) 81.46 ms (GPU 77.82 ms, 1233 dispatches) +strata metal: window pos 1994 token 271 -> 248069 (top 248069 34.4544, 2nd 510 14.2772, margin 20.1772) 61.85 ms (GPU 58.33 ms, 1233 dispatches) +strata metal: window pos 1997 token 71093 -> 2164 (top 2164 23.4538, 2nd 198 16.9753, margin 6.4785) 73.08 ms (GPU 70.44 ms, 1233 dispatches) +strata metal: window pos 2000 token 90 -> 198 (top 198 21.8618, 2nd 248046 16.6073, margin 5.2545) 81.77 ms (GPU 79.06 ms, 1233 dispatches) +strata metal: window pos 2003 token 328 -> 10083 (top 10083 21.7508, 2nd 7945 21.2813, margin 0.4695) 88.48 ms (GPU 86.00 ms, 1233 dispatches) +strata metal: window pos 2004 token 10083 -> 763 (top 763 23.1994, 2nd 3147 13.6225, margin 9.5770) 124.05 ms (GPU 120.97 ms, 1233 dispatches) +strata metal: window pos 2008 token 799 -> 1534 (top 1534 23.7638, 2nd 1483 16.7971, margin 6.9667) 126.35 ms (GPU 122.66 ms, 1233 dispatches) +strata metal: window pos 2012 token 318 -> 89478 (top 89478 23.2369, 2nd 13259 17.0454, margin 6.1915) 127.97 ms (GPU 124.18 ms, 1233 dispatches) +strata metal: window pos 2016 token 11 -> 4713 (top 4713 23.7222, 2nd 14408 19.8758, margin 3.8464) 95.19 ms (GPU 92.51 ms, 1233 dispatches) +strata metal: window pos 2020 token 11 -> 14408 (top 14408 26.5939, 2nd 321 17.7990, margin 8.7950) 96.53 ms (GPU 93.83 ms, 1233 dispatches) +strata metal: window pos 2024 token 11836 -> 681 (top 681 22.5685, 2nd 11 19.3141, margin 3.2543) 105.45 ms (GPU 102.51 ms, 1233 dispatches) +strata metal: window pos 2028 token 279 -> 3274 (top 3274 24.5338, 2nd 4434 15.0477, margin 9.4860) 105.12 ms (GPU 101.73 ms, 1233 dispatches) +strata metal: window pos 2032 token 198 -> 220 (top 220 23.5962, 2nd 1 16.3865, margin 7.2098) 98.22 ms (GPU 95.16 ms, 1233 dispatches) +strata metal: window pos 2036 token 763 -> 498 (top 498 23.6984, 2nd 9774 18.6071, margin 5.0914) 93.99 ms (GPU 91.04 ms, 1233 dispatches) +strata metal: window pos 2040 token 313 -> 198 (top 198 24.4607, 2nd 328 18.3975, margin 6.0631) 114.42 ms (GPU 110.41 ms, 1233 dispatches) +strata metal: window pos 2044 token 432 -> 582 (top 582 23.2588, 2nd 39903 12.5246, margin 10.7341) 145.59 ms (GPU 139.38 ms, 1233 dispatches) +strata metal: window pos 2045 token 582 -> 53188 (top 53188 28.0282, 2nd 7206 19.4190, margin 8.6091) 171.55 ms (GPU 123.87 ms, 1233 dispatches) +strata metal: window pos 2046 token 53188 -> 763 (top 763 24.6892, 2nd 3147 15.0103, margin 9.6789) 110.19 ms (GPU 106.69 ms, 1233 dispatches) +strata metal: window pos 2050 token 359 -> 8442 (top 8442 23.6402, 2nd 95429 16.1098, margin 7.5303) 115.34 ms (GPU 108.92 ms, 1269 dispatches) +strata metal: window pos 2054 token 31361 -> 3381 (top 3381 22.9197, 2nd 1981 14.6839, margin 8.2359) 121.30 ms (GPU 116.11 ms, 1269 dispatches) +strata metal: window pos 2058 token 2502 -> 1228 (top 1228 25.2811, 2nd 18997 18.3873, margin 6.8938) 122.70 ms (GPU 118.74 ms, 1269 dispatches) +strata metal: window pos 2062 token 593 -> 2911 (top 2911 25.1616, 2nd 41490 19.8592, margin 5.3024) 124.65 ms (GPU 120.87 ms, 1269 dispatches) +strata metal: window pos 2066 token 3778 -> 976 (top 976 23.8748, 2nd 248046 17.3181, margin 6.5567) 117.33 ms (GPU 113.47 ms, 1269 dispatches) +strata metal: window pos 2070 token 5454 -> 53294 (top 53294 23.8362, 2nd 248046 16.4544, margin 7.3819) 120.97 ms (GPU 116.87 ms, 1269 dispatches) +strata metal: window pos 2074 token 33684 -> 6 (top 6 24.3685, 2nd 318 16.9675, margin 7.4010) 109.19 ms (GPU 104.93 ms, 1269 dispatches) +strata metal: window pos 2078 token 7018 -> 1152 (top 1152 24.8221, 2nd 13259 17.0250, margin 7.7970) 96.81 ms (GPU 91.27 ms, 1269 dispatches) +strata metal: window pos 2080 token 1331 -> 25 (top 25 24.1473, 2nd 314 14.1257, margin 10.0216) 110.47 ms (GPU 105.27 ms, 1269 dispatches) +strata metal: window pos 2084 token 471 -> 66 (top 66 26.1635, 2nd 248046 16.1265, margin 10.0370) 112.17 ms (GPU 108.66 ms, 1269 dispatches) +strata metal: window pos 2088 token 2911 -> 80334 (top 80334 26.8637, 2nd 14 16.2198, margin 10.6438) 91.46 ms (GPU 86.03 ms, 1269 dispatches) +strata metal: window pos 2090 token 53294 -> 3778 (top 3778 24.5964, 2nd 13 13.7553, margin 10.8411) 90.94 ms (GPU 85.62 ms, 1269 dispatches) +strata metal: window pos 2092 token 8 -> 5639 (top 5639 25.4322, 2nd 248046 17.5445, margin 7.8877) 103.54 ms (GPU 99.90 ms, 1269 dispatches) +strata metal: window pos 2096 token 1627 -> 359 (top 359 23.6442, 2nd 7018 23.0138, margin 0.6305) 81.98 ms (GPU 78.01 ms, 1269 dispatches) +strata metal: window pos 2098 token 4277 -> 593 (top 593 22.6791, 2nd 7658 18.9420, margin 3.7372) 98.51 ms (GPU 94.95 ms, 1269 dispatches) +strata metal: window pos 2102 token 11836 -> 12091 (top 12091 23.7703, 2nd 318 17.5935, margin 6.1768) 98.64 ms (GPU 95.54 ms, 1269 dispatches) +strata metal: window pos 2106 token 18997 -> 471 (top 471 24.9212, 2nd 248046 14.8304, margin 10.0908) 82.69 ms (GPU 78.83 ms, 1269 dispatches) +strata metal: window pos 2108 token 16 -> 593 (top 593 24.9663, 2nd 49 15.3718, margin 9.5945) 85.29 ms (GPU 82.15 ms, 1269 dispatches) +STRATA_RESULT init_ms=3336.81 prompt_tokens=1984 prefill_ms=6534.59 prefill_tok_s=303.61 decode_tokens=127 decode_ms=4213.72 decode_tok_s=30.14 decode_ms_per_tok=33.18 +STRATA_SPEC rounds=39 drafted=99 accepted=88 accept_pct=88.9 tok_per_round=3.26 mtp_ms=256.98 +GEN_IDS: 16 248046 198 248045 248046 198 248045 74455 198 248068 271 248069 271 71093 2164 198 90 198 220 328 10083 763 328 6516 799 1534 57287 18944 318 89478 11 1331 11 4713 2552 52186 11 14408 314 7658 11836 681 1179 1804 279 3274 4434 10152 198 220 328 23508 763 498 198 262 313 198 406 328 432 582 53188 763 328 2949 359 8442 64697 43681 31361 3381 1981 6 2502 1228 471 69 593 2911 80334 53294 3778 976 1627 359 5454 53294 3778 25 33684 6 2502 1627 7018 1152 1331 25 4771 23277 471 66 361 593 2911 80334 53294 3778 8 5639 2037 2502 1627 359 4277 593 2911 76644 11836 12091 6 2502 18997 471 16 593 2911 76644 +expert stream 39 windows: 42631 distinct lookups (1093.1 per window), hits 0.8959, read-ahead hits 0.0405 (221 of them late), misses 0.0635 (69.5 per window); read-aheads 5239 (134.3 per window, 6 dropped); 297.1 MB read per window (7942 reads); host 28.70 ms per window, of it waiting 27.56 ms (+ 1.75 ms after the GPU went on) +stream windows T=1 1 (decode): 71.50 ms each (GPU incl. the host's turns 69.40 ms) +stream windows T=3 15 (decode): 93.60 ms each (GPU incl. the host's turns 87.17 ms) +stream windows T=4 23 (decode): 110.55 ms each (GPU incl. the host's turns 107.21 ms) diff --git a/results/round39/model-ab/2k-3-baseline.log b/results/round39/model-ab/2k-3-baseline.log new file mode 100644 index 0000000..823bc10 --- /dev/null +++ b/results/round39/model-ab/2k-3-baseline.log @@ -0,0 +1,68 @@ +strata metal stream: experts read from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf.experts (one request per expert) +strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): 290 matrices, 680 MB, made in 132 ms +strata metal stream: locked 2.27 GB of active GPU dense tensors (unpinned 0.34 GB CPU token_embd + 1.28 GB skipped BF16 HC in 3.90 GB file map) +strata metal stream: 1079 dense tensors wrapped zero-copy from file (3.90 GB clean file-backed map, 124.4 ms) +strata metal: decode window on Apple M5: /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf streamed from the SSD (experts) (3.9 GB), 32896 cells, PLE on, ready in 590 ms +strata metal: MTP map: 0.71 GB wired (mlock; 0.00 GB failed), 0.18 GB unused released +strata metal: MTP draft round in one command buffer (embedding type 23 on the GPU: max |GPU - host| 0 over 8 rows) +strata metal: MTP draft layer from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf (0.89 GB, 9 norms folded), draft head over 40525 tokens +strata metal stream: 0.00 GB of the host code's device weight copies given back (unused by the Metal window) +strata metal stream: 2.60 GB allocated before the arena (the process holds 2.24 GB, 5.40 GB resident); RAM 25.8 GB (11.9 GB available, 7.0 GB kept for macOS), GPU working set 19.1 GB +strata metal stream: planned resident footprint 9.6 GB decode (10.7 GB virtual peak during prefill with 1.12 GB transient prompt scratch) = arena 6.5 + 2.6 held before it + 0.5 reserved for the n-gram table; cap 18.8 GB = the lesser of the GPU working set 19.1 GB and the RAM 25.8 GB less 7.0 GB left to macOS and its file cache +strata metal stream: expert arena 6.50 GB, 4129 slots in 2 size classes, 4129 experts read from the profile's top (24576 ranked) in 2.0 s (3.09 GB/s); 8 I/O threads, read-ahead top 6; 1 buffers +strata metal stream: 6.5 GB of the arena wired (mlock) in 731 ms; macOS cannot compress it while the engine idles (STRATA_METAL_WIRE_ARENA=0 turns it off) +strata metal stream: resident footprint after the fill: the process holds 8.7 GB (10.1 GB resident), planned 9.6 GB, cap 18.8 GB +strata metal stream: prompt scratch 1.12 GB mapped in separate anonymous mmap +strata metal: prompts on the matrix units in chunks of 2048 (experts v3, QSA flash, matmul2d layout K4; 1.12 GB of buffers, in a separate transient mmap while a prompt runs) +strata metal: prompt chunk pos 0, 1983 rows: 6431 ms (embd 18 ms, ple 3 ms, graph build 61 ms) +strata metal stream: prompt scratch 1.12 GB unmapped after prefill (phys_footprint 8.89 -> 8.78 GB, resident RSS 9.54 -> 9.54 GB) +strata metal: window pos 1983 token 28 -> 16 (top 16 19.7624, 2nd 248046 19.6264, margin 0.1360) 161.33 ms (GPU 119.55 ms, 1197 dispatches) +expert stream prompts 1 chunks: 21669 distinct lookups (451.4 per layer), in the arena 0.1383, streamed ahead 0.7971, missed 0.0647; 18410 streamed in (28785.8 MB/chunk, 19600 reads); host 621.6 ms per chunk, of it waiting 566.9 ms +expert stream 1 windows: 480 distinct lookups (480.0 per window), hits 0.5146, read-ahead hits 0.3125 (11 of them late), misses 0.1729 (83.0 per window); read-aheads 228 (228.0 per window, 211 dropped); 455.8 MB read per window (311 reads); host 28.33 ms per window, of it waiting 27.40 ms (+ 1.33 ms after the GPU went on) +stream windows T=1 1 (decode): 161.04 ms each (GPU incl. the host's turns 119.55 ms) +strata metal: window pos 1984 token 16 -> 248046 (top 248046 20.9569, 2nd 15 19.4699, margin 1.4870) 72.36 ms (GPU 69.90 ms, 1197 dispatches) +strata metal: window pos 1985 token 248046 -> 198 (top 198 30.4136, 2nd 248044 13.3331, margin 17.0805) 104.85 ms (GPU 101.68 ms, 1197 dispatches) +strata metal: window pos 1988 token 248046 -> 198 (top 198 34.6470, 2nd 271 17.3126, margin 17.3344) 73.74 ms (GPU 70.79 ms, 1233 dispatches) +strata metal: window pos 1991 token 74455 -> 198 (top 198 28.0621, 2nd 248046 18.3070, margin 9.7551) 81.44 ms (GPU 78.31 ms, 1233 dispatches) +strata metal: window pos 1994 token 271 -> 248069 (top 248069 34.4544, 2nd 510 14.2772, margin 20.1772) 61.44 ms (GPU 58.16 ms, 1233 dispatches) +strata metal: window pos 1997 token 71093 -> 2164 (top 2164 23.4538, 2nd 198 16.9753, margin 6.4785) 72.60 ms (GPU 70.23 ms, 1233 dispatches) +strata metal: window pos 2000 token 90 -> 198 (top 198 21.8618, 2nd 248046 16.6073, margin 5.2545) 81.08 ms (GPU 78.72 ms, 1233 dispatches) +strata metal: window pos 2003 token 328 -> 10083 (top 10083 21.7508, 2nd 7945 21.2813, margin 0.4695) 88.80 ms (GPU 86.34 ms, 1233 dispatches) +strata metal: window pos 2004 token 10083 -> 763 (top 763 23.1994, 2nd 3147 13.6225, margin 9.5770) 104.85 ms (GPU 101.79 ms, 1233 dispatches) +strata metal: window pos 2008 token 799 -> 1534 (top 1534 23.7638, 2nd 1483 16.7971, margin 6.9667) 128.23 ms (GPU 124.92 ms, 1233 dispatches) +strata metal: window pos 2012 token 318 -> 89478 (top 89478 23.2369, 2nd 13259 17.0454, margin 6.1915) 132.64 ms (GPU 129.52 ms, 1233 dispatches) +strata metal: window pos 2016 token 11 -> 4713 (top 4713 23.7222, 2nd 14408 19.8758, margin 3.8464) 95.34 ms (GPU 92.49 ms, 1233 dispatches) +strata metal: window pos 2020 token 11 -> 14408 (top 14408 26.5939, 2nd 321 17.7990, margin 8.7950) 98.95 ms (GPU 96.25 ms, 1233 dispatches) +strata metal: window pos 2024 token 11836 -> 681 (top 681 22.5685, 2nd 11 19.3141, margin 3.2543) 125.28 ms (GPU 122.62 ms, 1233 dispatches) +strata metal: window pos 2028 token 279 -> 3274 (top 3274 24.5338, 2nd 4434 15.0477, margin 9.4860) 107.94 ms (GPU 105.42 ms, 1233 dispatches) +strata metal: window pos 2032 token 198 -> 220 (top 220 23.5962, 2nd 1 16.3865, margin 7.2098) 95.60 ms (GPU 92.91 ms, 1233 dispatches) +strata metal: window pos 2036 token 763 -> 498 (top 498 23.6984, 2nd 9774 18.6071, margin 5.0914) 94.51 ms (GPU 91.87 ms, 1233 dispatches) +strata metal: window pos 2040 token 313 -> 198 (top 198 24.4607, 2nd 328 18.3975, margin 6.0631) 106.29 ms (GPU 103.60 ms, 1233 dispatches) +strata metal: window pos 2044 token 432 -> 582 (top 582 23.2588, 2nd 39903 12.5246, margin 10.7341) 114.37 ms (GPU 110.08 ms, 1233 dispatches) +strata metal: window pos 2045 token 582 -> 53188 (top 53188 28.0282, 2nd 7206 19.4190, margin 8.6091) 112.81 ms (GPU 107.38 ms, 1233 dispatches) +strata metal: window pos 2046 token 53188 -> 763 (top 763 24.6892, 2nd 3147 15.0103, margin 9.6789) 111.09 ms (GPU 107.70 ms, 1233 dispatches) +strata metal: window pos 2050 token 359 -> 8442 (top 8442 23.6402, 2nd 95429 16.1098, margin 7.5303) 117.17 ms (GPU 112.61 ms, 1269 dispatches) +strata metal: window pos 2054 token 31361 -> 3381 (top 3381 22.9197, 2nd 1981 14.6839, margin 8.2359) 115.95 ms (GPU 111.52 ms, 1269 dispatches) +strata metal: window pos 2058 token 2502 -> 1228 (top 1228 25.2811, 2nd 18997 18.3873, margin 6.8938) 120.96 ms (GPU 117.79 ms, 1269 dispatches) +strata metal: window pos 2062 token 593 -> 2911 (top 2911 25.1616, 2nd 41490 19.8592, margin 5.3024) 120.46 ms (GPU 117.01 ms, 1269 dispatches) +strata metal: window pos 2066 token 3778 -> 976 (top 976 23.8748, 2nd 248046 17.3181, margin 6.5567) 115.94 ms (GPU 112.69 ms, 1269 dispatches) +strata metal: window pos 2070 token 5454 -> 53294 (top 53294 23.8362, 2nd 248046 16.4544, margin 7.3819) 127.24 ms (GPU 123.87 ms, 1269 dispatches) +strata metal: window pos 2074 token 33684 -> 6 (top 6 24.3685, 2nd 318 16.9675, margin 7.4010) 106.57 ms (GPU 103.19 ms, 1269 dispatches) +strata metal: window pos 2078 token 7018 -> 1152 (top 1152 24.8221, 2nd 13259 17.0250, margin 7.7970) 98.03 ms (GPU 90.41 ms, 1269 dispatches) +strata metal: window pos 2080 token 1331 -> 25 (top 25 24.1473, 2nd 314 14.1257, margin 10.0216) 111.49 ms (GPU 106.78 ms, 1269 dispatches) +strata metal: window pos 2084 token 471 -> 66 (top 66 26.1635, 2nd 248046 16.1265, margin 10.0370) 110.97 ms (GPU 107.45 ms, 1269 dispatches) +strata metal: window pos 2088 token 2911 -> 80334 (top 80334 26.8637, 2nd 14 16.2198, margin 10.6438) 94.51 ms (GPU 89.94 ms, 1269 dispatches) +strata metal: window pos 2090 token 53294 -> 3778 (top 3778 24.5964, 2nd 13 13.7553, margin 10.8411) 92.13 ms (GPU 87.57 ms, 1269 dispatches) +strata metal: window pos 2092 token 8 -> 5639 (top 5639 25.4322, 2nd 248046 17.5445, margin 7.8877) 107.87 ms (GPU 104.76 ms, 1269 dispatches) +strata metal: window pos 2096 token 1627 -> 359 (top 359 23.6442, 2nd 7018 23.0138, margin 0.6305) 83.20 ms (GPU 79.74 ms, 1269 dispatches) +strata metal: window pos 2098 token 4277 -> 593 (top 593 22.6791, 2nd 7658 18.9420, margin 3.7372) 100.32 ms (GPU 96.95 ms, 1269 dispatches) +strata metal: window pos 2102 token 11836 -> 12091 (top 12091 23.7703, 2nd 318 17.5935, margin 6.1768) 102.40 ms (GPU 99.21 ms, 1269 dispatches) +strata metal: window pos 2106 token 18997 -> 471 (top 471 24.9212, 2nd 248046 14.8304, margin 10.0908) 83.19 ms (GPU 79.56 ms, 1269 dispatches) +strata metal: window pos 2108 token 16 -> 593 (top 593 24.9663, 2nd 49 15.3718, margin 9.5945) 88.51 ms (GPU 85.41 ms, 1269 dispatches) +STRATA_RESULT init_ms=3886.19 prompt_tokens=1984 prefill_ms=6764.25 prefill_tok_s=293.31 decode_tokens=127 decode_ms=4143.13 decode_tok_s=30.65 decode_ms_per_tok=32.62 +STRATA_SPEC rounds=39 drafted=99 accepted=88 accept_pct=88.9 tok_per_round=3.26 mtp_ms=258.69 +GEN_IDS: 16 248046 198 248045 248046 198 248045 74455 198 248068 271 248069 271 71093 2164 198 90 198 220 328 10083 763 328 6516 799 1534 57287 18944 318 89478 11 1331 11 4713 2552 52186 11 14408 314 7658 11836 681 1179 1804 279 3274 4434 10152 198 220 328 23508 763 498 198 262 313 198 406 328 432 582 53188 763 328 2949 359 8442 64697 43681 31361 3381 1981 6 2502 1228 471 69 593 2911 80334 53294 3778 976 1627 359 5454 53294 3778 25 33684 6 2502 1627 7018 1152 1331 25 4771 23277 471 66 361 593 2911 80334 53294 3778 8 5639 2037 2502 1627 359 4277 593 2911 76644 11836 12091 6 2502 18997 471 16 593 2911 76644 +expert stream 39 windows: 42631 distinct lookups (1093.1 per window), hits 0.8959, read-ahead hits 0.0405 (188 of them late), misses 0.0635 (69.5 per window); read-aheads 5239 (134.3 per window, 2 dropped); 297.3 MB read per window (7946 reads); host 27.11 ms per window, of it waiting 26.17 ms (+ 1.61 ms after the GPU went on) +stream windows T=1 1 (decode): 72.11 ms each (GPU incl. the host's turns 69.90 ms) +stream windows T=3 15 (decode): 88.35 ms each (GPU incl. the host's turns 84.95 ms) +stream windows T=4 23 (decode): 110.87 ms each (GPU incl. the host's turns 107.95 ms) diff --git a/results/round39/model-ab/baseline-kernels/alu_bench.metal b/results/round39/model-ab/baseline-kernels/alu_bench.metal new file mode 100644 index 0000000..e34890a --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/alu_bench.metal @@ -0,0 +1,181 @@ +// metal/kernels/alu_bench.metal - instruction-rate probes (alu_bench.cpp): each thread runs ITER iterations of 8 +// independent chains of one operation; the result is written so nothing is eliminated. +#define ITER 1024 + +kernel void a_fma32(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + float a0 = gid, a1 = gid + 1, a2 = gid + 2, a3 = gid + 3, a4 = gid + 4, a5 = gid + 5, a6 = gid + 6, a7 = gid + 7; + for (int i = 0; i < ITER; ++i) { + a0 = fma(a0, k, 0.5f); a1 = fma(a1, k, 0.5f); a2 = fma(a2, k, 0.5f); a3 = fma(a3, k, 0.5f); + a4 = fma(a4, k, 0.5f); a5 = fma(a5, k, 0.5f); a6 = fma(a6, k, 0.5f); a7 = fma(a7, k, 0.5f); + } + out[gid] = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7; +} + +kernel void a_fma16(device float* out [[buffer(0)]], constant float& kf [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + const half k = half(kf); + half a0 = gid & 7, a1 = 1, a2 = 2, a3 = 3, a4 = 4, a5 = 5, a6 = 6, a7 = 7; + for (int i = 0; i < ITER; ++i) { + a0 = fma(a0, k, half(0.5)); a1 = fma(a1, k, half(0.5)); a2 = fma(a2, k, half(0.5)); a3 = fma(a3, k, half(0.5)); + a4 = fma(a4, k, half(0.5)); a5 = fma(a5, k, half(0.5)); a6 = fma(a6, k, half(0.5)); a7 = fma(a7, k, half(0.5)); + } + out[gid] = float(a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7); +} + +// float accumulator, half operands (does the conversion come free?) +kernel void a_fmamix(device float* out [[buffer(0)]], constant float& kf [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + half h0 = half(gid & 7), h1 = h0 + half(1); + float a0 = gid, a1 = 1, a2 = 2, a3 = 3, a4 = 4, a5 = 5, a6 = 6, a7 = 7; + for (int i = 0; i < ITER; ++i) { + a0 = fma(float(h0), kf, a0); a1 = fma(float(h1), kf, a1); a2 = fma(float(h0), kf, a2); a3 = fma(float(h1), kf, a3); + a4 = fma(float(h0), kf, a4); a5 = fma(float(h1), kf, a5); a6 = fma(float(h0), kf, a6); a7 = fma(float(h1), kf, a7); + h0 = h0 * half(0.999); h1 = h1 * half(0.999); + } + out[gid] = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7; +} + +kernel void a_iand(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint a0 = gid, a1 = gid + 1, a2 = gid + 2, a3 = gid + 3, a4 = gid + 4, a5 = gid + 5, a6 = gid + 6, a7 = gid + 7; + for (int i = 0; i < ITER; ++i) { + a0 = (a0 ^ k) + 1; a1 = (a1 ^ k) + 1; a2 = (a2 ^ k) + 1; a3 = (a3 ^ k) + 1; + a4 = (a4 ^ k) + 1; a5 = (a5 ^ k) + 1; a6 = (a6 ^ k) + 1; a7 = (a7 ^ k) + 1; + } + out[gid] = float(a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7); +} + +kernel void a_shift(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint a0 = gid, a1 = gid + 1, a2 = gid + 2, a3 = gid + 3, a4 = gid + 4, a5 = gid + 5, a6 = gid + 6, a7 = gid + 7; + for (int i = 0; i < ITER; ++i) { + a0 = (a0 >> k) | k; a1 = (a1 >> k) | k; a2 = (a2 >> k) | k; a3 = (a3 >> k) | k; + a4 = (a4 >> k) | k; a5 = (a5 >> k) | k; a6 = (a6 >> k) | k; a7 = (a7 >> k) | k; + } + out[gid] = float(a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7); +} + +// int -> float conversion + fma +kernel void a_cvt(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint u = gid; + float a0 = 0, a1 = 0, a2 = 0, a3 = 0, a4 = 0, a5 = 0, a6 = 0, a7 = 0; + for (int i = 0; i < ITER; ++i) { + u += k; + a0 += float(u & 3); a1 += float(u & 12); a2 += float(u & 48); a3 += float(u & 192); + a4 += float(u & 768); a5 += float(u & 3072); a6 += float(u & 12288); a7 += float(u & 49152); + } + out[gid] = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7; +} + +// threadgroup memory loads: random-ish 8-byte reads from a 4 KB table +kernel void a_tgload(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint2 tab[512]; + for (int i = tid; i < 512; i += 256) tab[i] = uint2(i, i * 3); + threadgroup_barrier(mem_flags::mem_threadgroup); + uint u = gid * 2654435761u; + uint2 s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += tab[(u >> 3) & 511]; s1 += tab[(u >> 7) & 511]; s2 += tab[(u >> 11) & 511]; s3 += tab[(u >> 15) & 511]; + u += k; + } + out[gid] = float(s0.x + s1.x + s2.x + s3.x + s0.y + s1.y + s2.y + s3.y); +} + +// select (sign application) +kernel void a_sel(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint u = gid; + float a0 = 1, a1 = 2, a2 = 3, a3 = 4, a4 = 5, a5 = 6, a6 = 7, a7 = 8; + for (int i = 0; i < ITER; ++i) { + u ^= k; + a0 = (u & 1) ? -a0 : a0; a1 = (u & 2) ? -a1 : a1; a2 = (u & 4) ? -a2 : a2; a3 = (u & 8) ? -a3 : a3; + a4 = (u & 16) ? -a4 : a4; a5 = (u & 32) ? -a5 : a5; a6 = (u & 64) ? -a6 : a6; a7 = (u & 128) ? -a7 : a7; + } + out[gid] = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7; +} + +// dot-shaped: acc_i += q_i * x_i with loop-carried accumulators only (q, x vary slowly) +kernel void a_dot32(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + float q[8], x[8], a[8]; + for (int j = 0; j < 8; ++j) { q[j] = float(gid + j); x[j] = k * j; a[j] = 0; } + for (int i = 0; i < ITER; ++i) { + for (int j = 0; j < 8; ++j) a[j] = fma(q[j], x[j], a[j]); + q[0] += 1.0f; + } + float s = 0; + for (int j = 0; j < 8; ++j) s += a[j]; + out[gid] = s; +} +kernel void a_dot16(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + half q[8], x[8], a[8]; + for (int j = 0; j < 8; ++j) { q[j] = half(gid + j); x[j] = half(k * j); a[j] = 0; } + for (int i = 0; i < ITER; ++i) { + for (int j = 0; j < 8; ++j) a[j] = fma(q[j], x[j], a[j]); + q[0] += half(1.0f); + } + float s = 0; + for (int j = 0; j < 8; ++j) s += float(a[j]); + out[gid] = s; +} +kernel void a_dotmix(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + half q[8]; + float x[8], a[8]; + for (int j = 0; j < 8; ++j) { q[j] = half(gid + j); x[j] = k * j; a[j] = 0; } + for (int i = 0; i < ITER; ++i) { + for (int j = 0; j < 8; ++j) a[j] = fma(float(q[j]), x[j], a[j]); + q[0] += half(1.0f); + } + float s = 0; + for (int j = 0; j < 8; ++j) s += a[j]; + out[gid] = s; +} +kernel void a_fma32x16(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + float a[16]; + for (int j = 0; j < 16; ++j) a[j] = gid + j; + for (int i = 0; i < ITER / 2; ++i) + for (int j = 0; j < 16; ++j) a[j] = fma(a[j], k, 0.5f); + float s = 0; + for (int j = 0; j < 16; ++j) s += a[j]; + out[gid] = s; +} + +// 8-byte gathers from a 4 KB table in device memory (L1-resident) and in the constant address space +kernel void a_devload(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], device const uint2* tab [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + uint u = gid * 2654435761u; + uint2 s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += tab[(u >> 3) & 511]; s1 += tab[(u >> 7) & 511]; s2 += tab[(u >> 11) & 511]; s3 += tab[(u >> 15) & 511]; + u += k; + } + out[gid] = float(s0.x + s1.x + s2.x + s3.x + s0.y + s1.y + s2.y + s3.y); +} +kernel void a_cstload(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], constant uint2* tab [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + uint u = gid * 2654435761u; + uint2 s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += tab[(u >> 3) & 511]; s1 += tab[(u >> 7) & 511]; s2 += tab[(u >> 11) & 511]; s3 += tab[(u >> 15) & 511]; + u += k; + } + out[gid] = float(s0.x + s1.x + s2.x + s3.x + s0.y + s1.y + s2.y + s3.y); +} +// 4-byte gathers from a 1 KB threadgroup table (IQ3_XXS's grid) +kernel void a_tgload4(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint tab[256]; + for (int i = tid; i < 256; i += 256) tab[i] = i * 3; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint u = gid * 2654435761u; + uint s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += tab[(u >> 3) & 255]; s1 += tab[(u >> 7) & 255]; s2 += tab[(u >> 11) & 255]; s3 += tab[(u >> 15) & 255]; + u += k; + } + out[gid] = float(s0 + s1 + s2 + s3); +} +// popcount parity (computed ksigns) +kernel void a_popc(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint u = gid, s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += popcount(u & 127); s1 += popcount((u >> 7) & 127); s2 += popcount((u >> 14) & 127); s3 += popcount((u >> 21) & 127); + u += k; + } + out[gid] = float(s0 + s1 + s2 + s3); +} diff --git a/results/round39/model-ab/baseline-kernels/dense_mm.metal b/results/round39/model-ab/baseline-kernels/dense_mm.metal new file mode 100644 index 0000000..ccf06f7 --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/dense_mm.metal @@ -0,0 +1,101 @@ +// metal/kernels/dense_mm.metal - dense projections for a PROMPT chunk (after quant/units/mma.metal): +// act_f16 X[t][k] -> half (the matrix units' operand) +// dq_f16 a quantized matrix -> half rows (ref_quant's values rounded to fp16), once per chunk +// gemm_f16 Y[t][n] = Xh[t] . Wh[n] (fp32) by MetalPerformancePrimitives matmul2d, TM x TN tiles, K dynamic; +// rows of Y past n_out / t past n_tok are not written (the tensors carry the extents) +#include +using namespace mpp::tensor_ops; + +kernel void act_f16(device const float* x [[buffer(0)]], device half* xh [[buffer(1)]], constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + const uint n4 = n >> 2; + if (gid < n4) { + ((device half4*)xh)[gid] = half4(((device const float4*)x)[gid]); + } else if (gid == n4) { + for (uint i = n4 << 2; i < n; ++i) xh[i] = half(x[i]); + } +} + +struct DqArgs { + uint n_in; // values per row + uint n_out; // rows + uint row_bytes; + uint perm; // 0: row r -> r; p > 0: row r = c * (n_out / p) + d -> d * p + c (the HC up matrix with the p + // streams of an output interleaved, for gemm_hc_up's epilogue) +}; + +// one thread per unit: grid ((units per row + nth.x - 1) / nth.x, n_out), nth.x threads +template +kernel void dq_f16(device const uchar* W [[buffer(0)]], device half* Wh [[buffer(1)]], constant DqArgs& a [[buffer(2)]], + uint3 gid [[thread_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort3 nth [[threads_per_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, H = UNIT / 2; + threadgroup uint tab[fmt::TAB]; + if (F == T_IQ4_XS || F == T_IQ4_NL) { + if (tid < 64) { + const uint r0 = tid & 15u, sh = (tid >> 4) * 8u; + const half h0 = half(as_type(uchar(((r0 & 8u) ? 0x7159453526190d01ul : 0xf6eaddcfbfad9881ul) >> ((r0 & 7u) * 8u)))); + tab[tid] = as_type(half2(h0, half(as_type(uchar(0xbfad9881u >> sh))))); + tab[tid + 64] = as_type(half2(h0, half(as_type(uchar(0xf6eaddcfu >> sh))))); + tab[tid + 128] = as_type(half2(h0, half(as_type(uchar(0x26190d01u >> sh))))); + tab[tid + 192] = as_type(half2(h0, half(as_type(uchar(0x71594535u >> sh))))); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } else if (fmt::TAB > 1) { + load_tab(tab, tid, nth.x); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int u = int(gid.x); + if (u * UNIT >= int(a.n_in)) return; + uint r = gid.y, ro = r; + if (a.perm) { + if (fmt::BB * 16 < fmt::QK * 17) { + r = a.perm == 4 ? (ro & 3u) * (a.n_out >> 2) + (ro >> 2) : (ro % a.perm) * (a.n_out / a.perm) + ro / a.perm; + } else { + const uint q = a.n_out / a.perm; + ro = (r % q) * a.perm + r / q; + } + } + float w[UNIT]; + int p0, p1; + DW::run(W + ulong(r) * a.row_bytes, u, 0, w, p0, p1, tab); + device half* o = Wh + ulong(ro) * a.n_in; +#pragma clang loop unroll(full) + for (int j = 0; j < H; j += 4) { + *(device half4*)(o + p0 + j) = half4(half(w[j]), half(w[j + 1]), half(w[j + 2]), half(w[j + 3])); + *(device half4*)(o + p1 + j) = half4(half(w[H + j]), half(w[H + j + 1]), half(w[H + j + 2]), half(w[H + j + 3])); + } +} + +struct GemmArgs { int M, N, K; int pad; }; // Y is M rows of N floats + +template +kernel void gemm_f16(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device float* Y [[buffer(2)]], + constant GemmArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]]) { + tensor, tensor_inline> tA(A, dextents(a.K, a.M)); + tensor, tensor_inline> tB(B, dextents(a.K, a.N)); + tensor, tensor_inline> tY(Y, dextents(a.N, a.M)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, int(tg.y) * TM); + auto mB = tB.slice(0, int(tg.x) * TN); + auto mY = tY.slice(int(tg.x) * TN, int(tg.y) * TM); + op.run(mA, mB, mY); +} + +#define DQ(F) template [[host_name("dq_f16_" #F)]] kernel decltype(dq_f16) dq_f16; +#define DQ_(F) DQ(F) +DQ_(T_F32) DQ_(T_BF16) DQ_(T_Q8_0) DQ_(T_Q4_0) DQ_(T_Q5_0) DQ_(T_Q5_1) DQ_(T_IQ4_NL) DQ_(T_IQ4_XS) DQ_(T_Q2_0) DQ_(T_IQ3_S) DQ_(T_Q3_K) +DQ_(T_Q4_K) DQ_(T_Q5_K) DQ_(T_Q6_K) +template [[host_name("gemm_f16_64_64_4")]] kernel decltype(gemm_f16<64, 64, 4>) gemm_f16<64, 64, 4>; +template [[host_name("gemm_f16_128_64_4")]] kernel decltype(gemm_f16<128, 64, 4>) gemm_f16<128, 64, 4>; +template [[host_name("gemm_f16_64_128_4")]] kernel decltype(gemm_f16<64, 128, 4>) gemm_f16<64, 128, 4>; +template [[host_name("gemm_f16_64_32_4")]] kernel decltype(gemm_f16<64, 32, 4>) gemm_f16<64, 32, 4>; +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>; +template [[host_name("gemm_f16_16_64_2")]] kernel decltype(gemm_f16<16, 64, 2>) gemm_f16<16, 64, 2>; +template [[host_name("gemm_f16_32_64_2")]] kernel decltype(gemm_f16<32, 64, 2>) gemm_f16<32, 64, 2>; +template [[host_name("gemm_f16_32_128_2")]] kernel decltype(gemm_f16<32, 128, 2>) gemm_f16<32, 128, 2>; +template [[host_name("gemm_f16_64_32_2")]] kernel decltype(gemm_f16<64, 32, 2>) gemm_f16<64, 32, 2>; +template [[host_name("gemm_f16_16_96_2")]] kernel decltype(gemm_f16<16, 96, 2>) gemm_f16<16, 96, 2>; diff --git a/results/round39/model-ab/baseline-kernels/gdn.metal b/results/round39/model-ab/baseline-kernels/gdn.metal new file mode 100644 index 0000000..c057067 --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/gdn.metal @@ -0,0 +1,197 @@ +// metal/kernels/gdn.metal - one GDN (gated delta-net) layer's recurrence for a window of tokens, fused end to end: +// +// conv (4 taps, history [channels][3] oldest first) -> SiLU, on q (16 heads), k (16), v (48 heads) of 128 +// q, k: L2 norm as native_gdn_l2_norm (rsqrt(sum/128 + eps/128), then x 1/sqrt(128), two roundings) +// gate = exp(softplus(alpha + dt) * a) (softplus = v > 20 ? v : log1p(exp(v))), beta = sigmoid(beta) +// kv_j = sum_i S[i][j] k[i]; delta_j = (v_j - gate * kv_j) * beta; S = gate * S + k[i] delta_j +// o_j = (sum_i S[i][j] q[i]) / sqrt(128); y_j = o_j * rsqrt(mean(o^2) + eps) * gamma_j * sigmoid(z_j) +// +// State layout as Strata's: S[(i * h_v + h) * 128 + j]. Value head h reads q/k head h % h_k. +// One threadgroup per value head, 512 threads = 16 simdgroups: simdgroup s holds columns j = 32 (s % 4) + lane and +// rows i = 32 (s / 4) .. + 32 of the head's 128 x 128 state in registers. GDN_NTOK tokens run in order; with +// `commit` the state and the conv history are written back after the last one (a verify window runs without commit, +// the fold of m accepted tokens is the same kernel with GDN_NTOK = m and commit). The conv history is read from +// `hist` and the committed one written to `hist_out` (ping-pong: the q/k channels are shared by three threadgroups). + +constant int GDN_NTOK [[function_constant(2)]]; + +struct GdnArgs { + uint h_k; // 16 + uint h_v; // 48 + uint channels; // 10240 = (2 h_k + h_v) * 128 + uint commit; // 0: verify (no write); 1: commit after the last token; 0x100 | k (1 <= k <= GDN_NTOK): commit + // the state and the conv history after the first k tokens, in place, while every token's y is + // still computed (a window whose first k rows are kept) - bitwise the commit of a k-token step + float eps; // 1e-6 + uint plane; // 0: y only; else y is the raw plane of prepared activations (xprep layout, n = h_v * 128, + // `plane` floats between planes) and the other planes and the per-16 sums are written too +}; + +constant float gdn_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, + 0.00006103515625f}; +constant float gdn_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + +// The kernel is launched with IB * 128 threads (512 for gdn_step, 1024 for gdn_step8). Without the attribute the compiler +// is free to use as many registers as the 128-row state needs and a pipeline can end up with +// maxTotalThreadsPerThreadgroup below the launch size: an M2 Max gave 448 for the GDN_NTOK >= 2 specializations, and +// Metal then skips the dispatch without an error (#2323). The attribute makes the compiler fit the launch size. +template +[[max_total_threads_per_threadgroup(IB * 128)]] +kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[buffer(1)]], + device const float* qkv [[buffer(2)]], device const float* alpha [[buffer(3)]], + device const float* beta [[buffer(4)]], device const float* z [[buffer(5)]], + device const float* conv_w [[buffer(6)]], device const float* dt [[buffer(7)]], + device const float* ssm_a [[buffer(8)]], device const float* gamma [[buffer(9)]], + device float* y [[buffer(10)]], constant GdnArgs& a [[buffer(11)]], + device float* hist_out [[buffer(12)]], + 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; + threadgroup float qs[S], ks[S], vs[S]; + constexpr int RPT = S / IB; // state rows per thread + // 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 - 1][S]; + threadgroup float part[8], part2[4]; + threadgroup float gb[2]; + 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; + const float scale = 1.0f / sqrt(float(S)); + + // the channel this thread convolves: 0..127 q of head hk, 128..255 k of head hk, 256..383 v of head h + int ch = -1; + if (tid < 128) ch = hk * S + tid; + 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) { + #pragma clang loop unroll(full) + for (int k = 0; k < 3; ++k) hst[k] = hist[ch * 3 + k]; + const float4 cwv = *(device const float4*)(conv_w + ch * 4); + cw[0] = cwv.x; cw[1] = cwv.y; cw[2] = cwv.z; cw[3] = cwv.w; + } + const float gam = ib == 0 ? gamma[j] : 0.0f; + const float dth = tid == 384 ? dt[h] : 0.0f; + const float ssm_ah = tid == 384 ? ssm_a[h] : 0.0f; + float st[RPT]; + const int st_stride = hv * S; + const int st_off = (i0 * hv + h) * S + j; + #pragma clang loop unroll(full) + for (int r = 0; r < RPT; ++r) st[r] = (state + st_off)[r * st_stride]; + + 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 vals[4] = {hst[0], hst[1], hst[2], in}; + float s = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 4; ++k) s = fma(vals[k], cw[k], s); + s = s + 0.0f; + cv = s / (1.0f + exp(-s)); + hst[0] = hst[1]; + hst[1] = hst[2]; + hst[2] = in; + if (t + 1 == commit_at && (tid >= 256 || h < int(a.h_k))) { + #pragma clang loop unroll(full) + for (int k = 0; k < 3; ++k) hist_out[ch * 3 + k] = hst[k]; + } + } + in = in_next; + // raw q/k/v to shared memory + L2 norm partials of q (sg 0..3) and k (sg 4..7); gate and beta on idle lanes + if (tid < 256) { + const float ssq = simd_sum(cv * cv); + if (lane == 0) part[sg] = ssq; + if (tid < 128) qs[tid] = cv; + else ks[tid - 128] = cv; + } else if (tid < 384) { + vs[tid - 256] = cv; + } else if (tid == 384) { + const float v = alpha[t * hv + h] + dth; + // 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_ah); + } else if (tid == 416) { + gb[1] = 1.0f / (1.0f + exp(-beta[t * hv + h])); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const float tot_k = (part[4] + part[5]) + (part[6] + part[7]); + const float rsk = rsqrt(tot_k / float(S) + a.eps / float(S)); + // the simdgroup's RPT rows of k and q one a lane, normalized on the fly and broadcast from the lane that holds them + const int qk_idx = i0 + (lane & (RPT - 1)); + const float kl = (rsk * ks[qk_idx]) * scale; + // kv_j + float kv = 0.0f; + #pragma clang loop unroll(full) + for (int r = 0; r < RPT; ++r) kv = fma(st[r], simd_broadcast(kl, ushort(r)), kv); + red[ib][j] = kv; + const float tot_q = (part[0] + part[1]) + (part[2] + part[3]); + const float rsq = rsqrt(tot_q / float(S) + a.eps / float(S)); + const float ql = (rsq * qs[qk_idx]) * scale; + const float g = gb[0], be = gb[1], vj = vs[j]; + threadgroup_barrier(mem_flags::mem_threadgroup); + kv = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < IB; ++k) kv += red[k][j]; + const float delta = (vj - g * kv) * be; + float o = 0.0f; + #pragma clang loop unroll(full) + for (int r = 0; r < RPT; ++r) { + st[r] = fma(g, st[r], simd_broadcast(kl, ushort(r)) * delta); + o = fma(st[r], simd_broadcast(ql, ushort(r)), o); + } + if (ib > 0) red2[ib - 1][j] = o; + if (t + 1 == commit_at) { + // the state and the conv history after token t (registers: st after this token's update, hst after its + // shift); the state is read only before the loop, so in place is safe. q/k history: the threadgroups + // h < h_k own their q/k head (h == hk); v history: every head its own - to hist_out, NOT in place: a q/k + // channel is read by the three threadgroups h % h_k at their start, and a threadgroup may start after its + // owner has finished (seen on the 10-core M4); the caller swaps the two + device float* st_base = state + st_off; + #pragma clang loop unroll(full) + for (int r = 0; r < RPT; ++r) st_base[r * st_stride] = st[r]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // the readout and the output norm: column j's o (only ib == 0 computes readout and writes y) + if (ib == 0) { + #pragma clang loop unroll(full) + for (int k = 0; k < IB - 1; ++k) o += red2[k][j]; + o = o * scale; + const float osq = simd_sum(o * o); + if (lane == 0) part2[cb] = osq; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ib == 0) { + const float tot = (part2[0] + part2[1]) + (part2[2] + part2[3]); + const float rs = rsqrt(tot / float(S) + a.eps); + const float yv = ((rs * o) * gam) * (1.0f / (1.0f + exp(-zz))); + y[t * hv * S + h * S + j] = yv; + if (a.plane) { + device float* raw = y + t * hv * S + h * S + j; + const int jj = j & 15; + raw[a.plane] = yv * gdn_q2s[jj & 7]; + raw[2 * a.plane] = yv * gdn_p4s[jj & 3]; + raw[3 * a.plane] = (jj & 1) ? yv * 0.00390625f : 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; + #pragma clang loop unroll(full) + 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; + } + } + } +} + +template [[host_name("gdn_step")]] kernel decltype(gdn_step_t<4>) gdn_step_t<4>; +template [[host_name("gdn_step8")]] kernel decltype(gdn_step_t<8>) gdn_step_t<8>; diff --git a/results/round39/model-ab/baseline-kernels/gdn_prompt.metal b/results/round39/model-ab/baseline-kernels/gdn_prompt.metal new file mode 100644 index 0000000..c2c4020 --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/gdn_prompt.metal @@ -0,0 +1,703 @@ +// metal/kernels/gdn_prompt.metal - a GDN layer's recurrence for a PROMPT chunk of T tokens, in three passes: the work +// that does not touch the state runs parallel over the tokens; only the state update is a scan, and the columns of a +// head's state are independent in it (column j needs k, q, v_j, the gate and beta), so a lane owns one column. +// +// gdn_pre per (token, head): conv (4 taps; history oldest first) -> SiLU; q, k L2 norm * 1/sqrt(128) (the decode +// step's arithmetic); kq = k . q per key head; gate = exp(softplus(alpha + dt) * a), beta = sigmoid(beta); +// the chunk's last 3 conv inputs -> hist_out +// gdn_scan per (value head, 32 columns x CB simdgroups): the state column in registers (128 floats a lane); +// per token a = S^T k, b = S^T q (old state), delta = (v - g a) beta, o = (g b + delta kq) / sqrt(128) +// (= S_new^T q / sqrt(128)), S = g S + k delta^T; q and k staged per TB tokens in threadgroup memory +// gdn_post per (token, head): y = o * rsqrt(mean(o^2) + eps) * gamma * sigmoid(z) -> fp16 (the out GEMM's operand) +// and, with `f32`, fp32 +// +// State layout as Strata's: S[(i * h_v + h) * 128 + j]. Value head h reads q/k head h % h_k. +// The kernels are new entries (nothing of gdn.metal changes). + +struct GdnPArgs { + uint n_tok; + uint h_k; // 16 + uint h_v; // 48 + uint channels; // 10240 + float eps; // 1e-6 + uint f32; // gdn_post: also write fp32 y +}; + +static inline float gp_conv(device const float* qkv, device const float* hist, device const float* cw, int ch, int t, + int channels) { + const float4 w = *(device const float4*)(cw + ch * 4); + float s = 0.0f; +#pragma clang loop unroll(full) + for (int k = 0; k < 4; ++k) { + const int src = t - 3 + k; + const float x = src >= 0 ? qkv[src * channels + ch] : hist[ch * 3 + (src + 3)]; + s = fma(x, w[k], s); + } + s = s + 0.0f; + return s / (1.0f + exp(-s)); +} + +// grid ((h_k + h_v) / 8, T) = (8, T) or (h_k + h_v / 2, T) = (40, T), 256 threads: when grid_x == (h_k + h_v) / 8, +// each simdgroup (32 lanes) owns one key head (x < h_k / 8: q, k and kq in registers) or one value head (else: vc and +// gb), with zero threadgroup memory and zero barriers; grid_x == h_k + h_v / 2 keeps the legacy grid. +kernel void gdn_pre(device const float* qkv [[buffer(0)]], device const float* hist [[buffer(1)]], + device const float* conv_w [[buffer(2)]], device const float* alpha [[buffer(3)]], + device const float* beta [[buffer(4)]], device const float* dt [[buffer(5)]], + device const float* ssm_a [[buffer(6)]], device float* qn [[buffer(7)]], device float* kn [[buffer(8)]], + device float* vc [[buffer(9)]], device float2* gb [[buffer(10)]], device float* kq [[buffer(11)]], + device float* hist_out [[buffer(12)]], constant GdnPArgs& a [[buffer(13)]], + uint3 tg [[threadgroup_position_in_grid]], uint3 tgs [[threadgroups_per_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128; + const int x = int(tg.x), t = int(tg.y), hk_n = int(a.h_k), hv = int(a.h_v), C = int(a.channels), T = int(a.n_tok); + const bool packed = int(tgs.x) * 8 == hk_n + hv; + // alpha and beta as the halves of one [T][2 h_v] GEMM (beta == alpha + h_v), else each [T][h_v] + const int ab = beta == alpha + hv ? 2 * hv : hv; + const int k_tgs = packed ? hk_n / 8 : hk_n; + if (x < k_tgs) { + const int kh = packed ? x * 8 + int(sg) : (sg == 0 ? x : hk_n); + if (kh >= hk_n) return; + float cv_q[4], cv_k[4]; + if (t >= 3) { + device const float* pq = qkv + (t - 3) * C + kh * S + int(lane); + device const float* pk = qkv + (t - 3) * C + (hk_n + kh) * S + int(lane); +#pragma clang loop unroll(full) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; + const float4 wq = *(device const float4*)(conv_w + ch_q * 4); + const float4 wk = *(device const float4*)(conv_w + ch_k * 4); + device const float* pqm = pq + m * 32; + device const float* pkm = pk + m * 32; + float sq = fma(pqm[0], wq.x, 0.0f); + sq = fma(pqm[C], wq.y, sq); + sq = fma(pqm[2 * C], wq.z, sq); + sq = fma(pqm[3 * C], wq.w, sq) + 0.0f; + cv_q[m] = sq / (1.0f + exp(-sq)); + float sk = fma(pkm[0], wk.x, 0.0f); + sk = fma(pkm[C], wk.y, sk); + sk = fma(pkm[2 * C], wk.z, sk); + sk = fma(pkm[3 * C], wk.w, sk) + 0.0f; + cv_k[m] = sk / (1.0f + exp(-sk)); + } + if (t == T - 1) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; + for (int k = 0; k < 3; ++k) { + hist_out[ch_q * 3 + k] = pq[m * 32 + (k + 1) * C]; + hist_out[ch_k * 3 + k] = pk[m * 32 + (k + 1) * C]; + } + } + } else { +#pragma clang loop unroll(full) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; + cv_q[m] = gp_conv(qkv, hist, conv_w, ch_q, t, C); + cv_k[m] = gp_conv(qkv, hist, conv_w, ch_k, t, C); + } + if (t == T - 1) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; + for (int k = 0; k < 3; ++k) { + const int src = T - 3 + k; + hist_out[ch_q * 3 + k] = src >= 0 ? qkv[src * C + ch_q] : hist[ch_q * 3 + (src + 3)]; + hist_out[ch_k * 3 + k] = src >= 0 ? qkv[src * C + ch_k] : hist[ch_k * 3 + (src + 3)]; + } + } + } + const float4 sq = simd_sum(float4(cv_q[0] * cv_q[0], cv_q[1] * cv_q[1], cv_q[2] * cv_q[2], cv_q[3] * cv_q[3])); + const float4 sk = simd_sum(float4(cv_k[0] * cv_k[0], cv_k[1] * cv_k[1], cv_k[2] * cv_k[2], cv_k[3] * cv_k[3])); + const float tot_q = (sq.x + sq.y) + (sq.z + sq.w), tot_k = (sk.x + sk.y) + (sk.z + sk.w); + const float rs_q = rsqrt(tot_q / float(S) + a.eps / float(S)); + const float rs_k = rsqrt(tot_k / float(S) + a.eps / float(S)); + const float scale = 1.0f / sqrt(float(S)); + const int base = (t * hk_n + kh) * S + int(lane); + float4 prod; +#pragma clang loop unroll(full) + for (int m = 0; m < 4; ++m) { + const float nq = (rs_q * cv_q[m]) * scale, nk = (rs_k * cv_k[m]) * scale; + qn[base + m * 32] = nq; + kn[base + m * 32] = nk; + prod[m] = nk * nq; + } + const float4 pv = simd_sum(prod); + if (lane == 0) kq[t * hk_n + kh] = (pv.x + pv.y) + (pv.z + pv.w); + } else if (packed) { + const int h = (x - k_tgs) * 8 + int(sg); + if (t >= 3) { + const int d = int(lane) * 4, ch = (2 * hk_n + h) * S + d; + device const float4* p = (device const float4*)(qkv + (t - 3) * C + ch); + device const float4* w = (device const float4*)(conv_w + ch * 4); + const float4 w0 = w[0], w1 = w[1], w2 = w[2], w3 = w[3]; + float4 cv = fma(p[0], float4(w0.x, w1.x, w2.x, w3.x), float4(0.0f)); + cv = fma(p[C / 4], float4(w0.y, w1.y, w2.y, w3.y), cv); + cv = fma(p[C / 2], float4(w0.z, w1.z, w2.z, w3.z), cv); + cv = fma(p[3 * C / 4], float4(w0.w, w1.w, w2.w, w3.w), cv); + cv = cv + 0.0f; + cv = cv / (1.0f + exp(-cv)); + *((device float4*)(vc + (t * hv + h) * S + d)) = cv; + if (t == T - 1) + for (int k = 0; k < 3; ++k) { + const float4 v = p[(k + 1) * C / 4]; + for (int m = 0; m < 4; ++m) hist_out[(ch + m) * 3 + k] = v[m]; + } + } else { +#pragma clang loop unroll(full) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch = 2 * hk_n * S + h * S + d; + const float cv = gp_conv(qkv, hist, conv_w, ch, t, C); + if (t == T - 1) + for (int k = 0; k < 3; ++k) { + const int src = T - 3 + k; + hist_out[ch * 3 + k] = src >= 0 ? qkv[src * C + ch] : hist[ch * 3 + (src + 3)]; + } + vc[(t * hv + h) * S + d] = cv; + } + } + if (lane == 0) { + const float v = alpha[t * ab + h] + dt[h]; + 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[t * hv + h] = float2(exp(sp * ssm_a[h]), 1.0f / (1.0f + exp(-beta[t * ab + h]))); + } + } else { + const int h = 2 * (x - k_tgs) + (tid >> 7), d = tid & 127; + const int ch = 2 * hk_n * S + h * S + d; + const float cv = gp_conv(qkv, hist, conv_w, ch, t, C); + if (t == T - 1) + for (int k = 0; k < 3; ++k) { + const int src = T - 3 + k; + hist_out[ch * 3 + k] = src >= 0 ? qkv[src * C + ch] : hist[ch * 3 + (src + 3)]; + } + vc[(t * hv + h) * S + d] = cv; + if (d == 0) { + const float v = alpha[t * ab + h] + dt[h]; + 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[t * hv + h] = float2(exp(sp * ssm_a[h]), 1.0f / (1.0f + exp(-beta[t * ab + h]))); + } + } +} + +// grid (h_v * 4 / CB), CB * 32 threads; TB tokens of q and k staged at a time +template +kernel void gdn_scan_t(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]], + device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]], + device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]], + device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]], + 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, NB = 4 / CB; + threadgroup float4 qs[TB][32], ks[TB][32]; + const int h = int(tg.x) / NB, hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k); + const int j = ((int(tg.x) % NB) * CB + sg) * 32 + lane; + const int T = int(a.n_tok); + const float scale = 1.0f / sqrt(float(S)); + float st[S]; +#pragma clang loop unroll(full) + for (int i = 0; i < S; ++i) st[i] = state[(i * hv + h) * S + j]; + for (int t0 = 0; t0 < T; t0 += TB) { + const int nb = min(TB, T - t0); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < nb * 32; idx += CB * 32) { + const int tt = idx >> 5, w = idx & 31; + qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w]; + ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int tt = 0; tt < nb; ++tt) { + const int t = t0 + tt; + const float2 g_b = gb[t * hv + h]; + const float kqv = kq[t * hkn + hk]; + const float v = vc[(t * hv + h) * S + j]; + float ka = 0.0f, kb = 0.0f, qa = 0.0f, qb = 0.0f; +#pragma clang loop unroll(full) + for (int w = 0; w < 32; w += 2) { + const float4 k0 = ks[tt][w], k1 = ks[tt][w + 1], q0 = qs[tt][w], q1 = qs[tt][w + 1]; + ka = fma(st[4 * w], k0.x, ka); ka = fma(st[4 * w + 1], k0.y, ka); + ka = fma(st[4 * w + 2], k0.z, ka); ka = fma(st[4 * w + 3], k0.w, ka); + kb = fma(st[4 * w + 4], k1.x, kb); kb = fma(st[4 * w + 5], k1.y, kb); + kb = fma(st[4 * w + 6], k1.z, kb); kb = fma(st[4 * w + 7], k1.w, kb); + qa = fma(st[4 * w], q0.x, qa); qa = fma(st[4 * w + 1], q0.y, qa); + qa = fma(st[4 * w + 2], q0.z, qa); qa = fma(st[4 * w + 3], q0.w, qa); + qb = fma(st[4 * w + 4], q1.x, qb); qb = fma(st[4 * w + 5], q1.y, qb); + qb = fma(st[4 * w + 6], q1.z, qb); qb = fma(st[4 * w + 7], q1.w, qb); + } + const float g = g_b.x, kv = ka + kb, qv = qa + qb; + const float delta = (v - g * kv) * g_b.y; + o[(t * hv + h) * S + j] = fma(g, qv, delta * kqv) * scale; +#pragma clang loop unroll(full) + for (int w = 0; w < 32; ++w) { + const float4 k4 = ks[tt][w]; + 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); + } + } + } +#pragma clang loop unroll(full) + for (int i = 0; i < S; ++i) state[(i * hv + h) * S + j] = st[i]; +} + +// a probe for the bench (function constant 6): 1 = no per-token exchange (wrong results), 2 = no state update +constant int GDN_PROBE [[function_constant(6)]]; +constant int gdn_probe = is_function_constant_defined(GDN_PROBE) ? GDN_PROBE : 0; + +// the rows of a column split over IB simdgroups (128 / IB rows a lane, fewer registers): grid (h_v * 4), IB * 32 +// threads; lane = column (32 columns a threadgroup); when TB == 0 (single-pass), each simdgroup loads its R rows of q +// and k across its lanes and broadcasts them with simd_broadcast (one initial dot pass at t = 0, one barrier a token, +// only red[2][IB][32] in threadgroup memory); TB > 0 keeps the TB-token threadgroup-staged reference. +template +kernel void gdn_scan_r(device float* state [[buffer(0)]], device const float* qn [[buffer(1)]], + device const float* kn [[buffer(2)]], device const float* vc [[buffer(3)]], + device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]], + device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]], + 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; + threadgroup float2 red[2][IB][32]; + const int g_idx = int(tg.x) / 4, hv = int(a.h_v), hkn = int(a.h_k); + const int hk = (hv == 3 * hkn) ? (g_idx / 3) : (g_idx % hkn); + const int h = (hv == 3 * hkn) ? (hk + (g_idx % 3) * hkn) : g_idx; + const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * R; + const int T = int(a.n_tok); + const float scale = 1.0f / sqrt(float(S)); + const int rs = hv * S, s0 = (i0 * hv + h) * S + j; + float st[R]; +#pragma clang loop unroll(full) + for (int i = 0; i < R; ++i) st[i] = (state + s0)[i * rs]; + if (TB == 0) { + constexpr int C = (R + 31) / 32, R1 = R / C; + const ushort lr = lane & (R1 - 1); + const int ks = hkn * S, k0 = hk * S + i0 + int(lr), v0 = h * S + j; + float kl[C], ql[C]; +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) { + kl[c] = kn[k0 + c * 32]; + ql[c] = qn[k0 + c * 32]; + } + float2 g_b = gb[h]; + float v = vc[v0]; + float ka = 0.0f, qa = 0.0f; +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) +#pragma clang loop unroll(full) + for (int r = 0; r < R1; ++r) { + ka = fma(st[c * 32 + r], simd_broadcast(kl[c], ushort(r)), ka); + qa = fma(st[c * 32 + r], simd_broadcast(ql[c], ushort(r)), qa); + } + for (int t = 0; t < T - 1; ++t) { + const int slot = t & 1, tn = t + 1; + const int kn_off = tn * ks + k0, vn_off = tn * rs + v0; + float kn_l[C], qn_l[C]; +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) { + kn_l[c] = kn[kn_off + c * 32]; + qn_l[c] = qn[kn_off + c * 32]; + } + const float2 g_b_n = gb[tn * hv + h]; + const float v_n = vc[vn_off]; + const float kqv = sg == 0 ? kq[t * hkn + hk] : 0.0f; + float2 s = float2(ka, qa); + if (gdn_probe != 1) { + red[slot][sg][lane] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + s = red[slot][0][lane]; +#pragma clang loop unroll(full) + for (int b = 1; b < IB; ++b) s += red[slot][b][lane]; + } + const float g = g_b.x; + const float delta = (v - g * s.x) * g_b.y; + if (sg == 0) o[t * rs + v0] = fma(g, s.y, delta * kqv) * scale; + ka = 0.0f; + qa = 0.0f; + if (gdn_probe != 2) +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) +#pragma clang loop unroll(full) + for (int r = 0; r < R1; ++r) { + const float kr = simd_broadcast(kl[c], ushort(r)); + const float knr = simd_broadcast(kn_l[c], ushort(r)), qnr = simd_broadcast(qn_l[c], ushort(r)); + st[c * 32 + r] = fma(g, st[c * 32 + r], kr * delta); + ka = fma(st[c * 32 + r], knr, ka); + qa = fma(st[c * 32 + r], qnr, qa); + } +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) kl[c] = kn_l[c]; + g_b = g_b_n; + v = v_n; + } + if (T > 0) { + const int t = T - 1, slot = t & 1; + const float kqv = sg == 0 ? kq[t * hkn + hk] : 0.0f; + float2 s = float2(ka, qa); + if (gdn_probe != 1) { + red[slot][sg][lane] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + s = red[slot][0][lane]; +#pragma clang loop unroll(full) + for (int b = 1; b < IB; ++b) s += red[slot][b][lane]; + } + const float g = g_b.x; + const float delta = (v - g * s.x) * g_b.y; + if (sg == 0) o[t * rs + v0] = fma(g, s.y, delta * kqv) * scale; + if (gdn_probe != 2) +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) +#pragma clang loop unroll(full) + for (int r = 0; r < R1; ++r) { + const float kr = simd_broadcast(kl[c], ushort(r)); + st[c * 32 + r] = fma(g, st[c * 32 + r], kr * delta); + } + } + } else { + constexpr int R4 = R / 4, TBS = TB > 0 ? TB : 1; + constexpr bool KC = R4 <= 8; + threadgroup float4 qs[TBS][32], ks[TBS][32]; + threadgroup float vs[TBS][32]; + threadgroup float4 gk[TBS]; + device const float4* qn4_p = (device const float4*)qn; + device const float4* kn4_p = (device const float4*)kn; + for (int t0 = 0; t0 < T; t0 += TBS) { + const int nb = min(TBS, T - t0); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < nb * 32; idx += IB * 32) { + const int tt = idx >> 5, w = idx & 31; + qs[tt][w] = qn4_p[((t0 + tt) * hkn + hk) * 32 + w]; + ks[tt][w] = kn4_p[((t0 + tt) * hkn + hk) * 32 + w]; + vs[tt][w] = vc[((t0 + tt) * hv + h) * S + (j - lane) + w]; + if (w == 0) { + const float2 g_b = gb[(t0 + tt) * hv + h]; + gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float ka = 0.0f, qa = 0.0f; + float4 kc[R4]; +#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; + const float4 g_b = gk[tt]; + const float kqv = g_b.z; + const float v = vs[tt][lane]; + float2 s = float2(ka, qa); + if (gdn_probe != 1) { + red[slot][sg][lane] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + s = red[slot][0][lane]; +#pragma clang loop unroll(full) + for (int b = 1; b < IB; ++b) s += red[slot][b][lane]; + } + 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 = 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); + } + } + } + } + device float* sp = state + s0; +#pragma clang loop unroll(full) + for (int i = 0; i < R; ++i) sp[i * rs] = st[i]; +} + +// gdn_scan_r<8, 16> with the 16 state rows of a lane in four float4 variables (no array): a register-allocation check +kernel void gdn_scan_v(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]], + device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]], + device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]], + device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]], + 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, IB = 8, TB = 16; + threadgroup float4 qs[TB][32], ks[TB][32]; + threadgroup float2 red[2][IB][32]; + threadgroup float vs[TB][32]; + threadgroup float4 gk[TB]; + const int h = int(tg.x) / 4, hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k); + const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * 16, w0 = sg * 4; + const int T = int(a.n_tok); + const float scale = 1.0f / sqrt(float(S)); + device float* sp = state + (i0 * hv + h) * S + j; + const int rs = hv * S; + float4 s0 = float4(sp[0], sp[rs], sp[2 * rs], sp[3 * rs]); + float4 s1 = float4(sp[4 * rs], sp[5 * rs], sp[6 * rs], sp[7 * rs]); + float4 s2 = float4(sp[8 * rs], sp[9 * rs], sp[10 * rs], sp[11 * rs]); + float4 s3 = float4(sp[12 * rs], sp[13 * rs], sp[14 * rs], sp[15 * rs]); + for (int t0 = 0; t0 < T; t0 += TB) { + const int nb = min(TB, T - t0); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < nb * 32; idx += IB * 32) { + const int tt = idx >> 5, w = idx & 31; + qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w]; + ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w]; + vs[tt][w] = vc[((t0 + tt) * hv + h) * S + (j - lane) + w]; + if (w == 0) { + const float2 g_b = gb[(t0 + tt) * hv + h]; + gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int tt = 0; tt < nb; ++tt) { + const int t = t0 + tt, slot = t & 1; + const float4 k0 = ks[tt][w0], k1 = ks[tt][w0 + 1], k2 = ks[tt][w0 + 2], k3 = ks[tt][w0 + 3]; + const float4 q0 = qs[tt][w0], q1 = qs[tt][w0 + 1], q2 = qs[tt][w0 + 2], q3 = qs[tt][w0 + 3]; + const float ka = (dot(s0, k0) + dot(s1, k1)) + (dot(s2, k2) + dot(s3, k3)); + const float qa = (dot(s0, q0) + dot(s1, q1)) + (dot(s2, q2) + dot(s3, q3)); + red[slot][sg][lane] = float2(ka, qa); + const float4 g_b = gk[tt]; + const float v = vs[tt][lane]; + threadgroup_barrier(mem_flags::mem_threadgroup); + float2 r = red[slot][0][lane]; + for (int b = 1; b < IB; ++b) r += red[slot][b][lane]; + const float g = g_b.x; + const float delta = (v - g * r.x) * g_b.y; + if (sg == 0) o[(t * hv + h) * S + j] = fma(g, r.y, delta * g_b.z) * scale; + s0 = fma(float4(g), s0, k0 * delta); + s1 = fma(float4(g), s1, k1 * delta); + s2 = fma(float4(g), s2, k2 * delta); + s3 = fma(float4(g), s3, k3 * delta); + } + } + sp[0] = s0.x; sp[rs] = s0.y; sp[2 * rs] = s0.z; sp[3 * rs] = s0.w; + sp[4 * rs] = s1.x; sp[5 * rs] = s1.y; sp[6 * rs] = s1.z; sp[7 * rs] = s1.w; + sp[8 * rs] = s2.x; sp[9 * rs] = s2.y; sp[10 * rs] = s2.z; sp[11 * rs] = s2.w; + sp[12 * rs] = s3.x; sp[13 * rs] = s3.y; sp[14 * rs] = s3.z; sp[15 * rs] = s3.w; +} + +// grid (h_v / 4, T) or (h_v, T), 128 threads: one simdgroup (32 lanes) per head h computes the RMSNorm and gated SiLU +// over the 128 channels (4 channels a lane) with zero threadgroup memory and zero barriers. +kernel void gdn_post(device const float* o [[buffer(0)]], device const float* z [[buffer(1)]], + device const float* gamma [[buffer(2)]], device half* yh [[buffer(3)]], device float* y [[buffer(4)]], + constant GdnPArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], uint3 tgs [[threadgroups_per_grid]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const int hv = int(a.h_v), t = int(tg.y); + const int h = int(tgs.x) == hv ? (sg == 0 ? int(tg.x) : hv) : int(tg.x) * 4 + int(sg); + if (h >= hv) return; + const uint base = uint(t * hv + h) * 128 + lane; + device const float* o_row = o + base; + device const float* z_row = z + base; + device const float* g_row = gamma + lane; + device half* yh_row = yh + base; + device float* y_row = y + base; + const float ov0 = o_row[0], ov1 = o_row[32], ov2 = o_row[64], ov3 = o_row[96]; + // the gate and gamma do not depend on the norm: computed while the simd_sums run + float sig[4]; +#pragma clang loop unroll(full) + for (uint k = 0; k < 4; ++k) sig[k] = 1.0f / (1.0f + exp(-z_row[k * 32])); + const float gv[4] = {g_row[0], g_row[32], g_row[64], g_row[96]}; + const float4 pv = simd_sum(float4(ov0 * ov0, ov1 * ov1, ov2 * ov2, ov3 * ov3)); + const float tot = (pv.x + pv.y) + (pv.z + pv.w); + const float rs = rsqrt(tot / 128.0f + a.eps); + const float ov[4] = {ov0, ov1, ov2, ov3}; + float yv[4]; +#pragma clang loop unroll(full) + for (uint k = 0; k < 4; ++k) { + yv[k] = ((rs * ov[k]) * gv[k]) * sig[k]; + yh_row[k * 32] = half(yv[k]); + } + if (a.f32) { +#pragma clang loop unroll(full) + for (uint k = 0; k < 4; ++k) y_row[k * 32] = yv[k]; + } +} + +template [[host_name("gdn_scan_1")]] kernel decltype(gdn_scan_t<1, 16>) gdn_scan_t<1, 16>; +template [[host_name("gdn_scan_2")]] kernel decltype(gdn_scan_t<2, 16>) gdn_scan_t<2, 16>; +template [[host_name("gdn_scan_4")]] kernel decltype(gdn_scan_t<4, 16>) gdn_scan_t<4, 16>; +template [[host_name("gdn_scanr_2")]] kernel decltype(gdn_scan_r<2, 0>) gdn_scan_r<2, 0>; +template [[host_name("gdn_scanr_4")]] kernel decltype(gdn_scan_r<4, 0>) gdn_scan_r<4, 0>; +template [[host_name("gdn_scanr_8")]] kernel decltype(gdn_scan_r<8, 0>) gdn_scan_r<8, 0>; +template [[host_name("gdn_scanr_4_8")]] kernel decltype(gdn_scan_r<4, 8>) gdn_scan_r<4, 8>; +template [[host_name("gdn_scanr_4_16")]] kernel decltype(gdn_scan_r<4, 16>) gdn_scan_r<4, 16>; + +// ---- the chunked form (chunks of GC = 32 tokens), the WY / UT form of the gated delta rule. With G_t the product +// of the chunk's gates up to t (log space: lg_t) and S0 the state entering the chunk: +// (I + A) D = B, A_ts = beta_t (k_t . k_s) G_t / G_s (s < t), B = [beta v | beta G k] -> D = [U | W] +// delta = U - W S0, o_t = G_t S0^T q_t + sum_{s <= t} (q_t . k_s) G_t / G_s delta_s, +// S_L = G_L S0 + sum_s G_L / G_s k_s delta_s^T +// (the sequential recurrence's algebra; the rounding differs). Three kernels: +// gdn_cprep per (head, chunk): lg, A, the forward substitution for U and W (a thread per column of B) +// gdn_cstate per (head, 32 state columns): the chunks in order - delta, the inter-chunk output, the state update +// gdn_cout per (head, chunk): o += the intra-chunk part, times 1/sqrt(128) +constant constexpr int GC = 32; + +// grid (h_v, chunks), 256 threads. U, W: [h][chunk][GC][128]; lgo: [h][chunk][GC] +kernel void gdn_cprep(device const float* qn [[buffer(0)]], device const float* kn [[buffer(1)]], + device const float* vc [[buffer(2)]], device const float2* gb [[buffer(3)]], + device float* U [[buffer(4)]], device float* W [[buffer(5)]], device float* lgo [[buffer(6)]], + constant GdnPArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int S = 128; + threadgroup float A[GC][GC]; + threadgroup float lg[GC], bt[GC]; + const int h = int(tg.x), c = int(tg.y), hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok); + const int NC = (T + GC - 1) / GC, t0 = c * GC, nb = min(GC, T - t0); + if (tid == 0) { + float s = 0.0f; + for (int t = 0; t < GC; ++t) { + if (t < nb) s += max(log(gb[(t0 + t) * hv + h].x), -80.0f); + lg[t] = s; + bt[t] = t < nb ? gb[(t0 + t) * hv + h].y : 0.0f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < GC * GC; idx += 256) { + const int t = idx / GC, s = idx % GC; + float v = 0.0f; + if (s < t && t < nb) { + device const float4* kt = (device const float4*)(kn + ((t0 + t) * hkn + hk) * S); + device const float4* ks = (device const float4*)(kn + ((t0 + s) * hkn + hk) * S); + float d = 0.0f; + for (int i = 0; i < S / 4; ++i) d += dot(kt[i], ks[i]); + v = bt[t] * d * exp(lg[t] - lg[s]); + } + A[t][s] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // forward substitution, column j of B a thread (threads 0..127: beta v; 128..255: beta G k) + const int j = tid & 127; + const bool isw = tid >= 128; + float x[GC]; + for (int t = 0; t < GC; ++t) { + float b = 0.0f; + if (t < nb) b = isw ? bt[t] * exp(lg[t]) * kn[((t0 + t) * hkn + hk) * S + j] : bt[t] * vc[((t0 + t) * hv + h) * S + j]; + for (int s = 0; s < t; ++s) b = fma(-A[t][s], x[s], b); + x[t] = b; + } + device float* o = (isw ? W : U) + (ulong(h) * NC + c) * GC * S + j; + for (int t = 0; t < GC; ++t) o[t * S] = x[t]; + if (tid < GC) lgo[(h * NC + c) * GC + tid] = lg[tid]; +} + +// grid (h_v * 4), 256 threads = 8 row blocks of 16 x 32 columns; the state's 128 x 32 block in registers (16 a thread) +kernel void gdn_cstate(device float* state [[buffer(0)]], device const float* qn [[buffer(1)]], + device const float* kn [[buffer(2)]], device const float* U [[buffer(3)]], + device const float* W [[buffer(4)]], device const float* lgo [[buffer(5)]], + device float* delta [[buffer(6)]], device float* o [[buffer(7)]], constant GdnPArgs& a [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128, RB = 16; + threadgroup float red[8][16][32]; + threadgroup float dl[GC][32]; + const int h = int(tg.x) / 4, hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok); + const int NC = (T + GC - 1) / GC; + const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * RB; + float st[RB]; + for (int r = 0; r < RB; ++r) st[r] = state[((i0 + r) * hv + h) * S + j]; + for (int c = 0; c < NC; ++c) { + const int t0 = c * GC, nb = min(GC, T - t0); + device const float* Uc = U + (ulong(h) * NC + c) * GC * S; + device const float* Wc = W + (ulong(h) * NC + c) * GC * S; + device const float* lg = lgo + (h * NC + c) * GC; + // pass 0: delta = U - W S0; pass 1: o_inter = G_t S0^T q_t; 16 tokens at a time + for (int pass = 0; pass < 2; ++pass) { + for (int hf = 0; hf < 2; ++hf) { + float p[16]; + for (int u = 0; u < 16; ++u) { + const int t = hf * 16 + u; + float s = 0.0f; + if (t < nb) { + device const float* v = pass == 0 ? Wc + t * S + i0 : qn + ((t0 + t) * hkn + hk) * S + i0; + for (int r = 0; r < RB; ++r) s = fma(v[r], st[r], s); + } + p[u] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int u = 0; u < 16; ++u) red[sg][u][lane] = p[u]; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int u = sg * 2; u < sg * 2 + 2; ++u) { + const int t = hf * 16 + u; + float s = 0.0f; + for (int b = 0; b < 8; ++b) s += red[b][u][lane]; + if (t < nb) { + if (pass == 0) { + const float d = Uc[t * S + j] - s; + dl[t][lane] = d; + delta[((ulong(h) * NC + c) * GC + t) * S + j] = d; + } else { + o[((t0 + t) * hv + h) * S + j] = exp(lg[t]) * s; + } + } else if (pass == 0) { + dl[t][lane] = 0.0f; + } + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // S = G_L S0 + sum_s G_L / G_s k_s delta_s^T + const float lgl = lg[nb - 1], gl = exp(lgl); + for (int r = 0; r < RB; ++r) st[r] *= gl; + for (int s = 0; s < nb; ++s) { + const float f = exp(lgl - lg[s]) * dl[s][lane]; + device const float* ks = kn + ((t0 + s) * hkn + hk) * S + i0; + for (int r = 0; r < RB; ++r) st[r] = fma(ks[r], f, st[r]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (int r = 0; r < RB; ++r) state[((i0 + r) * hv + h) * S + j] = st[r]; +} + +// grid (h_v, chunks), 128 threads (column j): o_t = (o_inter_t + sum_{s <= t} (q_t . k_s) G_t / G_s delta_s) / sqrt(128) +kernel void gdn_cout(device const float* qn [[buffer(0)]], device const float* kn [[buffer(1)]], + device const float* delta [[buffer(2)]], device const float* lgo [[buffer(3)]], + device float* o [[buffer(4)]], constant GdnPArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int S = 128; + threadgroup float M[GC][GC]; + const int h = int(tg.x), c = int(tg.y), hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok); + const int NC = (T + GC - 1) / GC, t0 = c * GC, nb = min(GC, T - t0); + device const float* lg = lgo + (h * NC + c) * GC; + for (int idx = tid; idx < GC * GC; idx += 128) { + const int t = idx / GC, s = idx % GC; + float v = 0.0f; + if (s <= t && t < nb) { + device const float4* qt = (device const float4*)(qn + ((t0 + t) * hkn + hk) * S); + device const float4* ks = (device const float4*)(kn + ((t0 + s) * hkn + hk) * S); + float d = 0.0f; + for (int i = 0; i < S / 4; ++i) d += dot(qt[i], ks[i]); + v = d * exp(lg[t] - lg[s]); + } + M[t][s] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int j = tid; + device const float* dc = delta + (ulong(h) * NC + c) * GC * S + j; + float dv[GC]; + for (int s = 0; s < GC; ++s) dv[s] = s < nb ? dc[s * S] : 0.0f; + const float scale = 1.0f / sqrt(float(S)); + for (int t = 0; t < nb; ++t) { + float s = o[((t0 + t) * hv + h) * S + j]; + for (int u = 0; u <= t; ++u) s = fma(M[t][u], dv[u], s); + o[((t0 + t) * hv + h) * S + j] = s * scale; + } +} diff --git a/results/round39/model-ab/baseline-kernels/gemv.metal b/results/round39/model-ab/baseline-kernels/gemv.metal new file mode 100644 index 0000000..49b7b62 --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/gemv.metal @@ -0,0 +1,500 @@ +// metal/kernels/gemv.metal - y[t][r] = W[r] . x[t] for a window of 1..8 tokens (dense projections, the shared +// expert, the LM head), any format of quant.metal/units.metal. +// +// Shape of a threadgroup: SG simdgroups; KS of them split the row's units (KS = 1: no split), the other SG/KS cover +// different rows. In a simdgroup, LPR lanes share a row and 32/LPR rows run side by side; each lane holds NR rows. +// The weights of a unit are decoded once and applied to every token. Partial sums of a K split are added in a +// fixed order (deterministic). + +struct GemvArgs { + uint n_in; + uint n_out; + uint n_tok; + uint row_bytes; + uint x_stride; // floats between tokens' x rows + uint y_stride; // floats between tokens' y rows +}; + +constant int NTOK [[function_constant(0)]]; // tokens of the window (1..8): the token loops are specialised +// a bench probe (function constant 9): 1 = tokens past the first reuse token 0's x (no load; wrong values), +// 2 = tokens past the first skip the combine (wrong values) +constant int GEMV_PROBE [[function_constant(9)]]; +constant int gemv_probe = is_function_constant_defined(GEMV_PROBE) ? GEMV_PROBE : 0; + +// Wider raw-dot windows and prepared Q2 at NR=1 did not improve on M5. +template +static inline bool gemv_direct_x() { + return F == T_Q8_0 && (NTOK == 1 || (NTOK == 2 && LPR == 16 && NR == 1)); +} + +// Keep unit_combine's four FMA chains without materialising the raw x slice. +template +static inline void gemv_raw_combine(device const float* xp, thread const float (&q)[NR][UNIT], + thread const float4* sc, thread float (&acc)[NR][8], int t) { + constexpr int H = UNIT / 2; + device const float4* x = (device const float4*)xp; + float s[NR][4]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) s[r][0] = s[r][1] = s[r][2] = s[r][3] = 0.0f; +#pragma clang loop unroll(full) + for (int j = 0; j < H; j += 4) { + const float4 a = x[j / 4], b = x[(H + j) / 4]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + s[r][0] = fma(q[r][j], a.x, s[r][0]); + s[r][1] = fma(q[r][j + 1], a.y, s[r][1]); + s[r][2] = fma(q[r][H + j], b.x, s[r][2]); + s[r][3] = fma(q[r][H + j + 1], b.y, s[r][3]); + s[r][0] = fma(q[r][j + 2], a.z, s[r][0]); + s[r][1] = fma(q[r][j + 3], a.w, s[r][1]); + s[r][2] = fma(q[r][H + j + 2], b.z, s[r][2]); + s[r][3] = fma(q[r][H + j + 3], b.w, s[r][3]); + } + } +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + acc[r][t] += sc[r].x * (s[r][0] + s[r][1]) + sc[r].y * (s[r][2] + s[r][3]) - + (sc[r].z * 0.0f + sc[r].w * 0.0f); +} + +static inline float gemv_q2_raw_combine(ushort2 pk, float d, thread const float* xs) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; j += 2) { + const uint2 q = uint2(pk) >> (2 * j); + const uint2 bits = ((q | (q << 14)) & 0x00030003u) | 0x64006400u; + const float4 v = float4(as_type(bits) - 1025.0h); + s0 = fma(v.x, xs[j], s0); + s1 = fma(v.y, xs[j + 1], s1); + s2 = fma(v.z, xs[8 + j], s2); + s3 = fma(v.w, xs[8 + j + 1], s3); + } + return d * (s0 + s1) + d * (s2 + s3) - 0.0f; +} + +// IQ4_XS kept as its 16 code bytes and scale (5 registers a row instead of 36), decoded inside the combine. M5, weights +// from DRAM: 2560 -> 640 at 16,2,4,1 -1.5..-5% (1 token), -2..-3% (2..4); 3..4 tokens -1..-2% on 2560 -> 6144..12288 +// (but 6144 -> 2560 at 16,2,4,1 3 tokens +1.4%); LPR 32 under 4 tokens +0.4..0.7%, so those keep the decoded unit. +// IQ4_NL (18-byte blocks, ushort reads) measured no change and keeps the decoded unit. +template +static inline bool gemv_iq4_packed() { return LPR <= 16 || NTOK >= 4; } + +static inline uint4 gemv_iq4xs_codes(device const uchar* row, int u, thread float& d) { + device const uchar* b = row + (u >> 3) * 136; + d = iq4xs_scale(b, u & 7); + return uint4(*(device const packed_uint4*)(b + 8 + 16 * (u & 7))); +} + +// unit_combine<32> of deq's q, decoding the codes as it goes: the same four FMA chains in the same order +static inline float gemv_iq4_combine(uint4 pk, float d, thread const float* xs, threadgroup const uint* tab) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const uint v = pk[j >> 1] >> (16 * (j & 1)); + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[(v >> 8) & 255]); + s0 = fma(float(w0.x), xs[2 * j], s0); + s1 = fma(float(w1.x), xs[2 * j + 1], s1); + s2 = fma(float(w0.y), xs[16 + 2 * j], s2); + s3 = fma(float(w1.y), xs[16 + 2 * j + 1], s3); + } + return d * (s0 + s1) + d * (s2 + s3) - 0.0f; +} + +template +kernel void gemv(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RS = SG / KS, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + threadgroup float red[KS > 1 ? KS * RS * RG * NR * 8 : 1]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int ks = sg / RS, rs = sg % RS; + const int li = lane % LPR; + const int lrow0 = (rs * RG + lane / LPR) * NR; + const int row0 = int(tg.x) * RS * RG * NR + lrow0; + const int U = int(a.n_in) / UNIT; + const int u0 = ks * U / KS, u1 = (ks + 1) * U / KS; + float acc[NR][8]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) rp[r] = W + (ulong)min(row0 + r, int(a.n_out) - 1) * a.row_bytes; + const int s16n = int(a.n_in) / 16; + if (gemv_probe == 3) { + // rows outer, tokens inner: one row's decoded unit in registers at a time (x re-read per row from L1) + for (int u = u0 + li; u < u1; u += LPR) +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 sc; + Dec::run(rp[r], u, q, sc, tab); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + acc[r][t] += unit_combine(q, sc, xs, xsum); + } + } + } else + for (int u = u0 + li; u < u1; u += LPR) { + if (NTOK == 1 && F == T_Q2_0R) { + ushort2 pk[NR]; + float d[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + device const uchar* b = rp[r] + (u >> 2) * 18; + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + pk[r] = ushort2(p[0], p[1]); + d[r] = rdh(b); + } + float xs[16]; + load16(xp + 16 * u, xs); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][0] += gemv_q2_raw_combine(pk[r], d[r], xs); + } else if (F == T_IQ4_XS && gemv_iq4_packed()) { + uint4 pk[NR]; + float d[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) pk[r] = gemv_iq4xs_codes(rp[r], u, d[r]); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[32]; + load16(xp + t * int(a.x_stride) + 32 * u, xs); + load16(xp + t * int(a.x_stride) + 32 * u + 16, xs + 16); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += gemv_iq4_combine(pk[r], d[r], xs, tab); + } + } else { + float q[NR][UNIT]; + float4 sc[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); + float xs[UNIT], xsum[2]; +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + if (gemv_direct_x()) { + if (gemv_probe != 2 || t == 0) + gemv_raw_combine(xp + (gemv_probe == 1 ? 0 : t) * int(a.x_stride) + u * UNIT, + q, sc, acc, t); + } else { + if (gemv_probe != 1 || t == 0) X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + if (gemv_probe != 2 || t == 0) +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + } + } + constexpr int RPT = RS * RG * NR; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + if (KS == 1) { + if (li == 0 && row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = v; + } else if (li == 0) { + red[(ks * RPT + lrow0 + r) * 8 + t] = v; + } + } + if (KS > 1) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ks == 0 && li == 0) +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + if (row0 + r < int(a.n_out)) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float v = 0.0f; +#pragma clang loop unroll(full) + for (int k = 0; k < KS; ++k) v += red[(k * RPT + lrow0 + r) * 8 + t]; + y[t * int(a.y_stride) + row0 + r] = v; + } + } +} + +#define GV(F, L, R, S, K) template [[host_name("gemv_" #F "_" #L "_" #R "_" #S "_" #K)]] kernel decltype(gemv) gemv; +#define GV_SET(F) GV(F, 32, 1, 4, 1) GV(F, 32, 2, 4, 1) GV(F, 16, 1, 4, 1) GV(F, 16, 2, 4, 1) GV(F, 8, 1, 4, 1) \ + GV(F, 8, 2, 4, 1) GV(F, 4, 2, 4, 1) GV(F, 4, 1, 4, 1) GV(F, 32, 1, 4, 4) GV(F, 32, 2, 4, 2) \ + GV(F, 32, 1, 8, 8) GV(F, 16, 1, 8, 4) GV(F, 32, 1, 8, 2) GV(F, 8, 1, 8, 1) GV(F, 1, 1, 4, 1) GV(F, 1, 2, 4, 1) GV(F, 2, 1, 4, 1) GV(F, 2, 2, 4, 1) GV(F, 1, 1, 2, 1) GV(F, 4, 4, 4, 1) +#ifndef GEMV_TYPES +#define GEMV_TYPES GV_SET(T_F32) GV_SET(T_BF16) GV_SET(T_Q8_0) GV_SET(T_Q4_0) GV_SET(T_Q5_0) GV_SET(T_IQ4_NL) GV_SET(T_IQ4_XS) \ + GV_SET(T_Q2_0) GV_SET(T_IQ3_S) GV_SET(T_Q3_K) GV_SET(T_Q4_K) GV_SET(T_Q5_K) GV_SET(T_Q6_K) +#endif +GEMV_TYPES + +// The shared expert's gate and up in one dispatch: h[t][r] = silu(Wg[r] . x[t]) * (Wu[r] . x[t]), each with its own +// format (and so its own prepared activation kind); h is written raw (its consumer reads raw activations). +// KS = 1; tiles as gemv. Buffers: Wg, xg (Wg's kind), Wu, xu (Wu's kind), s16, h, args (row_bytes = Wg's), rb_up. +template +kernel void gemv_pair(device const uchar* Wg [[buffer(0)]], device const float* xg [[buffer(1)]], + device const uchar* Wu [[buffer(2)]], device const float* xu [[buffer(3)]], + device const float* s16 [[buffer(4)]], device float* h [[buffer(5)]], + constant GemvArgs& a [[buffer(6)]], constant uint& rb_up [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int U1 = fmt::UNIT, U2 = fmt::UNIT, RG = 32 / LPR; + constexpr int TB = fmt::TAB > fmt::TAB ? fmt::TAB : fmt::TAB; + threadgroup uint tab1[fmt::TAB]; + threadgroup uint tab2[fmt::TAB]; + if (fmt::TAB > 1) load_tab(tab1, tid, SG * 32); + if (fmt::TAB > 1) load_tab(tab2, tid, SG * 32); + if (TB > 1) threadgroup_barrier(mem_flags::mem_threadgroup); + const int li = lane % LPR; + const int row = (int(tg.x) * SG + sg) * RG + lane / LPR; + const int rr = min(row, int(a.n_out) - 1); + device const uchar* rg = Wg + ulong(rr) * a.row_bytes; + device const uchar* ru = Wu + ulong(rr) * rb_up; + const int s16n = int(a.n_in) / 16; + float ag[8], au[8]; +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) ag[t] = au[t] = 0.0f; + for (int u = li; u < int(a.n_in) / U1; u += LPR) { + float q[U1]; + float4 sc; + Dec::run(rg, u, q, sc, tab1); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[U1], xsum[2]; + X::load(xg + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + ag[t] += unit_combine(q, sc, xs, xsum); + } + } + for (int u = li; u < int(a.n_in) / U2; u += LPR) { + float q[U2]; + float4 sc; + Dec::run(ru, u, q, sc, tab2); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[U2], xsum[2]; + X::load(xu + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + au[t] += unit_combine(q, sc, xs, xsum); + } + } +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float gv = lane_sum(ag[t]), uv = lane_sum(au[t]); + if (li == 0 && row < int(a.n_out)) h[t * int(a.y_stride) + row] = gv / (1.0f + exp(-gv)) * uv; + } +} + +#define GP_(F1, F2) template [[host_name("gemv_pair_" #F1 "_" #F2)]] kernel decltype(gemv_pair) gemv_pair; +#define GP(F1, F2) GP_(F1, F2) +#define GP_ROW(F1) GP(F1, T_IQ3_S) GP(F1, T_IQ4_XS) GP(F1, T_Q2_0) GP(F1, T_Q3_K) GP(F1, T_Q4_K) GP(F1, T_Q5_K) GP(F1, T_Q6_K) +GP_ROW(T_IQ3_S) GP_ROW(T_IQ4_XS) GP_ROW(T_Q2_0) GP_ROW(T_Q3_K) GP_ROW(T_Q4_K) GP_ROW(T_Q5_K) GP_ROW(T_Q6_K) +// Q2_0 with raw activations (the shared expert's down reading the pair's raw h) +GV_SET(T_Q2_0R) + +// gemv_cat: two matrices of the same format and n_in over the same x in one dispatch - rows [0, n_split) are W's +// (into y), rows [n_split, n_out) are W2's (into y2, from its row 0). Each row's arithmetic is gemv's (KS = 1), so the +// outputs are bitwise the two separate GEMVs'. (GDN's ssm_alpha + ssm_beta: 48 + 48 rows.) +// Buffers: W, xp, s16, y, args (n_out = both), W2, y2, n_split. +template +kernel void gemv_cat(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], + device const uchar* W2 [[buffer(5)]], device float* y2 [[buffer(6)]], constant uint& n_split [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int li = lane % LPR; + const int row0 = int(tg.x) * SG * RG * NR + (sg * RG + lane / LPR) * NR; + const int U = int(a.n_in) / UNIT, sp = int(n_split); + float acc[NR][8]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + const int row = min(row0 + r, int(a.n_out) - 1); + rp[r] = row < sp ? W + (ulong)row * a.row_bytes : W2 + (ulong)(row - sp) * a.row_bytes; + } + const int s16n = int(a.n_in) / 16; + for (int u = li; u < U; u += LPR) { + float q[NR][UNIT]; + float4 sc[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + if (gemv_direct_x()) { + gemv_raw_combine(xp + t * int(a.x_stride) + u * UNIT, q, sc, acc, t); + } else { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + } +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + const int row = row0 + r; + if (li == 0 && row < int(a.n_out)) { + if (row < sp) y[t * int(a.y_stride) + row] = v; + else y2[t * int(a.y_stride) + row - sp] = v; + } + } +} +#define GVC_(F, L, R, S) template [[host_name("gemvc_" #F "_" #L "_" #R "_" #S)]] kernel decltype(gemv_cat) gemv_cat; +#define GVC(F, L, R, S) GVC_(F, L, R, S) +GVC(T_BF16, 32, 1, 4) GVC(T_BF16, 32, 1, 8) GVC(T_BF16, 16, 1, 4) GVC(T_BF16, 32, 2, 4) +GVC(T_F32, 32, 1, 4) GVC(T_F32, 32, 1, 8) GVC(T_F32, 16, 1, 4) GVC(T_F32, 32, 2, 4) +GVC(T_Q8_0, 32, 1, 4) GVC(T_Q8_0, 32, 1, 8) GVC(T_Q8_0, 16, 1, 4) GVC(T_Q8_0, 32, 2, 4) GVC(T_Q8_0, 8, 1, 8) + +// gemv_catk: gemv_cat with gemv's K split (KS simdgroups share a row, partial sums added in gemv's fixed order), so +// the outputs are bitwise gemv_____'s. (The MTP draft step's k + v and shared gate + up: Q8_0 +// 2560 -> 512 + 512 / 640 + 640 at 32,1,4,4 on the M5.) Buffers as gemv_cat. +template +kernel void gemv_catk(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], + device const uchar* W2 [[buffer(5)]], device float* y2 [[buffer(6)]], constant uint& n_split [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RS = SG / KS, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + threadgroup float red[KS * RS * RG * NR * 8]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int ks = sg / RS, rs = sg % RS; + const int li = lane % LPR; + const int lrow0 = (rs * RG + lane / LPR) * NR; + const int row0 = int(tg.x) * RS * RG * NR + lrow0; + const int U = int(a.n_in) / UNIT, sp = int(n_split); + const int u0 = ks * U / KS, u1 = (ks + 1) * U / KS; + float acc[NR][8]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + const int row = min(row0 + r, int(a.n_out) - 1); + rp[r] = row < sp ? W + (ulong)row * a.row_bytes : W2 + (ulong)(row - sp) * a.row_bytes; + } + const int s16n = int(a.n_in) / 16; + for (int u = u0 + li; u < u1; u += LPR) { + float q[NR][UNIT]; + float4 sc[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + if (gemv_direct_x()) { + gemv_raw_combine(xp + t * int(a.x_stride) + u * UNIT, q, sc, acc, t); + } else { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + } + constexpr int RPT = RS * RG * NR; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + if (li == 0) red[(ks * RPT + lrow0 + r) * 8 + t] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ks == 0 && li == 0) +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + const int row = row0 + r; + if (row < int(a.n_out)) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float v = 0.0f; +#pragma clang loop unroll(full) + for (int k = 0; k < KS; ++k) v += red[(k * RPT + lrow0 + r) * 8 + t]; + if (row < sp) y[t * int(a.y_stride) + row] = v; + else y2[t * int(a.y_stride) + row - sp] = v; + } + } +} +#define GVCK_(F, L, R, S, K) template [[host_name("gemvck_" #F "_" #L "_" #R "_" #S "_" #K)]] kernel decltype(gemv_catk) gemv_catk; +#define GVCK(F, L, R, S, K) GVCK_(F, L, R, S, K) +GVCK(T_Q8_0, 32, 1, 4, 4) GVCK(T_Q8_0, 32, 2, 4, 2) GVCK(T_Q8_0, 32, 1, 8, 8) + +// gemv_x: gemv (KS = 1) with the window's prepared x (the format's plane, NTOK rows, and its per-16 sums) staged in +// threadgroup memory once per threadgroup, so the per-unit x reads of every token come from threadgroup memory +// instead of L1 (on the M4 the x traffic through L1 - 4 B a weight a token, shared by the rows in a simdgroup - is +// what a window token costs). Same tile meaning and the same arithmetic as gemv_____1: bitwise. +// Dynamic threadgroup memory 0: 4 * NTOK * (n_in + n_in / 16) bytes. +template +kernel void gemv_x(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], + threadgroup float* xt [[threadgroup(0)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1) load_tab(tab, tid, SG * 32); + const int n = int(a.n_in), s16n = n / 16; + threadgroup float* st = xt + NTOK * n; +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + device const float4* src = (device const float4*)(xp + t * int(a.x_stride)); + threadgroup float4* dst = (threadgroup float4*)(xt + t * n); + for (int i = tid; i < n / 4; i += SG * 32) dst[i] = src[i]; + for (int i = tid; i < s16n; i += SG * 32) st[t * s16n + i] = s16[t * s16n + i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int li = lane % LPR; + const int row0 = int(tg.x) * SG * RG * NR + (sg * RG + lane / LPR) * NR; + const int U = n / UNIT; + float acc[NR][8]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) rp[r] = W + (ulong)min(row0 + r, int(a.n_out) - 1) * a.row_bytes; + for (int u = li; u < U; u += LPR) { + float q[NR][UNIT]; + float4 sc[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[UNIT], xsum[2]; + X::load((threadgroup const float*)(xt + t * n), (threadgroup const float*)(st + t * s16n), u, xs, xsum); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + if (li == 0 && row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = v; + } +} +#define GVX(F, L, R, S) template [[host_name("gemvx_" #F "_" #L "_" #R "_" #S)]] kernel decltype(gemv_x) gemv_x; +#define GVX_SET(F) GVX(F, 32, 1, 4) GVX(F, 32, 2, 4) GVX(F, 16, 1, 4) GVX(F, 16, 2, 4) GVX(F, 8, 1, 4) GVX(F, 8, 2, 4) \ + GVX(F, 4, 2, 4) GVX(F, 16, 2, 8) GVX(F, 8, 2, 8) GVX(F, 16, 4, 4) GVX(F, 32, 2, 8) GVX(F, 2, 2, 4) +GVX_SET(T_BF16) GVX_SET(T_Q8_0) GVX_SET(T_Q4_0) GVX_SET(T_Q5_0) GVX_SET(T_IQ4_NL) GVX_SET(T_IQ4_XS) GVX_SET(T_Q2_0) +GVX_SET(T_IQ3_S) GVX_SET(T_Q3_K) GVX_SET(T_Q4_K) GVX_SET(T_Q5_K) GVX_SET(T_Q6_K) diff --git a/results/round39/model-ab/baseline-kernels/glue.metal b/results/round39/model-ab/baseline-kernels/glue.metal new file mode 100644 index 0000000..b34bb90 --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/glue.metal @@ -0,0 +1,72 @@ +// metal/kernels/glue.metal - small elementwise steps between the big kernels (after quant/units/moe.metal). +// +// swiglu_prep h = silu(g) * u for the shared expert, written as prepared activations (xprep layout) +// moe_combine_sh bo[t] = sum_k ye[entry(t, k)] + sigmoid(sg[t]) * ysh[t] (k in order) + +struct GlueArgs { + uint n; // vector width + uint n_tok; + uint n_embd; + uint topk; +}; + +constant float glue_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, + 0.00006103515625f}; +constant float glue_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + +// grid (n / 16, n_tok), 16 threads: one group of 16 values +kernel void swiglu_prep(device const float* g [[buffer(0)]], device const float* u [[buffer(1)]], + device float* xp [[buffer(2)]], constant GlueArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort j [[thread_index_in_threadgroup]]) { + const uint t = tg.y, d = tg.x * 16 + j, idx = t * a.n + d, plane = a.n_tok * a.n; + const float gv = g[idx]; + const float h = gv / (1.0f + exp(-gv)) * u[idx]; + device float* raw = xp + idx; + raw[0] = h; + raw[plane] = h * glue_q2s[j & 7]; + raw[2 * plane] = h * glue_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? h * 0.00390625f : h; + float s0 = 0.0f, s1 = 0.0f; +#pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(h, ushort(k)); + s1 += simd_shuffle_down(h, ushort(k + 1)); + } + if (j == 0) xp[4 * plane + t * (a.n / 16) + tg.x] = s0 + s1; +} + +kernel void moe_combine_sh(device const float* ye [[buffer(0)]], device const int* plan [[buffer(1)]], + device const float* ysh [[buffer(2)]], device const float* sg [[buffer(3)]], + device float* bo [[buffer(4)]], constant GlueArgs& a [[buffer(5)]], + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.n_embd, k = gid % a.n_embd; + if (t >= a.n_tok) return; + const int row_base = int(t * a.topk); + device const int* te = plan + P_TOKENT + row_base; + device const float* yep = ye + k; + float s = 0.0f; + if (a.topk == 10) { +#pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += yep[te[j] * int(a.n_embd)]; + } else { + for (uint j = 0; j < a.topk; ++j) s += yep[te[j] * int(a.n_embd)]; + } + const float gs = 1.0f / (1.0f + exp(-sg[t])); + bo[gid] = fma(gs, ysh[gid], s); +} + +// slc_touch: reads n16 16-byte words (a weight range the next kernels will need) so they sit in the SLC; the xor is +// written only if it equals a value it never takes (keeps the loads alive). grid (ceil(n16 / 1024)), 256 threads, +// 4 words a thread. +kernel void slc_touch(device const uint4* src [[buffer(0)]], device uint* sink [[buffer(1)]], constant uint& n16 [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + uint s = 0; + for (uint k = 0; k < 4; ++k) { + const uint j = (gid >> 8) * 1024 + k * 256 + (gid & 255); + if (j < n16) { + const uint4 v = src[j]; + s ^= v.x ^ v.y ^ v.z ^ v.w; + } + } + if (s == 0x9e3779b9u) sink[0] = s; +} diff --git a/results/round39/model-ab/baseline-kernels/gr.metal b/results/round39/model-ab/baseline-kernels/gr.metal new file mode 100644 index 0000000..c18d2ac --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/gr.metal @@ -0,0 +1,307 @@ +// metal/kernels/gr.metal - the hyper-connection (gated residual) read for a window of tokens, with the previous +// sublayer's write folded in. The math is fused_gr.hpp's (native contract: FP32 activations, BF16 weights): +// +// gr_norm R'[c] = R[c] + bo * 2 sigmoid(inj_prev[c] / hc) (when apply; written back in place) +// xn[c] = R'[c] * w_norm[c] * rsqrt(mean(R'[c]^2) + eps) +// gr_down lo[k] = silu((w_down[k] . xn) / hc) k < hc_lr; inject[c] = w_inject[c] . xn +// gr_up mixed[d] = mean_c xn[c][d] * sigmoid(w_up[c n + d] . lo), written as prepared activations (xprep's +// planes and per-16 sums) for the GEMVs that read it +// +// Shapes: n_embd 2560, hc 4, hc_lr 320. Per token t: R[t] (hc * n), bo[t] (n), inj_prev[t] (hc), xn[t] (hc * n), +// lo[t] (hc_lr), inject[t] (hc), XP planes of n_tok rows. + +constant int GR_NTOK [[function_constant(1)]]; + +// the prepared activations' per-position plane scales (units.metal's xprep): q2 planes 4^-(j%8), p4 planes 256^-(j%4) +constant float gr_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, + 0.00006103515625f}; +constant float gr_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + +struct GrArgs { + uint n_embd; + uint hc; + uint hc_lr; + uint n_tok; + uint apply; + float eps; +}; + +// one threadgroup per (stream c, token t), 256 threads +kernel void gr_norm(device float* R [[buffer(0)]], device const float* bo [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device float* xn [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint c = tg.x, t = tg.y, n = a.n_embd; + device float* r = R + (t * a.hc + c) * n; + const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; + float rv[10]; + float ss = 0.0f; + if (n == 2560) { + device float* rp = r + tid; + if (a.apply) { + device const float* bp = bo + t * 2560u + tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const float v = fma(bp[i * 256], g, rp[i * 256]); + rp[i * 256] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + rv[i] = rp[i * 256]; + ss = fma(rv[i], rv[i], ss); + } + } + } else { + for (uint d = tid; d < n; d += 256) { + float v = r[d]; + if (a.apply) { + v = fma(bo[t * n + d], g, v); + r[d] = v; + } + if (n <= 2560) rv[d / 256] = v; + ss = fma(v, v, ss); + } + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device float* o = xn + (t * a.hc + c) * n; + device const float* w = w_norm + c * n; + if (n == 2560) { + o += tid; + w += tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) o[i * 256] = rv[i] * w[i * 256] * rs; + } else { + for (uint d = tid; d < n; d += 256) + o[d] = (n <= 2560 ? rv[d / 256] : r[d]) * w[d] * rs; + } +} + +// rows: 0 .. hc_lr - 1 = w_down, hc_lr .. hc_lr + hc - 1 = w_inject (both bf16 rows of hc * n_embd). +// SG simdgroups split the row's units KS ways; LPR = 32 (a simdgroup per row); RS = SG / KS simdgroups per +// threadgroup, NR rows each. +// one token's dot with a decoded unit (gr_up): x is the token's slice of lo +template +static inline float gr_dot(device const float* x, thread const float* q, float4 sc) { + float xs[UNIT], xsum[2] = {0.0f, 0.0f}; + #pragma clang loop unroll(full) + for (int k = 0; k < UNIT; k += 16) load16(x + k, xs + k); + return unit_combine(q, sc, xs, xsum); +} + +// a unit's decoded rows for grd_acc, read four codes at a time +template +struct GrdUnit { + float q[NR][fmt::UNIT]; + float4 sc[NR]; + void load(device const uchar* w, uint rb, int u, threadgroup const uint* notab) thread { + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(w + ulong(r) * rb, u, q[r], sc[r], notab); + } + float4 v(int r, int i) const thread { return float4(q[r][4 * i], q[r][4 * i + 1], q[r][4 * i + 2], q[r][4 * i + 3]); } +}; + +// Q8_0 two rows a simdgroup: the codes stay packed (8 registers a row, not 32 floats) and are widened per token. +// Measured 10-16% faster past 6 tokens (where the token loop is rolled), no faster below, and slower for one row. +template <> +struct GrdUnit { + char4 q[2][8]; + float4 sc[2]; + void load(device const uchar* w, uint rb, int u, threadgroup const uint*) thread { + #pragma clang loop unroll(full) + for (int r = 0; r < 2; ++r) { + device const uchar* b = w + ulong(r) * rb + 34 * u; + device const ushort* p = (device const ushort*)(b + 2); + #pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) q[r][k] = as_type(uint(p[2 * k]) | (uint(p[2 * k + 1]) << 16)); + const float d = rdh(b); + sc[r] = float4(d, d, 0.0f, 0.0f); + } + } + float4 v(int r, int i) const thread { return float4(q[r][i]); } +}; + +// NR consecutive rows' K-split partials over xn (units of format F; the rows rb bytes apart), acc[r * 8 + t]. +// unit_combine's sums in its order, with each xn float4 read straight from memory once for all NR rows: xn is +// 15-30x the weights' bytes at 4-8 tokens, so NR = 2 halves the kernel's traffic. GR_NTOK is a function constant, +// which the compiler does not unroll over by itself: the pragma keeps acc in registers instead of the stack. Past 6 +// tokens the unrolled loop measured slower than the rolled one, so it stays rolled there; ROLL keeps it rolled at +// every count (the packed Q8_0 pair at KS = 4: 1-3% faster at 4-6 tokens, where KS = 8 measured 2-6% slower). +template +static inline void grd_acc(device const uchar* w, uint rb, int nin, int ks, int KS, ushort lane, device const float* xn, + thread float* acc, threadgroup const uint* notab) { + constexpr int UNIT = fmt::UNIT, H = UNIT / 2; + const int U = nin / UNIT, u0 = ks * U / KS, u1 = (ks + 1) * U / KS; + for (int u = u0 + lane; u < u1; u += 32) { + GrdUnit d; + d.load(w, rb, u, notab); + auto tok = [&](int t) { + device const float4* x = (device const float4*)(xn + t * nin + UNIT * u); + float s[NR][4]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) s[r][0] = s[r][1] = s[r][2] = s[r][3] = 0.0f; + #pragma clang loop unroll(full) + for (int j = 0; j < H; j += 4) { + const float4 a = x[j / 4], b = x[(H + j) / 4]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + const float4 qa = d.v(r, j / 4), qb = d.v(r, (H + j) / 4); + s[r][0] = fma(qa.x, a.x, s[r][0]); + s[r][1] = fma(qa.y, a.y, s[r][1]); + s[r][2] = fma(qb.x, b.x, s[r][2]); + s[r][3] = fma(qb.y, b.y, s[r][3]); + s[r][0] = fma(qa.z, a.z, s[r][0]); + s[r][1] = fma(qa.w, a.w, s[r][1]); + s[r][2] = fma(qb.z, b.z, s[r][2]); + s[r][3] = fma(qb.w, b.w, s[r][3]); + } + } + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + acc[r * 8 + t] += d.sc[r].x * (s[r][0] + s[r][1]) + d.sc[r].y * (s[r][2] + s[r][3]) - + (d.sc[r].z * 0.0f + d.sc[r].w * 0.0f); + }; + if (GR_NTOK <= 6 && !ROLL) { + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) tok(t); + } else { + for (int t = 0; t < GR_NTOK; ++t) tok(t); + } + } +} + +// FI: the inject rows' format (the Unsloth files: Q8_0 down, F32 inject); FI = F is the original kernel. NR must +// divide hc_lr and hc, so that a simdgroup's rows are all w_down or all w_inject. The sums do not depend on SG or NR: +// every SG / NR with the same KS is bitwise the same kernel. +template +kernel void gr_down_f(device const uchar* w_down [[buffer(0)]], device const uchar* w_inject [[buffer(1)]], + device const float* xn [[buffer(2)]], device float* lo [[buffer(3)]], + device float* inject [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + constexpr int RS = SG / KS; + threadgroup float red[KS * RS * NR * 8]; + threadgroup uint notab[1]; + const int ks = sg / RS, rs = sg % RS; + const int row0 = (int(tg.x) * RS + rs) * NR; + const int nrows = int(a.hc_lr + a.hc); + const int rr = min(row0, nrows - NR); + const int nin = int(a.hc * a.n_embd); + const uint rb = uint(nin / fmt::QK * fmt::BB), rbi = uint(nin / fmt::QK * fmt::BB); + float acc[NR * 8]; + #pragma clang loop unroll(full) + for (int t = 0; t < NR * 8; ++t) acc[t] = 0.0f; + constexpr bool roll = F == T_Q8_0 && NR == 2 && KS == 4; + if (rr < int(a.hc_lr)) grd_acc(w_down + ulong(rr) * rb, rb, nin, ks, KS, lane, xn, acc, notab); + else grd_acc(w_inject + ulong(rr - int(a.hc_lr)) * rbi, rbi, nin, ks, KS, lane, xn, acc, notab); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) { + const float v = simd_sum(acc[r * 8 + t]); + if (lane == 0) red[((ks * RS + rs) * NR + r) * 8 + t] = v; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ks == 0 && lane < NR * 8) { + const int r = lane / 8, t = lane % 8, row = row0 + r; + if (t < GR_NTOK && row < nrows) { + float v = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < KS; ++k) v += red[((k * RS + rs) * NR + r) * 8 + t]; + if (row < int(a.hc_lr)) { + const float z = v / float(a.hc); + lo[t * int(a.hc_lr) + row] = z / (1.0f + exp(-z)); + } else { + inject[t * int(a.hc) + row - int(a.hc_lr)] = v; + } + } + } +} + +// A threadgroup owns 16 consecutive d for all hc = 4 streams: 8 simdgroups, simdgroup s -> stream s / 2 and +// d0 + 8 (s % 2) .. + 8; LPR = 4 lanes per row (320 = 20 units of 16, 5 per lane). +// XP: [4 kinds][n_tok][n] then S16 [n_tok][n / 16] (units.metal's xprep layout). +template +kernel void gr_up_f(device const uchar* w_up [[buffer(0)]], device const float* lo [[buffer(1)]], + device const float* xn [[buffer(2)]], device float* xp [[buffer(3)]], + constant GrArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float gate[8][4][16]; + threadgroup uint notab[1]; + const int n = int(a.n_embd), c = sg >> 1, dl = 8 * (sg & 1) + lane / 4, li = lane & 3; + const int d = int(tg.x) * 16 + dl; + constexpr int UNIT = fmt::UNIT; + device const uchar* w = w_up + ulong(c * n + d) * (a.hc_lr / fmt::QK * fmt::BB); + const int U = int(a.hc_lr) / UNIT; + float acc[8]; + #pragma clang loop unroll(full) + for (int t = 0; t < 8; ++t) acc[t] = 0.0f; + for (int u = li; u < U; u += 4) { + float q[UNIT]; + float4 sc; + Dec::run(w, u, q, sc, notab); + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) acc[t] += gr_dot(lo + t * int(a.hc_lr) + UNIT * u, q, sc); + } + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) { + const float v = lane_sum<4>(acc[t]); + if (li == 0) gate[t][c][dl] = 1.0f / (1.0f + exp(-v)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // 16 threads per token: thread j -> d0 + j + const int t = tid / 16, j = tid & 15; + if (t < GR_NTOK) { + const int dd = int(tg.x) * 16 + j; + const int hc = int(a.hc); + float m = 0.0f; + if (hc == 4) { + #pragma clang loop unroll(full) + for (int cc = 0; cc < 4; ++cc) m = fma(xn[(t * 4 + cc) * n + dd], gate[t][cc][j], m); + } else { + for (int cc = 0; cc < hc; ++cc) m = fma(xn[(t * hc + cc) * n + dd], gate[t][cc][j], m); + } + m = m / float(hc); + const uint plane = a.n_tok * uint(n); + device float* raw = xp + t * n + dd; + raw[0] = m; + raw[plane] = m * gr_q2s[j & 7]; + raw[2 * plane] = m * gr_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? m * 0.00390625f : m; + float s0 = 0.0f, s1 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(m, ushort(k)); + s1 += simd_shuffle_down(m, ushort(k + 1)); + } + if (j == 0) xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1; + } +} + +#define GRD1(S, K, R, SFX) \ + template [[host_name("gr_down_" #S "_" #K SFX)]] kernel decltype(gr_down_f) gr_down_f; \ + template [[host_name("gr_down_q8_" #S "_" #K SFX)]] kernel decltype(gr_down_f) gr_down_f; \ + template [[host_name("gr_down_q8f_" #S "_" #K SFX)]] kernel decltype(gr_down_f) gr_down_f; +#define GRD(S, K) GRD1(S, K, 1, "") +GRD(8, 8) GRD(8, 4) GRD(4, 4) GRD(4, 2) GRD(16, 16) GRD(16, 8) +// two rows a simdgroup (gr_read past 4 tokens) +GRD1(8, 8, 2, "_r2") GRD1(4, 4, 2, "_r2") +template [[host_name("gr_up")]] kernel decltype(gr_up_f) gr_up_f; +template [[host_name("gr_up_q8")]] kernel decltype(gr_up_f) gr_up_f; + +template [[host_name("gr_down_q4_8_8")]] kernel decltype(gr_down_f<8, 8, 1, T_Q4_0, T_Q8_0>) gr_down_f<8, 8, 1, T_Q4_0, T_Q8_0>; +template [[host_name("gr_up_q4")]] kernel decltype(gr_up_f) gr_up_f; +template [[host_name("gr_down_q5_8_8")]] kernel decltype(gr_down_f<8, 8, 1, T_Q5_0, T_Q8_0>) gr_down_f<8, 8, 1, T_Q5_0, T_Q8_0>; +template [[host_name("gr_up_q5")]] kernel decltype(gr_up_f) gr_up_f; diff --git a/results/round39/model-ab/baseline-kernels/gr2.metal b/results/round39/model-ab/baseline-kernels/gr2.metal new file mode 100644 index 0000000..58342f0 --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/gr2.metal @@ -0,0 +1,484 @@ +// metal/kernels/gr2.metal - the hyper-connection read in TWO dispatches (after quant/units/gr.metal), for BF16 HC +// weights (the exact path) or a Q8_0 copy of them (the STRATA_HC_Q8 option): +// +// gr_down3 grid (row blocks, hc streams). The RMSNorm is per stream, so a threadgroup owns one stream c: it forms +// R'_c = R_c + bo * 2 sigmoid(inj_prev_c / hc) (10 KB a token), its sum of squares and rs_c itself, +// stages xn_c = (R'_c * w_norm_c) * rs_c in threadgroup memory once, and every simdgroup dots its rows' +// stream-c quarter with it: D[t][row][c]. Row block 0 also writes R'_c to R_out (ping-pong, never in +// place) and rs_c. Rows 0..hc_lr-1 are w_down, hc_lr.. the inject rows. +// gr_up3 lo[k] = silu(sum_c D[t][k][c] / hc) formed in threadgroup memory, then as gr_up with xn = R' * w_norm +// * rs on the fly, writing the prepared activations. +// inj_prev of the NEXT read = sum_c D[t][hc_lr + c'][c] of this one (gr_down3 sums it, in c order). +// +// Per element the arithmetic is gr_norm's (fma(bo, g, R), then (R' * w) * rs); the K order differs (per stream, then +// the 4 streams in order). + +#ifndef P_TOKENT +#define P_TOKENT 532 // moe.metal's plan layout +#endif + +struct GrArgs3 { + uint n_embd; + uint hc; + uint hc_lr; + uint n_tok; // the prepared planes' row count (the plane stride) + uint apply; + float eps; + uint rb_down; // bytes per row of w_down / w_inject + uint rb_up; // bytes per row of w_up + uint comb; // 1: bo is not given; it is the MoE combine (moe_combine_sh's arithmetic) of buffers 10..13 + uint topk; +}; + +// bo[t][d], given or the MoE combine: sum_k ye[entry(t, k)][d] (k in order) + sigmoid(sg[t]) * ysh[t][d] +static inline float bo_at(device const float* bo, device const int* plan, device const float* ye, + device const float* ysh, device const float* sg, constant GrArgs3& a, int t, int d) { + const int n = int(a.n_embd); + if (!a.comb) return bo[t * n + d]; + float s = 0.0f; + for (uint j = 0; j < a.topk; ++j) s += ye[plan[P_TOKENT + t * int(a.topk) + j] * n + d]; + const float gs = 1.0f / (1.0f + exp(-sg[t])); + return fma(gs, ysh[t * n + d], s); +} + +template +kernel void gr_down3(device const uchar* w_down [[buffer(0)]], device const uchar* w_inject [[buffer(1)]], + device const float* Rin [[buffer(2)]], device float* Rout [[buffer(3)]], + device const float* bo [[buffer(4)]], device const float* Dprev [[buffer(5)]], + device const float* wn [[buffer(6)]], device float* D [[buffer(7)]], + device float* rs_out [[buffer(8)]], constant GrArgs3& a [[buffer(9)]], + device const int* plan [[buffer(10)]], device const float* ye [[buffer(11)]], + device const float* ysh [[buffer(12)]], device const float* sgv [[buffer(13)]], + 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 UNIT = fmt::UNIT; + threadgroup float xs_tg[TT][2560]; + threadgroup float part[8][TT]; + threadgroup float g_tg[TT], rs_tg[TT]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, 256); + const int n = int(a.n_embd), hc = int(a.hc), NIN = hc * n, c = int(tg.y); + const int nrows = int(a.hc_lr) + hc; + for (int t0 = 0; t0 < GR_NTOK; t0 += TT) { + const int T = min(TT, GR_NTOK - t0); + if (tid < T) { + float g = 0.0f; + if (a.apply) { + device const float* dp = Dprev + ((t0 + tid) * nrows + int(a.hc_lr) + c) * hc; + const float ip = ((dp[0] + dp[1]) + dp[2]) + dp[3]; + g = 2.0f / (1.0f + exp(-ip / float(hc))); + } + g_tg[tid] = g; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int t = 0; t < T; ++t) { + device const float* rt = Rin + (t0 + t) * NIN + c * n; + float ss = 0.0f; + for (int d = tid; d < n; d += 256) { + float v = rt[d]; + if (a.apply) v = fma(bo_at(bo, plan, ye, ysh, sgv, a, t0 + t, d), g_tg[t], v); + xs_tg[t][d] = v; + ss = fma(v, v, ss); + if (tg.x == 0) Rout[(t0 + t) * NIN + c * n + d] = v; + } + ss = simd_sum(ss); + if (lane == 0) part[sg][t] = ss; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < T) { + float tot = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) tot += part[k][tid]; + const float rsv = rsqrt(tot / float(n) + a.eps); + rs_tg[tid] = rsv; + if (tg.x == 0) rs_out[(t0 + tid) * hc + c] = rsv; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int t = 0; t < T; ++t) + for (int d = tid; d < n; d += 256) xs_tg[t][d] = xs_tg[t][d] * wn[c * n + d] * rs_tg[t]; + threadgroup_barrier(mem_flags::mem_threadgroup); + // the rows: simdgroup sg holds RPS rows; lanes split the stream's units + int rows[RPS]; + device const uchar* wr[RPS]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + rows[k] = (int(tg.x) * 8 + sg) * RPS + k; + const int rr = min(rows[k], nrows - 1); + wr[k] = rr < int(a.hc_lr) ? w_down + ulong(rr) * a.rb_down : w_inject + ulong(rr - int(a.hc_lr)) * a.rb_down; + } + float acc[RPS][TT]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) acc[k][t] = 0.0f; + } + const int U = n / UNIT, u0 = c * U; + for (int u = lane; u < U; u += 32) { + float q[RPS][UNIT]; + float4 sc[RPS]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) Dec::run(wr[k], u0 + u, q[k], sc[k], tab); + // TT-bounded with a T guard, so acc[][t] is indexed by a constant (registers, not the stack) + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) { + if (t >= T) break; + float xs[UNIT], xsum[2] = {0.0f, 0.0f}; + #pragma clang loop unroll(full) + for (int e = 0; e < UNIT; ++e) xs[e] = xs_tg[t][u * UNIT + e]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) acc[k][t] += unit_combine(q[k], sc[k], xs, xsum); + } + } + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) { + if (t >= T) break; + const float v = simd_sum(acc[k][t]); + if (lane == 0 && rows[k] < nrows) D[((t0 + t) * nrows + rows[k]) * hc + c] = v; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +// grid (n_embd / 16), 256 threads. +template +kernel void gr_up3(device const uchar* w_up [[buffer(0)]], device const float* D [[buffer(1)]], + device const float* Rn [[buffer(2)]], device float* xp [[buffer(3)]], + constant GrArgs3& a [[buffer(4)]], device const float* wn [[buffer(5)]], + device const float* rs [[buffer(6)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT; + threadgroup float lo_tg[8][320]; + threadgroup float gate[8][4][16]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, 256); + const int n = int(a.n_embd), hc = int(a.hc), lr = int(a.hc_lr), nrows = lr + hc; + for (int e = tid; e < GR_NTOK * lr; e += 256) { + const int t = e / lr, k = e - t * lr; + device const float* dp = D + (t * nrows + k) * hc; + const float z = (((dp[0] + dp[1]) + dp[2]) + dp[3]) / float(hc); + lo_tg[t][k] = z / (1.0f + exp(-z)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int c = sg >> 1, dl = 8 * (sg & 1) + lane / 4, li = lane & 3; + const int d = int(tg.x) * 16 + dl; + device const uchar* w = w_up + ulong(c * n + d) * a.rb_up; + const int U = lr / UNIT; + float acc[8]; + #pragma clang loop unroll(full) + for (int t = 0; t < 8; ++t) acc[t] = 0.0f; + for (int u = li; u < U; u += 4) { + float q[UNIT]; + float4 sc; + Dec::run(w, u, q, sc, tab); + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) { + float xs[UNIT], xsum[2] = {0.0f, 0.0f}; + #pragma clang loop unroll(full) + for (int e = 0; e < UNIT; ++e) xs[e] = lo_tg[t][UNIT * u + e]; + acc[t] += unit_combine(q, sc, xs, xsum); + } + } + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) { + const float v = lane_sum<4>(acc[t]); + if (li == 0) gate[t][c][dl] = 1.0f / (1.0f + exp(-v)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int t = tid / 16, j = tid & 15; + if (t < GR_NTOK) { + const int dd = int(tg.x) * 16 + j; + float m = 0.0f; + if (hc == 4) { + #pragma clang loop unroll(full) + for (int cc = 0; cc < 4; ++cc) { + const float xn = Rn[(t * 4 + cc) * n + dd] * wn[cc * n + dd] * rs[t * 4 + cc]; + m = fma(xn, gate[t][cc][j], m); + } + } else { + for (int cc = 0; cc < hc; ++cc) { + const float xn = Rn[(t * hc + cc) * n + dd] * wn[cc * n + dd] * rs[t * hc + cc]; + m = fma(xn, gate[t][cc][j], m); + } + } + m = m / float(hc); + const uint plane = a.n_tok * uint(n); + device float* raw = xp + t * n + dd; + raw[0] = m; + raw[plane] = m * gr_q2s[j & 7]; + raw[2 * plane] = m * gr_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? m * 0.00390625f : m; + float s0 = 0.0f, s1 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(m, ushort(k)); + s1 += simd_shuffle_down(m, ushort(k + 1)); + } + if (j == 0) xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1; + } +} + +#define GD3(F, T, R) template [[host_name("gr_down3_" #F "_" #T "_" #R)]] kernel decltype(gr_down3) gr_down3; +GD3(30, 1, 1) GD3(30, 2, 1) GD3(30, 3, 1) GD3(30, 1, 2) GD3(30, 2, 2) GD3(30, 3, 2) +GD3(8, 1, 1) GD3(8, 2, 1) GD3(8, 3, 1) GD3(8, 1, 2) GD3(8, 2, 2) GD3(8, 3, 2) +template [[host_name("gr_up3_30")]] kernel decltype(gr_up3) gr_up3; +template [[host_name("gr_up3_8")]] kernel decltype(gr_up3) gr_up3; + +// BF16 only: as gr_down3, but each lane loads its rows' weights (5 units of 16 per row and lane: 2560 = 160 units) +// into registers before the norm phase, so the weight stream is in flight while the threadgroup forms xn. +template +kernel void gr_down3b(device const uchar* w_down [[buffer(0)]], device const uchar* w_inject [[buffer(1)]], + device const float* Rin [[buffer(2)]], device float* Rout [[buffer(3)]], + device const float* bo [[buffer(4)]], device const float* Dprev [[buffer(5)]], + device const float* wn [[buffer(6)]], device float* D [[buffer(7)]], + device float* rs_out [[buffer(8)]], constant GrArgs3& a [[buffer(9)]], + device const int* plan [[buffer(10)]], device const float* ye [[buffer(11)]], + device const float* ysh [[buffer(12)]], device const float* sgv [[buffer(13)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float xs_tg[TT][2560]; + threadgroup float part[8][TT]; + threadgroup float g_tg[TT], rs_tg[TT]; + const int n = int(a.n_embd), hc = int(a.hc), NIN = hc * n, c = int(tg.y); + const int nrows = int(a.hc_lr) + hc; + int rows[RPS]; + uint4 wr[RPS][5][2]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + rows[k] = (int(tg.x) * 8 + sg) * RPS + k; + const int rr = min(rows[k], nrows - 1); + device const uchar* w = rr < int(a.hc_lr) ? w_down + ulong(rr) * a.rb_down : w_inject + ulong(rr - int(a.hc_lr)) * a.rb_down; + device const uint4* p = (device const uint4*)(w + 2 * (c * n)); + #pragma clang loop unroll(full) + for (int i = 0; i < 5; ++i) { + wr[k][i][0] = p[2 * (lane + 32 * i)]; + wr[k][i][1] = p[2 * (lane + 32 * i) + 1]; + } + } + for (int t0 = 0; t0 < GR_NTOK; t0 += TT) { + const int T = min(TT, GR_NTOK - t0); + if (tid < T) { + float g = 0.0f; + if (a.apply) { + device const float* dp = Dprev + ((t0 + tid) * nrows + int(a.hc_lr) + c) * hc; + const float ip = ((dp[0] + dp[1]) + dp[2]) + dp[3]; + g = 2.0f / (1.0f + exp(-ip / float(hc))); + } + g_tg[tid] = g; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int t = 0; t < T; ++t) { + device const float* rt = Rin + (t0 + t) * NIN + c * n; + float ss = 0.0f; + for (int d = tid; d < n; d += 256) { + float v = rt[d]; + if (a.apply) v = fma(bo_at(bo, plan, ye, ysh, sgv, a, t0 + t, d), g_tg[t], v); + xs_tg[t][d] = v; + ss = fma(v, v, ss); + if (tg.x == 0) Rout[(t0 + t) * NIN + c * n + d] = v; + } + ss = simd_sum(ss); + if (lane == 0) part[sg][t] = ss; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < T) { + float tot = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) tot += part[k][tid]; + const float rsv = rsqrt(tot / float(n) + a.eps); + rs_tg[tid] = rsv; + if (tg.x == 0) rs_out[(t0 + tid) * hc + c] = rsv; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int t = 0; t < T; ++t) + for (int d = tid; d < n; d += 256) xs_tg[t][d] = xs_tg[t][d] * wn[c * n + d] * rs_tg[t]; + threadgroup_barrier(mem_flags::mem_threadgroup); + float acc[RPS][TT]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) acc[k][t] = 0.0f; + } + #pragma clang loop unroll(full) + for (int i = 0; i < 5; ++i) { + const int u = lane + 32 * i; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + const uint v[8] = {wr[k][i][0].x, wr[k][i][0].y, wr[k][i][0].z, wr[k][i][0].w, + wr[k][i][1].x, wr[k][i][1].y, wr[k][i][1].z, wr[k][i][1].w}; + float q[16]; + #pragma clang loop unroll(full) + for (int e = 0; e < 8; ++e) { + q[2 * e] = as_type(v[e] << 16); + q[2 * e + 1] = as_type(v[e] & 0xFFFF0000u); + } + const float4 sc = float4(1.0f, 1.0f, 0.0f, 0.0f); + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) { + if (t >= T) break; + float xs[16], xsum[2] = {0.0f, 0.0f}; + #pragma clang loop unroll(full) + for (int e = 0; e < 16; ++e) xs[e] = xs_tg[t][u * 16 + e]; + acc[k][t] += unit_combine<16>(q, sc, xs, xsum); + } + } + } + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) { + if (t >= T) break; + const float v = simd_sum(acc[k][t]); + if (lane == 0 && rows[k] < nrows) D[((t0 + t) * nrows + rows[k]) * hc + c] = v; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} +#define GD3B(T, R) template [[host_name("gr_down3b_" #T "_" #R)]] kernel decltype(gr_down3b) gr_down3b; +GD3B(1, 1) GD3B(2, 1) GD3B(3, 1) GD3B(1, 2) GD3B(2, 2) GD3B(3, 2) + +// gr_norm_comb: gr_norm (gr.metal) with the sublayer output bo formed on the fly as the MoE combine - bo[t][d] = +// sum_k ye[entry(t, k)][d] (k in order) + sigmoid(sg[t]) * ysh[t][d], moe_combine_sh's arithmetic - so the combine's +// dispatch goes away; R', xn bitwise gr_norm's after moe_combine_sh. One threadgroup per (stream c, token t). +kernel void gr_norm_comb(device float* R [[buffer(0)]], device const float* ye [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device float* xn [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + device const int* plan [[buffer(6)]], device const float* ysh [[buffer(7)]], + device const float* sgv [[buffer(8)]], constant uint& topk [[buffer(9)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint c = tg.x, t = tg.y, n = a.n_embd; + device float* r = R + (t * a.hc + c) * n; + const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + float rv[10]; + float ss = 0.0f; + if (n == 2560 && topk == 10) { + device const float* yep[10]; + if (a.apply) { + #pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) + yep[j] = ye + uint(plan[P_TOKENT + t * 10 + j]) * 2560u + tid; + } + device const float* yshp = ysh + t * 2560u + tid; + device float* rp = r + tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + float v = rp[i * 256]; + if (a.apply) { + float s = 0.0f; + #pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += yep[j][i * 256]; + v = fma(fma(gs, yshp[i * 256], s), g, v); + rp[i * 256] = v; + } + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid; d < n; d += 256) { + float v = r[d]; + if (a.apply) { + float s = 0.0f; + for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; + const float bo = fma(gs, ysh[t * n + d], s); + v = fma(bo, g, v); + r[d] = v; + } + if (n <= 2560) rv[d / 256] = v; + ss = fma(v, v, ss); + } + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device float* o = xn + (t * a.hc + c) * n; + device const float* w = w_norm + c * n; + if (n == 2560) { + o += tid; + w += tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) o[i * 256] = rv[i] * w[i * 256] * rs; + } else { + for (uint d = tid; d < n; d += 256) + o[d] = (n <= 2560 ? rv[d / 256] : r[d]) * w[d] * rs; + } +} + +// Fixed n_embd = 2560, hc = 4; grid (1, n_tok), 256 threads. Form bo once for all streams. +kernel void gr_norm_comb_4s(device float* R [[buffer(0)]], device const float* ye [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device float* xn [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + device const int* plan [[buffer(6)]], device const float* ysh [[buffer(7)]], + device const float* sgv [[buffer(8)]], constant uint& topk [[buffer(9)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + const uint t = tg.y, n = 2560; + device float* rp = R + t * 4 * n + tid; + device float* op = xn + t * 4 * n + tid; + device const float* wp = w_norm + tid; + float rv[4][10], ss[4] = {0.0f, 0.0f, 0.0f, 0.0f}, g[4]; + #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) + g[c] = a.apply ? 2.0f / (1.0f + exp(-inj[t * 4 + c] / 4.0f)) : 0.0f; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + device const float* yep[10]; + if (a.apply && topk == 10) { + #pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) + yep[j] = ye + uint(plan[P_TOKENT + t * 10 + j]) * 2560u + tid; + } + device const float* yshp = ysh + t * n + tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const uint d = tid + i * 256; + float bo = 0.0f; + if (a.apply) { + float s = 0.0f; + if (topk == 10) { + #pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += yep[j][i * 256]; + } else { + for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; + } + bo = fma(gs, yshp[i * 256], s); + } + #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) { + device float* r = rp + c * n; + float v = r[i * 256]; + if (a.apply) { v = fma(bo, g[c], v); r[i * 256] = v; } + rv[c][i] = v; + ss[c] = fma(v, v, ss[c]); + } + } + #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) { + const float s = simd_sum(ss[c]); + if (lane == 0) part[c][sg] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) { + float tot = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[c][i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device float* o = op + c * n; + device const float* w = wp + c * n; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + o[i * 256] = rv[c][i] * w[i * 256] * rs; + } + } +} diff --git a/results/round39/model-ab/baseline-kernels/gridbar.metal b/results/round39/model-ab/baseline-kernels/gridbar.metal new file mode 100644 index 0000000..5cd9a0f --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/gridbar.metal @@ -0,0 +1,54 @@ +// metal/kernels/gridbar.metal - feasibility probe for a persistent kernel: phases separated by an in-kernel grid +// barrier (a device atomic counter) instead of dispatch boundaries. Metal does not promise that all threadgroups of +// a grid are resident at once; the probe launches at most one threadgroup per GPU core. +// +// Each phase p: every thread streams its share of chunk p (n uint4), and threadgroup g writes a token value +// (p * 1000 + g) to flag[p][g]; after the barrier, phase p + 1 checks every flag of phase p (written by the OTHER +// threadgroups) and counts mismatches into bad[0]. + +static inline void grid_barrier(device atomic_uint* bar, uint target, ushort tid) { + threadgroup_barrier(mem_flags::mem_device); +#ifdef GB_FENCE + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); +#endif + if (tid == 0) { + atomic_fetch_add_explicit(bar, 1u, memory_order_relaxed); + while (atomic_load_explicit(bar, memory_order_relaxed) < target) {} + } + threadgroup_barrier(mem_flags::mem_device); +#ifdef GB_FENCE + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); +#endif +} + +kernel void k_phases(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + device atomic_uint* bar [[buffer(2)]], device volatile uint* flag [[buffer(3)]], + device atomic_uint* bad [[buffer(4)]], constant uint& n [[buffer(5)]], + constant uint& phases [[buffer(6)]], + uint tgpx [[threadgroup_position_in_grid]], uint ntgx [[threadgroups_per_grid]], + ushort tid [[thread_index_in_threadgroup]], uint gid [[thread_position_in_grid]], + uint nt [[threads_per_grid]]) { + const uint g = tgpx, ntg = ntgx; + uint4 acc = 0; + for (uint p = 0; p < phases; ++p) { + if (p > 0 && tid < ntg) { + const uint v = flag[(p - 1) * ntg + tid]; + if (v != (p - 1) * 1000 + tid) atomic_fetch_add_explicit(bad, 1u, memory_order_relaxed); + } + device const uint4* s = src + ulong(p) * n; + for (uint i = gid; i < n; i += nt) acc += s[i]; + if (tid == 0) flag[p * ntg + g] = p * 1000 + g; + grid_barrier(bar, (p + 1) * ntg, tid); + } + out[gid] = acc.x ^ acc.y ^ acc.z ^ acc.w; +} + +// one phase per dispatch, for comparison +kernel void k_phase1(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + constant uint& n [[buffer(2)]], constant uint& p [[buffer(3)]], + uint gid [[thread_position_in_grid]], uint nt [[threads_per_grid]]) { + uint4 acc = 0; + device const uint4* s = src + ulong(p) * n; + for (uint i = gid; i < n; i += nt) acc += s[i]; + out[gid] = acc.x ^ acc.y ^ acc.z ^ acc.w; +} diff --git a/results/round39/model-ab/baseline-kernels/mm_bench.metal b/results/round39/model-ab/baseline-kernels/mm_bench.metal new file mode 100644 index 0000000..c04659f --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/mm_bench.metal @@ -0,0 +1,85 @@ +// metal/kernels/mm_bench.metal - matrix throughput probes (mm_bench.cpp): C[M][N] = A[M][K] . B[N][K]^T, fp32 out. +// mpp_mm Metal 4 MetalPerformancePrimitives matmul2d (the M5's neural accelerators), a TM x TN +// tile of C per threadgroup, K dynamic, A and B straight from device memory +// sg_mm simdgroup_matrix 8x8, a 32 x 32 tile per simdgroup (4 x 4 accumulators), 4 simdgroups +#include +using namespace mpp::tensor_ops; + +struct MmArgs { int M, N, K; }; + +template +kernel void mpp_mm(device T* A [[buffer(0)]], device T* B [[buffer(1)]], device float* C [[buffer(2)]], + constant MmArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]]) { + // extents are (inner, outer): A is M rows of K, B is N rows of K, C is M rows of N + tensor, tensor_inline> tA(A, dextents(a.K, a.M)); + tensor, tensor_inline> tB(B, dextents(a.K, a.N)); + tensor, tensor_inline> tC(C, dextents(a.N, a.M)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, int(tg.y) * TM); + auto mB = tB.slice(0, int(tg.x) * TN); + auto mC = tC.slice(int(tg.x) * TN, int(tg.y) * TM); + op.run(mA, mB, mC); +} + +#define MPP(T, M, N, S) template [[host_name("mpp_mm_" #T "_" #M "_" #N "_" #S)]] kernel decltype(mpp_mm) mpp_mm; +MPP(half, 64, 32, 4) MPP(half, 64, 64, 4) MPP(half, 128, 64, 4) MPP(half, 32, 64, 4) MPP(half, 64, 128, 4) +MPP(half, 16, 64, 2) MPP(half, 128, 128, 4) +MPP(half, 64, 16, 2) MPP(half, 64, 16, 4) MPP(half, 128, 16, 4) MPP(half, 128, 32, 4) MPP(half, 64, 32, 2) MPP(half, 32, 16, 2) +MPP(half, 32, 32, 2) MPP(half, 16, 32, 1) MPP(half, 16, 16, 1) MPP(half, 128, 48, 4) MPP(half, 64, 48, 4) +MPP(float, 64, 32, 4) MPP(float, 64, 64, 4) MPP(float, 128, 64, 4) +MPP(bfloat, 64, 64, 4) MPP(bfloat, 128, 64, 4) + +// simdgroup_matrix: threadgroup tile 64 x 64 (4 simdgroups x 32 x 32), K in steps of 8, operands from device memory +template +kernel void sg_mm(device const T* A [[buffer(0)]], device const T* B [[buffer(1)]], device float* C [[buffer(2)]], + constant MmArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + const int m0 = int(tg.y) * 64 + (sg >> 1) * 32, n0 = int(tg.x) * 64 + (sg & 1) * 32; + simdgroup_matrix acc[4][4]; + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) acc[i][j] = make_filled_simdgroup_matrix(0.0f); + for (int k = 0; k < a.K; k += 8) { + simdgroup_matrix ma[4], mb[4]; + for (int i = 0; i < 4; ++i) simdgroup_load(ma[i], A + (m0 + 8 * i) * a.K + k, a.K); + for (int j = 0; j < 4; ++j) simdgroup_load(mb[j], B + (n0 + 8 * j) * a.K + k, a.K, ulong2(0, 0), true); + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) simdgroup_multiply_accumulate(acc[i][j], ma[i], mb[j], acc[i][j]); + } + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) simdgroup_store(acc[i][j], C + (m0 + 8 * i) * a.N + n0 + 8 * j, a.N); +} +template [[host_name("sg_mm_half")]] kernel decltype(sg_mm) sg_mm; +template [[host_name("sg_mm_float")]] kernel decltype(sg_mm) sg_mm; + + +// the element layout of single-simdgroup matmul2d cooperative tensors (right input, left input, destination) for a +// TM x TN x TK descriptor: out[((which * 32 + lane) * 128 + i) * 4 + (0, 1)] = the two indices, [2] = capacity +template +kernel void coop_layout(device int* out [[buffer(0)]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + auto r = op.template get_right_input_cooperative_tensor(); + auto l = op.template get_left_input_cooperative_tensor(); + auto dd = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < dd.get_capacity() && i < 128; ++i) { + auto ix = dd.get_multidimensional_index(i); + device int* o = out + ((2 * 32 + lane) * 128 + i) * 4; + o[0] = ix[0]; o[1] = ix[1]; o[2] = dd.get_capacity(); + } + for (uint16_t i = 0; i < r.get_capacity() && i < 128; ++i) { + auto ix = r.get_multidimensional_index(i); + device int* o = out + ((0 * 32 + lane) * 128 + i) * 4; + o[0] = ix[0]; o[1] = ix[1]; o[2] = r.get_capacity(); + } + for (uint16_t i = 0; i < l.get_capacity() && i < 128; ++i) { + auto ix = l.get_multidimensional_index(i); + device int* o = out + ((1 * 32 + lane) * 128 + i) * 4; + o[0] = ix[0]; o[1] = ix[1]; o[2] = l.get_capacity(); + } +} +template [[host_name("coop_layout_16_32_16")]] kernel decltype(coop_layout<16, 32, 16>) coop_layout<16, 32, 16>; +template [[host_name("coop_layout_32_32_16")]] kernel decltype(coop_layout<32, 32, 16>) coop_layout<32, 32, 16>; +template [[host_name("coop_layout_16_64_32")]] kernel decltype(coop_layout<16, 64, 32>) coop_layout<16, 64, 32>; +template [[host_name("coop_layout_16_32_32")]] kernel decltype(coop_layout<16, 32, 32>) coop_layout<16, 32, 32>; diff --git a/results/round39/model-ab/baseline-kernels/mma.metal b/results/round39/model-ab/baseline-kernels/mma.metal new file mode 100644 index 0000000..8fcc0ab --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/mma.metal @@ -0,0 +1,279 @@ +// metal/kernels/mma.metal - a GEMV for windows of tokens on the simdgroup matrix units (after quant/units.metal). +// +// y[t][r] = W[r] . x[t] for NTOK (1..8) tokens: each simdgroup owns 8 rows; the K axis runs in chunks of 128. Per +// chunk the simdgroup decodes its 8 x 128 weights - with ref_quant's arithmetic, so a tile element is exactly the +// dequantized weight - into threadgroup memory, then 16 simdgroup_multiply_accumulate steps take A = the 8 x 8 weight +// tile and B = the 8 (k) x 8 (token) slice of x, loaded transposed straight from device memory (x has 8 rows: the +// tokens past the window are padding, and an output column depends only on its own token). +// +// Activations are the raw plane (no pre-scaling); n_in must be a multiple of 128. Every format's units sit inside +// one 128-wide chunk: Q4_K/Q5_K units span 64 positions, Q6_K units 128, the others are contiguous. + +// ---- decoders to actual weights: w[UNIT] and the chunk offsets of its two halves (p0: w[0..H), p1: w[H..UNIT)) ---- +template struct DW { + // generic: w = s * q (quant.metal's deq is ref_quant's order: the scale times the code), contiguous + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT, UPB = fmt::QK / UNIT; + float q[UNIT]; + const float s = deq(row + (u / UPB) * fmt::BB, u % UPB, q, tab); +#pragma clang loop unroll(full) + for (int j = 0; j < UNIT; ++j) w[j] = q[j] * s; + p0 = u * UNIT - k0; + p1 = p0 + UNIT / 2; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint* tab) { + device const uchar* b = row + u * 18; + device const ushort* p = (device const ushort*)(b + 2); +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); + w[2 * j] = float(w0.x); + w[2 * j + 16] = float(w0.y); + w[2 * j + 1] = float(w1.x); + w[2 * j + 17] = float(w1.y); + } + const float d = rdh(b); +#pragma clang loop unroll(full) + for (int j = 0; j < 32; ++j) w[j] *= d; + p0 = 32 * u - k0; + p1 = p0 + 16; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint* tab) { + device const uchar* b = row + (u >> 3) * 136; + const int ui = u & 7; + device const ushort* p = (device const ushort*)(b + 8 + 16 * ui); +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); + w[2 * j] = float(w0.x); + w[2 * j + 16] = float(w0.y); + w[2 * j + 1] = float(w1.x); + w[2 * j + 17] = float(w1.y); + } + const uint sh = *(device const ushort*)(b + 2); + const int ls = int((b[4 + (ui >> 1)] >> (4 * (ui & 1))) & 15) | int(((sh >> (2 * ui)) & 3) << 4); + const float d = rdh(b) * float(ls - 32); +#pragma clang loop unroll(full) + for (int j = 0; j < 32; ++j) w[j] *= d; + p0 = 32 * u - k0; + p1 = p0 + 16; + } +}; +// Q2_0: (code - 1) * d +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 2) * 18; + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + const ushort a = p[0], c = p[1]; + const float d = rdh(b); +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + w[j] = float(int((a >> (2 * j)) & 3) - 1) * d; + w[8 + j] = float(int((c >> (2 * j)) & 3) - 1) * d; + } + p0 = 16 * u - k0; + p1 = p0 + 8; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint* tab) { + device const uchar* b = row + u * 24; + float q[32]; + const float d = deq(b, 0, q, tab), m = rdh(b + 2); +#pragma clang loop unroll(full) + for (int j = 0; j < 32; ++j) w[j] = q[j] * d + m; + p0 = 32 * u - k0; + p1 = p0 + 16; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const float* p = (device const float*)(row + 64 * u); +#pragma clang loop unroll(full) + for (int j = 0; j < 16; ++j) w[j] = p[j]; + p0 = 16 * u - k0; + p1 = p0 + 8; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uint4* p = (device const uint4*)(row + 32 * u); + const uint4 a = p[0], c = p[1]; + const uint v[8] = {a.x, a.y, a.z, a.w, c.x, c.y, c.z, c.w}; +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) { + w[2 * k] = as_type(v[k] << 16); + w[2 * k + 1] = as_type(v[k] & 0xFFFF0000u); + } + p0 = 16 * u - k0; + p1 = p0 + 8; + } +}; +// Q4_K / Q5_K: d1 * code - m1 (ref: d * sc, dmin * m), runs at 256b + 64c + 16m and + 32 +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 144; + const int c = (u >> 1) & 3, m = u & 1; + const float d = rdh(b), dm = rdh(b + 2); + float s1, m1, s2, m2; + scale_min_k4(2 * c, b + 4, s1, m1); + scale_min_k4(2 * c + 1, b + 4, s2, m2); + const float d1 = d * s1, mm1 = dm * m1, d2 = d * s2, mm2 = dm * m2; + device const uchar* q = b + 16 + 32 * c + 16 * m; +#pragma clang loop unroll(full) + for (int j = 0; j < 16; ++j) { + w[j] = d1 * float(q[j] & 15) - mm1; + w[16 + j] = d2 * float(q[j] >> 4) - mm2; + } + p0 = 256 * (u >> 3) + 64 * c + 16 * m - k0; + p1 = p0 + 32; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 176; + const int c = (u >> 1) & 3, m = u & 1; + const float d = rdh(b), dm = rdh(b + 2); + float s1, m1, s2, m2; + scale_min_k4(2 * c, b + 4, s1, m1); + scale_min_k4(2 * c + 1, b + 4, s2, m2); + const float d1 = d * s1, mm1 = dm * m1, d2 = d * s2, mm2 = dm * m2; + device const uchar* ql = b + 48 + 32 * c + 16 * m; + device const uchar* qh = b + 16 + 16 * m; + const uchar u1 = uchar(1 << (2 * c)), u2 = uchar(2 << (2 * c)); +#pragma clang loop unroll(full) + for (int j = 0; j < 16; ++j) { + w[j] = d1 * float((ql[j] & 15) + ((qh[j] & u1) ? 16 : 0)) - mm1; + w[16 + j] = d2 * float((ql[j] >> 4) + ((qh[j] & u2) ? 16 : 0)) - mm2; + } + p0 = 256 * (u >> 3) + 64 * c + 16 * m - k0; + p1 = p0 + 32; + } +}; +// Q6_K: d * sc * (code - 32), runs at 256b + 128n + 32p + 16m and + 64 +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 210; + const int n = (u >> 2) & 1, m = (u >> 1) & 1, p = u & 1; + device const uchar* ql = b + 64 * n + 32 * p + 16 * m; + device const uchar* qh = b + 128 + 32 * n + 16 * m; + device const char* sc = (device const char*)(b + 192); + const float d = rdh(b + 208); + const float ds1 = d * float(sc[8 * n + m + 2 * p]), ds2 = d * float(sc[8 * n + m + 2 * p + 4]); + const uint sh = uint(2 * p), sh4 = sh + 4u; +#pragma clang loop unroll(full) + for (int j = 0; j < 16; j += 4) { + const uint la = rd32(ql + j), ha = rd32(qh + j); + const uchar4 c1 = as_type((la & 0x0F0F0F0Fu) | (((ha >> sh) & 0x03030303u) << 4)); + const uchar4 c2 = as_type(((la >> 4) & 0x0F0F0F0Fu) | (((ha >> sh4) & 0x03030303u) << 4)); + w[j] = ds1 * float(int(c1.x) - 32); + w[j + 1] = ds1 * float(int(c1.y) - 32); + w[j + 2] = ds1 * float(int(c1.z) - 32); + w[j + 3] = ds1 * float(int(c1.w) - 32); + w[16 + j] = ds2 * float(int(c2.x) - 32); + w[16 + j + 1] = ds2 * float(int(c2.y) - 32); + w[16 + j + 2] = ds2 * float(int(c2.z) - 32); + w[16 + j + 3] = ds2 * float(int(c2.w) - 32); + } + p0 = 256 * (u >> 3) + 128 * n + 32 * p + 16 * m - k0; + p1 = p0 + 64; + } +}; +// Q3_K: d_all * (sc - 32) * (code - (hbit ? 0 : 4)), contiguous 32 +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 110; + const int n = (u >> 2) & 1, j = u & 3; + device const uchar* q = b + 32 + 32 * n; + device const uchar* hm = b; + device const uchar* s = b + 96; + const int i1 = 8 * n + 2 * j, i2 = i1 + 1; + const int l1 = (i1 < 8 ? s[i1] & 0xF : s[i1 - 8] >> 4) | (((s[8 + (i1 & 3)] >> (2 * (i1 >> 2))) & 3) << 4); + const int l2 = (i2 < 8 ? s[i2] & 0xF : s[i2 - 8] >> 4) | (((s[8 + (i2 & 3)] >> (2 * (i2 >> 2))) & 3) << 4); + const float d = rdh(b + 108); + const float dl1 = d * float(l1 - 32), dl2 = d * float(l2 - 32); + const uchar hb = uchar(1 << (4 * n + j)); +#pragma clang loop unroll(full) + for (int l = 0; l < 16; ++l) { + w[l] = dl1 * float(int((q[l] >> (2 * j)) & 3) - ((hm[l] & hb) ? 0 : 4)); + w[16 + l] = dl2 * float(int((q[l + 16] >> (2 * j)) & 3) - ((hm[l + 16] & hb) ? 0 : 4)); + } + p0 = 32 * u - k0; + p1 = p0 + 16; + } +}; + +struct MmaArgs { + uint n_in; + uint n_out; + uint row_bytes; + uint x_stride; // floats between the 8 token rows of x + uint y_stride; +}; +constant int MMA_NTOK [[function_constant(4)]]; + +template +kernel void gemv_mma(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]], + device float* y [[buffer(2)]], constant MmaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, UPC = 128 / UNIT, H = UNIT / 2; + threadgroup float wt[SG][8 * 128]; + threadgroup float ct[SG][64]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int row0 = (int(tg.x) * SG + sg) * 8; + simdgroup_float8x8 C = make_filled_simdgroup_matrix(0.0f); + threadgroup float* w8 = wt[sg]; + for (int k0 = 0; k0 < int(a.n_in); k0 += 128) { + for (int it = lane; it < 8 * UPC; it += 32) { + const int r = it / UPC, uu = it - r * UPC; + const int rr = min(row0 + r, int(a.n_out) - 1); + float w[UNIT]; + int p0, p1; + DW::run(W + ulong(rr) * a.row_bytes, k0 / UNIT + uu, k0, w, p0, p1, tab); +#pragma clang loop unroll(full) + for (int j = 0; j < H; ++j) { + w8[r * 128 + p0 + j] = w[j]; + w8[r * 128 + p1 + j] = w[H + j]; + } + } + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int k8 = 0; k8 < 128; k8 += 8) { + simdgroup_float8x8 A, B; + simdgroup_load(A, w8 + k8, 128); + simdgroup_load(B, x + k0 + k8, a.x_stride, ulong2(0, 0), true); + simdgroup_multiply_accumulate(C, A, B, C); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + } + simdgroup_store(C, ct[sg], 8); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int e = lane; e < 64; e += 32) { + const int r = e >> 3, t = e & 7; + if (t < MMA_NTOK && row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = ct[sg][e]; + } +} + +#define MM(F, S) template [[host_name("gemv_mma_" #F "_" #S)]] kernel decltype(gemv_mma) gemv_mma; +#define MM_SET(F) MM(F, 1) MM(F, 2) MM(F, 4) +MM_SET(T_BF16) MM_SET(T_Q8_0) MM_SET(T_Q4_0) MM_SET(T_Q5_0) MM_SET(T_IQ4_NL) MM_SET(T_IQ4_XS) MM_SET(T_Q2_0) +MM_SET(T_IQ3_S) MM_SET(T_Q3_K) MM_SET(T_Q4_K) MM_SET(T_Q5_K) MM_SET(T_Q6_K) diff --git a/results/round39/model-ab/baseline-kernels/moe.metal b/results/round39/model-ab/baseline-kernels/moe.metal new file mode 100644 index 0000000..878252e --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/moe.metal @@ -0,0 +1,765 @@ +// metal/kernels/moe.metal - routed experts of one layer for a verify window of 1..8 tokens, grouped by expert: +// +// moe_plan one threadgroup: the window's top-k ids -> groups (distinct experts, by id) with their entries +// (token, weight), each (token, k)'s entry index, and the indirect grid sizes of the next two kernels. +// Deterministic (no atomics). +// moe_gu per group, a block of FFN rows: gate and up rows read once, applied to every entry of the group +// (the shared input x is read from cache), h = silu(gate) * up per entry. +// moe_down per group, a block of output rows: ye[entry] = weight * (down . h[entry]). +// moe_combine y[token] = sum over k of ye[entry(token, k)], in k order. +// +// The plan is one int buffer (see P_* below). Expert addressing: MoeArgs (a pack blob or the GGUF tensors): +// [gate rows | up rows | down rows] in GGUF blocks (tools/iq_pack.py's experts.bin layout). + +#define P_NG 0 +#define P_IGU 4 +#define P_IDN 8 +#define P_MAXE 128 +#define P_GEXP 16 +#define P_GSTART (P_GEXP + P_MAXE) +#define P_ETOK (P_GSTART + P_MAXE + 4) +#define P_EW (P_ETOK + P_MAXE) +#define P_TOKENT (P_EW + P_MAXE) +#define P_SPLIT (P_TOKENT + P_MAXE) // function constant 26 only: [0] resident groups, then the indirect grids +#define P_IGU_R (P_SPLIT + 4) // ... of the resident groups (the GPU computes them during the host's turn) +#define P_IDN_R (P_SPLIT + 8) +#define P_IGU_M (P_SPLIT + 12) // ... and of the missing ones (after it) +#define P_IDN_M (P_SPLIT + 16) +#define P_GSLOT (P_SPLIT + 20) +#define P_INTS (P_GSLOT + P_MAXE) + +// Expert e's matrices: gate at W + gate_off + e * gu_stride, up at W + up_off + e * gu_stride, down at +// W + down_off + e * d_stride (a pack's blobs: gate_off 0, strides = the blob; a GGUF: the three tensors). +struct MoeArgs { + ulong gate_off; + ulong up_off; + ulong down_off; + ulong gu_stride; + ulong d_stride; + uint gu_row; // bytes per gate/up row + uint d_row; // bytes per down row + uint n_embd; + uint n_ff; + uint n_tok; + uint topk; + uint n_expert; + uint gx_gu; // threadgroups along the rows for moe_gu (n_ff / rows per threadgroup) + uint gx_dn; // ... for moe_down (n_embd / rows per threadgroup) + uint pad; +}; + +// The SSD-streaming arena (metal/src/engine/stream.cpp): with function constant 20 set, expert e of the layer lies in +// an arena slot, smap[e] = (chunk << 28) | its 16 KiB page in that chunk; the chunks are W (buffer 0) and buffers 7..9 +// (gate_off / up_off / down_off are then offsets inside the slot). The same block is in moe_mm.metal. +#ifndef STRATA_ARENA_METAL +#define STRATA_ARENA_METAL +constant bool kArenaFc [[function_constant(20)]]; +constant bool kArena = is_function_constant_defined(kArenaFc) && kArenaFc; +static inline device const uchar* arena_slot(uint v, device const uchar* W, device const uchar* W1, device const uchar* W2, + device const uchar* W3) { + const uint c = v >> 28; + device const uchar* b = c == 0 ? W : (c == 1 ? W1 : (c == 2 ? W2 : W3)); + return b + (ulong)(v & 0x0FFFFFFFu) * 16384ul; +} +#endif + +// Function constant 26 (the SSD-streaming decode window, with smap): the groups whose expert is in the arena when the +// routing is made come first, the missing ones after them, each part with its own indirect grids (P_SPLIT): the +// resident part runs while the host reads the missing experts. The entries keep their (token, k) mapping, so the +// combine (and every number) is the same; without it the plan is the old one, bit for bit. +constant bool kSplitResFc [[function_constant(26)]]; +constant bool kSplitRes = is_function_constant_defined(kSplitResFc) && kSplitResFc; + +// s_grp, s_ent: 64 each (the upper halves for the missing groups) +static inline void plan_body(device const int* ids, device const float* w, device int* plan, constant MoeArgs& a, + ushort tid, ushort nth, ushort sg, ushort lane, threadgroup int* s_grp, + threadgroup int* s_ent, device const uint* smap = nullptr) { + const int ne = int(a.n_tok * a.topk); + // one thread per expert id (nth == n_expert, at most 1024) + const int e = tid; + int cnt = 0; + if (e < int(a.n_expert)) { + #pragma clang loop unroll_count(10) + for (int i = 0; i < ne; ++i) cnt += ids[i] == e; + } + const int has = cnt > 0; + const bool split = kSplitRes && smap != nullptr; + const int miss = (split && has && smap[e] == 0xFFFFFFFFu) ? 1 : 0; + const int h0 = has - miss, c0 = miss ? 0 : cnt, c1 = miss ? cnt : 0; + const int pg = simd_prefix_exclusive_sum(h0), pe = simd_prefix_exclusive_sum(c0); + const int pg1 = split ? simd_prefix_exclusive_sum(miss) : 0; + const int pe1 = split ? simd_prefix_exclusive_sum(c1) : 0; + if (lane == 31) { + s_grp[sg] = pg + h0; s_ent[sg] = pe + c0; + if (split) { s_grp[32 + sg] = pg1 + miss; s_ent[32 + sg] = pe1 + c1; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + const int nsg = (nth + 31) / 32; + const int g0 = lane < nsg ? s_grp[lane] : 0, n0 = lane < nsg ? s_ent[lane] : 0; + const int og0 = simd_prefix_exclusive_sum(g0), oe0 = simd_prefix_exclusive_sum(n0); + if (lane < nsg) { s_grp[lane] = og0; s_ent[lane] = oe0; } + const int ng0 = simd_broadcast(og0 + g0, 31), ne0 = simd_broadcast(oe0 + n0, 31); + int og = ng0, oe = ne0; + if (split) { + const int g1 = lane < nsg ? s_grp[32 + lane] : 0, n1 = lane < nsg ? s_ent[32 + lane] : 0; + const int og1 = ng0 + simd_prefix_exclusive_sum(g1), oe1 = ne0 + simd_prefix_exclusive_sum(n1); + if (lane < nsg) { s_grp[32 + lane] = og1; s_ent[32 + lane] = oe1; } + og = simd_broadcast(og1 + g1, 31); + oe = simd_broadcast(oe1 + n1, 31); + } + if (lane == 0) { + plan[P_NG] = og; + plan[P_IGU] = int(a.gx_gu); plan[P_IGU + 1] = og; plan[P_IGU + 2] = 1; + plan[P_IDN] = int(a.gx_dn); plan[P_IDN + 1] = og; plan[P_IDN + 2] = 1; + plan[P_GSTART + og] = oe; + if (split) { + plan[P_SPLIT] = ng0; + plan[P_IGU_R] = int(a.gx_gu); plan[P_IGU_R + 1] = ng0; plan[P_IGU_R + 2] = 1; + plan[P_IDN_R] = int(a.gx_dn); plan[P_IDN_R + 1] = ng0; plan[P_IDN_R + 2] = 1; + plan[P_IGU_M] = int(a.gx_gu); plan[P_IGU_M + 1] = og - ng0; plan[P_IGU_M + 2] = 1; + plan[P_IDN_M] = int(a.gx_dn); plan[P_IDN_M + 1] = og - ng0; plan[P_IDN_M + 2] = 1; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (has) { + const int gi = (miss ? s_grp[32 + sg] + pg1 : s_grp[sg] + pg), off = miss ? s_ent[32 + sg] + pe1 : s_ent[sg] + pe; + plan[P_GEXP + gi] = e; + plan[P_GSTART + gi] = off; + if (split && !miss) plan[P_GSLOT + gi] = int(smap[e]); + const bool topk10 = (a.topk == 10u); + int j = 0; + #pragma clang loop unroll_count(10) + for (int i = 0; i < ne; ++i) + if (ids[i] == e) { + plan[P_ETOK + off + j] = topk10 ? ((i * 205) >> 11) : (i / int(a.topk)); + plan[P_EW + off + j] = as_type(w[i]); + plan[P_TOKENT + i] = off + j; + ++j; + } + } +} + +kernel void moe_plan(device const int* ids [[buffer(0)]], device const float* w [[buffer(1)]], + device int* plan [[buffer(2)]], constant MoeArgs& a [[buffer(3)]], + ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup int s_grp[64], s_ent[64]; + plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); +} + +// The router's tail for the window, then the plan, in one threadgroup of n_expert (512) threads: simdgroup t routes +// token t exactly as native_router.cu (pinned llama.cpp topk-moe): softmax over the 512 logits, top 10 by +// probability (equal values: the lower expert index), weights = p / max(sum of the selected, 2^-14), the sum taken +// by winning lane in selection order and reduced with the same xor tree. +#ifndef STRATA_MOCE_FC_METAL +#define STRATA_MOCE_FC_METAL +constant int kMoceProtectFc [[function_constant(24)]]; +constant int kMoceProtect = is_function_constant_defined(kMoceProtectFc) ? kMoceProtectFc : 3; +constant uint kMoceScaleR1BitsFc [[function_constant(25)]]; +constant float kMoceScaleR1 = is_function_constant_defined(kMoceScaleR1BitsFc) ? as_type(kMoceScaleR1BitsFc) : 1.0f; +#endif +kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]], + device float* w [[buffer(2)]], device int* plan [[buffer(3)]], constant MoeArgs& a [[buffer(4)]], + device const uint* smap [[buffer(5)]], + ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup int s_grp[64], s_ent[64]; + if (sg < a.n_tok) { + device const float* l = logits + sg * 512; + const bool e512 = (a.n_expert == 512u); + float v[16]; + float mx = -INFINITY; + float sum = 0.0f; + if (e512) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] = l[lane + 32 * i]; + mx = max(mx, v[i]); + } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] = exp(v[i] - mx); + sum += v[i]; + } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + } else { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + v[i] = idx < int(a.n_expert) ? l[idx] : -INFINITY; + mx = max(mx, v[i]); + } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] = (lane + 32 * i < int(a.n_expert)) ? exp(v[i] - mx) : 0.0f; + sum += v[i]; + } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + if (lane + 32 * i >= int(a.n_expert)) { + v[i] = -INFINITY; + } else { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + } + } + const bool draft_moce = (a.pad & 0x80000000u) != 0u && sg > 0; + const float base_moce = (a.pad & 0x7FFFFFFFu) != 0u ? as_type(a.pad & 0x7FFFFFFFu) : 1.0f; + const float moce_scale = draft_moce ? (base_moce * base_moce) : base_moce; + const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f; + const int r_prot = draft_moce ? 1 : kMoceProtect; + const float r1_scale = kMoceScaleR1; + const float rcp_r1 = r1_scale > 1.0f ? (1.0f / r1_scale) : 1.0f; + const float r_step_scale = (r1_scale > 1.0f && moce_scale > 1.0f) ? (moce_scale * rcp_r1) : moce_scale; + const bool has_r1 = (r1_scale > 1.0f && r_prot > 1); + const bool has_moce = (moce_scale > 1.0f || r1_scale > 1.0f) && smap != nullptr; + const int min_scaled_rank = has_r1 ? 1 : r_prot; + ushort res_mask = 0; + if (has_moce) { + if (e512) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) + if (smap[lane + 32 * i] != 0xFFFFFFFFu) res_mask |= ushort(1u << i); + } else { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if (idx < int(a.n_expert) && smap[idx] != 0xFFFFFFFFu) res_mask |= ushort(1u << i); + } + } + } + ushort scaled_mask = 0; + float selected = 0.0f, ssum = 0.0f; + if (!has_moce) { + for (int rank = 0; rank < 10; ++rank) { + float best = v[0]; + int ex = lane; + #pragma clang loop unroll(full) + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) v[i] = (ex == lane + 32 * i) ? -INFINITY : v[i]; + if ((ex & 31) == lane) { + ids[sg * 10 + rank] = ex; + ssum += best; + } + if (rank == lane) selected = best; + } + } else { + for (int rank = 0; rank < 10; ++rank) { + if (has_r1) { + if (rank == 1) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + if ((res_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) { + v[i] *= r1_scale; + scaled_mask |= ushort(1u << i); + } + } + } else if (rank == r_prot && moce_scale > 1.0f) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + if ((scaled_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) { + v[i] *= r_step_scale; + } + } + } + } else if (rank == r_prot && moce_scale > 1.0f) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + if ((res_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) { + v[i] *= moce_scale; + scaled_mask |= ushort(1u << i); + } + } + } + float best = v[0]; + int ex = lane; + #pragma clang loop unroll(full) + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + float true_prob = best; + if (rank >= min_scaled_rank && best > 0.0f) { + const ushort win_lane = ushort(ex & 31); + 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) { + // without a tail bonus, the rank-1 scale remains on unselected residents. + const bool r1_scaled = has_r1 && (rank < r_prot || moce_scale <= 1.0f); + true_prob = best * (r1_scaled ? rcp_r1 : rcp_moce); + } + } + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) v[i] = (ex == lane + 32 * i) ? -INFINITY : v[i]; + if ((ex & 31) == lane) { + ids[sg * 10 + rank] = ex; + ssum += true_prob; + } + if (rank == lane) selected = true_prob; + } + } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + ssum = max(ssum, 6.103515625e-5f); + const float inv = 1.0f / ssum; + if (lane < 10) w[sg * 10 + lane] = selected * inv; + } + threadgroup_barrier(mem_flags::mem_device); + plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent, smap); +} + +// ---- gate + up + SwiGLU ------------------------------------------------------------------------------------------- +// One pass over a group's rows for all its NE entries (1..8: an expert appears once per token); NE is uniform per +// threadgroup, so the kernel switches once into a body specialised for it. +// +// gu_fused: per row, gate and up decoded together and each entry's x slice loaded once for both (the 32-weight units, +// whose x load is two runs and their sums: the M5 Air, 2560 x 640, NE 3-4: Q4_K / Q5_K / Q6_K / Q8_0 -8..-25%, +// IQ4_XS -31..-36% from NE 2). The 16-weight units keep two x loads a row: the two decoded units held at once cost +// them more (IQ2_XS / IQ3_S at NR 1, NE 3: +15..+21%). The same sums in the same order either way. +template struct gu_fused { + enum { value = fmt::UNIT == 32 && (NE >= 3 || F == T_IQ4_XS) }; +}; + +template +static inline float q2_combine(ushort2 pk, float d, thread const float* xs, float xsum0) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; + #pragma clang loop unroll(full) + for (int j = 0; j < 8; j += 2) { + if (F == T_Q2_0R) { + const uint2 q = uint2(pk) >> (2 * j); + const uint2 bits = ((q | (q << 14)) & 0x00030003u) | 0x64006400u; + const float4 v = float4(as_type(bits) - 1025.0h); + s0 = fma(v.x, xs[j], s0); + s1 = fma(v.y, xs[j + 1], s1); + s2 = fma(v.z, xs[8 + j], s2); + s3 = fma(v.w, xs[8 + j + 1], s3); + } else { + s0 = fma(float(pk.x & ushort(3u << (2 * j))), xs[j], s0); + s1 = fma(float(pk.x & ushort(3u << (2 * j + 2))), xs[j + 1], s1); + s2 = fma(float(pk.y & ushort(3u << (2 * j))), xs[8 + j], s2); + s3 = fma(float(pk.y & ushort(3u << (2 * j + 2))), xs[8 + j + 1], s3); + } + } + const float dot = d * (s0 + s1) + d * (s2 + s3); + return F == T_Q2_0R ? (dot - 0.0f) : (dot - (d * xsum0 + 0.0f)); +} + +static inline float iq_combine(uint4 g, uint s16, float sc, thread const float* xs) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 2; ++k) { + const uchar4 v0 = as_type(g[k]), v1 = as_type(g[2 + k]); + const uint m0 = s16 >> (4 * k), m1 = s16 >> (8 + 4 * k); + #pragma clang loop unroll(full) + for (int j = 0; j < 4; j += 2) { + s0 = fma((m0 & (1u << j)) ? -float(v0[j]) : float(v0[j]), xs[4 * k + j], s0); + s1 = fma((m0 & (2u << j)) ? -float(v0[j + 1]) : float(v0[j + 1]), xs[4 * k + j + 1], s1); + s2 = fma((m1 & (1u << j)) ? -float(v1[j]) : float(v1[j]), xs[8 + 4 * k + j], s2); + s3 = fma((m1 & (2u << j)) ? -float(v1[j + 1]) : float(v1[j + 1]), xs[8 + 4 * k + j + 1], s3); + } + } + return sc * (s0 + s1) + sc * (s2 + s3) - 0.0f; +} + +template +static inline float deq_iq_packed(device const uchar* b, int ui, thread uint4& g, thread uint& s16, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + if (F == T_IQ2_S) { + const uint q2 = *(device const ushort*)(b + 2 + 2 * ui); + s16 = *(device const ushort*)(b + 34 + 2 * ui); + const uint qh = uint(b[66 + ib]) >> (4 * h); + const uint i0 = (q2 & 255) | ((qh << 8) & 0x300); + const uint i1 = (q2 >> 8) | ((qh << 6) & 0x300); + const uint2 g0 = ctab ? as_type(iq2s_grid[i0]) : ((threadgroup const uint2*)tab)[i0]; + const uint2 g1 = ctab ? as_type(iq2s_grid[i1]) : ((threadgroup const uint2*)tab)[i1]; + g = uint4(g0.x, g0.y, g1.x, g1.y); + const uint sc = b[74 + ib]; + return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); + } else if (F == T_IQ2_XXS) { + device const uchar* qs = b + 2 + 4 * ib + 2 * ui; + const uint a1 = rd32(b + 2 + 8 * ib + 4), s0 = a1 >> (14 * h); + const uint sa = s0 & 127, sb = (s0 >> 7) & 127; + s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12); + const uint2 g0 = ctab ? as_type(iq2xxs_grid[qs[0]]) : ((threadgroup const uint2*)tab)[qs[0]]; + const uint2 g1 = ctab ? as_type(iq2xxs_grid[qs[1]]) : ((threadgroup const uint2*)tab)[qs[1]]; + g = uint4(g0.x, g0.y, g1.x, g1.y); + return rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f); + } else if (F == T_IQ2_XS) { + device const ushort* qs = (device const ushort*)(b + 2) + 2 * ui; + const uint v0 = qs[0], v1 = qs[1]; + const uint sa = v0 >> 9, sb = v1 >> 9; + s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12); + const uint2 g0 = ctab ? as_type(iq2xs_grid[v0 & 511]) : ((threadgroup const uint2*)tab)[v0 & 511]; + const uint2 g1 = ctab ? as_type(iq2xs_grid[v1 & 511]) : ((threadgroup const uint2*)tab)[v1 & 511]; + g = uint4(g0.x, g0.y, g1.x, g1.y); + const uint sc = b[66 + ib]; + return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); + } else if (F == T_IQ3_S) { + const int l0 = 2 * h, l1 = 2 * h + 1; + const uint qh = b[66 + ib]; + s16 = *(device const ushort*)(b + 74 + 2 * ui); + const uint i0 = uint(b[2 + 8 * ib + 2 * l0]) | ((qh << (8 - 2 * l0)) & 256); + const uint i1 = uint(b[2 + 8 * ib + 2 * l0 + 1]) | ((qh << (7 - 2 * l0)) & 256); + const uint i2 = uint(b[2 + 8 * ib + 2 * l1]) | ((qh << (8 - 2 * l1)) & 256); + const uint i3 = uint(b[2 + 8 * ib + 2 * l1 + 1]) | ((qh << (7 - 2 * l1)) & 256); + g = uint4(g_iq3s(tab, i0), g_iq3s(tab, i1), g_iq3s(tab, i2), g_iq3s(tab, i3)); + const uint sc = b[106 + (ib >> 1)]; + return rdh(b) * fma(float((sc >> (4 * (ib & 1))) & 15), 2.0f, 1.0f); + } else { + const uint aux = rd32(b + 2 + 64 + 4 * ib), s0 = aux >> (14 * h); + const uint sa = s0 & 127, sb = (s0 >> 7) & 127; + s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12); + device const uchar* qs = b + 2 + 4 * ui; + g = uint4(g_iq3xxs(tab, qs[0]), g_iq3xxs(tab, qs[1]), g_iq3xxs(tab, qs[2]), g_iq3xxs(tab, qs[3])); + return rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f); + } +} + +template +static inline void gu_body(device const uchar* g0, device const uchar* u0, device const float* x, + device const float* xs16, device float* h, device const int* etok, int c, int row0, int li, + constant MoeArgs& a, threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT; + float ag[NR][NE], au[NR][NE]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) ag[r][t] = au[r][t] = 0.0f; + const int U = int(a.n_embd) / UNIT; + if (NE == 1) { + const int tok = etok[0]; + x += tok * int(a.n_embd); + xs16 += tok * int(a.n_embd / 16); + } + for (int u = li; u < U; u += LPR) { + if (NE == 1 && UNIT == 16 && F != T_Q2_0 && F != T_Q2_0R) { + float xs[UNIT], xsum[2]; + X::load(x, xs16, u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 c; + Dec::run(g0 + r * a.gu_row, u, q, c, tab); + ag[r][0] += unit_combine(q, c, xs, xsum); + Dec::run(u0 + r * a.gu_row, u, q, c, tab); + au[r][0] += unit_combine(q, c, xs, xsum); + } + } else if (NE == 1) { + float xs[UNIT], xsum[2]; + X::load(x, xs16, u, xs, xsum); + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 c; + Dec::run(g0 + r * a.gu_row, u, q, c, tab); + ag[r][0] += unit_combine(q, c, xs, xsum); + Dec::run(u0 + r * a.gu_row, u, q, c, tab); + au[r][0] += unit_combine(q, c, xs, xsum); + } + } else if (F == T_Q2_0 || F == T_Q2_0R) { + const int boff = (u >> 2) * 18, qoff = 2 + 4 * (u & 3); + ushort2 pkg[NR], pku[NR]; + float dg[NR], du[NR]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + device const uchar* bg = g0 + r * a.gu_row + boff; + device const uchar* bu = u0 + r * a.gu_row + boff; + device const ushort* pg = (device const ushort*)(bg + qoff); + device const ushort* pu = (device const ushort*)(bu + qoff); + pkg[r] = ushort2(pg[0], pg[1]); + pku[r] = ushort2(pu[0], pu[1]); + dg[r] = rdh(bg); + du[r] = rdh(bu); + } + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + ag[r][t] += q2_combine(pkg[r], dg[r], xs, xsum[0]); + au[r][t] += q2_combine(pku[r], du[r], xs, xsum[0]); + } + } + } else if (gu_fused::value) { + for (int r = 0; r < NR; ++r) { + float qg[UNIT], qu[UNIT]; + float4 cg, cu; + Dec::run(g0 + r * a.gu_row, u, qg, cg, tab); + Dec::run(u0 + r * a.gu_row, u, qu, cu, tab); + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + ag[r][t] += unit_combine(qg, cg, xs, xsum); + au[r][t] += unit_combine(qu, cu, xs, xsum); + } + } + } else if (NR == 1 && NE == 4 && (F == T_IQ2_S || F == T_IQ2_XXS || F == T_IQ2_XS || F == T_IQ3_S || F == T_IQ3_XXS)) { + const int boff = (u >> 4) * fmt::BB, ui = (LPR % 16 == 0) ? (li & 15) : (u & 15); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + uint4 gg, gu; + uint sg, su; + const float dg = deq_iq_packed(g0 + r * a.gu_row + boff, ui, gg, sg, tab); + const float du = deq_iq_packed(u0 + r * a.gu_row + boff, ui, gu, su, tab); + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + ag[r][t] += iq_combine(gg, sg, dg, xs); + au[r][t] += iq_combine(gu, su, du, xs); + } + } + } else { + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 c; + Dec::run(g0 + r * a.gu_row, u, q, c, tab); + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + ag[r][t] += unit_combine(q, c, xs, xsum); + } + Dec::run(u0 + r * a.gu_row, u, q, c, tab); + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + au[r][t] += unit_combine(q, c, xs, xsum); + } + } + } + } + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const float gv = lane_sum(ag[r][t]), uv = lane_sum(au[r][t]); + if (li == 0) h[(c + t) * int(a.n_ff) + row0 + r] = gv / (1.0f + exp(-gv)) * uv; + } +} + +// function constant 27, with 26's plan: 1 = only the groups resident at routing time (indirect grid P_IGU_R / P_IDN_R), +// 2 = only the missing ones (P_IGU_M / P_IDN_M); 0 or unset = all (P_IGU / P_IDN) +constant int kPhaseFc [[function_constant(27)]]; +constant int kPhase = is_function_constant_defined(kPhaseFc) ? kPhaseFc : 0; + +// 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; + +template +kernel void moe_gu(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]], + device float* h [[buffer(2)]], device const int* plan [[buffer(3)]], + constant MoeArgs& a [[buffer(4)]], device const float* xs16 [[buffer(5)]], + device const uint* smap [[buffer(6)]], device const uchar* W1 [[buffer(7)]], + device const uchar* W2 [[buffer(8)]], device const uchar* W3 [[buffer(9)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1 && moe_probe != 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int g = int(tg.y) + (kPhase == 2 ? plan[P_SPLIT] : 0); + if (g >= (kPhase == 1 ? plan[P_SPLIT] : plan[P_NG])) return; + const int li = lane % LPR; + const int row0 = ((int(tg.x) * SG + sg) * RG + lane / LPR) * NR; + const ulong ex = (ulong)plan[P_GEXP + g]; + const int c = plan[P_GSTART + g], ne = plan[P_GSTART + g + 1] - c; + device const int* etok = plan + P_ETOK + c; + const uint scode = (kArena && kPhase == 1) ? uint(plan[P_GSLOT + g]) : (kArena ? smap[ex] : 0u); + device const uchar* E = kArena ? arena_slot(scode, W, W1, W2, W3) : W + ex * a.gu_stride; + 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) { + case 1: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + case 2: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + case 3: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + case 4: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + default: + for (int k = 0; k < ne; k += 4) { + const int m = min(4, ne - k); + if (m == 4) gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + else if (m == 3) gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + else if (m == 2) gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + else gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + } + } +} + +// ---- down --------------------------------------------------------------------------------------------------------- +template +static inline void dn_body(device const uchar* d0, device const float* h, device const float* hs16, device float* ye, + device const int* plan, int c, int row0, int li, constant MoeArgs& a, + threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT; + float ac[NR][NE]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + #pragma clang loop unroll(full) + 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) { + if (F == T_Q2_0 || F == T_Q2_0R) { + device const uchar* bu = d0 + (u >> 2) * 18; + const int qoff = 2 + 4 * (u & 3); + ushort2 pk[NR]; + float d[NR]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + device const uchar* b = bu + r * a.d_row; + device const ushort* p = (device const ushort*)(b + qoff); + pk[r] = ushort2(p[0], p[1]); + d[r] = rdh(b); + } + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + float xs[UNIT], xsum[2]; + X::load(h + (c + t) * int(a.n_ff), hs16 + (c + t) * int(a.n_ff / 16), u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + ac[r][t] += q2_combine(pk[r], d[r], xs, xsum[0]); + } + } else if (NE == 1) { + float xs[UNIT], xsum[2]; + X::load(h + c * int(a.n_ff), hs16 + c * int(a.n_ff / 16), u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 sc; + Dec::run(d0 + r * a.d_row, u, q, sc, tab); + ac[r][0] += unit_combine(q, sc, xs, xsum); + } + } else { + float q[NR][UNIT]; + float4 sc[NR]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(d0 + r * a.d_row, u, q[r], sc[r], tab); + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + float xs[UNIT], xsum[2]; + X::load(h + (c + t) * int(a.n_ff), hs16 + (c + t) * int(a.n_ff / 16), u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) ac[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + } + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const float v = lane_sum(ac[r][t]); + if (li == 0) ye[(c + t) * int(a.n_embd) + row0 + r] = as_type(plan[P_EW + c + t]) * v; + } +} + +template +kernel void moe_down(device const uchar* W [[buffer(0)]], device const float* h [[buffer(1)]], + device float* ye [[buffer(2)]], device const int* plan [[buffer(3)]], + constant MoeArgs& a [[buffer(4)]], device const float* hs16 [[buffer(5)]], + device const uint* smap [[buffer(6)]], device const uchar* W1 [[buffer(7)]], + device const uchar* W2 [[buffer(8)]], device const uchar* W3 [[buffer(9)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int g = int(tg.y) + (kPhase == 2 ? plan[P_SPLIT] : 0); + if (g >= (kPhase == 1 ? plan[P_SPLIT] : plan[P_NG])) return; + const int li = lane % LPR; + const int row0 = ((int(tg.x) * SG + sg) * RG + lane / LPR) * NR; + const ulong ex = (ulong)plan[P_GEXP + g]; + const int c = plan[P_GSTART + g], ne = plan[P_GSTART + g + 1] - c; + const uint scode = (kArena && kPhase == 1) ? uint(plan[P_GSLOT + g]) : (kArena ? smap[ex] : 0u); + device const uchar* d0 = (kArena ? arena_slot(scode, W, W1, W2, W3) : W + ex * a.d_stride) + a.down_off + (ulong)row0 * a.d_row; + switch (ne) { + case 1: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + case 2: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + case 3: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + case 4: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + default: + for (int k = 0; k < ne; k += 4) { + const int m = min(4, ne - k); + if (m == 4) dn_body(d0, h, hs16, ye, plan, c + k, row0, li, a, tab); + else if (m == 3) dn_body(d0, h, hs16, ye, plan, c + k, row0, li, a, tab); + else if (m == 2) dn_body(d0, h, hs16, ye, plan, c + k, row0, li, a, tab); + else dn_body(d0, h, hs16, ye, plan, c + k, row0, li, a, tab); + } + } +} + +kernel void moe_combine(device const float* ye [[buffer(0)]], device float* y [[buffer(1)]], + device const int* plan [[buffer(2)]], constant MoeArgs& a [[buffer(3)]], + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.n_embd, k = gid % a.n_embd; + if (t >= a.n_tok) return; + device const int* te = plan + P_TOKENT + t * a.topk; + float s = 0.0f; + if (a.topk == 10) { +#pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += ye[te[j] * int(a.n_embd) + k]; + } else { + for (uint j = 0; j < a.topk; ++j) s += ye[te[j] * int(a.n_embd) + k]; + } + y[gid] = s; +} + +// ---- instantiations: moe_gu____, moe_down____ ------------------------------- +#define GU(F, L, R, S) template [[host_name("moe_gu_" #F "_" #L "_" #R "_" #S)]] kernel decltype(moe_gu) moe_gu; +#define DN(F, L, R, S) template [[host_name("moe_down_" #F "_" #L "_" #R "_" #S)]] kernel decltype(moe_down) moe_down; +#define GU_SET(F) GU(F, 32, 1, 4) GU(F, 32, 2, 4) GU(F, 32, 4, 4) GU(F, 32, 2, 2) GU(F, 32, 2, 8) GU(F, 32, 1, 8) \ + GU(F, 16, 1, 4) GU(F, 16, 2, 4) GU(F, 8, 1, 4) GU(F, 8, 2, 4) GU(F, 8, 1, 2) +#define DN_SET(F) DN(F, 8, 1, 4) DN(F, 8, 2, 4) DN(F, 8, 4, 4) DN(F, 8, 1, 8) DN(F, 8, 2, 8) DN(F, 4, 1, 4) DN(F, 4, 2, 4) \ + DN(F, 4, 4, 4) DN(F, 16, 1, 4) DN(F, 16, 2, 4) DN(F, 32, 1, 4) DN(F, 2, 2, 4) +GU_SET(T_Q2_0) +GU_SET(T_IQ3_XXS) +GU_SET(T_IQ2_XXS) +GU_SET(T_IQ2_XS) +GU_SET(T_IQ2_S) +GU_SET(T_IQ3_S) +GU_SET(T_IQ4_XS) // the IQ3_S pack's gate/up in one layer +GU_SET(T_Q4_K) GU_SET(T_Q5_K) GU_SET(T_Q6_K) GU_SET(T_Q8_0) // the UD-Q4/5/6_K_XL packs +DN_SET(T_Q5_1) DN_SET(T_Q8_0) +DN_SET(T_Q2_0) +DN_SET(T_Q2_0R) +DN_SET(T_IQ4_NL) diff --git a/results/round39/model-ab/baseline-kernels/moe_mm.metal b/results/round39/model-ab/baseline-kernels/moe_mm.metal new file mode 100644 index 0000000..03ac90a --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/moe_mm.metal @@ -0,0 +1,1312 @@ +// metal/kernels/moe_mm.metal - routed experts for a PROMPT chunk (after quant/units/mma.metal): a grouped GEMM on +// the matrix units through Metal 4's MetalPerformancePrimitives matmul2d (on the M5 its neural accelerators). +// +// Plan (deterministic, three passes): the chunk's T x topk (token, expert) entries are sorted by expert, then by +// entry index; each expert's run is padded to a multiple of TM rows; tile i = (expert, first padded row). +// mm_hist per block of 1024 entries: a histogram of experts -> bcnt[block][expert] +// mm_scan one threadgroup (thread = expert): block offsets, the padded expert offsets, the tile list +// mm_place per block: an entry's padded row = expert offset + earlier blocks' count + rank in its block; +// srt_tok[row] = token (-1 on padding), srt_w[row] = weight, pos[entry] = row +// GEMMs: a threadgroup takes a tile (TM rows of one expert) and TN output rows; per K chunk of TK it converts the +// tile's activations to fp16 and dequantizes TN x TK weights (ref_quant's arithmetic, then rounded to fp16) into +// threadgroup memory, and matmul2d multiplies them, fp32 accumulation. +// mm_gu H[row][f] = silu(x . gate_f) * (x . up_f), fp32, for the n_ff rows +// mm_down YE[row][d] = h . down_d, fp32 +// mm_combine y[t][d] = sum_k w(t, k) * YE[pos(t, k)][d], k in order +// +// The idea of feeding matmul2d tiles of dequantized weights from threadgroup memory follows yibie/strata-mlx's +// tile_matmul (MIT); the code is our own. +#include +using namespace mpp::tensor_ops; + +struct MmPlanArgs { + uint n_ent; // T * topk + uint topk; + uint n_expert; + uint tm; // the tile height (rows of one expert) + uint n_blocks; // ceil(n_ent / 1024) +}; + +// counts per block: bcnt[block * n_expert + e] +kernel void mm_hist(device const int* ids [[buffer(0)]], device uint* bcnt [[buffer(1)]], + constant MmPlanArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup atomic_uint h[1024]; + if (tid < a.n_expert) atomic_store_explicit(&h[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint i = tg.x * 1024 + tid; + if (i < a.n_ent && ids[i] >= 0) atomic_fetch_add_explicit(&h[ids[i]], 1u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < a.n_expert) + bcnt[tg.x * a.n_expert + tid] = atomic_load_explicit(&h[tid], memory_order_relaxed); +} + +// one threadgroup of n_expert (<= 1024) threads. boff[block][e] = expert e's rows before this block (within e); +// eoff[e] = expert e's first padded row; tiles[i] = (expert, row0); meta[0] = n_tiles, meta[1] = padded rows; +// indirect grids: meta[2..4] = (gx_gu, n_tiles, 1) for the gate/up GEMM, meta[5..7] = (gx_down, n_tiles, 1) for the +// down GEMM (gx_gu in meta[2], gx_down in meta[5] written by the host once), meta[8..10] = (rows, 1, 1) for mm_gather. +// Padding rows [r0 + tot, r0 + rows) of srt_tok (when bound at buffer 6) are cleared to -1 here so mm_clear is not needed. +kernel void mm_scan(device const uint* bcnt [[buffer(0)]], device uint* boff [[buffer(1)]], + device uint* eoff [[buffer(2)]], device int2* tiles [[buffer(3)]], device uint* meta [[buffer(4)]], + constant MmPlanArgs& a [[buffer(5)]], device int* srt_tok [[buffer(6)]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort nsg [[simdgroups_per_threadgroup]]) { + threadgroup uint s_rows[32], s_tiles[32]; + const uint e = tid; + uint tot = 0; + if (e < a.n_expert) + for (uint b = 0; b < a.n_blocks; ++b) { + boff[b * a.n_expert + e] = tot; + tot += bcnt[b * a.n_expert + e]; + } + const uint nt = (tot + a.tm - 1) / a.tm, rows = nt * a.tm; + const uint pr = simd_prefix_exclusive_sum(rows), pt = simd_prefix_exclusive_sum(nt); + if (lane == 31) { s_rows[sg] = pr + rows; s_tiles[sg] = pt + nt; } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint orow = 0, otile = 0, all_rows = 0, all_tiles = 0; + for (uint i = 0; i < nsg; ++i) { + if (i < sg) { orow += s_rows[i]; otile += s_tiles[i]; } + all_rows += s_rows[i]; + all_tiles += s_tiles[i]; + } + if (e < a.n_expert) { + const uint r0 = orow + pr, t0 = otile + pt; + eoff[e] = r0; + for (uint k = 0; k < nt; ++k) tiles[t0 + k] = int2(int(e), int(r0 + k * a.tm)); + if (srt_tok) + for (uint r = tot; r < rows; ++r) srt_tok[r0 + r] = -1; + } + if (tid == 0) { + meta[0] = all_tiles; + meta[1] = all_rows; + meta[3] = all_tiles; // indirect grid of the gate/up GEMM: (meta[2] given, tiles, 1) + meta[4] = 1; + meta[6] = all_tiles; // of the down GEMM: (meta[5] given, tiles, 1) + meta[7] = 1; + meta[8] = all_rows; // of the gather: (rows, 1, 1) + meta[9] = 1; + meta[10] = 1; + } +} + +static inline uint mm_sum_nibbles(uint c) { + c = (c & 0x0F0F0F0Fu) + ((c >> 4) & 0x0F0F0F0Fu); + c = (c + (c >> 8)) & 0x00FF00FFu; + return (c + (c >> 16)) & 0xFFFFu; +} + +// grid (n_blocks), 1024 threads. Two-level deterministic intra-block rank: 32-lane sub-block nibble counts in p_cnt +// plus simd_broadcast scan within the caller's 32-lane simdgroup (no s_id[1024] threadgroup array). +kernel void mm_place(device const int* ids [[buffer(0)]], device const float* w [[buffer(1)]], + device const uint* boff [[buffer(2)]], device const uint* eoff [[buffer(3)]], + device int* srt_tok [[buffer(4)]], device float* srt_w [[buffer(5)]], device int* pos [[buffer(6)]], + constant MmPlanArgs& a [[buffer(7)]], + uint tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint p_cnt[4096]; + const uint n4e = 4 * a.n_expert; + if (a.n_expert == 512) { + atomic_store_explicit(&p_cnt[tid], 0u, memory_order_relaxed); + atomic_store_explicit(&p_cnt[tid + 1024], 0u, memory_order_relaxed); + } else { + for (uint k = tid; k < n4e; k += 1024) atomic_store_explicit(&p_cnt[k], 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint i = tg * 1024 + tid; + const int e = i < a.n_ent ? ids[i] : -1; + const uint sub = tid >> 5; + if (e >= 0 && sub < 31) { + const uint idx = (sub >> 3) * a.n_expert + uint(e); + const uint sh = (sub & 7u) * 4u; + atomic_fetch_add_explicit(&p_cnt[idx], 1u << sh, memory_order_relaxed); + } + uint rank = 0; + #pragma clang loop unroll(full) + for (ushort l = 0; l < 32; ++l) { + const int el = simd_broadcast(e, l); + if (l < lane && el == e) ++rank; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (i >= a.n_ent) return; + if (e < 0) { pos[i] = -1; return; } + for (uint g = 0; g < (sub >> 3); ++g) { + const uint c = atomic_load_explicit(&p_cnt[g * a.n_expert + uint(e)], memory_order_relaxed); + rank += mm_sum_nibbles(c); + } + if (sub & 7u) { + const uint m = (1u << ((sub & 7u) * 4u)) - 1u; + const uint c = atomic_load_explicit(&p_cnt[(sub >> 3) * a.n_expert + uint(e)], memory_order_relaxed) & m; + rank += mm_sum_nibbles(c); + } + const uint row = eoff[e] + boff[tg * a.n_expert + e] + rank; + srt_tok[row] = int(i / a.topk); + srt_w[row] = w[i]; + pos[i] = int(row); +} + +kernel void mm_clear(device int* srt_tok [[buffer(0)]], constant uint& n [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + if (gid < n) srt_tok[gid] = -1; +} + +#ifndef MM_GATHER_ONLY +// ---- the GEMMs ----------------------------------------------------------------------------------------------------- +// weights of one expert: gate at W + gate_off + e * gu_stride (+ up_off for up), down at W + down_off + e * d_stride +struct MmArgs { + ulong gate_off; + ulong up_off; + ulong down_off; + ulong gu_stride; + ulong d_stride; + uint gu_row; // bytes per gate/up row (n_embd values) + uint d_row; // bytes per down row (n_ff values) + uint n_embd; + uint n_ff; +}; + +// TN x TK weights of `nrow`-strided rows into threadgroup half: thread-strided units of the format +template +static inline void stage_w(device const uchar* w0, uint row_bytes, int k0, threadgroup half* dst, + threadgroup const uint* tab, ushort tid, ushort nth) { + constexpr int UNIT = fmt::UNIT, UPC = TK / UNIT, H = UNIT / 2; + static_assert(TK % UNIT == 0, "the K chunk must hold whole units of the format"); + for (int it = tid; it < TN * UPC; it += nth) { + const int r = it / UPC, uu = it - r * UPC; + float wv[UNIT]; + int p0, p1; + DW::run(w0 + ulong(r) * row_bytes, k0 / UNIT + uu, k0, wv, p0, p1, tab); +#pragma clang loop unroll(full) + for (int j = 0; j < H; ++j) { + dst[r * TK + p0 + j] = half(wv[j]); + dst[r * TK + p1 + j] = half(wv[H + j]); + } + } +} + +#ifndef MM_V3_ONLY +template +kernel void mm_gu(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]], + device float* H [[buffer(2)]], device const int* srt_tok [[buffer(3)]], + device const int2* tiles [[buffer(4)]], constant MmArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int NTH = SG * 32; + threadgroup half xa[TM * TK]; + threadgroup half wg[TN * TK]; + threadgroup half wu[TN * TK]; + threadgroup uint tab[fmt::TAB]; + threadgroup int tok[TM]; + load_tab(tab, tid, NTH); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, f0 = int(tg.x) * TN; + for (int m = tid; m < TM; m += NTH) tok[m] = srt_tok[row0 + m]; + device const uchar* g0 = W + a.gate_off + ulong(e) * a.gu_stride + ulong(f0) * a.gu_row; + device const uchar* u0 = W + a.up_off + ulong(e) * a.gu_stride + ulong(f0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false); + matmul2d> op; + tensor, tensor_inline> tA(xa, dextents(TK, TM)); + tensor, tensor_inline> tG(wg, dextents(TK, TN)); + tensor, tensor_inline> tU(wu, dextents(TK, TN)); + auto accg = op.template get_destination_cooperative_tensor(); + auto accu = op.template get_destination_cooperative_tensor(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + for (int it = tid; it < TM * TK; it += NTH) { + const int m = it / TK, k = it - m * TK; + xa[it] = tok[m] >= 0 ? half(x[ulong(tok[m]) * a.n_embd + k0 + k]) : half(0.0f); + } + stage_w(g0, a.gu_row, k0, wg, tab, tid, NTH); + stage_w(u0, a.gu_row, k0, wu, tab, tid, NTH); + threadgroup_barrier(mem_flags::mem_threadgroup); + auto pg = op.template get_destination_cooperative_tensor(); + auto pu = op.template get_destination_cooperative_tensor(); + op.run(tA, tG, pg); + op.run(tA, tU, pu); + for (uint16_t i = 0; i < pg.get_capacity(); ++i) + if (pg.is_valid_element(i)) { accg[i] += pg[i]; accu[i] += pu[i]; } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = accg[i]; + accg[i] = g / (1.0f + exp(-g)) * accu[i]; + } + tensor, tensor_inline> tH(H, dextents(a.n_ff, 1 << 30)); + accg.store(tH.slice(f0, row0)); +} + +template +kernel void mm_down(device const uchar* W [[buffer(0)]], device const float* H [[buffer(1)]], + device float* YE [[buffer(2)]], device const int2* tiles [[buffer(3)]], + constant MmArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int NTH = SG * 32; + threadgroup half ha[TM * TK]; + threadgroup half wd[TN * TK]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, NTH); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, d0 = int(tg.x) * TN; + device const uchar* w0 = W + a.down_off + ulong(e) * a.d_stride + ulong(d0) * a.d_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false); + matmul2d> op; + tensor, tensor_inline> tA(ha, dextents(TK, TM)); + tensor, tensor_inline> tW(wd, dextents(TK, TN)); + auto acc = op.template get_destination_cooperative_tensor(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) { + for (int it = tid; it < TM * TK; it += NTH) { + const int m = it / TK, k = it - m * TK; + ha[it] = half(H[ulong(row0 + m) * a.n_ff + k0 + k]); + } + stage_w(w0, a.d_row, k0, wd, tab, tid, NTH); + threadgroup_barrier(mem_flags::mem_threadgroup); + auto p = op.template get_destination_cooperative_tensor(); + op.run(tA, tW, p); + for (uint16_t i = 0; i < p.get_capacity(); ++i) + if (p.is_valid_element(i)) acc[i] += p[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + tensor, tensor_inline> tY(YE, dextents(a.n_embd, 1 << 30)); + acc.store(tY.slice(d0, row0)); +} + +// grid (T * n_embd / 256), 256 threads +#endif +kernel void mm_combine(device const float* YE [[buffer(0)]], device const int* pos [[buffer(1)]], + device const float* srt_w [[buffer(2)]], device float* y [[buffer(3)]], + constant uint4& a [[buffer(4)]], // (T, n_embd, topk, -) + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.y, d = gid % a.y; + if (t >= a.x) return; + float s = 0.0f; + for (uint k = 0; k < a.z; ++k) { + const int r = pos[t * a.z + k]; + if (r >= 0) s = fma(srt_w[r], YE[ulong(r) * a.y + d], s); + } + y[gid] = s; +} + +#ifndef MM_V3_ONLY +#define MMGU_(F, M, N, K, S) template [[host_name("mm_gu_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu) mm_gu; +#define MMDN_(F, M, N, K, S) template [[host_name("mm_down_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down) mm_down; +#define MMGU(F, M, N, K, S) MMGU_(F, M, N, K, S) +#define MMDN(F, M, N, K, S) MMDN_(F, M, N, K, S) +#define MM_GU_SET(F) MMGU(F, 32, 64, 64, 4) MMGU(F, 64, 64, 64, 4) MMGU(F, 16, 64, 64, 2) MMGU(F, 32, 32, 64, 2) +#define MM_DN_SET(F) MMDN(F, 32, 64, 64, 4) MMDN(F, 64, 64, 64, 4) MMDN(F, 16, 64, 64, 2) MMDN(F, 32, 32, 64, 2) +MM_GU_SET(T_Q2_0) MM_GU_SET(T_IQ3_XXS) MM_GU_SET(T_IQ2_XXS) MM_GU_SET(T_IQ2_XS) MM_GU_SET(T_IQ2_S) MM_GU_SET(T_IQ3_S) +MM_DN_SET(T_Q2_0) MM_DN_SET(T_IQ4_NL) +#endif + +#endif +// ---- v2: activations gathered once per layer as fp16 rows, read by matmul2d from device memory --------------------- +// mm_gather: XS[row][k] = half(x[srt_tok[row]][k]) (0 on padding rows). grid (rows), 256 threads. +kernel void mm_gather(device const float* x [[buffer(0)]], device const int* srt_tok [[buffer(1)]], + device half* XS [[buffer(2)]], constant uint& n [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const int t = srt_tok[tg.x]; + device half* o = XS + ulong(tg.x) * n; + if (n == 2560) { + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const uint k = tid + i * 256; + o[k] = t >= 0 ? half(x[ulong(t) * n + k]) : half(0.0f); + } + } else { + for (uint k = tid; k < n; k += 256) o[k] = t >= 0 ? half(x[ulong(t) * n + k]) : half(0.0f); + } +} + +// mm_gather_h: gathers straight from fp16 rows (mixh) with 128-bit uint4 copies across nth threads per row. +kernel void mm_gather_h(device const half* xh [[buffer(0)]], device const int* srt_tok [[buffer(1)]], + device half* XS [[buffer(2)]], constant uint& n [[buffer(3)]], + uint tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort nth [[threads_per_threadgroup]]) { + const int t = srt_tok[tg]; + device uint4* o8 = (device uint4*)(XS + ulong(tg) * n); + const uint n8 = n >> 3; + if (n == 2560) { + device const uint4* x8 = (device const uint4*)(xh + ulong(max(t, 0)) * n); + if (nth == 32) { + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) o8[tid + i * 32] = t >= 0 ? x8[tid + i * 32] : uint4(0u); + return; + } + if (nth == 64) { + #pragma clang loop unroll(full) + for (uint i = 0; i < 5; ++i) o8[tid + i * 64] = t >= 0 ? x8[tid + i * 64] : uint4(0u); + return; + } + if (nth == 160) { + o8[tid] = t >= 0 ? x8[tid] : uint4(0u); + o8[tid + 160] = t >= 0 ? x8[tid + 160] : uint4(0u); + return; + } + if (nth == 256) { + o8[tid] = t >= 0 ? x8[tid] : uint4(0u); + if (tid < 64) o8[tid + 256] = t >= 0 ? x8[tid + 256] : uint4(0u); + return; + } + } + if (t >= 0) { + device const uint4* x8 = (device const uint4*)(xh + ulong(t) * n); + for (uint k = tid; k < n8; k += nth) o8[k] = x8[k]; + } else { + for (uint k = tid; k < n8; k += nth) o8[k] = uint4(0u); + } +} + +#ifndef MM_GATHER_ONLY +// The SSD-streaming arena (as in moe.metal): with function constant 20, expert e's slot is smap[e] in the chunks W, +// W1..W3 (buffers 5..8 of mm_gu2 / mm_down2 / mm_gu3 / mm_down3); gate_off / up_off / down_off are offsets in the slot. +#ifndef STRATA_ARENA_METAL +#define STRATA_ARENA_METAL +constant bool kArenaFc [[function_constant(20)]]; +constant bool kArena = is_function_constant_defined(kArenaFc) && kArenaFc; +static inline device const uchar* arena_slot(uint v, device const uchar* W, device const uchar* W1, device const uchar* W2, + device const uchar* W3) { + const uint c = v >> 28; + device const uchar* b = c == 0 ? W : (c == 1 ? W1 : (c == 2 ? W2 : W3)); + return b + (ulong)(v & 0x0FFFFFFFu) * 16384ul; +} +#endif + +// a probe for the bench (function constant 5): 1 = the weights not decoded (a cheap fill), 2 = no matrix multiply +constant int MM_PROBE [[function_constant(5)]]; +constant int mm_probe = is_function_constant_defined(MM_PROBE) ? MM_PROBE : 0; + +#ifndef MM_V3_ONLY +template +kernel void mm_gu2(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], + device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int NTH = SG * 32; + threadgroup half wg[TN * TK]; + threadgroup half wu[TN * TK]; + threadgroup float ot[TM * TN]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, NTH); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, f0 = int(tg.x) * TN; + device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride; + device const uchar* g0 = E + a.gate_off + ulong(f0) * a.gu_row; + device const uchar* u0 = E + a.up_off + ulong(f0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d> op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + tensor, tensor_inline> tG(wg, dextents(TK, TN)); + tensor, tensor_inline> tU(wu, dextents(TK, TN)); + auto tA = tX.template slice(0, row0); + auto accg = op.template get_destination_cooperative_tensor(); + auto accu = op.template get_destination_cooperative_tensor(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + if (mm_probe == 1) { + for (int it = tid; it < TN * TK; it += NTH) { wg[it] = half(k0 & 7); wu[it] = half(it & 7); } + } else { + stage_w(g0, a.gu_row, k0, wg, tab, tid, NTH); + stage_w(u0, a.gu_row, k0, wu, tab, tid, NTH); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + tA = tX.template slice(k0, row0); + if (mm_probe != 2) { + op.run(tA, tG, accg); + op.run(tA, tU, accu); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = accg[i]; + accg[i] = g / (1.0f + exp(-g)) * accu[i]; + } + tensor, tensor_inline> tO(ot, dextents(TN, TM)); + accg.store(tO); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int it = tid; it < TM * TN; it += NTH) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + m) * a.n_ff + f0 + f] = half(ot[it]); + } +} + +template +kernel void mm_down2(device const uchar* W [[buffer(0)]], device const half* HS [[buffer(1)]], + device float* YE [[buffer(2)]], device const int2* tiles [[buffer(3)]], + constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], + device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int NTH = SG * 32; + threadgroup half wd[TN * TK]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, NTH); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, d0 = int(tg.x) * TN; + device const uchar* w0 = (kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.d_stride) + a.down_off + ulong(d0) * a.d_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d> op; + tensor, tensor_inline> tH((device half*)HS, dextents(a.n_ff, 1 << 30)); + tensor, tensor_inline> tW(wd, dextents(TK, TN)); + auto tA = tH.template slice(0, row0); + auto acc = op.template get_destination_cooperative_tensor(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) { + stage_w(w0, a.d_row, k0, wd, tab, tid, NTH); + threadgroup_barrier(mem_flags::mem_threadgroup); + tA = tH.template slice(k0, row0); + op.run(tA, tW, acc); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + tensor, tensor_inline> tY(YE, dextents(a.n_embd, 1 << 30)); + acc.store(tY.slice(d0, row0)); +} + +#define MMGU2_(F, M, N, K, S) template [[host_name("mm_gu2_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu2) mm_gu2; +#define MMDN2_(F, M, N, K, S) template [[host_name("mm_down2_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down2) mm_down2; +#define MMGU2(F, M, N, K, S) MMGU2_(F, M, N, K, S) +#define MMDN2(F, M, N, K, S) MMDN2_(F, M, N, K, S) +#define MM_GU2_SET(F) MMGU2(F, 32, 64, 32, 2) MMGU2(F, 32, 32, 32, 2) MMGU2(F, 32, 64, 16, 2) MMGU2(F, 64, 32, 16, 2) MMGU2(F, 16, 32, 32, 2) MMGU2(F, 16, 64, 16, 2) MMGU2(F, 16, 128, 32, 4) MMGU2(F, 32, 64, 16, 4) MMGU2(F, 16, 32, 32, 1) MMGU2(F, 16, 64, 32, 4) MMGU2(F, 16, 64, 64, 2) MMGU2(F, 32, 64, 64, 4) MMGU2(F, 64, 64, 64, 4) MMGU2(F, 16, 32, 64, 2) MMGU2(F, 32, 64, 32, 4) MMGU2(F, 64, 64, 32, 4) MMGU2(F, 16, 64, 32, 2) +#define MM_DN2_SET(F) MMDN2(F, 32, 64, 32, 2) MMDN2(F, 32, 32, 32, 2) MMDN2(F, 32, 64, 16, 2) MMDN2(F, 64, 32, 16, 2) MMDN2(F, 16, 32, 32, 2) MMDN2(F, 16, 64, 16, 2) MMDN2(F, 16, 128, 32, 4) MMDN2(F, 32, 64, 16, 4) MMDN2(F, 16, 32, 32, 1) MMDN2(F, 16, 64, 32, 4) MMDN2(F, 16, 64, 64, 2) MMDN2(F, 32, 64, 64, 4) MMDN2(F, 64, 64, 64, 4) MMDN2(F, 16, 32, 64, 2) MMDN2(F, 32, 64, 32, 4) MMDN2(F, 64, 64, 32, 4) MMDN2(F, 16, 64, 32, 2) +#define MM_GU2_U32(F) MMGU2(F, 16, 64, 32, 2) MMGU2(F, 16, 64, 64, 2) MMGU2(F, 32, 64, 32, 2) MMGU2(F, 16, 32, 32, 2) \ + MMGU2(F, 64, 64, 32, 4) +#define MM_DN2_U32(F) MMDN2(F, 16, 64, 32, 2) MMDN2(F, 16, 64, 64, 2) MMDN2(F, 16, 32, 32, 2) MMDN2(F, 32, 32, 32, 2) \ + MMDN2(F, 16, 128, 32, 4) MMDN2(F, 32, 64, 32, 2) MMDN2(F, 64, 64, 32, 4) +// K-quants: the staged K chunk must hold whole units as DW places them (128 wide): TK 128 +#define MM_GU2_KQ(F) MMGU2(F, 16, 32, 128, 2) MMGU2(F, 16, 16, 128, 2) MMGU2(F, 32, 32, 128, 2) MMGU2(F, 16, 32, 128, 4) +MM_GU2_KQ(T_Q4_K) MM_GU2_KQ(T_Q5_K) MM_GU2_KQ(T_Q6_K) MM_GU2_U32(T_Q8_0) MM_DN2_U32(T_Q5_1) MM_DN2_U32(T_Q8_0) +// IQ4_XS (UNIT 32): the K chunk must be >= 32 +MMGU2(T_IQ4_XS, 16, 64, 32, 2) MMGU2(T_IQ4_XS, 16, 64, 64, 2) MMGU2(T_IQ4_XS, 32, 64, 32, 2) MMGU2(T_IQ4_XS, 16, 32, 32, 2) +MMGU2(T_IQ4_XS, 64, 64, 32, 4) +MM_GU2_SET(T_Q2_0) MM_GU2_SET(T_IQ3_XXS) MM_GU2_SET(T_IQ2_XXS) MM_GU2_SET(T_IQ2_XS) MM_GU2_SET(T_IQ2_S) MM_GU2_SET(T_IQ3_S) +MM_DN2_SET(T_Q2_0) +MMDN2(T_IQ4_NL, 16, 64, 32, 2) MMDN2(T_IQ4_NL, 16, 64, 64, 2) MMDN2(T_IQ4_NL, 16, 32, 32, 2) MMDN2(T_IQ4_NL, 32, 32, 32, 2) MMDN2(T_IQ4_NL, 16, 128, 32, 4) + +#endif + +// ---- v3: the weights decoded straight into matmul2d's right-input cooperative tensor (single-simdgroup scope), +// no threadgroup staging and no barriers in the K loop. Each simdgroup owns TN output rows; a threadgroup SG * TN. +// The right input's element layout (tests/coop_layout): runs of 4 consecutive k at one n, so a lane decodes 4 values +// at a time (W4). Activations: the gathered fp16 rows from device memory. +#ifndef STRATA_COOP_K4 +#define STRATA_COOP_K4 +// function constant 7: the right-input cooperative tensor holds runs of 4 consecutive k at one n per lane (the M5's +// layout; tests/coop_layout). false (the M4: runs of 2 n at one k): every element decoded on its own. +constant bool COOP_K4 [[function_constant(7)]]; +constant bool coop_k4 = is_function_constant_defined(COOP_K4) ? COOP_K4 : true; +#endif +template struct W4; +// the four weights k .. k + 3 (k % 4 == 0) of a row, in quant.metal's deq arithmetic (code * scale in fp32), as fp16 +// the threadgroup forms of quant.metal's field readers (blocks staged in threadgroup memory) +static inline uint rd32(threadgroup const uchar* p) { + threadgroup const ushort* s = (threadgroup const ushort*)p; + return uint(s[0]) | (uint(s[1]) << 16); +} +static inline float rdh(threadgroup const uchar* p) { return float(*(threadgroup const half*)p); } + +static inline half4 w4_signed(uint g, uint s, float sc) { + const uchar4 v = as_type(g); + const float4 q = float4((s & 1) ? -float(v[0]) : float(v[0]), (s & 2) ? -float(v[1]) : float(v[1]), + (s & 4) ? -float(v[2]) : float(v[2]), (s & 8) ? -float(v[3]) : float(v[3])); + return half4(q * sc); +} + +template <> struct W4 { // 64 values a block: fp16 d, 16 bytes of 2-bit codes (value v at byte v / 4) + // the codes c land in fp16 mantissas (1024 + c), so c - 1 and its product with d are exact in fp16: bit for bit the + // fp32 decode for every d and code byte (checked exhaustively), and 8-11% off mm_down3 at T 1024-2048 + static inline half4 dec(uint q, half d) { + const uint2 c = (uint2(q | (q << 14), (q >> 4) | (q << 10)) & 0x00030003u) | 0x64006400u; + return (as_type(c) - 1025.0h) * d; + } + template + static inline half4 run_q(device const uchar* b, uint kl4) { + return dec((b + kl4)[OFF], *(device const half*)b); + } + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + return run_q<2>(row + (k >> 6) * 18, uint(k & 63) >> 2); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 98; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint aux = rd32(b + 2 + 64 + 4 * ib); + const float sc = rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f); + return w4_signed(g_iq3xxs(tab, b[2 + ((k & 255) >> 2)]), sign4_iq2xs((aux >> (7 * l)) & 127, hh), sc); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 66; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint a0 = rd32(b + 2 + 8 * ib), a1 = rd32(b + 2 + 8 * ib + 4); + const uint idx = (a0 >> (8 * l)) & 255; + const float sc = rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f); + return w4_signed(g_iq2xxs(tab, idx, hh), sign4_iq2xs((a1 >> (7 * l)) & 127, hh), sc); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 74; + const int ib = (k & 255) >> 5, hh = (k & 7) >> 2; + const uint v = ((device const ushort*)(b + 2))[(k & 255) >> 3]; + const uint scb = b[66 + ib]; + const float sc = rdh(b) * fma(float((scb >> ((k & 16) >> 2)) & 15), 0.25f, 0.125f); + return w4_signed(g_iq2xs(tab, v & 511, hh), sign4_iq2xs(v >> 9, hh), sc); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 82; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint qh = b[66 + ib], o = (k & 255) >> 3; + const uint idx = uint(b[2 + o]) | ((qh << (8 - 2 * l)) & 0x300); + const uint s = b[34 + o]; + const uint scb = b[74 + ib]; + const float sc = rdh(b) * fma(float((scb >> ((k & 16) >> 2)) & 15), 0.25f, 0.125f); + return w4_signed(g_iq2s(tab, idx, hh), hh ? s >> 4 : s, sc); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 110; + const int ib = (k & 255) >> 5, hh = (k & 7) >> 2; + const uint qh = b[66 + ib]; + const uint i = uint(b[2 + ((k & 255) >> 2)]) | ((qh << (8 - ((k & 31) >> 2))) & 256u); + const uint s = b[74 + ((k & 255) >> 3)]; + const uint scb = b[106 + (ib >> 1)]; + const float sc = rdh(b) * float(1 + 2 * ((scb >> (4 * (ib & 1))) & 15)); + return w4_signed(g_iq3s(tab, i), hh ? s >> 4 : s, sc); + } +}; +template <> struct W4 { // 256 a block: d, scales_h, scales_l[4], qs[128]; 32-groups as IQ4_NL's + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 136; + const int j = k & 255, ib = j >> 5, o = j & 15, hi = (j >> 4) & 1; + const uint sh = *(device const ushort*)(b + 2); + const int ls = int((b[4 + (ib >> 1)] >> (4 * (ib & 1))) & 15) | int(((sh >> (2 * ib)) & 3) << 4); + const float s = rdh(b) * float(ls - 32); + device const uchar* qs = b + 8 + 16 * ib; + float4 q; + for (int c = 0; c < 4; ++c) { + const half2 w = as_type(tab[qs[o + c]]); + q[c] = float(hi ? w.y : w.x); + } + return half4(q * s); + } +}; +template <> struct W4 { // 32 a block: d, 16 bytes; low nibbles are elements 0..15, high 16..31 + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 5) * 18; + const int j = k & 31, hi = j >> 4, o = j & 15; + const float d = rdh(b); + float4 q; + for (int c = 0; c < 4; ++c) { + const half2 w = as_type(tab[b[2 + o + c]]); + q[c] = float(hi ? w.y : w.x); + } + return half4(q * d); + } +}; + +// W8: the runs k .. k + 3 and k + 16 .. k + 19 of a row (k % 32 < 16): one 32-group, so the per-group fields (scale, +// sign word) are read once. The generic one is two W4. +// ---- the K-quant / legacy-quant W4 decoders (the UD-Q4/5/6_K_XL packs' experts): values k .. k + 3 (k % 4 == 0) +// of a row in ref_quant's arithmetic (float), then fp16 +template <> struct W4 { // 32 a block: d, int8 qs[32] + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 5) * 34; + const float d = rdh(b); + device const char* q = (device const char*)(b + 2 + (k & 31)); + return half4(float4(float(q[0]), float(q[1]), float(q[2]), float(q[3])) * d); + } +}; +template <> struct W4 { // 32 a block: d, m, qh (32 bits), qs[16] (low nibbles 0..15, high 16..31) + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 5) * 24; + const float d = rdh(b), m = rdh(b + 2); + const uint hb = rd32(b + 4); + const int j = k & 31, o = j & 15, hi = j >> 4; + float4 w; + for (int c = 0; c < 4; ++c) { + const uint qv = b[8 + o + c]; + const int x = int(hi ? (qv >> 4) : (qv & 15)) | int(((hb >> (j + c)) & 1) << 4); + w[c] = float(x) * d + m; + } + return half4(w); + } +}; +template <> struct W4 { // 256 a block: d, dmin, scales[12], qs[128] (64-chunks: low nibbles then high) + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 8) * 144; + const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31; + float s, mm; + scale_min_k4(2 * c + hi, b + 4, s, mm); + const float d1 = rdh(b) * s, m1 = rdh(b + 2) * mm; + device const uchar* q = b + 16 + 32 * c + l; + float4 w; + for (int i = 0; i < 4; ++i) w[i] = d1 * float(hi ? (q[i] >> 4) : (q[i] & 15)) - m1; + return half4(w); + } +}; +template <> struct W4 { // 256 a block: d, dmin, scales[12], qh[32], qs[128] + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 8) * 176; + const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31; + float s, mm; + scale_min_k4(2 * c + hi, b + 4, s, mm); + const float d1 = rdh(b) * s, m1 = rdh(b + 2) * mm; + device const uchar* qh = b + 16 + l; + device const uchar* ql = b + 48 + 32 * c + l; + const uint bit = 1u << (2 * c + hi); + float4 w; + for (int i = 0; i < 4; ++i) + w[i] = d1 * float(int(hi ? (ql[i] >> 4) : (ql[i] & 15)) + ((qh[i] & bit) ? 16 : 0)) - m1; + return half4(w); + } +}; +template <> struct W4 { // 256 a block: ql[128], qh[64], int8 scales[16], d; halves of 128 + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 8) * 210; + const uint j = uint(k & 255), j4 = j >> 4; + const uint la = rd32(b + (((j >> 1) & 64u) | (j & 63u))); + const uint ha = rd32(b + 128u + (((j >> 2) & 32u) | (j & 31u))); + const float ds = rdh(b + 208) * float(((device const char*)(b + 192))[j4]); + const uint qa = ((la >> (j4 & 4u)) & 0x0F0F0F0Fu) | (((ha >> (j4 & 6u)) & 0x03030303u) << 4); + return half4(ds * (float4(as_type(qa)) - 32.0f)); + } +}; + +template struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + a = W4::run(row, k, tab); + b = W4::run(row, k + 16, tab); + } +}; +// (a W8 sharing the block scale was measured slower than the two W4 of the generic one: 5.5 vs 3.5-4.2 ms) +template <> struct W8 { + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + const P bl = row + (k >> 8) * 98; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint aux = rd32(bl + 2 + 64 + 4 * ib); + const float sc = rdh(bl) * fma(float(aux >> 28), 0.5f, 0.25f); + const P qs = bl + 2 + ((k & 255) >> 2); + if (mm_probe == 3) { // bench probe: no threadgroup-table lookups (wrong values) + a = w4_signed(qs[0] * 0x01010101u, aux, sc); + b = w4_signed(qs[4] * 0x01010101u, aux >> 7, sc); + return; + } + const uint j0 = (aux >> (7 * l)) & 127, j1 = (aux >> (7 * (l + 2))) & 127; + a = w4_signed(g_iq3xxs(tab, qs[0]), sign4_iq2xs(j0, hh), sc); + b = w4_signed(g_iq3xxs(tab, qs[4]), sign4_iq2xs(j1, hh), sc); + } +}; +template <> struct W8 { + template + static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { + const int ib = (k & 255) >> 5, l = (k & 31) >> 3; + const P p = bl + 2 + 8 * ib; + const uint a1 = rd32(p + 4); + const float sc = d * fma(float(a1 >> 28), 0.25f, 0.125f); + const uint i0 = p[l], i1 = p[l + 2]; + const uint j0 = (a1 >> (7 * l)) & 127, j1 = (a1 >> (7 * (l + 2))) & 127; + threadgroup const uint* th = tab + hh; + a = w4_signed(ctab ? uint(iq2xxs_grid[i0] >> (32 * hh)) : th[2 * i0], sign4_iq2xs(j0, hh), sc); + b = w4_signed(ctab ? uint(iq2xxs_grid[i1] >> (32 * hh)) : th[2 * i1], sign4_iq2xs(j1, hh), sc); + } + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + const P bl = row + (k >> 8) * 66; + block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b); + } +}; +template <> struct W8 { + static inline void block(device const uchar* bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { + const int ib = (k & 255) >> 5; + device const ushort* q = (device const ushort*)(bl + 2) + ((k & 255) >> 3); + const uint v0 = q[0], v1 = q[2]; + const uint scb = bl[66 + ib]; + const float sc0 = d * fma(float(scb & 15), 0.25f, 0.125f); // l < 2 (k % 32 < 16): runs l, l + 2 take nibbles 0, 1 + const float sc1 = d * fma(float(scb >> 4), 0.25f, 0.125f); + threadgroup const uint* th = tab + hh; + a = w4_signed(ctab ? uint(iq2xs_grid[v0 & 511] >> (32 * hh)) : th[2 * (v0 & 511)], sign4_iq2xs(v0 >> 9, hh), sc0); + b = w4_signed(ctab ? uint(iq2xs_grid[v1 & 511] >> (32 * hh)) : th[2 * (v1 & 511)], sign4_iq2xs(v1 >> 9, hh), sc1); + } + static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 74; + block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b); + } +}; +template <> struct W8 { + // grid words read through tab + hh and the sign byte as s >> 4 * hh (not a select): 8% off mm_gu3 at T 2048; + // the same for IQ2_XXS / IQ2_XS, whose sign4_iq2xs select stays (a shift form there measured 10% slower) + template + static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { + const int ib = (k & 255) >> 5, l = (k & 31) >> 3; + threadgroup const uint* th = tab + hh; + const uint sh = uint(hh) << 2; + const uint qh = bl[66 + ib], scb = bl[74 + ib]; + const P p = bl + 2 + ((k & 255) >> 3); + const uint i0 = uint(p[0]) | ((qh << (8 - 2 * l)) & 0x300); + const uint i1 = uint(p[2]) | ((qh << (4 - 2 * l)) & 0x300); + const uint s0 = p[32], s1 = p[34]; + const float sc0 = d * fma(float(scb & 15), 0.25f, 0.125f); + const float sc1 = d * fma(float(scb >> 4), 0.25f, 0.125f); + a = w4_signed(ctab ? uint(iq2s_grid[i0] >> (32 * hh)) : th[2 * i0], s0 >> sh, sc0); + b = w4_signed(ctab ? uint(iq2s_grid[i1] >> (32 * hh)) : th[2 * i1], s1 >> sh, sc1); + } + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + const P bl = row + (k >> 8) * 82; + block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b); + } +}; +template <> struct W8 { + template + static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { + const int ib = k >> 5; + const uint qh = bl[66 + ib], scb = bl[106 + (ib >> 1)]; + const float sc = d * fma(float((scb >> (4 * (ib & 1))) & 15), 2.0f, 1.0f); + const P qs = bl + 2 + (k >> 2); + const int sh0 = 8 - ((k & 31) >> 2); + const uint i0 = uint(qs[0]) | ((qh << sh0) & 256u); + const uint i1 = uint(qs[4]) | ((qh << (sh0 - 4)) & 256u); + const P sp = bl + 74 + (k >> 3); + const uint s0 = sp[0], s1 = sp[2]; + if (mm_probe == 5) { // bench probe: the loads and the scale, no grid lookup / sign math (wrong values) + a = half4(half(float(i0 ^ s0) * sc)); + b = half4(half(float(i1 ^ s1) * sc)); + return; + } + a = w4_signed(g_iq3s(tab, i0), s0 >> (4 * hh), sc); + b = w4_signed(g_iq3s(tab, i1), s1 >> (4 * hh), sc); + } + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + if (mm_probe == 7) { + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint4 r = *(device const uint4*)((ulong(row) + ulong(k >> 5) * 16) & ~ulong(15)); + const uint qhv = r.z & 0xFF, scv = (r.z >> 8) & 0xFF; + const float scr = float(as_type(ushort(r.w & 0xFFFF))) * float(1 + 2 * ((scv >> (4 * (ib & 1))) & 15)); + const uint j0 = ((r.x >> (8 * (l & 1))) & 0xFF) | ((qhv << (8 - 2 * l)) & 256); + const uint j1 = ((r.x >> (16 + 8 * (l & 1))) & 0xFF) | ((qhv << (7 - 2 * l)) & 256); + a = w4_signed(tab[j0], (r.y >> (8 * hh)) & 0xFF, scr); + b = w4_signed(tab[j1], (r.y >> (16 + 8 * hh)) & 0xFF, scr); + return; + } + const P bl = row + (k >> 8) * 110; + block(bl, k & 255, rdh(bl), tab, (k & 7) >> 2, a, b); + } +}; + +// W8 for the K-quants and Q5_1: the lane's runs k and k + 16 (k % 32 < 16) share the block's scales / mins / high +// bits; the 4-byte code runs read as one uint where the block layout keeps them 4-aligned (Q4_K 144-, Q5_K 176-, +// Q5_1 24-byte blocks: rows of whole blocks start 4-aligned). The same per-value arithmetic as W4. +static inline float4 nib4(uint v, int hi) { const uint s = hi ? (v >> 4) : v; return float4(as_type(s & 0x0F0F0F0Fu)); } +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 144; + const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31; + float s, mm; + scale_min_k4(2 * c + hi, bl + 4, s, mm); + const float d1 = rdh(bl) * s, m1 = rdh(bl + 2) * mm; + device const uint* q = (device const uint*)(bl + 16 + 32 * c + l); + const float4 qa = nib4(q[0], hi), qb = nib4(q[4], hi); + a = half4(d1 * qa - m1); + b = half4(d1 * qb - m1); + } +}; +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 176; + const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31; + float s, mm; + scale_min_k4(2 * c + hi, bl + 4, s, mm); + const float d1 = rdh(bl) * s, m1 = rdh(bl + 2) * mm; + device const uint* qh = (device const uint*)(bl + 16 + l); + device const uint* ql = (device const uint*)(bl + 48 + 32 * c + l); + const uint sh = uint(2 * c + hi); + const float4 ha = float4(as_type((qh[0] >> sh) & 0x01010101u)) * 16.0f; + const float4 hb = float4(as_type((qh[4] >> sh) & 0x01010101u)) * 16.0f; + a = half4(d1 * (nib4(ql[0], hi) + ha) - m1); + b = half4(d1 * (nib4(ql[4], hi) + hb) - m1); + } +}; +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 5) * 24; + const float d = rdh(bl), m = rdh(bl + 2); + const uint hb = rd32(bl + 4); + const int j = k & 15; // runs at j and j + 16: the low nibbles of qs[j..j+3], then the high ones + const uint v = *(device const uint*)(bl + 8 + j); + float4 wa, wb; + for (int c = 0; c < 4; ++c) { + const uint qv = (v >> (8 * c)) & 0xFFu; + wa[c] = float(int(qv & 15) | int(((hb >> (j + c)) & 1) << 4)) * d + m; + wb[c] = float(int(qv >> 4) | int(((hb >> (j + 16 + c)) & 1) << 4)) * d + m; + } + a = half4(wa); + b = half4(wb); + } +}; +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 210; + const uint j = uint(k & 255), j4 = j >> 4; // l = j & 31 < 16; the pair at l + 16 + device const uchar* ql = bl + (((j >> 1) & 64u) | (j & 63u)); + device const uchar* qh = bl + 128u + (((j >> 2) & 32u) | (j & 31u)); + const float d = rdh(bl + 208); + device const char* sc = (device const char*)(bl + 192u + j4); + const float dsa = d * float(sc[0]), dsb = d * float(sc[1]); + const uint la = rd32(ql), lb = rd32(ql + 16); + const uint ha = rd32(qh), hb = rd32(qh + 16); + const uint sl = j4 & 4u, shh = j4 & 6u; + const uint qa = ((la >> sl) & 0x0F0F0F0Fu) | (((ha >> shh) & 0x03030303u) << 4); + const uint qb = ((lb >> sl) & 0x0F0F0F0Fu) | (((hb >> shh) & 0x03030303u) << 4); + a = half4(dsa * (float4(as_type(qa)) - 32.0f)); + b = half4(dsb * (float4(as_type(qb)) - 32.0f)); + } +}; + +template +kernel void mm_gu3(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], + device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float ot[SG][TM * TN]; + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN; + device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride; + device const uchar* g0 = E + a.gate_off + ulong(n0) * a.gu_row; + device const uchar* u0 = E + a.up_off + ulong(n0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + auto rw = op.template get_right_input_cooperative_tensor(); + auto tA = tX.template slice(0, row0); + auto accg = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + auto accu = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + if (coop_k4 && TK == 32 && mm_probe != 1) { + const uint16_t half_cap = rw.get_capacity() / 2; + const auto ix0 = rw.get_multidimensional_index(uint16_t(0)); + const int kl = ix0[0], hh = (kl >> 2) & 1; + const ulong roff0 = ulong(ix0[1]) * a.gu_row; + device const uchar* g0_l = g0 + roff0; + device const uchar* u0_l = u0 + roff0; + const uint step8 = 8u * a.gu_row; + // IQ2 formats: a 256 super-block at a time, its d read once; the lane's hh is fixed for the whole K loop + constexpr bool sb = F == T_IQ2_S || F == T_IQ2_XXS || F == T_IQ2_XS; + for (int kb = 0; kb < int(a.n_embd); kb += (sb ? 256 : TK)) { + const uint sboff = uint(kb >> 8) * fmt::BB; + float dg[TN * TK / 256], du[TN * TK / 256]; + if (sb) { +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap / 4; ++i) { + dg[i] = rdh(g0_l + i * step8 + sboff); + du[i] = rdh(u0_l + i * step8 + sboff); + } + } + for (int ki = 0; ki < (sb ? 256 : TK); ki += TK) { + const int k0 = kb + ki; + tA = tX.template slice(k0, row0); + const uint boff = F == T_IQ2_XS ? sboff : (fmt::QK == 256 ? uint(k0 >> 8) * fmt::BB : (fmt::QK == 64 ? uint(k0 >> 6) * fmt::BB : 0u)); + const int j_in = fmt::QK == 256 ? ((k0 + kl) & 255) : (fmt::QK == 64 ? ((k0 + kl) & 63) : (k0 + kl)); + device const uchar* rgw = fmt::QK >= 64 ? g0_l + boff : g0_l; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, rgw += step8) { + half4 wg, wg2; + if (F == T_IQ2_S) W8::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2); + else if (F == T_IQ2_XXS) W8::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2); + else if (F == T_IQ2_XS) W8::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2); + else if (F == T_IQ3_S) W8::block(rgw, j_in, rdh(rgw), tab, hh, wg, wg2); + else W8::run(rgw, j_in, tab, wg, wg2); + rw[i] = wg.x; rw[i + 1] = wg.y; rw[i + 2] = wg.z; rw[i + 3] = wg.w; + const uint16_t j = i + half_cap; + rw[j] = wg2.x; rw[j + 1] = wg2.y; rw[j + 2] = wg2.z; rw[j + 3] = wg2.w; + } + op.run(tA, rw, accg); + device const uchar* ruw = fmt::QK >= 64 ? u0_l + boff : u0_l; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, ruw += step8) { + half4 wu, wu2; + if (F == T_IQ2_S) W8::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2); + else if (F == T_IQ2_XXS) W8::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2); + else if (F == T_IQ2_XS) W8::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2); + else if (F == T_IQ3_S) W8::block(ruw, j_in, rdh(ruw), tab, hh, wu, wu2); + else W8::run(ruw, j_in, tab, wu, wu2); + rw[i] = wu.x; rw[i + 1] = wu.y; rw[i + 2] = wu.z; rw[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rw[j] = wu2.x; rw[j + 1] = wu2.y; rw[j + 2] = wu2.z; rw[j + 3] = wu2.w; + } + op.run(tA, rw, accu); + } + } + } else + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + const uint16_t cap = rw.get_capacity(), half_cap = cap / 2; + tA = tX.template slice(k0, row0); + if (!coop_k4) { + for (uint16_t i = 0; i < cap; ++i) { + const auto ix = rw.get_multidimensional_index(i); + const int k = k0 + ix[0], n = ix[1]; + rw[i] = W4::run(g0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3]; + } + op.run(tA, rw, accg); + for (uint16_t i = 0; i < cap; ++i) { + const auto ix = rw.get_multidimensional_index(i); + const int k = k0 + ix[0], n = ix[1]; + rw[i] = W4::run(u0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3]; + } + op.run(tA, rw, accu); + } else { +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rw.get_multidimensional_index(i); + const auto iy = rw.get_multidimensional_index(i + half_cap); + const int k = k0 + ix[0], n = ix[1]; + device const uchar* rgw = g0 + ulong(n) * a.gu_row; + half4 wg, wg2; + if (mm_probe == 1) { + wg = half4(half(n & 7)); wg2 = half4(half(k & 7)); + } else if (TK == 32) { + W8::run(rgw, k, tab, wg, wg2); + } else { + const int k2 = k0 + iy[0]; + wg = W4::run(rgw, k, tab); + wg2 = W4::run(g0 + ulong(iy[1]) * a.gu_row, k2, tab); + } + rw[i] = wg.x; rw[i + 1] = wg.y; rw[i + 2] = wg.z; rw[i + 3] = wg.w; + const uint16_t j = i + half_cap; + rw[j] = wg2.x; rw[j + 1] = wg2.y; rw[j + 2] = wg2.z; rw[j + 3] = wg2.w; + } + op.run(tA, rw, accg); +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rw.get_multidimensional_index(i); + const auto iy = rw.get_multidimensional_index(i + half_cap); + const int k = k0 + ix[0], n = ix[1]; + device const uchar* ruw = u0 + ulong(n) * a.gu_row; + half4 wu, wu2; + if (mm_probe == 1) { + wu = half4(half(n & 7)); wu2 = half4(half(k & 7)); + } else if (TK == 32) { + W8::run(ruw, k, tab, wu, wu2); + } else { + const int k2 = k0 + iy[0]; + wu = W4::run(ruw, k, tab); + wu2 = W4::run(u0 + ulong(iy[1]) * a.gu_row, k2, tab); + } + rw[i] = wu.x; rw[i + 1] = wu.y; rw[i + 2] = wu.z; rw[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rw[j] = wu2.x; rw[j + 1] = wu2.y; rw[j + 2] = wu2.z; rw[j + 3] = wu2.w; + } + op.run(tA, rw, accu); + } + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = accg[i]; + accg[i] = g / (1.0f + exp(-g)) * accu[i]; + } + tensor, tensor_inline> tO(ot[sg], dextents(TN, TM)); + accg.store(tO); + simdgroup_barrier(mem_flags::mem_threadgroup); + if ((F == T_IQ2_XS || F == T_IQ3_XXS || F == T_IQ2_XXS || F == T_Q2_0 || F == T_IQ2_S || F == T_IQ3_S) && 128 % TN == 0 && (TM * TN) % 128 == 0) { + // a lane stores 4 adjacent outputs a pass, 128 a pass over the simdgroup + constexpr int RP = 128 / TN; // rows a pass + const int mr = lane / (TN / 4), f = (lane % (TN / 4)) * 4; + threadgroup const float4* ot_p = (threadgroup const float4*)(&ot[sg][mr * TN + f]); + device half* hs_p = HS + ulong(row0 + mr) * a.n_ff + n0 + f; + const uint hs_step = uint(RP) * a.n_ff; +#pragma clang loop unroll(full) + for (int m = 0; m < TM / RP; ++m, hs_p += hs_step) + *(device half4*)hs_p = half4(ot_p[m * (RP * TN / 4)]); + } else + for (int it = lane; it < TM * TN; it += 32) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + m) * a.n_ff + n0 + f] = half(ot[sg][it]); + } +} + +template +kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS [[buffer(1)]], + device float* YE [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], + device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN; + device const uchar* w0 = (kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.d_stride) + a.down_off + ulong(n0) * a.d_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tH((device half*)HS, dextents(a.n_ff, 1 << 30)); + auto rw = op.template get_right_input_cooperative_tensor(); + auto tA = tH.template slice(0, row0); + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + const auto ix0 = rw.get_multidimensional_index(uint16_t(0)); + const int kl = ix0[0]; + device const uchar* w0_l = w0 + ulong(ix0[1]) * a.d_row; + const uint step8 = 8u * a.d_row; + if (coop_k4 && F == T_Q2_0 && TK == 32) { + const uint16_t half_cap = rw.get_capacity() / 2; + const uint kl4 = uint(kl) >> 2; + for (int k0 = 0; k0 < int(a.n_ff); k0 += 64, w0_l += 18) { + device const uchar* rdw = w0_l; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) { + const half4 w = W4::run_q<2>(rdw, kl4); + const half4 w2 = W4::run_q<6>(rdw, kl4); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + const uint16_t j = i + half_cap; + rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; + } + tA = tH.template slice(k0, row0); + op.run(tA, rw, acc); + rdw = w0_l; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) { + const half4 w = W4::run_q<10>(rdw, kl4); + const half4 w2 = W4::run_q<14>(rdw, kl4); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + const uint16_t j = i + half_cap; + rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; + } + tA = tH.template slice(k0 + 32, row0); + op.run(tA, rw, acc); + } + } else + for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) { + if (!coop_k4) { + for (uint16_t i = 0; i < rw.get_capacity(); ++i) { + const auto ix = rw.get_multidimensional_index(i); + const int k = k0 + ix[0]; + rw[i] = W4::run(w0 + ulong(ix[1]) * a.d_row, k & ~3, tab)[k & 3]; + } + } else if (TK == 32) { + // the paired runs (element i + cap / 2 is (k + 16, n), the COOP_K4 layout): one W8 a pair + const uint16_t half_cap = rw.get_capacity() / 2; + const uint boff = fmt::QK == 64 ? uint(k0 >> 6) * fmt::BB : (fmt::QK == 32 ? uint(k0 >> 5) * fmt::BB : 0u); + const int j_in = fmt::QK == 64 ? ((k0 + kl) & 63) : (fmt::QK == 32 ? kl : (k0 + kl)); + device const uchar* rdw = w0_l + boff; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) { + half4 w, w2; + W8::run(rdw, j_in, tab, w, w2); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + const uint16_t j = i + half_cap; + rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rw.get_capacity(); i += 4) { + const auto ix = rw.get_multidimensional_index(i); + const half4 w = W4::run(w0 + ulong(ix[1]) * a.d_row, k0 + ix[0], tab); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + } + tA = tH.template slice(k0, row0); + op.run(tA, rw, acc); + } + tensor, tensor_inline> tY(YE, dextents(a.n_embd, 1 << 30)); + acc.store(tY.slice(n0, row0)); +} + +#define MMGU3_(F, M, N, K, S) template [[host_name("mm_gu3_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu3) mm_gu3; +#define MMDN3_(F, M, N, K, S) template [[host_name("mm_down3_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down3) mm_down3; +#define MMGU3(F, M, N, K, S) MMGU3_(F, M, N, K, S) +#define MMDN3(F, M, N, K, S) MMDN3_(F, M, N, K, S) +#define MM3_ONE(F, M, N, K, S) MMGU3(F, M, N, K, S) MMDN3(F, M, N, K, S) +#define MM3_GU(F) MMGU3(F, 16, 64, 32, 2) MMGU3(F, 16, 32, 32, 4) MMGU3(F, 32, 64, 32, 2) MMGU3(F, 16, 64, 64, 2) \ + MMGU3(F, 16, 32, 16, 2) MMGU3(F, 32, 32, 32, 2) MMGU3(F, 16, 64, 32, 1) MMGU3(F, 16, 64, 32, 4) MMGU3(F, 16, 32, 16, 4) MMGU3(F, 16, 32, 32, 5) MMGU3(F, 16, 32, 16, 5) MMGU3(F, 16, 16, 32, 4) MMGU3(F, 16, 16, 32, 8) MMGU3(F, 16, 16, 16, 8) +#define MM3_DN(F) MMDN3(F, 16, 64, 32, 2) MMDN3(F, 16, 32, 32, 4) MMDN3(F, 32, 64, 32, 2) MMDN3(F, 16, 64, 64, 2) \ + MMDN3(F, 16, 32, 16, 2) MMDN3(F, 32, 32, 32, 2) MMDN3(F, 16, 64, 32, 1) MMDN3(F, 16, 64, 32, 4) MMDN3(F, 16, 64, 16, 2) MMDN3(F, 16, 32, 16, 4) MMDN3(F, 16, 64, 32, 5) MMDN3(F, 16, 64, 16, 4) +MM3_GU(T_Q4_K) MM3_GU(T_Q5_K) MM3_GU(T_Q6_K) MM3_GU(T_Q8_0) MM3_GU(T_IQ4_XS) MM3_GU(T_Q2_0) MM3_GU(T_IQ3_XXS) MM3_GU(T_IQ2_XXS) MM3_GU(T_IQ2_XS) MM3_GU(T_IQ2_S) MM3_GU(T_IQ3_S) +MM3_DN(T_Q2_0) MM3_DN(T_IQ4_NL) MM3_DN(T_Q5_1) MM3_DN(T_Q8_0) + +#ifndef MM_V3_ONLY +// mm_gu4: mm_gu3 with each decoded right operand (TN x TK weights of gate and of up) applied to MS 16-row token +// tiles (a tile of 16 * MS rows from the plan, tm = 16 * MS): the decode, which bounds mm_gu3, is shared by MS times +// as many token rows (more padding rows per expert, but the matrix ops are cheap next to the decode). The same +// per-row arithmetic as mm_gu3 (each 16-row op is mm_gu3's): bitwise for the rows both compute. COOP_K4 layout only. +template +kernel void mm_gu4(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int TM = 16; + threadgroup float ot[SG][TM * TN]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN; + device const uchar* g0 = W + a.gate_off + ulong(e) * a.gu_stride + ulong(n0) * a.gu_row; + device const uchar* u0 = W + a.up_off + ulong(e) * a.gu_stride + ulong(n0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + auto rg = op.template get_right_input_cooperative_tensor(); + auto ru = op.template get_right_input_cooperative_tensor(); + auto tA0 = tX.template slice(0, row0); + auto tA1 = tX.template slice(0, row0 + TM); +#define MM4_DST op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>() + auto accg0 = MM4_DST; + auto accu0 = MM4_DST; + auto accg1 = MM4_DST; + auto accu1 = MM4_DST; +#undef MM4_DST + using DT = decltype(accg0); + for (uint16_t i = 0; i < accg0.get_capacity(); ++i) + if (accg0.is_valid_element(i)) { accg0[i] = 0.0f; accu0[i] = 0.0f; accg1[i] = 0.0f; accu1[i] = 0.0f; } + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + const uint16_t cap = rg.get_capacity(), half_cap = cap / 2; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rg.get_multidimensional_index(i); + const int k = k0 + ix[0], n = ix[1]; + device const uchar* rgw = g0 + ulong(n) * a.gu_row; + device const uchar* ruw = u0 + ulong(n) * a.gu_row; + half4 wg, wu, wg2, wu2; + W8::run(rgw, k, tab, wg, wg2); + W8::run(ruw, k, tab, wu, wu2); + rg[i] = wg.x; rg[i + 1] = wg.y; rg[i + 2] = wg.z; rg[i + 3] = wg.w; + ru[i] = wu.x; ru[i + 1] = wu.y; ru[i + 2] = wu.z; ru[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rg[j] = wg2.x; rg[j + 1] = wg2.y; rg[j + 2] = wg2.z; rg[j + 3] = wg2.w; + ru[j] = wu2.x; ru[j + 1] = wu2.y; ru[j + 2] = wu2.z; ru[j + 3] = wu2.w; + } + tA0 = tX.template slice(k0, row0); + op.run(tA0, rg, accg0); + op.run(tA0, ru, accu0); + if (MS > 1) { + tA1 = tX.template slice(k0, row0 + TM); + op.run(tA1, rg, accg1); + op.run(tA1, ru, accu1); + } + } + for (int s = 0; s < MS; ++s) { + thread DT& ag = s == 0 ? accg0 : accg1; + thread DT& au = s == 0 ? accu0 : accu1; + for (uint16_t i = 0; i < ag.get_capacity(); ++i) + if (ag.is_valid_element(i)) { + const float g = ag[i]; + ag[i] = g / (1.0f + exp(-g)) * au[i]; + } + tensor, tensor_inline> tO(ot[sg], dextents(TN, TM)); + simdgroup_barrier(mem_flags::mem_threadgroup); + ag.store(tO); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int it = lane; it < TM * TN; it += 32) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + s * TM + m) * a.n_ff + n0 + f] = half(ot[sg][it]); + } + } +} +#define MMGU4_(F, M, N, K, S) template [[host_name("mm_gu4_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu4) mm_gu4; +#define MMGU4(F, M, N, K, S) MMGU4_(F, M, N, K, S) +#define MM4_GU(F) MMGU4(F, 2, 32, 32, 4) MMGU4(F, 2, 32, 32, 2) MMGU4(F, 2, 16, 32, 4) MMGU4(F, 1, 32, 32, 4) +MM4_GU(T_Q2_0) MM4_GU(T_IQ3_XXS) MM4_GU(T_IQ2_XXS) MM4_GU(T_IQ2_XS) MM4_GU(T_IQ2_S) MM4_GU(T_IQ3_S) MM4_GU(T_IQ4_XS) + +// mm_gu5: mm_gu3 (COOP_K4 layout, TK 32) with the threadgroup's weight rows staged block by block (QK = 256 values a +// row, gate and up) into threadgroup memory by coalesced ushort loads, the lanes decoding from there: the decode's +// many small scattered device loads (2-6 a run) become threadgroup reads. Same decoded values and the same matrix +// ops as mm_gu3: bitwise. Formats with 256-value blocks and a templated W8 (IQ3_S, IQ3_XXS, IQ2_S). +template +kernel void mm_gu5(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], + device const uint* smap [[buffer(5)]], device const uchar* W1 [[buffer(6)]], + device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int TM = 16, TK = 32, BB = fmt::BB, ROWS = SG * TN, RU = BB / 2; + threadgroup float ot[SG][TM * TN]; + threadgroup uint tab[fmt::TAB]; + threadgroup ushort stg[2][ROWS * RU]; + load_tab(tab, tid, SG * 32); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, nt0 = int(tg.x) * ROWS, n0 = nt0 + sg * TN; + device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride; + device const uchar* gt = E + a.gate_off + ulong(nt0) * a.gu_row; + device const uchar* ut = E + a.up_off + ulong(nt0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + auto rg = op.template get_right_input_cooperative_tensor(); + auto ru = op.template get_right_input_cooperative_tensor(); + auto tA = tX.template slice(0, row0); + auto accg = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + auto accu = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + const int NB = int(a.n_embd) / 256; + for (int kb = 0; kb < NB; ++kb) { + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < ROWS * RU; idx += SG * 32) { + const int r = idx / RU, w = idx - r * RU; + stg[0][idx] = ((device const ushort*)(gt + ulong(r) * a.gu_row + kb * BB))[w]; + stg[1][idx] = ((device const ushort*)(ut + ulong(r) * a.gu_row + kb * BB))[w]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int ks = 0; ks < 8; ++ks) { + const int k0 = kb * 256 + ks * TK; + const uint16_t half_cap = rg.get_capacity() / 2; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rg.get_multidimensional_index(i); + const int kl = ks * TK + ix[0], n = sg * TN + ix[1]; + // the block's row in the stage; W8 indexes block (k >> 8) = 0 of it with k = kl + threadgroup const uchar* rgw = (threadgroup const uchar*)(stg[0] + n * RU); + threadgroup const uchar* ruw = (threadgroup const uchar*)(stg[1] + n * RU); + half4 wg, wu, wg2, wu2; + W8::run(rgw, kl, tab, wg, wg2); + W8::run(ruw, kl, tab, wu, wu2); + rg[i] = wg.x; rg[i + 1] = wg.y; rg[i + 2] = wg.z; rg[i + 3] = wg.w; + ru[i] = wu.x; ru[i + 1] = wu.y; ru[i + 2] = wu.z; ru[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rg[j] = wg2.x; rg[j + 1] = wg2.y; rg[j + 2] = wg2.z; rg[j + 3] = wg2.w; + ru[j] = wu2.x; ru[j + 1] = wu2.y; ru[j + 2] = wu2.z; ru[j + 3] = wu2.w; + } + tA = tX.template slice(k0, row0); + op.run(tA, rg, accg); + op.run(tA, ru, accu); + } + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = accg[i]; + accg[i] = g / (1.0f + exp(-g)) * accu[i]; + } + tensor, tensor_inline> tO(ot[sg], dextents(TN, TM)); + accg.store(tO); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int it = lane; it < TM * TN; it += 32) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + m) * a.n_ff + n0 + f] = half(ot[sg][it]); + } +} +#define MMGU5_(F, N, S) template [[host_name("mm_gu5_" #F "_" #N "_" #S)]] kernel decltype(mm_gu5) mm_gu5; +#define MMGU5(F, N, S) MMGU5_(F, N, S) +MMGU5(T_IQ3_S, 32, 2) MMGU5(T_IQ3_S, 32, 1) MMGU5(T_IQ3_S, 16, 4) MMGU5(T_IQ3_XXS, 32, 2) MMGU5(T_IQ3_XXS, 32, 1) +MMGU5(T_IQ3_XXS, 16, 4) MMGU5(T_IQ2_S, 32, 2) MMGU5(T_IQ2_S, 32, 1) MMGU5(T_IQ2_S, 16, 4) MMGU5(T_IQ2_XXS, 32, 2) MMGU5(T_IQ2_XXS, 16, 4) +#endif +#endif diff --git a/results/round39/model-ab/baseline-kernels/prompt.metal b/results/round39/model-ab/baseline-kernels/prompt.metal new file mode 100644 index 0000000..a3fd96b --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/prompt.metal @@ -0,0 +1,690 @@ +// metal/kernels/prompt.metal - the elementwise steps of a prompt chunk around the GEMMs (after quant/units.metal): +// +// route_topk per token: softmax over the 512 logits, top 10 (ties: the lower expert), weights renormalised +// over the selected with a 2^-14 floor - moe_route's rule; one simdgroup a token +// gr_epi_down after the HC down GEMM (V = xn . [w_down; w_inject], T x (hc_lr + hc)): lo_h = half(silu(v / hc)), +// inject = v +// gr_epi_up after the HC up GEMM (G = lo . w_up, T x hc*n): mixed = mean_c xn[c] * sigmoid(G[c]) -> float and +// half (the next GEMMs' operand) +// swiglu_h h = half(silu(g) * u) +// gemm_hc_down, gemm_swiglu_up (after gemm_hc_up) the HC down and shared-expert up GEMMs with gr_epi_down / +// swiglu_h as register epilogues +// moe_out bo = y_routed + sigmoid(sg) * y_shared (prompt_chain.inc; the prompt chunk folds it into the next +// read: gr_norm_h_moe, gr_apply_moe, ple_moe_gate_norm; with mm_combine folded in too: gr_norm_h_moec, +// gr_apply_moec) + +kernel void route_topk(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]], + device float* w [[buffer(2)]], constant uint& n_tok_packed [[buffer(3)]], + device const uint* smap [[buffer(4)]], constant float& prompt_moce [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + const uint n_tok = n_tok_packed & 0xFFFFu; + const int n_exp = (n_tok_packed >> 16) ? int(n_tok_packed >> 16) : 512; + const uint t = tg.x * 8 + sg; + if (t >= n_tok) return; + device const float* l = logits + ulong(t) * ulong(n_exp); + float v[16]; + float mx = -INFINITY; + float sum = 0.0f; + if (n_exp == 512) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] = l[lane + 32 * i]; + mx = max(mx, v[i]); + } + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] = exp(v[i] - mx); + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + } else { + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + v[i] = idx < n_exp ? l[idx] : -INFINITY; + } + for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + for (int i = 0; i < 16; ++i) { + v[i] = (lane + 32 * i < n_exp) ? exp(v[i] - mx) : 0.0f; + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + for (int i = 0; i < 16; ++i) { + if (lane + 32 * i >= n_exp) { + v[i] = -INFINITY; + } else { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + } + } + const int tail_exact = int(prompt_moce) / 10000; + const float rem_moce = prompt_moce - float(tail_exact * 10000); + const int body_prot = int(rem_moce) / 100; + const float body_bonus = rem_moce - float(body_prot * 100); + const bool in_tail = tail_exact > 0 && (t + uint(tail_exact) >= n_tok); + const int p_prot = in_tail ? 10 : body_prot; + const float p_bonus = in_tail ? 0.0f : body_bonus; + const float moce_scale = p_bonus > 0.0f ? exp(min(p_bonus, 40.0f)) : 1.0f; + const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f; + ushort scaled_mask = 0; + float selected = 0.0f, ssum = 0.0f; + for (int rank = 0; rank < 10; ++rank) { + if (rank == p_prot && moce_scale > 1.0f) { + if (n_exp == 512) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if (v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) { + v[i] *= moce_scale; + scaled_mask |= ushort(1u << i); + } + } + } else { + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if (idx < n_exp && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) { + v[i] *= moce_scale; + scaled_mask |= ushort(1u << i); + } + } + } + } + float best = v[0]; + int ex = lane; + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + float true_prob = best; + if (rank >= p_prot && moce_scale > 1.0f && best > 0.0f) { + const ushort win_lane = ushort(ex & 31); + 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 * rcp_moce; + } + if ((ex & 31) == lane) { + // v stays in registers: no runtime index into it + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) v[i] = i == (ex >> 5) ? -INFINITY : v[i]; + ids[t * 10 + rank] = ex; + ssum += true_prob; + } + if (rank == lane) selected = true_prob; + } + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + if (prompt_moce < 0.0f) { + const int skip_tail = int(-prompt_moce); + const float tau = (-prompt_moce) - float(skip_tail); + const bool skip_active = tau > 0.0f && !(skip_tail > 0 && (t + uint(skip_tail) >= n_tok)); + if (skip_active) { + float kept_prob = (lane < 10) ? selected : 0.0f; + if (lane >= 3 && lane < 10) { + const int ex_r = ids[t * 10 + lane]; + if (smap[ex_r] == 0xFFFFFFFFu && selected < tau * ssum) { + ids[t * 10 + lane] = -1; + kept_prob = 0.0f; + } + } + selected = kept_prob; + ssum = kept_prob; + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + } + } + ssum = max(ssum, 6.103515625e-5f); + if (lane < 10) w[t * 10 + lane] = selected * (1.0f / ssum); +} + +struct EpiArgs { + uint n_tok; + uint n_embd; // 2560 + uint hc; // 4 + uint hc_lr; // 320 +}; + +// grid (T * (hc_lr + hc) / 256) +kernel void gr_epi_down(device const float* V [[buffer(0)]], device half* lo_h [[buffer(1)]], + device float* inject [[buffer(2)]], constant EpiArgs& a [[buffer(3)]], + uint gid [[thread_position_in_grid]]) { + const uint nr = a.hc_lr + a.hc, t = gid / nr, r = gid % nr; + if (t >= a.n_tok) return; + const float v = V[gid]; + if (r < a.hc_lr) { + const float z = v / float(a.hc); + lo_h[t * a.hc_lr + r] = half(z / (1.0f + exp(-z))); + } else { + inject[t * a.hc + r - a.hc_lr] = v; + } +} + +// grid (T * n_embd / 256) +kernel void gr_epi_up(device const float* G [[buffer(0)]], device const float* xn [[buffer(1)]], + device float* mixed [[buffer(2)]], device half* mixed_h [[buffer(3)]], + constant EpiArgs& a [[buffer(4)]], uint gid [[thread_position_in_grid]]) { + const uint n = a.n_embd, t = gid / n, d = gid % n; + if (t >= a.n_tok) return; + float m = 0.0f; + for (uint c = 0; c < a.hc; ++c) { + const float g = G[ulong(t) * a.hc * n + c * n + d]; + m = fma(xn[ulong(t) * a.hc * n + c * n + d], 1.0f / (1.0f + exp(-g)), m); + } + m = m / float(a.hc); + mixed[gid] = m; + mixed_h[gid] = half(m); +} + +kernel void swiglu_h(device const float* g [[buffer(0)]], device const float* u [[buffer(1)]], device half* h [[buffer(2)]], + constant uint& n [[buffer(3)]], uint gid [[thread_position_in_grid]]) { + if (gid >= n) return; + const float gv = g[gid]; + h[gid] = half(gv / (1.0f + exp(-gv)) * u[gid]); +} + +// grid (T * n_embd / 256); sg: T values +kernel void moe_out(device const float* yr [[buffer(0)]], device const float* ys [[buffer(1)]], + device const float* sg [[buffer(2)]], device float* bo [[buffer(3)]], constant uint2& a [[buffer(4)]], + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.y; + if (t >= a.x) return; + bo[gid] = fma(1.0f / (1.0f + exp(-sg[t])), ys[gid], yr[gid]); +} + +// ---- the HC read with fewer passes (after dense_mm.metal: GemmArgs, MetalPerformancePrimitives) --------------------- +// gr_norm_h: gr_norm with xn written as fp16 only (the GEMM operand; gemm_hc_up's epilogue reads it too). +// one threadgroup per (stream c, token t), 256 threads; args as gr_norm's (GrArgs, gr.metal) +kernel void gr_norm_h(device float* R [[buffer(0)]], device const float* bo [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device half* xnh [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint c = tg.x, t = tg.y, n = a.n_embd; + device float* r = R + (t * a.hc + c) * n; + const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; + float rv[10], ss = 0.0f; + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + float v = r[d]; + if (a.apply) { + v = fma(bo[t * n + d], g, v); + r[d] = v; + } + rv[i] = v; + ss = fma(v, v, ss); + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; +#pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device half* o = xnh + (t * a.hc + c) * n; + device const float* w = w_norm + c * n; + for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[i] * w[d] * rs); +} + +// gr_norm_h's tail for one token's hc streams: the sums of squares ss[c] (per thread) to the norm, xnh = R' * w * rs +inline void gr_norm_h_out(thread const float (*rv)[10], device const float* w_norm, device half* xnh, uint t, uint hc, + uint n, float eps, thread const float* ss, threadgroup float (*part)[8], ushort tid, + ushort sg, ushort lane) { + for (uint c = 0; c < hc; ++c) { + const float s = simd_sum(ss[c]); + if (lane == 0) part[c][sg] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint c = 0; c < hc; ++c) { + float tot = 0.0f; +#pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[c][i]; + const float rs = rsqrt(tot / float(n) + eps); + device half* o = xnh + (t * hc + c) * n; + device const float* w = w_norm + c * n; + for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[c][i] * w[d] * rs); + } +} + +// stream-wise reduction keeps one residual stream plus bo live, rather than all four residual streams. +inline void gr_norm_h_stream_out(thread const float* rv, device const float* w, device half* o, uint n, float eps, + float ss, threadgroup float* part, ushort tid, ushort sg, ushort lane) { + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; +#pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + eps); + if (n == 2560) { +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const uint d = tid + i * 256; + o[d] = half(rv[i] * w[d] * rs); + } + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[i] * w[d] * rs); + } +} + +// gr_norm_h4: gr_norm_h with one threadgroup per token over all hc streams, so bo is read once a token rather than +// hc times; R', xnh bitwise gr_norm_h's. grid (1, T), 256 threads; hc <= 4 +kernel void gr_norm_h4(device float* R [[buffer(0)]], device const float* bo [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device half* xnh [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + const uint t = tg.y, n = a.n_embd, hc = a.hc; + device float* r = R + t * hc * n; + if (!a.apply) { + for (uint c = 0; c < hc; ++c) { + float rv[10], ss = 0.0f; + if (n == 2560) { + device const float* rc = r + c * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const float v = rc[i * 256]; + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + const float v = r[c * n + d]; + rv[i] = v; + ss = fma(v, v, ss); + } + } + gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); + } + return; + } + float bv[10]; + if (n == 2560) { + device const float* brow = bo + t * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) bv[i] = brow[i * 256]; + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) bv[i] = bo[t * n + d]; + } + for (uint c = 0; c < hc; ++c) { + const float g = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); + float rv[10], ss = 0.0f; + if (n == 2560) { + device float* rc = r + c * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const float v = fma(bv[i], g, rc[i * 256]); + rc[i * 256] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + const float v = fma(bv[i], g, r[c * n + d]); + r[c * n + d] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } + gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); + } +} + +// gr_norm_h_moe: gr_norm_h with the sublayer output bo formed on the fly as moe_out - bo = fma(sigmoid(sg[t]), ys, yr) +// - so the MoE's write never round-trips through bo; R', xnh bitwise moe_out + gr_norm_h. Always applies. +// One threadgroup per token t (all hc streams: bo is formed once a token), 256 threads; hc <= 4. +kernel void gr_norm_h_moe(device float* R [[buffer(0)]], device const float* yr [[buffer(1)]], + device const float* ys [[buffer(2)]], device const float* sgv [[buffer(3)]], + device const float* inj [[buffer(4)]], device const float* w_norm [[buffer(5)]], + device half* xnh [[buffer(6)]], constant GrArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + const uint t = tg.y, n = a.n_embd, hc = a.hc; + device float* r = R + t * hc * n; + float rv[4][10], g[4], ss[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (uint c = 0; c < hc; ++c) g[c] = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + const float bo = fma(gs, ys[t * n + d], yr[t * n + d]); + for (uint c = 0; c < hc; ++c) { + const float v = fma(bo, g[c], r[c * n + d]); + r[c * n + d] = v; + rv[c][i] = v; + ss[c] = fma(v, v, ss[c]); + } + } + gr_norm_h_out(rv, w_norm, xnh, t, hc, n, a.eps, ss, part, tid, sg, lane); +} + +// token t's routed slots for moe_comb: compacted valid YE row offsets (uint(r) * n) and their weights in threadgroup memory +inline uint moe_comb_slots(device const int* pos, device const float* srt_w, uint t, uint topk, uint n, + threadgroup uint* roff, threadgroup float* wk, ushort tid) { + int r = -1; + float w = 0.0f; + if (tid < topk) { + r = pos[t * topk + tid]; + if (r >= 0) w = srt_w[r]; + } + const bool valid = r >= 0; + const ushort prefix = simd_prefix_exclusive_sum(ushort(valid)); + const ushort cnt = simd_sum(ushort(valid)); + if (valid) { + roff[prefix] = uint(r) * n; + wk[prefix] = w; + } + if (tid == 0) roff[16] = cnt; + threadgroup_barrier(mem_flags::mem_threadgroup); + return roff[16]; +} + +// the routed experts' write at (t, d) as mm_combine (moe_mm.metal) forms it: sum_k w(t, k) * YE[pos(t, k)][d], k in order +inline float moe_comb(device const float* YE, threadgroup const uint* roff, threadgroup const float* wk, uint cnt, + uint d) { + float s = 0.0f; + for (uint k = 0; k < cnt; ++k) s = fma(wk[k], YE[roff[k] + d], s); + return s; +} + +// gr_norm_h_moec: gr_norm_h_moe with mm_combine folded in - yr formed from the expert rows YE (pos, srt_w: mm_place's) +// as it is read, so it is never written; R', xnh bitwise mm_combine + gr_norm_h_moe. grid (1, T), 256 threads; +// hc <= 4, topk <= 16 +kernel void gr_norm_h_moec(device float* R [[buffer(0)]], device const float* YE [[buffer(1)]], + device const int* pos [[buffer(2)]], device const float* srt_w [[buffer(3)]], + device const float* ys [[buffer(4)]], device const float* sgv [[buffer(5)]], + device const float* inj [[buffer(6)]], device const float* w_norm [[buffer(7)]], + device half* xnh [[buffer(8)]], constant GrArgs& a [[buffer(9)]], constant uint& topk [[buffer(10)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + threadgroup uint roff[17]; + threadgroup float wk[16]; + const uint t = tg.y, n = a.n_embd, hc = a.hc; + const uint cnt = moe_comb_slots(pos, srt_w, t, topk, n, roff, wk, tid); + device float* r = R + t * hc * n; + float bv[10]; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + if (n == 2560) { + // k outer: each slot's weight and row read once, its ten loads in flight together; s[i]'s fma order is moe_comb's + float s[10] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + for (uint k = 0; k < cnt; ++k) { + const float w = wk[k]; + device const float* yk = YE + roff[k] + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) s[i] = fma(w, yk[i * 256], s[i]); + } + device const float* yrow = ys + t * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) bv[i] = fma(gs, yrow[i * 256], s[i]); + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) bv[i] = fma(gs, ys[t * n + d], moe_comb(YE, roff, wk, cnt, d)); + } + for (uint c = 0; c < hc; ++c) { + const float g = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); + float rv[10], ss = 0.0f; + if (n == 2560) { + device float* rc = r + c * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const float v = fma(bv[i], g, rc[i * 256]); + rc[i * 256] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + const float v = fma(bv[i], g, r[c * n + d]); + r[c * n + d] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } + gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); + } +} + +// gr_apply_moe: the last layer's MoE write applied to R alone (R' of moe_out + gr_norm; its xn has no reader). +// grid (hc * n / 256, T), 256 threads +kernel void gr_apply_moe(device float* R [[buffer(0)]], device const float* yr [[buffer(1)]], + device const float* ys [[buffer(2)]], device const float* sgv [[buffer(3)]], + device const float* inj [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const uint ch = tg.x * 256 + tid, t = tg.y, n = a.n_embd; + if (ch >= a.hc * n) return; + const uint c = ch / n, d = ch - c * n; + const float g = 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))); + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + device float* r = R + t * a.hc * n + ch; + *r = fma(fma(gs, ys[t * n + d], yr[t * n + d]), g, *r); +} + +// gr_apply_moec: gr_apply_moe with mm_combine folded in (as gr_norm_h_moec); R' bitwise mm_combine + gr_apply_moe. +// grid (ceil(n / 256), T), 256 threads: a thread's d over all hc streams (yr formed once); topk <= 16 +kernel void gr_apply_moec(device float* R [[buffer(0)]], device const float* YE [[buffer(1)]], + device const int* pos [[buffer(2)]], device const float* srt_w [[buffer(3)]], + device const float* ys [[buffer(4)]], device const float* sgv [[buffer(5)]], + device const float* inj [[buffer(6)]], constant GrArgs& a [[buffer(7)]], + constant uint& topk [[buffer(8)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint roff[17]; + threadgroup float wk[16]; + const uint d = tg.x * 256 + tid, t = tg.y, n = a.n_embd, hc = a.hc; + const uint cnt = moe_comb_slots(pos, srt_w, t, topk, n, roff, wk, tid); + if (d >= n) return; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + const float bo = fma(gs, ys[t * n + d], moe_comb(YE, roff, wk, cnt, d)); + device float* r = R + t * hc * n + d; + for (uint c = 0; c < hc; ++c) r[c * n] = fma(bo, 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))), r[c * n]); +} + +// gemm_hc_up: G = lo . w_up^T with w_up's rows interleaved by dq_f16 (perm = hc: row d * hc + c), so a TN-column +// tile holds TN / hc outputs d with all their streams; the epilogue mixed[t][d] = mean_c xn[t][c][d] * +// sigmoid(G[t][d * hc + c]) -> fp16 and, if mixed is bound, fp32. G is never written. +// grid (hc * n / TN, ceil(T / TM)), 32 * SG threads. +template +inline void hc_up_tile(device half* A, device half* B, device const half* xnh, device float* mixed, + device half* mixh, constant GemmArgs& g, constant EpiArgs& e, + int t0, int c0, ushort tid, threadgroup float* ot) { + tensor, tensor_inline> tA(A, dextents(g.K, g.M)); + tensor, tensor_inline> tB(B, dextents(g.K, g.N)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, t0); + auto mB = tB.slice(0, c0); + if constexpr (REG) { + auto acc = op.template get_destination_cooperative_tensor(); + op.run(mA, mB, acc); + const int n = int(e.n_embd), d0 = c0 / 4; + // On M5, each lane owns four consecutive output columns: one output's HC streams. + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { + const auto ix = acc.get_multidimensional_index(i); + const int t = t0 + ix[1], d = d0 + ix[0] / 4; + if (t >= int(e.n_tok)) continue; + device const half* xnh_t = xnh + ulong(t) * 4 * n + d; + float s = 0.0f; + #pragma clang loop unroll(full) + for (int c = 0; c < 4; ++c) + s = fma(float(xnh_t[c * n]), 1.0f / (1.0f + exp(-acc[i + c])), s); + s *= 0.25f; + const ulong out_idx = ulong(t) * n + d; + if (mixed) mixed[out_idx] = s; + mixh[out_idx] = half(s); + } + } else { + tensor, tensor_inline> tO(ot, dextents(TN, TM)); + op.run(mA, mB, tO); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int hc = int(e.hc), n = int(e.n_embd), DPT = TN / hc; + const int d0 = c0 / hc; + for (int idx = tid; idx < TM * DPT; idx += SG * 32) { + const int m = idx / DPT, dd = idx - m * DPT, t = t0 + m; + if (t >= int(e.n_tok)) continue; + float s = 0.0f; + for (int c = 0; c < hc; ++c) + s = fma(float(xnh[ulong(t) * hc * n + c * n + d0 + dd]), 1.0f / (1.0f + exp(-ot[m * TN + dd * hc + c])), s); + s = s / float(hc); + if (mixed) mixed[ulong(t) * n + d0 + dd] = s; + mixh[ulong(t) * n + d0 + dd] = half(s); + } + } +} +template +kernel void gemm_hc_up(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device const half* xnh [[buffer(2)]], + device float* mixed [[buffer(3)]], device half* mixh [[buffer(4)]], constant GemmArgs& g [[buffer(5)]], + constant EpiArgs& e [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float ot[REG ? 1 : TM * TN]; + const int t0 = int(tg.y) * TM, c0 = int(tg.x) * TN; + if constexpr (REG && TM > 16) { + if (t0 + 16 >= int(e.n_tok)) { + hc_up_tile<16, TN, SG, REG>(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot); + return; + } + } + if constexpr (REG && TM > 32) { + if (t0 + 32 >= int(e.n_tok)) { + hc_up_tile<32, TN, SG, REG>(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot); + return; + } + } + hc_up_tile(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot); +} +template [[host_name("gemm_hc_up_64_64_4")]] kernel decltype(gemm_hc_up<64, 64, 4>) gemm_hc_up<64, 64, 4>; +template [[host_name("gemm_hc_up_64_32_4")]] kernel decltype(gemm_hc_up<64, 32, 4>) gemm_hc_up<64, 32, 4>; +template [[host_name("gemm_hc_up_32_64_4")]] kernel decltype(gemm_hc_up<32, 64, 4>) gemm_hc_up<32, 64, 4>; +template [[host_name("gemm_hc_up_128_64_4")]] kernel decltype(gemm_hc_up<128, 64, 4>) gemm_hc_up<128, 64, 4>; +template [[host_name("gemm_hc_up_64_128_4")]] kernel decltype(gemm_hc_up<64, 128, 4>) gemm_hc_up<64, 128, 4>; +template [[host_name("gemm_hc_up_reg_64_64_4")]] kernel decltype(gemm_hc_up<64, 64, 4, true>) gemm_hc_up<64, 64, 4, true>; +template [[host_name("gemm_hc_up_reg_128_64_4")]] kernel decltype(gemm_hc_up<128, 64, 4, true>) gemm_hc_up<128, 64, 4, true>; +template [[host_name("gemm_hc_up_reg_64_128_4")]] kernel decltype(gemm_hc_up<64, 128, 4, true>) gemm_hc_up<64, 128, 4, true>; +template [[host_name("gemm_hc_up_reg_64_32_4")]] kernel decltype(gemm_hc_up<64, 32, 4, true>) gemm_hc_up<64, 32, 4, true>; +template [[host_name("gemm_hc_up_reg_32_128_4")]] kernel decltype(gemm_hc_up<32, 128, 4, true>) gemm_hc_up<32, 128, 4, true>; + +// gemm_hc_down: the HC down GEMM V = xn . [w_down; w_inject]^T (T x (hc_lr + hc), K = hc * n_embd) with gr_epi_down +// as its register epilogue: lo_h = half(silu(v / hc)), inject = v. V is never written; lo_h, inject bitwise gemm_f16 + +// gr_epi_down. M5 layout only (coop_k4: a lane's elements 4j .. 4j + 3 are four consecutive columns of one row); +// hc_lr % 4 == 0, hc == 4. grid (ceil((hc_lr + hc) / TN), ceil(T / TM)), 32 * SG threads +template +inline void hc_down_tile(device half* A, device half* B, device half* lo_h, device float* inject, + constant EpiArgs& e, int t0, int c0) { + const int K = int(e.hc * e.n_embd), nr = int(e.hc_lr + e.hc); + tensor, tensor_inline> tA(A, dextents(K, int(e.n_tok))); + tensor, tensor_inline> tB(B, dextents(K, nr)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, t0); + auto mB = tB.slice(0, c0); + auto acc = op.template get_destination_cooperative_tensor(); + op.run(mA, mB, acc); + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { + const auto ix = acc.get_multidimensional_index(i); + const int r = c0 + ix[0], t = t0 + ix[1]; + if (t >= int(e.n_tok) || r >= nr) continue; + if (r < int(e.hc_lr)) { + half4 o; + #pragma clang loop unroll(full) + for (int c = 0; c < 4; ++c) { + const float z = acc[i + c] * 0.25f; // / hc, exact: hc == 4 + o[c] = half(z / (1.0f + exp(-z))); + } + *(device half4*)(lo_h + ulong(t) * e.hc_lr + r) = o; + } else { + *(device float4*)(inject + ulong(t) * e.hc) = float4(acc[i], acc[i + 1], acc[i + 2], acc[i + 3]); + } + } +} +template +inline void hc_down_rows(device half* A, device half* B, device half* lo_h, device float* inject, + constant EpiArgs& e, int t0, int c0) { + // Trimming 32-row parents regresses full tiles on M5. + if constexpr (TM > 32) { + if (t0 + 16 >= int(e.n_tok)) { + hc_down_tile<16, TN, SG>(A, B, lo_h, inject, e, t0, c0); + return; + } + if (t0 + 32 >= int(e.n_tok)) { + hc_down_tile<32, TN, SG>(A, B, lo_h, inject, e, t0, c0); + return; + } + } + hc_down_tile(A, B, lo_h, inject, e, t0, c0); +} +template +kernel void gemm_hc_down(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device half* lo_h [[buffer(2)]], + device float* inject [[buffer(3)]], constant EpiArgs& e [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int c0 = int(tg.x) * TN, t0 = int(tg.y) * TM; + if constexpr (TN == 128) { + // A single 96-column tail beats serial 64+16 on M5, including SG=2. + if (e.hc == 4 && c0 + 64 == int(e.hc_lr)) { + hc_down_rows(A, B, lo_h, inject, e, t0, c0); + return; + } + } + hc_down_rows(A, B, lo_h, inject, e, t0, c0); +} +template [[host_name("gemm_hc_down_16_64_2")]] kernel decltype(gemm_hc_down<16, 64, 2>) gemm_hc_down<16, 64, 2>; +template [[host_name("gemm_hc_down_32_64_2")]] kernel decltype(gemm_hc_down<32, 64, 2>) gemm_hc_down<32, 64, 2>; +template [[host_name("gemm_hc_down_32_128_4")]] kernel decltype(gemm_hc_down<32, 128, 4>) gemm_hc_down<32, 128, 4>; +template [[host_name("gemm_hc_down_64_128_4")]] kernel decltype(gemm_hc_down<64, 128, 4>) gemm_hc_down<64, 128, 4>; +template [[host_name("gemm_hc_down_32_128_2")]] kernel decltype(gemm_hc_down<32, 128, 2>) gemm_hc_down<32, 128, 2>; + +// gemm_swiglu_up: the shared expert's up GEMM u = x . w_up^T (M x N, K) with swiglu_h as its register epilogue against +// the gate GEMM's g: h = half(silu(g) * u). u is never written; h bitwise gemm_f16 + swiglu_h. One accumulator: the +// gate stays a plain gemm_f16 before it. M5 layout only (as gemm_hc_down); N % 4 == 0. +// grid (ceil(N / TN), ceil(M / TM)), 32 * SG threads +template +inline void swiglu_up_tile(device half* A, device half* B, device const float* g, device half* h, + constant GemmArgs& a, int t0, int c0) { + tensor, tensor_inline> tA(A, dextents(a.K, a.M)); + tensor, tensor_inline> tB(B, dextents(a.K, a.N)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, t0); + auto mB = tB.slice(0, c0); + auto acc = op.template get_destination_cooperative_tensor(); + op.run(mA, mB, acc); + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { + const auto ix = acc.get_multidimensional_index(i); + const int c = c0 + ix[0], t = t0 + ix[1]; + if (t >= a.M || c >= a.N) continue; + const ulong idx = ulong(t) * uint(a.N) + uint(c); + const float4 gv = *(device const float4*)(g + idx); + half4 o; + #pragma clang loop unroll(full) + for (int k = 0; k < 4; ++k) o[k] = half(gv[k] / (1.0f + exp(-gv[k])) * acc[i + k]); + *(device half4*)(h + idx) = o; + } +} +template +kernel void gemm_swiglu_up(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device const float* g [[buffer(2)]], + device half* h [[buffer(3)]], constant GemmArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int t0 = int(tg.y) * TM, c0 = int(tg.x) * TN; + if constexpr (TM > 32) { + if (t0 + 16 >= a.M) { + swiglu_up_tile<16, TN, SG>(A, B, g, h, a, t0, c0); + return; + } + if (t0 + 32 >= a.M) { + swiglu_up_tile<32, TN, SG>(A, B, g, h, a, t0, c0); + return; + } + } + swiglu_up_tile(A, B, g, h, a, t0, c0); +} +template [[host_name("gemm_swiglu_up_32_64_2")]] kernel decltype(gemm_swiglu_up<32, 64, 2>) gemm_swiglu_up<32, 64, 2>; +template [[host_name("gemm_swiglu_up_64_128_4")]] kernel decltype(gemm_swiglu_up<64, 128, 4>) gemm_swiglu_up<64, 128, 4>; diff --git a/results/round39/model-ab/baseline-kernels/qsa.metal b/results/round39/model-ab/baseline-kernels/qsa.metal new file mode 100644 index 0000000..63259b4 --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/qsa.metal @@ -0,0 +1,492 @@ +// metal/kernels/qsa.metal - QSA (sparse full attention) decode for a window of tokens, between the projections. +// Contract of include/strata/kernels/qsa.hpp / qsa_select.hpp / qsa_decode_attn.hpp / kv_q8.hpp (the engine's +// native path), in a Metal layout (no paging: unified memory): +// +// qsa_prep per token t (cell pos0 + t): q heads rms_norm * q_norm, NEOX RoPE on 64 dims -> qcur; +// k rms_norm * k_norm, RoPE -> int8 KV append (fp16 scale per 64 = fp16(amax/127), rint, clamp); +// v -> int8 append; indexer q rms_norm * iq_norm, RoPE -> qidx. The indexer KEYS (last +// threadgroup, all window tokens in order): raw tail [cell % 4]; on a block completion +// pooled[b] = rope(rms_norm(mean of the block's 4 raw keys) * ik_norm, pos of the block's first cell); +// cell 0 also sets dead = rms_norm(raw) * ik_norm (the tail block's key). +// qsa_score per token, block b <= n_bid: sum over the 4 indexer heads of relu(qidx . key_b); key_b = pooled[b], +// the tail block n_bid scores `dead` and gets +1e9 when it has cells. +// qsa_topk per token: width = min(n_kv, 2051) cells; n_kv <= 2051: all; else the width cells of largest +// score (a cell inherits its block's), ties to the lowest cell, emitted ascending (radix select). +// qsa_attn split-K: (chunk of CH cells, kv head, token): 12 query heads; s = q.k / 16, chunk max m, sum l, +// acc = p . v -> partials +// qsa_merge per (head, token): log-sum-exp merge, * sigmoid(gate), written as prepared activations (xprep +// planes) for attn_output. +// +// RoPE: cos/sin from a host table [pos][32] (double precision, as build_rope_table). + +#define QH 24 +#define QKV 2 +#define QD 256 +#define QROT 64 +#define IH 4 +#define ID 128 +#define QBUDGET 2051 +#define QCH 64 + +constant int QSA_NTOK [[function_constant(3)]]; +// QSA_MROPE (function constant 24): interleaved M-RoPE (the vision path; strata/kernels/mrope.hpp, ggml IMROPE, +// sections 11/11/10): buffer 20 of qsa_prep is a table int32 [cells][3] (t, h, w), and rotary pair i (0..31) of the row +// at (pos_base + cell) rotates by the cos/sin table's row tab[(pos_base + cell) * 3 + i % 3] instead of its own - +// the pooled indexer keys too (position pos_base + 4 b of a block). Off: exactly the plain RoPE. +constant bool QSA_MROPE [[function_constant(24)]]; +constant bool qsa_mrope = is_function_constant_defined(QSA_MROPE) && QSA_MROPE; +// QSA_Q_F16_FC (function constant 25): qsa_prep writes qcur directly as fp16 (device half*) for qsa_flash, skipping qsa_q_f16. +constant bool QSA_Q_F16_FC [[function_constant(25)]]; +constant bool qsa_qf16 = is_function_constant_defined(QSA_Q_F16_FC) ? QSA_Q_F16_FC : false; +// QSA_ROW0 (function constant 28): qsa_prep's threadgroup x is row QSA_ROW0 + x (24: the K/V rows only, for rows whose +// queries nobody reads). Undefined: 0. +constant int QSA_ROW0 [[function_constant(28)]]; +constant int qsa_row0 = is_function_constant_defined(QSA_ROW0) ? QSA_ROW0 : 0; + +struct QsaArgs { + uint pos0; // the window's first cell (relative to the sequence) + uint pos_base; // absolute position of cell 0 (RoPE) + uint max_cells; + uint n_tok; + float eps; +}; + +// cache layout per layer +// kq/vq [cell][kv head][256] int8, ks/vs [cell][kv head][4] half +// tail [4][128] f32 (slot = cell % 4), dead [128], pooled [max_cells / 4 + 1][128] + +static inline float tg_sum256(float v, threadgroup float* part, ushort sg, ushort lane, ushort nsg) { + v = simd_sum(v); + if (lane == 0) part[sg] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + float t = 0.0f; + for (int i = 0; i < nsg; ++i) t += part[i]; + return t; +} + +// one row (width w <= 256 = threads) : out = rope((rs * x) * g) (RoPE on the first 64 dims, NEOX pairs i, i+32) +static inline float norm_rope(float x, float xhi_unused, device const float* g, threadgroup float* row, + threadgroup float* part, ushort tid, ushort sg, ushort lane, int w, float eps, + device const float2* cs, uint absp, bool rope, device const int* mtab, + bool reuse = false) { + const float ss = tg_sum256(tid < w ? x * x : 0.0f, part, sg, lane, 8); + const float rs = rsqrt(ss / float(w) + eps); + const float y = tid < w ? (rs * x) * g[tid] : 0.0f; + float o = y; + if (rope) { + if (tid < QROT) row[tid] = y; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < QROT) { + const int p = tid & 31; + const uint pp = qsa_mrope ? uint(mtab[absp * 3 + uint(p % 3)]) : absp; + const float2 c = cs[pp * 32 + p]; + const float a = row[p], b = row[p + 32]; + o = tid < 32 ? a * c.x - b * c.y : a * c.y + b * c.x; + } + } + if (reuse) threadgroup_barrier(mem_flags::mem_threadgroup); + (void)xhi_unused; + return o; +} + +// QSA_KV_STRIDE (function constant 26): kcur, vcur, iqraw and ikraw are column slices of one [n_tok][QSA_KV_STRIDE] +// projection (prompt_chunk's fused K | V | indexer K | indexer Q GEMM, 1152 or 1664 columns). Undefined: their own +// [n_tok][rows] buffers. +constant int QSA_KV_STRIDE [[function_constant(26)]]; +constant int qsa_kv_s = is_function_constant_defined(QSA_KV_STRIDE) ? QSA_KV_STRIDE : 0; + +// grid (32 rows + 1, n_tok): rows 0..23 q heads, 24..25 k heads, 26..27 v heads, 28..31 indexer q heads; +// threadgroup 32 (x = 32, y = 0) is the indexer-key threadgroup. 256 threads. A grid of 28 (or of 4 from QSA_ROW0 = +// 24) writes no indexer state: the MTP layer's attention is a dense window, never scored. +kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float* kcur [[buffer(1)]], + device const float* vcur [[buffer(2)]], device const float* iqraw [[buffer(3)]], + device const float* ikraw [[buffer(4)]], device const float* qn [[buffer(5)]], + device const float* kn [[buffer(6)]], device const float* iqn [[buffer(7)]], + device const float* ikn [[buffer(8)]], device const float2* cs [[buffer(9)]], + device float* qcur [[buffer(10)]], device float* qidx [[buffer(11)]], + device char* kq [[buffer(12)]], device char* vq [[buffer(13)]], device half* ksc [[buffer(14)]], + device half* vsc [[buffer(15)]], device float* tail [[buffer(16)]], device float* dead [[buffer(17)]], + device float* pooled [[buffer(18)]], constant QsaArgs& a [[buffer(19)]], + device const int* mtab [[buffer(20)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + threadgroup float row[QROT]; + threadgroup float amax[8]; + const int r = int(tg.x) + qsa_row0; + // the row strides of kcur / vcur, ikraw, iqraw + const int s_kv = qsa_kv_s ? qsa_kv_s : QKV * QD, s_ik = qsa_kv_s ? qsa_kv_s : ID; + const int s_iq = qsa_kv_s ? qsa_kv_s : IH * ID; + if (r == 32) { + // the indexer keys: threadgroup (32, 0) handles t < 3 (boundary with prior tail + cell 0 dead) and final tail; + // every threadgroup (32, t >= 3) with (cell & 3) == 3 pools its 4 in-window keys in parallel. + const int t_self = int(tg.y); + if (t_self >= 3) { + const uint cell = a.pos0 + uint(t_self); + if ((cell & 3) == 3) { + float m = 0.0f; + if (tid < ID) { + m = (((ikraw[(t_self - 3) * s_ik + tid] + ikraw[(t_self - 2) * s_ik + tid]) + + ikraw[(t_self - 1) * s_ik + tid]) + ikraw[t_self * s_ik + tid]) * 0.25f; + } + const uint b = cell >> 2; + const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab); + if (tid < ID) pooled[b * ID + tid] = o; + } + return; + } + if (t_self != 0) return; + const int n_pre = min(int(a.n_tok), 3); + for (int t = 0; t < n_pre; ++t) { + const uint cell = a.pos0 + uint(t); + const float x = tid < ID ? ikraw[t * s_ik + tid] : 0.0f; + if (tid < ID) tail[(cell & 3) * ID + tid] = x; + threadgroup_barrier(mem_flags::mem_device); + if (cell == 0) { + const float o = norm_rope(x, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, 0, false, mtab, true); + if (tid < ID) dead[tid] = o; + } + if ((cell & 3) == 3) { + float m = 0.0f; + if (tid < ID) m = (((tail[tid] + tail[ID + tid]) + tail[2 * ID + tid]) + tail[3 * ID + tid]) * 0.25f; + const uint b = cell >> 2; + const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab, true); + if (tid < ID) pooled[b * ID + tid] = o; + } + threadgroup_barrier(mem_flags::mem_device); + } + const int t_tail0 = max(n_pre, int(a.n_tok) - 4); + for (int t = t_tail0; t < int(a.n_tok); ++t) { + const uint cell = a.pos0 + uint(t); + if (tid < ID) tail[(cell & 3) * ID + tid] = ikraw[t * s_ik + tid]; + } + return; + } + const int t = int(tg.y); + const uint cell = a.pos0 + uint(t), absp = a.pos_base + cell; + if (r < QH) { + const float x = qfull[(t * QH + r) * 2 * QD + tid]; + const float o = norm_rope(x, 0.0f, qn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); + if (qsa_qf16) ((device half*)qcur)[(t * QH + r) * QD + tid] = half(o); + else qcur[(t * QH + r) * QD + tid] = o; + } else if (r < QH + 2 * QKV) { + const bool isv = r >= QH + QKV; + const int h = r - QH - (isv ? QKV : 0); + float x = (isv ? vcur : kcur)[t * s_kv + h * QD + tid]; + if (!isv) x = norm_rope(x, 0.0f, kn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); + // int8 with an fp16 scale per 64 values: groups = simdgroup pairs (64 threads) + float m = simd_max(fabs(x)); + if (lane == 0) amax[sg] = m; + threadgroup_barrier(mem_flags::mem_threadgroup); + const int g = tid >> 6; + m = max(amax[2 * g], amax[2 * g + 1]); + const half sh = half(m / 127.0f); + const float sf = float(sh); + int q = 0; + if (sf > 0.0f) q = clamp(int(rint(x / sf)), -127, 127); + const uint rowi = cell * QKV + uint(h); + (isv ? vq : kq)[rowi * QD + tid] = char(q); + if ((tid & 63) == 0) (isv ? vsc : ksc)[rowi * 4 + g] = sh; + } else { + if (a.pos0 + a.n_tok <= uint(QBUDGET)) return; + const int h = r - QH - 2 * QKV; + const float x = tid < ID ? iqraw[t * s_iq + 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); + if (tid < ID) qidx[(t * IH + h) * ID + tid] = o; + } +} + +// scores [n_tok][max_blocks]; grid (ceil(max_blocks / 8), n_tok), 256 threads: a simdgroup per block. +kernel void qsa_score(device const float* qidx [[buffer(0)]], device const float* pooled [[buffer(1)]], + device const float* dead [[buffer(2)]], device float* scores [[buffer(3)]], + constant QsaArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + const int t = int(tg.y); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + const uint b = tg.x * 8 + sg; + if (b > n_bid) return; + device const float* key = b == n_bid ? dead : pooled + b * ID; + device const float* qp = qidx + (t * IH) * ID + lane; + float k[4], q[IH][4]; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) k[i] = key[lane + 32 * i]; + #pragma clang loop unroll(full) + for (int h = 0; h < IH; ++h) { + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) q[h][i] = qp[h * ID + 32 * i]; + } + float s = 0.0f; + #pragma clang loop unroll(full) + for (int h = 0; h < IH; ++h) { + float d = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) d = fma(q[h][i], k[i], d); + d = simd_sum(d); + s += max(d, 0.0f); + } + if (b == n_bid && (n_kv & 3) != 0) s += 1e9f; + if (lane == 0) scores[uint(t) * (a.max_cells / 4 + 2) + b] = s; +} + +// ids [n_tok][QBUDGET], ascending. One threadgroup of 1024 threads per token. Order key of a float: larger score = +// larger key; a radix select (4 passes of 8 bits) over the blocks with weights = their cell counts finds the +// threshold key T and how many cells with key == T to take (the lowest cells first). +static inline uint okey(float f) { + const uint u = as_type(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]], + constant QsaArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint hist[256]; + threadgroup uint sel[4]; + threadgroup uint2 pref[32]; + threadgroup uint2 totals; + const int t = int(tg.x); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2, nb = n_bid + 1; + device int* out = ids + t * QBUDGET; + device const float* sc = scores + uint(t) * (a.max_cells / 4 + 2); + if (n_kv <= QBUDGET) { + for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); + return; + } + const bool cached = nb <= 4096; + uint keys[4]; + if (cached) + for (uint b = tid, i = 0; b < nb; b += 1024, ++i) keys[i] = okey(sc[b]); + // radix select of the `need`-th largest cell key + uint prefix = 0, pmask = 0, need = QBUDGET; + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint b = tid; b < nb; b += 1024) { + const uint k = cached ? keys[b / 1024] : okey(sc[b]); + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + uint counts[8], sum = 0; + for (uint j = 0; j < 8; ++j) { + counts[j] = atomic_load_explicit(&hist[255 - (lane * 8 + j)], memory_order_relaxed); + sum += counts[j]; + } + uint acc = simd_prefix_exclusive_sum(sum); + if (acc < need && acc + sum >= need) { + for (uint j = 0; j < 8; ++j) { + if (acc + counts[j] >= need) { + sel[0] = 255 - (lane * 8 + j); + sel[1] = need - acc; + break; + } + acc += counts[j]; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; + } + // T = prefix; take every cell with key > T and the first `need` cells (lowest index) with key == T, ascending. + // Pass over blocks in order of index with a prefix scan of each block's selected cell count. + const uint T = prefix; + uint base = 0, eq_taken = 0; + for (uint b0 = 0; b0 < nb; b0 += 1024) { + const uint b = b0 + tid; + uint w = 0, weq = 0; + if (b < nb) { + const uint k = cached ? keys[b / 1024] : okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + if (k > T) w = cells; + else if (k == T) weq = cells; + } + const uint eq_before = simd_prefix_exclusive_sum(weq); + const uint gt_before = simd_prefix_exclusive_sum(w); + if (lane == 31) pref[sg] = uint2(eq_before + weq, gt_before + w); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + const uint2 count = pref[lane]; + const uint2 before = uint2(simd_prefix_exclusive_sum(count.x), + simd_prefix_exclusive_sum(count.y)); + pref[lane] = before; + if (lane == 31) totals = before + count; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint2 group_before = pref[sg], count = totals; + simdgroup_barrier(mem_flags::mem_threadgroup); + const uint eq0 = eq_taken + group_before.x + eq_before; + const uint take_eq = eq0 < need ? min(weq, need - eq0) : 0u; + const uint mine = w + take_eq; + // Clipping the equal-key prefix preserves the lowest-cell tie break. + const uint remaining = need - min(need, eq_taken); + const uint off = base + group_before.y + gt_before + + min(remaining, group_before.x + eq_before); + for (uint c = 0; c < mine; ++c) out[off + c] = int(4 * b + c); + base += count.y + min(remaining, count.x); + eq_taken += count.x; + } +} + +// partials: [n_tok][chunks][QH][QD + 2] (acc, then m, l). grid (chunks, QKV, n_tok), 256 threads. +kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + 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 HPG = QH / QKV; // 12 query heads per kv head + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH]; + threadgroup int cid[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; + const int nchunks = (QBUDGET + QCH - 1) / QCH; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + if (uint(c0) >= width) { + // an empty chunk: m = -inf, l = 0 (the merge skips it) + for (int h = 0; h < HPG; ++h) + if (tid == 0) { pp[(g * HPG + h) * (QD + 2) + QD] = -INFINITY; pp[(g * HPG + h) * (QD + 2) + QD + 1] = 0.0f; } + return; + } + const int nc = min(QCH, int(width) - c0); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + g * HPG + h) * QD + tid]; + if (tid < QCH) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + // scores: 12 x 64 dots of 256; simdgroup sg takes cells sg, sg + 8, ...; lanes split the 256 dims (8 each) + threadgroup const float* q_lane = &qs[0][lane * 8]; + for (int c = sg; c < nc; c += 8) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char* kr = kq + rowi * QD + lane * 8; + const char4 k0 = *(device const char4*)(kr); + const char4 k1 = *(device const char4*)(kr + 4); + const float scl = float(ksc[rowi * 4 + (lane >> 3)]); + const float kv[8] = { + float(k0.x) * scl, float(k0.y) * scl, float(k0.z) * scl, float(k0.w) * scl, + float(k1.x) * scl, float(k1.y) * scl, float(k1.z) * scl, float(k1.w) * scl + }; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + threadgroup const float* qh = q_lane + h * QD; + float d = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) d = fma(qh[i], kv[i], d); + d = simd_sum(d); + if (lane == 0) s[h][c] = d * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // per head: max, exp, sum (simdgroup h % 8 handles heads h, h + 8) + threadgroup float mh[HPG], lh[HPG]; + // nc <= QCH = 64: a lane holds cells lane and lane + 32 in registers (the same max and sum as a strided loop) + const bool in0 = int(lane) < nc, in1 = int(lane) + 32 < nc; + for (int h = sg; h < HPG; h += 8) { + const float s0 = in0 ? s[h][lane] : -INFINITY, s1 = in1 ? s[h][lane + 32] : -INFINITY; + const float m = simd_max(max(s0, s1)); + const float p0 = in0 ? exp(s0 - m) : 0.0f, p1 = in1 ? exp(s1 - m) : 0.0f; + if (in0) s[h][lane] = p0; + if (in1) s[h][lane + 32] = p1; + const float l = simd_sum(p0 + p1); + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // acc[h][j] = sum_c p[h][c] v[c][j]: thread = dim j + float acc[HPG]; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + // four cells' V loads in flight at once; the fma order per head stays c, c + 1, ... + auto vcell = [&](int c) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + return float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); + }; + int c = 0; + for (; c + 4 <= nc; c += 4) { + const float v0 = vcell(c), v1 = vcell(c + 1), v2 = vcell(c + 2), v3 = vcell(c + 3); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + float x = fma(s[h][c], v0, acc[h]); + x = fma(s[h][c + 1], v1, x); + x = fma(s[h][c + 2], v2, x); + acc[h] = fma(s[h][c + 3], v3, x); + } + } + for (; c < nc; ++c) { + const float v = vcell(c); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][c], v, acc[h]); + } + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (g * HPG + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} + +constant float qsa_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; +constant float qsa_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + +static inline void qsa_write_prepared(float y, device float* xp, constant QsaArgs& a, + int t, int h, ushort tid) { + // prepared planes (xprep layout) of the n = QH * QD vector + const uint n = QH * QD, plane = a.n_tok * n, d = uint(h) * QD + tid; + const int j = tid & 15; + device float* raw = xp + uint(t) * n + d; + raw[0] = y; + raw[plane] = y * qsa_q2s[j & 7]; + raw[2 * plane] = y * qsa_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? y * 0.00390625f : y; + float s0 = 0.0f, s1 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(y, ushort(k)); + s1 += simd_shuffle_down(y, ushort(k + 1)); + } + if (j == 0) xp[4 * plane + uint(t) * (n / 16) + d / 16] = s0 + s1; +} + +// grid (QH, n_tok), 256 threads: out = merge * sigmoid(gate), as prepared activations of n = QH * QD per token +kernel void qsa_merge(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], + device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const int h = int(tg.x), t = int(tg.y); + const int nchunks = (QBUDGET + QCH - 1) / QCH; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nc = int((width + QCH - 1) / QCH); + constexpr uint STRIDE = QH * (QD + 2); + device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2); + // the chunks' (m, l) and weights exp(m - M) once per threadgroup, not once per thread (nc <= 33 < 64) + threadgroup float mk[64], lk[64], ek[64], L_sh; + if (tid < uint(nc)) { + const float2 ml = *(device const float2*)(pp + tid * STRIDE + QD); + mk[tid] = ml.x; + lk[tid] = ml.y; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float M = -INFINITY; + for (int k = tid & 31; k < nc; k += 32) M = max(M, mk[k]); + M = simd_max(M); + if (tid < uint(nc)) ek[tid] = exp(mk[tid] - M); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + float L = 0.0f; + for (int k = 0; k < nc; ++k) L = fma(lk[k], ek[k], L); + L_sh = L; + } + float o = 0.0f; + device const float* q = pp + tid; + for (int k = 0; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); + threadgroup_barrier(mem_flags::mem_threadgroup); + o = o / L_sh; + const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; + const float y = o * (1.0f / (1.0f + exp(-gate))); + qsa_write_prepared(y, xp, a, t, h, tid); +} + diff --git a/results/round39/model-ab/baseline-kernels/qsa2.metal b/results/round39/model-ab/baseline-kernels/qsa2.metal new file mode 100644 index 0000000..d5b594c --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/qsa2.metal @@ -0,0 +1,338 @@ +// metal/kernels/qsa2.metal - QSA decode v2 (after qsa.metal): the same selection and attention math, more parallel. +// +// qsa_attn2 chunks of 32 cells, the 12 query heads of a kv head in 2 groups of 6: 4x the threadgroups of +// qsa_attn (at a 1K context 128 instead of 32); V loads 8 cells at a time. +// qsa_merge2 qsa_merge for that chunking. +// qsa_topk2 the digit of each radix pass found by a parallel suffix scan of the 256 bins (simdgroup prefix +// sums) instead of one thread; cells emitted from contiguous block ranges per thread after one scan. + +#define QCH2 32 +#define QHG 2 + +// partials: [n_tok][chunks2][QH][QD + 2]. grid (chunks2, QKV * QHG, n_tok), 256 threads. +kernel void qsa_attn2(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + 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 HPG = QH / QKV / QHG; // 6 + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH2]; + threadgroup float mh[HPG], lh[HPG]; + threadgroup int cid[QCH2]; + const int chunk = int(tg.x), g = int(tg.y) / QHG, hg = int(tg.y) % QHG, t = int(tg.z); + const int h0 = g * (QH / QKV) + hg * HPG; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH2; + const int nchunks = (QBUDGET + QCH2 - 1) / QCH2; + if (uint(c0) >= width) return; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + const int nc = min(QCH2, int(width) - c0); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + h0 + h) * QD + tid]; + if (tid < QCH2) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup const float* q_lane = &qs[0][lane * 8]; + for (int c = sg; c < nc; c += 8) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char4* kr = (device const char4*)(kq + rowi * QD + lane * 8); + const char4 k0 = kr[0], k1 = kr[1]; + const float scl = float(ksc[rowi * 4 + (lane >> 3)]); + const float kv[8] = {float(k0.x) * scl, float(k0.y) * scl, float(k0.z) * scl, float(k0.w) * scl, + float(k1.x) * scl, float(k1.y) * scl, float(k1.z) * scl, float(k1.w) * scl}; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + float d = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) d = fma(q_lane[h * QD + i], kv[i], d); + d = simd_sum(d); + if (lane == 0) s[h][c] = d * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg < HPG) { + const int h = sg; + const float sv = lane < nc ? s[h][lane] : -INFINITY; + const float m = simd_max(sv); + const float p = lane < nc ? exp(sv - m) : 0.0f; + const float l = simd_sum(p); + if (lane < nc) s[h][lane] = p; + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float acc[HPG]; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + const uint v_scale_off = tid >> 6; + int cb = 0; + for (; cb + 8 <= nc; cb += 8) { + float v[8]; + #pragma clang loop unroll(full) + for (int c = 0; c < 8; ++c) { + const uint rowi = uint(cid[cb + c]) * QKV + uint(g); + v[c] = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + v_scale_off]); + } + #pragma clang loop unroll(full) + for (int c = 0; c < 8; ++c) + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb + c], v[c], acc[h]); + } + for (; cb < nc; ++cb) { + const uint rowi = uint(cid[cb]) * QKV + uint(g); + const float v = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + v_scale_off]); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb], v, acc[h]); + } + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (h0 + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} + +kernel void qsa_merge2(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], + device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float mk[128], lk[128], ek[128], L_sh; + const int h = int(tg.x), t = int(tg.y); + const int nchunks = (QBUDGET + QCH2 - 1) / QCH2; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nc = int((width + QCH2 - 1) / QCH2); + constexpr uint STRIDE = QH * (QD + 2); + device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2); + if (tid < uint(nc)) { + const float2 ml = *(device const float2*)(pp + tid * STRIDE + QD); + mk[tid] = ml.x; + lk[tid] = ml.y; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float M = -INFINITY; + for (int k = tid & 31; k < nc; k += 32) M = max(M, mk[k]); + M = simd_max(M); + if (tid < uint(nc)) ek[tid] = exp(mk[tid] - M); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + float L = 0.0f; + for (int k = 0; k < nc; ++k) L = fma(lk[k], ek[k], L); + L_sh = L; + } + float o = 0.0f; + device const float* q = pp + tid; + for (int k = 0; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); + threadgroup_barrier(mem_flags::mem_threadgroup); + o = o / L_sh; + const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; + const float y = o * (1.0f / (1.0f + exp(-gate))); + qsa_write_prepared(y, xp, a, t, h, tid); +} + +kernel void qsa_topk2(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]], + constant QsaArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint hist[256]; + threadgroup uint sgs[32], sgs2[32]; + threadgroup uint sel[2]; + const int t = int(tg.x); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2, nb = n_bid + 1; + device int* out = ids + t * QBUDGET; + device const float* sc = scores + uint(t) * (a.max_cells / 4 + 2); + if (n_kv <= QBUDGET) { + for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); + return; + } + if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + uint prefix = 0, pmask = 0, need = QBUDGET; + if (nb <= 4096) { + uint keys[4]; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) { + const uint b = tid + 1024u * uint(i); + keys[i] = b < nb ? okey(sc[b]) : 0u; + } + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) { + const uint b = tid + 1024u * uint(i); + if (b >= nb) break; + const uint k = keys[i]; + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint cnt = 0, ex = 0; + if (tid < 256) { + cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); + if (pass < 3) atomic_store_explicit(&hist[255 - tid], 0u, memory_order_relaxed); + ex = simd_prefix_exclusive_sum(cnt); + if (lane == 31) sgs[sg] = ex + cnt; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 256) { + uint above = ex; + for (int i = 0; i < int(sg); ++i) above += sgs[i]; + if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; + } + } else { + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + for (uint b = tid; b < nb; b += 1024) { + const uint k = okey(sc[b]); + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint cnt = 0, ex = 0; + if (tid < 256) { + cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); + if (pass < 3) atomic_store_explicit(&hist[255 - tid], 0u, memory_order_relaxed); + ex = simd_prefix_exclusive_sum(cnt); + if (lane == 31) sgs[sg] = ex + cnt; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 256) { + uint above = ex; + for (int i = 0; i < int(sg); ++i) above += sgs[i]; + if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; + } + } + const uint T = prefix; + const uint per = (nb + 1023) / 1024, b0 = min(nb, uint(tid) * per), b1 = min(nb, b0 + per); + uint gt = 0, eq = 0; + for (uint b = b0; b < b1; ++b) { + const uint k = okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + if (k > T) gt += cells; + else if (k == T) eq += cells; + } + const uint eqb = simd_prefix_exclusive_sum(eq); + if (lane == 31) sgs2[sg] = eqb + eq; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint eq_before = eqb; + for (int i = 0; i < int(sg); ++i) eq_before += sgs2[i]; + const uint eq_take = eq_before >= need ? 0u : min(eq, need - eq_before); + const uint mine = gt + eq_take; + const uint mb = simd_prefix_exclusive_sum(mine); + if (lane == 31) sgs[sg] = mb + mine; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint off = mb; + for (int i = 0; i < int(sg); ++i) off += sgs[i]; + uint left = eq_take; + for (uint b = b0; b < b1; ++b) { + const uint k = okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + uint take = 0; + if (k > T) take = cells; + else if (k == T) { take = min(cells, left); left -= take; } + for (uint c = 0; c < take; ++c) out[off + c] = int(4 * b + c); + off += take; + } +} + +// qsa_attn3: qsa_attn's chunking and partial layout (qsa_merge reads it), without a reduction per (cell, head): +// for q.k a lane owns a cell (8 simdgroups = 2 halves of the 64 cells x 4 groups of 3 heads; q from threadgroup +// memory, read by every lane alike), for p.v a thread owns a dimension. grid (33, QKV, n_tok), 256 threads. +kernel void qsa_attn3(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + 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 HPG = QH / QKV; // 12 + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH]; + threadgroup float mh[HPG], lh[HPG]; + threadgroup int cid[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; + const int nchunks = (QBUDGET + QCH - 1) / QCH; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + if (uint(c0) >= width) return; + const int nc = min(QCH, int(width) - c0); + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + g * HPG + h) * QD + tid]; + if (tid < QCH) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + { + const int half_ = sg & 1, hgrp = sg >> 1; // cells 32 * half_ + lane, heads 3 * hgrp .. + 3 + const int c = 32 * half_ + lane; + if (c < nc) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char4* kr = (device const char4*)(kq + rowi * QD); + float d0 = 0.0f, d1 = 0.0f, d2 = 0.0f; + for (int gq = 0; gq < 4; ++gq) { + const float scl = float(ksc[rowi * 4 + gq]); + float e0 = 0.0f, e1 = 0.0f, e2 = 0.0f; + for (int i = 0; i < 16; ++i) { + const char4 k4 = kr[gq * 16 + i]; + const int base = gq * 64 + 4 * i; + const float kv[4] = {float(k4.x), float(k4.y), float(k4.z), float(k4.w)}; + for (int e = 0; e < 4; ++e) { + e0 = fma(qs[3 * hgrp][base + e], kv[e], e0); + e1 = fma(qs[3 * hgrp + 1][base + e], kv[e], e1); + e2 = fma(qs[3 * hgrp + 2][base + e], kv[e], e2); + } + } + d0 = fma(e0, scl, d0); + d1 = fma(e1, scl, d1); + d2 = fma(e2, scl, d2); + } + s[3 * hgrp][c] = d0 * 0.0625f; + s[3 * hgrp + 1][c] = d1 * 0.0625f; + s[3 * hgrp + 2][c] = d2 * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int h = sg; h < HPG; h += 8) { + float m = -INFINITY; + for (int c = lane; c < nc; c += 32) m = max(m, s[h][c]); + m = simd_max(m); + float l = 0.0f; + for (int c = lane; c < nc; c += 32) { + const float p = exp(s[h][c] - m); + s[h][c] = p; + l += p; + } + l = simd_sum(l); + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float acc[HPG]; + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + for (int cb = 0; cb < nc; cb += 8) { + float v[8]; + const int m = min(8, nc - cb); + for (int c = 0; c < 8; ++c) { + v[c] = 0.0f; + if (c < m) { + const uint rowi = uint(cid[cb + c]) * QKV + uint(g); + v[c] = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); + } + } + for (int c = 0; c < m; ++c) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb + c], v[c], acc[h]); + } + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (g * HPG + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} diff --git a/results/round39/model-ab/baseline-kernels/qsa_prompt.metal b/results/round39/model-ab/baseline-kernels/qsa_prompt.metal new file mode 100644 index 0000000..444e0c4 --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/qsa_prompt.metal @@ -0,0 +1,453 @@ +// metal/kernels/qsa_prompt.metal - QSA attention for a PROMPT chunk on the matrix units (after qsa.metal: QH, QKV, QD, +// QBUDGET, QCH, QsaArgs). The same split-K partials as qsa_attn (part [t][chunk][QH][QD + 2]: acc, then m, l), so +// qsa_merge is unchanged; per (chunk of QCH = 64 selected cells, kv head g, token t) one simdgroup: +// S[12 x 64] = Qh[12 x 256] . K^T / 16 single-simdgroup matmul2d, K (int8 * fp16 scale per 64) decoded into +// the right-input cooperative tensor, 32 dims a step +// per head: m = max, p = exp(s - m), l = sum p (fp32, through threadgroup memory), P -> fp16 +// acc[12 x 256] = P[12 x 64] . V[64 x 256] four 64-dim matmuls, V decoded into the right input +// The operands are fp16 (q rounded from fp32; k, v = int8 * scale rounded): to rounding of qsa_attn's fp32. +#include +using namespace mpp::tensor_ops; +#ifndef STRATA_COOP_K4 +#define STRATA_COOP_K4 +// function constant 7: the right-input cooperative tensor holds runs of 4 consecutive k at one n per lane (the M5's +// layout; tests/coop_layout). false (the M4: runs of 2 n at one k): every element decoded on its own. +constant bool COOP_K4 [[function_constant(7)]]; +constant bool coop_k4 = is_function_constant_defined(COOP_K4) ? COOP_K4 : true; +#endif + +// a bisection probe (function constant 8): stop at stage n and write m = n (1: after the ids, 2: after S, 3: after +// the softmax); 0 = normal +constant int QSA_STOP [[function_constant(8)]]; +constant int qsa_stop = is_function_constant_defined(QSA_STOP) ? QSA_STOP : 0; + +// grid (chunks, QKV, n_tok), 32 threads. qh: the q heads as fp16 [t][QH][QD] (qcur rounded). +kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int HPG = QH / QKV, NCH = (QBUDGET + QCH - 1) / QCH; + threadgroup float s[16 * QCH]; + threadgroup half p[16 * QCH]; + threadgroup int cid[QCH]; // coop_k4: byte offsets; otherwise cell ids + threadgroup half4 ks[QCH], vs[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; + device float* pp = part + ((uint(t) * NCH + chunk) * QH + g * HPG) * (QD + 2); + if (uint(c0) >= width) { + if (lane < HPG) { pp[lane * (QD + 2) + QD] = -INFINITY; pp[lane * (QD + 2) + QD + 1] = 0.0f; } + return; + } + const int nc = min(QCH, int(width) - c0); + for (int c = lane; c < QCH; c += 32) { + const int id = c < nc ? (n_kv <= QBUDGET ? c0 + c : ids[t * QBUDGET + c0 + c]) : -1; + cid[c] = coop_k4 ? int((uint(max(id, 0)) * QKV + g) * QD) : id; + if (id >= 0) { + const uint row4 = (uint(id) * QKV + g) * 4; + if (coop_k4) ks[c] = *(device const half4*)(ksc + row4); + vs[c] = *(device const half4*)(vsc + row4); + } else { + ks[c] = half4(0.0h); + vs[c] = half4(0.0h); + } + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (qsa_stop == 1) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 1.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } + // ---- S = Q K^T + { + constexpr auto desc = matmul2d_descriptor(16, QCH, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tQ((device half*)qh + (uint(t) * QH + g * HPG) * QD, + dextents(QD, HPG)); + auto rk = op.template get_right_input_cooperative_tensor(); + auto tA = tQ.template slice<32, 16>(0, 0); + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + // the second half repeats these eight cells, 16 k later; empty cells read row 0 at scale 0. + uint ko[8]; + float ks8[8][4]; + if (coop_k4) +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = rk.get_multidimensional_index(uint16_t(4 * j))[1]; + ko[j] = uint(cid[c]); +#pragma clang loop unroll(full) + for (int h = 0; h < 4; ++h) ks8[j][h] = float(ks[c][h]); + } + for (int k0 = 0; k0 < QD; k0 += 32) { + if (!coop_k4) { + for (uint16_t i = 0; i < rk.get_capacity(); ++i) { + const auto ix = rk.get_multidimensional_index(i); + const int k = k0 + ix[0], id = cid[ix[1]]; + const uint row = uint(id) * QKV + g; + rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)])) : half(0.0f); + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rk.get_capacity(); i += 4) { + const int k = k0 + rk.get_multidimensional_index(i)[0], j = (i >> 2) & 7; + const half4 w = half4(float4(*(device const char4*)(kq + (ko[j] + uint(k)))) * ks8[j][k >> 6]); + rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w; + } + tA = tQ.template slice<32, 16>(k0, 0); + op.run(tA, rk, acc); + } + tensor, tensor_inline> tS(s, dextents(QCH, 16)); + acc.store(tS); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (qsa_stop == 2) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 2.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } + // ---- softmax pieces per head (lane = head for the reductions; 12 heads) + float mh = -INFINITY, lh = 0.0f; + if (lane < HPG) { +#pragma clang loop unroll_count(4) + for (int c = 0; c < nc; ++c) mh = max(mh, s[lane * QCH + c] * 0.0625f); +#pragma clang loop unroll_count(4) + for (int c = 0; c < QCH; ++c) { + const float e = c < nc ? exp(s[lane * QCH + c] * 0.0625f - mh) : 0.0f; + lh += e; + p[lane * QCH + c] = half(e); + } + } else if (lane < 16) { +#pragma clang loop unroll_count(4) + for (int c = 0; c < QCH; ++c) p[lane * QCH + c] = half(0.0f); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (qsa_stop == 3) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 3.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } + // ---- acc = P V, 64 dims at a time + { + // K = 32 cells a run (two runs a 64-cell chunk): K = 64 runs compute garbage on the M4 + constexpr auto desc = matmul2d_descriptor(16, 64, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tP(p, dextents(QCH, 16)); + tensor, tensor_inline> tO(pp, dextents(QD + 2, HPG)); + for (int d0 = 0; d0 < QD; d0 += 64) { + auto rv = op.template get_right_input_cooperative_tensor(); + auto tPh = tP.template slice<32, 16>(0, 0); + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + for (int c0 = 0; c0 < nc; c0 += 32) { + if (!coop_k4) { +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); ++i) { + const auto ix = rv.get_multidimensional_index(i); // (cell, dim) + const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; + rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vs[c][d >> 6])) : half(0.0f); + } + } else { + uint vo[8]; + float vs8[8]; + const int cb = c0 + rv.get_multidimensional_index(uint16_t(0))[0]; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = cb + (j & 3) + 16 * (j >> 2); + vo[j] = uint(cid[c]) + uint(d0); + vs8[j] = float(vs[c][d0 >> 6]); + } +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); i += 4) { + const int d = rv.get_multidimensional_index(i)[1], r = i < rv.get_capacity() / 2 ? 0 : 4; +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) rv[i + j] = half(float(vq[vo[r + j] + uint(d)]) * vs8[r + j]); + } + } + tPh = tP.template slice<32, 16>(c0, 0); + op.run(tPh, rv, acc); + } + if (qsa_stop == 4) { // the bisection: the old direct store into the 258-float rows + auto mO = tO.slice(d0, 0); + acc.store(mO); + continue; + } + // through threadgroup memory (s is free now): a matmul2d store into rows of QD + 2 floats loses the + // whole threadgroup on the M4 + tensor, tensor_inline> tT(s, dextents(64, 16)); + simdgroup_barrier(mem_flags::mem_threadgroup); + acc.store(tT); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int it = lane; it < HPG * 16; it += 32) { + const int h = it >> 4, d = (it & 15) * 4; + const float4 v = *(threadgroup const float4*)(s + h * 64 + d); + // odd heads start on an 8-byte boundary, not a float4 boundary. + device float* dst = pp + h * (QD + 2) + d0 + d; + dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w; + } + } + } + if (lane < HPG) { pp[lane * (QD + 2) + QD] = mh; pp[lane * (QD + 2) + QD + 1] = lh; } +} + +// qcur (fp32) -> fp16 for qsa_attn_mm +kernel void qsa_q_f16(device const float* q [[buffer(0)]], device half* qh [[buffer(1)]], constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + if (gid < n) qh[gid] = half(q[gid]); +} + +// qsa_score as a GEMM (fp32 operands, so the selection sees qsa_score's values to fp32 rounding): S[(t, h)][b] = +// qidx[t][h] . pooled[b] for 16 tokens x IH heads (64 rows) by 64 blocks a threadgroup; the epilogue sums relu over +// the heads. The tail block n_bid (the `dead` key, + 1e9 when it has cells) and blocks past it are left to the +// scalar path: the tail is scored here by one simdgroup a token; blocks > n_bid are not written (qsa_topk reads +// b <= n_bid). grid (ceil(max blocks / 64), ceil(T / 16)), 128 threads. +kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const float* pooled [[buffer(1)]], + device const float* dead [[buffer(2)]], device float* scores [[buffer(3)]], + constant QsaArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float ot[64 * 64]; + const int b0 = int(tg.x) * 64, t0 = int(tg.y) * 16; + const uint max_t = min(uint(t0 + 15), a.n_tok - 1u); + const uint max_bid = (a.pos0 + max_t + 1u) >> 2; + // column zero still owns every token's tail score. + if (tg.x != 0 && uint(b0) >= max_bid) return; + const uint nb_all = (a.pos0 + a.n_tok) / 4 + 1; // blocks any token of the chunk can score + tensor, tensor_inline> tQ((device float*)qidx, dextents(ID, int(a.n_tok) * IH)); + tensor, tensor_inline> tK((device float*)pooled, dextents(ID, int(nb_all))); + tensor, tensor_inline> tO(ot, dextents(64, 64)); + constexpr auto desc = matmul2d_descriptor(64, 64, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tQ.slice(0, t0 * IH); + auto mB = tK.slice(0, b0); + op.run(mA, mB, tO); + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint stride = a.max_cells / 4 + 2; + const int j = tid & 63, b = b0 + j; + for (int tt = tid >> 6; tt < 16; tt += 2) { + const int t = t0 + tt; + if (t >= int(a.n_tok)) break; + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + if (uint(b) >= n_bid) continue; + float s = 0.0f; + for (int h = 0; h < IH; ++h) s += max(ot[(tt * IH + h) * 64 + j], 0.0f); + scores[uint(t) * stride + b] = s; + } + // the tail block of each token in this threadgroup's block range (the first column threadgroup does it) + if (tg.x == 0) { + float k[4]; + for (int i = 0; i < 4; ++i) k[i] = dead[lane + 32 * i]; + for (int tt = sg; tt < 16; tt += 4) { + const int t = t0 + tt; + if (t >= int(a.n_tok)) break; + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + float s = 0.0f; + for (int h = 0; h < IH; ++h) { + float d = 0.0f; + for (int i = 0; i < 4; ++i) d = fma(qidx[(t * IH + h) * ID + lane + 32 * i], k[i], d); + d = simd_sum(d); + s += max(d, 0.0f); + } + if ((n_kv & 3) != 0) s += 1e9f; + if (lane == 0) scores[uint(t) * stride + n_bid] = s; + } + } +} + +// qsa_flash: the whole attention of a (kv head g, token t) in one threadgroup (no split-K partials, no merge): the +// width selected cells in chunks of 64 with an online softmax; 4 simdgroups, simdgroup s owns dims 64 s .. 64 s + 63 +// (its slice of S's K sum and its 64 output dims). Out: xq = (acc / l) * sigmoid(gate) as fp16 (attn_output's GEMM +// operand). grid (QKV, n_tok), 128 threads. +kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device const float* qfull [[buffer(6)]], device half* xqh [[buffer(7)]], constant QsaArgs& a [[buffer(8)]], + 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 HPG = QH / QKV; + threadgroup float sp[4][HPG * QCH]; + threadgroup half* p = (threadgroup half*)sp[1]; // P [16][QCH] in fp16 over a dead S partial + threadgroup int cid[QCH]; // coop_k4: byte offsets; otherwise cell ids + threadgroup half4 ks_tg[QCH], vs_tg[QCH]; + threadgroup float fac[16], lsum[16], mrun[16]; + const int g = int(tg.x), t = int(tg.y); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nch = int((width + QCH - 1) / QCH); + if (tid < 16) { mrun[tid] = -INFINITY; lsum[tid] = 0.0f; } + constexpr auto dS = matmul2d_descriptor(16, QCH, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + constexpr auto dO = matmul2d_descriptor(16, 64, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d opS; + matmul2d opO; + tensor, tensor_inline> tQ((device half*)qh + (uint(t) * QH + g * HPG) * QD, + dextents(QD, HPG)); + tensor, tensor_inline> tP(p, dextents(QCH, 16)); + auto rv = opO.template get_right_input_cooperative_tensor(); + auto tPh = tP.template slice<32, 16>(0, 0); + auto acc = opO.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); + // (a second destination tensor, not a copy of acc: the MetalPerformancePrimitives headers of macOS 26.3 cannot copy a + // cooperative tensor - "cannot initialize a parameter of type 'void *' with an rvalue of type 'const void *'", #2323; + // tot is zeroed right below, so the copy was only there for its type) + auto tot = opO.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); // the running output; acc takes one chunk (fresh each chunk: on the M4 a run does not accumulate + // into elements rescaled in place) + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) tot[i] = 0.0f; + const int d0 = sg * 64; + if (tid < 16 - HPG) fac[HPG + tid] = 0.0f; // the padding rows' factor (P's rows are zeroed each chunk) + for (int ch = 0; ch < nch; ++ch) { + const int c0 = ch * QCH, nc = min(QCH, int(width) - c0); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < QCH) { + const int id = tid < nc ? (n_kv <= QBUDGET ? c0 + tid : ids[t * QBUDGET + c0 + tid]) : -1; + cid[tid] = coop_k4 ? int((uint(max(id, 0)) * QKV + g) * QD) : id; + if (id >= 0) { + const uint row4 = (uint(id) * QKV + g) * 4; + ks_tg[tid] = *(device const half4*)(ksc + row4); + vs_tg[tid] = *(device const half4*)(vsc + row4); + } else { + ks_tg[tid] = half4(0.0h); + vs_tg[tid] = half4(0.0h); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // Eight lanes per head; keep the denominator's original sequential sum. + const int h = sg * 4 + (lane >> 3), hlane = lane & 7; + float sc[8]; + // S partial over this simdgroup's 64 dims + { + auto rk = opS.template get_right_input_cooperative_tensor(); + auto tA = tQ.template slice<32, 16>(d0, 0); + auto s = opS.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < s.get_capacity(); ++i) + if (s.is_valid_element(i)) s[i] = 0.0f; + // the lane's eight cells (one a run of four k; the second half repeats them 16 k on): row pointers and + // scales once a chunk, unconditional loads (an empty cell reads row 0 at scale 0, and a +-0 operand sums to + // the same value as the skipped one; the branchy loads cost half the kernel) + uint ko[8]; + float ks8[8]; + if (coop_k4) { +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = rk.get_multidimensional_index(uint16_t(4 * j))[1]; + ko[j] = uint(cid[c]) + uint(d0 + rk.get_multidimensional_index(uint16_t(0))[0]); + ks8[j] = float(ks_tg[c][sg]); + } + } +#pragma clang loop unroll(full) + for (int k0 = d0; k0 < d0 + 64; k0 += 32) { + if (!coop_k4) { + for (uint16_t i = 0; i < rk.get_capacity(); ++i) { + const auto ix = rk.get_multidimensional_index(i); + const int k = k0 + ix[0], id = cid[ix[1]]; + const uint row = uint(id) * QKV + g; + rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ks_tg[ix[1]][k >> 6])) : half(0.0f); + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rk.get_capacity(); i += 4) { + const int k = k0 - d0 + (i >= rk.get_capacity() / 2 ? 16 : 0), j = (i >> 2) & 7; + const half4 w = half4(float4(*(device const char4*)(kq + ko[j] + uint(k))) * ks8[j]); + rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w; + } + tA = tQ.template slice<32, 16>(k0, 0); + opS.run(tA, rk, s); + } + // Preserve ((s0 + s1) + (s2 + s3)) while storing all partials together. + tensor, tensor_inline> tS(sp[sg], dextents(QCH, HPG)); + s.store(tS); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (h < HPG) +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) { + const int q = h * QCH + hlane + 8 * k; + sc[k] = ((sp[0][q] + sp[1][q]) + (sp[2][q] + sp[3][q])) * 0.0625f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); // P overwrites sp[1] + if (h < HPG) { + float m = mrun[h]; +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) + if (hlane + 8 * k < nc) m = max(m, sc[k]); + m = max(m, simd_shuffle_xor(m, 4)); + m = max(m, simd_shuffle_xor(m, 2)); + m = max(m, simd_shuffle_xor(m, 1)); +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) { + const int c = hlane + 8 * k, q = h * QCH + c; + const float e = c < nc ? exp(sc[k] - m) : 0.0f; + p[q] = half(e); + sp[0][q] = e; + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (hlane == 0) { + const float f = exp(mrun[h] - m); + float l = lsum[h] * f; +#pragma clang loop unroll(full) + for (int c = 0; c < QCH; c += 4) { // in order, four loads at a time + const float4 v = *(threadgroup const float4*)(sp[0] + h * QCH + c); + l += v.x; l += v.y; l += v.z; l += v.w; + } + mrun[h] = m; + lsum[h] = l; + fac[h] = f; + } + } else { + for (int c = hlane; c < QCH; c += 8) p[h * QCH + c] = half(0.0f); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // rescale this simdgroup's accumulator rows, then acc += P V over its 64 dims + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + for (int c0 = 0; c0 < nc; c0 += 32) { // K = 32 cells a run (K = 64 is wrong on the M4) + if (!coop_k4) { +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); ++i) { + const auto ix = rv.get_multidimensional_index(i); // (cell, dim) + const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; + rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vs_tg[c][sg])) : half(0.0f); + } + } else { + // a lane holds the same eight cells (two runs of four, 16 apart) at each of its dims: as for rk, row + // pointers and scales once a run, unconditional loads + uint vo[8]; + float vs8[8]; + const int cb = c0 + rv.get_multidimensional_index(uint16_t(0))[0]; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = cb + (j & 3) + 16 * (j >> 2); + vo[j] = uint(cid[c]) + uint(d0); + vs8[j] = float(vs_tg[c][sg]); + } +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); i += 4) { + const int d = rv.get_multidimensional_index(i)[1], r = i < rv.get_capacity() / 2 ? 0 : 4; +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) rv[i + j] = half(float(vq[vo[r + j] + uint(d)]) * vs8[r + j]); + } + } + tPh = tP.template slice<32, 16>(c0, 0); + opO.run(tPh, rv, acc); + } + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) tot[i] = fma(tot[i], fac[acc.get_multidimensional_index(i)[1]], acc[i]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (!coop_k4) { + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) { + const auto ix = acc.get_multidimensional_index(i); // (dim, row) + const int h = ix[1], d = d0 + ix[0]; + if (h < HPG) { + const int hh = g * HPG + h; + const float gate = qfull[(uint(t) * QH + hh) * 2 * QD + QD + d]; + xqh[(uint(t) * QH + hh) * QD + d] = half(tot[i] / lsum[h] * (1.0f / (1.0f + exp(-gate)))); + } + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { // runs of four dims of one row + const auto ix = acc.get_multidimensional_index(i); + const int h = ix[1], d = d0 + ix[0]; + if (h < HPG) { + const int hh = g * HPG + h; + const float l = lsum[h]; + const float4 gate = *(device const float4*)(qfull + (uint(t) * QH + hh) * 2 * QD + QD + d); + const float4 v = float4(tot[i], tot[i + 1], tot[i + 2], tot[i + 3]) / l * (1.0f / (1.0f + exp(-gate))); + *(device half4*)(xqh + (uint(t) * QH + hh) * QD + d) = half4(v); + } + } +} diff --git a/results/round39/model-ab/baseline-kernels/quant.metal b/results/round39/model-ab/baseline-kernels/quant.metal new file mode 100644 index 0000000..fe667be --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/quant.metal @@ -0,0 +1,316 @@ +// metal/kernels/quant.metal - the weight formats as "units": a unit is UNIT consecutive weights of one row that one +// SIMD lane decodes at a time into raw values q[0..UNIT) and one scale s, so that the unit's dot with x is +// s * sum(q[j] * x[j]). Every GEMV kernel (moe.metal, gemv.metal) is written once over these. +// +// Decode rules transcribed from ggml-quants.c dequantize_row_* (3cf03257) and Strata's Q2_0 contract +// (docs/q2_0-contract.md); the grids and sign tables are ggml-common.h's (prepended by the runtime). The tables a +// format needs are copied to threadgroup memory once per threadgroup (constant-memory lookups with divergent indices +// serialise on Apple GPUs). +// +// Blocks start at 2-byte aligned addresses (18, 66, 74, 82, 98, 110 bytes), so every multi-byte field is read as +// ushorts. + +#define T_Q4_0 2 +#define T_Q3_K 11 +#define T_BF16 30 +#define T_F32 0 +#define T_Q5_0 6 +#define T_Q5_1 7 +#define T_Q8_0 8 +#define T_Q4_K 12 +#define T_Q5_K 13 +#define T_Q6_K 14 +#define T_IQ2_XXS 16 +#define T_IQ2_XS 17 +#define T_IQ3_XXS 18 +#define T_IQ4_NL 20 +#define T_IQ3_S 21 +#define T_IQ2_S 22 +#define T_IQ4_XS 23 +#define T_Q2_0 42 +#define T_Q2_0R 142 // Q2_0 read with raw activations (codes - 1 in the decoder; no prepared x needed) + +template struct fmt; +// QK block elements, BB block bytes, UNIT weights per unit, TAB words of threadgroup table +template <> struct fmt { enum { QK = 64, BB = 18, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 64, BB = 18, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 98, UNIT = 16, TAB = 256 }; }; +template <> struct fmt { enum { QK = 256, BB = 66, UNIT = 16, TAB = 512 }; }; +template <> struct fmt { enum { QK = 256, BB = 74, UNIT = 16, TAB = 1024 }; }; +template <> struct fmt { enum { QK = 256, BB = 82, UNIT = 16, TAB = 2048 }; }; +template <> struct fmt { enum { QK = 256, BB = 110, UNIT = 16, TAB = 512 }; }; +template <> struct fmt { enum { QK = 32, BB = 18, UNIT = 32, TAB = 256 }; }; +template <> struct fmt { enum { QK = 256, BB = 136, UNIT = 32, TAB = 256 }; }; +template <> struct fmt { enum { QK = 256, BB = 110, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 144, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 176, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 210, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 16, BB = 32, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 16, BB = 64, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 32, BB = 34, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 32, BB = 18, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 32, BB = 22, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 32, BB = 24, UNIT = 32, TAB = 1 }; }; + +static inline uint rd32(device const uchar* p) { + device const ushort* s = (device const ushort*)p; + return uint(s[0]) | (uint(s[1]) << 16); +} +static inline float rdh(device const uchar* p) { return float(*(device const half*)p); } + +// ---- threadgroup tables ------------------------------------------------------------------------------------------- +// layout in `tab` (uint words): [0, G) the grid; IQ4_NL/XS: 256 half2. +template static inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth); + +// QUANT_CTAB (function constant 23): the i-quant grids read from ggml-common's constant arrays +// instead of a threadgroup copy (no table load and no barrier wait at the threadgroup's start); the same values. +constant bool QUANT_CTAB [[function_constant(23)]]; +constant bool ctab = is_function_constant_defined(QUANT_CTAB) && QUANT_CTAB; +// ksigns_iq2xs[j] = j | parity(j) << 7 (an even count of minus signs): the low 4 bits of run hh (0 or 1). +static inline uint sign4_iq2xs(uint j, int hh) { + return hh ? ((j >> 4) | ((popcount(j) & 1u) << 3)) : j; +} +static inline uint g_iq3xxs(threadgroup const uint* tab, uint i) { return ctab ? iq3xxs_grid[i] : tab[i]; } +static inline uint g_iq3s(threadgroup const uint* tab, uint i) { return ctab ? iq3s_grid[i] : tab[i]; } +static inline uint g_iq2xxs(threadgroup const uint* tab, uint i, int w) { return ctab ? uint(iq2xxs_grid[i] >> (32 * w)) : tab[2 * i + w]; } +static inline uint g_iq2xs(threadgroup const uint* tab, uint i, int w) { return ctab ? uint(iq2xs_grid[i] >> (32 * w)) : tab[2 * i + w]; } +static inline uint g_iq2s(threadgroup const uint* tab, uint i, int w) { return ctab ? uint(iq2s_grid[i] >> (32 * w)) : tab[2 * i + w]; } + +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} + +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 256; i += nth) tab[i] = iq3xxs_grid[i]; +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 512; i += nth) tab[i] = iq3s_grid[i]; +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 256; i += nth) { tab[2 * i] = uint(iq2xxs_grid[i]); tab[2 * i + 1] = uint(iq2xxs_grid[i] >> 32); } +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 512; i += nth) { tab[2 * i] = uint(iq2xs_grid[i]); tab[2 * i + 1] = uint(iq2xs_grid[i] >> 32); } +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 1024; i += nth) { tab[2 * i] = uint(iq2s_grid[i]); tab[2 * i + 1] = uint(iq2s_grid[i] >> 32); } +} +// IQ4: a byte -> (codebook[low nibble], codebook[high nibble]) as half2, so one lookup gives two weights +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + for (int i = tid; i < 256; i += nth) + tab[i] = as_type(half2(half(kvalues_iq4nl[i & 15]), half(kvalues_iq4nl[i >> 4]))); +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + load_tab(tab, tid, nth); +} + +// ---- unit decoders: return s, fill q[UNIT] ------------------------------------------------------------------------ +// ui = unit index inside the block (0 .. QK/UNIT - 1). +template static inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab); + +// Q2_0: 64 = 4 units of 16; code c (2 bits, LSB first) -> c - 1; scale d. +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint*) { + device const ushort* p = (device const ushort*)(b + 2 + 4 * ui); + const ushort a = p[0], c = p[1]; + for (int j = 0; j < 8; ++j) { + q[j] = float((a >> (2 * j)) & 3) - 1.0f; + q[8 + j] = float((c >> (2 * j)) & 3) - 1.0f; + } + return rdh(b); +} + +// 8 signed grid bytes from a 4-byte grid word pair helper +static inline void put_signed4(thread float* q, uint g, uint s) { + const uchar4 v = as_type(g); + q[0] = (s & 1) ? -float(v[0]) : float(v[0]); + q[1] = (s & 2) ? -float(v[1]) : float(v[1]); + q[2] = (s & 4) ? -float(v[2]) : float(v[2]); + q[3] = (s & 8) ? -float(v[3]) : float(v[3]); +} + +// IQ3_XXS: 256 = 8 x 32; unit = half of a 32 (two grid pairs). +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + const uint aux = rd32(b + 2 + 64 + 4 * ib), s0 = aux >> (14 * h); + device const uchar* qs = b + 2 + 4 * ui; + #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { + const uint s = (s0 >> (7 * l2)) & 127; + put_signed4(q + 8 * l2, g_iq3xxs(tab, qs[2 * l2]), s); + put_signed4(q + 8 * l2 + 4, g_iq3xxs(tab, qs[2 * l2 + 1]), sign4_iq2xs(s, 1)); + } + return rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f); +} + +// IQ2_XXS: per 32: two u32 (4 grid indices, then 4 x 7 sign bits + 4-bit scale). +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + device const uchar* qs = b + 2 + 4 * ib + 2 * ui; + const uint a1 = rd32(b + 2 + 8 * ib + 4), s0 = a1 >> (14 * h); + #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { + const uint idx = qs[l2], s = (s0 >> (7 * l2)) & 127; + const uint2 g = ctab ? as_type(iq2xxs_grid[idx]) : ((threadgroup const uint2*)tab)[idx]; + put_signed4(q + 8 * l2, g.x, s); + put_signed4(q + 8 * l2 + 4, g.y, sign4_iq2xs(s, 1)); + } + return rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f); +} + +// IQ2_XS: uint16 qs[32] (9-bit grid index + 7-bit sign index), scales[8] nibbles. +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + device const ushort* qs = (device const ushort*)(b + 2) + 2 * ui; + #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { + const uint v = qs[l2]; + const uint idx = v & 511, s = v >> 9; + const uint2 g = ctab ? as_type(iq2xs_grid[idx]) : ((threadgroup const uint2*)tab)[idx]; + put_signed4(q + 8 * l2, g.x, s); + put_signed4(q + 8 * l2 + 4, g.y, sign4_iq2xs(s, 1)); + } + const uint sc = b[66 + ib]; + return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); +} + +// IQ2_S: qs[32] low 8 bits of the 10-bit index, signs[32], qh[8] (2 high bits per group of 8), scales[8]. +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + const uint q2 = *(device const ushort*)(b + 2 + 2 * ui); + const uint s2 = *(device const ushort*)(b + 34 + 2 * ui); + const uint qh = uint(b[66 + ib]) >> (4 * h); + #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { + const uint idx = ((q2 >> (8 * l2)) & 255) | ((qh << (8 - 2 * l2)) & 0x300); + const uint s = (s2 >> (8 * l2)) & 255; + const uint2 g = ctab ? as_type(iq2s_grid[idx]) : ((threadgroup const uint2*)tab)[idx]; + put_signed4(q + 8 * l2, g.x, s); + put_signed4(q + 8 * l2 + 4, g.y, s >> 4); + } + const uint sc = b[74 + ib]; + return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); +} + +// IQ3_S: qs[64] low 8 bits of 9-bit indices, qh[8], signs[32], scales[4] (nibble per 32). +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + const uint s2 = *(device const ushort*)(b + 74 + 2 * ui); + const uint qh = b[66 + ib]; + #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { + const int l = 2 * h + l2; + const uint i1 = uint(b[2 + 8 * ib + 2 * l]) | ((qh << (8 - 2 * l)) & 256); + const uint i2 = uint(b[2 + 8 * ib + 2 * l + 1]) | ((qh << (7 - 2 * l)) & 256); + const uint s = (s2 >> (8 * l2)) & 255; + put_signed4(q + 8 * l2, g_iq3s(tab, i1), s); + put_signed4(q + 8 * l2 + 4, g_iq3s(tab, i2), s >> 4); + } + const uint sc = b[106 + (ib >> 1)]; + return rdh(b) * fma(float((sc >> (4 * (ib & 1))) & 15), 2.0f, 1.0f); +} + +// IQ4_NL: 32 = 1 unit; low nibbles are elements 0..15, high 16..31 (one half2 lookup per byte). +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint* tab) { + device const ushort* p = (device const ushort*)(b + 2); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); + q[2 * j] = float(w0.x); + q[2 * j + 16] = float(w0.y); + q[2 * j + 1] = float(w1.x); + q[2 * j + 17] = float(w1.y); + } + return rdh(b); +} + +// IQ4_XS: 256 = 8 units of 32; 6-bit scale from scales_l nibble + 2 bits of scales_h, minus 32. +static inline float iq4xs_scale(device const uchar* b, int ui) { + const uint sh = *(device const ushort*)(b + 2); + const int ls = int((b[4 + (ui >> 1)] >> (4 * (ui & 1))) & 15) | int(((sh >> (2 * ui)) & 3) << 4); + return rdh(b) * float(ls - 32); +} +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + device const ushort* p = (device const ushort*)(b + 8 + 16 * ui); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); + q[2 * j] = float(w0.x); + q[2 * j + 16] = float(w0.y); + q[2 * j + 1] = float(w1.x); + q[2 * j + 17] = float(w1.y); + } + return iq4xs_scale(b, ui); +} + +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { + device const ushort* p = (device const ushort*)(b + 2); + for (int j = 0; j < 16; ++j) { + const ushort v = p[j]; + q[2 * j] = float(as_type(uchar(v & 255))); + q[2 * j + 1] = float(as_type(uchar(v >> 8))); + } + return rdh(b); +} + +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { + device const ushort* p = (device const ushort*)(b + 2); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + q[2 * j] = float(v & 15) - 8.0f; + q[2 * j + 16] = float((v >> 4) & 15) - 8.0f; + q[2 * j + 1] = float((v >> 8) & 15) - 8.0f; + q[2 * j + 17] = float(v >> 12) - 8.0f; + } + return rdh(b); +} + +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { + const uint hb = rd32(b + 2); + device const ushort* p = (device const ushort*)(b + 6); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const int j0 = 2 * j, j1 = 2 * j + 1; + q[j0] = float(int(v & 15) | int(((hb >> j0) & 1) << 4)) - 16.0f; + q[j0 + 16] = float(int((v >> 4) & 15) | int(((hb >> (j0 + 16)) & 1) << 4)) - 16.0f; + q[j1] = float(int((v >> 8) & 15) | int(((hb >> j1) & 1) << 4)) - 16.0f; + q[j1 + 16] = float(int(v >> 12) | int(((hb >> (j1 + 16)) & 1) << 4)) - 16.0f; + } + return rdh(b); +} + +// Q5_1: 32 = 1 unit; q = the 5-bit code (no offset; the min m is applied through the x sums), returns d +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { + const uint hb = rd32(b + 4); + device const ushort* p = (device const ushort*)(b + 8); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const int j0 = 2 * j, j1 = 2 * j + 1; + q[j0] = float(int(v & 15) | int(((hb >> j0) & 1) << 4)); + q[j0 + 16] = float(int((v >> 4) & 15) | int(((hb >> (j0 + 16)) & 1) << 4)); + q[j1] = float(int((v >> 8) & 15) | int(((hb >> j1) & 1) << 4)); + q[j1 + 16] = float(int(v >> 12) | int(((hb >> (j1 + 16)) & 1) << 4)); + } + return rdh(b); +} + +// sum over groups of LPR consecutive lanes (LPR = 1, 2, 4, 8, 16 or 32); the result is valid in every lane +template +static inline float lane_sum(float v) { + if (LPR == 32) return simd_sum(v); + for (ushort o = LPR / 2; o > 0; o >>= 1) v += simd_shuffle_xor(v, o); + return v; +} diff --git a/results/round39/model-ab/baseline-kernels/router.metal b/results/round39/model-ab/baseline-kernels/router.metal new file mode 100644 index 0000000..5c6bf59 --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/router.metal @@ -0,0 +1,173 @@ +// metal/kernels/router.metal - the router GEMV and moe_route in ONE dispatch (after quant/units/moe/gemv.metal): the +// logits rows as gemv's (LPR = 32 lanes a row, KS = 1: bitwise gemv__32_1_*_1's), 16 rows a threadgroup of 512 +// threads; every threadgroup bumps a device counter after its rows (device-scope fences around it), and the LAST one +// runs moe_route's body (the softmax / top-10 / weights, then the plan) with its 512 threads (= n_expert) and resets +// the counter. No threadgroup waits for another. ids / w / plan are bitwise moe_route's. +// Buffers: W, xp, s16, logits, args (GemvArgs: n_out = n_expert), ids, w, plan, MoeArgs, counter (one uint, zero). + +#ifndef STRATA_MOCE_FC_METAL +#define STRATA_MOCE_FC_METAL +constant int kMoceProtectFc [[function_constant(24)]]; +constant int kMoceProtect = is_function_constant_defined(kMoceProtectFc) ? kMoceProtectFc : 3; +constant uint kMoceScaleR1BitsFc [[function_constant(25)]]; +constant float kMoceScaleR1 = is_function_constant_defined(kMoceScaleR1BitsFc) ? as_type(kMoceScaleR1BitsFc) : 1.0f; +#endif +static inline void route_body(device const float* logits, device int* ids, device float* w, device int* plan, + constant MoeArgs& a, device const uint* smap, ushort tid, ushort nth, ushort sg, ushort lane, + threadgroup int* s_grp, threadgroup int* s_ent) { + if (sg < a.n_tok) { + device const float* l = logits + sg * 512; + float v[16]; + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + v[i] = idx < int(a.n_expert) ? l[idx] : -INFINITY; + } + float mx = -INFINITY; + for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + float sum = 0.0f; + for (int i = 0; i < 16; ++i) { + v[i] = (lane + 32 * i < int(a.n_expert)) ? exp(v[i] - mx) : 0.0f; + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + for (int i = 0; i < 16; ++i) { + if (lane + 32 * i >= int(a.n_expert)) { + v[i] = -INFINITY; + } else { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + } + const bool draft_moce = (a.pad & 0x80000000u) != 0u && sg > 0; + const float base_moce = (a.pad & 0x7FFFFFFFu) != 0u ? as_type(a.pad & 0x7FFFFFFFu) : 1.0f; + const float moce_scale = draft_moce ? (base_moce * base_moce) : base_moce; + const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f; + const int r_prot = draft_moce ? 1 : kMoceProtect; + const float r1_scale = kMoceScaleR1; + const float rcp_r1 = r1_scale > 1.0f ? (1.0f / r1_scale) : 1.0f; + const float r_step_scale = (r1_scale > 1.0f && moce_scale > 1.0f) ? (moce_scale * rcp_r1) : moce_scale; + ushort scaled_mask = 0; + float selected = 0.0f, ssum = 0.0f; + for (int rank = 0; rank < 10; ++rank) { + if (r1_scale > 1.0f && r_prot > 1) { + if (rank == 1 && smap != nullptr) { + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) { + v[i] *= r1_scale; + scaled_mask |= ushort(1u << i); + } + } + } else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) { + for (int i = 0; i < 16; ++i) { + if ((scaled_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) { + v[i] *= r_step_scale; + } + } + } + } else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) { + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) { + v[i] *= moce_scale; + scaled_mask |= ushort(1u << i); + } + } + } + float best = v[0]; + int ex = lane; + #pragma clang loop unroll(full) + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + float true_prob = best; + const int min_scaled_rank = (r1_scale > 1.0f && r_prot > 1) ? 1 : r_prot; + if (rank >= min_scaled_rank && (moce_scale > 1.0f || r1_scale > 1.0f) && best > 0.0f) { + const ushort win_lane = ushort(ex & 31); + 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) { + // 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); + } + } + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) v[i] = (ex == lane + 32 * i) ? -INFINITY : v[i]; + if ((ex & 31) == lane) { + ids[sg * 10 + rank] = ex; + ssum += true_prob; + } + if (rank == lane) selected = true_prob; + } + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + ssum = max(ssum, 6.103515625e-5f); + const float inv = 1.0f / ssum; + if (lane < 10) w[sg * 10 + lane] = selected * inv; + } + threadgroup_barrier(mem_flags::mem_device); + plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent, smap); +} + +template +kernel void router_route(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + 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)]], + device atomic_uint* counter [[buffer(9)]], device const uint* smap [[buffer(10)]], + uint tgx [[threadgroup_position_in_grid]], uint ngx [[threadgroups_per_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort nth [[threads_per_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int UNIT = fmt::UNIT; + threadgroup uint tab[fmt::TAB]; + threadgroup int s_grp[64], s_ent[64]; + threadgroup uint last; + load_tab(tab, tid, nth); + threadgroup_barrier(mem_flags::mem_threadgroup); + // the GEMV: simdgroup sg -> row tg.x * 16 + sg, 32 lanes over its units + const int row = int(tgx) * 16 + sg; + const int rr = min(row, int(ga.n_out) - 1); + device const uchar* rp = W + (ulong)rr * ga.row_bytes; + const int U = int(ga.n_in) / UNIT, s16n = int(ga.n_in) / 16; + float acc[8]; +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[t] = 0.0f; + for (int u = lane; u < U; u += 32) { + float q[UNIT]; + float4 sc; + Dec::run(rp, u, q, sc, tab); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(ga.x_stride), s16 + t * s16n, u, xs, xsum); + acc[t] += unit_combine(q, sc, xs, xsum); + } + } +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum<32>(acc[t]); + if (lane == 0 && row < int(ga.n_out)) logits[t * int(ga.y_stride) + row] = v; + } + // the last threadgroup to finish routes + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + threadgroup_barrier(mem_flags::mem_device); + if (tid == 0) last = atomic_fetch_add_explicit(counter, 1u, memory_order_relaxed) == ngx - 1 ? 1u : 0u; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (!last) return; + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + route_body(logits, ids, w, plan, a, smap, tid, nth, sg, lane, s_grp, s_ent); + if (tid == 0) atomic_store_explicit(counter, 0u, memory_order_relaxed); +} + +template [[host_name("router_route_30")]] kernel decltype(router_route) router_route; +template [[host_name("router_route_0")]] kernel decltype(router_route) router_route; + +template [[host_name("router_route_8")]] kernel decltype(router_route) router_route; diff --git a/results/round39/model-ab/baseline-kernels/rt_bench.metal b/results/round39/model-ab/baseline-kernels/rt_bench.metal new file mode 100644 index 0000000..bdce780 --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/rt_bench.metal @@ -0,0 +1,67 @@ +// metal/kernels/rt_bench.metal - bandwidth and dispatch-cost probes (rt_bench.cpp). + +// Grid-stride streaming read: thread i reads src[i], src[i + nt], ... (coalesced across the grid). +kernel void k_read(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]], uint nt [[threads_per_grid]]) { + uint4 a = 0; + for (uint i = gid; i < n; i += nt) a += src[i]; + out[gid] = a.x ^ a.y ^ a.z ^ a.w; +} + +// The same, 4 independent loads in flight per iteration. +kernel void k_read_u4(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]], uint nt [[threads_per_grid]]) { + uint4 a = 0, b = 0, c = 0, d = 0; + uint i = gid; + for (; i + 3 * nt < n; i += 4 * nt) { + a += src[i]; + b += src[i + nt]; + c += src[i + 2 * nt]; + d += src[i + 3 * nt]; + } + for (; i < n; i += nt) a += src[i]; + a += b + c + d; + out[gid] = a.x ^ a.y ^ a.z ^ a.w; +} + +// Contiguous chunk per threadgroup (how a GEMV row block reads): threadgroup g reads [g*chunk, (g+1)*chunk). +kernel void k_read_chunk(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + constant uint& n [[buffer(2)]], constant uint& chunk [[buffer(3)]], + uint tid [[thread_position_in_threadgroup]], uint tgs [[threads_per_threadgroup]], + uint g [[threadgroup_position_in_grid]], uint gid [[thread_position_in_grid]]) { + uint4 a = 0, b = 0; + const uint base = g * chunk, end = min(base + chunk, n); + uint i = base + tid; + for (; i + tgs < end; i += 2 * tgs) { + a += src[i]; + b += src[i + tgs]; + } + for (; i < end; i += tgs) a += src[i]; + a += b; + out[gid] = a.x ^ a.y ^ a.z ^ a.w; +} + +// Read + write (a state update): dst[i] = src[i] * 1.0001 for float4. +kernel void k_copy(device const float4* src [[buffer(0)]], device float4* dst [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]], uint nt [[threads_per_grid]]) { + for (uint i = gid; i < n; i += nt) dst[i] = src[i] * 1.0001f; +} + +// Dispatch-cost probes. +kernel void k_empty(device uint* out [[buffer(0)]], uint gid [[thread_position_in_grid]]) { + if (gid == 0xFFFFFFFFu) out[0] = 1; +} + +// A dependent chain: each dispatch reads what the previous one wrote. +kernel void k_chain(device uint* buf [[buffer(0)]], uint gid [[thread_position_in_grid]]) { + if (gid == 0) buf[0] = buf[0] + 1; +} + +// A small but real kernel: 2560 floats read, 2560 written (an elementwise op on one token). +kernel void k_small(device const float* x [[buffer(0)]], device float* y [[buffer(1)]], + uint gid [[thread_position_in_grid]]) { + if (gid < 2560) y[gid] = x[gid] * 0.5f + y[gid]; +} diff --git a/results/round39/model-ab/baseline-kernels/sample.metal b/results/round39/model-ab/baseline-kernels/sample.metal new file mode 100644 index 0000000..4531ca2 --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/sample.metal @@ -0,0 +1,61 @@ +// metal/kernels/sample.metal - greedy sampling over the LM head's logits: the largest logit, ties to the lowest +// token id (NaN never wins). Two stages: argmax_part (grid (parts, n_tok), 256 threads) writes one (value, id) +// pair per threadgroup; argmax_final (grid (n_tok), 256 threads) reduces them in order. + +struct ArgArgs { + uint n_vocab; + uint parts; +}; + +static inline void better(thread float& bv, thread int& bi, float v, int i) { + if (v > bv || (v == bv && i < bi)) { bv = v; bi = i; } +} + +static inline void tg_argmax(thread float& bv, thread int& bi, threadgroup float* sv, threadgroup int* si, ushort sg, + ushort lane) { +#pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) { + const float ov = simd_shuffle_xor(bv, m); + const int oi = simd_shuffle_xor(bi, m); + better(bv, bi, ov, oi); + } + if (lane == 0) { sv[sg] = bv; si[sg] = bi; } + threadgroup_barrier(mem_flags::mem_threadgroup); + bv = -INFINITY; + bi = 0x7fffffff; +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) better(bv, bi, sv[k], si[k]); +} + +kernel void argmax_part(device const float* logits [[buffer(0)]], device float* pv [[buffer(1)]], + device int* pi [[buffer(2)]], constant ArgArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float sv[8]; + threadgroup int si[8]; + const uint t = tg.y, per = (a.n_vocab + a.parts - 1) / a.parts; + const uint b0 = tg.x * per, b1 = min(b0 + per, a.n_vocab); + device const float* l = logits + t * a.n_vocab; + float bv = -INFINITY; + int bi = 0x7fffffff; + for (uint i = b0 + tid; i < b1; i += 256) { + const float v = l[i]; + if (!isnan(v)) better(bv, bi, v, int(i)); + } + tg_argmax(bv, bi, sv, si, sg, lane); + if (tid == 0) { pv[t * a.parts + tg.x] = bv; pi[t * a.parts + tg.x] = bi; } +} + +kernel void argmax_final(device const float* pv [[buffer(0)]], device const int* pi [[buffer(1)]], + device int* out [[buffer(2)]], constant ArgArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float sv[8]; + threadgroup int si[8]; + const uint t = tg.x; + float bv = -INFINITY; + int bi = 0x7fffffff; + for (uint k = tid; k < a.parts; k += 256) better(bv, bi, pv[t * a.parts + k], pi[t * a.parts + k]); + tg_argmax(bv, bi, sv, si, sg, lane); + if (tid == 0) out[t] = bi; +} diff --git a/results/round39/model-ab/baseline-kernels/units.metal b/results/round39/model-ab/baseline-kernels/units.metal new file mode 100644 index 0000000..25d0761 --- /dev/null +++ b/results/round39/model-ab/baseline-kernels/units.metal @@ -0,0 +1,328 @@ +// metal/kernels/units.metal - the dot of one unit with the window's tokens (after quant.metal). +// +// X::load(xrow, u, xs, xsum) prepares unit u's x slice of one token (xrow = the token's row start), once per unit +// and shared by every weight row the lane reads: the UNIT x values in the order the decoder produces weights, and +// up to two partial sums (for the formats with an offset or a min). Where the codes are read with the mask trick +// (an AND keeps a field in place, so the code arrives multiplied by 2^k), the x values are pre-scaled by 2^-k here: +// exact in float. + +// the threadgroup-memory form (x staged once per threadgroup; the same values) +static inline void load16(threadgroup const float* p, thread float* xs) { + threadgroup const float4* v = (threadgroup const float4*)p; + for (int j = 0; j < 4; ++j) { + const float4 a = v[j]; + xs[4 * j] = a.x; + xs[4 * j + 1] = a.y; + xs[4 * j + 2] = a.z; + xs[4 * j + 3] = a.w; + } +} +static inline void load16(device const float* p, thread float* xs) { + device const float4* v = (device const float4*)p; + for (int j = 0; j < 4; ++j) { + const float4 a = v[j]; + xs[4 * j] = a.x; + xs[4 * j + 1] = a.y; + xs[4 * j + 2] = a.z; + xs[4 * j + 3] = a.w; + } +} + +// ---- prepared x (xprep): XP holds, per token row of n floats, the kinds +// XK_RAW x (every format without the mask trick) +// XK_Q2 x * 4^-(i % 8) (Q2_0: 2-bit fields in place in a 16-bit word) +// XK_P4 x * 2^-(8 (i % 4)) (Q4_K, Q5_K: byte fields in place in a 32-bit word) +// XK_P2 x * 2^-(8 (i % 2)) (Q6_K, Q3_K: byte fields in place in a 16-bit word) +// and S16 = the sum of every 16 consecutive raw x (the formats with an offset or a min). +#define XK_RAW 0 +#define XK_Q2 1 +#define XK_P4 2 +#define XK_P2 3 +template struct XKind { enum { K = XK_RAW }; }; +template <> struct XKind { enum { K = XK_Q2 }; }; +template <> struct XKind { enum { K = XK_P4 }; }; +template <> struct XKind { enum { K = XK_P4 }; }; +template <> struct XKind { enum { K = XK_P2 }; }; +template <> struct XKind { enum { K = XK_P2 }; }; + +// X::load(xp, s16, u, xs, xsum): unit u's prepared x (xp = the token's row of the format's kind, s16 = its sums) +template struct X { + template static inline void load(P xp, P, int u, thread float* xs, thread float* xsum) { + constexpr int UNIT = fmt::UNIT; + for (int k = 0; k < UNIT; k += 16) load16(xp + u * UNIT + k, xs + k); + xsum[0] = xsum[1] = 0.0f; + } +}; +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + load16(xp + 16 * u, xs); + xsum[0] = s16[u]; + xsum[1] = 0.0f; + } +}; +// Q4_K / Q5_K: unit (block b, chunk c, half m): run0 at 256b + 64c + 16m, run1 = run0 + 32 +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + const int p0 = 256 * (u >> 3) + 64 * ((u >> 1) & 3) + 16 * (u & 1); + load16(xp + p0, xs); + load16(xp + p0 + 32, xs + 16); + xsum[0] = s16[p0 >> 4]; + xsum[1] = s16[(p0 >> 4) + 2]; + } +}; +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + X::load(xp, s16, u, xs, xsum); + } +}; +// Q6_K: unit (block b, half n, m, p): run0 at 256b + 128n + 32p + 16m, run1 = run0 + 64 +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + const int p0 = 256 * (u >> 3) + 128 * ((u >> 2) & 1) + 32 * (u & 1) + 16 * ((u >> 1) & 1); + load16(xp + p0, xs); + load16(xp + p0 + 64, xs + 16); + xsum[0] = s16[p0 >> 4]; + xsum[1] = s16[(p0 >> 4) + 4]; + } +}; +// Q5_1: 32 contiguous raw x, the two per-16 sums formed from them (the GEMVs that read Q5_1 - the routed down - get +// h without prepared sums) +template <> struct X { + template static inline void load(P xp, P, int u, thread float* xs, thread float* xsum) { + load16(xp + 32 * u, xs); + load16(xp + 32 * u + 16, xs + 16); + float a = 0.0f, b = 0.0f; + for (int j = 0; j < 16; ++j) { a += xs[j]; b += xs[16 + j]; } + xsum[0] = a; + xsum[1] = b; + } +}; +// Q3_K: 32 contiguous, sums per 16 +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + load16(xp + 32 * u, xs); + load16(xp + 32 * u + 16, xs + 16); + xsum[0] = s16[2 * u]; + xsum[1] = s16[2 * u + 1]; + } +}; + +// xprep: one thread per 16 consecutive values of one token row; XP = [4 kinds][n_tok][n] then S16 [n_tok][n / 16]. +// raw and plane are 16-float aligned: the four planes are written as float4s (mtp.metal's mtp_prepare, bit for bit) +kernel void xprep(device const float* x [[buffer(0)]], device float* xp [[buffer(1)]], constant uint& n [[buffer(2)]], + constant uint& n_tok [[buffer(3)]], uint gid [[thread_position_in_grid]]) { + const uint per = n / 16; + if (gid >= per * n_tok) return; + const uint t = gid / per, i = gid % per; + device const float4* src = (device const float4*)(x + t * n + 16 * i); + const uint plane = n_tok * n; + device float* raw = xp + t * n + 16 * i; + device float4* r0 = (device float4*)raw; + device float4* r1 = (device float4*)(raw + plane); + device float4* r2 = (device float4*)(raw + 2 * plane); + device float4* r3 = (device float4*)(raw + 3 * plane); + const float4 q2lo = float4(1.0f, 0.25f, 0.0625f, 0.015625f), + q2hi = float4(0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f), + p4 = float4(1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f); + float s0 = 0, s1 = 0; +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) { + const float4 v = src[j]; + s0 += v.x; s0 += v.z; + s1 += v.y; s1 += v.w; + r0[j] = v; + r1[j] = v * ((j & 1) ? q2hi : q2lo); + r2[j] = v * p4; + r3[j] = float4(v.x, v.y * 0.00390625f, v.z, v.w * 0.00390625f); + } + xp[4 * plane + t * per + i] = s0 + s1; +} + +// ---- decoders: Dec::run(row, u, q, sc) decodes unit u of a row into q[UNIT] and sc = (d1, d2, m1, m2), so that +// for any token: dot = d1 * sum_{j < UNIT/2} q[j] xs[j] + d2 * sum_{j >= UNIT/2} q[j] xs[j] - (m1 xsum0 + m2 xsum1) +// with xs, xsum from X::load. A unit's weights are decoded once per row and applied to every token. +template struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT, UPB = fmt::QK / UNIT; + const float s = deq(row + (u / UPB) * fmt::BB, u % UPB, q, tab); + sc = float4(s, s, 0.0f, 0.0f); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 2) * 18; + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + const ushort a = p[0], c = p[1]; + for (int j = 0; j < 8; ++j) { + q[j] = float(a & ushort(3u << (2 * j))); + q[8 + j] = float(c & ushort(3u << (2 * j))); + } + const float d = rdh(b); + sc = float4(d, d, d, 0.0f); + } +}; + +// Q2_0 with raw x: q = c - 1 (shift, and, convert, subtract), one scale +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 2) * 18; + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + const ushort a = p[0], c = p[1]; + for (int j = 0; j < 8; ++j) { + q[j] = float((a >> (2 * j)) & 3) - 1.0f; + q[8 + j] = float((c >> (2 * j)) & 3) - 1.0f; + } + const float d = rdh(b); + sc = float4(d, d, 0.0f, 0.0f); + } +}; + +static inline void scale_min_k4(int j, device const uchar* q, thread float& d, thread float& m) { + if (j < 4) { + d = float(q[j] & 63); + m = float(q[j + 4] & 63); + } else { + d = float((q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4)); + m = float((q[j + 4] >> 4) | ((q[j] >> 6) << 4)); + } +} +static inline void bytes4(uint w, thread float* q) { + q[0] = float(w & 0xFFu); + q[1] = float(w & 0xFF00u); + q[2] = float(w & 0xFF0000u); + q[3] = float(w & 0xFF000000u); +} +static inline void bytes2(ushort w, thread float* q) { + q[0] = float(w & 0xFF); + q[1] = float(w & 0xFF00); +} + +template <> struct Dec { // w = d q + m: sum = d (q . x) + m (sum x) + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint* tab) { + device const uchar* b = row + u * 24; + const float d = deq(b, 0, q, tab), m = rdh(b + 2); + sc = float4(d, d, -m, -m); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 144; + const int c = (u >> 1) & 3, m = u & 1; + const float d = rdh(b), dm = rdh(b + 2); + float s1, m1, s2, m2; + scale_min_k4(2 * c, b + 4, s1, m1); + scale_min_k4(2 * c + 1, b + 4, s2, m2); + device const uint* qp = (device const uint*)(b + 16 + 32 * c + 16 * m); + for (int w = 0; w < 4; ++w) { + const uint v = qp[w]; + bytes4(v & 0x0F0F0F0Fu, q + 4 * w); + bytes4((v >> 4) & 0x0F0F0F0Fu, q + 16 + 4 * w); + } + sc = float4(d * s1, d * s2, dm * m1, dm * m2); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 176; + const int c = (u >> 1) & 3, m = u & 1; + const float d = rdh(b), dm = rdh(b + 2); + float s1, m1, s2, m2; + scale_min_k4(2 * c, b + 4, s1, m1); + scale_min_k4(2 * c + 1, b + 4, s2, m2); + device const uint* qp = (device const uint*)(b + 48 + 32 * c + 16 * m); + device const uint* hp = (device const uint*)(b + 16 + 16 * m); + for (int w = 0; w < 4; ++w) { + const uint v = qp[w], h = hp[w]; + bytes4((v & 0x0F0F0F0Fu) | (((h >> (2 * c)) & 0x01010101u) << 4), q + 4 * w); + bytes4(((v >> 4) & 0x0F0F0F0Fu) | (((h >> (2 * c + 1)) & 0x01010101u) << 4), q + 16 + 4 * w); + } + sc = float4(d * s1, d * s2, dm * m1, dm * m2); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 210; + const int n = (u >> 2) & 1, m = (u >> 1) & 1, p = u & 1; + device const ushort* ql = (device const ushort*)(b + 64 * n + 32 * p + 16 * m); + device const ushort* qh = (device const ushort*)(b + 128 + 32 * n + 16 * m); + device const char* scs = (device const char*)(b + 192); + const float d = rdh(b + 208); + for (int w = 0; w < 8; ++w) { + const ushort l = ql[w], h = qh[w]; + bytes2((l & 0x0F0F) | (((h >> (2 * p)) & 0x0303) << 4), q + 2 * w); + bytes2(((l >> 4) & 0x0F0F) | (((h >> (2 * p + 4)) & 0x0303) << 4), q + 16 + 2 * w); + } + const float d1 = d * float(scs[8 * n + m + 2 * p]), d2 = d * float(scs[8 * n + m + 2 * p + 4]); + sc = float4(d1, d2, 32.0f * d1, 32.0f * d2); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 110; + const int n = (u >> 2) & 1, j = u & 3; + // qs and hmask as packed (2-byte aligned) vectors: blocks are 110 bytes, so odd blocks are only 2-aligned + device const packed_ushort4* qv = (device const packed_ushort4*)(b + 32 + 32 * n); + device const packed_ushort4* hv = (device const packed_ushort4*)b; + ushort qq[16], hm[16]; + for (int v = 0; v < 4; ++v) { + const ushort4 a4 = ushort4(qv[v]), h4 = ushort4(hv[v]); + qq[4 * v] = a4.x; qq[4 * v + 1] = a4.y; qq[4 * v + 2] = a4.z; qq[4 * v + 3] = a4.w; + hm[4 * v] = h4.x; hm[4 * v + 1] = h4.y; hm[4 * v + 2] = h4.z; hm[4 * v + 3] = h4.w; + } + device const uchar* s = b + 96; + const int i1 = 8 * n + 2 * j, i2 = i1 + 1; + const int l1 = (i1 < 8 ? s[i1] & 0xF : s[i1 - 8] >> 4) | (((s[8 + (i1 & 3)] >> (2 * (i1 >> 2))) & 3) << 4); + const int l2 = (i2 < 8 ? s[i2] & 0xF : s[i2 - 8] >> 4) | (((s[8 + (i2 & 3)] >> (2 * (i2 >> 2))) & 3) << 4); + const float d = rdh(b + 108); + const int hs = 4 * n + j; + for (int w = 0; w < 8; ++w) { + bytes2(((qq[w] >> (2 * j)) & 0x0303) | (((hm[w] >> hs) & 0x0101) << 2), q + 2 * w); + bytes2(((qq[w + 8] >> (2 * j)) & 0x0303) | (((hm[w + 8] >> hs) & 0x0101) << 2), q + 16 + 2 * w); + } + const float d1 = d * float(l1 - 32), d2 = d * float(l2 - 32); + sc = float4(d1, d2, 4.0f * d1, 4.0f * d2); + } +}; + +template <> struct Dec { // 16 floats a unit (the Unsloth files' router / alpha / beta / inject rows) + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const float4* p = (device const float4*)(row + 64 * u); + for (int k = 0; k < 4; ++k) { + const float4 v = p[k]; + q[4 * k] = v.x; q[4 * k + 1] = v.y; q[4 * k + 2] = v.z; q[4 * k + 3] = v.w; + } + sc = float4(1.0f, 1.0f, 0.0f, 0.0f); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uint4* p = (device const uint4*)(row + 32 * u); + const uint4 a = p[0], c = p[1]; + const uint v[8] = {a.x, a.y, a.z, a.w, c.x, c.y, c.z, c.w}; + for (int k = 0; k < 8; ++k) { + q[2 * k] = as_type(v[k] << 16); + q[2 * k + 1] = as_type(v[k] & 0xFFFF0000u); + } + sc = float4(1.0f, 1.0f, 0.0f, 0.0f); + } +}; + +// acc += unit dot of decoded q / sc with one token's prepared slice +template +static inline float unit_combine(thread const float* q, float4 sc, thread const float* xs, thread const float* xsum) { + constexpr int H = UNIT / 2; + float s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int j = 0; j < H; j += 2) { + s0 = fma(q[j], xs[j], s0); + s1 = fma(q[j + 1], xs[j + 1], s1); + s2 = fma(q[H + j], xs[H + j], s2); + s3 = fma(q[H + j + 1], xs[H + j + 1], s3); + } + return sc.x * (s0 + s1) + sc.y * (s2 + s3) - (sc.z * xsum[0] + sc.w * xsum[1]); +} diff --git a/results/round39/model-ab/binary-sha256.json b/results/round39/model-ab/binary-sha256.json new file mode 100644 index 0000000..322c28a --- /dev/null +++ b/results/round39/model-ab/binary-sha256.json @@ -0,0 +1,4 @@ +{ + "baseline": "40de4f6d591745af62f289bd1e7614e0ab6b45de3eb7e8de33415533e22e1396", + "candidate": "0c704e710d289417bd26d204e458169aad21c9d0a812305f2a70e35576a05c89" +} diff --git a/results/round39/model-ab/build-commands.json b/results/round39/model-ab/build-commands.json new file mode 100644 index 0000000..d490f10 --- /dev/null +++ b/results/round39/model-ab/build-commands.json @@ -0,0 +1,67 @@ +{ + "compile": [ + "/usr/bin/clang++", + "-DSTRATA_GPU_METAL=1", + "-DSTRATA_NATIVE_EXPERTS=1", + "-DSTRATA_VERSION=\"0.1.44.1-metal\"", + "-I/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/include", + "-I/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/src/core", + "-I/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/third_party/ggml", + "-I/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/llama.cpp/ggml/include", + "-I/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/llama.cpp/ggml/src", + "-I/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/llama.cpp/ggml/src/ggml-cpu", + "-I/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/metal/include/metal_compat", + "-I/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/metal/include", + "-I/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/llama.cpp/ggml/src/../include", + "-isystem", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/third_party/metal-cpp", + "-O2", + "-std=gnu++20", + "-arch", + "arm64", + "-isysroot", + "/Library/Developer/CommandLineTools/SDKs/MacOSX.sdk", + "-mmacosx-version-min=14.0", + "-Wno-unused-parameter", + "-Wno-deprecated-declarations", + "-MD", + "-MT", + "CMakeFiles/strata_metal_host.dir/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/src/kernels/ngram.cpp.o", + "-MF", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/ngram.d", + "-o", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/ngram.cpp.o", + "-c", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/ngram-baseline.cpp" + ], + "link": [ + "/usr/bin/clang++", + "-O2", + "-arch", + "arm64", + "-isysroot", + "/Library/Developer/CommandLineTools/SDKs/MacOSX.sdk", + "-mmacosx-version-min=14.0", + "-Wl,-search_paths_first", + "-Wl,-headerpad_max_install_names", + "CMakeFiles/strata_window_cli.dir/src/strata_window_cli.cpp.o", + "-o", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/baseline-cli", + "libstrata_metal_kernels.a", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/libstrata_metal_host.a", + "libstrata_gpu_metal.a", + "-framework", + "QuartzCore", + "ggml/src/libggml-cpu.a", + "-framework", + "Accelerate", + "ggml/src/libggml-base.a", + "-lm", + "libmtl_rt.a", + "-framework", + "Metal", + "-framework", + "Foundation" + ], + "note": "identical current objects and link flags except replacing ngram.cpp.o with pre-round39 source; baseline shader snapshots restore all round39 shader changes" +} diff --git a/results/round39/model-ab/candidate-kernels/alu_bench.metal b/results/round39/model-ab/candidate-kernels/alu_bench.metal new file mode 100644 index 0000000..e34890a --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/alu_bench.metal @@ -0,0 +1,181 @@ +// metal/kernels/alu_bench.metal - instruction-rate probes (alu_bench.cpp): each thread runs ITER iterations of 8 +// independent chains of one operation; the result is written so nothing is eliminated. +#define ITER 1024 + +kernel void a_fma32(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + float a0 = gid, a1 = gid + 1, a2 = gid + 2, a3 = gid + 3, a4 = gid + 4, a5 = gid + 5, a6 = gid + 6, a7 = gid + 7; + for (int i = 0; i < ITER; ++i) { + a0 = fma(a0, k, 0.5f); a1 = fma(a1, k, 0.5f); a2 = fma(a2, k, 0.5f); a3 = fma(a3, k, 0.5f); + a4 = fma(a4, k, 0.5f); a5 = fma(a5, k, 0.5f); a6 = fma(a6, k, 0.5f); a7 = fma(a7, k, 0.5f); + } + out[gid] = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7; +} + +kernel void a_fma16(device float* out [[buffer(0)]], constant float& kf [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + const half k = half(kf); + half a0 = gid & 7, a1 = 1, a2 = 2, a3 = 3, a4 = 4, a5 = 5, a6 = 6, a7 = 7; + for (int i = 0; i < ITER; ++i) { + a0 = fma(a0, k, half(0.5)); a1 = fma(a1, k, half(0.5)); a2 = fma(a2, k, half(0.5)); a3 = fma(a3, k, half(0.5)); + a4 = fma(a4, k, half(0.5)); a5 = fma(a5, k, half(0.5)); a6 = fma(a6, k, half(0.5)); a7 = fma(a7, k, half(0.5)); + } + out[gid] = float(a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7); +} + +// float accumulator, half operands (does the conversion come free?) +kernel void a_fmamix(device float* out [[buffer(0)]], constant float& kf [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + half h0 = half(gid & 7), h1 = h0 + half(1); + float a0 = gid, a1 = 1, a2 = 2, a3 = 3, a4 = 4, a5 = 5, a6 = 6, a7 = 7; + for (int i = 0; i < ITER; ++i) { + a0 = fma(float(h0), kf, a0); a1 = fma(float(h1), kf, a1); a2 = fma(float(h0), kf, a2); a3 = fma(float(h1), kf, a3); + a4 = fma(float(h0), kf, a4); a5 = fma(float(h1), kf, a5); a6 = fma(float(h0), kf, a6); a7 = fma(float(h1), kf, a7); + h0 = h0 * half(0.999); h1 = h1 * half(0.999); + } + out[gid] = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7; +} + +kernel void a_iand(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint a0 = gid, a1 = gid + 1, a2 = gid + 2, a3 = gid + 3, a4 = gid + 4, a5 = gid + 5, a6 = gid + 6, a7 = gid + 7; + for (int i = 0; i < ITER; ++i) { + a0 = (a0 ^ k) + 1; a1 = (a1 ^ k) + 1; a2 = (a2 ^ k) + 1; a3 = (a3 ^ k) + 1; + a4 = (a4 ^ k) + 1; a5 = (a5 ^ k) + 1; a6 = (a6 ^ k) + 1; a7 = (a7 ^ k) + 1; + } + out[gid] = float(a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7); +} + +kernel void a_shift(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint a0 = gid, a1 = gid + 1, a2 = gid + 2, a3 = gid + 3, a4 = gid + 4, a5 = gid + 5, a6 = gid + 6, a7 = gid + 7; + for (int i = 0; i < ITER; ++i) { + a0 = (a0 >> k) | k; a1 = (a1 >> k) | k; a2 = (a2 >> k) | k; a3 = (a3 >> k) | k; + a4 = (a4 >> k) | k; a5 = (a5 >> k) | k; a6 = (a6 >> k) | k; a7 = (a7 >> k) | k; + } + out[gid] = float(a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7); +} + +// int -> float conversion + fma +kernel void a_cvt(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint u = gid; + float a0 = 0, a1 = 0, a2 = 0, a3 = 0, a4 = 0, a5 = 0, a6 = 0, a7 = 0; + for (int i = 0; i < ITER; ++i) { + u += k; + a0 += float(u & 3); a1 += float(u & 12); a2 += float(u & 48); a3 += float(u & 192); + a4 += float(u & 768); a5 += float(u & 3072); a6 += float(u & 12288); a7 += float(u & 49152); + } + out[gid] = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7; +} + +// threadgroup memory loads: random-ish 8-byte reads from a 4 KB table +kernel void a_tgload(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint2 tab[512]; + for (int i = tid; i < 512; i += 256) tab[i] = uint2(i, i * 3); + threadgroup_barrier(mem_flags::mem_threadgroup); + uint u = gid * 2654435761u; + uint2 s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += tab[(u >> 3) & 511]; s1 += tab[(u >> 7) & 511]; s2 += tab[(u >> 11) & 511]; s3 += tab[(u >> 15) & 511]; + u += k; + } + out[gid] = float(s0.x + s1.x + s2.x + s3.x + s0.y + s1.y + s2.y + s3.y); +} + +// select (sign application) +kernel void a_sel(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint u = gid; + float a0 = 1, a1 = 2, a2 = 3, a3 = 4, a4 = 5, a5 = 6, a6 = 7, a7 = 8; + for (int i = 0; i < ITER; ++i) { + u ^= k; + a0 = (u & 1) ? -a0 : a0; a1 = (u & 2) ? -a1 : a1; a2 = (u & 4) ? -a2 : a2; a3 = (u & 8) ? -a3 : a3; + a4 = (u & 16) ? -a4 : a4; a5 = (u & 32) ? -a5 : a5; a6 = (u & 64) ? -a6 : a6; a7 = (u & 128) ? -a7 : a7; + } + out[gid] = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7; +} + +// dot-shaped: acc_i += q_i * x_i with loop-carried accumulators only (q, x vary slowly) +kernel void a_dot32(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + float q[8], x[8], a[8]; + for (int j = 0; j < 8; ++j) { q[j] = float(gid + j); x[j] = k * j; a[j] = 0; } + for (int i = 0; i < ITER; ++i) { + for (int j = 0; j < 8; ++j) a[j] = fma(q[j], x[j], a[j]); + q[0] += 1.0f; + } + float s = 0; + for (int j = 0; j < 8; ++j) s += a[j]; + out[gid] = s; +} +kernel void a_dot16(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + half q[8], x[8], a[8]; + for (int j = 0; j < 8; ++j) { q[j] = half(gid + j); x[j] = half(k * j); a[j] = 0; } + for (int i = 0; i < ITER; ++i) { + for (int j = 0; j < 8; ++j) a[j] = fma(q[j], x[j], a[j]); + q[0] += half(1.0f); + } + float s = 0; + for (int j = 0; j < 8; ++j) s += float(a[j]); + out[gid] = s; +} +kernel void a_dotmix(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + half q[8]; + float x[8], a[8]; + for (int j = 0; j < 8; ++j) { q[j] = half(gid + j); x[j] = k * j; a[j] = 0; } + for (int i = 0; i < ITER; ++i) { + for (int j = 0; j < 8; ++j) a[j] = fma(float(q[j]), x[j], a[j]); + q[0] += half(1.0f); + } + float s = 0; + for (int j = 0; j < 8; ++j) s += a[j]; + out[gid] = s; +} +kernel void a_fma32x16(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + float a[16]; + for (int j = 0; j < 16; ++j) a[j] = gid + j; + for (int i = 0; i < ITER / 2; ++i) + for (int j = 0; j < 16; ++j) a[j] = fma(a[j], k, 0.5f); + float s = 0; + for (int j = 0; j < 16; ++j) s += a[j]; + out[gid] = s; +} + +// 8-byte gathers from a 4 KB table in device memory (L1-resident) and in the constant address space +kernel void a_devload(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], device const uint2* tab [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + uint u = gid * 2654435761u; + uint2 s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += tab[(u >> 3) & 511]; s1 += tab[(u >> 7) & 511]; s2 += tab[(u >> 11) & 511]; s3 += tab[(u >> 15) & 511]; + u += k; + } + out[gid] = float(s0.x + s1.x + s2.x + s3.x + s0.y + s1.y + s2.y + s3.y); +} +kernel void a_cstload(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], constant uint2* tab [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + uint u = gid * 2654435761u; + uint2 s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += tab[(u >> 3) & 511]; s1 += tab[(u >> 7) & 511]; s2 += tab[(u >> 11) & 511]; s3 += tab[(u >> 15) & 511]; + u += k; + } + out[gid] = float(s0.x + s1.x + s2.x + s3.x + s0.y + s1.y + s2.y + s3.y); +} +// 4-byte gathers from a 1 KB threadgroup table (IQ3_XXS's grid) +kernel void a_tgload4(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint tab[256]; + for (int i = tid; i < 256; i += 256) tab[i] = i * 3; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint u = gid * 2654435761u; + uint s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += tab[(u >> 3) & 255]; s1 += tab[(u >> 7) & 255]; s2 += tab[(u >> 11) & 255]; s3 += tab[(u >> 15) & 255]; + u += k; + } + out[gid] = float(s0 + s1 + s2 + s3); +} +// popcount parity (computed ksigns) +kernel void a_popc(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint u = gid, s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += popcount(u & 127); s1 += popcount((u >> 7) & 127); s2 += popcount((u >> 14) & 127); s3 += popcount((u >> 21) & 127); + u += k; + } + out[gid] = float(s0 + s1 + s2 + s3); +} diff --git a/results/round39/model-ab/candidate-kernels/dense_mm.metal b/results/round39/model-ab/candidate-kernels/dense_mm.metal new file mode 100644 index 0000000..ccf06f7 --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/dense_mm.metal @@ -0,0 +1,101 @@ +// metal/kernels/dense_mm.metal - dense projections for a PROMPT chunk (after quant/units/mma.metal): +// act_f16 X[t][k] -> half (the matrix units' operand) +// dq_f16 a quantized matrix -> half rows (ref_quant's values rounded to fp16), once per chunk +// gemm_f16 Y[t][n] = Xh[t] . Wh[n] (fp32) by MetalPerformancePrimitives matmul2d, TM x TN tiles, K dynamic; +// rows of Y past n_out / t past n_tok are not written (the tensors carry the extents) +#include +using namespace mpp::tensor_ops; + +kernel void act_f16(device const float* x [[buffer(0)]], device half* xh [[buffer(1)]], constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + const uint n4 = n >> 2; + if (gid < n4) { + ((device half4*)xh)[gid] = half4(((device const float4*)x)[gid]); + } else if (gid == n4) { + for (uint i = n4 << 2; i < n; ++i) xh[i] = half(x[i]); + } +} + +struct DqArgs { + uint n_in; // values per row + uint n_out; // rows + uint row_bytes; + uint perm; // 0: row r -> r; p > 0: row r = c * (n_out / p) + d -> d * p + c (the HC up matrix with the p + // streams of an output interleaved, for gemm_hc_up's epilogue) +}; + +// one thread per unit: grid ((units per row + nth.x - 1) / nth.x, n_out), nth.x threads +template +kernel void dq_f16(device const uchar* W [[buffer(0)]], device half* Wh [[buffer(1)]], constant DqArgs& a [[buffer(2)]], + uint3 gid [[thread_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort3 nth [[threads_per_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, H = UNIT / 2; + threadgroup uint tab[fmt::TAB]; + if (F == T_IQ4_XS || F == T_IQ4_NL) { + if (tid < 64) { + const uint r0 = tid & 15u, sh = (tid >> 4) * 8u; + const half h0 = half(as_type(uchar(((r0 & 8u) ? 0x7159453526190d01ul : 0xf6eaddcfbfad9881ul) >> ((r0 & 7u) * 8u)))); + tab[tid] = as_type(half2(h0, half(as_type(uchar(0xbfad9881u >> sh))))); + tab[tid + 64] = as_type(half2(h0, half(as_type(uchar(0xf6eaddcfu >> sh))))); + tab[tid + 128] = as_type(half2(h0, half(as_type(uchar(0x26190d01u >> sh))))); + tab[tid + 192] = as_type(half2(h0, half(as_type(uchar(0x71594535u >> sh))))); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } else if (fmt::TAB > 1) { + load_tab(tab, tid, nth.x); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int u = int(gid.x); + if (u * UNIT >= int(a.n_in)) return; + uint r = gid.y, ro = r; + if (a.perm) { + if (fmt::BB * 16 < fmt::QK * 17) { + r = a.perm == 4 ? (ro & 3u) * (a.n_out >> 2) + (ro >> 2) : (ro % a.perm) * (a.n_out / a.perm) + ro / a.perm; + } else { + const uint q = a.n_out / a.perm; + ro = (r % q) * a.perm + r / q; + } + } + float w[UNIT]; + int p0, p1; + DW::run(W + ulong(r) * a.row_bytes, u, 0, w, p0, p1, tab); + device half* o = Wh + ulong(ro) * a.n_in; +#pragma clang loop unroll(full) + for (int j = 0; j < H; j += 4) { + *(device half4*)(o + p0 + j) = half4(half(w[j]), half(w[j + 1]), half(w[j + 2]), half(w[j + 3])); + *(device half4*)(o + p1 + j) = half4(half(w[H + j]), half(w[H + j + 1]), half(w[H + j + 2]), half(w[H + j + 3])); + } +} + +struct GemmArgs { int M, N, K; int pad; }; // Y is M rows of N floats + +template +kernel void gemm_f16(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device float* Y [[buffer(2)]], + constant GemmArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]]) { + tensor, tensor_inline> tA(A, dextents(a.K, a.M)); + tensor, tensor_inline> tB(B, dextents(a.K, a.N)); + tensor, tensor_inline> tY(Y, dextents(a.N, a.M)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, int(tg.y) * TM); + auto mB = tB.slice(0, int(tg.x) * TN); + auto mY = tY.slice(int(tg.x) * TN, int(tg.y) * TM); + op.run(mA, mB, mY); +} + +#define DQ(F) template [[host_name("dq_f16_" #F)]] kernel decltype(dq_f16) dq_f16; +#define DQ_(F) DQ(F) +DQ_(T_F32) DQ_(T_BF16) DQ_(T_Q8_0) DQ_(T_Q4_0) DQ_(T_Q5_0) DQ_(T_Q5_1) DQ_(T_IQ4_NL) DQ_(T_IQ4_XS) DQ_(T_Q2_0) DQ_(T_IQ3_S) DQ_(T_Q3_K) +DQ_(T_Q4_K) DQ_(T_Q5_K) DQ_(T_Q6_K) +template [[host_name("gemm_f16_64_64_4")]] kernel decltype(gemm_f16<64, 64, 4>) gemm_f16<64, 64, 4>; +template [[host_name("gemm_f16_128_64_4")]] kernel decltype(gemm_f16<128, 64, 4>) gemm_f16<128, 64, 4>; +template [[host_name("gemm_f16_64_128_4")]] kernel decltype(gemm_f16<64, 128, 4>) gemm_f16<64, 128, 4>; +template [[host_name("gemm_f16_64_32_4")]] kernel decltype(gemm_f16<64, 32, 4>) gemm_f16<64, 32, 4>; +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>; +template [[host_name("gemm_f16_16_64_2")]] kernel decltype(gemm_f16<16, 64, 2>) gemm_f16<16, 64, 2>; +template [[host_name("gemm_f16_32_64_2")]] kernel decltype(gemm_f16<32, 64, 2>) gemm_f16<32, 64, 2>; +template [[host_name("gemm_f16_32_128_2")]] kernel decltype(gemm_f16<32, 128, 2>) gemm_f16<32, 128, 2>; +template [[host_name("gemm_f16_64_32_2")]] kernel decltype(gemm_f16<64, 32, 2>) gemm_f16<64, 32, 2>; +template [[host_name("gemm_f16_16_96_2")]] kernel decltype(gemm_f16<16, 96, 2>) gemm_f16<16, 96, 2>; diff --git a/results/round39/model-ab/candidate-kernels/gdn.metal b/results/round39/model-ab/candidate-kernels/gdn.metal new file mode 100644 index 0000000..c057067 --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/gdn.metal @@ -0,0 +1,197 @@ +// metal/kernels/gdn.metal - one GDN (gated delta-net) layer's recurrence for a window of tokens, fused end to end: +// +// conv (4 taps, history [channels][3] oldest first) -> SiLU, on q (16 heads), k (16), v (48 heads) of 128 +// q, k: L2 norm as native_gdn_l2_norm (rsqrt(sum/128 + eps/128), then x 1/sqrt(128), two roundings) +// gate = exp(softplus(alpha + dt) * a) (softplus = v > 20 ? v : log1p(exp(v))), beta = sigmoid(beta) +// kv_j = sum_i S[i][j] k[i]; delta_j = (v_j - gate * kv_j) * beta; S = gate * S + k[i] delta_j +// o_j = (sum_i S[i][j] q[i]) / sqrt(128); y_j = o_j * rsqrt(mean(o^2) + eps) * gamma_j * sigmoid(z_j) +// +// State layout as Strata's: S[(i * h_v + h) * 128 + j]. Value head h reads q/k head h % h_k. +// One threadgroup per value head, 512 threads = 16 simdgroups: simdgroup s holds columns j = 32 (s % 4) + lane and +// rows i = 32 (s / 4) .. + 32 of the head's 128 x 128 state in registers. GDN_NTOK tokens run in order; with +// `commit` the state and the conv history are written back after the last one (a verify window runs without commit, +// the fold of m accepted tokens is the same kernel with GDN_NTOK = m and commit). The conv history is read from +// `hist` and the committed one written to `hist_out` (ping-pong: the q/k channels are shared by three threadgroups). + +constant int GDN_NTOK [[function_constant(2)]]; + +struct GdnArgs { + uint h_k; // 16 + uint h_v; // 48 + uint channels; // 10240 = (2 h_k + h_v) * 128 + uint commit; // 0: verify (no write); 1: commit after the last token; 0x100 | k (1 <= k <= GDN_NTOK): commit + // the state and the conv history after the first k tokens, in place, while every token's y is + // still computed (a window whose first k rows are kept) - bitwise the commit of a k-token step + float eps; // 1e-6 + uint plane; // 0: y only; else y is the raw plane of prepared activations (xprep layout, n = h_v * 128, + // `plane` floats between planes) and the other planes and the per-16 sums are written too +}; + +constant float gdn_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, + 0.00006103515625f}; +constant float gdn_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + +// The kernel is launched with IB * 128 threads (512 for gdn_step, 1024 for gdn_step8). Without the attribute the compiler +// is free to use as many registers as the 128-row state needs and a pipeline can end up with +// maxTotalThreadsPerThreadgroup below the launch size: an M2 Max gave 448 for the GDN_NTOK >= 2 specializations, and +// Metal then skips the dispatch without an error (#2323). The attribute makes the compiler fit the launch size. +template +[[max_total_threads_per_threadgroup(IB * 128)]] +kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[buffer(1)]], + device const float* qkv [[buffer(2)]], device const float* alpha [[buffer(3)]], + device const float* beta [[buffer(4)]], device const float* z [[buffer(5)]], + device const float* conv_w [[buffer(6)]], device const float* dt [[buffer(7)]], + device const float* ssm_a [[buffer(8)]], device const float* gamma [[buffer(9)]], + device float* y [[buffer(10)]], constant GdnArgs& a [[buffer(11)]], + device float* hist_out [[buffer(12)]], + 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; + threadgroup float qs[S], ks[S], vs[S]; + constexpr int RPT = S / IB; // state rows per thread + // 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 - 1][S]; + threadgroup float part[8], part2[4]; + threadgroup float gb[2]; + 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; + const float scale = 1.0f / sqrt(float(S)); + + // the channel this thread convolves: 0..127 q of head hk, 128..255 k of head hk, 256..383 v of head h + int ch = -1; + if (tid < 128) ch = hk * S + tid; + 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) { + #pragma clang loop unroll(full) + for (int k = 0; k < 3; ++k) hst[k] = hist[ch * 3 + k]; + const float4 cwv = *(device const float4*)(conv_w + ch * 4); + cw[0] = cwv.x; cw[1] = cwv.y; cw[2] = cwv.z; cw[3] = cwv.w; + } + const float gam = ib == 0 ? gamma[j] : 0.0f; + const float dth = tid == 384 ? dt[h] : 0.0f; + const float ssm_ah = tid == 384 ? ssm_a[h] : 0.0f; + float st[RPT]; + const int st_stride = hv * S; + const int st_off = (i0 * hv + h) * S + j; + #pragma clang loop unroll(full) + for (int r = 0; r < RPT; ++r) st[r] = (state + st_off)[r * st_stride]; + + 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 vals[4] = {hst[0], hst[1], hst[2], in}; + float s = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 4; ++k) s = fma(vals[k], cw[k], s); + s = s + 0.0f; + cv = s / (1.0f + exp(-s)); + hst[0] = hst[1]; + hst[1] = hst[2]; + hst[2] = in; + if (t + 1 == commit_at && (tid >= 256 || h < int(a.h_k))) { + #pragma clang loop unroll(full) + for (int k = 0; k < 3; ++k) hist_out[ch * 3 + k] = hst[k]; + } + } + in = in_next; + // raw q/k/v to shared memory + L2 norm partials of q (sg 0..3) and k (sg 4..7); gate and beta on idle lanes + if (tid < 256) { + const float ssq = simd_sum(cv * cv); + if (lane == 0) part[sg] = ssq; + if (tid < 128) qs[tid] = cv; + else ks[tid - 128] = cv; + } else if (tid < 384) { + vs[tid - 256] = cv; + } else if (tid == 384) { + const float v = alpha[t * hv + h] + dth; + // 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_ah); + } else if (tid == 416) { + gb[1] = 1.0f / (1.0f + exp(-beta[t * hv + h])); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const float tot_k = (part[4] + part[5]) + (part[6] + part[7]); + const float rsk = rsqrt(tot_k / float(S) + a.eps / float(S)); + // the simdgroup's RPT rows of k and q one a lane, normalized on the fly and broadcast from the lane that holds them + const int qk_idx = i0 + (lane & (RPT - 1)); + const float kl = (rsk * ks[qk_idx]) * scale; + // kv_j + float kv = 0.0f; + #pragma clang loop unroll(full) + for (int r = 0; r < RPT; ++r) kv = fma(st[r], simd_broadcast(kl, ushort(r)), kv); + red[ib][j] = kv; + const float tot_q = (part[0] + part[1]) + (part[2] + part[3]); + const float rsq = rsqrt(tot_q / float(S) + a.eps / float(S)); + const float ql = (rsq * qs[qk_idx]) * scale; + const float g = gb[0], be = gb[1], vj = vs[j]; + threadgroup_barrier(mem_flags::mem_threadgroup); + kv = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < IB; ++k) kv += red[k][j]; + const float delta = (vj - g * kv) * be; + float o = 0.0f; + #pragma clang loop unroll(full) + for (int r = 0; r < RPT; ++r) { + st[r] = fma(g, st[r], simd_broadcast(kl, ushort(r)) * delta); + o = fma(st[r], simd_broadcast(ql, ushort(r)), o); + } + if (ib > 0) red2[ib - 1][j] = o; + if (t + 1 == commit_at) { + // the state and the conv history after token t (registers: st after this token's update, hst after its + // shift); the state is read only before the loop, so in place is safe. q/k history: the threadgroups + // h < h_k own their q/k head (h == hk); v history: every head its own - to hist_out, NOT in place: a q/k + // channel is read by the three threadgroups h % h_k at their start, and a threadgroup may start after its + // owner has finished (seen on the 10-core M4); the caller swaps the two + device float* st_base = state + st_off; + #pragma clang loop unroll(full) + for (int r = 0; r < RPT; ++r) st_base[r * st_stride] = st[r]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // the readout and the output norm: column j's o (only ib == 0 computes readout and writes y) + if (ib == 0) { + #pragma clang loop unroll(full) + for (int k = 0; k < IB - 1; ++k) o += red2[k][j]; + o = o * scale; + const float osq = simd_sum(o * o); + if (lane == 0) part2[cb] = osq; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ib == 0) { + const float tot = (part2[0] + part2[1]) + (part2[2] + part2[3]); + const float rs = rsqrt(tot / float(S) + a.eps); + const float yv = ((rs * o) * gam) * (1.0f / (1.0f + exp(-zz))); + y[t * hv * S + h * S + j] = yv; + if (a.plane) { + device float* raw = y + t * hv * S + h * S + j; + const int jj = j & 15; + raw[a.plane] = yv * gdn_q2s[jj & 7]; + raw[2 * a.plane] = yv * gdn_p4s[jj & 3]; + raw[3 * a.plane] = (jj & 1) ? yv * 0.00390625f : 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; + #pragma clang loop unroll(full) + 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; + } + } + } +} + +template [[host_name("gdn_step")]] kernel decltype(gdn_step_t<4>) gdn_step_t<4>; +template [[host_name("gdn_step8")]] kernel decltype(gdn_step_t<8>) gdn_step_t<8>; diff --git a/results/round39/model-ab/candidate-kernels/gdn_prompt.metal b/results/round39/model-ab/candidate-kernels/gdn_prompt.metal new file mode 100644 index 0000000..c2c4020 --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/gdn_prompt.metal @@ -0,0 +1,703 @@ +// metal/kernels/gdn_prompt.metal - a GDN layer's recurrence for a PROMPT chunk of T tokens, in three passes: the work +// that does not touch the state runs parallel over the tokens; only the state update is a scan, and the columns of a +// head's state are independent in it (column j needs k, q, v_j, the gate and beta), so a lane owns one column. +// +// gdn_pre per (token, head): conv (4 taps; history oldest first) -> SiLU; q, k L2 norm * 1/sqrt(128) (the decode +// step's arithmetic); kq = k . q per key head; gate = exp(softplus(alpha + dt) * a), beta = sigmoid(beta); +// the chunk's last 3 conv inputs -> hist_out +// gdn_scan per (value head, 32 columns x CB simdgroups): the state column in registers (128 floats a lane); +// per token a = S^T k, b = S^T q (old state), delta = (v - g a) beta, o = (g b + delta kq) / sqrt(128) +// (= S_new^T q / sqrt(128)), S = g S + k delta^T; q and k staged per TB tokens in threadgroup memory +// gdn_post per (token, head): y = o * rsqrt(mean(o^2) + eps) * gamma * sigmoid(z) -> fp16 (the out GEMM's operand) +// and, with `f32`, fp32 +// +// State layout as Strata's: S[(i * h_v + h) * 128 + j]. Value head h reads q/k head h % h_k. +// The kernels are new entries (nothing of gdn.metal changes). + +struct GdnPArgs { + uint n_tok; + uint h_k; // 16 + uint h_v; // 48 + uint channels; // 10240 + float eps; // 1e-6 + uint f32; // gdn_post: also write fp32 y +}; + +static inline float gp_conv(device const float* qkv, device const float* hist, device const float* cw, int ch, int t, + int channels) { + const float4 w = *(device const float4*)(cw + ch * 4); + float s = 0.0f; +#pragma clang loop unroll(full) + for (int k = 0; k < 4; ++k) { + const int src = t - 3 + k; + const float x = src >= 0 ? qkv[src * channels + ch] : hist[ch * 3 + (src + 3)]; + s = fma(x, w[k], s); + } + s = s + 0.0f; + return s / (1.0f + exp(-s)); +} + +// grid ((h_k + h_v) / 8, T) = (8, T) or (h_k + h_v / 2, T) = (40, T), 256 threads: when grid_x == (h_k + h_v) / 8, +// each simdgroup (32 lanes) owns one key head (x < h_k / 8: q, k and kq in registers) or one value head (else: vc and +// gb), with zero threadgroup memory and zero barriers; grid_x == h_k + h_v / 2 keeps the legacy grid. +kernel void gdn_pre(device const float* qkv [[buffer(0)]], device const float* hist [[buffer(1)]], + device const float* conv_w [[buffer(2)]], device const float* alpha [[buffer(3)]], + device const float* beta [[buffer(4)]], device const float* dt [[buffer(5)]], + device const float* ssm_a [[buffer(6)]], device float* qn [[buffer(7)]], device float* kn [[buffer(8)]], + device float* vc [[buffer(9)]], device float2* gb [[buffer(10)]], device float* kq [[buffer(11)]], + device float* hist_out [[buffer(12)]], constant GdnPArgs& a [[buffer(13)]], + uint3 tg [[threadgroup_position_in_grid]], uint3 tgs [[threadgroups_per_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128; + const int x = int(tg.x), t = int(tg.y), hk_n = int(a.h_k), hv = int(a.h_v), C = int(a.channels), T = int(a.n_tok); + const bool packed = int(tgs.x) * 8 == hk_n + hv; + // alpha and beta as the halves of one [T][2 h_v] GEMM (beta == alpha + h_v), else each [T][h_v] + const int ab = beta == alpha + hv ? 2 * hv : hv; + const int k_tgs = packed ? hk_n / 8 : hk_n; + if (x < k_tgs) { + const int kh = packed ? x * 8 + int(sg) : (sg == 0 ? x : hk_n); + if (kh >= hk_n) return; + float cv_q[4], cv_k[4]; + if (t >= 3) { + device const float* pq = qkv + (t - 3) * C + kh * S + int(lane); + device const float* pk = qkv + (t - 3) * C + (hk_n + kh) * S + int(lane); +#pragma clang loop unroll(full) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; + const float4 wq = *(device const float4*)(conv_w + ch_q * 4); + const float4 wk = *(device const float4*)(conv_w + ch_k * 4); + device const float* pqm = pq + m * 32; + device const float* pkm = pk + m * 32; + float sq = fma(pqm[0], wq.x, 0.0f); + sq = fma(pqm[C], wq.y, sq); + sq = fma(pqm[2 * C], wq.z, sq); + sq = fma(pqm[3 * C], wq.w, sq) + 0.0f; + cv_q[m] = sq / (1.0f + exp(-sq)); + float sk = fma(pkm[0], wk.x, 0.0f); + sk = fma(pkm[C], wk.y, sk); + sk = fma(pkm[2 * C], wk.z, sk); + sk = fma(pkm[3 * C], wk.w, sk) + 0.0f; + cv_k[m] = sk / (1.0f + exp(-sk)); + } + if (t == T - 1) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; + for (int k = 0; k < 3; ++k) { + hist_out[ch_q * 3 + k] = pq[m * 32 + (k + 1) * C]; + hist_out[ch_k * 3 + k] = pk[m * 32 + (k + 1) * C]; + } + } + } else { +#pragma clang loop unroll(full) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; + cv_q[m] = gp_conv(qkv, hist, conv_w, ch_q, t, C); + cv_k[m] = gp_conv(qkv, hist, conv_w, ch_k, t, C); + } + if (t == T - 1) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; + for (int k = 0; k < 3; ++k) { + const int src = T - 3 + k; + hist_out[ch_q * 3 + k] = src >= 0 ? qkv[src * C + ch_q] : hist[ch_q * 3 + (src + 3)]; + hist_out[ch_k * 3 + k] = src >= 0 ? qkv[src * C + ch_k] : hist[ch_k * 3 + (src + 3)]; + } + } + } + const float4 sq = simd_sum(float4(cv_q[0] * cv_q[0], cv_q[1] * cv_q[1], cv_q[2] * cv_q[2], cv_q[3] * cv_q[3])); + const float4 sk = simd_sum(float4(cv_k[0] * cv_k[0], cv_k[1] * cv_k[1], cv_k[2] * cv_k[2], cv_k[3] * cv_k[3])); + const float tot_q = (sq.x + sq.y) + (sq.z + sq.w), tot_k = (sk.x + sk.y) + (sk.z + sk.w); + const float rs_q = rsqrt(tot_q / float(S) + a.eps / float(S)); + const float rs_k = rsqrt(tot_k / float(S) + a.eps / float(S)); + const float scale = 1.0f / sqrt(float(S)); + const int base = (t * hk_n + kh) * S + int(lane); + float4 prod; +#pragma clang loop unroll(full) + for (int m = 0; m < 4; ++m) { + const float nq = (rs_q * cv_q[m]) * scale, nk = (rs_k * cv_k[m]) * scale; + qn[base + m * 32] = nq; + kn[base + m * 32] = nk; + prod[m] = nk * nq; + } + const float4 pv = simd_sum(prod); + if (lane == 0) kq[t * hk_n + kh] = (pv.x + pv.y) + (pv.z + pv.w); + } else if (packed) { + const int h = (x - k_tgs) * 8 + int(sg); + if (t >= 3) { + const int d = int(lane) * 4, ch = (2 * hk_n + h) * S + d; + device const float4* p = (device const float4*)(qkv + (t - 3) * C + ch); + device const float4* w = (device const float4*)(conv_w + ch * 4); + const float4 w0 = w[0], w1 = w[1], w2 = w[2], w3 = w[3]; + float4 cv = fma(p[0], float4(w0.x, w1.x, w2.x, w3.x), float4(0.0f)); + cv = fma(p[C / 4], float4(w0.y, w1.y, w2.y, w3.y), cv); + cv = fma(p[C / 2], float4(w0.z, w1.z, w2.z, w3.z), cv); + cv = fma(p[3 * C / 4], float4(w0.w, w1.w, w2.w, w3.w), cv); + cv = cv + 0.0f; + cv = cv / (1.0f + exp(-cv)); + *((device float4*)(vc + (t * hv + h) * S + d)) = cv; + if (t == T - 1) + for (int k = 0; k < 3; ++k) { + const float4 v = p[(k + 1) * C / 4]; + for (int m = 0; m < 4; ++m) hist_out[(ch + m) * 3 + k] = v[m]; + } + } else { +#pragma clang loop unroll(full) + for (int m = 0; m < 4; ++m) { + const int d = int(lane) + m * 32; + const int ch = 2 * hk_n * S + h * S + d; + const float cv = gp_conv(qkv, hist, conv_w, ch, t, C); + if (t == T - 1) + for (int k = 0; k < 3; ++k) { + const int src = T - 3 + k; + hist_out[ch * 3 + k] = src >= 0 ? qkv[src * C + ch] : hist[ch * 3 + (src + 3)]; + } + vc[(t * hv + h) * S + d] = cv; + } + } + if (lane == 0) { + const float v = alpha[t * ab + h] + dt[h]; + 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[t * hv + h] = float2(exp(sp * ssm_a[h]), 1.0f / (1.0f + exp(-beta[t * ab + h]))); + } + } else { + const int h = 2 * (x - k_tgs) + (tid >> 7), d = tid & 127; + const int ch = 2 * hk_n * S + h * S + d; + const float cv = gp_conv(qkv, hist, conv_w, ch, t, C); + if (t == T - 1) + for (int k = 0; k < 3; ++k) { + const int src = T - 3 + k; + hist_out[ch * 3 + k] = src >= 0 ? qkv[src * C + ch] : hist[ch * 3 + (src + 3)]; + } + vc[(t * hv + h) * S + d] = cv; + if (d == 0) { + const float v = alpha[t * ab + h] + dt[h]; + 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[t * hv + h] = float2(exp(sp * ssm_a[h]), 1.0f / (1.0f + exp(-beta[t * ab + h]))); + } + } +} + +// grid (h_v * 4 / CB), CB * 32 threads; TB tokens of q and k staged at a time +template +kernel void gdn_scan_t(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]], + device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]], + device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]], + device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]], + 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, NB = 4 / CB; + threadgroup float4 qs[TB][32], ks[TB][32]; + const int h = int(tg.x) / NB, hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k); + const int j = ((int(tg.x) % NB) * CB + sg) * 32 + lane; + const int T = int(a.n_tok); + const float scale = 1.0f / sqrt(float(S)); + float st[S]; +#pragma clang loop unroll(full) + for (int i = 0; i < S; ++i) st[i] = state[(i * hv + h) * S + j]; + for (int t0 = 0; t0 < T; t0 += TB) { + const int nb = min(TB, T - t0); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < nb * 32; idx += CB * 32) { + const int tt = idx >> 5, w = idx & 31; + qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w]; + ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int tt = 0; tt < nb; ++tt) { + const int t = t0 + tt; + const float2 g_b = gb[t * hv + h]; + const float kqv = kq[t * hkn + hk]; + const float v = vc[(t * hv + h) * S + j]; + float ka = 0.0f, kb = 0.0f, qa = 0.0f, qb = 0.0f; +#pragma clang loop unroll(full) + for (int w = 0; w < 32; w += 2) { + const float4 k0 = ks[tt][w], k1 = ks[tt][w + 1], q0 = qs[tt][w], q1 = qs[tt][w + 1]; + ka = fma(st[4 * w], k0.x, ka); ka = fma(st[4 * w + 1], k0.y, ka); + ka = fma(st[4 * w + 2], k0.z, ka); ka = fma(st[4 * w + 3], k0.w, ka); + kb = fma(st[4 * w + 4], k1.x, kb); kb = fma(st[4 * w + 5], k1.y, kb); + kb = fma(st[4 * w + 6], k1.z, kb); kb = fma(st[4 * w + 7], k1.w, kb); + qa = fma(st[4 * w], q0.x, qa); qa = fma(st[4 * w + 1], q0.y, qa); + qa = fma(st[4 * w + 2], q0.z, qa); qa = fma(st[4 * w + 3], q0.w, qa); + qb = fma(st[4 * w + 4], q1.x, qb); qb = fma(st[4 * w + 5], q1.y, qb); + qb = fma(st[4 * w + 6], q1.z, qb); qb = fma(st[4 * w + 7], q1.w, qb); + } + const float g = g_b.x, kv = ka + kb, qv = qa + qb; + const float delta = (v - g * kv) * g_b.y; + o[(t * hv + h) * S + j] = fma(g, qv, delta * kqv) * scale; +#pragma clang loop unroll(full) + for (int w = 0; w < 32; ++w) { + const float4 k4 = ks[tt][w]; + 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); + } + } + } +#pragma clang loop unroll(full) + for (int i = 0; i < S; ++i) state[(i * hv + h) * S + j] = st[i]; +} + +// a probe for the bench (function constant 6): 1 = no per-token exchange (wrong results), 2 = no state update +constant int GDN_PROBE [[function_constant(6)]]; +constant int gdn_probe = is_function_constant_defined(GDN_PROBE) ? GDN_PROBE : 0; + +// the rows of a column split over IB simdgroups (128 / IB rows a lane, fewer registers): grid (h_v * 4), IB * 32 +// threads; lane = column (32 columns a threadgroup); when TB == 0 (single-pass), each simdgroup loads its R rows of q +// and k across its lanes and broadcasts them with simd_broadcast (one initial dot pass at t = 0, one barrier a token, +// only red[2][IB][32] in threadgroup memory); TB > 0 keeps the TB-token threadgroup-staged reference. +template +kernel void gdn_scan_r(device float* state [[buffer(0)]], device const float* qn [[buffer(1)]], + device const float* kn [[buffer(2)]], device const float* vc [[buffer(3)]], + device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]], + device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]], + 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; + threadgroup float2 red[2][IB][32]; + const int g_idx = int(tg.x) / 4, hv = int(a.h_v), hkn = int(a.h_k); + const int hk = (hv == 3 * hkn) ? (g_idx / 3) : (g_idx % hkn); + const int h = (hv == 3 * hkn) ? (hk + (g_idx % 3) * hkn) : g_idx; + const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * R; + const int T = int(a.n_tok); + const float scale = 1.0f / sqrt(float(S)); + const int rs = hv * S, s0 = (i0 * hv + h) * S + j; + float st[R]; +#pragma clang loop unroll(full) + for (int i = 0; i < R; ++i) st[i] = (state + s0)[i * rs]; + if (TB == 0) { + constexpr int C = (R + 31) / 32, R1 = R / C; + const ushort lr = lane & (R1 - 1); + const int ks = hkn * S, k0 = hk * S + i0 + int(lr), v0 = h * S + j; + float kl[C], ql[C]; +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) { + kl[c] = kn[k0 + c * 32]; + ql[c] = qn[k0 + c * 32]; + } + float2 g_b = gb[h]; + float v = vc[v0]; + float ka = 0.0f, qa = 0.0f; +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) +#pragma clang loop unroll(full) + for (int r = 0; r < R1; ++r) { + ka = fma(st[c * 32 + r], simd_broadcast(kl[c], ushort(r)), ka); + qa = fma(st[c * 32 + r], simd_broadcast(ql[c], ushort(r)), qa); + } + for (int t = 0; t < T - 1; ++t) { + const int slot = t & 1, tn = t + 1; + const int kn_off = tn * ks + k0, vn_off = tn * rs + v0; + float kn_l[C], qn_l[C]; +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) { + kn_l[c] = kn[kn_off + c * 32]; + qn_l[c] = qn[kn_off + c * 32]; + } + const float2 g_b_n = gb[tn * hv + h]; + const float v_n = vc[vn_off]; + const float kqv = sg == 0 ? kq[t * hkn + hk] : 0.0f; + float2 s = float2(ka, qa); + if (gdn_probe != 1) { + red[slot][sg][lane] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + s = red[slot][0][lane]; +#pragma clang loop unroll(full) + for (int b = 1; b < IB; ++b) s += red[slot][b][lane]; + } + const float g = g_b.x; + const float delta = (v - g * s.x) * g_b.y; + if (sg == 0) o[t * rs + v0] = fma(g, s.y, delta * kqv) * scale; + ka = 0.0f; + qa = 0.0f; + if (gdn_probe != 2) +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) +#pragma clang loop unroll(full) + for (int r = 0; r < R1; ++r) { + const float kr = simd_broadcast(kl[c], ushort(r)); + const float knr = simd_broadcast(kn_l[c], ushort(r)), qnr = simd_broadcast(qn_l[c], ushort(r)); + st[c * 32 + r] = fma(g, st[c * 32 + r], kr * delta); + ka = fma(st[c * 32 + r], knr, ka); + qa = fma(st[c * 32 + r], qnr, qa); + } +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) kl[c] = kn_l[c]; + g_b = g_b_n; + v = v_n; + } + if (T > 0) { + const int t = T - 1, slot = t & 1; + const float kqv = sg == 0 ? kq[t * hkn + hk] : 0.0f; + float2 s = float2(ka, qa); + if (gdn_probe != 1) { + red[slot][sg][lane] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + s = red[slot][0][lane]; +#pragma clang loop unroll(full) + for (int b = 1; b < IB; ++b) s += red[slot][b][lane]; + } + const float g = g_b.x; + const float delta = (v - g * s.x) * g_b.y; + if (sg == 0) o[t * rs + v0] = fma(g, s.y, delta * kqv) * scale; + if (gdn_probe != 2) +#pragma clang loop unroll(full) + for (int c = 0; c < C; ++c) +#pragma clang loop unroll(full) + for (int r = 0; r < R1; ++r) { + const float kr = simd_broadcast(kl[c], ushort(r)); + st[c * 32 + r] = fma(g, st[c * 32 + r], kr * delta); + } + } + } else { + constexpr int R4 = R / 4, TBS = TB > 0 ? TB : 1; + constexpr bool KC = R4 <= 8; + threadgroup float4 qs[TBS][32], ks[TBS][32]; + threadgroup float vs[TBS][32]; + threadgroup float4 gk[TBS]; + device const float4* qn4_p = (device const float4*)qn; + device const float4* kn4_p = (device const float4*)kn; + for (int t0 = 0; t0 < T; t0 += TBS) { + const int nb = min(TBS, T - t0); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < nb * 32; idx += IB * 32) { + const int tt = idx >> 5, w = idx & 31; + qs[tt][w] = qn4_p[((t0 + tt) * hkn + hk) * 32 + w]; + ks[tt][w] = kn4_p[((t0 + tt) * hkn + hk) * 32 + w]; + vs[tt][w] = vc[((t0 + tt) * hv + h) * S + (j - lane) + w]; + if (w == 0) { + const float2 g_b = gb[(t0 + tt) * hv + h]; + gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float ka = 0.0f, qa = 0.0f; + float4 kc[R4]; +#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; + const float4 g_b = gk[tt]; + const float kqv = g_b.z; + const float v = vs[tt][lane]; + float2 s = float2(ka, qa); + if (gdn_probe != 1) { + red[slot][sg][lane] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + s = red[slot][0][lane]; +#pragma clang loop unroll(full) + for (int b = 1; b < IB; ++b) s += red[slot][b][lane]; + } + 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 = 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); + } + } + } + } + device float* sp = state + s0; +#pragma clang loop unroll(full) + for (int i = 0; i < R; ++i) sp[i * rs] = st[i]; +} + +// gdn_scan_r<8, 16> with the 16 state rows of a lane in four float4 variables (no array): a register-allocation check +kernel void gdn_scan_v(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]], + device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]], + device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]], + device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]], + 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, IB = 8, TB = 16; + threadgroup float4 qs[TB][32], ks[TB][32]; + threadgroup float2 red[2][IB][32]; + threadgroup float vs[TB][32]; + threadgroup float4 gk[TB]; + const int h = int(tg.x) / 4, hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k); + const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * 16, w0 = sg * 4; + const int T = int(a.n_tok); + const float scale = 1.0f / sqrt(float(S)); + device float* sp = state + (i0 * hv + h) * S + j; + const int rs = hv * S; + float4 s0 = float4(sp[0], sp[rs], sp[2 * rs], sp[3 * rs]); + float4 s1 = float4(sp[4 * rs], sp[5 * rs], sp[6 * rs], sp[7 * rs]); + float4 s2 = float4(sp[8 * rs], sp[9 * rs], sp[10 * rs], sp[11 * rs]); + float4 s3 = float4(sp[12 * rs], sp[13 * rs], sp[14 * rs], sp[15 * rs]); + for (int t0 = 0; t0 < T; t0 += TB) { + const int nb = min(TB, T - t0); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < nb * 32; idx += IB * 32) { + const int tt = idx >> 5, w = idx & 31; + qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w]; + ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w]; + vs[tt][w] = vc[((t0 + tt) * hv + h) * S + (j - lane) + w]; + if (w == 0) { + const float2 g_b = gb[(t0 + tt) * hv + h]; + gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int tt = 0; tt < nb; ++tt) { + const int t = t0 + tt, slot = t & 1; + const float4 k0 = ks[tt][w0], k1 = ks[tt][w0 + 1], k2 = ks[tt][w0 + 2], k3 = ks[tt][w0 + 3]; + const float4 q0 = qs[tt][w0], q1 = qs[tt][w0 + 1], q2 = qs[tt][w0 + 2], q3 = qs[tt][w0 + 3]; + const float ka = (dot(s0, k0) + dot(s1, k1)) + (dot(s2, k2) + dot(s3, k3)); + const float qa = (dot(s0, q0) + dot(s1, q1)) + (dot(s2, q2) + dot(s3, q3)); + red[slot][sg][lane] = float2(ka, qa); + const float4 g_b = gk[tt]; + const float v = vs[tt][lane]; + threadgroup_barrier(mem_flags::mem_threadgroup); + float2 r = red[slot][0][lane]; + for (int b = 1; b < IB; ++b) r += red[slot][b][lane]; + const float g = g_b.x; + const float delta = (v - g * r.x) * g_b.y; + if (sg == 0) o[(t * hv + h) * S + j] = fma(g, r.y, delta * g_b.z) * scale; + s0 = fma(float4(g), s0, k0 * delta); + s1 = fma(float4(g), s1, k1 * delta); + s2 = fma(float4(g), s2, k2 * delta); + s3 = fma(float4(g), s3, k3 * delta); + } + } + sp[0] = s0.x; sp[rs] = s0.y; sp[2 * rs] = s0.z; sp[3 * rs] = s0.w; + sp[4 * rs] = s1.x; sp[5 * rs] = s1.y; sp[6 * rs] = s1.z; sp[7 * rs] = s1.w; + sp[8 * rs] = s2.x; sp[9 * rs] = s2.y; sp[10 * rs] = s2.z; sp[11 * rs] = s2.w; + sp[12 * rs] = s3.x; sp[13 * rs] = s3.y; sp[14 * rs] = s3.z; sp[15 * rs] = s3.w; +} + +// grid (h_v / 4, T) or (h_v, T), 128 threads: one simdgroup (32 lanes) per head h computes the RMSNorm and gated SiLU +// over the 128 channels (4 channels a lane) with zero threadgroup memory and zero barriers. +kernel void gdn_post(device const float* o [[buffer(0)]], device const float* z [[buffer(1)]], + device const float* gamma [[buffer(2)]], device half* yh [[buffer(3)]], device float* y [[buffer(4)]], + constant GdnPArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], uint3 tgs [[threadgroups_per_grid]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const int hv = int(a.h_v), t = int(tg.y); + const int h = int(tgs.x) == hv ? (sg == 0 ? int(tg.x) : hv) : int(tg.x) * 4 + int(sg); + if (h >= hv) return; + const uint base = uint(t * hv + h) * 128 + lane; + device const float* o_row = o + base; + device const float* z_row = z + base; + device const float* g_row = gamma + lane; + device half* yh_row = yh + base; + device float* y_row = y + base; + const float ov0 = o_row[0], ov1 = o_row[32], ov2 = o_row[64], ov3 = o_row[96]; + // the gate and gamma do not depend on the norm: computed while the simd_sums run + float sig[4]; +#pragma clang loop unroll(full) + for (uint k = 0; k < 4; ++k) sig[k] = 1.0f / (1.0f + exp(-z_row[k * 32])); + const float gv[4] = {g_row[0], g_row[32], g_row[64], g_row[96]}; + const float4 pv = simd_sum(float4(ov0 * ov0, ov1 * ov1, ov2 * ov2, ov3 * ov3)); + const float tot = (pv.x + pv.y) + (pv.z + pv.w); + const float rs = rsqrt(tot / 128.0f + a.eps); + const float ov[4] = {ov0, ov1, ov2, ov3}; + float yv[4]; +#pragma clang loop unroll(full) + for (uint k = 0; k < 4; ++k) { + yv[k] = ((rs * ov[k]) * gv[k]) * sig[k]; + yh_row[k * 32] = half(yv[k]); + } + if (a.f32) { +#pragma clang loop unroll(full) + for (uint k = 0; k < 4; ++k) y_row[k * 32] = yv[k]; + } +} + +template [[host_name("gdn_scan_1")]] kernel decltype(gdn_scan_t<1, 16>) gdn_scan_t<1, 16>; +template [[host_name("gdn_scan_2")]] kernel decltype(gdn_scan_t<2, 16>) gdn_scan_t<2, 16>; +template [[host_name("gdn_scan_4")]] kernel decltype(gdn_scan_t<4, 16>) gdn_scan_t<4, 16>; +template [[host_name("gdn_scanr_2")]] kernel decltype(gdn_scan_r<2, 0>) gdn_scan_r<2, 0>; +template [[host_name("gdn_scanr_4")]] kernel decltype(gdn_scan_r<4, 0>) gdn_scan_r<4, 0>; +template [[host_name("gdn_scanr_8")]] kernel decltype(gdn_scan_r<8, 0>) gdn_scan_r<8, 0>; +template [[host_name("gdn_scanr_4_8")]] kernel decltype(gdn_scan_r<4, 8>) gdn_scan_r<4, 8>; +template [[host_name("gdn_scanr_4_16")]] kernel decltype(gdn_scan_r<4, 16>) gdn_scan_r<4, 16>; + +// ---- the chunked form (chunks of GC = 32 tokens), the WY / UT form of the gated delta rule. With G_t the product +// of the chunk's gates up to t (log space: lg_t) and S0 the state entering the chunk: +// (I + A) D = B, A_ts = beta_t (k_t . k_s) G_t / G_s (s < t), B = [beta v | beta G k] -> D = [U | W] +// delta = U - W S0, o_t = G_t S0^T q_t + sum_{s <= t} (q_t . k_s) G_t / G_s delta_s, +// S_L = G_L S0 + sum_s G_L / G_s k_s delta_s^T +// (the sequential recurrence's algebra; the rounding differs). Three kernels: +// gdn_cprep per (head, chunk): lg, A, the forward substitution for U and W (a thread per column of B) +// gdn_cstate per (head, 32 state columns): the chunks in order - delta, the inter-chunk output, the state update +// gdn_cout per (head, chunk): o += the intra-chunk part, times 1/sqrt(128) +constant constexpr int GC = 32; + +// grid (h_v, chunks), 256 threads. U, W: [h][chunk][GC][128]; lgo: [h][chunk][GC] +kernel void gdn_cprep(device const float* qn [[buffer(0)]], device const float* kn [[buffer(1)]], + device const float* vc [[buffer(2)]], device const float2* gb [[buffer(3)]], + device float* U [[buffer(4)]], device float* W [[buffer(5)]], device float* lgo [[buffer(6)]], + constant GdnPArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int S = 128; + threadgroup float A[GC][GC]; + threadgroup float lg[GC], bt[GC]; + const int h = int(tg.x), c = int(tg.y), hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok); + const int NC = (T + GC - 1) / GC, t0 = c * GC, nb = min(GC, T - t0); + if (tid == 0) { + float s = 0.0f; + for (int t = 0; t < GC; ++t) { + if (t < nb) s += max(log(gb[(t0 + t) * hv + h].x), -80.0f); + lg[t] = s; + bt[t] = t < nb ? gb[(t0 + t) * hv + h].y : 0.0f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < GC * GC; idx += 256) { + const int t = idx / GC, s = idx % GC; + float v = 0.0f; + if (s < t && t < nb) { + device const float4* kt = (device const float4*)(kn + ((t0 + t) * hkn + hk) * S); + device const float4* ks = (device const float4*)(kn + ((t0 + s) * hkn + hk) * S); + float d = 0.0f; + for (int i = 0; i < S / 4; ++i) d += dot(kt[i], ks[i]); + v = bt[t] * d * exp(lg[t] - lg[s]); + } + A[t][s] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // forward substitution, column j of B a thread (threads 0..127: beta v; 128..255: beta G k) + const int j = tid & 127; + const bool isw = tid >= 128; + float x[GC]; + for (int t = 0; t < GC; ++t) { + float b = 0.0f; + if (t < nb) b = isw ? bt[t] * exp(lg[t]) * kn[((t0 + t) * hkn + hk) * S + j] : bt[t] * vc[((t0 + t) * hv + h) * S + j]; + for (int s = 0; s < t; ++s) b = fma(-A[t][s], x[s], b); + x[t] = b; + } + device float* o = (isw ? W : U) + (ulong(h) * NC + c) * GC * S + j; + for (int t = 0; t < GC; ++t) o[t * S] = x[t]; + if (tid < GC) lgo[(h * NC + c) * GC + tid] = lg[tid]; +} + +// grid (h_v * 4), 256 threads = 8 row blocks of 16 x 32 columns; the state's 128 x 32 block in registers (16 a thread) +kernel void gdn_cstate(device float* state [[buffer(0)]], device const float* qn [[buffer(1)]], + device const float* kn [[buffer(2)]], device const float* U [[buffer(3)]], + device const float* W [[buffer(4)]], device const float* lgo [[buffer(5)]], + device float* delta [[buffer(6)]], device float* o [[buffer(7)]], constant GdnPArgs& a [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128, RB = 16; + threadgroup float red[8][16][32]; + threadgroup float dl[GC][32]; + const int h = int(tg.x) / 4, hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok); + const int NC = (T + GC - 1) / GC; + const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * RB; + float st[RB]; + for (int r = 0; r < RB; ++r) st[r] = state[((i0 + r) * hv + h) * S + j]; + for (int c = 0; c < NC; ++c) { + const int t0 = c * GC, nb = min(GC, T - t0); + device const float* Uc = U + (ulong(h) * NC + c) * GC * S; + device const float* Wc = W + (ulong(h) * NC + c) * GC * S; + device const float* lg = lgo + (h * NC + c) * GC; + // pass 0: delta = U - W S0; pass 1: o_inter = G_t S0^T q_t; 16 tokens at a time + for (int pass = 0; pass < 2; ++pass) { + for (int hf = 0; hf < 2; ++hf) { + float p[16]; + for (int u = 0; u < 16; ++u) { + const int t = hf * 16 + u; + float s = 0.0f; + if (t < nb) { + device const float* v = pass == 0 ? Wc + t * S + i0 : qn + ((t0 + t) * hkn + hk) * S + i0; + for (int r = 0; r < RB; ++r) s = fma(v[r], st[r], s); + } + p[u] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int u = 0; u < 16; ++u) red[sg][u][lane] = p[u]; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int u = sg * 2; u < sg * 2 + 2; ++u) { + const int t = hf * 16 + u; + float s = 0.0f; + for (int b = 0; b < 8; ++b) s += red[b][u][lane]; + if (t < nb) { + if (pass == 0) { + const float d = Uc[t * S + j] - s; + dl[t][lane] = d; + delta[((ulong(h) * NC + c) * GC + t) * S + j] = d; + } else { + o[((t0 + t) * hv + h) * S + j] = exp(lg[t]) * s; + } + } else if (pass == 0) { + dl[t][lane] = 0.0f; + } + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // S = G_L S0 + sum_s G_L / G_s k_s delta_s^T + const float lgl = lg[nb - 1], gl = exp(lgl); + for (int r = 0; r < RB; ++r) st[r] *= gl; + for (int s = 0; s < nb; ++s) { + const float f = exp(lgl - lg[s]) * dl[s][lane]; + device const float* ks = kn + ((t0 + s) * hkn + hk) * S + i0; + for (int r = 0; r < RB; ++r) st[r] = fma(ks[r], f, st[r]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (int r = 0; r < RB; ++r) state[((i0 + r) * hv + h) * S + j] = st[r]; +} + +// grid (h_v, chunks), 128 threads (column j): o_t = (o_inter_t + sum_{s <= t} (q_t . k_s) G_t / G_s delta_s) / sqrt(128) +kernel void gdn_cout(device const float* qn [[buffer(0)]], device const float* kn [[buffer(1)]], + device const float* delta [[buffer(2)]], device const float* lgo [[buffer(3)]], + device float* o [[buffer(4)]], constant GdnPArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int S = 128; + threadgroup float M[GC][GC]; + const int h = int(tg.x), c = int(tg.y), hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok); + const int NC = (T + GC - 1) / GC, t0 = c * GC, nb = min(GC, T - t0); + device const float* lg = lgo + (h * NC + c) * GC; + for (int idx = tid; idx < GC * GC; idx += 128) { + const int t = idx / GC, s = idx % GC; + float v = 0.0f; + if (s <= t && t < nb) { + device const float4* qt = (device const float4*)(qn + ((t0 + t) * hkn + hk) * S); + device const float4* ks = (device const float4*)(kn + ((t0 + s) * hkn + hk) * S); + float d = 0.0f; + for (int i = 0; i < S / 4; ++i) d += dot(qt[i], ks[i]); + v = d * exp(lg[t] - lg[s]); + } + M[t][s] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int j = tid; + device const float* dc = delta + (ulong(h) * NC + c) * GC * S + j; + float dv[GC]; + for (int s = 0; s < GC; ++s) dv[s] = s < nb ? dc[s * S] : 0.0f; + const float scale = 1.0f / sqrt(float(S)); + for (int t = 0; t < nb; ++t) { + float s = o[((t0 + t) * hv + h) * S + j]; + for (int u = 0; u <= t; ++u) s = fma(M[t][u], dv[u], s); + o[((t0 + t) * hv + h) * S + j] = s * scale; + } +} diff --git a/results/round39/model-ab/candidate-kernels/gemv.metal b/results/round39/model-ab/candidate-kernels/gemv.metal new file mode 100644 index 0000000..49b7b62 --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/gemv.metal @@ -0,0 +1,500 @@ +// metal/kernels/gemv.metal - y[t][r] = W[r] . x[t] for a window of 1..8 tokens (dense projections, the shared +// expert, the LM head), any format of quant.metal/units.metal. +// +// Shape of a threadgroup: SG simdgroups; KS of them split the row's units (KS = 1: no split), the other SG/KS cover +// different rows. In a simdgroup, LPR lanes share a row and 32/LPR rows run side by side; each lane holds NR rows. +// The weights of a unit are decoded once and applied to every token. Partial sums of a K split are added in a +// fixed order (deterministic). + +struct GemvArgs { + uint n_in; + uint n_out; + uint n_tok; + uint row_bytes; + uint x_stride; // floats between tokens' x rows + uint y_stride; // floats between tokens' y rows +}; + +constant int NTOK [[function_constant(0)]]; // tokens of the window (1..8): the token loops are specialised +// a bench probe (function constant 9): 1 = tokens past the first reuse token 0's x (no load; wrong values), +// 2 = tokens past the first skip the combine (wrong values) +constant int GEMV_PROBE [[function_constant(9)]]; +constant int gemv_probe = is_function_constant_defined(GEMV_PROBE) ? GEMV_PROBE : 0; + +// Wider raw-dot windows and prepared Q2 at NR=1 did not improve on M5. +template +static inline bool gemv_direct_x() { + return F == T_Q8_0 && (NTOK == 1 || (NTOK == 2 && LPR == 16 && NR == 1)); +} + +// Keep unit_combine's four FMA chains without materialising the raw x slice. +template +static inline void gemv_raw_combine(device const float* xp, thread const float (&q)[NR][UNIT], + thread const float4* sc, thread float (&acc)[NR][8], int t) { + constexpr int H = UNIT / 2; + device const float4* x = (device const float4*)xp; + float s[NR][4]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) s[r][0] = s[r][1] = s[r][2] = s[r][3] = 0.0f; +#pragma clang loop unroll(full) + for (int j = 0; j < H; j += 4) { + const float4 a = x[j / 4], b = x[(H + j) / 4]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + s[r][0] = fma(q[r][j], a.x, s[r][0]); + s[r][1] = fma(q[r][j + 1], a.y, s[r][1]); + s[r][2] = fma(q[r][H + j], b.x, s[r][2]); + s[r][3] = fma(q[r][H + j + 1], b.y, s[r][3]); + s[r][0] = fma(q[r][j + 2], a.z, s[r][0]); + s[r][1] = fma(q[r][j + 3], a.w, s[r][1]); + s[r][2] = fma(q[r][H + j + 2], b.z, s[r][2]); + s[r][3] = fma(q[r][H + j + 3], b.w, s[r][3]); + } + } +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + acc[r][t] += sc[r].x * (s[r][0] + s[r][1]) + sc[r].y * (s[r][2] + s[r][3]) - + (sc[r].z * 0.0f + sc[r].w * 0.0f); +} + +static inline float gemv_q2_raw_combine(ushort2 pk, float d, thread const float* xs) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; j += 2) { + const uint2 q = uint2(pk) >> (2 * j); + const uint2 bits = ((q | (q << 14)) & 0x00030003u) | 0x64006400u; + const float4 v = float4(as_type(bits) - 1025.0h); + s0 = fma(v.x, xs[j], s0); + s1 = fma(v.y, xs[j + 1], s1); + s2 = fma(v.z, xs[8 + j], s2); + s3 = fma(v.w, xs[8 + j + 1], s3); + } + return d * (s0 + s1) + d * (s2 + s3) - 0.0f; +} + +// IQ4_XS kept as its 16 code bytes and scale (5 registers a row instead of 36), decoded inside the combine. M5, weights +// from DRAM: 2560 -> 640 at 16,2,4,1 -1.5..-5% (1 token), -2..-3% (2..4); 3..4 tokens -1..-2% on 2560 -> 6144..12288 +// (but 6144 -> 2560 at 16,2,4,1 3 tokens +1.4%); LPR 32 under 4 tokens +0.4..0.7%, so those keep the decoded unit. +// IQ4_NL (18-byte blocks, ushort reads) measured no change and keeps the decoded unit. +template +static inline bool gemv_iq4_packed() { return LPR <= 16 || NTOK >= 4; } + +static inline uint4 gemv_iq4xs_codes(device const uchar* row, int u, thread float& d) { + device const uchar* b = row + (u >> 3) * 136; + d = iq4xs_scale(b, u & 7); + return uint4(*(device const packed_uint4*)(b + 8 + 16 * (u & 7))); +} + +// unit_combine<32> of deq's q, decoding the codes as it goes: the same four FMA chains in the same order +static inline float gemv_iq4_combine(uint4 pk, float d, thread const float* xs, threadgroup const uint* tab) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const uint v = pk[j >> 1] >> (16 * (j & 1)); + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[(v >> 8) & 255]); + s0 = fma(float(w0.x), xs[2 * j], s0); + s1 = fma(float(w1.x), xs[2 * j + 1], s1); + s2 = fma(float(w0.y), xs[16 + 2 * j], s2); + s3 = fma(float(w1.y), xs[16 + 2 * j + 1], s3); + } + return d * (s0 + s1) + d * (s2 + s3) - 0.0f; +} + +template +kernel void gemv(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RS = SG / KS, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + threadgroup float red[KS > 1 ? KS * RS * RG * NR * 8 : 1]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int ks = sg / RS, rs = sg % RS; + const int li = lane % LPR; + const int lrow0 = (rs * RG + lane / LPR) * NR; + const int row0 = int(tg.x) * RS * RG * NR + lrow0; + const int U = int(a.n_in) / UNIT; + const int u0 = ks * U / KS, u1 = (ks + 1) * U / KS; + float acc[NR][8]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) rp[r] = W + (ulong)min(row0 + r, int(a.n_out) - 1) * a.row_bytes; + const int s16n = int(a.n_in) / 16; + if (gemv_probe == 3) { + // rows outer, tokens inner: one row's decoded unit in registers at a time (x re-read per row from L1) + for (int u = u0 + li; u < u1; u += LPR) +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 sc; + Dec::run(rp[r], u, q, sc, tab); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + acc[r][t] += unit_combine(q, sc, xs, xsum); + } + } + } else + for (int u = u0 + li; u < u1; u += LPR) { + if (NTOK == 1 && F == T_Q2_0R) { + ushort2 pk[NR]; + float d[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + device const uchar* b = rp[r] + (u >> 2) * 18; + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + pk[r] = ushort2(p[0], p[1]); + d[r] = rdh(b); + } + float xs[16]; + load16(xp + 16 * u, xs); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][0] += gemv_q2_raw_combine(pk[r], d[r], xs); + } else if (F == T_IQ4_XS && gemv_iq4_packed()) { + uint4 pk[NR]; + float d[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) pk[r] = gemv_iq4xs_codes(rp[r], u, d[r]); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[32]; + load16(xp + t * int(a.x_stride) + 32 * u, xs); + load16(xp + t * int(a.x_stride) + 32 * u + 16, xs + 16); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += gemv_iq4_combine(pk[r], d[r], xs, tab); + } + } else { + float q[NR][UNIT]; + float4 sc[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); + float xs[UNIT], xsum[2]; +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + if (gemv_direct_x()) { + if (gemv_probe != 2 || t == 0) + gemv_raw_combine(xp + (gemv_probe == 1 ? 0 : t) * int(a.x_stride) + u * UNIT, + q, sc, acc, t); + } else { + if (gemv_probe != 1 || t == 0) X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + if (gemv_probe != 2 || t == 0) +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + } + } + constexpr int RPT = RS * RG * NR; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + if (KS == 1) { + if (li == 0 && row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = v; + } else if (li == 0) { + red[(ks * RPT + lrow0 + r) * 8 + t] = v; + } + } + if (KS > 1) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ks == 0 && li == 0) +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + if (row0 + r < int(a.n_out)) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float v = 0.0f; +#pragma clang loop unroll(full) + for (int k = 0; k < KS; ++k) v += red[(k * RPT + lrow0 + r) * 8 + t]; + y[t * int(a.y_stride) + row0 + r] = v; + } + } +} + +#define GV(F, L, R, S, K) template [[host_name("gemv_" #F "_" #L "_" #R "_" #S "_" #K)]] kernel decltype(gemv) gemv; +#define GV_SET(F) GV(F, 32, 1, 4, 1) GV(F, 32, 2, 4, 1) GV(F, 16, 1, 4, 1) GV(F, 16, 2, 4, 1) GV(F, 8, 1, 4, 1) \ + GV(F, 8, 2, 4, 1) GV(F, 4, 2, 4, 1) GV(F, 4, 1, 4, 1) GV(F, 32, 1, 4, 4) GV(F, 32, 2, 4, 2) \ + GV(F, 32, 1, 8, 8) GV(F, 16, 1, 8, 4) GV(F, 32, 1, 8, 2) GV(F, 8, 1, 8, 1) GV(F, 1, 1, 4, 1) GV(F, 1, 2, 4, 1) GV(F, 2, 1, 4, 1) GV(F, 2, 2, 4, 1) GV(F, 1, 1, 2, 1) GV(F, 4, 4, 4, 1) +#ifndef GEMV_TYPES +#define GEMV_TYPES GV_SET(T_F32) GV_SET(T_BF16) GV_SET(T_Q8_0) GV_SET(T_Q4_0) GV_SET(T_Q5_0) GV_SET(T_IQ4_NL) GV_SET(T_IQ4_XS) \ + GV_SET(T_Q2_0) GV_SET(T_IQ3_S) GV_SET(T_Q3_K) GV_SET(T_Q4_K) GV_SET(T_Q5_K) GV_SET(T_Q6_K) +#endif +GEMV_TYPES + +// The shared expert's gate and up in one dispatch: h[t][r] = silu(Wg[r] . x[t]) * (Wu[r] . x[t]), each with its own +// format (and so its own prepared activation kind); h is written raw (its consumer reads raw activations). +// KS = 1; tiles as gemv. Buffers: Wg, xg (Wg's kind), Wu, xu (Wu's kind), s16, h, args (row_bytes = Wg's), rb_up. +template +kernel void gemv_pair(device const uchar* Wg [[buffer(0)]], device const float* xg [[buffer(1)]], + device const uchar* Wu [[buffer(2)]], device const float* xu [[buffer(3)]], + device const float* s16 [[buffer(4)]], device float* h [[buffer(5)]], + constant GemvArgs& a [[buffer(6)]], constant uint& rb_up [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int U1 = fmt::UNIT, U2 = fmt::UNIT, RG = 32 / LPR; + constexpr int TB = fmt::TAB > fmt::TAB ? fmt::TAB : fmt::TAB; + threadgroup uint tab1[fmt::TAB]; + threadgroup uint tab2[fmt::TAB]; + if (fmt::TAB > 1) load_tab(tab1, tid, SG * 32); + if (fmt::TAB > 1) load_tab(tab2, tid, SG * 32); + if (TB > 1) threadgroup_barrier(mem_flags::mem_threadgroup); + const int li = lane % LPR; + const int row = (int(tg.x) * SG + sg) * RG + lane / LPR; + const int rr = min(row, int(a.n_out) - 1); + device const uchar* rg = Wg + ulong(rr) * a.row_bytes; + device const uchar* ru = Wu + ulong(rr) * rb_up; + const int s16n = int(a.n_in) / 16; + float ag[8], au[8]; +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) ag[t] = au[t] = 0.0f; + for (int u = li; u < int(a.n_in) / U1; u += LPR) { + float q[U1]; + float4 sc; + Dec::run(rg, u, q, sc, tab1); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[U1], xsum[2]; + X::load(xg + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + ag[t] += unit_combine(q, sc, xs, xsum); + } + } + for (int u = li; u < int(a.n_in) / U2; u += LPR) { + float q[U2]; + float4 sc; + Dec::run(ru, u, q, sc, tab2); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[U2], xsum[2]; + X::load(xu + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + au[t] += unit_combine(q, sc, xs, xsum); + } + } +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float gv = lane_sum(ag[t]), uv = lane_sum(au[t]); + if (li == 0 && row < int(a.n_out)) h[t * int(a.y_stride) + row] = gv / (1.0f + exp(-gv)) * uv; + } +} + +#define GP_(F1, F2) template [[host_name("gemv_pair_" #F1 "_" #F2)]] kernel decltype(gemv_pair) gemv_pair; +#define GP(F1, F2) GP_(F1, F2) +#define GP_ROW(F1) GP(F1, T_IQ3_S) GP(F1, T_IQ4_XS) GP(F1, T_Q2_0) GP(F1, T_Q3_K) GP(F1, T_Q4_K) GP(F1, T_Q5_K) GP(F1, T_Q6_K) +GP_ROW(T_IQ3_S) GP_ROW(T_IQ4_XS) GP_ROW(T_Q2_0) GP_ROW(T_Q3_K) GP_ROW(T_Q4_K) GP_ROW(T_Q5_K) GP_ROW(T_Q6_K) +// Q2_0 with raw activations (the shared expert's down reading the pair's raw h) +GV_SET(T_Q2_0R) + +// gemv_cat: two matrices of the same format and n_in over the same x in one dispatch - rows [0, n_split) are W's +// (into y), rows [n_split, n_out) are W2's (into y2, from its row 0). Each row's arithmetic is gemv's (KS = 1), so the +// outputs are bitwise the two separate GEMVs'. (GDN's ssm_alpha + ssm_beta: 48 + 48 rows.) +// Buffers: W, xp, s16, y, args (n_out = both), W2, y2, n_split. +template +kernel void gemv_cat(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], + device const uchar* W2 [[buffer(5)]], device float* y2 [[buffer(6)]], constant uint& n_split [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int li = lane % LPR; + const int row0 = int(tg.x) * SG * RG * NR + (sg * RG + lane / LPR) * NR; + const int U = int(a.n_in) / UNIT, sp = int(n_split); + float acc[NR][8]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + const int row = min(row0 + r, int(a.n_out) - 1); + rp[r] = row < sp ? W + (ulong)row * a.row_bytes : W2 + (ulong)(row - sp) * a.row_bytes; + } + const int s16n = int(a.n_in) / 16; + for (int u = li; u < U; u += LPR) { + float q[NR][UNIT]; + float4 sc[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + if (gemv_direct_x()) { + gemv_raw_combine(xp + t * int(a.x_stride) + u * UNIT, q, sc, acc, t); + } else { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + } +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + const int row = row0 + r; + if (li == 0 && row < int(a.n_out)) { + if (row < sp) y[t * int(a.y_stride) + row] = v; + else y2[t * int(a.y_stride) + row - sp] = v; + } + } +} +#define GVC_(F, L, R, S) template [[host_name("gemvc_" #F "_" #L "_" #R "_" #S)]] kernel decltype(gemv_cat) gemv_cat; +#define GVC(F, L, R, S) GVC_(F, L, R, S) +GVC(T_BF16, 32, 1, 4) GVC(T_BF16, 32, 1, 8) GVC(T_BF16, 16, 1, 4) GVC(T_BF16, 32, 2, 4) +GVC(T_F32, 32, 1, 4) GVC(T_F32, 32, 1, 8) GVC(T_F32, 16, 1, 4) GVC(T_F32, 32, 2, 4) +GVC(T_Q8_0, 32, 1, 4) GVC(T_Q8_0, 32, 1, 8) GVC(T_Q8_0, 16, 1, 4) GVC(T_Q8_0, 32, 2, 4) GVC(T_Q8_0, 8, 1, 8) + +// gemv_catk: gemv_cat with gemv's K split (KS simdgroups share a row, partial sums added in gemv's fixed order), so +// the outputs are bitwise gemv_____'s. (The MTP draft step's k + v and shared gate + up: Q8_0 +// 2560 -> 512 + 512 / 640 + 640 at 32,1,4,4 on the M5.) Buffers as gemv_cat. +template +kernel void gemv_catk(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], + device const uchar* W2 [[buffer(5)]], device float* y2 [[buffer(6)]], constant uint& n_split [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RS = SG / KS, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + threadgroup float red[KS * RS * RG * NR * 8]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int ks = sg / RS, rs = sg % RS; + const int li = lane % LPR; + const int lrow0 = (rs * RG + lane / LPR) * NR; + const int row0 = int(tg.x) * RS * RG * NR + lrow0; + const int U = int(a.n_in) / UNIT, sp = int(n_split); + const int u0 = ks * U / KS, u1 = (ks + 1) * U / KS; + float acc[NR][8]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + const int row = min(row0 + r, int(a.n_out) - 1); + rp[r] = row < sp ? W + (ulong)row * a.row_bytes : W2 + (ulong)(row - sp) * a.row_bytes; + } + const int s16n = int(a.n_in) / 16; + for (int u = u0 + li; u < u1; u += LPR) { + float q[NR][UNIT]; + float4 sc[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + if (gemv_direct_x()) { + gemv_raw_combine(xp + t * int(a.x_stride) + u * UNIT, q, sc, acc, t); + } else { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + } + constexpr int RPT = RS * RG * NR; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + if (li == 0) red[(ks * RPT + lrow0 + r) * 8 + t] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ks == 0 && li == 0) +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + const int row = row0 + r; + if (row < int(a.n_out)) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float v = 0.0f; +#pragma clang loop unroll(full) + for (int k = 0; k < KS; ++k) v += red[(k * RPT + lrow0 + r) * 8 + t]; + if (row < sp) y[t * int(a.y_stride) + row] = v; + else y2[t * int(a.y_stride) + row - sp] = v; + } + } +} +#define GVCK_(F, L, R, S, K) template [[host_name("gemvck_" #F "_" #L "_" #R "_" #S "_" #K)]] kernel decltype(gemv_catk) gemv_catk; +#define GVCK(F, L, R, S, K) GVCK_(F, L, R, S, K) +GVCK(T_Q8_0, 32, 1, 4, 4) GVCK(T_Q8_0, 32, 2, 4, 2) GVCK(T_Q8_0, 32, 1, 8, 8) + +// gemv_x: gemv (KS = 1) with the window's prepared x (the format's plane, NTOK rows, and its per-16 sums) staged in +// threadgroup memory once per threadgroup, so the per-unit x reads of every token come from threadgroup memory +// instead of L1 (on the M4 the x traffic through L1 - 4 B a weight a token, shared by the rows in a simdgroup - is +// what a window token costs). Same tile meaning and the same arithmetic as gemv_____1: bitwise. +// Dynamic threadgroup memory 0: 4 * NTOK * (n_in + n_in / 16) bytes. +template +kernel void gemv_x(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], + threadgroup float* xt [[threadgroup(0)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1) load_tab(tab, tid, SG * 32); + const int n = int(a.n_in), s16n = n / 16; + threadgroup float* st = xt + NTOK * n; +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + device const float4* src = (device const float4*)(xp + t * int(a.x_stride)); + threadgroup float4* dst = (threadgroup float4*)(xt + t * n); + for (int i = tid; i < n / 4; i += SG * 32) dst[i] = src[i]; + for (int i = tid; i < s16n; i += SG * 32) st[t * s16n + i] = s16[t * s16n + i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int li = lane % LPR; + const int row0 = int(tg.x) * SG * RG * NR + (sg * RG + lane / LPR) * NR; + const int U = n / UNIT; + float acc[NR][8]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) rp[r] = W + (ulong)min(row0 + r, int(a.n_out) - 1) * a.row_bytes; + for (int u = li; u < U; u += LPR) { + float q[NR][UNIT]; + float4 sc[NR]; +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[UNIT], xsum[2]; + X::load((threadgroup const float*)(xt + t * n), (threadgroup const float*)(st + t * s16n), u, xs, xsum); +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } +#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + if (li == 0 && row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = v; + } +} +#define GVX(F, L, R, S) template [[host_name("gemvx_" #F "_" #L "_" #R "_" #S)]] kernel decltype(gemv_x) gemv_x; +#define GVX_SET(F) GVX(F, 32, 1, 4) GVX(F, 32, 2, 4) GVX(F, 16, 1, 4) GVX(F, 16, 2, 4) GVX(F, 8, 1, 4) GVX(F, 8, 2, 4) \ + GVX(F, 4, 2, 4) GVX(F, 16, 2, 8) GVX(F, 8, 2, 8) GVX(F, 16, 4, 4) GVX(F, 32, 2, 8) GVX(F, 2, 2, 4) +GVX_SET(T_BF16) GVX_SET(T_Q8_0) GVX_SET(T_Q4_0) GVX_SET(T_Q5_0) GVX_SET(T_IQ4_NL) GVX_SET(T_IQ4_XS) GVX_SET(T_Q2_0) +GVX_SET(T_IQ3_S) GVX_SET(T_Q3_K) GVX_SET(T_Q4_K) GVX_SET(T_Q5_K) GVX_SET(T_Q6_K) diff --git a/results/round39/model-ab/candidate-kernels/glue.metal b/results/round39/model-ab/candidate-kernels/glue.metal new file mode 100644 index 0000000..b34bb90 --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/glue.metal @@ -0,0 +1,72 @@ +// metal/kernels/glue.metal - small elementwise steps between the big kernels (after quant/units/moe.metal). +// +// swiglu_prep h = silu(g) * u for the shared expert, written as prepared activations (xprep layout) +// moe_combine_sh bo[t] = sum_k ye[entry(t, k)] + sigmoid(sg[t]) * ysh[t] (k in order) + +struct GlueArgs { + uint n; // vector width + uint n_tok; + uint n_embd; + uint topk; +}; + +constant float glue_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, + 0.00006103515625f}; +constant float glue_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + +// grid (n / 16, n_tok), 16 threads: one group of 16 values +kernel void swiglu_prep(device const float* g [[buffer(0)]], device const float* u [[buffer(1)]], + device float* xp [[buffer(2)]], constant GlueArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort j [[thread_index_in_threadgroup]]) { + const uint t = tg.y, d = tg.x * 16 + j, idx = t * a.n + d, plane = a.n_tok * a.n; + const float gv = g[idx]; + const float h = gv / (1.0f + exp(-gv)) * u[idx]; + device float* raw = xp + idx; + raw[0] = h; + raw[plane] = h * glue_q2s[j & 7]; + raw[2 * plane] = h * glue_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? h * 0.00390625f : h; + float s0 = 0.0f, s1 = 0.0f; +#pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(h, ushort(k)); + s1 += simd_shuffle_down(h, ushort(k + 1)); + } + if (j == 0) xp[4 * plane + t * (a.n / 16) + tg.x] = s0 + s1; +} + +kernel void moe_combine_sh(device const float* ye [[buffer(0)]], device const int* plan [[buffer(1)]], + device const float* ysh [[buffer(2)]], device const float* sg [[buffer(3)]], + device float* bo [[buffer(4)]], constant GlueArgs& a [[buffer(5)]], + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.n_embd, k = gid % a.n_embd; + if (t >= a.n_tok) return; + const int row_base = int(t * a.topk); + device const int* te = plan + P_TOKENT + row_base; + device const float* yep = ye + k; + float s = 0.0f; + if (a.topk == 10) { +#pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += yep[te[j] * int(a.n_embd)]; + } else { + for (uint j = 0; j < a.topk; ++j) s += yep[te[j] * int(a.n_embd)]; + } + const float gs = 1.0f / (1.0f + exp(-sg[t])); + bo[gid] = fma(gs, ysh[gid], s); +} + +// slc_touch: reads n16 16-byte words (a weight range the next kernels will need) so they sit in the SLC; the xor is +// written only if it equals a value it never takes (keeps the loads alive). grid (ceil(n16 / 1024)), 256 threads, +// 4 words a thread. +kernel void slc_touch(device const uint4* src [[buffer(0)]], device uint* sink [[buffer(1)]], constant uint& n16 [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + uint s = 0; + for (uint k = 0; k < 4; ++k) { + const uint j = (gid >> 8) * 1024 + k * 256 + (gid & 255); + if (j < n16) { + const uint4 v = src[j]; + s ^= v.x ^ v.y ^ v.z ^ v.w; + } + } + if (s == 0x9e3779b9u) sink[0] = s; +} diff --git a/results/round39/model-ab/candidate-kernels/gr.metal b/results/round39/model-ab/candidate-kernels/gr.metal new file mode 100644 index 0000000..c18d2ac --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/gr.metal @@ -0,0 +1,307 @@ +// metal/kernels/gr.metal - the hyper-connection (gated residual) read for a window of tokens, with the previous +// sublayer's write folded in. The math is fused_gr.hpp's (native contract: FP32 activations, BF16 weights): +// +// gr_norm R'[c] = R[c] + bo * 2 sigmoid(inj_prev[c] / hc) (when apply; written back in place) +// xn[c] = R'[c] * w_norm[c] * rsqrt(mean(R'[c]^2) + eps) +// gr_down lo[k] = silu((w_down[k] . xn) / hc) k < hc_lr; inject[c] = w_inject[c] . xn +// gr_up mixed[d] = mean_c xn[c][d] * sigmoid(w_up[c n + d] . lo), written as prepared activations (xprep's +// planes and per-16 sums) for the GEMVs that read it +// +// Shapes: n_embd 2560, hc 4, hc_lr 320. Per token t: R[t] (hc * n), bo[t] (n), inj_prev[t] (hc), xn[t] (hc * n), +// lo[t] (hc_lr), inject[t] (hc), XP planes of n_tok rows. + +constant int GR_NTOK [[function_constant(1)]]; + +// the prepared activations' per-position plane scales (units.metal's xprep): q2 planes 4^-(j%8), p4 planes 256^-(j%4) +constant float gr_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, + 0.00006103515625f}; +constant float gr_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + +struct GrArgs { + uint n_embd; + uint hc; + uint hc_lr; + uint n_tok; + uint apply; + float eps; +}; + +// one threadgroup per (stream c, token t), 256 threads +kernel void gr_norm(device float* R [[buffer(0)]], device const float* bo [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device float* xn [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint c = tg.x, t = tg.y, n = a.n_embd; + device float* r = R + (t * a.hc + c) * n; + const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; + float rv[10]; + float ss = 0.0f; + if (n == 2560) { + device float* rp = r + tid; + if (a.apply) { + device const float* bp = bo + t * 2560u + tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const float v = fma(bp[i * 256], g, rp[i * 256]); + rp[i * 256] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + rv[i] = rp[i * 256]; + ss = fma(rv[i], rv[i], ss); + } + } + } else { + for (uint d = tid; d < n; d += 256) { + float v = r[d]; + if (a.apply) { + v = fma(bo[t * n + d], g, v); + r[d] = v; + } + if (n <= 2560) rv[d / 256] = v; + ss = fma(v, v, ss); + } + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device float* o = xn + (t * a.hc + c) * n; + device const float* w = w_norm + c * n; + if (n == 2560) { + o += tid; + w += tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) o[i * 256] = rv[i] * w[i * 256] * rs; + } else { + for (uint d = tid; d < n; d += 256) + o[d] = (n <= 2560 ? rv[d / 256] : r[d]) * w[d] * rs; + } +} + +// rows: 0 .. hc_lr - 1 = w_down, hc_lr .. hc_lr + hc - 1 = w_inject (both bf16 rows of hc * n_embd). +// SG simdgroups split the row's units KS ways; LPR = 32 (a simdgroup per row); RS = SG / KS simdgroups per +// threadgroup, NR rows each. +// one token's dot with a decoded unit (gr_up): x is the token's slice of lo +template +static inline float gr_dot(device const float* x, thread const float* q, float4 sc) { + float xs[UNIT], xsum[2] = {0.0f, 0.0f}; + #pragma clang loop unroll(full) + for (int k = 0; k < UNIT; k += 16) load16(x + k, xs + k); + return unit_combine(q, sc, xs, xsum); +} + +// a unit's decoded rows for grd_acc, read four codes at a time +template +struct GrdUnit { + float q[NR][fmt::UNIT]; + float4 sc[NR]; + void load(device const uchar* w, uint rb, int u, threadgroup const uint* notab) thread { + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(w + ulong(r) * rb, u, q[r], sc[r], notab); + } + float4 v(int r, int i) const thread { return float4(q[r][4 * i], q[r][4 * i + 1], q[r][4 * i + 2], q[r][4 * i + 3]); } +}; + +// Q8_0 two rows a simdgroup: the codes stay packed (8 registers a row, not 32 floats) and are widened per token. +// Measured 10-16% faster past 6 tokens (where the token loop is rolled), no faster below, and slower for one row. +template <> +struct GrdUnit { + char4 q[2][8]; + float4 sc[2]; + void load(device const uchar* w, uint rb, int u, threadgroup const uint*) thread { + #pragma clang loop unroll(full) + for (int r = 0; r < 2; ++r) { + device const uchar* b = w + ulong(r) * rb + 34 * u; + device const ushort* p = (device const ushort*)(b + 2); + #pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) q[r][k] = as_type(uint(p[2 * k]) | (uint(p[2 * k + 1]) << 16)); + const float d = rdh(b); + sc[r] = float4(d, d, 0.0f, 0.0f); + } + } + float4 v(int r, int i) const thread { return float4(q[r][i]); } +}; + +// NR consecutive rows' K-split partials over xn (units of format F; the rows rb bytes apart), acc[r * 8 + t]. +// unit_combine's sums in its order, with each xn float4 read straight from memory once for all NR rows: xn is +// 15-30x the weights' bytes at 4-8 tokens, so NR = 2 halves the kernel's traffic. GR_NTOK is a function constant, +// which the compiler does not unroll over by itself: the pragma keeps acc in registers instead of the stack. Past 6 +// tokens the unrolled loop measured slower than the rolled one, so it stays rolled there; ROLL keeps it rolled at +// every count (the packed Q8_0 pair at KS = 4: 1-3% faster at 4-6 tokens, where KS = 8 measured 2-6% slower). +template +static inline void grd_acc(device const uchar* w, uint rb, int nin, int ks, int KS, ushort lane, device const float* xn, + thread float* acc, threadgroup const uint* notab) { + constexpr int UNIT = fmt::UNIT, H = UNIT / 2; + const int U = nin / UNIT, u0 = ks * U / KS, u1 = (ks + 1) * U / KS; + for (int u = u0 + lane; u < u1; u += 32) { + GrdUnit d; + d.load(w, rb, u, notab); + auto tok = [&](int t) { + device const float4* x = (device const float4*)(xn + t * nin + UNIT * u); + float s[NR][4]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) s[r][0] = s[r][1] = s[r][2] = s[r][3] = 0.0f; + #pragma clang loop unroll(full) + for (int j = 0; j < H; j += 4) { + const float4 a = x[j / 4], b = x[(H + j) / 4]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + const float4 qa = d.v(r, j / 4), qb = d.v(r, (H + j) / 4); + s[r][0] = fma(qa.x, a.x, s[r][0]); + s[r][1] = fma(qa.y, a.y, s[r][1]); + s[r][2] = fma(qb.x, b.x, s[r][2]); + s[r][3] = fma(qb.y, b.y, s[r][3]); + s[r][0] = fma(qa.z, a.z, s[r][0]); + s[r][1] = fma(qa.w, a.w, s[r][1]); + s[r][2] = fma(qb.z, b.z, s[r][2]); + s[r][3] = fma(qb.w, b.w, s[r][3]); + } + } + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + acc[r * 8 + t] += d.sc[r].x * (s[r][0] + s[r][1]) + d.sc[r].y * (s[r][2] + s[r][3]) - + (d.sc[r].z * 0.0f + d.sc[r].w * 0.0f); + }; + if (GR_NTOK <= 6 && !ROLL) { + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) tok(t); + } else { + for (int t = 0; t < GR_NTOK; ++t) tok(t); + } + } +} + +// FI: the inject rows' format (the Unsloth files: Q8_0 down, F32 inject); FI = F is the original kernel. NR must +// divide hc_lr and hc, so that a simdgroup's rows are all w_down or all w_inject. The sums do not depend on SG or NR: +// every SG / NR with the same KS is bitwise the same kernel. +template +kernel void gr_down_f(device const uchar* w_down [[buffer(0)]], device const uchar* w_inject [[buffer(1)]], + device const float* xn [[buffer(2)]], device float* lo [[buffer(3)]], + device float* inject [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + constexpr int RS = SG / KS; + threadgroup float red[KS * RS * NR * 8]; + threadgroup uint notab[1]; + const int ks = sg / RS, rs = sg % RS; + const int row0 = (int(tg.x) * RS + rs) * NR; + const int nrows = int(a.hc_lr + a.hc); + const int rr = min(row0, nrows - NR); + const int nin = int(a.hc * a.n_embd); + const uint rb = uint(nin / fmt::QK * fmt::BB), rbi = uint(nin / fmt::QK * fmt::BB); + float acc[NR * 8]; + #pragma clang loop unroll(full) + for (int t = 0; t < NR * 8; ++t) acc[t] = 0.0f; + constexpr bool roll = F == T_Q8_0 && NR == 2 && KS == 4; + if (rr < int(a.hc_lr)) grd_acc(w_down + ulong(rr) * rb, rb, nin, ks, KS, lane, xn, acc, notab); + else grd_acc(w_inject + ulong(rr - int(a.hc_lr)) * rbi, rbi, nin, ks, KS, lane, xn, acc, notab); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) { + const float v = simd_sum(acc[r * 8 + t]); + if (lane == 0) red[((ks * RS + rs) * NR + r) * 8 + t] = v; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ks == 0 && lane < NR * 8) { + const int r = lane / 8, t = lane % 8, row = row0 + r; + if (t < GR_NTOK && row < nrows) { + float v = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < KS; ++k) v += red[((k * RS + rs) * NR + r) * 8 + t]; + if (row < int(a.hc_lr)) { + const float z = v / float(a.hc); + lo[t * int(a.hc_lr) + row] = z / (1.0f + exp(-z)); + } else { + inject[t * int(a.hc) + row - int(a.hc_lr)] = v; + } + } + } +} + +// A threadgroup owns 16 consecutive d for all hc = 4 streams: 8 simdgroups, simdgroup s -> stream s / 2 and +// d0 + 8 (s % 2) .. + 8; LPR = 4 lanes per row (320 = 20 units of 16, 5 per lane). +// XP: [4 kinds][n_tok][n] then S16 [n_tok][n / 16] (units.metal's xprep layout). +template +kernel void gr_up_f(device const uchar* w_up [[buffer(0)]], device const float* lo [[buffer(1)]], + device const float* xn [[buffer(2)]], device float* xp [[buffer(3)]], + constant GrArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float gate[8][4][16]; + threadgroup uint notab[1]; + const int n = int(a.n_embd), c = sg >> 1, dl = 8 * (sg & 1) + lane / 4, li = lane & 3; + const int d = int(tg.x) * 16 + dl; + constexpr int UNIT = fmt::UNIT; + device const uchar* w = w_up + ulong(c * n + d) * (a.hc_lr / fmt::QK * fmt::BB); + const int U = int(a.hc_lr) / UNIT; + float acc[8]; + #pragma clang loop unroll(full) + for (int t = 0; t < 8; ++t) acc[t] = 0.0f; + for (int u = li; u < U; u += 4) { + float q[UNIT]; + float4 sc; + Dec::run(w, u, q, sc, notab); + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) acc[t] += gr_dot(lo + t * int(a.hc_lr) + UNIT * u, q, sc); + } + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) { + const float v = lane_sum<4>(acc[t]); + if (li == 0) gate[t][c][dl] = 1.0f / (1.0f + exp(-v)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // 16 threads per token: thread j -> d0 + j + const int t = tid / 16, j = tid & 15; + if (t < GR_NTOK) { + const int dd = int(tg.x) * 16 + j; + const int hc = int(a.hc); + float m = 0.0f; + if (hc == 4) { + #pragma clang loop unroll(full) + for (int cc = 0; cc < 4; ++cc) m = fma(xn[(t * 4 + cc) * n + dd], gate[t][cc][j], m); + } else { + for (int cc = 0; cc < hc; ++cc) m = fma(xn[(t * hc + cc) * n + dd], gate[t][cc][j], m); + } + m = m / float(hc); + const uint plane = a.n_tok * uint(n); + device float* raw = xp + t * n + dd; + raw[0] = m; + raw[plane] = m * gr_q2s[j & 7]; + raw[2 * plane] = m * gr_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? m * 0.00390625f : m; + float s0 = 0.0f, s1 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(m, ushort(k)); + s1 += simd_shuffle_down(m, ushort(k + 1)); + } + if (j == 0) xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1; + } +} + +#define GRD1(S, K, R, SFX) \ + template [[host_name("gr_down_" #S "_" #K SFX)]] kernel decltype(gr_down_f) gr_down_f; \ + template [[host_name("gr_down_q8_" #S "_" #K SFX)]] kernel decltype(gr_down_f) gr_down_f; \ + template [[host_name("gr_down_q8f_" #S "_" #K SFX)]] kernel decltype(gr_down_f) gr_down_f; +#define GRD(S, K) GRD1(S, K, 1, "") +GRD(8, 8) GRD(8, 4) GRD(4, 4) GRD(4, 2) GRD(16, 16) GRD(16, 8) +// two rows a simdgroup (gr_read past 4 tokens) +GRD1(8, 8, 2, "_r2") GRD1(4, 4, 2, "_r2") +template [[host_name("gr_up")]] kernel decltype(gr_up_f) gr_up_f; +template [[host_name("gr_up_q8")]] kernel decltype(gr_up_f) gr_up_f; + +template [[host_name("gr_down_q4_8_8")]] kernel decltype(gr_down_f<8, 8, 1, T_Q4_0, T_Q8_0>) gr_down_f<8, 8, 1, T_Q4_0, T_Q8_0>; +template [[host_name("gr_up_q4")]] kernel decltype(gr_up_f) gr_up_f; +template [[host_name("gr_down_q5_8_8")]] kernel decltype(gr_down_f<8, 8, 1, T_Q5_0, T_Q8_0>) gr_down_f<8, 8, 1, T_Q5_0, T_Q8_0>; +template [[host_name("gr_up_q5")]] kernel decltype(gr_up_f) gr_up_f; diff --git a/results/round39/model-ab/candidate-kernels/gr2.metal b/results/round39/model-ab/candidate-kernels/gr2.metal new file mode 100644 index 0000000..58342f0 --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/gr2.metal @@ -0,0 +1,484 @@ +// metal/kernels/gr2.metal - the hyper-connection read in TWO dispatches (after quant/units/gr.metal), for BF16 HC +// weights (the exact path) or a Q8_0 copy of them (the STRATA_HC_Q8 option): +// +// gr_down3 grid (row blocks, hc streams). The RMSNorm is per stream, so a threadgroup owns one stream c: it forms +// R'_c = R_c + bo * 2 sigmoid(inj_prev_c / hc) (10 KB a token), its sum of squares and rs_c itself, +// stages xn_c = (R'_c * w_norm_c) * rs_c in threadgroup memory once, and every simdgroup dots its rows' +// stream-c quarter with it: D[t][row][c]. Row block 0 also writes R'_c to R_out (ping-pong, never in +// place) and rs_c. Rows 0..hc_lr-1 are w_down, hc_lr.. the inject rows. +// gr_up3 lo[k] = silu(sum_c D[t][k][c] / hc) formed in threadgroup memory, then as gr_up with xn = R' * w_norm +// * rs on the fly, writing the prepared activations. +// inj_prev of the NEXT read = sum_c D[t][hc_lr + c'][c] of this one (gr_down3 sums it, in c order). +// +// Per element the arithmetic is gr_norm's (fma(bo, g, R), then (R' * w) * rs); the K order differs (per stream, then +// the 4 streams in order). + +#ifndef P_TOKENT +#define P_TOKENT 532 // moe.metal's plan layout +#endif + +struct GrArgs3 { + uint n_embd; + uint hc; + uint hc_lr; + uint n_tok; // the prepared planes' row count (the plane stride) + uint apply; + float eps; + uint rb_down; // bytes per row of w_down / w_inject + uint rb_up; // bytes per row of w_up + uint comb; // 1: bo is not given; it is the MoE combine (moe_combine_sh's arithmetic) of buffers 10..13 + uint topk; +}; + +// bo[t][d], given or the MoE combine: sum_k ye[entry(t, k)][d] (k in order) + sigmoid(sg[t]) * ysh[t][d] +static inline float bo_at(device const float* bo, device const int* plan, device const float* ye, + device const float* ysh, device const float* sg, constant GrArgs3& a, int t, int d) { + const int n = int(a.n_embd); + if (!a.comb) return bo[t * n + d]; + float s = 0.0f; + for (uint j = 0; j < a.topk; ++j) s += ye[plan[P_TOKENT + t * int(a.topk) + j] * n + d]; + const float gs = 1.0f / (1.0f + exp(-sg[t])); + return fma(gs, ysh[t * n + d], s); +} + +template +kernel void gr_down3(device const uchar* w_down [[buffer(0)]], device const uchar* w_inject [[buffer(1)]], + device const float* Rin [[buffer(2)]], device float* Rout [[buffer(3)]], + device const float* bo [[buffer(4)]], device const float* Dprev [[buffer(5)]], + device const float* wn [[buffer(6)]], device float* D [[buffer(7)]], + device float* rs_out [[buffer(8)]], constant GrArgs3& a [[buffer(9)]], + device const int* plan [[buffer(10)]], device const float* ye [[buffer(11)]], + device const float* ysh [[buffer(12)]], device const float* sgv [[buffer(13)]], + 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 UNIT = fmt::UNIT; + threadgroup float xs_tg[TT][2560]; + threadgroup float part[8][TT]; + threadgroup float g_tg[TT], rs_tg[TT]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, 256); + const int n = int(a.n_embd), hc = int(a.hc), NIN = hc * n, c = int(tg.y); + const int nrows = int(a.hc_lr) + hc; + for (int t0 = 0; t0 < GR_NTOK; t0 += TT) { + const int T = min(TT, GR_NTOK - t0); + if (tid < T) { + float g = 0.0f; + if (a.apply) { + device const float* dp = Dprev + ((t0 + tid) * nrows + int(a.hc_lr) + c) * hc; + const float ip = ((dp[0] + dp[1]) + dp[2]) + dp[3]; + g = 2.0f / (1.0f + exp(-ip / float(hc))); + } + g_tg[tid] = g; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int t = 0; t < T; ++t) { + device const float* rt = Rin + (t0 + t) * NIN + c * n; + float ss = 0.0f; + for (int d = tid; d < n; d += 256) { + float v = rt[d]; + if (a.apply) v = fma(bo_at(bo, plan, ye, ysh, sgv, a, t0 + t, d), g_tg[t], v); + xs_tg[t][d] = v; + ss = fma(v, v, ss); + if (tg.x == 0) Rout[(t0 + t) * NIN + c * n + d] = v; + } + ss = simd_sum(ss); + if (lane == 0) part[sg][t] = ss; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < T) { + float tot = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) tot += part[k][tid]; + const float rsv = rsqrt(tot / float(n) + a.eps); + rs_tg[tid] = rsv; + if (tg.x == 0) rs_out[(t0 + tid) * hc + c] = rsv; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int t = 0; t < T; ++t) + for (int d = tid; d < n; d += 256) xs_tg[t][d] = xs_tg[t][d] * wn[c * n + d] * rs_tg[t]; + threadgroup_barrier(mem_flags::mem_threadgroup); + // the rows: simdgroup sg holds RPS rows; lanes split the stream's units + int rows[RPS]; + device const uchar* wr[RPS]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + rows[k] = (int(tg.x) * 8 + sg) * RPS + k; + const int rr = min(rows[k], nrows - 1); + wr[k] = rr < int(a.hc_lr) ? w_down + ulong(rr) * a.rb_down : w_inject + ulong(rr - int(a.hc_lr)) * a.rb_down; + } + float acc[RPS][TT]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) acc[k][t] = 0.0f; + } + const int U = n / UNIT, u0 = c * U; + for (int u = lane; u < U; u += 32) { + float q[RPS][UNIT]; + float4 sc[RPS]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) Dec::run(wr[k], u0 + u, q[k], sc[k], tab); + // TT-bounded with a T guard, so acc[][t] is indexed by a constant (registers, not the stack) + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) { + if (t >= T) break; + float xs[UNIT], xsum[2] = {0.0f, 0.0f}; + #pragma clang loop unroll(full) + for (int e = 0; e < UNIT; ++e) xs[e] = xs_tg[t][u * UNIT + e]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) acc[k][t] += unit_combine(q[k], sc[k], xs, xsum); + } + } + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) { + if (t >= T) break; + const float v = simd_sum(acc[k][t]); + if (lane == 0 && rows[k] < nrows) D[((t0 + t) * nrows + rows[k]) * hc + c] = v; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +// grid (n_embd / 16), 256 threads. +template +kernel void gr_up3(device const uchar* w_up [[buffer(0)]], device const float* D [[buffer(1)]], + device const float* Rn [[buffer(2)]], device float* xp [[buffer(3)]], + constant GrArgs3& a [[buffer(4)]], device const float* wn [[buffer(5)]], + device const float* rs [[buffer(6)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT; + threadgroup float lo_tg[8][320]; + threadgroup float gate[8][4][16]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, 256); + const int n = int(a.n_embd), hc = int(a.hc), lr = int(a.hc_lr), nrows = lr + hc; + for (int e = tid; e < GR_NTOK * lr; e += 256) { + const int t = e / lr, k = e - t * lr; + device const float* dp = D + (t * nrows + k) * hc; + const float z = (((dp[0] + dp[1]) + dp[2]) + dp[3]) / float(hc); + lo_tg[t][k] = z / (1.0f + exp(-z)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int c = sg >> 1, dl = 8 * (sg & 1) + lane / 4, li = lane & 3; + const int d = int(tg.x) * 16 + dl; + device const uchar* w = w_up + ulong(c * n + d) * a.rb_up; + const int U = lr / UNIT; + float acc[8]; + #pragma clang loop unroll(full) + for (int t = 0; t < 8; ++t) acc[t] = 0.0f; + for (int u = li; u < U; u += 4) { + float q[UNIT]; + float4 sc; + Dec::run(w, u, q, sc, tab); + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) { + float xs[UNIT], xsum[2] = {0.0f, 0.0f}; + #pragma clang loop unroll(full) + for (int e = 0; e < UNIT; ++e) xs[e] = lo_tg[t][UNIT * u + e]; + acc[t] += unit_combine(q, sc, xs, xsum); + } + } + #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) { + const float v = lane_sum<4>(acc[t]); + if (li == 0) gate[t][c][dl] = 1.0f / (1.0f + exp(-v)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int t = tid / 16, j = tid & 15; + if (t < GR_NTOK) { + const int dd = int(tg.x) * 16 + j; + float m = 0.0f; + if (hc == 4) { + #pragma clang loop unroll(full) + for (int cc = 0; cc < 4; ++cc) { + const float xn = Rn[(t * 4 + cc) * n + dd] * wn[cc * n + dd] * rs[t * 4 + cc]; + m = fma(xn, gate[t][cc][j], m); + } + } else { + for (int cc = 0; cc < hc; ++cc) { + const float xn = Rn[(t * hc + cc) * n + dd] * wn[cc * n + dd] * rs[t * hc + cc]; + m = fma(xn, gate[t][cc][j], m); + } + } + m = m / float(hc); + const uint plane = a.n_tok * uint(n); + device float* raw = xp + t * n + dd; + raw[0] = m; + raw[plane] = m * gr_q2s[j & 7]; + raw[2 * plane] = m * gr_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? m * 0.00390625f : m; + float s0 = 0.0f, s1 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(m, ushort(k)); + s1 += simd_shuffle_down(m, ushort(k + 1)); + } + if (j == 0) xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1; + } +} + +#define GD3(F, T, R) template [[host_name("gr_down3_" #F "_" #T "_" #R)]] kernel decltype(gr_down3) gr_down3; +GD3(30, 1, 1) GD3(30, 2, 1) GD3(30, 3, 1) GD3(30, 1, 2) GD3(30, 2, 2) GD3(30, 3, 2) +GD3(8, 1, 1) GD3(8, 2, 1) GD3(8, 3, 1) GD3(8, 1, 2) GD3(8, 2, 2) GD3(8, 3, 2) +template [[host_name("gr_up3_30")]] kernel decltype(gr_up3) gr_up3; +template [[host_name("gr_up3_8")]] kernel decltype(gr_up3) gr_up3; + +// BF16 only: as gr_down3, but each lane loads its rows' weights (5 units of 16 per row and lane: 2560 = 160 units) +// into registers before the norm phase, so the weight stream is in flight while the threadgroup forms xn. +template +kernel void gr_down3b(device const uchar* w_down [[buffer(0)]], device const uchar* w_inject [[buffer(1)]], + device const float* Rin [[buffer(2)]], device float* Rout [[buffer(3)]], + device const float* bo [[buffer(4)]], device const float* Dprev [[buffer(5)]], + device const float* wn [[buffer(6)]], device float* D [[buffer(7)]], + device float* rs_out [[buffer(8)]], constant GrArgs3& a [[buffer(9)]], + device const int* plan [[buffer(10)]], device const float* ye [[buffer(11)]], + device const float* ysh [[buffer(12)]], device const float* sgv [[buffer(13)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float xs_tg[TT][2560]; + threadgroup float part[8][TT]; + threadgroup float g_tg[TT], rs_tg[TT]; + const int n = int(a.n_embd), hc = int(a.hc), NIN = hc * n, c = int(tg.y); + const int nrows = int(a.hc_lr) + hc; + int rows[RPS]; + uint4 wr[RPS][5][2]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + rows[k] = (int(tg.x) * 8 + sg) * RPS + k; + const int rr = min(rows[k], nrows - 1); + device const uchar* w = rr < int(a.hc_lr) ? w_down + ulong(rr) * a.rb_down : w_inject + ulong(rr - int(a.hc_lr)) * a.rb_down; + device const uint4* p = (device const uint4*)(w + 2 * (c * n)); + #pragma clang loop unroll(full) + for (int i = 0; i < 5; ++i) { + wr[k][i][0] = p[2 * (lane + 32 * i)]; + wr[k][i][1] = p[2 * (lane + 32 * i) + 1]; + } + } + for (int t0 = 0; t0 < GR_NTOK; t0 += TT) { + const int T = min(TT, GR_NTOK - t0); + if (tid < T) { + float g = 0.0f; + if (a.apply) { + device const float* dp = Dprev + ((t0 + tid) * nrows + int(a.hc_lr) + c) * hc; + const float ip = ((dp[0] + dp[1]) + dp[2]) + dp[3]; + g = 2.0f / (1.0f + exp(-ip / float(hc))); + } + g_tg[tid] = g; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int t = 0; t < T; ++t) { + device const float* rt = Rin + (t0 + t) * NIN + c * n; + float ss = 0.0f; + for (int d = tid; d < n; d += 256) { + float v = rt[d]; + if (a.apply) v = fma(bo_at(bo, plan, ye, ysh, sgv, a, t0 + t, d), g_tg[t], v); + xs_tg[t][d] = v; + ss = fma(v, v, ss); + if (tg.x == 0) Rout[(t0 + t) * NIN + c * n + d] = v; + } + ss = simd_sum(ss); + if (lane == 0) part[sg][t] = ss; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < T) { + float tot = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) tot += part[k][tid]; + const float rsv = rsqrt(tot / float(n) + a.eps); + rs_tg[tid] = rsv; + if (tg.x == 0) rs_out[(t0 + tid) * hc + c] = rsv; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int t = 0; t < T; ++t) + for (int d = tid; d < n; d += 256) xs_tg[t][d] = xs_tg[t][d] * wn[c * n + d] * rs_tg[t]; + threadgroup_barrier(mem_flags::mem_threadgroup); + float acc[RPS][TT]; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) acc[k][t] = 0.0f; + } + #pragma clang loop unroll(full) + for (int i = 0; i < 5; ++i) { + const int u = lane + 32 * i; + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + const uint v[8] = {wr[k][i][0].x, wr[k][i][0].y, wr[k][i][0].z, wr[k][i][0].w, + wr[k][i][1].x, wr[k][i][1].y, wr[k][i][1].z, wr[k][i][1].w}; + float q[16]; + #pragma clang loop unroll(full) + for (int e = 0; e < 8; ++e) { + q[2 * e] = as_type(v[e] << 16); + q[2 * e + 1] = as_type(v[e] & 0xFFFF0000u); + } + const float4 sc = float4(1.0f, 1.0f, 0.0f, 0.0f); + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) { + if (t >= T) break; + float xs[16], xsum[2] = {0.0f, 0.0f}; + #pragma clang loop unroll(full) + for (int e = 0; e < 16; ++e) xs[e] = xs_tg[t][u * 16 + e]; + acc[k][t] += unit_combine<16>(q, sc, xs, xsum); + } + } + } + #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) { + if (t >= T) break; + const float v = simd_sum(acc[k][t]); + if (lane == 0 && rows[k] < nrows) D[((t0 + t) * nrows + rows[k]) * hc + c] = v; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} +#define GD3B(T, R) template [[host_name("gr_down3b_" #T "_" #R)]] kernel decltype(gr_down3b) gr_down3b; +GD3B(1, 1) GD3B(2, 1) GD3B(3, 1) GD3B(1, 2) GD3B(2, 2) GD3B(3, 2) + +// gr_norm_comb: gr_norm (gr.metal) with the sublayer output bo formed on the fly as the MoE combine - bo[t][d] = +// sum_k ye[entry(t, k)][d] (k in order) + sigmoid(sg[t]) * ysh[t][d], moe_combine_sh's arithmetic - so the combine's +// dispatch goes away; R', xn bitwise gr_norm's after moe_combine_sh. One threadgroup per (stream c, token t). +kernel void gr_norm_comb(device float* R [[buffer(0)]], device const float* ye [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device float* xn [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + device const int* plan [[buffer(6)]], device const float* ysh [[buffer(7)]], + device const float* sgv [[buffer(8)]], constant uint& topk [[buffer(9)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint c = tg.x, t = tg.y, n = a.n_embd; + device float* r = R + (t * a.hc + c) * n; + const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + float rv[10]; + float ss = 0.0f; + if (n == 2560 && topk == 10) { + device const float* yep[10]; + if (a.apply) { + #pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) + yep[j] = ye + uint(plan[P_TOKENT + t * 10 + j]) * 2560u + tid; + } + device const float* yshp = ysh + t * 2560u + tid; + device float* rp = r + tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + float v = rp[i * 256]; + if (a.apply) { + float s = 0.0f; + #pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += yep[j][i * 256]; + v = fma(fma(gs, yshp[i * 256], s), g, v); + rp[i * 256] = v; + } + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid; d < n; d += 256) { + float v = r[d]; + if (a.apply) { + float s = 0.0f; + for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; + const float bo = fma(gs, ysh[t * n + d], s); + v = fma(bo, g, v); + r[d] = v; + } + if (n <= 2560) rv[d / 256] = v; + ss = fma(v, v, ss); + } + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device float* o = xn + (t * a.hc + c) * n; + device const float* w = w_norm + c * n; + if (n == 2560) { + o += tid; + w += tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) o[i * 256] = rv[i] * w[i * 256] * rs; + } else { + for (uint d = tid; d < n; d += 256) + o[d] = (n <= 2560 ? rv[d / 256] : r[d]) * w[d] * rs; + } +} + +// Fixed n_embd = 2560, hc = 4; grid (1, n_tok), 256 threads. Form bo once for all streams. +kernel void gr_norm_comb_4s(device float* R [[buffer(0)]], device const float* ye [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device float* xn [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + device const int* plan [[buffer(6)]], device const float* ysh [[buffer(7)]], + device const float* sgv [[buffer(8)]], constant uint& topk [[buffer(9)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + const uint t = tg.y, n = 2560; + device float* rp = R + t * 4 * n + tid; + device float* op = xn + t * 4 * n + tid; + device const float* wp = w_norm + tid; + float rv[4][10], ss[4] = {0.0f, 0.0f, 0.0f, 0.0f}, g[4]; + #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) + g[c] = a.apply ? 2.0f / (1.0f + exp(-inj[t * 4 + c] / 4.0f)) : 0.0f; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + device const float* yep[10]; + if (a.apply && topk == 10) { + #pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) + yep[j] = ye + uint(plan[P_TOKENT + t * 10 + j]) * 2560u + tid; + } + device const float* yshp = ysh + t * n + tid; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const uint d = tid + i * 256; + float bo = 0.0f; + if (a.apply) { + float s = 0.0f; + if (topk == 10) { + #pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += yep[j][i * 256]; + } else { + for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; + } + bo = fma(gs, yshp[i * 256], s); + } + #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) { + device float* r = rp + c * n; + float v = r[i * 256]; + if (a.apply) { v = fma(bo, g[c], v); r[i * 256] = v; } + rv[c][i] = v; + ss[c] = fma(v, v, ss[c]); + } + } + #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) { + const float s = simd_sum(ss[c]); + if (lane == 0) part[c][sg] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + #pragma clang loop unroll(full) + for (uint c = 0; c < 4; ++c) { + float tot = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[c][i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device float* o = op + c * n; + device const float* w = wp + c * n; + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + o[i * 256] = rv[c][i] * w[i * 256] * rs; + } + } +} diff --git a/results/round39/model-ab/candidate-kernels/gridbar.metal b/results/round39/model-ab/candidate-kernels/gridbar.metal new file mode 100644 index 0000000..5cd9a0f --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/gridbar.metal @@ -0,0 +1,54 @@ +// metal/kernels/gridbar.metal - feasibility probe for a persistent kernel: phases separated by an in-kernel grid +// barrier (a device atomic counter) instead of dispatch boundaries. Metal does not promise that all threadgroups of +// a grid are resident at once; the probe launches at most one threadgroup per GPU core. +// +// Each phase p: every thread streams its share of chunk p (n uint4), and threadgroup g writes a token value +// (p * 1000 + g) to flag[p][g]; after the barrier, phase p + 1 checks every flag of phase p (written by the OTHER +// threadgroups) and counts mismatches into bad[0]. + +static inline void grid_barrier(device atomic_uint* bar, uint target, ushort tid) { + threadgroup_barrier(mem_flags::mem_device); +#ifdef GB_FENCE + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); +#endif + if (tid == 0) { + atomic_fetch_add_explicit(bar, 1u, memory_order_relaxed); + while (atomic_load_explicit(bar, memory_order_relaxed) < target) {} + } + threadgroup_barrier(mem_flags::mem_device); +#ifdef GB_FENCE + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); +#endif +} + +kernel void k_phases(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + device atomic_uint* bar [[buffer(2)]], device volatile uint* flag [[buffer(3)]], + device atomic_uint* bad [[buffer(4)]], constant uint& n [[buffer(5)]], + constant uint& phases [[buffer(6)]], + uint tgpx [[threadgroup_position_in_grid]], uint ntgx [[threadgroups_per_grid]], + ushort tid [[thread_index_in_threadgroup]], uint gid [[thread_position_in_grid]], + uint nt [[threads_per_grid]]) { + const uint g = tgpx, ntg = ntgx; + uint4 acc = 0; + for (uint p = 0; p < phases; ++p) { + if (p > 0 && tid < ntg) { + const uint v = flag[(p - 1) * ntg + tid]; + if (v != (p - 1) * 1000 + tid) atomic_fetch_add_explicit(bad, 1u, memory_order_relaxed); + } + device const uint4* s = src + ulong(p) * n; + for (uint i = gid; i < n; i += nt) acc += s[i]; + if (tid == 0) flag[p * ntg + g] = p * 1000 + g; + grid_barrier(bar, (p + 1) * ntg, tid); + } + out[gid] = acc.x ^ acc.y ^ acc.z ^ acc.w; +} + +// one phase per dispatch, for comparison +kernel void k_phase1(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + constant uint& n [[buffer(2)]], constant uint& p [[buffer(3)]], + uint gid [[thread_position_in_grid]], uint nt [[threads_per_grid]]) { + uint4 acc = 0; + device const uint4* s = src + ulong(p) * n; + for (uint i = gid; i < n; i += nt) acc += s[i]; + out[gid] = acc.x ^ acc.y ^ acc.z ^ acc.w; +} diff --git a/results/round39/model-ab/candidate-kernels/mm_bench.metal b/results/round39/model-ab/candidate-kernels/mm_bench.metal new file mode 100644 index 0000000..c04659f --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/mm_bench.metal @@ -0,0 +1,85 @@ +// metal/kernels/mm_bench.metal - matrix throughput probes (mm_bench.cpp): C[M][N] = A[M][K] . B[N][K]^T, fp32 out. +// mpp_mm Metal 4 MetalPerformancePrimitives matmul2d (the M5's neural accelerators), a TM x TN +// tile of C per threadgroup, K dynamic, A and B straight from device memory +// sg_mm simdgroup_matrix 8x8, a 32 x 32 tile per simdgroup (4 x 4 accumulators), 4 simdgroups +#include +using namespace mpp::tensor_ops; + +struct MmArgs { int M, N, K; }; + +template +kernel void mpp_mm(device T* A [[buffer(0)]], device T* B [[buffer(1)]], device float* C [[buffer(2)]], + constant MmArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]]) { + // extents are (inner, outer): A is M rows of K, B is N rows of K, C is M rows of N + tensor, tensor_inline> tA(A, dextents(a.K, a.M)); + tensor, tensor_inline> tB(B, dextents(a.K, a.N)); + tensor, tensor_inline> tC(C, dextents(a.N, a.M)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, int(tg.y) * TM); + auto mB = tB.slice(0, int(tg.x) * TN); + auto mC = tC.slice(int(tg.x) * TN, int(tg.y) * TM); + op.run(mA, mB, mC); +} + +#define MPP(T, M, N, S) template [[host_name("mpp_mm_" #T "_" #M "_" #N "_" #S)]] kernel decltype(mpp_mm) mpp_mm; +MPP(half, 64, 32, 4) MPP(half, 64, 64, 4) MPP(half, 128, 64, 4) MPP(half, 32, 64, 4) MPP(half, 64, 128, 4) +MPP(half, 16, 64, 2) MPP(half, 128, 128, 4) +MPP(half, 64, 16, 2) MPP(half, 64, 16, 4) MPP(half, 128, 16, 4) MPP(half, 128, 32, 4) MPP(half, 64, 32, 2) MPP(half, 32, 16, 2) +MPP(half, 32, 32, 2) MPP(half, 16, 32, 1) MPP(half, 16, 16, 1) MPP(half, 128, 48, 4) MPP(half, 64, 48, 4) +MPP(float, 64, 32, 4) MPP(float, 64, 64, 4) MPP(float, 128, 64, 4) +MPP(bfloat, 64, 64, 4) MPP(bfloat, 128, 64, 4) + +// simdgroup_matrix: threadgroup tile 64 x 64 (4 simdgroups x 32 x 32), K in steps of 8, operands from device memory +template +kernel void sg_mm(device const T* A [[buffer(0)]], device const T* B [[buffer(1)]], device float* C [[buffer(2)]], + constant MmArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + const int m0 = int(tg.y) * 64 + (sg >> 1) * 32, n0 = int(tg.x) * 64 + (sg & 1) * 32; + simdgroup_matrix acc[4][4]; + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) acc[i][j] = make_filled_simdgroup_matrix(0.0f); + for (int k = 0; k < a.K; k += 8) { + simdgroup_matrix ma[4], mb[4]; + for (int i = 0; i < 4; ++i) simdgroup_load(ma[i], A + (m0 + 8 * i) * a.K + k, a.K); + for (int j = 0; j < 4; ++j) simdgroup_load(mb[j], B + (n0 + 8 * j) * a.K + k, a.K, ulong2(0, 0), true); + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) simdgroup_multiply_accumulate(acc[i][j], ma[i], mb[j], acc[i][j]); + } + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) simdgroup_store(acc[i][j], C + (m0 + 8 * i) * a.N + n0 + 8 * j, a.N); +} +template [[host_name("sg_mm_half")]] kernel decltype(sg_mm) sg_mm; +template [[host_name("sg_mm_float")]] kernel decltype(sg_mm) sg_mm; + + +// the element layout of single-simdgroup matmul2d cooperative tensors (right input, left input, destination) for a +// TM x TN x TK descriptor: out[((which * 32 + lane) * 128 + i) * 4 + (0, 1)] = the two indices, [2] = capacity +template +kernel void coop_layout(device int* out [[buffer(0)]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + auto r = op.template get_right_input_cooperative_tensor(); + auto l = op.template get_left_input_cooperative_tensor(); + auto dd = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < dd.get_capacity() && i < 128; ++i) { + auto ix = dd.get_multidimensional_index(i); + device int* o = out + ((2 * 32 + lane) * 128 + i) * 4; + o[0] = ix[0]; o[1] = ix[1]; o[2] = dd.get_capacity(); + } + for (uint16_t i = 0; i < r.get_capacity() && i < 128; ++i) { + auto ix = r.get_multidimensional_index(i); + device int* o = out + ((0 * 32 + lane) * 128 + i) * 4; + o[0] = ix[0]; o[1] = ix[1]; o[2] = r.get_capacity(); + } + for (uint16_t i = 0; i < l.get_capacity() && i < 128; ++i) { + auto ix = l.get_multidimensional_index(i); + device int* o = out + ((1 * 32 + lane) * 128 + i) * 4; + o[0] = ix[0]; o[1] = ix[1]; o[2] = l.get_capacity(); + } +} +template [[host_name("coop_layout_16_32_16")]] kernel decltype(coop_layout<16, 32, 16>) coop_layout<16, 32, 16>; +template [[host_name("coop_layout_32_32_16")]] kernel decltype(coop_layout<32, 32, 16>) coop_layout<32, 32, 16>; +template [[host_name("coop_layout_16_64_32")]] kernel decltype(coop_layout<16, 64, 32>) coop_layout<16, 64, 32>; +template [[host_name("coop_layout_16_32_32")]] kernel decltype(coop_layout<16, 32, 32>) coop_layout<16, 32, 32>; diff --git a/results/round39/model-ab/candidate-kernels/mma.metal b/results/round39/model-ab/candidate-kernels/mma.metal new file mode 100644 index 0000000..8fcc0ab --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/mma.metal @@ -0,0 +1,279 @@ +// metal/kernels/mma.metal - a GEMV for windows of tokens on the simdgroup matrix units (after quant/units.metal). +// +// y[t][r] = W[r] . x[t] for NTOK (1..8) tokens: each simdgroup owns 8 rows; the K axis runs in chunks of 128. Per +// chunk the simdgroup decodes its 8 x 128 weights - with ref_quant's arithmetic, so a tile element is exactly the +// dequantized weight - into threadgroup memory, then 16 simdgroup_multiply_accumulate steps take A = the 8 x 8 weight +// tile and B = the 8 (k) x 8 (token) slice of x, loaded transposed straight from device memory (x has 8 rows: the +// tokens past the window are padding, and an output column depends only on its own token). +// +// Activations are the raw plane (no pre-scaling); n_in must be a multiple of 128. Every format's units sit inside +// one 128-wide chunk: Q4_K/Q5_K units span 64 positions, Q6_K units 128, the others are contiguous. + +// ---- decoders to actual weights: w[UNIT] and the chunk offsets of its two halves (p0: w[0..H), p1: w[H..UNIT)) ---- +template struct DW { + // generic: w = s * q (quant.metal's deq is ref_quant's order: the scale times the code), contiguous + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT, UPB = fmt::QK / UNIT; + float q[UNIT]; + const float s = deq(row + (u / UPB) * fmt::BB, u % UPB, q, tab); +#pragma clang loop unroll(full) + for (int j = 0; j < UNIT; ++j) w[j] = q[j] * s; + p0 = u * UNIT - k0; + p1 = p0 + UNIT / 2; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint* tab) { + device const uchar* b = row + u * 18; + device const ushort* p = (device const ushort*)(b + 2); +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); + w[2 * j] = float(w0.x); + w[2 * j + 16] = float(w0.y); + w[2 * j + 1] = float(w1.x); + w[2 * j + 17] = float(w1.y); + } + const float d = rdh(b); +#pragma clang loop unroll(full) + for (int j = 0; j < 32; ++j) w[j] *= d; + p0 = 32 * u - k0; + p1 = p0 + 16; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint* tab) { + device const uchar* b = row + (u >> 3) * 136; + const int ui = u & 7; + device const ushort* p = (device const ushort*)(b + 8 + 16 * ui); +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); + w[2 * j] = float(w0.x); + w[2 * j + 16] = float(w0.y); + w[2 * j + 1] = float(w1.x); + w[2 * j + 17] = float(w1.y); + } + const uint sh = *(device const ushort*)(b + 2); + const int ls = int((b[4 + (ui >> 1)] >> (4 * (ui & 1))) & 15) | int(((sh >> (2 * ui)) & 3) << 4); + const float d = rdh(b) * float(ls - 32); +#pragma clang loop unroll(full) + for (int j = 0; j < 32; ++j) w[j] *= d; + p0 = 32 * u - k0; + p1 = p0 + 16; + } +}; +// Q2_0: (code - 1) * d +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 2) * 18; + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + const ushort a = p[0], c = p[1]; + const float d = rdh(b); +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + w[j] = float(int((a >> (2 * j)) & 3) - 1) * d; + w[8 + j] = float(int((c >> (2 * j)) & 3) - 1) * d; + } + p0 = 16 * u - k0; + p1 = p0 + 8; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint* tab) { + device const uchar* b = row + u * 24; + float q[32]; + const float d = deq(b, 0, q, tab), m = rdh(b + 2); +#pragma clang loop unroll(full) + for (int j = 0; j < 32; ++j) w[j] = q[j] * d + m; + p0 = 32 * u - k0; + p1 = p0 + 16; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const float* p = (device const float*)(row + 64 * u); +#pragma clang loop unroll(full) + for (int j = 0; j < 16; ++j) w[j] = p[j]; + p0 = 16 * u - k0; + p1 = p0 + 8; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uint4* p = (device const uint4*)(row + 32 * u); + const uint4 a = p[0], c = p[1]; + const uint v[8] = {a.x, a.y, a.z, a.w, c.x, c.y, c.z, c.w}; +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) { + w[2 * k] = as_type(v[k] << 16); + w[2 * k + 1] = as_type(v[k] & 0xFFFF0000u); + } + p0 = 16 * u - k0; + p1 = p0 + 8; + } +}; +// Q4_K / Q5_K: d1 * code - m1 (ref: d * sc, dmin * m), runs at 256b + 64c + 16m and + 32 +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 144; + const int c = (u >> 1) & 3, m = u & 1; + const float d = rdh(b), dm = rdh(b + 2); + float s1, m1, s2, m2; + scale_min_k4(2 * c, b + 4, s1, m1); + scale_min_k4(2 * c + 1, b + 4, s2, m2); + const float d1 = d * s1, mm1 = dm * m1, d2 = d * s2, mm2 = dm * m2; + device const uchar* q = b + 16 + 32 * c + 16 * m; +#pragma clang loop unroll(full) + for (int j = 0; j < 16; ++j) { + w[j] = d1 * float(q[j] & 15) - mm1; + w[16 + j] = d2 * float(q[j] >> 4) - mm2; + } + p0 = 256 * (u >> 3) + 64 * c + 16 * m - k0; + p1 = p0 + 32; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 176; + const int c = (u >> 1) & 3, m = u & 1; + const float d = rdh(b), dm = rdh(b + 2); + float s1, m1, s2, m2; + scale_min_k4(2 * c, b + 4, s1, m1); + scale_min_k4(2 * c + 1, b + 4, s2, m2); + const float d1 = d * s1, mm1 = dm * m1, d2 = d * s2, mm2 = dm * m2; + device const uchar* ql = b + 48 + 32 * c + 16 * m; + device const uchar* qh = b + 16 + 16 * m; + const uchar u1 = uchar(1 << (2 * c)), u2 = uchar(2 << (2 * c)); +#pragma clang loop unroll(full) + for (int j = 0; j < 16; ++j) { + w[j] = d1 * float((ql[j] & 15) + ((qh[j] & u1) ? 16 : 0)) - mm1; + w[16 + j] = d2 * float((ql[j] >> 4) + ((qh[j] & u2) ? 16 : 0)) - mm2; + } + p0 = 256 * (u >> 3) + 64 * c + 16 * m - k0; + p1 = p0 + 32; + } +}; +// Q6_K: d * sc * (code - 32), runs at 256b + 128n + 32p + 16m and + 64 +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 210; + const int n = (u >> 2) & 1, m = (u >> 1) & 1, p = u & 1; + device const uchar* ql = b + 64 * n + 32 * p + 16 * m; + device const uchar* qh = b + 128 + 32 * n + 16 * m; + device const char* sc = (device const char*)(b + 192); + const float d = rdh(b + 208); + const float ds1 = d * float(sc[8 * n + m + 2 * p]), ds2 = d * float(sc[8 * n + m + 2 * p + 4]); + const uint sh = uint(2 * p), sh4 = sh + 4u; +#pragma clang loop unroll(full) + for (int j = 0; j < 16; j += 4) { + const uint la = rd32(ql + j), ha = rd32(qh + j); + const uchar4 c1 = as_type((la & 0x0F0F0F0Fu) | (((ha >> sh) & 0x03030303u) << 4)); + const uchar4 c2 = as_type(((la >> 4) & 0x0F0F0F0Fu) | (((ha >> sh4) & 0x03030303u) << 4)); + w[j] = ds1 * float(int(c1.x) - 32); + w[j + 1] = ds1 * float(int(c1.y) - 32); + w[j + 2] = ds1 * float(int(c1.z) - 32); + w[j + 3] = ds1 * float(int(c1.w) - 32); + w[16 + j] = ds2 * float(int(c2.x) - 32); + w[16 + j + 1] = ds2 * float(int(c2.y) - 32); + w[16 + j + 2] = ds2 * float(int(c2.z) - 32); + w[16 + j + 3] = ds2 * float(int(c2.w) - 32); + } + p0 = 256 * (u >> 3) + 128 * n + 32 * p + 16 * m - k0; + p1 = p0 + 64; + } +}; +// Q3_K: d_all * (sc - 32) * (code - (hbit ? 0 : 4)), contiguous 32 +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 110; + const int n = (u >> 2) & 1, j = u & 3; + device const uchar* q = b + 32 + 32 * n; + device const uchar* hm = b; + device const uchar* s = b + 96; + const int i1 = 8 * n + 2 * j, i2 = i1 + 1; + const int l1 = (i1 < 8 ? s[i1] & 0xF : s[i1 - 8] >> 4) | (((s[8 + (i1 & 3)] >> (2 * (i1 >> 2))) & 3) << 4); + const int l2 = (i2 < 8 ? s[i2] & 0xF : s[i2 - 8] >> 4) | (((s[8 + (i2 & 3)] >> (2 * (i2 >> 2))) & 3) << 4); + const float d = rdh(b + 108); + const float dl1 = d * float(l1 - 32), dl2 = d * float(l2 - 32); + const uchar hb = uchar(1 << (4 * n + j)); +#pragma clang loop unroll(full) + for (int l = 0; l < 16; ++l) { + w[l] = dl1 * float(int((q[l] >> (2 * j)) & 3) - ((hm[l] & hb) ? 0 : 4)); + w[16 + l] = dl2 * float(int((q[l + 16] >> (2 * j)) & 3) - ((hm[l + 16] & hb) ? 0 : 4)); + } + p0 = 32 * u - k0; + p1 = p0 + 16; + } +}; + +struct MmaArgs { + uint n_in; + uint n_out; + uint row_bytes; + uint x_stride; // floats between the 8 token rows of x + uint y_stride; +}; +constant int MMA_NTOK [[function_constant(4)]]; + +template +kernel void gemv_mma(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]], + device float* y [[buffer(2)]], constant MmaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, UPC = 128 / UNIT, H = UNIT / 2; + threadgroup float wt[SG][8 * 128]; + threadgroup float ct[SG][64]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int row0 = (int(tg.x) * SG + sg) * 8; + simdgroup_float8x8 C = make_filled_simdgroup_matrix(0.0f); + threadgroup float* w8 = wt[sg]; + for (int k0 = 0; k0 < int(a.n_in); k0 += 128) { + for (int it = lane; it < 8 * UPC; it += 32) { + const int r = it / UPC, uu = it - r * UPC; + const int rr = min(row0 + r, int(a.n_out) - 1); + float w[UNIT]; + int p0, p1; + DW::run(W + ulong(rr) * a.row_bytes, k0 / UNIT + uu, k0, w, p0, p1, tab); +#pragma clang loop unroll(full) + for (int j = 0; j < H; ++j) { + w8[r * 128 + p0 + j] = w[j]; + w8[r * 128 + p1 + j] = w[H + j]; + } + } + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int k8 = 0; k8 < 128; k8 += 8) { + simdgroup_float8x8 A, B; + simdgroup_load(A, w8 + k8, 128); + simdgroup_load(B, x + k0 + k8, a.x_stride, ulong2(0, 0), true); + simdgroup_multiply_accumulate(C, A, B, C); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + } + simdgroup_store(C, ct[sg], 8); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int e = lane; e < 64; e += 32) { + const int r = e >> 3, t = e & 7; + if (t < MMA_NTOK && row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = ct[sg][e]; + } +} + +#define MM(F, S) template [[host_name("gemv_mma_" #F "_" #S)]] kernel decltype(gemv_mma) gemv_mma; +#define MM_SET(F) MM(F, 1) MM(F, 2) MM(F, 4) +MM_SET(T_BF16) MM_SET(T_Q8_0) MM_SET(T_Q4_0) MM_SET(T_Q5_0) MM_SET(T_IQ4_NL) MM_SET(T_IQ4_XS) MM_SET(T_Q2_0) +MM_SET(T_IQ3_S) MM_SET(T_Q3_K) MM_SET(T_Q4_K) MM_SET(T_Q5_K) MM_SET(T_Q6_K) diff --git a/results/round39/model-ab/candidate-kernels/moe.metal b/results/round39/model-ab/candidate-kernels/moe.metal new file mode 100644 index 0000000..878252e --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/moe.metal @@ -0,0 +1,765 @@ +// metal/kernels/moe.metal - routed experts of one layer for a verify window of 1..8 tokens, grouped by expert: +// +// moe_plan one threadgroup: the window's top-k ids -> groups (distinct experts, by id) with their entries +// (token, weight), each (token, k)'s entry index, and the indirect grid sizes of the next two kernels. +// Deterministic (no atomics). +// moe_gu per group, a block of FFN rows: gate and up rows read once, applied to every entry of the group +// (the shared input x is read from cache), h = silu(gate) * up per entry. +// moe_down per group, a block of output rows: ye[entry] = weight * (down . h[entry]). +// moe_combine y[token] = sum over k of ye[entry(token, k)], in k order. +// +// The plan is one int buffer (see P_* below). Expert addressing: MoeArgs (a pack blob or the GGUF tensors): +// [gate rows | up rows | down rows] in GGUF blocks (tools/iq_pack.py's experts.bin layout). + +#define P_NG 0 +#define P_IGU 4 +#define P_IDN 8 +#define P_MAXE 128 +#define P_GEXP 16 +#define P_GSTART (P_GEXP + P_MAXE) +#define P_ETOK (P_GSTART + P_MAXE + 4) +#define P_EW (P_ETOK + P_MAXE) +#define P_TOKENT (P_EW + P_MAXE) +#define P_SPLIT (P_TOKENT + P_MAXE) // function constant 26 only: [0] resident groups, then the indirect grids +#define P_IGU_R (P_SPLIT + 4) // ... of the resident groups (the GPU computes them during the host's turn) +#define P_IDN_R (P_SPLIT + 8) +#define P_IGU_M (P_SPLIT + 12) // ... and of the missing ones (after it) +#define P_IDN_M (P_SPLIT + 16) +#define P_GSLOT (P_SPLIT + 20) +#define P_INTS (P_GSLOT + P_MAXE) + +// Expert e's matrices: gate at W + gate_off + e * gu_stride, up at W + up_off + e * gu_stride, down at +// W + down_off + e * d_stride (a pack's blobs: gate_off 0, strides = the blob; a GGUF: the three tensors). +struct MoeArgs { + ulong gate_off; + ulong up_off; + ulong down_off; + ulong gu_stride; + ulong d_stride; + uint gu_row; // bytes per gate/up row + uint d_row; // bytes per down row + uint n_embd; + uint n_ff; + uint n_tok; + uint topk; + uint n_expert; + uint gx_gu; // threadgroups along the rows for moe_gu (n_ff / rows per threadgroup) + uint gx_dn; // ... for moe_down (n_embd / rows per threadgroup) + uint pad; +}; + +// The SSD-streaming arena (metal/src/engine/stream.cpp): with function constant 20 set, expert e of the layer lies in +// an arena slot, smap[e] = (chunk << 28) | its 16 KiB page in that chunk; the chunks are W (buffer 0) and buffers 7..9 +// (gate_off / up_off / down_off are then offsets inside the slot). The same block is in moe_mm.metal. +#ifndef STRATA_ARENA_METAL +#define STRATA_ARENA_METAL +constant bool kArenaFc [[function_constant(20)]]; +constant bool kArena = is_function_constant_defined(kArenaFc) && kArenaFc; +static inline device const uchar* arena_slot(uint v, device const uchar* W, device const uchar* W1, device const uchar* W2, + device const uchar* W3) { + const uint c = v >> 28; + device const uchar* b = c == 0 ? W : (c == 1 ? W1 : (c == 2 ? W2 : W3)); + return b + (ulong)(v & 0x0FFFFFFFu) * 16384ul; +} +#endif + +// Function constant 26 (the SSD-streaming decode window, with smap): the groups whose expert is in the arena when the +// routing is made come first, the missing ones after them, each part with its own indirect grids (P_SPLIT): the +// resident part runs while the host reads the missing experts. The entries keep their (token, k) mapping, so the +// combine (and every number) is the same; without it the plan is the old one, bit for bit. +constant bool kSplitResFc [[function_constant(26)]]; +constant bool kSplitRes = is_function_constant_defined(kSplitResFc) && kSplitResFc; + +// s_grp, s_ent: 64 each (the upper halves for the missing groups) +static inline void plan_body(device const int* ids, device const float* w, device int* plan, constant MoeArgs& a, + ushort tid, ushort nth, ushort sg, ushort lane, threadgroup int* s_grp, + threadgroup int* s_ent, device const uint* smap = nullptr) { + const int ne = int(a.n_tok * a.topk); + // one thread per expert id (nth == n_expert, at most 1024) + const int e = tid; + int cnt = 0; + if (e < int(a.n_expert)) { + #pragma clang loop unroll_count(10) + for (int i = 0; i < ne; ++i) cnt += ids[i] == e; + } + const int has = cnt > 0; + const bool split = kSplitRes && smap != nullptr; + const int miss = (split && has && smap[e] == 0xFFFFFFFFu) ? 1 : 0; + const int h0 = has - miss, c0 = miss ? 0 : cnt, c1 = miss ? cnt : 0; + const int pg = simd_prefix_exclusive_sum(h0), pe = simd_prefix_exclusive_sum(c0); + const int pg1 = split ? simd_prefix_exclusive_sum(miss) : 0; + const int pe1 = split ? simd_prefix_exclusive_sum(c1) : 0; + if (lane == 31) { + s_grp[sg] = pg + h0; s_ent[sg] = pe + c0; + if (split) { s_grp[32 + sg] = pg1 + miss; s_ent[32 + sg] = pe1 + c1; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + const int nsg = (nth + 31) / 32; + const int g0 = lane < nsg ? s_grp[lane] : 0, n0 = lane < nsg ? s_ent[lane] : 0; + const int og0 = simd_prefix_exclusive_sum(g0), oe0 = simd_prefix_exclusive_sum(n0); + if (lane < nsg) { s_grp[lane] = og0; s_ent[lane] = oe0; } + const int ng0 = simd_broadcast(og0 + g0, 31), ne0 = simd_broadcast(oe0 + n0, 31); + int og = ng0, oe = ne0; + if (split) { + const int g1 = lane < nsg ? s_grp[32 + lane] : 0, n1 = lane < nsg ? s_ent[32 + lane] : 0; + const int og1 = ng0 + simd_prefix_exclusive_sum(g1), oe1 = ne0 + simd_prefix_exclusive_sum(n1); + if (lane < nsg) { s_grp[32 + lane] = og1; s_ent[32 + lane] = oe1; } + og = simd_broadcast(og1 + g1, 31); + oe = simd_broadcast(oe1 + n1, 31); + } + if (lane == 0) { + plan[P_NG] = og; + plan[P_IGU] = int(a.gx_gu); plan[P_IGU + 1] = og; plan[P_IGU + 2] = 1; + plan[P_IDN] = int(a.gx_dn); plan[P_IDN + 1] = og; plan[P_IDN + 2] = 1; + plan[P_GSTART + og] = oe; + if (split) { + plan[P_SPLIT] = ng0; + plan[P_IGU_R] = int(a.gx_gu); plan[P_IGU_R + 1] = ng0; plan[P_IGU_R + 2] = 1; + plan[P_IDN_R] = int(a.gx_dn); plan[P_IDN_R + 1] = ng0; plan[P_IDN_R + 2] = 1; + plan[P_IGU_M] = int(a.gx_gu); plan[P_IGU_M + 1] = og - ng0; plan[P_IGU_M + 2] = 1; + plan[P_IDN_M] = int(a.gx_dn); plan[P_IDN_M + 1] = og - ng0; plan[P_IDN_M + 2] = 1; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (has) { + const int gi = (miss ? s_grp[32 + sg] + pg1 : s_grp[sg] + pg), off = miss ? s_ent[32 + sg] + pe1 : s_ent[sg] + pe; + plan[P_GEXP + gi] = e; + plan[P_GSTART + gi] = off; + if (split && !miss) plan[P_GSLOT + gi] = int(smap[e]); + const bool topk10 = (a.topk == 10u); + int j = 0; + #pragma clang loop unroll_count(10) + for (int i = 0; i < ne; ++i) + if (ids[i] == e) { + plan[P_ETOK + off + j] = topk10 ? ((i * 205) >> 11) : (i / int(a.topk)); + plan[P_EW + off + j] = as_type(w[i]); + plan[P_TOKENT + i] = off + j; + ++j; + } + } +} + +kernel void moe_plan(device const int* ids [[buffer(0)]], device const float* w [[buffer(1)]], + device int* plan [[buffer(2)]], constant MoeArgs& a [[buffer(3)]], + ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup int s_grp[64], s_ent[64]; + plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); +} + +// The router's tail for the window, then the plan, in one threadgroup of n_expert (512) threads: simdgroup t routes +// token t exactly as native_router.cu (pinned llama.cpp topk-moe): softmax over the 512 logits, top 10 by +// probability (equal values: the lower expert index), weights = p / max(sum of the selected, 2^-14), the sum taken +// by winning lane in selection order and reduced with the same xor tree. +#ifndef STRATA_MOCE_FC_METAL +#define STRATA_MOCE_FC_METAL +constant int kMoceProtectFc [[function_constant(24)]]; +constant int kMoceProtect = is_function_constant_defined(kMoceProtectFc) ? kMoceProtectFc : 3; +constant uint kMoceScaleR1BitsFc [[function_constant(25)]]; +constant float kMoceScaleR1 = is_function_constant_defined(kMoceScaleR1BitsFc) ? as_type(kMoceScaleR1BitsFc) : 1.0f; +#endif +kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]], + device float* w [[buffer(2)]], device int* plan [[buffer(3)]], constant MoeArgs& a [[buffer(4)]], + device const uint* smap [[buffer(5)]], + ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup int s_grp[64], s_ent[64]; + if (sg < a.n_tok) { + device const float* l = logits + sg * 512; + const bool e512 = (a.n_expert == 512u); + float v[16]; + float mx = -INFINITY; + float sum = 0.0f; + if (e512) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] = l[lane + 32 * i]; + mx = max(mx, v[i]); + } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] = exp(v[i] - mx); + sum += v[i]; + } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + } else { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + v[i] = idx < int(a.n_expert) ? l[idx] : -INFINITY; + mx = max(mx, v[i]); + } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] = (lane + 32 * i < int(a.n_expert)) ? exp(v[i] - mx) : 0.0f; + sum += v[i]; + } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + if (lane + 32 * i >= int(a.n_expert)) { + v[i] = -INFINITY; + } else { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + } + } + const bool draft_moce = (a.pad & 0x80000000u) != 0u && sg > 0; + const float base_moce = (a.pad & 0x7FFFFFFFu) != 0u ? as_type(a.pad & 0x7FFFFFFFu) : 1.0f; + const float moce_scale = draft_moce ? (base_moce * base_moce) : base_moce; + const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f; + const int r_prot = draft_moce ? 1 : kMoceProtect; + const float r1_scale = kMoceScaleR1; + const float rcp_r1 = r1_scale > 1.0f ? (1.0f / r1_scale) : 1.0f; + const float r_step_scale = (r1_scale > 1.0f && moce_scale > 1.0f) ? (moce_scale * rcp_r1) : moce_scale; + const bool has_r1 = (r1_scale > 1.0f && r_prot > 1); + const bool has_moce = (moce_scale > 1.0f || r1_scale > 1.0f) && smap != nullptr; + const int min_scaled_rank = has_r1 ? 1 : r_prot; + ushort res_mask = 0; + if (has_moce) { + if (e512) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) + if (smap[lane + 32 * i] != 0xFFFFFFFFu) res_mask |= ushort(1u << i); + } else { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if (idx < int(a.n_expert) && smap[idx] != 0xFFFFFFFFu) res_mask |= ushort(1u << i); + } + } + } + ushort scaled_mask = 0; + float selected = 0.0f, ssum = 0.0f; + if (!has_moce) { + for (int rank = 0; rank < 10; ++rank) { + float best = v[0]; + int ex = lane; + #pragma clang loop unroll(full) + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) v[i] = (ex == lane + 32 * i) ? -INFINITY : v[i]; + if ((ex & 31) == lane) { + ids[sg * 10 + rank] = ex; + ssum += best; + } + if (rank == lane) selected = best; + } + } else { + for (int rank = 0; rank < 10; ++rank) { + if (has_r1) { + if (rank == 1) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + if ((res_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) { + v[i] *= r1_scale; + scaled_mask |= ushort(1u << i); + } + } + } else if (rank == r_prot && moce_scale > 1.0f) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + if ((scaled_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) { + v[i] *= r_step_scale; + } + } + } + } else if (rank == r_prot && moce_scale > 1.0f) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + if ((res_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) { + v[i] *= moce_scale; + scaled_mask |= ushort(1u << i); + } + } + } + float best = v[0]; + int ex = lane; + #pragma clang loop unroll(full) + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + float true_prob = best; + if (rank >= min_scaled_rank && best > 0.0f) { + const ushort win_lane = ushort(ex & 31); + 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) { + // without a tail bonus, the rank-1 scale remains on unselected residents. + const bool r1_scaled = has_r1 && (rank < r_prot || moce_scale <= 1.0f); + true_prob = best * (r1_scaled ? rcp_r1 : rcp_moce); + } + } + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) v[i] = (ex == lane + 32 * i) ? -INFINITY : v[i]; + if ((ex & 31) == lane) { + ids[sg * 10 + rank] = ex; + ssum += true_prob; + } + if (rank == lane) selected = true_prob; + } + } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + ssum = max(ssum, 6.103515625e-5f); + const float inv = 1.0f / ssum; + if (lane < 10) w[sg * 10 + lane] = selected * inv; + } + threadgroup_barrier(mem_flags::mem_device); + plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent, smap); +} + +// ---- gate + up + SwiGLU ------------------------------------------------------------------------------------------- +// One pass over a group's rows for all its NE entries (1..8: an expert appears once per token); NE is uniform per +// threadgroup, so the kernel switches once into a body specialised for it. +// +// gu_fused: per row, gate and up decoded together and each entry's x slice loaded once for both (the 32-weight units, +// whose x load is two runs and their sums: the M5 Air, 2560 x 640, NE 3-4: Q4_K / Q5_K / Q6_K / Q8_0 -8..-25%, +// IQ4_XS -31..-36% from NE 2). The 16-weight units keep two x loads a row: the two decoded units held at once cost +// them more (IQ2_XS / IQ3_S at NR 1, NE 3: +15..+21%). The same sums in the same order either way. +template struct gu_fused { + enum { value = fmt::UNIT == 32 && (NE >= 3 || F == T_IQ4_XS) }; +}; + +template +static inline float q2_combine(ushort2 pk, float d, thread const float* xs, float xsum0) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; + #pragma clang loop unroll(full) + for (int j = 0; j < 8; j += 2) { + if (F == T_Q2_0R) { + const uint2 q = uint2(pk) >> (2 * j); + const uint2 bits = ((q | (q << 14)) & 0x00030003u) | 0x64006400u; + const float4 v = float4(as_type(bits) - 1025.0h); + s0 = fma(v.x, xs[j], s0); + s1 = fma(v.y, xs[j + 1], s1); + s2 = fma(v.z, xs[8 + j], s2); + s3 = fma(v.w, xs[8 + j + 1], s3); + } else { + s0 = fma(float(pk.x & ushort(3u << (2 * j))), xs[j], s0); + s1 = fma(float(pk.x & ushort(3u << (2 * j + 2))), xs[j + 1], s1); + s2 = fma(float(pk.y & ushort(3u << (2 * j))), xs[8 + j], s2); + s3 = fma(float(pk.y & ushort(3u << (2 * j + 2))), xs[8 + j + 1], s3); + } + } + const float dot = d * (s0 + s1) + d * (s2 + s3); + return F == T_Q2_0R ? (dot - 0.0f) : (dot - (d * xsum0 + 0.0f)); +} + +static inline float iq_combine(uint4 g, uint s16, float sc, thread const float* xs) { + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 2; ++k) { + const uchar4 v0 = as_type(g[k]), v1 = as_type(g[2 + k]); + const uint m0 = s16 >> (4 * k), m1 = s16 >> (8 + 4 * k); + #pragma clang loop unroll(full) + for (int j = 0; j < 4; j += 2) { + s0 = fma((m0 & (1u << j)) ? -float(v0[j]) : float(v0[j]), xs[4 * k + j], s0); + s1 = fma((m0 & (2u << j)) ? -float(v0[j + 1]) : float(v0[j + 1]), xs[4 * k + j + 1], s1); + s2 = fma((m1 & (1u << j)) ? -float(v1[j]) : float(v1[j]), xs[8 + 4 * k + j], s2); + s3 = fma((m1 & (2u << j)) ? -float(v1[j + 1]) : float(v1[j + 1]), xs[8 + 4 * k + j + 1], s3); + } + } + return sc * (s0 + s1) + sc * (s2 + s3) - 0.0f; +} + +template +static inline float deq_iq_packed(device const uchar* b, int ui, thread uint4& g, thread uint& s16, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + if (F == T_IQ2_S) { + const uint q2 = *(device const ushort*)(b + 2 + 2 * ui); + s16 = *(device const ushort*)(b + 34 + 2 * ui); + const uint qh = uint(b[66 + ib]) >> (4 * h); + const uint i0 = (q2 & 255) | ((qh << 8) & 0x300); + const uint i1 = (q2 >> 8) | ((qh << 6) & 0x300); + const uint2 g0 = ctab ? as_type(iq2s_grid[i0]) : ((threadgroup const uint2*)tab)[i0]; + const uint2 g1 = ctab ? as_type(iq2s_grid[i1]) : ((threadgroup const uint2*)tab)[i1]; + g = uint4(g0.x, g0.y, g1.x, g1.y); + const uint sc = b[74 + ib]; + return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); + } else if (F == T_IQ2_XXS) { + device const uchar* qs = b + 2 + 4 * ib + 2 * ui; + const uint a1 = rd32(b + 2 + 8 * ib + 4), s0 = a1 >> (14 * h); + const uint sa = s0 & 127, sb = (s0 >> 7) & 127; + s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12); + const uint2 g0 = ctab ? as_type(iq2xxs_grid[qs[0]]) : ((threadgroup const uint2*)tab)[qs[0]]; + const uint2 g1 = ctab ? as_type(iq2xxs_grid[qs[1]]) : ((threadgroup const uint2*)tab)[qs[1]]; + g = uint4(g0.x, g0.y, g1.x, g1.y); + return rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f); + } else if (F == T_IQ2_XS) { + device const ushort* qs = (device const ushort*)(b + 2) + 2 * ui; + const uint v0 = qs[0], v1 = qs[1]; + const uint sa = v0 >> 9, sb = v1 >> 9; + s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12); + const uint2 g0 = ctab ? as_type(iq2xs_grid[v0 & 511]) : ((threadgroup const uint2*)tab)[v0 & 511]; + const uint2 g1 = ctab ? as_type(iq2xs_grid[v1 & 511]) : ((threadgroup const uint2*)tab)[v1 & 511]; + g = uint4(g0.x, g0.y, g1.x, g1.y); + const uint sc = b[66 + ib]; + return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); + } else if (F == T_IQ3_S) { + const int l0 = 2 * h, l1 = 2 * h + 1; + const uint qh = b[66 + ib]; + s16 = *(device const ushort*)(b + 74 + 2 * ui); + const uint i0 = uint(b[2 + 8 * ib + 2 * l0]) | ((qh << (8 - 2 * l0)) & 256); + const uint i1 = uint(b[2 + 8 * ib + 2 * l0 + 1]) | ((qh << (7 - 2 * l0)) & 256); + const uint i2 = uint(b[2 + 8 * ib + 2 * l1]) | ((qh << (8 - 2 * l1)) & 256); + const uint i3 = uint(b[2 + 8 * ib + 2 * l1 + 1]) | ((qh << (7 - 2 * l1)) & 256); + g = uint4(g_iq3s(tab, i0), g_iq3s(tab, i1), g_iq3s(tab, i2), g_iq3s(tab, i3)); + const uint sc = b[106 + (ib >> 1)]; + return rdh(b) * fma(float((sc >> (4 * (ib & 1))) & 15), 2.0f, 1.0f); + } else { + const uint aux = rd32(b + 2 + 64 + 4 * ib), s0 = aux >> (14 * h); + const uint sa = s0 & 127, sb = (s0 >> 7) & 127; + s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12); + device const uchar* qs = b + 2 + 4 * ui; + g = uint4(g_iq3xxs(tab, qs[0]), g_iq3xxs(tab, qs[1]), g_iq3xxs(tab, qs[2]), g_iq3xxs(tab, qs[3])); + return rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f); + } +} + +template +static inline void gu_body(device const uchar* g0, device const uchar* u0, device const float* x, + device const float* xs16, device float* h, device const int* etok, int c, int row0, int li, + constant MoeArgs& a, threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT; + float ag[NR][NE], au[NR][NE]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) ag[r][t] = au[r][t] = 0.0f; + const int U = int(a.n_embd) / UNIT; + if (NE == 1) { + const int tok = etok[0]; + x += tok * int(a.n_embd); + xs16 += tok * int(a.n_embd / 16); + } + for (int u = li; u < U; u += LPR) { + if (NE == 1 && UNIT == 16 && F != T_Q2_0 && F != T_Q2_0R) { + float xs[UNIT], xsum[2]; + X::load(x, xs16, u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 c; + Dec::run(g0 + r * a.gu_row, u, q, c, tab); + ag[r][0] += unit_combine(q, c, xs, xsum); + Dec::run(u0 + r * a.gu_row, u, q, c, tab); + au[r][0] += unit_combine(q, c, xs, xsum); + } + } else if (NE == 1) { + float xs[UNIT], xsum[2]; + X::load(x, xs16, u, xs, xsum); + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 c; + Dec::run(g0 + r * a.gu_row, u, q, c, tab); + ag[r][0] += unit_combine(q, c, xs, xsum); + Dec::run(u0 + r * a.gu_row, u, q, c, tab); + au[r][0] += unit_combine(q, c, xs, xsum); + } + } else if (F == T_Q2_0 || F == T_Q2_0R) { + const int boff = (u >> 2) * 18, qoff = 2 + 4 * (u & 3); + ushort2 pkg[NR], pku[NR]; + float dg[NR], du[NR]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + device const uchar* bg = g0 + r * a.gu_row + boff; + device const uchar* bu = u0 + r * a.gu_row + boff; + device const ushort* pg = (device const ushort*)(bg + qoff); + device const ushort* pu = (device const ushort*)(bu + qoff); + pkg[r] = ushort2(pg[0], pg[1]); + pku[r] = ushort2(pu[0], pu[1]); + dg[r] = rdh(bg); + du[r] = rdh(bu); + } + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + ag[r][t] += q2_combine(pkg[r], dg[r], xs, xsum[0]); + au[r][t] += q2_combine(pku[r], du[r], xs, xsum[0]); + } + } + } else if (gu_fused::value) { + for (int r = 0; r < NR; ++r) { + float qg[UNIT], qu[UNIT]; + float4 cg, cu; + Dec::run(g0 + r * a.gu_row, u, qg, cg, tab); + Dec::run(u0 + r * a.gu_row, u, qu, cu, tab); + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + ag[r][t] += unit_combine(qg, cg, xs, xsum); + au[r][t] += unit_combine(qu, cu, xs, xsum); + } + } + } else if (NR == 1 && NE == 4 && (F == T_IQ2_S || F == T_IQ2_XXS || F == T_IQ2_XS || F == T_IQ3_S || F == T_IQ3_XXS)) { + const int boff = (u >> 4) * fmt::BB, ui = (LPR % 16 == 0) ? (li & 15) : (u & 15); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + uint4 gg, gu; + uint sg, su; + const float dg = deq_iq_packed(g0 + r * a.gu_row + boff, ui, gg, sg, tab); + const float du = deq_iq_packed(u0 + r * a.gu_row + boff, ui, gu, su, tab); + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + ag[r][t] += iq_combine(gg, sg, dg, xs); + au[r][t] += iq_combine(gu, su, du, xs); + } + } + } else { + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 c; + Dec::run(g0 + r * a.gu_row, u, q, c, tab); + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + ag[r][t] += unit_combine(q, c, xs, xsum); + } + Dec::run(u0 + r * a.gu_row, u, q, c, tab); + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + au[r][t] += unit_combine(q, c, xs, xsum); + } + } + } + } + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const float gv = lane_sum(ag[r][t]), uv = lane_sum(au[r][t]); + if (li == 0) h[(c + t) * int(a.n_ff) + row0 + r] = gv / (1.0f + exp(-gv)) * uv; + } +} + +// function constant 27, with 26's plan: 1 = only the groups resident at routing time (indirect grid P_IGU_R / P_IDN_R), +// 2 = only the missing ones (P_IGU_M / P_IDN_M); 0 or unset = all (P_IGU / P_IDN) +constant int kPhaseFc [[function_constant(27)]]; +constant int kPhase = is_function_constant_defined(kPhaseFc) ? kPhaseFc : 0; + +// 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; + +template +kernel void moe_gu(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]], + device float* h [[buffer(2)]], device const int* plan [[buffer(3)]], + constant MoeArgs& a [[buffer(4)]], device const float* xs16 [[buffer(5)]], + device const uint* smap [[buffer(6)]], device const uchar* W1 [[buffer(7)]], + device const uchar* W2 [[buffer(8)]], device const uchar* W3 [[buffer(9)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1 && moe_probe != 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int g = int(tg.y) + (kPhase == 2 ? plan[P_SPLIT] : 0); + if (g >= (kPhase == 1 ? plan[P_SPLIT] : plan[P_NG])) return; + const int li = lane % LPR; + const int row0 = ((int(tg.x) * SG + sg) * RG + lane / LPR) * NR; + const ulong ex = (ulong)plan[P_GEXP + g]; + const int c = plan[P_GSTART + g], ne = plan[P_GSTART + g + 1] - c; + device const int* etok = plan + P_ETOK + c; + const uint scode = (kArena && kPhase == 1) ? uint(plan[P_GSLOT + g]) : (kArena ? smap[ex] : 0u); + device const uchar* E = kArena ? arena_slot(scode, W, W1, W2, W3) : W + ex * a.gu_stride; + 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) { + case 1: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + case 2: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + case 3: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + case 4: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + default: + for (int k = 0; k < ne; k += 4) { + const int m = min(4, ne - k); + if (m == 4) gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + else if (m == 3) gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + else if (m == 2) gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + else gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + } + } +} + +// ---- down --------------------------------------------------------------------------------------------------------- +template +static inline void dn_body(device const uchar* d0, device const float* h, device const float* hs16, device float* ye, + device const int* plan, int c, int row0, int li, constant MoeArgs& a, + threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT; + float ac[NR][NE]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + #pragma clang loop unroll(full) + 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) { + if (F == T_Q2_0 || F == T_Q2_0R) { + device const uchar* bu = d0 + (u >> 2) * 18; + const int qoff = 2 + 4 * (u & 3); + ushort2 pk[NR]; + float d[NR]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + device const uchar* b = bu + r * a.d_row; + device const ushort* p = (device const ushort*)(b + qoff); + pk[r] = ushort2(p[0], p[1]); + d[r] = rdh(b); + } + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + float xs[UNIT], xsum[2]; + X::load(h + (c + t) * int(a.n_ff), hs16 + (c + t) * int(a.n_ff / 16), u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + ac[r][t] += q2_combine(pk[r], d[r], xs, xsum[0]); + } + } else if (NE == 1) { + float xs[UNIT], xsum[2]; + X::load(h + c * int(a.n_ff), hs16 + c * int(a.n_ff / 16), u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 sc; + Dec::run(d0 + r * a.d_row, u, q, sc, tab); + ac[r][0] += unit_combine(q, sc, xs, xsum); + } + } else { + float q[NR][UNIT]; + float4 sc[NR]; + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(d0 + r * a.d_row, u, q[r], sc[r], tab); + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + float xs[UNIT], xsum[2]; + X::load(h + (c + t) * int(a.n_ff), hs16 + (c + t) * int(a.n_ff / 16), u, xs, xsum); + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) ac[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + } + #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) + #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const float v = lane_sum(ac[r][t]); + if (li == 0) ye[(c + t) * int(a.n_embd) + row0 + r] = as_type(plan[P_EW + c + t]) * v; + } +} + +template +kernel void moe_down(device const uchar* W [[buffer(0)]], device const float* h [[buffer(1)]], + device float* ye [[buffer(2)]], device const int* plan [[buffer(3)]], + constant MoeArgs& a [[buffer(4)]], device const float* hs16 [[buffer(5)]], + device const uint* smap [[buffer(6)]], device const uchar* W1 [[buffer(7)]], + device const uchar* W2 [[buffer(8)]], device const uchar* W3 [[buffer(9)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int g = int(tg.y) + (kPhase == 2 ? plan[P_SPLIT] : 0); + if (g >= (kPhase == 1 ? plan[P_SPLIT] : plan[P_NG])) return; + const int li = lane % LPR; + const int row0 = ((int(tg.x) * SG + sg) * RG + lane / LPR) * NR; + const ulong ex = (ulong)plan[P_GEXP + g]; + const int c = plan[P_GSTART + g], ne = plan[P_GSTART + g + 1] - c; + const uint scode = (kArena && kPhase == 1) ? uint(plan[P_GSLOT + g]) : (kArena ? smap[ex] : 0u); + device const uchar* d0 = (kArena ? arena_slot(scode, W, W1, W2, W3) : W + ex * a.d_stride) + a.down_off + (ulong)row0 * a.d_row; + switch (ne) { + case 1: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + case 2: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + case 3: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + case 4: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + default: + for (int k = 0; k < ne; k += 4) { + const int m = min(4, ne - k); + if (m == 4) dn_body(d0, h, hs16, ye, plan, c + k, row0, li, a, tab); + else if (m == 3) dn_body(d0, h, hs16, ye, plan, c + k, row0, li, a, tab); + else if (m == 2) dn_body(d0, h, hs16, ye, plan, c + k, row0, li, a, tab); + else dn_body(d0, h, hs16, ye, plan, c + k, row0, li, a, tab); + } + } +} + +kernel void moe_combine(device const float* ye [[buffer(0)]], device float* y [[buffer(1)]], + device const int* plan [[buffer(2)]], constant MoeArgs& a [[buffer(3)]], + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.n_embd, k = gid % a.n_embd; + if (t >= a.n_tok) return; + device const int* te = plan + P_TOKENT + t * a.topk; + float s = 0.0f; + if (a.topk == 10) { +#pragma clang loop unroll(full) + for (uint j = 0; j < 10; ++j) s += ye[te[j] * int(a.n_embd) + k]; + } else { + for (uint j = 0; j < a.topk; ++j) s += ye[te[j] * int(a.n_embd) + k]; + } + y[gid] = s; +} + +// ---- instantiations: moe_gu____, moe_down____ ------------------------------- +#define GU(F, L, R, S) template [[host_name("moe_gu_" #F "_" #L "_" #R "_" #S)]] kernel decltype(moe_gu) moe_gu; +#define DN(F, L, R, S) template [[host_name("moe_down_" #F "_" #L "_" #R "_" #S)]] kernel decltype(moe_down) moe_down; +#define GU_SET(F) GU(F, 32, 1, 4) GU(F, 32, 2, 4) GU(F, 32, 4, 4) GU(F, 32, 2, 2) GU(F, 32, 2, 8) GU(F, 32, 1, 8) \ + GU(F, 16, 1, 4) GU(F, 16, 2, 4) GU(F, 8, 1, 4) GU(F, 8, 2, 4) GU(F, 8, 1, 2) +#define DN_SET(F) DN(F, 8, 1, 4) DN(F, 8, 2, 4) DN(F, 8, 4, 4) DN(F, 8, 1, 8) DN(F, 8, 2, 8) DN(F, 4, 1, 4) DN(F, 4, 2, 4) \ + DN(F, 4, 4, 4) DN(F, 16, 1, 4) DN(F, 16, 2, 4) DN(F, 32, 1, 4) DN(F, 2, 2, 4) +GU_SET(T_Q2_0) +GU_SET(T_IQ3_XXS) +GU_SET(T_IQ2_XXS) +GU_SET(T_IQ2_XS) +GU_SET(T_IQ2_S) +GU_SET(T_IQ3_S) +GU_SET(T_IQ4_XS) // the IQ3_S pack's gate/up in one layer +GU_SET(T_Q4_K) GU_SET(T_Q5_K) GU_SET(T_Q6_K) GU_SET(T_Q8_0) // the UD-Q4/5/6_K_XL packs +DN_SET(T_Q5_1) DN_SET(T_Q8_0) +DN_SET(T_Q2_0) +DN_SET(T_Q2_0R) +DN_SET(T_IQ4_NL) diff --git a/results/round39/model-ab/candidate-kernels/moe_mm.metal b/results/round39/model-ab/candidate-kernels/moe_mm.metal new file mode 100644 index 0000000..03ac90a --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/moe_mm.metal @@ -0,0 +1,1312 @@ +// metal/kernels/moe_mm.metal - routed experts for a PROMPT chunk (after quant/units/mma.metal): a grouped GEMM on +// the matrix units through Metal 4's MetalPerformancePrimitives matmul2d (on the M5 its neural accelerators). +// +// Plan (deterministic, three passes): the chunk's T x topk (token, expert) entries are sorted by expert, then by +// entry index; each expert's run is padded to a multiple of TM rows; tile i = (expert, first padded row). +// mm_hist per block of 1024 entries: a histogram of experts -> bcnt[block][expert] +// mm_scan one threadgroup (thread = expert): block offsets, the padded expert offsets, the tile list +// mm_place per block: an entry's padded row = expert offset + earlier blocks' count + rank in its block; +// srt_tok[row] = token (-1 on padding), srt_w[row] = weight, pos[entry] = row +// GEMMs: a threadgroup takes a tile (TM rows of one expert) and TN output rows; per K chunk of TK it converts the +// tile's activations to fp16 and dequantizes TN x TK weights (ref_quant's arithmetic, then rounded to fp16) into +// threadgroup memory, and matmul2d multiplies them, fp32 accumulation. +// mm_gu H[row][f] = silu(x . gate_f) * (x . up_f), fp32, for the n_ff rows +// mm_down YE[row][d] = h . down_d, fp32 +// mm_combine y[t][d] = sum_k w(t, k) * YE[pos(t, k)][d], k in order +// +// The idea of feeding matmul2d tiles of dequantized weights from threadgroup memory follows yibie/strata-mlx's +// tile_matmul (MIT); the code is our own. +#include +using namespace mpp::tensor_ops; + +struct MmPlanArgs { + uint n_ent; // T * topk + uint topk; + uint n_expert; + uint tm; // the tile height (rows of one expert) + uint n_blocks; // ceil(n_ent / 1024) +}; + +// counts per block: bcnt[block * n_expert + e] +kernel void mm_hist(device const int* ids [[buffer(0)]], device uint* bcnt [[buffer(1)]], + constant MmPlanArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup atomic_uint h[1024]; + if (tid < a.n_expert) atomic_store_explicit(&h[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint i = tg.x * 1024 + tid; + if (i < a.n_ent && ids[i] >= 0) atomic_fetch_add_explicit(&h[ids[i]], 1u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < a.n_expert) + bcnt[tg.x * a.n_expert + tid] = atomic_load_explicit(&h[tid], memory_order_relaxed); +} + +// one threadgroup of n_expert (<= 1024) threads. boff[block][e] = expert e's rows before this block (within e); +// eoff[e] = expert e's first padded row; tiles[i] = (expert, row0); meta[0] = n_tiles, meta[1] = padded rows; +// indirect grids: meta[2..4] = (gx_gu, n_tiles, 1) for the gate/up GEMM, meta[5..7] = (gx_down, n_tiles, 1) for the +// down GEMM (gx_gu in meta[2], gx_down in meta[5] written by the host once), meta[8..10] = (rows, 1, 1) for mm_gather. +// Padding rows [r0 + tot, r0 + rows) of srt_tok (when bound at buffer 6) are cleared to -1 here so mm_clear is not needed. +kernel void mm_scan(device const uint* bcnt [[buffer(0)]], device uint* boff [[buffer(1)]], + device uint* eoff [[buffer(2)]], device int2* tiles [[buffer(3)]], device uint* meta [[buffer(4)]], + constant MmPlanArgs& a [[buffer(5)]], device int* srt_tok [[buffer(6)]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort nsg [[simdgroups_per_threadgroup]]) { + threadgroup uint s_rows[32], s_tiles[32]; + const uint e = tid; + uint tot = 0; + if (e < a.n_expert) + for (uint b = 0; b < a.n_blocks; ++b) { + boff[b * a.n_expert + e] = tot; + tot += bcnt[b * a.n_expert + e]; + } + const uint nt = (tot + a.tm - 1) / a.tm, rows = nt * a.tm; + const uint pr = simd_prefix_exclusive_sum(rows), pt = simd_prefix_exclusive_sum(nt); + if (lane == 31) { s_rows[sg] = pr + rows; s_tiles[sg] = pt + nt; } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint orow = 0, otile = 0, all_rows = 0, all_tiles = 0; + for (uint i = 0; i < nsg; ++i) { + if (i < sg) { orow += s_rows[i]; otile += s_tiles[i]; } + all_rows += s_rows[i]; + all_tiles += s_tiles[i]; + } + if (e < a.n_expert) { + const uint r0 = orow + pr, t0 = otile + pt; + eoff[e] = r0; + for (uint k = 0; k < nt; ++k) tiles[t0 + k] = int2(int(e), int(r0 + k * a.tm)); + if (srt_tok) + for (uint r = tot; r < rows; ++r) srt_tok[r0 + r] = -1; + } + if (tid == 0) { + meta[0] = all_tiles; + meta[1] = all_rows; + meta[3] = all_tiles; // indirect grid of the gate/up GEMM: (meta[2] given, tiles, 1) + meta[4] = 1; + meta[6] = all_tiles; // of the down GEMM: (meta[5] given, tiles, 1) + meta[7] = 1; + meta[8] = all_rows; // of the gather: (rows, 1, 1) + meta[9] = 1; + meta[10] = 1; + } +} + +static inline uint mm_sum_nibbles(uint c) { + c = (c & 0x0F0F0F0Fu) + ((c >> 4) & 0x0F0F0F0Fu); + c = (c + (c >> 8)) & 0x00FF00FFu; + return (c + (c >> 16)) & 0xFFFFu; +} + +// grid (n_blocks), 1024 threads. Two-level deterministic intra-block rank: 32-lane sub-block nibble counts in p_cnt +// plus simd_broadcast scan within the caller's 32-lane simdgroup (no s_id[1024] threadgroup array). +kernel void mm_place(device const int* ids [[buffer(0)]], device const float* w [[buffer(1)]], + device const uint* boff [[buffer(2)]], device const uint* eoff [[buffer(3)]], + device int* srt_tok [[buffer(4)]], device float* srt_w [[buffer(5)]], device int* pos [[buffer(6)]], + constant MmPlanArgs& a [[buffer(7)]], + uint tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint p_cnt[4096]; + const uint n4e = 4 * a.n_expert; + if (a.n_expert == 512) { + atomic_store_explicit(&p_cnt[tid], 0u, memory_order_relaxed); + atomic_store_explicit(&p_cnt[tid + 1024], 0u, memory_order_relaxed); + } else { + for (uint k = tid; k < n4e; k += 1024) atomic_store_explicit(&p_cnt[k], 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint i = tg * 1024 + tid; + const int e = i < a.n_ent ? ids[i] : -1; + const uint sub = tid >> 5; + if (e >= 0 && sub < 31) { + const uint idx = (sub >> 3) * a.n_expert + uint(e); + const uint sh = (sub & 7u) * 4u; + atomic_fetch_add_explicit(&p_cnt[idx], 1u << sh, memory_order_relaxed); + } + uint rank = 0; + #pragma clang loop unroll(full) + for (ushort l = 0; l < 32; ++l) { + const int el = simd_broadcast(e, l); + if (l < lane && el == e) ++rank; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (i >= a.n_ent) return; + if (e < 0) { pos[i] = -1; return; } + for (uint g = 0; g < (sub >> 3); ++g) { + const uint c = atomic_load_explicit(&p_cnt[g * a.n_expert + uint(e)], memory_order_relaxed); + rank += mm_sum_nibbles(c); + } + if (sub & 7u) { + const uint m = (1u << ((sub & 7u) * 4u)) - 1u; + const uint c = atomic_load_explicit(&p_cnt[(sub >> 3) * a.n_expert + uint(e)], memory_order_relaxed) & m; + rank += mm_sum_nibbles(c); + } + const uint row = eoff[e] + boff[tg * a.n_expert + e] + rank; + srt_tok[row] = int(i / a.topk); + srt_w[row] = w[i]; + pos[i] = int(row); +} + +kernel void mm_clear(device int* srt_tok [[buffer(0)]], constant uint& n [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + if (gid < n) srt_tok[gid] = -1; +} + +#ifndef MM_GATHER_ONLY +// ---- the GEMMs ----------------------------------------------------------------------------------------------------- +// weights of one expert: gate at W + gate_off + e * gu_stride (+ up_off for up), down at W + down_off + e * d_stride +struct MmArgs { + ulong gate_off; + ulong up_off; + ulong down_off; + ulong gu_stride; + ulong d_stride; + uint gu_row; // bytes per gate/up row (n_embd values) + uint d_row; // bytes per down row (n_ff values) + uint n_embd; + uint n_ff; +}; + +// TN x TK weights of `nrow`-strided rows into threadgroup half: thread-strided units of the format +template +static inline void stage_w(device const uchar* w0, uint row_bytes, int k0, threadgroup half* dst, + threadgroup const uint* tab, ushort tid, ushort nth) { + constexpr int UNIT = fmt::UNIT, UPC = TK / UNIT, H = UNIT / 2; + static_assert(TK % UNIT == 0, "the K chunk must hold whole units of the format"); + for (int it = tid; it < TN * UPC; it += nth) { + const int r = it / UPC, uu = it - r * UPC; + float wv[UNIT]; + int p0, p1; + DW::run(w0 + ulong(r) * row_bytes, k0 / UNIT + uu, k0, wv, p0, p1, tab); +#pragma clang loop unroll(full) + for (int j = 0; j < H; ++j) { + dst[r * TK + p0 + j] = half(wv[j]); + dst[r * TK + p1 + j] = half(wv[H + j]); + } + } +} + +#ifndef MM_V3_ONLY +template +kernel void mm_gu(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]], + device float* H [[buffer(2)]], device const int* srt_tok [[buffer(3)]], + device const int2* tiles [[buffer(4)]], constant MmArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int NTH = SG * 32; + threadgroup half xa[TM * TK]; + threadgroup half wg[TN * TK]; + threadgroup half wu[TN * TK]; + threadgroup uint tab[fmt::TAB]; + threadgroup int tok[TM]; + load_tab(tab, tid, NTH); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, f0 = int(tg.x) * TN; + for (int m = tid; m < TM; m += NTH) tok[m] = srt_tok[row0 + m]; + device const uchar* g0 = W + a.gate_off + ulong(e) * a.gu_stride + ulong(f0) * a.gu_row; + device const uchar* u0 = W + a.up_off + ulong(e) * a.gu_stride + ulong(f0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false); + matmul2d> op; + tensor, tensor_inline> tA(xa, dextents(TK, TM)); + tensor, tensor_inline> tG(wg, dextents(TK, TN)); + tensor, tensor_inline> tU(wu, dextents(TK, TN)); + auto accg = op.template get_destination_cooperative_tensor(); + auto accu = op.template get_destination_cooperative_tensor(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + for (int it = tid; it < TM * TK; it += NTH) { + const int m = it / TK, k = it - m * TK; + xa[it] = tok[m] >= 0 ? half(x[ulong(tok[m]) * a.n_embd + k0 + k]) : half(0.0f); + } + stage_w(g0, a.gu_row, k0, wg, tab, tid, NTH); + stage_w(u0, a.gu_row, k0, wu, tab, tid, NTH); + threadgroup_barrier(mem_flags::mem_threadgroup); + auto pg = op.template get_destination_cooperative_tensor(); + auto pu = op.template get_destination_cooperative_tensor(); + op.run(tA, tG, pg); + op.run(tA, tU, pu); + for (uint16_t i = 0; i < pg.get_capacity(); ++i) + if (pg.is_valid_element(i)) { accg[i] += pg[i]; accu[i] += pu[i]; } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = accg[i]; + accg[i] = g / (1.0f + exp(-g)) * accu[i]; + } + tensor, tensor_inline> tH(H, dextents(a.n_ff, 1 << 30)); + accg.store(tH.slice(f0, row0)); +} + +template +kernel void mm_down(device const uchar* W [[buffer(0)]], device const float* H [[buffer(1)]], + device float* YE [[buffer(2)]], device const int2* tiles [[buffer(3)]], + constant MmArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int NTH = SG * 32; + threadgroup half ha[TM * TK]; + threadgroup half wd[TN * TK]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, NTH); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, d0 = int(tg.x) * TN; + device const uchar* w0 = W + a.down_off + ulong(e) * a.d_stride + ulong(d0) * a.d_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false); + matmul2d> op; + tensor, tensor_inline> tA(ha, dextents(TK, TM)); + tensor, tensor_inline> tW(wd, dextents(TK, TN)); + auto acc = op.template get_destination_cooperative_tensor(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) { + for (int it = tid; it < TM * TK; it += NTH) { + const int m = it / TK, k = it - m * TK; + ha[it] = half(H[ulong(row0 + m) * a.n_ff + k0 + k]); + } + stage_w(w0, a.d_row, k0, wd, tab, tid, NTH); + threadgroup_barrier(mem_flags::mem_threadgroup); + auto p = op.template get_destination_cooperative_tensor(); + op.run(tA, tW, p); + for (uint16_t i = 0; i < p.get_capacity(); ++i) + if (p.is_valid_element(i)) acc[i] += p[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + tensor, tensor_inline> tY(YE, dextents(a.n_embd, 1 << 30)); + acc.store(tY.slice(d0, row0)); +} + +// grid (T * n_embd / 256), 256 threads +#endif +kernel void mm_combine(device const float* YE [[buffer(0)]], device const int* pos [[buffer(1)]], + device const float* srt_w [[buffer(2)]], device float* y [[buffer(3)]], + constant uint4& a [[buffer(4)]], // (T, n_embd, topk, -) + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.y, d = gid % a.y; + if (t >= a.x) return; + float s = 0.0f; + for (uint k = 0; k < a.z; ++k) { + const int r = pos[t * a.z + k]; + if (r >= 0) s = fma(srt_w[r], YE[ulong(r) * a.y + d], s); + } + y[gid] = s; +} + +#ifndef MM_V3_ONLY +#define MMGU_(F, M, N, K, S) template [[host_name("mm_gu_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu) mm_gu; +#define MMDN_(F, M, N, K, S) template [[host_name("mm_down_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down) mm_down; +#define MMGU(F, M, N, K, S) MMGU_(F, M, N, K, S) +#define MMDN(F, M, N, K, S) MMDN_(F, M, N, K, S) +#define MM_GU_SET(F) MMGU(F, 32, 64, 64, 4) MMGU(F, 64, 64, 64, 4) MMGU(F, 16, 64, 64, 2) MMGU(F, 32, 32, 64, 2) +#define MM_DN_SET(F) MMDN(F, 32, 64, 64, 4) MMDN(F, 64, 64, 64, 4) MMDN(F, 16, 64, 64, 2) MMDN(F, 32, 32, 64, 2) +MM_GU_SET(T_Q2_0) MM_GU_SET(T_IQ3_XXS) MM_GU_SET(T_IQ2_XXS) MM_GU_SET(T_IQ2_XS) MM_GU_SET(T_IQ2_S) MM_GU_SET(T_IQ3_S) +MM_DN_SET(T_Q2_0) MM_DN_SET(T_IQ4_NL) +#endif + +#endif +// ---- v2: activations gathered once per layer as fp16 rows, read by matmul2d from device memory --------------------- +// mm_gather: XS[row][k] = half(x[srt_tok[row]][k]) (0 on padding rows). grid (rows), 256 threads. +kernel void mm_gather(device const float* x [[buffer(0)]], device const int* srt_tok [[buffer(1)]], + device half* XS [[buffer(2)]], constant uint& n [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const int t = srt_tok[tg.x]; + device half* o = XS + ulong(tg.x) * n; + if (n == 2560) { + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const uint k = tid + i * 256; + o[k] = t >= 0 ? half(x[ulong(t) * n + k]) : half(0.0f); + } + } else { + for (uint k = tid; k < n; k += 256) o[k] = t >= 0 ? half(x[ulong(t) * n + k]) : half(0.0f); + } +} + +// mm_gather_h: gathers straight from fp16 rows (mixh) with 128-bit uint4 copies across nth threads per row. +kernel void mm_gather_h(device const half* xh [[buffer(0)]], device const int* srt_tok [[buffer(1)]], + device half* XS [[buffer(2)]], constant uint& n [[buffer(3)]], + uint tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort nth [[threads_per_threadgroup]]) { + const int t = srt_tok[tg]; + device uint4* o8 = (device uint4*)(XS + ulong(tg) * n); + const uint n8 = n >> 3; + if (n == 2560) { + device const uint4* x8 = (device const uint4*)(xh + ulong(max(t, 0)) * n); + if (nth == 32) { + #pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) o8[tid + i * 32] = t >= 0 ? x8[tid + i * 32] : uint4(0u); + return; + } + if (nth == 64) { + #pragma clang loop unroll(full) + for (uint i = 0; i < 5; ++i) o8[tid + i * 64] = t >= 0 ? x8[tid + i * 64] : uint4(0u); + return; + } + if (nth == 160) { + o8[tid] = t >= 0 ? x8[tid] : uint4(0u); + o8[tid + 160] = t >= 0 ? x8[tid + 160] : uint4(0u); + return; + } + if (nth == 256) { + o8[tid] = t >= 0 ? x8[tid] : uint4(0u); + if (tid < 64) o8[tid + 256] = t >= 0 ? x8[tid + 256] : uint4(0u); + return; + } + } + if (t >= 0) { + device const uint4* x8 = (device const uint4*)(xh + ulong(t) * n); + for (uint k = tid; k < n8; k += nth) o8[k] = x8[k]; + } else { + for (uint k = tid; k < n8; k += nth) o8[k] = uint4(0u); + } +} + +#ifndef MM_GATHER_ONLY +// The SSD-streaming arena (as in moe.metal): with function constant 20, expert e's slot is smap[e] in the chunks W, +// W1..W3 (buffers 5..8 of mm_gu2 / mm_down2 / mm_gu3 / mm_down3); gate_off / up_off / down_off are offsets in the slot. +#ifndef STRATA_ARENA_METAL +#define STRATA_ARENA_METAL +constant bool kArenaFc [[function_constant(20)]]; +constant bool kArena = is_function_constant_defined(kArenaFc) && kArenaFc; +static inline device const uchar* arena_slot(uint v, device const uchar* W, device const uchar* W1, device const uchar* W2, + device const uchar* W3) { + const uint c = v >> 28; + device const uchar* b = c == 0 ? W : (c == 1 ? W1 : (c == 2 ? W2 : W3)); + return b + (ulong)(v & 0x0FFFFFFFu) * 16384ul; +} +#endif + +// a probe for the bench (function constant 5): 1 = the weights not decoded (a cheap fill), 2 = no matrix multiply +constant int MM_PROBE [[function_constant(5)]]; +constant int mm_probe = is_function_constant_defined(MM_PROBE) ? MM_PROBE : 0; + +#ifndef MM_V3_ONLY +template +kernel void mm_gu2(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], + device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int NTH = SG * 32; + threadgroup half wg[TN * TK]; + threadgroup half wu[TN * TK]; + threadgroup float ot[TM * TN]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, NTH); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, f0 = int(tg.x) * TN; + device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride; + device const uchar* g0 = E + a.gate_off + ulong(f0) * a.gu_row; + device const uchar* u0 = E + a.up_off + ulong(f0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d> op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + tensor, tensor_inline> tG(wg, dextents(TK, TN)); + tensor, tensor_inline> tU(wu, dextents(TK, TN)); + auto tA = tX.template slice(0, row0); + auto accg = op.template get_destination_cooperative_tensor(); + auto accu = op.template get_destination_cooperative_tensor(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + if (mm_probe == 1) { + for (int it = tid; it < TN * TK; it += NTH) { wg[it] = half(k0 & 7); wu[it] = half(it & 7); } + } else { + stage_w(g0, a.gu_row, k0, wg, tab, tid, NTH); + stage_w(u0, a.gu_row, k0, wu, tab, tid, NTH); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + tA = tX.template slice(k0, row0); + if (mm_probe != 2) { + op.run(tA, tG, accg); + op.run(tA, tU, accu); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = accg[i]; + accg[i] = g / (1.0f + exp(-g)) * accu[i]; + } + tensor, tensor_inline> tO(ot, dextents(TN, TM)); + accg.store(tO); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int it = tid; it < TM * TN; it += NTH) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + m) * a.n_ff + f0 + f] = half(ot[it]); + } +} + +template +kernel void mm_down2(device const uchar* W [[buffer(0)]], device const half* HS [[buffer(1)]], + device float* YE [[buffer(2)]], device const int2* tiles [[buffer(3)]], + constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], + device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int NTH = SG * 32; + threadgroup half wd[TN * TK]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, NTH); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, d0 = int(tg.x) * TN; + device const uchar* w0 = (kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.d_stride) + a.down_off + ulong(d0) * a.d_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d> op; + tensor, tensor_inline> tH((device half*)HS, dextents(a.n_ff, 1 << 30)); + tensor, tensor_inline> tW(wd, dextents(TK, TN)); + auto tA = tH.template slice(0, row0); + auto acc = op.template get_destination_cooperative_tensor(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) { + stage_w(w0, a.d_row, k0, wd, tab, tid, NTH); + threadgroup_barrier(mem_flags::mem_threadgroup); + tA = tH.template slice(k0, row0); + op.run(tA, tW, acc); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + tensor, tensor_inline> tY(YE, dextents(a.n_embd, 1 << 30)); + acc.store(tY.slice(d0, row0)); +} + +#define MMGU2_(F, M, N, K, S) template [[host_name("mm_gu2_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu2) mm_gu2; +#define MMDN2_(F, M, N, K, S) template [[host_name("mm_down2_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down2) mm_down2; +#define MMGU2(F, M, N, K, S) MMGU2_(F, M, N, K, S) +#define MMDN2(F, M, N, K, S) MMDN2_(F, M, N, K, S) +#define MM_GU2_SET(F) MMGU2(F, 32, 64, 32, 2) MMGU2(F, 32, 32, 32, 2) MMGU2(F, 32, 64, 16, 2) MMGU2(F, 64, 32, 16, 2) MMGU2(F, 16, 32, 32, 2) MMGU2(F, 16, 64, 16, 2) MMGU2(F, 16, 128, 32, 4) MMGU2(F, 32, 64, 16, 4) MMGU2(F, 16, 32, 32, 1) MMGU2(F, 16, 64, 32, 4) MMGU2(F, 16, 64, 64, 2) MMGU2(F, 32, 64, 64, 4) MMGU2(F, 64, 64, 64, 4) MMGU2(F, 16, 32, 64, 2) MMGU2(F, 32, 64, 32, 4) MMGU2(F, 64, 64, 32, 4) MMGU2(F, 16, 64, 32, 2) +#define MM_DN2_SET(F) MMDN2(F, 32, 64, 32, 2) MMDN2(F, 32, 32, 32, 2) MMDN2(F, 32, 64, 16, 2) MMDN2(F, 64, 32, 16, 2) MMDN2(F, 16, 32, 32, 2) MMDN2(F, 16, 64, 16, 2) MMDN2(F, 16, 128, 32, 4) MMDN2(F, 32, 64, 16, 4) MMDN2(F, 16, 32, 32, 1) MMDN2(F, 16, 64, 32, 4) MMDN2(F, 16, 64, 64, 2) MMDN2(F, 32, 64, 64, 4) MMDN2(F, 64, 64, 64, 4) MMDN2(F, 16, 32, 64, 2) MMDN2(F, 32, 64, 32, 4) MMDN2(F, 64, 64, 32, 4) MMDN2(F, 16, 64, 32, 2) +#define MM_GU2_U32(F) MMGU2(F, 16, 64, 32, 2) MMGU2(F, 16, 64, 64, 2) MMGU2(F, 32, 64, 32, 2) MMGU2(F, 16, 32, 32, 2) \ + MMGU2(F, 64, 64, 32, 4) +#define MM_DN2_U32(F) MMDN2(F, 16, 64, 32, 2) MMDN2(F, 16, 64, 64, 2) MMDN2(F, 16, 32, 32, 2) MMDN2(F, 32, 32, 32, 2) \ + MMDN2(F, 16, 128, 32, 4) MMDN2(F, 32, 64, 32, 2) MMDN2(F, 64, 64, 32, 4) +// K-quants: the staged K chunk must hold whole units as DW places them (128 wide): TK 128 +#define MM_GU2_KQ(F) MMGU2(F, 16, 32, 128, 2) MMGU2(F, 16, 16, 128, 2) MMGU2(F, 32, 32, 128, 2) MMGU2(F, 16, 32, 128, 4) +MM_GU2_KQ(T_Q4_K) MM_GU2_KQ(T_Q5_K) MM_GU2_KQ(T_Q6_K) MM_GU2_U32(T_Q8_0) MM_DN2_U32(T_Q5_1) MM_DN2_U32(T_Q8_0) +// IQ4_XS (UNIT 32): the K chunk must be >= 32 +MMGU2(T_IQ4_XS, 16, 64, 32, 2) MMGU2(T_IQ4_XS, 16, 64, 64, 2) MMGU2(T_IQ4_XS, 32, 64, 32, 2) MMGU2(T_IQ4_XS, 16, 32, 32, 2) +MMGU2(T_IQ4_XS, 64, 64, 32, 4) +MM_GU2_SET(T_Q2_0) MM_GU2_SET(T_IQ3_XXS) MM_GU2_SET(T_IQ2_XXS) MM_GU2_SET(T_IQ2_XS) MM_GU2_SET(T_IQ2_S) MM_GU2_SET(T_IQ3_S) +MM_DN2_SET(T_Q2_0) +MMDN2(T_IQ4_NL, 16, 64, 32, 2) MMDN2(T_IQ4_NL, 16, 64, 64, 2) MMDN2(T_IQ4_NL, 16, 32, 32, 2) MMDN2(T_IQ4_NL, 32, 32, 32, 2) MMDN2(T_IQ4_NL, 16, 128, 32, 4) + +#endif + +// ---- v3: the weights decoded straight into matmul2d's right-input cooperative tensor (single-simdgroup scope), +// no threadgroup staging and no barriers in the K loop. Each simdgroup owns TN output rows; a threadgroup SG * TN. +// The right input's element layout (tests/coop_layout): runs of 4 consecutive k at one n, so a lane decodes 4 values +// at a time (W4). Activations: the gathered fp16 rows from device memory. +#ifndef STRATA_COOP_K4 +#define STRATA_COOP_K4 +// function constant 7: the right-input cooperative tensor holds runs of 4 consecutive k at one n per lane (the M5's +// layout; tests/coop_layout). false (the M4: runs of 2 n at one k): every element decoded on its own. +constant bool COOP_K4 [[function_constant(7)]]; +constant bool coop_k4 = is_function_constant_defined(COOP_K4) ? COOP_K4 : true; +#endif +template struct W4; +// the four weights k .. k + 3 (k % 4 == 0) of a row, in quant.metal's deq arithmetic (code * scale in fp32), as fp16 +// the threadgroup forms of quant.metal's field readers (blocks staged in threadgroup memory) +static inline uint rd32(threadgroup const uchar* p) { + threadgroup const ushort* s = (threadgroup const ushort*)p; + return uint(s[0]) | (uint(s[1]) << 16); +} +static inline float rdh(threadgroup const uchar* p) { return float(*(threadgroup const half*)p); } + +static inline half4 w4_signed(uint g, uint s, float sc) { + const uchar4 v = as_type(g); + const float4 q = float4((s & 1) ? -float(v[0]) : float(v[0]), (s & 2) ? -float(v[1]) : float(v[1]), + (s & 4) ? -float(v[2]) : float(v[2]), (s & 8) ? -float(v[3]) : float(v[3])); + return half4(q * sc); +} + +template <> struct W4 { // 64 values a block: fp16 d, 16 bytes of 2-bit codes (value v at byte v / 4) + // the codes c land in fp16 mantissas (1024 + c), so c - 1 and its product with d are exact in fp16: bit for bit the + // fp32 decode for every d and code byte (checked exhaustively), and 8-11% off mm_down3 at T 1024-2048 + static inline half4 dec(uint q, half d) { + const uint2 c = (uint2(q | (q << 14), (q >> 4) | (q << 10)) & 0x00030003u) | 0x64006400u; + return (as_type(c) - 1025.0h) * d; + } + template + static inline half4 run_q(device const uchar* b, uint kl4) { + return dec((b + kl4)[OFF], *(device const half*)b); + } + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + return run_q<2>(row + (k >> 6) * 18, uint(k & 63) >> 2); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 98; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint aux = rd32(b + 2 + 64 + 4 * ib); + const float sc = rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f); + return w4_signed(g_iq3xxs(tab, b[2 + ((k & 255) >> 2)]), sign4_iq2xs((aux >> (7 * l)) & 127, hh), sc); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 66; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint a0 = rd32(b + 2 + 8 * ib), a1 = rd32(b + 2 + 8 * ib + 4); + const uint idx = (a0 >> (8 * l)) & 255; + const float sc = rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f); + return w4_signed(g_iq2xxs(tab, idx, hh), sign4_iq2xs((a1 >> (7 * l)) & 127, hh), sc); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 74; + const int ib = (k & 255) >> 5, hh = (k & 7) >> 2; + const uint v = ((device const ushort*)(b + 2))[(k & 255) >> 3]; + const uint scb = b[66 + ib]; + const float sc = rdh(b) * fma(float((scb >> ((k & 16) >> 2)) & 15), 0.25f, 0.125f); + return w4_signed(g_iq2xs(tab, v & 511, hh), sign4_iq2xs(v >> 9, hh), sc); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 82; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint qh = b[66 + ib], o = (k & 255) >> 3; + const uint idx = uint(b[2 + o]) | ((qh << (8 - 2 * l)) & 0x300); + const uint s = b[34 + o]; + const uint scb = b[74 + ib]; + const float sc = rdh(b) * fma(float((scb >> ((k & 16) >> 2)) & 15), 0.25f, 0.125f); + return w4_signed(g_iq2s(tab, idx, hh), hh ? s >> 4 : s, sc); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 110; + const int ib = (k & 255) >> 5, hh = (k & 7) >> 2; + const uint qh = b[66 + ib]; + const uint i = uint(b[2 + ((k & 255) >> 2)]) | ((qh << (8 - ((k & 31) >> 2))) & 256u); + const uint s = b[74 + ((k & 255) >> 3)]; + const uint scb = b[106 + (ib >> 1)]; + const float sc = rdh(b) * float(1 + 2 * ((scb >> (4 * (ib & 1))) & 15)); + return w4_signed(g_iq3s(tab, i), hh ? s >> 4 : s, sc); + } +}; +template <> struct W4 { // 256 a block: d, scales_h, scales_l[4], qs[128]; 32-groups as IQ4_NL's + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 136; + const int j = k & 255, ib = j >> 5, o = j & 15, hi = (j >> 4) & 1; + const uint sh = *(device const ushort*)(b + 2); + const int ls = int((b[4 + (ib >> 1)] >> (4 * (ib & 1))) & 15) | int(((sh >> (2 * ib)) & 3) << 4); + const float s = rdh(b) * float(ls - 32); + device const uchar* qs = b + 8 + 16 * ib; + float4 q; + for (int c = 0; c < 4; ++c) { + const half2 w = as_type(tab[qs[o + c]]); + q[c] = float(hi ? w.y : w.x); + } + return half4(q * s); + } +}; +template <> struct W4 { // 32 a block: d, 16 bytes; low nibbles are elements 0..15, high 16..31 + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 5) * 18; + const int j = k & 31, hi = j >> 4, o = j & 15; + const float d = rdh(b); + float4 q; + for (int c = 0; c < 4; ++c) { + const half2 w = as_type(tab[b[2 + o + c]]); + q[c] = float(hi ? w.y : w.x); + } + return half4(q * d); + } +}; + +// W8: the runs k .. k + 3 and k + 16 .. k + 19 of a row (k % 32 < 16): one 32-group, so the per-group fields (scale, +// sign word) are read once. The generic one is two W4. +// ---- the K-quant / legacy-quant W4 decoders (the UD-Q4/5/6_K_XL packs' experts): values k .. k + 3 (k % 4 == 0) +// of a row in ref_quant's arithmetic (float), then fp16 +template <> struct W4 { // 32 a block: d, int8 qs[32] + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 5) * 34; + const float d = rdh(b); + device const char* q = (device const char*)(b + 2 + (k & 31)); + return half4(float4(float(q[0]), float(q[1]), float(q[2]), float(q[3])) * d); + } +}; +template <> struct W4 { // 32 a block: d, m, qh (32 bits), qs[16] (low nibbles 0..15, high 16..31) + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 5) * 24; + const float d = rdh(b), m = rdh(b + 2); + const uint hb = rd32(b + 4); + const int j = k & 31, o = j & 15, hi = j >> 4; + float4 w; + for (int c = 0; c < 4; ++c) { + const uint qv = b[8 + o + c]; + const int x = int(hi ? (qv >> 4) : (qv & 15)) | int(((hb >> (j + c)) & 1) << 4); + w[c] = float(x) * d + m; + } + return half4(w); + } +}; +template <> struct W4 { // 256 a block: d, dmin, scales[12], qs[128] (64-chunks: low nibbles then high) + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 8) * 144; + const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31; + float s, mm; + scale_min_k4(2 * c + hi, b + 4, s, mm); + const float d1 = rdh(b) * s, m1 = rdh(b + 2) * mm; + device const uchar* q = b + 16 + 32 * c + l; + float4 w; + for (int i = 0; i < 4; ++i) w[i] = d1 * float(hi ? (q[i] >> 4) : (q[i] & 15)) - m1; + return half4(w); + } +}; +template <> struct W4 { // 256 a block: d, dmin, scales[12], qh[32], qs[128] + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 8) * 176; + const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31; + float s, mm; + scale_min_k4(2 * c + hi, b + 4, s, mm); + const float d1 = rdh(b) * s, m1 = rdh(b + 2) * mm; + device const uchar* qh = b + 16 + l; + device const uchar* ql = b + 48 + 32 * c + l; + const uint bit = 1u << (2 * c + hi); + float4 w; + for (int i = 0; i < 4; ++i) + w[i] = d1 * float(int(hi ? (ql[i] >> 4) : (ql[i] & 15)) + ((qh[i] & bit) ? 16 : 0)) - m1; + return half4(w); + } +}; +template <> struct W4 { // 256 a block: ql[128], qh[64], int8 scales[16], d; halves of 128 + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 8) * 210; + const uint j = uint(k & 255), j4 = j >> 4; + const uint la = rd32(b + (((j >> 1) & 64u) | (j & 63u))); + const uint ha = rd32(b + 128u + (((j >> 2) & 32u) | (j & 31u))); + const float ds = rdh(b + 208) * float(((device const char*)(b + 192))[j4]); + const uint qa = ((la >> (j4 & 4u)) & 0x0F0F0F0Fu) | (((ha >> (j4 & 6u)) & 0x03030303u) << 4); + return half4(ds * (float4(as_type(qa)) - 32.0f)); + } +}; + +template struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + a = W4::run(row, k, tab); + b = W4::run(row, k + 16, tab); + } +}; +// (a W8 sharing the block scale was measured slower than the two W4 of the generic one: 5.5 vs 3.5-4.2 ms) +template <> struct W8 { + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + const P bl = row + (k >> 8) * 98; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint aux = rd32(bl + 2 + 64 + 4 * ib); + const float sc = rdh(bl) * fma(float(aux >> 28), 0.5f, 0.25f); + const P qs = bl + 2 + ((k & 255) >> 2); + if (mm_probe == 3) { // bench probe: no threadgroup-table lookups (wrong values) + a = w4_signed(qs[0] * 0x01010101u, aux, sc); + b = w4_signed(qs[4] * 0x01010101u, aux >> 7, sc); + return; + } + const uint j0 = (aux >> (7 * l)) & 127, j1 = (aux >> (7 * (l + 2))) & 127; + a = w4_signed(g_iq3xxs(tab, qs[0]), sign4_iq2xs(j0, hh), sc); + b = w4_signed(g_iq3xxs(tab, qs[4]), sign4_iq2xs(j1, hh), sc); + } +}; +template <> struct W8 { + template + static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { + const int ib = (k & 255) >> 5, l = (k & 31) >> 3; + const P p = bl + 2 + 8 * ib; + const uint a1 = rd32(p + 4); + const float sc = d * fma(float(a1 >> 28), 0.25f, 0.125f); + const uint i0 = p[l], i1 = p[l + 2]; + const uint j0 = (a1 >> (7 * l)) & 127, j1 = (a1 >> (7 * (l + 2))) & 127; + threadgroup const uint* th = tab + hh; + a = w4_signed(ctab ? uint(iq2xxs_grid[i0] >> (32 * hh)) : th[2 * i0], sign4_iq2xs(j0, hh), sc); + b = w4_signed(ctab ? uint(iq2xxs_grid[i1] >> (32 * hh)) : th[2 * i1], sign4_iq2xs(j1, hh), sc); + } + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + const P bl = row + (k >> 8) * 66; + block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b); + } +}; +template <> struct W8 { + static inline void block(device const uchar* bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { + const int ib = (k & 255) >> 5; + device const ushort* q = (device const ushort*)(bl + 2) + ((k & 255) >> 3); + const uint v0 = q[0], v1 = q[2]; + const uint scb = bl[66 + ib]; + const float sc0 = d * fma(float(scb & 15), 0.25f, 0.125f); // l < 2 (k % 32 < 16): runs l, l + 2 take nibbles 0, 1 + const float sc1 = d * fma(float(scb >> 4), 0.25f, 0.125f); + threadgroup const uint* th = tab + hh; + a = w4_signed(ctab ? uint(iq2xs_grid[v0 & 511] >> (32 * hh)) : th[2 * (v0 & 511)], sign4_iq2xs(v0 >> 9, hh), sc0); + b = w4_signed(ctab ? uint(iq2xs_grid[v1 & 511] >> (32 * hh)) : th[2 * (v1 & 511)], sign4_iq2xs(v1 >> 9, hh), sc1); + } + static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 74; + block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b); + } +}; +template <> struct W8 { + // grid words read through tab + hh and the sign byte as s >> 4 * hh (not a select): 8% off mm_gu3 at T 2048; + // the same for IQ2_XXS / IQ2_XS, whose sign4_iq2xs select stays (a shift form there measured 10% slower) + template + static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { + const int ib = (k & 255) >> 5, l = (k & 31) >> 3; + threadgroup const uint* th = tab + hh; + const uint sh = uint(hh) << 2; + const uint qh = bl[66 + ib], scb = bl[74 + ib]; + const P p = bl + 2 + ((k & 255) >> 3); + const uint i0 = uint(p[0]) | ((qh << (8 - 2 * l)) & 0x300); + const uint i1 = uint(p[2]) | ((qh << (4 - 2 * l)) & 0x300); + const uint s0 = p[32], s1 = p[34]; + const float sc0 = d * fma(float(scb & 15), 0.25f, 0.125f); + const float sc1 = d * fma(float(scb >> 4), 0.25f, 0.125f); + a = w4_signed(ctab ? uint(iq2s_grid[i0] >> (32 * hh)) : th[2 * i0], s0 >> sh, sc0); + b = w4_signed(ctab ? uint(iq2s_grid[i1] >> (32 * hh)) : th[2 * i1], s1 >> sh, sc1); + } + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + const P bl = row + (k >> 8) * 82; + block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b); + } +}; +template <> struct W8 { + template + static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { + const int ib = k >> 5; + const uint qh = bl[66 + ib], scb = bl[106 + (ib >> 1)]; + const float sc = d * fma(float((scb >> (4 * (ib & 1))) & 15), 2.0f, 1.0f); + const P qs = bl + 2 + (k >> 2); + const int sh0 = 8 - ((k & 31) >> 2); + const uint i0 = uint(qs[0]) | ((qh << sh0) & 256u); + const uint i1 = uint(qs[4]) | ((qh << (sh0 - 4)) & 256u); + const P sp = bl + 74 + (k >> 3); + const uint s0 = sp[0], s1 = sp[2]; + if (mm_probe == 5) { // bench probe: the loads and the scale, no grid lookup / sign math (wrong values) + a = half4(half(float(i0 ^ s0) * sc)); + b = half4(half(float(i1 ^ s1) * sc)); + return; + } + a = w4_signed(g_iq3s(tab, i0), s0 >> (4 * hh), sc); + b = w4_signed(g_iq3s(tab, i1), s1 >> (4 * hh), sc); + } + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + if (mm_probe == 7) { + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint4 r = *(device const uint4*)((ulong(row) + ulong(k >> 5) * 16) & ~ulong(15)); + const uint qhv = r.z & 0xFF, scv = (r.z >> 8) & 0xFF; + const float scr = float(as_type(ushort(r.w & 0xFFFF))) * float(1 + 2 * ((scv >> (4 * (ib & 1))) & 15)); + const uint j0 = ((r.x >> (8 * (l & 1))) & 0xFF) | ((qhv << (8 - 2 * l)) & 256); + const uint j1 = ((r.x >> (16 + 8 * (l & 1))) & 0xFF) | ((qhv << (7 - 2 * l)) & 256); + a = w4_signed(tab[j0], (r.y >> (8 * hh)) & 0xFF, scr); + b = w4_signed(tab[j1], (r.y >> (16 + 8 * hh)) & 0xFF, scr); + return; + } + const P bl = row + (k >> 8) * 110; + block(bl, k & 255, rdh(bl), tab, (k & 7) >> 2, a, b); + } +}; + +// W8 for the K-quants and Q5_1: the lane's runs k and k + 16 (k % 32 < 16) share the block's scales / mins / high +// bits; the 4-byte code runs read as one uint where the block layout keeps them 4-aligned (Q4_K 144-, Q5_K 176-, +// Q5_1 24-byte blocks: rows of whole blocks start 4-aligned). The same per-value arithmetic as W4. +static inline float4 nib4(uint v, int hi) { const uint s = hi ? (v >> 4) : v; return float4(as_type(s & 0x0F0F0F0Fu)); } +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 144; + const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31; + float s, mm; + scale_min_k4(2 * c + hi, bl + 4, s, mm); + const float d1 = rdh(bl) * s, m1 = rdh(bl + 2) * mm; + device const uint* q = (device const uint*)(bl + 16 + 32 * c + l); + const float4 qa = nib4(q[0], hi), qb = nib4(q[4], hi); + a = half4(d1 * qa - m1); + b = half4(d1 * qb - m1); + } +}; +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 176; + const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31; + float s, mm; + scale_min_k4(2 * c + hi, bl + 4, s, mm); + const float d1 = rdh(bl) * s, m1 = rdh(bl + 2) * mm; + device const uint* qh = (device const uint*)(bl + 16 + l); + device const uint* ql = (device const uint*)(bl + 48 + 32 * c + l); + const uint sh = uint(2 * c + hi); + const float4 ha = float4(as_type((qh[0] >> sh) & 0x01010101u)) * 16.0f; + const float4 hb = float4(as_type((qh[4] >> sh) & 0x01010101u)) * 16.0f; + a = half4(d1 * (nib4(ql[0], hi) + ha) - m1); + b = half4(d1 * (nib4(ql[4], hi) + hb) - m1); + } +}; +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 5) * 24; + const float d = rdh(bl), m = rdh(bl + 2); + const uint hb = rd32(bl + 4); + const int j = k & 15; // runs at j and j + 16: the low nibbles of qs[j..j+3], then the high ones + const uint v = *(device const uint*)(bl + 8 + j); + float4 wa, wb; + for (int c = 0; c < 4; ++c) { + const uint qv = (v >> (8 * c)) & 0xFFu; + wa[c] = float(int(qv & 15) | int(((hb >> (j + c)) & 1) << 4)) * d + m; + wb[c] = float(int(qv >> 4) | int(((hb >> (j + 16 + c)) & 1) << 4)) * d + m; + } + a = half4(wa); + b = half4(wb); + } +}; +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 210; + const uint j = uint(k & 255), j4 = j >> 4; // l = j & 31 < 16; the pair at l + 16 + device const uchar* ql = bl + (((j >> 1) & 64u) | (j & 63u)); + device const uchar* qh = bl + 128u + (((j >> 2) & 32u) | (j & 31u)); + const float d = rdh(bl + 208); + device const char* sc = (device const char*)(bl + 192u + j4); + const float dsa = d * float(sc[0]), dsb = d * float(sc[1]); + const uint la = rd32(ql), lb = rd32(ql + 16); + const uint ha = rd32(qh), hb = rd32(qh + 16); + const uint sl = j4 & 4u, shh = j4 & 6u; + const uint qa = ((la >> sl) & 0x0F0F0F0Fu) | (((ha >> shh) & 0x03030303u) << 4); + const uint qb = ((lb >> sl) & 0x0F0F0F0Fu) | (((hb >> shh) & 0x03030303u) << 4); + a = half4(dsa * (float4(as_type(qa)) - 32.0f)); + b = half4(dsb * (float4(as_type(qb)) - 32.0f)); + } +}; + +template +kernel void mm_gu3(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], + device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float ot[SG][TM * TN]; + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN; + device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride; + device const uchar* g0 = E + a.gate_off + ulong(n0) * a.gu_row; + device const uchar* u0 = E + a.up_off + ulong(n0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + auto rw = op.template get_right_input_cooperative_tensor(); + auto tA = tX.template slice(0, row0); + auto accg = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + auto accu = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + if (coop_k4 && TK == 32 && mm_probe != 1) { + const uint16_t half_cap = rw.get_capacity() / 2; + const auto ix0 = rw.get_multidimensional_index(uint16_t(0)); + const int kl = ix0[0], hh = (kl >> 2) & 1; + const ulong roff0 = ulong(ix0[1]) * a.gu_row; + device const uchar* g0_l = g0 + roff0; + device const uchar* u0_l = u0 + roff0; + const uint step8 = 8u * a.gu_row; + // IQ2 formats: a 256 super-block at a time, its d read once; the lane's hh is fixed for the whole K loop + constexpr bool sb = F == T_IQ2_S || F == T_IQ2_XXS || F == T_IQ2_XS; + for (int kb = 0; kb < int(a.n_embd); kb += (sb ? 256 : TK)) { + const uint sboff = uint(kb >> 8) * fmt::BB; + float dg[TN * TK / 256], du[TN * TK / 256]; + if (sb) { +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap / 4; ++i) { + dg[i] = rdh(g0_l + i * step8 + sboff); + du[i] = rdh(u0_l + i * step8 + sboff); + } + } + for (int ki = 0; ki < (sb ? 256 : TK); ki += TK) { + const int k0 = kb + ki; + tA = tX.template slice(k0, row0); + const uint boff = F == T_IQ2_XS ? sboff : (fmt::QK == 256 ? uint(k0 >> 8) * fmt::BB : (fmt::QK == 64 ? uint(k0 >> 6) * fmt::BB : 0u)); + const int j_in = fmt::QK == 256 ? ((k0 + kl) & 255) : (fmt::QK == 64 ? ((k0 + kl) & 63) : (k0 + kl)); + device const uchar* rgw = fmt::QK >= 64 ? g0_l + boff : g0_l; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, rgw += step8) { + half4 wg, wg2; + if (F == T_IQ2_S) W8::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2); + else if (F == T_IQ2_XXS) W8::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2); + else if (F == T_IQ2_XS) W8::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2); + else if (F == T_IQ3_S) W8::block(rgw, j_in, rdh(rgw), tab, hh, wg, wg2); + else W8::run(rgw, j_in, tab, wg, wg2); + rw[i] = wg.x; rw[i + 1] = wg.y; rw[i + 2] = wg.z; rw[i + 3] = wg.w; + const uint16_t j = i + half_cap; + rw[j] = wg2.x; rw[j + 1] = wg2.y; rw[j + 2] = wg2.z; rw[j + 3] = wg2.w; + } + op.run(tA, rw, accg); + device const uchar* ruw = fmt::QK >= 64 ? u0_l + boff : u0_l; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, ruw += step8) { + half4 wu, wu2; + if (F == T_IQ2_S) W8::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2); + else if (F == T_IQ2_XXS) W8::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2); + else if (F == T_IQ2_XS) W8::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2); + else if (F == T_IQ3_S) W8::block(ruw, j_in, rdh(ruw), tab, hh, wu, wu2); + else W8::run(ruw, j_in, tab, wu, wu2); + rw[i] = wu.x; rw[i + 1] = wu.y; rw[i + 2] = wu.z; rw[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rw[j] = wu2.x; rw[j + 1] = wu2.y; rw[j + 2] = wu2.z; rw[j + 3] = wu2.w; + } + op.run(tA, rw, accu); + } + } + } else + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + const uint16_t cap = rw.get_capacity(), half_cap = cap / 2; + tA = tX.template slice(k0, row0); + if (!coop_k4) { + for (uint16_t i = 0; i < cap; ++i) { + const auto ix = rw.get_multidimensional_index(i); + const int k = k0 + ix[0], n = ix[1]; + rw[i] = W4::run(g0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3]; + } + op.run(tA, rw, accg); + for (uint16_t i = 0; i < cap; ++i) { + const auto ix = rw.get_multidimensional_index(i); + const int k = k0 + ix[0], n = ix[1]; + rw[i] = W4::run(u0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3]; + } + op.run(tA, rw, accu); + } else { +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rw.get_multidimensional_index(i); + const auto iy = rw.get_multidimensional_index(i + half_cap); + const int k = k0 + ix[0], n = ix[1]; + device const uchar* rgw = g0 + ulong(n) * a.gu_row; + half4 wg, wg2; + if (mm_probe == 1) { + wg = half4(half(n & 7)); wg2 = half4(half(k & 7)); + } else if (TK == 32) { + W8::run(rgw, k, tab, wg, wg2); + } else { + const int k2 = k0 + iy[0]; + wg = W4::run(rgw, k, tab); + wg2 = W4::run(g0 + ulong(iy[1]) * a.gu_row, k2, tab); + } + rw[i] = wg.x; rw[i + 1] = wg.y; rw[i + 2] = wg.z; rw[i + 3] = wg.w; + const uint16_t j = i + half_cap; + rw[j] = wg2.x; rw[j + 1] = wg2.y; rw[j + 2] = wg2.z; rw[j + 3] = wg2.w; + } + op.run(tA, rw, accg); +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rw.get_multidimensional_index(i); + const auto iy = rw.get_multidimensional_index(i + half_cap); + const int k = k0 + ix[0], n = ix[1]; + device const uchar* ruw = u0 + ulong(n) * a.gu_row; + half4 wu, wu2; + if (mm_probe == 1) { + wu = half4(half(n & 7)); wu2 = half4(half(k & 7)); + } else if (TK == 32) { + W8::run(ruw, k, tab, wu, wu2); + } else { + const int k2 = k0 + iy[0]; + wu = W4::run(ruw, k, tab); + wu2 = W4::run(u0 + ulong(iy[1]) * a.gu_row, k2, tab); + } + rw[i] = wu.x; rw[i + 1] = wu.y; rw[i + 2] = wu.z; rw[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rw[j] = wu2.x; rw[j + 1] = wu2.y; rw[j + 2] = wu2.z; rw[j + 3] = wu2.w; + } + op.run(tA, rw, accu); + } + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = accg[i]; + accg[i] = g / (1.0f + exp(-g)) * accu[i]; + } + tensor, tensor_inline> tO(ot[sg], dextents(TN, TM)); + accg.store(tO); + simdgroup_barrier(mem_flags::mem_threadgroup); + if ((F == T_IQ2_XS || F == T_IQ3_XXS || F == T_IQ2_XXS || F == T_Q2_0 || F == T_IQ2_S || F == T_IQ3_S) && 128 % TN == 0 && (TM * TN) % 128 == 0) { + // a lane stores 4 adjacent outputs a pass, 128 a pass over the simdgroup + constexpr int RP = 128 / TN; // rows a pass + const int mr = lane / (TN / 4), f = (lane % (TN / 4)) * 4; + threadgroup const float4* ot_p = (threadgroup const float4*)(&ot[sg][mr * TN + f]); + device half* hs_p = HS + ulong(row0 + mr) * a.n_ff + n0 + f; + const uint hs_step = uint(RP) * a.n_ff; +#pragma clang loop unroll(full) + for (int m = 0; m < TM / RP; ++m, hs_p += hs_step) + *(device half4*)hs_p = half4(ot_p[m * (RP * TN / 4)]); + } else + for (int it = lane; it < TM * TN; it += 32) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + m) * a.n_ff + n0 + f] = half(ot[sg][it]); + } +} + +template +kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS [[buffer(1)]], + device float* YE [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], + device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint tab[fmt::TAB]; + if (fmt::TAB > 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN; + device const uchar* w0 = (kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.d_stride) + a.down_off + ulong(n0) * a.d_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tH((device half*)HS, dextents(a.n_ff, 1 << 30)); + auto rw = op.template get_right_input_cooperative_tensor(); + auto tA = tH.template slice(0, row0); + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + const auto ix0 = rw.get_multidimensional_index(uint16_t(0)); + const int kl = ix0[0]; + device const uchar* w0_l = w0 + ulong(ix0[1]) * a.d_row; + const uint step8 = 8u * a.d_row; + if (coop_k4 && F == T_Q2_0 && TK == 32) { + const uint16_t half_cap = rw.get_capacity() / 2; + const uint kl4 = uint(kl) >> 2; + for (int k0 = 0; k0 < int(a.n_ff); k0 += 64, w0_l += 18) { + device const uchar* rdw = w0_l; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) { + const half4 w = W4::run_q<2>(rdw, kl4); + const half4 w2 = W4::run_q<6>(rdw, kl4); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + const uint16_t j = i + half_cap; + rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; + } + tA = tH.template slice(k0, row0); + op.run(tA, rw, acc); + rdw = w0_l; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) { + const half4 w = W4::run_q<10>(rdw, kl4); + const half4 w2 = W4::run_q<14>(rdw, kl4); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + const uint16_t j = i + half_cap; + rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; + } + tA = tH.template slice(k0 + 32, row0); + op.run(tA, rw, acc); + } + } else + for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) { + if (!coop_k4) { + for (uint16_t i = 0; i < rw.get_capacity(); ++i) { + const auto ix = rw.get_multidimensional_index(i); + const int k = k0 + ix[0]; + rw[i] = W4::run(w0 + ulong(ix[1]) * a.d_row, k & ~3, tab)[k & 3]; + } + } else if (TK == 32) { + // the paired runs (element i + cap / 2 is (k + 16, n), the COOP_K4 layout): one W8 a pair + const uint16_t half_cap = rw.get_capacity() / 2; + const uint boff = fmt::QK == 64 ? uint(k0 >> 6) * fmt::BB : (fmt::QK == 32 ? uint(k0 >> 5) * fmt::BB : 0u); + const int j_in = fmt::QK == 64 ? ((k0 + kl) & 63) : (fmt::QK == 32 ? kl : (k0 + kl)); + device const uchar* rdw = w0_l + boff; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) { + half4 w, w2; + W8::run(rdw, j_in, tab, w, w2); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + const uint16_t j = i + half_cap; + rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rw.get_capacity(); i += 4) { + const auto ix = rw.get_multidimensional_index(i); + const half4 w = W4::run(w0 + ulong(ix[1]) * a.d_row, k0 + ix[0], tab); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + } + tA = tH.template slice(k0, row0); + op.run(tA, rw, acc); + } + tensor, tensor_inline> tY(YE, dextents(a.n_embd, 1 << 30)); + acc.store(tY.slice(n0, row0)); +} + +#define MMGU3_(F, M, N, K, S) template [[host_name("mm_gu3_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu3) mm_gu3; +#define MMDN3_(F, M, N, K, S) template [[host_name("mm_down3_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down3) mm_down3; +#define MMGU3(F, M, N, K, S) MMGU3_(F, M, N, K, S) +#define MMDN3(F, M, N, K, S) MMDN3_(F, M, N, K, S) +#define MM3_ONE(F, M, N, K, S) MMGU3(F, M, N, K, S) MMDN3(F, M, N, K, S) +#define MM3_GU(F) MMGU3(F, 16, 64, 32, 2) MMGU3(F, 16, 32, 32, 4) MMGU3(F, 32, 64, 32, 2) MMGU3(F, 16, 64, 64, 2) \ + MMGU3(F, 16, 32, 16, 2) MMGU3(F, 32, 32, 32, 2) MMGU3(F, 16, 64, 32, 1) MMGU3(F, 16, 64, 32, 4) MMGU3(F, 16, 32, 16, 4) MMGU3(F, 16, 32, 32, 5) MMGU3(F, 16, 32, 16, 5) MMGU3(F, 16, 16, 32, 4) MMGU3(F, 16, 16, 32, 8) MMGU3(F, 16, 16, 16, 8) +#define MM3_DN(F) MMDN3(F, 16, 64, 32, 2) MMDN3(F, 16, 32, 32, 4) MMDN3(F, 32, 64, 32, 2) MMDN3(F, 16, 64, 64, 2) \ + MMDN3(F, 16, 32, 16, 2) MMDN3(F, 32, 32, 32, 2) MMDN3(F, 16, 64, 32, 1) MMDN3(F, 16, 64, 32, 4) MMDN3(F, 16, 64, 16, 2) MMDN3(F, 16, 32, 16, 4) MMDN3(F, 16, 64, 32, 5) MMDN3(F, 16, 64, 16, 4) +MM3_GU(T_Q4_K) MM3_GU(T_Q5_K) MM3_GU(T_Q6_K) MM3_GU(T_Q8_0) MM3_GU(T_IQ4_XS) MM3_GU(T_Q2_0) MM3_GU(T_IQ3_XXS) MM3_GU(T_IQ2_XXS) MM3_GU(T_IQ2_XS) MM3_GU(T_IQ2_S) MM3_GU(T_IQ3_S) +MM3_DN(T_Q2_0) MM3_DN(T_IQ4_NL) MM3_DN(T_Q5_1) MM3_DN(T_Q8_0) + +#ifndef MM_V3_ONLY +// mm_gu4: mm_gu3 with each decoded right operand (TN x TK weights of gate and of up) applied to MS 16-row token +// tiles (a tile of 16 * MS rows from the plan, tm = 16 * MS): the decode, which bounds mm_gu3, is shared by MS times +// as many token rows (more padding rows per expert, but the matrix ops are cheap next to the decode). The same +// per-row arithmetic as mm_gu3 (each 16-row op is mm_gu3's): bitwise for the rows both compute. COOP_K4 layout only. +template +kernel void mm_gu4(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int TM = 16; + threadgroup float ot[SG][TM * TN]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN; + device const uchar* g0 = W + a.gate_off + ulong(e) * a.gu_stride + ulong(n0) * a.gu_row; + device const uchar* u0 = W + a.up_off + ulong(e) * a.gu_stride + ulong(n0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + auto rg = op.template get_right_input_cooperative_tensor(); + auto ru = op.template get_right_input_cooperative_tensor(); + auto tA0 = tX.template slice(0, row0); + auto tA1 = tX.template slice(0, row0 + TM); +#define MM4_DST op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>() + auto accg0 = MM4_DST; + auto accu0 = MM4_DST; + auto accg1 = MM4_DST; + auto accu1 = MM4_DST; +#undef MM4_DST + using DT = decltype(accg0); + for (uint16_t i = 0; i < accg0.get_capacity(); ++i) + if (accg0.is_valid_element(i)) { accg0[i] = 0.0f; accu0[i] = 0.0f; accg1[i] = 0.0f; accu1[i] = 0.0f; } + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + const uint16_t cap = rg.get_capacity(), half_cap = cap / 2; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rg.get_multidimensional_index(i); + const int k = k0 + ix[0], n = ix[1]; + device const uchar* rgw = g0 + ulong(n) * a.gu_row; + device const uchar* ruw = u0 + ulong(n) * a.gu_row; + half4 wg, wu, wg2, wu2; + W8::run(rgw, k, tab, wg, wg2); + W8::run(ruw, k, tab, wu, wu2); + rg[i] = wg.x; rg[i + 1] = wg.y; rg[i + 2] = wg.z; rg[i + 3] = wg.w; + ru[i] = wu.x; ru[i + 1] = wu.y; ru[i + 2] = wu.z; ru[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rg[j] = wg2.x; rg[j + 1] = wg2.y; rg[j + 2] = wg2.z; rg[j + 3] = wg2.w; + ru[j] = wu2.x; ru[j + 1] = wu2.y; ru[j + 2] = wu2.z; ru[j + 3] = wu2.w; + } + tA0 = tX.template slice(k0, row0); + op.run(tA0, rg, accg0); + op.run(tA0, ru, accu0); + if (MS > 1) { + tA1 = tX.template slice(k0, row0 + TM); + op.run(tA1, rg, accg1); + op.run(tA1, ru, accu1); + } + } + for (int s = 0; s < MS; ++s) { + thread DT& ag = s == 0 ? accg0 : accg1; + thread DT& au = s == 0 ? accu0 : accu1; + for (uint16_t i = 0; i < ag.get_capacity(); ++i) + if (ag.is_valid_element(i)) { + const float g = ag[i]; + ag[i] = g / (1.0f + exp(-g)) * au[i]; + } + tensor, tensor_inline> tO(ot[sg], dextents(TN, TM)); + simdgroup_barrier(mem_flags::mem_threadgroup); + ag.store(tO); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int it = lane; it < TM * TN; it += 32) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + s * TM + m) * a.n_ff + n0 + f] = half(ot[sg][it]); + } + } +} +#define MMGU4_(F, M, N, K, S) template [[host_name("mm_gu4_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu4) mm_gu4; +#define MMGU4(F, M, N, K, S) MMGU4_(F, M, N, K, S) +#define MM4_GU(F) MMGU4(F, 2, 32, 32, 4) MMGU4(F, 2, 32, 32, 2) MMGU4(F, 2, 16, 32, 4) MMGU4(F, 1, 32, 32, 4) +MM4_GU(T_Q2_0) MM4_GU(T_IQ3_XXS) MM4_GU(T_IQ2_XXS) MM4_GU(T_IQ2_XS) MM4_GU(T_IQ2_S) MM4_GU(T_IQ3_S) MM4_GU(T_IQ4_XS) + +// mm_gu5: mm_gu3 (COOP_K4 layout, TK 32) with the threadgroup's weight rows staged block by block (QK = 256 values a +// row, gate and up) into threadgroup memory by coalesced ushort loads, the lanes decoding from there: the decode's +// many small scattered device loads (2-6 a run) become threadgroup reads. Same decoded values and the same matrix +// ops as mm_gu3: bitwise. Formats with 256-value blocks and a templated W8 (IQ3_S, IQ3_XXS, IQ2_S). +template +kernel void mm_gu5(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], + device const uint* smap [[buffer(5)]], device const uchar* W1 [[buffer(6)]], + device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int TM = 16, TK = 32, BB = fmt::BB, ROWS = SG * TN, RU = BB / 2; + threadgroup float ot[SG][TM * TN]; + threadgroup uint tab[fmt::TAB]; + threadgroup ushort stg[2][ROWS * RU]; + load_tab(tab, tid, SG * 32); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, nt0 = int(tg.x) * ROWS, n0 = nt0 + sg * TN; + device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride; + device const uchar* gt = E + a.gate_off + ulong(nt0) * a.gu_row; + device const uchar* ut = E + a.up_off + ulong(nt0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + auto rg = op.template get_right_input_cooperative_tensor(); + auto ru = op.template get_right_input_cooperative_tensor(); + auto tA = tX.template slice(0, row0); + auto accg = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + auto accu = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + const int NB = int(a.n_embd) / 256; + for (int kb = 0; kb < NB; ++kb) { + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < ROWS * RU; idx += SG * 32) { + const int r = idx / RU, w = idx - r * RU; + stg[0][idx] = ((device const ushort*)(gt + ulong(r) * a.gu_row + kb * BB))[w]; + stg[1][idx] = ((device const ushort*)(ut + ulong(r) * a.gu_row + kb * BB))[w]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int ks = 0; ks < 8; ++ks) { + const int k0 = kb * 256 + ks * TK; + const uint16_t half_cap = rg.get_capacity() / 2; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rg.get_multidimensional_index(i); + const int kl = ks * TK + ix[0], n = sg * TN + ix[1]; + // the block's row in the stage; W8 indexes block (k >> 8) = 0 of it with k = kl + threadgroup const uchar* rgw = (threadgroup const uchar*)(stg[0] + n * RU); + threadgroup const uchar* ruw = (threadgroup const uchar*)(stg[1] + n * RU); + half4 wg, wu, wg2, wu2; + W8::run(rgw, kl, tab, wg, wg2); + W8::run(ruw, kl, tab, wu, wu2); + rg[i] = wg.x; rg[i + 1] = wg.y; rg[i + 2] = wg.z; rg[i + 3] = wg.w; + ru[i] = wu.x; ru[i + 1] = wu.y; ru[i + 2] = wu.z; ru[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rg[j] = wg2.x; rg[j + 1] = wg2.y; rg[j + 2] = wg2.z; rg[j + 3] = wg2.w; + ru[j] = wu2.x; ru[j + 1] = wu2.y; ru[j + 2] = wu2.z; ru[j + 3] = wu2.w; + } + tA = tX.template slice(k0, row0); + op.run(tA, rg, accg); + op.run(tA, ru, accu); + } + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = accg[i]; + accg[i] = g / (1.0f + exp(-g)) * accu[i]; + } + tensor, tensor_inline> tO(ot[sg], dextents(TN, TM)); + accg.store(tO); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int it = lane; it < TM * TN; it += 32) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + m) * a.n_ff + n0 + f] = half(ot[sg][it]); + } +} +#define MMGU5_(F, N, S) template [[host_name("mm_gu5_" #F "_" #N "_" #S)]] kernel decltype(mm_gu5) mm_gu5; +#define MMGU5(F, N, S) MMGU5_(F, N, S) +MMGU5(T_IQ3_S, 32, 2) MMGU5(T_IQ3_S, 32, 1) MMGU5(T_IQ3_S, 16, 4) MMGU5(T_IQ3_XXS, 32, 2) MMGU5(T_IQ3_XXS, 32, 1) +MMGU5(T_IQ3_XXS, 16, 4) MMGU5(T_IQ2_S, 32, 2) MMGU5(T_IQ2_S, 32, 1) MMGU5(T_IQ2_S, 16, 4) MMGU5(T_IQ2_XXS, 32, 2) MMGU5(T_IQ2_XXS, 16, 4) +#endif +#endif diff --git a/results/round39/model-ab/candidate-kernels/prompt.metal b/results/round39/model-ab/candidate-kernels/prompt.metal new file mode 100644 index 0000000..a3fd96b --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/prompt.metal @@ -0,0 +1,690 @@ +// metal/kernels/prompt.metal - the elementwise steps of a prompt chunk around the GEMMs (after quant/units.metal): +// +// route_topk per token: softmax over the 512 logits, top 10 (ties: the lower expert), weights renormalised +// over the selected with a 2^-14 floor - moe_route's rule; one simdgroup a token +// gr_epi_down after the HC down GEMM (V = xn . [w_down; w_inject], T x (hc_lr + hc)): lo_h = half(silu(v / hc)), +// inject = v +// gr_epi_up after the HC up GEMM (G = lo . w_up, T x hc*n): mixed = mean_c xn[c] * sigmoid(G[c]) -> float and +// half (the next GEMMs' operand) +// swiglu_h h = half(silu(g) * u) +// gemm_hc_down, gemm_swiglu_up (after gemm_hc_up) the HC down and shared-expert up GEMMs with gr_epi_down / +// swiglu_h as register epilogues +// moe_out bo = y_routed + sigmoid(sg) * y_shared (prompt_chain.inc; the prompt chunk folds it into the next +// read: gr_norm_h_moe, gr_apply_moe, ple_moe_gate_norm; with mm_combine folded in too: gr_norm_h_moec, +// gr_apply_moec) + +kernel void route_topk(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]], + device float* w [[buffer(2)]], constant uint& n_tok_packed [[buffer(3)]], + device const uint* smap [[buffer(4)]], constant float& prompt_moce [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + const uint n_tok = n_tok_packed & 0xFFFFu; + const int n_exp = (n_tok_packed >> 16) ? int(n_tok_packed >> 16) : 512; + const uint t = tg.x * 8 + sg; + if (t >= n_tok) return; + device const float* l = logits + ulong(t) * ulong(n_exp); + float v[16]; + float mx = -INFINITY; + float sum = 0.0f; + if (n_exp == 512) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] = l[lane + 32 * i]; + mx = max(mx, v[i]); + } + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] = exp(v[i] - mx); + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + } else { + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + v[i] = idx < n_exp ? l[idx] : -INFINITY; + } + for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + for (int i = 0; i < 16; ++i) { + v[i] = (lane + 32 * i < n_exp) ? exp(v[i] - mx) : 0.0f; + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + for (int i = 0; i < 16; ++i) { + if (lane + 32 * i >= n_exp) { + v[i] = -INFINITY; + } else { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + } + } + const int tail_exact = int(prompt_moce) / 10000; + const float rem_moce = prompt_moce - float(tail_exact * 10000); + const int body_prot = int(rem_moce) / 100; + const float body_bonus = rem_moce - float(body_prot * 100); + const bool in_tail = tail_exact > 0 && (t + uint(tail_exact) >= n_tok); + const int p_prot = in_tail ? 10 : body_prot; + const float p_bonus = in_tail ? 0.0f : body_bonus; + const float moce_scale = p_bonus > 0.0f ? exp(min(p_bonus, 40.0f)) : 1.0f; + const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f; + ushort scaled_mask = 0; + float selected = 0.0f, ssum = 0.0f; + for (int rank = 0; rank < 10; ++rank) { + if (rank == p_prot && moce_scale > 1.0f) { + if (n_exp == 512) { + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if (v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) { + v[i] *= moce_scale; + scaled_mask |= ushort(1u << i); + } + } + } else { + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if (idx < n_exp && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) { + v[i] *= moce_scale; + scaled_mask |= ushort(1u << i); + } + } + } + } + float best = v[0]; + int ex = lane; + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + float true_prob = best; + if (rank >= p_prot && moce_scale > 1.0f && best > 0.0f) { + const ushort win_lane = ushort(ex & 31); + 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 * rcp_moce; + } + if ((ex & 31) == lane) { + // v stays in registers: no runtime index into it + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) v[i] = i == (ex >> 5) ? -INFINITY : v[i]; + ids[t * 10 + rank] = ex; + ssum += true_prob; + } + if (rank == lane) selected = true_prob; + } + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + if (prompt_moce < 0.0f) { + const int skip_tail = int(-prompt_moce); + const float tau = (-prompt_moce) - float(skip_tail); + const bool skip_active = tau > 0.0f && !(skip_tail > 0 && (t + uint(skip_tail) >= n_tok)); + if (skip_active) { + float kept_prob = (lane < 10) ? selected : 0.0f; + if (lane >= 3 && lane < 10) { + const int ex_r = ids[t * 10 + lane]; + if (smap[ex_r] == 0xFFFFFFFFu && selected < tau * ssum) { + ids[t * 10 + lane] = -1; + kept_prob = 0.0f; + } + } + selected = kept_prob; + ssum = kept_prob; + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + } + } + ssum = max(ssum, 6.103515625e-5f); + if (lane < 10) w[t * 10 + lane] = selected * (1.0f / ssum); +} + +struct EpiArgs { + uint n_tok; + uint n_embd; // 2560 + uint hc; // 4 + uint hc_lr; // 320 +}; + +// grid (T * (hc_lr + hc) / 256) +kernel void gr_epi_down(device const float* V [[buffer(0)]], device half* lo_h [[buffer(1)]], + device float* inject [[buffer(2)]], constant EpiArgs& a [[buffer(3)]], + uint gid [[thread_position_in_grid]]) { + const uint nr = a.hc_lr + a.hc, t = gid / nr, r = gid % nr; + if (t >= a.n_tok) return; + const float v = V[gid]; + if (r < a.hc_lr) { + const float z = v / float(a.hc); + lo_h[t * a.hc_lr + r] = half(z / (1.0f + exp(-z))); + } else { + inject[t * a.hc + r - a.hc_lr] = v; + } +} + +// grid (T * n_embd / 256) +kernel void gr_epi_up(device const float* G [[buffer(0)]], device const float* xn [[buffer(1)]], + device float* mixed [[buffer(2)]], device half* mixed_h [[buffer(3)]], + constant EpiArgs& a [[buffer(4)]], uint gid [[thread_position_in_grid]]) { + const uint n = a.n_embd, t = gid / n, d = gid % n; + if (t >= a.n_tok) return; + float m = 0.0f; + for (uint c = 0; c < a.hc; ++c) { + const float g = G[ulong(t) * a.hc * n + c * n + d]; + m = fma(xn[ulong(t) * a.hc * n + c * n + d], 1.0f / (1.0f + exp(-g)), m); + } + m = m / float(a.hc); + mixed[gid] = m; + mixed_h[gid] = half(m); +} + +kernel void swiglu_h(device const float* g [[buffer(0)]], device const float* u [[buffer(1)]], device half* h [[buffer(2)]], + constant uint& n [[buffer(3)]], uint gid [[thread_position_in_grid]]) { + if (gid >= n) return; + const float gv = g[gid]; + h[gid] = half(gv / (1.0f + exp(-gv)) * u[gid]); +} + +// grid (T * n_embd / 256); sg: T values +kernel void moe_out(device const float* yr [[buffer(0)]], device const float* ys [[buffer(1)]], + device const float* sg [[buffer(2)]], device float* bo [[buffer(3)]], constant uint2& a [[buffer(4)]], + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.y; + if (t >= a.x) return; + bo[gid] = fma(1.0f / (1.0f + exp(-sg[t])), ys[gid], yr[gid]); +} + +// ---- the HC read with fewer passes (after dense_mm.metal: GemmArgs, MetalPerformancePrimitives) --------------------- +// gr_norm_h: gr_norm with xn written as fp16 only (the GEMM operand; gemm_hc_up's epilogue reads it too). +// one threadgroup per (stream c, token t), 256 threads; args as gr_norm's (GrArgs, gr.metal) +kernel void gr_norm_h(device float* R [[buffer(0)]], device const float* bo [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device half* xnh [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint c = tg.x, t = tg.y, n = a.n_embd; + device float* r = R + (t * a.hc + c) * n; + const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; + float rv[10], ss = 0.0f; + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + float v = r[d]; + if (a.apply) { + v = fma(bo[t * n + d], g, v); + r[d] = v; + } + rv[i] = v; + ss = fma(v, v, ss); + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; +#pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device half* o = xnh + (t * a.hc + c) * n; + device const float* w = w_norm + c * n; + for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[i] * w[d] * rs); +} + +// gr_norm_h's tail for one token's hc streams: the sums of squares ss[c] (per thread) to the norm, xnh = R' * w * rs +inline void gr_norm_h_out(thread const float (*rv)[10], device const float* w_norm, device half* xnh, uint t, uint hc, + uint n, float eps, thread const float* ss, threadgroup float (*part)[8], ushort tid, + ushort sg, ushort lane) { + for (uint c = 0; c < hc; ++c) { + const float s = simd_sum(ss[c]); + if (lane == 0) part[c][sg] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint c = 0; c < hc; ++c) { + float tot = 0.0f; +#pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[c][i]; + const float rs = rsqrt(tot / float(n) + eps); + device half* o = xnh + (t * hc + c) * n; + device const float* w = w_norm + c * n; + for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[c][i] * w[d] * rs); + } +} + +// stream-wise reduction keeps one residual stream plus bo live, rather than all four residual streams. +inline void gr_norm_h_stream_out(thread const float* rv, device const float* w, device half* o, uint n, float eps, + float ss, threadgroup float* part, ushort tid, ushort sg, ushort lane) { + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; +#pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + eps); + if (n == 2560) { +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const uint d = tid + i * 256; + o[d] = half(rv[i] * w[d] * rs); + } + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[i] * w[d] * rs); + } +} + +// gr_norm_h4: gr_norm_h with one threadgroup per token over all hc streams, so bo is read once a token rather than +// hc times; R', xnh bitwise gr_norm_h's. grid (1, T), 256 threads; hc <= 4 +kernel void gr_norm_h4(device float* R [[buffer(0)]], device const float* bo [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device half* xnh [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + const uint t = tg.y, n = a.n_embd, hc = a.hc; + device float* r = R + t * hc * n; + if (!a.apply) { + for (uint c = 0; c < hc; ++c) { + float rv[10], ss = 0.0f; + if (n == 2560) { + device const float* rc = r + c * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const float v = rc[i * 256]; + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + const float v = r[c * n + d]; + rv[i] = v; + ss = fma(v, v, ss); + } + } + gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); + } + return; + } + float bv[10]; + if (n == 2560) { + device const float* brow = bo + t * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) bv[i] = brow[i * 256]; + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) bv[i] = bo[t * n + d]; + } + for (uint c = 0; c < hc; ++c) { + const float g = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); + float rv[10], ss = 0.0f; + if (n == 2560) { + device float* rc = r + c * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const float v = fma(bv[i], g, rc[i * 256]); + rc[i * 256] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + const float v = fma(bv[i], g, r[c * n + d]); + r[c * n + d] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } + gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); + } +} + +// gr_norm_h_moe: gr_norm_h with the sublayer output bo formed on the fly as moe_out - bo = fma(sigmoid(sg[t]), ys, yr) +// - so the MoE's write never round-trips through bo; R', xnh bitwise moe_out + gr_norm_h. Always applies. +// One threadgroup per token t (all hc streams: bo is formed once a token), 256 threads; hc <= 4. +kernel void gr_norm_h_moe(device float* R [[buffer(0)]], device const float* yr [[buffer(1)]], + device const float* ys [[buffer(2)]], device const float* sgv [[buffer(3)]], + device const float* inj [[buffer(4)]], device const float* w_norm [[buffer(5)]], + device half* xnh [[buffer(6)]], constant GrArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + const uint t = tg.y, n = a.n_embd, hc = a.hc; + device float* r = R + t * hc * n; + float rv[4][10], g[4], ss[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (uint c = 0; c < hc; ++c) g[c] = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + const float bo = fma(gs, ys[t * n + d], yr[t * n + d]); + for (uint c = 0; c < hc; ++c) { + const float v = fma(bo, g[c], r[c * n + d]); + r[c * n + d] = v; + rv[c][i] = v; + ss[c] = fma(v, v, ss[c]); + } + } + gr_norm_h_out(rv, w_norm, xnh, t, hc, n, a.eps, ss, part, tid, sg, lane); +} + +// token t's routed slots for moe_comb: compacted valid YE row offsets (uint(r) * n) and their weights in threadgroup memory +inline uint moe_comb_slots(device const int* pos, device const float* srt_w, uint t, uint topk, uint n, + threadgroup uint* roff, threadgroup float* wk, ushort tid) { + int r = -1; + float w = 0.0f; + if (tid < topk) { + r = pos[t * topk + tid]; + if (r >= 0) w = srt_w[r]; + } + const bool valid = r >= 0; + const ushort prefix = simd_prefix_exclusive_sum(ushort(valid)); + const ushort cnt = simd_sum(ushort(valid)); + if (valid) { + roff[prefix] = uint(r) * n; + wk[prefix] = w; + } + if (tid == 0) roff[16] = cnt; + threadgroup_barrier(mem_flags::mem_threadgroup); + return roff[16]; +} + +// the routed experts' write at (t, d) as mm_combine (moe_mm.metal) forms it: sum_k w(t, k) * YE[pos(t, k)][d], k in order +inline float moe_comb(device const float* YE, threadgroup const uint* roff, threadgroup const float* wk, uint cnt, + uint d) { + float s = 0.0f; + for (uint k = 0; k < cnt; ++k) s = fma(wk[k], YE[roff[k] + d], s); + return s; +} + +// gr_norm_h_moec: gr_norm_h_moe with mm_combine folded in - yr formed from the expert rows YE (pos, srt_w: mm_place's) +// as it is read, so it is never written; R', xnh bitwise mm_combine + gr_norm_h_moe. grid (1, T), 256 threads; +// hc <= 4, topk <= 16 +kernel void gr_norm_h_moec(device float* R [[buffer(0)]], device const float* YE [[buffer(1)]], + device const int* pos [[buffer(2)]], device const float* srt_w [[buffer(3)]], + device const float* ys [[buffer(4)]], device const float* sgv [[buffer(5)]], + device const float* inj [[buffer(6)]], device const float* w_norm [[buffer(7)]], + device half* xnh [[buffer(8)]], constant GrArgs& a [[buffer(9)]], constant uint& topk [[buffer(10)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4][8]; + threadgroup uint roff[17]; + threadgroup float wk[16]; + const uint t = tg.y, n = a.n_embd, hc = a.hc; + const uint cnt = moe_comb_slots(pos, srt_w, t, topk, n, roff, wk, tid); + device float* r = R + t * hc * n; + float bv[10]; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + if (n == 2560) { + // k outer: each slot's weight and row read once, its ten loads in flight together; s[i]'s fma order is moe_comb's + float s[10] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + for (uint k = 0; k < cnt; ++k) { + const float w = wk[k]; + device const float* yk = YE + roff[k] + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) s[i] = fma(w, yk[i * 256], s[i]); + } + device const float* yrow = ys + t * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) bv[i] = fma(gs, yrow[i * 256], s[i]); + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) bv[i] = fma(gs, ys[t * n + d], moe_comb(YE, roff, wk, cnt, d)); + } + for (uint c = 0; c < hc; ++c) { + const float g = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); + float rv[10], ss = 0.0f; + if (n == 2560) { + device float* rc = r + c * 2560 + tid; +#pragma clang loop unroll(full) + for (uint i = 0; i < 10; ++i) { + const float v = fma(bv[i], g, rc[i * 256]); + rc[i * 256] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } else { + for (uint d = tid, i = 0; d < n; d += 256, ++i) { + const float v = fma(bv[i], g, r[c * n + d]); + r[c * n + d] = v; + rv[i] = v; + ss = fma(v, v, ss); + } + } + gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); + } +} + +// gr_apply_moe: the last layer's MoE write applied to R alone (R' of moe_out + gr_norm; its xn has no reader). +// grid (hc * n / 256, T), 256 threads +kernel void gr_apply_moe(device float* R [[buffer(0)]], device const float* yr [[buffer(1)]], + device const float* ys [[buffer(2)]], device const float* sgv [[buffer(3)]], + device const float* inj [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const uint ch = tg.x * 256 + tid, t = tg.y, n = a.n_embd; + if (ch >= a.hc * n) return; + const uint c = ch / n, d = ch - c * n; + const float g = 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))); + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + device float* r = R + t * a.hc * n + ch; + *r = fma(fma(gs, ys[t * n + d], yr[t * n + d]), g, *r); +} + +// gr_apply_moec: gr_apply_moe with mm_combine folded in (as gr_norm_h_moec); R' bitwise mm_combine + gr_apply_moe. +// grid (ceil(n / 256), T), 256 threads: a thread's d over all hc streams (yr formed once); topk <= 16 +kernel void gr_apply_moec(device float* R [[buffer(0)]], device const float* YE [[buffer(1)]], + device const int* pos [[buffer(2)]], device const float* srt_w [[buffer(3)]], + device const float* ys [[buffer(4)]], device const float* sgv [[buffer(5)]], + device const float* inj [[buffer(6)]], constant GrArgs& a [[buffer(7)]], + constant uint& topk [[buffer(8)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint roff[17]; + threadgroup float wk[16]; + const uint d = tg.x * 256 + tid, t = tg.y, n = a.n_embd, hc = a.hc; + const uint cnt = moe_comb_slots(pos, srt_w, t, topk, n, roff, wk, tid); + if (d >= n) return; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + const float bo = fma(gs, ys[t * n + d], moe_comb(YE, roff, wk, cnt, d)); + device float* r = R + t * hc * n + d; + for (uint c = 0; c < hc; ++c) r[c * n] = fma(bo, 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))), r[c * n]); +} + +// gemm_hc_up: G = lo . w_up^T with w_up's rows interleaved by dq_f16 (perm = hc: row d * hc + c), so a TN-column +// tile holds TN / hc outputs d with all their streams; the epilogue mixed[t][d] = mean_c xn[t][c][d] * +// sigmoid(G[t][d * hc + c]) -> fp16 and, if mixed is bound, fp32. G is never written. +// grid (hc * n / TN, ceil(T / TM)), 32 * SG threads. +template +inline void hc_up_tile(device half* A, device half* B, device const half* xnh, device float* mixed, + device half* mixh, constant GemmArgs& g, constant EpiArgs& e, + int t0, int c0, ushort tid, threadgroup float* ot) { + tensor, tensor_inline> tA(A, dextents(g.K, g.M)); + tensor, tensor_inline> tB(B, dextents(g.K, g.N)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, t0); + auto mB = tB.slice(0, c0); + if constexpr (REG) { + auto acc = op.template get_destination_cooperative_tensor(); + op.run(mA, mB, acc); + const int n = int(e.n_embd), d0 = c0 / 4; + // On M5, each lane owns four consecutive output columns: one output's HC streams. + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { + const auto ix = acc.get_multidimensional_index(i); + const int t = t0 + ix[1], d = d0 + ix[0] / 4; + if (t >= int(e.n_tok)) continue; + device const half* xnh_t = xnh + ulong(t) * 4 * n + d; + float s = 0.0f; + #pragma clang loop unroll(full) + for (int c = 0; c < 4; ++c) + s = fma(float(xnh_t[c * n]), 1.0f / (1.0f + exp(-acc[i + c])), s); + s *= 0.25f; + const ulong out_idx = ulong(t) * n + d; + if (mixed) mixed[out_idx] = s; + mixh[out_idx] = half(s); + } + } else { + tensor, tensor_inline> tO(ot, dextents(TN, TM)); + op.run(mA, mB, tO); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int hc = int(e.hc), n = int(e.n_embd), DPT = TN / hc; + const int d0 = c0 / hc; + for (int idx = tid; idx < TM * DPT; idx += SG * 32) { + const int m = idx / DPT, dd = idx - m * DPT, t = t0 + m; + if (t >= int(e.n_tok)) continue; + float s = 0.0f; + for (int c = 0; c < hc; ++c) + s = fma(float(xnh[ulong(t) * hc * n + c * n + d0 + dd]), 1.0f / (1.0f + exp(-ot[m * TN + dd * hc + c])), s); + s = s / float(hc); + if (mixed) mixed[ulong(t) * n + d0 + dd] = s; + mixh[ulong(t) * n + d0 + dd] = half(s); + } + } +} +template +kernel void gemm_hc_up(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device const half* xnh [[buffer(2)]], + device float* mixed [[buffer(3)]], device half* mixh [[buffer(4)]], constant GemmArgs& g [[buffer(5)]], + constant EpiArgs& e [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float ot[REG ? 1 : TM * TN]; + const int t0 = int(tg.y) * TM, c0 = int(tg.x) * TN; + if constexpr (REG && TM > 16) { + if (t0 + 16 >= int(e.n_tok)) { + hc_up_tile<16, TN, SG, REG>(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot); + return; + } + } + if constexpr (REG && TM > 32) { + if (t0 + 32 >= int(e.n_tok)) { + hc_up_tile<32, TN, SG, REG>(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot); + return; + } + } + hc_up_tile(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot); +} +template [[host_name("gemm_hc_up_64_64_4")]] kernel decltype(gemm_hc_up<64, 64, 4>) gemm_hc_up<64, 64, 4>; +template [[host_name("gemm_hc_up_64_32_4")]] kernel decltype(gemm_hc_up<64, 32, 4>) gemm_hc_up<64, 32, 4>; +template [[host_name("gemm_hc_up_32_64_4")]] kernel decltype(gemm_hc_up<32, 64, 4>) gemm_hc_up<32, 64, 4>; +template [[host_name("gemm_hc_up_128_64_4")]] kernel decltype(gemm_hc_up<128, 64, 4>) gemm_hc_up<128, 64, 4>; +template [[host_name("gemm_hc_up_64_128_4")]] kernel decltype(gemm_hc_up<64, 128, 4>) gemm_hc_up<64, 128, 4>; +template [[host_name("gemm_hc_up_reg_64_64_4")]] kernel decltype(gemm_hc_up<64, 64, 4, true>) gemm_hc_up<64, 64, 4, true>; +template [[host_name("gemm_hc_up_reg_128_64_4")]] kernel decltype(gemm_hc_up<128, 64, 4, true>) gemm_hc_up<128, 64, 4, true>; +template [[host_name("gemm_hc_up_reg_64_128_4")]] kernel decltype(gemm_hc_up<64, 128, 4, true>) gemm_hc_up<64, 128, 4, true>; +template [[host_name("gemm_hc_up_reg_64_32_4")]] kernel decltype(gemm_hc_up<64, 32, 4, true>) gemm_hc_up<64, 32, 4, true>; +template [[host_name("gemm_hc_up_reg_32_128_4")]] kernel decltype(gemm_hc_up<32, 128, 4, true>) gemm_hc_up<32, 128, 4, true>; + +// gemm_hc_down: the HC down GEMM V = xn . [w_down; w_inject]^T (T x (hc_lr + hc), K = hc * n_embd) with gr_epi_down +// as its register epilogue: lo_h = half(silu(v / hc)), inject = v. V is never written; lo_h, inject bitwise gemm_f16 + +// gr_epi_down. M5 layout only (coop_k4: a lane's elements 4j .. 4j + 3 are four consecutive columns of one row); +// hc_lr % 4 == 0, hc == 4. grid (ceil((hc_lr + hc) / TN), ceil(T / TM)), 32 * SG threads +template +inline void hc_down_tile(device half* A, device half* B, device half* lo_h, device float* inject, + constant EpiArgs& e, int t0, int c0) { + const int K = int(e.hc * e.n_embd), nr = int(e.hc_lr + e.hc); + tensor, tensor_inline> tA(A, dextents(K, int(e.n_tok))); + tensor, tensor_inline> tB(B, dextents(K, nr)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, t0); + auto mB = tB.slice(0, c0); + auto acc = op.template get_destination_cooperative_tensor(); + op.run(mA, mB, acc); + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { + const auto ix = acc.get_multidimensional_index(i); + const int r = c0 + ix[0], t = t0 + ix[1]; + if (t >= int(e.n_tok) || r >= nr) continue; + if (r < int(e.hc_lr)) { + half4 o; + #pragma clang loop unroll(full) + for (int c = 0; c < 4; ++c) { + const float z = acc[i + c] * 0.25f; // / hc, exact: hc == 4 + o[c] = half(z / (1.0f + exp(-z))); + } + *(device half4*)(lo_h + ulong(t) * e.hc_lr + r) = o; + } else { + *(device float4*)(inject + ulong(t) * e.hc) = float4(acc[i], acc[i + 1], acc[i + 2], acc[i + 3]); + } + } +} +template +inline void hc_down_rows(device half* A, device half* B, device half* lo_h, device float* inject, + constant EpiArgs& e, int t0, int c0) { + // Trimming 32-row parents regresses full tiles on M5. + if constexpr (TM > 32) { + if (t0 + 16 >= int(e.n_tok)) { + hc_down_tile<16, TN, SG>(A, B, lo_h, inject, e, t0, c0); + return; + } + if (t0 + 32 >= int(e.n_tok)) { + hc_down_tile<32, TN, SG>(A, B, lo_h, inject, e, t0, c0); + return; + } + } + hc_down_tile(A, B, lo_h, inject, e, t0, c0); +} +template +kernel void gemm_hc_down(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device half* lo_h [[buffer(2)]], + device float* inject [[buffer(3)]], constant EpiArgs& e [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int c0 = int(tg.x) * TN, t0 = int(tg.y) * TM; + if constexpr (TN == 128) { + // A single 96-column tail beats serial 64+16 on M5, including SG=2. + if (e.hc == 4 && c0 + 64 == int(e.hc_lr)) { + hc_down_rows(A, B, lo_h, inject, e, t0, c0); + return; + } + } + hc_down_rows(A, B, lo_h, inject, e, t0, c0); +} +template [[host_name("gemm_hc_down_16_64_2")]] kernel decltype(gemm_hc_down<16, 64, 2>) gemm_hc_down<16, 64, 2>; +template [[host_name("gemm_hc_down_32_64_2")]] kernel decltype(gemm_hc_down<32, 64, 2>) gemm_hc_down<32, 64, 2>; +template [[host_name("gemm_hc_down_32_128_4")]] kernel decltype(gemm_hc_down<32, 128, 4>) gemm_hc_down<32, 128, 4>; +template [[host_name("gemm_hc_down_64_128_4")]] kernel decltype(gemm_hc_down<64, 128, 4>) gemm_hc_down<64, 128, 4>; +template [[host_name("gemm_hc_down_32_128_2")]] kernel decltype(gemm_hc_down<32, 128, 2>) gemm_hc_down<32, 128, 2>; + +// gemm_swiglu_up: the shared expert's up GEMM u = x . w_up^T (M x N, K) with swiglu_h as its register epilogue against +// the gate GEMM's g: h = half(silu(g) * u). u is never written; h bitwise gemm_f16 + swiglu_h. One accumulator: the +// gate stays a plain gemm_f16 before it. M5 layout only (as gemm_hc_down); N % 4 == 0. +// grid (ceil(N / TN), ceil(M / TM)), 32 * SG threads +template +inline void swiglu_up_tile(device half* A, device half* B, device const float* g, device half* h, + constant GemmArgs& a, int t0, int c0) { + tensor, tensor_inline> tA(A, dextents(a.K, a.M)); + tensor, tensor_inline> tB(B, dextents(a.K, a.N)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, t0); + auto mB = tB.slice(0, c0); + auto acc = op.template get_destination_cooperative_tensor(); + op.run(mA, mB, acc); + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { + const auto ix = acc.get_multidimensional_index(i); + const int c = c0 + ix[0], t = t0 + ix[1]; + if (t >= a.M || c >= a.N) continue; + const ulong idx = ulong(t) * uint(a.N) + uint(c); + const float4 gv = *(device const float4*)(g + idx); + half4 o; + #pragma clang loop unroll(full) + for (int k = 0; k < 4; ++k) o[k] = half(gv[k] / (1.0f + exp(-gv[k])) * acc[i + k]); + *(device half4*)(h + idx) = o; + } +} +template +kernel void gemm_swiglu_up(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device const float* g [[buffer(2)]], + device half* h [[buffer(3)]], constant GemmArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]]) { + const int t0 = int(tg.y) * TM, c0 = int(tg.x) * TN; + if constexpr (TM > 32) { + if (t0 + 16 >= a.M) { + swiglu_up_tile<16, TN, SG>(A, B, g, h, a, t0, c0); + return; + } + if (t0 + 32 >= a.M) { + swiglu_up_tile<32, TN, SG>(A, B, g, h, a, t0, c0); + return; + } + } + swiglu_up_tile(A, B, g, h, a, t0, c0); +} +template [[host_name("gemm_swiglu_up_32_64_2")]] kernel decltype(gemm_swiglu_up<32, 64, 2>) gemm_swiglu_up<32, 64, 2>; +template [[host_name("gemm_swiglu_up_64_128_4")]] kernel decltype(gemm_swiglu_up<64, 128, 4>) gemm_swiglu_up<64, 128, 4>; diff --git a/results/round39/model-ab/candidate-kernels/qsa.metal b/results/round39/model-ab/candidate-kernels/qsa.metal new file mode 100644 index 0000000..63259b4 --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/qsa.metal @@ -0,0 +1,492 @@ +// metal/kernels/qsa.metal - QSA (sparse full attention) decode for a window of tokens, between the projections. +// Contract of include/strata/kernels/qsa.hpp / qsa_select.hpp / qsa_decode_attn.hpp / kv_q8.hpp (the engine's +// native path), in a Metal layout (no paging: unified memory): +// +// qsa_prep per token t (cell pos0 + t): q heads rms_norm * q_norm, NEOX RoPE on 64 dims -> qcur; +// k rms_norm * k_norm, RoPE -> int8 KV append (fp16 scale per 64 = fp16(amax/127), rint, clamp); +// v -> int8 append; indexer q rms_norm * iq_norm, RoPE -> qidx. The indexer KEYS (last +// threadgroup, all window tokens in order): raw tail [cell % 4]; on a block completion +// pooled[b] = rope(rms_norm(mean of the block's 4 raw keys) * ik_norm, pos of the block's first cell); +// cell 0 also sets dead = rms_norm(raw) * ik_norm (the tail block's key). +// qsa_score per token, block b <= n_bid: sum over the 4 indexer heads of relu(qidx . key_b); key_b = pooled[b], +// the tail block n_bid scores `dead` and gets +1e9 when it has cells. +// qsa_topk per token: width = min(n_kv, 2051) cells; n_kv <= 2051: all; else the width cells of largest +// score (a cell inherits its block's), ties to the lowest cell, emitted ascending (radix select). +// qsa_attn split-K: (chunk of CH cells, kv head, token): 12 query heads; s = q.k / 16, chunk max m, sum l, +// acc = p . v -> partials +// qsa_merge per (head, token): log-sum-exp merge, * sigmoid(gate), written as prepared activations (xprep +// planes) for attn_output. +// +// RoPE: cos/sin from a host table [pos][32] (double precision, as build_rope_table). + +#define QH 24 +#define QKV 2 +#define QD 256 +#define QROT 64 +#define IH 4 +#define ID 128 +#define QBUDGET 2051 +#define QCH 64 + +constant int QSA_NTOK [[function_constant(3)]]; +// QSA_MROPE (function constant 24): interleaved M-RoPE (the vision path; strata/kernels/mrope.hpp, ggml IMROPE, +// sections 11/11/10): buffer 20 of qsa_prep is a table int32 [cells][3] (t, h, w), and rotary pair i (0..31) of the row +// at (pos_base + cell) rotates by the cos/sin table's row tab[(pos_base + cell) * 3 + i % 3] instead of its own - +// the pooled indexer keys too (position pos_base + 4 b of a block). Off: exactly the plain RoPE. +constant bool QSA_MROPE [[function_constant(24)]]; +constant bool qsa_mrope = is_function_constant_defined(QSA_MROPE) && QSA_MROPE; +// QSA_Q_F16_FC (function constant 25): qsa_prep writes qcur directly as fp16 (device half*) for qsa_flash, skipping qsa_q_f16. +constant bool QSA_Q_F16_FC [[function_constant(25)]]; +constant bool qsa_qf16 = is_function_constant_defined(QSA_Q_F16_FC) ? QSA_Q_F16_FC : false; +// QSA_ROW0 (function constant 28): qsa_prep's threadgroup x is row QSA_ROW0 + x (24: the K/V rows only, for rows whose +// queries nobody reads). Undefined: 0. +constant int QSA_ROW0 [[function_constant(28)]]; +constant int qsa_row0 = is_function_constant_defined(QSA_ROW0) ? QSA_ROW0 : 0; + +struct QsaArgs { + uint pos0; // the window's first cell (relative to the sequence) + uint pos_base; // absolute position of cell 0 (RoPE) + uint max_cells; + uint n_tok; + float eps; +}; + +// cache layout per layer +// kq/vq [cell][kv head][256] int8, ks/vs [cell][kv head][4] half +// tail [4][128] f32 (slot = cell % 4), dead [128], pooled [max_cells / 4 + 1][128] + +static inline float tg_sum256(float v, threadgroup float* part, ushort sg, ushort lane, ushort nsg) { + v = simd_sum(v); + if (lane == 0) part[sg] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + float t = 0.0f; + for (int i = 0; i < nsg; ++i) t += part[i]; + return t; +} + +// one row (width w <= 256 = threads) : out = rope((rs * x) * g) (RoPE on the first 64 dims, NEOX pairs i, i+32) +static inline float norm_rope(float x, float xhi_unused, device const float* g, threadgroup float* row, + threadgroup float* part, ushort tid, ushort sg, ushort lane, int w, float eps, + device const float2* cs, uint absp, bool rope, device const int* mtab, + bool reuse = false) { + const float ss = tg_sum256(tid < w ? x * x : 0.0f, part, sg, lane, 8); + const float rs = rsqrt(ss / float(w) + eps); + const float y = tid < w ? (rs * x) * g[tid] : 0.0f; + float o = y; + if (rope) { + if (tid < QROT) row[tid] = y; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < QROT) { + const int p = tid & 31; + const uint pp = qsa_mrope ? uint(mtab[absp * 3 + uint(p % 3)]) : absp; + const float2 c = cs[pp * 32 + p]; + const float a = row[p], b = row[p + 32]; + o = tid < 32 ? a * c.x - b * c.y : a * c.y + b * c.x; + } + } + if (reuse) threadgroup_barrier(mem_flags::mem_threadgroup); + (void)xhi_unused; + return o; +} + +// QSA_KV_STRIDE (function constant 26): kcur, vcur, iqraw and ikraw are column slices of one [n_tok][QSA_KV_STRIDE] +// projection (prompt_chunk's fused K | V | indexer K | indexer Q GEMM, 1152 or 1664 columns). Undefined: their own +// [n_tok][rows] buffers. +constant int QSA_KV_STRIDE [[function_constant(26)]]; +constant int qsa_kv_s = is_function_constant_defined(QSA_KV_STRIDE) ? QSA_KV_STRIDE : 0; + +// grid (32 rows + 1, n_tok): rows 0..23 q heads, 24..25 k heads, 26..27 v heads, 28..31 indexer q heads; +// threadgroup 32 (x = 32, y = 0) is the indexer-key threadgroup. 256 threads. A grid of 28 (or of 4 from QSA_ROW0 = +// 24) writes no indexer state: the MTP layer's attention is a dense window, never scored. +kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float* kcur [[buffer(1)]], + device const float* vcur [[buffer(2)]], device const float* iqraw [[buffer(3)]], + device const float* ikraw [[buffer(4)]], device const float* qn [[buffer(5)]], + device const float* kn [[buffer(6)]], device const float* iqn [[buffer(7)]], + device const float* ikn [[buffer(8)]], device const float2* cs [[buffer(9)]], + device float* qcur [[buffer(10)]], device float* qidx [[buffer(11)]], + device char* kq [[buffer(12)]], device char* vq [[buffer(13)]], device half* ksc [[buffer(14)]], + device half* vsc [[buffer(15)]], device float* tail [[buffer(16)]], device float* dead [[buffer(17)]], + device float* pooled [[buffer(18)]], constant QsaArgs& a [[buffer(19)]], + device const int* mtab [[buffer(20)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + threadgroup float row[QROT]; + threadgroup float amax[8]; + const int r = int(tg.x) + qsa_row0; + // the row strides of kcur / vcur, ikraw, iqraw + const int s_kv = qsa_kv_s ? qsa_kv_s : QKV * QD, s_ik = qsa_kv_s ? qsa_kv_s : ID; + const int s_iq = qsa_kv_s ? qsa_kv_s : IH * ID; + if (r == 32) { + // the indexer keys: threadgroup (32, 0) handles t < 3 (boundary with prior tail + cell 0 dead) and final tail; + // every threadgroup (32, t >= 3) with (cell & 3) == 3 pools its 4 in-window keys in parallel. + const int t_self = int(tg.y); + if (t_self >= 3) { + const uint cell = a.pos0 + uint(t_self); + if ((cell & 3) == 3) { + float m = 0.0f; + if (tid < ID) { + m = (((ikraw[(t_self - 3) * s_ik + tid] + ikraw[(t_self - 2) * s_ik + tid]) + + ikraw[(t_self - 1) * s_ik + tid]) + ikraw[t_self * s_ik + tid]) * 0.25f; + } + const uint b = cell >> 2; + const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab); + if (tid < ID) pooled[b * ID + tid] = o; + } + return; + } + if (t_self != 0) return; + const int n_pre = min(int(a.n_tok), 3); + for (int t = 0; t < n_pre; ++t) { + const uint cell = a.pos0 + uint(t); + const float x = tid < ID ? ikraw[t * s_ik + tid] : 0.0f; + if (tid < ID) tail[(cell & 3) * ID + tid] = x; + threadgroup_barrier(mem_flags::mem_device); + if (cell == 0) { + const float o = norm_rope(x, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, 0, false, mtab, true); + if (tid < ID) dead[tid] = o; + } + if ((cell & 3) == 3) { + float m = 0.0f; + if (tid < ID) m = (((tail[tid] + tail[ID + tid]) + tail[2 * ID + tid]) + tail[3 * ID + tid]) * 0.25f; + const uint b = cell >> 2; + const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab, true); + if (tid < ID) pooled[b * ID + tid] = o; + } + threadgroup_barrier(mem_flags::mem_device); + } + const int t_tail0 = max(n_pre, int(a.n_tok) - 4); + for (int t = t_tail0; t < int(a.n_tok); ++t) { + const uint cell = a.pos0 + uint(t); + if (tid < ID) tail[(cell & 3) * ID + tid] = ikraw[t * s_ik + tid]; + } + return; + } + const int t = int(tg.y); + const uint cell = a.pos0 + uint(t), absp = a.pos_base + cell; + if (r < QH) { + const float x = qfull[(t * QH + r) * 2 * QD + tid]; + const float o = norm_rope(x, 0.0f, qn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); + if (qsa_qf16) ((device half*)qcur)[(t * QH + r) * QD + tid] = half(o); + else qcur[(t * QH + r) * QD + tid] = o; + } else if (r < QH + 2 * QKV) { + const bool isv = r >= QH + QKV; + const int h = r - QH - (isv ? QKV : 0); + float x = (isv ? vcur : kcur)[t * s_kv + h * QD + tid]; + if (!isv) x = norm_rope(x, 0.0f, kn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); + // int8 with an fp16 scale per 64 values: groups = simdgroup pairs (64 threads) + float m = simd_max(fabs(x)); + if (lane == 0) amax[sg] = m; + threadgroup_barrier(mem_flags::mem_threadgroup); + const int g = tid >> 6; + m = max(amax[2 * g], amax[2 * g + 1]); + const half sh = half(m / 127.0f); + const float sf = float(sh); + int q = 0; + if (sf > 0.0f) q = clamp(int(rint(x / sf)), -127, 127); + const uint rowi = cell * QKV + uint(h); + (isv ? vq : kq)[rowi * QD + tid] = char(q); + if ((tid & 63) == 0) (isv ? vsc : ksc)[rowi * 4 + g] = sh; + } else { + if (a.pos0 + a.n_tok <= uint(QBUDGET)) return; + const int h = r - QH - 2 * QKV; + const float x = tid < ID ? iqraw[t * s_iq + 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); + if (tid < ID) qidx[(t * IH + h) * ID + tid] = o; + } +} + +// scores [n_tok][max_blocks]; grid (ceil(max_blocks / 8), n_tok), 256 threads: a simdgroup per block. +kernel void qsa_score(device const float* qidx [[buffer(0)]], device const float* pooled [[buffer(1)]], + device const float* dead [[buffer(2)]], device float* scores [[buffer(3)]], + constant QsaArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + const int t = int(tg.y); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + const uint b = tg.x * 8 + sg; + if (b > n_bid) return; + device const float* key = b == n_bid ? dead : pooled + b * ID; + device const float* qp = qidx + (t * IH) * ID + lane; + float k[4], q[IH][4]; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) k[i] = key[lane + 32 * i]; + #pragma clang loop unroll(full) + for (int h = 0; h < IH; ++h) { + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) q[h][i] = qp[h * ID + 32 * i]; + } + float s = 0.0f; + #pragma clang loop unroll(full) + for (int h = 0; h < IH; ++h) { + float d = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) d = fma(q[h][i], k[i], d); + d = simd_sum(d); + s += max(d, 0.0f); + } + if (b == n_bid && (n_kv & 3) != 0) s += 1e9f; + if (lane == 0) scores[uint(t) * (a.max_cells / 4 + 2) + b] = s; +} + +// ids [n_tok][QBUDGET], ascending. One threadgroup of 1024 threads per token. Order key of a float: larger score = +// larger key; a radix select (4 passes of 8 bits) over the blocks with weights = their cell counts finds the +// threshold key T and how many cells with key == T to take (the lowest cells first). +static inline uint okey(float f) { + const uint u = as_type(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]], + constant QsaArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint hist[256]; + threadgroup uint sel[4]; + threadgroup uint2 pref[32]; + threadgroup uint2 totals; + const int t = int(tg.x); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2, nb = n_bid + 1; + device int* out = ids + t * QBUDGET; + device const float* sc = scores + uint(t) * (a.max_cells / 4 + 2); + if (n_kv <= QBUDGET) { + for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); + return; + } + const bool cached = nb <= 4096; + uint keys[4]; + if (cached) + for (uint b = tid, i = 0; b < nb; b += 1024, ++i) keys[i] = okey(sc[b]); + // radix select of the `need`-th largest cell key + uint prefix = 0, pmask = 0, need = QBUDGET; + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint b = tid; b < nb; b += 1024) { + const uint k = cached ? keys[b / 1024] : okey(sc[b]); + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + uint counts[8], sum = 0; + for (uint j = 0; j < 8; ++j) { + counts[j] = atomic_load_explicit(&hist[255 - (lane * 8 + j)], memory_order_relaxed); + sum += counts[j]; + } + uint acc = simd_prefix_exclusive_sum(sum); + if (acc < need && acc + sum >= need) { + for (uint j = 0; j < 8; ++j) { + if (acc + counts[j] >= need) { + sel[0] = 255 - (lane * 8 + j); + sel[1] = need - acc; + break; + } + acc += counts[j]; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; + } + // T = prefix; take every cell with key > T and the first `need` cells (lowest index) with key == T, ascending. + // Pass over blocks in order of index with a prefix scan of each block's selected cell count. + const uint T = prefix; + uint base = 0, eq_taken = 0; + for (uint b0 = 0; b0 < nb; b0 += 1024) { + const uint b = b0 + tid; + uint w = 0, weq = 0; + if (b < nb) { + const uint k = cached ? keys[b / 1024] : okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + if (k > T) w = cells; + else if (k == T) weq = cells; + } + const uint eq_before = simd_prefix_exclusive_sum(weq); + const uint gt_before = simd_prefix_exclusive_sum(w); + if (lane == 31) pref[sg] = uint2(eq_before + weq, gt_before + w); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + const uint2 count = pref[lane]; + const uint2 before = uint2(simd_prefix_exclusive_sum(count.x), + simd_prefix_exclusive_sum(count.y)); + pref[lane] = before; + if (lane == 31) totals = before + count; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint2 group_before = pref[sg], count = totals; + simdgroup_barrier(mem_flags::mem_threadgroup); + const uint eq0 = eq_taken + group_before.x + eq_before; + const uint take_eq = eq0 < need ? min(weq, need - eq0) : 0u; + const uint mine = w + take_eq; + // Clipping the equal-key prefix preserves the lowest-cell tie break. + const uint remaining = need - min(need, eq_taken); + const uint off = base + group_before.y + gt_before + + min(remaining, group_before.x + eq_before); + for (uint c = 0; c < mine; ++c) out[off + c] = int(4 * b + c); + base += count.y + min(remaining, count.x); + eq_taken += count.x; + } +} + +// partials: [n_tok][chunks][QH][QD + 2] (acc, then m, l). grid (chunks, QKV, n_tok), 256 threads. +kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + 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 HPG = QH / QKV; // 12 query heads per kv head + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH]; + threadgroup int cid[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; + const int nchunks = (QBUDGET + QCH - 1) / QCH; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + if (uint(c0) >= width) { + // an empty chunk: m = -inf, l = 0 (the merge skips it) + for (int h = 0; h < HPG; ++h) + if (tid == 0) { pp[(g * HPG + h) * (QD + 2) + QD] = -INFINITY; pp[(g * HPG + h) * (QD + 2) + QD + 1] = 0.0f; } + return; + } + const int nc = min(QCH, int(width) - c0); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + g * HPG + h) * QD + tid]; + if (tid < QCH) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + // scores: 12 x 64 dots of 256; simdgroup sg takes cells sg, sg + 8, ...; lanes split the 256 dims (8 each) + threadgroup const float* q_lane = &qs[0][lane * 8]; + for (int c = sg; c < nc; c += 8) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char* kr = kq + rowi * QD + lane * 8; + const char4 k0 = *(device const char4*)(kr); + const char4 k1 = *(device const char4*)(kr + 4); + const float scl = float(ksc[rowi * 4 + (lane >> 3)]); + const float kv[8] = { + float(k0.x) * scl, float(k0.y) * scl, float(k0.z) * scl, float(k0.w) * scl, + float(k1.x) * scl, float(k1.y) * scl, float(k1.z) * scl, float(k1.w) * scl + }; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + threadgroup const float* qh = q_lane + h * QD; + float d = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) d = fma(qh[i], kv[i], d); + d = simd_sum(d); + if (lane == 0) s[h][c] = d * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // per head: max, exp, sum (simdgroup h % 8 handles heads h, h + 8) + threadgroup float mh[HPG], lh[HPG]; + // nc <= QCH = 64: a lane holds cells lane and lane + 32 in registers (the same max and sum as a strided loop) + const bool in0 = int(lane) < nc, in1 = int(lane) + 32 < nc; + for (int h = sg; h < HPG; h += 8) { + const float s0 = in0 ? s[h][lane] : -INFINITY, s1 = in1 ? s[h][lane + 32] : -INFINITY; + const float m = simd_max(max(s0, s1)); + const float p0 = in0 ? exp(s0 - m) : 0.0f, p1 = in1 ? exp(s1 - m) : 0.0f; + if (in0) s[h][lane] = p0; + if (in1) s[h][lane + 32] = p1; + const float l = simd_sum(p0 + p1); + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // acc[h][j] = sum_c p[h][c] v[c][j]: thread = dim j + float acc[HPG]; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + // four cells' V loads in flight at once; the fma order per head stays c, c + 1, ... + auto vcell = [&](int c) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + return float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); + }; + int c = 0; + for (; c + 4 <= nc; c += 4) { + const float v0 = vcell(c), v1 = vcell(c + 1), v2 = vcell(c + 2), v3 = vcell(c + 3); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + float x = fma(s[h][c], v0, acc[h]); + x = fma(s[h][c + 1], v1, x); + x = fma(s[h][c + 2], v2, x); + acc[h] = fma(s[h][c + 3], v3, x); + } + } + for (; c < nc; ++c) { + const float v = vcell(c); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][c], v, acc[h]); + } + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (g * HPG + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} + +constant float qsa_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; +constant float qsa_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + +static inline void qsa_write_prepared(float y, device float* xp, constant QsaArgs& a, + int t, int h, ushort tid) { + // prepared planes (xprep layout) of the n = QH * QD vector + const uint n = QH * QD, plane = a.n_tok * n, d = uint(h) * QD + tid; + const int j = tid & 15; + device float* raw = xp + uint(t) * n + d; + raw[0] = y; + raw[plane] = y * qsa_q2s[j & 7]; + raw[2 * plane] = y * qsa_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? y * 0.00390625f : y; + float s0 = 0.0f, s1 = 0.0f; + #pragma clang loop unroll(full) + for (int k = 0; k < 16; k += 2) { + s0 += simd_shuffle_down(y, ushort(k)); + s1 += simd_shuffle_down(y, ushort(k + 1)); + } + if (j == 0) xp[4 * plane + uint(t) * (n / 16) + d / 16] = s0 + s1; +} + +// grid (QH, n_tok), 256 threads: out = merge * sigmoid(gate), as prepared activations of n = QH * QD per token +kernel void qsa_merge(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], + device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const int h = int(tg.x), t = int(tg.y); + const int nchunks = (QBUDGET + QCH - 1) / QCH; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nc = int((width + QCH - 1) / QCH); + constexpr uint STRIDE = QH * (QD + 2); + device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2); + // the chunks' (m, l) and weights exp(m - M) once per threadgroup, not once per thread (nc <= 33 < 64) + threadgroup float mk[64], lk[64], ek[64], L_sh; + if (tid < uint(nc)) { + const float2 ml = *(device const float2*)(pp + tid * STRIDE + QD); + mk[tid] = ml.x; + lk[tid] = ml.y; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float M = -INFINITY; + for (int k = tid & 31; k < nc; k += 32) M = max(M, mk[k]); + M = simd_max(M); + if (tid < uint(nc)) ek[tid] = exp(mk[tid] - M); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + float L = 0.0f; + for (int k = 0; k < nc; ++k) L = fma(lk[k], ek[k], L); + L_sh = L; + } + float o = 0.0f; + device const float* q = pp + tid; + for (int k = 0; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); + threadgroup_barrier(mem_flags::mem_threadgroup); + o = o / L_sh; + const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; + const float y = o * (1.0f / (1.0f + exp(-gate))); + qsa_write_prepared(y, xp, a, t, h, tid); +} + diff --git a/results/round39/model-ab/candidate-kernels/qsa2.metal b/results/round39/model-ab/candidate-kernels/qsa2.metal new file mode 100644 index 0000000..30439a1 --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/qsa2.metal @@ -0,0 +1,346 @@ +// metal/kernels/qsa2.metal - QSA decode v2 (after qsa.metal): the same selection and attention math, more parallel. +// +// qsa_attn2 chunks of 32 cells, the 12 query heads of a kv head in 2 groups of 6: 4x the threadgroups of +// qsa_attn (at a 1K context 128 instead of 32); V loads 8 cells at a time. +// qsa_merge2 qsa_merge for that chunking. +// qsa_topk2 the digit of each radix pass found by a parallel suffix scan of the 256 bins (simdgroup prefix +// sums) instead of one thread; cells emitted from contiguous block ranges per thread after one scan. + +#define QCH2 32 +#define QHG 2 + +// partials: [n_tok][chunks2][QH][QD + 2]. grid (chunks2, QKV * QHG, n_tok), 256 threads. +kernel void qsa_attn2(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + 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 HPG = QH / QKV / QHG; // 6 + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH2]; + threadgroup float mh[HPG], lh[HPG]; + threadgroup int cid[QCH2]; + const int chunk = int(tg.x), g = int(tg.y) / QHG, hg = int(tg.y) % QHG, t = int(tg.z); + const int h0 = g * (QH / QKV) + hg * HPG; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH2; + const int nchunks = (QBUDGET + QCH2 - 1) / QCH2; + if (uint(c0) >= width) return; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + const int nc = min(QCH2, int(width) - c0); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + h0 + h) * QD + tid]; + if (tid < QCH2) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup const float* q_lane = &qs[0][lane * 8]; + for (int c = sg; c < nc; c += 8) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char4* kr = (device const char4*)(kq + rowi * QD + lane * 8); + const char4 k0 = kr[0], k1 = kr[1]; + const float scl = float(ksc[rowi * 4 + (lane >> 3)]); + const float kv[8] = {float(k0.x) * scl, float(k0.y) * scl, float(k0.z) * scl, float(k0.w) * scl, + float(k1.x) * scl, float(k1.y) * scl, float(k1.z) * scl, float(k1.w) * scl}; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + float d = 0.0f; + #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) d = fma(q_lane[h * QD + i], kv[i], d); + d = simd_sum(d); + if (lane == 0) s[h][c] = d * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg < HPG) { + const int h = sg; + const float sv = lane < nc ? s[h][lane] : -INFINITY; + const float m = simd_max(sv); + const float p = lane < nc ? exp(sv - m) : 0.0f; + const float l = simd_sum(p); + if (lane < nc) s[h][lane] = p; + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float acc[HPG]; + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + const uint v_scale_off = tid >> 6; + int cb = 0; + for (; cb + 8 <= nc; cb += 8) { + float v[8]; + #pragma clang loop unroll(full) + for (int c = 0; c < 8; ++c) { + const uint rowi = uint(cid[cb + c]) * QKV + uint(g); + v[c] = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + v_scale_off]); + } + #pragma clang loop unroll(full) + for (int c = 0; c < 8; ++c) + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb + c], v[c], acc[h]); + } + for (; cb < nc; ++cb) { + const uint rowi = uint(cid[cb]) * QKV + uint(g); + const float v = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + v_scale_off]); + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb], v, acc[h]); + } + #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (h0 + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} + +kernel void qsa_merge2(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], + device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float mk[128], lk[128], ek[128], L_sh; + const int h = int(tg.x), t = int(tg.y); + const int nchunks = (QBUDGET + QCH2 - 1) / QCH2; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nc = int((width + QCH2 - 1) / QCH2); + constexpr uint STRIDE = QH * (QD + 2); + device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2); + if (tid < uint(nc)) { + const float2 ml = *(device const float2*)(pp + tid * STRIDE + QD); + mk[tid] = ml.x; + lk[tid] = ml.y; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float M = -INFINITY; + for (int k = tid & 31; k < nc; k += 32) M = max(M, mk[k]); + M = simd_max(M); + if (tid < uint(nc)) ek[tid] = exp(mk[tid] - M); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + float L = 0.0f; + for (int k = 0; k < nc; ++k) L = fma(lk[k], ek[k], L); + L_sh = L; + } + float o = 0.0f; + device const float* q = pp + tid; + int k = 0; + for (; k + 3 < nc; k += 4, q += 4 * STRIDE) { + const float p0 = q[0], p1 = q[STRIDE], p2 = q[2 * STRIDE], p3 = q[3 * STRIDE]; + o = fma(p0, ek[k], o); + o = fma(p1, ek[k + 1], o); + o = fma(p2, ek[k + 2], o); + o = fma(p3, ek[k + 3], o); + } + for (; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); + threadgroup_barrier(mem_flags::mem_threadgroup); + o = o / L_sh; + const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; + const float y = o * (1.0f / (1.0f + exp(-gate))); + qsa_write_prepared(y, xp, a, t, h, tid); +} + +kernel void qsa_topk2(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]], + constant QsaArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint hist[256]; + threadgroup uint sgs[32], sgs2[32]; + threadgroup uint sel[2]; + const int t = int(tg.x); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2, nb = n_bid + 1; + device int* out = ids + t * QBUDGET; + device const float* sc = scores + uint(t) * (a.max_cells / 4 + 2); + if (n_kv <= QBUDGET) { + for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); + return; + } + if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + uint prefix = 0, pmask = 0, need = QBUDGET; + if (nb <= 4096) { + uint keys[4]; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) { + const uint b = tid + 1024u * uint(i); + keys[i] = b < nb ? okey(sc[b]) : 0u; + } + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) { + const uint b = tid + 1024u * uint(i); + if (b >= nb) break; + const uint k = keys[i]; + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint cnt = 0, ex = 0; + if (tid < 256) { + cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); + if (pass < 3) atomic_store_explicit(&hist[255 - tid], 0u, memory_order_relaxed); + ex = simd_prefix_exclusive_sum(cnt); + if (lane == 31) sgs[sg] = ex + cnt; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 256) { + uint above = ex; + for (int i = 0; i < int(sg); ++i) above += sgs[i]; + if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; + } + } else { + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + for (uint b = tid; b < nb; b += 1024) { + const uint k = okey(sc[b]); + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint cnt = 0, ex = 0; + if (tid < 256) { + cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); + if (pass < 3) atomic_store_explicit(&hist[255 - tid], 0u, memory_order_relaxed); + ex = simd_prefix_exclusive_sum(cnt); + if (lane == 31) sgs[sg] = ex + cnt; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 256) { + uint above = ex; + for (int i = 0; i < int(sg); ++i) above += sgs[i]; + if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; + } + } + const uint T = prefix; + const uint per = (nb + 1023) / 1024, b0 = min(nb, uint(tid) * per), b1 = min(nb, b0 + per); + uint gt = 0, eq = 0; + for (uint b = b0; b < b1; ++b) { + const uint k = okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + if (k > T) gt += cells; + else if (k == T) eq += cells; + } + const uint eqb = simd_prefix_exclusive_sum(eq); + if (lane == 31) sgs2[sg] = eqb + eq; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint eq_before = eqb; + for (int i = 0; i < int(sg); ++i) eq_before += sgs2[i]; + const uint eq_take = eq_before >= need ? 0u : min(eq, need - eq_before); + const uint mine = gt + eq_take; + const uint mb = simd_prefix_exclusive_sum(mine); + if (lane == 31) sgs[sg] = mb + mine; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint off = mb; + for (int i = 0; i < int(sg); ++i) off += sgs[i]; + uint left = eq_take; + for (uint b = b0; b < b1; ++b) { + const uint k = okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + uint take = 0; + if (k > T) take = cells; + else if (k == T) { take = min(cells, left); left -= take; } + for (uint c = 0; c < take; ++c) out[off + c] = int(4 * b + c); + off += take; + } +} + +// qsa_attn3: qsa_attn's chunking and partial layout (qsa_merge reads it), without a reduction per (cell, head): +// for q.k a lane owns a cell (8 simdgroups = 2 halves of the 64 cells x 4 groups of 3 heads; q from threadgroup +// memory, read by every lane alike), for p.v a thread owns a dimension. grid (33, QKV, n_tok), 256 threads. +kernel void qsa_attn3(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + 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 HPG = QH / QKV; // 12 + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH]; + threadgroup float mh[HPG], lh[HPG]; + threadgroup int cid[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; + const int nchunks = (QBUDGET + QCH - 1) / QCH; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + if (uint(c0) >= width) return; + const int nc = min(QCH, int(width) - c0); + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + g * HPG + h) * QD + tid]; + if (tid < QCH) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + { + const int half_ = sg & 1, hgrp = sg >> 1; // cells 32 * half_ + lane, heads 3 * hgrp .. + 3 + const int c = 32 * half_ + lane; + if (c < nc) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char4* kr = (device const char4*)(kq + rowi * QD); + float d0 = 0.0f, d1 = 0.0f, d2 = 0.0f; + for (int gq = 0; gq < 4; ++gq) { + const float scl = float(ksc[rowi * 4 + gq]); + float e0 = 0.0f, e1 = 0.0f, e2 = 0.0f; + for (int i = 0; i < 16; ++i) { + const char4 k4 = kr[gq * 16 + i]; + const int base = gq * 64 + 4 * i; + const float kv[4] = {float(k4.x), float(k4.y), float(k4.z), float(k4.w)}; + for (int e = 0; e < 4; ++e) { + e0 = fma(qs[3 * hgrp][base + e], kv[e], e0); + e1 = fma(qs[3 * hgrp + 1][base + e], kv[e], e1); + e2 = fma(qs[3 * hgrp + 2][base + e], kv[e], e2); + } + } + d0 = fma(e0, scl, d0); + d1 = fma(e1, scl, d1); + d2 = fma(e2, scl, d2); + } + s[3 * hgrp][c] = d0 * 0.0625f; + s[3 * hgrp + 1][c] = d1 * 0.0625f; + s[3 * hgrp + 2][c] = d2 * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int h = sg; h < HPG; h += 8) { + float m = -INFINITY; + for (int c = lane; c < nc; c += 32) m = max(m, s[h][c]); + m = simd_max(m); + float l = 0.0f; + for (int c = lane; c < nc; c += 32) { + const float p = exp(s[h][c] - m); + s[h][c] = p; + l += p; + } + l = simd_sum(l); + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float acc[HPG]; + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + for (int cb = 0; cb < nc; cb += 8) { + float v[8]; + const int m = min(8, nc - cb); + for (int c = 0; c < 8; ++c) { + v[c] = 0.0f; + if (c < m) { + const uint rowi = uint(cid[cb + c]) * QKV + uint(g); + v[c] = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); + } + } + for (int c = 0; c < m; ++c) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb + c], v[c], acc[h]); + } + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (g * HPG + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} diff --git a/results/round39/model-ab/candidate-kernels/qsa_prompt.metal b/results/round39/model-ab/candidate-kernels/qsa_prompt.metal new file mode 100644 index 0000000..2ee4a4b --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/qsa_prompt.metal @@ -0,0 +1,461 @@ +// metal/kernels/qsa_prompt.metal - QSA attention for a PROMPT chunk on the matrix units (after qsa.metal: QH, QKV, QD, +// QBUDGET, QCH, QsaArgs). The same split-K partials as qsa_attn (part [t][chunk][QH][QD + 2]: acc, then m, l), so +// qsa_merge is unchanged; per (chunk of QCH = 64 selected cells, kv head g, token t) one simdgroup: +// S[12 x 64] = Qh[12 x 256] . K^T / 16 single-simdgroup matmul2d, K (int8 * fp16 scale per 64) decoded into +// the right-input cooperative tensor, 32 dims a step +// per head: m = max, p = exp(s - m), l = sum p (fp32, through threadgroup memory), P -> fp16 +// acc[12 x 256] = P[12 x 64] . V[64 x 256] four 64-dim matmuls, V decoded into the right input +// The operands are fp16 (q rounded from fp32; k, v = int8 * scale rounded): to rounding of qsa_attn's fp32. +#include +using namespace mpp::tensor_ops; +#ifndef STRATA_COOP_K4 +#define STRATA_COOP_K4 +// function constant 7: the right-input cooperative tensor holds runs of 4 consecutive k at one n per lane (the M5's +// layout; tests/coop_layout). false (the M4: runs of 2 n at one k): every element decoded on its own. +constant bool COOP_K4 [[function_constant(7)]]; +constant bool coop_k4 = is_function_constant_defined(COOP_K4) ? COOP_K4 : true; +#endif + +// a bisection probe (function constant 8): stop at stage n and write m = n (1: after the ids, 2: after S, 3: after +// the softmax); 0 = normal +constant int QSA_STOP [[function_constant(8)]]; +constant int qsa_stop = is_function_constant_defined(QSA_STOP) ? QSA_STOP : 0; + +// grid (chunks, QKV, n_tok), 32 threads. qh: the q heads as fp16 [t][QH][QD] (qcur rounded). +kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int HPG = QH / QKV, NCH = (QBUDGET + QCH - 1) / QCH; + threadgroup float s[16 * QCH]; + threadgroup half p[16 * QCH]; + threadgroup int cid[QCH]; // coop_k4: byte offsets; otherwise cell ids + threadgroup half4 ks[QCH], vs[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; + device float* pp = part + ((uint(t) * NCH + chunk) * QH + g * HPG) * (QD + 2); + if (uint(c0) >= width) { + if (lane < HPG) { pp[lane * (QD + 2) + QD] = -INFINITY; pp[lane * (QD + 2) + QD + 1] = 0.0f; } + return; + } + const int nc = min(QCH, int(width) - c0); + for (int c = lane; c < QCH; c += 32) { + const int id = c < nc ? (n_kv <= QBUDGET ? c0 + c : ids[t * QBUDGET + c0 + c]) : -1; + cid[c] = coop_k4 ? int((uint(max(id, 0)) * QKV + g) * QD) : id; + if (id >= 0) { + const uint row4 = (uint(id) * QKV + g) * 4; + if (coop_k4) ks[c] = *(device const half4*)(ksc + row4); + vs[c] = *(device const half4*)(vsc + row4); + } else { + ks[c] = half4(0.0h); + vs[c] = half4(0.0h); + } + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (qsa_stop == 1) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 1.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } + // ---- S = Q K^T + { + constexpr auto desc = matmul2d_descriptor(16, QCH, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tQ((device half*)qh + (uint(t) * QH + g * HPG) * QD, + dextents(QD, HPG)); + auto rk = op.template get_right_input_cooperative_tensor(); + auto tA = tQ.template slice<32, 16>(0, 0); + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + // the second half repeats these eight cells, 16 k later; empty cells read row 0 at scale 0. + uint ko[8]; + float ks8[8][4]; + if (coop_k4) +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = rk.get_multidimensional_index(uint16_t(4 * j))[1]; + ko[j] = uint(cid[c]); +#pragma clang loop unroll(full) + for (int h = 0; h < 4; ++h) ks8[j][h] = float(ks[c][h]); + } + for (int k0 = 0; k0 < QD; k0 += 32) { + if (!coop_k4) { + for (uint16_t i = 0; i < rk.get_capacity(); ++i) { + const auto ix = rk.get_multidimensional_index(i); + const int k = k0 + ix[0], id = cid[ix[1]]; + // the row is clamped and the select is on the value: a select with the device load in the arm + // loses the threadgroup on the M1 (macOS 26, measured); the discarded value cannot matter + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)]); + rk[i] = id >= 0 ? half(v) : half(0.0f); + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rk.get_capacity(); i += 4) { + const int k = k0 + rk.get_multidimensional_index(i)[0], j = (i >> 2) & 7; + const half4 w = half4(float4(*(device const char4*)(kq + (ko[j] + uint(k)))) * ks8[j][k >> 6]); + rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w; + } + tA = tQ.template slice<32, 16>(k0, 0); + op.run(tA, rk, acc); + } + tensor, tensor_inline> tS(s, dextents(QCH, 16)); + acc.store(tS); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (qsa_stop == 2) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 2.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } + // ---- softmax pieces per head (lane = head for the reductions; 12 heads) + float mh = -INFINITY, lh = 0.0f; + if (lane < HPG) { +#pragma clang loop unroll_count(4) + for (int c = 0; c < nc; ++c) mh = max(mh, s[lane * QCH + c] * 0.0625f); +#pragma clang loop unroll_count(4) + for (int c = 0; c < QCH; ++c) { + const float e = c < nc ? exp(s[lane * QCH + c] * 0.0625f - mh) : 0.0f; + lh += e; + p[lane * QCH + c] = half(e); + } + } else if (lane < 16) { +#pragma clang loop unroll_count(4) + for (int c = 0; c < QCH; ++c) p[lane * QCH + c] = half(0.0f); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (qsa_stop == 3) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 3.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } + // ---- acc = P V, 64 dims at a time + { + // K = 32 cells a run (two runs a 64-cell chunk): K = 64 runs compute garbage on the M4 + constexpr auto desc = matmul2d_descriptor(16, 64, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tP(p, dextents(QCH, 16)); + tensor, tensor_inline> tO(pp, dextents(QD + 2, HPG)); + for (int d0 = 0; d0 < QD; d0 += 64) { + auto rv = op.template get_right_input_cooperative_tensor(); + auto tPh = tP.template slice<32, 16>(0, 0); + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + for (int c0 = 0; c0 < nc; c0 += 32) { + if (!coop_k4) { +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); ++i) { + const auto ix = rv.get_multidimensional_index(i); // (cell, dim) + const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(vq[row * QD + d]) * float(vs[c][d >> 6]); + rv[i] = id >= 0 ? half(v) : half(0.0f); + } + } else { + uint vo[8]; + float vs8[8]; + const int cb = c0 + rv.get_multidimensional_index(uint16_t(0))[0]; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = cb + (j & 3) + 16 * (j >> 2); + vo[j] = uint(cid[c]) + uint(d0); + vs8[j] = float(vs[c][d0 >> 6]); + } +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); i += 4) { + const int d = rv.get_multidimensional_index(i)[1], r = i < rv.get_capacity() / 2 ? 0 : 4; +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) rv[i + j] = half(float(vq[vo[r + j] + uint(d)]) * vs8[r + j]); + } + } + tPh = tP.template slice<32, 16>(c0, 0); + op.run(tPh, rv, acc); + } + if (qsa_stop == 4) { // the bisection: the old direct store into the 258-float rows + auto mO = tO.slice(d0, 0); + acc.store(mO); + continue; + } + // through threadgroup memory (s is free now): a matmul2d store into rows of QD + 2 floats loses the + // whole threadgroup on the M4 + tensor, tensor_inline> tT(s, dextents(64, 16)); + simdgroup_barrier(mem_flags::mem_threadgroup); + acc.store(tT); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int it = lane; it < HPG * 16; it += 32) { + const int h = it >> 4, d = (it & 15) * 4; + const float4 v = *(threadgroup const float4*)(s + h * 64 + d); + // odd heads start on an 8-byte boundary, not a float4 boundary. + device float* dst = pp + h * (QD + 2) + d0 + d; + dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w; + } + } + } + if (lane < HPG) { pp[lane * (QD + 2) + QD] = mh; pp[lane * (QD + 2) + QD + 1] = lh; } +} + +// qcur (fp32) -> fp16 for qsa_attn_mm +kernel void qsa_q_f16(device const float* q [[buffer(0)]], device half* qh [[buffer(1)]], constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + if (gid < n) qh[gid] = half(q[gid]); +} + +// qsa_score as a GEMM (fp32 operands, so the selection sees qsa_score's values to fp32 rounding): S[(t, h)][b] = +// qidx[t][h] . pooled[b] for 16 tokens x IH heads (64 rows) by 64 blocks a threadgroup; the epilogue sums relu over +// the heads. The tail block n_bid (the `dead` key, + 1e9 when it has cells) and blocks past it are left to the +// scalar path: the tail is scored here by one simdgroup a token; blocks > n_bid are not written (qsa_topk reads +// b <= n_bid). grid (ceil(max blocks / 64), ceil(T / 16)), 128 threads. +kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const float* pooled [[buffer(1)]], + device const float* dead [[buffer(2)]], device float* scores [[buffer(3)]], + constant QsaArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float ot[64 * 64]; + const int b0 = int(tg.x) * 64, t0 = int(tg.y) * 16; + const uint max_t = min(uint(t0 + 15), a.n_tok - 1u); + const uint max_bid = (a.pos0 + max_t + 1u) >> 2; + // column zero still owns every token's tail score. + if (tg.x != 0 && uint(b0) >= max_bid) return; + const uint nb_all = (a.pos0 + a.n_tok) / 4 + 1; // blocks any token of the chunk can score + tensor, tensor_inline> tQ((device float*)qidx, dextents(ID, int(a.n_tok) * IH)); + tensor, tensor_inline> tK((device float*)pooled, dextents(ID, int(nb_all))); + tensor, tensor_inline> tO(ot, dextents(64, 64)); + constexpr auto desc = matmul2d_descriptor(64, 64, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tQ.slice(0, t0 * IH); + auto mB = tK.slice(0, b0); + op.run(mA, mB, tO); + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint stride = a.max_cells / 4 + 2; + const int j = tid & 63, b = b0 + j; + for (int tt = tid >> 6; tt < 16; tt += 2) { + const int t = t0 + tt; + if (t >= int(a.n_tok)) break; + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + if (uint(b) >= n_bid) continue; + float s = 0.0f; + for (int h = 0; h < IH; ++h) s += max(ot[(tt * IH + h) * 64 + j], 0.0f); + scores[uint(t) * stride + b] = s; + } + // the tail block of each token in this threadgroup's block range (the first column threadgroup does it) + if (tg.x == 0) { + float k[4]; + for (int i = 0; i < 4; ++i) k[i] = dead[lane + 32 * i]; + for (int tt = sg; tt < 16; tt += 4) { + const int t = t0 + tt; + if (t >= int(a.n_tok)) break; + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + float s = 0.0f; + for (int h = 0; h < IH; ++h) { + float d = 0.0f; + for (int i = 0; i < 4; ++i) d = fma(qidx[(t * IH + h) * ID + lane + 32 * i], k[i], d); + d = simd_sum(d); + s += max(d, 0.0f); + } + if ((n_kv & 3) != 0) s += 1e9f; + if (lane == 0) scores[uint(t) * stride + n_bid] = s; + } + } +} + +// qsa_flash: the whole attention of a (kv head g, token t) in one threadgroup (no split-K partials, no merge): the +// width selected cells in chunks of 64 with an online softmax; 4 simdgroups, simdgroup s owns dims 64 s .. 64 s + 63 +// (its slice of S's K sum and its 64 output dims). Out: xq = (acc / l) * sigmoid(gate) as fp16 (attn_output's GEMM +// operand). grid (QKV, n_tok), 128 threads. +kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device const float* qfull [[buffer(6)]], device half* xqh [[buffer(7)]], constant QsaArgs& a [[buffer(8)]], + 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 HPG = QH / QKV; + threadgroup float sp[4][HPG * QCH]; + threadgroup half* p = (threadgroup half*)sp[1]; // P [16][QCH] in fp16 over a dead S partial + threadgroup int cid[QCH]; // coop_k4: byte offsets; otherwise cell ids + threadgroup half4 ks_tg[QCH], vs_tg[QCH]; + threadgroup float fac[16], lsum[16], mrun[16]; + const int g = int(tg.x), t = int(tg.y); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nch = int((width + QCH - 1) / QCH); + if (tid < 16) { mrun[tid] = -INFINITY; lsum[tid] = 0.0f; } + constexpr auto dS = matmul2d_descriptor(16, QCH, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + constexpr auto dO = matmul2d_descriptor(16, 64, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d opS; + matmul2d opO; + tensor, tensor_inline> tQ((device half*)qh + (uint(t) * QH + g * HPG) * QD, + dextents(QD, HPG)); + tensor, tensor_inline> tP(p, dextents(QCH, 16)); + auto rv = opO.template get_right_input_cooperative_tensor(); + auto tPh = tP.template slice<32, 16>(0, 0); + auto acc = opO.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); + // (a second destination tensor, not a copy of acc: the MetalPerformancePrimitives headers of macOS 26.3 cannot copy a + // cooperative tensor - "cannot initialize a parameter of type 'void *' with an rvalue of type 'const void *'", #2323; + // tot is zeroed right below, so the copy was only there for its type) + auto tot = opO.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); // the running output; acc takes one chunk (fresh each chunk: on the M4 a run does not accumulate + // into elements rescaled in place) + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) tot[i] = 0.0f; + const int d0 = sg * 64; + if (tid < 16 - HPG) fac[HPG + tid] = 0.0f; // the padding rows' factor (P's rows are zeroed each chunk) + for (int ch = 0; ch < nch; ++ch) { + const int c0 = ch * QCH, nc = min(QCH, int(width) - c0); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < QCH) { + const int id = tid < nc ? (n_kv <= QBUDGET ? c0 + tid : ids[t * QBUDGET + c0 + tid]) : -1; + cid[tid] = coop_k4 ? int((uint(max(id, 0)) * QKV + g) * QD) : id; + if (id >= 0) { + const uint row4 = (uint(id) * QKV + g) * 4; + ks_tg[tid] = *(device const half4*)(ksc + row4); + vs_tg[tid] = *(device const half4*)(vsc + row4); + } else { + ks_tg[tid] = half4(0.0h); + vs_tg[tid] = half4(0.0h); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // Eight lanes per head; keep the denominator's original sequential sum. + const int h = sg * 4 + (lane >> 3), hlane = lane & 7; + float sc[8]; + // S partial over this simdgroup's 64 dims + { + auto rk = opS.template get_right_input_cooperative_tensor(); + auto tA = tQ.template slice<32, 16>(d0, 0); + auto s = opS.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < s.get_capacity(); ++i) + if (s.is_valid_element(i)) s[i] = 0.0f; + // the lane's eight cells (one a run of four k; the second half repeats them 16 k on): row pointers and + // scales once a chunk, unconditional loads (an empty cell reads row 0 at scale 0, and a +-0 operand sums to + // the same value as the skipped one; the branchy loads cost half the kernel) + uint ko[8]; + float ks8[8]; + if (coop_k4) { +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = rk.get_multidimensional_index(uint16_t(4 * j))[1]; + ko[j] = uint(cid[c]) + uint(d0 + rk.get_multidimensional_index(uint16_t(0))[0]); + ks8[j] = float(ks_tg[c][sg]); + } + } +#pragma clang loop unroll(full) + for (int k0 = d0; k0 < d0 + 64; k0 += 32) { + if (!coop_k4) { + for (uint16_t i = 0; i < rk.get_capacity(); ++i) { + const auto ix = rk.get_multidimensional_index(i); + const int k = k0 + ix[0], id = cid[ix[1]]; + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(kq[row * QD + k]) * float(ks_tg[ix[1]][k >> 6]); + rk[i] = id >= 0 ? half(v) : half(0.0f); + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rk.get_capacity(); i += 4) { + const int k = k0 - d0 + (i >= rk.get_capacity() / 2 ? 16 : 0), j = (i >> 2) & 7; + const half4 w = half4(float4(*(device const char4*)(kq + ko[j] + uint(k))) * ks8[j]); + rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w; + } + tA = tQ.template slice<32, 16>(k0, 0); + opS.run(tA, rk, s); + } + // Preserve ((s0 + s1) + (s2 + s3)) while storing all partials together. + tensor, tensor_inline> tS(sp[sg], dextents(QCH, HPG)); + s.store(tS); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (h < HPG) +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) { + const int q = h * QCH + hlane + 8 * k; + sc[k] = ((sp[0][q] + sp[1][q]) + (sp[2][q] + sp[3][q])) * 0.0625f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); // P overwrites sp[1] + if (h < HPG) { + float m = mrun[h]; +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) + if (hlane + 8 * k < nc) m = max(m, sc[k]); + m = max(m, simd_shuffle_xor(m, 4)); + m = max(m, simd_shuffle_xor(m, 2)); + m = max(m, simd_shuffle_xor(m, 1)); +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) { + const int c = hlane + 8 * k, q = h * QCH + c; + const float e = c < nc ? exp(sc[k] - m) : 0.0f; + p[q] = half(e); + sp[0][q] = e; + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (hlane == 0) { + const float f = exp(mrun[h] - m); + float l = lsum[h] * f; +#pragma clang loop unroll(full) + for (int c = 0; c < QCH; c += 4) { // in order, four loads at a time + const float4 v = *(threadgroup const float4*)(sp[0] + h * QCH + c); + l += v.x; l += v.y; l += v.z; l += v.w; + } + mrun[h] = m; + lsum[h] = l; + fac[h] = f; + } + } else { + for (int c = hlane; c < QCH; c += 8) p[h * QCH + c] = half(0.0f); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // rescale this simdgroup's accumulator rows, then acc += P V over its 64 dims + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + for (int c0 = 0; c0 < nc; c0 += 32) { // K = 32 cells a run (K = 64 is wrong on the M4) + if (!coop_k4) { +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); ++i) { + const auto ix = rv.get_multidimensional_index(i); // (cell, dim) + const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(vq[row * QD + d]) * float(vs_tg[c][sg]); + rv[i] = id >= 0 ? half(v) : half(0.0f); + } + } else { + // a lane holds the same eight cells (two runs of four, 16 apart) at each of its dims: as for rk, row + // pointers and scales once a run, unconditional loads + uint vo[8]; + float vs8[8]; + const int cb = c0 + rv.get_multidimensional_index(uint16_t(0))[0]; +#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + const int c = cb + (j & 3) + 16 * (j >> 2); + vo[j] = uint(cid[c]) + uint(d0); + vs8[j] = float(vs_tg[c][sg]); + } +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); i += 4) { + const int d = rv.get_multidimensional_index(i)[1], r = i < rv.get_capacity() / 2 ? 0 : 4; +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) rv[i + j] = half(float(vq[vo[r + j] + uint(d)]) * vs8[r + j]); + } + } + tPh = tP.template slice<32, 16>(c0, 0); + opO.run(tPh, rv, acc); + } + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) tot[i] = fma(tot[i], fac[acc.get_multidimensional_index(i)[1]], acc[i]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (!coop_k4) { + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) { + const auto ix = acc.get_multidimensional_index(i); // (dim, row) + const int h = ix[1], d = d0 + ix[0]; + if (h < HPG) { + const int hh = g * HPG + h; + const float gate = qfull[(uint(t) * QH + hh) * 2 * QD + QD + d]; + xqh[(uint(t) * QH + hh) * QD + d] = half(tot[i] / lsum[h] * (1.0f / (1.0f + exp(-gate)))); + } + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { // runs of four dims of one row + const auto ix = acc.get_multidimensional_index(i); + const int h = ix[1], d = d0 + ix[0]; + if (h < HPG) { + const int hh = g * HPG + h; + const float l = lsum[h]; + const float4 gate = *(device const float4*)(qfull + (uint(t) * QH + hh) * 2 * QD + QD + d); + const float4 v = float4(tot[i], tot[i + 1], tot[i + 2], tot[i + 3]) / l * (1.0f / (1.0f + exp(-gate))); + *(device half4*)(xqh + (uint(t) * QH + hh) * QD + d) = half4(v); + } + } +} diff --git a/results/round39/model-ab/candidate-kernels/quant.metal b/results/round39/model-ab/candidate-kernels/quant.metal new file mode 100644 index 0000000..fe667be --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/quant.metal @@ -0,0 +1,316 @@ +// metal/kernels/quant.metal - the weight formats as "units": a unit is UNIT consecutive weights of one row that one +// SIMD lane decodes at a time into raw values q[0..UNIT) and one scale s, so that the unit's dot with x is +// s * sum(q[j] * x[j]). Every GEMV kernel (moe.metal, gemv.metal) is written once over these. +// +// Decode rules transcribed from ggml-quants.c dequantize_row_* (3cf03257) and Strata's Q2_0 contract +// (docs/q2_0-contract.md); the grids and sign tables are ggml-common.h's (prepended by the runtime). The tables a +// format needs are copied to threadgroup memory once per threadgroup (constant-memory lookups with divergent indices +// serialise on Apple GPUs). +// +// Blocks start at 2-byte aligned addresses (18, 66, 74, 82, 98, 110 bytes), so every multi-byte field is read as +// ushorts. + +#define T_Q4_0 2 +#define T_Q3_K 11 +#define T_BF16 30 +#define T_F32 0 +#define T_Q5_0 6 +#define T_Q5_1 7 +#define T_Q8_0 8 +#define T_Q4_K 12 +#define T_Q5_K 13 +#define T_Q6_K 14 +#define T_IQ2_XXS 16 +#define T_IQ2_XS 17 +#define T_IQ3_XXS 18 +#define T_IQ4_NL 20 +#define T_IQ3_S 21 +#define T_IQ2_S 22 +#define T_IQ4_XS 23 +#define T_Q2_0 42 +#define T_Q2_0R 142 // Q2_0 read with raw activations (codes - 1 in the decoder; no prepared x needed) + +template struct fmt; +// QK block elements, BB block bytes, UNIT weights per unit, TAB words of threadgroup table +template <> struct fmt { enum { QK = 64, BB = 18, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 64, BB = 18, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 98, UNIT = 16, TAB = 256 }; }; +template <> struct fmt { enum { QK = 256, BB = 66, UNIT = 16, TAB = 512 }; }; +template <> struct fmt { enum { QK = 256, BB = 74, UNIT = 16, TAB = 1024 }; }; +template <> struct fmt { enum { QK = 256, BB = 82, UNIT = 16, TAB = 2048 }; }; +template <> struct fmt { enum { QK = 256, BB = 110, UNIT = 16, TAB = 512 }; }; +template <> struct fmt { enum { QK = 32, BB = 18, UNIT = 32, TAB = 256 }; }; +template <> struct fmt { enum { QK = 256, BB = 136, UNIT = 32, TAB = 256 }; }; +template <> struct fmt { enum { QK = 256, BB = 110, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 144, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 176, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 210, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 16, BB = 32, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 16, BB = 64, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 32, BB = 34, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 32, BB = 18, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 32, BB = 22, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 32, BB = 24, UNIT = 32, TAB = 1 }; }; + +static inline uint rd32(device const uchar* p) { + device const ushort* s = (device const ushort*)p; + return uint(s[0]) | (uint(s[1]) << 16); +} +static inline float rdh(device const uchar* p) { return float(*(device const half*)p); } + +// ---- threadgroup tables ------------------------------------------------------------------------------------------- +// layout in `tab` (uint words): [0, G) the grid; IQ4_NL/XS: 256 half2. +template static inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth); + +// QUANT_CTAB (function constant 23): the i-quant grids read from ggml-common's constant arrays +// instead of a threadgroup copy (no table load and no barrier wait at the threadgroup's start); the same values. +constant bool QUANT_CTAB [[function_constant(23)]]; +constant bool ctab = is_function_constant_defined(QUANT_CTAB) && QUANT_CTAB; +// ksigns_iq2xs[j] = j | parity(j) << 7 (an even count of minus signs): the low 4 bits of run hh (0 or 1). +static inline uint sign4_iq2xs(uint j, int hh) { + return hh ? ((j >> 4) | ((popcount(j) & 1u) << 3)) : j; +} +static inline uint g_iq3xxs(threadgroup const uint* tab, uint i) { return ctab ? iq3xxs_grid[i] : tab[i]; } +static inline uint g_iq3s(threadgroup const uint* tab, uint i) { return ctab ? iq3s_grid[i] : tab[i]; } +static inline uint g_iq2xxs(threadgroup const uint* tab, uint i, int w) { return ctab ? uint(iq2xxs_grid[i] >> (32 * w)) : tab[2 * i + w]; } +static inline uint g_iq2xs(threadgroup const uint* tab, uint i, int w) { return ctab ? uint(iq2xs_grid[i] >> (32 * w)) : tab[2 * i + w]; } +static inline uint g_iq2s(threadgroup const uint* tab, uint i, int w) { return ctab ? uint(iq2s_grid[i] >> (32 * w)) : tab[2 * i + w]; } + +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} + +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 256; i += nth) tab[i] = iq3xxs_grid[i]; +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 512; i += nth) tab[i] = iq3s_grid[i]; +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 256; i += nth) { tab[2 * i] = uint(iq2xxs_grid[i]); tab[2 * i + 1] = uint(iq2xxs_grid[i] >> 32); } +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 512; i += nth) { tab[2 * i] = uint(iq2xs_grid[i]); tab[2 * i + 1] = uint(iq2xs_grid[i] >> 32); } +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 1024; i += nth) { tab[2 * i] = uint(iq2s_grid[i]); tab[2 * i + 1] = uint(iq2s_grid[i] >> 32); } +} +// IQ4: a byte -> (codebook[low nibble], codebook[high nibble]) as half2, so one lookup gives two weights +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + for (int i = tid; i < 256; i += nth) + tab[i] = as_type(half2(half(kvalues_iq4nl[i & 15]), half(kvalues_iq4nl[i >> 4]))); +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + load_tab(tab, tid, nth); +} + +// ---- unit decoders: return s, fill q[UNIT] ------------------------------------------------------------------------ +// ui = unit index inside the block (0 .. QK/UNIT - 1). +template static inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab); + +// Q2_0: 64 = 4 units of 16; code c (2 bits, LSB first) -> c - 1; scale d. +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint*) { + device const ushort* p = (device const ushort*)(b + 2 + 4 * ui); + const ushort a = p[0], c = p[1]; + for (int j = 0; j < 8; ++j) { + q[j] = float((a >> (2 * j)) & 3) - 1.0f; + q[8 + j] = float((c >> (2 * j)) & 3) - 1.0f; + } + return rdh(b); +} + +// 8 signed grid bytes from a 4-byte grid word pair helper +static inline void put_signed4(thread float* q, uint g, uint s) { + const uchar4 v = as_type(g); + q[0] = (s & 1) ? -float(v[0]) : float(v[0]); + q[1] = (s & 2) ? -float(v[1]) : float(v[1]); + q[2] = (s & 4) ? -float(v[2]) : float(v[2]); + q[3] = (s & 8) ? -float(v[3]) : float(v[3]); +} + +// IQ3_XXS: 256 = 8 x 32; unit = half of a 32 (two grid pairs). +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + const uint aux = rd32(b + 2 + 64 + 4 * ib), s0 = aux >> (14 * h); + device const uchar* qs = b + 2 + 4 * ui; + #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { + const uint s = (s0 >> (7 * l2)) & 127; + put_signed4(q + 8 * l2, g_iq3xxs(tab, qs[2 * l2]), s); + put_signed4(q + 8 * l2 + 4, g_iq3xxs(tab, qs[2 * l2 + 1]), sign4_iq2xs(s, 1)); + } + return rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f); +} + +// IQ2_XXS: per 32: two u32 (4 grid indices, then 4 x 7 sign bits + 4-bit scale). +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + device const uchar* qs = b + 2 + 4 * ib + 2 * ui; + const uint a1 = rd32(b + 2 + 8 * ib + 4), s0 = a1 >> (14 * h); + #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { + const uint idx = qs[l2], s = (s0 >> (7 * l2)) & 127; + const uint2 g = ctab ? as_type(iq2xxs_grid[idx]) : ((threadgroup const uint2*)tab)[idx]; + put_signed4(q + 8 * l2, g.x, s); + put_signed4(q + 8 * l2 + 4, g.y, sign4_iq2xs(s, 1)); + } + return rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f); +} + +// IQ2_XS: uint16 qs[32] (9-bit grid index + 7-bit sign index), scales[8] nibbles. +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + device const ushort* qs = (device const ushort*)(b + 2) + 2 * ui; + #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { + const uint v = qs[l2]; + const uint idx = v & 511, s = v >> 9; + const uint2 g = ctab ? as_type(iq2xs_grid[idx]) : ((threadgroup const uint2*)tab)[idx]; + put_signed4(q + 8 * l2, g.x, s); + put_signed4(q + 8 * l2 + 4, g.y, sign4_iq2xs(s, 1)); + } + const uint sc = b[66 + ib]; + return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); +} + +// IQ2_S: qs[32] low 8 bits of the 10-bit index, signs[32], qh[8] (2 high bits per group of 8), scales[8]. +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + const uint q2 = *(device const ushort*)(b + 2 + 2 * ui); + const uint s2 = *(device const ushort*)(b + 34 + 2 * ui); + const uint qh = uint(b[66 + ib]) >> (4 * h); + #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { + const uint idx = ((q2 >> (8 * l2)) & 255) | ((qh << (8 - 2 * l2)) & 0x300); + const uint s = (s2 >> (8 * l2)) & 255; + const uint2 g = ctab ? as_type(iq2s_grid[idx]) : ((threadgroup const uint2*)tab)[idx]; + put_signed4(q + 8 * l2, g.x, s); + put_signed4(q + 8 * l2 + 4, g.y, s >> 4); + } + const uint sc = b[74 + ib]; + return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); +} + +// IQ3_S: qs[64] low 8 bits of 9-bit indices, qh[8], signs[32], scales[4] (nibble per 32). +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + const uint s2 = *(device const ushort*)(b + 74 + 2 * ui); + const uint qh = b[66 + ib]; + #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { + const int l = 2 * h + l2; + const uint i1 = uint(b[2 + 8 * ib + 2 * l]) | ((qh << (8 - 2 * l)) & 256); + const uint i2 = uint(b[2 + 8 * ib + 2 * l + 1]) | ((qh << (7 - 2 * l)) & 256); + const uint s = (s2 >> (8 * l2)) & 255; + put_signed4(q + 8 * l2, g_iq3s(tab, i1), s); + put_signed4(q + 8 * l2 + 4, g_iq3s(tab, i2), s >> 4); + } + const uint sc = b[106 + (ib >> 1)]; + return rdh(b) * fma(float((sc >> (4 * (ib & 1))) & 15), 2.0f, 1.0f); +} + +// IQ4_NL: 32 = 1 unit; low nibbles are elements 0..15, high 16..31 (one half2 lookup per byte). +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint* tab) { + device const ushort* p = (device const ushort*)(b + 2); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); + q[2 * j] = float(w0.x); + q[2 * j + 16] = float(w0.y); + q[2 * j + 1] = float(w1.x); + q[2 * j + 17] = float(w1.y); + } + return rdh(b); +} + +// IQ4_XS: 256 = 8 units of 32; 6-bit scale from scales_l nibble + 2 bits of scales_h, minus 32. +static inline float iq4xs_scale(device const uchar* b, int ui) { + const uint sh = *(device const ushort*)(b + 2); + const int ls = int((b[4 + (ui >> 1)] >> (4 * (ui & 1))) & 15) | int(((sh >> (2 * ui)) & 3) << 4); + return rdh(b) * float(ls - 32); +} +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + device const ushort* p = (device const ushort*)(b + 8 + 16 * ui); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); + q[2 * j] = float(w0.x); + q[2 * j + 16] = float(w0.y); + q[2 * j + 1] = float(w1.x); + q[2 * j + 17] = float(w1.y); + } + return iq4xs_scale(b, ui); +} + +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { + device const ushort* p = (device const ushort*)(b + 2); + for (int j = 0; j < 16; ++j) { + const ushort v = p[j]; + q[2 * j] = float(as_type(uchar(v & 255))); + q[2 * j + 1] = float(as_type(uchar(v >> 8))); + } + return rdh(b); +} + +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { + device const ushort* p = (device const ushort*)(b + 2); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + q[2 * j] = float(v & 15) - 8.0f; + q[2 * j + 16] = float((v >> 4) & 15) - 8.0f; + q[2 * j + 1] = float((v >> 8) & 15) - 8.0f; + q[2 * j + 17] = float(v >> 12) - 8.0f; + } + return rdh(b); +} + +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { + const uint hb = rd32(b + 2); + device const ushort* p = (device const ushort*)(b + 6); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const int j0 = 2 * j, j1 = 2 * j + 1; + q[j0] = float(int(v & 15) | int(((hb >> j0) & 1) << 4)) - 16.0f; + q[j0 + 16] = float(int((v >> 4) & 15) | int(((hb >> (j0 + 16)) & 1) << 4)) - 16.0f; + q[j1] = float(int((v >> 8) & 15) | int(((hb >> j1) & 1) << 4)) - 16.0f; + q[j1 + 16] = float(int(v >> 12) | int(((hb >> (j1 + 16)) & 1) << 4)) - 16.0f; + } + return rdh(b); +} + +// Q5_1: 32 = 1 unit; q = the 5-bit code (no offset; the min m is applied through the x sums), returns d +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { + const uint hb = rd32(b + 4); + device const ushort* p = (device const ushort*)(b + 8); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const int j0 = 2 * j, j1 = 2 * j + 1; + q[j0] = float(int(v & 15) | int(((hb >> j0) & 1) << 4)); + q[j0 + 16] = float(int((v >> 4) & 15) | int(((hb >> (j0 + 16)) & 1) << 4)); + q[j1] = float(int((v >> 8) & 15) | int(((hb >> j1) & 1) << 4)); + q[j1 + 16] = float(int(v >> 12) | int(((hb >> (j1 + 16)) & 1) << 4)); + } + return rdh(b); +} + +// sum over groups of LPR consecutive lanes (LPR = 1, 2, 4, 8, 16 or 32); the result is valid in every lane +template +static inline float lane_sum(float v) { + if (LPR == 32) return simd_sum(v); + for (ushort o = LPR / 2; o > 0; o >>= 1) v += simd_shuffle_xor(v, o); + return v; +} diff --git a/results/round39/model-ab/candidate-kernels/router.metal b/results/round39/model-ab/candidate-kernels/router.metal new file mode 100644 index 0000000..5c6bf59 --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/router.metal @@ -0,0 +1,173 @@ +// metal/kernels/router.metal - the router GEMV and moe_route in ONE dispatch (after quant/units/moe/gemv.metal): the +// logits rows as gemv's (LPR = 32 lanes a row, KS = 1: bitwise gemv__32_1_*_1's), 16 rows a threadgroup of 512 +// threads; every threadgroup bumps a device counter after its rows (device-scope fences around it), and the LAST one +// runs moe_route's body (the softmax / top-10 / weights, then the plan) with its 512 threads (= n_expert) and resets +// the counter. No threadgroup waits for another. ids / w / plan are bitwise moe_route's. +// Buffers: W, xp, s16, logits, args (GemvArgs: n_out = n_expert), ids, w, plan, MoeArgs, counter (one uint, zero). + +#ifndef STRATA_MOCE_FC_METAL +#define STRATA_MOCE_FC_METAL +constant int kMoceProtectFc [[function_constant(24)]]; +constant int kMoceProtect = is_function_constant_defined(kMoceProtectFc) ? kMoceProtectFc : 3; +constant uint kMoceScaleR1BitsFc [[function_constant(25)]]; +constant float kMoceScaleR1 = is_function_constant_defined(kMoceScaleR1BitsFc) ? as_type(kMoceScaleR1BitsFc) : 1.0f; +#endif +static inline void route_body(device const float* logits, device int* ids, device float* w, device int* plan, + constant MoeArgs& a, device const uint* smap, ushort tid, ushort nth, ushort sg, ushort lane, + threadgroup int* s_grp, threadgroup int* s_ent) { + if (sg < a.n_tok) { + device const float* l = logits + sg * 512; + float v[16]; + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + v[i] = idx < int(a.n_expert) ? l[idx] : -INFINITY; + } + float mx = -INFINITY; + for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + float sum = 0.0f; + for (int i = 0; i < 16; ++i) { + v[i] = (lane + 32 * i < int(a.n_expert)) ? exp(v[i] - mx) : 0.0f; + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + for (int i = 0; i < 16; ++i) { + if (lane + 32 * i >= int(a.n_expert)) { + v[i] = -INFINITY; + } else { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + } + const bool draft_moce = (a.pad & 0x80000000u) != 0u && sg > 0; + const float base_moce = (a.pad & 0x7FFFFFFFu) != 0u ? as_type(a.pad & 0x7FFFFFFFu) : 1.0f; + const float moce_scale = draft_moce ? (base_moce * base_moce) : base_moce; + const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f; + const int r_prot = draft_moce ? 1 : kMoceProtect; + const float r1_scale = kMoceScaleR1; + const float rcp_r1 = r1_scale > 1.0f ? (1.0f / r1_scale) : 1.0f; + const float r_step_scale = (r1_scale > 1.0f && moce_scale > 1.0f) ? (moce_scale * rcp_r1) : moce_scale; + ushort scaled_mask = 0; + float selected = 0.0f, ssum = 0.0f; + for (int rank = 0; rank < 10; ++rank) { + if (r1_scale > 1.0f && r_prot > 1) { + if (rank == 1 && smap != nullptr) { + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) { + v[i] *= r1_scale; + scaled_mask |= ushort(1u << i); + } + } + } else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) { + for (int i = 0; i < 16; ++i) { + if ((scaled_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) { + v[i] *= r_step_scale; + } + } + } + } else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) { + for (int i = 0; i < 16; ++i) { + const int idx = lane + 32 * i; + if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) { + v[i] *= moce_scale; + scaled_mask |= ushort(1u << i); + } + } + } + float best = v[0]; + int ex = lane; + #pragma clang loop unroll(full) + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + float true_prob = best; + const int min_scaled_rank = (r1_scale > 1.0f && r_prot > 1) ? 1 : r_prot; + if (rank >= min_scaled_rank && (moce_scale > 1.0f || r1_scale > 1.0f) && best > 0.0f) { + const ushort win_lane = ushort(ex & 31); + 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) { + // 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); + } + } + #pragma clang loop unroll(full) + for (int i = 0; i < 16; ++i) v[i] = (ex == lane + 32 * i) ? -INFINITY : v[i]; + if ((ex & 31) == lane) { + ids[sg * 10 + rank] = ex; + ssum += true_prob; + } + if (rank == lane) selected = true_prob; + } + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + ssum = max(ssum, 6.103515625e-5f); + const float inv = 1.0f / ssum; + if (lane < 10) w[sg * 10 + lane] = selected * inv; + } + threadgroup_barrier(mem_flags::mem_device); + plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent, smap); +} + +template +kernel void router_route(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + 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)]], + device atomic_uint* counter [[buffer(9)]], device const uint* smap [[buffer(10)]], + uint tgx [[threadgroup_position_in_grid]], uint ngx [[threadgroups_per_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort nth [[threads_per_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int UNIT = fmt::UNIT; + threadgroup uint tab[fmt::TAB]; + threadgroup int s_grp[64], s_ent[64]; + threadgroup uint last; + load_tab(tab, tid, nth); + threadgroup_barrier(mem_flags::mem_threadgroup); + // the GEMV: simdgroup sg -> row tg.x * 16 + sg, 32 lanes over its units + const int row = int(tgx) * 16 + sg; + const int rr = min(row, int(ga.n_out) - 1); + device const uchar* rp = W + (ulong)rr * ga.row_bytes; + const int U = int(ga.n_in) / UNIT, s16n = int(ga.n_in) / 16; + float acc[8]; +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[t] = 0.0f; + for (int u = lane; u < U; u += 32) { + float q[UNIT]; + float4 sc; + Dec::run(rp, u, q, sc, tab); +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(ga.x_stride), s16 + t * s16n, u, xs, xsum); + acc[t] += unit_combine(q, sc, xs, xsum); + } + } +#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum<32>(acc[t]); + if (lane == 0 && row < int(ga.n_out)) logits[t * int(ga.y_stride) + row] = v; + } + // the last threadgroup to finish routes + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + threadgroup_barrier(mem_flags::mem_device); + if (tid == 0) last = atomic_fetch_add_explicit(counter, 1u, memory_order_relaxed) == ngx - 1 ? 1u : 0u; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (!last) return; + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + route_body(logits, ids, w, plan, a, smap, tid, nth, sg, lane, s_grp, s_ent); + if (tid == 0) atomic_store_explicit(counter, 0u, memory_order_relaxed); +} + +template [[host_name("router_route_30")]] kernel decltype(router_route) router_route; +template [[host_name("router_route_0")]] kernel decltype(router_route) router_route; + +template [[host_name("router_route_8")]] kernel decltype(router_route) router_route; diff --git a/results/round39/model-ab/candidate-kernels/rt_bench.metal b/results/round39/model-ab/candidate-kernels/rt_bench.metal new file mode 100644 index 0000000..bdce780 --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/rt_bench.metal @@ -0,0 +1,67 @@ +// metal/kernels/rt_bench.metal - bandwidth and dispatch-cost probes (rt_bench.cpp). + +// Grid-stride streaming read: thread i reads src[i], src[i + nt], ... (coalesced across the grid). +kernel void k_read(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]], uint nt [[threads_per_grid]]) { + uint4 a = 0; + for (uint i = gid; i < n; i += nt) a += src[i]; + out[gid] = a.x ^ a.y ^ a.z ^ a.w; +} + +// The same, 4 independent loads in flight per iteration. +kernel void k_read_u4(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]], uint nt [[threads_per_grid]]) { + uint4 a = 0, b = 0, c = 0, d = 0; + uint i = gid; + for (; i + 3 * nt < n; i += 4 * nt) { + a += src[i]; + b += src[i + nt]; + c += src[i + 2 * nt]; + d += src[i + 3 * nt]; + } + for (; i < n; i += nt) a += src[i]; + a += b + c + d; + out[gid] = a.x ^ a.y ^ a.z ^ a.w; +} + +// Contiguous chunk per threadgroup (how a GEMV row block reads): threadgroup g reads [g*chunk, (g+1)*chunk). +kernel void k_read_chunk(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + constant uint& n [[buffer(2)]], constant uint& chunk [[buffer(3)]], + uint tid [[thread_position_in_threadgroup]], uint tgs [[threads_per_threadgroup]], + uint g [[threadgroup_position_in_grid]], uint gid [[thread_position_in_grid]]) { + uint4 a = 0, b = 0; + const uint base = g * chunk, end = min(base + chunk, n); + uint i = base + tid; + for (; i + tgs < end; i += 2 * tgs) { + a += src[i]; + b += src[i + tgs]; + } + for (; i < end; i += tgs) a += src[i]; + a += b; + out[gid] = a.x ^ a.y ^ a.z ^ a.w; +} + +// Read + write (a state update): dst[i] = src[i] * 1.0001 for float4. +kernel void k_copy(device const float4* src [[buffer(0)]], device float4* dst [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]], uint nt [[threads_per_grid]]) { + for (uint i = gid; i < n; i += nt) dst[i] = src[i] * 1.0001f; +} + +// Dispatch-cost probes. +kernel void k_empty(device uint* out [[buffer(0)]], uint gid [[thread_position_in_grid]]) { + if (gid == 0xFFFFFFFFu) out[0] = 1; +} + +// A dependent chain: each dispatch reads what the previous one wrote. +kernel void k_chain(device uint* buf [[buffer(0)]], uint gid [[thread_position_in_grid]]) { + if (gid == 0) buf[0] = buf[0] + 1; +} + +// A small but real kernel: 2560 floats read, 2560 written (an elementwise op on one token). +kernel void k_small(device const float* x [[buffer(0)]], device float* y [[buffer(1)]], + uint gid [[thread_position_in_grid]]) { + if (gid < 2560) y[gid] = x[gid] * 0.5f + y[gid]; +} diff --git a/results/round39/model-ab/candidate-kernels/sample.metal b/results/round39/model-ab/candidate-kernels/sample.metal new file mode 100644 index 0000000..4531ca2 --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/sample.metal @@ -0,0 +1,61 @@ +// metal/kernels/sample.metal - greedy sampling over the LM head's logits: the largest logit, ties to the lowest +// token id (NaN never wins). Two stages: argmax_part (grid (parts, n_tok), 256 threads) writes one (value, id) +// pair per threadgroup; argmax_final (grid (n_tok), 256 threads) reduces them in order. + +struct ArgArgs { + uint n_vocab; + uint parts; +}; + +static inline void better(thread float& bv, thread int& bi, float v, int i) { + if (v > bv || (v == bv && i < bi)) { bv = v; bi = i; } +} + +static inline void tg_argmax(thread float& bv, thread int& bi, threadgroup float* sv, threadgroup int* si, ushort sg, + ushort lane) { +#pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) { + const float ov = simd_shuffle_xor(bv, m); + const int oi = simd_shuffle_xor(bi, m); + better(bv, bi, ov, oi); + } + if (lane == 0) { sv[sg] = bv; si[sg] = bi; } + threadgroup_barrier(mem_flags::mem_threadgroup); + bv = -INFINITY; + bi = 0x7fffffff; +#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) better(bv, bi, sv[k], si[k]); +} + +kernel void argmax_part(device const float* logits [[buffer(0)]], device float* pv [[buffer(1)]], + device int* pi [[buffer(2)]], constant ArgArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float sv[8]; + threadgroup int si[8]; + const uint t = tg.y, per = (a.n_vocab + a.parts - 1) / a.parts; + const uint b0 = tg.x * per, b1 = min(b0 + per, a.n_vocab); + device const float* l = logits + t * a.n_vocab; + float bv = -INFINITY; + int bi = 0x7fffffff; + for (uint i = b0 + tid; i < b1; i += 256) { + const float v = l[i]; + if (!isnan(v)) better(bv, bi, v, int(i)); + } + tg_argmax(bv, bi, sv, si, sg, lane); + if (tid == 0) { pv[t * a.parts + tg.x] = bv; pi[t * a.parts + tg.x] = bi; } +} + +kernel void argmax_final(device const float* pv [[buffer(0)]], device const int* pi [[buffer(1)]], + device int* out [[buffer(2)]], constant ArgArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float sv[8]; + threadgroup int si[8]; + const uint t = tg.x; + float bv = -INFINITY; + int bi = 0x7fffffff; + for (uint k = tid; k < a.parts; k += 256) better(bv, bi, pv[t * a.parts + k], pi[t * a.parts + k]); + tg_argmax(bv, bi, sv, si, sg, lane); + if (tid == 0) out[t] = bi; +} diff --git a/results/round39/model-ab/candidate-kernels/units.metal b/results/round39/model-ab/candidate-kernels/units.metal new file mode 100644 index 0000000..25d0761 --- /dev/null +++ b/results/round39/model-ab/candidate-kernels/units.metal @@ -0,0 +1,328 @@ +// metal/kernels/units.metal - the dot of one unit with the window's tokens (after quant.metal). +// +// X::load(xrow, u, xs, xsum) prepares unit u's x slice of one token (xrow = the token's row start), once per unit +// and shared by every weight row the lane reads: the UNIT x values in the order the decoder produces weights, and +// up to two partial sums (for the formats with an offset or a min). Where the codes are read with the mask trick +// (an AND keeps a field in place, so the code arrives multiplied by 2^k), the x values are pre-scaled by 2^-k here: +// exact in float. + +// the threadgroup-memory form (x staged once per threadgroup; the same values) +static inline void load16(threadgroup const float* p, thread float* xs) { + threadgroup const float4* v = (threadgroup const float4*)p; + for (int j = 0; j < 4; ++j) { + const float4 a = v[j]; + xs[4 * j] = a.x; + xs[4 * j + 1] = a.y; + xs[4 * j + 2] = a.z; + xs[4 * j + 3] = a.w; + } +} +static inline void load16(device const float* p, thread float* xs) { + device const float4* v = (device const float4*)p; + for (int j = 0; j < 4; ++j) { + const float4 a = v[j]; + xs[4 * j] = a.x; + xs[4 * j + 1] = a.y; + xs[4 * j + 2] = a.z; + xs[4 * j + 3] = a.w; + } +} + +// ---- prepared x (xprep): XP holds, per token row of n floats, the kinds +// XK_RAW x (every format without the mask trick) +// XK_Q2 x * 4^-(i % 8) (Q2_0: 2-bit fields in place in a 16-bit word) +// XK_P4 x * 2^-(8 (i % 4)) (Q4_K, Q5_K: byte fields in place in a 32-bit word) +// XK_P2 x * 2^-(8 (i % 2)) (Q6_K, Q3_K: byte fields in place in a 16-bit word) +// and S16 = the sum of every 16 consecutive raw x (the formats with an offset or a min). +#define XK_RAW 0 +#define XK_Q2 1 +#define XK_P4 2 +#define XK_P2 3 +template struct XKind { enum { K = XK_RAW }; }; +template <> struct XKind { enum { K = XK_Q2 }; }; +template <> struct XKind { enum { K = XK_P4 }; }; +template <> struct XKind { enum { K = XK_P4 }; }; +template <> struct XKind { enum { K = XK_P2 }; }; +template <> struct XKind { enum { K = XK_P2 }; }; + +// X::load(xp, s16, u, xs, xsum): unit u's prepared x (xp = the token's row of the format's kind, s16 = its sums) +template struct X { + template static inline void load(P xp, P, int u, thread float* xs, thread float* xsum) { + constexpr int UNIT = fmt::UNIT; + for (int k = 0; k < UNIT; k += 16) load16(xp + u * UNIT + k, xs + k); + xsum[0] = xsum[1] = 0.0f; + } +}; +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + load16(xp + 16 * u, xs); + xsum[0] = s16[u]; + xsum[1] = 0.0f; + } +}; +// Q4_K / Q5_K: unit (block b, chunk c, half m): run0 at 256b + 64c + 16m, run1 = run0 + 32 +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + const int p0 = 256 * (u >> 3) + 64 * ((u >> 1) & 3) + 16 * (u & 1); + load16(xp + p0, xs); + load16(xp + p0 + 32, xs + 16); + xsum[0] = s16[p0 >> 4]; + xsum[1] = s16[(p0 >> 4) + 2]; + } +}; +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + X::load(xp, s16, u, xs, xsum); + } +}; +// Q6_K: unit (block b, half n, m, p): run0 at 256b + 128n + 32p + 16m, run1 = run0 + 64 +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + const int p0 = 256 * (u >> 3) + 128 * ((u >> 2) & 1) + 32 * (u & 1) + 16 * ((u >> 1) & 1); + load16(xp + p0, xs); + load16(xp + p0 + 64, xs + 16); + xsum[0] = s16[p0 >> 4]; + xsum[1] = s16[(p0 >> 4) + 4]; + } +}; +// Q5_1: 32 contiguous raw x, the two per-16 sums formed from them (the GEMVs that read Q5_1 - the routed down - get +// h without prepared sums) +template <> struct X { + template static inline void load(P xp, P, int u, thread float* xs, thread float* xsum) { + load16(xp + 32 * u, xs); + load16(xp + 32 * u + 16, xs + 16); + float a = 0.0f, b = 0.0f; + for (int j = 0; j < 16; ++j) { a += xs[j]; b += xs[16 + j]; } + xsum[0] = a; + xsum[1] = b; + } +}; +// Q3_K: 32 contiguous, sums per 16 +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + load16(xp + 32 * u, xs); + load16(xp + 32 * u + 16, xs + 16); + xsum[0] = s16[2 * u]; + xsum[1] = s16[2 * u + 1]; + } +}; + +// xprep: one thread per 16 consecutive values of one token row; XP = [4 kinds][n_tok][n] then S16 [n_tok][n / 16]. +// raw and plane are 16-float aligned: the four planes are written as float4s (mtp.metal's mtp_prepare, bit for bit) +kernel void xprep(device const float* x [[buffer(0)]], device float* xp [[buffer(1)]], constant uint& n [[buffer(2)]], + constant uint& n_tok [[buffer(3)]], uint gid [[thread_position_in_grid]]) { + const uint per = n / 16; + if (gid >= per * n_tok) return; + const uint t = gid / per, i = gid % per; + device const float4* src = (device const float4*)(x + t * n + 16 * i); + const uint plane = n_tok * n; + device float* raw = xp + t * n + 16 * i; + device float4* r0 = (device float4*)raw; + device float4* r1 = (device float4*)(raw + plane); + device float4* r2 = (device float4*)(raw + 2 * plane); + device float4* r3 = (device float4*)(raw + 3 * plane); + const float4 q2lo = float4(1.0f, 0.25f, 0.0625f, 0.015625f), + q2hi = float4(0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f), + p4 = float4(1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f); + float s0 = 0, s1 = 0; +#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) { + const float4 v = src[j]; + s0 += v.x; s0 += v.z; + s1 += v.y; s1 += v.w; + r0[j] = v; + r1[j] = v * ((j & 1) ? q2hi : q2lo); + r2[j] = v * p4; + r3[j] = float4(v.x, v.y * 0.00390625f, v.z, v.w * 0.00390625f); + } + xp[4 * plane + t * per + i] = s0 + s1; +} + +// ---- decoders: Dec::run(row, u, q, sc) decodes unit u of a row into q[UNIT] and sc = (d1, d2, m1, m2), so that +// for any token: dot = d1 * sum_{j < UNIT/2} q[j] xs[j] + d2 * sum_{j >= UNIT/2} q[j] xs[j] - (m1 xsum0 + m2 xsum1) +// with xs, xsum from X::load. A unit's weights are decoded once per row and applied to every token. +template struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT, UPB = fmt::QK / UNIT; + const float s = deq(row + (u / UPB) * fmt::BB, u % UPB, q, tab); + sc = float4(s, s, 0.0f, 0.0f); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 2) * 18; + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + const ushort a = p[0], c = p[1]; + for (int j = 0; j < 8; ++j) { + q[j] = float(a & ushort(3u << (2 * j))); + q[8 + j] = float(c & ushort(3u << (2 * j))); + } + const float d = rdh(b); + sc = float4(d, d, d, 0.0f); + } +}; + +// Q2_0 with raw x: q = c - 1 (shift, and, convert, subtract), one scale +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 2) * 18; + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + const ushort a = p[0], c = p[1]; + for (int j = 0; j < 8; ++j) { + q[j] = float((a >> (2 * j)) & 3) - 1.0f; + q[8 + j] = float((c >> (2 * j)) & 3) - 1.0f; + } + const float d = rdh(b); + sc = float4(d, d, 0.0f, 0.0f); + } +}; + +static inline void scale_min_k4(int j, device const uchar* q, thread float& d, thread float& m) { + if (j < 4) { + d = float(q[j] & 63); + m = float(q[j + 4] & 63); + } else { + d = float((q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4)); + m = float((q[j + 4] >> 4) | ((q[j] >> 6) << 4)); + } +} +static inline void bytes4(uint w, thread float* q) { + q[0] = float(w & 0xFFu); + q[1] = float(w & 0xFF00u); + q[2] = float(w & 0xFF0000u); + q[3] = float(w & 0xFF000000u); +} +static inline void bytes2(ushort w, thread float* q) { + q[0] = float(w & 0xFF); + q[1] = float(w & 0xFF00); +} + +template <> struct Dec { // w = d q + m: sum = d (q . x) + m (sum x) + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint* tab) { + device const uchar* b = row + u * 24; + const float d = deq(b, 0, q, tab), m = rdh(b + 2); + sc = float4(d, d, -m, -m); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 144; + const int c = (u >> 1) & 3, m = u & 1; + const float d = rdh(b), dm = rdh(b + 2); + float s1, m1, s2, m2; + scale_min_k4(2 * c, b + 4, s1, m1); + scale_min_k4(2 * c + 1, b + 4, s2, m2); + device const uint* qp = (device const uint*)(b + 16 + 32 * c + 16 * m); + for (int w = 0; w < 4; ++w) { + const uint v = qp[w]; + bytes4(v & 0x0F0F0F0Fu, q + 4 * w); + bytes4((v >> 4) & 0x0F0F0F0Fu, q + 16 + 4 * w); + } + sc = float4(d * s1, d * s2, dm * m1, dm * m2); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 176; + const int c = (u >> 1) & 3, m = u & 1; + const float d = rdh(b), dm = rdh(b + 2); + float s1, m1, s2, m2; + scale_min_k4(2 * c, b + 4, s1, m1); + scale_min_k4(2 * c + 1, b + 4, s2, m2); + device const uint* qp = (device const uint*)(b + 48 + 32 * c + 16 * m); + device const uint* hp = (device const uint*)(b + 16 + 16 * m); + for (int w = 0; w < 4; ++w) { + const uint v = qp[w], h = hp[w]; + bytes4((v & 0x0F0F0F0Fu) | (((h >> (2 * c)) & 0x01010101u) << 4), q + 4 * w); + bytes4(((v >> 4) & 0x0F0F0F0Fu) | (((h >> (2 * c + 1)) & 0x01010101u) << 4), q + 16 + 4 * w); + } + sc = float4(d * s1, d * s2, dm * m1, dm * m2); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 210; + const int n = (u >> 2) & 1, m = (u >> 1) & 1, p = u & 1; + device const ushort* ql = (device const ushort*)(b + 64 * n + 32 * p + 16 * m); + device const ushort* qh = (device const ushort*)(b + 128 + 32 * n + 16 * m); + device const char* scs = (device const char*)(b + 192); + const float d = rdh(b + 208); + for (int w = 0; w < 8; ++w) { + const ushort l = ql[w], h = qh[w]; + bytes2((l & 0x0F0F) | (((h >> (2 * p)) & 0x0303) << 4), q + 2 * w); + bytes2(((l >> 4) & 0x0F0F) | (((h >> (2 * p + 4)) & 0x0303) << 4), q + 16 + 2 * w); + } + const float d1 = d * float(scs[8 * n + m + 2 * p]), d2 = d * float(scs[8 * n + m + 2 * p + 4]); + sc = float4(d1, d2, 32.0f * d1, 32.0f * d2); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 110; + const int n = (u >> 2) & 1, j = u & 3; + // qs and hmask as packed (2-byte aligned) vectors: blocks are 110 bytes, so odd blocks are only 2-aligned + device const packed_ushort4* qv = (device const packed_ushort4*)(b + 32 + 32 * n); + device const packed_ushort4* hv = (device const packed_ushort4*)b; + ushort qq[16], hm[16]; + for (int v = 0; v < 4; ++v) { + const ushort4 a4 = ushort4(qv[v]), h4 = ushort4(hv[v]); + qq[4 * v] = a4.x; qq[4 * v + 1] = a4.y; qq[4 * v + 2] = a4.z; qq[4 * v + 3] = a4.w; + hm[4 * v] = h4.x; hm[4 * v + 1] = h4.y; hm[4 * v + 2] = h4.z; hm[4 * v + 3] = h4.w; + } + device const uchar* s = b + 96; + const int i1 = 8 * n + 2 * j, i2 = i1 + 1; + const int l1 = (i1 < 8 ? s[i1] & 0xF : s[i1 - 8] >> 4) | (((s[8 + (i1 & 3)] >> (2 * (i1 >> 2))) & 3) << 4); + const int l2 = (i2 < 8 ? s[i2] & 0xF : s[i2 - 8] >> 4) | (((s[8 + (i2 & 3)] >> (2 * (i2 >> 2))) & 3) << 4); + const float d = rdh(b + 108); + const int hs = 4 * n + j; + for (int w = 0; w < 8; ++w) { + bytes2(((qq[w] >> (2 * j)) & 0x0303) | (((hm[w] >> hs) & 0x0101) << 2), q + 2 * w); + bytes2(((qq[w + 8] >> (2 * j)) & 0x0303) | (((hm[w + 8] >> hs) & 0x0101) << 2), q + 16 + 2 * w); + } + const float d1 = d * float(l1 - 32), d2 = d * float(l2 - 32); + sc = float4(d1, d2, 4.0f * d1, 4.0f * d2); + } +}; + +template <> struct Dec { // 16 floats a unit (the Unsloth files' router / alpha / beta / inject rows) + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const float4* p = (device const float4*)(row + 64 * u); + for (int k = 0; k < 4; ++k) { + const float4 v = p[k]; + q[4 * k] = v.x; q[4 * k + 1] = v.y; q[4 * k + 2] = v.z; q[4 * k + 3] = v.w; + } + sc = float4(1.0f, 1.0f, 0.0f, 0.0f); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uint4* p = (device const uint4*)(row + 32 * u); + const uint4 a = p[0], c = p[1]; + const uint v[8] = {a.x, a.y, a.z, a.w, c.x, c.y, c.z, c.w}; + for (int k = 0; k < 8; ++k) { + q[2 * k] = as_type(v[k] << 16); + q[2 * k + 1] = as_type(v[k] & 0xFFFF0000u); + } + sc = float4(1.0f, 1.0f, 0.0f, 0.0f); + } +}; + +// acc += unit dot of decoded q / sc with one token's prepared slice +template +static inline float unit_combine(thread const float* q, float4 sc, thread const float* xs, thread const float* xsum) { + constexpr int H = UNIT / 2; + float s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int j = 0; j < H; j += 2) { + s0 = fma(q[j], xs[j], s0); + s1 = fma(q[j + 1], xs[j + 1], s1); + s2 = fma(q[H + j], xs[H + j], s2); + s3 = fma(q[H + j + 1], xs[H + j + 1], s3); + } + return sc.x * (s0 + s1) + sc.y * (s2 + s3) - (sc.z * xsum[0] + sc.w * xsum[1]); +} diff --git a/results/round39/model-ab/ngram-baseline.cpp b/results/round39/model-ab/ngram-baseline.cpp new file mode 100644 index 0000000..91de370 --- /dev/null +++ b/results/round39/model-ab/ngram-baseline.cpp @@ -0,0 +1,647 @@ +// src/kernels/ngram.cpp - P2.S4: the PLE n-gram hash and the IQ4_NL table read. +// +// See include/strata/kernels/ngram.hpp for the semantics, the rival readings, and the note on MADV_RANDOM. +#include "strata/kernels/ngram.hpp" +#include "strata/artifact/gguf_reader.hpp" +#include "strata/artifact/dequant.hpp" +#include "strata/kernels/f16_bits.hpp" +#include "strata/ngram/ple_reader.hpp" +#include "strata/platform/memory.hpp" + +#include +#include +#include +#include +#include +#if !defined(_WIN32) +#include +#endif +#include +#include +#include +#include +#include + +#if defined(_WIN32) +// `PrefetchVirtualMemory` (memoryapi.h, Windows 8+) is the whole point of the change in `gather` below. +#define WIN32_LEAN_AND_MEAN +#define NOMINMAX +#include +#endif + +namespace strata::kernels { + +namespace { +/// The A/B arm. Host-token-path only, so it needs no atomics; see the note on `ple_prefetch_enable`. +bool g_ple_prefetch = true; +} // namespace + +void ple_prefetch_enable(bool on) { g_ple_prefetch = on; } +bool ple_prefetch_enabled() { return g_ple_prefetch; } + +PleConsts ple_artifact_consts() { + // docs/gguf-dump-shard1.txt, verbatim. Written once here rather than derived, because `head_offsets` is + // ALSO in the metadata as its own array - and deriving one from the other would hide a mismatch between + // them instead of surfacing it. The parity test checks the two agree. + PleConsts c{}; + c.mult[0] = 23703573157769ull; + c.mult[1] = 20109073645365ull; + c.mult[2] = 8052911324071ull; + const uint64_t vocab[PLE_N_HEADS] = { + 20000003, 20000023, 20000033, 20000047, 20000059, 20000063, 20000069, 20000077, + 20000081, 20000093, 20000107, 20000147, 20000153, 20000159, 20000161, 20000171}; + const uint64_t offset[PLE_N_HEADS] = { + 0, 20000003, 40000026, 60000059, 80000106, 100000165, 120000228, 140000297, + 160000374, 180000455, 200000548, 220000655, 240000802, 260000955, 280001114, 300001275}; + for (int i = 0; i < PLE_N_HEADS; ++i) { + c.vocab[i] = vocab[i]; + c.offset[i] = offset[i]; + } + return c; +} + +uint64_t ngram_mixed(const int64_t* ctx, const uint64_t* mult, int n) { + // The first term is an ASSIGNMENT and the rest are XORed into it, which is how the source writes it + // (`uint64_t mixed = ctx[0]*m[0]; for j=1.. mixed ^= ctx[j]*m[j];`). Every product wraps mod 2^64, + // which is what the `(uint64_t)` casts in the source make explicit. + uint64_t mixed = (uint64_t) ctx[0] * mult[0]; + for (int j = 1; j < n; ++j) mixed ^= (uint64_t) ctx[j] * mult[j]; + return mixed; +} + +void ngram_rows(const int32_t* tokens, const int32_t* prev, int n_tokens, const PleConsts& c, uint32_t* out) { + const int n_prev = NGRAM_SIZE - 1; + for (int i = 0; i < n_tokens; ++i) { + int64_t ctx[NGRAM_SIZE]; + ctx[0] = tokens[i]; + bool cut = false; + for (int s = 1; s < NGRAM_SIZE; ++s) { + // `prev` is OLDEST FIRST, so predecessor `s` positions back is entry (n_prev - s): s=1 reads the + // NEWEST. Reading index (s-1) instead walks the window backwards, which still produces indices + // in range and so cannot be caught by a range check - only by an oracle. + const int32_t t = cut ? TOKEN_NULL : prev[i * n_prev + (n_prev - s)]; + // The cut is evaluated BEFORE the value is stored, so the position whose predecessor was EOS is + // itself EOS. Storing first and then cutting would leave position s holding the real token while + // position s+1 became EOS - one token of history too much. + cut = cut || t < 0 || t == PLE_EOS_TOKEN_ID; + ctx[s] = cut ? PLE_EOS_TOKEN_ID : t; + } + for (int n = 2; n <= NGRAM_SIZE; ++n) { + const uint64_t mixed = ngram_mixed(ctx, c.mult, n); + const int base = (n - 2) * HEADS_PER_NGRAM; + for (int g = 0; g < HEADS_PER_NGRAM; ++g) { + const int h = base + g; + out[i * PLE_N_HEADS + h] = (uint32_t) (mixed % c.vocab[h] + c.offset[h]); + } + } + } +} + +namespace { +const int8_t kIq4Nl[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113}; +} + +int iq4nl_code(int code) { return kIq4Nl[code & 15]; } + +void iq4nl_dequant_row(const uint8_t* row, float* out160) { + for (int b = 0; b < PLE_HEAD_DIM / 32; ++b) { + const uint8_t* blk = row + (size_t) b * 18; + uint16_t dbits; + std::memcpy(&dbits, blk, 2); + const float d = f32_from_f16(dbits); + const uint8_t* qs = blk + 2; + // SPLIT HALVES: qs[j] holds elements j and j+16, not 2j and 2j+1. + for (int j = 0; j < 16; ++j) { + out160[b * 32 + j] = d * (float) kIq4Nl[qs[j] & 0x0F]; + out160[b * 32 + j + 16] = d * (float) kIq4Nl[qs[j] >> 4]; + } + } +} + +namespace { +struct Fp8Table { + float v[256]; + Fp8Table() { + for (int x = 0; x < 256; ++x) { + const int s = x >> 7, e = (x >> 3) & 15, m = x & 7; + float f; + if (e == 15 && m == 7) f = 0.0f; // NaN in e4m3fn; never in a real table, read as 0 + else if (e == 0) f = std::ldexp((float) m / 8.0f, -6); + else f = std::ldexp(1.0f + (float) m / 8.0f, e - 7); + v[x] = s ? -f : f; + } + } +}; +const Fp8Table kFp8; +} // namespace + +void fp8_e4m3_dequant_row(const uint8_t* row, float scale, float* out160) { + for (int j = 0; j < PLE_HEAD_DIM; ++j) out160[j] = kFp8.v[row[j]] * scale; +} + +void bf16_dequant_row(const uint8_t* row, float* out160) { + // bfloat16 IS the top 16 bits of a float32, so widening is exact for every value, normal or not - no rounding, + // no clamp, no scale. Bytes are assembled little-endian explicitly rather than reinterpreted, so this is the + // same on a big-endian host. + for (int j = 0; j < PLE_HEAD_DIM; ++j) { + const uint32_t bits = (uint32_t) row[2 * j] | ((uint32_t) row[2 * j + 1] << 8); + const uint32_t wide = bits << 16; + float v; + std::memcpy(&v, &wide, sizeof v); + out160[j] = v; + } +} + +// --------------------------------------------------------------------------------------------------- +// THE FORMAT TABLE (see PleFormatInfo). One entry per table type; the block formats decode 5 blocks of 32. +namespace { +template +void dequant_blocks(const uint8_t* row, float /*scale*/, float* out160) { + for (int b = 0; b < PLE_HEAD_DIM / 32; ++b) Block(row + (size_t) b * BlockBytes, out160 + b * 32); +} +void dequant_iq4_nl(const uint8_t* row, float, float* out160) { iq4nl_dequant_row(row, out160); } +void dequant_fp8(const uint8_t* row, float scale, float* out160) { fp8_e4m3_dequant_row(row, scale, out160); } +void dequant_bf16(const uint8_t* row, float, float* out160) { bf16_dequant_row(row, out160); } + +const PleFormatInfo kPleFormats[] = { + {PleFormat::IQ4_NL, "IQ4_NL", "IQ4_NL", PLE_ROW_BYTES, false, dequant_iq4_nl}, + {PleFormat::Q5_0, "Q5_0", "Q5_0", (PLE_HEAD_DIM / 32) * 22, false, dequant_blocks}, + {PleFormat::F8_E4M3, "F8_E4M3", "I8", PLE_ROW_BYTES_FP8, true, dequant_fp8}, + // Ordinary GGUFs ship the table in other types: Unsloth's UD-Q6_K_XL and Swift-1.5 Q4_K_L as Q8_0, a Q5_K_M finetune as Q5_1 + {PleFormat::Q5_1, "Q5_1", "Q5_1", PLE_ROW_BYTES_Q5_1, false, dequant_blocks}, + {PleFormat::Q8_0, "Q8_0", "Q8_0", PLE_ROW_BYTES_Q8_0, false, dequant_blocks}, + // Q4_0 (plain llama-quantize Q4_0 files, #599): 90-byte rows like IQ4_NL (18-byte blocks), a linear 4-bit grid + {PleFormat::Q4_0, "Q4_0", "Q4_0", PLE_ROW_BYTES, false, dequant_blocks}, + // BF16 (#586): the checkpoint's own table at full precision, 320-byte rows; self-describing, so no scale and no + // metadata to trust (tools/ple_fp8_pack.py writes the FP8 form of the same table) + {PleFormat::BF16, "BF16", "BF16", PLE_ROW_BYTES_BF16, false, dequant_bf16}, +}; +constexpr int kPleFormatCount = (int) (sizeof kPleFormats / sizeof kPleFormats[0]); +} // namespace + +const PleFormatInfo* ple_formats() { return kPleFormats; } +int ple_format_count() { return kPleFormatCount; } +const PleFormatInfo& ple_format_info(PleFormat f) { return kPleFormats[(int) f]; } +const PleFormatInfo* ple_format_for_type(const char* gguf_type_name) { + for (const PleFormatInfo& f : kPleFormats) + if (std::strcmp(f.gguf_type, gguf_type_name) == 0) return &f; + return nullptr; +} +std::string ple_format_list() { + std::string out; + for (int i = 0; i < kPleFormatCount; ++i) { + const PleFormatInfo& f = kPleFormats[i]; + if (i > 0) out += i + 1 == kPleFormatCount ? " or " : ", "; + out += std::strcmp(f.gguf_type, "I8") == 0 ? "FP8 (I8)" : f.name; + } + return out; +} + +static_assert(PLE_ROW_BYTES_MAX >= PLE_ROW_BYTES && PLE_ROW_BYTES_MAX >= PLE_ROW_BYTES_FP8, "row buffers too small"); + +struct PleTable::Impl { + GgufFile* file = nullptr; + const uint8_t* data = nullptr; + uint64_t n_rows = 0; + const PleFormatInfo* fmt = &ple_format_info(PleFormat::IQ4_NL); // the table's format (#296: Q5_0 is one of them) + mutable uint64_t bytes_read = 0; + // Direct mode (plan v0.3 P2): the mapping above is released after the header parse and every row comes + // from an unbuffered SSD read into `raw`. + PleIo mode = PleIo::Mmap; + strata::ngram::PleReader reader; + strata::ngram::PleReader::Ticket ticket; + bool pending = false; + bool locked = false; + uint32_t rows[PLE_N_HEADS] = {}; + uint8_t raw[PLE_N_HEADS * PLE_ROW_BYTES_MAX] = {}; + static constexpr size_t kMaxPrefetch = 16; + size_t n_prefetch = 0; + strata::ngram::PleReader::Ticket prefetch_tickets[kMaxPrefetch] = {}; + uint32_t prefetch_keys[kMaxPrefetch][PLE_N_HEADS] = {}; + uint8_t prefetch_raw[kMaxPrefetch][PLE_N_HEADS * PLE_ROW_BYTES_MAX] = {}; + float scale = 1.0f; // the FP8 table's one scale + uint32_t rb = PLE_ROW_BYTES; // bytes per row + void decode(const uint8_t* row, float* out160) const { fmt->dequant(row, scale, out160); } +}; + +PleTable::PleTable() : impl_(new Impl) {} +PleTable::~PleTable() { close(); delete impl_; } + +bool PleTable::open(const std::string& gguf_path, std::string& err) { + return open(gguf_path, err, PleIoOptions{}); +} + +bool PleTable::open(const std::string& gguf_path, std::string& err, const PleIoOptions& io) { + close(); + try { + impl_->file = new GgufFile(gguf_path); + } catch (const std::exception& e) { + err = e.what(); + return false; + } + const TensorInfo* t = impl_->file->find("per_layer_token_embd.weight"); + if (t == nullptr) { + err = "per_layer_token_embd.weight is not in " + gguf_path; + close(); + return false; + } + // [160, 320001536]: ne0 = 160 is the FAST axis, so the ROW index is shape[1] and a row is contiguous. + if (t->shape.size() != 2 || t->shape[0] != (uint64_t) PLE_HEAD_DIM) { + err = "per_layer_token_embd.weight has an unexpected shape"; + close(); + return false; + } + // The table's type is looked up in the format list (ple_formats()): IQ4_NL (ISTA-DASLab's shard 2, the original's + // own GGUF), Q5_0 (#296: OrcaRouter's GGUF), or the FP8 table as shipped: I8 bytes marked strata.ple.format = + // f8_e4m3 with strata.ple.scale (tools/ple_fp8_pack.py) + const PleFormatInfo* fmt = ple_format_for_type(t->type_name()); + if (fmt == nullptr) { + err = std::string("per_layer_token_embd.weight is ") + t->type_name() + ", not " + ple_format_list(); + close(); + return false; + } + if (fmt->needs_scale) { + const MetaValue* f = impl_->file->get("strata.ple.format"); + const MetaValue* s = impl_->file->get("strata.ple.scale"); + if (f == nullptr || f->s != "f8_e4m3" || s == nullptr || !(s->num() > 0.0)) { + err = "per_layer_token_embd.weight is I8 without strata.ple.format = f8_e4m3 and a positive strata.ple.scale"; + close(); + return false; + } + impl_->scale = (float) s->num(); + } + impl_->fmt = fmt; + impl_->rb = fmt->row_bytes; + // PleReader's row_bytes has been a runtime parameter since the FP8 table (160 B rows) needed it; Q5_0's + // 110 B rows go through the exact same generic path (ple_reader_test --selftest covers both 90 and 110 B + // rows: straddling, caching, in-flight tickets, keep-alive). This refusal was stale. + impl_->n_rows = t->shape[1]; + + // THE CHECK THAT MAKES THE OFFSET FALSIFIABLE. The manifest's `shard2_tensor.offset` is 0, but that is + // the offset within the GGUF's DATA SECTION: the file's first 192 bytes are a header, and reading at 0 + // would decode the header plus 192 bytes of shifted rows - still plausible IQ4_NL, and wrong for every + // row. `GgufFile` parses the header, so `tensor_data` is already correct; this asserts the tensor + // exactly fills the file from there, which is what makes the whole arrangement checkable rather than + // assumed. A wrong data offset would leave a different remainder. + // A shard may hold other tensors too (Swift 1.5's shard 1 holds layers 0-12 and the table): the table must + // then fit inside the file at its own offset; alone in its shard (the original's shard 2) it fills it exactly. + // The size arithmetic is checked before it is used (#865): a header that claims enough rows to wrap the 64-bit + // product, or a data section that starts past the end of the file, must be refused, not compared after wrapping. + if (impl_->n_rows > std::numeric_limits::max() / impl_->rb || + impl_->file->data_start() > impl_->file->file_size()) { + err = "PLE table size overflow or invalid data offset in " + gguf_path; + close(); + return false; + } + const uint64_t need = impl_->n_rows * (uint64_t) impl_->rb; + const uint64_t have = impl_->file->file_size() - impl_->file->data_start(); + const bool alone = impl_->file->tensors().size() == 1; + if (t->offset > have || need > have - t->offset || (alone && (t->offset != 0 || need != have))) { + char buf[256]; + std::snprintf(buf, sizeof buf, + "PLE table size mismatch: %llu rows x %d B = %llu at offset %llu, but the file holds %llu from " + "data_start %llu", + (unsigned long long) impl_->n_rows, (int) impl_->rb, (unsigned long long) need, + (unsigned long long) t->offset, (unsigned long long) have, + (unsigned long long) impl_->file->data_start()); + err = buf; + if (alone && have > need) // #657: extra bytes after the table: a damaged or wrong file, not a format question + err += "; the file is longer than its table - delete it and its .done mark and run setup again"; + close(); + return false; + } + impl_->data = impl_->file->tensor_data(*t); // only a table that fits the file is mapped + if (io.mode == PleIo::Direct) { + // The parse above is the validated source of the offset; the mapping itself is not kept, so no page of + // the table can enter this process's working set or the file cache through it. + const uint64_t table_offset = impl_->file->data_start() + t->offset; + const uint64_t n_rows = impl_->n_rows; + delete impl_->file; + impl_->file = nullptr; + impl_->data = nullptr; + if (!impl_->reader.open(gguf_path, table_offset, n_rows, io.max_inflight, io.cache_rows, err, io.io_thread, + impl_->rb)) { + close(); + return false; + } + impl_->reader.set_keepalive(io.keepalive_ms, io.keepalive_window_s); + impl_->n_rows = n_rows; + } + if (io.mode == PleIo::Mmap && io.lock && impl_->data != nullptr) { +#if defined(_WIN32) + // Windows has no mlock: the platform helper raises the process's minimum working set and + // VirtualLocks the mapped table instead (ordinary accounts hold the privilege). A failure + // falls back to touching the pages, i.e. plain mmap behaviour, with a warning - same as POSIX. + // A table that is more than half of the RAM leaves the rest of the process little (the locked pages cannot be + // paged out): say so, and go on (recommend, never force). + MEMORYSTATUSEX ms; + ms.dwLength = sizeof(ms); + if (GlobalMemoryStatusEx(&ms) && need > ms.ullTotalPhys / 2) + std::fprintf(stderr, "strata: warning: --ple-io ram locks %.1f GiB, more than half of this PC's %.1f GiB of RAM\n", + (double) need / (1ull << 30), (double) ms.ullTotalPhys / (1ull << 30)); + const strata::platform::LockResult lr = + strata::platform::lock_resident((void*) impl_->data, need); + if (lr.ok && lr.locked_bytes >= need) { + impl_->locked = true; + } else { + // a partial lock is not a lock: the unlocked rest would fault on the token path, so give it back + if (lr.locked_bytes > 0) strata::platform::unlock_resident((void*) impl_->data, lr.locked_bytes); + std::fprintf(stderr, "strata: PLE table lock failed (%s): touching its pages instead\n", + lr.note.c_str()); + volatile uint8_t sink = 0; + for (uint64_t off = 0; off < need; off += 4096) sink = sink + impl_->data[off]; + (void) sink; + } +#else + const uint64_t page = 4096; + const uintptr_t a0 = (uintptr_t) impl_->data & ~(uintptr_t) (page - 1); + const uintptr_t a1 = (uintptr_t) impl_->data + (uintptr_t) need; + madvise((void*) a0, a1 - a0, MADV_WILLNEED); + // Fault the table in with several threads first: mlock (and a single toucher) brings it in one page at a time + // from one thread, ~0.5 GB/s from a cold file - 106 s for a 54 GB Q8_0 table. The pages are then resident + // and mlock only pins them. + { + constexpr uintptr_t kPiece = 64ull << 20; + const uintptr_t pieces = (a1 - a0 + kPiece - 1) / kPiece; + const unsigned threads = std::max(1u, std::min(16u, std::thread::hardware_concurrency())); + std::atomic next{0}; + auto touch = [&] { + volatile uint8_t sink = 0; + for (uintptr_t i; (i = next.fetch_add(1)) < pieces;) + for (uintptr_t p = a0 + i * kPiece, e = std::min(a1, p + kPiece); p < e; p += page) + sink = sink + *(const volatile uint8_t*) p; + (void) sink; + }; + std::vector pool; + for (unsigned t = 0; t < threads; ++t) pool.emplace_back(touch); + for (auto& t : pool) t.join(); + } + if (mlock((const void*) a0, a1 - a0) == 0) { + impl_->locked = true; + } else { + std::fprintf(stderr, "strata: PLE table mlock failed (%s; raise `ulimit -l`): its pages stay faulted in " + "but may be reclaimed\n", std::strerror(errno)); + } +#endif + } + impl_->mode = io.mode; + return true; +} + +void PleTable::wait_prefetches() { + if (impl_->n_prefetch == 0) return; + std::string dummy; + for (size_t i = 0; i < impl_->n_prefetch; ++i) (void) impl_->reader.collect(impl_->prefetch_tickets[i], dummy); + impl_->n_prefetch = 0; +} + +void PleTable::prefetch_rows(const uint32_t* rows16) { + if (impl_->mode == PleIo::Direct && impl_->reader.is_open()) { + for (size_t i = 0; i < impl_->n_prefetch; ++i) + if (std::memcmp(impl_->prefetch_keys[i], rows16, sizeof(uint32_t) * PLE_N_HEADS) == 0) return; + if (impl_->n_prefetch >= Impl::kMaxPrefetch) wait_prefetches(); + const size_t slot = impl_->n_prefetch++; + std::memcpy(impl_->prefetch_keys[slot], rows16, sizeof(uint32_t) * PLE_N_HEADS); + impl_->prefetch_tickets[slot] = impl_->reader.issue(impl_->prefetch_keys[slot], PLE_N_HEADS, impl_->prefetch_raw[slot]); + return; + } +#if defined(_WIN32) + if (impl_->data != nullptr && g_ple_prefetch) { + WIN32_MEMORY_RANGE_ENTRY ranges[PLE_N_HEADS]; + ULONG_PTR n = 0; + for (int h = 0; h < PLE_N_HEADS; ++h) { + if (rows16[h] >= impl_->n_rows) continue; + ranges[n].VirtualAddress = (PVOID) (impl_->data + (size_t) rows16[h] * impl_->rb); + ranges[n].NumberOfBytes = impl_->rb; + ++n; + } + if (n > 0) (void) PrefetchVirtualMemory(GetCurrentProcess(), n, ranges, 0); + } +#endif +} + +void PleTable::close() { + wait_prefetches(); + impl_->reader.close(); + impl_->pending = false; + impl_->locked = false; // the unmap below releases the lock + impl_->mode = PleIo::Mmap; + delete impl_->file; + impl_->file = nullptr; + impl_->data = nullptr; + impl_->n_rows = 0; + impl_->rb = PLE_ROW_BYTES; + impl_->fmt = &ple_format_info(PleFormat::IQ4_NL); + impl_->scale = 1.0f; + impl_->bytes_read = 0; +} + +bool PleTable::is_open() const { return impl_->data != nullptr || impl_->reader.is_open(); } +bool PleTable::locked() const { return impl_->locked; } +const char* PleTable::format() const { return impl_->fmt->name; } +PleIo PleTable::mode() const { return impl_->mode; } +uint64_t PleTable::rows() const { return impl_->n_rows; } +uint64_t PleTable::bytes_read() const { return impl_->bytes_read; } + +void PleTable::read_row(uint32_t row, float* out160) const { + if (impl_->mode == PleIo::Direct && impl_->reader.is_open()) { + uint8_t raw[PLE_ROW_BYTES_MAX]; + std::string err; + const auto t = impl_->reader.issue(&row, 1, raw); + if (!impl_->reader.collect(t, err)) { + std::memset(out160, 0, (size_t) PLE_HEAD_DIM * sizeof(float)); + return; + } + impl_->decode(raw, out160); + impl_->bytes_read += impl_->rb; + return; + } + if (impl_->data == nullptr || row >= impl_->n_rows) { + std::memset(out160, 0, (size_t) PLE_HEAD_DIM * sizeof(float)); + return; + } + impl_->decode(impl_->data + (size_t) row * impl_->rb, out160); + impl_->bytes_read += impl_->rb; +} + +bool PleTable::issue(const uint32_t* rows16) { + wait_prefetches(); + std::memcpy(impl_->rows, rows16, sizeof impl_->rows); + if (impl_->mode == PleIo::Direct) { + if (impl_->pending) return false; // one token in flight per table + impl_->ticket = impl_->reader.issue(impl_->rows, PLE_N_HEADS, impl_->raw); + impl_->pending = true; + return true; + } +#if defined(_WIN32) + if (impl_->data != nullptr && g_ple_prefetch) { + WIN32_MEMORY_RANGE_ENTRY ranges[PLE_N_HEADS]; + ULONG_PTR n = 0; + for (int h = 0; h < PLE_N_HEADS; ++h) { + if (rows16[h] >= impl_->n_rows) continue; + ranges[n].VirtualAddress = (PVOID) (impl_->data + (size_t) rows16[h] * impl_->rb); + ranges[n].NumberOfBytes = impl_->rb; + ++n; + } + if (n > 0) (void) PrefetchVirtualMemory(GetCurrentProcess(), n, ranges, 0); + } +#endif + impl_->pending = true; + return true; +} + +bool PleTable::collect(float* out2560, std::string& err) { + if (!impl_->pending) { err = "PleTable::collect without issue"; return false; } + impl_->pending = false; + if (impl_->mode == PleIo::Direct) { + if (!impl_->reader.collect(impl_->ticket, err)) return false; + for (int h = 0; h < PLE_N_HEADS; ++h) + impl_->decode(impl_->raw + (size_t) h * impl_->rb, out2560 + (size_t) h * PLE_HEAD_DIM); + impl_->bytes_read += (uint64_t) PLE_N_HEADS * impl_->rb; + return true; + } + for (int h = 0; h < PLE_N_HEADS; ++h) read_row(impl_->rows[h], out2560 + (size_t) h * PLE_HEAD_DIM); + return true; +} + +bool PleTable::gather_batch(const uint32_t* rows, size_t n_tokens, float* out, std::string& err) { + if (impl_->pending) { err = "PleTable::gather_batch while a token is in flight"; return false; } + const size_t n = n_tokens * (size_t) PLE_N_HEADS; + if (impl_->mode == PleIo::Direct) { + // Fast path: if all requested tokens match our in-flight/completed prefetch slots, collect them directly + // without issuing a second PleReader ticket or copying through RowCache. + if (impl_->n_prefetch >= n_tokens && n_tokens <= Impl::kMaxPrefetch) { + bool exact_match = true; + for (size_t t = 0; t < n_tokens; ++t) { + if (std::memcmp(impl_->prefetch_keys[t], rows + t * PLE_N_HEADS, sizeof(uint32_t) * PLE_N_HEADS) != 0) { + exact_match = false; + break; + } + } + if (exact_match) { + for (size_t i = 0; i < impl_->n_prefetch; ++i) { + if (!impl_->reader.collect(impl_->prefetch_tickets[i], err)) { + impl_->n_prefetch = 0; + return false; + } + } + impl_->n_prefetch = 0; + for (size_t t = 0; t < n_tokens; ++t) { + for (int h = 0; h < PLE_N_HEADS; ++h) { + impl_->decode(impl_->prefetch_raw[t] + (size_t) h * impl_->rb, + out + (t * PLE_N_HEADS + (size_t) h) * PLE_HEAD_DIM); + } + } + impl_->bytes_read += (uint64_t) n * impl_->rb; + return true; + } + } + wait_prefetches(); + std::vector raw(n * impl_->rb); + const auto ticket = impl_->reader.issue(rows, n, raw.data()); + if (!impl_->reader.collect(ticket, err)) return false; + for (size_t i = 0; i < n; ++i) impl_->decode(raw.data() + i * impl_->rb, out + i * PLE_HEAD_DIM); + impl_->bytes_read += (uint64_t) n * impl_->rb; + return true; + } + wait_prefetches(); + // #1984 #1971: a mapped table that is not in the file cache served this loop as serial demand faults, one random + // 4 KiB page at a time (16 rows a token, 16 x the chunk's tokens for a prompt), and the first chunk of a prompt has + // only layer 0 to hide it behind: a cold 27 GB table at ~0.2-1 ms a fault is minutes with no progress mark, which + // is what the stall watchdog stops. Ask the OS for the next window's pages in one go (PrefetchVirtualMemory / + // MADV_WILLNEED, a hint that cannot change a value: an unfetched page is faulted by the read) while this window is read. + constexpr size_t kWin = 4096; // rows a window + auto hint = [&](size_t a, size_t b) { + if (impl_->data == nullptr || !g_ple_prefetch) return; +#if defined(_WIN32) + WIN32_MEMORY_RANGE_ENTRY ranges[kWin]; + ULONG_PTR k = 0; + for (size_t i = a; i < b; ++i) { + if (rows[i] >= impl_->n_rows) continue; + ranges[k].VirtualAddress = (PVOID) (impl_->data + (size_t) rows[i] * impl_->rb); + ranges[k].NumberOfBytes = impl_->rb; + ++k; + } + if (k > 0) (void) PrefetchVirtualMemory(GetCurrentProcess(), k, ranges, 0); +#else + for (size_t i = a; i < b; ++i) { + if (rows[i] >= impl_->n_rows) continue; + const uintptr_t r0 = (uintptr_t) (impl_->data + (size_t) rows[i] * impl_->rb); + const uintptr_t p0 = r0 & ~(uintptr_t) 4095; + (void) madvise((void*) p0, (size_t) (r0 + impl_->rb - p0), MADV_WILLNEED); + } +#endif + }; + hint(0, std::min(n, kWin)); + for (size_t w = 0; w < n; w += kWin) { + if (w + kWin < n) hint(w + kWin, std::min(n, w + 2 * kWin)); + const size_t e = std::min(n, w + kWin); + for (size_t i = w; i < e; ++i) read_row(rows[i], out + i * PLE_HEAD_DIM); + } + return true; +} + +void PleTable::set_injected_delay_us(double us) { impl_->reader.set_injected_delay_us(us); } + +std::string PleTable::io_report() const { + if (impl_->mode != PleIo::Direct || !impl_->reader.is_open()) return {}; + const strata::ngram::ReaderStats s = impl_->reader.snapshot(); + char buf[400]; + int n = std::snprintf(buf, sizeof buf, + "ple io: %llu rows, %.1f%% row-cache hits, %llu SSD reads (%.1f MB), read p50 %.0f us p99 %.0f us, " + "blocked %.3f ms total (submit %.3f ms), cache %llu/%llu rows", + (unsigned long long) s.requests, s.requests ? 100.0 * (double) s.cache_hits / (double) s.requests : 0.0, + (unsigned long long) s.reads, (double) s.bytes / 1e6, s.percentile(0.5), s.percentile(0.99), + s.wait_us / 1000.0, s.submit_us / 1000.0, (unsigned long long) impl_->reader.cache_size(), + (unsigned long long) impl_->reader.cache_capacity()); + if (s.keepalive_reads > 0 && n > 0 && n < (int) sizeof buf) + std::snprintf(buf + n, sizeof buf - (size_t) n, ", SSD kept awake by %llu reads (slowest %.1f ms)", + (unsigned long long) s.keepalive_reads, s.keepalive_us_max / 1000.0); + return buf; +} + +void PleTable::gather(const uint32_t* rows16, float* out2560) const { + if (impl_->mode == PleIo::Direct) { + // `gather` stays const for its existing callers; the reader's state is the table's I/O state. + PleTable* self = const_cast(this); + std::string err; + if (!self->issue(rows16) || !self->collect(out2560, err)) { + std::fprintf(stderr, "PleTable::gather: %s\n", err.empty() ? "a token is already in flight" : err.c_str()); + std::memset(out2560, 0, (size_t) NG_N_EMBD * sizeof(float)); + } + return; + } + // ================================ SIXTEEN SERIAL PAGE FAULTS, MEASURED ================================ + // + // **THIS COST 2.10-2.61 ms PER TOKEN AND HAD NEVER BEEN IN THE PLAN'S BUDGET AT ALL.** The round-309 + // `token host phases` line put it second behind the layer loop among avoidable terms, and the arithmetic + // says why: the table is 320,001,536 rows of `PLE_ROW_BYTES` = 90 B in a 26.8 GB mapping, so the sixteen + // rows a token needs are 1,440 B - **0.5 MB/s**. That is not bandwidth, it is latency: sixteen reads into + // sixteen different 4 KB pages scattered across 26.8 GB, taken ONE AT A TIME, and on this machine the PLE + // shard is 26.8 GB against 63 GB of RAM that the 31.6 GB expert arena is also competing for, so they are + // not in the OS cache. Sixteen serial NVMe reads at ~150 us is 2.4 ms, which is the measurement. + // + // `PrefetchVirtualMemory` issues all sixteen in ONE call and lets them complete in parallel. It is a hint + // and cannot change the answer - a range it does not fetch is simply faulted in by the read that follows - + // so the only risk is that it does nothing. +#if defined(_WIN32) + if (impl_->data != nullptr && g_ple_prefetch) { + WIN32_MEMORY_RANGE_ENTRY ranges[PLE_N_HEADS]; + ULONG_PTR n = 0; + for (int h = 0; h < PLE_N_HEADS; ++h) { + // Out-of-range rows are handled by `read_row` as zeros and have no address to prefetch. + if (rows16[h] >= impl_->n_rows) continue; + ranges[n].VirtualAddress = (PVOID) (impl_->data + (size_t) rows16[h] * impl_->rb); + ranges[n].NumberOfBytes = impl_->rb; + ++n; + } + if (n > 0) (void) PrefetchVirtualMemory(GetCurrentProcess(), n, ranges, 0); + } +#endif + // HEAD-SLOWEST, which is what `ggml_get_rows` does and what the source's own comment says: head h's 160 + // values occupy [h*160, (h+1)*160). A head-fastest layout would put element (d, h) at d*16 + h and needs + // a real transpose - a reshape of the same flat buffer compares equal and would make the check vacuous. + for (int h = 0; h < PLE_N_HEADS; ++h) read_row(rows16[h], out2560 + (size_t) h * PLE_HEAD_DIM); +} + +} // namespace strata::kernels diff --git a/results/round39/model-ab/plan.json b/results/round39/model-ab/plan.json new file mode 100644 index 0000000..eda1d3c --- /dev/null +++ b/results/round39/model-ab/plan.json @@ -0,0 +1,38 @@ +{ + "question": "does round39 regress CLI decode vs reconstructed round38 on the same machine now?", + "baseline": "same linked objects with pre-round39 ngram.cpp.o and baseline runtime shader directory", + "candidate": "current CLI copied and current shader directory snapshot", + "cases": [ + { + "prompt": "data/m5_val_prompts/prompt_albedo_11k.txt", + "n_prompt": 10180, + "n_gen": 32, + "warmup_order": [ + "baseline", + "candidate" + ], + "measured_order": [ + "baseline", + "candidate", + "candidate", + "baseline", + "baseline", + "candidate" + ] + }, + { + "prompt": "data/m5_val_prompts/prompt_2048tok_ops.txt", + "n_prompt": 1984, + "n_gen": 128, + "measured_order": [ + "baseline", + "candidate", + "candidate", + "baseline" + ] + } + ], + "decision": "report adjacent-pair ratios and medians; repeated >2% slowdown is regression evidence; isolated/noisy differences remain inconclusive; no broad zero-regression claim", + "conditions": "all agent cpu suites complete; live service stopped; one model process; identical env except shader path; no source writes during runs", + "limits": "local only, 12 model launches, per-run timeout180s; no servo/telemetry overhead claim" +} diff --git a/results/round39/model-ab/runs.json b/results/round39/model-ab/runs.json new file mode 100644 index 0000000..d1f58a1 --- /dev/null +++ b/results/round39/model-ab/runs.json @@ -0,0 +1,842 @@ +[ + { + "case": "10k-warm", + "index": 0, + "variant": "baseline", + "exit": 0, + "seconds": 37.28526037497795, + "argv": [ + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/baseline-cli", + "--metal-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf", + "--ple-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K256_reap_l2_sum/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf", + "--mtp-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf", + "--n-expert", + "512", + "--ctx", + "32768", + "--spec", + "3", + "--spec-min-p", + "0.35", + "--suffix-draft", + "3", + "--tokens-file", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/data/m5_val_prompts/prompt_albedo_11k.txt", + "--n-prompt", + "10180", + "--n-gen", + "32" + ], + "env": { + "STRATA_METAL_ARENA_GB": "6.5", + "STRATA_METAL_MLOCK_DENSE": "1", + "STRATA_METAL_WIRE_ARENA": "1", + "STRATA_METAL_EXPERT_PROFILE": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/data/expert-profile.bin", + "STRATA_METAL_MOCE_BONUS": "0.85", + "STRATA_METAL_EXPERT_SPLIT": "1", + "STRATA_METAL_SHEXP_OVERLAP": "1", + "STRATA_METAL_SPLIT_TURN": "1", + "STRATA_METAL_PRED_Q8": "1", + "STRATA_METAL_ROUTER_BF16": "1", + "STRATA_METAL_MOE_DN_TILE": "8,2,4", + "STRATA_METAL_PREFETCH": "6", + "STRATA_METAL_PREFETCH2": "2", + "STRATA_METAL_IO_THREADS": "8", + "STRATA_METAL_IO_URGENT": "2", + "STRATA_ANTI_SLOP_PENALTY": "9.0", + "STRATA_ANTI_SLOP_HONEST": "6.0", + "STRATA_ANTI_SLOP_MD": "4.0", + "STRATA_ANTI_SLOP_COLON": "3.0", + "STRATA_METAL_PROMPT_CHUNK": "2048", + "STRATA_METAL_PREFETCH_DRAFT": "0", + "STRATA_METAL_PUBLISH_PREFETCH": "1", + "STRATA_METAL_MTP_SLOP": "1", + "STRATA_METAL_DRAFT_MOCE": "1", + "STRATA_METAL_TRACE": "1", + "STRATA_METAL_KERNELS": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/baseline-kernels" + }, + "metrics": { + "init_ms": "14139.98", + "prompt_tokens": "10180", + "prefill_ms": "21641.96", + "prefill_tok_s": "470.38", + "decode_tokens": "31", + "decode_ms": "888.16", + "decode_tok_s": "34.90", + "decode_ms_per_tok": "28.65" + } + }, + { + "case": "10k-warm", + "index": 1, + "variant": "candidate", + "exit": 0, + "seconds": 26.902124457992613, + "argv": [ + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/candidate-cli", + "--metal-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf", + "--ple-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K256_reap_l2_sum/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf", + "--mtp-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf", + "--n-expert", + "512", + "--ctx", + "32768", + "--spec", + "3", + "--spec-min-p", + "0.35", + "--suffix-draft", + "3", + "--tokens-file", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/data/m5_val_prompts/prompt_albedo_11k.txt", + "--n-prompt", + "10180", + "--n-gen", + "32" + ], + "env": { + "STRATA_METAL_ARENA_GB": "6.5", + "STRATA_METAL_MLOCK_DENSE": "1", + "STRATA_METAL_WIRE_ARENA": "1", + "STRATA_METAL_EXPERT_PROFILE": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/data/expert-profile.bin", + "STRATA_METAL_MOCE_BONUS": "0.85", + "STRATA_METAL_EXPERT_SPLIT": "1", + "STRATA_METAL_SHEXP_OVERLAP": "1", + "STRATA_METAL_SPLIT_TURN": "1", + "STRATA_METAL_PRED_Q8": "1", + "STRATA_METAL_ROUTER_BF16": "1", + "STRATA_METAL_MOE_DN_TILE": "8,2,4", + "STRATA_METAL_PREFETCH": "6", + "STRATA_METAL_PREFETCH2": "2", + "STRATA_METAL_IO_THREADS": "8", + "STRATA_METAL_IO_URGENT": "2", + "STRATA_ANTI_SLOP_PENALTY": "9.0", + "STRATA_ANTI_SLOP_HONEST": "6.0", + "STRATA_ANTI_SLOP_MD": "4.0", + "STRATA_ANTI_SLOP_COLON": "3.0", + "STRATA_METAL_PROMPT_CHUNK": "2048", + "STRATA_METAL_PREFETCH_DRAFT": "0", + "STRATA_METAL_PUBLISH_PREFETCH": "1", + "STRATA_METAL_MTP_SLOP": "1", + "STRATA_METAL_DRAFT_MOCE": "1", + "STRATA_METAL_TRACE": "1", + "STRATA_METAL_KERNELS": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/candidate-kernels" + }, + "metrics": { + "init_ms": "3039.06", + "prompt_tokens": "10180", + "prefill_ms": "22208.75", + "prefill_tok_s": "458.38", + "decode_tokens": "31", + "decode_ms": "891.67", + "decode_tok_s": "34.77", + "decode_ms_per_tok": "28.76" + } + }, + { + "case": "10k", + "index": 0, + "variant": "baseline", + "exit": 0, + "seconds": 27.85658016600064, + "argv": [ + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/baseline-cli", + "--metal-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf", + "--ple-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K256_reap_l2_sum/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf", + "--mtp-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf", + "--n-expert", + "512", + "--ctx", + "32768", + "--spec", + "3", + "--spec-min-p", + "0.35", + "--suffix-draft", + "3", + "--tokens-file", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/data/m5_val_prompts/prompt_albedo_11k.txt", + "--n-prompt", + "10180", + "--n-gen", + "32" + ], + "env": { + "STRATA_METAL_ARENA_GB": "6.5", + "STRATA_METAL_MLOCK_DENSE": "1", + "STRATA_METAL_WIRE_ARENA": "1", + "STRATA_METAL_EXPERT_PROFILE": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/data/expert-profile.bin", + "STRATA_METAL_MOCE_BONUS": "0.85", + "STRATA_METAL_EXPERT_SPLIT": "1", + "STRATA_METAL_SHEXP_OVERLAP": "1", + "STRATA_METAL_SPLIT_TURN": "1", + "STRATA_METAL_PRED_Q8": "1", + "STRATA_METAL_ROUTER_BF16": "1", + "STRATA_METAL_MOE_DN_TILE": "8,2,4", + "STRATA_METAL_PREFETCH": "6", + "STRATA_METAL_PREFETCH2": "2", + "STRATA_METAL_IO_THREADS": "8", + "STRATA_METAL_IO_URGENT": "2", + "STRATA_ANTI_SLOP_PENALTY": "9.0", + "STRATA_ANTI_SLOP_HONEST": "6.0", + "STRATA_ANTI_SLOP_MD": "4.0", + "STRATA_ANTI_SLOP_COLON": "3.0", + "STRATA_METAL_PROMPT_CHUNK": "2048", + "STRATA_METAL_PREFETCH_DRAFT": "0", + "STRATA_METAL_PUBLISH_PREFETCH": "1", + "STRATA_METAL_MTP_SLOP": "1", + "STRATA_METAL_DRAFT_MOCE": "1", + "STRATA_METAL_TRACE": "1", + "STRATA_METAL_KERNELS": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/baseline-kernels" + }, + "metrics": { + "init_ms": "3103.75", + "prompt_tokens": "10180", + "prefill_ms": "23415.08", + "prefill_tok_s": "434.76", + "decode_tokens": "31", + "decode_ms": "908.29", + "decode_tok_s": "34.13", + "decode_ms_per_tok": "29.30" + } + }, + { + "case": "10k", + "index": 1, + "variant": "candidate", + "exit": 0, + "seconds": 29.545378999988316, + "argv": [ + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/candidate-cli", + "--metal-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf", + "--ple-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K256_reap_l2_sum/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf", + "--mtp-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf", + "--n-expert", + "512", + "--ctx", + "32768", + "--spec", + "3", + "--spec-min-p", + "0.35", + "--suffix-draft", + "3", + "--tokens-file", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/data/m5_val_prompts/prompt_albedo_11k.txt", + "--n-prompt", + "10180", + "--n-gen", + "32" + ], + "env": { + "STRATA_METAL_ARENA_GB": "6.5", + "STRATA_METAL_MLOCK_DENSE": "1", + "STRATA_METAL_WIRE_ARENA": "1", + "STRATA_METAL_EXPERT_PROFILE": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/data/expert-profile.bin", + "STRATA_METAL_MOCE_BONUS": "0.85", + "STRATA_METAL_EXPERT_SPLIT": "1", + "STRATA_METAL_SHEXP_OVERLAP": "1", + "STRATA_METAL_SPLIT_TURN": "1", + "STRATA_METAL_PRED_Q8": "1", + "STRATA_METAL_ROUTER_BF16": "1", + "STRATA_METAL_MOE_DN_TILE": "8,2,4", + "STRATA_METAL_PREFETCH": "6", + "STRATA_METAL_PREFETCH2": "2", + "STRATA_METAL_IO_THREADS": "8", + "STRATA_METAL_IO_URGENT": "2", + "STRATA_ANTI_SLOP_PENALTY": "9.0", + "STRATA_ANTI_SLOP_HONEST": "6.0", + "STRATA_ANTI_SLOP_MD": "4.0", + "STRATA_ANTI_SLOP_COLON": "3.0", + "STRATA_METAL_PROMPT_CHUNK": "2048", + "STRATA_METAL_PREFETCH_DRAFT": "0", + "STRATA_METAL_PUBLISH_PREFETCH": "1", + "STRATA_METAL_MTP_SLOP": "1", + "STRATA_METAL_DRAFT_MOCE": "1", + "STRATA_METAL_TRACE": "1", + "STRATA_METAL_KERNELS": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/candidate-kernels" + }, + "metrics": { + "init_ms": "4106.98", + "prompt_tokens": "10180", + "prefill_ms": "24080.38", + "prefill_tok_s": "422.75", + "decode_tokens": "31", + "decode_ms": "920.70", + "decode_tok_s": "33.67", + "decode_ms_per_tok": "29.70" + } + }, + { + "case": "10k", + "index": 2, + "variant": "candidate", + "exit": 0, + "seconds": 28.710599082987756, + "argv": [ + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/candidate-cli", + "--metal-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf", + "--ple-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K256_reap_l2_sum/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf", + "--mtp-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf", + "--n-expert", + "512", + "--ctx", + "32768", + "--spec", + "3", + "--spec-min-p", + "0.35", + "--suffix-draft", + "3", + "--tokens-file", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/data/m5_val_prompts/prompt_albedo_11k.txt", + "--n-prompt", + "10180", + "--n-gen", + "32" + ], + "env": { + "STRATA_METAL_ARENA_GB": "6.5", + "STRATA_METAL_MLOCK_DENSE": "1", + "STRATA_METAL_WIRE_ARENA": "1", + "STRATA_METAL_EXPERT_PROFILE": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/data/expert-profile.bin", + "STRATA_METAL_MOCE_BONUS": "0.85", + "STRATA_METAL_EXPERT_SPLIT": "1", + "STRATA_METAL_SHEXP_OVERLAP": "1", + "STRATA_METAL_SPLIT_TURN": "1", + "STRATA_METAL_PRED_Q8": "1", + "STRATA_METAL_ROUTER_BF16": "1", + "STRATA_METAL_MOE_DN_TILE": "8,2,4", + "STRATA_METAL_PREFETCH": "6", + "STRATA_METAL_PREFETCH2": "2", + "STRATA_METAL_IO_THREADS": "8", + "STRATA_METAL_IO_URGENT": "2", + "STRATA_ANTI_SLOP_PENALTY": "9.0", + "STRATA_ANTI_SLOP_HONEST": "6.0", + "STRATA_ANTI_SLOP_MD": "4.0", + "STRATA_ANTI_SLOP_COLON": "3.0", + "STRATA_METAL_PROMPT_CHUNK": "2048", + "STRATA_METAL_PREFETCH_DRAFT": "0", + "STRATA_METAL_PUBLISH_PREFETCH": "1", + "STRATA_METAL_MTP_SLOP": "1", + "STRATA_METAL_DRAFT_MOCE": "1", + "STRATA_METAL_TRACE": "1", + "STRATA_METAL_KERNELS": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/candidate-kernels" + }, + "metrics": { + "init_ms": "3115.54", + "prompt_tokens": "10180", + "prefill_ms": "24248.45", + "prefill_tok_s": "419.82", + "decode_tokens": "31", + "decode_ms": "915.71", + "decode_tok_s": "33.85", + "decode_ms_per_tok": "29.54" + } + }, + { + "case": "10k", + "index": 3, + "variant": "baseline", + "exit": 0, + "seconds": 32.19512254098663, + "argv": [ + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/baseline-cli", + "--metal-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf", + "--ple-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K256_reap_l2_sum/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf", + "--mtp-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf", + "--n-expert", + "512", + "--ctx", + "32768", + "--spec", + "3", + "--spec-min-p", + "0.35", + "--suffix-draft", + "3", + "--tokens-file", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/data/m5_val_prompts/prompt_albedo_11k.txt", + "--n-prompt", + "10180", + "--n-gen", + "32" + ], + "env": { + "STRATA_METAL_ARENA_GB": "6.5", + "STRATA_METAL_MLOCK_DENSE": "1", + "STRATA_METAL_WIRE_ARENA": "1", + "STRATA_METAL_EXPERT_PROFILE": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/data/expert-profile.bin", + "STRATA_METAL_MOCE_BONUS": "0.85", + "STRATA_METAL_EXPERT_SPLIT": "1", + "STRATA_METAL_SHEXP_OVERLAP": "1", + "STRATA_METAL_SPLIT_TURN": "1", + "STRATA_METAL_PRED_Q8": "1", + "STRATA_METAL_ROUTER_BF16": "1", + "STRATA_METAL_MOE_DN_TILE": "8,2,4", + "STRATA_METAL_PREFETCH": "6", + "STRATA_METAL_PREFETCH2": "2", + "STRATA_METAL_IO_THREADS": "8", + "STRATA_METAL_IO_URGENT": "2", + "STRATA_ANTI_SLOP_PENALTY": "9.0", + "STRATA_ANTI_SLOP_HONEST": "6.0", + "STRATA_ANTI_SLOP_MD": "4.0", + "STRATA_ANTI_SLOP_COLON": "3.0", + "STRATA_METAL_PROMPT_CHUNK": "2048", + "STRATA_METAL_PREFETCH_DRAFT": "0", + "STRATA_METAL_PUBLISH_PREFETCH": "1", + "STRATA_METAL_MTP_SLOP": "1", + "STRATA_METAL_DRAFT_MOCE": "1", + "STRATA_METAL_TRACE": "1", + "STRATA_METAL_KERNELS": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/baseline-kernels" + }, + "metrics": { + "init_ms": "4189.01", + "prompt_tokens": "10180", + "prefill_ms": "26637.98", + "prefill_tok_s": "382.16", + "decode_tokens": "31", + "decode_ms": "949.61", + "decode_tok_s": "32.64", + "decode_ms_per_tok": "30.63" + } + }, + { + "case": "10k", + "index": 4, + "variant": "baseline", + "exit": 0, + "seconds": 32.90031333299703, + "argv": [ + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/baseline-cli", + "--metal-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf", + "--ple-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K256_reap_l2_sum/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf", + "--mtp-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf", + "--n-expert", + "512", + "--ctx", + "32768", + "--spec", + "3", + "--spec-min-p", + "0.35", + "--suffix-draft", + "3", + "--tokens-file", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/data/m5_val_prompts/prompt_albedo_11k.txt", + "--n-prompt", + "10180", + "--n-gen", + "32" + ], + "env": { + "STRATA_METAL_ARENA_GB": "6.5", + "STRATA_METAL_MLOCK_DENSE": "1", + "STRATA_METAL_WIRE_ARENA": "1", + "STRATA_METAL_EXPERT_PROFILE": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/data/expert-profile.bin", + "STRATA_METAL_MOCE_BONUS": "0.85", + "STRATA_METAL_EXPERT_SPLIT": "1", + "STRATA_METAL_SHEXP_OVERLAP": "1", + "STRATA_METAL_SPLIT_TURN": "1", + "STRATA_METAL_PRED_Q8": "1", + "STRATA_METAL_ROUTER_BF16": "1", + "STRATA_METAL_MOE_DN_TILE": "8,2,4", + "STRATA_METAL_PREFETCH": "6", + "STRATA_METAL_PREFETCH2": "2", + "STRATA_METAL_IO_THREADS": "8", + "STRATA_METAL_IO_URGENT": "2", + "STRATA_ANTI_SLOP_PENALTY": "9.0", + "STRATA_ANTI_SLOP_HONEST": "6.0", + "STRATA_ANTI_SLOP_MD": "4.0", + "STRATA_ANTI_SLOP_COLON": "3.0", + "STRATA_METAL_PROMPT_CHUNK": "2048", + "STRATA_METAL_PREFETCH_DRAFT": "0", + "STRATA_METAL_PUBLISH_PREFETCH": "1", + "STRATA_METAL_MTP_SLOP": "1", + "STRATA_METAL_DRAFT_MOCE": "1", + "STRATA_METAL_TRACE": "1", + "STRATA_METAL_KERNELS": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/baseline-kernels" + }, + "metrics": { + "init_ms": "2935.91", + "prompt_tokens": "10180", + "prefill_ms": "28519.16", + "prefill_tok_s": "356.95", + "decode_tokens": "31", + "decode_ms": "996.64", + "decode_tok_s": "31.10", + "decode_ms_per_tok": "32.15" + } + }, + { + "case": "10k", + "index": 5, + "variant": "candidate", + "exit": 0, + "seconds": 35.647112042002846, + "argv": [ + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/candidate-cli", + "--metal-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf", + "--ple-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K256_reap_l2_sum/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf", + "--mtp-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf", + "--n-expert", + "512", + "--ctx", + "32768", + "--spec", + "3", + "--spec-min-p", + "0.35", + "--suffix-draft", + "3", + "--tokens-file", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/data/m5_val_prompts/prompt_albedo_11k.txt", + "--n-prompt", + "10180", + "--n-gen", + "32" + ], + "env": { + "STRATA_METAL_ARENA_GB": "6.5", + "STRATA_METAL_MLOCK_DENSE": "1", + "STRATA_METAL_WIRE_ARENA": "1", + "STRATA_METAL_EXPERT_PROFILE": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/data/expert-profile.bin", + "STRATA_METAL_MOCE_BONUS": "0.85", + "STRATA_METAL_EXPERT_SPLIT": "1", + "STRATA_METAL_SHEXP_OVERLAP": "1", + "STRATA_METAL_SPLIT_TURN": "1", + "STRATA_METAL_PRED_Q8": "1", + "STRATA_METAL_ROUTER_BF16": "1", + "STRATA_METAL_MOE_DN_TILE": "8,2,4", + "STRATA_METAL_PREFETCH": "6", + "STRATA_METAL_PREFETCH2": "2", + "STRATA_METAL_IO_THREADS": "8", + "STRATA_METAL_IO_URGENT": "2", + "STRATA_ANTI_SLOP_PENALTY": "9.0", + "STRATA_ANTI_SLOP_HONEST": "6.0", + "STRATA_ANTI_SLOP_MD": "4.0", + "STRATA_ANTI_SLOP_COLON": "3.0", + "STRATA_METAL_PROMPT_CHUNK": "2048", + "STRATA_METAL_PREFETCH_DRAFT": "0", + "STRATA_METAL_PUBLISH_PREFETCH": "1", + "STRATA_METAL_MTP_SLOP": "1", + "STRATA_METAL_DRAFT_MOCE": "1", + "STRATA_METAL_TRACE": "1", + "STRATA_METAL_KERNELS": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/candidate-kernels" + }, + "metrics": { + "init_ms": "3740.62", + "prompt_tokens": "10180", + "prefill_ms": "30410.22", + "prefill_tok_s": "334.76", + "decode_tokens": "31", + "decode_ms": "1024.55", + "decode_tok_s": "30.26", + "decode_ms_per_tok": "33.05" + } + }, + { + "case": "2k", + "index": 0, + "variant": "baseline", + "exit": 0, + "seconds": 14.231617666024249, + "argv": [ + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/baseline-cli", + "--metal-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf", + "--ple-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K256_reap_l2_sum/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf", + "--mtp-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf", + "--n-expert", + "512", + "--ctx", + "32768", + "--spec", + "3", + "--spec-min-p", + "0.35", + "--suffix-draft", + "3", + "--tokens-file", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/data/m5_val_prompts/prompt_2048tok_ops.txt", + "--n-prompt", + "1984", + "--n-gen", + "128" + ], + "env": { + "STRATA_METAL_ARENA_GB": "6.5", + "STRATA_METAL_MLOCK_DENSE": "1", + "STRATA_METAL_WIRE_ARENA": "1", + "STRATA_METAL_EXPERT_PROFILE": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/data/expert-profile.bin", + "STRATA_METAL_MOCE_BONUS": "0.85", + "STRATA_METAL_EXPERT_SPLIT": "1", + "STRATA_METAL_SHEXP_OVERLAP": "1", + "STRATA_METAL_SPLIT_TURN": "1", + "STRATA_METAL_PRED_Q8": "1", + "STRATA_METAL_ROUTER_BF16": "1", + "STRATA_METAL_MOE_DN_TILE": "8,2,4", + "STRATA_METAL_PREFETCH": "6", + "STRATA_METAL_PREFETCH2": "2", + "STRATA_METAL_IO_THREADS": "8", + "STRATA_METAL_IO_URGENT": "2", + "STRATA_ANTI_SLOP_PENALTY": "9.0", + "STRATA_ANTI_SLOP_HONEST": "6.0", + "STRATA_ANTI_SLOP_MD": "4.0", + "STRATA_ANTI_SLOP_COLON": "3.0", + "STRATA_METAL_PROMPT_CHUNK": "2048", + "STRATA_METAL_PREFETCH_DRAFT": "0", + "STRATA_METAL_PUBLISH_PREFETCH": "1", + "STRATA_METAL_MTP_SLOP": "1", + "STRATA_METAL_DRAFT_MOCE": "1", + "STRATA_METAL_TRACE": "1", + "STRATA_METAL_KERNELS": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/baseline-kernels" + }, + "metrics": { + "init_ms": "3523.84", + "prompt_tokens": "1984", + "prefill_ms": "6219.43", + "prefill_tok_s": "319.00", + "decode_tokens": "127", + "decode_ms": "3997.56", + "decode_tok_s": "31.77", + "decode_ms_per_tok": "31.48" + } + }, + { + "case": "2k", + "index": 1, + "variant": "candidate", + "exit": 0, + "seconds": 15.059506499994313, + "argv": [ + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/candidate-cli", + "--metal-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf", + "--ple-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K256_reap_l2_sum/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf", + "--mtp-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf", + "--n-expert", + "512", + "--ctx", + "32768", + "--spec", + "3", + "--spec-min-p", + "0.35", + "--suffix-draft", + "3", + "--tokens-file", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/data/m5_val_prompts/prompt_2048tok_ops.txt", + "--n-prompt", + "1984", + "--n-gen", + "128" + ], + "env": { + "STRATA_METAL_ARENA_GB": "6.5", + "STRATA_METAL_MLOCK_DENSE": "1", + "STRATA_METAL_WIRE_ARENA": "1", + "STRATA_METAL_EXPERT_PROFILE": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/data/expert-profile.bin", + "STRATA_METAL_MOCE_BONUS": "0.85", + "STRATA_METAL_EXPERT_SPLIT": "1", + "STRATA_METAL_SHEXP_OVERLAP": "1", + "STRATA_METAL_SPLIT_TURN": "1", + "STRATA_METAL_PRED_Q8": "1", + "STRATA_METAL_ROUTER_BF16": "1", + "STRATA_METAL_MOE_DN_TILE": "8,2,4", + "STRATA_METAL_PREFETCH": "6", + "STRATA_METAL_PREFETCH2": "2", + "STRATA_METAL_IO_THREADS": "8", + "STRATA_METAL_IO_URGENT": "2", + "STRATA_ANTI_SLOP_PENALTY": "9.0", + "STRATA_ANTI_SLOP_HONEST": "6.0", + "STRATA_ANTI_SLOP_MD": "4.0", + "STRATA_ANTI_SLOP_COLON": "3.0", + "STRATA_METAL_PROMPT_CHUNK": "2048", + "STRATA_METAL_PREFETCH_DRAFT": "0", + "STRATA_METAL_PUBLISH_PREFETCH": "1", + "STRATA_METAL_MTP_SLOP": "1", + "STRATA_METAL_DRAFT_MOCE": "1", + "STRATA_METAL_TRACE": "1", + "STRATA_METAL_KERNELS": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/candidate-kernels" + }, + "metrics": { + "init_ms": "3630.76", + "prompt_tokens": "1984", + "prefill_ms": "6868.23", + "prefill_tok_s": "288.87", + "decode_tokens": "127", + "decode_ms": "4056.08", + "decode_tok_s": "31.31", + "decode_ms_per_tok": "31.94" + } + }, + { + "case": "2k", + "index": 2, + "variant": "candidate", + "exit": 0, + "seconds": 14.592543375008972, + "argv": [ + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/candidate-cli", + "--metal-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf", + "--ple-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K256_reap_l2_sum/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf", + "--mtp-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf", + "--n-expert", + "512", + "--ctx", + "32768", + "--spec", + "3", + "--spec-min-p", + "0.35", + "--suffix-draft", + "3", + "--tokens-file", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/data/m5_val_prompts/prompt_2048tok_ops.txt", + "--n-prompt", + "1984", + "--n-gen", + "128" + ], + "env": { + "STRATA_METAL_ARENA_GB": "6.5", + "STRATA_METAL_MLOCK_DENSE": "1", + "STRATA_METAL_WIRE_ARENA": "1", + "STRATA_METAL_EXPERT_PROFILE": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/data/expert-profile.bin", + "STRATA_METAL_MOCE_BONUS": "0.85", + "STRATA_METAL_EXPERT_SPLIT": "1", + "STRATA_METAL_SHEXP_OVERLAP": "1", + "STRATA_METAL_SPLIT_TURN": "1", + "STRATA_METAL_PRED_Q8": "1", + "STRATA_METAL_ROUTER_BF16": "1", + "STRATA_METAL_MOE_DN_TILE": "8,2,4", + "STRATA_METAL_PREFETCH": "6", + "STRATA_METAL_PREFETCH2": "2", + "STRATA_METAL_IO_THREADS": "8", + "STRATA_METAL_IO_URGENT": "2", + "STRATA_ANTI_SLOP_PENALTY": "9.0", + "STRATA_ANTI_SLOP_HONEST": "6.0", + "STRATA_ANTI_SLOP_MD": "4.0", + "STRATA_ANTI_SLOP_COLON": "3.0", + "STRATA_METAL_PROMPT_CHUNK": "2048", + "STRATA_METAL_PREFETCH_DRAFT": "0", + "STRATA_METAL_PUBLISH_PREFETCH": "1", + "STRATA_METAL_MTP_SLOP": "1", + "STRATA_METAL_DRAFT_MOCE": "1", + "STRATA_METAL_TRACE": "1", + "STRATA_METAL_KERNELS": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/candidate-kernels" + }, + "metrics": { + "init_ms": "3336.81", + "prompt_tokens": "1984", + "prefill_ms": "6534.59", + "prefill_tok_s": "303.61", + "decode_tokens": "127", + "decode_ms": "4213.72", + "decode_tok_s": "30.14", + "decode_ms_per_tok": "33.18" + } + }, + { + "case": "2k", + "index": 3, + "variant": "baseline", + "exit": 0, + "seconds": 15.292481250013225, + "argv": [ + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/baseline-cli", + "--metal-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf", + "--ple-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K256_reap_l2_sum/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf", + "--mtp-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf", + "--n-expert", + "512", + "--ctx", + "32768", + "--spec", + "3", + "--spec-min-p", + "0.35", + "--suffix-draft", + "3", + "--tokens-file", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/data/m5_val_prompts/prompt_2048tok_ops.txt", + "--n-prompt", + "1984", + "--n-gen", + "128" + ], + "env": { + "STRATA_METAL_ARENA_GB": "6.5", + "STRATA_METAL_MLOCK_DENSE": "1", + "STRATA_METAL_WIRE_ARENA": "1", + "STRATA_METAL_EXPERT_PROFILE": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/data/expert-profile.bin", + "STRATA_METAL_MOCE_BONUS": "0.85", + "STRATA_METAL_EXPERT_SPLIT": "1", + "STRATA_METAL_SHEXP_OVERLAP": "1", + "STRATA_METAL_SPLIT_TURN": "1", + "STRATA_METAL_PRED_Q8": "1", + "STRATA_METAL_ROUTER_BF16": "1", + "STRATA_METAL_MOE_DN_TILE": "8,2,4", + "STRATA_METAL_PREFETCH": "6", + "STRATA_METAL_PREFETCH2": "2", + "STRATA_METAL_IO_THREADS": "8", + "STRATA_METAL_IO_URGENT": "2", + "STRATA_ANTI_SLOP_PENALTY": "9.0", + "STRATA_ANTI_SLOP_HONEST": "6.0", + "STRATA_ANTI_SLOP_MD": "4.0", + "STRATA_ANTI_SLOP_COLON": "3.0", + "STRATA_METAL_PROMPT_CHUNK": "2048", + "STRATA_METAL_PREFETCH_DRAFT": "0", + "STRATA_METAL_PUBLISH_PREFETCH": "1", + "STRATA_METAL_MTP_SLOP": "1", + "STRATA_METAL_DRAFT_MOCE": "1", + "STRATA_METAL_TRACE": "1", + "STRATA_METAL_KERNELS": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/results/round39/model-ab/baseline-kernels" + }, + "metrics": { + "init_ms": "3886.19", + "prompt_tokens": "1984", + "prefill_ms": "6764.25", + "prefill_tok_s": "293.31", + "decode_tokens": "127", + "decode_ms": "4143.13", + "decode_tok_s": "30.65", + "decode_ms_per_tok": "32.62" + } + } +] diff --git a/results/round39/model-ab/summary.json b/results/round39/model-ab/summary.json new file mode 100644 index 0000000..0efeec8 --- /dev/null +++ b/results/round39/model-ab/summary.json @@ -0,0 +1,59 @@ +{ + "10k": { + "pairs": [ + { + "tokens_equal": true, + "decode_ms": 1.0136630371357167, + "prefill_ms": 1.0284133131298292 + }, + { + "tokens_equal": true, + "decode_ms": 0.9643011341498089, + "prefill_ms": 0.9102961260576065 + }, + { + "tokens_equal": true, + "decode_ms": 1.0280040937550168, + "prefill_ms": 1.0663084045953668 + } + ], + "medians": { + "baseline": { + "decode_ms": 949.61, + "decode_tok_s": 32.64, + "prefill_ms": 26637.98 + }, + "candidate": { + "decode_ms": 920.7, + "decode_tok_s": 33.67, + "prefill_ms": 24248.45 + } + } + }, + "2k": { + "pairs": [ + { + "tokens_equal": true, + "decode_ms": 1.0146389297471456, + "prefill_ms": 1.1043182413822488 + }, + { + "tokens_equal": true, + "decode_ms": 1.0170378433696263, + "prefill_ms": 0.9660479727981669 + } + ], + "medians": { + "baseline": { + "decode_ms": 4070.3450000000003, + "decode_tok_s": 31.21, + "prefill_ms": 6491.84 + }, + "candidate": { + "decode_ms": 4134.9, + "decode_tok_s": 30.725, + "prefill_ms": 6701.41 + } + } + } +} diff --git a/results/round39/model-smoke-10k.log b/results/round39/model-smoke-10k.log new file mode 100644 index 0000000..230d7eb --- /dev/null +++ b/results/round39/model-smoke-10k.log @@ -0,0 +1,46 @@ +strata metal stream: experts read from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf.experts (one request per expert) +strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): 290 matrices, 680 MB, made in 135 ms +strata metal stream: locked 2.27 GB of active GPU dense tensors (unpinned 0.34 GB CPU token_embd + 1.28 GB skipped BF16 HC in 3.90 GB file map) +strata metal stream: 1079 dense tensors wrapped zero-copy from file (3.90 GB clean file-backed map, 135.8 ms) +strata metal: decode window on Apple M5: /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf streamed from the SSD (experts) (3.9 GB), 32896 cells, PLE on, ready in 406 ms +strata metal: MTP map: 0.71 GB wired (mlock; 0.00 GB failed), 0.18 GB unused released +strata metal: MTP draft round in one command buffer (embedding type 23 on the GPU: max |GPU - host| 0 over 8 rows) +strata metal: MTP draft layer from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf (0.89 GB, 9 norms folded), draft head over 40525 tokens +strata metal stream: 0.00 GB of the host code's device weight copies given back (unused by the Metal window) +strata metal stream: 2.60 GB allocated before the arena (the process holds 2.24 GB, 5.40 GB resident); RAM 25.8 GB (11.7 GB available, 7.0 GB kept for macOS), GPU working set 19.1 GB +strata metal stream: planned resident footprint 9.6 GB decode (10.7 GB virtual peak during prefill with 1.12 GB transient prompt scratch) = arena 6.5 + 2.6 held before it + 0.5 reserved for the n-gram table; cap 18.8 GB = the lesser of the GPU working set 19.1 GB and the RAM 25.8 GB less 7.0 GB left to macOS and its file cache +strata metal stream: expert arena 6.50 GB, 4129 slots in 2 size classes, 4129 experts read from the profile's top (24576 ranked) in 1.8 s (3.49 GB/s); 8 I/O threads, read-ahead top 6; 1 buffers +strata metal stream: 6.5 GB of the arena wired (mlock) in 629 ms; macOS cannot compress it while the engine idles (STRATA_METAL_WIRE_ARENA=0 turns it off) +strata metal stream: resident footprint after the fill: the process holds 8.7 GB (9.8 GB resident), planned 9.6 GB, cap 18.8 GB +strata metal stream: prompt scratch 1.12 GB mapped in separate anonymous mmap +strata metal: prompts on the matrix units in chunks of 2048 (experts v3, QSA flash, matmul2d layout K4; 1.12 GB of buffers, in a separate transient mmap while a prompt runs) +strata metal: prompt chunk pos 0, 2048 rows: 4709 ms (embd 36 ms, ple 0 ms, graph build 36 ms) +strata metal: prompt chunk pos 2048, 2048 rows: 4371 ms (embd 7 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 4096, 2048 rows: 4627 ms (embd 4 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 6144, 2048 rows: 4794 ms (embd 4 ms, ple 0 ms, graph build 1 ms) +strata metal: prompt chunk pos 8192, 1987 rows: 4684 ms (embd 5 ms, ple 0 ms, graph build 2 ms) +strata metal stream: prompt scratch 1.12 GB unmapped after prefill (phys_footprint 8.89 -> 8.79 GB, resident RSS 9.62 -> 9.53 GB) +strata metal: window pos 10179 token 198 -> 760 (top 760 23.8239, 2nd 35 23.7743, margin 0.0496) 131.94 ms (GPU 100.64 ms, 1233 dispatches) +expert stream prompts 5 chunks: 82802 distinct lookups (345.0 per layer), in the arena 0.1981, streamed ahead 0.6932, missed 0.1088; 66194 streamed in (21901.5 MB/chunk, 74674 reads); host 572.1 ms per chunk, of it waiting 535.6 ms +expert stream 1 windows: 480 distinct lookups (480.0 per window), hits 0.5417, read-ahead hits 0.2875 (5 of them late), misses 0.1708 (82.0 per window); read-aheads 176 (176.0 per window, 525 dropped); 375.2 MB read per window (258 reads); host 24.35 ms per window, of it waiting 23.63 ms (+ 1.14 ms after the GPU went on) +stream windows T=1 1 (decode): 131.62 ms each (GPU incl. the host's turns 100.64 ms) +strata metal: window pos 10180 token 760 -> 1156 (top 1156 29.0186, 2nd 6353 15.3183, margin 13.7003) 86.10 ms (GPU 83.08 ms, 1233 dispatches) +strata metal: window pos 10183 token 33362 -> 11 (top 11 27.0843, 2nd 13 19.1119, margin 7.9724) 83.87 ms (GPU 81.28 ms, 1269 dispatches) +strata metal: window pos 10186 token 8523 -> 728 (top 728 26.5949, 2nd 279 17.2821, margin 9.3128) 77.28 ms (GPU 74.53 ms, 1269 dispatches) +strata metal: window pos 10188 token 13 -> 2838 (top 2838 22.9993, 2nd 2302 20.3597, margin 2.6395) 50.87 ms (GPU 48.62 ms, 1269 dispatches) +strata metal: window pos 10189 token 2838 -> 579 (top 579 24.7598, 2nd 42315 19.9055, margin 4.8543) 60.95 ms (GPU 58.91 ms, 1233 dispatches) +strata metal: window pos 10191 token 5315 -> 1292 (top 1292 22.6356, 2nd 15131 21.4028, margin 1.2328) 60.85 ms (GPU 58.17 ms, 1269 dispatches) +strata metal: window pos 10193 token 37444 -> 685 (top 685 25.6466, 2nd 321 15.4891, margin 10.1575) 66.09 ms (GPU 63.49 ms, 1269 dispatches) +strata metal: window pos 10196 token 9859 -> 1092 (top 1092 23.8715, 2nd 62431 16.8946, margin 6.9769) 62.29 ms (GPU 60.37 ms, 1269 dispatches) +strata metal: window pos 10199 token 2908 -> 978 (top 978 27.4026, 2nd 2725 16.9151, margin 10.4874) 57.68 ms (GPU 55.99 ms, 1269 dispatches) +strata metal: window pos 10202 token 13 -> 1061 (top 1061 21.7433, 2nd 2532 20.4921, margin 1.2513) 63.55 ms (GPU 61.42 ms, 1269 dispatches) +strata metal: window pos 10205 token 264 -> 15855 (top 15855 26.3293, 2nd 3471 19.5190, margin 6.8104) 74.14 ms (GPU 72.00 ms, 1269 dispatches) +strata metal: window pos 10207 token 40719 -> 13 (top 13 21.3241, 2nd 11 20.2257, margin 1.0984) 59.48 ms (GPU 56.98 ms, 1269 dispatches) +strata metal: window pos 10208 token 13 -> 353 (top 353 23.4751, 2nd 6558 22.0544, margin 1.4206) 62.46 ms (GPU 59.79 ms, 1269 dispatches) +STRATA_RESULT init_ms=3072.87 prompt_tokens=10180 prefill_ms=23498.33 prefill_tok_s=433.22 decode_tokens=31 decode_ms=905.63 decode_tok_s=34.23 decode_ms_per_tok=29.21 +STRATA_SPEC rounds=13 drafted=23 accepted=18 accept_pct=78.3 tok_per_round=2.38 mtp_ms=106.58 +GEN_IDS: 760 1156 369 33362 11 852 8523 728 13 2838 579 5315 1292 37444 685 321 9859 1092 353 2908 978 3604 13 1061 369 264 15855 40719 13 353 1220 5707 +expert stream 13 windows: 11692 distinct lookups (899.4 per window), hits 0.8885, read-ahead hits 0.0457 (62 of them late), misses 0.0659 (59.2 per window); read-aheads 1303 (100.2 per window, 0 dropped); 232.8 MB read per window (2073 reads); host 17.63 ms per window, of it waiting 17.08 ms (+ 0.82 ms after the GPU went on) +stream windows T=1 1 (decode): 50.65 ms each (GPU incl. the host's turns 48.62 ms) +stream windows T=2 1 (decode): 59.25 ms each (GPU incl. the host's turns 56.98 ms) +stream windows T=3 11 (decode): 68.45 ms each (GPU incl. the host's turns 66.28 ms) diff --git a/results/round39/model-smoke-2k.log b/results/round39/model-smoke-2k.log new file mode 100644 index 0000000..7707b5b --- /dev/null +++ b/results/round39/model-smoke-2k.log @@ -0,0 +1,68 @@ +strata metal stream: experts read from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf.experts (one request per expert) +strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): 290 matrices, 680 MB, made in 215 ms +strata metal stream: locked 2.27 GB of active GPU dense tensors (unpinned 0.34 GB CPU token_embd + 1.28 GB skipped BF16 HC in 3.90 GB file map) +strata metal stream: 1079 dense tensors wrapped zero-copy from file (3.90 GB clean file-backed map, 103.1 ms) +strata metal: decode window on Apple M5: /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf streamed from the SSD (experts) (3.9 GB), 32896 cells, PLE on, ready in 519 ms +strata metal: MTP map: 0.71 GB wired (mlock; 0.00 GB failed), 0.18 GB unused released +strata metal: MTP draft round in one command buffer (embedding type 23 on the GPU: max |GPU - host| 0 over 8 rows) +strata metal: MTP draft layer from /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf (0.89 GB, 9 norms folded), draft head over 40525 tokens +strata metal stream: 0.00 GB of the host code's device weight copies given back (unused by the Metal window) +strata metal stream: 2.60 GB allocated before the arena (the process holds 2.24 GB, 5.40 GB resident); RAM 25.8 GB (11.8 GB available, 7.0 GB kept for macOS), GPU working set 19.1 GB +strata metal stream: planned resident footprint 9.6 GB decode (10.7 GB virtual peak during prefill with 1.12 GB transient prompt scratch) = arena 6.5 + 2.6 held before it + 0.5 reserved for the n-gram table; cap 18.8 GB = the lesser of the GPU working set 19.1 GB and the RAM 25.8 GB less 7.0 GB left to macOS and its file cache +strata metal stream: expert arena 6.50 GB, 4129 slots in 2 size classes, 4129 experts read from the profile's top (24576 ranked) in 0.9 s (6.63 GB/s); 8 I/O threads, read-ahead top 6; 1 buffers +strata metal stream: 6.5 GB of the arena wired (mlock) in 176 ms; macOS cannot compress it while the engine idles (STRATA_METAL_WIRE_ARENA=0 turns it off) +strata metal stream: resident footprint after the fill: the process holds 8.7 GB (11.9 GB resident), planned 9.6 GB, cap 18.8 GB +strata metal stream: prompt scratch 1.12 GB mapped in separate anonymous mmap +strata metal: prompts on the matrix units in chunks of 2048 (experts v3, QSA flash, matmul2d layout K4; 1.12 GB of buffers, in a separate transient mmap while a prompt runs) +strata metal: prompt chunk pos 0, 1983 rows: 4567 ms (embd 6 ms, ple 0 ms, graph build 203 ms) +strata metal stream: prompt scratch 1.12 GB unmapped after prefill (phys_footprint 8.89 -> 8.78 GB, resident RSS 10.28 -> 10.20 GB) +strata metal: window pos 1983 token 28 -> 16 (top 16 19.7624, 2nd 248046 19.6264, margin 0.1360) 119.79 ms (GPU 90.66 ms, 1197 dispatches) +expert stream prompts 1 chunks: 21669 distinct lookups (451.4 per layer), in the arena 0.1383, streamed ahead 0.7971, missed 0.0647; 18410 streamed in (28273.9 MB/chunk, 19256 reads); host 929.7 ms per chunk, of it waiting 853.0 ms +expert stream 1 windows: 480 distinct lookups (480.0 per window), hits 0.5146, read-ahead hits 0.3125 (24 of them late), misses 0.1729 (83.0 per window); read-aheads 228 (228.0 per window, 555 dropped); 455.8 MB read per window (311 reads); host 26.94 ms per window, of it waiting 26.31 ms (+ 0.88 ms after the GPU went on) +stream windows T=1 1 (decode): 119.57 ms each (GPU incl. the host's turns 90.66 ms) +strata metal: window pos 1984 token 16 -> 248046 (top 248046 20.9569, 2nd 15 19.4699, margin 1.4870) 55.93 ms (GPU 53.96 ms, 1197 dispatches) +strata metal: window pos 1985 token 248046 -> 198 (top 198 30.4136, 2nd 248044 13.3331, margin 17.0805) 91.06 ms (GPU 88.94 ms, 1197 dispatches) +strata metal: window pos 1988 token 248046 -> 198 (top 198 34.6470, 2nd 271 17.3126, margin 17.3344) 65.74 ms (GPU 63.66 ms, 1233 dispatches) +strata metal: window pos 1991 token 74455 -> 198 (top 198 28.0621, 2nd 248046 18.3070, margin 9.7551) 71.70 ms (GPU 69.51 ms, 1233 dispatches) +strata metal: window pos 1994 token 271 -> 248069 (top 248069 34.4544, 2nd 510 14.2772, margin 20.1772) 52.40 ms (GPU 50.29 ms, 1233 dispatches) +strata metal: window pos 1997 token 71093 -> 2164 (top 2164 23.4538, 2nd 198 16.9753, margin 6.4785) 63.13 ms (GPU 61.30 ms, 1233 dispatches) +strata metal: window pos 2000 token 90 -> 198 (top 198 21.8618, 2nd 248046 16.6073, margin 5.2545) 69.32 ms (GPU 67.82 ms, 1233 dispatches) +strata metal: window pos 2003 token 328 -> 10083 (top 10083 21.7508, 2nd 7945 21.2813, margin 0.4695) 75.21 ms (GPU 73.70 ms, 1233 dispatches) +strata metal: window pos 2004 token 10083 -> 763 (top 763 23.1994, 2nd 3147 13.6225, margin 9.5770) 85.24 ms (GPU 83.33 ms, 1233 dispatches) +strata metal: window pos 2008 token 799 -> 1534 (top 1534 23.7638, 2nd 1483 16.7971, margin 6.9667) 94.10 ms (GPU 92.27 ms, 1233 dispatches) +strata metal: window pos 2012 token 318 -> 89478 (top 89478 23.2369, 2nd 13259 17.0454, margin 6.1915) 100.33 ms (GPU 98.51 ms, 1233 dispatches) +strata metal: window pos 2016 token 11 -> 4713 (top 4713 23.7222, 2nd 14408 19.8758, margin 3.8464) 76.78 ms (GPU 75.64 ms, 1233 dispatches) +strata metal: window pos 2020 token 11 -> 14408 (top 14408 26.5939, 2nd 321 17.7990, margin 8.7950) 78.86 ms (GPU 77.66 ms, 1233 dispatches) +strata metal: window pos 2024 token 11836 -> 681 (top 681 22.5685, 2nd 11 19.3141, margin 3.2543) 81.50 ms (GPU 80.23 ms, 1233 dispatches) +strata metal: window pos 2028 token 279 -> 3274 (top 3274 24.5338, 2nd 4434 15.0477, margin 9.4860) 80.87 ms (GPU 79.68 ms, 1233 dispatches) +strata metal: window pos 2032 token 198 -> 220 (top 220 23.5962, 2nd 1 16.3865, margin 7.2098) 72.00 ms (GPU 70.84 ms, 1233 dispatches) +strata metal: window pos 2036 token 763 -> 498 (top 498 23.6984, 2nd 9774 18.6071, margin 5.0914) 64.76 ms (GPU 63.49 ms, 1233 dispatches) +strata metal: window pos 2040 token 313 -> 198 (top 198 24.4607, 2nd 328 18.3975, margin 6.0631) 72.70 ms (GPU 71.38 ms, 1233 dispatches) +strata metal: window pos 2044 token 432 -> 582 (top 582 23.2588, 2nd 39903 12.5246, margin 10.7341) 87.52 ms (GPU 85.16 ms, 1233 dispatches) +strata metal: window pos 2045 token 582 -> 53188 (top 53188 28.0282, 2nd 7206 19.4190, margin 8.6091) 75.64 ms (GPU 73.30 ms, 1233 dispatches) +strata metal: window pos 2046 token 53188 -> 763 (top 763 24.6892, 2nd 3147 15.0103, margin 9.6789) 65.51 ms (GPU 64.18 ms, 1233 dispatches) +strata metal: window pos 2050 token 359 -> 8442 (top 8442 23.6402, 2nd 95429 16.1098, margin 7.5303) 73.94 ms (GPU 71.88 ms, 1269 dispatches) +strata metal: window pos 2054 token 31361 -> 3381 (top 3381 22.9197, 2nd 1981 14.6839, margin 8.2359) 75.32 ms (GPU 73.52 ms, 1269 dispatches) +strata metal: window pos 2058 token 2502 -> 1228 (top 1228 25.2811, 2nd 18997 18.3873, margin 6.8938) 75.47 ms (GPU 74.27 ms, 1269 dispatches) +strata metal: window pos 2062 token 593 -> 2911 (top 2911 25.1616, 2nd 41490 19.8592, margin 5.3024) 76.86 ms (GPU 75.67 ms, 1269 dispatches) +strata metal: window pos 2066 token 3778 -> 976 (top 976 23.8748, 2nd 248046 17.3181, margin 6.5567) 94.00 ms (GPU 92.73 ms, 1269 dispatches) +strata metal: window pos 2070 token 5454 -> 53294 (top 53294 23.8362, 2nd 248046 16.4544, margin 7.3819) 80.44 ms (GPU 79.18 ms, 1269 dispatches) +strata metal: window pos 2074 token 33684 -> 6 (top 6 24.3685, 2nd 318 16.9675, margin 7.4010) 70.57 ms (GPU 69.36 ms, 1269 dispatches) +strata metal: window pos 2078 token 7018 -> 1152 (top 1152 24.8221, 2nd 13259 17.0250, margin 7.7970) 65.39 ms (GPU 63.24 ms, 1269 dispatches) +strata metal: window pos 2080 token 1331 -> 25 (top 25 24.1473, 2nd 314 14.1257, margin 10.0216) 73.43 ms (GPU 71.83 ms, 1269 dispatches) +strata metal: window pos 2084 token 471 -> 66 (top 66 26.1635, 2nd 248046 16.1265, margin 10.0370) 73.87 ms (GPU 72.61 ms, 1269 dispatches) +strata metal: window pos 2088 token 2911 -> 80334 (top 80334 26.8637, 2nd 14 16.2198, margin 10.6438) 64.68 ms (GPU 62.46 ms, 1269 dispatches) +strata metal: window pos 2090 token 53294 -> 3778 (top 3778 24.5964, 2nd 13 13.7553, margin 10.8411) 69.00 ms (GPU 66.85 ms, 1269 dispatches) +strata metal: window pos 2092 token 8 -> 5639 (top 5639 25.4322, 2nd 248046 17.5445, margin 7.8877) 76.00 ms (GPU 74.68 ms, 1269 dispatches) +strata metal: window pos 2096 token 1627 -> 359 (top 359 23.6442, 2nd 7018 23.0138, margin 0.6305) 63.31 ms (GPU 61.27 ms, 1269 dispatches) +strata metal: window pos 2098 token 4277 -> 593 (top 593 22.6791, 2nd 7658 18.9420, margin 3.7372) 69.95 ms (GPU 68.36 ms, 1269 dispatches) +strata metal: window pos 2102 token 11836 -> 12091 (top 12091 23.7703, 2nd 318 17.5935, margin 6.1768) 66.93 ms (GPU 65.58 ms, 1269 dispatches) +strata metal: window pos 2106 token 18997 -> 471 (top 471 24.9212, 2nd 248046 14.8304, margin 10.0908) 56.20 ms (GPU 54.41 ms, 1269 dispatches) +strata metal: window pos 2108 token 16 -> 593 (top 593 24.9663, 2nd 49 15.3718, margin 9.5945) 59.56 ms (GPU 58.28 ms, 1269 dispatches) +STRATA_RESULT init_ms=2430.79 prompt_tokens=1984 prefill_ms=4941.54 prefill_tok_s=401.49 decode_tokens=127 decode_ms=2991.30 decode_tok_s=42.46 decode_ms_per_tok=23.55 +STRATA_SPEC rounds=39 drafted=99 accepted=88 accept_pct=88.9 tok_per_round=3.26 mtp_ms=189.56 +GEN_IDS: 16 248046 198 248045 248046 198 248045 74455 198 248068 271 248069 271 71093 2164 198 90 198 220 328 10083 763 328 6516 799 1534 57287 18944 318 89478 11 1331 11 4713 2552 52186 11 14408 314 7658 11836 681 1179 1804 279 3274 4434 10152 198 220 328 23508 763 498 198 262 313 198 406 328 432 582 53188 763 328 2949 359 8442 64697 43681 31361 3381 1981 6 2502 1228 471 69 593 2911 80334 53294 3778 976 1627 359 5454 53294 3778 25 33684 6 2502 1627 7018 1152 1331 25 4771 23277 471 66 361 593 2911 80334 53294 3778 8 5639 2037 2502 1627 359 4277 593 2911 76644 11836 12091 6 2502 18997 471 16 593 2911 76644 +expert stream 39 windows: 42631 distinct lookups (1093.1 per window), hits 0.8959, read-ahead hits 0.0405 (268 of them late), misses 0.0635 (69.5 per window); read-aheads 5239 (134.3 per window, 0 dropped); 297.4 MB read per window (7948 reads); host 22.54 ms per window, of it waiting 22.05 ms (+ 0.75 ms after the GPU went on) +stream windows T=1 1 (decode): 55.76 ms each (GPU incl. the host's turns 53.96 ms) +stream windows T=3 15 (decode): 68.46 ms each (GPU incl. the host's turns 66.68 ms) +stream windows T=4 23 (decode): 77.21 ms each (GPU incl. the host's turns 75.95 ms) diff --git a/results/round39/model-smoke-plan.json b/results/round39/model-smoke-plan.json new file mode 100644 index 0000000..32f22c1 --- /dev/null +++ b/results/round39/model-smoke-plan.json @@ -0,0 +1,62 @@ +{ + "purpose": "integration smoke only; no paired baseline and no speedup claim", + "env": { + "STRATA_METAL_ARENA_GB": "6.5", + "STRATA_METAL_MLOCK_DENSE": "1", + "STRATA_METAL_WIRE_ARENA": "1", + "STRATA_METAL_EXPERT_PROFILE": "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/data/expert-profile.bin", + "STRATA_METAL_MOCE_BONUS": "0.85", + "STRATA_METAL_EXPERT_SPLIT": "1", + "STRATA_METAL_SHEXP_OVERLAP": "1", + "STRATA_METAL_SPLIT_TURN": "1", + "STRATA_METAL_PRED_Q8": "1", + "STRATA_METAL_ROUTER_BF16": "1", + "STRATA_METAL_MOE_DN_TILE": "8,2,4", + "STRATA_METAL_PREFETCH": "6", + "STRATA_METAL_PREFETCH2": "2", + "STRATA_METAL_IO_THREADS": "8", + "STRATA_METAL_IO_URGENT": "2", + "STRATA_ANTI_SLOP_PENALTY": "9.0", + "STRATA_ANTI_SLOP_HONEST": "6.0", + "STRATA_ANTI_SLOP_MD": "4.0", + "STRATA_ANTI_SLOP_COLON": "3.0", + "STRATA_METAL_PROMPT_CHUNK": "2048", + "STRATA_METAL_PREFETCH_DRAFT": "0", + "STRATA_METAL_PUBLISH_PREFETCH": "1", + "STRATA_METAL_MTP_SLOP": "1", + "STRATA_METAL_DRAFT_MOCE": "1", + "STRATA_METAL_TRACE": "1" + }, + "base_argv": [ + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/build-metal/strata_window_cli", + "--metal-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/k512-iq2_xs-metal.gguf", + "--ple-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K256_reap_l2_sum/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf", + "--mtp-gguf", + "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/models/K512_iq2_xs/mtp-q2_0.gguf", + "--n-expert", + "512", + "--ctx", + "32768", + "--spec", + "3", + "--spec-min-p", + "0.35", + "--suffix-draft", + "3" + ], + "cases": [ + [ + "data/m5_val_prompts/prompt_2048tok_ops.txt", + 1984, + 128 + ], + [ + "data/m5_val_prompts/prompt_albedo_11k.txt", + 10180, + 32 + ] + ], + "concurrency": "one model process; stop existing service before smoke and restart afterwards" +} diff --git a/results/round39/model-smoke-results.json b/results/round39/model-smoke-results.json new file mode 100644 index 0000000..b492a02 --- /dev/null +++ b/results/round39/model-smoke-results.json @@ -0,0 +1,12 @@ +[ + { + "case": "2k", + "exit": 0, + "seconds": 10.854593625001144 + }, + { + "case": "10k", + "exit": 0, + "seconds": 27.82232933299383 + } +] diff --git a/results/round39/parent-bias-image.log b/results/round39/parent-bias-image.log new file mode 100644 index 0000000..4e564ee --- /dev/null +++ b/results/round39/parent-bias-image.log @@ -0,0 +1,8 @@ +..[strata] 400 invalid request: logit_bias: expected an object or a list of [token_id, bias] pairs +[strata] 400 invalid request: logit_bias: this engine does not support token biases; use a supporting CUDA/HIP build +[strata] 400 invalid request: logit_bias: biases must be finite numbers between -100 and 100; false is a pair-list ban +............... +---------------------------------------------------------------------- +Ran 17 tests in 0.872s + +OK diff --git a/results/round39/parent-hardening.log b/results/round39/parent-hardening.log new file mode 100644 index 0000000..14e81e9 --- /dev/null +++ b/results/round39/parent-hardening.log @@ -0,0 +1,16 @@ +........[strata] done: 8 tokens in 0 s (177536.7 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.....[strata] done: 32 tokens in 0 s (273355.9 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (363733.7 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (360800.3 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (372827.0 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (351355.3 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (375960.0 tok/s) (stop, cancel=False) +..............[strata] 400 invalid request: logit_bias: expected an object or a list of [token_id, bias] pairs +[strata] 400 invalid request: logit_bias: this engine does not support token biases; use a supporting CUDA/HIP build +[strata] 400 invalid request: logit_bias: biases must be finite numbers between -100 and 100; false is a pair-list ban +............ +---------------------------------------------------------------------- +Ran 39 tests in 3.500s + +OK diff --git a/results/round39/parent-metal-telemetry.json b/results/round39/parent-metal-telemetry.json new file mode 100644 index 0000000..3eb82a4 --- /dev/null +++ b/results/round39/parent-metal-telemetry.json @@ -0,0 +1 @@ +{"ok": true, "name": "Apple M5", "budget": 19069665280, "sample": {"util": 0, "mem_used": 15040512, "mem_total": 19069665280}} diff --git a/results/round39/parent-moe-default.log b/results/round39/parent-moe-default.log new file mode 100644 index 0000000..66eac20 --- /dev/null +++ b/results/round39/parent-moe-default.log @@ -0,0 +1,3 @@ +device Apple M5; Q2_0/Q2_0, T 256, 64 experts, tile TM 16 TN 64 TK 64 SG 2 +plan: 190 tiles of 16, 3040 padded rows for 2560 entries +parity: plan ok; y vs fp16-operand reference rel L2 2.03e-06, vs exact float reference 2.37e-04 PASS diff --git a/results/round39/parent-moe-v3.log b/results/round39/parent-moe-v3.log new file mode 100644 index 0000000..f6f610d --- /dev/null +++ b/results/round39/parent-moe-v3.log @@ -0,0 +1,4 @@ +device Apple M5; IQ2_S/Q2_0, T 256, 64 experts, tile TM 16 TN 32 TK 32 SG 4 +right-input layout: runs of 4 k +plan: 190 tiles of 16, 3040 padded rows for 2560 entries +parity: plan ok; y vs fp16-operand reference rel L2 1.60e-05, vs exact float reference 4.38e-04 PASS diff --git a/results/round39/parent-qsa-merge2-bitwise.log b/results/round39/parent-qsa-merge2-bitwise.log new file mode 100644 index 0000000..baac0f8 --- /dev/null +++ b/results/round39/parent-qsa-merge2-bitwise.log @@ -0,0 +1,8 @@ +bitwise: 0 of 40 (pos0, ntok) cases differ PASS +device Apple M5: GPU us per dispatch, median of 9 paired rounds of 32 (ref / this) + pos0 512 ntok 1 attn 67.41 / 67.15 merge 18.21 / 12.17 + pos0 512 ntok 4 attn 116.69 / 116.68 merge 17.55 / 13.90 + pos0 2048 ntok 1 attn 51.01 / 51.02 merge 9.94 / 6.21 + pos0 2048 ntok 4 attn 192.52 / 192.52 merge 45.81 / 43.33 + pos0 8192 ntok 1 attn 51.26 / 51.33 merge 10.15 / 6.23 + pos0 8192 ntok 4 attn 195.43 / 195.61 merge 45.72 / 43.73 diff --git a/results/round39/parent-qsa-prompt-noncoop.log b/results/round39/parent-qsa-prompt-noncoop.log new file mode 100644 index 0000000..38a5495 --- /dev/null +++ b/results/round39/parent-qsa-prompt-noncoop.log @@ -0,0 +1,13 @@ +device Apple M5 +pipelines: qsa_attn_mm max threads 1024, tg mem 6912; qsa_flash max threads 1024, tg mem 13760 +pos0 0, T 96: qsa_attn_mm vs qsa_attn (after merge) rel L2 3.09e-04 PASS +pos0 0, T 96: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 3.72e-04 PASS +pos0 1984, T 96: qsa_attn_mm vs qsa_attn (after merge) rel L2 3.61e-04 PASS +pos0 1984, T 96: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 4.27e-04 PASS +pos0 2050, T 96: qsa_attn_mm vs qsa_attn (after merge) rel L2 3.62e-04 PASS +pos0 2050, T 96: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 4.28e-04 PASS +pos0 30000, T 96: qsa_attn_mm vs qsa_attn (after merge) rel L2 3.57e-04 PASS +pos0 30000, T 96: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 4.23e-04 PASS +score pos0 30000, T 96: qsa_score_mm vs qsa_score worst rel 2.76e-06 PASS +qsa_prep(QSA_Q_F16_FC) vs qsa_prep+qsa_q_f16: qcurh memcmp PASS, qsa_flash memcmp PASS +ALL PASS diff --git a/results/round39/parent-qsa-prompt.log b/results/round39/parent-qsa-prompt.log new file mode 100644 index 0000000..32edaa0 --- /dev/null +++ b/results/round39/parent-qsa-prompt.log @@ -0,0 +1,13 @@ +device Apple M5 +pipelines: qsa_attn_mm max threads 1024, tg mem 7424; qsa_flash max threads 1024, tg mem 13760 +pos0 0, T 96: qsa_attn_mm vs qsa_attn (after merge) rel L2 3.09e-04 PASS +pos0 0, T 96: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 3.72e-04 PASS +pos0 1984, T 96: qsa_attn_mm vs qsa_attn (after merge) rel L2 3.61e-04 PASS +pos0 1984, T 96: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 4.27e-04 PASS +pos0 2050, T 96: qsa_attn_mm vs qsa_attn (after merge) rel L2 3.62e-04 PASS +pos0 2050, T 96: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 4.28e-04 PASS +pos0 30000, T 96: qsa_attn_mm vs qsa_attn (after merge) rel L2 3.57e-04 PASS +pos0 30000, T 96: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 4.23e-04 PASS +score pos0 30000, T 96: qsa_score_mm vs qsa_score worst rel 2.76e-06 PASS +qsa_prep(QSA_Q_F16_FC) vs qsa_prep+qsa_q_f16: qcurh memcmp PASS, qsa_flash memcmp PASS +ALL PASS diff --git a/results/round39/parent-serve-full.log b/results/round39/parent-serve-full.log new file mode 100644 index 0000000..9bad430 --- /dev/null +++ b/results/round39/parent-serve-full.log @@ -0,0 +1,1617 @@ +.........[strata] done: 12 tokens in 0 s (151601.3 tok/s) (stop, cancel=False) +..[strata] other apps now use the Chat settings: temperature=0.5 +.[strata] done: 12 tokens in 0 s (285975.3 tok/s) (stop, cancel=False) +.....[strata] done: 12 tokens in 0 s (278075.4 tok/s) (stop, cancel=False) +.[strata] other apps now use the Chat settings: temperature=0.5 +[strata] other apps now use the Chat settings: temperature=0.5 +[strata] other apps now use the Chat settings: temperature=0.5 +..[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +.[strata] done: 1 tokens in 0 s (28532.7 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) + max tokens (1) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] done: 1 tokens in 0 s (16578.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] done: 1 tokens in 0 s (26214.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) + max tokens (1) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] done: 1 tokens in 0 s (35848.8 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] done: 1 tokens in 0 s (33288.1 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 1 tokens in 0 s (45590.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +....sssss....[strata] 400 invalid request: unknown reasoning effort 'huge': use none, low, medium or high +.[strata] done: 24 tokens in 0 s (106297.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (281970.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (304118.7 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 500) +[strata] done: 24 tokens in 0 s (337796.3 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 2047) +[strata] done: 24 tokens in 0 s (189216.7 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 2048) +[strata] done: 24 tokens in 0 s (333322.2 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 8191) +[strata] done: 24 tokens in 0 s (226209.7 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 8192) +[strata] done: 24 tokens in 0 s (244922.9 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 31999) +[strata] done: 24 tokens in 0 s (324720.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (252922.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (303202.7 tok/s) (stop, cancel=False) +.[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (328965.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (267011.4 tok/s) (stop, cancel=False) +[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (198546.9 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (275789.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (289262.3 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (235194.6 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (323676.2 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (320583.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (240821.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (303202.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (278845.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (261463.1 tok/s) (stop, cancel=False) +.[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (304118.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (300487.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (285975.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (189216.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (153684.4 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (291777.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (308783.1 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (307838.8 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (312618.9 tok/s) (stop, cancel=False) +[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (319566.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (289262.3 tok/s) (stop, cancel=False) +.[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (324720.3 tok/s) (stop, cancel=False) +[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (328965.0 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (319566.0 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (315558.9 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (285975.3 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (333322.2 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (311651.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (306900.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (343560.7 tok/s) (stop, cancel=False) +.[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (232478.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (260785.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (240821.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (275789.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (275789.9 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (200524.5 tok/s) (stop, cancel=False) +.[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (332222.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (260785.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (264208.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (311651.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (319566.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (267011.4 tok/s) (stop, cancel=False) +.[strata] thinking limited to 5964 of max_tokens 7951 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (308783.1 tok/s) (stop, cancel=False) +[strata] thinking limited to 6073 of max_tokens 8097 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (338933.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (328965.0 tok/s) (stop, cancel=False) +.[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (348315.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (324720.3 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 20000) +[strata] done: 24 tokens in 0 s (315558.9 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (260785.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (343560.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (296068.5 tok/s) (stop, cancel=False) +.[strata] done: 24 tokens in 0 s (90524.5 tok/s) (stop, cancel=False) +[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (177850.3 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (112097.2 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (173857.2 tok/s) (stop, cancel=False) +.[strata] thinking limited to 5266 of max_tokens 7021 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (328965.0 tok/s) (stop, cancel=False) +[strata] thinking limited to 5213 of max_tokens 6950 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (334429.6 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (323676.2 tok/s) (stop, cancel=False) +.................[strata] the engine stopped unexpectedly (exit code -15). Unrecoverable native verification failure: verify: timed out at layer 18; its GPU waits were released and the GPU finished (#267) The next request starts the engine again. +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -15); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 1 tokens in 0 s (14315.0 tok/s) (length, cancel=False) +...[strata] the engine stopped unexpectedly (exit code -15). Unrecoverable native verification failure: verify: timed out at layer 18; its GPU waits were released and the GPU finished (#267) The next request starts the engine again. +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -15); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 1 tokens in 0 s (23831.3 tok/s) (length, cancel=False) +...................................[strata] image: 512x768, n=2, 8 steps, seed 11 +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5726: ResourceWarning: unclosed file <_io.TextIOWrapper name=11 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 0.3 s (keep) +.[strata] 400 invalid request: size 999x999: each side a multiple of 32 from 256 to 2048 +[strata] image: 1024x1024, n=1, 8 steps, seed 533030141 +[strata] image job failed: sd-cli exited with code 1 +[INFO ] loading tensors completed, taking 0.10s +[ERROR] ggml_vulkan: ErrorDeviceLost +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5730: ResourceWarning: unclosed file <_io.TextIOWrapper name=11 encoding='UTF-8'> + return +ResourceWarning: Enable tracemalloc to get the object allocation traceback +.[strata] image: 1024x1024, n=1, 8 steps, seed 11558204 +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5726: ResourceWarning: unclosed file <_io.TextIOWrapper name=11 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 2.0 s (keep) +[strata] done: 8 tokens in 0 s (122910.0 tok/s) (length, cancel=False) +...[strata] image: 1024x1024, n=1, 8 steps, seed 0 +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5726: ResourceWarning: unclosed file <_io.TextIOWrapper name=11 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 0.3 s (keep) +[strata] image: 1024x1024, n=1, 8 steps, seed 1 +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5726: ResourceWarning: unclosed file <_io.TextIOWrapper name=9 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 0.3 s (keep) +[strata] image: 1024x1024, n=1, 8 steps, seed 2 +[strata] image done in 0.2 s (keep) +[strata] image: 1024x1024, n=1, 8 steps, seed 3 +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5726: ResourceWarning: unclosed file <_io.TextIOWrapper name=9 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 0.3 s (keep) +.[strata] image job: the conversation (123 tokens) saved; unloading the model +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] the conversation from before the image job is back (123 tokens) +...[strata] image job: the engine's expert cache gave VRAM back (9000 MiB free) +.[strata] image job: the conversation (123 tokens) saved; unloading the model +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] the conversation from before the image job is back (123 tokens) +.........[strata] done: 3 tokens in 0 s (64198.5 tok/s) (stop, cancel=False) +[strata] done: 1 tokens in 0 s (15.8 tok/s) (disconnect, cancel=False) +.[strata] loading the model again (it was unloaded) ... +.[strata] model unloaded; the next request loads it again +.....[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 7 tokens in 0 s (117912.2 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (173728.6 tok/s) (stop, cancel=False) +[strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 7 tokens in 0 s (67339.7 tok/s) (stop, cancel=False) +.[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +....[strata] 400 invalid request: logit_bias: expected an object or a list of [token_id, bias] pairs +[strata] 400 invalid request: logit_bias: this engine does not support token biases; use a supporting CUDA/HIP build +[strata] 400 invalid request: logit_bias: biases must be finite numbers between -100 and 100; false is a pair-list ban +...........[strata] MCP server 'missing' did not start: could not start 'strata-no-such-program': [Errno 2] No such file or directory: 'strata-no-such-program'. Its tools are left out; the chat works without them. +[strata] MCP server 'paged': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +...[strata] MCP server 'solo': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'solo' had stopped; starting it again +[strata] MCP server 'solo': 6 tools (echo, fail, sleep, big, die, add) +........[strata] MCP server 'down' did not start: could not reach http://127.0.0.1:9/mcp: . Its tools are left out; the chat works without them. +[strata] MCP server 'web': 6 tools (echo, fail, sleep, big, die, add) +..[strata] thinking limited to 5955 of max_tokens 7939 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (153162.8 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 83 tokens in 0 s (58994.6 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 120 tokens in 0 s (79374.9 tok/s) (stop, cancel=False) +[strata] tool fake__echo: ok, 5 characters in 0.0 s +[strata] thinking limited to 4520 of max_tokens 6026 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 120 tokens in 0 s (110692.0 tok/s) (stop, cancel=False) +[strata] tool fake__echo: ok, 5 characters in 0.0 s +[strata] thinking limited to 4380 of max_tokens 5839 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 120 tokens in 0 s (89414.9 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 141 tokens in 0 s (97913.4 tok/s) (stop, cancel=False) +[strata] tool fake__add: ok, 1 characters in 0.0 s +[strata] thinking limited to 4507 of max_tokens 6009 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (194042.5 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] refused an API request from the web page 'http://evil.example' (no API key; add its host to "allowed_hosts" or its origin to "cors_origins" in the config) +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (78794.7 tok/s) (stop, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4590 of max_tokens 6119 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 109 tokens in 0 s (75804.9 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 5949 of max_tokens 7932 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 116 tokens in 0 s (38614.2 tok/s) (stop, cancel=False) +[strata] thinking limited to 5949 of max_tokens 7932 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 16 tokens in 0 s (77314.4 tok/s) (stop, cancel=False) +[strata] thinking limited to 5949 of max_tokens 7932 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 16 tokens in 0 s (35544.9 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4665 of max_tokens 6220 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 120 tokens in 0 s (73177.7 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 134 tokens in 0 s (83178.4 tok/s) (stop, cancel=False) +[strata] tool fake__echo: ok, 19 characters in 0.0 s +[strata] thinking limited to 4499 of max_tokens 5998 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 31 tokens in 0 s (90671.8 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 84 tokens in 0 s (53301.3 tok/s) (stop, cancel=False) +[strata] tool fake__fail: error, 27 characters in 0.0 s +[strata] thinking limited to 4530 of max_tokens 6040 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 21 tokens in 0 s (100091.3 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 115 tokens in 0 s (73104.7 tok/s) (stop, cancel=False) +[strata] tool fake__big: ok, 3,000 characters in 0.0 s (truncated for the model) +[strata] thinking limited to 4090 of max_tokens 5453 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 16 tokens in 0 s (131328.5 tok/s) (stop, cancel=False) +.[strata] done: 8 tokens in 0 s (98689.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 8 tokens in 0 s (157532.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 8 tokens in 0 s (119410.8 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] done: 8 tokens in 0 s (163680.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] done: 8 tokens in 0 s (173857.2 tok/s) (length, cancel=False) +.[strata] done: 8 tokens in 0 s (173857.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 8 tokens in 0 s (166111.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +..[strata] done: 8 tokens in 0 s (166111.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (208412.6 tok/s) (stop, cancel=False) +....[strata] 400 invalid request: request exceeds the context window: prompt (183 tokens) + max tokens (20000) exceeds the context (16384); requests are never truncated. Send a smaller max_tokens (at most 16193 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] thinking capped at 128 of max_tokens 256 to leave room for the answer (budget 1024) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] thinking capped at 128 of max_tokens 256 to leave room for the answer (budget 1024) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] thinking limited to 12098 of max_tokens 16130 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 50 tokens in 0 s (232500.2 tok/s) (stop, cancel=False) +.[strata] done: 50 tokens in 0 s (249661.0 tok/s) (stop, cancel=False) +[strata] done: 50 tokens in 0 s (107712.0 tok/s) (stop, cancel=False) +[strata] thinking capped at 128 of max_tokens 256 to leave room for the answer (budget 1024) +[strata] done: 50 tokens in 0 s (251457.1 tok/s) (stop, cancel=False) +[strata] thinking capped at 128 of max_tokens 256 to leave room for the answer (budget 1024) +[strata] done: 50 tokens in 0 s (104180.4 tok/s) (stop, cancel=False) +..[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 50 tokens in 0 s (74393.5 tok/s) (stop, cancel=False) +[strata] done: 50 tokens in 0 s (81856.0 tok/s) (stop, cancel=False) +.[strata] thinking limited to 2879 of max_tokens 3838 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 7 tokens in 0 s (170698.4 tok/s) (stop, cancel=False) +..[strata] "parallel" must be a whole number of requests (2..8), not '4': ignored +[strata] "parallel" must be a whole number of requests (2..8), not True: ignored +[strata] "parallel": 12 - the engine runs at most 8 requests together; it will use 8 +.[strata] answering: 5 of max 8 tokens, 6.6 tok/s, 1 s +[strata] answering: 6 of max 8 tokens, 5.3 tok/s, 1 s +[strata] done: 8 tokens in 1 s (5.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 1 s so far +[strata] done: 8 tokens in 2 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 2 s so far +[strata] answering: 7 of max 8 tokens, 6.9 tok/s, 2 s +[strata] done: 8 tokens in 2 s (7.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 3 s so far +[strata] answering: 8 of max 8 tokens, 7.1 tok/s, 3 s +[strata] done: 8 tokens in 3 s (7.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 3 s so far +[strata] answering: 6 of max 8 tokens, 5.6 tok/s, 4 s +[strata] done: 8 tokens in 4 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 4 s so far +[strata] answering: 8 of max 8 tokens, 7.0 tok/s, 4 s +[strata] done: 8 tokens in 4 s (7.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 4 s so far +[strata] answering: 8 of max 8 tokens, 7.4 tok/s, 5 s +[strata] done: 8 tokens in 5 s (7.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 5 s so far +[strata] answering: 8 of max 8 tokens, 6.4 tok/s, 5 s +[strata] done: 8 tokens in 5 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 5 s so far +[strata] answering: 6 of max 8 tokens, 5.6 tok/s, 6 s +[strata] done: 8 tokens in 6 s (5.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 6 s so far +[strata] answering: 8 of max 8 tokens, 6.1 tok/s, 7 s +[strata] done: 8 tokens in 7 s (6.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 7 s so far +[strata] answering: 7 of max 8 tokens, 6.5 tok/s, 8 s +[strata] done: 8 tokens in 8 s (6.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 8 s so far +[strata] answering: 8 of max 8 tokens, 6.1 tok/s, 8 s +[strata] done: 8 tokens in 8 s (6.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 8 s so far +[strata] answering: 8 of max 8 tokens, 7.2 tok/s, 9 s +[strata] done: 8 tokens in 9 s (7.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 9 s so far +[strata] done: 8 tokens in 9 s (8.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 9 s so far +[strata] answering: 7 of max 8 tokens, 6.7 tok/s, 10 s +[strata] done: 8 tokens in 10 s (7.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 10 s so far +[strata] answering: 8 of max 8 tokens, 6.7 tok/s, 10 s +[strata] done: 8 tokens in 10 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 10 s so far +[strata] answering: 8 of max 8 tokens, 7.6 tok/s, 11 s +[strata] done: 8 tokens in 11 s (7.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 11 s so far +[strata] done: 8 tokens in 11 s (8.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 11 s so far +[strata] done: 8 tokens in 12 s (8.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 12 s so far +[strata] answering: 8 of max 8 tokens, 6.5 tok/s, 13 s +[strata] done: 8 tokens in 13 s (6.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 13 s so far +[strata] answering: 6 of max 8 tokens, 5.8 tok/s, 13 s +[strata] done: 8 tokens in 14 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 14 s so far +[strata] answering: 8 of max 8 tokens, 6.6 tok/s, 14 s +[strata] done: 8 tokens in 14 s (6.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 14 s so far +[strata] answering: 7 of max 8 tokens, 6.8 tok/s, 15 s +[strata] done: 8 tokens in 15 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 15 s so far +[strata] answering: 8 of max 8 tokens, 6.7 tok/s, 15 s +[strata] done: 8 tokens in 15 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 15 s so far +[strata] answering: 8 of max 8 tokens, 7.9 tok/s, 16 s +[strata] done: 8 tokens in 16 s (7.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 16 s so far +[strata] answering: 8 of max 8 tokens, 7.2 tok/s, 16 s +[strata] done: 8 tokens in 16 s (7.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 16 s so far +[strata] answering: 7 of max 8 tokens, 6.5 tok/s, 17 s +[strata] done: 8 tokens in 17 s (6.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 17 s so far +[strata] answering: 8 of max 8 tokens, 7.3 tok/s, 17 s +[strata] done: 8 tokens in 17 s (7.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 17 s so far +[strata] done: 8 tokens in 18 s (8.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 18 s so far +[strata] answering: 8 of max 8 tokens, 7.7 tok/s, 18 s +[strata] done: 8 tokens in 18 s (7.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 19 s so far +[strata] answering: 8 of max 8 tokens, 7.1 tok/s, 19 s +[strata] done: 8 tokens in 19 s (7.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 19 s so far +[strata] answering: 8 of max 8 tokens, 7.1 tok/s, 20 s +[strata] done: 8 tokens in 20 s (7.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 20 s so far +[strata] answering: 6 of max 8 tokens, 5.8 tok/s, 21 s +[strata] answering: 7 of max 8 tokens, 6.5 tok/s, 21 s +[strata] done: 8 tokens in 21 s (6.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 21 s so far +[strata] done: 8 tokens in 21 s (5.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 21 s so far +[strata] answering: 7 of max 8 tokens, 6.8 tok/s, 22 s +[strata] done: 8 tokens in 22 s (6.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 22 s so far +[strata] answering: 8 of max 8 tokens, 6.6 tok/s, 22 s +[strata] done: 8 tokens in 22 s (6.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 22 s so far +[strata] done: 8 tokens in 23 s (8.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 23 s so far +[strata] done: 8 tokens in 23 s (9.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 23 s so far +[strata] done: 8 tokens in 24 s (8.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 24 s so far +[strata] done: 8 tokens in 24 s (8.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 24 s so far +[strata] answering: 7 of max 8 tokens, 6.8 tok/s, 25 s +[strata] done: 8 tokens in 25 s (6.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 25 s so far +[strata] answering: 8 of max 8 tokens, 6.1 tok/s, 26 s +[strata] done: 8 tokens in 26 s (6.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 26 s so far +[strata] answering: 7 of max 8 tokens, 6.9 tok/s, 26 s +[strata] done: 8 tokens in 27 s (6.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 27 s so far +[strata] answering: 8 of max 8 tokens, 7.4 tok/s, 27 s +[strata] done: 8 tokens in 27 s (7.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 27 s so far +[strata] answering: 8 of max 8 tokens, 7.1 tok/s, 28 s +[strata] done: 8 tokens in 28 s (7.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 28 s so far +[strata] answering: 8 of max 8 tokens, 7.1 tok/s, 28 s +[strata] done: 8 tokens in 28 s (7.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 28 s so far +[strata] done: 8 tokens in 29 s (9.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 29 s so far +[strata] done: 8 tokens in 29 s (8.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 29 s so far +[strata] answering: 8 of max 8 tokens, 7.2 tok/s, 30 s +[strata] done: 8 tokens in 30 s (7.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 30 s so far +[strata] answering: 8 of max 8 tokens, 6.7 tok/s, 30 s +[strata] done: 8 tokens in 30 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 30 s so far +[strata] answering: 6 of max 8 tokens, 5.4 tok/s, 31 s +[strata] done: 8 tokens in 31 s (6.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 31 s so far +[strata] answering: 8 of max 8 tokens, 6.4 tok/s, 31 s +[strata] done: 8 tokens in 31 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 32 s so far +[strata] done: 8 tokens in 32 s (8.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 32 s so far +[strata] done: 8 tokens in 32 s (8.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 32 s so far +[strata] answering: 7 of max 8 tokens, 6.6 tok/s, 33 s +[strata] done: 8 tokens in 33 s (6.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 33 s so far +[strata] answering: 8 of max 8 tokens, 6.7 tok/s, 34 s +[strata] done: 8 tokens in 34 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 34 s so far +[strata] answering: 7 of max 8 tokens, 6.4 tok/s, 35 s +[strata] done: 8 tokens in 35 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 35 s so far +[strata] answering: 8 of max 8 tokens, 6.4 tok/s, 35 s +[strata] done: 8 tokens in 35 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 35 s so far +[strata] answering: 6 of max 8 tokens, 5.4 tok/s, 36 s +[strata] done: 8 tokens in 36 s (6.0 tok/s) (length, cancel=False) +[strata] answering: 8 of max 8 tokens, 7.0 tok/s, 36 s +[strata] done: 8 tokens in 36 s (7.0 tok/s) (length, cancel=False) +.[strata] thinking: 1 of max 200 tokens, 1.9 tok/s, 1 s +[strata] reading the prompt: 32 of 249 tokens, 1 s so far +[strata] thinking: 4 of max 200 tokens, 2.3 tok/s, 2 s +[strata] thinking: 5 of max 200 tokens, 8.1 tok/s, 2 s +[strata] done: 10 tokens in 3 s (7.6 tok/s) (stop, cancel=False) +[strata] thinking: 13 of max 200 tokens, 4.6 tok/s, 3 s +[strata] thinking: 18 of max 200 tokens, 4.6 tok/s, 4 s +[strata] thinking: 29 of max 200 tokens, 5.8 tok/s, 5 s +[strata] thinking: 37 of max 200 tokens, 6.2 tok/s, 7 s +[strata] thinking: 44 of max 200 tokens, 6.3 tok/s, 8 s +[strata] thinking: 52 of max 200 tokens, 6.4 tok/s, 9 s +[strata] done: 55 tokens in 9 s (6.3 tok/s) (stop, cancel=False) +.[strata] the engine stopped unexpectedly (exit code 1). The engine exited (code 1) after it reported: verify batch: layer 34 never rang (graph finished) - please report it at github.com/Niko1221/Strata/issues with the log. The next request starts the engine again. +[strata] done: 3 tokens in 1 s (3.7 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] the engine stopped unexpectedly (exit code 1). The engine exited (code 1) after it reported: verify batch: layer 34 never rang (graph finished) - please report it at github.com/Niko1221/Strata/issues with the log. The next request starts the engine again. +[strata] done: 1 tokens in 1 s (4.6 tok/s) (error, cancel=False) +.[strata] parallel: slot(s) 0, 1 were marked busy with no request using them (issue #1603); freed, the waiting request goes on +[strata] reading the prompt: 32 of 52 tokens, 1 s so far +[strata] done: 10 tokens in 2 s (8.8 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 224 of 3,037 tokens, 1 s so far +[strata] answering: 5 of max 64 tokens, 4.7 tok/s, 1 s +[strata] reading the prompt: 800 of 3,037 tokens, 2 s so far +[strata] done: 10 tokens in 2 s (4.3 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,280 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 1,792 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 2,272 of 3,037 tokens, 5 s so far +[strata] reading the prompt: 2,720 of 3,037 tokens, 6 s so far +[strata] answering: 7 of max 64 tokens, 10.8 tok/s, 7 s +[strata] done: 10 tokens in 8 s (11.3 tok/s) (stop, cancel=False) +.[strata] answering: 8 of max 64 tokens, 7.5 tok/s, 1 s +[strata] done: 10 tokens in 2 s (6.6 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 1 s (10.4 tok/s) (stop, cancel=False) +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] answering: 9 of max 16 tokens, 8.6 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.7 tok/s) (stop, cancel=False) +.[strata] answering: 6 of max 200 tokens, 5.7 tok/s, 1 s +[strata] answering: 4 of max 200 tokens, 7.9 tok/s, 1 s +[strata] answering: 14 of max 200 tokens, 6.7 tok/s, 2 s +[strata] done: 10 tokens in 2 s (7.5 tok/s) (stop, cancel=False) +[strata] answering: 22 of max 200 tokens, 6.8 tok/s, 3 s +[strata] answering: 29 of max 200 tokens, 6.7 tok/s, 4 s +[strata] answering: 38 of max 200 tokens, 7.0 tok/s, 5 s +[strata] answering: 47 of max 200 tokens, 7.1 tok/s, 7 s +[strata] done: 55 tokens in 8 s (7.0 tok/s) (stop, cancel=False) +.[strata] answering: 6 of max 200 tokens, 6.2 tok/s, 1 s +[strata] answering: 5 of max 200 tokens, 6.8 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.1 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 2 s (7.5 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 224 of 3,037 tokens, 1 s so far +[strata] answering: 5 of max 60 tokens, 4.1 tok/s, 1 s +[strata] answering: 4 of max 64 tokens, 3.9 tok/s, 1 s +[strata] reading the prompt: 672 of 3,037 tokens, 2 s so far +[strata] answering: 9 of max 60 tokens, 3.7 tok/s, 3 s +[strata] answering: 8 of max 64 tokens, 3.5 tok/s, 2 s +[strata] done: 10 tokens in 3 s (3.7 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,152 of 3,037 tokens, 3 s so far +[strata] done: 10 tokens in 3 s (3.6 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,664 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 2,208 of 3,037 tokens, 5 s so far +[strata] reading the prompt: 2,688 of 3,037 tokens, 6 s so far +[strata] answering: 4 of max 64 tokens, 8.3 tok/s, 7 s +[strata] done: 10 tokens in 8 s (7.9 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 2,560 of 3,037 tokens, 6 s so far +[strata] answering: 1 of max 64 tokens, 167772.2 tok/s, 5 s +[strata] answering: 2 of max 64 tokens, 0.5 tok/s, 10 s +[strata] reading the prompt: 2,560 of 3,037 tokens, 10 s so far +[strata] answering: 4 of max 60 tokens, 0.4 tok/s, 10 s +[strata] answering: 3 of max 64 tokens, 0.5 tok/s, 11 s +[strata] answering: 2 of max 64 tokens, 6.5 tok/s, 11 s +[strata] answering: 5 of max 60 tokens, 0.4 tok/s, 11 s +[strata] done: 10 tokens in 12 s (1.6 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 12 s (8.2 tok/s) (stop, cancel=False) +[strata] answering: 14 of max 60 tokens, 1.1 tok/s, 13 s +[strata] answering: 24 of max 60 tokens, 1.8 tok/s, 14 s +[strata] answering: 34 of max 60 tokens, 2.3 tok/s, 15 s +[strata] answering: 44 of max 60 tokens, 2.8 tok/s, 16 s +[strata] answering: 53 of max 60 tokens, 3.1 tok/s, 17 s +[strata] done: 55 tokens in 17 s (3.2 tok/s) (stop, cancel=False) +.[strata] parallel: slot(s) 0 were marked busy with no request using them (issue #1603); freed, the waiting request goes on +..[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (1/6) (issue #1994) +[strata] answering: 7 of max 64 tokens, 6.6 tok/s, 1 s +[strata] answering: 3 of max 64 tokens, 7.2 tok/s, 1 s +[strata] answering: 15 of max 64 tokens, 7.1 tok/s, 2 s +[strata] done: 10 tokens in 2 s (7.2 tok/s) (stop, cancel=False) +[strata] answering: 24 of max 64 tokens, 7.5 tok/s, 3 s +[strata] answering: 31 of max 64 tokens, 7.3 tok/s, 4 s +[strata] answering: 40 of max 64 tokens, 7.5 tok/s, 5 s +[strata] answering: 49 of max 64 tokens, 7.6 tok/s, 7 s +[strata] done: 55 tokens in 8 s (7.4 tok/s) (stop, cancel=False) +[strata] answering: 9 of max 64 tokens, 8.8 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.2 tok/s) (stop, cancel=False) +.[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (1/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (2/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (3/6) (issue #1994) +[strata] answering: 8 of max 64 tokens, 8.1 tok/s, 1 s +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (4/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (5/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (6/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (7/6) (issue #1994) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (length, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] answering: 13 of max 64 tokens, 6.3 tok/s, 2 s +[strata] answering: 21 of max 64 tokens, 6.8 tok/s, 3 s +[strata] answering: 28 of max 64 tokens, 6.8 tok/s, 4 s +[strata] answering: 35 of max 64 tokens, 6.7 tok/s, 5 s +[strata] answering: 45 of max 64 tokens, 7.1 tok/s, 6 s +[strata] done: 55 tokens in 7 s (7.6 tok/s) (stop, cancel=False) +[strata] answering: 7 of max 64 tokens, 7.4 tok/s, 1 s +[strata] done: 10 tokens in 2 s (6.7 tok/s) (stop, cancel=False) +.[strata] answering: 7 of max 64 tokens, 6.6 tok/s, 1 s +[strata] answering: 4 of max 64 tokens, 6.5 tok/s, 1 s +[strata] done: 10 tokens in 2 s (7.0 tok/s) (stop, cancel=False) +[strata] reading the prompt: 32 of 66 tokens, 2 s so far +[strata] answering: 15 of max 64 tokens, 6.4 tok/s, 2 s +[strata] done: 10 tokens in 3 s (10.4 tok/s) (stop, cancel=False) +[strata] answering: 25 of max 64 tokens, 7.4 tok/s, 3 s +[strata] answering: 34 of max 64 tokens, 7.6 tok/s, 5 s +[strata] answering: 42 of max 64 tokens, 7.6 tok/s, 6 s +[strata] answering: 49 of max 64 tokens, 7.5 tok/s, 7 s +[strata] done: 55 tokens in 7 s (7.5 tok/s) (stop, cancel=False) +.[strata] parallel requests: 4 asked, the engine runs 2 (its log says why) +.[strata] parallel requests: 4 asked, the engine runs one at a time (its log says why) +[strata] answering: 8 of max 64 tokens, 8.0 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.4 tok/s) (stop, cancel=False) +.[strata] answering: 6 of max 200 tokens, 5.3 tok/s, 1 s +[strata] answering: 4 of max 200 tokens, 7.6 tok/s, 1 s +[strata] done: 10 tokens in 2 s (7.7 tok/s) (stop, cancel=False) +[strata] answering: 15 of max 200 tokens, 6.6 tok/s, 2 s +[strata] answering: 22 of max 200 tokens, 6.6 tok/s, 3 s +[strata] answering: 29 of max 200 tokens, 6.5 tok/s, 4 s +[strata] answering: 38 of max 200 tokens, 6.9 tok/s, 6 s +[strata] answering: 48 of max 200 tokens, 7.3 tok/s, 7 s +[strata] done: 55 tokens in 8 s (7.1 tok/s) (stop, cancel=False) +.[strata] answering: 3 of max 64 tokens, 6.7 tok/s, 1 s +[strata] answering: 5 of max 200 tokens, 4.9 tok/s, 1 s +[strata] done: 10 tokens in 2 s (7.4 tok/s) (stop, cancel=False) +[strata] answering: 13 of max 200 tokens, 6.2 tok/s, 2 s +[strata] answering: 23 of max 200 tokens, 7.3 tok/s, 3 s +[strata] answering: 30 of max 200 tokens, 7.3 tok/s, 4 s +[strata] answering: 38 of max 200 tokens, 7.3 tok/s, 5 s +[strata] answering: 47 of max 200 tokens, 7.6 tok/s, 6 s +[strata] done: 55 tokens in 7 s (7.6 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 640 of 3,037 tokens, 1 s so far +[strata] reading the prompt: 1,152 of 3,037 tokens, 2 s so far +[strata] reading the prompt: 1,728 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 2,272 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 2,784 of 3,037 tokens, 5 s so far +[strata] answering: 1 of max 64 tokens, 1.8 tok/s, 6 s +[strata] answering: 1 of max 64 tokens, 0.6 tok/s, 7 s +[strata] answering: 1 of max 64 tokens, 0.4 tok/s, 8 s +[strata] answering: 1 of max 64 tokens, 0.3 tok/s, 9 s +[strata] answering: 1 of max 64 tokens, 0.2 tok/s, 10 s +[strata] reading the prompt: 2,464 of 3,038 tokens, 11 s so far +[strata] answering: 1 of max 64 tokens, 0.2 tok/s, 11 s +[strata] reading the prompt: 32 of 3,037 tokens, 12 s so far +[strata] answering: 4 of max 64 tokens, 0.6 tok/s, 12 s +[strata] reading the prompt: 608 of 3,037 tokens, 13 s so far +[strata] answering: 9 of max 64 tokens, 1.1 tok/s, 14 s +[strata] done: 10 tokens in 14 s (1.2 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,280 of 3,037 tokens, 14 s so far +[strata] reading the prompt: 1,792 of 3,037 tokens, 15 s so far +[strata] reading the prompt: 2,336 of 3,037 tokens, 16 s so far +[strata] reading the prompt: 2,880 of 3,037 tokens, 17 s so far +[strata] answering: 7 of max 64 tokens, 8.1 tok/s, 18 s +[strata] done: 10 tokens in 18 s (8.6 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 2,560 of 3,037 tokens, 5 s so far +[strata] answering: 1 of max 64 tokens, 0.2 tok/s, 10 s +[strata] reading the prompt: 2,560 of 3,038 tokens, 10 s so far +[strata] answering: 3 of max 64 tokens, 0.3 tok/s, 16 s +[strata] reading the prompt: 2,560 of 3,037 tokens, 15 s so far +[strata] answering: 2 of max 64 tokens, 7.7 tok/s, 16 s +[strata] answering: 4 of max 64 tokens, 0.4 tok/s, 17 s +[strata] done: 10 tokens in 17 s (0.8 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 17 s (8.3 tok/s) (stop, cancel=False) +.[strata] the engine's logits became non-finite before the first token (#879); it dropped its cached state: sending the request again once +[strata] done: 10 tokens in 1 s (8.7 tok/s) (stop, cancel=False) +[strata] the engine's logits became non-finite before the first token (#879); it dropped its cached state: sending the request again once +[strata] answering: 7 of max 64 tokens, 6.9 tok/s, 1 s +[strata] answering: 6 of max 64 tokens, 7.9 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.8 tok/s) (stop, cancel=False) +[strata] answering: 16 of max 64 tokens, 7.8 tok/s, 2 s +[strata] answering: 23 of max 64 tokens, 7.4 tok/s, 3 s +[strata] answering: 31 of max 64 tokens, 7.5 tok/s, 4 s +[strata] answering: 39 of max 64 tokens, 7.6 tok/s, 5 s +[strata] answering: 51 of max 64 tokens, 8.2 tok/s, 6 s +[strata] done: 55 tokens in 7 s (8.3 tok/s) (stop, cancel=False) +[strata] answering: 7 of max 64 tokens, 7.3 tok/s, 1 s +[strata] the engine's logits became non-finite after 2 tokens (#879): the reply ends here as "length"; the engine dropped its cached state, so the next request reads its prompt from the start +[strata] done: 2 tokens in 1 s (8.9 tok/s) (length, cancel=False) +[strata] answering: 16 of max 64 tokens, 7.8 tok/s, 2 s +[strata] answering: 25 of max 64 tokens, 8.2 tok/s, 3 s +[strata] answering: 34 of max 64 tokens, 8.3 tok/s, 4 s +[strata] answering: 44 of max 64 tokens, 8.4 tok/s, 5 s +[strata] answering: 52 of max 64 tokens, 8.2 tok/s, 6 s +[strata] done: 55 tokens in 7 s (8.2 tok/s) (stop, cancel=False) +.[strata] done: 9 tokens in 1 s (8.8 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 1 s (12.3 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 2,560 of 3,037 tokens, 4 s so far +[strata] answering: 1 of max 64 tokens, 322638.8 tok/s, 4 s +[strata] reading the prompt: 2,560 of 3,037 tokens, 9 s so far +[strata] answering: 2 of max 64 tokens, 0.4 tok/s, 9 s +[strata] answering: 4 of max 64 tokens, 9.8 tok/s, 10 s +[strata] answering: 5 of max 64 tokens, 0.8 tok/s, 10 s +[strata] done: 10 tokens in 11 s (1.6 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 11 s (9.1 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 32 of 3,037 tokens, 1 s so far +[strata] reading the prompt: 416 of 3,037 tokens, 2 s so far +[strata] reading the prompt: 928 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 1,600 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 2,336 of 3,037 tokens, 5 s so far +[strata] reading the prompt: 2,976 of 3,037 tokens, 6 s so far +[strata] reading the prompt: 32 of 3,037 tokens, 6 s so far +[strata] answering: 5 of max 64 tokens, 5.1 tok/s, 7 s +[strata] reading the prompt: 576 of 3,037 tokens, 7 s so far +[strata] answering: 9 of max 64 tokens, 4.2 tok/s, 8 s +[strata] reading the prompt: 1,056 of 3,037 tokens, 8 s so far +[strata] done: 10 tokens in 9 s (4.1 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,632 of 3,037 tokens, 9 s so far +[strata] reading the prompt: 2,176 of 3,037 tokens, 10 s so far +[strata] reading the prompt: 2,688 of 3,037 tokens, 11 s so far +[strata] reading the prompt: 32 of 43 tokens, 12 s so far +[strata] answering: 4 of max 64 tokens, 6.7 tok/s, 12 s +[strata] answering: 9 of max 64 tokens, 8.1 tok/s, 13 s +[strata] done: 10 tokens in 13 s (7.8 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 13 s (6.6 tok/s) (stop, cancel=False) +...............[strata] answering: 11 of max 64 tokens, 12.4 tok/s, 1 s +[strata] reading the prompt: 10 of max 64 tokens, 10.2 tok/s, 1 s +[strata] answering: 11 of max 64 tokens, 9.8 tok/s, 1 s +[strata] answering: 19 of max 64 tokens, 9.6 tok/s, 2 s +[strata] answering: 19 of max 64 tokens, 9.1 tok/s, 2 s +[strata] answering: 20 of max 64 tokens, 8.9 tok/s, 2 s +[strata] answering: 29 of max 64 tokens, 9.3 tok/s, 3 s +[strata] answering: 27 of max 64 tokens, 8.7 tok/s, 3 s +[strata] answering: 29 of max 64 tokens, 8.9 tok/s, 3 s +[strata] answering: 37 of max 64 tokens, 8.8 tok/s, 4 s +[strata] answering: 35 of max 64 tokens, 8.3 tok/s, 4 s +[strata] answering: 39 of max 64 tokens, 8.9 tok/s, 4 s +[strata] answering: 44 of max 64 tokens, 8.4 tok/s, 5 s +[strata] answering: 46 of max 64 tokens, 8.8 tok/s, 5 s +[strata] answering: 47 of max 64 tokens, 8.6 tok/s, 5 s +[strata] done: 51 tokens in 6 s (8.8 tok/s) (stop, cancel=False) +[strata] done: 51 tokens in 6 s (8.6 tok/s) (stop, cancel=False) +[strata] done: 51 tokens in 6 s (8.5 tok/s) (stop, cancel=False) +.[strata] done: 22 tokens in 0 s (265921.3 tok/s) (stop, cancel=False) +............[strata] the config has a "vision" section but no --vision in its args: added --vision (run setup again with --vision to write them) +....[strata] thinking limited to 24386 of max_tokens 32514 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 3 tokens in 0 s (30615.4 tok/s) (stop, cancel=False) +...[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (7358.4 tok/s) (length, cancel=False) +.[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] answering: 1 of max 1 tokens, 139810.1 tok/s, 1 s +[strata] done: 1 tokens in 1 s (4539.3 tok/s) (length, cancel=False) +.[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (7884.0 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] answering: 1 of max 1 tokens, 262144.0 tok/s, 1 s +[strata] done: 1 tokens in 1 s (7884.0 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (12336.2 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (6403.5 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (7463.2 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (5108.8 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 2 s (0.0 tok/s) (cancel, cancel=True) +[strata] answering: 1 of max 1 tokens, 144631.2 tok/s, 1 s +[strata] done: 1 tokens in 1 s (4350.9 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (7145.3 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (6413.3 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (6887.2 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] answering: 1 of max 1 tokens, 262144.0 tok/s, 1 s +[strata] done: 1 tokens in 1 s (5817.3 tok/s) (length, cancel=False) +...[strata] loading the model again (it was unloaded) ... +[strata] the engine did not start (try 1 of 3); again in 0 s +[strata] the engine did not start (try 2 of 3); again in 0 s +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +.[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 1 tokens in 0 s (27776.8 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine did not start (try 1 of 3); again in 0 s +[strata] the engine did not start (try 2 of 3); again in 0 s +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +..[strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2086 tokens in 0 s (244886.9 tok/s) (stop, cancel=False) +.[strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2071 tokens in 0 s (412459.8 tok/s) (stop, cancel=False) +.[strata] done: 1 tokens in 0 s (27594.1 tok/s) (cancel, cancel=True) +.[strata] done: 16384 tokens in 0 s (664264.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] done: 28001 tokens in 0 s (787537.7 tok/s) (stop, cancel=False) +.[strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2082 tokens in 0 s (695210.6 tok/s) (stop, cancel=False) +.[strata] done: 28001 tokens in 0 s (761145.2 tok/s) (stop, cancel=False) +.[strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2106 tokens in 0 s (730439.4 tok/s) (stop, cancel=False) +.[strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2086 tokens in 0 s (742915.7 tok/s) (stop, cancel=False) +..[strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 16384 tokens in 0 s (779070.6 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] done: 28001 tokens in 0 s (803359.3 tok/s) (stop, cancel=False) +..[strata] the thinking repeated the same passages (coverage=1.000) at 2560 tokens: ended as "length" (reasoning_loop_recovery: "stop" in strata-.json; remove it to turn this off) +[strata] done: 2560 tokens in 0 s (754933.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] the thinking repeated the same passages (coverage=1.000) at 2048 tokens: ended as "length" (reasoning_loop_recovery: "stop" in strata-.json; remove it to turn this off) +[strata] done: 2048 tokens in 0 s (763889.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +..................[strata] done: 116 tokens in 0 s (210988.4 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (165489.5 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (220552.7 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (166224.6 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (220053.9 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (156393.2 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +.[strata] done: 133 tokens in 0 s (163590.2 tok/s) (stop, cancel=False) +[strata] done: 133 tokens in 0 s (96179.7 tok/s) (stop, cancel=False) +[strata] done: 133 tokens in 0 s (175256.8 tok/s) (stop, cancel=False) +[strata] done: 133 tokens in 0 s (105372.6 tok/s) (stop, cancel=False) +[strata] done: 133 tokens in 0 s (175201.8 tok/s) (stop, cancel=False) +[strata] done: 133 tokens in 0 s (77683.1 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (171825.5 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (100683.9 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (181955.3 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (98739.1 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (185743.2 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (104079.2 tok/s) (stop, cancel=False) +[strata] done: 150 tokens in 0 s (140717.0 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 150 tokens in 0 s (117751.4 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 150 tokens in 0 s (163541.9 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 150 tokens in 0 s (121974.7 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 150 tokens in 0 s (168220.7 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 150 tokens in 0 s (125955.1 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (184290.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (149204.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (198503.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (160462.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (212071.6 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (166172.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (202512.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (209068.6 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (218031.9 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (200813.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (215263.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (199925.1 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +.[strata] done: 134 tokens in 0 s (183014.2 tok/s) (stop, cancel=False) +[strata] done: 134 tokens in 0 s (98533.8 tok/s) (stop, cancel=False) +[strata] done: 134 tokens in 0 s (189813.1 tok/s) (stop, cancel=False) +[strata] done: 134 tokens in 0 s (106607.9 tok/s) (stop, cancel=False) +[strata] done: 134 tokens in 0 s (178651.2 tok/s) (stop, cancel=False) +[strata] done: 134 tokens in 0 s (98828.3 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (164189.8 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (108604.9 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (186174.2 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (111516.3 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (179205.8 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (103425.9 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (162450.5 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (128036.6 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (186128.3 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (140253.5 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (183807.8 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (131213.4 tok/s) (stop, cancel=False) +.....[strata] done: 121 tokens in 0 s (209282.8 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 121 tokens in 0 s (161730.7 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 121 tokens in 0 s (213778.8 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 121 tokens in 0 s (159394.1 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 121 tokens in 0 s (141171.3 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +.............[strata] done: 119 tokens in 0 s (170756.8 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (134752.2 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (188918.3 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (137234.6 tok/s) (length, cancel=False) +.[strata] done: 120 tokens in 0 s (189003.6 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (144589.6 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (203032.1 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (150738.7 tok/s) (stop, cancel=False) +.[strata] done: 147 tokens in 0 s (149723.8 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (121394.5 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (161912.5 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (136861.9 tok/s) (stop, cancel=False) +........[strata] reading the prompt: 512 of 523 tokens, 1 s so far +[strata] answering: 8 of max 64 tokens, 8.0 tok/s, 2 s +[strata] done: 10 tokens in 2 s (7.1 tok/s) (stop, cancel=False) +[strata] reading the prompt: 480 of 581 tokens, 1 s so far +[strata] done: 10 tokens in 2 s (11.2 tok/s) (stop, cancel=False) +.[strata] answering: 9 of max 64 tokens, 8.2 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.1 tok/s) (stop, cancel=False) +[strata] answering: 8 of max 64 tokens, 8.4 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.5 tok/s) (stop, cancel=False) +[strata] answering: 8 of max 64 tokens, 8.0 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.9 tok/s) (stop, cancel=False) +.[strata] answering: 8 of max 64 tokens, 7.8 tok/s, 1 s +[strata] answering: 9 of max 64 tokens, 8.4 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.7 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 1 s (7.3 tok/s) (stop, cancel=False) +..[strata] done: 10 tokens in 1 s (8.7 tok/s) (stop, cancel=False) +[strata] reading the prompt: 192 of 355 tokens, 1 s so far +[strata] answering: 1 of max 64 tokens, 2.3 tok/s, 1 s +[strata] answering: 3 of max 64 tokens, 4.4 tok/s, 1 s +[strata] reading the prompt: 32 of 355 tokens, 1 s so far +[strata] reading the prompt: 32 of 355 tokens, 2 s so far +[strata] answering: 5 of max 64 tokens, 6.4 tok/s, 2 s +[strata] reading the prompt: 32 of 355 tokens, 2 s so far +[strata] answering: 9 of max 64 tokens, 4.9 tok/s, 2 s +[strata] done: 10 tokens in 2 s (5.4 tok/s) (stop, cancel=False) +[strata] answering: 5 of max 64 tokens, 2.9 tok/s, 2 s +[strata] answering: 7 of max 64 tokens, 14.8 tok/s, 2 s +[strata] done: 10 tokens in 3 s (8.2 tok/s) (stop, cancel=False) +[strata] answering: 3 of max 64 tokens, 5.0 tok/s, 3 s +[strata] done: 10 tokens in 3 s (12.4 tok/s) (stop, cancel=False) +[strata] answering: 4 of max 64 tokens, 11.1 tok/s, 3 s +[strata] done: 10 tokens in 3 s (4.0 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 3 s (6.6 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 4 s (10.8 tok/s) (stop, cancel=False) +........[strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (163566.2 tok/s) (stop, cancel=False), expert cache 88.9% hit +.[strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (98855.6 tok/s) (stop, cancel=False), expert cache 88.9% hit +.[strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 98 tokens in 0 s (80834.2 tok/s) (cancel, cancel=True), expert cache 88.9% hit +.[strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] the continuation ended without a DONE line. The next request starts the engine again. +[strata] done: 95 tokens in 0 s (87554.1 tok/s) (error, cancel=False), expert cache 90.0% hit +.[strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (155808.3 tok/s) (stop, cancel=False), expert cache 88.9% hit +.[strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (90617.7 tok/s) (stop, cancel=False), expert cache 88.9% hit +.[strata] thinking capped at 200 of max_tokens 400 to leave room for the answer (budget 10000) +[strata] done: 211 tokens in 0 s (104141.9 tok/s) (stop, cancel=False), expert cache 90.0% hit +....[strata] done: 8 tokens in 0 s (189573.1 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +...[strata] done: 8 tokens in 0 s (266305.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 8 tokens in 0 s (277309.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 8 tokens in 0 s (182361.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] done: 8 tokens in 0 s (277309.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +...[strata] done: 8 tokens in 0 s (266305.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +...............s......[strata] 400 invalid request: Expecting property name enclosed in double quotes: line 1 column 2 (char 1) +.....................[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +.[strata] the engine reported an error: refused: images are not enabled +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] 400 invalid request: refused: images are not enabled +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] done: 10 tokens in 1 s (9.0 tok/s) (stop, cancel=False) +.[strata] the engine reported an error: refused: images are not enabled +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] 400 invalid request: refused: images are not enabled +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] done: 10 tokens in 1 s (8.5 tok/s) (stop, cancel=False) +..[strata] the engine stopped unexpectedly (exit code -9). The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 10 tokens in 1 s (10.0 tok/s) (stop, cancel=False) +.[strata] the engine stopped unexpectedly (exit code -9). The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine stopped unexpectedly (exit code -9). The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 2 tokens in 1 s (4.4 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] answering: 9 of max 64 tokens, 8.7 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.9 tok/s) (stop, cancel=False) +.[strata] parallel requests: 4 asked, the engine runs one at a time (its log says why) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] parallel requests: 4 asked, the engine runs one at a time (its log says why) +[strata] the engine is running again +[strata] answering: 8 of max 64 tokens, 7.8 tok/s, 1 s +[strata] done: 10 tokens in 2 s (6.6 tok/s) (stop, cancel=False) +.[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 10 tokens in 1 s (9.7 tok/s) (stop, cancel=False) +.[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 10 tokens in 1 s (9.1 tok/s) (stop, cancel=False) +.[strata] answering: 2 of max 64 tokens, 7.8 tok/s, 1 s +[strata] answering: 2 of max 64 tokens, 3.9 tok/s, 1 s +[strata] done: 10 tokens in 2 s (6.2 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 2 s (7.4 tok/s) (stop, cancel=False) +[strata] reading the prompt: 32 of 48 tokens, 2 s so far +[strata] answering: 9 of max 64 tokens, 8.7 tok/s, 3 s +[strata] done: 10 tokens in 3 s (8.6 tok/s) (stop, cancel=False) +.[strata] answering: 2 of max 64 tokens, 3.5 tok/s, 1 s +[strata] answering: 2 of max 64 tokens, 6.6 tok/s, 1 s +[strata] answering: 2 of max 64 tokens, 5.2 tok/s, 1 s +[strata] done: 10 tokens in 2 s (7.2 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 2 s (7.6 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 2 s (6.4 tok/s) (stop, cancel=False) +....................................[strata] model unloaded; the next request loads it again +..[strata] refused a request for Host 'rebind.example.com' from 127.0.0.1: not a name this server answers to (add it to "allowed_hosts" in the config or STRATA_ALLOWED_HOSTS, or set an API key) +....[strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +...........[strata] done: 10 tokens in 0 s (169125.2 tok/s) (length, cancel=False) +[strata] done: 10 tokens in 0 s (248183.7 tok/s) (length, cancel=False) +[strata] done: 10 tokens in 0 s (272357.4 tok/s) (length, cancel=False) +............[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +.[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +.[strata] 400 invalid request: tools[0] (f): "parameters" must be an object (the JSON schema of its parameters), not str +[strata] 400 invalid request: tools[0] (f): "parameters" must be an object (the JSON schema of its parameters), not list +[strata] 400 invalid request: tools[0] (f): "parameters" must be an object (the JSON schema of its parameters), not int +[strata] 400 invalid request: tools[0] (f): "input_schema" must be an object (the JSON schema of its parameters), not str +[strata] 400 invalid request: tools[0] (f): "input_schema" must be an object (the JSON schema of its parameters), not list +[strata] 400 invalid request: tools[0] (f): "input_schema" must be an object (the JSON schema of its parameters), not bool +[strata] done: 8 tokens in 0 s (268435.5 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (360800.3 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (399457.5 tok/s) (length, cancel=False) +.[strata] done: 12 tokens in 0 s (384211.1 tok/s) (stop, cancel=False) +..[strata] 400 invalid request: tools must be a list of tool objects, each {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools must be a list of tool objects, each {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools must be a list of tool objects, each {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools must be a list of tool objects, each {"name": ..., "input_schema": {...}} +[strata] 400 invalid request: tools must be a list of tool objects, each {"name": ..., "input_schema": {...}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"name": ..., "input_schema": {...}} +[strata] 400 invalid request: tools must be a list of tool objects, each {"name": ..., "input_schema": {...}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"name": ..., "input_schema": {...}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"name": ..., "input_schema": {...}} +[strata] done: 8 tokens in 0 s (419430.4 tok/s) (length, cancel=False) +.[strata] done: 12 tokens in 0 s (457560.4 tok/s) (stop, cancel=False) +[strata] 400 invalid request: messages must be a list of objects +[strata] done: 12 tokens in 0 s (461758.2 tok/s) (stop, cancel=False) +[strata] 400 invalid request: messages must be a list of objects +..[strata] 400 invalid request: No user query found in messages. +[strata] done: 12 tokens in 0 s (399457.5 tok/s) (stop, cancel=False) +.[strata] done: 12 tokens in 0 s (441505.7 tok/s) (stop, cancel=False) +.[strata] done: 12 tokens in 0 s (474826.9 tok/s) (stop, cancel=False) +[strata] done: 12 tokens in 0 s (518883.0 tok/s) (stop, cancel=False) +[strata] done: 12 tokens in 0 s (441505.7 tok/s) (stop, cancel=False) +...[strata] done: 8 tokens in 0 s (364722.1 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (399457.5 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (447392.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (466033.8 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (466033.8 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (399457.5 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (466033.8 tok/s) (length, cancel=False) +........[strata] done: 13 tokens in 0 s (302922.0 tok/s) (stop, cancel=False) +.[strata] done: 13 tokens in 0 s (245612.4 tok/s) (stop, cancel=False) +....[strata] done: 10 tokens in 0 s (169125.2 tok/s) (length, cancel=False) +[strata] done: 10 tokens in 0 s (233016.9 tok/s) (length, cancel=False) +[strata] done: 10 tokens in 0 s (235635.1 tok/s) (length, cancel=False), expert cache 90.0% hit +.........[strata] the engine reported an error: verify: layer 31 never rang (an illegal memory access was encountered) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] the engine reported an error: verify: layer 31 never rang (an illegal memory access was encountered) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +.[strata] the engine stopped unexpectedly (exit code -9). The next request starts the engine again. +[strata] done: 5 tokens in 0 s (105916.8 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code None); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 50 tokens in 0 s (239400.9 tok/s) (length, cancel=False) +.[strata] done: 50 tokens in 0 s (251155.9 tok/s) (length, cancel=False) +.[strata] done: 1 tokens in 0 s (123361.9 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (4088 tokens) leaves no room to answer in the context (4096); requests are never truncated +.[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19181.8 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19088.0 tok/s) (stop, cancel=False) +.[strata] done: 5 tokens in 0 s (179243.8 tok/s) (length, cancel=False) +.[strata] done: 120 tokens in 0 s (131586.0 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (107799.6 tok/s) (stop, cancel=False) +.[strata] done: 272 tokens in 0 s (95805.4 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (78070.9 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (107081.9 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (85533.9 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (145845.1 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (94238.4 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (147436.1 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (95485.4 tok/s) (stop, cancel=False) +[strata] tool_choice {"type": "tool", "name": "nope"} is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (147640.7 tok/s) (stop, cancel=False) +[strata] tool_choice {"type": "tool", "name": "nope"} is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (92659.0 tok/s) (stop, cancel=False) +.[strata] done: 203 tokens in 0 s (118305.4 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (121843.7 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (134530.5 tok/s) (stop, cancel=False) +.[strata] done: 272 tokens in 0 s (91678.8 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (71276.4 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (97868.3 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (74800.1 tok/s) (stop, cancel=False) +.[strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 181 tokens in 0 s (113072.5 tok/s) (stop, cancel=False) +.[strata] done: 272 tokens in 0 s (92804.9 tok/s) (stop, cancel=False) +.[strata] tool_choice {"type": "function", "function": {"name": "nope"}} is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (128558.6 tok/s) (stop, cancel=False) +[strata] tool_choice {"type": "function"} is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (133727.6 tok/s) (stop, cancel=False) +[strata] tool_choice "required" is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (143051.7 tok/s) (stop, cancel=False) +[strata] tool_choice {"type": "banana"} is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (140386.4 tok/s) (stop, cancel=False) +[strata] tool_choice "sometimes" is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (133121.3 tok/s) (stop, cancel=False) +.[strata] done: 64 tokens in 0 s (240533.6 tok/s) (stop, cancel=False) +[strata] done: 64 tokens in 0 s (187193.5 tok/s) (stop, cancel=False) +[strata] done: 56 tokens in 0 s (280287.6 tok/s) (stop, cancel=False) +[strata] done: 56 tokens in 0 s (218697.4 tok/s) (stop, cancel=False) +.[strata] layer split: card order 2,0 (the faster card last; "gpu_order": "as_given" keeps 0,2, #1352) +[strata] layer split: card order 2,0 (the faster card last; "gpu_order": "as_given" keeps 0,2, #1352) +[strata] layer split: card order 2,0 (the faster card last; "gpu_order": "as_given" keeps 0,2, #1352) +[strata] layer split: card order 0,2,1 (the faster card last; "gpu_order": "as_given" keeps 0,1,2, #1352) +[strata] layer split: keeping card order 0,1 - the faster card has less VRAM than the one the split ends on (12288 vs 22528 MiB, #1576) +[strata] layer split: card order 1,0 (the faster card last; "gpu_order": "as_given" keeps 0,1, #1352) +[strata] layer split: card order 1,0 (the faster card last; "gpu_order": "as_given" keeps 0,1, #1352) +[strata] layer split: card order 1,0 (the faster card last; "gpu_order": "as_given" keeps 0,1, #1352) +[strata] layer split: keeping card order 0,1,2 - the faster card has less VRAM than the one the split ends on (12288 vs 24576 MiB, #1576) +.../Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:2795: ResourceWarning: unclosed file <_io.TextIOWrapper name='/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpyedlx3cc/0/properties' mode='r' encoding='UTF-8'> + try: +ResourceWarning: Enable tracemalloc to get the object allocation traceback +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:2795: ResourceWarning: unclosed file <_io.TextIOWrapper name='/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpyedlx3cc/1/properties' mode='r' encoding='UTF-8'> + try: +ResourceWarning: Enable tracemalloc to get the object allocation traceback +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:2795: ResourceWarning: unclosed file <_io.TextIOWrapper name='/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpyedlx3cc/2/properties' mode='r' encoding='UTF-8'> + try: +ResourceWarning: Enable tracemalloc to get the object allocation traceback +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:2795: ResourceWarning: unclosed file <_io.TextIOWrapper name='/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpyedlx3cc/3/properties' mode='r' encoding='UTF-8'> + try: +ResourceWarning: Enable tracemalloc to get the object allocation traceback +..[strata] done: 54 tokens in 0 s (182508.0 tok/s) (stop, cancel=False) +.[strata] done: 3 tokens in 0 s (151601.3 tok/s) (stop, cancel=False) +..strata serve: AMD: HIP_PATH=Z:\definitely\not\here does not exist: not passed to the engine +.strata serve: AMD: HIP_PATH=/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpoaoya9dn/TheRock/build does not exist: not passed to the engine +strata serve: AMD: HIP_DEVICE_LIB_PATH=/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpoaoya9dn/TheRock/build does not exist: not passed to the engine +.strata serve: AMD: TEMP=/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpukk3nz26/no/such/dir is not writable: using /private/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpukk3nz26/tmp +.[strata] a picture a tool returned was left out of the prompt: this server has no image encoder +[strata] a picture a tool returned was left out of the prompt: this server has no image encoder +[strata] a picture a tool returned was left out of the prompt: this image format needs Pillow +[strata] a picture a tool returned was left out of the prompt: the image URL could not be read +....[strata] refused an API request from the web page 'https://evil.example' (no API key; add its host to "allowed_hosts" or its origin to "cors_origins" in the config) +[strata] 400 invalid request: a web page of another origin cannot have a file on this computer read as an image: send it as a data: URL (or add the page's origin to the config's trusted_origins) +[strata] 400 invalid request: a web page of another origin cannot have a file on this computer read as an image: send it as a data: URL (or add the page's origin to the config's trusted_origins) +[strata] 400 invalid request: a web page of another origin cannot have a file on this computer read as an image: send it as a data: URL (or add the page's origin to the config's trusted_origins) +[strata] 400 invalid request: a web page of another origin cannot have a file on this computer read as an image: send it as a data: URL (or add the page's origin to the config's trusted_origins) +[strata] done: 8 tokens in 0 s (250406.2 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (419430.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (447392.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (360800.3 tok/s) (length, cancel=False) +../opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +.....[strata] done: 8 tokens in 0 s (217885.9 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +....[strata] done: 32 tokens in 0 s (329773.3 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (362750.6 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (381300.4 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (371794.3 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (385683.1 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (344148.0 tok/s) (stop, cancel=False) +............................[strata] the engine stopped unexpectedly (exit code -9). The next request starts the engine again. +[strata] done: 3 tokens in 0 s (38956.4 tok/s) (error, cancel=False) +..[strata] thinking: 8 of max 30 tokens, 8.4 tok/s, 1 s +[strata] thinking: 16 of max 30 tokens, 7.9 tok/s, 2 s +[strata] thinking: 24 of max 30 tokens, 7.6 tok/s, 3 s +[strata] done: 30 tokens in 4 s (7.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] the engine said nothing for 1 s during the request; the server ended the engine. The engine and the server lost step (issue #481; a very slow PC can raise "engine_silence_s" in the config, 0 = wait forever). If it happens again, please add the end of the engine log to github.com/Niko1221/Strata/issues/481. The next request starts the engine again. +[strata] done: 1 tokens in 1 s (2.0 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 2 tokens in 0 s (24105.2 tok/s) (length, cancel=False) +.[strata] the engine said nothing for 1 s during the request; the server ended the engine. The engine and the server lost step (issue #481; a very slow PC can raise "engine_silence_s" in the config, 0 = wait forever). If it happens again, please add the end of the engine log to github.com/Niko1221/Strata/issues/481. The next request starts the engine again. +[strata] done: 1 tokens in 1 s (1.7 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 2 tokens in 0 s (23497.5 tok/s) (length, cancel=False) +.[strata] the engine did not finish the request after it was stopped (STOP) within 0 s; the server ended the engine. The engine and the server lost step (issue #481; a very slow PC can raise "engine_silence_s" in the config, 0 = wait forever). If it happens again, please add the end of the engine log to github.com/Niko1221/Strata/issues/481. The next request starts the engine again. +[strata] done: 2 tokens in 1 s (3.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 2 tokens in 0 s (23237.1 tok/s) (length, cancel=False) +......[strata] done: 8 tokens in 0 s (250406.2 tok/s) (length, cancel=False) +..[strata] done: 50 tokens in 0 s (314415.6 tok/s) (length, cancel=False) +[strata] done: 50 tokens in 0 s (337705.6 tok/s) (length, cancel=False) +[strata] done: 50 tokens in 0 s (335544.3 tok/s) (length, cancel=False) +[strata] done: 50 tokens in 0 s (324636.5 tok/s) (length, cancel=False) +[strata] thinking limited to 988 of max_tokens 1500 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 1500 tokens in 0 s (25110.0 tok/s) (length, cancel=False) +[strata] thinking limited to 988 of max_tokens 1500 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 1500 tokens in 0 s (25104.7 tok/s) (length, cancel=False) +.[strata] 400 invalid request: request exceeds the context window: prompt (246 tokens) + max tokens (4096) exceeds the context (4096); requests are never truncated. Send a smaller max_tokens (at most 3842 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (246 tokens) + max tokens (4096) exceeds the context (4096); requests are never truncated. Send a smaller max_tokens (at most 3842 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +..[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19029.8 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19109.8 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19067.6 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19173.6 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19211.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19180.5 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19077.4 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (18932.4 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19069.5 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19247.0 tok/s) (stop, cancel=False) +.[strata] done: 1 tokens in 0 s (127100.1 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 1 tokens in 0 s (144631.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 5 tokens in 0 s (381300.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (4088 tokens) leaves no room to answer in the context (4096); requests are never truncated +[strata] done: 1 tokens in 0 s (139810.1 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 5 tokens in 0 s (381300.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (4088 tokens) leaves no room to answer in the context (4096); requests are never truncated +..[strata] done: 8 tokens in 0 s (250406.2 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (335544.3 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (319566.0 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (399457.5 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (404270.3 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (424739.6 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (424739.6 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (419430.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (381300.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (424739.6 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (399457.5 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (381300.4 tok/s) (length, cancel=False) +.........[strata] the config has a "vision" section but no --vision in its args: added --vision and --vram-reserve-mib 700 (run setup again with --vision to write them) +[strata] the config has a "vision" section but no --vision in its args: added --vision (run setup again with --vision to write them) +[strata] the config has a "vision" section but no --vision in its args: added --vision (run setup again with --vision to write them) +........[strata] done: 10 tokens in 0 s (192399.3 tok/s) (length, cancel=False) +.[strata] done: 10 tokens in 0 s (181571.6 tok/s) (length, cancel=False) +[strata] done: 10 tokens in 0 s (238312.7 tok/s) (length, cancel=False) +.[strata] done: 40 tokens in 0 s (95271.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] done: 40 tokens in 0 s (101250.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +....[strata] done: 18 tokens in 0 s (236669.2 tok/s) (stop, cancel=False) +.[strata] done: 28 tokens in 0 s (274393.7 tok/s) (stop, cancel=False) +.[strata] done: 36 tokens in 0 s (277054.9 tok/s) (stop, cancel=False) +.[strata] the reply ended inside its thinking with no answer: closing the thinking once and continuing (reasoning_close_retry) +[strata] done: 64 tokens in 0 s (226146.1 tok/s) (stop, cancel=False) +....[strata] done: 13 tokens in 0 s (245612.4 tok/s) (stop, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/unittest/case.py:615: ResourceWarning: unclosed + result = method() +ResourceWarning: Enable tracemalloc to get the object allocation traceback +..[strata] the engine did not start (try 1 of 3); again in 0 s +............[strata] other apps now use the Chat settings: temperature=0.3 +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 16 tokens in 0 s (285569.6 tok/s) (stop, cancel=False) +.[strata] other apps use their own settings again +.[strata] other apps now use the Chat settings: temperature=0.3 +[strata] other apps use their own settings again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 16 tokens in 0 s (370767.2 tok/s) (stop, cancel=False) +.[strata] other apps now use the Chat settings: temperature=1.0 +.[strata] other apps now use the Chat settings: temperature=0.3, top_p=0.9, top_k=10, seed=7, max_tokens=77, reasoning_effort=low, experimental_speed_projection=False +[strata] done: 16 tokens in 0 s (411710.8 tok/s) (stop, cancel=False) +[strata] done: 5 tokens in 0 s (355449.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +..[strata] other apps now use the Chat settings: temperature=1.0 +.[strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(1) +[strata] the before_load command failed (it exited with code 1); loading anyway +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (265980.3 tok/s) (stop, cancel=False) +.[strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c open('/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpa8qxdouz/ran', 'w').close() +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (259647.4 tok/s) (stop, cancel=False) +.[strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(7) +[strata] the before_load command failed (it exited with code 7); the model stays unloaded and this request is answered 503 - try again (before_load_strict) +.[strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import time; time.sleep(30) +[strata] the before_load command failed (it did not finish in 0.5 s (before_load_timeout_s)); the model stays unloaded and this request is answered 503 - try again (before_load_strict) +.[strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(7) +[strata] the before_load command failed (it exited with code 7); checking the free VRAM +[strata] the before_load command failed (it exited with code 7), but 20000 MiB of VRAM is free: loading +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (260889.7 tok/s) (stop, cancel=False) +.[strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(3) +[strata] the before_load command failed (it exited with code 3); checking the free VRAM +[strata] the before_load command failed (it exited with code 3) and only 2000 MiB of VRAM is free (the model needs 8000): not loading; 503 +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(3) +[strata] the before_load command failed (it exited with code 3); checking the free VRAM +[strata] the before_load command failed (it exited with code 3), but 20000 MiB of VRAM is free: loading +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (209715.2 tok/s) (stop, cancel=False) +.[strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(7) +[strata] the before_load command failed (it exited with code 7); loading anyway (set min_free_vram_mib to have the free VRAM checked, or before_load_strict to refuse) +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (250119.0 tok/s) (stop, cancel=False) +.[strata] the model unloads after 1 s without requests +[strata] model unloaded after 1 s idle; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (109270.4 tok/s) (stop, cancel=False) +.[strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +.[strata] model unloaded after 1 s idle; the next request loads it again +[strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (152733.8 tok/s) (stop, cancel=False) +.[strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (250119.0 tok/s) (stop, cancel=False) +...[strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (248976.9 tok/s) (stop, cancel=False) +..[strata] model unloaded; the next request loads it again +.[strata] model unloaded; the next request loads it again +[strata] starting the vision encoder ... +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (290031.7 tok/s) (stop, cancel=False) +.[strata] the engine has said nothing for 1 s but is still working; it is not ended (the limit is engine_silence_s). If it is stuck, "engine_silence_s" can be set lower. +...[strata] the engine has said nothing for 1 s but is still working; it is not ended (the limit is engine_silence_s). If it is stuck, "engine_silence_s" can be set lower. +...................[strata] done: 15 tokens in 0 s (57247.1 tok/s) (disconnect, cancel=False) +[strata] done: 20 tokens in 0 s (157090.0 tok/s) (length, cancel=False) +..[strata] done: 36 tokens in 0 s (216635.5 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (125411.1 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (262599.9 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (126780.0 tok/s) (stop, cancel=False) +.[strata] 400 invalid request: stop takes at most 4 strings +[strata] 400 invalid request: stop must be a string or a list of strings +[strata] 400 invalid request: stop_sequences must be a list of strings +[strata] 400 invalid request: stop_sequences must be a list of strings +.[strata] done: 19 tokens in 0 s (322638.8 tok/s) (stop, cancel=False) +[strata] done: 19 tokens in 0 s (197746.3 tok/s) (stop, cancel=False) +[strata] done: 19 tokens in 0 s (351065.1 tok/s) (stop, cancel=False) +[strata] done: 19 tokens in 0 s (197746.3 tok/s) (stop, cancel=False) +..[strata] done: 36 tokens in 0 s (346318.7 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (342392.2 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (313267.5 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (130055.9 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (154549.6 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (153919.4 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (336291.6 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (329683.3 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (294912.0 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (121966.8 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (126355.6 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (115793.7 tok/s) (stop, cancel=False) +.[strata] done: 29 tokens in 0 s (358804.8 tok/s) (stop, cancel=False) +.[strata] done: 24 tokens in 0 s (363405.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (51150.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (356962.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (135664.8 tok/s) (stop, cancel=False) +...[strata] 400 invalid request: reasoning_budget_tokens='lots': expected a whole number of tokens (0: no budget) +[strata] 400 invalid request: reasoning_budget_tokens=2.5: expected a whole number of tokens (0: no budget) +[strata] 400 invalid request: reasoning_budget_tokens=True: expected a whole number of tokens (0: no budget) +[strata] 400 invalid request: reasoning_budget_tokens=[1]: expected a whole number of tokens (0: no budget) +..[strata] thinking capped at 200 of max_tokens 400 to leave room for the answer (budget 10000) +[strata] done: 211 tokens in 0 s (118807.6 tok/s) (stop, cancel=False) +..[strata] thinking limited to 512 of max_tokens 1024 so the answer has room ("answer_room": false in the config turns this off) +[strata] thinking budget reached (512 tokens): wrapping up the thinking +[strata] done: 604 tokens in 0 s (54581.8 tok/s) (stop, cancel=False) +...[strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (166111.0 tok/s) (stop, cancel=False) +[strata] done: 211 tokens in 0 s (120016.0 tok/s) (stop, cancel=False) +.[strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 108 tokens in 0 s (149796.6 tok/s) (stop, cancel=False) +.[strata] thinking budget reached (30 tokens): wrapping up the thinking +[strata] done: 122 tokens in 0 s (84888.0 tok/s) (stop, cancel=False) +.[strata] thinking budget reached (30 tokens): wrapping up the thinking +[strata] done: 122 tokens in 0 s (160158.1 tok/s) (stop, cancel=False) +.[strata] done: 20 tokens in 0 s (188508.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] done: 211 tokens in 0 s (117420.5 tok/s) (stop, cancel=False) +.[strata] done: 300 tokens in 0 s (92081.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 1024 tokens in 0 s (36032.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 17 tokens in 0 s (396128.7 tok/s) (stop, cancel=False) +.[strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (177805.5 tok/s) (stop, cancel=False) +.[strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (94901.4 tok/s) (stop, cancel=False) +[strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (161652.5 tok/s) (stop, cancel=False) +.[strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (178075.1 tok/s) (stop, cancel=False) +[strata] done: 211 tokens in 0 s (125265.1 tok/s) (stop, cancel=False) +..[strata] done: 17 tokens in 0 s (270087.8 tok/s) (stop, cancel=False) +......................[strata] done: 119 tokens in 0 s (183974.3 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (122423.9 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (200611.8 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (130762.9 tok/s) (length, cancel=False) +.[strata] done: 120 tokens in 0 s (162569.9 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (122193.9 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (186068.9 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (145467.2 tok/s) (stop, cancel=False) +.[strata] done: 147 tokens in 0 s (148641.0 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (122382.4 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (158948.9 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (133138.1 tok/s) (stop, cancel=False) +.......[strata] done: 13 tokens in 0 s (182973.0 tok/s) (stop, cancel=False), expert cache 90.0% hit +...[strata] done: 5 tokens in 0 s (227951.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] done: 13 tokens in 0 s (42900.0 tok/s) (stop, cancel=False), expert cache 90.0% hit +.[strata] done: 13 tokens in 0 s (24728.3 tok/s) (stop, cancel=False), expert cache 90.0% hit +..[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (39626.4 tok/s) (stop, cancel=False), expert cache 90.0% hit +.[strata] done: 13 tokens in 0 s (90876.6 tok/s) (stop, cancel=False), expert cache 90.0% hit +.........../opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +./opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +.[strata] done: 5 tokens in 0 s (220752.8 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +../opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +..........[strata] SAVE waiting for RAM after cache reclaim: need=1125 MiB, limit=30s; floor unchanged +.[strata] SAVE waiting for RAM after cache reclaim: need=1125 MiB, limit=30s; floor unchanged +................................................[strata] done: 255 tokens in 0 s (107362.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 10991 of max_tokens 14654 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 282 tokens in 0 s (104376.4 tok/s) (stop, cancel=False) +.ss[strata] 400 invalid request: response_format must be an object +[strata] 400 invalid request: response_format.type must be text, json_object or json_schema +[strata] 400 invalid request: response_format.json_schema needs a name +[strata] 400 invalid request: response_format.json_schema.strict must be boolean +[strata] 400 invalid request: structured response_format with tools/MCP is not supported +[strata] 400 invalid request: structured response_format with tools/MCP is not supported +[strata] 400 invalid request: send a JSON object +.[strata] done: 13 tokens in 0 s (219862.7 tok/s) (stop, cancel=False) +[strata] done: 3 tokens in 0 s (187804.7 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 0 s (313007.8 tok/s) (stop, cancel=False) +[strata] done: 14 tokens in 0 s (331752.9 tok/s) (stop, cancel=False) +[strata] done: 12 tokens in 0 s (362098.2 tok/s) (stop, cancel=False) +[strata] done: 15 tokens in 0 s (332881.3 tok/s) (stop, cancel=False) +[strata] done: 6 tokens in 0 s (314572.8 tok/s) (stop, cancel=False) +[strata] done: 15 tokens in 0 s (332881.3 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (326563.8 tok/s) (stop, cancel=False) +[strata] done: 5 tokens in 0 s (265462.3 tok/s) (length, cancel=False) +[strata] done: 7 tokens in 0 s (319131.8 tok/s) (stop, cancel=False) +.[strata] 400 invalid request: response_format.json_schema needs a schema object +[strata] done: 21 tokens in 0 s (229974.9 tok/s) (stop, cancel=False) +.[strata] done: 40 tokens in 0 s (256532.4 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (271853.0 tok/s) (stop, cancel=False) +[strata] done: 40 tokens in 0 s (327041.2 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (333637.8 tok/s) (stop, cancel=False) +[strata] done: 40 tokens in 0 s (325771.2 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (315700.3 tok/s) (stop, cancel=False) +[strata] done: 40 tokens in 0 s (336216.8 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (349525.3 tok/s) (stop, cancel=False) +[strata] done: 40 tokens in 0 s (93466.4 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (349525.3 tok/s) (stop, cancel=False) +.[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +..[strata] done: 255 tokens in 0 s (98612.2 tok/s) (stop, cancel=False) +[strata] done: 26 tokens in 0 s (361099.0 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (279620.3 tok/s) (stop, cancel=False) +[strata] done: 13 tokens in 0 s (370924.8 tok/s) (stop, cancel=False) +.......[strata] VRAM: 4000 MiB free, expert cache 2000 of 6000 MiB (100 experts) +[strata] images: the expert cache keeps 2048 MiB free for the image encoder (500 of 6000 MiB left) +[strata] images: the image encoder started in 0.0 s, 0 MiB of VRAM +[strata] VRAM: 3048 MiB free, expert cache 2952 of 6000 MiB (952 experts) +[strata] images: the image encoder stopped after 60 s without a picture; the expert cache took its VRAM back (5000 of 6000 MiB) +.[strata] images: the expert cache did not make room for the image encoder (VRAM needs an engine started with --vram-elastic (one NVIDIA GPU, --serve)); starting it anyway +[strata] images: the image encoder started in 0.0 s, 0 MiB of VRAM +.[strata] images: the expert cache keeps 2048 MiB free for the image encoder (3952 of 6000 MiB left) +[strata] images: the image encoder started in 0.0 s, 0 MiB of VRAM +[strata] images: the image encoder stopped after 60 s without a picture; the expert cache took its VRAM back (6000 of 6000 MiB) +[strata] images: the expert cache keeps 2048 MiB free for the image encoder (3952 of 6000 MiB left) +[strata] images: the image encoder started in 0.0 s, 0 MiB of VRAM +.[strata] images: the expert cache keeps 2048 MiB free for the image encoder (3952 of 6000 MiB left) +[strata] images: the image encoder started in 0.0 s, 0 MiB of VRAM +[strata] model unloaded; the next request loads it again +..[strata] images: the expert cache keeps 2048 MiB free for the image encoder (3952 of 6000 MiB left) +[strata] images: the image encoder started in 0.0 s, 0 MiB of VRAM +.....[strata] 400 invalid request: VRAM needs an engine started with --vram-elastic (one NVIDIA GPU, --serve) +.[strata] 400 invalid request: reserve_mib: a whole number of MiB (0 or more), or null for the start's +[strata] 400 invalid request: reserve_mib: a whole number of MiB (0 or more), or null for the start's +[strata] 400 invalid request: reserve_mib: a whole number of MiB (0 or more), or null for the start's +[strata] 400 invalid request: reserve_mib: a whole number of MiB (0 or more), or null for the start's +[strata] refused an API request from the web page 'https://other.example' (no API key; add its host to "allowed_hosts" or its origin to "cors_origins" in the config) +.[strata] VRAM: 4000 MiB free, expert cache 2000 of 6000 MiB (100 experts) +[strata] VRAM: 700 MiB free, expert cache 5300 of 6000 MiB (3300 experts) +.[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +.....sss +---------------------------------------------------------------------- +Ran 852 tests in 517.636s + +OK (skipped=11) diff --git a/results/round39/parent-tests.json b/results/round39/parent-tests.json new file mode 100644 index 0000000..1342f3f --- /dev/null +++ b/results/round39/parent-tests.json @@ -0,0 +1,27 @@ +[ + { + "test": "qsa-merge2-bitwise", + "exit": 0, + "seconds": 1.4664330839877948 + }, + { + "test": "qsa-prompt", + "exit": 0, + "seconds": 0.5909264170040842 + }, + { + "test": "qsa-prompt-noncoop", + "exit": 0, + "seconds": 0.4592962079914287 + }, + { + "test": "moe-default", + "exit": 0, + "seconds": 24.26980204100255 + }, + { + "test": "moe-v3", + "exit": 0, + "seconds": 22.33838270799606 + } +] diff --git a/results/round39/parent-tokenizer-real.log b/results/round39/parent-tokenizer-real.log new file mode 100644 index 0000000..b762612 --- /dev/null +++ b/results/round39/parent-tokenizer-real.log @@ -0,0 +1,5 @@ +.........s........ +---------------------------------------------------------------------- +Ran 18 tests in 0.784s + +OK (skipped=1) diff --git a/results/round39/ple-baseline-build.log b/results/round39/ple-baseline-build.log new file mode 100644 index 0000000..e69de29 diff --git a/results/round39/ple-baseline.log b/results/round39/ple-baseline.log new file mode 100644 index 0000000..7d31d3f --- /dev/null +++ b/results/round39/ple-baseline.log @@ -0,0 +1,7 @@ +FAIL vendor/strata/src/ngram/ple_reader_test.cpp:343: keep-alive: 2 reads in 300 ms at a 20 ms period +FAIL vendor/strata/src/ngram/ple_reader_test.cpp:356: the keep-alive did not start again (3 -> 5) +FAIL vendor/strata/src/ngram/ple_reader_test.cpp:343: keep-alive: 4 reads in 300 ms at a 20 ms period +FAIL vendor/strata/src/ngram/ple_reader_test.cpp:356: the keep-alive did not start again (5 -> 9) +ple formats, both readers vs their dequantizers: IQ4_NL Q5_0 F8_E4M3 Q5_1 Q8_0 Q4_0 BF16: OK +ple_reader selftest (row_bytes=90): FAILED +ple_reader selftest (row_bytes=110): FAILED diff --git a/results/round39/ple-candidate-build.log b/results/round39/ple-candidate-build.log new file mode 100644 index 0000000..e69de29 diff --git a/results/round39/ple-candidate.log b/results/round39/ple-candidate.log new file mode 100644 index 0000000..03c3adf --- /dev/null +++ b/results/round39/ple-candidate.log @@ -0,0 +1,8 @@ +FAIL vendor/strata/src/ngram/ple_reader_test.cpp:423: keep-alive: 4 reads in 300 ms at a 20 ms period +FAIL vendor/strata/src/ngram/ple_reader_test.cpp:436: the keep-alive did not start again (4 -> 7) +FAIL vendor/strata/src/ngram/ple_reader_test.cpp:423: keep-alive: 3 reads in 300 ms at a 20 ms period +FAIL vendor/strata/src/ngram/ple_reader_test.cpp:436: the keep-alive did not start again (3 -> 6) +ple prefetched-key reuse (reordered, duplicated, subset), Direct vs Mmap bytes: OK +ple formats, both readers vs their dequantizers: IQ4_NL Q5_0 F8_E4M3 Q5_1 Q8_0 Q4_0 BF16: OK +ple_reader selftest (row_bytes=90): FAILED +ple_reader selftest (row_bytes=110): FAILED diff --git a/results/round39/ple-targeted.log b/results/round39/ple-targeted.log new file mode 100644 index 0000000..a099ddc --- /dev/null +++ b/results/round39/ple-targeted.log @@ -0,0 +1 @@ +ple prefetched-key reuse (reordered, duplicated, subset), Direct vs Mmap bytes: OK diff --git a/results/round39/ple-verification.json b/results/round39/ple-verification.json new file mode 100644 index 0000000..68e570e --- /dev/null +++ b/results/round39/ple-verification.json @@ -0,0 +1,51 @@ +{ + "baseline_build": [ + "c++", + "-std=c++20", + "-O2", + "-pthread", + "-Ivendor/strata/include", + "-Ivendor/strata/third_party/ggml", + "vendor/strata/src/ngram/ple_reader_test.cpp", + "vendor/strata/src/kernels/ngram.cpp", + "vendor/strata/src/ngram/ple_reader.cpp", + "vendor/strata/src/platform/direct_file.cpp", + "vendor/strata/src/platform/memory.cpp", + "-o", + "results/round39/ple_reader_baseline" + ], + "candidate_build": [ + "c++", + "-std=c++20", + "-O2", + "-pthread", + "-Ivendor/strata/include", + "-Ivendor/strata/third_party/ggml", + "vendor/strata/src/ngram/ple_reader_test.cpp", + "vendor/strata/src/kernels/ngram.cpp", + "vendor/strata/src/ngram/ple_reader.cpp", + "vendor/strata/src/platform/direct_file.cpp", + "vendor/strata/src/platform/memory.cpp", + "-o", + "results/round39/ple_reader_candidate" + ], + "baseline_selftest_exit": 1, + "candidate_selftest_exit": 1, + "targeted_build": [ + "c++", + "-std=c++20", + "-O2", + "-pthread", + "-Ivendor/strata/include", + "-Ivendor/strata/third_party/ggml", + "results/round39/ple_reuse_target.cpp", + "vendor/strata/src/kernels/ngram.cpp", + "vendor/strata/src/ngram/ple_reader.cpp", + "vendor/strata/src/platform/direct_file.cpp", + "vendor/strata/src/platform/memory.cpp", + "-o", + "results/round39/ple_reuse_target" + ], + "targeted_exit": 0, + "limitation": "full suites fail the same pre-existing keepalive timing assertions; no thresholds changed" +} diff --git a/results/round39/ple.patch b/results/round39/ple.patch new file mode 100644 index 0000000..fcb0d6f --- /dev/null +++ b/results/round39/ple.patch @@ -0,0 +1,141 @@ +diff --git a/src/kernels/ngram.cpp b/src/kernels/ngram.cpp +index 91de3705..1fd7a3a5 100644 +--- a/src/kernels/ngram.cpp ++++ b/src/kernels/ngram.cpp +@@ -508,27 +508,31 @@ bool PleTable::gather_batch(const uint32_t* rows, size_t n_tokens, float* out, s + if (impl_->pending) { err = "PleTable::gather_batch while a token is in flight"; return false; } + const size_t n = n_tokens * (size_t) PLE_N_HEADS; + if (impl_->mode == PleIo::Direct) { +- // Fast path: if all requested tokens match our in-flight/completed prefetch slots, collect them directly ++ // Fast path: match each requested token to any in-flight/completed prefetch slot, then collect directly + // without issuing a second PleReader ticket or copying through RowCache. +- if (impl_->n_prefetch >= n_tokens && n_tokens <= Impl::kMaxPrefetch) { ++ if (n_tokens <= Impl::kMaxPrefetch) { ++ size_t matched[Impl::kMaxPrefetch]; + bool exact_match = true; + for (size_t t = 0; t < n_tokens; ++t) { +- if (std::memcmp(impl_->prefetch_keys[t], rows + t * PLE_N_HEADS, sizeof(uint32_t) * PLE_N_HEADS) != 0) { +- exact_match = false; +- break; +- } ++ size_t i = 0; ++ while (i < impl_->n_prefetch && ++ std::memcmp(impl_->prefetch_keys[i], rows + t * PLE_N_HEADS, ++ sizeof(uint32_t) * PLE_N_HEADS) != 0) ++i; ++ if (i == impl_->n_prefetch) { exact_match = false; break; } ++ matched[t] = i; // duplicate requests may share one immutable prefetch slot + } + if (exact_match) { ++ bool collect_ok = true; + for (size_t i = 0; i < impl_->n_prefetch; ++i) { +- if (!impl_->reader.collect(impl_->prefetch_tickets[i], err)) { +- impl_->n_prefetch = 0; +- return false; +- } ++ std::string ignored; ++ if (!impl_->reader.collect(impl_->prefetch_tickets[i], collect_ok ? err : ignored)) ++ collect_ok = false; // retain the first error, but drain every output buffer + } + impl_->n_prefetch = 0; ++ if (!collect_ok) return false; + for (size_t t = 0; t < n_tokens; ++t) { + for (int h = 0; h < PLE_N_HEADS; ++h) { +- impl_->decode(impl_->prefetch_raw[t] + (size_t) h * impl_->rb, ++ impl_->decode(impl_->prefetch_raw[matched[t]] + (size_t) h * impl_->rb, + out + (t * PLE_N_HEADS + (size_t) h) * PLE_HEAD_DIM); + } + } +diff --git a/src/ngram/ple_reader_test.cpp b/src/ngram/ple_reader_test.cpp +index 4663455f..670b29cc 100644 +--- a/src/ngram/ple_reader_test.cpp ++++ b/src/ngram/ple_reader_test.cpp +@@ -242,9 +242,89 @@ void bad_sizes_are_refused(const std::string& dir) { + } + } + ++// #2622 (reported by @hotschmoe): gather_batch matched the prefetched tokens by position, so a prefetch of A..F ++// gathered as D,E,F,A,B,C (or A gathered as A,A,A) issued a second ticket for rows that were already in flight. ++// Every case: the bytes equal the mapped reader's, and the requested-row counter says how many rows were re-read. ++void prefetched_keys_are_reused(const std::string& dir) { ++ const k::PleFormatInfo* fmt = nullptr; ++ for (int i = 0; i < k::ple_format_count(); ++i) ++ if (std::string(k::ple_formats()[i].name) == "IQ4_NL") fmt = &k::ple_formats()[i]; ++ if (fmt == nullptr) { CHECK(false, "no IQ4_NL format"); return; } ++ const std::string path = write_format_table(dir, *fmt, 4096); ++ if (path.empty()) { CHECK(false, "cannot write the prefetch table"); return; } ++ struct Case { const char* name; std::vector pre, want; uint64_t extra_rows; }; ++ std::vector seventeen; ++ for (int i = 0; i < 17; ++i) seventeen.push_back(i); ++ const std::vector cases = { ++ {"ordered", {0, 1}, {0, 1}, 0}, ++ {"permutation", {0, 1, 2, 3, 4, 5}, {3, 4, 5, 0, 1, 2}, 0}, ++ {"subset", {0, 1, 2, 3, 4, 5}, {4, 1}, 0}, ++ {"duplicates", {0}, {0, 0, 0}, 0}, ++ {"duplicate prefetch", {0, 0, 1}, {0, 1}, 0}, ++ {"missing key", {0, 1}, {0, 2}, 32}, ++ {"extra prefetch", {0, 1, 2}, {1}, 0}, ++ {"none", {0}, {}, 0}, ++ {"over the window", {0, 1}, seventeen, 17u * k::PLE_N_HEADS}, ++ {"window overflow", seventeen, {16}, 0}, ++ }; ++ const auto keys = [](const std::vector& ids) { ++ std::vector r; ++ for (int id : ids) for (int h = 0; h < k::PLE_N_HEADS; ++h) r.push_back((uint32_t) (id * 32 + h)); ++ return r; ++ }; ++ const auto requests = [](const k::PleTable& t) { ++ unsigned long long n = 0; ++ const std::string s = t.io_report(); ++ return std::sscanf(s.c_str(), "ple io: %llu rows", &n) == 1 ? (uint64_t) n : ~0ull; ++ }; ++ const size_t tok = (size_t) k::PLE_N_HEADS * k::PLE_HEAD_DIM; ++ std::vector> refs(cases.size()); ++ { ++ k::PleTable mapped; ++ k::PleIoOptions mio; ++ mio.mode = k::PleIo::Mmap; ++ std::string err; ++ if (!mapped.open(path, err, mio)) { CHECK(false, "mapped open: %s", err.c_str()); return; } ++ for (size_t c = 0; c < cases.size(); ++c) { ++ const auto r = keys(cases[c].want); ++ refs[c].resize(cases[c].want.size() * tok); ++ CHECK(mapped.gather_batch(r.data(), cases[c].want.size(), refs[c].data(), err), "mapped gather: %s", err.c_str()); ++ } ++ } ++ for (uint64_t cache : {uint64_t(0), uint64_t(256)}) for (bool worker : {false, true}) { ++ for (size_t c = 0; c < cases.size(); ++c) { ++ k::PleTable t; ++ k::PleIoOptions io; ++ io.mode = k::PleIo::Direct; ++ io.cache_rows = cache; ++ io.io_thread = worker; ++ io.max_inflight = 8; ++ io.keepalive_ms = 0; ++ std::string err; ++ if (!t.open(path, err, io)) { CHECK(false, "direct open: %s", err.c_str()); return; } ++ const auto pr = keys(cases[c].pre), wr = keys(cases[c].want); ++ for (size_t i = 0; i < cases[c].pre.size(); ++i) t.prefetch_rows(pr.data() + i * k::PLE_N_HEADS); ++ const uint64_t rq = requests(t), by = t.bytes_read(); ++ const float sentinel = -12345.5f; ++ std::vector out(refs[c].size() + 2, sentinel); ++ const bool ok = t.gather_batch(wr.data(), cases[c].want.size(), out.data() + 1, err); ++ CHECK(ok, "%s (cache %llu worker %d): %s", cases[c].name, (unsigned long long) cache, (int) worker, err.c_str()); ++ CHECK(requests(t) - rq == cases[c].extra_rows, "%s (cache %llu worker %d): %llu rows re-requested, want %llu", ++ cases[c].name, (unsigned long long) cache, (int) worker, (unsigned long long) (requests(t) - rq), ++ (unsigned long long) cases[c].extra_rows); ++ CHECK(t.bytes_read() - by == cases[c].want.size() * k::PLE_N_HEADS * fmt->row_bytes, "%s: bytes_read", cases[c].name); ++ CHECK(refs[c].empty() || std::memcmp(out.data() + 1, refs[c].data(), refs[c].size() * sizeof(float)) == 0, ++ "%s (cache %llu worker %d): bytes differ from the mapped reader", cases[c].name, (unsigned long long) cache, (int) worker); ++ CHECK(out.front() == sentinel && out.back() == sentinel, "%s: wrote outside the output", cases[c].name); ++ } ++ } ++ if (g_fail == 0) std::printf("ple prefetched-key reuse (reordered, duplicated, subset), Direct vs Mmap bytes: OK\n"); ++} ++ + void all_formats_round_trip(const std::string& dir) { + bf16_widening_is_exact(); // BF16's own claim: a shift, exact for every pattern + bad_sizes_are_refused(dir); ++ prefetched_keys_are_reused(dir); + CHECK(k::ple_format_count() > 0, "no formats"); + for (int i = 0; i < k::ple_format_count(); ++i) { + const k::PleFormatInfo& f = k::ple_formats()[i]; diff --git a/results/round39/ple_reuse_target.cpp b/results/round39/ple_reuse_target.cpp new file mode 100644 index 0000000..192642c --- /dev/null +++ b/results/round39/ple_reuse_target.cpp @@ -0,0 +1,9 @@ +#define main ple_full_suite_main +#include "../../vendor/strata/src/ngram/ple_reader_test.cpp" +#undef main + +int main(int argc, char** argv) { + if (argc != 2) return 2; + prefetched_keys_are_reused(argv[1]); + return g_fail ? 1 : 0; +} diff --git a/results/round39/serve-baseline-full.log b/results/round39/serve-baseline-full.log new file mode 100644 index 0000000..5a71b01 --- /dev/null +++ b/results/round39/serve-baseline-full.log @@ -0,0 +1,2318 @@ +test_default_off (test_cache_drop.CacheDrop.test_default_off) ... ok +test_files_and_advise (test_cache_drop.CacheDrop.test_files_and_advise) ... ok +test_not_while_engine_runs_or_not_opted_in (test_cache_drop.CacheDrop.test_not_while_engine_runs_or_not_opted_in) ... ok +test_on (test_cache_drop.CacheDrop.test_on) ... ok +test_a_bad_chunk_size_is_a_400 (test_chunked_body.ChunkedBody.test_a_bad_chunk_size_is_a_400) ... ok +test_a_chunked_chat_request_is_read (test_chunked_body.ChunkedBody.test_a_chunked_chat_request_is_read) ... [strata] done: 12 tokens in 0 s (199728.8 tok/s) (stop, cancel=False) +ok +test_a_chunked_control_body_is_consumed (test_chunked_body.ChunkedBody.test_a_chunked_control_body_is_consumed) ... ok +test_a_chunked_settings_body_is_read_too (test_chunked_body.ChunkedBody.test_a_chunked_settings_body_is_read_too) ... [strata] other apps now use the Chat settings: temperature=0.5 +ok +test_a_content_length_request_is_unchanged (test_chunked_body.ChunkedBody.test_a_content_length_request_is_unchanged) ... [strata] done: 12 tokens in 0 s (278075.4 tok/s) (stop, cancel=False) +ok +test_a_truncated_chunked_body_is_a_400 (test_chunked_body.ChunkedBody.test_a_truncated_chunked_body_is_a_400) ... ok +test_an_overlong_chunked_trailer_does_not_change_settings (test_chunked_body.ChunkedBody.test_an_overlong_chunked_trailer_does_not_change_settings) ... ok +test_an_oversized_chunked_body_is_a_413_and_is_never_allocated (test_chunked_body.ChunkedBody.test_an_oversized_chunked_body_is_a_413_and_is_never_allocated) ... ok +test_chunked_trailer_limit_does_not_change_settings (test_chunked_body.ChunkedBody.test_chunked_trailer_limit_does_not_change_settings) ... ok +test_chunked_wins_over_content_length (test_chunked_body.ChunkedBody.test_chunked_wins_over_content_length) ... [strata] done: 12 tokens in 0 s (306900.3 tok/s) (stop, cancel=False) +ok +test_complete_settings_bodies_still_apply (test_chunked_body.ChunkedBody.test_complete_settings_bodies_still_apply) ... [strata] other apps now use the Chat settings: temperature=0.5 +[strata] other apps now use the Chat settings: temperature=0.5 +[strata] other apps now use the Chat settings: temperature=0.5 +ok +test_incomplete_chunked_trailers_do_not_change_settings (test_chunked_body.ChunkedBody.test_incomplete_chunked_trailers_do_not_change_settings) ... ok +test_both_overflow_errors_are_recognizable_before_streaming (test_context_overflow.ContextOverflow.test_both_overflow_errors_are_recognizable_before_streaming) ... [strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +ok +test_one_token_of_room_still_answers_but_zero_is_rejected (test_context_overflow.ContextOverflow.test_one_token_of_room_still_answers_but_zero_is_rejected) ... [strata] done: 1 tokens in 0 s (50533.8 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) + max tokens (1) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] done: 1 tokens in 0 s (62601.6 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] done: 1 tokens in 0 s (65536.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) + max tokens (1) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] done: 1 tokens in 0 s (62601.6 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] done: 1 tokens in 0 s (39945.8 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 1 tokens in 0 s (43690.7 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_prepare_preserves_prompt_and_budget_limits (test_context_overflow.ContextOverflow.test_prepare_preserves_prompt_and_budget_limits) ... ok +test_no_stop_without_cancel (test_control_cancel.ControlCancelTest.test_no_stop_without_cancel) ... ok +test_stop_during_prompt_read (test_control_cancel.ControlCancelTest.test_stop_during_prompt_read) ... ok +test_constant_cost (test_detok.Detok.test_constant_cost) ... skipped 'no pack tokenizer here' +test_random_ids_and_broken_bytes (test_detok.Detok.test_random_ids_and_broken_bytes) ... skipped 'no pack tokenizer here' +test_text_with_split_characters (test_detok.Detok.test_text_with_split_characters) ... skipped 'no pack tokenizer here' +test_pack_vocab (test_detok.HeapBpe.test_pack_vocab) ... skipped 'no pack tokenizer here' +test_speed_pack (test_detok.HeapBpe.test_speed_pack) ... skipped 'no pack tokenizer here' +test_speed_synthetic (test_detok.HeapBpe.test_speed_synthetic) ... ok +test_synthetic_vocab (test_detok.HeapBpe.test_synthetic_vocab) ... ok +test_words_directly (test_detok.HeapBpe.test_words_directly) ... ok +test_threads_match_serial (test_detok.ThreadedEncode.test_threads_match_serial) ... ok +test_an_unknown_level_is_a_400_not_the_default (test_effort_paths.EffortPaths.test_an_unknown_level_is_a_400_not_the_default) ... [strata] 400 invalid request: unknown reasoning effort 'huge': use none, low, medium or high +ok +test_anthropic_output_config_and_thinking (test_effort_paths.EffortPaths.test_anthropic_output_config_and_thinking) ... [strata] done: 24 tokens in 0 s (278845.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (333322.2 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (354448.2 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (348315.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (370085.6 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (381300.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (368730.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (379861.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (336666.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (351969.6 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (368730.0 tok/s) (stop, cancel=False) +ok +test_config_level_reaches_every_api (test_effort_paths.EffortPaths.test_config_level_reaches_every_api) ... [strata] done: 24 tokens in 0 s (353204.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (379861.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (262828.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (399457.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (370085.6 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (285975.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (387166.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (374213.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (292625.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (343560.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (348315.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (316551.2 tok/s) (stop, cancel=False) +ok +test_openai_reasoning_effort (test_effort_paths.EffortPaths.test_openai_reasoning_effort) ... [strata] done: 24 tokens in 0 s (387166.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (363405.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (363405.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (368730.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (385683.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (393216.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (387166.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (393216.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (353204.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (399457.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (375609.3 tok/s) (stop, cancel=False) +ok +test_openai_reasoning_object_and_template_kwargs (test_effort_paths.EffortPaths.test_openai_reasoning_object_and_template_kwargs) ... [strata] done: 24 tokens in 0 s (407543.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (415964.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (412554.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (381300.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (399457.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (381300.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (356962.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (358232.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (362098.2 tok/s) (stop, cancel=False) +ok +test_responses_reasoning_effort (test_effort_paths.EffortPaths.test_responses_reasoning_effort) ... [strata] done: 24 tokens in 0 s (285165.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (299593.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (312618.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (307838.8 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (299593.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (315558.9 tok/s) (stop, cancel=False) +ok +test_the_chat_settings_shared_with_apps_beat_the_config (test_effort_paths.EffortPaths.test_the_chat_settings_shared_with_apps_beat_the_config) ... [strata] done: 24 tokens in 0 s (407543.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (387166.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (311651.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (363405.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (359511.8 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (310689.2 tok/s) (stop, cancel=False) +ok +test_the_level_follows_a_conversation (test_effort_paths.EffortPaths.test_the_level_follows_a_conversation) ... [strata] done: 24 tokens in 0 s (393216.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (356962.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (356962.0 tok/s) (stop, cancel=False) +ok +test_the_requests_own_level_beats_the_config (test_effort_paths.EffortPaths.test_the_requests_own_level_beats_the_config) ... [strata] done: 24 tokens in 0 s (393216.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (368730.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (381300.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (304118.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (363405.4 tok/s) (stop, cancel=False) +ok +test_the_web_apps_request_shape (test_effort_paths.EffortPaths.test_the_web_apps_request_shape) ... [strata] done: 24 tokens in 0 s (107202.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (120123.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (144631.2 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (152059.4 tok/s) (stop, cancel=False) +ok +test_tools_and_a_system_message_keep_the_level (test_effort_paths.EffortPaths.test_tools_and_a_system_message_keep_the_level) ... [strata] done: 24 tokens in 0 s (348315.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (354448.2 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (358232.4 tok/s) (stop, cancel=False) +ok +test_anthropic_requests_too (test_env_to_user.EnvBlockToUser.test_anthropic_requests_too) ... ok +test_content_parts_get_a_leading_text_part (test_env_to_user.EnvBlockToUser.test_content_parts_get_a_leading_text_part) ... ok +test_moved_to_the_first_user_message (test_env_to_user.EnvBlockToUser.test_moved_to_the_first_user_message) ... ok +test_off_and_no_block_change_nothing (test_env_to_user.EnvBlockToUser.test_off_and_no_block_change_nothing) ... ok +test_text_parts_are_read_as_one_text (test_env_to_user.EnvBlockToUser.test_text_parts_are_read_as_one_text) ... ok +test_two_folders_share_the_system_turn (test_env_to_user.EnvBlockToUser.test_two_folders_share_the_system_turn) ... ok +test_a_broken_jpeg_is_left_as_it_was_sent (test_exif_orientation.ExifOrientation.test_a_broken_jpeg_is_left_as_it_was_sent) ... ok +test_a_jpeg_without_the_tag_or_with_tag_1_passes_through_unchanged (test_exif_orientation.ExifOrientation.test_a_jpeg_without_the_tag_or_with_tag_1_passes_through_unchanged) ... ok +test_a_turned_jpeg_comes_out_upright_as_a_png (test_exif_orientation.ExifOrientation.test_a_turned_jpeg_comes_out_upright_as_a_png) ... ok +test_png_passes_through (test_exif_orientation.ExifOrientation.test_png_passes_through) ... ok +test_config_name_wins (test_family_name.FamilyName.test_config_name_wins) ... ok +test_old_configs (test_family_name.FamilyName.test_old_configs) ... ok +test_per_family (test_family_name.FamilyName.test_per_family) ... ok +test_an_earlier_window_never_finished_is_fatal (test_fatal_recovery.FatalRecovery.test_an_earlier_window_never_finished_is_fatal) ... ok +test_batch_timeout_signature_also_invalidates_engine (test_fatal_recovery.FatalRecovery.test_batch_timeout_signature_also_invalidates_engine) ... ok +test_closed_serial_engine_rejects_without_accessing_closed_pipe (test_fatal_recovery.FatalRecovery.test_closed_serial_engine_rejects_without_accessing_closed_pipe) ... ok +test_fatal_error_is_not_a_success_and_next_request_reloads (test_fatal_recovery.FatalRecovery.test_fatal_error_is_not_a_success_and_next_request_reloads) ... [strata] the engine stopped unexpectedly (exit code -15). Unrecoverable native verification failure: verify: timed out at layer 18; its GPU waits were released and the GPU finished (#267) The next request starts the engine again. +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -15); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 1 tokens in 0 s (22671.9 tok/s) (length, cancel=False) +ok +test_fatal_lines_are_checked_before_batch_routing (test_fatal_recovery.FatalRecovery.test_fatal_lines_are_checked_before_batch_routing) ... ok +test_generate_batched_refuses_a_dead_engine (test_fatal_recovery.FatalRecovery.test_generate_batched_refuses_a_dead_engine) ... ok +test_http_health_and_next_request_recovery (test_fatal_recovery.FatalRecovery.test_http_health_and_next_request_recovery) ... [strata] the engine stopped unexpectedly (exit code -15). Unrecoverable native verification failure: verify: timed out at layer 18; its GPU waits were released and the GPU finished (#267) The next request starts the engine again. +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -15); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 1 tokens in 0 s (27060.0 tok/s) (length, cancel=False) +ok +test_validation_error_does_not_kill_engine (test_fatal_recovery.FatalRecovery.test_validation_error_does_not_kill_engine) ... ok +test_anthropic_thinking_only_replies_are_not_rendered (test_frontend.EmptyAssistantHistory.test_anthropic_thinking_only_replies_are_not_rendered) ... ok +test_empty_user_and_tool_turns_are_kept (test_frontend.EmptyAssistantHistory.test_empty_user_and_tool_turns_are_kept) ... ok +test_final_empty_assistant_turn_is_kept (test_frontend.EmptyAssistantHistory.test_final_empty_assistant_turn_is_kept) ... ok +test_image_only_assistant_turn_is_kept (test_frontend.EmptyAssistantHistory.test_image_only_assistant_turn_is_kept) ... ok +test_openai_empty_replies_are_not_rendered (test_frontend.EmptyAssistantHistory.test_openai_empty_replies_are_not_rendered) ... ok +test_text_and_tool_calls_are_kept (test_frontend.EmptyAssistantHistory.test_text_and_tool_calls_are_kept) ... ok +test_keep_empty_turns_restores_the_old_prompt (test_frontend.EmptyTurnsSwitch.test_keep_empty_turns_restores_the_old_prompt) ... ok +test_fenced_and_bare_agree (test_frontend.FencedJsonCallInWrapper.test_fenced_and_bare_agree) ... ok +test_off_without_recovery_or_for_an_undeclared_tool (test_frontend.FencedJsonCallInWrapper.test_off_without_recovery_or_for_an_undeclared_tool) ... ok +test_other_fences_and_prose_stay_text (test_frontend.FencedJsonCallInWrapper.test_other_fences_and_prose_stay_text) ... ok +test_required_with_one_tool_names_it (test_frontend.ForcedCallOpening.test_required_with_one_tool_names_it) ... ok +test_a_user_asking_for_a_closer_verbatim_gets_text (test_frontend.QuotedToolCallCloser.test_a_user_asking_for_a_closer_verbatim_gets_text) ... ok +test_assistant_replay_keeps_every_tool_tag (test_frontend.QuotedToolCallCloser.test_assistant_replay_keeps_every_tool_tag) ... ok +test_assistant_replay_keeps_its_token (test_frontend.QuotedToolCallCloser.test_assistant_replay_keeps_its_token) ... ok +test_every_tool_tag_a_user_quotes_stays_text (test_frontend.QuotedToolCallCloser.test_every_tool_tag_a_user_quotes_stays_text) ... ok +test_template_keeps_seeing_a_tool_result_turn (test_frontend.QuotedToolCallCloser.test_template_keeps_seeing_a_tool_result_turn) ... ok +test_user_text_stays_text (test_frontend.QuotedToolCallCloser.test_user_text_stays_text) ... ok +test_user_tool_result_keeps_its_outer_markers (test_frontend.QuotedToolCallCloser.test_user_tool_result_keeps_its_outer_markers) ... ok +test_a_system_prompt_with_an_image_part_is_read_as_text_as_before (test_frontend.UnreadContentParts.test_a_system_prompt_with_an_image_part_is_read_as_text_as_before) ... ok +test_chat_file_and_audio_parts_are_left_out (test_frontend.UnreadContentParts.test_chat_file_and_audio_parts_are_left_out) ... ok +test_chat_refusal_part_is_read_as_its_text (test_frontend.UnreadContentParts.test_chat_refusal_part_is_read_as_its_text) ... ok +test_messages_blocks_real_clients_send_are_left_out (test_frontend.UnreadContentParts.test_messages_blocks_real_clients_send_are_left_out) ... ok +test_messages_redacted_thinking_is_dropped (test_frontend.UnreadContentParts.test_messages_redacted_thinking_is_dropped) ... ok +test_messages_tool_reference_in_a_tool_result_is_left_out (test_frontend.UnreadContentParts.test_messages_tool_reference_in_a_tool_result_is_left_out) ... ok +test_text_and_image_requests_render_as_before (test_frontend.UnreadContentParts.test_text_and_image_requests_render_as_before) ... ok +test_config_checked (test_images.Command.test_config_checked) ... ok +test_env (test_images.Command.test_env) ... ok +test_native_switch (test_images.Command.test_native_switch) ... ok +test_other_steps_and_batch (test_images.Command.test_other_steps_and_batch) ... ok +test_the_model_cards_command (test_images.Command.test_the_model_cards_command) +Unsloth's stable-diffusion.cpp command for Turbo: 8 steps, CFG 1, euler, the Turbo sigmas, --diffusion-fa. ... ok +test_b64 (test_images.OverHttp.test_b64) ... [strata] image: 512x768, n=2, 8 steps, seed 11 +/private/tmp/r39-serve-baseline/serve/server.py:5603: ResourceWarning: unclosed file <_io.TextIOWrapper name=9 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 0.3 s (keep) +ok +test_bad_request_and_failure (test_images.OverHttp.test_bad_request_and_failure) ... [strata] 400 invalid request: size 999x999: each side a multiple of 32 from 256 to 2048 +[strata] image: 1024x1024, n=1, 8 steps, seed 409967793 +[strata] image job failed: sd-cli exited with code 1 +[INFO ] loading tensors completed, taking 0.10s +[ERROR] ggml_vulkan: ErrorDeviceLost +/private/tmp/r39-serve-baseline/serve/server.py:5607: ResourceWarning: unclosed file <_io.TextIOWrapper name=9 encoding='UTF-8'> + return +ResourceWarning: Enable tracemalloc to get the object allocation traceback +ok +test_chat_waits_for_the_image (test_images.OverHttp.test_chat_waits_for_the_image) +One job at a time on the GPU: a chat request sent during an image job is answered after it. ... [strata] image: 1024x1024, n=1, 8 steps, seed 1458926076 +[strata] image done in 1.9 s (keep) +[strata] done: 8 tokens in 0 s (157532.5 tok/s) (length, cancel=False) +ok +test_listed (test_images.OverHttp.test_listed) ... ok +test_not_set_up (test_images.OverHttp.test_not_set_up) ... ok +test_url (test_images.OverHttp.test_url) ... [strata] image: 1024x1024, n=1, 8 steps, seed 0 +[strata] image done in 0.3 s (keep) +[strata] image: 1024x1024, n=1, 8 steps, seed 1 +[strata] image done in 0.3 s (keep) +[strata] image: 1024x1024, n=1, 8 steps, seed 2 +[strata] image done in 0.2 s (keep) +[strata] image: 1024x1024, n=1, 8 steps, seed 3 +[strata] image done in 0.2 s (keep) +ok +test_eager_reload (test_images.Policy.test_eager_reload) ... [strata] image job: the conversation (123 tokens) saved; unloading the model +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] the conversation from before the image job is back (123 tokens) +ok +test_keep_touches_nothing (test_images.Policy.test_keep_touches_nothing) ... ok +test_resolve (test_images.Policy.test_resolve) ... ok +test_shrink_gives_the_cache_back (test_images.Policy.test_shrink_gives_the_cache_back) ... [strata] image job: the engine's expert cache gave VRAM back (9000 MiB free) +ok +test_unload_saves_and_the_next_start_restores (test_images.Policy.test_unload_saves_and_the_next_start_restores) ... [strata] image job: the conversation (123 tokens) saved; unloading the model +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] the conversation from before the image job is back (123 tokens) +ok +test_defaults (test_images.Request.test_defaults) ... ok +test_fields (test_images.Request.test_fields) ... ok +test_refused (test_images.Request.test_refused) ... ok +test_a_failing_call_never_breaks_a_request (test_keepawake.KeepAwakeTests.test_a_failing_call_never_breaks_a_request) ... ok +test_default_service_has_none (test_keepawake.KeepAwakeTests.test_default_service_has_none) ... ok +test_exit_clears_a_held_state (test_keepawake.KeepAwakeTests.test_exit_clears_a_held_state) ... ok +test_held_while_busy_and_released_when_idle (test_keepawake.KeepAwakeTests.test_held_while_busy_and_released_when_idle) ... ok +test_off_and_non_windows (test_keepawake.KeepAwakeTests.test_off_and_non_windows) ... ok +test_service_holds_for_the_whole_request_including_early_close (test_keepawake.KeepAwakeTests.test_service_holds_for_the_whole_request_including_early_close) ... [strata] done: 3 tokens in 0 s (62291.6 tok/s) (stop, cancel=False) +FAIL +test_an_engine_that_died_on_load_is_reported_as_such (test_lifecycle.Lifecycle.test_an_engine_that_died_on_load_is_reported_as_such) ... [strata] loading the model again (it was unloaded) ... +ok +test_busy_auth_and_foreign_origin_cannot_unload (test_lifecycle.Lifecycle.test_busy_auth_and_foreign_origin_cannot_unload) ... [strata] model unloaded; the next request loads it again +ok +test_close_cleans_up_after_a_broken_stdin_pipe (test_lifecycle.Lifecycle.test_close_cleans_up_after_a_broken_stdin_pipe) ... ok +test_close_does_not_forget_a_process_still_exiting (test_lifecycle.Lifecycle.test_close_does_not_forget_a_process_still_exiting) ... ok +test_close_sends_eof_before_waiting_for_windows_reader (test_lifecycle.Lifecycle.test_close_sends_eof_before_waiting_for_windows_reader) ... ok +test_close_terminates_before_it_kills (test_lifecycle.Lifecycle.test_close_terminates_before_it_kills) ... ok +test_discover_chat_unload_and_stream_reload (test_lifecycle.Lifecycle.test_discover_chat_unload_and_stream_reload) ... [strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 7 tokens in 0 s (240656.8 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (257545.0 tok/s) (stop, cancel=False) +[strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 7 tokens in 0 s (59796.6 tok/s) (stop, cancel=False) +ok +test_explicit_load_and_model_matching (test_lifecycle.Lifecycle.test_explicit_load_and_model_matching) ... [strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +ok +test_native_lazy_constructor_does_not_start_a_process (test_lifecycle.Lifecycle.test_native_lazy_constructor_does_not_start_a_process) ... ok +test_capability_and_batch (test_logit_bias.LogitBias.test_capability_and_batch) ... ok +test_forms_and_transport (test_logit_bias.LogitBias.test_forms_and_transport) ... ok +test_http_rejects_before_stream (test_logit_bias.LogitBias.test_http_rejects_before_stream) ... [strata] 400 invalid request: logit_bias: expected an object or a list of [token_id, bias] pairs +[strata] 400 invalid request: logit_bias: this engine does not support token biases; use a supporting CUDA/HIP build +[strata] 400 invalid request: logit_bias: biases must be finite numbers between -100 and 100; false is a pair-list ban +ok +test_invalid (test_logit_bias.LogitBias.test_invalid) ... ok +test_large_list (test_logit_bias.LogitBias.test_large_list) ... ok +[strata] MCP server 'missing' did not start: could not start 'strata-no-such-program': [Errno 2] No such file or directory: 'strata-no-such-program'. Its tools are left out; the chat works without them. +[strata] MCP server 'paged': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +test_a_server_that_fails_is_left_out (test_mcp.Client.test_a_server_that_fails_is_left_out) ... ok +test_call (test_mcp.Client.test_call) ... ok +test_cancel (test_mcp.Client.test_cancel) ... ok +test_crash_mid_call_then_restart (test_mcp.Client.test_crash_mid_call_then_restart) ... [strata] MCP server 'solo': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'solo' had stopped; starting it again +[strata] MCP server 'solo': 6 tools (echo, fail, sleep, big, die, add) +ok +test_errors_become_text (test_mcp.Client.test_errors_become_text) ... ok +test_initialize_and_list_every_page (test_mcp.Client.test_initialize_and_list_every_page) ... ok +test_namespaced_openai_tools (test_mcp.Client.test_namespaced_openai_tools) ... ok +test_timeout (test_mcp.Client.test_timeout) ... ok +test_truncation (test_mcp.Client.test_truncation) ... ok +test_bad_entries_stop_the_start (test_mcp.Config.test_bad_entries_stop_the_start) ... ok +test_both_spellings_and_the_file (test_mcp.Config.test_both_spellings_and_the_file) ... ok +[strata] MCP server 'down' did not start: could not reach http://127.0.0.1:9/mcp: . Its tools are left out; the chat works without them. +[strata] MCP server 'web': 6 tools (echo, fail, sleep, big, die, add) +test_list_and_call (test_mcp.HttpClient.test_list_and_call) ... ok +test_timeout_and_unreachable (test_mcp.HttpClient.test_timeout_and_unreachable) ... ok +test_opt_in_without_servers_is_a_plain_chat (test_mcp.NoServers.test_opt_in_without_servers_is_a_plain_chat) ... [strata] done: 13 tokens in 0 s (199728.8 tok/s) (stop, cancel=False) +ok +test_a_call_the_output_ends_inside_is_not_run (test_mcp.ToolLoop.test_a_call_the_output_ends_inside_is_not_run) +#211: the model's turn ends inside an MCP call: nothing runs (it used to run with no arguments). ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 83 tokens in 0 s (80121.3 tok/s) (stop, cancel=False) +ok +test_get_mcp (test_mcp.ToolLoop.test_get_mcp) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +ok +test_max_rounds (test_mcp.ToolLoop.test_max_rounds) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 120 tokens in 0 s (111205.6 tok/s) (stop, cancel=False) +[strata] tool fake__echo: ok, 5 characters in 0.0 s +[strata] done: 120 tokens in 0 s (105274.3 tok/s) (stop, cancel=False) +[strata] tool fake__echo: ok, 5 characters in 0.0 s +[strata] done: 120 tokens in 0 s (98592.8 tok/s) (stop, cancel=False) +ok +test_non_stream (test_mcp.ToolLoop.test_non_stream) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 141 tokens in 0 s (116923.1 tok/s) (stop, cancel=False) +[strata] tool fake__add: ok, 1 characters in 0.0 s +[strata] done: 13 tokens in 0 s (232025.3 tok/s) (stop, cancel=False) +ok +test_only_from_the_app_s_own_page (test_mcp.ToolLoop.test_only_from_the_app_s_own_page) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] refused an API request from the web page 'http://evil.example' (no API key; add its host to "allowed_hosts" or its origin to "cors_origins" in the config) +[strata] done: 13 tokens in 0 s (55300.2 tok/s) (stop, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +ok +test_own_tools_still_go_to_the_client (test_mcp.ToolLoop.test_own_tools_still_go_to_the_client) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 109 tokens in 0 s (127632.4 tok/s) (stop, cancel=False) +ok +test_plain_requests_get_no_mcp_tools (test_mcp.ToolLoop.test_plain_requests_get_no_mcp_tools) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 116 tokens in 0 s (89701.2 tok/s) (stop, cancel=False) +[strata] done: 16 tokens in 0 s (120482.7 tok/s) (stop, cancel=False) +[strata] done: 16 tokens in 0 s (113551.4 tok/s) (stop, cancel=False) +ok +test_stop_during_a_tool (test_mcp.ToolLoop.test_stop_during_a_tool) +Closing the connection while a slow tool runs stops the tool (notifications/cancelled) and the loop. ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 120 tokens in 0 s (116266.2 tok/s) (stop, cancel=False) +ok +test_tool_call_then_answer (test_mcp.ToolLoop.test_tool_call_then_answer) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 134 tokens in 0 s (115621.6 tok/s) (stop, cancel=False) +[strata] tool fake__echo: ok, 19 characters in 0.0 s +[strata] done: 31 tokens in 0 s (125991.7 tok/s) (stop, cancel=False) +ok +test_tool_error_is_a_result (test_mcp.ToolLoop.test_tool_error_is_a_result) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 84 tokens in 0 s (153249.9 tok/s) (stop, cancel=False) +[strata] tool fake__fail: error, 27 characters in 0.0 s +[strata] done: 21 tokens in 0 s (164023.1 tok/s) (stop, cancel=False) +ok +test_truncated_for_the_model (test_mcp.ToolLoop.test_truncated_for_the_model) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 115 tokens in 0 s (125708.9 tok/s) (stop, cancel=False) +[strata] tool fake__big: ok, 3,000 characters in 0.0 s (truncated for the model) +[strata] done: 16 tokens in 0 s (192841.6 tok/s) (stop, cancel=False) +ok +test_metrics (test_monitor.ConversationCacheCard.test_metrics) ... [strata] done: 24 tokens in 0 s (300487.5 tok/s) (stop, cancel=False) +ok +test_the_card_is_on_the_page (test_monitor.ConversationCacheCard.test_the_card_is_on_the_page) ... ok +test_the_log (test_monitor.ConversationCacheCard.test_the_log) ... ok +test_capture_is_bounded_and_eviction_returns_404 (test_monitor.Monitor.test_capture_is_bounded_and_eviction_returns_404) ... ok +test_error_and_engine_death_are_not_successes (test_monitor.Monitor.test_error_and_engine_death_are_not_successes) ... [strata] 400 invalid request: request exceeds the context window: prompt (183 tokens) + max tokens (20000) exceeds the context (16384); requests are never truncated. Send a smaller max_tokens (at most 16193 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_history_is_authorized_and_does_not_capture_headers (test_monitor.Monitor.test_history_is_authorized_and_does_not_capture_headers) ... [strata] done: 50 tokens in 0 s (225258.0 tok/s) (stop, cancel=False) +ok +test_output_and_reasoning_for_both_dialects_and_stream_modes (test_monitor.Monitor.test_output_and_reasoning_for_both_dialects_and_stream_modes) ... [strata] done: 50 tokens in 0 s (134778.4 tok/s) (stop, cancel=False) +[strata] done: 50 tokens in 0 s (56299.4 tok/s) (stop, cancel=False) +[strata] done: 50 tokens in 0 s (156737.8 tok/s) (stop, cancel=False) +[strata] done: 50 tokens in 0 s (46937.2 tok/s) (stop, cancel=False) +ok +test_page_and_relative_assets_are_served (test_monitor.Monitor.test_page_and_relative_assets_are_served) ... ok +test_queue_wait_and_loading_are_attributed_separately (test_monitor.Monitor.test_queue_wait_and_loading_are_attributed_separately) ... [strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 50 tokens in 0 s (110434.5 tok/s) (stop, cancel=False) +[strata] done: 50 tokens in 0 s (147479.0 tok/s) (stop, cancel=False) +ok +test_off_by_default (test_monitor.MonitorOff.test_off_by_default) ... [strata] done: 7 tokens in 0 s (170698.4 tok/s) (stop, cancel=False) +ok +test_parallel_above_the_window_says_so_unless_batch_mtp (test_parallel.ParallelArgs.test_parallel_above_the_window_says_so_unless_batch_mtp) ... ok +test_parallel_config_key (test_parallel.ParallelArgs.test_parallel_config_key) ... [strata] "parallel" must be a whole number of requests (2..8), not '4': ignored +[strata] "parallel" must be a whole number of requests (2..8), not True: ignored +[strata] "parallel": 12 - the engine runs at most 8 requests together; it will use 8 +ok +test_a_burst_of_connections_is_not_reset (test_parallel.ParallelService.test_a_burst_of_connections_is_not_reset) +30-40 clients at once got "connection reset by peer" with the listen backlog of 5 (4 x R9700 burst): the ... [strata] done: 8 tokens in 1 s (10.1 tok/s) (length, cancel=False) +[strata] answering: 8 of max 8 tokens, 11.4 tok/s, 1 s +[strata] done: 8 tokens in 1 s (11.4 tok/s) (length, cancel=False) +[strata] answering: 2 of max 8 tokens, 15.3 tok/s, 1 s +[strata] reading the prompt: 32 of 46 tokens, 1 s so far +[strata] done: 8 tokens in 2 s (8.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 2 s so far +[strata] answering: 8 of max 8 tokens, 7.7 tok/s, 2 s +[strata] done: 8 tokens in 2 s (7.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 2 s so far +[strata] done: 8 tokens in 3 s (8.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 3 s so far +[strata] answering: 8 of max 8 tokens, 7.7 tok/s, 3 s +[strata] done: 8 tokens in 3 s (7.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 3 s so far +[strata] answering: 8 of max 8 tokens, 7.3 tok/s, 4 s +[strata] done: 8 tokens in 4 s (7.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 4 s so far +[strata] answering: 8 of max 8 tokens, 6.8 tok/s, 4 s +[strata] done: 8 tokens in 4 s (6.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 4 s so far +[strata] answering: 7 of max 8 tokens, 6.0 tok/s, 5 s +[strata] done: 8 tokens in 5 s (6.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 5 s so far +[strata] answering: 8 of max 8 tokens, 6.9 tok/s, 6 s +[strata] done: 8 tokens in 6 s (6.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 6 s so far +[strata] done: 8 tokens in 6 s (8.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 6 s so far +[strata] done: 8 tokens in 6 s (9.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 7 s so far +[strata] done: 8 tokens in 7 s (8.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 7 s so far +[strata] answering: 8 of max 8 tokens, 7.2 tok/s, 8 s +[strata] done: 8 tokens in 8 s (7.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 8 s so far +[strata] answering: 7 of max 8 tokens, 7.0 tok/s, 8 s +[strata] done: 8 tokens in 9 s (6.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 9 s so far +[strata] answering: 8 of max 8 tokens, 6.6 tok/s, 9 s +[strata] done: 8 tokens in 9 s (6.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 9 s so far +[strata] answering: 6 of max 8 tokens, 5.4 tok/s, 10 s +[strata] done: 8 tokens in 10 s (5.8 tok/s) (length, cancel=False) +[strata] answering: 7 of max 8 tokens, 6.8 tok/s, 10 s +[strata] reading the prompt: 32 of 46 tokens, 10 s so far +[strata] done: 8 tokens in 10 s (6.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 10 s so far +[strata] answering: 8 of max 8 tokens, 7.7 tok/s, 11 s +[strata] done: 8 tokens in 11 s (7.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 11 s so far +[strata] answering: 8 of max 8 tokens, 6.8 tok/s, 11 s +[strata] done: 8 tokens in 11 s (6.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 11 s so far +[strata] answering: 6 of max 8 tokens, 5.9 tok/s, 12 s +[strata] done: 8 tokens in 12 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 12 s so far +[strata] answering: 8 of max 8 tokens, 7.2 tok/s, 13 s +[strata] done: 8 tokens in 13 s (7.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 13 s so far +[strata] answering: 7 of max 8 tokens, 7.0 tok/s, 13 s +[strata] done: 8 tokens in 14 s (6.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 14 s so far +[strata] answering: 8 of max 8 tokens, 6.9 tok/s, 14 s +[strata] done: 8 tokens in 14 s (6.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 14 s so far +[strata] answering: 7 of max 8 tokens, 6.8 tok/s, 15 s +[strata] done: 8 tokens in 15 s (7.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 15 s so far +[strata] answering: 8 of max 8 tokens, 6.9 tok/s, 15 s +[strata] done: 8 tokens in 15 s (6.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 15 s so far +[strata] answering: 7 of max 8 tokens, 6.8 tok/s, 16 s +[strata] done: 8 tokens in 16 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 16 s so far +[strata] answering: 8 of max 8 tokens, 7.2 tok/s, 16 s +[strata] done: 8 tokens in 16 s (7.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 16 s so far +[strata] answering: 7 of max 8 tokens, 6.3 tok/s, 17 s +[strata] answering: 7 of max 8 tokens, 6.9 tok/s, 17 s +[strata] done: 8 tokens in 17 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 17 s so far +[strata] done: 8 tokens in 17 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 18 s so far +[strata] answering: 8 of max 8 tokens, 7.0 tok/s, 18 s +[strata] done: 8 tokens in 18 s (7.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 19 s so far +[strata] answering: 8 of max 8 tokens, 7.0 tok/s, 19 s +[strata] done: 8 tokens in 19 s (7.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 19 s so far +[strata] done: 8 tokens in 19 s (9.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 19 s so far +[strata] done: 8 tokens in 20 s (8.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 20 s so far +[strata] answering: 7 of max 8 tokens, 6.2 tok/s, 21 s +[strata] done: 8 tokens in 21 s (6.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 21 s so far +[strata] answering: 8 of max 8 tokens, 6.1 tok/s, 21 s +[strata] done: 8 tokens in 21 s (6.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 21 s so far +[strata] answering: 7 of max 8 tokens, 6.3 tok/s, 22 s +[strata] done: 8 tokens in 22 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 22 s so far +[strata] answering: 8 of max 8 tokens, 6.5 tok/s, 22 s +[strata] done: 8 tokens in 22 s (6.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 22 s so far +[strata] answering: 7 of max 8 tokens, 6.1 tok/s, 23 s +[strata] done: 8 tokens in 23 s (6.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 23 s so far +[strata] answering: 8 of max 8 tokens, 6.1 tok/s, 24 s +[strata] done: 8 tokens in 24 s (6.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 24 s so far +[strata] answering: 6 of max 8 tokens, 5.5 tok/s, 25 s +[strata] done: 8 tokens in 25 s (5.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 25 s so far +[strata] answering: 8 of max 8 tokens, 6.5 tok/s, 25 s +[strata] done: 8 tokens in 25 s (6.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 25 s so far +[strata] done: 8 tokens in 26 s (8.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 26 s so far +[strata] answering: 8 of max 8 tokens, 7.3 tok/s, 26 s +[strata] done: 8 tokens in 26 s (7.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 26 s so far +[strata] answering: 6 of max 8 tokens, 5.8 tok/s, 27 s +[strata] done: 8 tokens in 27 s (6.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 27 s so far +[strata] answering: 8 of max 8 tokens, 6.9 tok/s, 27 s +[strata] done: 8 tokens in 27 s (6.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 28 s so far +[strata] answering: 6 of max 8 tokens, 5.5 tok/s, 28 s +[strata] answering: 7 of max 8 tokens, 6.8 tok/s, 29 s +[strata] done: 8 tokens in 29 s (5.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 29 s so far +[strata] done: 8 tokens in 29 s (6.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 29 s so far +[strata] answering: 8 of max 8 tokens, 7.5 tok/s, 30 s +[strata] done: 8 tokens in 30 s (7.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 30 s so far +[strata] answering: 8 of max 8 tokens, 6.6 tok/s, 30 s +[strata] done: 8 tokens in 30 s (6.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 30 s so far +[strata] answering: 5 of max 8 tokens, 4.7 tok/s, 31 s +[strata] done: 8 tokens in 31 s (6.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 31 s so far +[strata] answering: 8 of max 8 tokens, 6.6 tok/s, 31 s +[strata] done: 8 tokens in 31 s (6.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 31 s so far +[strata] answering: 7 of max 8 tokens, 6.6 tok/s, 32 s +[strata] done: 8 tokens in 32 s (6.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 32 s so far +[strata] answering: 8 of max 8 tokens, 6.5 tok/s, 33 s +[strata] done: 8 tokens in 33 s (6.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 33 s so far +[strata] done: 8 tokens in 33 s (8.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 33 s so far +[strata] done: 8 tokens in 34 s (8.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 34 s so far +[strata] answering: 7 of max 8 tokens, 6.7 tok/s, 34 s +[strata] done: 8 tokens in 34 s (6.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 35 s so far +[strata] answering: 8 of max 8 tokens, 7.1 tok/s, 35 s +[strata] done: 8 tokens in 35 s (7.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 35 s so far +[strata] answering: 7 of max 8 tokens, 6.1 tok/s, 36 s +[strata] done: 8 tokens in 36 s (6.3 tok/s) (length, cancel=False) +[strata] answering: 8 of max 8 tokens, 7.4 tok/s, 36 s +[strata] done: 8 tokens in 36 s (7.4 tok/s) (length, cancel=False) +ok +test_a_continued_request_reports_its_own_prompt_reuse (test_parallel.ParallelService.test_a_continued_request_reports_its_own_prompt_reuse) +A request STOPped on the solo path and continued in a slot (or back on the solo path) is sent again as its ... [strata] thinking: 1 of max 200 tokens, 2.0 tok/s, 1 s +[strata] reading the prompt: 32 of 249 tokens, 1 s so far +[strata] thinking: 6 of max 200 tokens, 3.8 tok/s, 2 s +[strata] thinking: 6 of max 200 tokens, 7.3 tok/s, 2 s +[strata] done: 10 tokens in 3 s (6.9 tok/s) (stop, cancel=False) +[strata] thinking: 13 of max 200 tokens, 4.9 tok/s, 3 s +[strata] thinking: 21 of max 200 tokens, 5.8 tok/s, 4 s +[strata] thinking: 30 of max 200 tokens, 6.4 tok/s, 5 s +[strata] thinking: 38 of max 200 tokens, 6.7 tok/s, 6 s +[strata] thinking: 47 of max 200 tokens, 6.9 tok/s, 7 s +[strata] done: 55 tokens in 8 s (7.0 tok/s) (stop, cancel=False) +ok +test_a_failed_window_names_its_error (test_parallel.ParallelService.test_a_failed_window_names_its_error) ... [strata] the engine stopped unexpectedly (exit code 1). The engine exited (code 1) after it reported: verify batch: layer 34 never rang (graph finished) - please report it at github.com/Niko1221/Strata/issues with the log. The next request starts the engine again. +[strata] done: 3 tokens in 1 s (4.3 tok/s) (error, cancel=False) +[strata] the engine stopped unexpectedly (exit code 1). The engine exited (code 1) after it reported: verify batch: layer 34 never rang (graph finished) - please report it at github.com/Niko1221/Strata/issues with the log. The next request starts the engine again. +[strata] done: 1 tokens in 1 s (5.8 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +ok +test_a_leaked_busy_slot_does_not_hang_the_next_request (test_parallel.ParallelService.test_a_leaked_busy_slot_does_not_hang_the_next_request) +#1603: a slot whose busy flag has no owner (no request holds it, no BDONE wait is running) kept every later ... [strata] parallel: slot(s) 0, 1 were marked busy with no request using them (issue #1603); freed, the waiting request goes on +[strata] reading the prompt: 32 of 52 tokens, 1 s so far +[strata] done: 10 tokens in 2 s (9.3 tok/s) (stop, cancel=False) +ok +test_a_long_prompt_gives_way_to_a_short_one (test_parallel.ParallelService.test_a_long_prompt_gives_way_to_a_short_one) +#656: a long solo prompt read gives way at a chunk boundary; the short request is answered first, then the ... [strata] reading the prompt: 352 of 3,037 tokens, 1 s so far +[strata] answering: 5 of max 64 tokens, 5.3 tok/s, 1 s +[strata] reading the prompt: 960 of 3,037 tokens, 2 s so far +[strata] done: 10 tokens in 2 s (4.9 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,504 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 2,048 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 2,624 of 3,037 tokens, 5 s so far +[strata] answering: 2 of max 64 tokens, 7.1 tok/s, 6 s +[strata] done: 10 tokens in 7 s (8.2 tok/s) (stop, cancel=False) +ok +test_a_parallel_request_does_not_wait_for_the_fifo_on_a_loaded_engine (test_parallel.ParallelService.test_a_parallel_request_does_not_wait_for_the_fifo_on_a_loaded_engine) +#1603: parallel requests used to take the FIFO (held by an image encode, /v1/vram, load) just to find the ... [strata] answering: 8 of max 64 tokens, 8.1 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.2 tok/s) (stop, cancel=False) +[strata] answering: 8 of max 64 tokens, 8.3 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.2 tok/s) (stop, cancel=False) +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 10 tokens in 1 s (9.7 tok/s) (stop, cancel=False) +ok +test_a_request_moved_between_solo_and_slot_keeps_its_totals (test_parallel.ParallelService.test_a_request_moved_between_solo_and_slot_keeps_its_totals) +#1869/#1870: a request that read its prompt and decoded on the solo path, was STOPped when another arrived and ... [strata] answering: 5 of max 200 tokens, 5.2 tok/s, 1 s +[strata] answering: 6 of max 200 tokens, 9.5 tok/s, 1 s +[strata] done: 10 tokens in 2 s (8.6 tok/s) (stop, cancel=False) +[strata] answering: 14 of max 200 tokens, 6.7 tok/s, 2 s +[strata] answering: 23 of max 200 tokens, 7.1 tok/s, 3 s +[strata] answering: 31 of max 200 tokens, 7.1 tok/s, 4 s +[strata] answering: 40 of max 200 tokens, 7.4 tok/s, 5 s +[strata] answering: 51 of max 200 tokens, 7.9 tok/s, 7 s +[strata] done: 55 tokens in 7 s (7.8 tok/s) (stop, cancel=False) +ok +test_a_slot_request_ended_by_its_stop_token_reports_the_slot_leg (test_parallel.ParallelService.test_a_slot_request_ended_by_its_stop_token_reports_the_slot_leg) +#1869: the consumer leaves at the stop token, before the BDONE: the request kept the admission's one token and ... [strata] answering: 5 of max 200 tokens, 7.6 tok/s, 1 s +[strata] answering: 6 of max 200 tokens, 5.5 tok/s, 1 s +[strata] done: 10 tokens in 2 s (5.8 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 2 s (6.8 tok/s) (stop, cancel=False) +ok +test_an_admission_gives_way_too (test_parallel.ParallelService.test_an_admission_gives_way_too) +The same while another request decodes in a slot: the long one is being admitted, a short one waits. ... [strata] reading the prompt: 128 of 3,037 tokens, 1 s so far +[strata] answering: 4 of max 60 tokens, 3.8 tok/s, 1 s +[strata] answering: 4 of max 64 tokens, 4.4 tok/s, 1 s +[strata] reading the prompt: 608 of 3,037 tokens, 2 s so far +[strata] answering: 8 of max 60 tokens, 3.6 tok/s, 2 s +[strata] answering: 8 of max 64 tokens, 3.9 tok/s, 2 s +[strata] done: 10 tokens in 3 s (3.7 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,088 of 3,037 tokens, 3 s so far +[strata] done: 10 tokens in 3 s (3.8 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,632 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 2,112 of 3,037 tokens, 5 s so far +[strata] reading the prompt: 2,688 of 3,037 tokens, 6 s so far +[strata] answering: 4 of max 64 tokens, 7.4 tok/s, 7 s +[strata] done: 10 tokens in 8 s (8.5 tok/s) (stop, cancel=False) +ok +test_an_admission_is_asked_between_far_apart_progress_lines (test_parallel.ParallelService.test_an_admission_is_asked_between_far_apart_progress_lines) +As above for a read admitted beside a decoding slot (BGEN, `ask()`): with one PP line per 80 chunks the ... [strata] reading the prompt: 2,560 of 3,037 tokens, 6 s so far +[strata] answering: 1 of max 64 tokens, 161319.4 tok/s, 6 s +[strata] answering: 4 of max 60 tokens, 0.4 tok/s, 11 s +[strata] answering: 2 of max 64 tokens, 0.4 tok/s, 11 s +[strata] reading the prompt: 2,560 of 3,037 tokens, 11 s so far +[strata] answering: 3 of max 64 tokens, 0.5 tok/s, 12 s +[strata] answering: 2 of max 64 tokens, 6.1 tok/s, 12 s +[strata] answering: 6 of max 60 tokens, 0.5 tok/s, 12 s +[strata] answering: 9 of max 64 tokens, 6.4 tok/s, 13 s +[strata] answering: 13 of max 60 tokens, 1.0 tok/s, 13 s +[strata] done: 10 tokens in 13 s (1.4 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 13 s (6.4 tok/s) (stop, cancel=False) +[strata] answering: 22 of max 60 tokens, 1.6 tok/s, 14 s +[strata] answering: 30 of max 60 tokens, 2.0 tok/s, 15 s +[strata] answering: 38 of max 60 tokens, 2.3 tok/s, 16 s +[strata] answering: 45 of max 60 tokens, 2.6 tok/s, 18 s +[strata] answering: 54 of max 60 tokens, 2.9 tok/s, 19 s +[strata] done: 55 tokens in 19 s (3.0 tok/s) (stop, cancel=False) +ok +test_an_owned_busy_slot_is_never_taken (test_parallel.ParallelService.test_an_owned_busy_slot_is_never_taken) +The leak check frees only slots with no owner: a slot a request holds stays busy however long it waits. ... [strata] parallel: slot(s) 0 were marked busy with no request using them (issue #1603); freed, the waiting request goes on +ok +test_batch_groups_auto_follows_the_engine (test_parallel.ParallelService.test_batch_groups_auto_follows_the_engine) +--batch-groups auto: the engine picks the groups and says so (INFO batch_groups=G); the server spreads the ... ok +test_bgen_refused_by_the_engine_is_retried_not_a_400 (test_parallel.ParallelService.test_bgen_refused_by_the_engine_is_retried_not_a_400) +#1994: the engine refuses a BGEN ("no such free slot") although the server saw the slot free: the request ... [strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (1/6) (issue #1994) +[strata] answering: 2 of max 64 tokens, 6.1 tok/s, 1 s +[strata] answering: 7 of max 64 tokens, 5.2 tok/s, 1 s +[strata] answering: 15 of max 64 tokens, 6.0 tok/s, 3 s +[strata] done: 10 tokens in 2 s (6.8 tok/s) (stop, cancel=False) +[strata] answering: 22 of max 64 tokens, 6.3 tok/s, 4 s +[strata] answering: 33 of max 64 tokens, 7.2 tok/s, 5 s +[strata] answering: 40 of max 64 tokens, 7.0 tok/s, 6 s +[strata] answering: 49 of max 64 tokens, 7.2 tok/s, 7 s +[strata] done: 55 tokens in 8 s (7.2 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 1 s (8.7 tok/s) (stop, cancel=False) +ok +test_bgen_refused_for_good_is_a_retryable_503_and_frees_the_slots (test_parallel.ParallelService.test_bgen_refused_for_good_is_a_retryable_503_and_frees_the_slots) +#1994: when the engine keeps refusing, the answer is a 503 (retryable), never a 400, and no slot stays busy. ... [strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (1/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (2/6) (issue #1994) +[strata] answering: 5 of max 64 tokens, 5.4 tok/s, 1 s +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (3/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (4/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (5/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (6/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (7/6) (issue #1994) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (length, cancel=False) +[strata] answering: 12 of max 64 tokens, 6.1 tok/s, 2 s +[strata] answering: 20 of max 64 tokens, 6.5 tok/s, 3 s +[strata] answering: 28 of max 64 tokens, 6.8 tok/s, 4 s +[strata] answering: 35 of max 64 tokens, 6.8 tok/s, 5 s +[strata] answering: 44 of max 64 tokens, 7.0 tok/s, 6 s +[strata] answering: 52 of max 64 tokens, 7.1 tok/s, 7 s +[strata] done: 55 tokens in 8 s (7.0 tok/s) (stop, cancel=False) +[strata] answering: 9 of max 64 tokens, 8.9 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.5 tok/s) (stop, cancel=False) +ok +test_concurrent_requests_share_the_slots (test_parallel.ParallelService.test_concurrent_requests_share_the_slots) ... [strata] answering: 5 of max 64 tokens, 5.1 tok/s, 1 s +[strata] answering: 4 of max 64 tokens, 7.6 tok/s, 1 s +[strata] done: 10 tokens in 2 s (8.6 tok/s) (stop, cancel=False) +[strata] reading the prompt: 32 of 66 tokens, 2 s so far +[strata] answering: 13 of max 64 tokens, 6.3 tok/s, 2 s +[strata] answering: 9 of max 64 tokens, 8.5 tok/s, 3 s +[strata] done: 10 tokens in 3 s (9.1 tok/s) (stop, cancel=False) +[strata] answering: 23 of max 64 tokens, 7.2 tok/s, 3 s +[strata] answering: 33 of max 64 tokens, 7.7 tok/s, 4 s +[strata] answering: 44 of max 64 tokens, 8.2 tok/s, 5 s +[strata] answering: 54 of max 64 tokens, 8.4 tok/s, 7 s +[strata] done: 55 tokens in 7 s (8.4 tok/s) (stop, cancel=False) +ok +test_engine_reports_fewer_slots (test_parallel.ParallelService.test_engine_reports_fewer_slots) ... [strata] parallel requests: 4 asked, the engine runs 2 (its log says why) +ok +test_engine_turns_batching_off (test_parallel.ParallelService.test_engine_turns_batching_off) ... [strata] parallel requests: 4 asked, the engine runs one at a time (its log says why) +[strata] answering: 9 of max 64 tokens, 9.6 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.7 tok/s) (stop, cancel=False) +ok +test_left_alone_it_goes_back_to_the_solo_path (test_parallel.ParallelService.test_left_alone_it_goes_back_to_the_solo_path) +A request decoding in a slot whose neighbour finished goes back to the solo path (MTP drafts), continued ... [strata] answering: 7 of max 200 tokens, 6.7 tok/s, 1 s +[strata] answering: 6 of max 200 tokens, 8.8 tok/s, 1 s +[strata] done: 10 tokens in 2 s (8.2 tok/s) (stop, cancel=False) +[strata] answering: 16 of max 200 tokens, 7.8 tok/s, 2 s +[strata] answering: 25 of max 200 tokens, 8.2 tok/s, 3 s +[strata] answering: 35 of max 200 tokens, 8.4 tok/s, 4 s +[strata] answering: 45 of max 200 tokens, 8.6 tok/s, 5 s +[strata] answering: 53 of max 200 tokens, 8.4 tok/s, 6 s +[strata] done: 55 tokens in 7 s (8.4 tok/s) (stop, cancel=False) +ok +test_next_turn_goes_back_to_its_slot (test_parallel.ParallelService.test_next_turn_goes_back_to_its_slot) ... [strata] answering: 7 of max 200 tokens, 6.8 tok/s, 1 s +[strata] answering: 5 of max 64 tokens, 9.4 tok/s, 1 s +[strata] done: 10 tokens in 2 s (9.6 tok/s) (stop, cancel=False) +[strata] answering: 16 of max 200 tokens, 7.6 tok/s, 2 s +[strata] answering: 23 of max 200 tokens, 7.2 tok/s, 3 s +[strata] answering: 31 of max 200 tokens, 7.3 tok/s, 4 s +[strata] answering: 39 of max 200 tokens, 7.4 tok/s, 5 s +[strata] answering: 49 of max 200 tokens, 7.8 tok/s, 6 s +[strata] done: 55 tokens in 7 s (8.1 tok/s) (stop, cancel=False) +ok +test_no_way_given_to_a_long_one (test_parallel.ParallelService.test_no_way_given_to_a_long_one) +Two prompts of about the same length: no BYIELD (it only goes to a much shorter one). ... [strata] reading the prompt: 544 of 3,037 tokens, 1 s so far +[strata] reading the prompt: 1,024 of 3,037 tokens, 2 s so far +[strata] reading the prompt: 1,504 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 1,984 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 2,464 of 3,037 tokens, 5 s so far +[strata] reading the prompt: 2,976 of 3,037 tokens, 6 s so far +[strata] answering: 1 of max 64 tokens, 1.1 tok/s, 7 s +[strata] answering: 1 of max 64 tokens, 0.5 tok/s, 8 s +[strata] answering: 1 of max 64 tokens, 0.3 tok/s, 9 s +[strata] answering: 1 of max 64 tokens, 0.2 tok/s, 10 s +[strata] reading the prompt: 2,048 of 3,038 tokens, 10 s so far +[strata] answering: 1 of max 64 tokens, 0.2 tok/s, 11 s +[strata] answering: 2 of max 64 tokens, 0.3 tok/s, 12 s +[strata] reading the prompt: 32 of 3,037 tokens, 12 s so far +[strata] answering: 5 of max 64 tokens, 0.7 tok/s, 13 s +[strata] reading the prompt: 544 of 3,037 tokens, 13 s so far +[strata] done: 10 tokens in 15 s (1.2 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,088 of 3,037 tokens, 14 s so far +[strata] reading the prompt: 1,632 of 3,037 tokens, 15 s so far +[strata] reading the prompt: 2,208 of 3,037 tokens, 16 s so far +[strata] reading the prompt: 2,784 of 3,037 tokens, 18 s so far +[strata] answering: 4 of max 64 tokens, 6.9 tok/s, 19 s +[strata] done: 10 tokens in 19 s (7.0 tok/s) (stop, cancel=False) +ok +test_no_way_given_to_a_long_one_between_far_apart_progress_lines (test_parallel.ParallelService.test_no_way_given_to_a_long_one_between_far_apart_progress_lines) +The ask between PP lines keeps the rule: a waiting request about as long gets no BYIELD there either. ... [strata] reading the prompt: 2,560 of 3,037 tokens, 4 s so far +[strata] answering: 1 of max 64 tokens, 0.2 tok/s, 10 s +[strata] reading the prompt: 2,560 of 3,038 tokens, 10 s so far +[strata] reading the prompt: 2,560 of 3,037 tokens, 15 s so far +[strata] answering: 3 of max 64 tokens, 0.3 tok/s, 16 s +[strata] answering: 2 of max 64 tokens, 6.0 tok/s, 16 s +[strata] answering: 4 of max 64 tokens, 0.3 tok/s, 17 s +[strata] done: 10 tokens in 18 s (0.8 tok/s) (stop, cancel=False) +[strata] answering: 9 of max 64 tokens, 6.3 tok/s, 17 s +[strata] done: 10 tokens in 18 s (6.5 tok/s) (stop, cancel=False) +ok +test_nonfinite_in_batch_mode (test_parallel.ParallelService.test_nonfinite_in_batch_mode) +#879 in --batch mode: DONE ... nonfinite on the solo path or at an admission (nothing emitted yet) sends the ... [strata] the engine's logits became non-finite before the first token (#879); it dropped its cached state: sending the request again once +[strata] answering: 9 of max 64 tokens, 8.7 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.2 tok/s) (stop, cancel=False) +[strata] the engine's logits became non-finite before the first token (#879); it dropped its cached state: sending the request again once +[strata] answering: 7 of max 64 tokens, 7.1 tok/s, 1 s +[strata] answering: 6 of max 64 tokens, 7.5 tok/s, 1 s +[strata] done: 10 tokens in 2 s (7.0 tok/s) (stop, cancel=False) +[strata] answering: 14 of max 64 tokens, 6.8 tok/s, 2 s +[strata] answering: 22 of max 64 tokens, 6.8 tok/s, 3 s +[strata] answering: 31 of max 64 tokens, 7.3 tok/s, 4 s +[strata] answering: 40 of max 64 tokens, 7.4 tok/s, 5 s +[strata] answering: 47 of max 64 tokens, 7.3 tok/s, 7 s +[strata] answering: 54 of max 64 tokens, 7.2 tok/s, 8 s +[strata] done: 55 tokens in 8 s (7.2 tok/s) (stop, cancel=False) +[strata] the engine's logits became non-finite after 2 tokens (#879): the reply ends here as "length"; the engine dropped its cached state, so the next request reads its prompt from the start +[strata] done: 2 tokens in 1 s (5.4 tok/s) (length, cancel=False) +[strata] answering: 7 of max 64 tokens, 6.5 tok/s, 1 s +[strata] answering: 16 of max 64 tokens, 7.6 tok/s, 2 s +[strata] answering: 26 of max 64 tokens, 8.2 tok/s, 3 s +[strata] answering: 36 of max 64 tokens, 8.6 tok/s, 4 s +[strata] answering: 44 of max 64 tokens, 8.4 tok/s, 5 s +[strata] answering: 53 of max 64 tokens, 8.4 tok/s, 6 s +[strata] done: 55 tokens in 6 s (8.5 tok/s) (stop, cancel=False) +ok +test_stop_strings_in_a_batch_slot (test_parallel.ParallelService.test_stop_strings_in_a_batch_slot) +#454: a stop string cuts the answer in --batch mode too, and the slot is freed for the next request. ... [strata] done: 9 tokens in 1 s (9.6 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 1 s (9.6 tok/s) (stop, cancel=False) +ok +test_the_way_is_asked_between_far_apart_progress_lines (test_parallel.ParallelService.test_the_way_is_asked_between_far_apart_progress_lines) +#656 when the PP lines come far apart (the pipelined groups' read in pieces of one chunk: one PP per piece, ... [strata] reading the prompt: 2,560 of 3,037 tokens, 5 s so far +[strata] answering: 1 of max 64 tokens, 1000000.0 tok/s, 5 s +[strata] answering: 2 of max 64 tokens, 0.4 tok/s, 10 s +[strata] reading the prompt: 2,560 of 3,037 tokens, 10 s so far +[strata] answering: 3 of max 64 tokens, 7.2 tok/s, 11 s +[strata] answering: 4 of max 64 tokens, 0.7 tok/s, 11 s +[strata] done: 10 tokens in 11 s (1.5 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 12 s (9.1 tok/s) (stop, cancel=False) +ok +test_three_requests_long_long_short_do_not_deadlock (test_parallel.ParallelService.test_three_requests_long_long_short_do_not_deadlock) +ENGINE_REVIEW finding 1 (0.1.41 check): two long prompts fill both slots, a short one arrives a moment later ... [strata] reading the prompt: 96 of 3,037 tokens, 1 s so far +[strata] reading the prompt: 800 of 3,037 tokens, 2 s so far +[strata] reading the prompt: 1,600 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 2,272 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 2,912 of 3,037 tokens, 5 s so far +[strata] reading the prompt: 32 of 3,037 tokens, 5 s so far +[strata] answering: 4 of max 64 tokens, 4.9 tok/s, 6 s +[strata] reading the prompt: 608 of 3,037 tokens, 6 s so far +[strata] answering: 9 of max 64 tokens, 4.3 tok/s, 7 s +[strata] reading the prompt: 1,088 of 3,037 tokens, 7 s so far +[strata] done: 10 tokens in 8 s (4.1 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,504 of 3,037 tokens, 8 s so far +[strata] reading the prompt: 2,048 of 3,037 tokens, 9 s so far +[strata] reading the prompt: 2,656 of 3,037 tokens, 10 s so far +[strata] reading the prompt: 32 of 43 tokens, 11 s so far +[strata] answering: 3 of max 64 tokens, 5.3 tok/s, 12 s +[strata] answering: 8 of max 64 tokens, 7.5 tok/s, 12 s +[strata] done: 10 tokens in 13 s (6.4 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 12 s (7.2 tok/s) (stop, cancel=False) +ok +test_batch_groups_spread_requests_over_the_groups (test_parallel.PickSlot.test_batch_groups_spread_requests_over_the_groups) ... ok +test_slots_view (test_parallel.PickSlot.test_slots_view) ... ok +test_the_slot_that_holds_the_conversation (test_parallel.PickSlot.test_the_slot_that_holds_the_conversation) ... ok +test_compaction_boundaries_reads_and_line_replacement (test_parser_track.TrackTests.test_compaction_boundaries_reads_and_line_replacement) ... ok +test_feed_events_with_long_lines_and_quoted_tool_tags (test_parser_track.TrackTests.test_feed_events_with_long_lines_and_quoted_tool_tags) ... ok +test_fragment_state_from_varied_initial_states (test_parser_track.TrackTests.test_fragment_state_from_varied_initial_states) ... ok +test_long_line_snapshots_restore_and_append (test_parser_track.TrackTests.test_long_line_snapshots_restore_and_append) ... ok +test_random_streaming_boundaries (test_parser_track.TrackTests.test_random_streaming_boundaries) ... ok +test_every_concurrent_request_is_counted (test_prometheus.Batched.test_every_concurrent_request_is_counted) ... [strata] answering: 12 of max 64 tokens, 12.7 tok/s, 1 s +[strata] reading the prompt: 10 of max 64 tokens, 10.6 tok/s, 1 s +[strata] answering: 12 of max 64 tokens, 12.0 tok/s, 1 s +[strata] answering: 22 of max 64 tokens, 11.3 tok/s, 2 s +[strata] answering: 20 of max 64 tokens, 10.3 tok/s, 2 s +[strata] answering: 23 of max 64 tokens, 11.0 tok/s, 2 s +[strata] answering: 29 of max 64 tokens, 9.8 tok/s, 3 s +[strata] answering: 33 of max 64 tokens, 11.2 tok/s, 3 s +[strata] answering: 32 of max 64 tokens, 10.0 tok/s, 3 s +[strata] answering: 38 of max 64 tokens, 9.5 tok/s, 4 s +[strata] answering: 44 of max 64 tokens, 10.7 tok/s, 4 s +[strata] answering: 41 of max 64 tokens, 9.7 tok/s, 4 s +[strata] answering: 46 of max 64 tokens, 9.1 tok/s, 5 s +[strata] done: 51 tokens in 5 s (10.1 tok/s) (stop, cancel=False) +[strata] answering: 50 of max 64 tokens, 9.4 tok/s, 5 s +[strata] done: 51 tokens in 6 s (9.4 tok/s) (stop, cancel=False) +[strata] done: 51 tokens in 6 s (8.7 tok/s) (stop, cancel=False) +ok +test_a_request_shows_in_both_formats (test_prometheus.Endpoint.test_a_request_shows_in_both_formats) ... [strata] done: 22 tokens in 0 s (150775.6 tok/s) (stop, cancel=False) +ok +test_histograms_are_cumulative (test_prometheus.Units.test_histograms_are_cumulative) ... ok +test_render_reads_the_server_record (test_prometheus.Units.test_render_reads_the_server_record) ... ok +test_strata_names_follow_the_json (test_prometheus.Units.test_strata_names_follow_the_json) ... ok +test_who_gets_the_text (test_prometheus.Units.test_who_gets_the_text) ... ok +test_a_queued_stream_gets_keep_alives_then_its_answer (test_queue_keepalive.QueueKeepAlive.test_a_queued_stream_gets_keep_alives_then_its_answer) ... [strata] done: 3 tokens in 0 s (37900.3 tok/s) (stop, cancel=False) +ok +test_a_request_cancelled_while_queued_never_takes_the_turn (test_queue_keepalive.QueueKeepAlive.test_a_request_cancelled_while_queued_never_takes_the_turn) ... ok +test_closing_a_queued_request_is_no_longer_queued (test_queue_keepalive.QueueKeepAlive.test_closing_a_queued_request_is_no_longer_queued) ... ok +test_anthropic_stream (test_quiet_cancel.QuietPromptCancel.test_anthropic_stream) ... [strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (25420.0 tok/s) (length, cancel=False) +ok +test_batch_solo_stream (test_quiet_cancel.QuietPromptCancel.test_batch_solo_stream) ... [strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] answering: 1 of max 1 tokens, 1000000.0 tok/s, 1 s +[strata] done: 1 tokens in 1 s (27776.8 tok/s) (length, cancel=False) +ok +test_other_http_modes (test_quiet_cancel.QuietPromptCancel.test_other_http_modes) ... [strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (43240.2 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (12018.1 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 2 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (38479.9 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] answering: 1 of max 1 tokens, 144631.2 tok/s, 1 s +[strata] done: 1 tokens in 1 s (6096.4 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (6754.1 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (5343.1 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (6326.3 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (7307.1 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] answering: 1 of max 1 tokens, 349525.3 tok/s, 1 s +[strata] done: 1 tokens in 1 s (14315.0 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (6710.9 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] answering: 1 of max 1 tokens, 110376.4 tok/s, 1 s +[strata] done: 1 tokens in 1 s (5841.6 tok/s) (length, cancel=False) +ok +test_expired_silence_is_not_a_heartbeat (test_quiet_cancel.QuietWaitPolicy.test_expired_silence_is_not_a_heartbeat) ... ok +test_polling_keeps_ten_second_heartbeats (test_quiet_cancel.QuietWaitPolicy.test_polling_keeps_ten_second_heartbeats) ... ok +test_ready_timeout_on_lazy_load_answers_the_request (test_ready_timeout.ReadyTimeout.test_ready_timeout_on_lazy_load_answers_the_request) ... [strata] loading the model again (it was unloaded) ... +[strata] the engine did not start (try 1 of 3); again in 0 s +[strata] the engine did not start (try 2 of 3); again in 0 s +ok +test_ready_timeout_on_restart_answers_the_request (test_ready_timeout.ReadyTimeout.test_ready_timeout_on_restart_answers_the_request) ... [strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 1 tokens in 0 s (15887.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine did not start (try 1 of 3); again in 0 s +[strata] the engine did not start (try 2 of 3); again in 0 s +ok +test_a_loop_written_as_one_long_word (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_a_loop_written_as_one_long_word) +#1753: a 24k-digit string cycling an 11-digit pattern is one word to the split and never one token twice. ... ok +test_a_lower_top_p_and_top_k_are_never_raised_or_lowered (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_a_lower_top_p_and_top_k_are_never_raised_or_lowered) ... [strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2086 tokens in 0 s (297950.6 tok/s) (stop, cancel=False) +ok +test_batch_status_and_history_record_recovery (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_batch_status_and_history_record_recovery) ... [strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2071 tokens in 0 s (373946.5 tok/s) (stop, cancel=False) +ok +test_cancel_during_reasoning_does_not_start_recovery (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_cancel_during_reasoning_does_not_start_recovery) ... [strata] done: 1 tokens in 0 s (34663.7 tok/s) (cancel, cancel=True) +ok +test_default_single_token_guard_does_not_detect_passage_loop (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_default_single_token_guard_does_not_detect_passage_loop) ... [strata] done: 16384 tokens in 0 s (664778.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_disabled_is_one_pass (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_disabled_is_one_pass) ... [strata] done: 28001 tokens in 0 s (706931.8 tok/s) (stop, cancel=False) +ok +test_frozen_cohort_sampler_is_unchanged_during_recovery (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_frozen_cohort_sampler_is_unchanged_during_recovery) ... [strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2082 tokens in 0 s (665590.0 tok/s) (stop, cancel=False) +ok +test_missing_native_high_hint_skips_recovery (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_missing_native_high_hint_skips_recovery) ... [strata] done: 28001 tokens in 0 s (626940.2 tok/s) (stop, cancel=False) +ok +test_natural_answer_with_shared_budget_and_preserved_prefix (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_natural_answer_with_shared_budget_and_preserved_prefix) ... [strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2106 tokens in 0 s (629235.2 tok/s) (stop, cancel=False) +ok +test_old_explicit_sampler_recovers_without_mutating_request (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_old_explicit_sampler_recovers_without_mutating_request) ... [strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2086 tokens in 0 s (619376.9 tok/s) (stop, cancel=False) +ok +test_preserves_task_and_exact_generated_tokens (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_preserves_task_and_exact_generated_tokens) ... ok +test_recovery_is_once_and_never_extends_output_limit (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_recovery_is_once_and_never_extends_output_limit) ... [strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 16384 tokens in 0 s (653500.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_repeated_final_text_does_not_change_sampler_or_prompt (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_repeated_final_text_does_not_change_sampler_or_prompt) ... [strata] done: 28001 tokens in 0 s (552824.1 tok/s) (stop, cancel=False) +ok +test_repetition_and_novel_work (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_repetition_and_novel_work) ... ok +test_stop_mode_ends_a_loop_written_as_one_long_word (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_stop_mode_ends_a_loop_written_as_one_long_word) +#1753: the word count never sees the digits, so only the tail's period ends the reply. ... [strata] the thinking repeated the same passages (coverage=1.000) at 2560 tokens: ended as "length" (reasoning_loop_recovery: "stop" in strata-.json; remove it to turn this off) +[strata] done: 2560 tokens in 0 s (644657.7 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_stop_mode_ends_the_reply_at_the_loop (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_stop_mode_ends_the_reply_at_the_loop) ... [strata] the thinking repeated the same passages (coverage=1.000) at 2048 tokens: ended as "length" (reasoning_loop_recovery: "stop" in strata-.json; remove it to turn this off) +[strata] done: 2048 tokens in 0 s (626545.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_the_splice_keeps_a_literal_think_tag_as_text (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_the_splice_keeps_a_literal_think_tag_as_text) +#537: the prompt is spliced as token ids, never decoded and encoded again. ... ok +test_user_instruction_is_never_rewritten (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_user_instruction_is_never_rewritten) ... ok +test_every_specimen_whole_and_in_chunks (test_reasoning_rescue.Corpus.test_every_specimen_whole_and_in_chunks) ... ok +test_the_corpus_has_the_reported_split (test_reasoning_rescue.Corpus.test_the_corpus_has_the_reported_split) ... ok +test_a_call_waits_and_is_not_streamed_early (test_reasoning_rescue.Gates.test_a_call_waits_and_is_not_streamed_early) ... ok +test_a_fence_needs_the_same_marker_to_close (test_reasoning_rescue.Gates.test_a_fence_needs_the_same_marker_to_close) ... ok +test_a_length_cut_keeps_even_a_clean_stranded_call_as_text (test_reasoning_rescue.Gates.test_a_length_cut_keeps_even_a_clean_stranded_call_as_text) ... ok +test_a_stranded_call_is_recovered (test_reasoning_rescue.Gates.test_a_stranded_call_is_recovered) ... ok +test_after_a_closed_fence_a_call_counts_again (test_reasoning_rescue.Gates.test_after_a_closed_fence_a_call_counts_again) ... ok +test_an_undeclared_name_stays_reasoning (test_reasoning_rescue.Gates.test_an_undeclared_name_stays_reasoning) ... ok +test_fence_split_across_chunks (test_reasoning_rescue.Gates.test_fence_split_across_chunks) ... ok +test_no_tools_declared_is_unchanged (test_reasoning_rescue.Gates.test_no_tools_declared_is_unchanged) ... ok +test_opener_split_across_chunks (test_reasoning_rescue.Gates.test_opener_split_across_chunks) ... ok +test_recovered_call_keeps_the_text_before_it (test_reasoning_rescue.Gates.test_recovered_call_keeps_the_text_before_it) ... ok +test_the_counters (test_reasoning_rescue.Gates.test_the_counters) ... ok +test_the_end_of_the_turn_splits_a_held_tail (test_reasoning_rescue.Gates.test_the_end_of_the_turn_splits_a_held_tail) ... ok +test_the_reporters_negative_controls (test_reasoning_rescue.Gates.test_the_reporters_negative_controls) ... ok +test_a_stranded_call_is_a_call_in_every_api (test_reasoning_rescue.OverHttp.test_a_stranded_call_is_a_call_in_every_api) ... [strata] done: 116 tokens in 0 s (222164.0 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (169821.7 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (217593.6 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (164371.4 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (221355.4 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (172592.9 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +ok +test_quotes_and_cuts_are_not (test_reasoning_rescue.OverHttp.test_quotes_and_cuts_are_not) ... [strata] done: 133 tokens in 0 s (77188.7 tok/s) (stop, cancel=False) +[strata] done: 133 tokens in 0 s (49188.1 tok/s) (stop, cancel=False) +[strata] done: 133 tokens in 0 s (89440.8 tok/s) (stop, cancel=False) +[strata] done: 133 tokens in 0 s (57403.0 tok/s) (stop, cancel=False) +[strata] done: 133 tokens in 0 s (100223.2 tok/s) (stop, cancel=False) +[strata] done: 133 tokens in 0 s (59030.9 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (95417.8 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (57917.3 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (108234.1 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (68778.4 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (107240.7 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (71944.5 tok/s) (stop, cancel=False) +[strata] done: 150 tokens in 0 s (114933.4 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 150 tokens in 0 s (85668.0 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 150 tokens in 0 s (115481.9 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 150 tokens in 0 s (92250.1 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 150 tokens in 0 s (123847.6 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 150 tokens in 0 s (101491.5 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (221292.1 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (180904.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (231232.7 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (175916.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (216739.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (186797.9 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (233820.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (216523.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (240169.7 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (198268.8 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (229339.1 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (211349.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +ok +test_visible_answer_fences_on_all_three_apis (test_reasoning_rescue.OverHttp.test_visible_answer_fences_on_all_three_apis) ... [strata] done: 134 tokens in 0 s (182835.6 tok/s) (stop, cancel=False) +[strata] done: 134 tokens in 0 s (102993.7 tok/s) (stop, cancel=False) +[strata] done: 134 tokens in 0 s (182538.7 tok/s) (stop, cancel=False) +[strata] done: 134 tokens in 0 s (112497.3 tok/s) (stop, cancel=False) +[strata] done: 134 tokens in 0 s (169186.3 tok/s) (stop, cancel=False) +[strata] done: 134 tokens in 0 s (99704.9 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (108814.5 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (65266.3 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (129137.6 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (84136.4 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (120835.9 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (77258.5 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (123393.2 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (90620.1 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (128442.3 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (96650.6 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (127166.6 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (99557.9 tok/s) (stop, cancel=False) +ok +test_a_closed_call_in_a_length_cut_answer_stays_a_call (test_reasoning_rescue.VisibleAnswer.test_a_closed_call_in_a_length_cut_answer_stays_a_call) ... ok +test_quoted_calls_stay_text (test_reasoning_rescue.VisibleAnswer.test_quoted_calls_stay_text) ... ok +test_real_calls_still_fire (test_reasoning_rescue.VisibleAnswer.test_real_calls_still_fire) ... ok +test_without_thinking (test_reasoning_rescue.VisibleAnswer.test_without_thinking) ... ok +test_finish_reasons_and_block_order (test_reasoning_tools.OverHttp.test_finish_reasons_and_block_order) ... [strata] done: 121 tokens in 0 s (158994.6 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 121 tokens in 0 s (126027.0 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 121 tokens in 0 s (157759.0 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 121 tokens in 0 s (130097.6 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 121 tokens in 0 s (123964.5 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +ok +test_a_call_in_the_answer_is_unchanged (test_reasoning_tools.Parser.test_a_call_in_the_answer_is_unchanged) ... ok +test_a_declared_call_in_the_reasoning_is_a_call (test_reasoning_tools.Parser.test_a_declared_call_in_the_reasoning_is_a_call) ... ok +test_a_partial_tag_is_held (test_reasoning_tools.Parser.test_a_partial_tag_is_held) ... ok +test_a_tag_that_never_closes_does_not_hold_the_thinking_back (test_reasoning_tools.Parser.test_a_tag_that_never_closes_does_not_hold_the_thinking_back) ... ok +test_an_undeclared_name_stays_reasoning (test_reasoning_tools.Parser.test_an_undeclared_name_stays_reasoning) ... ok +test_an_unfinished_call_is_never_a_call (test_reasoning_tools.Parser.test_an_unfinished_call_is_never_a_call) ... ok +test_no_tools_declared_is_unchanged (test_reasoning_tools.Parser.test_no_tools_declared_is_unchanged) ... ok +test_the_thinking_ending_inside_a_call_makes_it_text (test_reasoning_tools.Parser.test_the_thinking_ending_inside_a_call_makes_it_text) ... ok +test_a_call_quoted_after_it_is_not_run (test_reasoning_tools.QuotedThinkEnd.test_a_call_quoted_after_it_is_not_run) ... ok +test_a_quoted_tag_stays_reasoning (test_reasoning_tools.QuotedThinkEnd.test_a_quoted_tag_stays_reasoning) ... ok +test_an_output_that_ends_on_the_tag (test_reasoning_tools.QuotedThinkEnd.test_an_output_that_ends_on_the_tag) ... ok +test_the_end_after_a_quote_still_ends_the_thinking (test_reasoning_tools.QuotedThinkEnd.test_the_end_after_a_quote_still_ends_the_thinking) ... ok +test_a_call_cut_at_the_token_limit (serve.test_server.UnfinishedToolCall.test_a_call_cut_at_the_token_limit) +The same cut by max_tokens: "length" / "max_tokens", and the whole (non-streamed) answers leave the call ... [strata] done: 119 tokens in 0 s (92174.0 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (70052.2 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (97067.7 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (85305.4 tok/s) (length, cancel=False) +ok +test_a_cut_call (serve.test_server.UnfinishedToolCall.test_a_cut_call) ... [strata] done: 120 tokens in 0 s (128758.4 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (93675.1 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (136808.0 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (102906.7 tok/s) (stop, cancel=False) +ok +test_a_whole_call (serve.test_server.UnfinishedToolCall.test_a_whole_call) ... [strata] done: 147 tokens in 0 s (166279.0 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (124810.3 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (171936.1 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (136740.4 tok/s) (stop, cancel=False) +ok +test_collect_keeps_calls_whose_arguments_parse (serve.test_server.UnfinishedToolCall.test_collect_keeps_calls_whose_arguments_parse) ... ok +test_parser (serve.test_server.UnfinishedToolCall.test_parser) ... ok +test_a_number_cuts_the_gpu_list (test_replicas.Config.test_a_number_cuts_the_gpu_list) ... ok +test_cpu_sets_off (test_replicas.Config.test_cpu_sets_off) ... ok +test_mistakes_are_said (test_replicas.Config.test_mistakes_are_said) ... ok +test_named_groups (test_replicas.Config.test_named_groups) ... ok +test_replica_cfg (test_replicas.Config.test_replica_cfg) ... ok +test_a_conversation_goes_back_to_its_replica (test_replicas.ReplicaService.test_a_conversation_goes_back_to_its_replica) ... [strata] reading the prompt: 512 of 523 tokens, 1 s so far +[strata] done: 10 tokens in 2 s (8.8 tok/s) (stop, cancel=False) +[strata] reading the prompt: 512 of 581 tokens, 1 s so far +[strata] answering: 9 of max 16 tokens, 9.9 tok/s, 2 s +[strata] done: 10 tokens in 2 s (9.1 tok/s) (stop, cancel=False) +ok +test_a_replica_that_died_is_skipped (test_replicas.ReplicaService.test_a_replica_that_died_is_skipped) ... [strata] answering: 9 of max 64 tokens, 8.6 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.1 tok/s) (stop, cancel=False) +[strata] answering: 9 of max 64 tokens, 8.3 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.3 tok/s) (stop, cancel=False) +[strata] answering: 7 of max 64 tokens, 7.3 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.5 tok/s) (stop, cancel=False) +ok +test_engines_without_slots_run_one_each_at_once (test_replicas.ReplicaService.test_engines_without_slots_run_one_each_at_once) ... [strata] answering: 9 of max 64 tokens, 9.8 tok/s, 1 s +[strata] done: 10 tokens in 1 s (9.8 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 1 s (9.5 tok/s) (stop, cancel=False) +ok +test_metrics_and_status_see_all_slots (test_replicas.ReplicaService.test_metrics_and_status_see_all_slots) ... ok +test_requests_spread_over_both_and_answers_are_the_solo_answer (test_replicas.ReplicaService.test_requests_spread_over_both_and_answers_are_the_solo_answer) ... [strata] done: 10 tokens in 1 s (9.8 tok/s) (stop, cancel=False) +[strata] answering: 1 of max 64 tokens, 2.0 tok/s, 1 s +[strata] answering: 1 of max 64 tokens, 2.3 tok/s, 1 s +[strata] reading the prompt: 32 of 355 tokens, 1 s so far +[strata] reading the prompt: 32 of 355 tokens, 1 s so far +[strata] reading the prompt: 32 of 355 tokens, 2 s so far +[strata] reading the prompt: 32 of 355 tokens, 2 s so far +[strata] answering: 5 of max 64 tokens, 3.1 tok/s, 2 s +[strata] answering: 5 of max 64 tokens, 2.8 tok/s, 2 s +[strata] answering: 3 of max 64 tokens, 5.5 tok/s, 2 s +[strata] answering: 3 of max 64 tokens, 4.8 tok/s, 3 s +[strata] answering: 3 of max 64 tokens, 6.8 tok/s, 3 s +[strata] answering: 9 of max 64 tokens, 3.3 tok/s, 3 s +[strata] answering: 5 of max 64 tokens, 9.9 tok/s, 3 s +[strata] done: 10 tokens in 3 s (3.6 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 3 s (3.5 tok/s) (stop, cancel=False) +[strata] answering: 8 of max 64 tokens, 5.2 tok/s, 4 s +[strata] done: 10 tokens in 4 s (5.8 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 4 s (6.0 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 4 s (7.9 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 4 s (8.5 tok/s) (stop, cancel=False) +ok +test_a_dead_replica_is_skipped (test_replicas.RouteChoice.test_a_dead_replica_is_skipped) ... ok +test_a_long_prefix_outweighs_load (test_replicas.RouteChoice.test_a_long_prefix_outweighs_load) ... ok +test_batch_is_the_sum (test_replicas.RouteChoice.test_batch_is_the_sum) ... ok +test_least_loaded_then_affinity (test_replicas.RouteChoice.test_least_loaded_then_affinity) ... ok +test_nonfinite_forgets_the_replica (test_replicas.RouteChoice.test_nonfinite_forgets_the_replica) +#879: after DONE ... nonfinite the engine holds no conversation: the router forgets that replica. ... ok +test_matched_prefix_per_replica (test_replicas.Routing.test_matched_prefix_per_replica) ... ok +test_the_index_is_bounded (test_replicas.Routing.test_the_index_is_bounded) ... ok +test_anthropic_nonstream_classifies_only_original_input_as_cached (test_request_accounting.RequestAccounting.test_anthropic_nonstream_classifies_only_original_input_as_cached) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (159783.0 tok/s) (stop, cancel=False), expert cache 88.9% hit +ok +test_anthropic_stream_classifies_only_original_input_as_cached (test_request_accounting.RequestAccounting.test_anthropic_stream_classifies_only_original_input_as_cached) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (50001.3 tok/s) (stop, cancel=False), expert cache 88.9% hit +ok +test_cancelled_continuation_keeps_original_input_and_counts_each_done_once (test_request_accounting.RequestAccounting.test_cancelled_continuation_keeps_original_input_and_counts_each_done_once) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 98 tokens in 0 s (61514.8 tok/s) (cancel, cancel=True), expert cache 88.9% hit +ok +test_failed_continuation_does_not_count_the_first_done_twice (test_request_accounting.RequestAccounting.test_failed_continuation_does_not_count_the_first_done_twice) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] the continuation ended without a DONE line. The next request starts the engine again. +[strata] done: 95 tokens in 0 s (102589.8 tok/s) (error, cancel=False), expert cache 90.0% hit +ok +test_openai_nonstream_keeps_original_cache_and_all_segment_work (test_request_accounting.RequestAccounting.test_openai_nonstream_keeps_original_cache_and_all_segment_work) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (180747.2 tok/s) (stop, cancel=False), expert cache 88.9% hit +ok +test_openai_stream_keeps_original_cache_and_all_segment_work (test_request_accounting.RequestAccounting.test_openai_stream_keeps_original_cache_and_all_segment_work) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (48948.8 tok/s) (stop, cancel=False), expert cache 88.9% hit +ok +test_single_segment_keeps_its_own_usage_and_timings (test_request_accounting.RequestAccounting.test_single_segment_keeps_its_own_usage_and_timings) ... [strata] thinking capped at 200 of max_tokens 400 to leave room for the answer (budget 10000) +[strata] done: 211 tokens in 0 s (49449.5 tok/s) (stop, cancel=False), expert cache 90.0% hit +ok +test_prompt_n_is_what_was_read (test_request_accounting.SlotLegAccounting.test_prompt_n_is_what_was_read) ... ok +test_solo_leg_adds_to_the_slot_leg (test_request_accounting.SlotLegAccounting.test_solo_leg_adds_to_the_slot_leg) ... ok +test_two_hand_backs_sum (test_request_accounting.SlotLegAccounting.test_two_hand_backs_sum) ... ok +test_a_good_request_still_answers (test_request_hardening.Hardening.test_a_good_request_still_answers) ... [strata] done: 8 tokens in 0 s (258111.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_a_wrong_type_is_a_400_that_names_the_field (test_request_hardening.Hardening.test_a_wrong_type_is_a_400_that_names_the_field) ... ok +test_default_limit_is_generous (test_request_hardening.Hardening.test_default_limit_is_generous) ... ok +test_in_type_values_and_nulls_still_answer (test_request_hardening.Hardening.test_in_type_values_and_nulls_still_answer) ... [strata] done: 8 tokens in 0 s (286790.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 8 tokens in 0 s (260111.9 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 8 tokens in 0 s (228261.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_llama_cpp_names_for_the_repetition_penalty (test_request_hardening.Hardening.test_llama_cpp_names_for_the_repetition_penalty) +#1819: repeat_penalty / repeat_last_n are repetition_penalty / penalty_last_n; the OpenAI-style name wins. ... [strata] done: 8 tokens in 0 s (319566.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_negative_or_bad_content_length_is_400 (test_request_hardening.Hardening.test_negative_or_bad_content_length_is_400) ... ok +test_not_json_is_400 (test_request_hardening.Hardening.test_not_json_is_400) ... ok +test_numeric_strings_are_read_and_used (test_request_hardening.Hardening.test_numeric_strings_are_read_and_used) ... [strata] done: 8 tokens in 0 s (305040.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_oversized_content_length_is_413_unread (test_request_hardening.Hardening.test_oversized_content_length_is_413_unread) ... ok +test_wrong_shapes_are_400 (test_request_hardening.Hardening.test_wrong_shapes_are_400) ... ok +test_a_compaction_prompt_starts_as_its_conversations_last_prompt (test_responses.CodexCompaction.test_a_compaction_prompt_starts_as_its_conversations_last_prompt) ... ok +test_a_compaction_that_sends_tools_keeps_them (test_responses.CodexCompaction.test_a_compaction_that_sends_tools_keeps_them) ... ok +test_a_custom_compact_prompt_is_a_compaction_when_codex_says_so (test_responses.CodexCompaction.test_a_custom_compact_prompt_is_a_compaction_when_codex_says_so) ... ok +test_a_thread_title_turn_does_not_replace_the_kept_tools (test_responses.CodexCompaction.test_a_thread_title_turn_does_not_replace_the_kept_tools) ... ok +test_another_conversation_never_gets_these_tools (test_responses.CodexCompaction.test_another_conversation_never_gets_these_tools) ... ok +test_malformed_metadata_is_no_error_and_no_reuse (test_responses.CodexCompaction.test_malformed_metadata_is_no_error_and_no_reuse) ... ok +test_requests_that_are_not_a_codex_compaction_are_rendered_as_sent (test_responses.CodexCompaction.test_requests_that_are_not_a_codex_compaction_are_rendered_as_sent) ... ok +test_the_compact_prompt_in_a_message_is_not_a_compaction (test_responses.CodexCompaction.test_the_compact_prompt_in_a_message_is_not_a_compaction) ... ok +test_without_its_conversation_nothing_is_kept_or_used (test_responses.CodexCompaction.test_without_its_conversation_nothing_is_kept_or_used) ... ok +test_a_compaction_too_long_with_the_kept_tools_is_rendered_as_sent (test_responses.CodexCompactionOverHttp.test_a_compaction_too_long_with_the_kept_tools_is_rendered_as_sent) ... ok +test_a_request_that_fails_leaves_the_kept_tools_as_they_were (test_responses.CodexCompactionOverHttp.test_a_request_that_fails_leaves_the_kept_tools_as_they_were) ... ok +test_the_kept_tools_shape_only_the_prompt (test_responses.CodexCompactionOverHttp.test_the_kept_tools_shape_only_the_prompt) ... ok +test_a_text_answer_that_breaks_the_schema_fails_as_without_tools (test_responses.JsonFormatWithTools.test_a_text_answer_that_breaks_the_schema_fails_as_without_tools) ... skipped 'the schema check needs jsonschema' +test_a_tool_call_passes_through (test_responses.JsonFormatWithTools.test_a_tool_call_passes_through) ... ok +test_the_final_text_answer_is_the_schema_object (test_responses.JsonFormatWithTools.test_the_final_text_answer_is_the_schema_object) ... ok +test_server_close_ends_the_sampler (test_responses.NoLeakedSampler.test_server_close_ends_the_sampler) ... ok +test_a_foreign_page_cannot_name_a_file_as_an_image (test_responses.OverHttp.test_a_foreign_page_cannot_name_a_file_as_an_image) +#553: the chat routes refuse a file image from another origin; so does /v1/responses. ... ok +test_api_key_and_foreign_pages (test_responses.OverHttp.test_api_key_and_foreign_pages) ... ok +test_effort_none_does_not_think (test_responses.OverHttp.test_effort_none_does_not_think) ... ok +test_errors_use_the_responses_format (test_responses.OverHttp.test_errors_use_the_responses_format) ... [strata] 400 invalid request: Expecting property name enclosed in double quotes: line 1 column 2 (char 1) +ok +test_json_schema_text_format (test_responses.OverHttp.test_json_schema_text_format) ... ok +test_max_output_tokens_ends_incomplete (test_responses.OverHttp.test_max_output_tokens_ends_incomplete) ... ok +test_monitor_and_status (test_responses.OverHttp.test_monitor_and_status) ... ok +test_non_streaming_text_with_reasoning (test_responses.OverHttp.test_non_streaming_text_with_reasoning) ... ok +test_stream_event_order (test_responses.OverHttp.test_stream_event_order) ... ok +test_too_long_for_the_context (test_responses.OverHttp.test_too_long_for_the_context) ... ok +test_a_string_input_is_one_user_message (test_responses.Parsing.test_a_string_input_is_one_user_message) ... ok +test_additional_tools_input_items_are_accepted (test_responses.Parsing.test_additional_tools_input_items_are_accepted) +#782: Codex sends an `additional_tools` input item; it used to be refused with a 400. ... ok +test_bad_arguments_from_a_cut_off_call_do_not_refuse_the_conversation (test_responses.Parsing.test_bad_arguments_from_a_cut_off_call_do_not_refuse_the_conversation) ... ok +test_codex_history_becomes_one_assistant_turn_per_answer (test_responses.Parsing.test_codex_history_becomes_one_assistant_turn_per_answer) ... ok +test_effort_and_text_format (test_responses.Parsing.test_effort_and_text_format) ... ok +test_leading_developer_messages_join_the_system_message_later_ones_become_user (test_responses.Parsing.test_leading_developer_messages_join_the_system_message_later_ones_become_user) ... ok +test_reasoning_comes_back_from_our_encrypted_content_and_foreign_tokens_are_ignored (test_responses.Parsing.test_reasoning_comes_back_from_our_encrypted_content_and_foreign_tokens_are_ignored) ... ok +test_tools_namespaces_custom_and_hosted (test_responses.Parsing.test_tools_namespaces_custom_and_hosted) ... ok +test_unsupported_items_are_named (test_responses.Parsing.test_unsupported_items_are_named) ... ok +test_a_thread_title_is_answered_without_the_engine (test_responses.ThreadTitle.test_a_thread_title_is_answered_without_the_engine) ... ok +test_off_by_default_the_engine_answers (test_responses.ThreadTitle.test_off_by_default_the_engine_answers) ... ok +test_function_call_round_trip (test_responses.ToolRoundTrip.test_function_call_round_trip) ... ok +test_namespace_and_custom_calls_come_back_under_their_own_names (test_responses.ToolRoundTrip.test_namespace_and_custom_calls_come_back_under_their_own_names) ... ok +test_namespace_call_written_with_double_underscore (test_responses.ToolRoundTrip.test_namespace_call_written_with_double_underscore) ... ok +test_a_restart_that_fails_wakes_the_waiters (test_restart_waiters.RestartWaiters.test_a_restart_that_fails_wakes_the_waiters) ... [strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +ok +test_an_engine_err_answers_at_once_not_batched (test_restart_waiters.RestartWaiters.test_an_engine_err_answers_at_once_not_batched) ... [strata] the engine reported an error: refused: images are not enabled +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] 400 invalid request: refused: images are not enabled +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] done: 10 tokens in 1 s (10.5 tok/s) (stop, cancel=False) +ok +test_an_engine_err_answers_at_once_solo_loop (test_restart_waiters.RestartWaiters.test_an_engine_err_answers_at_once_solo_loop) ... [strata] the engine reported an error: refused: images are not enabled +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] 400 invalid request: refused: images are not enabled +[strata] done: 10 tokens in 1 s (10.5 tok/s) (stop, cancel=False) +ok +test_an_old_request_meets_a_restart_in_progress (test_restart_waiters.RestartWaiters.test_an_old_request_meets_a_restart_in_progress) +A request that learns its engine died while restart() has already taken the process (self.proc is None): ... ok +test_killed_while_reading_alone (test_restart_waiters.RestartWaiters.test_killed_while_reading_alone) ... [strata] the engine stopped unexpectedly (exit code -9). The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] answering: 9 of max 64 tokens, 8.6 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.5 tok/s) (stop, cancel=False) +ok +test_killed_while_reading_with_others_arriving (test_restart_waiters.RestartWaiters.test_killed_while_reading_with_others_arriving) ... [strata] the engine stopped unexpectedly (exit code -9). The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine stopped unexpectedly (exit code -9). The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 1 tokens in 0 s (3.9 tok/s) (error, cancel=False) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] answering: 8 of max 64 tokens, 8.6 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.3 tok/s) (stop, cancel=False) +ok +test_solo_engine_is_unchanged (test_restart_waiters.RestartWaiters.test_solo_engine_is_unchanged) ... [strata] parallel requests: 4 asked, the engine runs one at a time (its log says why) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] parallel requests: 4 asked, the engine runs one at a time (its log says why) +[strata] the engine is running again +[strata] answering: 9 of max 64 tokens, 9.4 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.1 tok/s) (stop, cancel=False) +ok +test_waiters_get_a_clean_503_when_the_engine_stays_down_2_slots (test_restart_waiters.RestartWaiters.test_waiters_get_a_clean_503_when_the_engine_stays_down_2_slots) ... [strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 10 tokens in 1 s (9.3 tok/s) (stop, cancel=False) +ok +test_waiters_get_a_clean_503_when_the_engine_stays_down_4_slots (test_restart_waiters.RestartWaiters.test_waiters_get_a_clean_503_when_the_engine_stays_down_4_slots) ... [strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 10 tokens in 1 s (10.1 tok/s) (stop, cancel=False) +ok +test_waiters_go_on_with_the_new_engine_2_slots (test_restart_waiters.RestartWaiters.test_waiters_go_on_with_the_new_engine_2_slots) ... [strata] answering: 4 of max 64 tokens, 9.9 tok/s, 1 s +[strata] answering: 4 of max 64 tokens, 6.7 tok/s, 1 s +[strata] done: 10 tokens in 2 s (11.1 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 2 s (9.2 tok/s) (stop, cancel=False) +[strata] reading the prompt: 32 of 48 tokens, 2 s so far +[strata] answering: 9 of max 64 tokens, 8.8 tok/s, 3 s +[strata] done: 10 tokens in 3 s (8.4 tok/s) (stop, cancel=False) +ok +test_waiters_go_on_with_the_new_engine_4_slots (test_restart_waiters.RestartWaiters.test_waiters_go_on_with_the_new_engine_4_slots) ... [strata] answering: 2 of max 64 tokens, 3.5 tok/s, 1 s +[strata] answering: 2 of max 64 tokens, 7.1 tok/s, 1 s +[strata] answering: 2 of max 64 tokens, 4.4 tok/s, 1 s +[strata] done: 10 tokens in 2 s (6.3 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 2 s (7.7 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 2 s (6.8 tok/s) (stop, cancel=False) +ok +test_null_is_the_default_again (test_runconfig.Apply.test_null_is_the_default_again) ... ok +test_only_the_named_keys_change (test_runconfig.Apply.test_only_the_named_keys_change) ... ok +test_refused (test_runconfig.Apply.test_refused) ... ok +test_view (test_runconfig.Apply.test_view) ... ok +test_only_from_strata_s_own_page (test_runconfig.Http.test_only_from_strata_s_own_page) ... ok +test_read_and_write_from_the_page (test_runconfig.Http.test_read_and_write_from_the_page) ... ok +test_the_api_key (test_runconfig.Http.test_the_api_key) ... ok +test_the_page_has_the_view (test_runconfig.Http.test_the_page_has_the_view) ... ok +test_without_a_run_config (test_runconfig.Http.test_without_a_run_config) ... ok +test_accepted (test_security.HostCheck.test_accepted) ... ok +test_any_host (test_security.HostCheck.test_any_host) ... ok +test_refused (test_security.HostCheck.test_refused) ... ok +test_allowed_hosts_of (test_security.HostNames.test_allowed_hosts_of) ... ok +test_host_name (test_security.HostNames.test_host_name) ... ok +test_names_from_the_bind_address (test_security.HostNames.test_names_from_the_bind_address) ... ok +test_accepted (test_security.OriginCheck.test_accepted) ... ok +test_refused (test_security.OriginCheck.test_refused) ... ok +test_a_trusted_origin_s_host_is_allowed (test_security.OverHttp.test_a_trusted_origin_s_host_is_allowed) ... ok +test_allowed_hosts_and_wildcard (test_security.OverHttp.test_allowed_hosts_and_wildcard) ... ok +test_configured_origins_pass (test_security.OverHttp.test_configured_origins_pass) ... ok +test_control_body_arrives_before_the_operation_and_reply (test_security.OverHttp.test_control_body_arrives_before_the_operation_and_reply) ... ok +test_cross_site_page_is_refused (test_security.OverHttp.test_cross_site_page_is_refused) ... ok +test_curl_and_sdks_without_origin (test_security.OverHttp.test_curl_and_sdks_without_origin) ... ok +test_local_names_and_ips_pass (test_security.OverHttp.test_local_names_and_ips_pass) ... ok +test_no_host_header_passes (test_security.OverHttp.test_no_host_header_passes) ... ok +test_own_and_local_pages_pass (test_security.OverHttp.test_own_and_local_pages_pass) ... ok +test_rebinding_name_is_refused_everywhere (test_security.OverHttp.test_rebinding_name_is_refused_everywhere) ... ok +test_refusal_names_both_ways_out (test_security.OverHttp.test_refusal_names_both_ways_out) ... ok +test_slots_need_json_from_the_own_page (test_security.OverHttp.test_slots_need_json_from_the_own_page) ... ok +test_unload_and_load_need_json_from_the_own_page (test_security.OverHttp.test_unload_and_load_need_json_from_the_own_page) ... ok +test_with_a_key_the_host_check_is_off (test_security.OverHttp.test_with_a_key_the_host_check_is_off) ... ok +test_with_a_key_the_key_decides (test_security.OverHttp.test_with_a_key_the_key_decides) ... ok +test_free_vram_on_hip (test_server.AmdTelemetry.test_free_vram_on_hip) ... ok +test_pcie_link_is_the_narrowest_slowest_hop (test_server.AmdTelemetry.test_pcie_link_is_the_narrowest_slowest_hop) +#1733: the link is the slowest/narrowest hop between the root port and the card; the bottleneck is only ... ok +test_readings (test_server.AmdTelemetry.test_readings) ... ok +test_a_body_a_handler_read_is_not_waited_for (test_server.AnswerBeforeTheBody.test_a_body_a_handler_read_is_not_waited_for) +#594: /config and /load read their body themselves: the close must not wait for the drain's 5 s. ... [strata] model unloaded; the next request loads it again +ok +test_a_body_that_comes_in_drops_does_not_hold_the_connection (test_server.AnswerBeforeTheBody.test_a_body_that_comes_in_drops_does_not_hold_the_connection) +A client that announces a body and sends it a byte at a time is let go after DRAIN_SECONDS, with its ... ok +test_a_host_the_server_does_not_answer_to (test_server.AnswerBeforeTheBody.test_a_host_the_server_does_not_answer_to) ... [strata] refused a request for Host 'rebind.example.com' from 127.0.0.1: not a name this server answers to (add it to "allowed_hosts" in the config or STRATA_ALLOWED_HOSTS, or set an API key) +ok +test_a_method_with_no_handler (test_server.AnswerBeforeTheBody.test_a_method_with_no_handler) ... ok +test_a_wrong_key (test_server.AnswerBeforeTheBody.test_a_wrong_key) ... ok +test_a_wrong_key_and_a_body_of_megabytes (test_server.AnswerBeforeTheBody.test_a_wrong_key_and_a_body_of_megabytes) +An agent client's conversation, or one screenshot, is several MiB, and a rotated key is when the 401 matters. ... ok +test_load_and_unload (test_server.AnswerBeforeTheBody.test_load_and_unload) ... [strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +ok +test_not_the_apps_own_page (test_server.AnswerBeforeTheBody.test_not_the_apps_own_page) ... ok +test_a_blank_key_is_an_error_not_no_key (test_server.ApiKeyForms.test_a_blank_key_is_an_error_not_no_key) ... ok +test_a_key_outside_ascii_matches_as_utf8_and_as_latin1 (test_server.ApiKeyForms.test_a_key_outside_ascii_matches_as_utf8_and_as_latin1) ... ok +test_a_utf8_key_over_http (test_server.ApiKeyForms.test_a_utf8_key_over_http) ... ok +test_several_keys_1344 (test_server.ApiKeyForms.test_several_keys_1344) ... ok +test_the_key_loses_what_a_header_cannot_carry (test_server.ApiKeyForms.test_the_key_loses_what_a_header_cannot_carry) ... ok +test_a_budget_that_fits_is_quiet (test_server.BudgetOverRam.test_a_budget_that_fits_is_quiet) ... ok +test_no_budget_argument_is_quiet (test_server.BudgetOverRam.test_no_budget_argument_is_quiet) ... ok +test_over_free_ram_warns (test_server.BudgetOverRam.test_over_free_ram_warns) ... ok +test_over_total_less_headroom_warns (test_server.BudgetOverRam.test_over_total_less_headroom_warns) ... ok +test_history_and_totals (test_server.CancelledRead.test_history_and_totals) ... [strata] done: 10 tokens in 0 s (117487.5 tok/s) (length, cancel=False) +[strata] done: 10 tokens in 0 s (143150.3 tok/s) (length, cancel=False) +[strata] done: 10 tokens in 0 s (138425.9 tok/s) (length, cancel=False) +ok +test_parse_done_read_field (test_server.CancelledRead.test_parse_done_read_field) ... ok +test_prompt_tokens_seen (test_server.CancelledRead.test_prompt_tokens_seen) ... ok +test_a_plain_version_is_kept (test_server.ClaudeCodeBillingStamp.test_a_plain_version_is_kept) ... ok +test_bounds (test_server.ClaudeCodeBillingStamp.test_bounds) ... ok +test_cc_version_tail_is_pinned (test_server.ClaudeCodeBillingStamp.test_cc_version_tail_is_pinned) ... ok +test_cch_is_pinned (test_server.ClaudeCodeBillingStamp.test_cch_is_pinned) ... ok +test_only_at_the_start_of_the_system_prompt (test_server.ClaudeCodeBillingStamp.test_only_at_the_start_of_the_system_prompt) ... ok +test_string_system_is_pinned (test_server.ClaudeCodeBillingStamp.test_string_system_is_pinned) ... ok +test_two_turns_share_the_whole_system_prompt (test_server.ClaudeCodeBillingStamp.test_two_turns_share_the_whole_system_prompt) ... ok +test_non_streamed (test_server.ClientHangUp.test_non_streamed) ... [strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +ok +test_streamed (test_server.ClientHangUp.test_streamed) ... [strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +ok +test_a_tool_schema_that_is_not_an_object_is_a_400 (test_server.ClientShapes.test_a_tool_schema_that_is_not_an_object_is_a_400) ... [strata] 400 invalid request: tools[0] (f): "parameters" must be an object (the JSON schema of its parameters), not str +[strata] 400 invalid request: tools[0] (f): "parameters" must be an object (the JSON schema of its parameters), not list +[strata] 400 invalid request: tools[0] (f): "parameters" must be an object (the JSON schema of its parameters), not int +[strata] 400 invalid request: tools[0] (f): "input_schema" must be an object (the JSON schema of its parameters), not str +[strata] 400 invalid request: tools[0] (f): "input_schema" must be an object (the JSON schema of its parameters), not list +[strata] 400 invalid request: tools[0] (f): "input_schema" must be an object (the JSON schema of its parameters), not bool +[strata] done: 8 tokens in 0 s (294337.1 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (349525.3 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (364722.1 tok/s) (length, cancel=False) +ok +test_anthropic_mid_conversation_system (test_server.ClientShapes.test_anthropic_mid_conversation_system) ... [strata] done: 12 tokens in 0 s (351969.6 tok/s) (stop, cancel=False) +ok +test_leading_system_unchanged (test_server.ClientShapes.test_leading_system_unchanged) ... ok +test_malformed_tools_are_a_400 (test_server.ClientShapes.test_malformed_tools_are_a_400) ... [strata] 400 invalid request: tools must be a list of tool objects, each {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools must be a list of tool objects, each {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools must be a list of tool objects, each {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools must be a list of tool objects, each {"name": ..., "input_schema": {...}} +[strata] 400 invalid request: tools must be a list of tool objects, each {"name": ..., "input_schema": {...}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"name": ..., "input_schema": {...}} +[strata] 400 invalid request: tools must be a list of tool objects, each {"name": ..., "input_schema": {...}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"name": ..., "input_schema": {...}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"name": ..., "input_schema": {...}} +[strata] done: 8 tokens in 0 s (399457.5 tok/s) (length, cancel=False) +ok +test_messages_as_a_json_string_over_http (test_server.ClientShapes.test_messages_as_a_json_string_over_http) ... [strata] done: 12 tokens in 0 s (430185.0 tok/s) (stop, cancel=False) +[strata] 400 invalid request: messages must be a list of objects +[strata] done: 12 tokens in 0 s (430185.0 tok/s) (stop, cancel=False) +[strata] 400 invalid request: messages must be a list of objects +ok +test_messages_sent_as_a_json_string (test_server.ClientShapes.test_messages_sent_as_a_json_string) ... ok +test_no_user_turn_is_a_400 (test_server.ClientShapes.test_no_user_turn_is_a_400) ... [strata] 400 invalid request: No user query found in messages. +[strata] done: 12 tokens in 0 s (384211.1 tok/s) (stop, cancel=False) +ok +test_openai_late_developer_and_system (test_server.ClientShapes.test_openai_late_developer_and_system) ... [strata] done: 12 tokens in 0 s (445412.8 tok/s) (stop, cancel=False) +ok +test_query_string (test_server.ClientShapes.test_query_string) ... [strata] done: 12 tokens in 0 s (441505.7 tok/s) (stop, cancel=False) +[strata] done: 12 tokens in 0 s (479349.0 tok/s) (stop, cancel=False) +[strata] done: 12 tokens in 0 s (461758.2 tok/s) (stop, cancel=False) +ok +test_tool_call_malformed_arguments_fallback (test_server.ClientShapes.test_tool_call_malformed_arguments_fallback) ... ok +test_vision_temp_image_removed_when_the_pipe_fails (test_server.ClientShapes.test_vision_temp_image_removed_when_the_pipe_fails) ... ok +test_well_formed_tools_still_work (test_server.ClientShapes.test_well_formed_tools_still_work) ... [strata] done: 8 tokens in 0 s (322638.8 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (381300.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (364722.1 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (419430.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (364722.1 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (381300.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (424739.6 tok/s) (length, cancel=False) +ok +test_a_blank_key_is_refused (test_server.ConfigApiKey.test_a_blank_key_is_refused) ... ok +test_a_key_and_no_key (test_server.ConfigApiKey.test_a_key_and_no_key) ... ok +test_an_empty_key_is_a_warning (test_server.ConfigApiKey.test_an_empty_key_is_a_warning) ... ok +test_config_level_fills_in_what_a_request_leaves_out (test_server.ConfigReasoningEffort.test_config_level_fills_in_what_a_request_leaves_out) ... ok +test_default_changes_nothing (test_server.ConfigReasoningEffort.test_default_changes_nothing) ... ok +test_shared_settings_win (test_server.ConfigReasoningEffort.test_shared_settings_win) ... ok +test_the_flag_in_args_becomes_the_key (test_server.ConfigReasoningEffort.test_the_flag_in_args_becomes_the_key) ... ok +test_a_non_streamed_answer_says_the_connection_closes (test_server.ConnectionClose.test_a_non_streamed_answer_says_the_connection_closes) ... [strata] done: 13 tokens in 0 s (361099.0 tok/s) (stop, cancel=False) +ok +test_a_streamed_answer_says_the_connection_closes (test_server.ConnectionClose.test_a_streamed_answer_says_the_connection_closes) ... [strata] done: 13 tokens in 0 s (294734.9 tok/s) (stop, cancel=False) +ok +test_desktop_detection (test_server.DesktopVramNote.test_desktop_detection) ... ok +test_when_it_applies (test_server.DesktopVramNote.test_when_it_applies) ... ok +test_when_it_does_not (test_server.DesktopVramNote.test_when_it_does_not) ... ok +test_drafts_in_history_and_totals (test_server.DraftCounts.test_drafts_in_history_and_totals) ... [strata] done: 10 tokens in 0 s (178481.0 tok/s) (length, cancel=False) +[strata] done: 10 tokens in 0 s (192399.3 tok/s) (length, cancel=False) +[strata] done: 10 tokens in 0 s (204600.2 tok/s) (length, cancel=False), expert cache 90.0% hit +ok +test_an_older_engine_gets_the_advice_in_words (test_server.DraftHeadHint.test_an_older_engine_gets_the_advice_in_words) ... ok +test_other_failures_and_earlier_starts_add_nothing (test_server.DraftHeadHint.test_other_failures_and_earlier_starts_add_nothing) ... ok +test_the_engines_hint_is_relayed (test_server.DraftHeadHint.test_the_engines_hint_is_relayed) ... ok +test_changing_the_effort_keeps_the_conversation (test_server.EffortAtTheEnd.test_changing_the_effort_keeps_the_conversation) ... ok +test_the_config (test_server.EffortAtTheEnd.test_the_config) ... ok +test_the_default_prompt_is_unchanged (test_server.EffortAtTheEnd.test_the_default_prompt_is_unchanged) ... ok +test_the_next_turn_reuses_the_checkpoint (test_server.EffortAtTheEnd.test_the_next_turn_reuses_the_checkpoint) ... ok +test_the_trailing_turn (test_server.EffortAtTheEnd.test_the_trailing_turn) ... ok +test_engine_err_mid_stream (test_server.EngineDeath.test_engine_err_mid_stream) +The engine's ERR line after the stream started reaches the client as an error event (it used to be a ... [strata] the engine reported an error: verify: layer 31 never rang (an illegal memory access was encountered) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] the engine reported an error: verify: layer 31 never rang (an illegal memory access was encountered) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +ok +test_error_then_restart (test_server.EngineDeath.test_error_then_restart) ... [strata] the engine stopped unexpectedly (exit code -9). The next request starts the engine again. +[strata] done: 5 tokens in 0 s (151967.5 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code None); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 50 tokens in 0 s (279620.3 tok/s) (length, cancel=False) +ok +test_a_budget_that_fits_is_unchanged (test_server.FitMaxTokens.test_a_budget_that_fits_is_unchanged) ... [strata] done: 50 tokens in 0 s (142083.5 tok/s) (length, cancel=False) +ok +test_no_room_is_still_a_400 (test_server.FitMaxTokens.test_no_room_is_still_a_400) ... [strata] done: 1 tokens in 0 s (66576.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (4088 tokens) leaves no room to answer in the context (4096); requests are never truncated +ok +test_overshoot_is_clamped_to_the_room (test_server.FitMaxTokens.test_overshoot_is_clamped_to_the_room) ... [strata] done: 2011 tokens in 0 s (16933.2 tok/s) (stop, cancel=False) +[strata] done: 2011 tokens in 0 s (19184.9 tok/s) (stop, cancel=False) +ok +test_a_call_cut_by_max_tokens_is_not_a_tool_call (test_server.ForcedToolChoice.test_a_call_cut_by_max_tokens_is_not_a_tool_call) ... [strata] done: 5 tokens in 0 s (160087.9 tok/s) (length, cancel=False) +ok +test_a_thinking_that_ends_right_after_a_quote_opens_the_call (test_server.ForcedToolChoice.test_a_thinking_that_ends_right_after_a_quote_opens_the_call) ... [strata] done: 120 tokens in 0 s (67495.8 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (60033.0 tok/s) (stop, cancel=False) +ok +test_anthropic_any_tool_none (test_server.ForcedToolChoice.test_anthropic_any_tool_none) ... [strata] done: 272 tokens in 0 s (59400.7 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (51906.4 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (70249.4 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (61027.6 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (102534.2 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (72897.6 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (115937.3 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (81353.3 tok/s) (stop, cancel=False) +[strata] tool_choice {"type": "tool", "name": "nope"} is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (124845.1 tok/s) (stop, cancel=False) +[strata] tool_choice {"type": "tool", "name": "nope"} is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (83877.8 tok/s) (stop, cancel=False) +ok +test_auto_absent_and_none_leave_the_tools_to_the_model (test_server.ForcedToolChoice.test_auto_absent_and_none_leave_the_tools_to_the_model) ... [strata] done: 203 tokens in 0 s (60240.8 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (69381.0 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (72897.6 tok/s) (stop, cancel=False) +ok +test_required_and_named_after_the_thinking (test_server.ForcedToolChoice.test_required_and_named_after_the_thinking) ... [strata] done: 272 tokens in 0 s (52390.3 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (43132.4 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (61359.2 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (51371.2 tok/s) (stop, cancel=False) +ok +test_the_budget_wrap_up_opens_the_call (test_server.ForcedToolChoice.test_the_budget_wrap_up_opens_the_call) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 181 tokens in 0 s (131253.3 tok/s) (stop, cancel=False) +ok +test_the_flat_shape_names_a_function (test_server.ForcedToolChoice.test_the_flat_shape_names_a_function) ... [strata] done: 272 tokens in 0 s (102493.1 tok/s) (stop, cancel=False) +ok +test_values_it_cannot_honour_act_as_auto (test_server.ForcedToolChoice.test_values_it_cannot_honour_act_as_auto) +Not a 400: an odd tool_choice is logged and the model decides. ... [strata] tool_choice {"type": "function", "function": {"name": "nope"}} is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (131720.9 tok/s) (stop, cancel=False) +[strata] tool_choice {"type": "function"} is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (127195.1 tok/s) (stop, cancel=False) +[strata] tool_choice "required" is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (138649.0 tok/s) (stop, cancel=False) +[strata] tool_choice {"type": "banana"} is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (131558.1 tok/s) (stop, cancel=False) +[strata] tool_choice "sometimes" is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (132934.2 tok/s) (stop, cancel=False) +ok +test_without_thinking_the_prompt_ends_with_the_opening (test_server.ForcedToolChoice.test_without_thinking_the_prompt_ends_with_the_opening) ... [strata] done: 64 tokens in 0 s (106142.9 tok/s) (stop, cancel=False) +[strata] done: 64 tokens in 0 s (72063.2 tok/s) (stop, cancel=False) +[strata] done: 56 tokens in 0 s (118149.4 tok/s) (stop, cancel=False) +[strata] done: 56 tokens in 0 s (91500.2 tok/s) (stop, cancel=False) +ok +test_card_order (test_server.GpuChoice.test_card_order) +#1352: with an auto layer split the faster card (SMs x clock) goes last; ties, a manual split, a ... [strata] layer split: card order 2,0 (the faster card last; "gpu_order": "as_given" keeps 0,2, #1352) +[strata] layer split: card order 2,0 (the faster card last; "gpu_order": "as_given" keeps 0,2, #1352) +[strata] layer split: card order 2,0 (the faster card last; "gpu_order": "as_given" keeps 0,2, #1352) +[strata] layer split: card order 0,2,1 (the faster card last; "gpu_order": "as_given" keeps 0,1,2, #1352) +[strata] layer split: keeping card order 0,1 - the faster card has less VRAM than the one the split ends on (12288 vs 22528 MiB, #1576) +[strata] layer split: card order 1,0 (the faster card last; "gpu_order": "as_given" keeps 0,1, #1352) +[strata] layer split: card order 1,0 (the faster card last; "gpu_order": "as_given" keeps 0,1, #1352) +[strata] layer split: card order 1,0 (the faster card last; "gpu_order": "as_given" keeps 0,1, #1352) +[strata] layer split: keeping card order 0,1,2 - the faster card has less VRAM than the one the split ends on (12288 vs 24576 MiB, #1576) +ok +test_env (test_server.GpuChoice.test_env) ... ok +test_hip_ordinal (test_server.GpuChoice.test_hip_ordinal) +#325: on Windows the HIP ordinal setup resolved wins over the config's "gpu" (an iGPU takes HIP's 0). ... ok +test_kfd_scores (test_server.GpuChoice.test_kfd_scores) +AMD on Linux: the GPU nodes of the KFD topology, SIMDs x clock, node order = the config's index. ... /private/tmp/r39-serve-baseline/serve/server.py:2725: ResourceWarning: unclosed file <_io.TextIOWrapper name='/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpguybtoxu/0/properties' mode='r' encoding='UTF-8'> + props = dict(line.split() for line in open(os.path.join(root, str(n), "properties")) if len(line.split()) == 2) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +/private/tmp/r39-serve-baseline/serve/server.py:2725: ResourceWarning: unclosed file <_io.TextIOWrapper name='/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpguybtoxu/1/properties' mode='r' encoding='UTF-8'> + props = dict(line.split() for line in open(os.path.join(root, str(n), "properties")) if len(line.split()) == 2) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +/private/tmp/r39-serve-baseline/serve/server.py:2725: ResourceWarning: unclosed file <_io.TextIOWrapper name='/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpguybtoxu/2/properties' mode='r' encoding='UTF-8'> + props = dict(line.split() for line in open(os.path.join(root, str(n), "properties")) if len(line.split()) == 2) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +/private/tmp/r39-serve-baseline/serve/server.py:2725: ResourceWarning: unclosed file <_io.TextIOWrapper name='/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpguybtoxu/3/properties' mode='r' encoding='UTF-8'> + props = dict(line.split() for line in open(os.path.join(root, str(n), "properties")) if len(line.split()) == 2) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +ok +test_vision_device (test_server.GpuChoice.test_vision_device) ... ok +test_a_held_back_call_still_pings (test_server.HeldBackArgKeepalive.test_a_held_back_call_still_pings) ... [strata] done: 54 tokens in 0 s (150393.4 tok/s) (stop, cancel=False) +ok +test_an_ordinary_answer_does_not_ping (test_server.HeldBackArgKeepalive.test_an_ordinary_answer_does_not_ping) ... [strata] done: 3 tokens in 0 s (142987.6 tok/s) (stop, cancel=False) +ok +test_healthy_untouched (test_server.HipEnvGuard.test_healthy_untouched) ... ok +test_only_for_hip (test_server.HipEnvGuard.test_only_for_hip) ... strata serve: AMD: HIP_PATH=Z:\definitely\not\here does not exist: not passed to the engine +ok +test_stale_hip_path_dropped (test_server.HipEnvGuard.test_stale_hip_path_dropped) ... strata serve: AMD: HIP_PATH=/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmp4ogc_pen/TheRock/build does not exist: not passed to the engine +strata serve: AMD: HIP_DEVICE_LIB_PATH=/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmp4ogc_pen/TheRock/build does not exist: not passed to the engine +ok +test_unwritable_temp_replaced (test_server.HipEnvGuard.test_unwritable_temp_replaced) ... strata serve: AMD: TEMP=/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmp7etngxgi/no/such/dir is not writable: using /private/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmp7etngxgi/tmp +ok +test_anthropic_tool_result_image (test_server.ImageMarkers.test_anthropic_tool_result_image) +Claude Code's Read of a picture returns it inside a tool_result: it reaches the encoder, in place. ... [strata] a picture a tool returned was left out of the prompt: this server has no image encoder +[strata] a picture a tool returned was left out of the prompt: this server has no image encoder +[strata] a picture a tool returned was left out of the prompt: this image format needs Pillow +[strata] a picture a tool returned was left out of the prompt: the image URL could not be read +ok +test_literal_marker_with_an_image (test_server.ImageMarkers.test_literal_marker_with_an_image) ... ok +test_pictures_are_keyed_by_their_bytes_for_the_engine (test_server.ImageMarkers.test_pictures_are_keyed_by_their_bytes_for_the_engine) +#2228: each picture's key (its normalized bytes' hash: the encoder's file name) goes to the engine as ... ok +test_the_whole_marker_in_text_before_images (test_server.ImageMarkers.test_the_whole_marker_in_text_before_images) ... ok +test_a_page_of_another_site_cannot_name_a_file (test_server.ImageSources.test_a_page_of_another_site_cannot_name_a_file) ... [strata] refused an API request from the web page 'https://evil.example' (no API key; add its host to "allowed_hosts" or its origin to "cors_origins" in the config) +[strata] 400 invalid request: a web page of another origin cannot have a file on this computer read as an image: send it as a data: URL (or add the page's origin to the config's trusted_origins) +[strata] 400 invalid request: a web page of another origin cannot have a file on this computer read as an image: send it as a data: URL (or add the page's origin to the config's trusted_origins) +[strata] 400 invalid request: a web page of another origin cannot have a file on this computer read as an image: send it as a data: URL (or add the page's origin to the config's trusted_origins) +[strata] 400 invalid request: a web page of another origin cannot have a file on this computer read as an image: send it as a data: URL (or add the page's origin to the config's trusted_origins) +[strata] done: 8 tokens in 0 s (332222.1 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (364722.1 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (424739.6 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (399457.5 tok/s) (length, cancel=False) +ok +test_a_url_is_downloaded_outside_the_fifo (test_server.ImageSources.test_a_url_is_downloaded_outside_the_fifo) ... ok +test_a_url_is_read_up_to_the_cap (test_server.ImageSources.test_a_url_is_read_up_to_the_cap) ... /opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +ok +test_network_paths_are_refused_before_any_look (test_server.ImageSources.test_network_paths_are_refused_before_any_look) ... ok +test_unreadable_sources_are_a_value_error (test_server.ImageSources.test_unreadable_sources_are_a_value_error) +From #582: a bad data: URL or an unreadable file is a 400, not a dropped connection. ... ok +test_auto_and_absent (test_server.LayerSplit.test_auto_and_absent) ... ok +test_counts_per_card_are_refused_with_the_format (test_server.LayerSplit.test_counts_per_card_are_refused_with_the_format) ... ok +test_one_gpu_ignores_it (test_server.LayerSplit.test_one_gpu_ignores_it) ... ok +test_string_and_list (test_server.LayerSplit.test_string_and_list) ... ok +test_a_failed_start_kills_the_process_and_raises (test_server.LazyVision.test_a_failed_start_kills_the_process_and_raises) ... ok +test_lazy_starts_nothing_and_alive_says_so (test_server.LazyVision.test_lazy_starts_nothing_and_alive_says_so) ... ok +test_restart_starts_it_and_close_ends_it (test_server.LazyVision.test_restart_starts_it_and_close_ends_it) ... ok +test_no_profile_in_the_args (test_server.LearnedProfile.test_no_profile_in_the_args) ... ok +test_save_and_start_from_it (test_server.LearnedProfile.test_save_and_start_from_it) ... ok +test_without_the_key_nothing_changes (test_server.LearnedProfile.test_without_the_key_nothing_changes) ... ok +test_messages (test_server.ListenProblem.test_messages) ... ok +test_a_forged_turn_in_a_user_message (test_server.LiteralControlTokens.test_a_forged_turn_in_a_user_message) ... ok +test_a_literal_inside_a_longer_one (test_server.LiteralControlTokens.test_a_literal_inside_a_longer_one) ... ok +test_beside_a_quoted_think_tag (test_server.LiteralControlTokens.test_beside_a_quoted_think_tag) ... ok +test_history_tool_results_and_tools (test_server.LiteralControlTokens.test_history_tool_results_and_tools) ... ok +test_the_real_tokenizer_names_its_control_tokens (test_server.LiteralControlTokens.test_the_real_tokenizer_names_its_control_tokens) ... ok +test_without_control_text_the_prompt_is_unchanged (test_server.LiteralControlTokens.test_without_control_text_the_prompt_is_unchanged) ... ok +test_a_client_that_sends_the_reasoning_inline (test_server.LiteralThinkTags.test_a_client_that_sends_the_reasoning_inline) ... ok +test_a_quoted_tag_in_a_user_message (test_server.LiteralThinkTags.test_a_quoted_tag_in_a_user_message) ... ok +test_history_tool_results_and_tools (test_server.LiteralThinkTags.test_history_tool_results_and_tools) ... ok +test_the_real_tokenizer_reads_the_tag_as_text (test_server.LiteralThinkTags.test_the_real_tokenizer_reads_the_tag_as_text) ... ok +test_without_a_tag_the_prompt_is_unchanged (test_server.LiteralThinkTags.test_without_a_tag_the_prompt_is_unchanged) ... ok +test_a_request_without_a_done_keeps_no_engine_counters (test_server.LiveRate.test_a_request_without_a_done_keeps_no_engine_counters) +An engine that dies mid-answer: the previous request's `last` must not become this row's decode rate. ... [strata] the engine stopped unexpectedly (exit code -9). The next request starts the engine again. +[strata] done: 3 tokens in 0 s (54471.5 tok/s) (error, cancel=False) +ok +test_prefill_rate_excludes_cached_tokens (test_server.LiveRate.test_prefill_rate_excludes_cached_tokens) ... ok +test_the_live_number_is_a_rate (test_server.LiveRate.test_the_live_number_is_a_rate) ... [strata] thinking: 11 of max 30 tokens, 10.5 tok/s, 1 s +[strata] thinking: 19 of max 30 tokens, 9.3 tok/s, 2 s +[strata] thinking: 28 of max 30 tokens, 9.1 tok/s, 3 s +[strata] done: 30 tokens in 3 s (9.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_silent_mid_answer (test_server.LostStep.test_silent_mid_answer) ... [strata] the engine said nothing for 1 s during the request; the server ended the engine. The engine and the server lost step (issue #481; a very slow PC can raise "engine_silence_s" in the config, 0 = wait forever). If it happens again, please add the end of the engine log to github.com/Niko1221/Strata/issues/481. The next request starts the engine again. +[strata] done: 1 tokens in 1 s (1.7 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 2 tokens in 0 s (22550.0 tok/s) (length, cancel=False) +ok +test_silent_mid_stream (test_server.LostStep.test_silent_mid_stream) ... [strata] the engine said nothing for 1 s during the request; the server ended the engine. The engine and the server lost step (issue #481; a very slow PC can raise "engine_silence_s" in the config, 0 = wait forever). If it happens again, please add the end of the engine log to github.com/Niko1221/Strata/issues/481. The next request starts the engine again. +[strata] done: 1 tokens in 1 s (1.6 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 2 tokens in 0 s (17810.2 tok/s) (length, cancel=False) +ok +test_stop_never_acknowledged (test_server.LostStep.test_stop_never_acknowledged) ... [strata] the engine did not finish the request after it was stopped (STOP) within 0 s; the server ended the engine. The engine and the server lost step (issue #481; a very slow PC can raise "engine_silence_s" in the config, 0 = wait forever). If it happens again, please add the end of the engine log to github.com/Niko1221/Strata/issues/481. The next request starts the engine again. +[strata] done: 2 tokens in 1 s (3.5 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 2 tokens in 0 s (29330.8 tok/s) (length, cancel=False) +ok +test_anthropic_thinks_only_when_asked (test_server.MaxTokens.test_anthropic_thinks_only_when_asked) ... ok +test_count_tokens_is_the_prompt_messages_reads (test_server.MaxTokens.test_count_tokens_is_the_prompt_messages_reads) ... [strata] done: 8 tokens in 0 s (258111.0 tok/s) (length, cancel=False) +ok +test_debug_log_shows_the_resolved_budget (test_server.MaxTokens.test_debug_log_shows_the_resolved_budget) ... ok +test_explicit_budget_is_honoured (test_server.MaxTokens.test_explicit_budget_is_honoured) ... [strata] done: 50 tokens in 0 s (305262.3 tok/s) (length, cancel=False) +[strata] done: 50 tokens in 0 s (307050.1 tok/s) (length, cancel=False) +[strata] done: 50 tokens in 0 s (308404.7 tok/s) (length, cancel=False) +[strata] done: 50 tokens in 0 s (311150.1 tok/s) (length, cancel=False) +[strata] done: 1500 tokens in 0 s (25163.2 tok/s) (length, cancel=False) +[strata] done: 1500 tokens in 0 s (25177.0 tok/s) (length, cancel=False) +ok +test_explicit_budget_over_the_context_is_rejected (test_server.MaxTokens.test_explicit_budget_over_the_context_is_rejected) ... [strata] 400 invalid request: request exceeds the context window: prompt (246 tokens) + max tokens (4096) exceeds the context (4096); requests are never truncated. Send a smaller max_tokens (at most 3842 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (246 tokens) + max tokens (4096) exceeds the context (4096); requests are never truncated. Send a smaller max_tokens (at most 3842 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +ok +test_request_line_parses_the_engine_summary (test_server.MaxTokens.test_request_line_parses_the_engine_summary) ... ok +test_unset_budget_is_the_rest_of_the_context (test_server.MaxTokens.test_unset_budget_is_the_rest_of_the_context) ... [strata] done: 2011 tokens in 0 s (18638.8 tok/s) (stop, cancel=False) +[strata] done: 2011 tokens in 0 s (18976.5 tok/s) (stop, cancel=False) +[strata] done: 2011 tokens in 0 s (19321.3 tok/s) (stop, cancel=False) +[strata] done: 2011 tokens in 0 s (19284.8 tok/s) (stop, cancel=False) +[strata] done: 2011 tokens in 0 s (19231.9 tok/s) (stop, cancel=False) +[strata] done: 2011 tokens in 0 s (18894.2 tok/s) (stop, cancel=False) +[strata] done: 2011 tokens in 0 s (18700.4 tok/s) (stop, cancel=False) +[strata] done: 2011 tokens in 0 s (19177.6 tok/s) (stop, cancel=False) +[strata] done: 2011 tokens in 0 s (19204.7 tok/s) (stop, cancel=False) +[strata] done: 2011 tokens in 0 s (19227.5 tok/s) (stop, cancel=False) +ok +test_unset_budget_with_a_near_full_prompt (test_server.MaxTokens.test_unset_budget_with_a_near_full_prompt) ... [strata] done: 1 tokens in 0 s (144631.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 1 tokens in 0 s (107546.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 5 tokens in 0 s (308404.7 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (4088 tokens) leaves no room to answer in the context (4096); requests are never truncated +[strata] done: 1 tokens in 0 s (144631.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 5 tokens in 0 s (317750.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (4088 tokens) leaves no room to answer in the context (4096); requests are never truncated +ok +test_models_lists_them (test_server.ModelAliases.test_models_lists_them) ... ok +test_requests_are_answered_under_the_alias (test_server.ModelAliases.test_requests_are_answered_under_the_alias) ... [strata] done: 8 tokens in 0 s (138084.1 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (153919.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (178481.0 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (204600.2 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (199728.8 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (205855.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (209715.2 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (216480.2 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (215092.5 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (228261.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (220752.8 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (228261.4 tok/s) (length, cancel=False) +ok +test_set_aliases (test_server.ModelAliases.test_set_aliases) ... ok +test_without_aliases_nothing_changes (test_server.ModelAliases.test_without_aliases_nothing_changes) ... ok +test_engine_done_line_raises (test_server.NonFiniteLogits.test_engine_done_line_raises) +StrataEngine.generate: `DONE ... nonfinite` raises EngineNonFinite after the tokens before it. ... ok +test_mid_reply_ends_without_retry (test_server.NonFiniteLogits.test_mid_reply_ends_without_retry) ... ok +test_nothing_emitted_is_retried_once (test_server.NonFiniteLogits.test_nothing_emitted_is_retried_once) ... ok +test_twice_is_an_error (test_server.NonFiniteLogits.test_twice_is_an_error) ... ok +test_twice_over_http (test_server.NonFiniteLogits.test_twice_over_http) ... ok +test_history (test_server.PcieShare.test_history) ... ok +test_a_vision_section_starts_the_engine_with_images_1322 (test_server.PeerDevice.test_a_vision_section_starts_the_engine_with_images_1322) ... [strata] the config has a "vision" section but no --vision in its args: added --vision and --vram-reserve-mib 700 (run setup again with --vision to write them) +[strata] the config has a "vision" section but no --vision in its args: added --vision (run setup again with --vision to write them) +[strata] the config has a "vision" section but no --vision in its args: added --vision (run setup again with --vision to write them) +ok +test_both_cards_visible (test_server.PeerDevice.test_both_cards_visible) ... ok +test_no_split_with_a_peer (test_server.PeerDevice.test_no_split_with_a_peer) ... ok +test_split_added_for_several_gpus (test_server.PeerDevice.test_split_added_for_several_gpus) ... ok +test_cancel_before_the_encode_runs_none (test_server.PrepareCancelled.test_cancel_before_the_encode_runs_none) ... ok +test_cancel_between_pictures_stops_the_next_encode (test_server.PrepareCancelled.test_cancel_between_pictures_stops_the_next_encode) ... ok +test_cancel_while_waiting_for_the_turn (test_server.PrepareCancelled.test_cancel_while_waiting_for_the_turn) ... ok +test_without_cancel_nothing_changes (test_server.PrepareCancelled.test_without_cancel_nothing_changes) ... ok +test_totals (test_server.PrometheusMetrics.test_totals) ... [strata] done: 10 tokens in 0 s (270600.3 tok/s) (length, cancel=False) +[strata] done: 10 tokens in 0 s (285326.8 tok/s) (length, cancel=False) +ok +test_a_client_that_asks_gets_one_progress_per_chunk (test_server.PromptProgress.test_a_client_that_asks_gets_one_progress_per_chunk) ... [strata] done: 40 tokens in 0 s (109870.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_a_client_that_does_not_ask_gets_nothing (test_server.PromptProgress.test_a_client_that_does_not_ask_gets_nothing) ... [strata] done: 40 tokens in 0 s (120439.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_nothing_before_the_first_chunk (test_server.PromptProgress.test_nothing_before_the_first_chunk) ... ok +test_the_line_that_says_the_prompt_is_read_is_dropped (test_server.PromptProgress.test_the_line_that_says_the_prompt_is_read_is_dropped) ... ok +test_nothing_under_batch (test_server.PromptProgressBatch.test_nothing_under_batch) ... ok +test_a_reply_that_answered_is_not_touched (test_server.ReasoningCloseRetry.test_a_reply_that_answered_is_not_touched) ... [strata] done: 18 tokens in 0 s (153762.7 tok/s) (stop, cancel=False) +ok +test_a_reply_that_closed_its_thinking_right_after_a_quote_is_not_touched (test_server.ReasoningCloseRetry.test_a_reply_that_closed_its_thinking_right_after_a_quote_is_not_touched) ... [strata] done: 28 tokens in 0 s (201441.7 tok/s) (stop, cancel=False) +ok +test_off_by_default_the_reply_stays_empty (test_server.ReasoningCloseRetry.test_off_by_default_the_reply_stays_empty) ... [strata] done: 36 tokens in 0 s (142313.8 tok/s) (stop, cancel=False) +ok +test_on_it_closes_the_thinking_once_and_answers (test_server.ReasoningCloseRetry.test_on_it_closes_the_thinking_once_and_answers) ... [strata] the reply ended inside its thinking with no answer: closing the thinking once and continuing (reasoning_close_retry) +[strata] done: 64 tokens in 0 s (106437.5 tok/s) (stop, cancel=False) +ok +test_a_long_run_is_ended (test_server.RepeatStop.test_a_long_run_is_ended) ... ok +test_short_runs_and_off (test_server.RepeatStop.test_short_runs_and_off) ... ok +test_context_zero_is_503 (test_server.RestartWindow.test_context_zero_is_503) ... ok +test_failed_restart_keeps_the_known_context (test_server.RestartWindow.test_failed_restart_keeps_the_known_context) +After a restart that failed, max_context is 0 but the engine is not starting: a request is checked against ... [strata] done: 13 tokens in 0 s (191319.1 tok/s) (stop, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/unittest/case.py:615: ResourceWarning: unclosed + result = method() +ResourceWarning: Enable tracemalloc to get the object allocation traceback +ok +test_not_alive_until_ready (test_server.RestartWindow.test_not_alive_until_ready) ... ok +test_restart_retries_a_start_that_fails (test_server.RestartWindow.test_restart_retries_a_start_that_fails) +A dead engine's VRAM is freed only when its process is gone, so a new engine started at once can exit ... [strata] the engine did not start (try 1 of 3); again in 0 s +ok +test_checkpoint_key (test_server.SamplingKeys.test_checkpoint_key) ... ok +test_penalty_window (test_server.SamplingKeys.test_penalty_window) ... ok +test_top_k (test_server.SamplingKeys.test_top_k) ... ok +test_tune_aux_cpus (test_server.SamplingKeys.test_tune_aux_cpus) ... ok +test_tune_keys (test_server.SamplingKeys.test_tune_keys) ... ok +test_chars_end_inside_a_message (test_server.SharedPrefix.test_chars_end_inside_a_message) ... ok +test_messages_resolve_to_the_boundary_token (test_server.SharedPrefix.test_messages_resolve_to_the_boundary_token) ... ok +test_pin_key (test_server.SharedPrefix.test_pin_key) ... ok +test_tokens_and_unusable_prefixes (test_server.SharedPrefix.test_tokens_and_unusable_prefixes) ... ok +test_without_the_field_nothing_changes (test_server.SharedPrefix.test_without_the_field_nothing_changes) ... ok +test_a_body_without_the_defaults_wrapper_is_rejected (test_server.SharedSettings.test_a_body_without_the_defaults_wrapper_is_rejected) +A body with no "defaults" key is a mistake, not the documented "clear them" ({"defaults": null}). ... [strata] other apps now use the Chat settings: temperature=0.3 +[strata] done: 16 tokens in 0 s (321095.0 tok/s) (stop, cancel=False) +ok +test_a_trusted_origin_is_strata_s_own_page (test_server.SharedSettings.test_a_trusted_origin_is_strata_s_own_page) ... [strata] other apps use their own settings again +ok +test_off_again (test_server.SharedSettings.test_off_again) ... [strata] other apps now use the Chat settings: temperature=0.3 +[strata] other apps use their own settings again +[strata] done: 16 tokens in 0 s (345922.0 tok/s) (stop, cancel=False) +ok +test_only_strata_s_own_page_may_set_them (test_server.SharedSettings.test_only_strata_s_own_page_may_set_them) ... [strata] other apps now use the Chat settings: temperature=1.0 +ok +test_other_apps_get_the_chat_settings (test_server.SharedSettings.test_other_apps_get_the_chat_settings) ... [strata] other apps now use the Chat settings: temperature=0.3, top_p=0.9, top_k=10, seed=7, max_tokens=77, reasoning_effort=low, experimental_speed_projection=False +[strata] done: 16 tokens in 0 s (409200.4 tok/s) (stop, cancel=False) +[strata] done: 5 tokens in 0 s (327680.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_proxy_headers_do_not_make_a_page_strata_s_own (test_server.SharedSettings.test_proxy_headers_do_not_make_a_page_strata_s_own) ... ok +test_they_need_the_key_when_one_is_set (test_server.SharedSettings.test_they_need_the_key_when_one_is_set) ... [strata] other apps now use the Chat settings: temperature=1.0 +ok +test_before_load_failure_can_be_ignored_by_config (test_server.SharingTheGpu.test_before_load_failure_can_be_ignored_by_config) ... [strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(1) +[strata] the before_load command failed (it exited with code 1); loading anyway +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 13 tokens in 0 s (217234.9 tok/s) (stop, cancel=False) +ok +test_before_load_runs_first (test_server.SharingTheGpu.test_before_load_runs_first) ... [strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c open('/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpuhv_gxzy/ran', 'w').close() +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 13 tokens in 0 s (184833.7 tok/s) (stop, cancel=False) +ok +test_before_load_strict_refuses_even_with_free_vram (test_server.SharingTheGpu.test_before_load_strict_refuses_even_with_free_vram) ... [strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(7) +[strata] the before_load command failed (it exited with code 7); the model stays unloaded and this request is answered 503 - try again (before_load_strict) +ok +test_before_load_timeout_is_configurable_and_a_timeout_refuses (test_server.SharingTheGpu.test_before_load_timeout_is_configurable_and_a_timeout_refuses) ... [strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import time; time.sleep(30) +[strata] the before_load command failed (it did not finish in 0.5 s (before_load_timeout_s)); the model stays unloaded and this request is answered 503 - try again (before_load_strict) +ok +test_failing_before_load_with_free_vram_still_loads (test_server.SharingTheGpu.test_failing_before_load_with_free_vram_still_loads) +#1997: a peer that is not running makes `curl` exit 7; the card has room, so the model loads. ... [strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(7) +[strata] the before_load command failed (it exited with code 7); checking the free VRAM +[strata] the before_load command failed (it exited with code 7), but 20000 MiB of VRAM is free: loading +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 13 tokens in 0 s (216372.8 tok/s) (stop, cancel=False) +ok +test_failing_before_load_with_full_vram_refuses_with_503_and_retry_after (test_server.SharingTheGpu.test_failing_before_load_with_full_vram_refuses_with_503_and_retry_after) ... [strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(3) +[strata] the before_load command failed (it exited with code 3); checking the free VRAM +[strata] the before_load command failed (it exited with code 3) and only 2000 MiB of VRAM is free (the model needs 8000): not loading; 503 +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(3) +[strata] the before_load command failed (it exited with code 3); checking the free VRAM +[strata] the before_load command failed (it exited with code 3), but 20000 MiB of VRAM is free: loading +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 13 tokens in 0 s (154903.3 tok/s) (stop, cancel=False) +ok +test_failing_before_load_without_a_vram_limit_loads (test_server.SharingTheGpu.test_failing_before_load_without_a_vram_limit_loads) ... [strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(7) +[strata] the before_load command failed (it exited with code 7); loading anyway (set min_free_vram_mib to have the free VRAM checked, or before_load_strict to refuse) +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 13 tokens in 0 s (170927.7 tok/s) (stop, cancel=False) +ok +test_idle_unload (test_server.SharingTheGpu.test_idle_unload) ... [strata] the model unloads after 1 s without requests +[strata] model unloaded after 1 s idle; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 13 tokens in 0 s (127695.4 tok/s) (stop, cancel=False) +ok +test_load_endpoint (test_server.SharingTheGpu.test_load_endpoint) ... [strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +ok +test_min_free_vram_refuses_to_load (test_server.SharingTheGpu.test_min_free_vram_refuses_to_load) ... [strata] model unloaded; the next request loads it again +[strata] model unloaded after 1 s idle; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 13 tokens in 0 s (57760.5 tok/s) (stop, cancel=False) +ok +test_min_free_vram_unreadable_loads (test_server.SharingTheGpu.test_min_free_vram_unreadable_loads) ... [strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 13 tokens in 0 s (149386.2 tok/s) (stop, cancel=False) +ok +test_off_by_default (test_server.SharingTheGpu.test_off_by_default) ... ok +test_unload_refused_while_a_request_runs (test_server.SharingTheGpu.test_unload_refused_while_a_request_runs) ... ok +test_unload_then_the_next_request_loads (test_server.SharingTheGpu.test_unload_then_the_next_request_loads) ... [strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 13 tokens in 0 s (178189.4 tok/s) (stop, cancel=False) +ok +test_unload_wait_gives_up_after_the_given_seconds (test_server.SharingTheGpu.test_unload_wait_gives_up_after_the_given_seconds) ... ok +test_unload_wait_waits_for_the_running_turn (test_server.SharingTheGpu.test_unload_wait_waits_for_the_running_turn) +#1997: POST /unload answers busy at once (409); with ?wait=1 it waits for the turn and then unloads. ... [strata] model unloaded; the next request loads it again +ok +test_vision_encoder_unloads_and_starts_first (test_server.SharingTheGpu.test_vision_encoder_unloads_and_starts_first) ... [strata] model unloaded; the next request loads it again +[strata] starting the vision encoder ... +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 13 tokens in 0 s (358723.4 tok/s) (stop, cancel=False) +ok +test_a_busy_gpu_is_not_a_frozen_engine (test_server.SilentEngine.test_a_busy_gpu_is_not_a_frozen_engine) ... [strata] the engine has said nothing for 1 s but is still working; it is not ended (the limit is engine_silence_s). If it is stuck, "engine_silence_s" can be set lower. +ok +test_a_long_first_chunk_is_allowed (test_server.SilentEngine.test_a_long_first_chunk_is_allowed) ... ok +test_a_stop_never_acknowledged (test_server.SilentEngine.test_a_stop_never_acknowledged) ... ok +test_a_working_silent_engine_is_not_ended_by_the_stall_watchdog (test_server.SilentEngine.test_a_working_silent_engine_is_not_ended_by_the_stall_watchdog) ... [strata] the engine has said nothing for 1 s but is still working; it is not ended (the limit is engine_silence_s). If it is stuck, "engine_silence_s" can be set lower. +ok +test_config (test_server.SilentEngine.test_config) ... ok +test_death_note_finds_the_watchdog_line_under_a_long_trace (test_server.SilentEngine.test_death_note_finds_the_watchdog_line_under_a_long_trace) ... ok +test_engine_frozen_samples (test_server.SilentEngine.test_engine_frozen_samples) ... ok +test_frozen_engine_is_ended_by_the_stall_watchdog (test_server.SilentEngine.test_frozen_engine_is_ended_by_the_stall_watchdog) ... ok +test_no_reading_means_no_kill_and_zero_is_off (test_server.SilentEngine.test_no_reading_means_no_kill_and_zero_is_off) ... ok +test_prompt_chunks_set_the_wait (test_server.SilentEngine.test_prompt_chunks_set_the_wait) ... ok +test_silence_mid_answer_ends_the_engine (test_server.SilentEngine.test_silence_mid_answer_ends_the_engine) ... ok +test_zero_waits_as_before (test_server.SilentEngine.test_zero_waits_as_before) ... ok +test_a_crash_is_not_a_transient_start (test_server.StartCrashReported.test_a_crash_is_not_a_transient_start) ... ok +test_a_silent_exit_is_retried_without_claiming_the_gpu_runtime (test_server.StartCrashReported.test_a_silent_exit_is_retried_without_claiming_the_gpu_runtime) ... ok +test_sigill_is_reported_not_retried (test_server.StartCrashReported.test_sigill_is_reported_not_retried) ... ok +test_status_words (test_server.StartCrashReported.test_status_words) ... ok +test_the_runtimes_own_failure_is_named (test_server.StartCrashReported.test_the_runtimes_own_failure_is_named) ... ok +test_the_last_lines_of_this_start (test_server.StartFailureLog.test_the_last_lines_of_this_start) ... ok +test_the_start_error_has_them (test_server.StartFailureLog.test_the_start_error_has_them) ... ok +test_a_mapped_start (test_server.StartNarrator.test_a_mapped_start) ... ok +test_an_arena_start (test_server.StartNarrator.test_an_arena_start) ... ok +test_the_words_follow_the_flags (test_server.StartNarrator.test_the_words_follow_the_flags) ... ok +test_abort_then_the_next_request (test_server.StatusHandover.test_abort_then_the_next_request) ... [strata] done: 15 tokens in 0 s (93068.9 tok/s) (disconnect, cancel=False) +[strata] done: 20 tokens in 0 s (232371.4 tok/s) (length, cancel=False) +ok +test_status (test_server.StatusNeedsTheKey.test_status) ... ok +test_a_prefix_that_is_not_a_stop_is_sent (test_server.StopStrings.test_a_prefix_that_is_not_a_stop_is_sent) +"ENDX" begins like the answer's "END" but is not in it: the held tail goes out and nothing is lost. ... [strata] done: 36 tokens in 0 s (166111.0 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (81092.9 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (173159.3 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (92014.0 tok/s) (stop, cancel=False) +ok +test_bad_values_are_a_400 (test_server.StopStrings.test_bad_values_are_a_400) ... [strata] 400 invalid request: stop takes at most 4 strings +[strata] 400 invalid request: stop must be a string or a list of strings +[strata] 400 invalid request: stop_sequences must be a list of strings +[strata] 400 invalid request: stop_sequences must be a list of strings +ok +test_cut_at_the_stop_string (test_server.StopStrings.test_cut_at_the_stop_string) ... [strata] done: 19 tokens in 0 s (195802.9 tok/s) (stop, cancel=False) +[strata] done: 19 tokens in 0 s (111769.7 tok/s) (stop, cancel=False) +[strata] done: 19 tokens in 0 s (190195.2 tok/s) (stop, cancel=False) +[strata] done: 19 tokens in 0 s (143848.0 tok/s) (stop, cancel=False) +ok +test_matcher_holds_back_only_a_possible_prefix (test_server.StopStrings.test_matcher_holds_back_only_a_possible_prefix) ... ok +test_no_stop_and_an_empty_list (test_server.StopStrings.test_no_stop_and_an_empty_list) ... [strata] done: 36 tokens in 0 s (226719.1 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (209134.3 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (205715.2 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (91180.5 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (87787.8 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (87584.1 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (224028.1 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (224028.1 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (222378.4 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (90470.3 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (104639.6 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (103492.1 tok/s) (stop, cancel=False) +ok +test_openai_takes_a_string (test_server.StopStrings.test_openai_takes_a_string) ... [strata] done: 29 tokens in 0 s (234363.8 tok/s) (stop, cancel=False) +ok +test_the_first_of_several (test_server.StopStrings.test_the_first_of_several) ... [strata] done: 24 tokens in 0 s (258111.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (120554.8 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (249784.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (119410.8 tok/s) (stop, cancel=False) +ok +test_sampler_runs_without_disk_counters (test_server.TelemetryNoDisks.test_sampler_runs_without_disk_counters) ... ok +test_probe_errors_are_logged_and_swallowed (test_server.TemplateCaps.test_probe_errors_are_logged_and_swallowed) ... ok +test_a_bad_value_is_a_400 (test_server.ThinkingBudget.test_a_bad_value_is_a_400) ... [strata] 400 invalid request: reasoning_budget_tokens='lots': expected a whole number of tokens (0: no budget) +[strata] 400 invalid request: reasoning_budget_tokens=2.5: expected a whole number of tokens (0: no budget) +[strata] 400 invalid request: reasoning_budget_tokens=True: expected a whole number of tokens (0: no budget) +[strata] 400 invalid request: reasoning_budget_tokens=[1]: expected a whole number of tokens (0: no budget) +ok +test_a_budget_above_max_tokens_is_capped_to_leave_room_for_the_answer (test_server.ThinkingBudget.test_a_budget_above_max_tokens_is_capped_to_leave_room_for_the_answer) +#984: max_tokens 300 and a budget of 5000 never fired: the reply ended while still thinking. The thinking ... ok +test_a_budget_the_thinking_stays_under (test_server.ThinkingBudget.test_a_budget_the_thinking_stays_under) ... [strata] thinking capped at 200 of max_tokens 400 to leave room for the answer (budget 10000) +[strata] done: 211 tokens in 0 s (132524.4 tok/s) (stop, cancel=False) +ok +test_a_budget_well_below_max_tokens_is_not_capped (test_server.ThinkingBudget.test_a_budget_well_below_max_tokens_is_not_capped) ... ok +test_a_reply_cut_while_thinking_is_named_in_the_log (test_server.ThinkingBudget.test_a_reply_cut_while_thinking_is_named_in_the_log) +#530: max tokens reached inside the thinking gives an empty answer; the server log says what helps. ... ok +test_a_shared_budget_must_be_a_whole_number_of_tokens (test_server.ThinkingBudget.test_a_shared_budget_must_be_a_whole_number_of_tokens) ... ok +test_a_shared_budget_reaches_a_request_that_sets_none (test_server.ThinkingBudget.test_a_shared_budget_reaches_a_request_that_sets_none) +The Chat settings shared with other apps may carry a thinking budget too, like max_tokens and the effort. ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (156639.6 tok/s) (stop, cancel=False) +[strata] done: 211 tokens in 0 s (119610.5 tok/s) (stop, cancel=False) +ok +test_a_stop_string_is_found_in_the_answer_after_the_wrap_up (test_server.ThinkingBudget.test_a_stop_string_is_found_in_the_answer_after_the_wrap_up) +#454: the continuation after the thinking budget is cut at a stop string too. ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 108 tokens in 0 s (86200.7 tok/s) (stop, cancel=False) +ok +test_anthropic_stream (test_server.ThinkingBudget.test_anthropic_stream) ... [strata] thinking budget reached (30 tokens): wrapping up the thinking +[strata] done: 122 tokens in 0 s (60693.3 tok/s) (stop, cancel=False) +ok +test_no_room_left_to_answer (test_server.ThinkingBudget.test_no_room_left_to_answer) ... [strata] done: 20 tokens in 0 s (161319.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_off_by_default (test_server.ThinkingBudget.test_off_by_default) ... [strata] done: 211 tokens in 0 s (103026.6 tok/s) (stop, cancel=False) +ok +test_the_budget_wraps_up_the_thinking (test_server.ThinkingBudget.test_the_budget_wraps_up_the_thinking) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (196800.2 tok/s) (stop, cancel=False) +ok +test_the_config_budget_reaches_the_anthropic_and_effort_requests (test_server.ThinkingBudget.test_the_config_budget_reaches_the_anthropic_and_effort_requests) +#984: a config budget applies to a request that sets no budget of its own, whatever its max_tokens and ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (78925.1 tok/s) (stop, cancel=False) +[strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (138736.6 tok/s) (stop, cancel=False) +ok +test_the_config_default_and_a_request_that_turns_it_off (test_server.ThinkingBudget.test_the_config_default_and_a_request_that_turns_it_off) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (111344.4 tok/s) (stop, cancel=False) +[strata] done: 211 tokens in 0 s (79981.8 tok/s) (stop, cancel=False) +ok +test_the_thinking_progress_line_shows_the_budget_not_max_tokens (test_server.ThinkingBudget.test_the_thinking_progress_line_shows_the_budget_not_max_tokens) +#984: 'thinking: x of max 6400' read as if the budget was ignored; while a budget applies the line shows it. ... ok +test_without_thinking_there_is_nothing_to_limit (test_server.ThinkingBudget.test_without_thinking_there_is_nothing_to_limit) ... [strata] done: 17 tokens in 0 s (146413.1 tok/s) (stop, cancel=False) +ok +test_draft_fields (test_server.TimingsDrafts.test_draft_fields) ... ok +test_both_frontends_render_it (test_server.ToolArguments.test_both_frontends_render_it) ... ok +test_helper (test_server.ToolArguments.test_helper) ... ok +test_bare_function_stays_text_where_it_is_not_a_call (test_server.ToolCallRecovery.test_bare_function_stays_text_where_it_is_not_a_call) ... ok +test_bare_function_without_the_wrapper (test_server.ToolCallRecovery.test_bare_function_without_the_wrapper) ... ok +test_broken_json_is_content_not_an_error (test_server.ToolCallRecovery.test_broken_json_is_content_not_an_error) ... ok +test_json_in_the_wrapper (test_server.ToolCallRecovery.test_json_in_the_wrapper) ... ok +test_json_never_closed_is_content (test_server.ToolCallRecovery.test_json_never_closed_is_content) ... ok +test_json_of_an_unknown_tool_is_content (test_server.ToolCallRecovery.test_json_of_an_unknown_tool_is_content) ... ok +test_parameter_as_the_opener (test_server.ToolCallRecovery.test_parameter_as_the_opener) ... ok +test_parameter_opener_that_is_not_a_tool_is_content (test_server.ToolCallRecovery.test_parameter_opener_that_is_not_a_tool_is_content) ... ok +test_two_calls_in_one_wrapper (test_server.ToolCallRecovery.test_two_calls_in_one_wrapper) ... ok +test_a_batch_is_one_call (test_server.ToolCallRecoveryOff.test_a_batch_is_one_call) ... ok +test_drifted_forms_stay_text (test_server.ToolCallRecoveryOff.test_drifted_forms_stay_text) ... ok +test_switch (test_server.ToolCallRecoverySwitch.test_switch) ... ok +test_named_tag_is_content (test_server.ToolCallTagInProse.test_named_tag_is_content) ... ok +test_tag_alone_at_the_end_is_content (test_server.ToolCallTagInProse.test_tag_alone_at_the_end_is_content) ... ok +test_values_keep_the_terminators (test_server.ToolCallTerminators.test_values_keep_the_terminators) ... ok +test_anthropic_shape_and_a_different_set (test_server.ToolOrder.test_anthropic_shape_and_a_different_set) ... ok +test_duplicate_names_and_one_tool_are_left_alone (test_server.ToolOrder.test_duplicate_names_and_one_tool_are_left_alone) ... ok +test_openai_shape_keeps_the_first_order (test_server.ToolOrder.test_openai_shape_keeps_the_first_order) ... ok +test_a_call_cut_at_the_token_limit (test_server.UnfinishedToolCall.test_a_call_cut_at_the_token_limit) +The same cut by max_tokens: "length" / "max_tokens", and the whole (non-streamed) answers leave the call ... [strata] done: 119 tokens in 0 s (147103.5 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (110400.8 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (154574.8 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (105880.8 tok/s) (length, cancel=False) +ok +test_a_cut_call (test_server.UnfinishedToolCall.test_a_cut_call) ... [strata] done: 120 tokens in 0 s (88956.6 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (65365.8 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (97391.0 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (82483.9 tok/s) (stop, cancel=False) +ok +test_a_whole_call (test_server.UnfinishedToolCall.test_a_whole_call) ... [strata] done: 147 tokens in 0 s (73017.8 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (66850.6 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (85621.8 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (68920.5 tok/s) (stop, cancel=False) +ok +test_collect_keeps_calls_whose_arguments_parse (test_server.UnfinishedToolCall.test_collect_keeps_calls_whose_arguments_parse) ... ok +test_parser (test_server.UnfinishedToolCall.test_parser) ... ok +test_a_vision_answer_in_time_is_unchanged (test_server.UntimedReads.test_a_vision_answer_in_time_is_unchanged) ... ok +test_engine_that_never_says_ready_is_ended (test_server.UntimedReads.test_engine_that_never_says_ready_is_ended) ... ok +test_vision_encode_read_times_out (test_server.UntimedReads.test_vision_encode_read_times_out) ... ok +test_vision_ready_read_times_out (test_server.UntimedReads.test_vision_ready_read_times_out) ... ok +test_anthropic (test_server.UsageAndStatus.test_anthropic) ... [strata] done: 13 tokens in 0 s (106496.0 tok/s) (stop, cancel=False), expert cache 90.0% hit +ok +test_cors_for_the_configured_origins_on_the_api_only (test_server.UsageAndStatus.test_cors_for_the_configured_origins_on_the_api_only) ... ok +test_cors_is_off_by_default (test_server.UsageAndStatus.test_cors_is_off_by_default) ... ok +test_no_clock (test_server.UsageAndStatus.test_no_clock) +An engine without a clock (MockEngine): no timings, nothing cached. ... [strata] done: 5 tokens in 0 s (146654.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_openai (test_server.UsageAndStatus.test_openai) ... [strata] done: 13 tokens in 0 s (136999.9 tok/s) (stop, cancel=False), expert cache 90.0% hit +ok +test_openai_stream (test_server.UsageAndStatus.test_openai_stream) ... [strata] done: 13 tokens in 0 s (117260.1 tok/s) (stop, cancel=False), expert cache 90.0% hit +ok +test_origin_lists_are_checked (test_server.UsageAndStatus.test_origin_lists_are_checked) ... ok +test_sse_is_not_buffered_by_proxies (test_server.UsageAndStatus.test_sse_is_not_buffered_by_proxies) ... [strata] done: 13 tokens in 0 s (125926.0 tok/s) (stop, cancel=False), expert cache 90.0% hit +ok +test_v1_status (test_server.UsageAndStatus.test_v1_status) ... [strata] done: 13 tokens in 0 s (144248.6 tok/s) (stop, cancel=False), expert cache 90.0% hit +ok +test_min_tokens_is_passed_only_when_set (test_server.VisionArgs.test_min_tokens_is_passed_only_when_set) ... ok +test_one_request_with_more_than_64_images_keeps_all_its_files (test_server.VisionCacheEviction.test_one_request_with_more_than_64_images_keeps_all_its_files) ... ok +test_plain_encode_still_evicts_the_oldest (test_server.VisionCacheEviction.test_plain_encode_still_evicts_the_oldest) ... ok +test_shutdown_removes_the_work_dir_but_close_keeps_it (test_server.VisionShutdown.test_shutdown_removes_the_work_dir_but_close_keeps_it) ... ok +test_combined_file_is_written_whole_and_removable (test_server.VisionTempFiles.test_combined_file_is_written_whole_and_removable) ... ok +test_combined_file_stays_off_a_small_shm (test_server.VisionTempFiles.test_combined_file_stays_off_a_small_shm) ... ok +test_enc_line_names_files_relative_to_the_encoder_dir (test_server.VisionTempFiles.test_enc_line_names_files_relative_to_the_encoder_dir) ... ok +test_smart_app_control_is_said_in_words (test_server.VisionTempFiles.test_smart_app_control_is_said_in_words) ... ok +test_work_dir_has_no_space_when_temp_has_one (test_server.VisionTempFiles.test_work_dir_has_no_space_when_temp_has_one) ... ok +test_build_model_path_and_slot_status (test_server.WebApp.test_build_model_path_and_slot_status) ... ok +test_discovery_does_not_restart_a_dead_engine (test_server.WebApp.test_discovery_does_not_restart_a_dead_engine) ... ok +test_discovery_needs_the_api_key (test_server.WebApp.test_discovery_needs_the_api_key) ... ok +test_metrics (test_server.WebApp.test_metrics) ... [strata] done: 5 tokens in 0 s (249661.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_metrics_need_the_key_when_one_is_set (test_server.WebApp.test_metrics_need_the_key_when_one_is_set) ... ok +test_model_discovery_and_props (test_server.WebApp.test_model_discovery_and_props) ... ok +test_only_the_app_files_are_served (test_server.WebApp.test_only_the_app_files_are_served) ... ok +test_page_and_files (test_server.WebApp.test_page_and_files) ... ok +test_props_caps_are_ready_for_concurrent_requests (test_server.WebApp.test_props_caps_are_ready_for_concurrent_requests) ... ok +test_props_caps_follow_the_active_template (test_server.WebApp.test_props_caps_follow_the_active_template) ... ok +test_props_total_slots_follows_the_batch_slots (test_server.WebApp.test_props_total_slots_follows_the_batch_slots) ... ok +test_slot_reports_the_context_in_use (test_server.WebApp.test_slot_reports_the_context_in_use) +/slots in llama.cpp's names: a front-end's context meter divides n_prompt_tokens by n_ctx, so a server ... ok +test_slots_lists_every_batch_slot (test_server.WebApp.test_slots_lists_every_batch_slot) ... ok +test_the_endpoints_keep_the_context_while_the_engine_restarts (test_server.WebApp.test_the_endpoints_keep_the_context_while_the_engine_restarts) +#351: max_context is 0 until READY; the endpoints report the last known one. ... ok +test_default_refusal_is_immediate_and_byte_identical (test_session_save_retry.HTTP.test_default_refusal_is_immediate_and_byte_identical) ... ok +test_fifo_blocks_restore_while_save_waits_and_releases_afterwards (test_session_save_retry.HTTP.test_fifo_blocks_restore_while_save_waits_and_releases_afterwards) ... [strata] SAVE waiting for RAM after cache reclaim: need=1125 MiB, limit=30s; floor unchanged +ok +test_opt_in_recovers_native_refusal_and_keeps_response_and_process (test_session_save_retry.HTTP.test_opt_in_recovers_native_refusal_and_keeps_response_and_process) ... [strata] SAVE waiting for RAM after cache reclaim: need=1125 MiB, limit=30s; floor unchanged +ok +test_admission_races_are_bounded_and_keep_latest_error (test_session_save_retry.Retry.test_admission_races_are_bounded_and_keep_latest_error) ... ok +test_delayed_ram_reuses_same_process_path_and_floor (test_session_save_retry.Retry.test_delayed_ram_reuses_same_process_path_and_floor) ... ok +test_memory_probe_is_portable_and_unknown_fails_closed (test_session_save_retry.Retry.test_memory_probe_is_portable_and_unknown_fails_closed) ... ok +test_new_refusal_rechecks_its_larger_estimate (test_session_save_retry.Retry.test_new_refusal_rechecks_its_larger_estimate) ... ok +test_non_admission_error_during_retry_is_returned (test_session_save_retry.Retry.test_non_admission_error_during_retry_is_returned) ... ok +test_persistent_pressure_expires_without_another_save (test_session_save_retry.Retry.test_persistent_pressure_expires_without_another_save) ... ok +test_process_death_or_replacement_ends_wait (test_session_save_retry.Retry.test_process_death_or_replacement_ends_wait) ... ok +test_restore_other_categories_and_published_files_never_retry (test_session_save_retry.Retry.test_restore_other_categories_and_published_files_never_retry) ... ok +test_success_does_not_poll_or_wait (test_session_save_retry.Retry.test_success_does_not_poll_or_wait) ... ok +test_unknown_telemetry_or_estimate_never_retries (test_session_save_retry.Retry.test_unknown_telemetry_or_estimate_never_retries) ... ok +test_reuse_then_least_recently_used (test_slot_alloc.OneGroup.test_reuse_then_least_recently_used) ... ok +test_fills_a_leading_hole_first (test_slot_alloc.Pipelined.test_fills_a_leading_hole_first) ... ok +test_held_tokens_are_ignored (test_slot_alloc.Pipelined.test_held_tokens_are_ignored) ... ok +test_never_opens_a_hole_below_the_pick (test_slot_alloc.Pipelined.test_never_opens_a_hole_below_the_pick) ... ok +test_spreads_over_the_groups_first (test_slot_alloc.Pipelined.test_spreads_over_the_groups_first) ... ok +test_a_blocking_step_gets_its_announced_allowance (test_slots.Slots.test_a_blocking_step_gets_its_announced_allowance) +B: SWAIT gives one non-streaming step (a flush, the device transfer) its own bounded wait. ... ok +test_a_blocking_step_that_never_ends_is_ended (test_slots.Slots.test_a_blocking_step_that_never_ends_is_ended) +B: the allowance is a bound, not a heartbeat: a step that outlives it ends the engine. ... ok +test_a_foreign_page_or_a_form_never_reaches_the_engine (test_slots.Slots.test_a_foreign_page_or_a_form_never_reaches_the_engine) ... ok +test_a_silent_engine_is_ended_not_waited_for (test_slots.Slots.test_a_silent_engine_is_ended_not_waited_for) ... ok +test_a_small_save_reports_its_only_partial_block (test_slots.Slots.test_a_small_save_reports_its_only_partial_block) +B: a file below one 16 MiB block still reports (the start, then its last partial block). ... ok +test_disabled_without_a_slot_save_path (test_slots.Slots.test_disabled_without_a_slot_save_path) ... ok +test_disk_and_ram_refusals_have_their_own_codes (test_slots.Slots.test_disk_and_ram_refusals_have_their_own_codes) ... ok +test_engine_death_is_a_server_error (test_slots.Slots.test_engine_death_is_a_server_error) ... ok +test_engine_gets_an_absolute_path_on_one_line (test_slots.Slots.test_engine_gets_an_absolute_path_on_one_line) ... ok +test_failed_restore_transfer_is_a_server_error (test_slots.Slots.test_failed_restore_transfer_is_a_server_error) ... ok +test_filename_cannot_leave_the_directory (test_slots.Slots.test_filename_cannot_leave_the_directory) ... ok +test_log_lines_are_not_answers (test_slots.Slots.test_log_lines_are_not_answers) ... ok +test_malformed_allowances_end_the_engine (test_slots.Slots.test_malformed_allowances_end_the_engine) ... ok +test_malformed_or_mismatched_reply_is_a_server_error (test_slots.Slots.test_malformed_or_mismatched_reply_is_a_server_error) ... ok +test_missing_file_is_not_found (test_slots.Slots.test_missing_file_is_not_found) ... ok +test_only_slot_zero_and_known_actions (test_slots.Slots.test_only_slot_zero_and_known_actions) ... ok +test_ordinary_names_are_accepted (test_slots.Slots.test_ordinary_names_are_accepted) ... ok +test_out_of_protocol_lines_end_the_engine (test_slots.Slots.test_out_of_protocol_lines_end_the_engine) +R4: a line this exchange does not allow is not a heartbeat, and an answer must hold valid numbers. ... ok +test_progress_ends_the_allowance (test_slots.Slots.test_progress_ends_the_allowance) +B: after a SESSION line the ordinary silence limit applies again (the allowance was for one step). ... ok +test_progress_lines_keep_a_long_save_alive (test_slots.Slots.test_progress_lines_keep_a_long_save_alive) ... ok +test_published_failure_is_reported (test_slots.Slots.test_published_failure_is_reported) +R3: the folder flush failed after the rename - the new file replaced the old one; the answer says so. ... ok +test_refused_file_is_an_error_not_a_success (test_slots.Slots.test_refused_file_is_an_error_not_a_success) ... ok +test_refused_with_parallel_requests (test_slots.Slots.test_refused_with_parallel_requests) +#465: batch slots do not hold the FIFO, so a session file is refused before anything reaches the engine. ... ok +test_relative_save_dir_becomes_absolute (test_slots.Slots.test_relative_save_dir_becomes_absolute) ... ok +test_restore_answers_like_llama_server (test_slots.Slots.test_restore_answers_like_llama_server) ... ok +test_restore_of_a_link_reaches_the_engine_which_decides (test_slots.Slots.test_restore_of_a_link_reaches_the_engine_which_decides) +The server only checks the name exists (lstat); the engine opens it without following a link. ... ok +test_save_writes_into_the_slot_save_path (test_slots.Slots.test_save_writes_into_the_slot_save_path) ... ok +test_slot_work_counts_as_activity_and_shows_in_status (test_slots.Slots.test_slot_work_counts_as_activity_and_shows_in_status) ... ok +test_typed_refusals_have_their_own_codes (test_slots.Slots.test_typed_refusals_have_their_own_codes) +R5: the engine's category decides the status, never the words of the message. ... ok +test_waits_for_a_running_request (test_slots.Slots.test_waits_for_a_running_request) +A save never interleaves with a generation: it takes the same FIFO. ... ok +test_h3_schema_success_and_stream_with_reasoning (test_structured.Structured.test_h3_schema_success_and_stream_with_reasoning) ... [strata] done: 255 tokens in 0 s (58432.4 tok/s) (stop, cancel=False) +[strata] done: 282 tokens in 0 s (60721.5 tok/s) (stop, cancel=False) +ok +test_invalid_generation_is_an_error_and_never_streams_content (test_structured.Structured.test_invalid_generation_is_an_error_and_never_streams_content) ... skipped 'jsonschema is not installed' +test_invalid_schema_is_400_and_local_refs_work (test_structured.Structured.test_invalid_schema_is_400_and_local_refs_work) ... skipped 'jsonschema is not installed' +test_malformed_formats_and_tools_are_rejected_before_loading (test_structured.Structured.test_malformed_formats_and_tools_are_rejected_before_loading) ... [strata] 400 invalid request: response_format must be an object +[strata] 400 invalid request: response_format.type must be text, json_object or json_schema +[strata] 400 invalid request: response_format.json_schema needs a name +[strata] 400 invalid request: response_format.json_schema.strict must be boolean +[strata] 400 invalid request: structured response_format with tools/MCP is not supported +[strata] 400 invalid request: structured response_format with tools/MCP is not supported +[strata] 400 invalid request: send a JSON object +ok +test_object_mode_strict_json_and_truncation (test_structured.Structured.test_object_mode_strict_json_and_truncation) ... [strata] done: 13 tokens in 0 s (132344.5 tok/s) (stop, cancel=False) +[strata] done: 3 tokens in 0 s (99864.4 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 0 s (163840.0 tok/s) (stop, cancel=False) +[strata] done: 14 tokens in 0 s (168736.4 tok/s) (stop, cancel=False) +[strata] done: 12 tokens in 0 s (166661.1 tok/s) (stop, cancel=False) +[strata] done: 15 tokens in 0 s (189501.7 tok/s) (stop, cancel=False) +[strata] done: 6 tokens in 0 s (142987.6 tok/s) (stop, cancel=False) +[strata] done: 15 tokens in 0 s (178734.5 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (192841.6 tok/s) (stop, cancel=False) +[strata] done: 5 tokens in 0 s (161319.4 tok/s) (length, cancel=False) +[strata] done: 7 tokens in 0 s (194437.9 tok/s) (stop, cancel=False) +ok +test_rejected_schema_does_not_load_the_engine_or_retry (test_structured.Structured.test_rejected_schema_does_not_load_the_engine_or_retry) ... [strata] 400 invalid request: response_format.json_schema needs a schema object +[strata] done: 21 tokens in 0 s (328658.1 tok/s) (stop, cancel=False) +ok +test_root_unions_of_object_shapes_are_accepted (test_structured.Structured.test_root_unions_of_object_shapes_are_accepted) ... [strata] done: 40 tokens in 0 s (127777.7 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (127100.1 tok/s) (stop, cancel=False) +[strata] done: 40 tokens in 0 s (143272.6 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (149036.2 tok/s) (stop, cancel=False) +[strata] done: 40 tokens in 0 s (144382.2 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (152125.0 tok/s) (stop, cancel=False) +[strata] done: 40 tokens in 0 s (166771.5 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (180123.5 tok/s) (stop, cancel=False) +[strata] done: 40 tokens in 0 s (156358.0 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (175809.1 tok/s) (stop, cancel=False) +ok +test_schemas_that_allow_non_objects_are_still_rejected (test_structured.Structured.test_schemas_that_allow_non_objects_are_still_rejected) ... [strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +ok +test_the_server_does_not_import_jsonschema_at_start (test_structured.Structured.test_the_server_does_not_import_jsonschema_at_start) ... ok +test_without_jsonschema_an_object_is_still_required (test_structured.Structured.test_without_jsonschema_an_object_is_still_required) ... [strata] done: 255 tokens in 0 s (52858.9 tok/s) (stop, cancel=False) +[strata] done: 26 tokens in 0 s (182973.0 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (146800.6 tok/s) (stop, cancel=False) +[strata] done: 13 tokens in 0 s (194735.5 tok/s) (stop, cancel=False) +ok +test_after_the_thinking (test_tool_name_space.SpacedName.test_after_the_thinking) ... ok +test_inside_the_thinking (test_tool_name_space.SpacedName.test_inside_the_thinking) ... ok +test_message_names_the_cause_from_the_line (test_vision_busy.Busy.test_message_names_the_cause_from_the_line) ... ok +test_message_names_the_cause_from_the_log_tail (test_vision_busy.Busy.test_message_names_the_cause_from_the_log_tail) ... ok +test_other_failures_are_unchanged (test_vision_busy.Busy.test_other_failures_are_unchanged) ... ok +test_start_raises_it (test_vision_busy.Busy.test_start_raises_it) ... ok +test_a_vram_reserve_stacks_on_the_encoders_share (test_vision_elastic.Elastic.test_a_vram_reserve_stacks_on_the_encoders_share) ... [strata] VRAM: 4000 MiB free, expert cache 2000 of 6000 MiB (100 experts) +[strata] images: the expert cache keeps 2048 MiB free for the image encoder (500 of 6000 MiB left) +[strata] images: the image encoder started in 0.0 s +[strata] VRAM: 3048 MiB free, expert cache 2952 of 6000 MiB (952 experts) +[strata] images: the image encoder stopped after 60 s without a picture; the expert cache took its VRAM back (5000 of 6000 MiB) +ok +test_an_engine_that_cannot_shrink_still_gets_its_encoder (test_vision_elastic.Elastic.test_an_engine_that_cannot_shrink_still_gets_its_encoder) ... [strata] images: the expert cache did not make room for the image encoder (VRAM needs an engine started with --vram-elastic (one NVIDIA GPU, --serve)); starting it anyway +[strata] images: the image encoder started in 0.0 s +ok +test_an_idle_encoder_stops_and_the_cache_grows_back_whole (test_vision_elastic.Elastic.test_an_idle_encoder_stops_and_the_cache_grows_back_whole) ... [strata] images: the expert cache keeps 2048 MiB free for the image encoder (3952 of 6000 MiB left) +[strata] images: the image encoder started in 0.0 s +[strata] images: the image encoder stopped after 60 s without a picture; the expert cache took its VRAM back (6000 of 6000 MiB) +[strata] images: the expert cache keeps 2048 MiB free for the image encoder (3952 of 6000 MiB left) +[strata] images: the image encoder started in 0.0 s +ok +test_an_unload_lets_go_of_the_share (test_vision_elastic.Elastic.test_an_unload_lets_go_of_the_share) ... [strata] images: the expert cache keeps 2048 MiB free for the image encoder (3952 of 6000 MiB left) +[strata] images: the image encoder started in 0.0 s +[strata] model unloaded; the next request loads it again +ok +test_the_config_keys (test_vision_elastic.Elastic.test_the_config_keys) ... ok +test_the_first_picture_makes_room_then_starts_the_encoder (test_vision_elastic.Elastic.test_the_first_picture_makes_room_then_starts_the_encoder) ... [strata] images: the expert cache keeps 2048 MiB free for the image encoder (3952 of 6000 MiB left) +[strata] images: the image encoder started in 0.0 s +ok +test_the_load_never_starts_an_elastic_encoder (test_vision_elastic.Elastic.test_the_load_never_starts_an_elastic_encoder) ... ok +test_config_key (test_vram.EngineCommand.test_config_key) ... ok +test_refused_and_died (test_vram.EngineCommand.test_refused_and_died) ... ok +test_the_line_and_the_answer (test_vram.EngineCommand.test_the_line_and_the_answer) ... ok +test_an_engine_without_the_flag_says_so (test_vram.Route.test_an_engine_without_the_flag_says_so) ... [strata] 400 invalid request: VRAM needs an engine started with --vram-elastic (one NVIDIA GPU, --serve) +ok +test_bad_input_busy_and_foreign_pages (test_vram.Route.test_bad_input_busy_and_foreign_pages) ... [strata] 400 invalid request: reserve_mib: a whole number of MiB (0 or more), or null for the start's +[strata] 400 invalid request: reserve_mib: a whole number of MiB (0 or more), or null for the start's +[strata] 400 invalid request: reserve_mib: a whole number of MiB (0 or more), or null for the start's +[strata] 400 invalid request: reserve_mib: a whole number of MiB (0 or more), or null for the start's +[strata] refused an API request from the web page 'https://other.example' (no API key; add its host to "allowed_hosts" or its origin to "cors_origins" in the config) +ok +test_resize_and_status (test_vram.Route.test_resize_and_status) ... [strata] VRAM: 4000 MiB free, expert cache 2000 of 6000 MiB (100 experts) +[strata] VRAM: 700 MiB free, expert cache 5300 of 6000 MiB (3300 experts) +ok +test_unloaded_keeps_the_reserve_for_the_next_load (test_vram.Route.test_unloaded_keeps_the_reserve_for_the_next_load) ... [strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +ok +test_a_contained_child_is_opted_out_of_power_throttling (test_winjob.TestNoThrottle.test_a_contained_child_is_opted_out_of_power_throttling) ... skipped 'power throttling is a Windows setting' +test_children_end_with_a_killed_server (test_winjob.TestWinJob.test_children_end_with_a_killed_server) ... skipped 'job objects are Windows-only' +test_elsewhere_or_without_a_process_it_is_a_no_op (test_winjob.TestWinJob.test_elsewhere_or_without_a_process_it_is_a_no_op) ... skipped 'job objects are Windows-only' + +====================================================================== +FAIL: test_service_holds_for_the_whole_request_including_early_close (test_keepawake.KeepAwakeTests.test_service_holds_for_the_whole_request_including_early_close) +---------------------------------------------------------------------- +Traceback (most recent call last): + File "/private/tmp/r39-serve-baseline/serve/test_keepawake.py", line 87, in test_service_holds_for_the_whole_request_including_early_close + self.assertTrue(wait_for(lambda: rec.calls == [ON, ES_CONTINUOUS])) + ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +AssertionError: False is not true + +---------------------------------------------------------------------- +Ran 785 tests in 494.236s + +FAILED (failures=1, skipped=11) diff --git a/results/round39/serve-baseline.log b/results/round39/serve-baseline.log new file mode 100644 index 0000000..6404327 --- /dev/null +++ b/results/round39/serve-baseline.log @@ -0,0 +1,752 @@ +.....[strata] done: 12 tokens in 0 s (184365.0 tok/s) (stop, cancel=False) +..[strata] other apps now use the Chat settings: temperature=0.5 +.[strata] done: 12 tokens in 0 s (316551.2 tok/s) (stop, cancel=False) +.....[strata] done: 12 tokens in 0 s (316551.2 tok/s) (stop, cancel=False) +.[strata] other apps now use the Chat settings: temperature=0.5 +[strata] other apps now use the Chat settings: temperature=0.5 +[strata] other apps now use the Chat settings: temperature=0.5 +..[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +.[strata] done: 1 tokens in 0 s (72315.6 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) + max tokens (1) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] done: 1 tokens in 0 s (107546.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] done: 1 tokens in 0 s (89240.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) + max tokens (1) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] done: 1 tokens in 0 s (91180.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] done: 1 tokens in 0 s (55188.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 1 tokens in 0 s (47127.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +....sssss....[strata] 400 invalid request: unknown reasoning effort 'huge': use none, low, medium or high +.[strata] done: 24 tokens in 0 s (249784.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (295200.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (328965.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (337796.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (343560.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (342392.2 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (351969.6 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (348315.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (348315.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (315558.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (353204.5 tok/s) (stop, cancel=False) +.[strata] done: 24 tokens in 0 s (359511.8 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (353204.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (264208.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (363405.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (351969.6 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (285975.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (337796.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (337796.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (289262.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (324720.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (328965.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (282762.1 tok/s) (stop, cancel=False) +.[strata] done: 24 tokens in 0 s (328965.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (328965.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (338933.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (332222.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (315558.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (343560.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (353204.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (356962.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (363405.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (363405.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (328965.0 tok/s) (stop, cancel=False) +.[strata] done: 24 tokens in 0 s (358232.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (356962.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (363405.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (368730.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (363405.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (332222.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (332222.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (332222.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (324720.3 tok/s) (stop, cancel=False) +.[strata] done: 24 tokens in 0 s (266305.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (275789.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (289262.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (285975.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (290096.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (278845.7 tok/s) (stop, cancel=False) +.[strata] done: 24 tokens in 0 s (358232.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (342392.2 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (289262.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (334429.6 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (333322.2 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (289262.3 tok/s) (stop, cancel=False) +.[strata] done: 24 tokens in 0 s (358232.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (348315.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (332222.1 tok/s) (stop, cancel=False) +.[strata] done: 24 tokens in 0 s (363405.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (359511.8 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (356962.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (290096.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (337796.3 tok/s) (stop, cancel=False) +.[strata] done: 24 tokens in 0 s (117734.8 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (135664.8 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (152983.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (127421.9 tok/s) (stop, cancel=False) +.[strata] done: 24 tokens in 0 s (319566.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (334429.6 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (337796.3 tok/s) (stop, cancel=False) +.................[strata] the engine stopped unexpectedly (exit code -15). Unrecoverable native verification failure: verify: timed out at layer 18; its GPU waits were released and the GPU finished (#267) The next request starts the engine again. +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -15); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 1 tokens in 0 s (18558.9 tok/s) (length, cancel=False) +...[strata] the engine stopped unexpectedly (exit code -15). Unrecoverable native verification failure: verify: timed out at layer 18; its GPU waits were released and the GPU finished (#267) The next request starts the engine again. +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -15); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 1 tokens in 0 s (24385.5 tok/s) (length, cancel=False) +................................[strata] image: 512x768, n=2, 8 steps, seed 11 +/private/tmp/r39-serve-baseline/serve/server.py:5603: ResourceWarning: unclosed file <_io.TextIOWrapper name=9 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 0.2 s (keep) +.[strata] 400 invalid request: size 999x999: each side a multiple of 32 from 256 to 2048 +[strata] image: 1024x1024, n=1, 8 steps, seed 1394238321 +[strata] image job failed: sd-cli exited with code 1 +[INFO ] loading tensors completed, taking 0.10s +[ERROR] ggml_vulkan: ErrorDeviceLost +/private/tmp/r39-serve-baseline/serve/server.py:5607: ResourceWarning: unclosed file <_io.TextIOWrapper name=9 encoding='UTF-8'> + return +ResourceWarning: Enable tracemalloc to get the object allocation traceback +.[strata] image: 1024x1024, n=1, 8 steps, seed 2022972244 +[strata] image done in 1.7 s (keep) +[strata] done: 8 tokens in 0 s (204600.2 tok/s) (length, cancel=False) +...[strata] image: 1024x1024, n=1, 8 steps, seed 0 +[strata] image done in 0.3 s (keep) +[strata] image: 1024x1024, n=1, 8 steps, seed 1 +[strata] image done in 0.3 s (keep) +[strata] image: 1024x1024, n=1, 8 steps, seed 2 +[strata] image done in 0.3 s (keep) +[strata] image: 1024x1024, n=1, 8 steps, seed 3 +[strata] image done in 0.4 s (keep) +.[strata] image job: the conversation (123 tokens) saved; unloading the model +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] the conversation from before the image job is back (123 tokens) +...[strata] image job: the engine's expert cache gave VRAM back (9000 MiB free) +.[strata] image job: the conversation (123 tokens) saved; unloading the model +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] the conversation from before the image job is back (123 tokens) +.........[strata] done: 3 tokens in 0 s (103991.0 tok/s) (stop, cancel=False) +F[strata] loading the model again (it was unloaded) ... +.[strata] model unloaded; the next request loads it again +.....[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 7 tokens in 0 s (166818.9 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (233016.9 tok/s) (stop, cancel=False) +[strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 7 tokens in 0 s (81330.0 tok/s) (stop, cancel=False) +.[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +....[strata] 400 invalid request: logit_bias: expected an object or a list of [token_id, bias] pairs +[strata] 400 invalid request: logit_bias: this engine does not support token biases; use a supporting CUDA/HIP build +[strata] 400 invalid request: logit_bias: biases must be finite numbers between -100 and 100; false is a pair-list ban +...[strata] MCP server 'missing' did not start: could not start 'strata-no-such-program': [Errno 2] No such file or directory: 'strata-no-such-program'. Its tools are left out; the chat works without them. +[strata] MCP server 'paged': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +...[strata] MCP server 'solo': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'solo' had stopped; starting it again +[strata] MCP server 'solo': 6 tools (echo, fail, sleep, big, die, add) +........[strata] MCP server 'down' did not start: could not reach http://127.0.0.1:9/mcp: . Its tools are left out; the chat works without them. +[strata] MCP server 'web': 6 tools (echo, fail, sleep, big, die, add) +..[strata] done: 13 tokens in 0 s (212992.0 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 83 tokens in 0 s (162070.4 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 120 tokens in 0 s (128857.3 tok/s) (stop, cancel=False) +[strata] tool fake__echo: ok, 5 characters in 0.0 s +[strata] done: 120 tokens in 0 s (153450.1 tok/s) (stop, cancel=False) +[strata] tool fake__echo: ok, 5 characters in 0.0 s +[strata] done: 120 tokens in 0 s (156796.4 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 141 tokens in 0 s (157035.8 tok/s) (stop, cancel=False) +[strata] tool fake__add: ok, 1 characters in 0.0 s +[strata] done: 13 tokens in 0 s (351780.3 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] refused an API request from the web page 'http://evil.example' (no API key; add its host to "allowed_hosts" or its origin to "cors_origins" in the config) +[strata] done: 13 tokens in 0 s (200463.1 tok/s) (stop, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 109 tokens in 0 s (147144.9 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 116 tokens in 0 s (130334.7 tok/s) (stop, cancel=False) +[strata] done: 16 tokens in 0 s (207767.4 tok/s) (stop, cancel=False) +[strata] done: 16 tokens in 0 s (208412.6 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 120 tokens in 0 s (127648.1 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 134 tokens in 0 s (146592.8 tok/s) (stop, cancel=False) +[strata] tool fake__echo: ok, 19 characters in 0.0 s +[strata] done: 31 tokens in 0 s (148259.3 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 84 tokens in 0 s (168979.2 tok/s) (stop, cancel=False) +[strata] tool fake__fail: error, 27 characters in 0.0 s +[strata] done: 21 tokens in 0 s (190650.2 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] done: 115 tokens in 0 s (104290.8 tok/s) (stop, cancel=False) +[strata] tool fake__big: ok, 3,000 characters in 0.0 s (truncated for the model) +[strata] done: 16 tokens in 0 s (175677.7 tok/s) (stop, cancel=False) +.[strata] done: 24 tokens in 0 s (319566.0 tok/s) (stop, cancel=False) +....[strata] 400 invalid request: request exceeds the context window: prompt (183 tokens) + max tokens (20000) exceeds the context (16384); requests are never truncated. Send a smaller max_tokens (at most 16193 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] done: 50 tokens in 0 s (249958.5 tok/s) (stop, cancel=False) +.[strata] done: 50 tokens in 0 s (114410.9 tok/s) (stop, cancel=False) +[strata] done: 50 tokens in 0 s (57614.1 tok/s) (stop, cancel=False) +[strata] done: 50 tokens in 0 s (123725.8 tok/s) (stop, cancel=False) +[strata] done: 50 tokens in 0 s (61284.4 tok/s) (stop, cancel=False) +..[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 50 tokens in 0 s (36766.3 tok/s) (stop, cancel=False) +[strata] done: 50 tokens in 0 s (39741.4 tok/s) (stop, cancel=False) +.[strata] done: 7 tokens in 0 s (107941.6 tok/s) (stop, cancel=False) +..[strata] "parallel" must be a whole number of requests (2..8), not '4': ignored +[strata] "parallel" must be a whole number of requests (2..8), not True: ignored +[strata] "parallel": 12 - the engine runs at most 8 requests together; it will use 8 +.[strata] answering: 5 of max 8 tokens, 5.1 tok/s, 1 s +[strata] answering: 4 of max 8 tokens, 6.2 tok/s, 1 s +[strata] done: 8 tokens in 1 s (6.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 1 s so far +[strata] done: 8 tokens in 2 s (6.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 2 s so far +[strata] answering: 8 of max 8 tokens, 7.4 tok/s, 3 s +[strata] done: 8 tokens in 3 s (7.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 3 s so far +[strata] answering: 8 of max 8 tokens, 7.8 tok/s, 3 s +[strata] done: 8 tokens in 3 s (7.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 3 s so far +[strata] done: 8 tokens in 3 s (8.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 4 s so far +[strata] done: 8 tokens in 4 s (8.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 4 s so far +[strata] answering: 7 of max 8 tokens, 6.2 tok/s, 5 s +[strata] done: 8 tokens in 5 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 5 s so far +[strata] answering: 8 of max 8 tokens, 6.4 tok/s, 5 s +[strata] done: 8 tokens in 5 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 5 s so far +[strata] answering: 8 of max 8 tokens, 7.3 tok/s, 6 s +[strata] done: 8 tokens in 6 s (7.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 6 s so far +[strata] answering: 8 of max 8 tokens, 7.6 tok/s, 6 s +[strata] done: 8 tokens in 6 s (7.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 6 s so far +[strata] answering: 8 of max 8 tokens, 7.5 tok/s, 7 s +[strata] done: 8 tokens in 7 s (7.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 7 s so far +[strata] answering: 8 of max 8 tokens, 6.6 tok/s, 7 s +[strata] done: 8 tokens in 7 s (6.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 7 s so far +[strata] answering: 6 of max 8 tokens, 5.8 tok/s, 8 s +[strata] done: 8 tokens in 8 s (6.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 8 s so far +[strata] answering: 8 of max 8 tokens, 7.2 tok/s, 9 s +[strata] done: 8 tokens in 9 s (7.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 9 s so far +[strata] answering: 8 of max 8 tokens, 7.2 tok/s, 10 s +[strata] done: 8 tokens in 10 s (7.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 10 s so far +[strata] answering: 8 of max 8 tokens, 6.2 tok/s, 10 s +[strata] done: 8 tokens in 10 s (6.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 10 s so far +[strata] answering: 7 of max 8 tokens, 6.7 tok/s, 11 s +[strata] done: 8 tokens in 11 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 11 s so far +[strata] answering: 8 of max 8 tokens, 6.2 tok/s, 11 s +[strata] done: 8 tokens in 11 s (6.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 11 s so far +[strata] answering: 6 of max 8 tokens, 5.9 tok/s, 12 s +[strata] done: 8 tokens in 12 s (6.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 12 s so far +[strata] answering: 8 of max 8 tokens, 6.6 tok/s, 13 s +[strata] done: 8 tokens in 13 s (6.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 13 s so far +[strata] done: 8 tokens in 13 s (9.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 13 s so far +[strata] done: 8 tokens in 13 s (9.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 13 s so far +[strata] answering: 8 of max 8 tokens, 7.0 tok/s, 14 s +[strata] done: 8 tokens in 14 s (7.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 14 s so far +[strata] answering: 8 of max 8 tokens, 7.0 tok/s, 15 s +[strata] done: 8 tokens in 15 s (7.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 15 s so far +[strata] answering: 7 of max 8 tokens, 6.3 tok/s, 15 s +[strata] done: 8 tokens in 16 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 16 s so far +[strata] answering: 8 of max 8 tokens, 6.1 tok/s, 16 s +[strata] done: 8 tokens in 16 s (6.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 16 s so far +[strata] answering: 6 of max 8 tokens, 5.7 tok/s, 17 s +[strata] done: 8 tokens in 17 s (5.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 17 s so far +[strata] answering: 8 of max 8 tokens, 6.4 tok/s, 17 s +[strata] done: 8 tokens in 17 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 17 s so far +[strata] answering: 8 of max 8 tokens, 7.9 tok/s, 18 s +[strata] done: 8 tokens in 18 s (7.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 18 s so far +[strata] done: 8 tokens in 18 s (8.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 18 s so far +[strata] answering: 8 of max 8 tokens, 7.0 tok/s, 19 s +[strata] done: 8 tokens in 19 s (7.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 19 s so far +[strata] answering: 8 of max 8 tokens, 6.6 tok/s, 20 s +[strata] done: 8 tokens in 20 s (6.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 20 s so far +[strata] answering: 7 of max 8 tokens, 6.6 tok/s, 20 s +[strata] done: 8 tokens in 21 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 21 s so far +[strata] answering: 8 of max 8 tokens, 6.9 tok/s, 21 s +[strata] done: 8 tokens in 21 s (6.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 21 s so far +[strata] answering: 7 of max 8 tokens, 6.2 tok/s, 22 s +[strata] done: 8 tokens in 22 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 22 s so far +[strata] answering: 8 of max 8 tokens, 6.4 tok/s, 22 s +[strata] done: 8 tokens in 22 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 22 s so far +[strata] answering: 6 of max 8 tokens, 6.0 tok/s, 23 s +[strata] done: 8 tokens in 23 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 23 s so far +[strata] answering: 8 of max 8 tokens, 6.1 tok/s, 24 s +[strata] done: 8 tokens in 24 s (6.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 24 s so far +[strata] answering: 6 of max 8 tokens, 5.4 tok/s, 24 s +[strata] done: 8 tokens in 25 s (5.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 25 s so far +[strata] answering: 8 of max 8 tokens, 6.4 tok/s, 25 s +[strata] done: 8 tokens in 25 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 25 s so far +[strata] answering: 6 of max 8 tokens, 6.0 tok/s, 26 s +[strata] done: 8 tokens in 26 s (6.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 26 s so far +[strata] answering: 8 of max 8 tokens, 6.8 tok/s, 26 s +[strata] done: 8 tokens in 26 s (6.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 26 s so far +[strata] answering: 6 of max 8 tokens, 5.6 tok/s, 27 s +[strata] done: 8 tokens in 27 s (6.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 27 s so far +[strata] answering: 8 of max 8 tokens, 5.9 tok/s, 28 s +[strata] done: 8 tokens in 28 s (5.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 28 s so far +[strata] answering: 6 of max 8 tokens, 5.5 tok/s, 28 s +[strata] answering: 7 of max 8 tokens, 6.8 tok/s, 29 s +[strata] done: 8 tokens in 29 s (5.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 29 s so far +[strata] done: 8 tokens in 29 s (5.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 29 s so far +[strata] answering: 6 of max 8 tokens, 5.8 tok/s, 30 s +[strata] done: 8 tokens in 30 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 30 s so far +[strata] answering: 8 of max 8 tokens, 7.3 tok/s, 30 s +[strata] done: 8 tokens in 30 s (7.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 30 s so far +[strata] answering: 6 of max 8 tokens, 5.8 tok/s, 31 s +[strata] done: 8 tokens in 31 s (5.9 tok/s) (length, cancel=False) +[strata] answering: 7 of max 8 tokens, 6.7 tok/s, 31 s +[strata] reading the prompt: 32 of 46 tokens, 31 s so far +[strata] done: 8 tokens in 32 s (6.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 32 s so far +[strata] answering: 7 of max 8 tokens, 6.7 tok/s, 32 s +[strata] done: 8 tokens in 33 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 33 s so far +[strata] answering: 8 of max 8 tokens, 6.0 tok/s, 33 s +[strata] done: 8 tokens in 33 s (6.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 33 s so far +[strata] answering: 8 of max 8 tokens, 7.4 tok/s, 34 s +[strata] done: 8 tokens in 34 s (7.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 34 s so far +[strata] answering: 8 of max 8 tokens, 7.9 tok/s, 34 s +[strata] done: 8 tokens in 34 s (7.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 34 s so far +[strata] answering: 7 of max 8 tokens, 6.8 tok/s, 35 s +[strata] done: 8 tokens in 35 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 35 s so far +[strata] answering: 8 of max 8 tokens, 6.3 tok/s, 35 s +[strata] done: 8 tokens in 35 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 35 s so far +[strata] answering: 6 of max 8 tokens, 5.7 tok/s, 36 s +[strata] done: 8 tokens in 36 s (5.9 tok/s) (length, cancel=False) +[strata] answering: 7 of max 8 tokens, 6.7 tok/s, 36 s +[strata] reading the prompt: 32 of 46 tokens, 37 s so far +[strata] done: 8 tokens in 37 s (5.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 37 s so far +[strata] answering: 8 of max 8 tokens, 7.2 tok/s, 38 s +[strata] done: 8 tokens in 38 s (7.2 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 38 s (8.5 tok/s) (length, cancel=False) +.[strata] thinking: 3 of max 200 tokens, 4.4 tok/s, 1 s +[strata] reading the prompt: 160 of 249 tokens, 1 s so far +[strata] thinking: 12 of max 200 tokens, 6.6 tok/s, 2 s +[strata] done: 10 tokens in 2 s (10.4 tok/s) (stop, cancel=False) +[strata] thinking: 19 of max 200 tokens, 6.5 tok/s, 3 s +[strata] thinking: 26 of max 200 tokens, 6.4 tok/s, 5 s +[strata] thinking: 33 of max 200 tokens, 6.5 tok/s, 6 s +[strata] thinking: 41 of max 200 tokens, 6.6 tok/s, 7 s +[strata] thinking: 50 of max 200 tokens, 6.8 tok/s, 8 s +[strata] done: 55 tokens in 9 s (6.7 tok/s) (stop, cancel=False) +.[strata] the engine stopped unexpectedly (exit code 1). The engine exited (code 1) after it reported: verify batch: layer 34 never rang (graph finished) - please report it at github.com/Niko1221/Strata/issues with the log. The next request starts the engine again. +[strata] done: 3 tokens in 1 s (3.1 tok/s) (error, cancel=False) +[strata] the engine stopped unexpectedly (exit code 1). The engine exited (code 1) after it reported: verify batch: layer 34 never rang (graph finished) - please report it at github.com/Niko1221/Strata/issues with the log. The next request starts the engine again. +[strata] done: 1 tokens in 1 s (3.0 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +.[strata] parallel: slot(s) 0, 1 were marked busy with no request using them (issue #1603); freed, the waiting request goes on +[strata] reading the prompt: 32 of 52 tokens, 1 s so far +[strata] done: 10 tokens in 2 s (8.8 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 224 of 3,037 tokens, 1 s so far +[strata] answering: 4 of max 64 tokens, 4.3 tok/s, 1 s +[strata] reading the prompt: 768 of 3,037 tokens, 2 s so far +[strata] answering: 9 of max 64 tokens, 4.2 tok/s, 2 s +[strata] done: 10 tokens in 2 s (4.3 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,312 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 1,792 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 2,272 of 3,037 tokens, 5 s so far +[strata] reading the prompt: 2,784 of 3,037 tokens, 6 s so far +[strata] answering: 6 of max 64 tokens, 8.4 tok/s, 7 s +[strata] done: 10 tokens in 8 s (7.7 tok/s) (stop, cancel=False) +.[strata] answering: 9 of max 64 tokens, 8.4 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.1 tok/s) (stop, cancel=False) +[strata] answering: 9 of max 64 tokens, 8.5 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.5 tok/s) (stop, cancel=False) +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] answering: 9 of max 16 tokens, 9.2 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.3 tok/s) (stop, cancel=False) +.[strata] answering: 6 of max 200 tokens, 6.0 tok/s, 1 s +[strata] answering: 5 of max 200 tokens, 6.4 tok/s, 1 s +[strata] answering: 13 of max 200 tokens, 6.2 tok/s, 2 s +[strata] done: 10 tokens in 2 s (6.7 tok/s) (stop, cancel=False) +[strata] answering: 24 of max 200 tokens, 7.6 tok/s, 3 s +[strata] answering: 34 of max 200 tokens, 8.1 tok/s, 4 s +[strata] answering: 41 of max 200 tokens, 7.7 tok/s, 5 s +[strata] answering: 50 of max 200 tokens, 7.8 tok/s, 6 s +[strata] done: 55 tokens in 7 s (7.8 tok/s) (stop, cancel=False) +.[strata] answering: 4 of max 200 tokens, 6.0 tok/s, 1 s +[strata] answering: 5 of max 200 tokens, 4.6 tok/s, 1 s +[strata] done: 10 tokens in 2 s (5.8 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 2 s (7.0 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 64 of 3,037 tokens, 1 s so far +[strata] answering: 5 of max 60 tokens, 3.9 tok/s, 1 s +[strata] answering: 3 of max 64 tokens, 4.5 tok/s, 1 s +[strata] reading the prompt: 544 of 3,037 tokens, 2 s so far +[strata] answering: 9 of max 60 tokens, 3.9 tok/s, 2 s +[strata] done: 10 tokens in 3 s (4.0 tok/s) (stop, cancel=False) +[strata] answering: 8 of max 64 tokens, 4.5 tok/s, 2 s +[strata] done: 10 tokens in 2 s (4.9 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,280 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 1,792 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 2,432 of 3,037 tokens, 5 s so far +[strata] reading the prompt: 3,008 of 3,037 tokens, 6 s so far +[strata] answering: 7 of max 64 tokens, 6.8 tok/s, 7 s +[strata] done: 10 tokens in 8 s (7.6 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 2,560 of 3,037 tokens, 6 s so far +[strata] answering: 1 of max 64 tokens, 167772.2 tok/s, 5 s +[strata] answering: 4 of max 60 tokens, 0.4 tok/s, 11 s +[strata] answering: 2 of max 64 tokens, 0.4 tok/s, 11 s +[strata] reading the prompt: 2,560 of 3,037 tokens, 11 s so far +[strata] answering: 6 of max 60 tokens, 0.5 tok/s, 12 s +[strata] answering: 2 of max 64 tokens, 10.9 tok/s, 12 s +[strata] answering: 3 of max 64 tokens, 0.5 tok/s, 12 s +[strata] answering: 13 of max 60 tokens, 1.0 tok/s, 14 s +[strata] answering: 9 of max 64 tokens, 6.8 tok/s, 13 s +[strata] done: 10 tokens in 13 s (1.3 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 14 s (6.7 tok/s) (stop, cancel=False) +[strata] answering: 20 of max 60 tokens, 1.4 tok/s, 15 s +[strata] answering: 28 of max 60 tokens, 1.8 tok/s, 16 s +[strata] answering: 36 of max 60 tokens, 2.2 tok/s, 17 s +[strata] answering: 43 of max 60 tokens, 2.4 tok/s, 18 s +[strata] answering: 50 of max 60 tokens, 2.6 tok/s, 19 s +[strata] done: 55 tokens in 19 s (2.8 tok/s) (stop, cancel=False) +.[strata] parallel: slot(s) 0 were marked busy with no request using them (issue #1603); freed, the waiting request goes on +..[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (1/6) (issue #1994) +[strata] answering: 2 of max 64 tokens, 8.3 tok/s, 1 s +[strata] answering: 8 of max 64 tokens, 6.7 tok/s, 1 s +[strata] done: 10 tokens in 2 s (8.4 tok/s) (stop, cancel=False) +[strata] answering: 17 of max 64 tokens, 7.3 tok/s, 2 s +[strata] answering: 25 of max 64 tokens, 7.3 tok/s, 3 s +[strata] answering: 34 of max 64 tokens, 7.7 tok/s, 4 s +[strata] answering: 45 of max 64 tokens, 8.1 tok/s, 6 s +[strata] answering: 53 of max 64 tokens, 8.0 tok/s, 7 s +[strata] done: 55 tokens in 7 s (8.1 tok/s) (stop, cancel=False) +[strata] answering: 9 of max 64 tokens, 9.2 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.1 tok/s) (stop, cancel=False) +.[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (1/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (2/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (3/6) (issue #1994) +[strata] answering: 8 of max 64 tokens, 7.9 tok/s, 1 s +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (4/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (5/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (6/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (7/6) (issue #1994) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (length, cancel=False) +[strata] answering: 14 of max 64 tokens, 6.7 tok/s, 2 s +[strata] answering: 22 of max 64 tokens, 6.9 tok/s, 3 s +[strata] answering: 29 of max 64 tokens, 6.9 tok/s, 4 s +[strata] answering: 36 of max 64 tokens, 6.9 tok/s, 5 s +[strata] answering: 43 of max 64 tokens, 6.7 tok/s, 6 s +[strata] answering: 50 of max 64 tokens, 6.6 tok/s, 8 s +[strata] done: 55 tokens in 8 s (6.7 tok/s) (stop, cancel=False) +[strata] answering: 9 of max 64 tokens, 8.4 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.2 tok/s) (stop, cancel=False) +.[strata] answering: 8 of max 64 tokens, 8.1 tok/s, 1 s +[strata] answering: 5 of max 64 tokens, 9.6 tok/s, 1 s +[strata] done: 10 tokens in 2 s (9.8 tok/s) (stop, cancel=False) +[strata] reading the prompt: 32 of 66 tokens, 1 s so far +[strata] answering: 16 of max 64 tokens, 7.5 tok/s, 2 s +[strata] answering: 8 of max 64 tokens, 7.9 tok/s, 2 s +[strata] done: 10 tokens in 3 s (7.8 tok/s) (stop, cancel=False) +[strata] answering: 24 of max 64 tokens, 7.5 tok/s, 3 s +[strata] answering: 33 of max 64 tokens, 7.6 tok/s, 4 s +[strata] answering: 41 of max 64 tokens, 7.6 tok/s, 5 s +[strata] answering: 48 of max 64 tokens, 7.4 tok/s, 6 s +[strata] done: 55 tokens in 7 s (7.5 tok/s) (stop, cancel=False) +.[strata] parallel requests: 4 asked, the engine runs 2 (its log says why) +.[strata] parallel requests: 4 asked, the engine runs one at a time (its log says why) +[strata] done: 10 tokens in 1 s (10.6 tok/s) (stop, cancel=False) +.[strata] answering: 5 of max 200 tokens, 4.9 tok/s, 1 s +[strata] answering: 5 of max 200 tokens, 8.5 tok/s, 1 s +[strata] done: 10 tokens in 2 s (7.9 tok/s) (stop, cancel=False) +[strata] answering: 14 of max 200 tokens, 6.5 tok/s, 2 s +[strata] answering: 21 of max 200 tokens, 6.6 tok/s, 3 s +[strata] answering: 31 of max 200 tokens, 7.4 tok/s, 4 s +[strata] answering: 40 of max 200 tokens, 7.7 tok/s, 5 s +[strata] answering: 48 of max 200 tokens, 7.6 tok/s, 6 s +[strata] done: 55 tokens in 8 s (7.3 tok/s) (stop, cancel=False) +.[strata] answering: 5 of max 200 tokens, 5.3 tok/s, 1 s +[strata] answering: 5 of max 64 tokens, 7.7 tok/s, 1 s +[strata] done: 10 tokens in 2 s (8.7 tok/s) (stop, cancel=False) +[strata] answering: 14 of max 200 tokens, 7.1 tok/s, 2 s +[strata] answering: 22 of max 200 tokens, 7.2 tok/s, 3 s +[strata] answering: 31 of max 200 tokens, 7.3 tok/s, 4 s +[strata] answering: 38 of max 200 tokens, 7.2 tok/s, 5 s +[strata] answering: 46 of max 200 tokens, 7.3 tok/s, 6 s +[strata] done: 55 tokens in 7 s (7.5 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 576 of 3,037 tokens, 1 s so far +[strata] reading the prompt: 1,344 of 3,037 tokens, 2 s so far +[strata] reading the prompt: 2,016 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 2,624 of 3,037 tokens, 4 s so far +[strata] answering: 1 of max 64 tokens, 3.5 tok/s, 5 s +[strata] answering: 1 of max 64 tokens, 0.8 tok/s, 6 s +[strata] answering: 1 of max 64 tokens, 0.4 tok/s, 7 s +[strata] answering: 1 of max 64 tokens, 0.3 tok/s, 8 s +[strata] answering: 1 of max 64 tokens, 0.2 tok/s, 9 s +[strata] reading the prompt: 32 of 3,037 tokens, 10 s so far +[strata] answering: 4 of max 64 tokens, 0.7 tok/s, 10 s +[strata] reading the prompt: 672 of 3,037 tokens, 11 s so far +[strata] answering: 9 of max 64 tokens, 1.4 tok/s, 11 s +[strata] done: 10 tokens in 12 s (1.5 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,280 of 3,037 tokens, 12 s so far +[strata] reading the prompt: 1,920 of 3,037 tokens, 13 s so far +[strata] reading the prompt: 2,528 of 3,037 tokens, 14 s so far +[strata] answering: 4 of max 64 tokens, 11.4 tok/s, 15 s +[strata] done: 10 tokens in 16 s (8.7 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 2,560 of 3,037 tokens, 4 s so far +[strata] answering: 1 of max 64 tokens, 0.2 tok/s, 9 s +[strata] answering: 3 of max 64 tokens, 0.3 tok/s, 14 s +[strata] reading the prompt: 2,560 of 3,037 tokens, 13 s so far +[strata] answering: 5 of max 64 tokens, 0.5 tok/s, 15 s +[strata] answering: 3 of max 64 tokens, 9.6 tok/s, 14 s +[strata] done: 10 tokens in 15 s (0.9 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 15 s (8.7 tok/s) (stop, cancel=False) +.[strata] the engine's logits became non-finite before the first token (#879); it dropped its cached state: sending the request again once +[strata] answering: 9 of max 64 tokens, 9.4 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.0 tok/s) (stop, cancel=False) +[strata] the engine's logits became non-finite before the first token (#879); it dropped its cached state: sending the request again once +[strata] answering: 7 of max 64 tokens, 7.0 tok/s, 1 s +[strata] answering: 6 of max 64 tokens, 9.5 tok/s, 1 s +[strata] done: 10 tokens in 1 s (9.5 tok/s) (stop, cancel=False) +[strata] answering: 16 of max 64 tokens, 7.8 tok/s, 2 s +[strata] answering: 26 of max 64 tokens, 8.1 tok/s, 3 s +[strata] answering: 35 of max 64 tokens, 8.2 tok/s, 4 s +[strata] answering: 44 of max 64 tokens, 8.2 tok/s, 5 s +[strata] answering: 54 of max 64 tokens, 8.4 tok/s, 7 s +[strata] done: 55 tokens in 7 s (8.3 tok/s) (stop, cancel=False) +[strata] the engine's logits became non-finite after 2 tokens (#879): the reply ends here as "length"; the engine dropped its cached state, so the next request reads its prompt from the start +[strata] done: 2 tokens in 1 s (6.3 tok/s) (length, cancel=False) +[strata] answering: 8 of max 64 tokens, 7.6 tok/s, 1 s +[strata] answering: 16 of max 64 tokens, 7.4 tok/s, 2 s +[strata] answering: 24 of max 64 tokens, 7.6 tok/s, 3 s +[strata] answering: 36 of max 64 tokens, 8.5 tok/s, 4 s +[strata] answering: 46 of max 64 tokens, 8.7 tok/s, 5 s +[strata] done: 55 tokens in 6 s (8.6 tok/s) (stop, cancel=False) +.[strata] answering: 9 of max 64 tokens, 8.8 tok/s, 1 s +[strata] done: 9 tokens in 1 s (8.5 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 1 s (10.6 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 2,560 of 3,037 tokens, 5 s so far +[strata] answering: 1 of max 64 tokens, 161319.4 tok/s, 5 s +[strata] answering: 2 of max 64 tokens, 0.4 tok/s, 10 s +[strata] reading the prompt: 2,560 of 3,037 tokens, 10 s so far +[strata] answering: 2 of max 64 tokens, 6.8 tok/s, 11 s +[strata] answering: 3 of max 64 tokens, 0.5 tok/s, 11 s +[strata] done: 10 tokens in 12 s (1.4 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 13 s (7.2 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 32 of 3,037 tokens, 1 s so far +[strata] reading the prompt: 64 of 3,037 tokens, 1 s so far +[strata] reading the prompt: 416 of 3,037 tokens, 2 s so far +[strata] reading the prompt: 928 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 1,504 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 2,112 of 3,037 tokens, 5 s so far +[strata] reading the prompt: 2,624 of 3,037 tokens, 6 s so far +[strata] reading the prompt: 32 of 3,037 tokens, 7 s so far +[strata] answering: 2 of max 64 tokens, 5.1 tok/s, 8 s +[strata] reading the prompt: 672 of 3,037 tokens, 8 s so far +[strata] answering: 7 of max 64 tokens, 5.0 tok/s, 9 s +[strata] done: 10 tokens in 9 s (4.4 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,184 of 3,037 tokens, 9 s so far +[strata] reading the prompt: 1,696 of 3,037 tokens, 10 s so far +[strata] reading the prompt: 2,176 of 3,037 tokens, 11 s so far +[strata] reading the prompt: 2,688 of 3,037 tokens, 12 s so far +[strata] reading the prompt: 32 of 43 tokens, 13 s so far +[strata] answering: 2 of max 64 tokens, 3.8 tok/s, 13 s +[strata] answering: 7 of max 64 tokens, 6.4 tok/s, 14 s +[strata] answering: 9 of max 64 tokens, 5.5 tok/s, 15 s +[strata] done: 10 tokens in 14 s (6.2 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 15 s (5.6 tok/s) (stop, cancel=False) +.........[strata] reading the prompt: 9 of max 64 tokens, 9.3 tok/s, 1 s +[strata] reading the prompt: 9 of max 64 tokens, 9.0 tok/s, 1 s +[strata] reading the prompt: 8 of max 64 tokens, 8.1 tok/s, 1 s +[strata] answering: 16 of max 64 tokens, 8.0 tok/s, 2 s +[strata] answering: 16 of max 64 tokens, 7.9 tok/s, 2 s +[strata] answering: 16 of max 64 tokens, 7.6 tok/s, 2 s +[strata] answering: 26 of max 64 tokens, 8.3 tok/s, 3 s +[strata] answering: 24 of max 64 tokens, 7.7 tok/s, 3 s +[strata] answering: 25 of max 64 tokens, 7.6 tok/s, 3 s +[strata] answering: 35 of max 64 tokens, 8.4 tok/s, 4 s +[strata] answering: 32 of max 64 tokens, 7.6 tok/s, 4 s +[strata] answering: 35 of max 64 tokens, 8.1 tok/s, 4 s +[strata] answering: 40 of max 64 tokens, 7.7 tok/s, 5 s +[strata] answering: 44 of max 64 tokens, 8.4 tok/s, 5 s +[strata] answering: 42 of max 64 tokens, 7.8 tok/s, 6 s +[strata] done: 51 tokens in 6 s (8.2 tok/s) (stop, cancel=False) +[strata] answering: 47 of max 64 tokens, 7.4 tok/s, 6 s +[strata] answering: 50 of max 64 tokens, 7.8 tok/s, 7 s +[strata] done: 51 tokens in 7 s (7.7 tok/s) (stop, cancel=False) +[strata] done: 51 tokens in 7 s (7.5 tok/s) (stop, cancel=False) +.[strata] done: 22 tokens in 0 s (149796.6 tok/s) (stop, cancel=False) +.....[strata] done: 3 tokens in 0 s (11305.4 tok/s) (stop, cancel=False) +...[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (7884.0 tok/s) (length, cancel=False) +.[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] answering: 1 of max 1 tokens, 167772.2 tok/s, 1 s +[strata] done: 1 tokens in 1 s (4675.9 tok/s) (length, cancel=False) +.[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (10838.0 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] answering: 1 of max 1 tokens, 322638.8 tok/s, 1 s +[strata] done: 1 tokens in 1 s (12633.4 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (40329.8 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (5468.5 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (6978.9 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (9258.9 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (6754.1 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (6492.7 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (6141.0 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (30393.5 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (29330.8 tok/s) (length, cancel=False) +...[strata] loading the model again (it was unloaded) ... +[strata] the engine did not start (try 1 of 3); again in 0 s +[strata] the engine did not start (try 2 of 3); again in 0 s +.[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 1 tokens in 0 s (13706.9 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine did not start (try 1 of 3); again in 0 s +[strata] the engine did not start (try 2 of 3); again in 0 s +..[strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2086 tokens in 0 s (215830.0 tok/s) (stop, cancel=False) +.[strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2071 tokens in 0 s (361556.9 tok/s) (stop, cancel=False) +.[strata] done: 1 tokens in 0 s (27776.8 tok/s) (cancel, cancel=True) +.[strata] done: 16384 tokens in 0 s (596777.1 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] done: 28001 tokens in 0 s (692256.1 tok/s) (stop, cancel=False) +.[strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2082 tokens in 0 s (639840.3 tok/s) (stop, cancel=False) +.[strata] done: 28001 tokens in 0 s (675210.2 tok/s) (stop, cancel=False) +.[strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2106 tokens in 0 s (640876.7 tok/s) (stop, cancel=False) +.[strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2086 tokens in 0 s (630581.5 tok/s) (stop, cancel=False) +..[strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 16384 tokens in 0 s (674646.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] done: 28001 tokens in 0 s (566856.4 tok/s) (stop, cancel=False) +..[strata] the thinking repeated the same passages (coverage=1.000) at 2560 tokens: ended as "length" (reasoning_loop_recovery: "stop" in strata-.json; remove it to turn this off) +[strata] done: 2560 tokens in 0 s (642190.1 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] the thinking repeated the same passages (coverage=1.000) at 2048 tokens: ended as "length" (reasoning_loop_recovery: "stop" in strata-.json; remove it to turn this off) +[strata] done: 2048 tokens in 0 s (645423.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +..................[strata] done: 116 tokens in 0 s (214428.9 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (161265.9 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (216432.1 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (163817.9 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (218081.2 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (171558.3 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) + +[deadline exceeded after 300s; process group 91073 terminated after SIGTERM] diff --git a/results/round39/serve-final-full.log b/results/round39/serve-final-full.log new file mode 100644 index 0000000..9a78f5b --- /dev/null +++ b/results/round39/serve-final-full.log @@ -0,0 +1,2550 @@ +test_a_continuation_is_not_held_back (test_admit_order.AdmitOrderTest.test_a_continuation_is_not_held_back) +A request already admitted once (new=False: promoted from the solo path) is not part of the order. ... ok +test_a_long_request_that_waited_too_long_goes_first (test_admit_order.AdmitOrderTest.test_a_long_request_that_waited_too_long_goes_first) ... ok +test_off_by_default (test_admit_order.AdmitOrderTest.test_off_by_default) ... ok +test_short_first_puts_the_chat_before_the_agent_reads (test_admit_order.AdmitOrderTest.test_short_first_puts_the_chat_before_the_agent_reads) ... ok +test_default_off (test_cache_drop.CacheDrop.test_default_off) ... ok +test_files_and_advise (test_cache_drop.CacheDrop.test_files_and_advise) ... ok +test_not_while_engine_runs_or_not_opted_in (test_cache_drop.CacheDrop.test_not_while_engine_runs_or_not_opted_in) ... ok +test_on (test_cache_drop.CacheDrop.test_on) ... ok +test_a_bad_chunk_size_is_a_400 (test_chunked_body.ChunkedBody.test_a_bad_chunk_size_is_a_400) ... ok +test_a_chunked_chat_request_is_read (test_chunked_body.ChunkedBody.test_a_chunked_chat_request_is_read) ... [strata] done: 12 tokens in 0 s (202950.2 tok/s) (stop, cancel=False) +ok +test_a_chunked_control_body_is_consumed (test_chunked_body.ChunkedBody.test_a_chunked_control_body_is_consumed) ... ok +test_a_chunked_settings_body_is_read_too (test_chunked_body.ChunkedBody.test_a_chunked_settings_body_is_read_too) ... [strata] other apps now use the Chat settings: temperature=0.5 +ok +test_a_content_length_request_is_unchanged (test_chunked_body.ChunkedBody.test_a_content_length_request_is_unchanged) ... [strata] done: 12 tokens in 0 s (316551.2 tok/s) (stop, cancel=False) +ok +test_a_truncated_chunked_body_is_a_400 (test_chunked_body.ChunkedBody.test_a_truncated_chunked_body_is_a_400) ... ok +test_an_overlong_chunked_trailer_does_not_change_settings (test_chunked_body.ChunkedBody.test_an_overlong_chunked_trailer_does_not_change_settings) ... ok +test_an_oversized_chunked_body_is_a_413_and_is_never_allocated (test_chunked_body.ChunkedBody.test_an_oversized_chunked_body_is_a_413_and_is_never_allocated) ... ok +test_chunked_trailer_limit_does_not_change_settings (test_chunked_body.ChunkedBody.test_chunked_trailer_limit_does_not_change_settings) ... ok +test_chunked_wins_over_content_length (test_chunked_body.ChunkedBody.test_chunked_wins_over_content_length) ... [strata] done: 12 tokens in 0 s (356962.0 tok/s) (stop, cancel=False) +ok +test_complete_settings_bodies_still_apply (test_chunked_body.ChunkedBody.test_complete_settings_bodies_still_apply) ... [strata] other apps now use the Chat settings: temperature=0.5 +[strata] other apps now use the Chat settings: temperature=0.5 +[strata] other apps now use the Chat settings: temperature=0.5 +ok +test_incomplete_chunked_trailers_do_not_change_settings (test_chunked_body.ChunkedBody.test_incomplete_chunked_trailers_do_not_change_settings) ... ok +test_both_overflow_errors_are_recognizable_before_streaming (test_context_overflow.ContextOverflow.test_both_overflow_errors_are_recognizable_before_streaming) ... [strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +ok +test_one_token_of_room_still_answers_but_zero_is_rejected (test_context_overflow.ContextOverflow.test_one_token_of_room_still_answers_but_zero_is_rejected) ... [strata] done: 1 tokens in 0 s (59074.7 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) + max tokens (1) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] done: 1 tokens in 0 s (91180.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] done: 1 tokens in 0 s (91180.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) + max tokens (1) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] done: 1 tokens in 0 s (82241.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] done: 1 tokens in 0 s (47662.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 1 tokens in 0 s (53092.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_prepare_preserves_prompt_and_budget_limits (test_context_overflow.ContextOverflow.test_prepare_preserves_prompt_and_budget_limits) ... ok +test_no_stop_without_cancel (test_control_cancel.ControlCancelTest.test_no_stop_without_cancel) ... ok +test_stop_during_prompt_read (test_control_cancel.ControlCancelTest.test_stop_during_prompt_read) ... ok +test_constant_cost (test_detok.Detok.test_constant_cost) ... ok +test_random_ids_and_broken_bytes (test_detok.Detok.test_random_ids_and_broken_bytes) ... ok +test_text_with_split_characters (test_detok.Detok.test_text_with_split_characters) ... ok +test_pack_vocab (test_detok.HeapBpe.test_pack_vocab) ... ok +test_speed_pack (test_detok.HeapBpe.test_speed_pack) ... ok +test_speed_synthetic (test_detok.HeapBpe.test_speed_synthetic) ... ok +test_synthetic_vocab (test_detok.HeapBpe.test_synthetic_vocab) ... ok +test_words_directly (test_detok.HeapBpe.test_words_directly) ... ok +test_threads_match_serial (test_detok.ThreadedEncode.test_threads_match_serial) ... ok +test_an_unknown_level_is_a_400_not_the_default (test_effort_paths.EffortPaths.test_an_unknown_level_is_a_400_not_the_default) ... [strata] 400 invalid request: unknown reasoning effort 'huge': use none, low, medium or high +ok +test_anthropic_output_config_and_thinking (test_effort_paths.EffortPaths.test_anthropic_output_config_and_thinking) ... [strata] done: 24 tokens in 0 s (255490.6 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (343560.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (358232.4 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 500) +[strata] done: 24 tokens in 0 s (343560.7 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 2047) +[strata] done: 24 tokens in 0 s (363405.4 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 2048) +[strata] done: 24 tokens in 0 s (337796.3 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 8191) +[strata] done: 24 tokens in 0 s (368730.0 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 8192) +[strata] done: 24 tokens in 0 s (368730.0 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 31999) +[strata] done: 24 tokens in 0 s (381300.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (370085.6 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (401049.0 tok/s) (stop, cancel=False) +ok +test_config_level_reaches_every_api (test_effort_paths.EffortPaths.test_config_level_reaches_every_api) ... [strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (379861.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (401049.0 tok/s) (stop, cancel=False) +[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (272063.0 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (394758.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (374213.0 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (311651.1 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (356962.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (397878.6 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (292625.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (381300.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (385683.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (328965.0 tok/s) (stop, cancel=False) +ok +test_openai_reasoning_effort (test_effort_paths.EffortPaths.test_openai_reasoning_effort) ... [strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (385683.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (375609.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (388661.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (379861.5 tok/s) (stop, cancel=False) +[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (414252.2 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (407543.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (407543.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (342392.2 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (393216.0 tok/s) (stop, cancel=False) +[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (394758.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (388661.4 tok/s) (stop, cancel=False) +ok +test_openai_reasoning_object_and_template_kwargs (test_effort_paths.EffortPaths.test_openai_reasoning_object_and_template_kwargs) ... [strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (399457.5 tok/s) (stop, cancel=False) +[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (393216.0 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (407543.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (353204.5 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (407543.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (405900.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (387166.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (381300.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (393216.0 tok/s) (stop, cancel=False) +ok +test_responses_reasoning_effort (test_effort_paths.EffortPaths.test_responses_reasoning_effort) ... [strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (307838.8 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (307838.8 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (320583.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (323676.2 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (296068.5 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (306900.3 tok/s) (stop, cancel=False) +ok +test_the_chat_settings_shared_with_apps_beat_the_config (test_effort_paths.EffortPaths.test_the_chat_settings_shared_with_apps_beat_the_config) ... [strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (407543.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (393216.0 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (285975.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (394758.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (374213.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (338933.7 tok/s) (stop, cancel=False) +ok +test_the_level_follows_a_conversation (test_effort_paths.EffortPaths.test_the_level_follows_a_conversation) ... [strata] thinking limited to 5964 of max_tokens 7951 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (399457.5 tok/s) (stop, cancel=False) +[strata] thinking limited to 6073 of max_tokens 8097 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (405900.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (394758.0 tok/s) (stop, cancel=False) +ok +test_the_requests_own_level_beats_the_config (test_effort_paths.EffortPaths.test_the_requests_own_level_beats_the_config) ... [strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (407543.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (399457.5 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 20000) +[strata] done: 24 tokens in 0 s (391686.0 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (300487.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (399457.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (394758.0 tok/s) (stop, cancel=False) +ok +test_the_web_apps_request_shape (test_effort_paths.EffortPaths.test_the_web_apps_request_shape) ... [strata] done: 24 tokens in 0 s (141182.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (140199.6 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (143804.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (137894.9 tok/s) (stop, cancel=False) +ok +test_tools_and_a_system_message_keep_the_level (test_effort_paths.EffortPaths.test_tools_and_a_system_message_keep_the_level) ... [strata] thinking limited to 5266 of max_tokens 7021 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (370085.6 tok/s) (stop, cancel=False) +[strata] thinking limited to 5213 of max_tokens 6950 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (342392.2 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (387166.5 tok/s) (stop, cancel=False) +ok +test_anthropic_requests_too (test_env_to_user.EnvBlockToUser.test_anthropic_requests_too) ... ok +test_content_parts_get_a_leading_text_part (test_env_to_user.EnvBlockToUser.test_content_parts_get_a_leading_text_part) ... ok +test_moved_to_the_first_user_message (test_env_to_user.EnvBlockToUser.test_moved_to_the_first_user_message) ... ok +test_off_and_no_block_change_nothing (test_env_to_user.EnvBlockToUser.test_off_and_no_block_change_nothing) ... ok +test_text_parts_are_read_as_one_text (test_env_to_user.EnvBlockToUser.test_text_parts_are_read_as_one_text) ... ok +test_two_folders_share_the_system_turn (test_env_to_user.EnvBlockToUser.test_two_folders_share_the_system_turn) ... ok +test_a_broken_jpeg_is_left_as_it_was_sent (test_exif_orientation.ExifOrientation.test_a_broken_jpeg_is_left_as_it_was_sent) ... ok +test_a_jpeg_without_the_tag_or_with_tag_1_passes_through_unchanged (test_exif_orientation.ExifOrientation.test_a_jpeg_without_the_tag_or_with_tag_1_passes_through_unchanged) ... ok +test_a_turned_jpeg_comes_out_upright_as_a_png (test_exif_orientation.ExifOrientation.test_a_turned_jpeg_comes_out_upright_as_a_png) ... ok +test_png_passes_through (test_exif_orientation.ExifOrientation.test_png_passes_through) ... ok +test_config_name_wins (test_family_name.FamilyName.test_config_name_wins) ... ok +test_old_configs (test_family_name.FamilyName.test_old_configs) ... ok +test_per_family (test_family_name.FamilyName.test_per_family) ... ok +test_an_earlier_window_never_finished_is_fatal (test_fatal_recovery.FatalRecovery.test_an_earlier_window_never_finished_is_fatal) ... ok +test_batch_timeout_signature_also_invalidates_engine (test_fatal_recovery.FatalRecovery.test_batch_timeout_signature_also_invalidates_engine) ... ok +test_closed_serial_engine_rejects_without_accessing_closed_pipe (test_fatal_recovery.FatalRecovery.test_closed_serial_engine_rejects_without_accessing_closed_pipe) ... ok +test_fatal_error_is_not_a_success_and_next_request_reloads (test_fatal_recovery.FatalRecovery.test_fatal_error_is_not_a_success_and_next_request_reloads) ... [strata] the engine stopped unexpectedly (exit code -15). Unrecoverable native verification failure: verify: timed out at layer 18; its GPU waits were released and the GPU finished (#267) The next request starts the engine again. +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -15); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 1 tokens in 0 s (21620.1 tok/s) (length, cancel=False) +ok +test_fatal_lines_are_checked_before_batch_routing (test_fatal_recovery.FatalRecovery.test_fatal_lines_are_checked_before_batch_routing) ... ok +test_generate_batched_refuses_a_dead_engine (test_fatal_recovery.FatalRecovery.test_generate_batched_refuses_a_dead_engine) ... ok +test_http_health_and_next_request_recovery (test_fatal_recovery.FatalRecovery.test_http_health_and_next_request_recovery) ... [strata] the engine stopped unexpectedly (exit code -15). Unrecoverable native verification failure: verify: timed out at layer 18; its GPU waits were released and the GPU finished (#267) The next request starts the engine again. +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -15); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 1 tokens in 0 s (28532.7 tok/s) (length, cancel=False) +ok +test_validation_error_does_not_kill_engine (test_fatal_recovery.FatalRecovery.test_validation_error_does_not_kill_engine) ... ok +test_anthropic_thinking_only_replies_are_not_rendered (test_frontend.EmptyAssistantHistory.test_anthropic_thinking_only_replies_are_not_rendered) ... ok +test_empty_user_and_tool_turns_are_kept (test_frontend.EmptyAssistantHistory.test_empty_user_and_tool_turns_are_kept) ... ok +test_final_empty_assistant_turn_is_kept (test_frontend.EmptyAssistantHistory.test_final_empty_assistant_turn_is_kept) ... ok +test_image_only_assistant_turn_is_kept (test_frontend.EmptyAssistantHistory.test_image_only_assistant_turn_is_kept) ... ok +test_openai_empty_replies_are_not_rendered (test_frontend.EmptyAssistantHistory.test_openai_empty_replies_are_not_rendered) ... ok +test_text_and_tool_calls_are_kept (test_frontend.EmptyAssistantHistory.test_text_and_tool_calls_are_kept) ... ok +test_keep_empty_turns_restores_the_old_prompt (test_frontend.EmptyTurnsSwitch.test_keep_empty_turns_restores_the_old_prompt) ... ok +test_fenced_and_bare_agree (test_frontend.FencedJsonCallInWrapper.test_fenced_and_bare_agree) ... ok +test_off_without_recovery_or_for_an_undeclared_tool (test_frontend.FencedJsonCallInWrapper.test_off_without_recovery_or_for_an_undeclared_tool) ... ok +test_other_fences_and_prose_stay_text (test_frontend.FencedJsonCallInWrapper.test_other_fences_and_prose_stay_text) ... ok +test_required_with_one_tool_names_it (test_frontend.ForcedCallOpening.test_required_with_one_tool_names_it) ... ok +test_a_user_asking_for_a_closer_verbatim_gets_text (test_frontend.QuotedToolCallCloser.test_a_user_asking_for_a_closer_verbatim_gets_text) ... ok +test_assistant_replay_keeps_every_tool_tag (test_frontend.QuotedToolCallCloser.test_assistant_replay_keeps_every_tool_tag) ... ok +test_assistant_replay_keeps_its_token (test_frontend.QuotedToolCallCloser.test_assistant_replay_keeps_its_token) ... ok +test_every_tool_tag_a_user_quotes_stays_text (test_frontend.QuotedToolCallCloser.test_every_tool_tag_a_user_quotes_stays_text) ... ok +test_template_keeps_seeing_a_tool_result_turn (test_frontend.QuotedToolCallCloser.test_template_keeps_seeing_a_tool_result_turn) ... ok +test_user_text_stays_text (test_frontend.QuotedToolCallCloser.test_user_text_stays_text) ... ok +test_user_tool_result_keeps_its_outer_markers (test_frontend.QuotedToolCallCloser.test_user_tool_result_keeps_its_outer_markers) ... ok +test_a_system_prompt_with_an_image_part_is_read_as_text_as_before (test_frontend.UnreadContentParts.test_a_system_prompt_with_an_image_part_is_read_as_text_as_before) ... ok +test_chat_file_and_audio_parts_are_left_out (test_frontend.UnreadContentParts.test_chat_file_and_audio_parts_are_left_out) ... ok +test_chat_refusal_part_is_read_as_its_text (test_frontend.UnreadContentParts.test_chat_refusal_part_is_read_as_its_text) ... ok +test_messages_blocks_real_clients_send_are_left_out (test_frontend.UnreadContentParts.test_messages_blocks_real_clients_send_are_left_out) ... ok +test_messages_redacted_thinking_is_dropped (test_frontend.UnreadContentParts.test_messages_redacted_thinking_is_dropped) ... ok +test_messages_tool_reference_in_a_tool_result_is_left_out (test_frontend.UnreadContentParts.test_messages_tool_reference_in_a_tool_result_is_left_out) ... ok +test_text_and_image_requests_render_as_before (test_frontend.UnreadContentParts.test_text_and_image_requests_render_as_before) ... ok +test_configured_vision_is_not_omitted (test_image_omission.ImageOmission.test_configured_vision_is_not_omitted) ... ok +test_disabled_switch_keeps_the_normal_vision_error (test_image_omission.ImageOmission.test_disabled_switch_keeps_the_normal_vision_error) ... ok +test_omitted_images_are_cached_as_text_and_retries_are_stable (test_image_omission.ImageOmission.test_omitted_images_are_cached_as_text_and_retries_are_stable) ... ok +test_config_checked (test_images.Command.test_config_checked) ... ok +test_env (test_images.Command.test_env) ... ok +test_native_switch (test_images.Command.test_native_switch) ... ok +test_other_steps_and_batch (test_images.Command.test_other_steps_and_batch) ... ok +test_the_model_cards_command (test_images.Command.test_the_model_cards_command) +Unsloth's stable-diffusion.cpp command for Turbo: 8 steps, CFG 1, euler, the Turbo sigmas, --diffusion-fa. ... ok +test_b64 (test_images.OverHttp.test_b64) ... [strata] image: 512x768, n=2, 8 steps, seed 11 +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5726: ResourceWarning: unclosed file <_io.TextIOWrapper name=11 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 0.3 s (keep) +ok +test_bad_request_and_failure (test_images.OverHttp.test_bad_request_and_failure) ... [strata] 400 invalid request: size 999x999: each side a multiple of 32 from 256 to 2048 +[strata] image: 1024x1024, n=1, 8 steps, seed 1909908448 +[strata] image job failed: sd-cli exited with code 1 +[INFO ] loading tensors completed, taking 0.10s +[ERROR] ggml_vulkan: ErrorDeviceLost +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5730: ResourceWarning: unclosed file <_io.TextIOWrapper name=11 encoding='UTF-8'> + return +ResourceWarning: Enable tracemalloc to get the object allocation traceback +ok +test_chat_waits_for_the_image (test_images.OverHttp.test_chat_waits_for_the_image) +One job at a time on the GPU: a chat request sent during an image job is answered after it. ... [strata] image: 1024x1024, n=1, 8 steps, seed 1669014356 +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5726: ResourceWarning: unclosed file <_io.TextIOWrapper name=11 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 1.6 s (keep) +[strata] done: 8 tokens in 0 s (89003.8 tok/s) (length, cancel=False) +ok +test_listed (test_images.OverHttp.test_listed) ... ok +test_not_set_up (test_images.OverHttp.test_not_set_up) ... ok +test_url (test_images.OverHttp.test_url) ... [strata] image: 1024x1024, n=1, 8 steps, seed 0 +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5726: ResourceWarning: unclosed file <_io.TextIOWrapper name=11 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 0.3 s (keep) +[strata] image: 1024x1024, n=1, 8 steps, seed 1 +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5726: ResourceWarning: unclosed file <_io.TextIOWrapper name=9 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 0.3 s (keep) +[strata] image: 1024x1024, n=1, 8 steps, seed 2 +[strata] image done in 0.4 s (keep) +[strata] image: 1024x1024, n=1, 8 steps, seed 3 +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5726: ResourceWarning: unclosed file <_io.TextIOWrapper name=9 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 0.3 s (keep) +ok +test_eager_reload (test_images.Policy.test_eager_reload) ... [strata] image job: the conversation (123 tokens) saved; unloading the model +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] the conversation from before the image job is back (123 tokens) +ok +test_keep_touches_nothing (test_images.Policy.test_keep_touches_nothing) ... ok +test_resolve (test_images.Policy.test_resolve) ... ok +test_shrink_gives_the_cache_back (test_images.Policy.test_shrink_gives_the_cache_back) ... [strata] image job: the engine's expert cache gave VRAM back (9000 MiB free) +ok +test_unload_saves_and_the_next_start_restores (test_images.Policy.test_unload_saves_and_the_next_start_restores) ... [strata] image job: the conversation (123 tokens) saved; unloading the model +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] the conversation from before the image job is back (123 tokens) +ok +test_defaults (test_images.Request.test_defaults) ... ok +test_fields (test_images.Request.test_fields) ... ok +test_refused (test_images.Request.test_refused) ... ok +test_a_failing_call_never_breaks_a_request (test_keepawake.KeepAwakeTests.test_a_failing_call_never_breaks_a_request) ... ok +test_default_service_has_none (test_keepawake.KeepAwakeTests.test_default_service_has_none) ... ok +test_exit_clears_a_held_state (test_keepawake.KeepAwakeTests.test_exit_clears_a_held_state) ... ok +test_held_while_busy_and_released_when_idle (test_keepawake.KeepAwakeTests.test_held_while_busy_and_released_when_idle) ... ok +test_off_and_non_windows (test_keepawake.KeepAwakeTests.test_off_and_non_windows) ... ok +test_service_holds_for_the_whole_request_including_early_close (test_keepawake.KeepAwakeTests.test_service_holds_for_the_whole_request_including_early_close) ... [strata] done: 3 tokens in 0 s (68015.7 tok/s) (stop, cancel=False) +FAIL +test_an_engine_that_died_on_load_is_reported_as_such (test_lifecycle.Lifecycle.test_an_engine_that_died_on_load_is_reported_as_such) ... [strata] loading the model again (it was unloaded) ... +ok +test_busy_auth_and_foreign_origin_cannot_unload (test_lifecycle.Lifecycle.test_busy_auth_and_foreign_origin_cannot_unload) ... [strata] model unloaded; the next request loads it again +ok +test_close_cleans_up_after_a_broken_stdin_pipe (test_lifecycle.Lifecycle.test_close_cleans_up_after_a_broken_stdin_pipe) ... ok +test_close_does_not_forget_a_process_still_exiting (test_lifecycle.Lifecycle.test_close_does_not_forget_a_process_still_exiting) ... ok +test_close_sends_eof_before_waiting_for_windows_reader (test_lifecycle.Lifecycle.test_close_sends_eof_before_waiting_for_windows_reader) ... ok +test_close_terminates_before_it_kills (test_lifecycle.Lifecycle.test_close_terminates_before_it_kills) ... ok +test_discover_chat_unload_and_stream_reload (test_lifecycle.Lifecycle.test_discover_chat_unload_and_stream_reload) ... [strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 7 tokens in 0 s (177940.2 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (233016.9 tok/s) (stop, cancel=False) +[strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 7 tokens in 0 s (60287.7 tok/s) (stop, cancel=False) +ok +test_explicit_load_and_model_matching (test_lifecycle.Lifecycle.test_explicit_load_and_model_matching) ... [strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +ok +test_native_lazy_constructor_does_not_start_a_process (test_lifecycle.Lifecycle.test_native_lazy_constructor_does_not_start_a_process) ... ok +test_capability_and_batch (test_logit_bias.LogitBias.test_capability_and_batch) ... ok +test_forms_and_transport (test_logit_bias.LogitBias.test_forms_and_transport) ... ok +test_http_rejects_before_stream (test_logit_bias.LogitBias.test_http_rejects_before_stream) ... [strata] 400 invalid request: logit_bias: expected an object or a list of [token_id, bias] pairs +[strata] 400 invalid request: logit_bias: this engine does not support token biases; use a supporting CUDA/HIP build +[strata] 400 invalid request: logit_bias: biases must be finite numbers between -100 and 100; false is a pair-list ban +ok +test_invalid (test_logit_bias.LogitBias.test_invalid) ... ok +test_large_list (test_logit_bias.LogitBias.test_large_list) ... ok +test_a_bad_list_is_never_stored (test_logit_bias.LogitBiasCache.test_a_bad_list_is_never_stored) ... ok +test_an_equal_list_in_a_later_request_is_a_hit (test_logit_bias.LogitBiasCache.test_an_equal_list_in_a_later_request_is_a_hit) ... ok +test_false_is_not_zero (test_logit_bias.LogitBiasCache.test_false_is_not_zero) ... ok +test_same_answers_as_the_plain_function (test_logit_bias.LogitBiasCache.test_same_answers_as_the_plain_function) ... ok +test_same_length_mutations_are_validated (test_logit_bias.LogitBiasCache.test_same_length_mutations_are_validated) ... ok +test_the_cache_is_bounded (test_logit_bias.LogitBiasCache.test_the_cache_is_bounded) ... ok +test_the_result_is_a_copy (test_logit_bias.LogitBiasCache.test_the_result_is_a_copy) ... ok +test_vocabulary_size_is_part_of_the_key (test_logit_bias.LogitBiasCache.test_vocabulary_size_is_part_of_the_key) ... ok +[strata] MCP server 'missing' did not start: could not start 'strata-no-such-program': [Errno 2] No such file or directory: 'strata-no-such-program'. Its tools are left out; the chat works without them. +[strata] MCP server 'paged': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +test_a_server_that_fails_is_left_out (test_mcp.Client.test_a_server_that_fails_is_left_out) ... ok +test_call (test_mcp.Client.test_call) ... ok +test_cancel (test_mcp.Client.test_cancel) ... ok +test_crash_mid_call_then_restart (test_mcp.Client.test_crash_mid_call_then_restart) ... [strata] MCP server 'solo': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'solo' had stopped; starting it again +[strata] MCP server 'solo': 6 tools (echo, fail, sleep, big, die, add) +ok +test_errors_become_text (test_mcp.Client.test_errors_become_text) ... ok +test_initialize_and_list_every_page (test_mcp.Client.test_initialize_and_list_every_page) ... ok +test_namespaced_openai_tools (test_mcp.Client.test_namespaced_openai_tools) ... ok +test_timeout (test_mcp.Client.test_timeout) ... ok +test_truncation (test_mcp.Client.test_truncation) ... ok +test_bad_entries_stop_the_start (test_mcp.Config.test_bad_entries_stop_the_start) ... ok +test_both_spellings_and_the_file (test_mcp.Config.test_both_spellings_and_the_file) ... ok +[strata] MCP server 'down' did not start: could not reach http://127.0.0.1:9/mcp: . Its tools are left out; the chat works without them. +[strata] MCP server 'web': 6 tools (echo, fail, sleep, big, die, add) +test_list_and_call (test_mcp.HttpClient.test_list_and_call) ... ok +test_timeout_and_unreachable (test_mcp.HttpClient.test_timeout_and_unreachable) ... ok +test_opt_in_without_servers_is_a_plain_chat (test_mcp.NoServers.test_opt_in_without_servers_is_a_plain_chat) ... [strata] thinking limited to 5955 of max_tokens 7939 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (264689.1 tok/s) (stop, cancel=False) +ok +test_a_call_the_output_ends_inside_is_not_run (test_mcp.ToolLoop.test_a_call_the_output_ends_inside_is_not_run) +#211: the model's turn ends inside an MCP call: nothing runs (it used to run with no arguments). ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 83 tokens in 0 s (129271.2 tok/s) (stop, cancel=False) +ok +test_get_mcp (test_mcp.ToolLoop.test_get_mcp) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +ok +test_max_rounds (test_mcp.ToolLoop.test_max_rounds) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 120 tokens in 0 s (149086.6 tok/s) (stop, cancel=False) +[strata] tool fake__echo: ok, 5 characters in 0.0 s +[strata] thinking limited to 4520 of max_tokens 6026 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 120 tokens in 0 s (159327.8 tok/s) (stop, cancel=False) +[strata] tool fake__echo: ok, 5 characters in 0.0 s +[strata] thinking limited to 4380 of max_tokens 5839 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 120 tokens in 0 s (174399.3 tok/s) (stop, cancel=False) +ok +test_non_stream (test_mcp.ToolLoop.test_non_stream) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 141 tokens in 0 s (174556.3 tok/s) (stop, cancel=False) +[strata] tool fake__add: ok, 1 characters in 0.0 s +[strata] thinking limited to 4507 of max_tokens 6009 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (349525.3 tok/s) (stop, cancel=False) +ok +test_only_from_the_app_s_own_page (test_mcp.ToolLoop.test_only_from_the_app_s_own_page) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] refused an API request from the web page 'http://evil.example' (no API key; add its host to "allowed_hosts" or its origin to "cors_origins" in the config) +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (102108.5 tok/s) (stop, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +ok +test_own_tools_still_go_to_the_client (test_mcp.ToolLoop.test_own_tools_still_go_to_the_client) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4590 of max_tokens 6119 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 109 tokens in 0 s (140627.2 tok/s) (stop, cancel=False) +ok +test_plain_requests_get_no_mcp_tools (test_mcp.ToolLoop.test_plain_requests_get_no_mcp_tools) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 5949 of max_tokens 7932 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 116 tokens in 0 s (136171.1 tok/s) (stop, cancel=False) +[strata] thinking limited to 5949 of max_tokens 7932 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 16 tokens in 0 s (234646.4 tok/s) (stop, cancel=False) +[strata] thinking limited to 5949 of max_tokens 7932 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 16 tokens in 0 s (211034.2 tok/s) (stop, cancel=False) +ok +test_stop_during_a_tool (test_mcp.ToolLoop.test_stop_during_a_tool) +Closing the connection while a slow tool runs stops the tool (notifications/cancelled) and the loop. ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4665 of max_tokens 6220 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 120 tokens in 0 s (144423.7 tok/s) (stop, cancel=False) +ok +test_tool_call_then_answer (test_mcp.ToolLoop.test_tool_call_then_answer) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 134 tokens in 0 s (146478.2 tok/s) (stop, cancel=False) +[strata] tool fake__echo: ok, 19 characters in 0.0 s +[strata] thinking limited to 4499 of max_tokens 5998 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 31 tokens in 0 s (143040.1 tok/s) (stop, cancel=False) +ok +test_tool_error_is_a_result (test_mcp.ToolLoop.test_tool_error_is_a_result) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 84 tokens in 0 s (153249.9 tok/s) (stop, cancel=False) +[strata] tool fake__fail: error, 27 characters in 0.0 s +[strata] thinking limited to 4530 of max_tokens 6040 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 21 tokens in 0 s (170698.4 tok/s) (stop, cancel=False) +ok +test_truncated_for_the_model (test_mcp.ToolLoop.test_truncated_for_the_model) ... [strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 115 tokens in 0 s (150310.1 tok/s) (stop, cancel=False) +[strata] tool fake__big: ok, 3,000 characters in 0.0 s (truncated for the model) +[strata] thinking limited to 4090 of max_tokens 5453 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 16 tokens in 0 s (128807.8 tok/s) (stop, cancel=False) +ok +test_config_block_and_shared_settings_win_over_the_card (test_model_card_sampling.ModelCardSampling.test_config_block_and_shared_settings_win_over_the_card) ... [strata] done: 8 tokens in 0 s (147816.9 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 8 tokens in 0 s (228261.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 8 tokens in 0 s (250406.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_no_family_no_defaults (test_model_card_sampling.ModelCardSampling.test_no_family_no_defaults) ... [strata] done: 8 tokens in 0 s (228261.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_non_thinking_default (test_model_card_sampling.ModelCardSampling.test_non_thinking_default) ... [strata] done: 8 tokens in 0 s (252289.0 tok/s) (length, cancel=False) +ok +test_request_wins_and_zero_is_greedy (test_model_card_sampling.ModelCardSampling.test_request_wins_and_zero_is_greedy) ... [strata] done: 8 tokens in 0 s (275036.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 8 tokens in 0 s (195083.9 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_tables (test_model_card_sampling.ModelCardSampling.test_tables) ... ok +test_thinking_default (test_model_card_sampling.ModelCardSampling.test_thinking_default) ... [strata] done: 8 tokens in 0 s (284359.6 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_metrics (test_monitor.ConversationCacheCard.test_metrics) ... [strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (288433.5 tok/s) (stop, cancel=False) +ok +test_the_card_is_on_the_page (test_monitor.ConversationCacheCard.test_the_card_is_on_the_page) ... ok +test_the_log (test_monitor.ConversationCacheCard.test_the_log) ... ok +test_capture_is_bounded_and_eviction_returns_404 (test_monitor.Monitor.test_capture_is_bounded_and_eviction_returns_404) ... ok +test_error_and_engine_death_are_not_successes (test_monitor.Monitor.test_error_and_engine_death_are_not_successes) ... [strata] 400 invalid request: request exceeds the context window: prompt (183 tokens) + max tokens (20000) exceeds the context (16384); requests are never truncated. Send a smaller max_tokens (at most 16193 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] thinking capped at 128 of max_tokens 256 to leave room for the answer (budget 1024) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] thinking capped at 128 of max_tokens 256 to leave room for the answer (budget 1024) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_history_is_authorized_and_does_not_capture_headers (test_monitor.Monitor.test_history_is_authorized_and_does_not_capture_headers) ... [strata] thinking limited to 12098 of max_tokens 16130 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 50 tokens in 0 s (231473.7 tok/s) (stop, cancel=False) +ok +test_output_and_reasoning_for_both_dialects_and_stream_modes (test_monitor.Monitor.test_output_and_reasoning_for_both_dialects_and_stream_modes) ... [strata] done: 50 tokens in 0 s (222155.9 tok/s) (stop, cancel=False) +[strata] done: 50 tokens in 0 s (125277.9 tok/s) (stop, cancel=False) +[strata] thinking capped at 128 of max_tokens 256 to leave room for the answer (budget 1024) +[strata] done: 50 tokens in 0 s (240223.6 tok/s) (stop, cancel=False) +[strata] thinking capped at 128 of max_tokens 256 to leave room for the answer (budget 1024) +[strata] done: 50 tokens in 0 s (102250.2 tok/s) (stop, cancel=False) +ok +test_page_and_relative_assets_are_served (test_monitor.Monitor.test_page_and_relative_assets_are_served) ... ok +test_queue_wait_and_loading_are_attributed_separately (test_monitor.Monitor.test_queue_wait_and_loading_are_attributed_separately) ... [strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 50 tokens in 0 s (174762.7 tok/s) (stop, cancel=False) +[strata] done: 50 tokens in 0 s (182997.6 tok/s) (stop, cancel=False) +ok +test_off_by_default (test_monitor.MonitorOff.test_off_by_default) ... [strata] thinking limited to 2879 of max_tokens 3838 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 7 tokens in 0 s (131659.8 tok/s) (stop, cancel=False) +ok +test_parallel_above_the_window_says_so_unless_batch_mtp (test_parallel.ParallelArgs.test_parallel_above_the_window_says_so_unless_batch_mtp) ... ok +test_parallel_config_key (test_parallel.ParallelArgs.test_parallel_config_key) ... [strata] "parallel" must be a whole number of requests (2..8), not '4': ignored +[strata] "parallel" must be a whole number of requests (2..8), not True: ignored +[strata] "parallel": 12 - the engine runs at most 8 requests together; it will use 8 +ok +test_a_burst_of_connections_is_not_reset (test_parallel.ParallelService.test_a_burst_of_connections_is_not_reset) +30-40 clients at once got "connection reset by peer" with the listen backlog of 5 (4 x R9700 burst): the ... [strata] answering: 7 of max 8 tokens, 6.6 tok/s, 1 s +[strata] answering: 6 of max 8 tokens, 8.9 tok/s, 1 s +[strata] done: 8 tokens in 1 s (6.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 1 s so far +[strata] done: 8 tokens in 1 s (7.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 2 s so far +[strata] answering: 7 of max 8 tokens, 6.5 tok/s, 2 s +[strata] done: 8 tokens in 2 s (6.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 3 s so far +[strata] answering: 8 of max 8 tokens, 6.3 tok/s, 3 s +[strata] done: 8 tokens in 3 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 3 s so far +[strata] answering: 6 of max 8 tokens, 5.7 tok/s, 4 s +[strata] answering: 7 of max 8 tokens, 6.9 tok/s, 4 s +[strata] done: 8 tokens in 4 s (6.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 4 s so far +[strata] done: 8 tokens in 4 s (5.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 4 s so far +[strata] answering: 6 of max 8 tokens, 5.9 tok/s, 5 s +[strata] done: 8 tokens in 5 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 5 s so far +[strata] answering: 8 of max 8 tokens, 6.8 tok/s, 5 s +[strata] done: 8 tokens in 5 s (6.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 6 s so far +[strata] answering: 7 of max 8 tokens, 6.4 tok/s, 6 s +[strata] done: 8 tokens in 7 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 7 s so far +[strata] answering: 8 of max 8 tokens, 6.6 tok/s, 7 s +[strata] done: 8 tokens in 7 s (6.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 7 s so far +[strata] answering: 8 of max 8 tokens, 7.2 tok/s, 8 s +[strata] done: 8 tokens in 8 s (7.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 8 s so far +[strata] answering: 8 of max 8 tokens, 7.6 tok/s, 8 s +[strata] done: 8 tokens in 8 s (7.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 8 s so far +[strata] answering: 8 of max 8 tokens, 7.0 tok/s, 9 s +[strata] done: 8 tokens in 9 s (7.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 9 s so far +[strata] answering: 8 of max 8 tokens, 7.0 tok/s, 9 s +[strata] done: 8 tokens in 9 s (7.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 9 s so far +[strata] answering: 8 of max 8 tokens, 7.4 tok/s, 10 s +[strata] done: 8 tokens in 10 s (7.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 10 s so far +[strata] answering: 8 of max 8 tokens, 7.2 tok/s, 10 s +[strata] done: 8 tokens in 10 s (7.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 10 s so far +[strata] done: 8 tokens in 11 s (8.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 11 s so far +[strata] done: 8 tokens in 11 s (10.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 11 s so far +[strata] done: 8 tokens in 12 s (8.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 12 s so far +[strata] done: 8 tokens in 12 s (8.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 12 s so far +[strata] done: 8 tokens in 13 s (9.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 13 s so far +[strata] done: 8 tokens in 13 s (8.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 13 s so far +[strata] answering: 7 of max 8 tokens, 6.0 tok/s, 14 s +[strata] done: 8 tokens in 14 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 14 s so far +[strata] answering: 8 of max 8 tokens, 6.0 tok/s, 15 s +[strata] done: 8 tokens in 15 s (6.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 15 s so far +[strata] answering: 6 of max 8 tokens, 5.7 tok/s, 15 s +[strata] done: 8 tokens in 16 s (6.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 16 s so far +[strata] answering: 8 of max 8 tokens, 6.2 tok/s, 16 s +[strata] done: 8 tokens in 16 s (6.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 16 s so far +[strata] answering: 6 of max 8 tokens, 6.0 tok/s, 17 s +[strata] done: 8 tokens in 17 s (6.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 17 s so far +[strata] answering: 8 of max 8 tokens, 5.8 tok/s, 17 s +[strata] done: 8 tokens in 17 s (5.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 17 s so far +[strata] answering: 7 of max 8 tokens, 6.3 tok/s, 18 s +[strata] done: 8 tokens in 18 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 18 s so far +[strata] answering: 8 of max 8 tokens, 6.7 tok/s, 19 s +[strata] done: 8 tokens in 19 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 19 s so far +[strata] answering: 7 of max 8 tokens, 6.3 tok/s, 20 s +[strata] done: 8 tokens in 20 s (6.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 20 s so far +[strata] answering: 8 of max 8 tokens, 6.3 tok/s, 20 s +[strata] done: 8 tokens in 20 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 20 s so far +[strata] answering: 7 of max 8 tokens, 6.2 tok/s, 21 s +[strata] done: 8 tokens in 21 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 21 s so far +[strata] answering: 8 of max 8 tokens, 7.3 tok/s, 21 s +[strata] done: 8 tokens in 21 s (7.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 21 s so far +[strata] answering: 7 of max 8 tokens, 6.8 tok/s, 22 s +[strata] done: 8 tokens in 22 s (6.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 22 s so far +[strata] answering: 8 of max 8 tokens, 6.9 tok/s, 22 s +[strata] done: 8 tokens in 22 s (6.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 22 s so far +[strata] answering: 8 of max 8 tokens, 7.4 tok/s, 23 s +[strata] done: 8 tokens in 23 s (7.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 23 s so far +[strata] done: 8 tokens in 23 s (8.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 23 s so far +[strata] done: 8 tokens in 24 s (8.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 24 s so far +[strata] done: 8 tokens in 24 s (8.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 24 s so far +[strata] done: 8 tokens in 25 s (9.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 25 s so far +[strata] done: 8 tokens in 25 s (8.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 25 s so far +[strata] answering: 7 of max 8 tokens, 7.0 tok/s, 26 s +[strata] done: 8 tokens in 26 s (7.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 26 s so far +[strata] answering: 8 of max 8 tokens, 7.3 tok/s, 27 s +[strata] done: 8 tokens in 27 s (7.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 27 s so far +[strata] done: 8 tokens in 27 s (9.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 27 s so far +[strata] answering: 8 of max 8 tokens, 7.7 tok/s, 28 s +[strata] done: 8 tokens in 28 s (7.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 28 s so far +[strata] answering: 8 of max 8 tokens, 7.4 tok/s, 28 s +[strata] done: 8 tokens in 28 s (7.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 28 s so far +[strata] done: 8 tokens in 29 s (8.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 29 s so far +[strata] done: 8 tokens in 29 s (11.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 29 s so far +[strata] done: 8 tokens in 30 s (8.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 30 s so far +[strata] answering: 8 of max 8 tokens, 7.4 tok/s, 30 s +[strata] done: 8 tokens in 30 s (7.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 30 s so far +[strata] answering: 8 of max 8 tokens, 7.6 tok/s, 31 s +[strata] done: 8 tokens in 31 s (7.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 31 s so far +[strata] answering: 6 of max 8 tokens, 5.9 tok/s, 31 s +[strata] done: 8 tokens in 32 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 32 s so far +[strata] answering: 8 of max 8 tokens, 6.3 tok/s, 32 s +[strata] done: 8 tokens in 32 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 32 s so far +[strata] answering: 7 of max 8 tokens, 6.1 tok/s, 33 s +[strata] answering: 7 of max 8 tokens, 6.7 tok/s, 33 s +[strata] done: 8 tokens in 33 s (6.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 33 s so far +[strata] done: 8 tokens in 33 s (5.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 33 s so far +[strata] answering: 6 of max 8 tokens, 5.9 tok/s, 34 s +[strata] done: 8 tokens in 34 s (7.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 34 s so far +[strata] answering: 8 of max 8 tokens, 6.8 tok/s, 35 s +[strata] done: 8 tokens in 35 s (6.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 35 s so far +[strata] answering: 7 of max 8 tokens, 6.9 tok/s, 35 s +[strata] done: 8 tokens in 35 s (6.9 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 36 s (8.2 tok/s) (length, cancel=False) +ok +test_a_continued_request_reports_its_own_prompt_reuse (test_parallel.ParallelService.test_a_continued_request_reports_its_own_prompt_reuse) +A request STOPped on the solo path and continued in a slot (or back on the solo path) is sent again as its ... [strata] thinking: 3 of max 200 tokens, 4.0 tok/s, 1 s +[strata] reading the prompt: 192 of 249 tokens, 1 s so far +[strata] thinking: 10 of max 200 tokens, 5.6 tok/s, 2 s +[strata] thinking: 9 of max 200 tokens, 8.9 tok/s, 2 s +[strata] done: 10 tokens in 2 s (8.5 tok/s) (stop, cancel=False) +[strata] thinking: 14 of max 200 tokens, 5.0 tok/s, 3 s +[strata] thinking: 24 of max 200 tokens, 6.2 tok/s, 4 s +[strata] thinking: 36 of max 200 tokens, 7.2 tok/s, 5 s +[strata] thinking: 45 of max 200 tokens, 7.5 tok/s, 6 s +[strata] thinking: 53 of max 200 tokens, 7.4 tok/s, 8 s +[strata] done: 55 tokens in 8 s (7.3 tok/s) (stop, cancel=False) +ok +test_a_failed_window_names_its_error (test_parallel.ParallelService.test_a_failed_window_names_its_error) ... [strata] the engine stopped unexpectedly (exit code 1). The engine exited (code 1) after it reported: verify batch: layer 34 never rang (graph finished) - please report it at github.com/Niko1221/Strata/issues with the log. The next request starts the engine again. +[strata] done: 4 tokens in 1 s (4.5 tok/s) (error, cancel=False) +[strata] the engine stopped unexpectedly (exit code 1). The engine exited (code 1) after it reported: verify batch: layer 34 never rang (graph finished) - please report it at github.com/Niko1221/Strata/issues with the log. The next request starts the engine again. +[strata] done: 1 tokens in 1 s (3.0 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +ok +test_a_leaked_busy_slot_does_not_hang_the_next_request (test_parallel.ParallelService.test_a_leaked_busy_slot_does_not_hang_the_next_request) +#1603: a slot whose busy flag has no owner (no request holds it, no BDONE wait is running) kept every later ... [strata] parallel: slot(s) 0, 1 were marked busy with no request using them (issue #1603); freed, the waiting request goes on +[strata] reading the prompt: 32 of 52 tokens, 1 s so far +[strata] answering: 8 of max 16 tokens, 7.4 tok/s, 2 s +[strata] done: 10 tokens in 3 s (7.4 tok/s) (stop, cancel=False) +ok +test_a_long_prompt_gives_way_to_a_short_one (test_parallel.ParallelService.test_a_long_prompt_gives_way_to_a_short_one) +#656: a long solo prompt read gives way at a chunk boundary; the short request is answered first, then the ... [strata] reading the prompt: 352 of 3,037 tokens, 1 s so far +[strata] answering: 5 of max 64 tokens, 5.4 tok/s, 1 s +[strata] reading the prompt: 896 of 3,037 tokens, 2 s so far +[strata] done: 10 tokens in 2 s (4.8 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,440 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 2,016 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 2,688 of 3,037 tokens, 5 s so far +[strata] answering: 4 of max 64 tokens, 8.8 tok/s, 6 s +[strata] done: 10 tokens in 7 s (8.9 tok/s) (stop, cancel=False) +ok +test_a_parallel_request_does_not_wait_for_the_fifo_on_a_loaded_engine (test_parallel.ParallelService.test_a_parallel_request_does_not_wait_for_the_fifo_on_a_loaded_engine) +#1603: parallel requests used to take the FIFO (held by an image encode, /v1/vram, load) just to find the ... [strata] answering: 9 of max 64 tokens, 8.9 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.7 tok/s) (stop, cancel=False) +[strata] answering: 9 of max 64 tokens, 8.4 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.8 tok/s) (stop, cancel=False) +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 10 tokens in 1 s (8.8 tok/s) (stop, cancel=False) +ok +test_a_request_moved_between_solo_and_slot_keeps_its_totals (test_parallel.ParallelService.test_a_request_moved_between_solo_and_slot_keeps_its_totals) +#1869/#1870: a request that read its prompt and decoded on the solo path, was STOPped when another arrived and ... [strata] answering: 5 of max 200 tokens, 4.4 tok/s, 1 s +[strata] answering: 4 of max 200 tokens, 8.2 tok/s, 1 s +[strata] done: 10 tokens in 2 s (8.8 tok/s) (stop, cancel=False) +[strata] answering: 14 of max 200 tokens, 6.3 tok/s, 2 s +[strata] answering: 23 of max 200 tokens, 6.9 tok/s, 3 s +[strata] answering: 32 of max 200 tokens, 7.2 tok/s, 5 s +[strata] answering: 40 of max 200 tokens, 7.3 tok/s, 6 s +[strata] answering: 49 of max 200 tokens, 7.6 tok/s, 7 s +[strata] done: 55 tokens in 7 s (7.5 tok/s) (stop, cancel=False) +ok +test_a_slot_request_ended_by_its_stop_token_reports_the_slot_leg (test_parallel.ParallelService.test_a_slot_request_ended_by_its_stop_token_reports_the_slot_leg) +#1869: the consumer leaves at the stop token, before the BDONE: the request kept the admission's one token and ... [strata] answering: 6 of max 200 tokens, 5.9 tok/s, 1 s +[strata] answering: 5 of max 200 tokens, 7.2 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.2 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 2 s (8.4 tok/s) (stop, cancel=False) +ok +test_an_admission_gives_way_too (test_parallel.ParallelService.test_an_admission_gives_way_too) +The same while another request decodes in a slot: the long one is being admitted, a short one waits. ... [strata] answering: 4 of max 60 tokens, 4.0 tok/s, 1 s +[strata] reading the prompt: 192 of 3,037 tokens, 1 s so far +[strata] answering: 4 of max 64 tokens, 5.2 tok/s, 1 s +[strata] answering: 9 of max 60 tokens, 4.2 tok/s, 2 s +[strata] reading the prompt: 768 of 3,037 tokens, 2 s so far +[strata] done: 10 tokens in 2 s (4.2 tok/s) (stop, cancel=False) +[strata] answering: 9 of max 64 tokens, 4.7 tok/s, 2 s +[strata] done: 10 tokens in 2 s (4.7 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,376 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 1,952 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 2,560 of 3,037 tokens, 5 s so far +[strata] answering: 2 of max 64 tokens, 9.5 tok/s, 6 s +[strata] done: 10 tokens in 7 s (8.6 tok/s) (stop, cancel=False) +ok +test_an_admission_is_asked_between_far_apart_progress_lines (test_parallel.ParallelService.test_an_admission_is_asked_between_far_apart_progress_lines) +As above for a read admitted beside a decoding slot (BGEN, `ask()`): with one PP line per 80 chunks the ... [strata] reading the prompt: 2,560 of 3,037 tokens, 5 s so far +[strata] answering: 1 of max 64 tokens, 167772.2 tok/s, 5 s +[strata] answering: 2 of max 64 tokens, 0.4 tok/s, 10 s +[strata] reading the prompt: 2,560 of 3,037 tokens, 10 s so far +[strata] answering: 4 of max 60 tokens, 0.4 tok/s, 10 s +[strata] answering: 2 of max 64 tokens, 7.5 tok/s, 11 s +[strata] answering: 5 of max 60 tokens, 0.4 tok/s, 12 s +[strata] answering: 3 of max 64 tokens, 0.5 tok/s, 11 s +[strata] done: 10 tokens in 12 s (1.3 tok/s) (stop, cancel=False) +[strata] answering: 13 of max 60 tokens, 1.0 tok/s, 13 s +[strata] done: 10 tokens in 12 s (7.3 tok/s) (stop, cancel=False) +[strata] answering: 21 of max 60 tokens, 1.5 tok/s, 14 s +[strata] answering: 28 of max 60 tokens, 1.9 tok/s, 15 s +[strata] answering: 38 of max 60 tokens, 2.4 tok/s, 16 s +[strata] answering: 50 of max 60 tokens, 2.9 tok/s, 17 s +[strata] done: 55 tokens in 18 s (3.1 tok/s) (stop, cancel=False) +ok +test_an_owned_busy_slot_is_never_taken (test_parallel.ParallelService.test_an_owned_busy_slot_is_never_taken) +The leak check frees only slots with no owner: a slot a request holds stays busy however long it waits. ... [strata] parallel: slot(s) 0 were marked busy with no request using them (issue #1603); freed, the waiting request goes on +ok +test_batch_groups_auto_follows_the_engine (test_parallel.ParallelService.test_batch_groups_auto_follows_the_engine) +--batch-groups auto: the engine picks the groups and says so (INFO batch_groups=G); the server spreads the ... ok +test_bgen_refused_by_the_engine_is_retried_not_a_400 (test_parallel.ParallelService.test_bgen_refused_by_the_engine_is_retried_not_a_400) +#1994: the engine refuses a BGEN ("no such free slot") although the server saw the slot free: the request ... [strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (1/6) (issue #1994) +[strata] answering: 8 of max 64 tokens, 8.7 tok/s, 1 s +[strata] answering: 2 of max 64 tokens, 6.0 tok/s, 1 s +[strata] answering: 15 of max 64 tokens, 7.8 tok/s, 2 s +[strata] done: 10 tokens in 2 s (8.2 tok/s) (stop, cancel=False) +[strata] answering: 24 of max 64 tokens, 8.0 tok/s, 3 s +[strata] answering: 34 of max 64 tokens, 8.3 tok/s, 4 s +[strata] answering: 42 of max 64 tokens, 8.0 tok/s, 5 s +[strata] answering: 50 of max 64 tokens, 8.0 tok/s, 6 s +[strata] done: 55 tokens in 7 s (7.9 tok/s) (stop, cancel=False) +[strata] answering: 9 of max 64 tokens, 9.3 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.7 tok/s) (stop, cancel=False) +ok +test_bgen_refused_for_good_is_a_retryable_503_and_frees_the_slots (test_parallel.ParallelService.test_bgen_refused_for_good_is_a_retryable_503_and_frees_the_slots) +#1994: when the engine keeps refusing, the answer is a 503 (retryable), never a 400, and no slot stays busy. ... [strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (1/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (2/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (3/6) (issue #1994) +[strata] answering: 5 of max 64 tokens, 4.9 tok/s, 1 s +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (4/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (5/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (6/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (7/6) (issue #1994) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (length, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] answering: 12 of max 64 tokens, 5.6 tok/s, 2 s +[strata] answering: 21 of max 64 tokens, 6.4 tok/s, 3 s +[strata] answering: 29 of max 64 tokens, 6.5 tok/s, 5 s +[strata] answering: 37 of max 64 tokens, 6.8 tok/s, 6 s +[strata] answering: 48 of max 64 tokens, 7.3 tok/s, 7 s +[strata] done: 55 tokens in 8 s (7.2 tok/s) (stop, cancel=False) +[strata] answering: 8 of max 64 tokens, 8.6 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.1 tok/s) (stop, cancel=False) +ok +test_concurrent_requests_share_the_slots (test_parallel.ParallelService.test_concurrent_requests_share_the_slots) ... [strata] answering: 6 of max 64 tokens, 6.4 tok/s, 1 s +[strata] answering: 5 of max 64 tokens, 8.1 tok/s, 1 s +[strata] done: 10 tokens in 2 s (7.9 tok/s) (stop, cancel=False) +[strata] reading the prompt: 32 of 66 tokens, 2 s so far +[strata] answering: 13 of max 64 tokens, 6.1 tok/s, 2 s +[strata] answering: 8 of max 64 tokens, 7.4 tok/s, 3 s +[strata] done: 10 tokens in 3 s (8.1 tok/s) (stop, cancel=False) +[strata] answering: 22 of max 64 tokens, 7.0 tok/s, 3 s +[strata] answering: 29 of max 64 tokens, 6.8 tok/s, 4 s +[strata] answering: 39 of max 64 tokens, 7.2 tok/s, 5 s +[strata] answering: 47 of max 64 tokens, 7.3 tok/s, 7 s +[strata] done: 55 tokens in 8 s (7.3 tok/s) (stop, cancel=False) +ok +test_engine_reports_fewer_slots (test_parallel.ParallelService.test_engine_reports_fewer_slots) ... [strata] parallel requests: 4 asked, the engine runs 2 (its log says why) +ok +test_engine_turns_batching_off (test_parallel.ParallelService.test_engine_turns_batching_off) ... [strata] parallel requests: 4 asked, the engine runs one at a time (its log says why) +[strata] answering: 7 of max 64 tokens, 7.1 tok/s, 1 s +[strata] done: 10 tokens in 2 s (6.6 tok/s) (stop, cancel=False) +ok +test_left_alone_it_goes_back_to_the_solo_path (test_parallel.ParallelService.test_left_alone_it_goes_back_to_the_solo_path) +A request decoding in a slot whose neighbour finished goes back to the solo path (MTP drafts), continued ... [strata] answering: 7 of max 200 tokens, 6.6 tok/s, 1 s +[strata] answering: 5 of max 200 tokens, 6.8 tok/s, 1 s +[strata] done: 10 tokens in 2 s (7.7 tok/s) (stop, cancel=False) +[strata] answering: 15 of max 200 tokens, 7.1 tok/s, 2 s +[strata] answering: 23 of max 200 tokens, 7.2 tok/s, 3 s +[strata] answering: 30 of max 200 tokens, 7.1 tok/s, 4 s +[strata] answering: 39 of max 200 tokens, 7.3 tok/s, 5 s +[strata] answering: 46 of max 200 tokens, 7.2 tok/s, 6 s +[strata] done: 55 tokens in 8 s (7.2 tok/s) (stop, cancel=False) +ok +test_next_turn_goes_back_to_its_slot (test_parallel.ParallelService.test_next_turn_goes_back_to_its_slot) ... [strata] answering: 6 of max 200 tokens, 5.9 tok/s, 1 s +[strata] answering: 5 of max 64 tokens, 6.8 tok/s, 1 s +[strata] done: 10 tokens in 2 s (8.0 tok/s) (stop, cancel=False) +[strata] answering: 16 of max 200 tokens, 7.5 tok/s, 2 s +[strata] answering: 23 of max 200 tokens, 7.2 tok/s, 3 s +[strata] answering: 33 of max 200 tokens, 7.7 tok/s, 4 s +[strata] answering: 43 of max 200 tokens, 8.0 tok/s, 5 s +[strata] answering: 51 of max 200 tokens, 8.0 tok/s, 6 s +[strata] done: 55 tokens in 7 s (8.0 tok/s) (stop, cancel=False) +ok +test_no_way_given_to_a_long_one (test_parallel.ParallelService.test_no_way_given_to_a_long_one) +Two prompts of about the same length: no BYIELD (it only goes to a much shorter one). ... [strata] reading the prompt: 608 of 3,037 tokens, 1 s so far +[strata] reading the prompt: 1,120 of 3,037 tokens, 2 s so far +[strata] reading the prompt: 1,728 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 2,400 of 3,037 tokens, 4 s so far +[strata] answering: 1 of max 64 tokens, 6.3 tok/s, 5 s +[strata] answering: 1 of max 64 tokens, 0.8 tok/s, 6 s +[strata] answering: 1 of max 64 tokens, 0.4 tok/s, 7 s +[strata] answering: 1 of max 64 tokens, 0.3 tok/s, 8 s +[strata] answering: 1 of max 64 tokens, 0.2 tok/s, 9 s +[strata] reading the prompt: 3,008 of 3,038 tokens, 10 s so far +[strata] answering: 3 of max 64 tokens, 0.5 tok/s, 11 s +[strata] reading the prompt: 416 of 3,037 tokens, 11 s so far +[strata] answering: 7 of max 64 tokens, 1.1 tok/s, 12 s +[strata] reading the prompt: 960 of 3,037 tokens, 12 s so far +[strata] done: 10 tokens in 13 s (1.3 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,536 of 3,037 tokens, 13 s so far +[strata] reading the prompt: 2,144 of 3,037 tokens, 14 s so far +[strata] reading the prompt: 2,752 of 3,037 tokens, 15 s so far +[strata] answering: 5 of max 64 tokens, 8.6 tok/s, 16 s +[strata] done: 10 tokens in 17 s (8.6 tok/s) (stop, cancel=False) +ok +test_no_way_given_to_a_long_one_between_far_apart_progress_lines (test_parallel.ParallelService.test_no_way_given_to_a_long_one_between_far_apart_progress_lines) +The ask between PP lines keeps the rule: a waiting request about as long gets no BYIELD there either. ... [strata] reading the prompt: 2,560 of 3,037 tokens, 5 s so far +[strata] reading the prompt: 2,560 of 3,037 tokens, 10 s so far +[strata] answering: 1 of max 64 tokens, 0.2 tok/s, 11 s +[strata] answering: 3 of max 64 tokens, 0.3 tok/s, 17 s +[strata] reading the prompt: 2,560 of 3,037 tokens, 17 s so far +[strata] answering: 2 of max 64 tokens, 6.4 tok/s, 18 s +[strata] answering: 4 of max 64 tokens, 0.3 tok/s, 18 s +[strata] done: 10 tokens in 18 s (0.8 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 19 s (7.6 tok/s) (stop, cancel=False) +ok +test_nonfinite_in_batch_mode (test_parallel.ParallelService.test_nonfinite_in_batch_mode) +#879 in --batch mode: DONE ... nonfinite on the solo path or at an admission (nothing emitted yet) sends the ... [strata] the engine's logits became non-finite before the first token (#879); it dropped its cached state: sending the request again once +[strata] done: 10 tokens in 1 s (9.6 tok/s) (stop, cancel=False) +[strata] the engine's logits became non-finite before the first token (#879); it dropped its cached state: sending the request again once +[strata] answering: 5 of max 64 tokens, 4.5 tok/s, 1 s +[strata] answering: 3 of max 64 tokens, 6.2 tok/s, 1 s +[strata] answering: 13 of max 64 tokens, 6.0 tok/s, 2 s +[strata] done: 10 tokens in 2 s (7.3 tok/s) (stop, cancel=False) +[strata] answering: 21 of max 64 tokens, 6.4 tok/s, 3 s +[strata] answering: 29 of max 64 tokens, 6.7 tok/s, 4 s +[strata] answering: 37 of max 64 tokens, 6.9 tok/s, 5 s +[strata] answering: 48 of max 64 tokens, 7.4 tok/s, 7 s +[strata] done: 55 tokens in 8 s (7.3 tok/s) (stop, cancel=False) +[strata] the engine's logits became non-finite after 2 tokens (#879): the reply ends here as "length"; the engine dropped its cached state, so the next request reads its prompt from the start +[strata] done: 2 tokens in 1 s (6.8 tok/s) (length, cancel=False) +[strata] answering: 6 of max 64 tokens, 6.0 tok/s, 1 s +[strata] answering: 16 of max 64 tokens, 7.4 tok/s, 2 s +[strata] answering: 24 of max 64 tokens, 7.6 tok/s, 3 s +[strata] answering: 32 of max 64 tokens, 7.6 tok/s, 4 s +[strata] answering: 39 of max 64 tokens, 7.4 tok/s, 5 s +[strata] answering: 47 of max 64 tokens, 7.4 tok/s, 6 s +[strata] done: 55 tokens in 7 s (7.9 tok/s) (stop, cancel=False) +ok +test_stop_strings_in_a_batch_slot (test_parallel.ParallelService.test_stop_strings_in_a_batch_slot) +#454: a stop string cuts the answer in --batch mode too, and the slot is freed for the next request. ... [strata] answering: 9 of max 64 tokens, 8.7 tok/s, 1 s +[strata] done: 9 tokens in 1 s (7.8 tok/s) (stop, cancel=False) +[strata] answering: 8 of max 64 tokens, 8.4 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.4 tok/s) (stop, cancel=False) +ok +test_the_way_is_asked_between_far_apart_progress_lines (test_parallel.ParallelService.test_the_way_is_asked_between_far_apart_progress_lines) +#656 when the PP lines come far apart (the pipelined groups' read in pieces of one chunk: one PP per piece, ... [strata] reading the prompt: 2,560 of 3,037 tokens, 5 s so far +[strata] answering: 1 of max 64 tokens, 161319.4 tok/s, 5 s +[strata] answering: 2 of max 64 tokens, 0.4 tok/s, 10 s +[strata] reading the prompt: 2,560 of 3,037 tokens, 10 s so far +[strata] answering: 3 of max 64 tokens, 0.5 tok/s, 11 s +[strata] answering: 2 of max 64 tokens, 6.0 tok/s, 11 s +[strata] done: 10 tokens in 12 s (1.5 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 12 s (7.6 tok/s) (stop, cancel=False) +ok +test_three_requests_long_long_short_do_not_deadlock (test_parallel.ParallelService.test_three_requests_long_long_short_do_not_deadlock) +ENGINE_REVIEW finding 1 (0.1.41 check): two long prompts fill both slots, a short one arrives a moment later ... [strata] reading the prompt: 32 of 3,037 tokens, 1 s so far +[strata] reading the prompt: 480 of 3,037 tokens, 2 s so far +[strata] reading the prompt: 1,056 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 1,568 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 2,176 of 3,037 tokens, 5 s so far +[strata] reading the prompt: 2,720 of 3,037 tokens, 6 s so far +[strata] reading the prompt: 32 of 3,037 tokens, 7 s so far +[strata] answering: 3 of max 64 tokens, 5.6 tok/s, 7 s +[strata] reading the prompt: 672 of 3,037 tokens, 8 s so far +[strata] answering: 8 of max 64 tokens, 5.1 tok/s, 8 s +[strata] done: 10 tokens in 9 s (4.9 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,216 of 3,037 tokens, 9 s so far +[strata] reading the prompt: 1,760 of 3,037 tokens, 10 s so far +[strata] reading the prompt: 2,272 of 3,037 tokens, 11 s so far +[strata] reading the prompt: 2,784 of 3,037 tokens, 12 s so far +[strata] reading the prompt: 32 of 43 tokens, 12 s so far +[strata] answering: 2 of max 64 tokens, 3.4 tok/s, 13 s +[strata] answering: 8 of max 64 tokens, 7.4 tok/s, 13 s +[strata] done: 10 tokens in 14 s (6.4 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 13 s (7.7 tok/s) (stop, cancel=False) +ok +test_a_rewritten_next_turn_goes_to_the_slot_of_its_chat (test_parallel.PickSlot.test_a_rewritten_next_turn_goes_to_the_slot_of_its_chat) ... ok +test_batch_groups_spread_requests_over_the_groups (test_parallel.PickSlot.test_batch_groups_spread_requests_over_the_groups) ... ok +test_slots_view (test_parallel.PickSlot.test_slots_view) ... ok +test_the_slot_that_holds_the_conversation (test_parallel.PickSlot.test_the_slot_that_holds_the_conversation) ... ok +test_chunk_boundaries_literal_tags_multiple_calls_and_recovery (test_parser_call_scan.CallScanTests.test_chunk_boundaries_literal_tags_multiple_calls_and_recovery) ... ok +test_no_structural_rescans_until_closer_and_reset_for_next_call (test_parser_call_scan.CallScanTests.test_no_structural_rescans_until_closer_and_reset_for_next_call) ... ok +test_every_real_tag_split (test_parser_hold.HoldTests.test_every_real_tag_split) ... ok +test_overlapping_and_empty_tags (test_parser_hold.HoldTests.test_overlapping_and_empty_tags) ... ok +test_streamed_events_match_reference (test_parser_hold.HoldTests.test_streamed_events_match_reference) ... ok +test_compaction_boundaries_reads_and_line_replacement (test_parser_track.TrackTests.test_compaction_boundaries_reads_and_line_replacement) ... ok +test_feed_events_with_long_lines_and_quoted_tool_tags (test_parser_track.TrackTests.test_feed_events_with_long_lines_and_quoted_tool_tags) ... ok +test_fragment_state_from_varied_initial_states (test_parser_track.TrackTests.test_fragment_state_from_varied_initial_states) ... ok +test_long_line_snapshots_restore_and_append (test_parser_track.TrackTests.test_long_line_snapshots_restore_and_append) ... ok +test_random_streaming_boundaries (test_parser_track.TrackTests.test_random_streaming_boundaries) ... ok +test_every_concurrent_request_is_counted (test_prometheus.Batched.test_every_concurrent_request_is_counted) ... [strata] reading the prompt: 10 of max 64 tokens, 9.7 tok/s, 1 s +[strata] reading the prompt: 10 of max 64 tokens, 10.4 tok/s, 1 s +[strata] reading the prompt: 10 of max 64 tokens, 9.4 tok/s, 1 s +[strata] answering: 20 of max 64 tokens, 9.8 tok/s, 2 s +[strata] answering: 19 of max 64 tokens, 8.7 tok/s, 2 s +[strata] answering: 18 of max 64 tokens, 8.6 tok/s, 2 s +[strata] answering: 28 of max 64 tokens, 9.0 tok/s, 3 s +[strata] answering: 25 of max 64 tokens, 8.1 tok/s, 3 s +[strata] answering: 28 of max 64 tokens, 8.4 tok/s, 3 s +[strata] answering: 38 of max 64 tokens, 8.9 tok/s, 4 s +[strata] answering: 35 of max 64 tokens, 8.2 tok/s, 4 s +[strata] answering: 37 of max 64 tokens, 8.3 tok/s, 5 s +[strata] answering: 47 of max 64 tokens, 8.7 tok/s, 5 s +[strata] answering: 46 of max 64 tokens, 8.6 tok/s, 5 s +[strata] answering: 47 of max 64 tokens, 8.4 tok/s, 6 s +[strata] done: 51 tokens in 6 s (8.5 tok/s) (stop, cancel=False) +[strata] done: 51 tokens in 6 s (8.6 tok/s) (stop, cancel=False) +[strata] done: 51 tokens in 6 s (8.4 tok/s) (stop, cancel=False) +ok +test_a_request_shows_in_both_formats (test_prometheus.Endpoint.test_a_request_shows_in_both_formats) ... [strata] done: 22 tokens in 0 s (270600.3 tok/s) (stop, cancel=False) +ok +test_histograms_are_cumulative (test_prometheus.Units.test_histograms_are_cumulative) ... ok +test_render_reads_the_server_record (test_prometheus.Units.test_render_reads_the_server_record) ... ok +test_strata_names_follow_the_json (test_prometheus.Units.test_strata_names_follow_the_json) ... ok +test_who_gets_the_text (test_prometheus.Units.test_who_gets_the_text) ... ok +test_a_refused_prompt_is_not_kept (test_prompt_encoder.Served.test_a_refused_prompt_is_not_kept) ... ok +test_another_conversation_does_not_evict_the_first (test_prompt_encoder.Served.test_another_conversation_does_not_evict_the_first) +A subagent beside the main conversation: every rendered prompt starts with the template's header, which ... ok +test_quoted_think_across_turns (test_prompt_encoder.Served.test_quoted_think_across_turns) ... ok +test_random_messages_equal_a_full_encode (test_prompt_encoder.Served.test_random_messages_equal_a_full_encode) +Messages quoting control-token texts and tags at random, growing, edited and cut back between ... ok +test_a_turn_reuses_the_previous_one (test_prompt_encoder.Units.test_a_turn_reuses_the_previous_one) ... ok +test_plain_agree_until (test_prompt_encoder.Units.test_plain_agree_until) ... ok +test_same_ids_as_a_full_encode (test_prompt_encoder.Units.test_same_ids_as_a_full_encode) ... ok +test_a_vision_section_starts_the_engine_with_pictures (test_q36_config.Q36Config.test_a_vision_section_starts_the_engine_with_pictures) ... [strata] the config has a "vision" section but no --vision in its args: added --vision (run setup again with --vision to write them) +ok +test_an_explicit_batch_in_the_args_wins (test_q36_config.Q36Config.test_an_explicit_batch_in_the_args_wins) ... ok +test_more_than_the_window_is_said (test_q36_config.Q36Config.test_more_than_the_window_is_said) ... ok +test_parallel_becomes_batch (test_q36_config.Q36Config.test_parallel_becomes_batch) ... ok +test_a_queued_stream_gets_keep_alives_then_its_answer (test_queue_keepalive.QueueKeepAlive.test_a_queued_stream_gets_keep_alives_then_its_answer) ... [strata] thinking limited to 24386 of max_tokens 32514 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 3 tokens in 0 s (9642.1 tok/s) (stop, cancel=False) +ok +test_a_request_cancelled_while_queued_never_takes_the_turn (test_queue_keepalive.QueueKeepAlive.test_a_request_cancelled_while_queued_never_takes_the_turn) ... ok +test_closing_a_queued_request_is_no_longer_queued (test_queue_keepalive.QueueKeepAlive.test_closing_a_queued_request_is_no_longer_queued) ... ok +test_anthropic_stream (test_quiet_cancel.QuietPromptCancel.test_anthropic_stream) ... [strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (39945.8 tok/s) (length, cancel=False) +ok +test_batch_solo_stream (test_quiet_cancel.QuietPromptCancel.test_batch_solo_stream) ... [strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (5777.3 tok/s) (length, cancel=False) +ok +test_other_http_modes (test_quiet_cancel.QuietPromptCancel.test_other_http_modes) ... [strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (5433.0 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] answering: 1 of max 1 tokens, 167772.2 tok/s, 1 s +[strata] done: 1 tokens in 1 s (4718.0 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (6955.7 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (6250.8 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (5817.3 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (16644.1 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (33288.1 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (14122.2 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (5991.9 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (7358.4 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (3663.1 tok/s) (length, cancel=False) +ok +test_expired_silence_is_not_a_heartbeat (test_quiet_cancel.QuietWaitPolicy.test_expired_silence_is_not_a_heartbeat) ... ok +test_polling_keeps_ten_second_heartbeats (test_quiet_cancel.QuietWaitPolicy.test_polling_keeps_ten_second_heartbeats) ... ok +test_ready_timeout_on_lazy_load_answers_the_request (test_ready_timeout.ReadyTimeout.test_ready_timeout_on_lazy_load_answers_the_request) ... [strata] loading the model again (it was unloaded) ... +[strata] the engine did not start (try 1 of 3); again in 0 s +[strata] the engine did not start (try 2 of 3); again in 0 s +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +ok +test_ready_timeout_on_restart_answers_the_request (test_ready_timeout.ReadyTimeout.test_ready_timeout_on_restart_answers_the_request) ... [strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 1 tokens in 0 s (14463.1 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine did not start (try 1 of 3); again in 0 s +[strata] the engine did not start (try 2 of 3); again in 0 s +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +ok +test_a_loop_written_as_one_long_word (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_a_loop_written_as_one_long_word) +#1753: a 24k-digit string cycling an 11-digit pattern is one word to the split and never one token twice. ... ok +test_a_lower_top_p_and_top_k_are_never_raised_or_lowered (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_a_lower_top_p_and_top_k_are_never_raised_or_lowered) ... [strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2086 tokens in 0 s (244026.3 tok/s) (stop, cancel=False) +ok +test_batch_status_and_history_record_recovery (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_batch_status_and_history_record_recovery) ... [strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2071 tokens in 0 s (397674.5 tok/s) (stop, cancel=False) +ok +test_cancel_during_reasoning_does_not_start_recovery (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_cancel_during_reasoning_does_not_start_recovery) ... [strata] done: 1 tokens in 0 s (25575.0 tok/s) (cancel, cancel=True) +ok +test_default_single_token_guard_does_not_detect_passage_loop (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_default_single_token_guard_does_not_detect_passage_loop) ... [strata] done: 16384 tokens in 0 s (672408.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_disabled_is_one_pass (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_disabled_is_one_pass) ... [strata] done: 28001 tokens in 0 s (795254.0 tok/s) (stop, cancel=False) +ok +test_frozen_cohort_sampler_is_unchanged_during_recovery (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_frozen_cohort_sampler_is_unchanged_during_recovery) ... [strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2082 tokens in 0 s (723971.2 tok/s) (stop, cancel=False) +ok +test_missing_native_high_hint_skips_recovery (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_missing_native_high_hint_skips_recovery) ... [strata] done: 28001 tokens in 0 s (538895.4 tok/s) (stop, cancel=False) +ok +test_natural_answer_with_shared_budget_and_preserved_prefix (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_natural_answer_with_shared_budget_and_preserved_prefix) ... [strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2106 tokens in 0 s (708413.2 tok/s) (stop, cancel=False) +ok +test_old_explicit_sampler_recovers_without_mutating_request (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_old_explicit_sampler_recovers_without_mutating_request) ... [strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 2086 tokens in 0 s (709711.1 tok/s) (stop, cancel=False) +ok +test_preserves_task_and_exact_generated_tokens (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_preserves_task_and_exact_generated_tokens) ... ok +test_recovery_is_once_and_never_extends_output_limit (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_recovery_is_once_and_never_extends_output_limit) ... [strata] repeated reasoning detected at 2048 tokens (coverage=1.000); resuming the same output with the low-effort instruction (reasoning_loop_recovery) +[strata] done: 16384 tokens in 0 s (760146.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_repeated_final_text_does_not_change_sampler_or_prompt (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_repeated_final_text_does_not_change_sampler_or_prompt) ... [strata] done: 28001 tokens in 0 s (790234.9 tok/s) (stop, cancel=False) +ok +test_repetition_and_novel_work (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_repetition_and_novel_work) ... ok +test_stop_mode_ends_a_loop_written_as_one_long_word (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_stop_mode_ends_a_loop_written_as_one_long_word) +#1753: the word count never sees the digits, so only the tail's period ends the reply. ... [strata] the thinking repeated the same passages (coverage=1.000) at 2560 tokens: ended as "length" (reasoning_loop_recovery: "stop" in strata-.json; remove it to turn this off) +[strata] done: 2560 tokens in 0 s (747470.8 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_stop_mode_ends_the_reply_at_the_loop (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_stop_mode_ends_the_reply_at_the_loop) ... [strata] the thinking repeated the same passages (coverage=1.000) at 2048 tokens: ended as "length" (reasoning_loop_recovery: "stop" in strata-.json; remove it to turn this off) +[strata] done: 2048 tokens in 0 s (740384.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_the_splice_keeps_a_literal_think_tag_as_text (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_the_splice_keeps_a_literal_think_tag_as_text) +#537: the prompt is spliced as token ids, never decoded and encoded again. ... ok +test_user_instruction_is_never_rewritten (test_reasoning_loop_recovery.ReasoningLoopRecovery.test_user_instruction_is_never_rewritten) ... ok +test_every_specimen_whole_and_in_chunks (test_reasoning_rescue.Corpus.test_every_specimen_whole_and_in_chunks) ... ok +test_the_corpus_has_the_reported_split (test_reasoning_rescue.Corpus.test_the_corpus_has_the_reported_split) ... ok +test_a_call_waits_and_is_not_streamed_early (test_reasoning_rescue.Gates.test_a_call_waits_and_is_not_streamed_early) ... ok +test_a_fence_needs_the_same_marker_to_close (test_reasoning_rescue.Gates.test_a_fence_needs_the_same_marker_to_close) ... ok +test_a_length_cut_keeps_even_a_clean_stranded_call_as_text (test_reasoning_rescue.Gates.test_a_length_cut_keeps_even_a_clean_stranded_call_as_text) ... ok +test_a_stranded_call_is_recovered (test_reasoning_rescue.Gates.test_a_stranded_call_is_recovered) ... ok +test_after_a_closed_fence_a_call_counts_again (test_reasoning_rescue.Gates.test_after_a_closed_fence_a_call_counts_again) ... ok +test_an_undeclared_name_stays_reasoning (test_reasoning_rescue.Gates.test_an_undeclared_name_stays_reasoning) ... ok +test_fence_split_across_chunks (test_reasoning_rescue.Gates.test_fence_split_across_chunks) ... ok +test_no_tools_declared_is_unchanged (test_reasoning_rescue.Gates.test_no_tools_declared_is_unchanged) ... ok +test_opener_split_across_chunks (test_reasoning_rescue.Gates.test_opener_split_across_chunks) ... ok +test_recovered_call_keeps_the_text_before_it (test_reasoning_rescue.Gates.test_recovered_call_keeps_the_text_before_it) ... ok +test_the_counters (test_reasoning_rescue.Gates.test_the_counters) ... ok +test_the_end_of_the_turn_splits_a_held_tail (test_reasoning_rescue.Gates.test_the_end_of_the_turn_splits_a_held_tail) ... ok +test_the_reporters_negative_controls (test_reasoning_rescue.Gates.test_the_reporters_negative_controls) ... ok +test_a_stranded_call_is_a_call_in_every_api (test_reasoning_rescue.OverHttp.test_a_stranded_call_is_a_call_in_every_api) ... [strata] done: 116 tokens in 0 s (197619.5 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (153047.9 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (225249.7 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (163377.9 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (215569.0 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 116 tokens in 0 s (159783.0 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +ok +test_quotes_and_cuts_are_not (test_reasoning_rescue.OverHttp.test_quotes_and_cuts_are_not) ... [strata] done: 133 tokens in 0 s (165679.4 tok/s) (stop, cancel=False) +[strata] done: 133 tokens in 0 s (98890.7 tok/s) (stop, cancel=False) +[strata] done: 133 tokens in 0 s (176924.3 tok/s) (stop, cancel=False) +[strata] done: 133 tokens in 0 s (103823.3 tok/s) (stop, cancel=False) +[strata] done: 133 tokens in 0 s (179717.3 tok/s) (stop, cancel=False) +[strata] done: 133 tokens in 0 s (102000.8 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (175136.9 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (97348.5 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (184833.7 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (111176.6 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (182225.6 tok/s) (stop, cancel=False) +[strata] done: 117 tokens in 0 s (100766.6 tok/s) (stop, cancel=False) +[strata] done: 150 tokens in 0 s (157247.1 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 150 tokens in 0 s (119723.2 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 150 tokens in 0 s (169901.6 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 150 tokens in 0 s (127100.1 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 150 tokens in 0 s (173413.9 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 150 tokens in 0 s (124583.3 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (193831.8 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (151045.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (199289.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (162418.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (209715.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 108 tokens in 0 s (169530.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (210911.1 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (179703.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (218031.9 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (207258.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (220872.7 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 97 tokens in 0 s (205894.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] tool calls inside the thinking: 0 read as calls, 1 kept as reasoning (quoted, or the turn did not end on a stop) +ok +test_visible_answer_fences_on_all_three_apis (test_reasoning_rescue.OverHttp.test_visible_answer_fences_on_all_three_apis) ... [strata] done: 134 tokens in 0 s (160124.4 tok/s) (stop, cancel=False) +[strata] done: 134 tokens in 0 s (94222.4 tok/s) (stop, cancel=False) +[strata] done: 134 tokens in 0 s (178708.0 tok/s) (stop, cancel=False) +[strata] done: 134 tokens in 0 s (112972.2 tok/s) (stop, cancel=False) +[strata] done: 134 tokens in 0 s (166826.0 tok/s) (stop, cancel=False) +[strata] done: 134 tokens in 0 s (100381.6 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (198945.8 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (111737.3 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (204641.4 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (129137.6 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (196481.1 tok/s) (stop, cancel=False) +[strata] done: 121 tokens in 0 s (108096.0 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (188873.9 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (138457.4 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (207479.4 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (148697.8 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (187058.5 tok/s) (stop, cancel=False) +[strata] done: 116 tokens in 0 s (122894.5 tok/s) (stop, cancel=False) +ok +test_a_closed_call_in_a_length_cut_answer_stays_a_call (test_reasoning_rescue.VisibleAnswer.test_a_closed_call_in_a_length_cut_answer_stays_a_call) ... ok +test_quoted_calls_stay_text (test_reasoning_rescue.VisibleAnswer.test_quoted_calls_stay_text) ... ok +test_real_calls_still_fire (test_reasoning_rescue.VisibleAnswer.test_real_calls_still_fire) ... ok +test_without_thinking (test_reasoning_rescue.VisibleAnswer.test_without_thinking) ... ok +test_finish_reasons_and_block_order (test_reasoning_tools.OverHttp.test_finish_reasons_and_block_order) ... [strata] done: 121 tokens in 0 s (186174.2 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 121 tokens in 0 s (144548.8 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 121 tokens in 0 s (192969.9 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 121 tokens in 0 s (154587.5 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +[strata] done: 121 tokens in 0 s (149796.6 tok/s) (stop, cancel=False) +[strata] tool calls inside the thinking: 1 read as calls, 0 kept as reasoning (quoted, or the turn did not end on a stop) +ok +test_a_call_in_the_answer_is_unchanged (test_reasoning_tools.Parser.test_a_call_in_the_answer_is_unchanged) ... ok +test_a_declared_call_in_the_reasoning_is_a_call (test_reasoning_tools.Parser.test_a_declared_call_in_the_reasoning_is_a_call) ... ok +test_a_partial_tag_is_held (test_reasoning_tools.Parser.test_a_partial_tag_is_held) ... ok +test_a_tag_that_never_closes_does_not_hold_the_thinking_back (test_reasoning_tools.Parser.test_a_tag_that_never_closes_does_not_hold_the_thinking_back) ... ok +test_an_undeclared_name_stays_reasoning (test_reasoning_tools.Parser.test_an_undeclared_name_stays_reasoning) ... ok +test_an_unfinished_call_is_never_a_call (test_reasoning_tools.Parser.test_an_unfinished_call_is_never_a_call) ... ok +test_no_tools_declared_is_unchanged (test_reasoning_tools.Parser.test_no_tools_declared_is_unchanged) ... ok +test_the_thinking_ending_inside_a_call_makes_it_text (test_reasoning_tools.Parser.test_the_thinking_ending_inside_a_call_makes_it_text) ... ok +test_a_call_quoted_after_it_is_not_run (test_reasoning_tools.QuotedThinkEnd.test_a_call_quoted_after_it_is_not_run) ... ok +test_a_quoted_tag_stays_reasoning (test_reasoning_tools.QuotedThinkEnd.test_a_quoted_tag_stays_reasoning) ... ok +test_an_output_that_ends_on_the_tag (test_reasoning_tools.QuotedThinkEnd.test_an_output_that_ends_on_the_tag) ... ok +test_the_end_after_a_quote_still_ends_the_thinking (test_reasoning_tools.QuotedThinkEnd.test_the_end_after_a_quote_still_ends_the_thinking) ... ok +test_a_call_cut_at_the_token_limit (serve.test_server.UnfinishedToolCall.test_a_call_cut_at_the_token_limit) +The same cut by max_tokens: "length" / "max_tokens", and the whole (non-streamed) answers leave the call ... [strata] done: 119 tokens in 0 s (184791.6 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (126712.9 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (206932.9 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (132674.7 tok/s) (length, cancel=False) +ok +test_a_cut_call (serve.test_server.UnfinishedToolCall.test_a_cut_call) ... [strata] done: 120 tokens in 0 s (166385.6 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (130123.2 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (189859.1 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (135921.3 tok/s) (stop, cancel=False) +ok +test_a_whole_call (serve.test_server.UnfinishedToolCall.test_a_whole_call) ... [strata] done: 147 tokens in 0 s (144258.9 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (124684.1 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (168597.9 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (138397.9 tok/s) (stop, cancel=False) +ok +test_collect_keeps_calls_whose_arguments_parse (serve.test_server.UnfinishedToolCall.test_collect_keeps_calls_whose_arguments_parse) ... ok +test_parser (serve.test_server.UnfinishedToolCall.test_parser) ... ok +test_a_number_cuts_the_gpu_list (test_replicas.Config.test_a_number_cuts_the_gpu_list) ... ok +test_cpu_sets_off (test_replicas.Config.test_cpu_sets_off) ... ok +test_mistakes_are_said (test_replicas.Config.test_mistakes_are_said) ... ok +test_named_groups (test_replicas.Config.test_named_groups) ... ok +test_replica_cfg (test_replicas.Config.test_replica_cfg) ... ok +test_a_conversation_goes_back_to_its_replica (test_replicas.ReplicaService.test_a_conversation_goes_back_to_its_replica) ... [strata] answering: 2 of max 64 tokens, 15.0 tok/s, 1 s +[strata] done: 10 tokens in 2 s (8.3 tok/s) (stop, cancel=False) +[strata] reading the prompt: 576 of 581 tokens, 1 s so far +[strata] done: 10 tokens in 2 s (8.5 tok/s) (stop, cancel=False) +ok +test_a_replica_that_died_is_skipped (test_replicas.ReplicaService.test_a_replica_that_died_is_skipped) ... [strata] answering: 8 of max 64 tokens, 7.2 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.1 tok/s) (stop, cancel=False) +[strata] answering: 8 of max 64 tokens, 8.6 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.6 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 1 s (11.5 tok/s) (stop, cancel=False) +ok +test_engines_without_slots_run_one_each_at_once (test_replicas.ReplicaService.test_engines_without_slots_run_one_each_at_once) ... [strata] answering: 9 of max 64 tokens, 8.7 tok/s, 1 s +[strata] answering: 9 of max 64 tokens, 9.0 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.8 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 1 s (7.6 tok/s) (stop, cancel=False) +ok +test_metrics_and_status_see_all_slots (test_replicas.ReplicaService.test_metrics_and_status_see_all_slots) ... ok +test_requests_spread_over_both_and_answers_are_the_solo_answer (test_replicas.ReplicaService.test_requests_spread_over_both_and_answers_are_the_solo_answer) ... [strata] done: 10 tokens in 1 s (11.0 tok/s) (stop, cancel=False) +[strata] answering: 1 of max 64 tokens, 3.4 tok/s, 1 s +[strata] answering: 1 of max 64 tokens, 1.8 tok/s, 1 s +[strata] reading the prompt: 32 of 355 tokens, 1 s so far +[strata] reading the prompt: 32 of 355 tokens, 2 s so far +[strata] reading the prompt: 32 of 355 tokens, 2 s so far +[strata] reading the prompt: 32 of 355 tokens, 2 s so far +[strata] answering: 5 of max 64 tokens, 2.7 tok/s, 2 s +[strata] answering: 5 of max 64 tokens, 2.8 tok/s, 2 s +[strata] answering: 3 of max 64 tokens, 5.2 tok/s, 3 s +[strata] answering: 4 of max 64 tokens, 4.8 tok/s, 3 s +[strata] answering: 4 of max 64 tokens, 8.5 tok/s, 3 s +[strata] done: 10 tokens in 3 s (3.9 tok/s) (stop, cancel=False) +[strata] answering: 6 of max 64 tokens, 10.0 tok/s, 3 s +[strata] done: 10 tokens in 3 s (3.5 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 4 s (7.1 tok/s) (stop, cancel=False) +[strata] answering: 9 of max 64 tokens, 5.7 tok/s, 4 s +[strata] done: 10 tokens in 4 s (5.9 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 4 s (10.0 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 4 s (8.5 tok/s) (stop, cancel=False) +ok +test_a_dead_replica_is_skipped (test_replicas.RouteChoice.test_a_dead_replica_is_skipped) ... ok +test_a_long_prefix_outweighs_load (test_replicas.RouteChoice.test_a_long_prefix_outweighs_load) ... ok +test_batch_is_the_sum (test_replicas.RouteChoice.test_batch_is_the_sum) ... ok +test_least_loaded_then_affinity (test_replicas.RouteChoice.test_least_loaded_then_affinity) ... ok +test_nonfinite_forgets_the_replica (test_replicas.RouteChoice.test_nonfinite_forgets_the_replica) +#879: after DONE ... nonfinite the engine holds no conversation: the router forgets that replica. ... ok +test_matched_prefix_per_replica (test_replicas.Routing.test_matched_prefix_per_replica) ... ok +test_the_index_is_bounded (test_replicas.Routing.test_the_index_is_bounded) ... ok +test_anthropic_nonstream_classifies_only_original_input_as_cached (test_request_accounting.RequestAccounting.test_anthropic_nonstream_classifies_only_original_input_as_cached) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (174244.1 tok/s) (stop, cancel=False), expert cache 88.9% hit +ok +test_anthropic_stream_classifies_only_original_input_as_cached (test_request_accounting.RequestAccounting.test_anthropic_stream_classifies_only_original_input_as_cached) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (86752.0 tok/s) (stop, cancel=False), expert cache 88.9% hit +ok +test_cancelled_continuation_keeps_original_input_and_counts_each_done_once (test_request_accounting.RequestAccounting.test_cancelled_continuation_keeps_original_input_and_counts_each_done_once) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 98 tokens in 0 s (90498.0 tok/s) (cancel, cancel=True), expert cache 88.9% hit +ok +test_failed_continuation_does_not_count_the_first_done_twice (test_request_accounting.RequestAccounting.test_failed_continuation_does_not_count_the_first_done_twice) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] the continuation ended without a DONE line. The next request starts the engine again. +[strata] done: 95 tokens in 0 s (93954.0 tok/s) (error, cancel=False), expert cache 90.0% hit +ok +test_openai_nonstream_keeps_original_cache_and_all_segment_work (test_request_accounting.RequestAccounting.test_openai_nonstream_keeps_original_cache_and_all_segment_work) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (159565.9 tok/s) (stop, cancel=False), expert cache 88.9% hit +ok +test_openai_stream_keeps_original_cache_and_all_segment_work (test_request_accounting.RequestAccounting.test_openai_stream_keeps_original_cache_and_all_segment_work) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (92020.0 tok/s) (stop, cancel=False), expert cache 88.9% hit +ok +test_single_segment_keeps_its_own_usage_and_timings (test_request_accounting.RequestAccounting.test_single_segment_keeps_its_own_usage_and_timings) ... [strata] thinking capped at 200 of max_tokens 400 to leave room for the answer (budget 10000) +[strata] done: 211 tokens in 0 s (115309.2 tok/s) (stop, cancel=False), expert cache 90.0% hit +ok +test_prompt_n_is_what_was_read (test_request_accounting.SlotLegAccounting.test_prompt_n_is_what_was_read) ... ok +test_solo_leg_adds_to_the_slot_leg (test_request_accounting.SlotLegAccounting.test_solo_leg_adds_to_the_slot_leg) ... ok +test_two_hand_backs_sum (test_request_accounting.SlotLegAccounting.test_two_hand_backs_sum) ... ok +test_a_good_request_still_answers (test_request_hardening.Hardening.test_a_good_request_still_answers) ... [strata] done: 8 tokens in 0 s (195083.9 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_a_wrong_type_is_a_400_that_names_the_field (test_request_hardening.Hardening.test_a_wrong_type_is_a_400_that_names_the_field) ... ok +test_default_limit_is_generous (test_request_hardening.Hardening.test_default_limit_is_generous) ... ok +test_in_type_values_and_nulls_still_answer (test_request_hardening.Hardening.test_in_type_values_and_nulls_still_answer) ... [strata] done: 8 tokens in 0 s (296941.9 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 8 tokens in 0 s (258111.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 8 tokens in 0 s (217885.9 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_llama_cpp_names_for_the_repetition_penalty (test_request_hardening.Hardening.test_llama_cpp_names_for_the_repetition_penalty) +#1819: repeat_penalty / repeat_last_n are repetition_penalty / penalty_last_n; the OpenAI-style name wins. ... [strata] done: 8 tokens in 0 s (307838.8 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_negative_or_bad_content_length_is_400 (test_request_hardening.Hardening.test_negative_or_bad_content_length_is_400) ... ok +test_not_json_is_400 (test_request_hardening.Hardening.test_not_json_is_400) ... ok +test_numeric_strings_are_read_and_used (test_request_hardening.Hardening.test_numeric_strings_are_read_and_used) ... [strata] done: 8 tokens in 0 s (322638.8 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_oversized_content_length_is_413_unread (test_request_hardening.Hardening.test_oversized_content_length_is_413_unread) ... ok +test_wrong_shapes_are_400 (test_request_hardening.Hardening.test_wrong_shapes_are_400) ... ok +test_a_compaction_prompt_starts_as_its_conversations_last_prompt (test_responses.CodexCompaction.test_a_compaction_prompt_starts_as_its_conversations_last_prompt) ... ok +test_a_compaction_that_sends_tools_keeps_them (test_responses.CodexCompaction.test_a_compaction_that_sends_tools_keeps_them) ... ok +test_a_custom_compact_prompt_is_a_compaction_when_codex_says_so (test_responses.CodexCompaction.test_a_custom_compact_prompt_is_a_compaction_when_codex_says_so) ... ok +test_a_thread_title_turn_does_not_replace_the_kept_tools (test_responses.CodexCompaction.test_a_thread_title_turn_does_not_replace_the_kept_tools) ... ok +test_another_conversation_never_gets_these_tools (test_responses.CodexCompaction.test_another_conversation_never_gets_these_tools) ... ok +test_malformed_metadata_is_no_error_and_no_reuse (test_responses.CodexCompaction.test_malformed_metadata_is_no_error_and_no_reuse) ... ok +test_requests_that_are_not_a_codex_compaction_are_rendered_as_sent (test_responses.CodexCompaction.test_requests_that_are_not_a_codex_compaction_are_rendered_as_sent) ... ok +test_the_compact_prompt_in_a_message_is_not_a_compaction (test_responses.CodexCompaction.test_the_compact_prompt_in_a_message_is_not_a_compaction) ... ok +test_without_its_conversation_nothing_is_kept_or_used (test_responses.CodexCompaction.test_without_its_conversation_nothing_is_kept_or_used) ... ok +test_a_compaction_too_long_with_the_kept_tools_is_rendered_as_sent (test_responses.CodexCompactionOverHttp.test_a_compaction_too_long_with_the_kept_tools_is_rendered_as_sent) ... ok +test_a_request_that_fails_leaves_the_kept_tools_as_they_were (test_responses.CodexCompactionOverHttp.test_a_request_that_fails_leaves_the_kept_tools_as_they_were) ... ok +test_the_kept_tools_shape_only_the_prompt (test_responses.CodexCompactionOverHttp.test_the_kept_tools_shape_only_the_prompt) ... ok +test_a_text_answer_that_breaks_the_schema_fails_as_without_tools (test_responses.JsonFormatWithTools.test_a_text_answer_that_breaks_the_schema_fails_as_without_tools) ... skipped 'the schema check needs jsonschema' +test_a_tool_call_passes_through (test_responses.JsonFormatWithTools.test_a_tool_call_passes_through) ... ok +test_the_final_text_answer_is_the_schema_object (test_responses.JsonFormatWithTools.test_the_final_text_answer_is_the_schema_object) ... ok +test_server_close_ends_the_sampler (test_responses.NoLeakedSampler.test_server_close_ends_the_sampler) ... ok +test_a_foreign_page_cannot_name_a_file_as_an_image (test_responses.OverHttp.test_a_foreign_page_cannot_name_a_file_as_an_image) +#553: the chat routes refuse a file image from another origin; so does /v1/responses. ... ok +test_api_key_and_foreign_pages (test_responses.OverHttp.test_api_key_and_foreign_pages) ... ok +test_effort_none_does_not_think (test_responses.OverHttp.test_effort_none_does_not_think) ... ok +test_errors_use_the_responses_format (test_responses.OverHttp.test_errors_use_the_responses_format) ... [strata] 400 invalid request: Expecting property name enclosed in double quotes: line 1 column 2 (char 1) +ok +test_json_schema_text_format (test_responses.OverHttp.test_json_schema_text_format) ... ok +test_max_output_tokens_ends_incomplete (test_responses.OverHttp.test_max_output_tokens_ends_incomplete) ... ok +test_monitor_and_status (test_responses.OverHttp.test_monitor_and_status) ... ok +test_non_streaming_text_with_reasoning (test_responses.OverHttp.test_non_streaming_text_with_reasoning) ... ok +test_stream_event_order (test_responses.OverHttp.test_stream_event_order) ... ok +test_too_long_for_the_context (test_responses.OverHttp.test_too_long_for_the_context) ... ok +test_a_string_input_is_one_user_message (test_responses.Parsing.test_a_string_input_is_one_user_message) ... ok +test_additional_tools_input_items_are_accepted (test_responses.Parsing.test_additional_tools_input_items_are_accepted) +#782: Codex sends an `additional_tools` input item; it used to be refused with a 400. ... ok +test_bad_arguments_from_a_cut_off_call_do_not_refuse_the_conversation (test_responses.Parsing.test_bad_arguments_from_a_cut_off_call_do_not_refuse_the_conversation) ... ok +test_codex_history_becomes_one_assistant_turn_per_answer (test_responses.Parsing.test_codex_history_becomes_one_assistant_turn_per_answer) ... ok +test_effort_and_text_format (test_responses.Parsing.test_effort_and_text_format) ... ok +test_leading_developer_messages_join_the_system_message_later_ones_become_user (test_responses.Parsing.test_leading_developer_messages_join_the_system_message_later_ones_become_user) ... ok +test_reasoning_comes_back_from_our_encrypted_content_and_foreign_tokens_are_ignored (test_responses.Parsing.test_reasoning_comes_back_from_our_encrypted_content_and_foreign_tokens_are_ignored) ... ok +test_tools_namespaces_custom_and_hosted (test_responses.Parsing.test_tools_namespaces_custom_and_hosted) ... ok +test_unsupported_items_are_named (test_responses.Parsing.test_unsupported_items_are_named) ... ok +test_a_thread_title_is_answered_without_the_engine (test_responses.ThreadTitle.test_a_thread_title_is_answered_without_the_engine) ... ok +test_off_by_default_the_engine_answers (test_responses.ThreadTitle.test_off_by_default_the_engine_answers) ... ok +test_function_call_round_trip (test_responses.ToolRoundTrip.test_function_call_round_trip) ... ok +test_namespace_and_custom_calls_come_back_under_their_own_names (test_responses.ToolRoundTrip.test_namespace_and_custom_calls_come_back_under_their_own_names) ... ok +test_namespace_call_written_with_double_underscore (test_responses.ToolRoundTrip.test_namespace_call_written_with_double_underscore) ... ok +test_a_restart_that_fails_wakes_the_waiters (test_restart_waiters.RestartWaiters.test_a_restart_that_fails_wakes_the_waiters) ... [strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +ok +test_an_engine_err_answers_at_once_not_batched (test_restart_waiters.RestartWaiters.test_an_engine_err_answers_at_once_not_batched) ... [strata] the engine reported an error: refused: images are not enabled +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] 400 invalid request: refused: images are not enabled +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] answering: 9 of max 64 tokens, 8.9 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.9 tok/s) (stop, cancel=False) +ok +test_an_engine_err_answers_at_once_solo_loop (test_restart_waiters.RestartWaiters.test_an_engine_err_answers_at_once_solo_loop) ... [strata] the engine reported an error: refused: images are not enabled +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] 400 invalid request: refused: images are not enabled +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] answering: 9 of max 64 tokens, 8.6 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.4 tok/s) (stop, cancel=False) +ok +test_an_old_request_meets_a_restart_in_progress (test_restart_waiters.RestartWaiters.test_an_old_request_meets_a_restart_in_progress) +A request that learns its engine died while restart() has already taken the process (self.proc is None): ... ok +test_killed_while_reading_alone (test_restart_waiters.RestartWaiters.test_killed_while_reading_alone) ... [strata] the engine stopped unexpectedly (exit code -9). The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 10 tokens in 1 s (8.1 tok/s) (stop, cancel=False) +ok +test_killed_while_reading_with_others_arriving (test_restart_waiters.RestartWaiters.test_killed_while_reading_with_others_arriving) ... [strata] the engine stopped unexpectedly (exit code -9). The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] the engine stopped unexpectedly (exit code -9). The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 1 tokens in 1 s (2.7 tok/s) (error, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 10 tokens in 1 s (8.8 tok/s) (stop, cancel=False) +ok +test_solo_engine_is_unchanged (test_restart_waiters.RestartWaiters.test_solo_engine_is_unchanged) ... [strata] parallel requests: 4 asked, the engine runs one at a time (its log says why) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] parallel requests: 4 asked, the engine runs one at a time (its log says why) +[strata] the engine is running again +[strata] done: 10 tokens in 1 s (9.7 tok/s) (stop, cancel=False) +ok +test_waiters_get_a_clean_503_when_the_engine_stays_down_2_slots (test_restart_waiters.RestartWaiters.test_waiters_get_a_clean_503_when_the_engine_stays_down_2_slots) ... [strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 10 tokens in 1 s (10.9 tok/s) (stop, cancel=False) +ok +test_waiters_get_a_clean_503_when_the_engine_stays_down_4_slots (test_restart_waiters.RestartWaiters.test_waiters_get_a_clean_503_when_the_engine_stays_down_4_slots) ... [strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] the engine stopped while this request waited; it was not sent. The usual cause is running out of RAM: Linux then ends the biggest program (check: sudo dmesg | grep -i -E 'killed process|out of memory'); Windows slows down instead. Close other programs or use a smaller model (Q2_0 / IQ2_XS). The next request starts the engine again. +[strata] done: 0 tokens in 1 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 10 tokens in 1 s (10.6 tok/s) (stop, cancel=False) +ok +test_waiters_go_on_with_the_new_engine_2_slots (test_restart_waiters.RestartWaiters.test_waiters_go_on_with_the_new_engine_2_slots) ... [strata] answering: 2 of max 64 tokens, 4.3 tok/s, 1 s +[strata] answering: 2 of max 64 tokens, 9.2 tok/s, 1 s +[strata] done: 10 tokens in 2 s (7.7 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 2 s (6.5 tok/s) (stop, cancel=False) +[strata] reading the prompt: 32 of 48 tokens, 2 s so far +[strata] answering: 8 of max 64 tokens, 7.7 tok/s, 3 s +[strata] done: 10 tokens in 3 s (8.1 tok/s) (stop, cancel=False) +ok +test_waiters_go_on_with_the_new_engine_4_slots (test_restart_waiters.RestartWaiters.test_waiters_go_on_with_the_new_engine_4_slots) ... [strata] answering: 2 of max 64 tokens, 8.4 tok/s, 1 s +[strata] answering: 2 of max 64 tokens, 3.0 tok/s, 1 s +[strata] answering: 2 of max 64 tokens, 4.5 tok/s, 1 s +[strata] done: 10 tokens in 2 s (11.5 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 2 s (7.7 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 2 s (9.3 tok/s) (stop, cancel=False) +ok +test_null_is_the_default_again (test_runconfig.Apply.test_null_is_the_default_again) ... ok +test_only_the_named_keys_change (test_runconfig.Apply.test_only_the_named_keys_change) ... ok +test_refused (test_runconfig.Apply.test_refused) ... ok +test_view (test_runconfig.Apply.test_view) ... ok +test_only_from_strata_s_own_page (test_runconfig.Http.test_only_from_strata_s_own_page) ... ok +test_read_and_write_from_the_page (test_runconfig.Http.test_read_and_write_from_the_page) ... ok +test_the_api_key (test_runconfig.Http.test_the_api_key) ... ok +test_the_page_has_the_view (test_runconfig.Http.test_the_page_has_the_view) ... ok +test_without_a_run_config (test_runconfig.Http.test_without_a_run_config) ... ok +test_accepted (test_security.HostCheck.test_accepted) ... ok +test_any_host (test_security.HostCheck.test_any_host) ... ok +test_refused (test_security.HostCheck.test_refused) ... ok +test_allowed_hosts_of (test_security.HostNames.test_allowed_hosts_of) ... ok +test_host_name (test_security.HostNames.test_host_name) ... ok +test_names_from_the_bind_address (test_security.HostNames.test_names_from_the_bind_address) ... ok +test_accepted (test_security.OriginCheck.test_accepted) ... ok +test_refused (test_security.OriginCheck.test_refused) ... ok +test_a_trusted_origin_s_host_is_allowed (test_security.OverHttp.test_a_trusted_origin_s_host_is_allowed) ... ok +test_allowed_hosts_and_wildcard (test_security.OverHttp.test_allowed_hosts_and_wildcard) ... ok +test_configured_origins_pass (test_security.OverHttp.test_configured_origins_pass) ... ok +test_control_body_arrives_before_the_operation_and_reply (test_security.OverHttp.test_control_body_arrives_before_the_operation_and_reply) ... ok +test_cross_site_page_is_refused (test_security.OverHttp.test_cross_site_page_is_refused) ... ok +test_curl_and_sdks_without_origin (test_security.OverHttp.test_curl_and_sdks_without_origin) ... ok +test_local_names_and_ips_pass (test_security.OverHttp.test_local_names_and_ips_pass) ... ok +test_no_host_header_passes (test_security.OverHttp.test_no_host_header_passes) ... ok +test_own_and_local_pages_pass (test_security.OverHttp.test_own_and_local_pages_pass) ... ok +test_rebinding_name_is_refused_everywhere (test_security.OverHttp.test_rebinding_name_is_refused_everywhere) ... ok +test_refusal_names_both_ways_out (test_security.OverHttp.test_refusal_names_both_ways_out) ... ok +test_slots_need_json_from_the_own_page (test_security.OverHttp.test_slots_need_json_from_the_own_page) ... ok +test_unload_and_load_need_json_from_the_own_page (test_security.OverHttp.test_unload_and_load_need_json_from_the_own_page) ... ok +test_with_a_key_the_host_check_is_off (test_security.OverHttp.test_with_a_key_the_host_check_is_off) ... ok +test_with_a_key_the_key_decides (test_security.OverHttp.test_with_a_key_the_key_decides) ... ok +test_free_vram_on_hip (test_server.AmdTelemetry.test_free_vram_on_hip) ... ok +test_pcie_link_is_the_narrowest_slowest_hop (test_server.AmdTelemetry.test_pcie_link_is_the_narrowest_slowest_hop) +#1733: the link is the slowest/narrowest hop between the root port and the card; the bottleneck is only ... ok +test_readings (test_server.AmdTelemetry.test_readings) ... ok +test_a_body_a_handler_read_is_not_waited_for (test_server.AnswerBeforeTheBody.test_a_body_a_handler_read_is_not_waited_for) +#594: /config and /load read their body themselves: the close must not wait for the drain's 5 s. ... [strata] model unloaded; the next request loads it again +ok +test_a_body_that_comes_in_drops_does_not_hold_the_connection (test_server.AnswerBeforeTheBody.test_a_body_that_comes_in_drops_does_not_hold_the_connection) +A client that announces a body and sends it a byte at a time is let go after DRAIN_SECONDS, with its ... ok +test_a_host_the_server_does_not_answer_to (test_server.AnswerBeforeTheBody.test_a_host_the_server_does_not_answer_to) ... [strata] refused a request for Host 'rebind.example.com' from 127.0.0.1: not a name this server answers to (add it to "allowed_hosts" in the config or STRATA_ALLOWED_HOSTS, or set an API key) +ok +test_a_method_with_no_handler (test_server.AnswerBeforeTheBody.test_a_method_with_no_handler) ... ok +test_a_wrong_key (test_server.AnswerBeforeTheBody.test_a_wrong_key) ... ok +test_a_wrong_key_and_a_body_of_megabytes (test_server.AnswerBeforeTheBody.test_a_wrong_key_and_a_body_of_megabytes) +An agent client's conversation, or one screenshot, is several MiB, and a rotated key is when the 401 matters. ... ok +test_load_and_unload (test_server.AnswerBeforeTheBody.test_load_and_unload) ... [strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +ok +test_not_the_apps_own_page (test_server.AnswerBeforeTheBody.test_not_the_apps_own_page) ... ok +test_a_blank_key_is_an_error_not_no_key (test_server.ApiKeyForms.test_a_blank_key_is_an_error_not_no_key) ... ok +test_a_key_outside_ascii_matches_as_utf8_and_as_latin1 (test_server.ApiKeyForms.test_a_key_outside_ascii_matches_as_utf8_and_as_latin1) ... ok +test_a_utf8_key_over_http (test_server.ApiKeyForms.test_a_utf8_key_over_http) ... ok +test_several_keys_1344 (test_server.ApiKeyForms.test_several_keys_1344) ... ok +test_the_key_loses_what_a_header_cannot_carry (test_server.ApiKeyForms.test_the_key_loses_what_a_header_cannot_carry) ... ok +test_a_budget_that_fits_is_quiet (test_server.BudgetOverRam.test_a_budget_that_fits_is_quiet) ... ok +test_no_budget_argument_is_quiet (test_server.BudgetOverRam.test_no_budget_argument_is_quiet) ... ok +test_over_free_ram_warns (test_server.BudgetOverRam.test_over_free_ram_warns) ... ok +test_over_total_less_headroom_warns (test_server.BudgetOverRam.test_over_total_less_headroom_warns) ... ok +test_history_and_totals (test_server.CancelledRead.test_history_and_totals) ... [strata] done: 10 tokens in 0 s (120526.0 tok/s) (length, cancel=False) +[strata] done: 10 tokens in 0 s (182361.0 tok/s) (length, cancel=False) +[strata] done: 10 tokens in 0 s (263792.7 tok/s) (length, cancel=False) +ok +test_parse_done_chunk_fields (test_server.CancelledRead.test_parse_done_chunk_fields) +The prompt's first chunk and its streamed expert blobs come as key=value fields after the positional ones, ... ok +test_parse_done_read_field (test_server.CancelledRead.test_parse_done_read_field) ... ok +test_prompt_tokens_seen (test_server.CancelledRead.test_prompt_tokens_seen) ... ok +test_request_stats_chunk_fields (test_server.CancelledRead.test_request_stats_chunk_fields) +Across a request's continuations the streamed blobs add up; the chunk is the request's own read's. ... ok +test_a_plain_version_is_kept (test_server.ClaudeCodeBillingStamp.test_a_plain_version_is_kept) ... ok +test_bounds (test_server.ClaudeCodeBillingStamp.test_bounds) ... ok +test_cc_version_tail_is_pinned (test_server.ClaudeCodeBillingStamp.test_cc_version_tail_is_pinned) ... ok +test_cch_is_pinned (test_server.ClaudeCodeBillingStamp.test_cch_is_pinned) ... ok +test_only_at_the_start_of_the_system_prompt (test_server.ClaudeCodeBillingStamp.test_only_at_the_start_of_the_system_prompt) ... ok +test_string_system_is_pinned (test_server.ClaudeCodeBillingStamp.test_string_system_is_pinned) ... ok +test_two_turns_share_the_whole_system_prompt (test_server.ClaudeCodeBillingStamp.test_two_turns_share_the_whole_system_prompt) ... ok +test_non_streamed (test_server.ClientHangUp.test_non_streamed) ... [strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +ok +test_streamed (test_server.ClientHangUp.test_streamed) ... [strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +ok +test_a_tool_schema_that_is_not_an_object_is_a_400 (test_server.ClientShapes.test_a_tool_schema_that_is_not_an_object_is_a_400) ... [strata] 400 invalid request: tools[0] (f): "parameters" must be an object (the JSON schema of its parameters), not str +[strata] 400 invalid request: tools[0] (f): "parameters" must be an object (the JSON schema of its parameters), not list +[strata] 400 invalid request: tools[0] (f): "parameters" must be an object (the JSON schema of its parameters), not int +[strata] 400 invalid request: tools[0] (f): "input_schema" must be an object (the JSON schema of its parameters), not str +[strata] 400 invalid request: tools[0] (f): "input_schema" must be an object (the JSON schema of its parameters), not list +[strata] 400 invalid request: tools[0] (f): "input_schema" must be an object (the JSON schema of its parameters), not bool +[strata] done: 8 tokens in 0 s (236298.8 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (364722.1 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (332222.1 tok/s) (length, cancel=False) +ok +test_anthropic_mid_conversation_system (test_server.ClientShapes.test_anthropic_mid_conversation_system) ... [strata] done: 12 tokens in 0 s (342392.2 tok/s) (stop, cancel=False) +ok +test_leading_system_unchanged (test_server.ClientShapes.test_leading_system_unchanged) ... ok +test_malformed_tools_are_a_400 (test_server.ClientShapes.test_malformed_tools_are_a_400) ... [strata] 400 invalid request: tools must be a list of tool objects, each {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools must be a list of tool objects, each {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools must be a list of tool objects, each {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"type": "function", "function": {"name": ..., "parameters": {...}}} +[strata] 400 invalid request: tools must be a list of tool objects, each {"name": ..., "input_schema": {...}} +[strata] 400 invalid request: tools must be a list of tool objects, each {"name": ..., "input_schema": {...}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"name": ..., "input_schema": {...}} +[strata] 400 invalid request: tools must be a list of tool objects, each {"name": ..., "input_schema": {...}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"name": ..., "input_schema": {...}} +[strata] 400 invalid request: tools[0] has no name: each tool must be {"name": ..., "input_schema": {...}} +[strata] done: 8 tokens in 0 s (419430.4 tok/s) (length, cancel=False) +ok +test_messages_as_a_json_string_over_http (test_server.ClientShapes.test_messages_as_a_json_string_over_http) ... [strata] done: 12 tokens in 0 s (430185.0 tok/s) (stop, cancel=False) +[strata] 400 invalid request: messages must be a list of objects +[strata] done: 12 tokens in 0 s (466033.8 tok/s) (stop, cancel=False) +[strata] 400 invalid request: messages must be a list of objects +ok +test_messages_sent_as_a_json_string (test_server.ClientShapes.test_messages_sent_as_a_json_string) ... ok +test_no_user_turn_is_a_400 (test_server.ClientShapes.test_no_user_turn_is_a_400) ... [strata] 400 invalid request: No user query found in messages. +[strata] done: 12 tokens in 0 s (430185.0 tok/s) (stop, cancel=False) +ok +test_openai_late_developer_and_system (test_server.ClientShapes.test_openai_late_developer_and_system) ... [strata] done: 12 tokens in 0 s (387166.5 tok/s) (stop, cancel=False) +ok +test_query_string (test_server.ClientShapes.test_query_string) ... [strata] done: 12 tokens in 0 s (445412.8 tok/s) (stop, cancel=False) +[strata] done: 12 tokens in 0 s (461758.2 tok/s) (stop, cancel=False) +[strata] done: 12 tokens in 0 s (483958.2 tok/s) (stop, cancel=False) +ok +test_tool_call_malformed_arguments_fallback (test_server.ClientShapes.test_tool_call_malformed_arguments_fallback) ... ok +test_vision_temp_image_removed_when_the_pipe_fails (test_server.ClientShapes.test_vision_temp_image_removed_when_the_pipe_fails) ... ok +test_well_formed_tools_still_work (test_server.ClientShapes.test_well_formed_tools_still_work) ... [strata] done: 8 tokens in 0 s (419430.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (399457.5 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (441505.7 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (466033.8 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (447392.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (466033.8 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (466033.8 tok/s) (length, cancel=False) +ok +test_a_blank_key_is_refused (test_server.ConfigApiKey.test_a_blank_key_is_refused) ... ok +test_a_key_and_no_key (test_server.ConfigApiKey.test_a_key_and_no_key) ... ok +test_an_empty_key_is_a_warning (test_server.ConfigApiKey.test_an_empty_key_is_a_warning) ... ok +test_config_level_fills_in_what_a_request_leaves_out (test_server.ConfigReasoningEffort.test_config_level_fills_in_what_a_request_leaves_out) ... ok +test_default_changes_nothing (test_server.ConfigReasoningEffort.test_default_changes_nothing) ... ok +test_shared_settings_win (test_server.ConfigReasoningEffort.test_shared_settings_win) ... ok +test_the_flag_in_args_becomes_the_key (test_server.ConfigReasoningEffort.test_the_flag_in_args_becomes_the_key) ... ok +test_a_non_streamed_answer_says_the_connection_closes (test_server.ConnectionClose.test_a_non_streamed_answer_says_the_connection_closes) ... [strata] done: 13 tokens in 0 s (271273.4 tok/s) (stop, cancel=False) +ok +test_a_streamed_answer_says_the_connection_closes (test_server.ConnectionClose.test_a_streamed_answer_says_the_connection_closes) ... [strata] done: 13 tokens in 0 s (224386.6 tok/s) (stop, cancel=False) +ok +test_desktop_detection (test_server.DesktopVramNote.test_desktop_detection) ... ok +test_when_it_applies (test_server.DesktopVramNote.test_when_it_applies) ... ok +test_when_it_does_not (test_server.DesktopVramNote.test_when_it_does_not) ... ok +test_drafts_in_history_and_totals (test_server.DraftCounts.test_drafts_in_history_and_totals) ... [strata] done: 10 tokens in 0 s (179243.8 tok/s) (length, cancel=False) +[strata] done: 10 tokens in 0 s (249661.0 tok/s) (length, cancel=False) +[strata] done: 10 tokens in 0 s (257319.3 tok/s) (length, cancel=False), expert cache 90.0% hit +ok +test_an_older_engine_gets_the_advice_in_words (test_server.DraftHeadHint.test_an_older_engine_gets_the_advice_in_words) ... ok +test_other_failures_and_earlier_starts_add_nothing (test_server.DraftHeadHint.test_other_failures_and_earlier_starts_add_nothing) ... ok +test_the_engines_hint_is_relayed (test_server.DraftHeadHint.test_the_engines_hint_is_relayed) ... ok +test_changing_the_effort_keeps_the_conversation (test_server.EffortAtTheEnd.test_changing_the_effort_keeps_the_conversation) ... ok +test_the_config (test_server.EffortAtTheEnd.test_the_config) ... ok +test_the_default_prompt_is_unchanged (test_server.EffortAtTheEnd.test_the_default_prompt_is_unchanged) ... ok +test_the_next_turn_reuses_the_checkpoint (test_server.EffortAtTheEnd.test_the_next_turn_reuses_the_checkpoint) ... ok +test_the_trailing_turn (test_server.EffortAtTheEnd.test_the_trailing_turn) ... ok +test_engine_err_mid_stream (test_server.EngineDeath.test_engine_err_mid_stream) +The engine's ERR line after the stream started reaches the client as an error event (it used to be a ... [strata] the engine reported an error: verify: layer 31 never rang (an illegal memory access was encountered) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] the engine reported an error: verify: layer 31 never rang (an illegal memory access was encountered) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +ok +test_error_then_restart (test_server.EngineDeath.test_error_then_restart) ... [strata] the engine stopped unexpectedly (exit code -9). The next request starts the engine again. +[strata] done: 5 tokens in 0 s (102300.1 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code None); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 50 tokens in 0 s (247597.6 tok/s) (length, cancel=False) +ok +test_a_budget_that_fits_is_unchanged (test_server.FitMaxTokens.test_a_budget_that_fits_is_unchanged) ... [strata] done: 50 tokens in 0 s (243007.2 tok/s) (length, cancel=False) +ok +test_no_room_is_still_a_400 (test_server.FitMaxTokens.test_no_room_is_still_a_400) ... [strata] done: 1 tokens in 0 s (123361.9 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (4088 tokens) leaves no room to answer in the context (4096); requests are never truncated +ok +test_overshoot_is_clamped_to_the_room (test_server.FitMaxTokens.test_overshoot_is_clamped_to_the_room) ... [strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19399.9 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19253.8 tok/s) (stop, cancel=False) +ok +test_a_call_cut_by_max_tokens_is_not_a_tool_call (test_server.ForcedToolChoice.test_a_call_cut_by_max_tokens_is_not_a_tool_call) ... [strata] done: 5 tokens in 0 s (185588.7 tok/s) (length, cancel=False) +ok +test_a_thinking_that_ends_right_after_a_quote_opens_the_call (test_server.ForcedToolChoice.test_a_thinking_that_ends_right_after_a_quote_opens_the_call) ... [strata] done: 120 tokens in 0 s (131003.8 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (109203.0 tok/s) (stop, cancel=False) +ok +test_anthropic_any_tool_none (test_server.ForcedToolChoice.test_anthropic_any_tool_none) ... [strata] done: 272 tokens in 0 s (94441.3 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (76061.8 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (105390.4 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (76752.6 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (141647.6 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (84018.5 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (143946.5 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (81846.0 tok/s) (stop, cancel=False) +[strata] tool_choice {"type": "tool", "name": "nope"} is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (139833.1 tok/s) (stop, cancel=False) +[strata] tool_choice {"type": "tool", "name": "nope"} is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (84872.8 tok/s) (stop, cancel=False) +ok +test_auto_absent_and_none_leave_the_tools_to_the_model (test_server.ForcedToolChoice.test_auto_absent_and_none_leave_the_tools_to_the_model) ... [strata] done: 203 tokens in 0 s (122739.5 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (128481.0 tok/s) (stop, cancel=False) +[strata] done: 203 tokens in 0 s (132335.1 tok/s) (stop, cancel=False) +ok +test_required_and_named_after_the_thinking (test_server.ForcedToolChoice.test_required_and_named_after_the_thinking) ... [strata] done: 272 tokens in 0 s (97592.0 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (72173.8 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (108714.6 tok/s) (stop, cancel=False) +[strata] done: 272 tokens in 0 s (79813.3 tok/s) (stop, cancel=False) +ok +test_the_budget_wrap_up_opens_the_call (test_server.ForcedToolChoice.test_the_budget_wrap_up_opens_the_call) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 181 tokens in 0 s (119403.7 tok/s) (stop, cancel=False) +ok +test_the_flat_shape_names_a_function (test_server.ForcedToolChoice.test_the_flat_shape_names_a_function) ... [strata] done: 272 tokens in 0 s (97633.8 tok/s) (stop, cancel=False) +ok +test_values_it_cannot_honour_act_as_auto (test_server.ForcedToolChoice.test_values_it_cannot_honour_act_as_auto) +Not a 400: an odd tool_choice is logged and the model decides. ... [strata] tool_choice {"type": "function", "function": {"name": "nope"}} is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (123112.2 tok/s) (stop, cancel=False) +[strata] tool_choice {"type": "function"} is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (138671.6 tok/s) (stop, cancel=False) +[strata] tool_choice "required" is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (141271.6 tok/s) (stop, cancel=False) +[strata] tool_choice {"type": "banana"} is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (141365.4 tok/s) (stop, cancel=False) +[strata] tool_choice "sometimes" is not supported here (or names no tool of the request): the model decides, as with "auto" +[strata] done: 203 tokens in 0 s (141765.5 tok/s) (stop, cancel=False) +ok +test_without_thinking_the_prompt_ends_with_the_opening (test_server.ForcedToolChoice.test_without_thinking_the_prompt_ends_with_the_opening) ... [strata] done: 64 tokens in 0 s (217180.8 tok/s) (stop, cancel=False) +[strata] done: 64 tokens in 0 s (198841.1 tok/s) (stop, cancel=False) +[strata] done: 56 tokens in 0 s (282989.2 tok/s) (stop, cancel=False) +[strata] done: 56 tokens in 0 s (195083.9 tok/s) (stop, cancel=False) +ok +test_card_order (test_server.GpuChoice.test_card_order) +#1352: with an auto layer split the faster card (SMs x clock) goes last; ties, a manual split, a ... [strata] layer split: card order 2,0 (the faster card last; "gpu_order": "as_given" keeps 0,2, #1352) +[strata] layer split: card order 2,0 (the faster card last; "gpu_order": "as_given" keeps 0,2, #1352) +[strata] layer split: card order 2,0 (the faster card last; "gpu_order": "as_given" keeps 0,2, #1352) +[strata] layer split: card order 0,2,1 (the faster card last; "gpu_order": "as_given" keeps 0,1,2, #1352) +[strata] layer split: keeping card order 0,1 - the faster card has less VRAM than the one the split ends on (12288 vs 22528 MiB, #1576) +[strata] layer split: card order 1,0 (the faster card last; "gpu_order": "as_given" keeps 0,1, #1352) +[strata] layer split: card order 1,0 (the faster card last; "gpu_order": "as_given" keeps 0,1, #1352) +[strata] layer split: card order 1,0 (the faster card last; "gpu_order": "as_given" keeps 0,1, #1352) +[strata] layer split: keeping card order 0,1,2 - the faster card has less VRAM than the one the split ends on (12288 vs 24576 MiB, #1576) +ok +test_env (test_server.GpuChoice.test_env) ... ok +test_hip_ordinal (test_server.GpuChoice.test_hip_ordinal) +#325: on Windows the HIP ordinal setup resolved wins over the config's "gpu" (an iGPU takes HIP's 0). ... ok +test_kfd_scores (test_server.GpuChoice.test_kfd_scores) +AMD on Linux: the GPU nodes of the KFD topology, SIMDs x clock, node order = the config's index. ... /Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:2795: ResourceWarning: unclosed file <_io.TextIOWrapper name='/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpj8a9ucay/0/properties' mode='r' encoding='UTF-8'> + try: +ResourceWarning: Enable tracemalloc to get the object allocation traceback +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:2795: ResourceWarning: unclosed file <_io.TextIOWrapper name='/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpj8a9ucay/1/properties' mode='r' encoding='UTF-8'> + try: +ResourceWarning: Enable tracemalloc to get the object allocation traceback +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:2795: ResourceWarning: unclosed file <_io.TextIOWrapper name='/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpj8a9ucay/2/properties' mode='r' encoding='UTF-8'> + try: +ResourceWarning: Enable tracemalloc to get the object allocation traceback +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:2795: ResourceWarning: unclosed file <_io.TextIOWrapper name='/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpj8a9ucay/3/properties' mode='r' encoding='UTF-8'> + try: +ResourceWarning: Enable tracemalloc to get the object allocation traceback +ok +test_vision_device (test_server.GpuChoice.test_vision_device) ... ok +test_a_held_back_call_still_pings (test_server.HeldBackArgKeepalive.test_a_held_back_call_still_pings) ... [strata] done: 54 tokens in 0 s (205528.5 tok/s) (stop, cancel=False) +ok +test_an_ordinary_answer_does_not_ping (test_server.HeldBackArgKeepalive.test_an_ordinary_answer_does_not_ping) ... [strata] done: 3 tokens in 0 s (174762.7 tok/s) (stop, cancel=False) +ok +test_healthy_untouched (test_server.HipEnvGuard.test_healthy_untouched) ... ok +test_only_for_hip (test_server.HipEnvGuard.test_only_for_hip) ... strata serve: AMD: HIP_PATH=Z:\definitely\not\here does not exist: not passed to the engine +ok +test_stale_hip_path_dropped (test_server.HipEnvGuard.test_stale_hip_path_dropped) ... strata serve: AMD: HIP_PATH=/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpzsqao3tx/TheRock/build does not exist: not passed to the engine +strata serve: AMD: HIP_DEVICE_LIB_PATH=/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpzsqao3tx/TheRock/build does not exist: not passed to the engine +ok +test_unwritable_temp_replaced (test_server.HipEnvGuard.test_unwritable_temp_replaced) ... strata serve: AMD: TEMP=/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmp9xuos_hs/no/such/dir is not writable: using /private/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmp9xuos_hs/tmp +ok +test_anthropic_tool_result_image (test_server.ImageMarkers.test_anthropic_tool_result_image) +Claude Code's Read of a picture returns it inside a tool_result: it reaches the encoder, in place. ... [strata] a picture a tool returned was left out of the prompt: this server has no image encoder +[strata] a picture a tool returned was left out of the prompt: this server has no image encoder +[strata] a picture a tool returned was left out of the prompt: this image format needs Pillow +[strata] a picture a tool returned was left out of the prompt: the image URL could not be read +ok +test_literal_marker_with_an_image (test_server.ImageMarkers.test_literal_marker_with_an_image) ... ok +test_pictures_are_keyed_by_their_bytes_for_the_engine (test_server.ImageMarkers.test_pictures_are_keyed_by_their_bytes_for_the_engine) +#2228: each picture's key (its normalized bytes' hash: the encoder's file name) goes to the engine as ... ok +test_the_whole_marker_in_text_before_images (test_server.ImageMarkers.test_the_whole_marker_in_text_before_images) ... ok +test_a_page_of_another_site_cannot_name_a_file (test_server.ImageSources.test_a_page_of_another_site_cannot_name_a_file) ... [strata] refused an API request from the web page 'https://evil.example' (no API key; add its host to "allowed_hosts" or its origin to "cors_origins" in the config) +[strata] 400 invalid request: a web page of another origin cannot have a file on this computer read as an image: send it as a data: URL (or add the page's origin to the config's trusted_origins) +[strata] 400 invalid request: a web page of another origin cannot have a file on this computer read as an image: send it as a data: URL (or add the page's origin to the config's trusted_origins) +[strata] 400 invalid request: a web page of another origin cannot have a file on this computer read as an image: send it as a data: URL (or add the page's origin to the config's trusted_origins) +[strata] 400 invalid request: a web page of another origin cannot have a file on this computer read as an image: send it as a data: URL (or add the page's origin to the config's trusted_origins) +[strata] done: 8 tokens in 0 s (258111.0 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (399457.5 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (364722.1 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (447392.4 tok/s) (length, cancel=False) +ok +test_a_url_is_downloaded_outside_the_fifo (test_server.ImageSources.test_a_url_is_downloaded_outside_the_fifo) ... ok +test_a_url_is_read_up_to_the_cap (test_server.ImageSources.test_a_url_is_read_up_to_the_cap) ... /opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +ok +test_network_paths_are_refused_before_any_look (test_server.ImageSources.test_network_paths_are_refused_before_any_look) ... ok +test_unreadable_sources_are_a_value_error (test_server.ImageSources.test_unreadable_sources_are_a_value_error) +From #582: a bad data: URL or an unreadable file is a 400, not a dropped connection. ... ok +test_a_lazy_engine_keeps_the_prompts_its_arguments_ask_for (test_server.IncrementalPrompts.test_a_lazy_engine_keeps_the_prompts_its_arguments_ask_for) ... ok +test_a_tokenizer_without_resume_points_encodes_in_full (test_server.IncrementalPrompts.test_a_tokenizer_without_resume_points_encodes_in_full) ... ok +test_counting_tokens_leaves_the_kept_prompts_alone (test_server.IncrementalPrompts.test_counting_tokens_leaves_the_kept_prompts_alone) ... [strata] done: 8 tokens in 0 s (181375.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_same_ids_as_a_full_encode_on_varied_conversations (test_server.IncrementalPrompts.test_same_ids_as_a_full_encode_on_varied_conversations) ... ok +test_the_kept_prompts_follow_the_engines_conversation_slots (test_server.IncrementalPrompts.test_the_kept_prompts_follow_the_engines_conversation_slots) ... ok +test_the_slot_argument_wins_over_the_info_line_and_a_bad_one_falls_back (test_server.IncrementalPrompts.test_the_slot_argument_wins_over_the_info_line_and_a_bad_one_falls_back) ... ok +test_turns (test_server.IncrementalPrompts.test_turns) ... [strata] done: 32 tokens in 0 s (345032.7 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (377016.1 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (366715.1 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (394758.0 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (400649.9 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (364722.1 tok/s) (stop, cancel=False) +ok +test_without_a_slot_setting_four_prompts_are_kept (test_server.IncrementalPrompts.test_without_a_slot_setting_four_prompts_are_kept) ... ok +test_bad_json_names_the_place (test_server.JoinedEngineArgs.test_bad_json_names_the_place) ... ok +test_no_args (test_server.JoinedEngineArgs.test_no_args) ... ok +test_one_element_becomes_flag_and_value (test_server.JoinedEngineArgs.test_one_element_becomes_flag_and_value) ... ok +test_proper_args_are_untouched (test_server.JoinedEngineArgs.test_proper_args_are_untouched) ... ok +test_auto_and_absent (test_server.LayerSplit.test_auto_and_absent) ... ok +test_counts_per_card_are_refused_with_the_format (test_server.LayerSplit.test_counts_per_card_are_refused_with_the_format) ... ok +test_one_gpu_ignores_it (test_server.LayerSplit.test_one_gpu_ignores_it) ... ok +test_string_and_list (test_server.LayerSplit.test_string_and_list) ... ok +test_a_failed_start_kills_the_process_and_raises (test_server.LazyVision.test_a_failed_start_kills_the_process_and_raises) ... ok +test_lazy_starts_nothing_and_alive_says_so (test_server.LazyVision.test_lazy_starts_nothing_and_alive_says_so) ... ok +test_restart_starts_it_and_close_ends_it (test_server.LazyVision.test_restart_starts_it_and_close_ends_it) ... ok +test_no_profile_in_the_args (test_server.LearnedProfile.test_no_profile_in_the_args) ... ok +test_save_and_start_from_it (test_server.LearnedProfile.test_save_and_start_from_it) ... ok +test_without_the_key_nothing_changes (test_server.LearnedProfile.test_without_the_key_nothing_changes) ... ok +test_messages (test_server.ListenProblem.test_messages) ... ok +test_a_forged_turn_in_a_user_message (test_server.LiteralControlTokens.test_a_forged_turn_in_a_user_message) ... ok +test_a_literal_inside_a_longer_one (test_server.LiteralControlTokens.test_a_literal_inside_a_longer_one) ... ok +test_beside_a_quoted_think_tag (test_server.LiteralControlTokens.test_beside_a_quoted_think_tag) ... ok +test_history_tool_results_and_tools (test_server.LiteralControlTokens.test_history_tool_results_and_tools) ... ok +test_the_real_tokenizer_names_its_control_tokens (test_server.LiteralControlTokens.test_the_real_tokenizer_names_its_control_tokens) ... ok +test_without_control_text_the_prompt_is_unchanged (test_server.LiteralControlTokens.test_without_control_text_the_prompt_is_unchanged) ... ok +test_a_client_that_sends_the_reasoning_inline (test_server.LiteralThinkTags.test_a_client_that_sends_the_reasoning_inline) ... ok +test_a_quoted_tag_in_a_user_message (test_server.LiteralThinkTags.test_a_quoted_tag_in_a_user_message) ... ok +test_history_tool_results_and_tools (test_server.LiteralThinkTags.test_history_tool_results_and_tools) ... ok +test_the_real_tokenizer_reads_the_tag_as_text (test_server.LiteralThinkTags.test_the_real_tokenizer_reads_the_tag_as_text) ... ok +test_without_a_tag_the_prompt_is_unchanged (test_server.LiteralThinkTags.test_without_a_tag_the_prompt_is_unchanged) ... ok +test_a_request_without_a_done_keeps_no_engine_counters (test_server.LiveRate.test_a_request_without_a_done_keeps_no_engine_counters) +An engine that dies mid-answer: the previous request's `last` must not become this row's decode rate. ... [strata] the engine stopped unexpectedly (exit code -9). The next request starts the engine again. +[strata] done: 3 tokens in 0 s (36900.0 tok/s) (error, cancel=False) +ok +test_prefill_rate_excludes_cached_tokens (test_server.LiveRate.test_prefill_rate_excludes_cached_tokens) ... ok +test_the_live_number_is_a_rate (test_server.LiveRate.test_the_live_number_is_a_rate) ... [strata] thinking: 9 of max 30 tokens, 9.5 tok/s, 1 s +[strata] thinking: 18 of max 30 tokens, 8.8 tok/s, 2 s +[strata] thinking: 26 of max 30 tokens, 8.4 tok/s, 3 s +[strata] done: 30 tokens in 4 s (8.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_silent_mid_answer (test_server.LostStep.test_silent_mid_answer) ... [strata] the engine said nothing for 1 s during the request; the server ended the engine. The engine and the server lost step (issue #481; a very slow PC can raise "engine_silence_s" in the config, 0 = wait forever). If it happens again, please add the end of the engine log to github.com/Niko1221/Strata/issues/481. The next request starts the engine again. +[strata] done: 1 tokens in 1 s (1.5 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 2 tokens in 0 s (51150.0 tok/s) (length, cancel=False) +ok +test_silent_mid_stream (test_server.LostStep.test_silent_mid_stream) ... [strata] the engine said nothing for 1 s during the request; the server ended the engine. The engine and the server lost step (issue #481; a very slow PC can raise "engine_silence_s" in the config, 0 = wait forever). If it happens again, please add the end of the engine log to github.com/Niko1221/Strata/issues/481. The next request starts the engine again. +[strata] done: 1 tokens in 1 s (1.6 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 2 tokens in 0 s (16644.1 tok/s) (length, cancel=False) +ok +test_stop_never_acknowledged (test_server.LostStep.test_stop_never_acknowledged) ... [strata] the engine did not finish the request after it was stopped (STOP) within 0 s; the server ended the engine. The engine and the server lost step (issue #481; a very slow PC can raise "engine_silence_s" in the config, 0 = wait forever). If it happens again, please add the end of the engine log to github.com/Niko1221/Strata/issues/481. The next request starts the engine again. +[strata] done: 2 tokens in 1 s (3.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -9); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 2 tokens in 0 s (24036.1 tok/s) (length, cancel=False) +ok +test_bad_process (test_server.MacTelemetry.test_bad_process) ... ok +test_no_registry (test_server.MacTelemetry.test_no_registry) ... ok +test_note_reaches_the_monitor (test_server.MacTelemetry.test_note_reaches_the_monitor) ... ok +test_readings (test_server.MacTelemetry.test_readings) ... ok +test_anthropic_thinks_only_when_asked (test_server.MaxTokens.test_anthropic_thinks_only_when_asked) ... ok +test_count_tokens_is_the_prompt_messages_reads (test_server.MaxTokens.test_count_tokens_is_the_prompt_messages_reads) ... [strata] done: 8 tokens in 0 s (236298.8 tok/s) (length, cancel=False) +ok +test_debug_log_shows_the_resolved_budget (test_server.MaxTokens.test_debug_log_shows_the_resolved_budget) ... ok +test_explicit_budget_is_honoured (test_server.MaxTokens.test_explicit_budget_is_honoured) ... [strata] done: 50 tokens in 0 s (318716.1 tok/s) (length, cancel=False) +[strata] done: 50 tokens in 0 s (330781.1 tok/s) (length, cancel=False) +[strata] done: 50 tokens in 0 s (327168.8 tok/s) (length, cancel=False) +[strata] done: 50 tokens in 0 s (320175.9 tok/s) (length, cancel=False) +[strata] thinking limited to 988 of max_tokens 1500 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 1500 tokens in 0 s (25557.2 tok/s) (length, cancel=False) +[strata] thinking limited to 988 of max_tokens 1500 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 1500 tokens in 0 s (25535.0 tok/s) (length, cancel=False) +ok +test_explicit_budget_over_the_context_is_rejected (test_server.MaxTokens.test_explicit_budget_over_the_context_is_rejected) ... [strata] 400 invalid request: request exceeds the context window: prompt (246 tokens) + max tokens (4096) exceeds the context (4096); requests are never truncated. Send a smaller max_tokens (at most 3842 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (246 tokens) + max tokens (4096) exceeds the context (4096); requests are never truncated. Send a smaller max_tokens (at most 3842 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +ok +test_request_line_parses_the_engine_summary (test_server.MaxTokens.test_request_line_parses_the_engine_summary) ... ok +test_unset_budget_is_the_rest_of_the_context (test_server.MaxTokens.test_unset_budget_is_the_rest_of_the_context) ... [strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19341.8 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19369.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19191.3 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19286.8 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19224.1 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19348.6 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19124.9 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19274.6 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (18730.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 2011 tokens in 0 s (19083.3 tok/s) (stop, cancel=False) +ok +test_unset_budget_with_a_near_full_prompt (test_server.MaxTokens.test_unset_budget_with_a_near_full_prompt) ... [strata] done: 1 tokens in 0 s (127100.1 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 1 tokens in 0 s (174762.7 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 5 tokens in 0 s (332881.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (4088 tokens) leaves no room to answer in the context (4096); requests are never truncated +[strata] done: 1 tokens in 0 s (144631.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 5 tokens in 0 s (388361.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (4088 tokens) leaves no room to answer in the context (4096); requests are never truncated +ok +test_models_lists_them (test_server.ModelAliases.test_models_lists_them) ... ok +test_requests_are_answered_under_the_alias (test_server.ModelAliases.test_requests_are_answered_under_the_alias) ... [strata] done: 8 tokens in 0 s (248551.3 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (345922.0 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (381300.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (399457.5 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (381300.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (447392.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (447392.4 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (414252.2 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (404270.3 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (414252.2 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (424739.6 tok/s) (length, cancel=False) +[strata] done: 8 tokens in 0 s (364722.1 tok/s) (length, cancel=False) +ok +test_set_aliases (test_server.ModelAliases.test_set_aliases) ... ok +test_without_aliases_nothing_changes (test_server.ModelAliases.test_without_aliases_nothing_changes) ... ok +test_engine_done_line_raises (test_server.NonFiniteLogits.test_engine_done_line_raises) +StrataEngine.generate: `DONE ... nonfinite` raises EngineNonFinite after the tokens before it. ... ok +test_mid_reply_ends_without_retry (test_server.NonFiniteLogits.test_mid_reply_ends_without_retry) ... ok +test_nothing_emitted_is_retried_once (test_server.NonFiniteLogits.test_nothing_emitted_is_retried_once) ... ok +test_twice_is_an_error (test_server.NonFiniteLogits.test_twice_is_an_error) ... ok +test_twice_over_http (test_server.NonFiniteLogits.test_twice_over_http) ... ok +test_history (test_server.PcieShare.test_history) ... ok +test_a_vision_section_starts_the_engine_with_images_1322 (test_server.PeerDevice.test_a_vision_section_starts_the_engine_with_images_1322) ... [strata] the config has a "vision" section but no --vision in its args: added --vision and --vram-reserve-mib 700 (run setup again with --vision to write them) +[strata] the config has a "vision" section but no --vision in its args: added --vision (run setup again with --vision to write them) +[strata] the config has a "vision" section but no --vision in its args: added --vision (run setup again with --vision to write them) +ok +test_both_cards_visible (test_server.PeerDevice.test_both_cards_visible) ... ok +test_no_split_with_a_peer (test_server.PeerDevice.test_no_split_with_a_peer) ... ok +test_split_added_for_several_gpus (test_server.PeerDevice.test_split_added_for_several_gpus) ... ok +test_cancel_before_the_encode_runs_none (test_server.PrepareCancelled.test_cancel_before_the_encode_runs_none) ... ok +test_cancel_between_pictures_stops_the_next_encode (test_server.PrepareCancelled.test_cancel_between_pictures_stops_the_next_encode) ... ok +test_cancel_while_waiting_for_the_turn (test_server.PrepareCancelled.test_cancel_while_waiting_for_the_turn) ... ok +test_without_cancel_nothing_changes (test_server.PrepareCancelled.test_without_cancel_nothing_changes) ... ok +test_throughput_gauges_cover_the_window_between_scrapes (test_server.PrometheusMetrics.test_throughput_gauges_cover_the_window_between_scrapes) +#2341: llama.cpp reads these two from a bucket it resets on every scrape, so a scrape that saw no ... [strata] done: 10 tokens in 0 s (185588.7 tok/s) (length, cancel=False) +ok +test_totals (test_server.PrometheusMetrics.test_totals) ... [strata] done: 10 tokens in 0 s (129854.6 tok/s) (length, cancel=False) +[strata] done: 10 tokens in 0 s (243854.9 tok/s) (length, cancel=False) +ok +test_a_client_that_asks_gets_one_progress_per_chunk (test_server.PromptProgress.test_a_client_that_asks_gets_one_progress_per_chunk) ... [strata] done: 40 tokens in 0 s (99039.1 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_a_client_that_does_not_ask_gets_nothing (test_server.PromptProgress.test_a_client_that_does_not_ask_gets_nothing) ... [strata] done: 40 tokens in 0 s (97598.7 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_nothing_before_the_first_chunk (test_server.PromptProgress.test_nothing_before_the_first_chunk) ... ok +test_the_line_that_says_the_prompt_is_read_is_dropped (test_server.PromptProgress.test_the_line_that_says_the_prompt_is_read_is_dropped) ... ok +test_nothing_under_batch (test_server.PromptProgressBatch.test_nothing_under_batch) ... ok +test_a_reply_that_answered_is_not_touched (test_server.ReasoningCloseRetry.test_a_reply_that_answered_is_not_touched) ... [strata] done: 18 tokens in 0 s (276547.5 tok/s) (stop, cancel=False) +ok +test_a_reply_that_closed_its_thinking_right_after_a_quote_is_not_touched (test_server.ReasoningCloseRetry.test_a_reply_that_closed_its_thinking_right_after_a_quote_is_not_touched) ... [strata] done: 28 tokens in 0 s (277637.1 tok/s) (stop, cancel=False) +ok +test_off_by_default_the_reply_stays_empty (test_server.ReasoningCloseRetry.test_off_by_default_the_reply_stays_empty) ... [strata] done: 36 tokens in 0 s (276547.5 tok/s) (stop, cancel=False) +ok +test_on_it_closes_the_thinking_once_and_answers (test_server.ReasoningCloseRetry.test_on_it_closes_the_thinking_once_and_answers) ... [strata] the reply ended inside its thinking with no answer: closing the thinking once and continuing (reasoning_close_retry) +[strata] done: 64 tokens in 0 s (196224.7 tok/s) (stop, cancel=False) +ok +test_a_long_run_is_ended (test_server.RepeatStop.test_a_long_run_is_ended) ... ok +test_short_runs_and_off (test_server.RepeatStop.test_short_runs_and_off) ... ok +test_context_zero_is_503 (test_server.RestartWindow.test_context_zero_is_503) ... ok +test_failed_restart_keeps_the_known_context (test_server.RestartWindow.test_failed_restart_keeps_the_known_context) +After a restart that failed, max_context is 0 but the engine is not starting: a request is checked against ... [strata] done: 13 tokens in 0 s (236043.1 tok/s) (stop, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/unittest/case.py:615: ResourceWarning: unclosed + result = method() +ResourceWarning: Enable tracemalloc to get the object allocation traceback +ok +test_not_alive_until_ready (test_server.RestartWindow.test_not_alive_until_ready) ... ok +test_restart_retries_a_start_that_fails (test_server.RestartWindow.test_restart_retries_a_start_that_fails) +A dead engine's VRAM is freed only when its process is gone, so a new engine started at once can exit ... [strata] the engine did not start (try 1 of 3); again in 0 s +ok +test_checkpoint_key (test_server.SamplingKeys.test_checkpoint_key) ... ok +test_penalty_window (test_server.SamplingKeys.test_penalty_window) ... ok +test_top_k (test_server.SamplingKeys.test_top_k) ... ok +test_tune_aux_cpus (test_server.SamplingKeys.test_tune_aux_cpus) ... ok +test_tune_keys (test_server.SamplingKeys.test_tune_keys) ... ok +test_tune_prefill_pipe (test_server.SamplingKeys.test_tune_prefill_pipe) ... ok +test_chars_end_inside_a_message (test_server.SharedPrefix.test_chars_end_inside_a_message) ... ok +test_messages_resolve_to_the_boundary_token (test_server.SharedPrefix.test_messages_resolve_to_the_boundary_token) ... ok +test_pin_key (test_server.SharedPrefix.test_pin_key) ... ok +test_tokens_and_unusable_prefixes (test_server.SharedPrefix.test_tokens_and_unusable_prefixes) ... ok +test_without_the_field_nothing_changes (test_server.SharedPrefix.test_without_the_field_nothing_changes) ... ok +test_a_body_without_the_defaults_wrapper_is_rejected (test_server.SharedSettings.test_a_body_without_the_defaults_wrapper_is_rejected) +A body with no "defaults" key is a mistake, not the documented "clear them" ({"defaults": null}). ... [strata] other apps now use the Chat settings: temperature=0.3 +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 16 tokens in 0 s (271695.8 tok/s) (stop, cancel=False) +ok +test_a_trusted_origin_is_strata_s_own_page (test_server.SharedSettings.test_a_trusted_origin_is_strata_s_own_page) ... [strata] other apps use their own settings again +ok +test_off_again (test_server.SharedSettings.test_off_again) ... [strata] other apps now use the Chat settings: temperature=0.3 +[strata] other apps use their own settings again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 16 tokens in 0 s (390167.8 tok/s) (stop, cancel=False) +ok +test_only_strata_s_own_page_may_set_them (test_server.SharedSettings.test_only_strata_s_own_page_may_set_them) ... [strata] other apps now use the Chat settings: temperature=1.0 +ok +test_other_apps_get_the_chat_settings (test_server.SharedSettings.test_other_apps_get_the_chat_settings) ... [strata] other apps now use the Chat settings: temperature=0.3, top_p=0.9, top_k=10, seed=7, max_tokens=77, reasoning_effort=low, experimental_speed_projection=False +[strata] done: 16 tokens in 0 s (419430.4 tok/s) (stop, cancel=False) +[strata] done: 5 tokens in 0 s (279620.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_proxy_headers_do_not_make_a_page_strata_s_own (test_server.SharedSettings.test_proxy_headers_do_not_make_a_page_strata_s_own) ... ok +test_they_need_the_key_when_one_is_set (test_server.SharedSettings.test_they_need_the_key_when_one_is_set) ... [strata] other apps now use the Chat settings: temperature=1.0 +ok +test_before_load_failure_can_be_ignored_by_config (test_server.SharingTheGpu.test_before_load_failure_can_be_ignored_by_config) ... [strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(1) +[strata] the before_load command failed (it exited with code 1); loading anyway +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (288497.1 tok/s) (stop, cancel=False) +ok +test_before_load_runs_first (test_server.SharingTheGpu.test_before_load_runs_first) ... [strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c open('/var/folders/32/rbk61lf917s0ps54b9c4lmfc0000gn/T/tmpbb50noit/ran', 'w').close() +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (275383.6 tok/s) (stop, cancel=False) +ok +test_before_load_strict_refuses_even_with_free_vram (test_server.SharingTheGpu.test_before_load_strict_refuses_even_with_free_vram) ... [strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(7) +[strata] the before_load command failed (it exited with code 7); the model stays unloaded and this request is answered 503 - try again (before_load_strict) +ok +test_before_load_timeout_is_configurable_and_a_timeout_refuses (test_server.SharingTheGpu.test_before_load_timeout_is_configurable_and_a_timeout_refuses) ... [strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import time; time.sleep(30) +[strata] the before_load command failed (it did not finish in 0.5 s (before_load_timeout_s)); the model stays unloaded and this request is answered 503 - try again (before_load_strict) +ok +test_failing_before_load_with_free_vram_still_loads (test_server.SharingTheGpu.test_failing_before_load_with_free_vram_still_loads) +#1997: a peer that is not running makes `curl` exit 7; the card has room, so the model loads. ... [strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(7) +[strata] the before_load command failed (it exited with code 7); checking the free VRAM +[strata] the before_load command failed (it exited with code 7), but 20000 MiB of VRAM is free: loading +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (236043.1 tok/s) (stop, cancel=False) +ok +test_failing_before_load_with_full_vram_refuses_with_503_and_retry_after (test_server.SharingTheGpu.test_failing_before_load_with_full_vram_refuses_with_503_and_retry_after) ... [strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(3) +[strata] the before_load command failed (it exited with code 3); checking the free VRAM +[strata] the before_load command failed (it exited with code 3) and only 2000 MiB of VRAM is free (the model needs 8000): not loading; 503 +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(3) +[strata] the before_load command failed (it exited with code 3); checking the free VRAM +[strata] the before_load command failed (it exited with code 3), but 20000 MiB of VRAM is free: loading +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (227191.5 tok/s) (stop, cancel=False) +ok +test_failing_before_load_without_a_vram_limit_loads (test_server.SharingTheGpu.test_failing_before_load_without_a_vram_limit_loads) ... [strata] model unloaded; the next request loads it again +[strata] before loading: /private/tmp/r39-serve-venv/bin/python -c import sys; sys.exit(7) +[strata] the before_load command failed (it exited with code 7); loading anyway (set min_free_vram_mib to have the free VRAM checked, or before_load_strict to refuse) +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (264689.1 tok/s) (stop, cancel=False) +ok +test_idle_unload (test_server.SharingTheGpu.test_idle_unload) ... [strata] the model unloads after 1 s without requests +[strata] model unloaded after 1 s idle; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (232025.3 tok/s) (stop, cancel=False) +ok +test_load_endpoint (test_server.SharingTheGpu.test_load_endpoint) ... [strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] model unloaded after 1 s idle; the next request loads it again +ok +test_min_free_vram_refuses_to_load (test_server.SharingTheGpu.test_min_free_vram_refuses_to_load) ... [strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (212163.2 tok/s) (stop, cancel=False) +ok +test_min_free_vram_unreadable_loads (test_server.SharingTheGpu.test_min_free_vram_unreadable_loads) ... [strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (245612.4 tok/s) (stop, cancel=False) +ok +test_off_by_default (test_server.SharingTheGpu.test_off_by_default) ... ok +test_unload_refused_while_a_request_runs (test_server.SharingTheGpu.test_unload_refused_while_a_request_runs) ... ok +test_unload_then_the_next_request_loads (test_server.SharingTheGpu.test_unload_then_the_next_request_loads) ... [strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (264689.1 tok/s) (stop, cancel=False) +ok +test_unload_wait_gives_up_after_the_given_seconds (test_server.SharingTheGpu.test_unload_wait_gives_up_after_the_given_seconds) ... ok +test_unload_wait_waits_for_the_running_turn (test_server.SharingTheGpu.test_unload_wait_waits_for_the_running_turn) +#1997: POST /unload answers busy at once (409); with ?wait=1 it waits for the turn and then unloads. ... [strata] model unloaded; the next request loads it again +ok +test_vision_encoder_unloads_and_starts_first (test_server.SharingTheGpu.test_vision_encoder_unloads_and_starts_first) ... [strata] model unloaded; the next request loads it again +[strata] starting the vision encoder ... +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (264689.1 tok/s) (stop, cancel=False) +ok +test_a_busy_gpu_is_not_a_frozen_engine (test_server.SilentEngine.test_a_busy_gpu_is_not_a_frozen_engine) ... [strata] the engine has said nothing for 1 s but is still working; it is not ended (the limit is engine_silence_s). If it is stuck, "engine_silence_s" can be set lower. +ok +test_a_long_first_chunk_is_allowed (test_server.SilentEngine.test_a_long_first_chunk_is_allowed) ... ok +test_a_stop_never_acknowledged (test_server.SilentEngine.test_a_stop_never_acknowledged) ... ok +test_a_working_silent_engine_is_not_ended_by_the_stall_watchdog (test_server.SilentEngine.test_a_working_silent_engine_is_not_ended_by_the_stall_watchdog) ... [strata] the engine has said nothing for 1 s but is still working; it is not ended (the limit is engine_silence_s). If it is stuck, "engine_silence_s" can be set lower. +ok +test_config (test_server.SilentEngine.test_config) ... ok +test_death_note_finds_the_watchdog_line_under_a_long_trace (test_server.SilentEngine.test_death_note_finds_the_watchdog_line_under_a_long_trace) ... ok +test_engine_frozen_samples (test_server.SilentEngine.test_engine_frozen_samples) ... ok +test_frozen_engine_is_ended_by_the_stall_watchdog (test_server.SilentEngine.test_frozen_engine_is_ended_by_the_stall_watchdog) ... ok +test_no_reading_means_no_kill_and_zero_is_off (test_server.SilentEngine.test_no_reading_means_no_kill_and_zero_is_off) ... ok +test_prompt_chunks_set_the_wait (test_server.SilentEngine.test_prompt_chunks_set_the_wait) ... ok +test_silence_mid_answer_ends_the_engine (test_server.SilentEngine.test_silence_mid_answer_ends_the_engine) ... ok +test_zero_waits_as_before (test_server.SilentEngine.test_zero_waits_as_before) ... ok +test_a_crash_is_not_a_transient_start (test_server.StartCrashReported.test_a_crash_is_not_a_transient_start) ... ok +test_a_silent_exit_is_retried_without_claiming_the_gpu_runtime (test_server.StartCrashReported.test_a_silent_exit_is_retried_without_claiming_the_gpu_runtime) ... ok +test_sigill_is_reported_not_retried (test_server.StartCrashReported.test_sigill_is_reported_not_retried) ... ok +test_status_words (test_server.StartCrashReported.test_status_words) ... ok +test_the_runtimes_own_failure_is_named (test_server.StartCrashReported.test_the_runtimes_own_failure_is_named) ... ok +test_the_last_lines_of_this_start (test_server.StartFailureLog.test_the_last_lines_of_this_start) ... ok +test_the_start_error_has_them (test_server.StartFailureLog.test_the_start_error_has_them) ... ok +test_a_mapped_start (test_server.StartNarrator.test_a_mapped_start) ... ok +test_an_arena_start (test_server.StartNarrator.test_an_arena_start) ... ok +test_the_words_follow_the_flags (test_server.StartNarrator.test_the_words_follow_the_flags) ... ok +test_abort_then_the_next_request (test_server.StatusHandover.test_abort_then_the_next_request) ... [strata] done: 15 tokens in 0 s (3549.3 tok/s) (disconnect, cancel=False) +[strata] done: 20 tokens in 0 s (122640.5 tok/s) (length, cancel=False) +ok +test_status (test_server.StatusNeedsTheKey.test_status) ... ok +test_a_prefix_that_is_not_a_stop_is_sent (test_server.StopStrings.test_a_prefix_that_is_not_a_stop_is_sent) +"ENDX" begins like the answer's "END" but is not in it: the held tail goes out and nothing is lost. ... [strata] done: 36 tokens in 0 s (228089.0 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (118520.4 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (279103.4 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (174762.7 tok/s) (stop, cancel=False) +ok +test_bad_values_are_a_400 (test_server.StopStrings.test_bad_values_are_a_400) ... [strata] 400 invalid request: stop takes at most 4 strings +[strata] 400 invalid request: stop must be a string or a list of strings +[strata] 400 invalid request: stop_sequences must be a list of strings +[strata] 400 invalid request: stop_sequences must be a list of strings +ok +test_cut_at_the_stop_string (test_server.StopStrings.test_cut_at_the_stop_string) ... [strata] done: 19 tokens in 0 s (339113.9 tok/s) (stop, cancel=False) +[strata] done: 19 tokens in 0 s (218333.6 tok/s) (stop, cancel=False) +[strata] done: 19 tokens in 0 s (351065.1 tok/s) (stop, cancel=False) +[strata] done: 19 tokens in 0 s (209164.8 tok/s) (stop, cancel=False) +ok +test_matcher_holds_back_only_a_possible_prefix (test_server.StopStrings.test_matcher_holds_back_only_a_possible_prefix) ... ok +test_no_stop_and_an_empty_list (test_server.StopStrings.test_no_stop_and_an_empty_list) ... [strata] done: 36 tokens in 0 s (349525.3 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (336291.6 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (310689.2 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (164482.5 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (163591.5 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (105076.5 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (379384.3 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (382265.7 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (343170.3 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (169087.3 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (146882.2 tok/s) (stop, cancel=False) +[strata] done: 36 tokens in 0 s (167400.2 tok/s) (stop, cancel=False) +ok +test_openai_takes_a_string (test_server.StopStrings.test_openai_takes_a_string) ... [strata] done: 29 tokens in 0 s (357749.5 tok/s) (stop, cancel=False) +ok +test_the_first_of_several (test_server.StopStrings.test_the_first_of_several) ... [strata] done: 24 tokens in 0 s (342392.2 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (210592.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (333322.2 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (203360.2 tok/s) (stop, cancel=False) +ok +test_sampler_runs_without_disk_counters (test_server.TelemetryNoDisks.test_sampler_runs_without_disk_counters) ... ok +test_probe_errors_are_logged_and_swallowed (test_server.TemplateCaps.test_probe_errors_are_logged_and_swallowed) ... ok +test_a_bad_value_is_a_400 (test_server.ThinkingBudget.test_a_bad_value_is_a_400) ... [strata] 400 invalid request: reasoning_budget_tokens='lots': expected a whole number of tokens (0: no budget) +[strata] 400 invalid request: reasoning_budget_tokens=2.5: expected a whole number of tokens (0: no budget) +[strata] 400 invalid request: reasoning_budget_tokens=True: expected a whole number of tokens (0: no budget) +[strata] 400 invalid request: reasoning_budget_tokens=[1]: expected a whole number of tokens (0: no budget) +ok +test_a_budget_above_max_tokens_is_capped_to_leave_room_for_the_answer (test_server.ThinkingBudget.test_a_budget_above_max_tokens_is_capped_to_leave_room_for_the_answer) +#984: max_tokens 300 and a budget of 5000 never fired: the reply ended while still thinking. The thinking ... ok +test_a_budget_the_thinking_stays_under (test_server.ThinkingBudget.test_a_budget_the_thinking_stays_under) ... [strata] thinking capped at 200 of max_tokens 400 to leave room for the answer (budget 10000) +[strata] done: 211 tokens in 0 s (123689.5 tok/s) (stop, cancel=False) +ok +test_a_budget_well_below_max_tokens_is_not_capped (test_server.ThinkingBudget.test_a_budget_well_below_max_tokens_is_not_capped) ... ok +test_a_cap_that_the_thinking_would_eat_leaves_room_to_answer (test_server.ThinkingBudget.test_a_cap_that_the_thinking_would_eat_leaves_room_to_answer) +0.1.45: max_tokens counts the thinking; a thinking that would fill it gets closed with a quarter left (the ... [strata] thinking limited to 512 of max_tokens 1024 so the answer has room ("answer_room": false in the config turns this off) +[strata] thinking budget reached (512 tokens): wrapping up the thinking +[strata] done: 604 tokens in 0 s (58165.9 tok/s) (stop, cancel=False) +ok +test_a_reply_cut_while_thinking_is_named_in_the_log (test_server.ThinkingBudget.test_a_reply_cut_while_thinking_is_named_in_the_log) +#530: max tokens reached inside the thinking gives an empty answer; the server log says what helps. ... ok +test_a_shared_budget_must_be_a_whole_number_of_tokens (test_server.ThinkingBudget.test_a_shared_budget_must_be_a_whole_number_of_tokens) ... ok +test_a_shared_budget_reaches_a_request_that_sets_none (test_server.ThinkingBudget.test_a_shared_budget_reaches_a_request_that_sets_none) +The Chat settings shared with other apps may carry a thinking budget too, like max_tokens and the effort. ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (188281.4 tok/s) (stop, cancel=False) +[strata] done: 211 tokens in 0 s (141328.4 tok/s) (stop, cancel=False) +ok +test_a_stop_string_is_found_in_the_answer_after_the_wrap_up (test_server.ThinkingBudget.test_a_stop_string_is_found_in_the_answer_after_the_wrap_up) +#454: the continuation after the thinking budget is cut at a stop string too. ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 108 tokens in 0 s (171715.3 tok/s) (stop, cancel=False) +ok +test_anthropic_stream (test_server.ThinkingBudget.test_anthropic_stream) ... [strata] thinking budget reached (30 tokens): wrapping up the thinking +[strata] done: 122 tokens in 0 s (95538.7 tok/s) (stop, cancel=False) +ok +test_anthropic_thinking_budget_tokens_is_a_budget (test_server.ThinkingBudget.test_anthropic_thinking_budget_tokens_is_a_budget) +0.1.45: thinking {type: enabled, budget_tokens: N} closes the thinking after N tokens (it was only a level). ... [strata] thinking budget reached (30 tokens): wrapping up the thinking +[strata] done: 122 tokens in 0 s (186006.9 tok/s) (stop, cancel=False) +ok +test_no_room_left_to_answer (test_server.ThinkingBudget.test_no_room_left_to_answer) ... [strata] done: 20 tokens in 0 s (220173.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_off_by_default (test_server.ThinkingBudget.test_off_by_default) ... [strata] done: 211 tokens in 0 s (134990.6 tok/s) (stop, cancel=False) +ok +test_the_answer_room_has_its_limits (test_server.ThinkingBudget.test_the_answer_room_has_its_limits) ... [strata] done: 300 tokens in 0 s (106725.3 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 1024 tokens in 0 s (36058.8 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 17 tokens in 0 s (377265.4 tok/s) (stop, cancel=False) +ok +test_the_budget_wraps_up_the_thinking (test_server.ThinkingBudget.test_the_budget_wraps_up_the_thinking) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (194760.4 tok/s) (stop, cancel=False) +ok +test_the_config_budget_reaches_the_anthropic_and_effort_requests (test_server.ThinkingBudget.test_the_config_budget_reaches_the_anthropic_and_effort_requests) +#984: a config budget applies to a request that sets no budget of its own, whatever its max_tokens and ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (101812.3 tok/s) (stop, cancel=False) +[strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (205945.7 tok/s) (stop, cancel=False) +ok +test_the_config_default_and_a_request_that_turns_it_off (test_server.ThinkingBudget.test_the_config_default_and_a_request_that_turns_it_off) ... [strata] thinking budget reached (20 tokens): wrapping up the thinking +[strata] done: 112 tokens in 0 s (170512.5 tok/s) (stop, cancel=False) +[strata] done: 211 tokens in 0 s (122609.9 tok/s) (stop, cancel=False) +ok +test_the_thinking_progress_line_shows_the_budget_not_max_tokens (test_server.ThinkingBudget.test_the_thinking_progress_line_shows_the_budget_not_max_tokens) +#984: 'thinking: x of max 6400' read as if the budget was ignored; while a budget applies the line shows it. ... ok +test_without_thinking_there_is_nothing_to_limit (test_server.ThinkingBudget.test_without_thinking_there_is_nothing_to_limit) ... [strata] done: 17 tokens in 0 s (265067.5 tok/s) (stop, cancel=False) +ok +test_draft_fields (test_server.TimingsDrafts.test_draft_fields) ... ok +test_both_frontends_render_it (test_server.ToolArguments.test_both_frontends_render_it) ... ok +test_helper (test_server.ToolArguments.test_helper) ... ok +test_bare_function_stays_text_where_it_is_not_a_call (test_server.ToolCallRecovery.test_bare_function_stays_text_where_it_is_not_a_call) ... ok +test_bare_function_without_the_wrapper (test_server.ToolCallRecovery.test_bare_function_without_the_wrapper) ... ok +test_broken_json_is_content_not_an_error (test_server.ToolCallRecovery.test_broken_json_is_content_not_an_error) ... ok +test_json_in_the_wrapper (test_server.ToolCallRecovery.test_json_in_the_wrapper) ... ok +test_json_never_closed_is_content (test_server.ToolCallRecovery.test_json_never_closed_is_content) ... ok +test_json_of_an_unknown_tool_is_content (test_server.ToolCallRecovery.test_json_of_an_unknown_tool_is_content) ... ok +test_parameter_as_the_opener (test_server.ToolCallRecovery.test_parameter_as_the_opener) ... ok +test_parameter_opener_that_is_not_a_tool_is_content (test_server.ToolCallRecovery.test_parameter_opener_that_is_not_a_tool_is_content) ... ok +test_two_calls_in_one_wrapper (test_server.ToolCallRecovery.test_two_calls_in_one_wrapper) ... ok +test_a_batch_is_one_call (test_server.ToolCallRecoveryOff.test_a_batch_is_one_call) ... ok +test_drifted_forms_stay_text (test_server.ToolCallRecoveryOff.test_drifted_forms_stay_text) ... ok +test_switch (test_server.ToolCallRecoverySwitch.test_switch) ... ok +test_named_tag_is_content (test_server.ToolCallTagInProse.test_named_tag_is_content) ... ok +test_tag_alone_at_the_end_is_content (test_server.ToolCallTagInProse.test_tag_alone_at_the_end_is_content) ... ok +test_values_keep_the_terminators (test_server.ToolCallTerminators.test_values_keep_the_terminators) ... ok +test_anthropic_shape_and_a_different_set (test_server.ToolOrder.test_anthropic_shape_and_a_different_set) ... ok +test_duplicate_names_and_one_tool_are_left_alone (test_server.ToolOrder.test_duplicate_names_and_one_tool_are_left_alone) ... ok +test_openai_shape_keeps_the_first_order (test_server.ToolOrder.test_openai_shape_keeps_the_first_order) ... ok +test_a_call_cut_at_the_token_limit (test_server.UnfinishedToolCall.test_a_call_cut_at_the_token_limit) +The same cut by max_tokens: "length" / "max_tokens", and the whole (non-streamed) answers leave the call ... [strata] done: 119 tokens in 0 s (191675.2 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (123057.7 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (208750.4 tok/s) (length, cancel=False) +[strata] done: 119 tokens in 0 s (138837.9 tok/s) (length, cancel=False) +ok +test_a_cut_call (test_server.UnfinishedToolCall.test_a_cut_call) ... [strata] done: 120 tokens in 0 s (193881.5 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (125797.7 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (208326.4 tok/s) (stop, cancel=False) +[strata] done: 120 tokens in 0 s (131861.8 tok/s) (stop, cancel=False) +ok +test_a_whole_call (test_server.UnfinishedToolCall.test_a_whole_call) ... [strata] done: 147 tokens in 0 s (167817.8 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (133773.6 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (171505.6 tok/s) (stop, cancel=False) +[strata] done: 147 tokens in 0 s (151081.3 tok/s) (stop, cancel=False) +ok +test_collect_keeps_calls_whose_arguments_parse (test_server.UnfinishedToolCall.test_collect_keeps_calls_whose_arguments_parse) ... ok +test_parser (test_server.UnfinishedToolCall.test_parser) ... ok +test_a_vision_answer_in_time_is_unchanged (test_server.UntimedReads.test_a_vision_answer_in_time_is_unchanged) ... ok +test_engine_that_never_says_ready_is_ended (test_server.UntimedReads.test_engine_that_never_says_ready_is_ended) ... ok +test_vision_encode_read_times_out (test_server.UntimedReads.test_vision_encode_read_times_out) ... ok +test_vision_ready_read_times_out (test_server.UntimedReads.test_vision_ready_read_times_out) ... ok +test_anthropic (test_server.UsageAndStatus.test_anthropic) ... [strata] done: 13 tokens in 0 s (170927.7 tok/s) (stop, cancel=False), expert cache 90.0% hit +ok +test_cors_for_the_configured_origins_on_the_api_only (test_server.UsageAndStatus.test_cors_for_the_configured_origins_on_the_api_only) ... ok +test_cors_is_off_by_default (test_server.UsageAndStatus.test_cors_is_off_by_default) ... ok +test_no_clock (test_server.UsageAndStatus.test_no_clock) +An engine without a clock (MockEngine): no timings, nothing cached. ... [strata] done: 5 tokens in 0 s (201649.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_openai (test_server.UsageAndStatus.test_openai) ... [strata] done: 13 tokens in 0 s (37474.9 tok/s) (stop, cancel=False), expert cache 90.0% hit +ok +test_openai_stream (test_server.UsageAndStatus.test_openai_stream) ... [strata] done: 13 tokens in 0 s (25092.5 tok/s) (stop, cancel=False), expert cache 90.0% hit +ok +test_origin_lists_are_checked (test_server.UsageAndStatus.test_origin_lists_are_checked) ... ok +test_sse_is_not_buffered_by_proxies (test_server.UsageAndStatus.test_sse_is_not_buffered_by_proxies) ... [strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (61961.3 tok/s) (stop, cancel=False), expert cache 90.0% hit +ok +test_v1_status (test_server.UsageAndStatus.test_v1_status) ... [strata] done: 13 tokens in 0 s (93526.5 tok/s) (stop, cancel=False), expert cache 90.0% hit +ok +test_min_tokens_is_passed_only_when_set (test_server.VisionArgs.test_min_tokens_is_passed_only_when_set) ... ok +test_one_request_with_more_than_64_images_keeps_all_its_files (test_server.VisionCacheEviction.test_one_request_with_more_than_64_images_keeps_all_its_files) ... ok +test_plain_encode_still_evicts_the_oldest (test_server.VisionCacheEviction.test_plain_encode_still_evicts_the_oldest) ... ok +test_shutdown_removes_the_work_dir_but_close_keeps_it (test_server.VisionShutdown.test_shutdown_removes_the_work_dir_but_close_keeps_it) ... ok +test_combined_file_is_written_whole_and_removable (test_server.VisionTempFiles.test_combined_file_is_written_whole_and_removable) ... ok +test_combined_file_stays_off_a_small_shm (test_server.VisionTempFiles.test_combined_file_stays_off_a_small_shm) ... ok +test_enc_line_names_files_relative_to_the_encoder_dir (test_server.VisionTempFiles.test_enc_line_names_files_relative_to_the_encoder_dir) ... ok +test_smart_app_control_is_said_in_words (test_server.VisionTempFiles.test_smart_app_control_is_said_in_words) ... ok +test_work_dir_has_no_space_when_temp_has_one (test_server.VisionTempFiles.test_work_dir_has_no_space_when_temp_has_one) ... ok +test_build_model_path_and_slot_status (test_server.WebApp.test_build_model_path_and_slot_status) ... ok +test_discovery_does_not_restart_a_dead_engine (test_server.WebApp.test_discovery_does_not_restart_a_dead_engine) ... /opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +ok +test_discovery_needs_the_api_key (test_server.WebApp.test_discovery_needs_the_api_key) ... /opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +ok +test_metrics (test_server.WebApp.test_metrics) ... [strata] done: 5 tokens in 0 s (216201.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_metrics_need_the_key_when_one_is_set (test_server.WebApp.test_metrics_need_the_key_when_one_is_set) ... ok +test_model_discovery_and_props (test_server.WebApp.test_model_discovery_and_props) ... /opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +ok +test_only_the_app_files_are_served (test_server.WebApp.test_only_the_app_files_are_served) ... ok +test_page_and_files (test_server.WebApp.test_page_and_files) ... ok +test_props_caps_are_ready_for_concurrent_requests (test_server.WebApp.test_props_caps_are_ready_for_concurrent_requests) ... ok +test_props_caps_follow_the_active_template (test_server.WebApp.test_props_caps_follow_the_active_template) ... ok +test_props_total_slots_follows_the_batch_slots (test_server.WebApp.test_props_total_slots_follows_the_batch_slots) ... ok +test_slot_reports_the_context_in_use (test_server.WebApp.test_slot_reports_the_context_in_use) +/slots in llama.cpp's names: a front-end's context meter divides n_prompt_tokens by n_ctx, so a server ... ok +test_slots_lists_every_batch_slot (test_server.WebApp.test_slots_lists_every_batch_slot) ... ok +test_the_endpoints_keep_the_context_while_the_engine_restarts (test_server.WebApp.test_the_endpoints_keep_the_context_while_the_engine_restarts) +#351: max_context is 0 until READY; the endpoints report the last known one. ... ok +test_default_refusal_is_immediate_and_byte_identical (test_session_save_retry.HTTP.test_default_refusal_is_immediate_and_byte_identical) ... ok +test_fifo_blocks_restore_while_save_waits_and_releases_afterwards (test_session_save_retry.HTTP.test_fifo_blocks_restore_while_save_waits_and_releases_afterwards) ... [strata] SAVE waiting for RAM after cache reclaim: need=1125 MiB, limit=30s; floor unchanged +ok +test_opt_in_recovers_native_refusal_and_keeps_response_and_process (test_session_save_retry.HTTP.test_opt_in_recovers_native_refusal_and_keeps_response_and_process) ... [strata] SAVE waiting for RAM after cache reclaim: need=1125 MiB, limit=30s; floor unchanged +ok +test_admission_races_are_bounded_and_keep_latest_error (test_session_save_retry.Retry.test_admission_races_are_bounded_and_keep_latest_error) ... ok +test_delayed_ram_reuses_same_process_path_and_floor (test_session_save_retry.Retry.test_delayed_ram_reuses_same_process_path_and_floor) ... ok +test_memory_probe_is_portable_and_unknown_fails_closed (test_session_save_retry.Retry.test_memory_probe_is_portable_and_unknown_fails_closed) ... ok +test_new_refusal_rechecks_its_larger_estimate (test_session_save_retry.Retry.test_new_refusal_rechecks_its_larger_estimate) ... ok +test_non_admission_error_during_retry_is_returned (test_session_save_retry.Retry.test_non_admission_error_during_retry_is_returned) ... ok +test_persistent_pressure_expires_without_another_save (test_session_save_retry.Retry.test_persistent_pressure_expires_without_another_save) ... ok +test_process_death_or_replacement_ends_wait (test_session_save_retry.Retry.test_process_death_or_replacement_ends_wait) ... ok +test_restore_other_categories_and_published_files_never_retry (test_session_save_retry.Retry.test_restore_other_categories_and_published_files_never_retry) ... ok +test_success_does_not_poll_or_wait (test_session_save_retry.Retry.test_success_does_not_poll_or_wait) ... ok +test_unknown_telemetry_or_estimate_never_retries (test_session_save_retry.Retry.test_unknown_telemetry_or_estimate_never_retries) ... ok +test_reuse_then_least_recently_used (test_slot_alloc.OneGroup.test_reuse_then_least_recently_used) ... ok +test_groups_of_several_slots_still_ignore_held_tokens (test_slot_alloc.OneSlotGroups.test_groups_of_several_slots_still_ignore_held_tokens) ... ok +test_the_slot_of_the_chat (test_slot_alloc.OneSlotGroups.test_the_slot_of_the_chat) ... ok +test_fills_a_leading_hole_first (test_slot_alloc.Pipelined.test_fills_a_leading_hole_first) ... ok +test_held_tokens_are_ignored (test_slot_alloc.Pipelined.test_held_tokens_are_ignored) ... ok +test_never_opens_a_hole_below_the_pick (test_slot_alloc.Pipelined.test_never_opens_a_hole_below_the_pick) ... ok +test_spreads_over_the_groups_first (test_slot_alloc.Pipelined.test_spreads_over_the_groups_first) ... ok +test_a_blocking_step_gets_its_announced_allowance (test_slots.Slots.test_a_blocking_step_gets_its_announced_allowance) +B: SWAIT gives one non-streaming step (a flush, the device transfer) its own bounded wait. ... ok +test_a_blocking_step_that_never_ends_is_ended (test_slots.Slots.test_a_blocking_step_that_never_ends_is_ended) +B: the allowance is a bound, not a heartbeat: a step that outlives it ends the engine. ... ok +test_a_foreign_page_or_a_form_never_reaches_the_engine (test_slots.Slots.test_a_foreign_page_or_a_form_never_reaches_the_engine) ... ok +test_a_silent_engine_is_ended_not_waited_for (test_slots.Slots.test_a_silent_engine_is_ended_not_waited_for) ... ok +test_a_small_save_reports_its_only_partial_block (test_slots.Slots.test_a_small_save_reports_its_only_partial_block) +B: a file below one 16 MiB block still reports (the start, then its last partial block). ... ok +test_disabled_without_a_slot_save_path (test_slots.Slots.test_disabled_without_a_slot_save_path) ... ok +test_disk_and_ram_refusals_have_their_own_codes (test_slots.Slots.test_disk_and_ram_refusals_have_their_own_codes) ... ok +test_engine_death_is_a_server_error (test_slots.Slots.test_engine_death_is_a_server_error) ... ok +test_engine_gets_an_absolute_path_on_one_line (test_slots.Slots.test_engine_gets_an_absolute_path_on_one_line) ... ok +test_failed_restore_transfer_is_a_server_error (test_slots.Slots.test_failed_restore_transfer_is_a_server_error) ... ok +test_filename_cannot_leave_the_directory (test_slots.Slots.test_filename_cannot_leave_the_directory) ... ok +test_log_lines_are_not_answers (test_slots.Slots.test_log_lines_are_not_answers) ... ok +test_malformed_allowances_end_the_engine (test_slots.Slots.test_malformed_allowances_end_the_engine) ... ok +test_malformed_or_mismatched_reply_is_a_server_error (test_slots.Slots.test_malformed_or_mismatched_reply_is_a_server_error) ... ok +test_missing_file_is_not_found (test_slots.Slots.test_missing_file_is_not_found) ... ok +test_only_slot_zero_and_known_actions (test_slots.Slots.test_only_slot_zero_and_known_actions) ... ok +test_ordinary_names_are_accepted (test_slots.Slots.test_ordinary_names_are_accepted) ... ok +test_out_of_protocol_lines_end_the_engine (test_slots.Slots.test_out_of_protocol_lines_end_the_engine) +R4: a line this exchange does not allow is not a heartbeat, and an answer must hold valid numbers. ... ok +test_progress_ends_the_allowance (test_slots.Slots.test_progress_ends_the_allowance) +B: after a SESSION line the ordinary silence limit applies again (the allowance was for one step). ... ok +test_progress_lines_keep_a_long_save_alive (test_slots.Slots.test_progress_lines_keep_a_long_save_alive) ... ok +test_published_failure_is_reported (test_slots.Slots.test_published_failure_is_reported) +R3: the folder flush failed after the rename - the new file replaced the old one; the answer says so. ... ok +test_refused_file_is_an_error_not_a_success (test_slots.Slots.test_refused_file_is_an_error_not_a_success) ... ok +test_refused_with_parallel_requests (test_slots.Slots.test_refused_with_parallel_requests) +#465: batch slots do not hold the FIFO, so a session file is refused before anything reaches the engine. ... ok +test_relative_save_dir_becomes_absolute (test_slots.Slots.test_relative_save_dir_becomes_absolute) ... ok +test_restore_answers_like_llama_server (test_slots.Slots.test_restore_answers_like_llama_server) ... ok +test_restore_of_a_link_reaches_the_engine_which_decides (test_slots.Slots.test_restore_of_a_link_reaches_the_engine_which_decides) +The server only checks the name exists (lstat); the engine opens it without following a link. ... ok +test_save_writes_into_the_slot_save_path (test_slots.Slots.test_save_writes_into_the_slot_save_path) ... ok +test_slot_work_counts_as_activity_and_shows_in_status (test_slots.Slots.test_slot_work_counts_as_activity_and_shows_in_status) ... ok +test_typed_refusals_have_their_own_codes (test_slots.Slots.test_typed_refusals_have_their_own_codes) +R5: the engine's category decides the status, never the words of the message. ... ok +test_waits_for_a_running_request (test_slots.Slots.test_waits_for_a_running_request) +A save never interleaves with a generation: it takes the same FIFO. ... ok +test_h3_schema_success_and_stream_with_reasoning (test_structured.Structured.test_h3_schema_success_and_stream_with_reasoning) ... [strata] done: 255 tokens in 0 s (105197.9 tok/s) (stop, cancel=False) +[strata] thinking limited to 10991 of max_tokens 14654 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 282 tokens in 0 s (104137.5 tok/s) (stop, cancel=False) +ok +test_invalid_generation_is_an_error_and_never_streams_content (test_structured.Structured.test_invalid_generation_is_an_error_and_never_streams_content) ... skipped 'jsonschema is not installed' +test_invalid_schema_is_400_and_local_refs_work (test_structured.Structured.test_invalid_schema_is_400_and_local_refs_work) ... skipped 'jsonschema is not installed' +test_malformed_formats_and_tools_are_rejected_before_loading (test_structured.Structured.test_malformed_formats_and_tools_are_rejected_before_loading) ... [strata] 400 invalid request: response_format must be an object +[strata] 400 invalid request: response_format.type must be text, json_object or json_schema +[strata] 400 invalid request: response_format.json_schema needs a name +[strata] 400 invalid request: response_format.json_schema.strict must be boolean +[strata] 400 invalid request: structured response_format with tools/MCP is not supported +[strata] 400 invalid request: structured response_format with tools/MCP is not supported +[strata] 400 invalid request: send a JSON object +ok +test_object_mode_strict_json_and_truncation (test_structured.Structured.test_object_mode_strict_json_and_truncation) ... [strata] done: 13 tokens in 0 s (188020.5 tok/s) (stop, cancel=False) +[strata] done: 3 tokens in 0 s (213269.7 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 0 s (293308.0 tok/s) (stop, cancel=False) +[strata] done: 14 tokens in 0 s (376411.9 tok/s) (stop, cancel=False) +[strata] done: 12 tokens in 0 s (364722.1 tok/s) (stop, cancel=False) +[strata] done: 15 tokens in 0 s (374491.4 tok/s) (stop, cancel=False) +[strata] done: 6 tokens in 0 s (314572.8 tok/s) (stop, cancel=False) +[strata] done: 15 tokens in 0 s (349525.3 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (309971.7 tok/s) (stop, cancel=False) +[strata] done: 5 tokens in 0 s (299593.1 tok/s) (length, cancel=False) +[strata] done: 7 tokens in 0 s (205315.6 tok/s) (stop, cancel=False) +ok +test_rejected_schema_does_not_load_the_engine_or_retry (test_structured.Structured.test_rejected_schema_does_not_load_the_engine_or_retry) ... [strata] 400 invalid request: response_format.json_schema needs a schema object +[strata] done: 21 tokens in 0 s (241316.1 tok/s) (stop, cancel=False) +ok +test_root_unions_of_object_shapes_are_accepted (test_structured.Structured.test_root_unions_of_object_shapes_are_accepted) ... [strata] done: 40 tokens in 0 s (242445.3 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (269359.0 tok/s) (stop, cancel=False) +[strata] done: 40 tokens in 0 s (300667.0 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (319131.8 tok/s) (stop, cancel=False) +[strata] done: 40 tokens in 0 s (325771.2 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (333637.8 tok/s) (stop, cancel=False) +[strata] done: 40 tokens in 0 s (322638.8 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (345413.3 tok/s) (stop, cancel=False) +[strata] done: 40 tokens in 0 s (42506.2 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (353736.5 tok/s) (stop, cancel=False) +ok +test_schemas_that_allow_non_objects_are_still_rejected (test_structured.Structured.test_schemas_that_allow_non_objects_are_still_rejected) ... [strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +[strata] 400 invalid request: response_format schema must accept only JSON objects at its root (type object, or anyOf/oneOf of object schemas) +ok +test_the_server_does_not_import_jsonschema_at_start (test_structured.Structured.test_the_server_does_not_import_jsonschema_at_start) ... ok +test_without_jsonschema_an_object_is_still_required (test_structured.Structured.test_without_jsonschema_an_object_is_still_required) ... [strata] done: 255 tokens in 0 s (105374.1 tok/s) (stop, cancel=False) +[strata] done: 26 tokens in 0 s (333492.1 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (302681.7 tok/s) (stop, cancel=False) +[strata] done: 13 tokens in 0 s (332475.3 tok/s) (stop, cancel=False) +ok +test_after_the_thinking (test_tool_name_space.SpacedName.test_after_the_thinking) ... ok +test_inside_the_thinking (test_tool_name_space.SpacedName.test_inside_the_thinking) ... ok +test_message_names_the_cause_from_the_line (test_vision_busy.Busy.test_message_names_the_cause_from_the_line) ... ok +test_message_names_the_cause_from_the_log_tail (test_vision_busy.Busy.test_message_names_the_cause_from_the_log_tail) ... ok +test_other_failures_are_unchanged (test_vision_busy.Busy.test_other_failures_are_unchanged) ... ok +test_start_raises_it (test_vision_busy.Busy.test_start_raises_it) ... ok +test_a_vram_reserve_stacks_on_the_encoders_share (test_vision_elastic.Elastic.test_a_vram_reserve_stacks_on_the_encoders_share) ... [strata] VRAM: 4000 MiB free, expert cache 2000 of 6000 MiB (100 experts) +[strata] images: the expert cache keeps 2048 MiB free for the image encoder (500 of 6000 MiB left) +[strata] images: the image encoder started in 0.0 s, 0 MiB of VRAM +[strata] VRAM: 3048 MiB free, expert cache 2952 of 6000 MiB (952 experts) +[strata] images: the image encoder stopped after 60 s without a picture; the expert cache took its VRAM back (5000 of 6000 MiB) +ok +test_an_engine_that_cannot_shrink_still_gets_its_encoder (test_vision_elastic.Elastic.test_an_engine_that_cannot_shrink_still_gets_its_encoder) ... [strata] images: the expert cache did not make room for the image encoder (VRAM needs an engine started with --vram-elastic (one NVIDIA GPU, --serve)); starting it anyway +[strata] images: the image encoder started in 0.0 s, 0 MiB of VRAM +ok +test_an_idle_encoder_stops_and_the_cache_grows_back_whole (test_vision_elastic.Elastic.test_an_idle_encoder_stops_and_the_cache_grows_back_whole) ... [strata] images: the expert cache keeps 2048 MiB free for the image encoder (3952 of 6000 MiB left) +[strata] images: the image encoder started in 0.0 s, -55 MiB of VRAM +[strata] images: the image encoder stopped after 60 s without a picture; the expert cache took its VRAM back (6000 of 6000 MiB) +[strata] images: the expert cache keeps 2048 MiB free for the image encoder (3952 of 6000 MiB left) +[strata] images: the image encoder started in 0.0 s, 0 MiB of VRAM +ok +test_an_unload_lets_go_of_the_share (test_vision_elastic.Elastic.test_an_unload_lets_go_of_the_share) ... [strata] images: the expert cache keeps 2048 MiB free for the image encoder (3952 of 6000 MiB left) +[strata] images: the image encoder started in 0.0 s, 0 MiB of VRAM +[strata] model unloaded; the next request loads it again +ok +test_the_config_keys (test_vision_elastic.Elastic.test_the_config_keys) ... ok +test_the_first_picture_makes_room_then_starts_the_encoder (test_vision_elastic.Elastic.test_the_first_picture_makes_room_then_starts_the_encoder) ... [strata] images: the expert cache keeps 2048 MiB free for the image encoder (3952 of 6000 MiB left) +[strata] images: the image encoder started in 0.0 s, 0 MiB of VRAM +ok +test_the_load_never_starts_an_elastic_encoder (test_vision_elastic.Elastic.test_the_load_never_starts_an_elastic_encoder) ... ok +test_config_key (test_vram.EngineCommand.test_config_key) ... ok +test_refused_and_died (test_vram.EngineCommand.test_refused_and_died) ... ok +test_the_line_and_the_answer (test_vram.EngineCommand.test_the_line_and_the_answer) ... ok +test_an_engine_without_the_flag_says_so (test_vram.Route.test_an_engine_without_the_flag_says_so) ... [strata] 400 invalid request: VRAM needs an engine started with --vram-elastic (one NVIDIA GPU, --serve) +ok +test_bad_input_busy_and_foreign_pages (test_vram.Route.test_bad_input_busy_and_foreign_pages) ... [strata] 400 invalid request: reserve_mib: a whole number of MiB (0 or more), or null for the start's +[strata] 400 invalid request: reserve_mib: a whole number of MiB (0 or more), or null for the start's +[strata] 400 invalid request: reserve_mib: a whole number of MiB (0 or more), or null for the start's +[strata] 400 invalid request: reserve_mib: a whole number of MiB (0 or more), or null for the start's +[strata] refused an API request from the web page 'https://other.example' (no API key; add its host to "allowed_hosts" or its origin to "cors_origins" in the config) +ok +test_resize_and_status (test_vram.Route.test_resize_and_status) ... [strata] VRAM: 4000 MiB free, expert cache 2000 of 6000 MiB (100 experts) +[strata] VRAM: 700 MiB free, expert cache 5300 of 6000 MiB (3300 experts) +ok +test_unloaded_keeps_the_reserve_for_the_next_load (test_vram.Route.test_unloaded_keeps_the_reserve_for_the_next_load) ... [strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +ok +test_changed_values_are_sent_whole (test_web_sampling.WebSampling.test_changed_values_are_sent_whole) ... ok +test_the_old_default_is_shown_as_the_new_one (test_web_sampling.WebSampling.test_the_old_default_is_shown_as_the_new_one) ... ok +test_untouched_sends_nothing (test_web_sampling.WebSampling.test_untouched_sends_nothing) ... ok +test_zero_is_greedy (test_web_sampling.WebSampling.test_zero_is_greedy) ... ok +test_a_contained_child_is_opted_out_of_power_throttling (test_winjob.TestNoThrottle.test_a_contained_child_is_opted_out_of_power_throttling) ... skipped 'power throttling is a Windows setting' +test_children_end_with_a_killed_server (test_winjob.TestWinJob.test_children_end_with_a_killed_server) ... skipped 'job objects are Windows-only' +test_elsewhere_or_without_a_process_it_is_a_no_op (test_winjob.TestWinJob.test_elsewhere_or_without_a_process_it_is_a_no_op) ... skipped 'job objects are Windows-only' + +====================================================================== +FAIL: test_service_holds_for_the_whole_request_including_early_close (test_keepawake.KeepAwakeTests.test_service_holds_for_the_whole_request_including_early_close) +---------------------------------------------------------------------- +Traceback (most recent call last): + File "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/test_keepawake.py", line 87, in test_service_holds_for_the_whole_request_including_early_close + self.assertTrue(wait_for(lambda: rec.calls == [ON, ES_CONTINUOUS])) + ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +AssertionError: False is not true + +---------------------------------------------------------------------- +Ran 852 tests in 516.658s + +FAILED (failures=1, skipped=6) diff --git a/results/round39/serve-hardening-affected.log b/results/round39/serve-hardening-affected.log new file mode 100644 index 0000000..0e81ba6 --- /dev/null +++ b/results/round39/serve-hardening-affected.log @@ -0,0 +1,57 @@ +test_bad_process (serve.test_server.MacTelemetry.test_bad_process) ... ok +test_device_is_released_on_success_and_unsupported_selector (serve.test_server.MacTelemetry.test_device_is_released_on_success_and_unsupported_selector) ... ok +test_no_registry (serve.test_server.MacTelemetry.test_no_registry) ... ok +test_note_reaches_the_monitor (serve.test_server.MacTelemetry.test_note_reaches_the_monitor) ... ok +test_readings (serve.test_server.MacTelemetry.test_readings) ... ok +test_working_set_is_read_once_for_samples_and_static (serve.test_server.MacTelemetry.test_working_set_is_read_once_for_samples_and_static) ... ok +test_a_lazy_engine_keeps_the_prompts_its_arguments_ask_for (serve.test_server.IncrementalPrompts.test_a_lazy_engine_keeps_the_prompts_its_arguments_ask_for) ... ok +test_a_tokenizer_without_resume_points_encodes_in_full (serve.test_server.IncrementalPrompts.test_a_tokenizer_without_resume_points_encodes_in_full) ... ok +test_counting_tokens_leaves_the_kept_prompts_alone (serve.test_server.IncrementalPrompts.test_counting_tokens_leaves_the_kept_prompts_alone) ... [strata] done: 8 tokens in 0 s (153919.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_prompt_encoder_slot_count_is_capped (serve.test_server.IncrementalPrompts.test_prompt_encoder_slot_count_is_capped) ... ok +test_same_ids_as_a_full_encode_on_varied_conversations (serve.test_server.IncrementalPrompts.test_same_ids_as_a_full_encode_on_varied_conversations) ... ok +test_the_kept_prompts_follow_the_engines_conversation_slots (serve.test_server.IncrementalPrompts.test_the_kept_prompts_follow_the_engines_conversation_slots) ... ok +test_the_slot_argument_wins_over_the_info_line_and_a_bad_one_falls_back (serve.test_server.IncrementalPrompts.test_the_slot_argument_wins_over_the_info_line_and_a_bad_one_falls_back) ... ok +test_turns (serve.test_server.IncrementalPrompts.test_turns) ... [strata] done: 32 tokens in 0 s (298925.9 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (347714.3 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (380220.2 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (381300.4 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (390167.8 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (394758.0 tok/s) (stop, cancel=False) +ok +test_without_a_slot_setting_four_prompts_are_kept (serve.test_server.IncrementalPrompts.test_without_a_slot_setting_four_prompts_are_kept) ... ok +test_a_refused_prompt_is_not_kept (serve.test_prompt_encoder.Served.test_a_refused_prompt_is_not_kept) ... ok +test_another_conversation_does_not_evict_the_first (serve.test_prompt_encoder.Served.test_another_conversation_does_not_evict_the_first) +A subagent beside the main conversation: every rendered prompt starts with the template's header, which ... ok +test_quoted_think_across_turns (serve.test_prompt_encoder.Served.test_quoted_think_across_turns) ... ok +test_random_messages_equal_a_full_encode (serve.test_prompt_encoder.Served.test_random_messages_equal_a_full_encode) +Messages quoting control-token texts and tags at random, growing, edited and cut back between ... ok +test_a_turn_reuses_the_previous_one (serve.test_prompt_encoder.Units.test_a_turn_reuses_the_previous_one) ... ok +test_plain_agree_until (serve.test_prompt_encoder.Units.test_plain_agree_until) ... ok +test_same_ids_as_a_full_encode (serve.test_prompt_encoder.Units.test_same_ids_as_a_full_encode) ... ok +test_configured_vision_is_not_omitted (serve.test_image_omission.ImageOmission.test_configured_vision_is_not_omitted) ... ok +test_disabled_switch_keeps_the_normal_vision_error (serve.test_image_omission.ImageOmission.test_disabled_switch_keeps_the_normal_vision_error) ... ok +test_omitted_images_are_cached_as_text_and_retries_are_stable (serve.test_image_omission.ImageOmission.test_omitted_images_are_cached_as_text_and_retries_are_stable) ... ok +test_capability_and_batch (serve.test_logit_bias.LogitBias.test_capability_and_batch) ... ok +test_forms_and_transport (serve.test_logit_bias.LogitBias.test_forms_and_transport) ... ok +test_http_rejects_before_stream (serve.test_logit_bias.LogitBias.test_http_rejects_before_stream) ... [strata] 400 invalid request: logit_bias: expected an object or a list of [token_id, bias] pairs +[strata] 400 invalid request: logit_bias: this engine does not support token biases; use a supporting CUDA/HIP build +[strata] 400 invalid request: logit_bias: biases must be finite numbers between -100 and 100; false is a pair-list ban +ok +test_invalid (serve.test_logit_bias.LogitBias.test_invalid) ... ok +test_large_list (serve.test_logit_bias.LogitBias.test_large_list) ... ok +test_a_bad_list_is_never_stored (serve.test_logit_bias.LogitBiasCache.test_a_bad_list_is_never_stored) ... ok +test_an_equal_list_in_a_later_request_is_a_hit (serve.test_logit_bias.LogitBiasCache.test_an_equal_list_in_a_later_request_is_a_hit) ... ok +test_false_is_not_zero (serve.test_logit_bias.LogitBiasCache.test_false_is_not_zero) ... ok +test_json_equivalent_pairs_and_token_keys (serve.test_logit_bias.LogitBiasCache.test_json_equivalent_pairs_and_token_keys) ... ok +test_same_answers_as_the_plain_function (serve.test_logit_bias.LogitBiasCache.test_same_answers_as_the_plain_function) ... ok +test_same_length_mutations_are_validated (serve.test_logit_bias.LogitBiasCache.test_same_length_mutations_are_validated) ... ok +test_the_cache_is_bounded (serve.test_logit_bias.LogitBiasCache.test_the_cache_is_bounded) ... ok +test_the_result_is_a_copy (serve.test_logit_bias.LogitBiasCache.test_the_result_is_a_copy) ... ok +test_vocabulary_size_is_part_of_the_key (serve.test_logit_bias.LogitBiasCache.test_vocabulary_size_is_part_of_the_key) ... ok + +---------------------------------------------------------------------- +Ran 39 tests in 3.786s + +OK diff --git a/results/round39/serve-hardening-final.log b/results/round39/serve-hardening-final.log new file mode 100644 index 0000000..fbbb30f --- /dev/null +++ b/results/round39/serve-hardening-final.log @@ -0,0 +1,57 @@ +test_bad_process (serve.test_server.MacTelemetry.test_bad_process) ... ok +test_device_is_released_on_success_and_unsupported_selector (serve.test_server.MacTelemetry.test_device_is_released_on_success_and_unsupported_selector) ... ok +test_no_registry (serve.test_server.MacTelemetry.test_no_registry) ... ok +test_note_reaches_the_monitor (serve.test_server.MacTelemetry.test_note_reaches_the_monitor) ... ok +test_readings (serve.test_server.MacTelemetry.test_readings) ... ok +test_working_set_is_read_once_for_samples_and_static (serve.test_server.MacTelemetry.test_working_set_is_read_once_for_samples_and_static) ... ok +test_a_lazy_engine_keeps_the_prompts_its_arguments_ask_for (serve.test_server.IncrementalPrompts.test_a_lazy_engine_keeps_the_prompts_its_arguments_ask_for) ... ok +test_a_tokenizer_without_resume_points_encodes_in_full (serve.test_server.IncrementalPrompts.test_a_tokenizer_without_resume_points_encodes_in_full) ... ok +test_counting_tokens_leaves_the_kept_prompts_alone (serve.test_server.IncrementalPrompts.test_counting_tokens_leaves_the_kept_prompts_alone) ... [strata] done: 8 tokens in 0 s (156796.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +ok +test_prompt_encoder_slot_count_is_capped (serve.test_server.IncrementalPrompts.test_prompt_encoder_slot_count_is_capped) ... ok +test_same_ids_as_a_full_encode_on_varied_conversations (serve.test_server.IncrementalPrompts.test_same_ids_as_a_full_encode_on_varied_conversations) ... ok +test_the_kept_prompts_follow_the_engines_conversation_slots (serve.test_server.IncrementalPrompts.test_the_kept_prompts_follow_the_engines_conversation_slots) ... ok +test_the_slot_argument_wins_over_the_info_line_and_a_bad_one_falls_back (serve.test_server.IncrementalPrompts.test_the_slot_argument_wins_over_the_info_line_and_a_bad_one_falls_back) ... ok +test_turns (serve.test_server.IncrementalPrompts.test_turns) ... [strata] done: 32 tokens in 0 s (309971.7 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (366715.1 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (367719.8 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (385683.1 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (384578.0 tok/s) (stop, cancel=False) +[strata] done: 32 tokens in 0 s (394758.0 tok/s) (stop, cancel=False) +ok +test_without_a_slot_setting_four_prompts_are_kept (serve.test_server.IncrementalPrompts.test_without_a_slot_setting_four_prompts_are_kept) ... ok +test_a_refused_prompt_is_not_kept (serve.test_prompt_encoder.Served.test_a_refused_prompt_is_not_kept) ... ok +test_another_conversation_does_not_evict_the_first (serve.test_prompt_encoder.Served.test_another_conversation_does_not_evict_the_first) +A subagent beside the main conversation: every rendered prompt starts with the template's header, which ... ok +test_quoted_think_across_turns (serve.test_prompt_encoder.Served.test_quoted_think_across_turns) ... ok +test_random_messages_equal_a_full_encode (serve.test_prompt_encoder.Served.test_random_messages_equal_a_full_encode) +Messages quoting control-token texts and tags at random, growing, edited and cut back between ... ok +test_a_turn_reuses_the_previous_one (serve.test_prompt_encoder.Units.test_a_turn_reuses_the_previous_one) ... ok +test_plain_agree_until (serve.test_prompt_encoder.Units.test_plain_agree_until) ... ok +test_same_ids_as_a_full_encode (serve.test_prompt_encoder.Units.test_same_ids_as_a_full_encode) ... ok +test_configured_vision_is_not_omitted (serve.test_image_omission.ImageOmission.test_configured_vision_is_not_omitted) ... ok +test_disabled_switch_keeps_the_normal_vision_error (serve.test_image_omission.ImageOmission.test_disabled_switch_keeps_the_normal_vision_error) ... ok +test_omitted_images_are_cached_as_text_and_retries_are_stable (serve.test_image_omission.ImageOmission.test_omitted_images_are_cached_as_text_and_retries_are_stable) ... ok +test_capability_and_batch (serve.test_logit_bias.LogitBias.test_capability_and_batch) ... ok +test_forms_and_transport (serve.test_logit_bias.LogitBias.test_forms_and_transport) ... ok +test_http_rejects_before_stream (serve.test_logit_bias.LogitBias.test_http_rejects_before_stream) ... [strata] 400 invalid request: logit_bias: expected an object or a list of [token_id, bias] pairs +[strata] 400 invalid request: logit_bias: this engine does not support token biases; use a supporting CUDA/HIP build +[strata] 400 invalid request: logit_bias: biases must be finite numbers between -100 and 100; false is a pair-list ban +ok +test_invalid (serve.test_logit_bias.LogitBias.test_invalid) ... ok +test_large_list (serve.test_logit_bias.LogitBias.test_large_list) ... ok +test_a_bad_list_is_never_stored (serve.test_logit_bias.LogitBiasCache.test_a_bad_list_is_never_stored) ... ok +test_an_equal_list_in_a_later_request_is_a_hit (serve.test_logit_bias.LogitBiasCache.test_an_equal_list_in_a_later_request_is_a_hit) ... ok +test_false_is_not_zero (serve.test_logit_bias.LogitBiasCache.test_false_is_not_zero) ... ok +test_json_equivalent_pairs_and_token_keys (serve.test_logit_bias.LogitBiasCache.test_json_equivalent_pairs_and_token_keys) ... ok +test_same_answers_as_the_plain_function (serve.test_logit_bias.LogitBiasCache.test_same_answers_as_the_plain_function) ... ok +test_same_length_mutations_are_validated (serve.test_logit_bias.LogitBiasCache.test_same_length_mutations_are_validated) ... ok +test_the_cache_is_bounded (serve.test_logit_bias.LogitBiasCache.test_the_cache_is_bounded) ... ok +test_the_result_is_a_copy (serve.test_logit_bias.LogitBiasCache.test_the_result_is_a_copy) ... ok +test_vocabulary_size_is_part_of_the_key (serve.test_logit_bias.LogitBiasCache.test_vocabulary_size_is_part_of_the_key) ... ok + +---------------------------------------------------------------------- +Ran 39 tests in 3.698s + +OK diff --git a/results/round39/serve-image-omission.log b/results/round39/serve-image-omission.log new file mode 100644 index 0000000..5346403 --- /dev/null +++ b/results/round39/serve-image-omission.log @@ -0,0 +1,8 @@ +test_configured_vision_is_not_omitted (serve.test_image_omission.ImageOmission.test_configured_vision_is_not_omitted) ... ok +test_disabled_switch_keeps_the_normal_vision_error (serve.test_image_omission.ImageOmission.test_disabled_switch_keeps_the_normal_vision_error) ... ok +test_omitted_images_are_cached_as_text_and_retries_are_stable (serve.test_image_omission.ImageOmission.test_omitted_images_are_cached_as_text_and_retries_are_stable) ... ok + +---------------------------------------------------------------------- +Ran 3 tests in 0.047s + +OK diff --git a/results/round39/serve-logit-bias-final.log b/results/round39/serve-logit-bias-final.log new file mode 100644 index 0000000..bb320a7 --- /dev/null +++ b/results/round39/serve-logit-bias-final.log @@ -0,0 +1,22 @@ +test_capability_and_batch (serve.test_logit_bias.LogitBias.test_capability_and_batch) ... ok +test_forms_and_transport (serve.test_logit_bias.LogitBias.test_forms_and_transport) ... ok +test_http_rejects_before_stream (serve.test_logit_bias.LogitBias.test_http_rejects_before_stream) ... [strata] 400 invalid request: logit_bias: expected an object or a list of [token_id, bias] pairs +[strata] 400 invalid request: logit_bias: this engine does not support token biases; use a supporting CUDA/HIP build +[strata] 400 invalid request: logit_bias: biases must be finite numbers between -100 and 100; false is a pair-list ban +ok +test_invalid (serve.test_logit_bias.LogitBias.test_invalid) ... ok +test_large_list (serve.test_logit_bias.LogitBias.test_large_list) ... ok +test_a_bad_list_is_never_stored (serve.test_logit_bias.LogitBiasCache.test_a_bad_list_is_never_stored) ... ok +test_an_equal_list_in_a_later_request_is_a_hit (serve.test_logit_bias.LogitBiasCache.test_an_equal_list_in_a_later_request_is_a_hit) ... ok +test_false_is_not_zero (serve.test_logit_bias.LogitBiasCache.test_false_is_not_zero) ... ok +test_json_equivalent_pairs_and_token_keys (serve.test_logit_bias.LogitBiasCache.test_json_equivalent_pairs_and_token_keys) ... ok +test_same_answers_as_the_plain_function (serve.test_logit_bias.LogitBiasCache.test_same_answers_as_the_plain_function) ... ok +test_same_length_mutations_are_validated (serve.test_logit_bias.LogitBiasCache.test_same_length_mutations_are_validated) ... ok +test_the_cache_is_bounded (serve.test_logit_bias.LogitBiasCache.test_the_cache_is_bounded) ... ok +test_the_result_is_a_copy (serve.test_logit_bias.LogitBiasCache.test_the_result_is_a_copy) ... ok +test_vocabulary_size_is_part_of_the_key (serve.test_logit_bias.LogitBiasCache.test_vocabulary_size_is_part_of_the_key) ... ok + +---------------------------------------------------------------------- +Ran 14 tests in 1.394s + +OK diff --git a/results/round39/serve-logit-bias-upstream-repro.log b/results/round39/serve-logit-bias-upstream-repro.log new file mode 100644 index 0000000..e1ab28f --- /dev/null +++ b/results/round39/serve-logit-bias-upstream-repro.log @@ -0,0 +1,26 @@ +test_same_length_mutations_are_validated (serve.test_logit_bias.LogitBiasCache.test_same_length_mutations_are_validated) ... + test_same_length_mutations_are_validated (serve.test_logit_bias.LogitBiasCache.test_same_length_mutations_are_validated) (form='list') ... FAIL + test_same_length_mutations_are_validated (serve.test_logit_bias.LogitBiasCache.test_same_length_mutations_are_validated) (form='dict') ... FAIL + +====================================================================== +FAIL: test_same_length_mutations_are_validated (serve.test_logit_bias.LogitBiasCache.test_same_length_mutations_are_validated) (form='list') +---------------------------------------------------------------------- +Traceback (most recent call last): + File "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/test_logit_bias.py", line 146, in test_same_length_mutations_are_validated + self.assertEqual(normalize(value, 1000)[0], -100) + ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +AssertionError: 0.0 != -100 + +====================================================================== +FAIL: test_same_length_mutations_are_validated (serve.test_logit_bias.LogitBiasCache.test_same_length_mutations_are_validated) (form='dict') +---------------------------------------------------------------------- +Traceback (most recent call last): + File "/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/test_logit_bias.py", line 146, in test_same_length_mutations_are_validated + self.assertEqual(normalize(value, 1000)[0], -100) + ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +AssertionError: 0.0 != -100 + +---------------------------------------------------------------------- +Ran 1 test in 0.001s + +FAILED (failures=2) diff --git a/results/round39/serve-logit-bias.log b/results/round39/serve-logit-bias.log new file mode 100644 index 0000000..66b4008 --- /dev/null +++ b/results/round39/serve-logit-bias.log @@ -0,0 +1,21 @@ +test_capability_and_batch (serve.test_logit_bias.LogitBias.test_capability_and_batch) ... ok +test_forms_and_transport (serve.test_logit_bias.LogitBias.test_forms_and_transport) ... ok +test_http_rejects_before_stream (serve.test_logit_bias.LogitBias.test_http_rejects_before_stream) ... [strata] 400 invalid request: logit_bias: expected an object or a list of [token_id, bias] pairs +[strata] 400 invalid request: logit_bias: this engine does not support token biases; use a supporting CUDA/HIP build +[strata] 400 invalid request: logit_bias: biases must be finite numbers between -100 and 100; false is a pair-list ban +ok +test_invalid (serve.test_logit_bias.LogitBias.test_invalid) ... ok +test_large_list (serve.test_logit_bias.LogitBias.test_large_list) ... ok +test_a_bad_list_is_never_stored (serve.test_logit_bias.LogitBiasCache.test_a_bad_list_is_never_stored) ... ok +test_an_equal_list_in_a_later_request_is_a_hit (serve.test_logit_bias.LogitBiasCache.test_an_equal_list_in_a_later_request_is_a_hit) ... ok +test_false_is_not_zero (serve.test_logit_bias.LogitBiasCache.test_false_is_not_zero) ... ok +test_same_answers_as_the_plain_function (serve.test_logit_bias.LogitBiasCache.test_same_answers_as_the_plain_function) ... ok +test_same_length_mutations_are_validated (serve.test_logit_bias.LogitBiasCache.test_same_length_mutations_are_validated) ... ok +test_the_cache_is_bounded (serve.test_logit_bias.LogitBiasCache.test_the_cache_is_bounded) ... ok +test_the_result_is_a_copy (serve.test_logit_bias.LogitBiasCache.test_the_result_is_a_copy) ... ok +test_vocabulary_size_is_part_of_the_key (serve.test_logit_bias.LogitBiasCache.test_vocabulary_size_is_part_of_the_key) ... ok + +---------------------------------------------------------------------- +Ran 13 tests in 0.751s + +OK diff --git a/results/round39/serve-owned-final.patch b/results/round39/serve-owned-final.patch new file mode 100644 index 0000000..bc4defd --- /dev/null +++ b/results/round39/serve-owned-final.patch @@ -0,0 +1,2932 @@ +diff --git a/serve/frontend.py b/serve/frontend.py +index c36a44da..3fda9c56 100644 +--- a/serve/frontend.py ++++ b/serve/frontend.py +@@ -637,6 +637,11 @@ def anthropic_to_messages(req: dict, think_unasked: bool = True) -> tuple[list[d + # "output_config": {"effort": "low" | "medium" | "high"} + thinking = req.get("thinking") + effort = (req.get("output_config") or {}).get("effort") if isinstance(req.get("output_config"), dict) else None ++ # 0.1.45: Anthropic's budget_tokens is a real budget, not only a level: the thinking is closed (wrap-up, , ++ # the answer) after that many tokens. A reasoning_budget_tokens of the request's own wins. ++ bt = thinking.get("budget_tokens") if isinstance(thinking, dict) and thinking.get("type") != "disabled" else None ++ if isinstance(bt, int) and not isinstance(bt, bool) and bt > 0 and req.get("reasoning_budget_tokens") is None: ++ req["reasoning_budget_tokens"] = bt + if isinstance(thinking, dict) and thinking.get("type") == "disabled": + kwargs["enable_thinking"] = False + elif effort: +@@ -1039,6 +1044,7 @@ class OutputParser: + return out + + def _reset_scan(self): ++ self._call_end_checked = 0 # closing-tag search progress for the current wrapped call + self.sp = 0 # how much of self.buf (the call body) the scanner has consumed + self.ss = "name" # name -> between -> str|raw -> ... -> done + self.scall = None # the ToolCall being streamed (its id is reused by the final event) +@@ -1046,6 +1052,20 @@ class OutputParser: + self.sval_started = False + self.sdeclared = {} + ++ def _wrapped_call_end(self) -> int: ++ # Without a closing tag the structural walk cannot finish. Search only ++ # newly appended text plus enough overlap for a tag split across feeds. ++ # A literal closing tag inside a value needs structural disambiguation: ++ # once any candidate appears, keep the original full walk for this call. ++ # Recovery can rewrite the buffer, so it always uses the full walk. ++ if not self.recover and self._call_end_checked >= 0: ++ start = max(0, self._call_end_checked - len(CALL_END) + 1) ++ if self.buf.find(CALL_END, start) < 0: ++ self._call_end_checked = len(self.buf) ++ return -1 ++ self._call_end_checked = -1 ++ return call_end(self.buf) ++ + def _scan(self) -> list[Event]: + """Advance the streaming view of the call body in self.buf (see stream_tools).""" + out = [] +@@ -1152,10 +1172,19 @@ class OutputParser: + def _hold(self, text: str, tags: tuple[str, ...]) -> int: + """Length of the longest suffix of `text` that is a proper prefix of one of `tags`.""" + best = 0 ++ size = len(text) + for tag in tags: +- for n in range(1, len(tag)): +- if text.endswith(tag[:n]): +- best = max(best, n) ++ if len(tag) <= 1: ++ continue ++ # Only a suffix beginning with the tag's first character can match. ++ # Ordinary token deltas contain no '<': one C-level search replaces ++ # checking every possible prefix. Earliest candidate is the longest. ++ start = text.find(tag[0], max(0, size - len(tag) + 1)) ++ while start >= 0 and size - start > best: ++ if tag.startswith(text[start:]): ++ best = size - start ++ break ++ start = text.find(tag[0], start + 1) + return best + + def feed(self, delta: str) -> list[Event]: +@@ -1351,7 +1380,7 @@ class OutputParser: + elif after and not CALL_END.startswith(after): + drop = 0 + else: +- i = call_end(self.buf) ++ i = self._wrapped_call_end() + if i >= 0: + drop = len(CALL_END) + if self.recover: # a batch in one wrapper: the next call follows this one's +diff --git a/serve/logit_bias.py b/serve/logit_bias.py +index 59873a77..700511c0 100644 +--- a/serve/logit_bias.py ++++ b/serve/logit_bias.py +@@ -1,9 +1,12 @@ + """Request-scoped token biases. Normalize before HTTP streaming starts.""" ++import json + import math + import re ++import threading ++from collections import OrderedDict + + +-def normalize(value, vocab_size=None): ++def _normalize(value, vocab_size=None): + if value is None: + return {} + if isinstance(value, dict): +@@ -36,9 +39,52 @@ def normalize(value, vocab_size=None): + return out + + ++# A long list (a standing ban of a whole script is tens of thousands of IDs) is sent with every request and used to be ++# normalized five times per request: in the OpenAI route, twice in validate_request and twice in engine_key (sampling_keys ++# runs for the cache key and again for the GEN line), about 170 ms for 65,929 entries. JSON text is an exact cache key ++# that tells False from 0 and "1" from 1, and catches in-place edits to the same request object. ++# Only a list that normalized without an error is stored, so a bad list still raises on every call. Short lists skip the ++# cache: they cost less than the key. ++_CACHE_MIN = 64 ++_CACHE_MAX = 4 ++_cache = OrderedDict() # (vocab_size, JSON text) -> [normalized dict, engine key or None] ++_lock = threading.Lock() ++ ++ ++def _entry(value, vocab_size): ++ if not isinstance(value, (dict, list)) or len(value) < _CACHE_MIN: ++ return [_normalize(value, vocab_size), None] ++ try: ++ key = (vocab_size, json.dumps(value, separators=(",", ":"))) ++ except (TypeError, ValueError): ++ key = None # something JSON cannot write: _normalize raises for it ++ entry = None ++ if key is not None: ++ with _lock: ++ entry = _cache.get(key) ++ if entry is not None: ++ _cache.move_to_end(key) ++ if entry is None: ++ entry = [_normalize(value, vocab_size), None] ++ if key is not None: ++ with _lock: ++ _cache[key] = entry ++ while len(_cache) > _CACHE_MAX: ++ _cache.popitem(last=False) ++ return entry ++ ++ ++def normalize(value, vocab_size=None): ++ return dict(_entry(value, vocab_size)[0]) ++ ++ + def engine_key(value): +- entries = normalize(value) +- return " logit_bias=" + ",".join(f"{i}:{b:g}" for i, b in sorted(entries.items())) if entries else "" ++ entry = _entry(value, None) ++ if not entry[0]: ++ return "" ++ if entry[1] is None: ++ entry[1] = " logit_bias=" + ",".join(f"{i}:{b:g}" for i, b in sorted(entry[0].items())) ++ return entry[1] + + + def validate_request(req, engine, vocab_size): +diff --git a/serve/runconfig.py b/serve/runconfig.py +index 472b0021..842554ab 100644 +--- a/serve/runconfig.py ++++ b/serve/runconfig.py +@@ -18,7 +18,7 @@ from pathlib import Path + # of the "sampling" block, ("arg", flag) for an engine option kept in "args". + EDITABLE = [ + ("sampling.temperature", ("sampling", "num>=0"), +- "Default temperature for requests that send none (0 = greedy, the default without a sampling block)"), ++ "Default temperature for requests that send none (0 = greedy; without it the model card's: 1.0 thinking, 0.7 not)"), + ("sampling.top_p", ("sampling", "0, with no answer +@@ -254,6 +256,15 @@ def _timeout_env(name: str, default: float) -> float | None: + # still working is never ended by this, only reported once (the recommendation: engine_silence_s, or this variable). + # STRATA_ENGINE_STALL_S overrides it (0 = off). Needs psutil (setup installs it); without it nothing is ended. + ENGINE_STALL_S = _timeout_env("STRATA_ENGINE_STALL_S", 90.0) ++# --batch: new requests take the control lines (one prompt read at a time) short prompts first: a prompt of at most ++# STRATA_ADMIT_SHORT_FIRST tokens goes before longer ones (opt-in; unset or 0: whoever gets the lock first), and a ++# longer one that waited STRATA_ADMIT_MAX_WAIT_S (default 30) counts as short. A chat question then waits for the one ++# read in progress, not for every agent read that came before it. ++try: ++ ADMIT_SHORT_FIRST = max(0, int(os.environ.get("STRATA_ADMIT_SHORT_FIRST") or 0)) ++except ValueError: ++ ADMIT_SHORT_FIRST = 0 ++ADMIT_MAX_WAIT_S = _timeout_env("STRATA_ADMIT_MAX_WAIT_S", 30.0) or 0.0 + STALL_CPU_EPS_S = 0.5 # CPU seconds the engine may use in the whole stall window and still count as frozen + STALL_IO_EPS_B = 1 << 20 # bytes read or written in that window + +@@ -705,6 +716,19 @@ class ConvCacheLog: + return dict(self.state) + + ++def conversation_slots(engine) -> int | None: ++ """How many conversations the engine's cache holds: its --conversation-cache-slots argument, which is known ++ before the engine has started (a lazy start has reported no INFO line yet), else the INFO line it reported.""" ++ args = engine.spawn[1] if hasattr(engine, "spawn") else [] ++ if "--conversation-cache-slots" in args[:-1]: ++ try: ++ return int(args[args.index("--conversation-cache-slots") + 1]) ++ except ValueError: ++ pass ++ slots = (getattr(engine, "info", None) or {}).get("conversation_cache_slots") ++ return slots if isinstance(slots, int) else None ++ ++ + def conversation_cache_view(info: dict, hist: list, totals: dict, parked: dict) -> dict: + """#596: the Monitor's Conversation cache card: the parked conversations (the engine's opt-in + --conversation-cache-mib: budget, slots, what its log says) and how much of the prompts the cache gave back.""" +@@ -1096,6 +1120,13 @@ class StrataEngine: + self.last.update(prompt_read=int(f[14])) + if len(f) >= 16: # #588 (engine 0.1.39+): routed experts read over PCIe + self.last.update(offloaded=int(f[15])) ++ for tok in f[16:]: # key=value fields after the positional ones (the prompt's ++ key, eq, value = tok.partition("=") # first chunk and its streamed expert blobs) ++ if eq and key in ("chunk", "experts_streamed"): ++ try: ++ self.last[key] = int(value) ++ except ValueError: ++ pass + + def vram(self, reserve_mib: int | None, timeout: float = 120.0) -> dict: + """#533: `VRAM ` between requests (the caller holds the service's FIFO): the engine shrinks its +@@ -1173,6 +1204,10 @@ class StrataEngine: + v = tune.get(k) + if isinstance(v, (int, float)) and not isinstance(v, bool) and 0.0 <= float(v) <= 1.0: + keys += f" {k}={float(v)!r}" ++ # a layer split's prompt chunks (--prefill-pipe): 0 off, 1 evened out, > 1 the rig's b/a in tokens ++ v = tune.get("prefill_pipe") ++ if isinstance(v, (int, float)) and not isinstance(v, bool) and 0.0 <= float(v) <= 262144.0: ++ keys += f" pipe={float(v)!r}" + # --aux-cpus for this request: 1 puts the engine's threads that are neither pool workers nor the host on the + # spare CPUs, 0 leaves them where they were created (true/false or 1/0; the engine started without a spare CPU + # ignores a 1). Anything else is not sent. +@@ -1416,10 +1451,11 @@ class StrataEngine: + + YIELDS_MAX = 2 # #656: how often one request's prompt read gives way to a shorter waiting one + +- def _take_control(self, cancel, plen: int, after_epoch: int | None = None): ++ def _take_control(self, cancel, plen: int, after_epoch: int | None = None, new: bool = False): + """Waits for the control lines (one prompt read at a time), yielding None heartbeats; False when cancelled. +- `after_epoch`: a request whose read gave way lets the requests waiting then go first.""" +- entry = [plen] ++ `after_epoch`: a request whose read gave way lets the requests waiting then go first. ++ `new`: the request's first admission (STRATA_ADMIT_SHORT_FIRST orders only these).""" ++ entry = [plen, time.monotonic(), new] # (length, arrival, ordered: STRATA_ADMIT_SHORT_FIRST) + with self.slot_cv: + self.waiting += 1 + self.wait_lens.append(entry) +@@ -1433,7 +1469,17 @@ class StrataEngine: + raise EngineDied("the engine stopped while this request waited; it was not sent") + with self.slot_cv: + turn = after_epoch is None or self.ctl_epoch > after_epoch or self.waiting <= 1 ++ if turn and new and ADMIT_SHORT_FIRST > 0: ++ turn = self._admit_next() is entry + if turn and self.ctl.acquire(timeout=0.5): ++ if new and ADMIT_SHORT_FIRST > 0: ++ with self.slot_cv: # a shorter one came while this one waited in acquire() ++ first = self._admit_next() is entry ++ if not first: ++ self.ctl.release() ++ with self.slot_cv: ++ self.slot_cv.notify_all() ++ continue + if self.alive(): + break + self.ctl.release() # it died (or is being restarted) while this one was getting +@@ -1456,6 +1502,18 @@ class StrataEngine: + self.slot_cv.notify_all() + return True + ++ def _admit_next(self): ++ """STRATA_ADMIT_SHORT_FIRST= (the caller holds slot_cv): the new request that takes the control lines ++ next - one whose prompt is at most that long (it reads little) before a longer one, in arrival order within ++ each; a longer one that has waited STRATA_ADMIT_MAX_WAIT_S counts as short (it is never held back for good). ++ Off (unset or 0): whoever gets the lock first, as before. Only the order changes, never what is read.""" ++ now = time.monotonic() ++ ordered = [e for e in self.wait_lens if len(e) > 2 and e[2]] ++ if not ordered: ++ return None ++ return min(ordered, key=lambda e: (0 if e[0] <= ADMIT_SHORT_FIRST or now - e[1] >= ADMIT_MAX_WAIT_S else 1, ++ e[1])) ++ + SOLO_AGAIN_MAX = 2 # how often a request left alone in a slot goes back to the solo path + SOLO_AGAIN_MIN_LEFT = 32 # ... only with at least this many tokens still allowed (max_tokens) + +@@ -1493,7 +1551,7 @@ class StrataEngine: + prompt, left = list(ids), int(max_new) + reused0 = None # the first read's reused tokens: a continued leg's count includes this request's own output + carry = {} # #1869: the figures of the legs that ended in a hand-back to the solo path (the final DONE adds to them) +- ok = yield from self._take_control(cancel, len(prompt)) ++ ok = yield from self._take_control(cancel, len(prompt), new=True) + if not ok: + return + holding, born = True, self.gen # (the engine it now has the control lines of) +@@ -1803,6 +1861,7 @@ class StrataEngine: + + slot_groups = 1 # --batch-groups (set when the engine starts) + SLOT_PREFIX_MIN = 512 # a held prefix shorter than this is not worth unbalancing the pipeline's groups ++ TURN_ID = 248045 # <|im_start|>: the engine's default --turn-token (its turn checkpoint sits at the last one) + + def pick_slot(self, prompt: list[int]) -> int | None: + """A free slot for `prompt` (the caller holds slot_cv): the one whose held tokens are the longest start of the +@@ -1812,18 +1871,30 @@ class StrataEngine: + if not free: + return None + gs = getattr(self, "slot_gs", 0) +- if gs: +- # pipelined groups: the engine keeps no pipelined slot as a conversation cache, so nothing is worth keeping +- # free for a next turn. A group's window spans its slots up to its last busy one, so first fill a hole +- # below a busy slot of the same group (it costs no extra row), then follow the spreading order ++ if gs > 1: ++ # pipelined groups of several slots: the engine keeps no idle slot below a busy one as a conversation cache ++ # (its pad row writes the slot), so nothing is worth keeping free for a next turn. A group's window spans ++ # its slots up to its last busy one, so first fill a hole below a busy slot of the same group (it costs no ++ # extra row), then follow the spreading order + # (reported on #793: least-recently-used picks left leading holes, [_ B], that cost a whole row each) + def fills_hole(b): + g0 = b - b % gs + return any(self.slot_busy[x] for x in range(b + 1, g0 + gs)) + return min(free, key=lambda b: (not fills_hole(b), self.slot_order.index(b))) + def held_prefix(b): ++ # what of `prompt` the slot's cache serves: the slot's whole held stream (the prompt and what it decoded) ++ # when this prompt continues it, else its turn checkpoint - the engine takes one at the last <|im_start|> of ++ # a prompt (the next request of the chat renders the history again, but not the reply it decoded or the ++ # assistant header), so a prompt that agrees with the held tokens up to there is that chat's next turn ++ # too (#2518: it was only found by the engine's own source search, after pick_slot had sent it to another ++ # slot and the request in that slot had overwritten the chat whose cache the next one needed) + h = self.slot_held[b] +- return len(h) if h and len(h) < len(prompt) and prompt[:len(h)] == h else 0 ++ if not h: ++ return 0 ++ if len(h) < len(prompt) and prompt[:len(h)] == h: ++ return len(h) ++ k = len(h) - 1 - h[::-1].index(self.TURN_ID) if self.TURN_ID in h else 0 ++ return k if 0 < k < len(prompt) and prompt[:k] == h[:k] else 0 + best = max(free, key=held_prefix) + # without --batch-groups a held prefix of any length is taken, as before + if held_prefix(best) >= (self.SLOT_PREFIX_MIN if self.slot_groups > 1 else 1): +@@ -2868,23 +2939,32 @@ class ByteTokenizer: + + ALWAYS = () # specials matched without parse_special (type 4, as ) + ++ @property ++ def max_special_len(self): ++ return max(len(s) for s in self.SPECIALS) ++ + @property + def control_tokens(self): + return [s for s in self.SPECIALS if s not in self.ALWAYS] + + def encode(self, text, parse_special=False, plain=()): +- out, i = [], 0 ++ return self.encode_marked(text, parse_special, plain)[0] ++ ++ def encode_marked(self, text, parse_special=False, plain=()): ++ """encode() and its resume points (after each special), as strata_tokenizer's for PromptEncoder.""" ++ out, marks, i = [], [], 0 + while i < len(text): + for k, s in enumerate(self.SPECIALS): + if (parse_special or s in self.ALWAYS) and text.startswith(s, i) and not any( + a <= i < b for a, b in plain): + out.append(256 + k) + i += len(s) ++ marks.append((i, len(out))) + break + else: + out.extend(text[i].encode("utf-8")) + i += 1 +- return out ++ return out, marks + + def decode(self, ids, errors="replace"): + raw = bytearray() +@@ -3031,10 +3111,12 @@ class Service: + self.literals = literal_tags(getattr(tokenizer, "control_tokens", ())) # texts that stay text inside a message + self.fit_max_tokens = fit_max_tokens # --fit-max-tokens: clamp the output cap instead of 400 + self.aliases: list[str] = [] # #297: other names of the model (the config's `aliases`) ++ self.answer_room = True # 0.1.45: see ANSWER_ROOM_MIN_CAP (the config's "answer_room": false turns it off) + self.reasoning_close_retry = False # #1053 (opt-in): close the thinking once when a reply ends inside it + self.codex_compaction_cache = False # #924 (opt-in): a Codex compaction is rendered with its conversation's tools + self.codex_thread_titles = False # #923 (opt-in): answer Codex's thread-title turns without the engine + self.sampling_defaults = dict(sampling_defaults or {}) # the run config's `sampling` block ++ self.model_arch = None # the config's family ("qwen4exp", "qwen35moe"): None = no model-card defaults + self.shared = {} # the web app's Chat settings for every client (POST /settings) + self.shared_path = None # where they are kept between starts (next to the config) + self.config_effort = None # #1641: the config's "reasoning_effort": the level for requests with none +@@ -3075,6 +3157,9 @@ class Service: + "prompt_ms": 0.0, "decode_ms": 0.0, + "drafts_offered": 0, "drafts_accepted": 0} # #457: the MTP drafts, summed where reported + self.last_timings = None # the last finished request's, llama.cpp's names (/v1/status) ++ # /metrics/prometheus' two throughput gauges need the totals as they stood at the last scrape, because ++ # llama.cpp averages them over the window between two scrapes (#2341). ++ self.scrape_bucket = (0, 0.0, 0, 0.0) + from serve.prometheus import Latencies + self.latencies = Latencies() # GET /metrics in Prometheus text: the latency histograms + self.last_request_at = None # when a request last started or finished +@@ -3104,6 +3189,18 @@ class Service: + self.restore_after_load = None # a session file an image job saved before it unloaded + self.stop_ids = set(tokenizer.encode(IM_END, parse_special=True) + + tokenizer.encode("<|endoftext|>", parse_special=True)) ++ # #567: a prompt re-encodes only what follows the last special token it shares with a recent prompt (the ++ # same ids as a full encode: tools/strata_tokenizer.py PromptEncoder). Tokenizers without resume points ++ # encode in full. ++ self.prompts = None ++ if hasattr(tokenizer, "encode_marked") and hasattr(tokenizer, "max_special_len"): ++ from strata_tokenizer import PromptEncoder ++ # Keep at most one prompt per engine conversation slot, bounded independently of engine config. ++ slots = conversation_slots(engine) ++ if slots is not None and slots > 0: ++ self.prompts = PromptEncoder(tokenizer, keep=min(slots, PROMPT_ENCODER_KEEP_MAX)) ++ else: ++ self.prompts = PromptEncoder(tokenizer) + + def loaded(self) -> bool: + return not hasattr(self.engine, "alive") or self.engine.alive() +@@ -3296,20 +3393,17 @@ class Service: + threading.Thread(target=loop, daemon=True).start() + + def free_vram_mib(self) -> int | None: +- """Free VRAM on the engine's (first) GPU, from NVML (AMD backend: amdgpu's sysfs files, #301); None when it can't +- be read (then nothing is refused).""" ++ """Free VRAM on the engine's (first) GPU, from the backend's reader - NVML, the amdgpu sysfs files (#301) or ++ IOKit on macOS (Metal backend); None when it can't be read (then nothing is refused).""" + try: +- if getattr(self, "backend", None) == "hip": +- from serve.telemetry import free_vram_mib +- return free_vram_mib(int(getattr(self, "gpu_index", 0) or 0), amd=True) +- from serve.telemetry import _Nvml +- nv = _Nvml(int(getattr(self, "gpu_index", 0) or 0)) +- if not nv.ok(): ++ from serve.telemetry import gpu_reader ++ g = gpu_reader(int(getattr(self, "gpu_index", 0) or 0), amd=getattr(self, "backend", None) == "hip") ++ if not g.ok(): + return None +- m = nv.Mem() +- if nv.lib.nvmlDeviceGetMemoryInfo(nv.dev, ctypes.byref(m)) != 0: ++ r = g.read() ++ if r.get("mem_total") is None or r.get("mem_used") is None: + return None +- return int(m.free >> 20) ++ return int((r["mem_total"] - r["mem_used"]) >> 20) + except Exception: + return None + +@@ -3552,7 +3646,7 @@ class Service: + if effort: + if api == "openai": + ctk = req.get("chat_template_kwargs") if isinstance(req.get("chat_template_kwargs"), dict) else {} +- if not req.get("reasoning_effort") and not req.get("reasoning") and \ ++ if req.get("reasoning_effort") in (None, "") and not req.get("reasoning") and \ + "enable_thinking" not in ctk and "reasoning_effort" not in ctk: + req["reasoning_effort"] = effort + elif not req.get("thinking") and not req.get("output_config"): +@@ -3703,8 +3797,15 @@ class Service: + with self.status_lock: + s, t = dict(self.status), dict(self.totals) + running = len(self.live_reqs) or int(bool(s.get("busy"))) +- prompt_n = t["prompt_tokens"] - t["reused"] # what was read, as llama.cpp counts it +- prompt_s, decode_s = t["prompt_ms"] / 1000, t["decode_ms"] / 1000 ++ prompt_n = t["prompt_tokens"] - t["reused"] # what was read, as llama.cpp counts it ++ prompt_s, decode_s = t["prompt_ms"] / 1000, t["decode_ms"] / 1000 ++ # llama.cpp reads these two gauges from a bucket it resets on every scrape (server_metrics::bucket, ++ # "the gauges are averaged over the window between two scrapes"), so an idle server reports 0 rather ++ # than its lifetime average. Keep the same window here, or the names lie to llama.cpp dashboards. ++ was_prompt_n, was_prompt_s, was_predicted_n, was_decode_s = self.scrape_bucket ++ self.scrape_bucket = (prompt_n, prompt_s, t["output_tokens"], decode_s) ++ d_prompt_n, d_prompt_s = prompt_n - was_prompt_n, prompt_s - was_prompt_s ++ d_predicted_n, d_decode_s = t["output_tokens"] - was_predicted_n, decode_s - was_decode_s + rows = [ + ("llamacpp:prompt_tokens_total", "counter", "Prompt tokens read", prompt_n), + ("llamacpp:prompt_tokens_cached_total", "counter", "Prompt tokens the conversation cache already held", +@@ -3712,10 +3813,10 @@ class Service: + ("llamacpp:prompt_seconds_total", "counter", "Prompt read time", prompt_s), + ("llamacpp:tokens_predicted_total", "counter", "Tokens generated", t["output_tokens"]), + ("llamacpp:tokens_predicted_seconds_total", "counter", "Generation time", decode_s), +- ("llamacpp:prompt_tokens_seconds", "gauge", "Average prompt read throughput in tokens/s", +- prompt_n / prompt_s if prompt_s else 0), +- ("llamacpp:predicted_tokens_seconds", "gauge", "Average generation throughput in tokens/s", +- t["output_tokens"] / decode_s if decode_s else 0), ++ ("llamacpp:prompt_tokens_seconds", "gauge", "Average prompt read throughput in tokens/s, since the last scrape", ++ d_prompt_n / d_prompt_s if d_prompt_s > 0 else 0), ++ ("llamacpp:predicted_tokens_seconds", "gauge", "Average generation throughput in tokens/s, since the last scrape", ++ d_predicted_n / d_decode_s if d_decode_s > 0 else 0), + ("llamacpp:requests_processing", "gauge", "Requests running", running), + ("llamacpp:requests_deferred", "gauge", "Requests waiting their turn", s.get("queued") or 0), + ("strata:requests_total", "counter", "Requests finished", t["requests"]), +@@ -3787,17 +3888,21 @@ class Service: + tail = "" if off else EFFORT_TURN.format(EFFORT_TEXT.get(effort, EFFORT_TEXT["medium"])) + return history + tail + full[len(mine):] + +- def encode_prompt(self, messages, tools, kwargs) -> list[int]: ++ def encode_prompt(self, messages, tools, kwargs, store: bool = True) -> list[int]: + """The request's prompt: the template rendered and tokenized. #537: a / written inside a + message's text is encoded as the text it is, not as the model's reasoning markers (the template's own are). + A control token's text (<|im_start|>, <|im_end|>, <|endoftext|>, ...) inside a message is text as well: only +- the control tokens the template writes are control tokens.""" ++ the control tokens the template writes are control tokens. `store=False` (a token count) reuses the kept ++ prompts without adding or replacing one.""" + marked, marked_tools, changed = mark_think_literals(messages, tools, self.literals) + prompt = self.render_prompt(marked, marked_tools, kwargs) +- if not changed: +- return self.tok.encode(prompt, parse_special=True) +- prompt, plain = unmark_think_literals(prompt, self.literals) +- return self.tok.encode(prompt, parse_special=True, plain=plain) ++ plain = () ++ if changed: ++ prompt, plain = unmark_think_literals(prompt, self.literals) ++ if self.prompts is not None: # #567: only what follows the last special shared with a recent prompt ++ return self.prompts.encode(prompt, plain, store) ++ return self.tok.encode(prompt, parse_special=True, plain=plain) if plain else \ ++ self.tok.encode(prompt, parse_special=True) + + @contextlib.contextmanager + def _turn(self, cancel=None): +@@ -3845,14 +3950,26 @@ class Service: + content[n] = {"type": "text", "text": f"[image omitted: {why}]"} + return fetched + ++ def forget_prompt(self, ids): ++ if self.prompts is not None: ++ self.prompts.forget(ids) ++ ++ + def prepare(self, messages, tools, kwargs, max_new=None, force=None, req=None, cancel=None): + """-> (ids, thinking, max_new). An unset or non-positive max_new (some clients send -1) means "unlimited": + the rest of the context. `force` (forced_call): without thinking the reply starts with it, so it ends the + prompt; with thinking, Service.run writes it once the thinking is over.""" + # `cancel` (#1999): set when the client hangs up; the downloads, the wait for the turn and the encodes stop + # between pictures and raise ClientGone ++ if self.vision is None and os.environ.get("STRATA_OMIT_IMAGES_WHEN_NO_VISION") == "1": ++ for m in messages: ++ if isinstance(m.get("content"), list): ++ m["content"] = [ ++ {"type": "text", "text": "[image omitted on text-only engine]"} if isinstance(it, dict) and it.get("type") == "image" else it ++ for it in m["content"] ++ ] + fetched = self._note_unreadable_tool_images(messages, cancel) +- ids = self.encode_prompt(messages, tools, kwargs) ++ ids = encoded_ids = self.encode_prompt(messages, tools, kwargs) + if force and kwargs.get("enable_thinking", True) is False: + ids = ids + self.tok.encode(force, parse_special=True) + self.embeddings.path = None +@@ -3914,12 +4031,16 @@ class Service: + room = ctx - CTX_SLACK - len(ids) + if max_new is None or max_new <= 0 or (self.fit_max_tokens and room < 1): + if room < 1: ++ # A prompt refused for the context is not kept (#567): the client cannot resend it unchanged. The ++ # other refusals (engine starting, images) keep it: the same prompt comes back and its encode is reused. ++ self.forget_prompt(encoded_ids) + raise ValueError("request exceeds the context window: " + f"prompt ({len(ids)} tokens) leaves no room to answer in the context " + f"({ctx}); requests are never truncated") + max_new = room + elif max_new > room: + if not self.fit_max_tokens: ++ self.forget_prompt(encoded_ids) + raise ValueError("request exceeds the context window: " + f"prompt ({len(ids)} tokens) + max tokens ({max_new}) exceeds the context " + f"({ctx}); requests are never truncated. Send a smaller " +@@ -4076,12 +4197,23 @@ class Service: + def _run(self, ids, thinking, tools, max_new, sampling, cancel, force=None) -> Iterator[tuple[str, object]]: + """Yields ("event", Event) as text arrives, then ("done", {"finish": .., "completion_tokens": ..}). + `force` (forced_call): the opening of the call the reply must make - see prepare().""" +- defaults = {**self.sampling_defaults, **self.shared} # the config's, then the Chat settings shared with apps ++ # the model card's sampling for this thinking state, then the config's block, then the Chat settings shared with ++ # apps; the request's own fields win over all of them ++ defaults = {**model_card_sampling(self.model_arch, thinking), **self.sampling_defaults, **self.shared} + if defaults: # the request's own fields win (explicit 0 stays greedy) + req_values = {k: v for k, v in (sampling or {}).items() if v is not None} + sampling = {**defaults, **req_values} + # #123: read after the merge, so a budget shared through POST /settings is seen like the other keys + budget = self.reasoning_budget(sampling) if thinking else None ++ if thinking and not budget and self.answer_room and max_new and max_new >= ANSWER_ROOM_MIN_CAP: ++ # 0.1.45: max_tokens counts the thinking too. A client that caps the reply at 8192 (the desktop agent's ++ # harness sends max_completion_tokens: 8192 on every request) and a task the model thinks long about got ++ # the cut-off in the middle of the thinking and NO answer. So thinking is closed early enough to leave ++ # the answer its quarter of the cap (the budget's wrap-up: a short closing sentence, , the answer). ++ # A request with a budget of its own, or max_tokens below ANSWER_ROOM_MIN_CAP, is left as it was. ++ budget = max_new - min(max(ANSWER_RESERVE_MIN, max_new // 4), max_new // 2) ++ if max_new <= 32768: print(f"[strata] thinking limited to {budget} of max_tokens {max_new} so the answer has room " ++ "(\"answer_room\": false in the config turns this off)", flush=True) + if budget and max_new and max_new > 0: + # #984: a budget at or near max_tokens never fires before the reply is cut off while still thinking, so + # it has no answer. Close the thinking early enough to leave room for one. +@@ -4586,13 +4718,13 @@ def request_stats(segments: list[dict]) -> dict: + if not segments: + return {} + result = dict(segments[-1]) +- for key in ("reused", "prompt_read"): ++ for key in ("reused", "prompt_read", "chunk"): + if key in segments[0]: + result[key] = segments[0][key] + else: + result.pop(key, None) + for key in ("prompt_ms", "decode_ms", "generated", "drafts_offered", "drafts_accepted", +- "hits", "lookups", "offloaded", "ram_blobs", "file_blobs", "file_mb"): ++ "hits", "lookups", "offloaded", "ram_blobs", "file_blobs", "file_mb", "experts_streamed"): + if any(segment.get(key) is not None for segment in segments): + result[key] = sum(segment.get(key) or 0 for segment in segments) + return result +@@ -5612,7 +5744,7 @@ def make_handler(svc: Service): + if not svc.loaded() and not getattr(svc.engine, "unloaded", False): + self._json(503, {"error": {"message": "the engine is not running"}}) + return +- defaults = {**svc.sampling_defaults, **svc.shared} ++ defaults = {**model_card_sampling(svc.model_arch, True), **svc.sampling_defaults, **svc.shared} + names = {"repetition_penalty": "repeat_penalty", "penalty_last_n": "repeat_last_n"} + params = {names.get(k, k): v for k, v in defaults.items() + if k in ("temperature", "top_p", "top_k", "min_p", "seed", "repetition_penalty", +@@ -6045,7 +6177,7 @@ def make_handler(svc: Service): + messages, tools, kw = anthropic_to_messages(req, svc.anthropic_think_unasked) + stop_strings(req) # the same 400 as the request itself would get + check_request_sampling(req) # ... likewise a sampling field of the wrong type +- self._json(200, {"input_tokens": len(svc.encode_prompt(messages, tools, kw))}) ++ self._json(200, {"input_tokens": len(svc.encode_prompt(messages, tools, kw, store=False))}) + + def _anthropic(self, req): + svc.load() +@@ -6433,6 +6565,43 @@ def clean_shared_defaults(d) -> dict: + return out + + ++def split_joined_args(cfg: dict) -> None: ++ """#2267: an engine option written as ONE element of the config's "args" - "--ple-io mmap" or "--ple-io=mmap" - reaches ++ the engine as a single argument ("unknown argument: --ple-io mmap"; the usual form is two elements, "--ple-io", ++ "mmap"). Recommend, never force: such an element is split into the flag and its value, and said at startup. Only an ++ element that starts with a flag name (--letters-digits-dashes) followed by a space or "=" is touched, and never one ++ that sits in the value place of the flag before it.""" ++ args = cfg.get("args") ++ if not isinstance(args, list): ++ return ++ out, changed = [], [] ++ for a in args: ++ m = re.match(r"^(--[A-Za-z0-9][A-Za-z0-9-]*)(?:\s+|=)(\S.*)$", a) if isinstance(a, str) else None ++ prev = out[-1] if out else None ++ in_value_place = isinstance(prev, str) and re.match(r"^--[A-Za-z0-9][A-Za-z0-9-]*$", prev) is not None ++ if m is None or in_value_place: ++ out.append(a) ++ continue ++ out += [m.group(1), m.group(2).strip()] ++ changed.append(a) ++ if changed: ++ cfg["args"] = out ++ print("[strata] the config's args hold " + ", ".join(repr(c) for c in changed) + " as one element each: split into the " ++ 'flag and its value (write them as two elements, e.g. "--ple-io", "mmap")', flush=True) ++ ++ ++def read_config(path) -> dict: ++ """The run config's JSON. A mistake in it (a missing comma between two args, a trailing comma) is said with its ++ place in the file instead of a traceback (#2267: editing "args" by hand is where it happens).""" ++ text = Path(path).read_text(encoding="utf-8-sig") # Notepad adds a BOM ++ try: ++ return json.loads(text) ++ except json.JSONDecodeError as e: ++ raise SystemExit(f"[strata] {path} is not valid JSON: {e.msg} at line {e.lineno}, column {e.colno}. Check the commas " ++ 'between the entries of "args" (each option and each value is its own quoted entry, e.g. ' ++ '"--ple-io", "mmap",) and that no entry is followed by a comma at the end of a list.') from None ++ ++ + def lift_reasoning_effort_arg(cfg: dict) -> None: + """#1641: "--reasoning-effort high" in a config's "args" is meant for the thinking level, which is not an engine + option (the engine stopped at the unknown flag). It is the config's "reasoning_effort" (none, low, medium or high: +@@ -6615,6 +6784,25 @@ def check_request_sampling(req) -> None: + FAMILY_MODEL_NAMES = {"qwen4exp": "qwen3.8-flash-next", "qwen35moe": "qwen3.6-35b-a3b"} + + ++# The sampling each model card recommends ("Best Practices"), used when neither the request, the Chat settings nor the ++# run config's `sampling` block says otherwise. Greedy is only for a request or config that asks for temperature 0. ++# min_p 0, repetition_penalty 1 and presence_penalty 0 are the engine's own "off", so they are not listed. ++MODEL_CARD_SAMPLING = { ++ # Qwen3.8-Flash-Next (and Gyro, the same base model) ++ "qwen4exp": {"thinking": {"temperature": 1.0, "top_p": 0.95, "top_k": 20}, ++ "instruct": {"temperature": 0.7, "top_p": 0.8, "top_k": 20, "presence_penalty": 1.5}}, ++ # Qwen3.6-35B-A3B: the general-task sets (the card's precise-coding set is temperature 0.6, presence_penalty 0) ++ "qwen35moe": {"thinking": {"temperature": 1.0, "top_p": 0.95, "top_k": 20, "presence_penalty": 1.5}, ++ "instruct": {"temperature": 0.7, "top_p": 0.8, "top_k": 20, "presence_penalty": 1.5}}, ++} ++ ++ ++def model_card_sampling(arch, thinking) -> dict: ++ """The model card's sampling for the family `arch` and the thinking state ({} for an unknown or unset family).""" ++ sets = MODEL_CARD_SAMPLING.get(arch) if arch else None ++ return dict(sets["thinking" if thinking else "instruct"]) if sets else {} ++ ++ + def default_model_name(cfg: dict) -> str: + """The config's "model_name", else its family's name ("arch": qwen4exp, qwen35moe), else Flash-Next's.""" + return cfg.get("model_name") or FAMILY_MODEL_NAMES.get(cfg.get("arch") or "qwen4exp", "qwen3.8-flash-next") +@@ -6667,7 +6855,8 @@ def main() -> int: + "each, requests spread over them (also \"replicas\" in the config: a number, or " + "[{\"gpus\": [0, 1]}, {\"gpus\": [2, 3]}])") + a = ap.parse_args() +- cfg = json.loads(Path(a.config).read_text(encoding="utf-8-sig")) if a.config else {} # Notepad adds a BOM ++ cfg = read_config(a.config) if a.config else {} ++ split_joined_args(cfg) + lift_reasoning_effort_arg(cfg) + if a.replicas is not None: + cfg["replicas"] = a.replicas +@@ -6768,6 +6957,9 @@ def main() -> int: + model_name=default_model_name(cfg), vision=vision, + sampling_defaults=sampling_defaults, + fit_max_tokens=a.fit_max_tokens or cfg.get("fit_max_tokens") is True) ++ if a.engine == "strata": # the model card's sampling when nothing else is said (see MODEL_CARD_SAMPLING) ++ svc.model_arch = cfg.get("arch") or "qwen4exp" ++ svc.answer_room = cfg.get("answer_room") is not False + svc.reasoning_close_retry = cfg.get("reasoning_close_retry") is True # #1053: opt-in, off by default + if cfg.get("images"): # image generation (serve/images.py), set up by setup --images + icfg = {k: (os.path.abspath(os.path.join(cfg.get("cwd") or ".", v)) +diff --git a/serve/telemetry.py b/serve/telemetry.py +index 3232c63c..697eeec7 100644 +--- a/serve/telemetry.py ++++ b/serve/telemetry.py +@@ -3,7 +3,9 @@ + A background thread samples once a second and keeps the last 60 readings of each series for the sparklines: + - GPU: NVIDIA's own NVML library (nvml.dll / libnvidia-ml.so.1, installed with every driver) through ctypes, so no + pip package is needed: load, VRAM, temperature, power, PCIe link and throughput. With the AMD backend (#301): the +- amdgpu driver's Linux sysfs files - load, VRAM, temperature and power. ++ amdgpu driver's Linux sysfs files - load, VRAM, temperature and power. On macOS (Metal backend): IOKit's ++ IOAccelerator registry - load and VRAM; the temperature and power sensors need root (`powermetrics`), and an Apple ++ Silicon GPU sits on the die, so there is no PCIe link. + - CPU, RAM, disk: `psutil` when it is installed (setup installs it); without it the CPU and RAM readings fall back to + the OS (Windows GlobalMemoryStatusEx / GetSystemTimes, Linux /proc) and the disk rate is absent. + Anything that cannot be read is None; nothing here can stop the server. +@@ -14,6 +16,8 @@ import collections + import ctypes + import os + import platform ++import re ++import subprocess + import sys + import threading + import time +@@ -246,9 +250,114 @@ class _Amd: + return out + + ++def metal_working_set(): ++ """The Apple GPU's memory budget in bytes: MTLCreateSystemDefaultDevice's recommendedMaxWorkingSetSize, through ++ the Metal framework and objc_msgSend (no PyObjC) - the same property metal/src/mtl_rt.mm reads for the engine ++ (its "GPU working set" log line). None off macOS, or when it cannot be read: a selector the device does not ++ answer would raise an Objective-C exception no ctypes call can catch, so the class is asked first.""" ++ if sys.platform != "darwin": ++ return None ++ try: ++ metal = ctypes.CDLL("/System/Library/Frameworks/Metal.framework/Metal") ++ objc = ctypes.CDLL("/usr/lib/libobjc.dylib") ++ metal.MTLCreateSystemDefaultDevice.restype = ctypes.c_void_p ++ objc.sel_registerName.restype = ctypes.c_void_p ++ objc.object_getClass.argtypes = [ctypes.c_void_p] ++ objc.object_getClass.restype = ctypes.c_void_p ++ objc.class_respondsToSelector.argtypes = [ctypes.c_void_p, ctypes.c_void_p] ++ objc.class_respondsToSelector.restype = ctypes.c_ubyte ++ objc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] ++ objc.objc_msgSend.restype = ctypes.c_ulonglong ++ objc.objc_release.argtypes = [ctypes.c_void_p] ++ objc.objc_release.restype = None ++ dev = metal.MTLCreateSystemDefaultDevice() ++ if not dev: ++ return None ++ try: ++ sel = objc.sel_registerName(b"recommendedMaxWorkingSetSize") ++ if not objc.class_respondsToSelector(objc.object_getClass(dev), sel): ++ return None ++ return int(objc.objc_msgSend(ctypes.c_void_p(dev), ctypes.c_void_p(sel))) or None ++ finally: ++ objc.objc_release(dev) # MTLDevice.h: MTLCreateSystemDefaultDevice is NS_RETURNS_RETAINED ++ except (OSError, AttributeError, ValueError): ++ return None ++ ++ ++class _Metal: ++ """macOS: an Apple GPU's readings from the IOKit registry (`ioreg -l -w0 -r -c IOAccelerator`, 22 ms), with ++ _Nvml's interface: the load and the memory are the driver's PerformanceStatistics of the index'th accelerator ++ node - Device Utilization % and the in-use system memory in bytes. The total is the device's Metal working set, ++ the budget one app gets of the unified memory. No temperature or power reading exists without root ++ (`powermetrics`), and an Apple Silicon GPU sits on the die, so the PCIe cards have nothing to show.""" ++ note = ("no GPU temperature or power reading without root (powermetrics), and an Apple Silicon GPU sits on the " ++ "die, so there is no PCIe link to show; the engine's own figures are in its log") ++ cmd = ("ioreg", "-l", "-w0", "-r", "-c", "IOAccelerator") ++ ++ def __init__(self, index=0, run=None): ++ self.index = index ++ self._run = run or self._ioreg ++ self.nodes = len(self._blocks()) ++ self.working_set = metal_working_set() ++ self._name = None ++ b = self._block() ++ if b: ++ m = re.search(r'"model" = "([^"]+)"', b) ++ self._name = m.group(1) if m else None ++ ++ @staticmethod ++ def _ioreg(): ++ try: ++ return subprocess.run(list(_Metal.cmd), capture_output=True, text=True, timeout=10).stdout or "" ++ except (OSError, subprocess.SubprocessError): ++ return "" ++ ++ def _blocks(self): ++ """The accelerator nodes, in the registry's order: every node that carries a PerformanceStatistics.""" ++ try: ++ out = self._run() or "" ++ except OSError: # ioreg missing, or a caller's runner that raises ++ return [] ++ return [b for b in re.split(r"^\s*\+-o ", out, flags=re.M) if '"PerformanceStatistics"' in b] ++ ++ def _block(self): ++ b = self._blocks() ++ return b[self.index] if 0 <= self.index < len(b) else None ++ ++ def ok(self): ++ return self.index < self.nodes ++ ++ def name(self): ++ return self._name ++ ++ def read(self): ++ b = self._block() ++ if b is None: ++ return {} ++ ++ def num(key): ++ m = re.search(r'"%s"=([0-9.]+)' % re.escape(key), b) ++ return float(m.group(1)) if m else None ++ ++ out = {} ++ u = num("Device Utilization %") ++ if u is None: # older chips report only the renderer and the tiler ++ r, t = num("Renderer Utilization %"), num("Tiler Utilization %") ++ u = max(v for v in (r, t) if v is not None) if (r is not None or t is not None) else None ++ out["util"] = int(u) if u is not None else None ++ used = num("In use system memory") ++ out["mem_used"] = int(used) if used is not None else None ++ out["mem_total"] = self.working_set ++ return out ++ ++ + def gpu_reader(index=0, amd=False): +- """The card's readings: NVML (NVIDIA), or the amdgpu sysfs files with the AMD backend (#301).""" +- return _Amd(index) if amd else _Nvml(index) ++ """The card's readings: NVML (NVIDIA), the amdgpu sysfs files with the AMD backend (#301), or IOKit's ++ IOAccelerator registry on macOS (Metal backend). `amd` comes first: no Mac runs the HIP backend, and the ++ tests' AMD trees are read on a Mac too.""" ++ if amd: ++ return _Amd(index) ++ return _Metal(index) if sys.platform == "darwin" else _Nvml(index) + + + _DISPLAY_CLASS = r"SYSTEM\CurrentControlSet\Control\Class\{4d36e968-e325-11ce-bfc1-08002be10318}" +@@ -413,15 +522,21 @@ class Telemetry: + cards = windows_amd_cards() + win_amd = [cards[i] for i in idx if 0 <= i < len(cards)] or cards[:1] + gpu_name = " + ".join(c["name"] for c in win_amd) if win_amd else None ++ # the macOS reader knows the device's working set (the VRAM budget of the unified memory); None elsewhere ++ vram_total = None ++ if win_amd is None and isinstance(self.gpu, _Metal) and self.gpu.ok(): ++ ws = self.gpu.working_set ++ vram_total = ws >> 20 if ws else None + self.static = { + "gpu_name": gpu_name, + "gpu_vram_total_mib": (sum(c["vram_bytes"] for c in win_amd) >> 20 +- if win_amd and all(c["vram_bytes"] for c in win_amd) else None), ++ if win_amd and all(c["vram_bytes"] for c in win_amd) else vram_total), + "gpu_count": len(self.gpus), + # #1380: the AMD readings are the amdgpu driver's Linux sysfs files; a Windows AMD card has none yet, and the + # dashboard said "not readable (NVML)" or showed empty tiles with no word why + "gpu_note": ("no GPU load or VRAM readings for AMD cards on Windows yet (Linux reads them from the amdgpu " +- "driver); the engine's own VRAM figures are in its log" if amd and not self.gpu.ok() else None), ++ "driver); the engine's own VRAM figures are in its log" if amd and not self.gpu.ok() ++ else getattr(self.gpu, "note", None)), + "cpu_name": _cpu_name(), + "cores": (self.ps.cpu_count(logical=False) if self.ps else None) or None, + "threads": os.cpu_count(), +diff --git a/serve/test_effort_paths.py b/serve/test_effort_paths.py +index dbc1ae6f..15525f6b 100644 +--- a/serve/test_effort_paths.py ++++ b/serve/test_effort_paths.py +@@ -130,6 +130,7 @@ class EffortPaths(unittest.TestCase): + "thinking": {"type": "enabled", "budget_tokens": 20000}}), "xhigh") + self.assertEqual(self.level("/v1/responses", {"input": "hello", "reasoning": {"effort": "high"}}), "xhigh") + self.assertEqual(self.level("/v1/chat/completions", {"messages": MSG, "reasoning_effort": "none"}), "none") ++ self.assertEqual(self.level("/v1/chat/completions", {"messages": MSG, "reasoning_effort": False}), "none") + + def test_the_chat_settings_shared_with_apps_beat_the_config(self): + self.svc.config_effort = "low" +diff --git a/serve/test_logit_bias.py b/serve/test_logit_bias.py +index ee3751b1..8543a9c2 100644 +--- a/serve/test_logit_bias.py ++++ b/serve/test_logit_bias.py +@@ -59,3 +59,116 @@ class LogitBias(unittest.TestCase): + finally: + httpd.shutdown() + httpd.server_close() ++ ++ ++class LogitBiasCache(unittest.TestCase): ++ """A long list is normalized once; the cache must not change what a call returns or raises.""" ++ ++ def setUp(self): ++ from serve import logit_bias ++ logit_bias._cache.clear() ++ ++ def test_same_answers_as_the_plain_function(self): ++ from serve.logit_bias import _normalize ++ forms = [{str(i): -100 for i in range(200)}, ++ [[i, False] for i in range(200)], ++ {str(i): (i % 7) - 3.5 for i in range(200)}, ++ [[str(i), 2] for i in range(200)]] ++ for vocab in (None, 1000): ++ for form in forms: ++ want = _normalize(form, vocab) ++ for _ in range(3): # the first call fills the cache, the others read it ++ self.assertEqual(normalize(form, vocab), want) ++ self.assertEqual(normalize(json.loads(json.dumps(form)), vocab), want) # an equal list, another object ++ for form in forms: ++ ref = " logit_bias=" + ",".join(f"{i}:{b:g}" for i, b in sorted(_normalize(form).items())) ++ self.assertEqual(engine_key(form), ref) ++ self.assertEqual(engine_key(form), ref) ++ ++ def test_false_is_not_zero(self): ++ ban = [[i, False] for i in range(100)] ++ zero = [[i, 0] for i in range(100)] ++ self.assertTrue(ban == zero) # Python calls them equal, so an == key would mix them up ++ for first, second in [(ban, zero), (zero, ban)]: ++ self.setUp() ++ self.assertEqual(set(normalize(first).values()), {-100.0 if first is ban else 0.0}) ++ self.assertEqual(set(normalize(second).values()), {-100.0 if second is ban else 0.0}) ++ ++ def test_a_bad_list_is_never_stored(self): ++ good = {str(i): -1 for i in range(100)} ++ normalize(good, 1000) ++ for bad in [{**good, "5": True}, {**good, "5": 101}, {**good, "5": float("nan")}, {**good, "5": "x"}]: ++ for _ in range(2): ++ with self.assertRaises(ValueError): ++ normalize(bad, 1000) ++ pairs = [[i, 1] for i in range(100)] ++ normalize(pairs) ++ with self.assertRaises(ValueError): # True == 1, but a bool is not a bias ++ normalize([[i, 1] for i in range(99)] + [[99, True]]) ++ with self.assertRaises(ValueError): # a duplicate ID ++ normalize([[i, 1] for i in range(100)] + [[0, 1]]) ++ ++ def test_vocabulary_size_is_part_of_the_key(self): ++ ids = {str(i): -1 for i in range(100)} ++ self.assertEqual(len(normalize(ids, 1000)), 100) ++ with self.assertRaises(ValueError): ++ normalize(ids, 50) ++ self.assertEqual(len(normalize(ids)), 100) ++ with self.assertRaises(ValueError): ++ normalize(ids, 50) ++ ++ def test_the_result_is_a_copy(self): ++ ids = {str(i): -1 for i in range(100)} ++ first = normalize(ids) ++ first[0] = 99.0 ++ self.assertEqual(normalize(ids)[0], -1.0) ++ self.assertEqual(validate_request({"logit_bias": ids}, SimpleNamespace(info={"logit_bias": 1}), 1000)[0], -1.0) ++ ++ def test_an_equal_list_in_a_later_request_is_a_hit(self): ++ from serve import logit_bias ++ a = {str(i): -100 for i in range(100)} ++ b = json.loads(json.dumps(a)) ++ self.assertIsNot(a, b) ++ key_a = engine_key(a) ++ key_b = engine_key(b) ++ self.assertIs(key_a, key_b) # one entry, one key string ++ self.assertEqual(len(logit_bias._cache), 1) ++ ++ def test_json_equivalent_pairs_and_token_keys(self): ++ pairs = [[i, 0] for i in range(64)] ++ tuples = [(i, 0) for i in range(64)] ++ numeric = {i: 0 for i in range(64)} ++ strings = {str(i): 0 for i in range(64)} ++ for first, second in ((pairs, tuples), (numeric, strings)): ++ self.assertEqual(normalize(first, 1000), normalize(second, 1000)) ++ self.assertEqual(engine_key(first), engine_key(second)) ++ duplicates = {**numeric, "0": 0} ++ for _ in range(2): ++ with self.assertRaisesRegex(ValueError, "duplicate"): ++ normalize(duplicates, 1000) ++ ++ def test_same_length_mutations_are_validated(self): ++ for value in ([[i, 0] for i in range(64)], {str(i): 0 for i in range(64)}): ++ with self.subTest(form=type(value).__name__): ++ self.assertEqual(normalize(value, 1000)[0], 0) ++ engine_key(value) ++ if isinstance(value, list): ++ value[0][1] = -100 ++ else: ++ value["0"] = -100 ++ self.assertEqual(normalize(value, 1000)[0], -100) ++ self.assertIn("0:-100", engine_key(value)) ++ if isinstance(value, list): ++ value[0][1] = True ++ else: ++ value["0"] = True ++ with self.assertRaises(ValueError): ++ normalize(value, 1000) ++ with self.assertRaises(ValueError): ++ engine_key(value) ++ ++ def test_the_cache_is_bounded(self): ++ from serve import logit_bias ++ for n in range(logit_bias._CACHE_MAX + 3): ++ normalize({str(i): -(n + 1) for i in range(100)}) ++ self.assertEqual(len(logit_bias._cache), logit_bias._CACHE_MAX) +diff --git a/serve/test_parallel.py b/serve/test_parallel.py +index 1a638919..bde52084 100644 +--- a/serve/test_parallel.py ++++ b/serve/test_parallel.py +@@ -248,6 +248,34 @@ class PickSlot(unittest.TestCase): + e.slot_busy = [False] * 3 + self.assertEqual(e.pick_slot([1, 2, 3]), 1) # a held prompt is never the WHOLE prompt + ++ def test_a_rewritten_next_turn_goes_to_the_slot_of_its_chat(self): ++ # #2518: two chats in two slots; each next turn renders the history again but not the reply it decoded or the ++ # assistant header, so the held stream (prompt + reply) is no start of it. The engine's turn checkpoint sits at ++ # the last <|im_start|> of the old prompt: the prompt that agrees up to there belongs to that slot. Before, both ++ # went to the least recently used slot in turn, and the first one overwrote the other chat's cache. ++ T = StrataEngine.TURN_ID ++ hist_a, hist_b = list(range(1000, 1600)), list(range(2000, 2600)) ++ held_a = hist_a + [T, 7, 8, 9] + [50, 51, 52] # prompt (ends with the assistant header) + reply ++ held_b = hist_b + [T, 7, 8, 9] + [60, 61] ++ e = self.engine(2) ++ e.slot_held = [held_a, held_b] ++ e.slot_used = [10.0, 20.0] # slot 0 was used longer ago ++ next_a = hist_a + [T, 5, 6, 7, 8, 9] + list(range(3000, 3100)) # the history again, then new turns ++ next_b = hist_b + [T, 5, 6, 7, 8, 9] + list(range(3000, 3100)) ++ self.assertEqual(e.pick_slot(next_b), 1) # the LRU slot 0 holds the other chat ++ self.assertEqual(e.pick_slot(next_a), 0) ++ e.slot_busy[0] = True # a busy slot is never picked ++ self.assertEqual(e.pick_slot(next_a), 1) ++ e.slot_busy[0] = False ++ # a chat that agrees with neither (a new one, or one that rewrote its history) keeps the old order ++ self.assertEqual(e.pick_slot(list(range(5000, 5200))), 0) ++ self.assertEqual(e.pick_slot(hist_a[:300] + [T, 5, 6] + list(range(3000, 3100))), 0) ++ e.slot_held[0] = [] # an empty slot still comes before evicting one ++ self.assertEqual(e.pick_slot(list(range(5000, 5200))), 0) ++ # the whole held stream continued is still the longer match ++ e.slot_held = [held_a, held_b + [T, 1, 2]] ++ self.assertEqual(e.pick_slot(held_b + [T, 1, 2, 3, 4]), 1) ++ + def test_batch_groups_spread_requests_over_the_groups(self): + # #1249: with --batch-groups 2 over 4 slots (groups {0,1} and {2,3}) a second request goes to the other group, + # and a short held prefix does not pull it back into the busy one +diff --git a/serve/test_reasoning_loop_recovery.py b/serve/test_reasoning_loop_recovery.py +index b53ee56a..c00fc1be 100644 +--- a/serve/test_reasoning_loop_recovery.py ++++ b/serve/test_reasoning_loop_recovery.py +@@ -58,6 +58,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + engine = RecordingEngine(self.tok, [self.loop, + 'I have verified the solution.\n\nComplete solution.']) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = enabled + events = list(service.run(self.ids, True, [], 40000, {}, threading.Event())) + return engine, events +@@ -66,6 +67,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + engine = RecordingEngine(self.tok, [self.loop, + 'Verified.\n\nComplete solution.']) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + request = {'temperature': .2, 'presence_penalty': 0, 'top_p': .95, 'top_k': 64, 'seed': 2} + list(service.run(self.ids, True, [], 40000, request, threading.Event())) +@@ -81,6 +83,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_frozen_cohort_sampler_is_unchanged_during_recovery(self): + engine = RecordingEngine(self.tok, [self.loop, 'Done.\n\nComplete solution.']) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + sampler = {'temperature': 1.0, 'top_p': .95, 'top_k': 20, 'presence_penalty': 1.5, 'seed': 2} + list(service.run(self.ids, True, [], 40000, sampler, threading.Event())) +@@ -91,6 +94,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_a_lower_top_p_and_top_k_are_never_raised_or_lowered(self): + engine = RecordingEngine(self.tok, [self.loop, 'Verified.\n\nComplete solution.']) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = "recover" + list(service.run(self.ids, True, [], 40000, {'temperature': 0, 'top_p': .5, 'top_k': 5}, threading.Event())) + self.assertEqual((engine.sampling_calls[1]['top_p'], engine.sampling_calls[1]['top_k']), (.5, 5)) +@@ -110,6 +114,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_stop_mode_ends_the_reply_at_the_loop(self): + engine = RecordingEngine(self.tok, [self.loop]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = "stop" + events = list(service.run(self.ids, True, [], 40000, {}, threading.Event())) + self.assertEqual(len(engine.calls), 1) # no second pass +@@ -121,6 +126,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + """#1753: the word count never sees the digits, so only the tail's period ends the reply.""" + engine = RecordingEngine(self.tok, ['9^104 = ' + '58747693387' * 3000]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = "stop" + events = list(service.run(self.ids, True, [], 40000, {}, threading.Event())) + self.assertEqual(len(engine.calls), 1) # no second pass +@@ -136,6 +142,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_default_single_token_guard_does_not_detect_passage_loop(self): + engine = RecordingEngine(self.tok, [self.loop]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + events = list(service.run(self.ids, True, [], 16384, {}, threading.Event())) + self.assertEqual(service.repeat_stop_tokens, 256) + self.assertEqual(events[-1][1]['finish'], 'length') +@@ -145,6 +152,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_recovery_is_once_and_never_extends_output_limit(self): + engine = RecordingEngine(self.tok, [self.loop, self.loop]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + events = list(service.run(self.ids, True, [], 16384, {}, threading.Event())) + self.assertEqual(len(engine.calls), 2) +@@ -155,6 +163,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_repeated_final_text_does_not_change_sampler_or_prompt(self): + engine = RecordingEngine(self.tok, [self.loop]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + events = list(service.run(self.ids, False, [], 40000, {}, threading.Event())) + self.assertEqual(len(engine.calls), 1) +@@ -163,6 +172,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_missing_native_high_hint_skips_recovery(self): + engine = RecordingEngine(self.tok, [self.loop]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + ids = self.tok.encode('<|im_start|>user\nTask<|im_end|>\n', parse_special=True) + events = list(service.run(ids, True, [], 40000, {}, threading.Event())) +@@ -172,6 +182,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_cancel_during_reasoning_does_not_start_recovery(self): + engine = RecordingEngine(self.tok, [self.loop, 'Unused']) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + cancel = threading.Event() + events = [] +@@ -187,6 +198,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + engine = RecordingEngine(self.tok, [self.loop, 'Done.\n\nAnswer.']) + engine.batch = 2 + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + events = list(service.run(self.ids, True, [], 40000, {}, threading.Event())) + self.assertEqual(events[-1][1]['reasoning_recoveries'], 1) +diff --git a/serve/test_server.py b/serve/test_server.py +index 3a49ff25..5f58d09d 100644 +--- a/serve/test_server.py ++++ b/serve/test_server.py +@@ -26,7 +26,7 @@ from serve.server import (CTX_SLACK, ByteTokenizer, EngineDied, EngineNonFinite, + PP_DONE_TAIL, Service, + StrataEngine, api_key_of, engine_args, key_matches, layer_split_value, prompt_progress, + prompt_tokens_seen, +- request_timings, serve, start_failure_hint) ++ request_stats, request_timings, serve, start_failure_hint) + from types import SimpleNamespace # noqa: E402 + + ROOT = Path(__file__).resolve().parents[1] +@@ -1548,6 +1548,13 @@ class SamplingKeys(unittest.TestCase): + bad = self.keys(strata_tune={"pcie_frac": 3, "spec_min_p": True, "pool_workers": 2}) + self.assertFalse([x for x in bad if x.split("=")[0] in ("pcie_frac", "spec_min_p", "pool_workers")]) + ++ def test_tune_prefill_pipe(self): ++ self.assertIn("pipe=1536.0", self.keys(temperature=0, strata_tune={"prefill_pipe": 1536})) ++ self.assertIn("pipe=0.0", self.keys(temperature=0, strata_tune={"prefill_pipe": 0})) ++ for bad in (-1, True, "800", 1e9): ++ self.assertFalse([x for x in self.keys(temperature=0, strata_tune={"prefill_pipe": bad}) ++ if x.startswith("pipe=")], bad) ++ + def test_tune_aux_cpus(self): + for val, want in ((True, "aux_cpus=1"), (1, "aux_cpus=1"), (False, "aux_cpus=0"), (0, "aux_cpus=0"), + (1.0, "aux_cpus=1"), (0.0, "aux_cpus=0")): +@@ -1771,6 +1778,150 @@ class RecordingPrompt(MockEngine): + yield from super().generate(ids, max_new, sampling, cancel, embeddings) + + ++class IncrementalPrompts(unittest.TestCase): ++ """The prompt encoder: every request's ids are those of a full encode, and a turn reuses the previous one's.""" ++ ++ @classmethod ++ def setUpClass(cls): ++ tok = ByteTokenizer() ++ cls.engine = RecordingPrompt(tok, "Thinking.\n\n\nThe answer.", max_context=CTX) ++ cls.svc = Service(cls.engine, tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) ++ cls.httpd = serve(cls.svc, port=0) ++ cls.base = f"http://127.0.0.1:{cls.httpd.server_address[1]}" ++ ++ @classmethod ++ def tearDownClass(cls): ++ cls.httpd.shutdown() ++ cls.httpd.server_close() ++ ++ post = ClientShapes.post ++ ++ def test_turns(self): ++ self.assertIsNotNone(self.svc.prompts) ++ full = Service(self.engine, self.svc.tok, self.svc.template) ++ full.prompts = None # the reference: the same prompt, encoded in full ++ msgs = [{"role": "system", "content": "Be brief."}] ++ for turn in range(6): ++ msgs.append({"role": "user", "content": f"question {turn} <|im_end|> é你 " * (turn + 1)}) ++ status, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 64, "messages": msgs}) ++ self.assertEqual(status, 200, b) ++ prompt = self.svc.template.render(msgs) ++ self.assertEqual(self.engine.last_ids, full.encode_prompt(msgs, None, {})) ++ if turn: ++ self.assertGreater(self.svc.prompts.last_reused, len(prompt) // 3) ++ msgs.append({"role": "assistant", "content": b["choices"][0]["message"]["content"]}) ++ ++ def conversation(self, n): ++ """Messages of the n-th conversation: distinct from every other one from the number on.""" ++ return [{"role": "user", "content": f"conversation {n} says: " + f"topic {n} sentence {n} " * 20}] ++ ++ def service(self, **info): ++ engine = RecordingPrompt(self.svc.tok, "The answer.", max_context=CTX) ++ engine.info = info ++ return Service(engine, self.svc.tok, self.svc.template) ++ ++ def test_the_kept_prompts_follow_the_engines_conversation_slots(self): ++ svc = self.service(conversation_cache_slots=8) ++ texts = [] ++ for n in range(8): ++ msgs = self.conversation(n) ++ svc.encode_prompt(msgs, None, {}) ++ texts.append(svc.render_prompt(msgs, None, {})) ++ self.assertEqual(len(svc.prompts.entries), 8) ++ self.assertEqual(sorted(e[0] for e in svc.prompts.entries), sorted(texts)) ++ msgs = self.conversation(8) ++ svc.encode_prompt(msgs, None, {}) ++ kept = sorted(e[0] for e in svc.prompts.entries) ++ self.assertEqual(len(kept), 8) ++ self.assertEqual(kept, sorted(texts[1:] + [svc.render_prompt(msgs, None, {})])) ++ ++ def lazy_service(self, args, **info): ++ """A Service over the real StrataEngine class started lazily: no process yet, so no INFO line either.""" ++ engine = StrataEngine("strata", args, lazy=True) ++ engine.info.update(info) ++ return Service(engine, self.svc.tok, self.svc.template) ++ ++ def test_a_lazy_engine_keeps_the_prompts_its_arguments_ask_for(self): ++ svc = self.lazy_service(["--conversation-cache-slots", "8"]) ++ for n in range(10): ++ svc.encode_prompt(self.conversation(n), None, {}) ++ self.assertEqual(len(svc.prompts.entries), 8) ++ ++ def test_the_slot_argument_wins_over_the_info_line_and_a_bad_one_falls_back(self): ++ self.assertEqual(self.lazy_service(["--conversation-cache-slots", "6"], conversation_cache_slots=3) ++ .prompts.keep, 6) ++ self.assertEqual(self.lazy_service(["--conversation-cache-slots", "x"], conversation_cache_slots=3) ++ .prompts.keep, 3) ++ self.assertEqual(self.lazy_service(["--conversation-cache-slots"]).prompts.keep, 4) ++ self.assertEqual(self.lazy_service([]).prompts.keep, 4) ++ ++ def test_prompt_encoder_slot_count_is_capped(self): ++ from serve.server import PROMPT_ENCODER_KEEP_MAX ++ self.assertEqual(PROMPT_ENCODER_KEEP_MAX, 16) ++ self.assertEqual(self.service(conversation_cache_slots=10**9).prompts.keep, 16) ++ self.assertEqual(self.lazy_service(["--conversation-cache-slots", str(10**9)]).prompts.keep, 16) ++ ++ def test_without_a_slot_setting_four_prompts_are_kept(self): ++ svc = self.service() ++ for n in range(6): ++ svc.encode_prompt(self.conversation(n), None, {}) ++ self.assertEqual(len(svc.prompts.entries), 4) ++ ++ def test_counting_tokens_leaves_the_kept_prompts_alone(self): ++ svc = self.service(conversation_cache_slots=2) ++ httpd = serve(svc, port=0) ++ self.addCleanup(httpd.server_close) ++ self.addCleanup(httpd.shutdown) ++ self.base = f"http://127.0.0.1:{httpd.server_address[1]}" ++ main = self.conversation(0) ++ status, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 8, "messages": main}) ++ self.assertEqual(status, 200, b) ++ before = list(svc.prompts.entries) ++ full = self.service() ++ full.prompts = None ++ for n in range(1, 6): ++ other = self.conversation(n) ++ status, b = self.post("/v1/messages/count_tokens", {"model": "m", "messages": other}) ++ self.assertEqual(status, 200, b) ++ self.assertEqual(b["input_tokens"], len(full.encode_prompt(other, None, {}))) ++ self.assertTrue(svc.prompts.entries == before, "the count wrote into the kept prompts") ++ main.append({"role": "assistant", "content": "ok"}) ++ main.append({"role": "user", "content": "and then?"}) ++ svc.encode_prompt(main, None, {}) ++ self.assertGreater(svc.prompts.last_reused, len(svc.render_prompt(main, None, {})) // 2) ++ ++ def test_same_ids_as_a_full_encode_on_varied_conversations(self): ++ import random ++ sys.path.insert(0, str(ROOT / "tools")) ++ from test_strata_tokenizer import conversation_prompts, load_tokenizer ++ toks = [("byte", ByteTokenizer())] ++ qwen = load_tokenizer() ++ if qwen is not None: ++ toks.append(("qwen35", qwen)) ++ for name, tok in toks: ++ svc = Service(self.engine, tok, self.svc.template) ++ for seed in range(3): ++ for what, prompt in conversation_prompts(self.svc.template, random.Random(seed)): ++ with self.subTest(tokenizer=name, seed=seed, what=what): ++ self.assertEqual(svc.prompts.encode(prompt), tok.encode(prompt, parse_special=True)) ++ ++ def test_a_tokenizer_without_resume_points_encodes_in_full(self): ++ class Plain: ++ """ByteTokenizer without the resume points.""" ++ def __init__(self): ++ self.base = ByteTokenizer() ++ ++ def __getattr__(self, name): ++ if name in ("encode_marked", "max_special_len"): ++ raise AttributeError(name) ++ return getattr(self.base, name) ++ svc = Service(self.engine, Plain(), self.svc.template) ++ self.assertIsNone(svc.prompts) ++ msgs = [{"role": "user", "content": "hi"}] ++ self.assertEqual(svc.encode_prompt(msgs, None, {}), ++ ByteTokenizer().encode(svc.render_prompt(msgs, None, {}), parse_special=True)) ++ ++ + class DyingEngine(MockEngine): + """Issue #27: an engine that dies after a few tokens of its first answer, and comes back when restarted.""" + +@@ -1989,6 +2140,39 @@ class PrometheusMetrics(unittest.TestCase): + self.assertEqual((m["strata:requests_total"], m["strata:drafts_offered_total"], + m["strata:drafts_accepted_total"]), (2, 17, 10)) + ++ def test_throughput_gauges_cover_the_window_between_scrapes(self): ++ """#2341: llama.cpp reads these two from a bucket it resets on every scrape, so a scrape that saw no ++ traffic reports 0 - not the lifetime average, which froze at the last request's speed.""" ++ tok = ByteTokenizer() ++ engine = DoneLineEngine(tok, "\n\nok", max_context=CTX, ++ done_lines=["DONE 4 20 400.0 200.0 stop 7 12 5"]) ++ svc = Service(engine, tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) ++ httpd = serve(svc, port=0) ++ base = f"http://127.0.0.1:{httpd.server_address[1]}" ++ ++ def scrape(): ++ with urllib.request.urlopen(base + "/metrics/prometheus", timeout=10) as r: ++ return {line.split()[0]: float(line.split()[1]) ++ for line in r.read().decode().splitlines() if not line.startswith("#")} ++ ++ try: ++ body = json.dumps({"model": "m", "max_tokens": 10, "messages": [{"role": "user", "content": "hi"}]}) ++ with urllib.request.urlopen(urllib.request.Request(base + "/v1/chat/completions", data=body.encode(), ++ headers={"Content-Type": "application/json"}), ++ timeout=30) as r: ++ self.assertEqual(r.status, 200) ++ first = scrape() ++ self.assertGreater(first["llamacpp:prompt_tokens_seconds"], 0) ++ self.assertGreater(first["llamacpp:predicted_tokens_seconds"], 0) ++ idle = scrape() # nothing ran since the scrape above, so the window is empty ++ self.assertEqual((idle["llamacpp:prompt_tokens_seconds"], ++ idle["llamacpp:predicted_tokens_seconds"]), (0, 0)) ++ self.assertEqual(idle["llamacpp:tokens_predicted_total"], # counters stay cumulative ++ first["llamacpp:tokens_predicted_total"]) ++ finally: ++ httpd.shutdown() ++ httpd.server_close() ++ + + class PcieShare(unittest.TestCase): + """#588: the hit rate stays the VRAM share of the lookups; the routed experts the GPU read over PCIe (the DONE +@@ -2507,6 +2691,30 @@ class CancelledRead(unittest.TestCase): + StrataEngine._parse_done(e, "DONE 0 98179 17565.0 0.0 cancel 0 0 0 0 0 0 0 0.0") + self.assertNotIn("prompt_read", e.last) + ++ def test_parse_done_chunk_fields(self): ++ """The prompt's first chunk and its streamed expert blobs come as key=value fields after the positional ones, ++ with or without the five --lookup-chain fields; an older line has neither.""" ++ e = SimpleNamespace() ++ StrataEngine._parse_done(e, "DONE 5 12000 900.0 80.0 stop 3 4 0 7 9 0 0 0.0 12000 2 chunk=1536 " ++ "experts_streamed=24576") ++ self.assertEqual((e.last["chunk"], e.last["experts_streamed"], e.last["offloaded"]), (1536, 24576, 2)) ++ StrataEngine._parse_done(e, "DONE 5 12000 900.0 80.0 stop 3 4 0 7 9 0 0 0.0 12000 2 1 2 3 4 5 chunk=8192 " ++ "experts_streamed=0") ++ self.assertEqual((e.last["chunk"], e.last["experts_streamed"]), (8192, 0)) ++ StrataEngine._parse_done(e, "DONE 5 12000 900.0 80.0 stop 3 4 0 7 9 0 0 0.0 12000 2") ++ self.assertNotIn("chunk", e.last) ++ self.assertNotIn("experts_streamed", e.last) ++ StrataEngine._parse_done(e, "DONE 5 12000 900.0 80.0 stop 3 4 0 7 9 0 0 0.0 12000 2 chunk=x other=1") ++ self.assertNotIn("chunk", e.last) ++ self.assertNotIn("other", e.last) ++ ++ def test_request_stats_chunk_fields(self): ++ """Across a request's continuations the streamed blobs add up; the chunk is the request's own read's.""" ++ stats = request_stats([{"prompt_ms": 10.0, "chunk": 1536, "experts_streamed": 100}, ++ {"prompt_ms": 2.0, "chunk": 64, "experts_streamed": 5}]) ++ self.assertEqual((stats["chunk"], stats["experts_streamed"], stats["prompt_ms"]), (1536, 105, 12.0)) ++ self.assertNotIn("chunk", request_stats([{"prompt_ms": 1.0}, {"prompt_ms": 1.0, "chunk": 64}])) ++ + def test_prompt_tokens_seen(self): + last = {"finish": "cancel", "reused": 1000, "prompt_read": 2000} + self.assertEqual(prompt_tokens_seen(98179, last), 3000) +@@ -2822,6 +3030,46 @@ class ConfigReasoningEffort(unittest.TestCase): + self.assertEqual(cfg, {"args": ["--ctx", "4096"]}) + + ++class JoinedEngineArgs(unittest.TestCase): ++ """#2267: "--ple-io mmap" written as one element of the config's args reached the engine as one argument.""" ++ ++ def split(self, args): ++ from serve.server import split_joined_args ++ cfg = {"args": list(args)} ++ with mock.patch("builtins.print") as pr: ++ split_joined_args(cfg) ++ return cfg["args"], pr.called ++ ++ def test_one_element_becomes_flag_and_value(self): ++ self.assertEqual(self.split(["--spec", "4", "--ple-io mmap"]), (["--spec", "4", "--ple-io", "mmap"], True)) ++ self.assertEqual(self.split(["--ple-io=mmap", "--kv int8"]), (["--ple-io", "mmap", "--kv", "int8"], True)) ++ ++ def test_proper_args_are_untouched(self): ++ args = ["--ple-io", "mmap", "--native", "/data/My Models/a b.gguf", "--max-context", "131072", "--vision"] ++ self.assertEqual(self.split(args), (args, False)) ++ # a value in a flag's value place stays whole, whatever it looks like ++ args = ["--system-prompt", "--be brief please"] ++ self.assertEqual(self.split(args), (args, False)) ++ ++ def test_no_args(self): ++ from serve.server import split_joined_args ++ cfg = {} ++ split_joined_args(cfg) ++ self.assertEqual(cfg, {}) ++ ++ def test_bad_json_names_the_place(self): ++ import tempfile ++ from serve.server import read_config ++ with tempfile.TemporaryDirectory() as d: ++ p = Path(d) / "c.json" ++ p.write_text('{"args": ["--ple-io" "mmap"]}', encoding="utf-8") ++ with self.assertRaises(SystemExit) as cm: ++ read_config(p) ++ self.assertIn("line 1", str(cm.exception)) ++ p.write_text('{"args": ["--ple-io", "mmap"]}', encoding="utf-8") ++ self.assertEqual(read_config(p), {"args": ["--ple-io", "mmap"]}) ++ ++ + class SharedSettings(unittest.TestCase): + """The web app's "Use for other apps too": POST /settings makes its Chat settings every client's defaults.""" + +@@ -3776,6 +4024,45 @@ class ThinkingBudget(unittest.TestCase): + self.assertEqual(text, ThinkingEngine.ANSWER) + self.assertEqual(evs[-2]["delta"]["stop_reason"], "end_turn") + ++ def test_a_cap_that_the_thinking_would_eat_leaves_room_to_answer(self): ++ """0.1.45: max_tokens counts the thinking; a thinking that would fill it gets closed with a quarter left (the ++ agent harness sends max_completion_tokens 8192; the pagoda's thinking took all of it and no answer came).""" ++ self.engine.THOUGHT = "Let me think step by step about two plus two. " * 40 # 1,800 reasoning tokens ++ code, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 1024, ++ "messages": [{"role": "user", "content": "2+2?"}]}) ++ self.assertEqual(code, 200, b) ++ msg = b["choices"][0]["message"] ++ self.assertEqual(msg["content"], ThinkingEngine.ANSWER) ++ self.assertEqual(b["choices"][0]["finish_reason"], "stop") ++ self.assertLess(len(msg["reasoning_content"]), 1024 - 256) ++ self.assertEqual(len(self.engine.prompts), 2) # the wrap-up continuation ++ ++ def test_the_answer_room_has_its_limits(self): ++ self.engine.THOUGHT = "Let me think step by step about two plus two. " * 40 ++ code, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 300, # below the floor: as before ++ "messages": [{"role": "user", "content": "2+2?"}]}) ++ self.assertEqual(b["choices"][0]["finish_reason"], "length") ++ self.assertFalse(b["choices"][0]["message"]["content"]) ++ self.svc.answer_room = False # the config's "answer_room": false ++ code, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 1024, ++ "messages": [{"role": "user", "content": "2+2?"}]}) ++ self.assertEqual(b["choices"][0]["finish_reason"], "length") ++ self.svc.answer_room = True ++ code, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 1024, "reasoning_effort": "none", ++ "messages": [{"role": "user", "content": "2+2?"}]}) ++ self.assertEqual(len(self.engine.prompts), 3) # thinking off: no budget, one prompt each ++ ++ def test_anthropic_thinking_budget_tokens_is_a_budget(self): ++ """0.1.45: thinking {type: enabled, budget_tokens: N} closes the thinking after N tokens (it was only a level).""" ++ from serve.server import REASONING_WRAP_UP ++ code, b = self.post("/v1/messages", {"model": "m", "max_tokens": 400, "thinking": {"type": "enabled", "budget_tokens": 30}, ++ "messages": [{"role": "user", "content": "2+2?"}]}) ++ self.assertEqual(code, 200, b) ++ thinking = "".join(c.get("thinking", "") for c in b["content"] if c["type"] == "thinking") ++ self.assertEqual(thinking, ThinkingEngine.THOUGHT[:30] + REASONING_WRAP_UP.split("")[0]) ++ self.assertEqual("".join(c.get("text", "") for c in b["content"] if c["type"] == "text"), ThinkingEngine.ANSWER) ++ self.assertEqual(len(self.engine.prompts), 2) ++ + def test_a_budget_the_thinking_stays_under(self): + code, b = self.openai(reasoning_budget_tokens=10_000) + self.assertEqual(b["choices"][0]["message"]["reasoning_content"], ThinkingEngine.THOUGHT) +@@ -4481,6 +4768,114 @@ class AmdTelemetry(unittest.TestCase): + self.assertEqual(svc.free_vram_mib(), 26 << 10) + + ++IOREG_FIXTURE = ''' ++ +-o AGXAcceleratorG13X ++ { ++ "IOClass" = "AGXAcceleratorG13X" ++ "model" = "Apple M1 Max" ++ "PerformanceStatistics" = {"In use system memory (driver)"=0,"Alloc system memory"=22991831040,"Tiler Utilization %"=36,"recoveryCount"=44,"lastRecoveryTime"=216970984728,"Renderer Utilization %"=35,"TiledSceneBytes"=1671168,"Device Utilization %"=76,"SplitSceneCount"=0,"Allocated PB Size"=76152832,"In use system memory"=13384105984} ++ +-o IOUserClientCreator ++ { ++ "IOUserClientCreator" = "pid 402, WindowServer" ++ } ++ } ++ ++ +-o AGXAccelerator ++ { ++ "IOClass" = "AGXAccelerator" ++ "model" = "Apple M2 Ultra" ++ "PerformanceStatistics" = {"Renderer Utilization %"=3,"Tiler Utilization %"=2,"In use system memory"=1234567890} ++ } ++''' ++ ++WORKING_SET = 26800603136 # bytes; the M1 Max's device in the fixture above ++ ++ ++class MacTelemetry(unittest.TestCase): ++ """The macOS (Metal backend) reader: IOKit's IOAccelerator nodes and the device's Metal working set.""" ++ ++ def _reader(self, index=0, text=IOREG_FIXTURE): ++ from serve import telemetry ++ with mock.patch.object(telemetry, "metal_working_set", return_value=WORKING_SET): ++ return telemetry._Metal(index, run=lambda: text) ++ ++ def test_readings(self): ++ g = self._reader() ++ self.assertTrue(g.ok()) ++ self.assertEqual(g.name(), "Apple M1 Max") ++ with mock.patch("serve.telemetry.metal_working_set", lambda: WORKING_SET): ++ self.assertEqual(g.read(), {"util": 76, "mem_used": 13384105984, "mem_total": WORKING_SET}) ++ # the second accelerator is index 1: the older chips' renderer/tiler-only utilization is the fallback ++ g1 = self._reader(1) ++ self.assertTrue(g1.ok()) ++ with mock.patch("serve.telemetry.metal_working_set", lambda: WORKING_SET): ++ self.assertEqual(g1.read()["util"], 3) # the renderer, the higher of the two ++ self.assertEqual(g1.name(), "Apple M2 Ultra") ++ self.assertFalse(self._reader(5).ok()) ++ ++ def test_working_set_is_read_once_for_samples_and_static(self): ++ from serve import telemetry ++ with mock.patch.object(telemetry, "metal_working_set", return_value=WORKING_SET) as budget: ++ g = telemetry._Metal(run=lambda: IOREG_FIXTURE) ++ for _ in range(3): ++ self.assertEqual(g.read()["mem_total"], WORKING_SET) ++ with mock.patch.object(telemetry, "gpu_reader", return_value=g): ++ t = telemetry.Telemetry() ++ try: ++ self.assertEqual(t.static["gpu_vram_total_mib"], WORKING_SET >> 20) ++ finally: ++ t.close() ++ budget.assert_called_once_with() ++ ++ def test_device_is_released_on_success_and_unsupported_selector(self): ++ from serve import telemetry ++ metal, objc = mock.MagicMock(), mock.MagicMock() ++ metal.MTLCreateSystemDefaultDevice.return_value = 1234 ++ objc.sel_registerName.return_value = 5678 ++ objc.objc_msgSend.return_value = WORKING_SET ++ for supported in (True, False): ++ with self.subTest(supported=supported): ++ objc.objc_release.reset_mock() ++ objc.class_respondsToSelector.return_value = supported ++ with mock.patch.object(telemetry.sys, "platform", "darwin"), \ ++ mock.patch.object(telemetry.ctypes, "CDLL", side_effect=[metal, objc]): ++ self.assertEqual(telemetry.metal_working_set(), WORKING_SET if supported else None) ++ objc.objc_release.assert_called_once_with(1234) ++ ++ def test_no_registry(self): ++ g = self._reader(text="") ++ self.assertFalse(g.ok()) ++ self.assertEqual(g.read(), {}) ++ self.assertIsNone(g.name()) ++ ++ def test_bad_process(self): ++ from serve import telemetry ++ g = telemetry._Metal(0, run=lambda: (_ for _ in ()).throw(OSError("no ioreg"))) ++ self.assertFalse(g.ok()) ++ self.assertEqual(g.read(), {}) ++ ++ def test_note_reaches_the_monitor(self): ++ from serve import telemetry ++ t = telemetry.Telemetry(extra=lambda: {"tok_s": 1.0}) ++ try: ++ with mock.patch.object(telemetry, "gpu_reader", lambda index=0, amd=False: self._reader()): ++ t2 = telemetry.Telemetry(extra=lambda: {"tok_s": 1.0}) ++ try: ++ self.assertEqual(t2.static["gpu_name"], "Apple M1 Max") ++ self.assertIn("powermetrics", t2.static["gpu_note"]) ++ with mock.patch("serve.telemetry.metal_working_set", lambda: WORKING_SET): ++ s = t2.sample() ++ self.assertEqual(s["gpu_util"], 76) ++ self.assertEqual(s["gpu_mem_used"], 13384105984) ++ self.assertEqual(s["gpu_mem_total"], WORKING_SET) ++ self.assertIsNone(s.get("gpu_temp")) # no reading without root; the cards show a dash ++ self.assertIsNone(s.get("gpu_pcie_gen")) ++ finally: ++ t2.close() ++ finally: ++ t.close() ++ ++ + class TelemetryNoDisks(unittest.TestCase): + """psutil.disk_io_counters() returns None where it finds no disk (a gVisor container; Windows with its disk counters + off): the sampler goes on without a disk reading instead of its thread ending at the first sample, which left +diff --git a/serve/test_slot_alloc.py b/serve/test_slot_alloc.py +index 474b4f78..b2b4d35b 100644 +--- a/serve/test_slot_alloc.py ++++ b/serve/test_slot_alloc.py +@@ -79,5 +79,30 @@ class OneGroup(unittest.TestCase): + self.assertEqual(e.pick_slot([1, 2, 3]), 3) # else the least recently used + + ++class OneSlotGroups(unittest.TestCase): ++ """--batch 2 on a two-GPU layer split: two pipelined groups of ONE slot. No slot sits below another in its group, ++ so no pad row writes an idle slot and each one keeps its chat (#2518): the next turn goes back to it.""" ++ ++ def test_the_slot_of_the_chat(self): ++ T = StrataEngine.TURN_ID ++ e = engine(2, 2) ++ self.assertEqual(e.slot_gs, 1) ++ hist_a, hist_b = list(range(1000, 1700)), list(range(2000, 2700)) ++ e.slot_held = [hist_a + [T, 7, 8] + [50, 51], hist_b + [T, 7, 8] + [60]] ++ e.slot_used = [10.0, 20.0] ++ again = lambda h: h + [T, 5, 6] + list(range(3000, 3050)) ++ self.assertEqual(e.pick_slot(again(hist_b)), 1) ++ self.assertEqual(e.pick_slot(again(hist_a)), 0) ++ self.assertEqual(e.pick_slot(list(range(5000, 5100))), 0) # a new chat: the least recently used one ++ ++ def test_groups_of_several_slots_still_ignore_held_tokens(self): ++ T = StrataEngine.TURN_ID ++ e = engine(4, 2) ++ self.assertEqual(e.slot_gs, 2) ++ hist = list(range(1000, 1700)) ++ e.slot_held[3] = hist + [T, 7, 8] + [50] ++ self.assertEqual(e.pick_slot(hist + [T, 5, 6] + list(range(3000, 3050))), 0) ++ ++ + if __name__ == "__main__": + unittest.main() +diff --git a/serve/web/app.js b/serve/web/app.js +index d3ad404f..4bc2a46f 100644 +--- a/serve/web/app.js ++++ b/serve/web/app.js +@@ -560,8 +560,18 @@ function markdown(text) { + } + + // ------------------------------------------------------------------ Chat +-const DEFAULTS = {thinking: "high", temperature: 0.6, top_p: 0.95, top_k: 20, max: "", seed: "", show: true, esp: true, mcp: true}; ++// The sampling shown until the user changes it is the model card's thinking set. Untouched, the page sends NO sampling ++// fields: the server picks the card's set for the request's thinking state (thinking 1.0 / 0.95 / 20, not thinking ++// 0.7 / 0.8 / 20 with a presence penalty). Before 0.1.45 the page sent 0.6 itself; a stored 0.6 with the other two ++// values untouched counts as "not set" too. ++const DEFAULTS = {thinking: "high", temperature: 1.0, top_p: 0.95, top_k: 20, max: "", seed: "", show: true, esp: true, mcp: true}; ++const OLD_DEFAULT_TEMPERATURE = 0.6; ++function samplingUntouched(s) { ++ const t = +s.temperature; ++ return (t === DEFAULTS.temperature || t === OLD_DEFAULT_TEMPERATURE) && +s.top_p === DEFAULTS.top_p && +s.top_k === DEFAULTS.top_k; ++} + let settings = {...DEFAULTS, ...store.get("sampling", {})}; ++if (samplingUntouched(settings)) settings.temperature = DEFAULTS.temperature; + let messages = store.get("chat", []); + let attachments = []; // {name, url} + let busy = null; // {controller, msg} +@@ -790,10 +800,10 @@ async function send() { + + const body = {model: health.model, messages: apiMessages(), stream: true, + reasoning_effort: settings.thinking}; +- if (settings.temperature > 0) { +- Object.assign(body, {temperature: +settings.temperature, top_p: +settings.top_p, top_k: +settings.top_k}); +- } else { ++ if (!(settings.temperature > 0)) { + body.temperature = 0; ++ } else if (!samplingUntouched(settings)) { ++ Object.assign(body, {temperature: +settings.temperature, top_p: +settings.top_p, top_k: +settings.top_k}); + } + if (settings.seed) body.seed = +settings.seed; + if (settings.max) body.max_tokens = +settings.max; +@@ -1011,8 +1021,9 @@ async function loadShared() { + $("s-share").setAttribute("aria-checked", String(sharedOn)); + } + function sharedDefaults(s) { +- const d = {reasoning_effort: s.thinking, temperature: +s.temperature}; +- if (+s.temperature > 0) Object.assign(d, {top_p: +s.top_p, top_k: +s.top_k}); ++ const d = {reasoning_effort: s.thinking}; ++ if (!(+s.temperature > 0)) d.temperature = 0; ++ else if (!samplingUntouched(s)) Object.assign(d, {temperature: +s.temperature, top_p: +s.top_p, top_k: +s.top_k}); + if (s.seed) d.seed = +s.seed; + if (s.max) d.max_tokens = +s.max; + if (projectionLoaded()) d.experimental_speed_projection = s.esp !== false; +diff --git a/tools/strata_tokenizer.py b/tools/strata_tokenizer.py +index c9e84ee7..42a27fed 100644 +--- a/tools/strata_tokenizer.py ++++ b/tools/strata_tokenizer.py +@@ -15,9 +15,11 @@ emoji (4-byte), combining marks, whitespace runs, and C0 control bytes. + """ + from __future__ import annotations + ++import bisect + import json + import pathlib + import sys ++import threading + + import regex + +@@ -115,6 +117,8 @@ class Tokenizer: + # through __dict__.setdefault on each encoded segment or decoded token. + self._piece_ids: dict[str, list[int]] = {} + self._bytes_cache: dict[int, bytes] = {} ++ # How far a special-token match can look ahead of where it starts (PromptEncoder's margin). ++ self.max_special_len = max((len(t) for t in self.special_tokens), default=1) + + @staticmethod + def _alt(literals: list[str]): +@@ -234,13 +238,16 @@ class Tokenizer: + out.extend(got) + return out + +- def _encode_matching(self, text: str, pat, plain=()) -> list[int]: ++ def _encode_matching(self, text: str, pat, plain=(), marks: list | None = None) -> list[int]: + """Encode `text`, emitting any literal `pat` matches as single tokens and BPE-ing the rest. + + The split happens on the RAW text, before the byte mapping, because a special token's string is a + literal to match rather than bytes to decompose. Everything between the matches is tokenized + normally - which is why a near-miss like `<|im_star` still costs ordinary tokens. A match that starts + inside one of the `plain` (start, end) spans is left to the text around it (#537). ++ ++ `marks`, when given, receives (end of the match in `text`, ids so far) for every match: the points where ++ the encoding of a longer text with the same beginning can resume (PromptEncoder). + """ + if pat is None: + return self._encode_plain(text) +@@ -269,6 +276,8 @@ class Tokenizer: + out.extend(self._encode_plain(text[pos:start])) + out.append(self.special_tokens[m.group(0)]) + pos = m.end() ++ if marks is not None: ++ marks.append((pos, len(out))) + if pos < len(text): + out.extend(self._encode_plain(text[pos:])) + return out +@@ -283,6 +292,11 @@ class Tokenizer: + """ + return self._encode_matching(text, self._special_re if parse_special else self._always_re, plain) + ++ def encode_marked(self, text: str, parse_special: bool = False, plain=()) -> tuple[list[int], list[tuple[int, int]]]: ++ """encode() and its resume points: (end offset in `text`, ids so far) after every special-token match.""" ++ marks: list[tuple[int, int]] = [] ++ return self._encode_matching(text, self._special_re if parse_special else self._always_re, plain, marks), marks ++ + def token_bytes(self, i: int) -> bytes: + """The raw bytes of one token (a multi-byte character can be split across tokens).""" + cache = self._bytes_cache +@@ -303,6 +317,111 @@ class Tokenizer: + return b"".join(self.token_bytes(i) for i in ids).decode("utf-8", errors=errors) + + ++# ------------------------------------------------------------------ incremental prompts (#567) ++def common_prefix_len(a: str, b: str) -> int: ++ """Length of the longest common prefix, by slices (C speed) rather than a loop over characters.""" ++ n = min(len(a), len(b)) ++ lo, step = 0, 4096 ++ while lo < n: # whole chunks first, doubling ++ hi = min(n, lo + step) ++ if a[lo:hi] != b[lo:hi]: ++ break ++ lo, step = hi, min(step * 2, 1 << 20) ++ else: ++ return n ++ while hi - lo > 1: # the first difference is in [lo, hi): bisect it ++ mid = (lo + hi) // 2 ++ if a[lo:mid] == b[lo:mid]: ++ lo = mid ++ else: ++ hi = mid ++ return lo ++ ++ ++def plain_agree_until(a, b, limit: int) -> int: ++ """The first position below `limit` where the `plain` spans (#537) of two texts differ - a position inside a ++ span in one and not in the other - or `limit`. Spans are (start, end), sorted and disjoint.""" ++ def clip(spans): ++ return [(x, min(y, limit)) for x, y in spans if x < limit and y > x] ++ ca, cb = clip(a), clip(b) ++ for (xa, ya), (xb, yb) in zip(ca, cb): ++ if (xa, ya) != (xb, yb): ++ return min(xa, xb) if xa != xb else min(ya, yb) ++ if len(ca) != len(cb): ++ return (ca if len(ca) > len(cb) else cb)[min(len(ca), len(cb))][0] ++ return limit ++ ++ ++class PromptEncoder: ++ """`tok.encode(text, parse_special=True, plain=...)` for chat prompts, reusing the ids of an earlier prompt. ++ ++ A chat client sends the whole conversation every turn, and the rendered prompt of turn n+1 starts with most of ++ turn n's. Running BPE over all of it again costs hundreds of milliseconds at 100K tokens, even when the engine ++ reuses the whole prefix. So the ids up to a SPECIAL-TOKEN BOUNDARY both prompts share are taken as they were, ++ and only the rest is encoded. ++ ++ WHY THAT IS EXACT. Encoding splits the text at special-token matches and BPE-encodes each stretch between two ++ matches on its own (`_encode_matching`): nothing crosses a match. Whether a match is taken at position q ++ depends only on text[q : q + max_special_len] and on whether q lies inside a `plain` span (#537: a ++ quoted in a message is text). So if the two texts agree up to L and their plain spans agree below L, every ++ decision the scan makes at a position before c = the end of some match, with c + max_special_len - 1 <= L, is ++ the same for both - the same matches, the same stretches, the same ids up to c - and from c the scan starts ++ afresh, exactly as it does on text[c:] with the spans shifted by c. Hence encode(new) == ids_old[:ids at c] + ++ encode(new[c:]). The margin matters only for a vocabulary where a literal overlaps another's end; it costs ++ re-encoding one short stretch. ++ ++ A few recent prompts are kept (one per conversation: a prompt replaces the one it continues, that is, one whose ++ text it reuses for more than half; a prompt that shares only the template's head with every entry is a new ++ entry), so a second client does not evict the first. `encode(..., store=False)` reads the kept prompts without ++ writing one (a token count is not a conversation's turn). The tokenizer needs `encode_marked` and ++ `max_special_len`. ++ """ ++ ++ def __init__(self, tok, keep: int = 4): ++ self.tok, self.keep = tok, keep ++ # (text, plain spans, ids, mark ends, mark counts) ++ self.entries: list[tuple[str, tuple, list[int], list[int], list[int]]] = [] ++ self.lock = threading.Lock() ++ self.last_reused = 0 # characters taken from an earlier prompt (for tests) ++ ++ def encode(self, text: str, plain=(), store: bool = True) -> list[int]: ++ plain = tuple(sorted(plain)) ++ with self.lock: ++ entries = self.entries # replaced, never changed in place ++ margin = self.tok.max_special_len - 1 ++ src, cut_k = None, -1 ++ for e in entries: ++ same = plain_agree_until(e[1], plain, common_prefix_len(e[0], text)) ++ k = bisect.bisect_right(e[3], same - margin) - 1 # the last boundary c with c <= same - margin ++ if k >= 0 and (src is None or e[3][k] > src[3][cut_k]): ++ src, cut_k = e, k ++ if src is None: ++ ids, marks = self.tok.encode_marked(text, parse_special=True, plain=plain) ++ ends, counts = [m[0] for m in marks], [m[1] for m in marks] ++ self.last_reused = 0 ++ else: ++ c, n = src[3][cut_k], src[4][cut_k] ++ rest = tuple((max(x - c, 0), y - c) for x, y in plain if y > c) ++ tail, marks = self.tok.encode_marked(text[c:], parse_special=True, plain=rest) ++ ids = src[2][:n] + tail ++ ends = src[3][:cut_k + 1] + [c + m[0] for m in marks] ++ counts = src[4][:cut_k + 1] + [n + m[1] for m in marks] ++ self.last_reused = c ++ if not store: ++ return ids ++ # every rendered prompt starts with the template's header, so sharing a special boundary does not make two ++ # prompts one conversation: only the entry whose text is mostly reused is replaced ++ old = src if src is not None and 2 * src[3][cut_k] > len(src[0]) else None ++ with self.lock: ++ self.entries = [(text, plain, ids, ends, counts)] + [e for e in self.entries if e is not old][:self.keep - 1] ++ return list(ids) ++ ++ def forget(self, ids: list[int]) -> None: ++ """The prompt with these ids was refused: it is not kept.""" ++ with self.lock: ++ self.entries = [e for e in self.entries if e[2] != ids] ++ ++ + # ------------------------------------------------------------------ the pack's tokenizer/ directory + def extract(gguf_path, out_dir) -> dict: + """Write the tokenizer into `/tokenizer/` so the engine never opens the weight shards for it.""" +diff --git a/tools/test_strata_tokenizer.py b/tools/test_strata_tokenizer.py +index 17c3f698..e3cb498c 100644 +--- a/tools/test_strata_tokenizer.py ++++ b/tools/test_strata_tokenizer.py +@@ -1,9 +1,18 @@ +-"""The tokenizer's piece cache (0.1.41): the ids are the ones the merge loop gives, cache cold, warm or switched off. ++"""Tests for tools/strata_tokenizer.py: llama.cpp's reference vectors, the piece cache (0.1.41: the ids are the ones the ++merge loop gives, cache cold, warm or switched off), and the incremental prompt encoder, which must give exactly the ++ids of a plain encode. + +- python -m unittest tools.test_strata_tokenizer ++ python -m unittest tools.test_strata_tokenizer -v ++ ++The real vocabulary comes from $STRATA_TOKENIZER (a pack's tokenizer/ directory) when set, else from llama.cpp's ++`models/ggml-vocab-qwen35.gguf` next to the gguf-py the tools use ($STRATA_GGUF_PY, third_party/llama.cpp); the ++tests that need it are skipped without one. The chat template's USER_DEFINED tokens (``, ``, ++...) are added to the llama.cpp vocabulary, which does not have them, as the pack's has. + """ + from __future__ import annotations + ++import json ++import os + import random + import sys + import unittest +@@ -11,8 +20,295 @@ from pathlib import Path + + ROOT = Path(__file__).resolve().parents[1] + sys.path.insert(0, str(ROOT / "tools")) ++sys.path.insert(0, str(ROOT)) + import strata_tokenizer as ST # noqa: E402 + ++ADDED = ["", "", "", "", "", "", ++ "<|vision_start|>", "<|image_pad|>", "<|vision_end|>"] ++ ++ ++def vocab_gguf() -> Path | None: ++ cands = [os.environ.get("STRATA_GGUF_PY"), ROOT / "third_party" / "llama.cpp" / "gguf-py", "/opt/llama.cpp/gguf-py"] ++ for c in cands: ++ if c and (Path(c).parent / "models" / "ggml-vocab-qwen35.gguf").is_file(): ++ return Path(c).parent / "models" / "ggml-vocab-qwen35.gguf" ++ return None ++ ++ ++def load_tokenizer(): ++ pack = os.environ.get("STRATA_TOKENIZER") ++ if pack and (Path(pack) / "vocab.json").is_file(): ++ t = Path(pack) ++ vocab = json.loads((t / "vocab.json").read_text(encoding="utf-8")) ++ tokens = [None] * len(vocab) ++ for s, i in vocab.items(): ++ tokens[i] = s ++ return ST.Tokenizer(tokens, (t / "merges.txt").read_text(encoding="utf-8").split("\n"), ++ json.loads((t / "token_type.json").read_text())) ++ path = vocab_gguf() ++ return None if path is None else gguf_tokenizer(path) ++ ++ ++def gguf_tokenizer(path): ++ """llama.cpp's qwen35 vocabulary with the template's added tokens: the one its test vectors were made with.""" ++ from gguf_reader import GGUFFile ++ md = GGUFFile(path).metadata ++ tokens, types = list(md["tokenizer.ggml.tokens"]), list(md["tokenizer.ggml.token_type"]) ++ for s in ADDED: ++ if s not in tokens: ++ tokens.append(s) ++ types.append(4) ++ return ST.Tokenizer(tokens, list(md["tokenizer.ggml.merges"]), types) ++ ++ ++def byte_tokenizer(specials: list[str]) -> ST.Tokenizer: ++ """A vocabulary of the 256 byte tokens and `specials` (CONTROL), no merges: small enough to build per test.""" ++ tokens = [ST.BYTE_TO_UNICODE[b] for b in range(256)] + specials ++ return ST.Tokenizer(tokens, [], [1] * 256 + [3] * len(specials)) ++ ++ ++WORDS = ["the", "tokenizer", " ", " ", "\n", "\n\n", "\t", "def", "f(x):", "return", "x", "==", "->", "café", ++ "你好", "\U0001f600", "é", "1234", "don't", "I'm", "<|im_", "", ++ " ", "\r\n", "Hello,", "world!", "<", ">", "|", " ", "x" * 70] ++ ++ ++def random_text(rnd: random.Random, n: int, specials=()) -> str: ++ parts = [] ++ for _ in range(n): ++ parts.append(rnd.choice(WORDS) if not specials or rnd.random() > 0.05 else rnd.choice(specials)) ++ if rnd.random() < 0.5: ++ parts.append(" ") ++ return "".join(parts) ++ ++ ++def conversation_prompts(template, rnd: random.Random) -> list[tuple[str, str]]: ++ """Rendered prompts, in the order a server would see them, from clients that grow a conversation turn after turn, ++ use tools, send images, change the start of the prompt, edit an earlier turn, send a shorter conversation, and ++ interleave with each other: (what happens, prompt).""" ++ out: list[tuple[str, str]] = [] ++ specials = ["<|im_start|>", "<|im_end|>", "", "", "", ""] ++ weather = [{"type": "function", "function": {"name": "get_weather", "description": "Get the weather", ++ "parameters": {"type": "object", "properties": { ++ "city": {"type": "string"}}, "required": ["city"]}}}] ++ ++ def add(what, msgs, tools=None, **kw): ++ out.append((what, template.render(msgs, tools=tools, **kw))) ++ ++ # several turns, with unicode, a literal special token inside the text, and reasoning kept or not ++ msgs = [{"role": "system", "content": "Be brief. Réponds en français si on te parle en français."}] ++ for turn in range(6): ++ msgs.append({"role": "user", "content": random_text(rnd, 60, specials) + " é你\U0001f600 " * turn}) ++ add(f"turn {turn}", msgs) ++ answer = {"role": "assistant", "content": random_text(rnd, 30)} ++ if turn % 2: ++ answer["reasoning_content"] = random_text(rnd, 30) ++ msgs.append(answer) ++ grown = list(msgs) ++ ++ # tools: definitions, a call, its response, then a following turn ++ tmsgs = [{"role": "user", "content": "Weather in Oslo?"}] ++ add("tools 1", tmsgs, weather) ++ tmsgs += [{"role": "assistant", "content": "", "tool_calls": [{"type": "function", "function": { ++ "name": "get_weather", "arguments": {"city": "Oslo"}}}]}, ++ {"role": "tool", "content": "{\"temp\": 3}", "name": "get_weather"}] ++ add("tools 2", tmsgs, weather) ++ tmsgs += [{"role": "assistant", "content": "3 degrees."}, {"role": "user", "content": "And in Bergen? " * 20}] ++ add("tools 3", tmsgs, weather) ++ ++ # images: one per user turn, the placeholder tokens are specials ++ imsgs = [] ++ for turn in range(3): ++ imsgs.append({"role": "user", "content": [{"type": "image", "source": f"img{turn}"}, ++ {"type": "text", "text": f"What is on picture {turn}?"}]}) ++ add(f"images {turn}", imsgs) ++ imsgs.append({"role": "assistant", "content": "A cat."}) ++ ++ # the start of the prompt changes (a thinking level, a system message), then the old conversation grows again ++ add("effort changed", grown, reasoning_effort="low") ++ add("system changed", [dict(grown[0], content=grown[0]["content"] + " Be kind.")] + grown[1:]) ++ grown.append({"role": "user", "content": "one more question"}) ++ add("grown again", grown) ++ ++ # an earlier turn edited, then a conversation that got shorter, then the full one again ++ edited = list(grown) ++ edited[3] = dict(edited[3], content=edited[3]["content"] + " (edited)") ++ add("edited", edited) ++ add("shorter", grown[:5]) ++ add("empty user", [{"role": "user", "content": ""}]) ++ add("full again", grown) ++ return out ++ ++ ++class Oracle(unittest.TestCase): ++ """llama.cpp's own tokenizer test vectors for qwen35.""" ++ ++ def test_llama_cpp_vectors(self): ++ path = vocab_gguf() ++ if path is None or not path.with_suffix(".gguf.inp").is_file(): ++ self.skipTest("llama.cpp's ggml-vocab-qwen35.gguf test vectors are not here") ++ tk = gguf_tokenizer(path) # not $STRATA_TOKENIZER: the vectors belong to this vocabulary ++ inp = path.with_suffix(".gguf.inp").read_text(encoding="utf-8").split("\n__ggml_vocab_test__\n") ++ out = path.with_suffix(".gguf.out").read_text(encoding="utf-8").split("\n") ++ n = 0 ++ for text, want in zip(inp, out): ++ if text.endswith("\n__ggml_vocab_test__"): ++ text = text[:-len("\n__ggml_vocab_test__")] ++ self.assertEqual(tk.encode(text), [int(x) for x in want.split()], repr(text[:40])) ++ n += 1 ++ self.assertGreater(n, 40) ++ ++ ++class Marks(unittest.TestCase): ++ @classmethod ++ def setUpClass(cls): ++ cls.tk = load_tokenizer() ++ if cls.tk is None: ++ raise unittest.SkipTest("no tokenizer (set STRATA_TOKENIZER or STRATA_GGUF_PY)") ++ ++ def test_marks(self): ++ text = "a<|im_start|>user\nhi<|im_end|>\n\n" ++ ids, marks = self.tk.encode_marked(text, parse_special=True) ++ self.assertEqual(ids, self.tk.encode(text, parse_special=True)) ++ self.assertEqual([m[0] for m in marks], [text.index("user"), text.index("<|im_end|>") + len("<|im_end|>"), ++ len(text) - 1]) ++ for end, n in marks: # each mark: the ids of the text up to it ++ self.assertEqual(ids[:n], self.tk.encode(text[:end], parse_special=True)) ++ ++ ++class Prompts(unittest.TestCase): ++ """PromptEncoder: the ids of a plain encode, whatever the earlier prompts were.""" ++ ++ @classmethod ++ def setUpClass(cls): ++ cls.tk = load_tokenizer() ++ if cls.tk is None: ++ raise unittest.SkipTest("no tokenizer (set STRATA_TOKENIZER or STRATA_GGUF_PY)") ++ from serve.frontend import ChatTemplate ++ cls.template = ChatTemplate(ROOT / "serve" / "chat_template.jinja") ++ ++ def check(self, enc, text): ++ self.assertEqual(enc.encode(text), self.tk.encode(text, parse_special=True)) ++ ++ def test_chat_golden(self): ++ cases = json.loads((ROOT / "serve" / "chat_golden.json").read_text(encoding="utf-8")) ++ enc = ST.PromptEncoder(self.tk) ++ for c in cases + cases: # the second round finds its own earlier prompts ++ if c.get("rendered"): ++ with self.subTest(case=c["name"]): ++ self.check(enc, c["rendered"]) ++ ++ def test_conversations(self): ++ for seed in range(3): ++ prompts = conversation_prompts(self.template, random.Random(seed)) ++ wants = [self.tk.encode(prompt, parse_special=True) for _, prompt in prompts] ++ for keep in (1, 2, 4): ++ enc = ST.PromptEncoder(self.tk, keep=keep) ++ for (what, prompt), want in zip(prompts, wants): ++ with self.subTest(seed=seed, keep=keep, what=what): ++ self.assertEqual(enc.encode(prompt), want) ++ ++ def test_growing_turn_reuses_the_previous_prompt(self): ++ enc = ST.PromptEncoder(self.tk) ++ msgs = [{"role": "user", "content": "first question " * 50}] ++ prev = "" ++ for turn in range(4): ++ prompt = self.template.render(msgs) ++ self.check(enc, prompt) ++ if turn: ++ self.assertGreater(enc.last_reused, len(prev) - 100) ++ prev = prompt ++ msgs += [{"role": "assistant", "content": f"answer {turn}"}, {"role": "user", "content": f"question {turn}"}] ++ ++ def test_prefix_of_an_earlier_prompt_and_unrelated_prompts(self): ++ enc = ST.PromptEncoder(self.tk, keep=2) ++ a = self.template.render([{"role": "user", "content": "first question " * 50}]) ++ b = self.template.render([{"role": "user", "content": "another conversation"}]) ++ c = self.template.render([{"role": "user", "content": "a third one"}]) ++ for text in (a, a[:len(a) // 2], b, c, a, a + "tail", "", "<|im_end|>", "<|im_end|><|im_end|>x"): ++ with self.subTest(text=text[:30]): ++ self.check(enc, text) ++ ++ def test_a_read_only_encode_gives_the_ids_and_stores_nothing(self): ++ enc = ST.PromptEncoder(self.tk) ++ a = self.template.render([{"role": "user", "content": "first question " * 50}]) ++ enc.encode(a) ++ kept = list(enc.entries) ++ for text in (a + "<|im_start|>user\nand a follow-up<|im_end|>\n", a[:len(a) // 2], "<|im_end|>x", ++ self.template.render([{"role": "user", "content": "another conversation"}])): ++ with self.subTest(text=text[:30]): ++ self.assertEqual(enc.encode(text, store=False), self.tk.encode(text, parse_special=True)) ++ self.assertTrue(enc.entries == kept, "a read-only encode wrote an entry") ++ self.assertEqual(enc.encode(a + "tail", store=False), self.tk.encode(a + "tail", parse_special=True)) ++ self.assertGreater(enc.last_reused, len(a) // 2) ++ ++ def test_the_returned_ids_are_the_callers(self): ++ enc = ST.PromptEncoder(self.tk) ++ text = "<|im_start|>user\nhi<|im_end|>\n" + "and a tail longer than the look-ahead margin. " * 3 ++ want = self.tk.encode(text, parse_special=True) ++ got = enc.encode(text) ++ got[0] = -1 # inside the part the next call takes from this one ++ got.append(-1) ++ self.assertEqual(enc.encode(text), want) ++ self.assertGreater(enc.last_reused, 0) ++ ++ ++class Margin(unittest.TestCase): ++ """A special literal that is a prefix of a longer one: cutting at the shorter one's end without the look-ahead ++ margin would give different ids.""" ++ ++ def test_longer_literal_past_the_common_prefix(self): ++ tk = byte_tokenizer(["<|a|>", "<|a|>zz"]) ++ enc = ST.PromptEncoder(tk) ++ first = "x<|a|>y" + "<|a|>" * 3 + "q" ++ self.assertEqual(enc.encode(first), tk.encode(first, parse_special=True)) ++ second = "x<|a|>y" + "<|a|>" * 3 + "zz" # shares everything up to the last <|a|> ++ got = enc.encode(second) ++ self.assertEqual(got, tk.encode(second, parse_special=True)) ++ self.assertEqual(got[-1], tk.special_tokens["<|a|>zz"]) ++ self.assertGreater(enc.last_reused, 0) # the earlier boundaries were still reused ++ ++ def test_literal_ending_inside_another(self): ++ tk = byte_tokenizer(["<|a|>", "b<|a|>c"]) ++ enc = ST.PromptEncoder(tk) ++ for text in ("zzb<|a|>d", "zzb<|a|>c", "zzb<|a|>cz<|a|>", "zzb<|a|>cz<|a|>b<|a|>c"): ++ with self.subTest(text=text): ++ self.assertEqual(enc.encode(text), tk.encode(text, parse_special=True)) ++ ++ def test_random_specials(self): ++ """Literals that extend or end inside one another, texts that change right after one of them.""" ++ rnd = random.Random(3) ++ ++ def word(lo, hi): ++ return "".join(rnd.choice("<|ab>") for _ in range(rnd.randint(lo, hi))) ++ for trial in range(60): ++ a, b = word(2, 4), word(2, 4) ++ lits = sorted({a, a + word(1, 3), b, word(1, 2) + b + word(0, 2)}) ++ tk = byte_tokenizer(lits) ++ enc = ST.PromptEncoder(tk, keep=3) ++ pieces = lits + [x[:rnd.randint(1, len(x))] for x in lits] + ["x", "y", "a", "|"] ++ base = "".join(rnd.choice(pieces) for _ in range(60)) ++ for _ in range(30): ++ cut = rnd.randint(0, len(base)) ++ text = base[:cut] + "".join(rnd.choice(pieces) for _ in range(rnd.randint(0, 6))) ++ self.assertEqual(enc.encode(text), tk.encode(text, parse_special=True), (lits, text)) ++ if rnd.random() < 0.3: ++ base = text + base[cut:] ++ ++ ++class CommonPrefix(unittest.TestCase): ++ def test_against_a_loop(self): ++ rnd = random.Random(5) ++ for _ in range(500): ++ a = "".join(rnd.choice("ab") for _ in range(rnd.randint(0, 9000))) ++ k = rnd.randint(0, len(a)) ++ b = a[:k] + rnd.choice(["", "c", "cd"]) + a[k:k + rnd.randint(0, 50)] ++ want = 0 ++ while want < min(len(a), len(b)) and a[want] == b[want]: ++ want += 1 ++ self.assertEqual(ST.common_prefix_len(a, b), want) ++ self.assertEqual(ST.common_prefix_len(b, a), want) ++ ++ + + def toy() -> ST.Tokenizer: + base = [ST.BYTE_TO_UNICODE[b] for b in range(256)] +diff --git a/serve/test_admit_order.py b/serve/test_admit_order.py +new file mode 100644 +--- /dev/null ++++ b/serve/test_admit_order.py +@@ -0,0 +1,119 @@ ++"""serve/test_admit_order.py - STRATA_ADMIT_SHORT_FIRST: new requests take the control lines short prompts first ++(StrataEngine._take_control / _admit_next; no process, no GPU). ++ ++ python -m unittest serve.test_admit_order -v ++""" ++from __future__ import annotations ++ ++import sys ++import threading ++import time ++import unittest ++from pathlib import Path ++from unittest import mock ++ ++sys.path.insert(0, str(Path(__file__).resolve().parents[1])) ++import serve.server as server # noqa: E402 ++from serve.server import StrataEngine # noqa: E402 ++ ++ ++def engine_without_process(n=2): ++ """A StrataEngine with the admission state of a started one and no process: alive() is `e.up`.""" ++ e = StrataEngine("missing-executable", [], lazy=True) ++ e.batch = n ++ e.slot_busy = [False] * n ++ e.slot_claims, e.slot_releasing, e._orphan_seen = set(), set(), {} ++ e.slot_cv = threading.Condition() ++ e.waiting, e.admitting, e.wait_lens, e.ctl_epoch = 0, 0, [], 0 ++ e.ctl = threading.Lock() ++ e.gen = 1 ++ e.up = True ++ e.alive = lambda: e.up ++ return e ++ ++ ++class AdmitOrderTest(unittest.TestCase): ++ def setUp(self): ++ self.e = engine_without_process(16) ++ self.order, self.cancels = [], [] ++ self.order_lock = threading.Lock() ++ ++ def tearDown(self): ++ for c in self.cancels: ++ c.set() ++ ++ def take(self, name, plen, **kw): ++ """A request waiting for the control lines; once it has them it notes its name and gives them back.""" ++ c = threading.Event() ++ self.cancels.append(c) ++ ++ def go(): ++ gen = self.e._take_control(c, plen, **kw) ++ try: ++ while True: ++ next(gen) ++ except StopIteration as s: ++ if s.value: ++ with self.order_lock: ++ self.order.append(name) ++ time.sleep(0.05) # its read ++ self.e.ctl.release() ++ threading.Thread(target=go, daemon=True).start() ++ ++ def until(self, cond, s=5.0): ++ end = time.time() + s ++ while time.time() < end and not cond(): ++ time.sleep(0.01) ++ self.assertTrue(cond()) ++ ++ def queue(self, *reqs): ++ """While a long read holds the control lines, the requests arrive in this order; then the read ends.""" ++ self.assertTrue(self.e.ctl.acquire(timeout=1)) ++ for name, plen in reqs: ++ self.take(name, plen, new=True) ++ self.until(lambda n=len(self.order) + sum(1 for _ in self.e.wait_lens): self.e.waiting >= 1) ++ time.sleep(0.02) # (distinct arrival times) ++ self.until(lambda: self.e.waiting == len(reqs)) ++ self.e.ctl.release() ++ self.until(lambda: len(self.order) == len(reqs), s=15.0) ++ return self.order ++ ++ def test_short_first_puts_the_chat_before_the_agent_reads(self): ++ with mock.patch.object(server, "ADMIT_SHORT_FIRST", 8192), mock.patch.object(server, "ADMIT_MAX_WAIT_S", 60.0): ++ order = self.queue(("agent1", 120000), ("agent2", 190000), ("chat", 500), ("agent3", 150000)) ++ self.assertEqual(order[0], "chat") ++ self.assertEqual(order[1:], ["agent1", "agent2", "agent3"]) # arrival order among the long ones ++ ++ def test_a_long_request_that_waited_too_long_goes_first(self): ++ with mock.patch.object(server, "ADMIT_SHORT_FIRST", 8192), mock.patch.object(server, "ADMIT_MAX_WAIT_S", 0.2): ++ self.assertTrue(self.e.ctl.acquire(timeout=1)) ++ self.take("agent", 120000, new=True) ++ self.until(lambda: self.e.waiting == 1) ++ time.sleep(0.4) # past STRATA_ADMIT_MAX_WAIT_S ++ self.take("chat", 500, new=True) ++ self.until(lambda: self.e.waiting == 2) ++ self.e.ctl.release() ++ self.until(lambda: len(self.order) == 2, s=15.0) ++ self.assertEqual(self.order, ["agent", "chat"]) ++ ++ def test_off_by_default(self): ++ self.assertEqual(server.ADMIT_SHORT_FIRST if "STRATA_ADMIT_SHORT_FIRST" not in __import__("os").environ ++ else 0, 0) ++ with mock.patch.object(server, "ADMIT_SHORT_FIRST", 0): ++ order = self.queue(("agent1", 120000), ("chat", 500)) ++ self.assertEqual(sorted(order), ["agent1", "chat"]) # both get their turn, in any order ++ ++ def test_a_continuation_is_not_held_back(self): ++ """A request already admitted once (new=False: promoted from the solo path) is not part of the order.""" ++ with mock.patch.object(server, "ADMIT_SHORT_FIRST", 8192), mock.patch.object(server, "ADMIT_MAX_WAIT_S", 60.0): ++ self.assertTrue(self.e.ctl.acquire(timeout=1)) ++ self.take("chat", 500, new=True) ++ self.take("continuation", 190000) ++ self.until(lambda: self.e.waiting == 2) ++ self.e.ctl.release() ++ self.until(lambda: len(self.order) == 2, s=15.0) ++ self.assertEqual(sorted(self.order), ["chat", "continuation"]) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_model_card_sampling.py b/serve/test_model_card_sampling.py +new file mode 100644 +--- /dev/null ++++ b/serve/test_model_card_sampling.py +@@ -0,0 +1,108 @@ ++"""serve/test_model_card_sampling.py - a request that sends no sampling gets the model card's, not greedy (mock engine). ++ ++Qwen's "Best Practices": thinking mode temperature 1.0 / top_p 0.95 / top_k 20, non-thinking 0.7 / 0.8 / 20 with ++presence_penalty 1.5. The request's own fields win, then the Chat settings shared with apps, then the config's ++`sampling` block, then the card; temperature 0 asked for is still greedy. ++ ++ python -m unittest serve.test_model_card_sampling -v ++""" ++from __future__ import annotations ++ ++import json ++import sys ++import unittest ++import urllib.request ++from pathlib import Path ++ ++sys.path.insert(0, str(Path(__file__).resolve().parents[1])) ++from serve.frontend import ChatTemplate # noqa: E402 ++from serve.server import ByteTokenizer, MockEngine, Service, StrataEngine, model_card_sampling, serve # noqa: E402 ++ ++ROOT = Path(__file__).resolve().parents[1] ++ ++ ++class Recording(MockEngine): ++ def generate(self, ids, max_new, sampling, cancel, embeddings=None): ++ self.sampling = dict(sampling or {}) ++ yield from super().generate(ids, max_new, sampling, cancel, embeddings) ++ ++ ++class ModelCardSampling(unittest.TestCase): ++ @classmethod ++ def setUpClass(cls): ++ cls.tok = ByteTokenizer() ++ cls.engine = Recording(cls.tok, "\nhm\n\n\nok", max_context=8192) ++ cls.svc = Service(cls.engine, cls.tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) ++ cls.httpd = serve(cls.svc, port=0) ++ cls.base = f"http://127.0.0.1:{cls.httpd.server_address[1]}" ++ ++ @classmethod ++ def tearDownClass(cls): ++ cls.httpd.shutdown() ++ cls.httpd.server_close() ++ ++ def setUp(self): ++ self.svc.model_arch = "qwen4exp" ++ self.svc.sampling_defaults = {} ++ self.svc.shared = {} ++ self.svc.config_effort = None ++ ++ def tearDown(self): ++ self.svc.model_arch = None ++ ++ def keys(self, **body): ++ body = {"model": "m", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 8, **body} ++ self.engine.sampling = {} ++ req = urllib.request.Request(self.base + "/v1/chat/completions", data=json.dumps(body).encode(), ++ headers={"Content-Type": "application/json"}) ++ urllib.request.urlopen(req, timeout=10).read() ++ return StrataEngine.sampling_keys(self.engine.sampling) ++ ++ def test_tables(self): ++ self.assertEqual(model_card_sampling("qwen4exp", True), {"temperature": 1.0, "top_p": 0.95, "top_k": 20}) ++ self.assertEqual(model_card_sampling("qwen4exp", False), ++ {"temperature": 0.7, "top_p": 0.8, "top_k": 20, "presence_penalty": 1.5}) ++ self.assertEqual(model_card_sampling("qwen35moe", True)["presence_penalty"], 1.5) ++ self.assertEqual(model_card_sampling(None, True), {}) ++ self.assertEqual(model_card_sampling("unknown", False), {}) ++ ++ def test_thinking_default(self): ++ k = self.keys() ++ self.assertIn("temperature=1.0", k) ++ self.assertIn("top_p=0.95", k) ++ self.assertIn("top_k=20", k) ++ self.assertNotIn("penalty_present", k) ++ ++ def test_non_thinking_default(self): ++ k = self.keys(reasoning_effort="none") ++ self.assertIn("temperature=0.7", k) ++ self.assertIn("top_p=0.8", k) ++ self.assertIn("penalty_present=1.5", k) ++ ++ def test_request_wins_and_zero_is_greedy(self): ++ k = self.keys(temperature=0) ++ self.assertNotIn("temperature", k) ++ self.assertIn("top_k=20", k) ++ k = self.keys(temperature=0.3, top_p=0.5) ++ self.assertIn("temperature=0.3", k) ++ self.assertIn("top_p=0.5", k) ++ ++ def test_config_block_and_shared_settings_win_over_the_card(self): ++ self.svc.sampling_defaults = {"temperature": 0.6, "top_k": 40} ++ k = self.keys() ++ self.assertIn("temperature=0.6", k) ++ self.assertIn("top_k=40", k) ++ self.assertIn("top_p=0.95", k) ++ self.svc.shared = {"temperature": 0.2} ++ self.assertIn("temperature=0.2", self.keys()) ++ self.svc.sampling_defaults = {"temperature": 0} ++ self.svc.shared = {} ++ self.assertNotIn("temperature", self.keys()) # a config that asks for greedy gets it ++ ++ def test_no_family_no_defaults(self): ++ self.svc.model_arch = None ++ self.assertEqual(self.keys(), "") ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_parser_call_scan.py b/serve/test_parser_call_scan.py +new file mode 100644 +--- /dev/null ++++ b/serve/test_parser_call_scan.py +@@ -0,0 +1,61 @@ ++"""Wrapped-call precheck parity against the original structural walk.""" ++import itertools ++import random ++import unittest ++from unittest.mock import patch ++ ++from serve import frontend ++ ++ ++class ReferenceParser(frontend.OutputParser): ++ def _wrapped_call_end(self): ++ return frontend.call_end(self.buf) ++ ++ ++def events_without_ids(events): ++ return [(e.kind, e.text, (e.call.name, e.call.arguments) if e.call else None) for e in events] ++ ++ ++class CallScanTests(unittest.TestCase): ++ def test_chunk_boundaries_literal_tags_multiple_calls_and_recovery(self): ++ def call(value): ++ return "" + value + "" ++ bodies = [call("ordinary"), call("literal text here"), ++ call("literal prose"), call("a") + call("b"), ++ call("unfinished")[:-5], call("incomplete prose")[:-30], ++ "", ++ "bare", ++ "a" ++ "b"] ++ tools = [{"name": "write", "parameters": {"properties": {"text": {"type": "string"}}}}] ++ rng = random.Random(1030) ++ for text, stream, recover in itertools.product(bodies, (False, True), (False, True)): ++ # Every fixed chunk width, plus random partitions; compare after EACH feed. ++ for width in range(1, 16): ++ a = frontend.OutputParser(thinking=False, tools=tools, stream_tools=stream, recover=recover) ++ b = ReferenceParser(thinking=False, tools=tools, stream_tools=stream, recover=recover) ++ pos = 0 ++ while pos < len(text): ++ end = pos + (rng.randint(1, 20) if width == 15 else width) ++ delta = text[pos:end] ++ self.assertEqual(events_without_ids(a.feed(delta)), events_without_ids(b.feed(delta))) ++ pos = end ++ self.assertEqual(events_without_ids(a.finish()), events_without_ids(b.finish())) ++ ++ def test_no_structural_rescans_until_closer_and_reset_for_next_call(self): ++ parser = frontend.OutputParser(thinking=False) ++ with patch.object(frontend, "call_end", wraps=frontend.call_end) as walk: ++ for _ in range(2): ++ parser.feed("") ++ before = walk.call_count ++ for _ in range(100): ++ parser.feed("abcdefgh") ++ self.assertEqual(walk.call_count, before) ++ for char in "": ++ parser.feed(char) ++ self.assertEqual(walk.call_count, before + 1) ++ self.assertEqual(parser._call_end_checked, 0) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_parser_hold.py b/serve/test_parser_hold.py +new file mode 100644 +--- /dev/null ++++ b/serve/test_parser_hold.py +@@ -0,0 +1,60 @@ ++"""Partial-tag parity, including overlapping prefixes and arbitrary chunk boundaries.""" ++import itertools ++import random ++import unittest ++ ++from serve.frontend import OutputParser ++ ++ ++def reference_hold(text, tags): ++ return max((n for tag in tags for n in range(1, len(tag)) if text.endswith(tag[:n])), default=0) ++ ++ ++class ReferenceParser(OutputParser): ++ def _hold(self, text, tags): ++ return reference_hold(text, tags) ++ ++ ++class HoldTests(unittest.TestCase): ++ def test_overlapping_and_empty_tags(self): ++ parser = OutputParser() ++ tags = ("", "a", "aaa", "abab", "baab", "", "") ++ for size in range(8): ++ for chars in itertools.product("ab<", repeat=size): ++ text = "".join(chars) ++ self.assertEqual(parser._hold(text, tags), reference_hold(text, tags), text) ++ ++ def test_every_real_tag_split(self): ++ tags = ("", "", "") ++ parser = OutputParser() ++ for tag in tags: ++ for n in range(len(tag) + 1): ++ for prefix in ("", "prose\n", "<", "<<<", "🙂日本語"): ++ text = prefix + tag[:n] ++ self.assertEqual(parser._hold(text, tags), reference_hold(text, tags), text) ++ ++ def test_streamed_events_match_reference(self): ++ texts = ["Reasoning.\n\nAnswer café 🙂\n", ++ "Quoted `` text and prose.", ++ "\n```xml\n\n```\nAnswer\nabc\nDone", ++ "abcDone"] ++ rng = random.Random(1028) ++ for text, thinking, recover, stream_tools in itertools.product(texts, (False, True), (False, True), (False, True)): ++ for _ in range(5): ++ args = dict(thinking=thinking, tools=[{"name": "lookup"}], recover=recover, stream_tools=stream_tools) ++ actual, expected = OutputParser(**args), ReferenceParser(**args) ++ def normalized(events): ++ # Tool call IDs are deliberately random; compare their semantic contents. ++ return [(e.kind, e.text, (e.call.name, e.call.arguments) if e.call else None) for e in events] ++ pos = 0 ++ while pos < len(text): ++ end = pos + rng.randint(1, 9) ++ delta = text[pos:end] ++ self.assertEqual(normalized(actual.feed(delta)), normalized(expected.feed(delta))) ++ pos = end ++ self.assertEqual(normalized(actual.finish()), normalized(expected.finish())) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_prompt_encoder.py b/serve/test_prompt_encoder.py +new file mode 100644 +--- /dev/null ++++ b/serve/test_prompt_encoder.py +@@ -0,0 +1,162 @@ ++"""serve/test_prompt_encoder.py - #567's incremental prompt encoding with #537's plain spans: every prompt's ids ++are those of a full encode, and a turn reuses the previous one's. ++ ++ python -m unittest serve.test_prompt_encoder -v ++""" ++from __future__ import annotations ++ ++import random ++import re ++import sys ++import unittest ++from pathlib import Path ++ ++ROOT = Path(__file__).resolve().parents[1] ++sys.path.insert(0, str(ROOT)) ++sys.path.insert(0, str(ROOT / "tools")) ++from serve.frontend import ChatTemplate # noqa: E402 ++from serve.server import ByteTokenizer, MockEngine, Service # noqa: E402 ++from strata_tokenizer import PromptEncoder, plain_agree_until # noqa: E402 ++ ++ ++class ThinkTokenizer(ByteTokenizer): ++ """ByteTokenizer whose / are specials matched even without parse_special (type 4), as in the ++ Qwen vocabulary: the case #537's plain spans exist for.""" ++ SPECIALS = ByteTokenizer.SPECIALS + ["", ""] ++ ALWAYS = ("", "") ++ ++ ++PIECES = ["<|im_start|>user\n", "<|im_end|>\n", "", "", "hello ", "é你 ", "<|im_sta", "x" * 7, ++ "<|im_start|>assistant\n", "", "<|vision_start|>", "<|image_pad|>"] ++ ++ ++def spans_of(text, rng): ++ """Random plain spans over the specials' texts (, , <|im_start|>, ...), as unmark_think_literals ++ marks the ones quoted inside a message.""" ++ matches = re.finditer("|".join(map(re.escape, ThinkTokenizer.SPECIALS)), text) ++ return tuple((m.start(), m.end()) for m in matches if rng.random() < 0.5) # finditer: left to right, no overlap ++ ++ ++class Units(unittest.TestCase): ++ def test_plain_agree_until(self): ++ self.assertEqual(plain_agree_until((), (), 50), 50) ++ self.assertEqual(plain_agree_until(((5, 13),), ((5, 13),), 50), 50) ++ self.assertEqual(plain_agree_until(((5, 13),), (), 50), 5) ++ self.assertEqual(plain_agree_until(((5, 13),), ((5, 9),), 50), 9) ++ self.assertEqual(plain_agree_until(((60, 70),), (), 50), 50) # beyond the shared text: no matter ++ self.assertEqual(plain_agree_until(((5, 13), (20, 28)), ((5, 13),), 50), 20) ++ ++ def test_same_ids_as_a_full_encode(self): ++ tok = ThinkTokenizer() ++ rng = random.Random(7) ++ for conv in range(300): ++ enc = PromptEncoder(tok) ++ text, plain = "", () ++ for turn in range(8): ++ text += "".join(rng.choice(PIECES) for _ in range(rng.randint(1, 12))) ++ if rng.random() < 0.6: ++ plain = spans_of(text, rng) # the spans may change in the shared part too ++ with self.subTest(conv=conv, turn=turn): ++ self.assertEqual(enc.encode(text, plain), tok.encode(text, True, plain)) ++ if rng.random() < 0.2 and len(text) > 10: # an edit: a client re-rendering its history ++ cut = rng.randrange(len(text)) ++ text = text[:cut] ++ plain = tuple((a, b) for a, b in plain if b <= cut) ++ ++ def test_a_turn_reuses_the_previous_one(self): ++ tok = ThinkTokenizer() ++ enc = PromptEncoder(tok) ++ # as a chat turn: the earlier prompt goes on past its last special (the answer being written), and the ++ # next one shares it up to there - the resume point needs max_special_len characters after it ++ base = "<|im_start|>user\n" + "quoted here " * 200 + "<|im_end|>\n<|im_start|>assistant\nan answer" ++ plain = tuple((i, i + 8) for i in range(len(base)) if base.startswith("", i)) ++ enc.encode(base, plain) ++ more = base + " that goes on<|im_end|>\n<|im_start|>user\nnext<|im_end|>\n" ++ self.assertEqual(enc.encode(more, plain), tok.encode(more, True, plain)) ++ self.assertGreater(enc.last_reused, len(base) // 2) ++ # the same text with the quoted no longer plain: nothing before it may be reused past its start ++ self.assertEqual(enc.encode(more, ()), tok.encode(more, True, ())) ++ self.assertLessEqual(enc.last_reused, plain[0][0]) ++ ++ ++class Served(unittest.TestCase): ++ def test_quoted_think_across_turns(self): ++ tok = ThinkTokenizer() ++ template = ChatTemplate(ROOT / "serve/chat_template.jinja") ++ svc = Service(MockEngine(tok, "\n\nok", max_context=1 << 20), tok, template) ++ full = Service(MockEngine(tok, "\n\nok", max_context=1 << 20), tok, template) ++ full.prompts = None # the reference: a full encode every time ++ self.assertIsNotNone(svc.prompts) ++ msgs = [{"role": "system", "content": "Explain tags."}] ++ for turn in range(5): ++ msgs.append({"role": "user", "content": f"turn {turn}: what does mean? too " * 3}) ++ got, want = svc.encode_prompt(msgs, None, {}), full.encode_prompt(msgs, None, {}) ++ self.assertEqual(got, want) ++ if turn: ++ self.assertGreater(svc.prompts.last_reused, 0) ++ msgs.append({"role": "assistant", "content": "It closes the reasoning: is a marker."}) ++ ++ def service(self, max_context=1 << 22): ++ tok = ByteTokenizer() ++ return Service(MockEngine(tok, "ok", max_context=max_context), tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) ++ ++ def test_another_conversation_does_not_evict_the_first(self): ++ """A subagent beside the main conversation: every rendered prompt starts with the template's header, which ++ must not make the second client's prompt a continuation of the first's entry.""" ++ svc = self.service() ++ turn2 = [{"role": "user", "content": "a long question " * 500}, {"role": "assistant", "content": "answer"}, ++ {"role": "user", "content": "next"}] ++ svc.encode_prompt(turn2, None, {}) ++ svc.encode_prompt([{"role": "user", "content": "one unrelated prompt " * 300}], None, {}) ++ self.assertEqual(len(svc.prompts.entries), 2) ++ turn3 = turn2 + [{"role": "assistant", "content": "x"}, {"role": "user", "content": "more"}] ++ svc.encode_prompt(turn3, None, {}) ++ self.assertGreater(svc.prompts.last_reused, len(turn2[0]["content"])) ++ self.assertEqual(len(svc.prompts.entries), 2) # turn 3 replaced turn 2, the other one stayed ++ ++ def test_a_refused_prompt_is_not_kept(self): ++ svc = self.service(max_context=4096) ++ msgs = [{"role": "user", "content": "kept " * 20}] ++ svc.prepare(msgs, None, {}, 16) ++ refused = msgs + [{"role": "assistant", "content": "a"}, {"role": "user", "content": "more " * 20}] ++ full = self.service(max_context=4096) ++ full.prompts = None ++ refused_ids = full.encode_prompt(refused, None, {}) ++ with self.assertRaisesRegex(ValueError, "exceeds the context"): ++ svc.prepare(refused, None, {}, 1 << 20) ++ self.assertNotIn(refused_ids, [e[2] for e in svc.prompts.entries]) ++ self.assertNotIn(full.render_prompt(refused, None, {}), [e[0] for e in svc.prompts.entries]) ++ ++ def test_random_messages_equal_a_full_encode(self): ++ """Messages quoting control-token texts and tags at random, growing, edited and cut back between ++ requests: the incremental encoder's ids are the full encode's, 2400 prompts.""" ++ tok = ThinkTokenizer() ++ template = ChatTemplate(ROOT / "serve/chat_template.jinja") ++ svc = Service(MockEngine(tok, "ok", max_context=1 << 20), tok, template) ++ full = Service(MockEngine(tok, "ok", max_context=1 << 20), tok, template) ++ full.prompts = None ++ rng = random.Random(567) ++ quoted = ThinkTokenizer.SPECIALS + ["<|im_sta", " 2: # a shorter conversation ++ msgs = msgs[:rng.randrange(1, len(msgs))] + msgs[-1:] ++ with self.subTest(conv=conv, turn=turn): ++ self.assertEqual(svc.encode_prompt(msgs, None, {}), full.encode_prompt(msgs, None, {})) ++ reused += svc.prompts.last_reused > 0 ++ msgs.append({"role": "assistant", "content": text()}) ++ self.assertGreater(reused, 1000) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_q36_config.py b/serve/test_q36_config.py +new file mode 100644 +--- /dev/null ++++ b/serve/test_q36_config.py +@@ -0,0 +1,48 @@ ++"""serve/test_q36_config.py - what the server hands the Qwen3.6-35B-A3B engine (arch qwen35moe): the config's pictures ++("vision" section -> --vision) and parallel requests ("parallel": N -> --batch N), as for Flash-Next's engine. ++ ++ python -m unittest serve.test_q36_config ++""" ++from __future__ import annotations ++ ++import unittest ++from unittest import mock ++ ++from serve.server import engine_args, parallel_args ++ ++ ++def cfg(**kw): ++ base = {"arch": "qwen35moe", "exe": "strata", "args": ["--native", "/m/Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf", ++ "--max-context", "32768", "--prefill", "auto", "--mtp", "2"]} ++ base.update(kw) ++ return base ++ ++ ++class Q36Config(unittest.TestCase): ++ def test_parallel_becomes_batch(self): ++ a = engine_args(cfg(parallel=4)) ++ self.assertEqual(a[a.index("--batch") + 1], "4") ++ self.assertNotIn("--batch", engine_args(cfg())) ++ self.assertNotIn("--batch", engine_args(cfg(parallel=1))) ++ ++ def test_more_than_the_window_is_said(self): ++ with mock.patch("builtins.print") as pr: ++ a = engine_args(cfg(parallel=12)) ++ self.assertEqual(a[a.index("--batch") + 1], "12") # the engine caps it at 8 and says so ++ self.assertTrue(any("at most 8" in str(c) for c in pr.call_args_list)) ++ ++ def test_a_vision_section_starts_the_engine_with_pictures(self): ++ v = {"exe": "strata-vision", "mmproj": "/m/mmproj-F16.gguf", "model": "/m/Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf", ++ "gpu": False, "max_tokens": 300} ++ self.assertIn("--vision", engine_args(cfg(vision=v))) ++ self.assertIn("--vision", engine_args(cfg(vision=v, args=cfg()["args"] + ["--vision"]))) ++ self.assertEqual(engine_args(cfg(vision=v, args=cfg()["args"] + ["--vision"])).count("--vision"), 1) ++ self.assertNotIn("--vision", engine_args(cfg())) ++ ++ def test_an_explicit_batch_in_the_args_wins(self): ++ c = cfg(parallel=4, args=cfg()["args"] + ["--batch", "2"]) ++ self.assertEqual(parallel_args(c, c["args"]), []) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_web_sampling.py b/serve/test_web_sampling.py +new file mode 100644 +--- /dev/null ++++ b/serve/test_web_sampling.py +@@ -0,0 +1,77 @@ ++"""serve/test_web_sampling.py - the web page sends no sampling of its own until the user changes it (0.1.45, #2685). ++ ++The page used to send temperature 0.6 / top_p 0.95 / top_k 20 with every request, so the server's model-card defaults ++(thinking 1.0, not thinking 0.7 + presence penalty) never applied to it. Pins, with node, the page's own functions: ++untouched (the shown defaults, or the old stored 0.6) sends nothing, a changed value is sent whole, 0 is greedy. ++Skipped without node. ++ ++ python -m unittest serve.test_web_sampling -v ++""" ++from __future__ import annotations ++ ++import json ++import re ++import shutil ++import subprocess ++import unittest ++from pathlib import Path ++ ++APP = Path(__file__).resolve().parent / "web" / "app.js" ++ ++ ++def slice_between(src: str, start: str, end: str) -> str: ++ a = src.index(start) ++ return src[a:src.index(end, a)] ++ ++ ++@unittest.skipUnless(shutil.which("node"), "node is not installed") ++class WebSampling(unittest.TestCase): ++ @classmethod ++ def setUpClass(cls): ++ src = APP.read_text(encoding="utf-8") ++ cls.defaults = slice_between(src, "const DEFAULTS", "let settings") ++ cls.shared = slice_between(src, "function sharedDefaults", "async function saveShared") ++ cls.body = slice_between(src, " if (!(settings.temperature > 0)) {", " if (settings.seed)") ++ ++ def run_js(self, settings: dict, pieces: str): ++ js = (self.defaults + "function projectionLoaded() { return false; }\n" + self.shared + ++ f"\nconst settings = {json.dumps(settings)};\nconst body = {{}};\n{pieces}\n" ++ "console.log(JSON.stringify({body, shared: sharedDefaults(settings)}));") ++ out = subprocess.run(["node", "-e", js], capture_output=True, text=True, timeout=30) ++ self.assertEqual(out.returncode, 0, out.stderr) ++ return json.loads(out.stdout) ++ ++ def sampling(self, **kw): ++ s = {"thinking": "high", "temperature": 1.0, "top_p": 0.95, "top_k": 20, "max": "", "seed": "", "esp": True} ++ s.update(kw) ++ r = self.run_js(s, self.body) ++ return r["body"], r["shared"] ++ ++ def test_untouched_sends_nothing(self): ++ for t in (1.0, 0.6): # the shown default, and the old stored default ++ body, shared = self.sampling(temperature=t) ++ self.assertEqual(body, {}) ++ self.assertEqual(shared, {"reasoning_effort": "high"}) ++ ++ def test_changed_values_are_sent_whole(self): ++ body, shared = self.sampling(temperature=0.3) ++ self.assertEqual(body, {"temperature": 0.3, "top_p": 0.95, "top_k": 20}) ++ self.assertEqual(shared["temperature"], 0.3) ++ body, _ = self.sampling(top_k=40) ++ self.assertEqual(body, {"temperature": 1.0, "top_p": 0.95, "top_k": 40}) ++ body, _ = self.sampling(temperature=0.6, top_p=0.9) ++ self.assertEqual(body["temperature"], 0.6) ++ ++ def test_zero_is_greedy(self): ++ body, shared = self.sampling(temperature=0) ++ self.assertEqual(body, {"temperature": 0}) ++ self.assertEqual(shared["temperature"], 0) ++ ++ def test_the_old_default_is_shown_as_the_new_one(self): ++ src = APP.read_text(encoding="utf-8") ++ self.assertRegex(src, r"if \(samplingUntouched\(settings\)\) settings\.temperature = DEFAULTS\.temperature;") ++ self.assertIsNone(re.search(r"temperature: 0\.6", src)) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_image_omission.py b/serve/test_image_omission.py +new file mode 100644 +--- /dev/null ++++ b/serve/test_image_omission.py +@@ -0,0 +1,69 @@ ++"""The local text-only image switch must run before incremental prompt encoding.""" ++import copy ++import os ++import unittest ++from pathlib import Path ++from unittest import mock ++ ++from serve.frontend import ChatTemplate ++from serve.server import ByteTokenizer, MockEngine, Service ++ ++ ++ROOT = Path(__file__).resolve().parents[1] ++ ++ ++class ImageOmission(unittest.TestCase): ++ def service(self): ++ tok = ByteTokenizer() ++ return Service(MockEngine(tok, "ok", max_context=1 << 20), tok, ++ ChatTemplate(ROOT / "serve/chat_template.jinja")) ++ ++ def messages(self): ++ return [{"role": "user", "content": [ ++ {"type": "text", "text": "describe literally"}, ++ {"type": "image", "source": "unused.png"}, ++ ]}] ++ ++ @mock.patch.dict(os.environ, {"STRATA_OMIT_IMAGES_WHEN_NO_VISION": "1"}) ++ def test_omitted_images_are_cached_as_text_and_retries_are_stable(self): ++ svc = self.service() ++ original = self.messages() ++ messages = copy.deepcopy(original) ++ text = messages[0]["content"][0] ++ first, _, _ = svc.prepare(messages, None, {}, 16) ++ self.assertIs(messages[0]["content"][0], text) ++ self.assertEqual(messages[0]["content"][1], { ++ "type": "text", "text": "[image omitted on text-only engine]"}) ++ self.assertEqual(svc.prompts.entries[0][2], first) ++ full = self.service() ++ full.prompts = None ++ self.assertEqual(first, full.encode_prompt(messages, None, {})) ++ again, _, _ = svc.prepare(copy.deepcopy(original), None, {}, 16) ++ self.assertEqual(first, again) ++ self.assertGreater(svc.prompts.last_reused, 0) ++ same_object, _, _ = svc.prepare(messages, None, {}, 16) ++ self.assertEqual(first, same_object) ++ ++ @mock.patch.dict(os.environ, {"STRATA_OMIT_IMAGES_WHEN_NO_VISION": "0"}) ++ def test_disabled_switch_keeps_the_normal_vision_error(self): ++ svc = self.service() ++ messages = self.messages() ++ before = copy.deepcopy(messages) ++ with self.assertRaisesRegex(ValueError, "without the vision encoder"): ++ svc.prepare(messages, None, {}, 16) ++ self.assertEqual(messages, before) ++ ++ @mock.patch.dict(os.environ, {"STRATA_OMIT_IMAGES_WHEN_NO_VISION": "1"}) ++ def test_configured_vision_is_not_omitted(self): ++ svc = self.service() ++ svc.vision = object() ++ messages = self.messages() ++ before = copy.deepcopy(messages) ++ with mock.patch.object(svc, "_note_unreadable_tool_images", side_effect=RuntimeError("before encoding")): ++ with self.assertRaisesRegex(RuntimeError, "before encoding"): ++ svc.prepare(messages, None, {}, 16) ++ self.assertEqual(messages, before) ++ ++ ++if __name__ == "__main__": ++ unittest.main() diff --git a/results/round39/serve-port.patch b/results/round39/serve-port.patch new file mode 100644 index 0000000..d2b1f3d --- /dev/null +++ b/results/round39/serve-port.patch @@ -0,0 +1,2798 @@ +diff --git a/serve/frontend.py b/serve/frontend.py +index c36a44da..3fda9c56 100644 +--- a/serve/frontend.py ++++ b/serve/frontend.py +@@ -637,6 +637,11 @@ def anthropic_to_messages(req: dict, think_unasked: bool = True) -> tuple[list[d + # "output_config": {"effort": "low" | "medium" | "high"} + thinking = req.get("thinking") + effort = (req.get("output_config") or {}).get("effort") if isinstance(req.get("output_config"), dict) else None ++ # 0.1.45: Anthropic's budget_tokens is a real budget, not only a level: the thinking is closed (wrap-up, , ++ # the answer) after that many tokens. A reasoning_budget_tokens of the request's own wins. ++ bt = thinking.get("budget_tokens") if isinstance(thinking, dict) and thinking.get("type") != "disabled" else None ++ if isinstance(bt, int) and not isinstance(bt, bool) and bt > 0 and req.get("reasoning_budget_tokens") is None: ++ req["reasoning_budget_tokens"] = bt + if isinstance(thinking, dict) and thinking.get("type") == "disabled": + kwargs["enable_thinking"] = False + elif effort: +@@ -1039,6 +1044,7 @@ class OutputParser: + return out + + def _reset_scan(self): ++ self._call_end_checked = 0 # closing-tag search progress for the current wrapped call + self.sp = 0 # how much of self.buf (the call body) the scanner has consumed + self.ss = "name" # name -> between -> str|raw -> ... -> done + self.scall = None # the ToolCall being streamed (its id is reused by the final event) +@@ -1046,6 +1052,20 @@ class OutputParser: + self.sval_started = False + self.sdeclared = {} + ++ def _wrapped_call_end(self) -> int: ++ # Without a closing tag the structural walk cannot finish. Search only ++ # newly appended text plus enough overlap for a tag split across feeds. ++ # A literal closing tag inside a value needs structural disambiguation: ++ # once any candidate appears, keep the original full walk for this call. ++ # Recovery can rewrite the buffer, so it always uses the full walk. ++ if not self.recover and self._call_end_checked >= 0: ++ start = max(0, self._call_end_checked - len(CALL_END) + 1) ++ if self.buf.find(CALL_END, start) < 0: ++ self._call_end_checked = len(self.buf) ++ return -1 ++ self._call_end_checked = -1 ++ return call_end(self.buf) ++ + def _scan(self) -> list[Event]: + """Advance the streaming view of the call body in self.buf (see stream_tools).""" + out = [] +@@ -1152,10 +1172,19 @@ class OutputParser: + def _hold(self, text: str, tags: tuple[str, ...]) -> int: + """Length of the longest suffix of `text` that is a proper prefix of one of `tags`.""" + best = 0 ++ size = len(text) + for tag in tags: +- for n in range(1, len(tag)): +- if text.endswith(tag[:n]): +- best = max(best, n) ++ if len(tag) <= 1: ++ continue ++ # Only a suffix beginning with the tag's first character can match. ++ # Ordinary token deltas contain no '<': one C-level search replaces ++ # checking every possible prefix. Earliest candidate is the longest. ++ start = text.find(tag[0], max(0, size - len(tag) + 1)) ++ while start >= 0 and size - start > best: ++ if tag.startswith(text[start:]): ++ best = size - start ++ break ++ start = text.find(tag[0], start + 1) + return best + + def feed(self, delta: str) -> list[Event]: +@@ -1351,7 +1380,7 @@ class OutputParser: + elif after and not CALL_END.startswith(after): + drop = 0 + else: +- i = call_end(self.buf) ++ i = self._wrapped_call_end() + if i >= 0: + drop = len(CALL_END) + if self.recover: # a batch in one wrapper: the next call follows this one's +diff --git a/serve/logit_bias.py b/serve/logit_bias.py +index 59873a77..f9765bcf 100644 +--- a/serve/logit_bias.py ++++ b/serve/logit_bias.py +@@ -1,9 +1,12 @@ + """Request-scoped token biases. Normalize before HTTP streaming starts.""" ++import json + import math + import re ++import threading ++from collections import OrderedDict + + +-def normalize(value, vocab_size=None): ++def _normalize(value, vocab_size=None): + if value is None: + return {} + if isinstance(value, dict): +@@ -36,9 +39,67 @@ def normalize(value, vocab_size=None): + return out + + ++# A long list (a standing ban of a whole script is tens of thousands of IDs) is sent with every request and used to be ++# normalized five times per request: in the OpenAI route, twice in validate_request and twice in engine_key (sampling_keys ++# runs for the cache key and again for the GEN line), about 170 ms for 65,929 entries. A list is now normalized once: ++# - the same object again in one request is answered from `_last`; ++# - an equal list in a later request is found by its JSON text, an exact key that tells False from 0 and "1" from 1. ++# Only a list that normalized without an error is stored, so a bad list still raises on every call. Short lists skip the ++# cache: they cost less than the key. ++_CACHE_MIN = 64 ++_CACHE_MAX = 4 ++_cache = OrderedDict() # (vocab_size, JSON text) -> [normalized dict, engine key or None] ++_last = None # [raw value, its length, {vocab_size: entry}]: the object the previous call was given ++_lock = threading.Lock() ++ ++ ++def _entry(value, vocab_size): ++ global _last ++ if not isinstance(value, (dict, list)) or len(value) < _CACHE_MIN: ++ return [_normalize(value, vocab_size), None] ++ last = _last ++ if last is not None and last[0] is value and last[1] == len(value): ++ hit = last[2].get(vocab_size) ++ if hit is not None: ++ return hit ++ else: ++ last = None ++ try: ++ key = (vocab_size, json.dumps(value, separators=(",", ":"))) ++ except (TypeError, ValueError): ++ key = None # something JSON cannot write: _normalize raises for it ++ entry = None ++ if key is not None: ++ with _lock: ++ entry = _cache.get(key) ++ if entry is not None: ++ _cache.move_to_end(key) ++ if entry is None: ++ entry = [_normalize(value, vocab_size), None] ++ if key is not None: ++ with _lock: ++ _cache[key] = entry ++ while len(_cache) > _CACHE_MAX: ++ _cache.popitem(last=False) ++ if key is not None: ++ if last is None: ++ last = [value, len(value), {}] ++ _last = last ++ last[2][vocab_size] = entry ++ return entry ++ ++ ++def normalize(value, vocab_size=None): ++ return dict(_entry(value, vocab_size)[0]) ++ ++ + def engine_key(value): +- entries = normalize(value) +- return " logit_bias=" + ",".join(f"{i}:{b:g}" for i, b in sorted(entries.items())) if entries else "" ++ entry = _entry(value, None) ++ if not entry[0]: ++ return "" ++ if entry[1] is None: ++ entry[1] = " logit_bias=" + ",".join(f"{i}:{b:g}" for i, b in sorted(entry[0].items())) ++ return entry[1] + + + def validate_request(req, engine, vocab_size): +diff --git a/serve/runconfig.py b/serve/runconfig.py +index 472b0021..842554ab 100644 +--- a/serve/runconfig.py ++++ b/serve/runconfig.py +@@ -18,7 +18,7 @@ from pathlib import Path + # of the "sampling" block, ("arg", flag) for an engine option kept in "args". + EDITABLE = [ + ("sampling.temperature", ("sampling", "num>=0"), +- "Default temperature for requests that send none (0 = greedy, the default without a sampling block)"), ++ "Default temperature for requests that send none (0 = greedy; without it the model card's: 1.0 thinking, 0.7 not)"), + ("sampling.top_p", ("sampling", "0, with no answer +@@ -254,6 +255,15 @@ def _timeout_env(name: str, default: float) -> float | None: + # still working is never ended by this, only reported once (the recommendation: engine_silence_s, or this variable). + # STRATA_ENGINE_STALL_S overrides it (0 = off). Needs psutil (setup installs it); without it nothing is ended. + ENGINE_STALL_S = _timeout_env("STRATA_ENGINE_STALL_S", 90.0) ++# --batch: new requests take the control lines (one prompt read at a time) short prompts first: a prompt of at most ++# STRATA_ADMIT_SHORT_FIRST tokens goes before longer ones (opt-in; unset or 0: whoever gets the lock first), and a ++# longer one that waited STRATA_ADMIT_MAX_WAIT_S (default 30) counts as short. A chat question then waits for the one ++# read in progress, not for every agent read that came before it. ++try: ++ ADMIT_SHORT_FIRST = max(0, int(os.environ.get("STRATA_ADMIT_SHORT_FIRST") or 0)) ++except ValueError: ++ ADMIT_SHORT_FIRST = 0 ++ADMIT_MAX_WAIT_S = _timeout_env("STRATA_ADMIT_MAX_WAIT_S", 30.0) or 0.0 + STALL_CPU_EPS_S = 0.5 # CPU seconds the engine may use in the whole stall window and still count as frozen + STALL_IO_EPS_B = 1 << 20 # bytes read or written in that window + +@@ -705,6 +715,19 @@ class ConvCacheLog: + return dict(self.state) + + ++def conversation_slots(engine) -> int | None: ++ """How many conversations the engine's cache holds: its --conversation-cache-slots argument, which is known ++ before the engine has started (a lazy start has reported no INFO line yet), else the INFO line it reported.""" ++ args = engine.spawn[1] if hasattr(engine, "spawn") else [] ++ if "--conversation-cache-slots" in args[:-1]: ++ try: ++ return int(args[args.index("--conversation-cache-slots") + 1]) ++ except ValueError: ++ pass ++ slots = (getattr(engine, "info", None) or {}).get("conversation_cache_slots") ++ return slots if isinstance(slots, int) else None ++ ++ + def conversation_cache_view(info: dict, hist: list, totals: dict, parked: dict) -> dict: + """#596: the Monitor's Conversation cache card: the parked conversations (the engine's opt-in + --conversation-cache-mib: budget, slots, what its log says) and how much of the prompts the cache gave back.""" +@@ -1096,6 +1119,13 @@ class StrataEngine: + self.last.update(prompt_read=int(f[14])) + if len(f) >= 16: # #588 (engine 0.1.39+): routed experts read over PCIe + self.last.update(offloaded=int(f[15])) ++ for tok in f[16:]: # key=value fields after the positional ones (the prompt's ++ key, eq, value = tok.partition("=") # first chunk and its streamed expert blobs) ++ if eq and key in ("chunk", "experts_streamed"): ++ try: ++ self.last[key] = int(value) ++ except ValueError: ++ pass + + def vram(self, reserve_mib: int | None, timeout: float = 120.0) -> dict: + """#533: `VRAM ` between requests (the caller holds the service's FIFO): the engine shrinks its +@@ -1173,6 +1203,10 @@ class StrataEngine: + v = tune.get(k) + if isinstance(v, (int, float)) and not isinstance(v, bool) and 0.0 <= float(v) <= 1.0: + keys += f" {k}={float(v)!r}" ++ # a layer split's prompt chunks (--prefill-pipe): 0 off, 1 evened out, > 1 the rig's b/a in tokens ++ v = tune.get("prefill_pipe") ++ if isinstance(v, (int, float)) and not isinstance(v, bool) and 0.0 <= float(v) <= 262144.0: ++ keys += f" pipe={float(v)!r}" + # --aux-cpus for this request: 1 puts the engine's threads that are neither pool workers nor the host on the + # spare CPUs, 0 leaves them where they were created (true/false or 1/0; the engine started without a spare CPU + # ignores a 1). Anything else is not sent. +@@ -1416,10 +1450,11 @@ class StrataEngine: + + YIELDS_MAX = 2 # #656: how often one request's prompt read gives way to a shorter waiting one + +- def _take_control(self, cancel, plen: int, after_epoch: int | None = None): ++ def _take_control(self, cancel, plen: int, after_epoch: int | None = None, new: bool = False): + """Waits for the control lines (one prompt read at a time), yielding None heartbeats; False when cancelled. +- `after_epoch`: a request whose read gave way lets the requests waiting then go first.""" +- entry = [plen] ++ `after_epoch`: a request whose read gave way lets the requests waiting then go first. ++ `new`: the request's first admission (STRATA_ADMIT_SHORT_FIRST orders only these).""" ++ entry = [plen, time.monotonic(), new] # (length, arrival, ordered: STRATA_ADMIT_SHORT_FIRST) + with self.slot_cv: + self.waiting += 1 + self.wait_lens.append(entry) +@@ -1433,7 +1468,17 @@ class StrataEngine: + raise EngineDied("the engine stopped while this request waited; it was not sent") + with self.slot_cv: + turn = after_epoch is None or self.ctl_epoch > after_epoch or self.waiting <= 1 ++ if turn and new and ADMIT_SHORT_FIRST > 0: ++ turn = self._admit_next() is entry + if turn and self.ctl.acquire(timeout=0.5): ++ if new and ADMIT_SHORT_FIRST > 0: ++ with self.slot_cv: # a shorter one came while this one waited in acquire() ++ first = self._admit_next() is entry ++ if not first: ++ self.ctl.release() ++ with self.slot_cv: ++ self.slot_cv.notify_all() ++ continue + if self.alive(): + break + self.ctl.release() # it died (or is being restarted) while this one was getting +@@ -1456,6 +1501,18 @@ class StrataEngine: + self.slot_cv.notify_all() + return True + ++ def _admit_next(self): ++ """STRATA_ADMIT_SHORT_FIRST= (the caller holds slot_cv): the new request that takes the control lines ++ next - one whose prompt is at most that long (it reads little) before a longer one, in arrival order within ++ each; a longer one that has waited STRATA_ADMIT_MAX_WAIT_S counts as short (it is never held back for good). ++ Off (unset or 0): whoever gets the lock first, as before. Only the order changes, never what is read.""" ++ now = time.monotonic() ++ ordered = [e for e in self.wait_lens if len(e) > 2 and e[2]] ++ if not ordered: ++ return None ++ return min(ordered, key=lambda e: (0 if e[0] <= ADMIT_SHORT_FIRST or now - e[1] >= ADMIT_MAX_WAIT_S else 1, ++ e[1])) ++ + SOLO_AGAIN_MAX = 2 # how often a request left alone in a slot goes back to the solo path + SOLO_AGAIN_MIN_LEFT = 32 # ... only with at least this many tokens still allowed (max_tokens) + +@@ -1493,7 +1550,7 @@ class StrataEngine: + prompt, left = list(ids), int(max_new) + reused0 = None # the first read's reused tokens: a continued leg's count includes this request's own output + carry = {} # #1869: the figures of the legs that ended in a hand-back to the solo path (the final DONE adds to them) +- ok = yield from self._take_control(cancel, len(prompt)) ++ ok = yield from self._take_control(cancel, len(prompt), new=True) + if not ok: + return + holding, born = True, self.gen # (the engine it now has the control lines of) +@@ -1803,6 +1860,7 @@ class StrataEngine: + + slot_groups = 1 # --batch-groups (set when the engine starts) + SLOT_PREFIX_MIN = 512 # a held prefix shorter than this is not worth unbalancing the pipeline's groups ++ TURN_ID = 248045 # <|im_start|>: the engine's default --turn-token (its turn checkpoint sits at the last one) + + def pick_slot(self, prompt: list[int]) -> int | None: + """A free slot for `prompt` (the caller holds slot_cv): the one whose held tokens are the longest start of the +@@ -1812,18 +1870,30 @@ class StrataEngine: + if not free: + return None + gs = getattr(self, "slot_gs", 0) +- if gs: +- # pipelined groups: the engine keeps no pipelined slot as a conversation cache, so nothing is worth keeping +- # free for a next turn. A group's window spans its slots up to its last busy one, so first fill a hole +- # below a busy slot of the same group (it costs no extra row), then follow the spreading order ++ if gs > 1: ++ # pipelined groups of several slots: the engine keeps no idle slot below a busy one as a conversation cache ++ # (its pad row writes the slot), so nothing is worth keeping free for a next turn. A group's window spans ++ # its slots up to its last busy one, so first fill a hole below a busy slot of the same group (it costs no ++ # extra row), then follow the spreading order + # (reported on #793: least-recently-used picks left leading holes, [_ B], that cost a whole row each) + def fills_hole(b): + g0 = b - b % gs + return any(self.slot_busy[x] for x in range(b + 1, g0 + gs)) + return min(free, key=lambda b: (not fills_hole(b), self.slot_order.index(b))) + def held_prefix(b): ++ # what of `prompt` the slot's cache serves: the slot's whole held stream (the prompt and what it decoded) ++ # when this prompt continues it, else its turn checkpoint - the engine takes one at the last <|im_start|> of ++ # a prompt (the next request of the chat renders the history again, but not the reply it decoded or the ++ # assistant header), so a prompt that agrees with the held tokens up to there is that chat's next turn ++ # too (#2518: it was only found by the engine's own source search, after pick_slot had sent it to another ++ # slot and the request in that slot had overwritten the chat whose cache the next one needed) + h = self.slot_held[b] +- return len(h) if h and len(h) < len(prompt) and prompt[:len(h)] == h else 0 ++ if not h: ++ return 0 ++ if len(h) < len(prompt) and prompt[:len(h)] == h: ++ return len(h) ++ k = len(h) - 1 - h[::-1].index(self.TURN_ID) if self.TURN_ID in h else 0 ++ return k if 0 < k < len(prompt) and prompt[:k] == h[:k] else 0 + best = max(free, key=held_prefix) + # without --batch-groups a held prefix of any length is taken, as before + if held_prefix(best) >= (self.SLOT_PREFIX_MIN if self.slot_groups > 1 else 1): +@@ -2868,23 +2938,32 @@ class ByteTokenizer: + + ALWAYS = () # specials matched without parse_special (type 4, as ) + ++ @property ++ def max_special_len(self): ++ return max(len(s) for s in self.SPECIALS) ++ + @property + def control_tokens(self): + return [s for s in self.SPECIALS if s not in self.ALWAYS] + + def encode(self, text, parse_special=False, plain=()): +- out, i = [], 0 ++ return self.encode_marked(text, parse_special, plain)[0] ++ ++ def encode_marked(self, text, parse_special=False, plain=()): ++ """encode() and its resume points (after each special), as strata_tokenizer's for PromptEncoder.""" ++ out, marks, i = [], [], 0 + while i < len(text): + for k, s in enumerate(self.SPECIALS): + if (parse_special or s in self.ALWAYS) and text.startswith(s, i) and not any( + a <= i < b for a, b in plain): + out.append(256 + k) + i += len(s) ++ marks.append((i, len(out))) + break + else: + out.extend(text[i].encode("utf-8")) + i += 1 +- return out ++ return out, marks + + def decode(self, ids, errors="replace"): + raw = bytearray() +@@ -3031,10 +3110,12 @@ class Service: + self.literals = literal_tags(getattr(tokenizer, "control_tokens", ())) # texts that stay text inside a message + self.fit_max_tokens = fit_max_tokens # --fit-max-tokens: clamp the output cap instead of 400 + self.aliases: list[str] = [] # #297: other names of the model (the config's `aliases`) ++ self.answer_room = True # 0.1.45: see ANSWER_ROOM_MIN_CAP (the config's "answer_room": false turns it off) + self.reasoning_close_retry = False # #1053 (opt-in): close the thinking once when a reply ends inside it + self.codex_compaction_cache = False # #924 (opt-in): a Codex compaction is rendered with its conversation's tools + self.codex_thread_titles = False # #923 (opt-in): answer Codex's thread-title turns without the engine + self.sampling_defaults = dict(sampling_defaults or {}) # the run config's `sampling` block ++ self.model_arch = None # the config's family ("qwen4exp", "qwen35moe"): None = no model-card defaults + self.shared = {} # the web app's Chat settings for every client (POST /settings) + self.shared_path = None # where they are kept between starts (next to the config) + self.config_effort = None # #1641: the config's "reasoning_effort": the level for requests with none +@@ -3075,6 +3156,9 @@ class Service: + "prompt_ms": 0.0, "decode_ms": 0.0, + "drafts_offered": 0, "drafts_accepted": 0} # #457: the MTP drafts, summed where reported + self.last_timings = None # the last finished request's, llama.cpp's names (/v1/status) ++ # /metrics/prometheus' two throughput gauges need the totals as they stood at the last scrape, because ++ # llama.cpp averages them over the window between two scrapes (#2341). ++ self.scrape_bucket = (0, 0.0, 0, 0.0) + from serve.prometheus import Latencies + self.latencies = Latencies() # GET /metrics in Prometheus text: the latency histograms + self.last_request_at = None # when a request last started or finished +@@ -3104,6 +3188,17 @@ class Service: + self.restore_after_load = None # a session file an image job saved before it unloaded + self.stop_ids = set(tokenizer.encode(IM_END, parse_special=True) + + tokenizer.encode("<|endoftext|>", parse_special=True)) ++ # #567: a prompt re-encodes only what follows the last special token it shares with a recent prompt (the ++ # same ids as a full encode: tools/strata_tokenizer.py PromptEncoder). Tokenizers without resume points ++ # encode in full. ++ self.prompts = None ++ if hasattr(tokenizer, "encode_marked") and hasattr(tokenizer, "max_special_len"): ++ from strata_tokenizer import PromptEncoder ++ # one kept prompt per conversation the engine's cache can hold (--conversation-cache-slots); the ++ # encoder's own default without that setting ++ slots = conversation_slots(engine) ++ self.prompts = PromptEncoder(tokenizer, keep=slots) if slots is not None and slots > 0 \ ++ else PromptEncoder(tokenizer) + + def loaded(self) -> bool: + return not hasattr(self.engine, "alive") or self.engine.alive() +@@ -3296,20 +3391,17 @@ class Service: + threading.Thread(target=loop, daemon=True).start() + + def free_vram_mib(self) -> int | None: +- """Free VRAM on the engine's (first) GPU, from NVML (AMD backend: amdgpu's sysfs files, #301); None when it can't +- be read (then nothing is refused).""" ++ """Free VRAM on the engine's (first) GPU, from the backend's reader - NVML, the amdgpu sysfs files (#301) or ++ IOKit on macOS (Metal backend); None when it can't be read (then nothing is refused).""" + try: +- if getattr(self, "backend", None) == "hip": +- from serve.telemetry import free_vram_mib +- return free_vram_mib(int(getattr(self, "gpu_index", 0) or 0), amd=True) +- from serve.telemetry import _Nvml +- nv = _Nvml(int(getattr(self, "gpu_index", 0) or 0)) +- if not nv.ok(): ++ from serve.telemetry import gpu_reader ++ g = gpu_reader(int(getattr(self, "gpu_index", 0) or 0), amd=getattr(self, "backend", None) == "hip") ++ if not g.ok(): + return None +- m = nv.Mem() +- if nv.lib.nvmlDeviceGetMemoryInfo(nv.dev, ctypes.byref(m)) != 0: ++ r = g.read() ++ if r.get("mem_total") is None or r.get("mem_used") is None: + return None +- return int(m.free >> 20) ++ return int((r["mem_total"] - r["mem_used"]) >> 20) + except Exception: + return None + +@@ -3552,7 +3644,7 @@ class Service: + if effort: + if api == "openai": + ctk = req.get("chat_template_kwargs") if isinstance(req.get("chat_template_kwargs"), dict) else {} +- if not req.get("reasoning_effort") and not req.get("reasoning") and \ ++ if req.get("reasoning_effort") in (None, "") and not req.get("reasoning") and \ + "enable_thinking" not in ctk and "reasoning_effort" not in ctk: + req["reasoning_effort"] = effort + elif not req.get("thinking") and not req.get("output_config"): +@@ -3703,8 +3795,15 @@ class Service: + with self.status_lock: + s, t = dict(self.status), dict(self.totals) + running = len(self.live_reqs) or int(bool(s.get("busy"))) +- prompt_n = t["prompt_tokens"] - t["reused"] # what was read, as llama.cpp counts it +- prompt_s, decode_s = t["prompt_ms"] / 1000, t["decode_ms"] / 1000 ++ prompt_n = t["prompt_tokens"] - t["reused"] # what was read, as llama.cpp counts it ++ prompt_s, decode_s = t["prompt_ms"] / 1000, t["decode_ms"] / 1000 ++ # llama.cpp reads these two gauges from a bucket it resets on every scrape (server_metrics::bucket, ++ # "the gauges are averaged over the window between two scrapes"), so an idle server reports 0 rather ++ # than its lifetime average. Keep the same window here, or the names lie to llama.cpp dashboards. ++ was_prompt_n, was_prompt_s, was_predicted_n, was_decode_s = self.scrape_bucket ++ self.scrape_bucket = (prompt_n, prompt_s, t["output_tokens"], decode_s) ++ d_prompt_n, d_prompt_s = prompt_n - was_prompt_n, prompt_s - was_prompt_s ++ d_predicted_n, d_decode_s = t["output_tokens"] - was_predicted_n, decode_s - was_decode_s + rows = [ + ("llamacpp:prompt_tokens_total", "counter", "Prompt tokens read", prompt_n), + ("llamacpp:prompt_tokens_cached_total", "counter", "Prompt tokens the conversation cache already held", +@@ -3712,10 +3811,10 @@ class Service: + ("llamacpp:prompt_seconds_total", "counter", "Prompt read time", prompt_s), + ("llamacpp:tokens_predicted_total", "counter", "Tokens generated", t["output_tokens"]), + ("llamacpp:tokens_predicted_seconds_total", "counter", "Generation time", decode_s), +- ("llamacpp:prompt_tokens_seconds", "gauge", "Average prompt read throughput in tokens/s", +- prompt_n / prompt_s if prompt_s else 0), +- ("llamacpp:predicted_tokens_seconds", "gauge", "Average generation throughput in tokens/s", +- t["output_tokens"] / decode_s if decode_s else 0), ++ ("llamacpp:prompt_tokens_seconds", "gauge", "Average prompt read throughput in tokens/s, since the last scrape", ++ d_prompt_n / d_prompt_s if d_prompt_s > 0 else 0), ++ ("llamacpp:predicted_tokens_seconds", "gauge", "Average generation throughput in tokens/s, since the last scrape", ++ d_predicted_n / d_decode_s if d_decode_s > 0 else 0), + ("llamacpp:requests_processing", "gauge", "Requests running", running), + ("llamacpp:requests_deferred", "gauge", "Requests waiting their turn", s.get("queued") or 0), + ("strata:requests_total", "counter", "Requests finished", t["requests"]), +@@ -3787,17 +3886,21 @@ class Service: + tail = "" if off else EFFORT_TURN.format(EFFORT_TEXT.get(effort, EFFORT_TEXT["medium"])) + return history + tail + full[len(mine):] + +- def encode_prompt(self, messages, tools, kwargs) -> list[int]: ++ def encode_prompt(self, messages, tools, kwargs, store: bool = True) -> list[int]: + """The request's prompt: the template rendered and tokenized. #537: a / written inside a + message's text is encoded as the text it is, not as the model's reasoning markers (the template's own are). + A control token's text (<|im_start|>, <|im_end|>, <|endoftext|>, ...) inside a message is text as well: only +- the control tokens the template writes are control tokens.""" ++ the control tokens the template writes are control tokens. `store=False` (a token count) reuses the kept ++ prompts without adding or replacing one.""" + marked, marked_tools, changed = mark_think_literals(messages, tools, self.literals) + prompt = self.render_prompt(marked, marked_tools, kwargs) +- if not changed: +- return self.tok.encode(prompt, parse_special=True) +- prompt, plain = unmark_think_literals(prompt, self.literals) +- return self.tok.encode(prompt, parse_special=True, plain=plain) ++ plain = () ++ if changed: ++ prompt, plain = unmark_think_literals(prompt, self.literals) ++ if self.prompts is not None: # #567: only what follows the last special shared with a recent prompt ++ return self.prompts.encode(prompt, plain, store) ++ return self.tok.encode(prompt, parse_special=True, plain=plain) if plain else \ ++ self.tok.encode(prompt, parse_special=True) + + @contextlib.contextmanager + def _turn(self, cancel=None): +@@ -3845,6 +3948,11 @@ class Service: + content[n] = {"type": "text", "text": f"[image omitted: {why}]"} + return fetched + ++ def forget_prompt(self, ids): ++ if self.prompts is not None: ++ self.prompts.forget(ids) ++ ++ + def prepare(self, messages, tools, kwargs, max_new=None, force=None, req=None, cancel=None): + """-> (ids, thinking, max_new). An unset or non-positive max_new (some clients send -1) means "unlimited": + the rest of the context. `force` (forced_call): without thinking the reply starts with it, so it ends the +@@ -3852,7 +3960,7 @@ class Service: + # `cancel` (#1999): set when the client hangs up; the downloads, the wait for the turn and the encodes stop + # between pictures and raise ClientGone + fetched = self._note_unreadable_tool_images(messages, cancel) +- ids = self.encode_prompt(messages, tools, kwargs) ++ ids = encoded_ids = self.encode_prompt(messages, tools, kwargs) + if force and kwargs.get("enable_thinking", True) is False: + ids = ids + self.tok.encode(force, parse_special=True) + self.embeddings.path = None +@@ -3914,12 +4022,16 @@ class Service: + room = ctx - CTX_SLACK - len(ids) + if max_new is None or max_new <= 0 or (self.fit_max_tokens and room < 1): + if room < 1: ++ # A prompt refused for the context is not kept (#567): the client cannot resend it unchanged. The ++ # other refusals (engine starting, images) keep it: the same prompt comes back and its encode is reused. ++ self.forget_prompt(encoded_ids) + raise ValueError("request exceeds the context window: " + f"prompt ({len(ids)} tokens) leaves no room to answer in the context " + f"({ctx}); requests are never truncated") + max_new = room + elif max_new > room: + if not self.fit_max_tokens: ++ self.forget_prompt(encoded_ids) + raise ValueError("request exceeds the context window: " + f"prompt ({len(ids)} tokens) + max tokens ({max_new}) exceeds the context " + f"({ctx}); requests are never truncated. Send a smaller " +@@ -4076,12 +4188,23 @@ class Service: + def _run(self, ids, thinking, tools, max_new, sampling, cancel, force=None) -> Iterator[tuple[str, object]]: + """Yields ("event", Event) as text arrives, then ("done", {"finish": .., "completion_tokens": ..}). + `force` (forced_call): the opening of the call the reply must make - see prepare().""" +- defaults = {**self.sampling_defaults, **self.shared} # the config's, then the Chat settings shared with apps ++ # the model card's sampling for this thinking state, then the config's block, then the Chat settings shared with ++ # apps; the request's own fields win over all of them ++ defaults = {**model_card_sampling(self.model_arch, thinking), **self.sampling_defaults, **self.shared} + if defaults: # the request's own fields win (explicit 0 stays greedy) + req_values = {k: v for k, v in (sampling or {}).items() if v is not None} + sampling = {**defaults, **req_values} + # #123: read after the merge, so a budget shared through POST /settings is seen like the other keys + budget = self.reasoning_budget(sampling) if thinking else None ++ if thinking and not budget and self.answer_room and max_new and max_new >= ANSWER_ROOM_MIN_CAP: ++ # 0.1.45: max_tokens counts the thinking too. A client that caps the reply at 8192 (the desktop agent's ++ # harness sends max_completion_tokens: 8192 on every request) and a task the model thinks long about got ++ # the cut-off in the middle of the thinking and NO answer. So thinking is closed early enough to leave ++ # the answer its quarter of the cap (the budget's wrap-up: a short closing sentence, , the answer). ++ # A request with a budget of its own, or max_tokens below ANSWER_ROOM_MIN_CAP, is left as it was. ++ budget = max_new - min(max(ANSWER_RESERVE_MIN, max_new // 4), max_new // 2) ++ if max_new <= 32768: print(f"[strata] thinking limited to {budget} of max_tokens {max_new} so the answer has room " ++ "(\"answer_room\": false in the config turns this off)", flush=True) + if budget and max_new and max_new > 0: + # #984: a budget at or near max_tokens never fires before the reply is cut off while still thinking, so + # it has no answer. Close the thinking early enough to leave room for one. +@@ -4586,13 +4709,13 @@ def request_stats(segments: list[dict]) -> dict: + if not segments: + return {} + result = dict(segments[-1]) +- for key in ("reused", "prompt_read"): ++ for key in ("reused", "prompt_read", "chunk"): + if key in segments[0]: + result[key] = segments[0][key] + else: + result.pop(key, None) + for key in ("prompt_ms", "decode_ms", "generated", "drafts_offered", "drafts_accepted", +- "hits", "lookups", "offloaded", "ram_blobs", "file_blobs", "file_mb"): ++ "hits", "lookups", "offloaded", "ram_blobs", "file_blobs", "file_mb", "experts_streamed"): + if any(segment.get(key) is not None for segment in segments): + result[key] = sum(segment.get(key) or 0 for segment in segments) + return result +@@ -5612,7 +5735,7 @@ def make_handler(svc: Service): + if not svc.loaded() and not getattr(svc.engine, "unloaded", False): + self._json(503, {"error": {"message": "the engine is not running"}}) + return +- defaults = {**svc.sampling_defaults, **svc.shared} ++ defaults = {**model_card_sampling(svc.model_arch, True), **svc.sampling_defaults, **svc.shared} + names = {"repetition_penalty": "repeat_penalty", "penalty_last_n": "repeat_last_n"} + params = {names.get(k, k): v for k, v in defaults.items() + if k in ("temperature", "top_p", "top_k", "min_p", "seed", "repetition_penalty", +@@ -6045,7 +6168,7 @@ def make_handler(svc: Service): + messages, tools, kw = anthropic_to_messages(req, svc.anthropic_think_unasked) + stop_strings(req) # the same 400 as the request itself would get + check_request_sampling(req) # ... likewise a sampling field of the wrong type +- self._json(200, {"input_tokens": len(svc.encode_prompt(messages, tools, kw))}) ++ self._json(200, {"input_tokens": len(svc.encode_prompt(messages, tools, kw, store=False))}) + + def _anthropic(self, req): + svc.load() +@@ -6433,6 +6556,43 @@ def clean_shared_defaults(d) -> dict: + return out + + ++def split_joined_args(cfg: dict) -> None: ++ """#2267: an engine option written as ONE element of the config's "args" - "--ple-io mmap" or "--ple-io=mmap" - reaches ++ the engine as a single argument ("unknown argument: --ple-io mmap"; the usual form is two elements, "--ple-io", ++ "mmap"). Recommend, never force: such an element is split into the flag and its value, and said at startup. Only an ++ element that starts with a flag name (--letters-digits-dashes) followed by a space or "=" is touched, and never one ++ that sits in the value place of the flag before it.""" ++ args = cfg.get("args") ++ if not isinstance(args, list): ++ return ++ out, changed = [], [] ++ for a in args: ++ m = re.match(r"^(--[A-Za-z0-9][A-Za-z0-9-]*)(?:\s+|=)(\S.*)$", a) if isinstance(a, str) else None ++ prev = out[-1] if out else None ++ in_value_place = isinstance(prev, str) and re.match(r"^--[A-Za-z0-9][A-Za-z0-9-]*$", prev) is not None ++ if m is None or in_value_place: ++ out.append(a) ++ continue ++ out += [m.group(1), m.group(2).strip()] ++ changed.append(a) ++ if changed: ++ cfg["args"] = out ++ print("[strata] the config's args hold " + ", ".join(repr(c) for c in changed) + " as one element each: split into the " ++ 'flag and its value (write them as two elements, e.g. "--ple-io", "mmap")', flush=True) ++ ++ ++def read_config(path) -> dict: ++ """The run config's JSON. A mistake in it (a missing comma between two args, a trailing comma) is said with its ++ place in the file instead of a traceback (#2267: editing "args" by hand is where it happens).""" ++ text = Path(path).read_text(encoding="utf-8-sig") # Notepad adds a BOM ++ try: ++ return json.loads(text) ++ except json.JSONDecodeError as e: ++ raise SystemExit(f"[strata] {path} is not valid JSON: {e.msg} at line {e.lineno}, column {e.colno}. Check the commas " ++ 'between the entries of "args" (each option and each value is its own quoted entry, e.g. ' ++ '"--ple-io", "mmap",) and that no entry is followed by a comma at the end of a list.') from None ++ ++ + def lift_reasoning_effort_arg(cfg: dict) -> None: + """#1641: "--reasoning-effort high" in a config's "args" is meant for the thinking level, which is not an engine + option (the engine stopped at the unknown flag). It is the config's "reasoning_effort" (none, low, medium or high: +@@ -6615,6 +6775,25 @@ def check_request_sampling(req) -> None: + FAMILY_MODEL_NAMES = {"qwen4exp": "qwen3.8-flash-next", "qwen35moe": "qwen3.6-35b-a3b"} + + ++# The sampling each model card recommends ("Best Practices"), used when neither the request, the Chat settings nor the ++# run config's `sampling` block says otherwise. Greedy is only for a request or config that asks for temperature 0. ++# min_p 0, repetition_penalty 1 and presence_penalty 0 are the engine's own "off", so they are not listed. ++MODEL_CARD_SAMPLING = { ++ # Qwen3.8-Flash-Next (and Gyro, the same base model) ++ "qwen4exp": {"thinking": {"temperature": 1.0, "top_p": 0.95, "top_k": 20}, ++ "instruct": {"temperature": 0.7, "top_p": 0.8, "top_k": 20, "presence_penalty": 1.5}}, ++ # Qwen3.6-35B-A3B: the general-task sets (the card's precise-coding set is temperature 0.6, presence_penalty 0) ++ "qwen35moe": {"thinking": {"temperature": 1.0, "top_p": 0.95, "top_k": 20, "presence_penalty": 1.5}, ++ "instruct": {"temperature": 0.7, "top_p": 0.8, "top_k": 20, "presence_penalty": 1.5}}, ++} ++ ++ ++def model_card_sampling(arch, thinking) -> dict: ++ """The model card's sampling for the family `arch` and the thinking state ({} for an unknown or unset family).""" ++ sets = MODEL_CARD_SAMPLING.get(arch) if arch else None ++ return dict(sets["thinking" if thinking else "instruct"]) if sets else {} ++ ++ + def default_model_name(cfg: dict) -> str: + """The config's "model_name", else its family's name ("arch": qwen4exp, qwen35moe), else Flash-Next's.""" + return cfg.get("model_name") or FAMILY_MODEL_NAMES.get(cfg.get("arch") or "qwen4exp", "qwen3.8-flash-next") +@@ -6667,7 +6846,8 @@ def main() -> int: + "each, requests spread over them (also \"replicas\" in the config: a number, or " + "[{\"gpus\": [0, 1]}, {\"gpus\": [2, 3]}])") + a = ap.parse_args() +- cfg = json.loads(Path(a.config).read_text(encoding="utf-8-sig")) if a.config else {} # Notepad adds a BOM ++ cfg = read_config(a.config) if a.config else {} ++ split_joined_args(cfg) + lift_reasoning_effort_arg(cfg) + if a.replicas is not None: + cfg["replicas"] = a.replicas +@@ -6768,6 +6948,9 @@ def main() -> int: + model_name=default_model_name(cfg), vision=vision, + sampling_defaults=sampling_defaults, + fit_max_tokens=a.fit_max_tokens or cfg.get("fit_max_tokens") is True) ++ if a.engine == "strata": # the model card's sampling when nothing else is said (see MODEL_CARD_SAMPLING) ++ svc.model_arch = cfg.get("arch") or "qwen4exp" ++ svc.answer_room = cfg.get("answer_room") is not False + svc.reasoning_close_retry = cfg.get("reasoning_close_retry") is True # #1053: opt-in, off by default + if cfg.get("images"): # image generation (serve/images.py), set up by setup --images + icfg = {k: (os.path.abspath(os.path.join(cfg.get("cwd") or ".", v)) +diff --git a/serve/telemetry.py b/serve/telemetry.py +index 3232c63c..2e2a28ad 100644 +--- a/serve/telemetry.py ++++ b/serve/telemetry.py +@@ -3,7 +3,9 @@ + A background thread samples once a second and keeps the last 60 readings of each series for the sparklines: + - GPU: NVIDIA's own NVML library (nvml.dll / libnvidia-ml.so.1, installed with every driver) through ctypes, so no + pip package is needed: load, VRAM, temperature, power, PCIe link and throughput. With the AMD backend (#301): the +- amdgpu driver's Linux sysfs files - load, VRAM, temperature and power. ++ amdgpu driver's Linux sysfs files - load, VRAM, temperature and power. On macOS (Metal backend): IOKit's ++ IOAccelerator registry - load and VRAM; the temperature and power sensors need root (`powermetrics`), and an Apple ++ Silicon GPU sits on the die, so there is no PCIe link. + - CPU, RAM, disk: `psutil` when it is installed (setup installs it); without it the CPU and RAM readings fall back to + the OS (Windows GlobalMemoryStatusEx / GetSystemTimes, Linux /proc) and the disk rate is absent. + Anything that cannot be read is None; nothing here can stop the server. +@@ -14,6 +16,8 @@ import collections + import ctypes + import os + import platform ++import re ++import subprocess + import sys + import threading + import time +@@ -246,9 +250,108 @@ class _Amd: + return out + + ++def metal_working_set(): ++ """The Apple GPU's memory budget in bytes: MTLCreateSystemDefaultDevice's recommendedMaxWorkingSetSize, through ++ the Metal framework and objc_msgSend (no PyObjC) - the same property metal/src/mtl_rt.mm reads for the engine ++ (its "GPU working set" log line). None off macOS, or when it cannot be read: a selector the device does not ++ answer would raise an Objective-C exception no ctypes call can catch, so the class is asked first.""" ++ if sys.platform != "darwin": ++ return None ++ try: ++ metal = ctypes.CDLL("/System/Library/Frameworks/Metal.framework/Metal") ++ objc = ctypes.CDLL("/usr/lib/libobjc.dylib") ++ metal.MTLCreateSystemDefaultDevice.restype = ctypes.c_void_p ++ objc.sel_registerName.restype = ctypes.c_void_p ++ objc.object_getClass.argtypes = [ctypes.c_void_p] ++ objc.object_getClass.restype = ctypes.c_void_p ++ objc.class_respondsToSelector.argtypes = [ctypes.c_void_p, ctypes.c_void_p] ++ objc.class_respondsToSelector.restype = ctypes.c_ubyte ++ objc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] ++ objc.objc_msgSend.restype = ctypes.c_ulonglong ++ dev = metal.MTLCreateSystemDefaultDevice() ++ if not dev: ++ return None ++ sel = objc.sel_registerName(b"recommendedMaxWorkingSetSize") ++ if not objc.class_respondsToSelector(objc.object_getClass(dev), sel): ++ return None ++ return int(objc.objc_msgSend(ctypes.c_void_p(dev), ctypes.c_void_p(sel))) or None ++ except (OSError, AttributeError, ValueError): ++ return None ++ ++ ++class _Metal: ++ """macOS: an Apple GPU's readings from the IOKit registry (`ioreg -l -w0 -r -c IOAccelerator`, 22 ms), with ++ _Nvml's interface: the load and the memory are the driver's PerformanceStatistics of the index'th accelerator ++ node - Device Utilization % and the in-use system memory in bytes. The total is the device's Metal working set, ++ the budget one app gets of the unified memory. No temperature or power reading exists without root ++ (`powermetrics`), and an Apple Silicon GPU sits on the die, so the PCIe cards have nothing to show.""" ++ note = ("no GPU temperature or power reading without root (powermetrics), and an Apple Silicon GPU sits on the " ++ "die, so there is no PCIe link to show; the engine's own figures are in its log") ++ cmd = ("ioreg", "-l", "-w0", "-r", "-c", "IOAccelerator") ++ ++ def __init__(self, index=0, run=None): ++ self.index = index ++ self._run = run or self._ioreg ++ self.nodes = len(self._blocks()) ++ self._name = None ++ b = self._block() ++ if b: ++ m = re.search(r'"model" = "([^"]+)"', b) ++ self._name = m.group(1) if m else None ++ ++ @staticmethod ++ def _ioreg(): ++ try: ++ return subprocess.run(list(_Metal.cmd), capture_output=True, text=True, timeout=10).stdout or "" ++ except (OSError, subprocess.SubprocessError): ++ return "" ++ ++ def _blocks(self): ++ """The accelerator nodes, in the registry's order: every node that carries a PerformanceStatistics.""" ++ try: ++ out = self._run() or "" ++ except OSError: # ioreg missing, or a caller's runner that raises ++ return [] ++ return [b for b in re.split(r"^\s*\+-o ", out, flags=re.M) if '"PerformanceStatistics"' in b] ++ ++ def _block(self): ++ b = self._blocks() ++ return b[self.index] if 0 <= self.index < len(b) else None ++ ++ def ok(self): ++ return self.index < self.nodes ++ ++ def name(self): ++ return self._name ++ ++ def read(self): ++ b = self._block() ++ if b is None: ++ return {} ++ ++ def num(key): ++ m = re.search(r'"%s"=([0-9.]+)' % re.escape(key), b) ++ return float(m.group(1)) if m else None ++ ++ out = {} ++ u = num("Device Utilization %") ++ if u is None: # older chips report only the renderer and the tiler ++ r, t = num("Renderer Utilization %"), num("Tiler Utilization %") ++ u = max(v for v in (r, t) if v is not None) if (r is not None or t is not None) else None ++ out["util"] = int(u) if u is not None else None ++ used = num("In use system memory") ++ out["mem_used"] = int(used) if used is not None else None ++ out["mem_total"] = metal_working_set() ++ return out ++ ++ + def gpu_reader(index=0, amd=False): +- """The card's readings: NVML (NVIDIA), or the amdgpu sysfs files with the AMD backend (#301).""" +- return _Amd(index) if amd else _Nvml(index) ++ """The card's readings: NVML (NVIDIA), the amdgpu sysfs files with the AMD backend (#301), or IOKit's ++ IOAccelerator registry on macOS (Metal backend). `amd` comes first: no Mac runs the HIP backend, and the ++ tests' AMD trees are read on a Mac too.""" ++ if amd: ++ return _Amd(index) ++ return _Metal(index) if sys.platform == "darwin" else _Nvml(index) + + + _DISPLAY_CLASS = r"SYSTEM\CurrentControlSet\Control\Class\{4d36e968-e325-11ce-bfc1-08002be10318}" +@@ -413,15 +516,21 @@ class Telemetry: + cards = windows_amd_cards() + win_amd = [cards[i] for i in idx if 0 <= i < len(cards)] or cards[:1] + gpu_name = " + ".join(c["name"] for c in win_amd) if win_amd else None ++ # the macOS reader knows the device's working set (the VRAM budget of the unified memory); None elsewhere ++ vram_total = None ++ if win_amd is None and isinstance(self.gpu, _Metal) and self.gpu.ok(): ++ ws = metal_working_set() ++ vram_total = ws >> 20 if ws else None + self.static = { + "gpu_name": gpu_name, + "gpu_vram_total_mib": (sum(c["vram_bytes"] for c in win_amd) >> 20 +- if win_amd and all(c["vram_bytes"] for c in win_amd) else None), ++ if win_amd and all(c["vram_bytes"] for c in win_amd) else vram_total), + "gpu_count": len(self.gpus), + # #1380: the AMD readings are the amdgpu driver's Linux sysfs files; a Windows AMD card has none yet, and the + # dashboard said "not readable (NVML)" or showed empty tiles with no word why + "gpu_note": ("no GPU load or VRAM readings for AMD cards on Windows yet (Linux reads them from the amdgpu " +- "driver); the engine's own VRAM figures are in its log" if amd and not self.gpu.ok() else None), ++ "driver); the engine's own VRAM figures are in its log" if amd and not self.gpu.ok() ++ else getattr(self.gpu, "note", None)), + "cpu_name": _cpu_name(), + "cores": (self.ps.cpu_count(logical=False) if self.ps else None) or None, + "threads": os.cpu_count(), +diff --git a/serve/test_admit_order.py b/serve/test_admit_order.py +new file mode 100644 +index 00000000..67766e55 +--- /dev/null ++++ b/serve/test_admit_order.py +@@ -0,0 +1,119 @@ ++"""serve/test_admit_order.py - STRATA_ADMIT_SHORT_FIRST: new requests take the control lines short prompts first ++(StrataEngine._take_control / _admit_next; no process, no GPU). ++ ++ python -m unittest serve.test_admit_order -v ++""" ++from __future__ import annotations ++ ++import sys ++import threading ++import time ++import unittest ++from pathlib import Path ++from unittest import mock ++ ++sys.path.insert(0, str(Path(__file__).resolve().parents[1])) ++import serve.server as server # noqa: E402 ++from serve.server import StrataEngine # noqa: E402 ++ ++ ++def engine_without_process(n=2): ++ """A StrataEngine with the admission state of a started one and no process: alive() is `e.up`.""" ++ e = StrataEngine("missing-executable", [], lazy=True) ++ e.batch = n ++ e.slot_busy = [False] * n ++ e.slot_claims, e.slot_releasing, e._orphan_seen = set(), set(), {} ++ e.slot_cv = threading.Condition() ++ e.waiting, e.admitting, e.wait_lens, e.ctl_epoch = 0, 0, [], 0 ++ e.ctl = threading.Lock() ++ e.gen = 1 ++ e.up = True ++ e.alive = lambda: e.up ++ return e ++ ++ ++class AdmitOrderTest(unittest.TestCase): ++ def setUp(self): ++ self.e = engine_without_process(16) ++ self.order, self.cancels = [], [] ++ self.order_lock = threading.Lock() ++ ++ def tearDown(self): ++ for c in self.cancels: ++ c.set() ++ ++ def take(self, name, plen, **kw): ++ """A request waiting for the control lines; once it has them it notes its name and gives them back.""" ++ c = threading.Event() ++ self.cancels.append(c) ++ ++ def go(): ++ gen = self.e._take_control(c, plen, **kw) ++ try: ++ while True: ++ next(gen) ++ except StopIteration as s: ++ if s.value: ++ with self.order_lock: ++ self.order.append(name) ++ time.sleep(0.05) # its read ++ self.e.ctl.release() ++ threading.Thread(target=go, daemon=True).start() ++ ++ def until(self, cond, s=5.0): ++ end = time.time() + s ++ while time.time() < end and not cond(): ++ time.sleep(0.01) ++ self.assertTrue(cond()) ++ ++ def queue(self, *reqs): ++ """While a long read holds the control lines, the requests arrive in this order; then the read ends.""" ++ self.assertTrue(self.e.ctl.acquire(timeout=1)) ++ for name, plen in reqs: ++ self.take(name, plen, new=True) ++ self.until(lambda n=len(self.order) + sum(1 for _ in self.e.wait_lens): self.e.waiting >= 1) ++ time.sleep(0.02) # (distinct arrival times) ++ self.until(lambda: self.e.waiting == len(reqs)) ++ self.e.ctl.release() ++ self.until(lambda: len(self.order) == len(reqs), s=15.0) ++ return self.order ++ ++ def test_short_first_puts_the_chat_before_the_agent_reads(self): ++ with mock.patch.object(server, "ADMIT_SHORT_FIRST", 8192), mock.patch.object(server, "ADMIT_MAX_WAIT_S", 60.0): ++ order = self.queue(("agent1", 120000), ("agent2", 190000), ("chat", 500), ("agent3", 150000)) ++ self.assertEqual(order[0], "chat") ++ self.assertEqual(order[1:], ["agent1", "agent2", "agent3"]) # arrival order among the long ones ++ ++ def test_a_long_request_that_waited_too_long_goes_first(self): ++ with mock.patch.object(server, "ADMIT_SHORT_FIRST", 8192), mock.patch.object(server, "ADMIT_MAX_WAIT_S", 0.2): ++ self.assertTrue(self.e.ctl.acquire(timeout=1)) ++ self.take("agent", 120000, new=True) ++ self.until(lambda: self.e.waiting == 1) ++ time.sleep(0.4) # past STRATA_ADMIT_MAX_WAIT_S ++ self.take("chat", 500, new=True) ++ self.until(lambda: self.e.waiting == 2) ++ self.e.ctl.release() ++ self.until(lambda: len(self.order) == 2, s=15.0) ++ self.assertEqual(self.order, ["agent", "chat"]) ++ ++ def test_off_by_default(self): ++ self.assertEqual(server.ADMIT_SHORT_FIRST if "STRATA_ADMIT_SHORT_FIRST" not in __import__("os").environ ++ else 0, 0) ++ with mock.patch.object(server, "ADMIT_SHORT_FIRST", 0): ++ order = self.queue(("agent1", 120000), ("chat", 500)) ++ self.assertEqual(sorted(order), ["agent1", "chat"]) # both get their turn, in any order ++ ++ def test_a_continuation_is_not_held_back(self): ++ """A request already admitted once (new=False: promoted from the solo path) is not part of the order.""" ++ with mock.patch.object(server, "ADMIT_SHORT_FIRST", 8192), mock.patch.object(server, "ADMIT_MAX_WAIT_S", 60.0): ++ self.assertTrue(self.e.ctl.acquire(timeout=1)) ++ self.take("chat", 500, new=True) ++ self.take("continuation", 190000) ++ self.until(lambda: self.e.waiting == 2) ++ self.e.ctl.release() ++ self.until(lambda: len(self.order) == 2, s=15.0) ++ self.assertEqual(sorted(self.order), ["chat", "continuation"]) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_effort_paths.py b/serve/test_effort_paths.py +index dbc1ae6f..15525f6b 100644 +--- a/serve/test_effort_paths.py ++++ b/serve/test_effort_paths.py +@@ -130,6 +130,7 @@ class EffortPaths(unittest.TestCase): + "thinking": {"type": "enabled", "budget_tokens": 20000}}), "xhigh") + self.assertEqual(self.level("/v1/responses", {"input": "hello", "reasoning": {"effort": "high"}}), "xhigh") + self.assertEqual(self.level("/v1/chat/completions", {"messages": MSG, "reasoning_effort": "none"}), "none") ++ self.assertEqual(self.level("/v1/chat/completions", {"messages": MSG, "reasoning_effort": False}), "none") + + def test_the_chat_settings_shared_with_apps_beat_the_config(self): + self.svc.config_effort = "low" +diff --git a/serve/test_logit_bias.py b/serve/test_logit_bias.py +index ee3751b1..02915522 100644 +--- a/serve/test_logit_bias.py ++++ b/serve/test_logit_bias.py +@@ -59,3 +59,85 @@ class LogitBias(unittest.TestCase): + finally: + httpd.shutdown() + httpd.server_close() ++ ++ ++class LogitBiasCache(unittest.TestCase): ++ """A long list is normalized once; the cache must not change what a call returns or raises.""" ++ ++ def setUp(self): ++ from serve import logit_bias ++ logit_bias._cache.clear() ++ logit_bias._last = None ++ ++ def test_same_answers_as_the_plain_function(self): ++ from serve.logit_bias import _normalize ++ forms = [{str(i): -100 for i in range(200)}, ++ [[i, False] for i in range(200)], ++ {str(i): (i % 7) - 3.5 for i in range(200)}, ++ [[str(i), 2] for i in range(200)]] ++ for vocab in (None, 1000): ++ for form in forms: ++ want = _normalize(form, vocab) ++ for _ in range(3): # the first call fills the cache, the others read it ++ self.assertEqual(normalize(form, vocab), want) ++ self.assertEqual(normalize(json.loads(json.dumps(form)), vocab), want) # an equal list, another object ++ for form in forms: ++ ref = " logit_bias=" + ",".join(f"{i}:{b:g}" for i, b in sorted(_normalize(form).items())) ++ self.assertEqual(engine_key(form), ref) ++ self.assertEqual(engine_key(form), ref) ++ ++ def test_false_is_not_zero(self): ++ ban = [[i, False] for i in range(100)] ++ zero = [[i, 0] for i in range(100)] ++ self.assertTrue(ban == zero) # Python calls them equal, so an == key would mix them up ++ for first, second in [(ban, zero), (zero, ban)]: ++ self.setUp() ++ self.assertEqual(set(normalize(first).values()), {-100.0 if first is ban else 0.0}) ++ self.assertEqual(set(normalize(second).values()), {-100.0 if second is ban else 0.0}) ++ ++ def test_a_bad_list_is_never_stored(self): ++ good = {str(i): -1 for i in range(100)} ++ normalize(good, 1000) ++ for bad in [{**good, "5": True}, {**good, "5": 101}, {**good, "5": float("nan")}, {**good, "5": "x"}]: ++ for _ in range(2): ++ with self.assertRaises(ValueError): ++ normalize(bad, 1000) ++ pairs = [[i, 1] for i in range(100)] ++ normalize(pairs) ++ with self.assertRaises(ValueError): # True == 1, but a bool is not a bias ++ normalize([[i, 1] for i in range(99)] + [[99, True]]) ++ with self.assertRaises(ValueError): # a duplicate ID ++ normalize([[i, 1] for i in range(100)] + [[0, 1]]) ++ ++ def test_vocabulary_size_is_part_of_the_key(self): ++ ids = {str(i): -1 for i in range(100)} ++ self.assertEqual(len(normalize(ids, 1000)), 100) ++ with self.assertRaises(ValueError): ++ normalize(ids, 50) ++ self.assertEqual(len(normalize(ids)), 100) ++ with self.assertRaises(ValueError): ++ normalize(ids, 50) ++ ++ def test_the_result_is_a_copy(self): ++ ids = {str(i): -1 for i in range(100)} ++ first = normalize(ids) ++ first[0] = 99.0 ++ self.assertEqual(normalize(ids)[0], -1.0) ++ self.assertEqual(validate_request({"logit_bias": ids}, SimpleNamespace(info={"logit_bias": 1}), 1000)[0], -1.0) ++ ++ def test_an_equal_list_in_a_later_request_is_a_hit(self): ++ from serve import logit_bias ++ a = {str(i): -100 for i in range(100)} ++ b = json.loads(json.dumps(a)) ++ self.assertIsNot(a, b) ++ key_a = engine_key(a) ++ logit_bias._last = None # a later request: the previous object is gone ++ key_b = engine_key(b) ++ self.assertIs(key_a, key_b) # one entry, one key string ++ self.assertEqual(len(logit_bias._cache), 1) ++ ++ def test_the_cache_is_bounded(self): ++ from serve import logit_bias ++ for n in range(logit_bias._CACHE_MAX + 3): ++ normalize({str(i): -(n + 1) for i in range(100)}) ++ self.assertEqual(len(logit_bias._cache), logit_bias._CACHE_MAX) +diff --git a/serve/test_model_card_sampling.py b/serve/test_model_card_sampling.py +new file mode 100644 +index 00000000..ee491506 +--- /dev/null ++++ b/serve/test_model_card_sampling.py +@@ -0,0 +1,108 @@ ++"""serve/test_model_card_sampling.py - a request that sends no sampling gets the model card's, not greedy (mock engine). ++ ++Qwen's "Best Practices": thinking mode temperature 1.0 / top_p 0.95 / top_k 20, non-thinking 0.7 / 0.8 / 20 with ++presence_penalty 1.5. The request's own fields win, then the Chat settings shared with apps, then the config's ++`sampling` block, then the card; temperature 0 asked for is still greedy. ++ ++ python -m unittest serve.test_model_card_sampling -v ++""" ++from __future__ import annotations ++ ++import json ++import sys ++import unittest ++import urllib.request ++from pathlib import Path ++ ++sys.path.insert(0, str(Path(__file__).resolve().parents[1])) ++from serve.frontend import ChatTemplate # noqa: E402 ++from serve.server import ByteTokenizer, MockEngine, Service, StrataEngine, model_card_sampling, serve # noqa: E402 ++ ++ROOT = Path(__file__).resolve().parents[1] ++ ++ ++class Recording(MockEngine): ++ def generate(self, ids, max_new, sampling, cancel, embeddings=None): ++ self.sampling = dict(sampling or {}) ++ yield from super().generate(ids, max_new, sampling, cancel, embeddings) ++ ++ ++class ModelCardSampling(unittest.TestCase): ++ @classmethod ++ def setUpClass(cls): ++ cls.tok = ByteTokenizer() ++ cls.engine = Recording(cls.tok, "\nhm\n\n\nok", max_context=8192) ++ cls.svc = Service(cls.engine, cls.tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) ++ cls.httpd = serve(cls.svc, port=0) ++ cls.base = f"http://127.0.0.1:{cls.httpd.server_address[1]}" ++ ++ @classmethod ++ def tearDownClass(cls): ++ cls.httpd.shutdown() ++ cls.httpd.server_close() ++ ++ def setUp(self): ++ self.svc.model_arch = "qwen4exp" ++ self.svc.sampling_defaults = {} ++ self.svc.shared = {} ++ self.svc.config_effort = None ++ ++ def tearDown(self): ++ self.svc.model_arch = None ++ ++ def keys(self, **body): ++ body = {"model": "m", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 8, **body} ++ self.engine.sampling = {} ++ req = urllib.request.Request(self.base + "/v1/chat/completions", data=json.dumps(body).encode(), ++ headers={"Content-Type": "application/json"}) ++ urllib.request.urlopen(req, timeout=10).read() ++ return StrataEngine.sampling_keys(self.engine.sampling) ++ ++ def test_tables(self): ++ self.assertEqual(model_card_sampling("qwen4exp", True), {"temperature": 1.0, "top_p": 0.95, "top_k": 20}) ++ self.assertEqual(model_card_sampling("qwen4exp", False), ++ {"temperature": 0.7, "top_p": 0.8, "top_k": 20, "presence_penalty": 1.5}) ++ self.assertEqual(model_card_sampling("qwen35moe", True)["presence_penalty"], 1.5) ++ self.assertEqual(model_card_sampling(None, True), {}) ++ self.assertEqual(model_card_sampling("unknown", False), {}) ++ ++ def test_thinking_default(self): ++ k = self.keys() ++ self.assertIn("temperature=1.0", k) ++ self.assertIn("top_p=0.95", k) ++ self.assertIn("top_k=20", k) ++ self.assertNotIn("penalty_present", k) ++ ++ def test_non_thinking_default(self): ++ k = self.keys(reasoning_effort="none") ++ self.assertIn("temperature=0.7", k) ++ self.assertIn("top_p=0.8", k) ++ self.assertIn("penalty_present=1.5", k) ++ ++ def test_request_wins_and_zero_is_greedy(self): ++ k = self.keys(temperature=0) ++ self.assertNotIn("temperature", k) ++ self.assertIn("top_k=20", k) ++ k = self.keys(temperature=0.3, top_p=0.5) ++ self.assertIn("temperature=0.3", k) ++ self.assertIn("top_p=0.5", k) ++ ++ def test_config_block_and_shared_settings_win_over_the_card(self): ++ self.svc.sampling_defaults = {"temperature": 0.6, "top_k": 40} ++ k = self.keys() ++ self.assertIn("temperature=0.6", k) ++ self.assertIn("top_k=40", k) ++ self.assertIn("top_p=0.95", k) ++ self.svc.shared = {"temperature": 0.2} ++ self.assertIn("temperature=0.2", self.keys()) ++ self.svc.sampling_defaults = {"temperature": 0} ++ self.svc.shared = {} ++ self.assertNotIn("temperature", self.keys()) # a config that asks for greedy gets it ++ ++ def test_no_family_no_defaults(self): ++ self.svc.model_arch = None ++ self.assertEqual(self.keys(), "") ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_parallel.py b/serve/test_parallel.py +index 1a638919..bde52084 100644 +--- a/serve/test_parallel.py ++++ b/serve/test_parallel.py +@@ -248,6 +248,34 @@ class PickSlot(unittest.TestCase): + e.slot_busy = [False] * 3 + self.assertEqual(e.pick_slot([1, 2, 3]), 1) # a held prompt is never the WHOLE prompt + ++ def test_a_rewritten_next_turn_goes_to_the_slot_of_its_chat(self): ++ # #2518: two chats in two slots; each next turn renders the history again but not the reply it decoded or the ++ # assistant header, so the held stream (prompt + reply) is no start of it. The engine's turn checkpoint sits at ++ # the last <|im_start|> of the old prompt: the prompt that agrees up to there belongs to that slot. Before, both ++ # went to the least recently used slot in turn, and the first one overwrote the other chat's cache. ++ T = StrataEngine.TURN_ID ++ hist_a, hist_b = list(range(1000, 1600)), list(range(2000, 2600)) ++ held_a = hist_a + [T, 7, 8, 9] + [50, 51, 52] # prompt (ends with the assistant header) + reply ++ held_b = hist_b + [T, 7, 8, 9] + [60, 61] ++ e = self.engine(2) ++ e.slot_held = [held_a, held_b] ++ e.slot_used = [10.0, 20.0] # slot 0 was used longer ago ++ next_a = hist_a + [T, 5, 6, 7, 8, 9] + list(range(3000, 3100)) # the history again, then new turns ++ next_b = hist_b + [T, 5, 6, 7, 8, 9] + list(range(3000, 3100)) ++ self.assertEqual(e.pick_slot(next_b), 1) # the LRU slot 0 holds the other chat ++ self.assertEqual(e.pick_slot(next_a), 0) ++ e.slot_busy[0] = True # a busy slot is never picked ++ self.assertEqual(e.pick_slot(next_a), 1) ++ e.slot_busy[0] = False ++ # a chat that agrees with neither (a new one, or one that rewrote its history) keeps the old order ++ self.assertEqual(e.pick_slot(list(range(5000, 5200))), 0) ++ self.assertEqual(e.pick_slot(hist_a[:300] + [T, 5, 6] + list(range(3000, 3100))), 0) ++ e.slot_held[0] = [] # an empty slot still comes before evicting one ++ self.assertEqual(e.pick_slot(list(range(5000, 5200))), 0) ++ # the whole held stream continued is still the longer match ++ e.slot_held = [held_a, held_b + [T, 1, 2]] ++ self.assertEqual(e.pick_slot(held_b + [T, 1, 2, 3, 4]), 1) ++ + def test_batch_groups_spread_requests_over_the_groups(self): + # #1249: with --batch-groups 2 over 4 slots (groups {0,1} and {2,3}) a second request goes to the other group, + # and a short held prefix does not pull it back into the busy one +diff --git a/serve/test_parser_call_scan.py b/serve/test_parser_call_scan.py +new file mode 100644 +index 00000000..a4484514 +--- /dev/null ++++ b/serve/test_parser_call_scan.py +@@ -0,0 +1,61 @@ ++"""Wrapped-call precheck parity against the original structural walk.""" ++import itertools ++import random ++import unittest ++from unittest.mock import patch ++ ++from serve import frontend ++ ++ ++class ReferenceParser(frontend.OutputParser): ++ def _wrapped_call_end(self): ++ return frontend.call_end(self.buf) ++ ++ ++def events_without_ids(events): ++ return [(e.kind, e.text, (e.call.name, e.call.arguments) if e.call else None) for e in events] ++ ++ ++class CallScanTests(unittest.TestCase): ++ def test_chunk_boundaries_literal_tags_multiple_calls_and_recovery(self): ++ def call(value): ++ return "" + value + "" ++ bodies = [call("ordinary"), call("literal text here"), ++ call("literal prose"), call("a") + call("b"), ++ call("unfinished")[:-5], call("incomplete prose")[:-30], ++ "", ++ "bare", ++ "a" ++ "b"] ++ tools = [{"name": "write", "parameters": {"properties": {"text": {"type": "string"}}}}] ++ rng = random.Random(1030) ++ for text, stream, recover in itertools.product(bodies, (False, True), (False, True)): ++ # Every fixed chunk width, plus random partitions; compare after EACH feed. ++ for width in range(1, 16): ++ a = frontend.OutputParser(thinking=False, tools=tools, stream_tools=stream, recover=recover) ++ b = ReferenceParser(thinking=False, tools=tools, stream_tools=stream, recover=recover) ++ pos = 0 ++ while pos < len(text): ++ end = pos + (rng.randint(1, 20) if width == 15 else width) ++ delta = text[pos:end] ++ self.assertEqual(events_without_ids(a.feed(delta)), events_without_ids(b.feed(delta))) ++ pos = end ++ self.assertEqual(events_without_ids(a.finish()), events_without_ids(b.finish())) ++ ++ def test_no_structural_rescans_until_closer_and_reset_for_next_call(self): ++ parser = frontend.OutputParser(thinking=False) ++ with patch.object(frontend, "call_end", wraps=frontend.call_end) as walk: ++ for _ in range(2): ++ parser.feed("") ++ before = walk.call_count ++ for _ in range(100): ++ parser.feed("abcdefgh") ++ self.assertEqual(walk.call_count, before) ++ for char in "": ++ parser.feed(char) ++ self.assertEqual(walk.call_count, before + 1) ++ self.assertEqual(parser._call_end_checked, 0) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_parser_hold.py b/serve/test_parser_hold.py +new file mode 100644 +index 00000000..519f8192 +--- /dev/null ++++ b/serve/test_parser_hold.py +@@ -0,0 +1,60 @@ ++"""Partial-tag parity, including overlapping prefixes and arbitrary chunk boundaries.""" ++import itertools ++import random ++import unittest ++ ++from serve.frontend import OutputParser ++ ++ ++def reference_hold(text, tags): ++ return max((n for tag in tags for n in range(1, len(tag)) if text.endswith(tag[:n])), default=0) ++ ++ ++class ReferenceParser(OutputParser): ++ def _hold(self, text, tags): ++ return reference_hold(text, tags) ++ ++ ++class HoldTests(unittest.TestCase): ++ def test_overlapping_and_empty_tags(self): ++ parser = OutputParser() ++ tags = ("", "a", "aaa", "abab", "baab", "", "") ++ for size in range(8): ++ for chars in itertools.product("ab<", repeat=size): ++ text = "".join(chars) ++ self.assertEqual(parser._hold(text, tags), reference_hold(text, tags), text) ++ ++ def test_every_real_tag_split(self): ++ tags = ("", "", "") ++ parser = OutputParser() ++ for tag in tags: ++ for n in range(len(tag) + 1): ++ for prefix in ("", "prose\n", "<", "<<<", "🙂日本語"): ++ text = prefix + tag[:n] ++ self.assertEqual(parser._hold(text, tags), reference_hold(text, tags), text) ++ ++ def test_streamed_events_match_reference(self): ++ texts = ["Reasoning.\n\nAnswer café 🙂\n", ++ "Quoted `` text and prose.", ++ "\n```xml\n\n```\nAnswer\nabc\nDone", ++ "abcDone"] ++ rng = random.Random(1028) ++ for text, thinking, recover, stream_tools in itertools.product(texts, (False, True), (False, True), (False, True)): ++ for _ in range(5): ++ args = dict(thinking=thinking, tools=[{"name": "lookup"}], recover=recover, stream_tools=stream_tools) ++ actual, expected = OutputParser(**args), ReferenceParser(**args) ++ def normalized(events): ++ # Tool call IDs are deliberately random; compare their semantic contents. ++ return [(e.kind, e.text, (e.call.name, e.call.arguments) if e.call else None) for e in events] ++ pos = 0 ++ while pos < len(text): ++ end = pos + rng.randint(1, 9) ++ delta = text[pos:end] ++ self.assertEqual(normalized(actual.feed(delta)), normalized(expected.feed(delta))) ++ pos = end ++ self.assertEqual(normalized(actual.finish()), normalized(expected.finish())) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_prompt_encoder.py b/serve/test_prompt_encoder.py +new file mode 100644 +index 00000000..36cc06f2 +--- /dev/null ++++ b/serve/test_prompt_encoder.py +@@ -0,0 +1,162 @@ ++"""serve/test_prompt_encoder.py - #567's incremental prompt encoding with #537's plain spans: every prompt's ids ++are those of a full encode, and a turn reuses the previous one's. ++ ++ python -m unittest serve.test_prompt_encoder -v ++""" ++from __future__ import annotations ++ ++import random ++import re ++import sys ++import unittest ++from pathlib import Path ++ ++ROOT = Path(__file__).resolve().parents[1] ++sys.path.insert(0, str(ROOT)) ++sys.path.insert(0, str(ROOT / "tools")) ++from serve.frontend import ChatTemplate # noqa: E402 ++from serve.server import ByteTokenizer, MockEngine, Service # noqa: E402 ++from strata_tokenizer import PromptEncoder, plain_agree_until # noqa: E402 ++ ++ ++class ThinkTokenizer(ByteTokenizer): ++ """ByteTokenizer whose / are specials matched even without parse_special (type 4), as in the ++ Qwen vocabulary: the case #537's plain spans exist for.""" ++ SPECIALS = ByteTokenizer.SPECIALS + ["", ""] ++ ALWAYS = ("", "") ++ ++ ++PIECES = ["<|im_start|>user\n", "<|im_end|>\n", "", "", "hello ", "é你 ", "<|im_sta", "x" * 7, ++ "<|im_start|>assistant\n", "", "<|vision_start|>", "<|image_pad|>"] ++ ++ ++def spans_of(text, rng): ++ """Random plain spans over the specials' texts (, , <|im_start|>, ...), as unmark_think_literals ++ marks the ones quoted inside a message.""" ++ matches = re.finditer("|".join(map(re.escape, ThinkTokenizer.SPECIALS)), text) ++ return tuple((m.start(), m.end()) for m in matches if rng.random() < 0.5) # finditer: left to right, no overlap ++ ++ ++class Units(unittest.TestCase): ++ def test_plain_agree_until(self): ++ self.assertEqual(plain_agree_until((), (), 50), 50) ++ self.assertEqual(plain_agree_until(((5, 13),), ((5, 13),), 50), 50) ++ self.assertEqual(plain_agree_until(((5, 13),), (), 50), 5) ++ self.assertEqual(plain_agree_until(((5, 13),), ((5, 9),), 50), 9) ++ self.assertEqual(plain_agree_until(((60, 70),), (), 50), 50) # beyond the shared text: no matter ++ self.assertEqual(plain_agree_until(((5, 13), (20, 28)), ((5, 13),), 50), 20) ++ ++ def test_same_ids_as_a_full_encode(self): ++ tok = ThinkTokenizer() ++ rng = random.Random(7) ++ for conv in range(300): ++ enc = PromptEncoder(tok) ++ text, plain = "", () ++ for turn in range(8): ++ text += "".join(rng.choice(PIECES) for _ in range(rng.randint(1, 12))) ++ if rng.random() < 0.6: ++ plain = spans_of(text, rng) # the spans may change in the shared part too ++ with self.subTest(conv=conv, turn=turn): ++ self.assertEqual(enc.encode(text, plain), tok.encode(text, True, plain)) ++ if rng.random() < 0.2 and len(text) > 10: # an edit: a client re-rendering its history ++ cut = rng.randrange(len(text)) ++ text = text[:cut] ++ plain = tuple((a, b) for a, b in plain if b <= cut) ++ ++ def test_a_turn_reuses_the_previous_one(self): ++ tok = ThinkTokenizer() ++ enc = PromptEncoder(tok) ++ # as a chat turn: the earlier prompt goes on past its last special (the answer being written), and the ++ # next one shares it up to there - the resume point needs max_special_len characters after it ++ base = "<|im_start|>user\n" + "quoted here " * 200 + "<|im_end|>\n<|im_start|>assistant\nan answer" ++ plain = tuple((i, i + 8) for i in range(len(base)) if base.startswith("", i)) ++ enc.encode(base, plain) ++ more = base + " that goes on<|im_end|>\n<|im_start|>user\nnext<|im_end|>\n" ++ self.assertEqual(enc.encode(more, plain), tok.encode(more, True, plain)) ++ self.assertGreater(enc.last_reused, len(base) // 2) ++ # the same text with the quoted no longer plain: nothing before it may be reused past its start ++ self.assertEqual(enc.encode(more, ()), tok.encode(more, True, ())) ++ self.assertLessEqual(enc.last_reused, plain[0][0]) ++ ++ ++class Served(unittest.TestCase): ++ def test_quoted_think_across_turns(self): ++ tok = ThinkTokenizer() ++ template = ChatTemplate(ROOT / "serve/chat_template.jinja") ++ svc = Service(MockEngine(tok, "\n\nok", max_context=1 << 20), tok, template) ++ full = Service(MockEngine(tok, "\n\nok", max_context=1 << 20), tok, template) ++ full.prompts = None # the reference: a full encode every time ++ self.assertIsNotNone(svc.prompts) ++ msgs = [{"role": "system", "content": "Explain tags."}] ++ for turn in range(5): ++ msgs.append({"role": "user", "content": f"turn {turn}: what does mean? too " * 3}) ++ got, want = svc.encode_prompt(msgs, None, {}), full.encode_prompt(msgs, None, {}) ++ self.assertEqual(got, want) ++ if turn: ++ self.assertGreater(svc.prompts.last_reused, 0) ++ msgs.append({"role": "assistant", "content": "It closes the reasoning: is a marker."}) ++ ++ def service(self, max_context=1 << 22): ++ tok = ByteTokenizer() ++ return Service(MockEngine(tok, "ok", max_context=max_context), tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) ++ ++ def test_another_conversation_does_not_evict_the_first(self): ++ """A subagent beside the main conversation: every rendered prompt starts with the template's header, which ++ must not make the second client's prompt a continuation of the first's entry.""" ++ svc = self.service() ++ turn2 = [{"role": "user", "content": "a long question " * 500}, {"role": "assistant", "content": "answer"}, ++ {"role": "user", "content": "next"}] ++ svc.encode_prompt(turn2, None, {}) ++ svc.encode_prompt([{"role": "user", "content": "one unrelated prompt " * 300}], None, {}) ++ self.assertEqual(len(svc.prompts.entries), 2) ++ turn3 = turn2 + [{"role": "assistant", "content": "x"}, {"role": "user", "content": "more"}] ++ svc.encode_prompt(turn3, None, {}) ++ self.assertGreater(svc.prompts.last_reused, len(turn2[0]["content"])) ++ self.assertEqual(len(svc.prompts.entries), 2) # turn 3 replaced turn 2, the other one stayed ++ ++ def test_a_refused_prompt_is_not_kept(self): ++ svc = self.service(max_context=4096) ++ msgs = [{"role": "user", "content": "kept " * 20}] ++ svc.prepare(msgs, None, {}, 16) ++ refused = msgs + [{"role": "assistant", "content": "a"}, {"role": "user", "content": "more " * 20}] ++ full = self.service(max_context=4096) ++ full.prompts = None ++ refused_ids = full.encode_prompt(refused, None, {}) ++ with self.assertRaisesRegex(ValueError, "exceeds the context"): ++ svc.prepare(refused, None, {}, 1 << 20) ++ self.assertNotIn(refused_ids, [e[2] for e in svc.prompts.entries]) ++ self.assertNotIn(full.render_prompt(refused, None, {}), [e[0] for e in svc.prompts.entries]) ++ ++ def test_random_messages_equal_a_full_encode(self): ++ """Messages quoting control-token texts and tags at random, growing, edited and cut back between ++ requests: the incremental encoder's ids are the full encode's, 2400 prompts.""" ++ tok = ThinkTokenizer() ++ template = ChatTemplate(ROOT / "serve/chat_template.jinja") ++ svc = Service(MockEngine(tok, "ok", max_context=1 << 20), tok, template) ++ full = Service(MockEngine(tok, "ok", max_context=1 << 20), tok, template) ++ full.prompts = None ++ rng = random.Random(567) ++ quoted = ThinkTokenizer.SPECIALS + ["<|im_sta", " 2: # a shorter conversation ++ msgs = msgs[:rng.randrange(1, len(msgs))] + msgs[-1:] ++ with self.subTest(conv=conv, turn=turn): ++ self.assertEqual(svc.encode_prompt(msgs, None, {}), full.encode_prompt(msgs, None, {})) ++ reused += svc.prompts.last_reused > 0 ++ msgs.append({"role": "assistant", "content": text()}) ++ self.assertGreater(reused, 1000) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_q36_config.py b/serve/test_q36_config.py +new file mode 100644 +index 00000000..981ae5f3 +--- /dev/null ++++ b/serve/test_q36_config.py +@@ -0,0 +1,48 @@ ++"""serve/test_q36_config.py - what the server hands the Qwen3.6-35B-A3B engine (arch qwen35moe): the config's pictures ++("vision" section -> --vision) and parallel requests ("parallel": N -> --batch N), as for Flash-Next's engine. ++ ++ python -m unittest serve.test_q36_config ++""" ++from __future__ import annotations ++ ++import unittest ++from unittest import mock ++ ++from serve.server import engine_args, parallel_args ++ ++ ++def cfg(**kw): ++ base = {"arch": "qwen35moe", "exe": "strata", "args": ["--native", "/m/Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf", ++ "--max-context", "32768", "--prefill", "auto", "--mtp", "2"]} ++ base.update(kw) ++ return base ++ ++ ++class Q36Config(unittest.TestCase): ++ def test_parallel_becomes_batch(self): ++ a = engine_args(cfg(parallel=4)) ++ self.assertEqual(a[a.index("--batch") + 1], "4") ++ self.assertNotIn("--batch", engine_args(cfg())) ++ self.assertNotIn("--batch", engine_args(cfg(parallel=1))) ++ ++ def test_more_than_the_window_is_said(self): ++ with mock.patch("builtins.print") as pr: ++ a = engine_args(cfg(parallel=12)) ++ self.assertEqual(a[a.index("--batch") + 1], "12") # the engine caps it at 8 and says so ++ self.assertTrue(any("at most 8" in str(c) for c in pr.call_args_list)) ++ ++ def test_a_vision_section_starts_the_engine_with_pictures(self): ++ v = {"exe": "strata-vision", "mmproj": "/m/mmproj-F16.gguf", "model": "/m/Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf", ++ "gpu": False, "max_tokens": 300} ++ self.assertIn("--vision", engine_args(cfg(vision=v))) ++ self.assertIn("--vision", engine_args(cfg(vision=v, args=cfg()["args"] + ["--vision"]))) ++ self.assertEqual(engine_args(cfg(vision=v, args=cfg()["args"] + ["--vision"])).count("--vision"), 1) ++ self.assertNotIn("--vision", engine_args(cfg())) ++ ++ def test_an_explicit_batch_in_the_args_wins(self): ++ c = cfg(parallel=4, args=cfg()["args"] + ["--batch", "2"]) ++ self.assertEqual(parallel_args(c, c["args"]), []) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_reasoning_loop_recovery.py b/serve/test_reasoning_loop_recovery.py +index b53ee56a..c00fc1be 100644 +--- a/serve/test_reasoning_loop_recovery.py ++++ b/serve/test_reasoning_loop_recovery.py +@@ -58,6 +58,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + engine = RecordingEngine(self.tok, [self.loop, + 'I have verified the solution.\n\nComplete solution.']) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = enabled + events = list(service.run(self.ids, True, [], 40000, {}, threading.Event())) + return engine, events +@@ -66,6 +67,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + engine = RecordingEngine(self.tok, [self.loop, + 'Verified.\n\nComplete solution.']) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + request = {'temperature': .2, 'presence_penalty': 0, 'top_p': .95, 'top_k': 64, 'seed': 2} + list(service.run(self.ids, True, [], 40000, request, threading.Event())) +@@ -81,6 +83,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_frozen_cohort_sampler_is_unchanged_during_recovery(self): + engine = RecordingEngine(self.tok, [self.loop, 'Done.\n\nComplete solution.']) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + sampler = {'temperature': 1.0, 'top_p': .95, 'top_k': 20, 'presence_penalty': 1.5, 'seed': 2} + list(service.run(self.ids, True, [], 40000, sampler, threading.Event())) +@@ -91,6 +94,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_a_lower_top_p_and_top_k_are_never_raised_or_lowered(self): + engine = RecordingEngine(self.tok, [self.loop, 'Verified.\n\nComplete solution.']) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = "recover" + list(service.run(self.ids, True, [], 40000, {'temperature': 0, 'top_p': .5, 'top_k': 5}, threading.Event())) + self.assertEqual((engine.sampling_calls[1]['top_p'], engine.sampling_calls[1]['top_k']), (.5, 5)) +@@ -110,6 +114,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_stop_mode_ends_the_reply_at_the_loop(self): + engine = RecordingEngine(self.tok, [self.loop]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = "stop" + events = list(service.run(self.ids, True, [], 40000, {}, threading.Event())) + self.assertEqual(len(engine.calls), 1) # no second pass +@@ -121,6 +126,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + """#1753: the word count never sees the digits, so only the tail's period ends the reply.""" + engine = RecordingEngine(self.tok, ['9^104 = ' + '58747693387' * 3000]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = "stop" + events = list(service.run(self.ids, True, [], 40000, {}, threading.Event())) + self.assertEqual(len(engine.calls), 1) # no second pass +@@ -136,6 +142,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_default_single_token_guard_does_not_detect_passage_loop(self): + engine = RecordingEngine(self.tok, [self.loop]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + events = list(service.run(self.ids, True, [], 16384, {}, threading.Event())) + self.assertEqual(service.repeat_stop_tokens, 256) + self.assertEqual(events[-1][1]['finish'], 'length') +@@ -145,6 +152,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_recovery_is_once_and_never_extends_output_limit(self): + engine = RecordingEngine(self.tok, [self.loop, self.loop]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + events = list(service.run(self.ids, True, [], 16384, {}, threading.Event())) + self.assertEqual(len(engine.calls), 2) +@@ -155,6 +163,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_repeated_final_text_does_not_change_sampler_or_prompt(self): + engine = RecordingEngine(self.tok, [self.loop]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + events = list(service.run(self.ids, False, [], 40000, {}, threading.Event())) + self.assertEqual(len(engine.calls), 1) +@@ -163,6 +172,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_missing_native_high_hint_skips_recovery(self): + engine = RecordingEngine(self.tok, [self.loop]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + ids = self.tok.encode('<|im_start|>user\nTask<|im_end|>\n', parse_special=True) + events = list(service.run(ids, True, [], 40000, {}, threading.Event())) +@@ -172,6 +182,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_cancel_during_reasoning_does_not_start_recovery(self): + engine = RecordingEngine(self.tok, [self.loop, 'Unused']) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + cancel = threading.Event() + events = [] +@@ -187,6 +198,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + engine = RecordingEngine(self.tok, [self.loop, 'Done.\n\nAnswer.']) + engine.batch = 2 + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + events = list(service.run(self.ids, True, [], 40000, {}, threading.Event())) + self.assertEqual(events[-1][1]['reasoning_recoveries'], 1) +diff --git a/serve/test_server.py b/serve/test_server.py +index 3a49ff25..e68b44ac 100644 +--- a/serve/test_server.py ++++ b/serve/test_server.py +@@ -26,7 +26,7 @@ from serve.server import (CTX_SLACK, ByteTokenizer, EngineDied, EngineNonFinite, + PP_DONE_TAIL, Service, + StrataEngine, api_key_of, engine_args, key_matches, layer_split_value, prompt_progress, + prompt_tokens_seen, +- request_timings, serve, start_failure_hint) ++ request_stats, request_timings, serve, start_failure_hint) + from types import SimpleNamespace # noqa: E402 + + ROOT = Path(__file__).resolve().parents[1] +@@ -1548,6 +1548,13 @@ class SamplingKeys(unittest.TestCase): + bad = self.keys(strata_tune={"pcie_frac": 3, "spec_min_p": True, "pool_workers": 2}) + self.assertFalse([x for x in bad if x.split("=")[0] in ("pcie_frac", "spec_min_p", "pool_workers")]) + ++ def test_tune_prefill_pipe(self): ++ self.assertIn("pipe=1536.0", self.keys(temperature=0, strata_tune={"prefill_pipe": 1536})) ++ self.assertIn("pipe=0.0", self.keys(temperature=0, strata_tune={"prefill_pipe": 0})) ++ for bad in (-1, True, "800", 1e9): ++ self.assertFalse([x for x in self.keys(temperature=0, strata_tune={"prefill_pipe": bad}) ++ if x.startswith("pipe=")], bad) ++ + def test_tune_aux_cpus(self): + for val, want in ((True, "aux_cpus=1"), (1, "aux_cpus=1"), (False, "aux_cpus=0"), (0, "aux_cpus=0"), + (1.0, "aux_cpus=1"), (0.0, "aux_cpus=0")): +@@ -1771,6 +1778,144 @@ class RecordingPrompt(MockEngine): + yield from super().generate(ids, max_new, sampling, cancel, embeddings) + + ++class IncrementalPrompts(unittest.TestCase): ++ """The prompt encoder: every request's ids are those of a full encode, and a turn reuses the previous one's.""" ++ ++ @classmethod ++ def setUpClass(cls): ++ tok = ByteTokenizer() ++ cls.engine = RecordingPrompt(tok, "Thinking.\n\n\nThe answer.", max_context=CTX) ++ cls.svc = Service(cls.engine, tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) ++ cls.httpd = serve(cls.svc, port=0) ++ cls.base = f"http://127.0.0.1:{cls.httpd.server_address[1]}" ++ ++ @classmethod ++ def tearDownClass(cls): ++ cls.httpd.shutdown() ++ cls.httpd.server_close() ++ ++ post = ClientShapes.post ++ ++ def test_turns(self): ++ self.assertIsNotNone(self.svc.prompts) ++ full = Service(self.engine, self.svc.tok, self.svc.template) ++ full.prompts = None # the reference: the same prompt, encoded in full ++ msgs = [{"role": "system", "content": "Be brief."}] ++ for turn in range(6): ++ msgs.append({"role": "user", "content": f"question {turn} <|im_end|> é你 " * (turn + 1)}) ++ status, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 64, "messages": msgs}) ++ self.assertEqual(status, 200, b) ++ prompt = self.svc.template.render(msgs) ++ self.assertEqual(self.engine.last_ids, full.encode_prompt(msgs, None, {})) ++ if turn: ++ self.assertGreater(self.svc.prompts.last_reused, len(prompt) // 3) ++ msgs.append({"role": "assistant", "content": b["choices"][0]["message"]["content"]}) ++ ++ def conversation(self, n): ++ """Messages of the n-th conversation: distinct from every other one from the number on.""" ++ return [{"role": "user", "content": f"conversation {n} says: " + f"topic {n} sentence {n} " * 20}] ++ ++ def service(self, **info): ++ engine = RecordingPrompt(self.svc.tok, "The answer.", max_context=CTX) ++ engine.info = info ++ return Service(engine, self.svc.tok, self.svc.template) ++ ++ def test_the_kept_prompts_follow_the_engines_conversation_slots(self): ++ svc = self.service(conversation_cache_slots=8) ++ texts = [] ++ for n in range(8): ++ msgs = self.conversation(n) ++ svc.encode_prompt(msgs, None, {}) ++ texts.append(svc.render_prompt(msgs, None, {})) ++ self.assertEqual(len(svc.prompts.entries), 8) ++ self.assertEqual(sorted(e[0] for e in svc.prompts.entries), sorted(texts)) ++ msgs = self.conversation(8) ++ svc.encode_prompt(msgs, None, {}) ++ kept = sorted(e[0] for e in svc.prompts.entries) ++ self.assertEqual(len(kept), 8) ++ self.assertEqual(kept, sorted(texts[1:] + [svc.render_prompt(msgs, None, {})])) ++ ++ def lazy_service(self, args, **info): ++ """A Service over the real StrataEngine class started lazily: no process yet, so no INFO line either.""" ++ engine = StrataEngine("strata", args, lazy=True) ++ engine.info.update(info) ++ return Service(engine, self.svc.tok, self.svc.template) ++ ++ def test_a_lazy_engine_keeps_the_prompts_its_arguments_ask_for(self): ++ svc = self.lazy_service(["--conversation-cache-slots", "8"]) ++ for n in range(10): ++ svc.encode_prompt(self.conversation(n), None, {}) ++ self.assertEqual(len(svc.prompts.entries), 8) ++ ++ def test_the_slot_argument_wins_over_the_info_line_and_a_bad_one_falls_back(self): ++ self.assertEqual(self.lazy_service(["--conversation-cache-slots", "6"], conversation_cache_slots=3) ++ .prompts.keep, 6) ++ self.assertEqual(self.lazy_service(["--conversation-cache-slots", "x"], conversation_cache_slots=3) ++ .prompts.keep, 3) ++ self.assertEqual(self.lazy_service(["--conversation-cache-slots"]).prompts.keep, 4) ++ self.assertEqual(self.lazy_service([]).prompts.keep, 4) ++ ++ def test_without_a_slot_setting_four_prompts_are_kept(self): ++ svc = self.service() ++ for n in range(6): ++ svc.encode_prompt(self.conversation(n), None, {}) ++ self.assertEqual(len(svc.prompts.entries), 4) ++ ++ def test_counting_tokens_leaves_the_kept_prompts_alone(self): ++ svc = self.service(conversation_cache_slots=2) ++ httpd = serve(svc, port=0) ++ self.addCleanup(httpd.server_close) ++ self.addCleanup(httpd.shutdown) ++ self.base = f"http://127.0.0.1:{httpd.server_address[1]}" ++ main = self.conversation(0) ++ status, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 8, "messages": main}) ++ self.assertEqual(status, 200, b) ++ before = list(svc.prompts.entries) ++ full = self.service() ++ full.prompts = None ++ for n in range(1, 6): ++ other = self.conversation(n) ++ status, b = self.post("/v1/messages/count_tokens", {"model": "m", "messages": other}) ++ self.assertEqual(status, 200, b) ++ self.assertEqual(b["input_tokens"], len(full.encode_prompt(other, None, {}))) ++ self.assertTrue(svc.prompts.entries == before, "the count wrote into the kept prompts") ++ main.append({"role": "assistant", "content": "ok"}) ++ main.append({"role": "user", "content": "and then?"}) ++ svc.encode_prompt(main, None, {}) ++ self.assertGreater(svc.prompts.last_reused, len(svc.render_prompt(main, None, {})) // 2) ++ ++ def test_same_ids_as_a_full_encode_on_varied_conversations(self): ++ import random ++ sys.path.insert(0, str(ROOT / "tools")) ++ from test_strata_tokenizer import conversation_prompts, load_tokenizer ++ toks = [("byte", ByteTokenizer())] ++ qwen = load_tokenizer() ++ if qwen is not None: ++ toks.append(("qwen35", qwen)) ++ for name, tok in toks: ++ svc = Service(self.engine, tok, self.svc.template) ++ for seed in range(3): ++ for what, prompt in conversation_prompts(self.svc.template, random.Random(seed)): ++ with self.subTest(tokenizer=name, seed=seed, what=what): ++ self.assertEqual(svc.prompts.encode(prompt), tok.encode(prompt, parse_special=True)) ++ ++ def test_a_tokenizer_without_resume_points_encodes_in_full(self): ++ class Plain: ++ """ByteTokenizer without the resume points.""" ++ def __init__(self): ++ self.base = ByteTokenizer() ++ ++ def __getattr__(self, name): ++ if name in ("encode_marked", "max_special_len"): ++ raise AttributeError(name) ++ return getattr(self.base, name) ++ svc = Service(self.engine, Plain(), self.svc.template) ++ self.assertIsNone(svc.prompts) ++ msgs = [{"role": "user", "content": "hi"}] ++ self.assertEqual(svc.encode_prompt(msgs, None, {}), ++ ByteTokenizer().encode(svc.render_prompt(msgs, None, {}), parse_special=True)) ++ ++ + class DyingEngine(MockEngine): + """Issue #27: an engine that dies after a few tokens of its first answer, and comes back when restarted.""" + +@@ -1989,6 +2134,39 @@ class PrometheusMetrics(unittest.TestCase): + self.assertEqual((m["strata:requests_total"], m["strata:drafts_offered_total"], + m["strata:drafts_accepted_total"]), (2, 17, 10)) + ++ def test_throughput_gauges_cover_the_window_between_scrapes(self): ++ """#2341: llama.cpp reads these two from a bucket it resets on every scrape, so a scrape that saw no ++ traffic reports 0 - not the lifetime average, which froze at the last request's speed.""" ++ tok = ByteTokenizer() ++ engine = DoneLineEngine(tok, "\n\nok", max_context=CTX, ++ done_lines=["DONE 4 20 400.0 200.0 stop 7 12 5"]) ++ svc = Service(engine, tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) ++ httpd = serve(svc, port=0) ++ base = f"http://127.0.0.1:{httpd.server_address[1]}" ++ ++ def scrape(): ++ with urllib.request.urlopen(base + "/metrics/prometheus", timeout=10) as r: ++ return {line.split()[0]: float(line.split()[1]) ++ for line in r.read().decode().splitlines() if not line.startswith("#")} ++ ++ try: ++ body = json.dumps({"model": "m", "max_tokens": 10, "messages": [{"role": "user", "content": "hi"}]}) ++ with urllib.request.urlopen(urllib.request.Request(base + "/v1/chat/completions", data=body.encode(), ++ headers={"Content-Type": "application/json"}), ++ timeout=30) as r: ++ self.assertEqual(r.status, 200) ++ first = scrape() ++ self.assertGreater(first["llamacpp:prompt_tokens_seconds"], 0) ++ self.assertGreater(first["llamacpp:predicted_tokens_seconds"], 0) ++ idle = scrape() # nothing ran since the scrape above, so the window is empty ++ self.assertEqual((idle["llamacpp:prompt_tokens_seconds"], ++ idle["llamacpp:predicted_tokens_seconds"]), (0, 0)) ++ self.assertEqual(idle["llamacpp:tokens_predicted_total"], # counters stay cumulative ++ first["llamacpp:tokens_predicted_total"]) ++ finally: ++ httpd.shutdown() ++ httpd.server_close() ++ + + class PcieShare(unittest.TestCase): + """#588: the hit rate stays the VRAM share of the lookups; the routed experts the GPU read over PCIe (the DONE +@@ -2507,6 +2685,30 @@ class CancelledRead(unittest.TestCase): + StrataEngine._parse_done(e, "DONE 0 98179 17565.0 0.0 cancel 0 0 0 0 0 0 0 0.0") + self.assertNotIn("prompt_read", e.last) + ++ def test_parse_done_chunk_fields(self): ++ """The prompt's first chunk and its streamed expert blobs come as key=value fields after the positional ones, ++ with or without the five --lookup-chain fields; an older line has neither.""" ++ e = SimpleNamespace() ++ StrataEngine._parse_done(e, "DONE 5 12000 900.0 80.0 stop 3 4 0 7 9 0 0 0.0 12000 2 chunk=1536 " ++ "experts_streamed=24576") ++ self.assertEqual((e.last["chunk"], e.last["experts_streamed"], e.last["offloaded"]), (1536, 24576, 2)) ++ StrataEngine._parse_done(e, "DONE 5 12000 900.0 80.0 stop 3 4 0 7 9 0 0 0.0 12000 2 1 2 3 4 5 chunk=8192 " ++ "experts_streamed=0") ++ self.assertEqual((e.last["chunk"], e.last["experts_streamed"]), (8192, 0)) ++ StrataEngine._parse_done(e, "DONE 5 12000 900.0 80.0 stop 3 4 0 7 9 0 0 0.0 12000 2") ++ self.assertNotIn("chunk", e.last) ++ self.assertNotIn("experts_streamed", e.last) ++ StrataEngine._parse_done(e, "DONE 5 12000 900.0 80.0 stop 3 4 0 7 9 0 0 0.0 12000 2 chunk=x other=1") ++ self.assertNotIn("chunk", e.last) ++ self.assertNotIn("other", e.last) ++ ++ def test_request_stats_chunk_fields(self): ++ """Across a request's continuations the streamed blobs add up; the chunk is the request's own read's.""" ++ stats = request_stats([{"prompt_ms": 10.0, "chunk": 1536, "experts_streamed": 100}, ++ {"prompt_ms": 2.0, "chunk": 64, "experts_streamed": 5}]) ++ self.assertEqual((stats["chunk"], stats["experts_streamed"], stats["prompt_ms"]), (1536, 105, 12.0)) ++ self.assertNotIn("chunk", request_stats([{"prompt_ms": 1.0}, {"prompt_ms": 1.0, "chunk": 64}])) ++ + def test_prompt_tokens_seen(self): + last = {"finish": "cancel", "reused": 1000, "prompt_read": 2000} + self.assertEqual(prompt_tokens_seen(98179, last), 3000) +@@ -2822,6 +3024,46 @@ class ConfigReasoningEffort(unittest.TestCase): + self.assertEqual(cfg, {"args": ["--ctx", "4096"]}) + + ++class JoinedEngineArgs(unittest.TestCase): ++ """#2267: "--ple-io mmap" written as one element of the config's args reached the engine as one argument.""" ++ ++ def split(self, args): ++ from serve.server import split_joined_args ++ cfg = {"args": list(args)} ++ with mock.patch("builtins.print") as pr: ++ split_joined_args(cfg) ++ return cfg["args"], pr.called ++ ++ def test_one_element_becomes_flag_and_value(self): ++ self.assertEqual(self.split(["--spec", "4", "--ple-io mmap"]), (["--spec", "4", "--ple-io", "mmap"], True)) ++ self.assertEqual(self.split(["--ple-io=mmap", "--kv int8"]), (["--ple-io", "mmap", "--kv", "int8"], True)) ++ ++ def test_proper_args_are_untouched(self): ++ args = ["--ple-io", "mmap", "--native", "/data/My Models/a b.gguf", "--max-context", "131072", "--vision"] ++ self.assertEqual(self.split(args), (args, False)) ++ # a value in a flag's value place stays whole, whatever it looks like ++ args = ["--system-prompt", "--be brief please"] ++ self.assertEqual(self.split(args), (args, False)) ++ ++ def test_no_args(self): ++ from serve.server import split_joined_args ++ cfg = {} ++ split_joined_args(cfg) ++ self.assertEqual(cfg, {}) ++ ++ def test_bad_json_names_the_place(self): ++ import tempfile ++ from serve.server import read_config ++ with tempfile.TemporaryDirectory() as d: ++ p = Path(d) / "c.json" ++ p.write_text('{"args": ["--ple-io" "mmap"]}', encoding="utf-8") ++ with self.assertRaises(SystemExit) as cm: ++ read_config(p) ++ self.assertIn("line 1", str(cm.exception)) ++ p.write_text('{"args": ["--ple-io", "mmap"]}', encoding="utf-8") ++ self.assertEqual(read_config(p), {"args": ["--ple-io", "mmap"]}) ++ ++ + class SharedSettings(unittest.TestCase): + """The web app's "Use for other apps too": POST /settings makes its Chat settings every client's defaults.""" + +@@ -3776,6 +4018,45 @@ class ThinkingBudget(unittest.TestCase): + self.assertEqual(text, ThinkingEngine.ANSWER) + self.assertEqual(evs[-2]["delta"]["stop_reason"], "end_turn") + ++ def test_a_cap_that_the_thinking_would_eat_leaves_room_to_answer(self): ++ """0.1.45: max_tokens counts the thinking; a thinking that would fill it gets closed with a quarter left (the ++ agent harness sends max_completion_tokens 8192; the pagoda's thinking took all of it and no answer came).""" ++ self.engine.THOUGHT = "Let me think step by step about two plus two. " * 40 # 1,800 reasoning tokens ++ code, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 1024, ++ "messages": [{"role": "user", "content": "2+2?"}]}) ++ self.assertEqual(code, 200, b) ++ msg = b["choices"][0]["message"] ++ self.assertEqual(msg["content"], ThinkingEngine.ANSWER) ++ self.assertEqual(b["choices"][0]["finish_reason"], "stop") ++ self.assertLess(len(msg["reasoning_content"]), 1024 - 256) ++ self.assertEqual(len(self.engine.prompts), 2) # the wrap-up continuation ++ ++ def test_the_answer_room_has_its_limits(self): ++ self.engine.THOUGHT = "Let me think step by step about two plus two. " * 40 ++ code, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 300, # below the floor: as before ++ "messages": [{"role": "user", "content": "2+2?"}]}) ++ self.assertEqual(b["choices"][0]["finish_reason"], "length") ++ self.assertFalse(b["choices"][0]["message"]["content"]) ++ self.svc.answer_room = False # the config's "answer_room": false ++ code, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 1024, ++ "messages": [{"role": "user", "content": "2+2?"}]}) ++ self.assertEqual(b["choices"][0]["finish_reason"], "length") ++ self.svc.answer_room = True ++ code, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 1024, "reasoning_effort": "none", ++ "messages": [{"role": "user", "content": "2+2?"}]}) ++ self.assertEqual(len(self.engine.prompts), 3) # thinking off: no budget, one prompt each ++ ++ def test_anthropic_thinking_budget_tokens_is_a_budget(self): ++ """0.1.45: thinking {type: enabled, budget_tokens: N} closes the thinking after N tokens (it was only a level).""" ++ from serve.server import REASONING_WRAP_UP ++ code, b = self.post("/v1/messages", {"model": "m", "max_tokens": 400, "thinking": {"type": "enabled", "budget_tokens": 30}, ++ "messages": [{"role": "user", "content": "2+2?"}]}) ++ self.assertEqual(code, 200, b) ++ thinking = "".join(c.get("thinking", "") for c in b["content"] if c["type"] == "thinking") ++ self.assertEqual(thinking, ThinkingEngine.THOUGHT[:30] + REASONING_WRAP_UP.split("")[0]) ++ self.assertEqual("".join(c.get("text", "") for c in b["content"] if c["type"] == "text"), ThinkingEngine.ANSWER) ++ self.assertEqual(len(self.engine.prompts), 2) ++ + def test_a_budget_the_thinking_stays_under(self): + code, b = self.openai(reasoning_budget_tokens=10_000) + self.assertEqual(b["choices"][0]["message"]["reasoning_content"], ThinkingEngine.THOUGHT) +@@ -4481,6 +4762,84 @@ class AmdTelemetry(unittest.TestCase): + self.assertEqual(svc.free_vram_mib(), 26 << 10) + + ++IOREG_FIXTURE = ''' ++ +-o AGXAcceleratorG13X ++ { ++ "IOClass" = "AGXAcceleratorG13X" ++ "model" = "Apple M1 Max" ++ "PerformanceStatistics" = {"In use system memory (driver)"=0,"Alloc system memory"=22991831040,"Tiler Utilization %"=36,"recoveryCount"=44,"lastRecoveryTime"=216970984728,"Renderer Utilization %"=35,"TiledSceneBytes"=1671168,"Device Utilization %"=76,"SplitSceneCount"=0,"Allocated PB Size"=76152832,"In use system memory"=13384105984} ++ +-o IOUserClientCreator ++ { ++ "IOUserClientCreator" = "pid 402, WindowServer" ++ } ++ } ++ ++ +-o AGXAccelerator ++ { ++ "IOClass" = "AGXAccelerator" ++ "model" = "Apple M2 Ultra" ++ "PerformanceStatistics" = {"Renderer Utilization %"=3,"Tiler Utilization %"=2,"In use system memory"=1234567890} ++ } ++''' ++ ++WORKING_SET = 26800603136 # bytes; the M1 Max's device in the fixture above ++ ++ ++class MacTelemetry(unittest.TestCase): ++ """The macOS (Metal backend) reader: IOKit's IOAccelerator nodes and the device's Metal working set.""" ++ ++ def _reader(self, index=0, text=IOREG_FIXTURE): ++ from serve import telemetry ++ return telemetry._Metal(index, run=lambda: text) ++ ++ def test_readings(self): ++ g = self._reader() ++ self.assertTrue(g.ok()) ++ self.assertEqual(g.name(), "Apple M1 Max") ++ with mock.patch("serve.telemetry.metal_working_set", lambda: WORKING_SET): ++ self.assertEqual(g.read(), {"util": 76, "mem_used": 13384105984, "mem_total": WORKING_SET}) ++ # the second accelerator is index 1: the older chips' renderer/tiler-only utilization is the fallback ++ g1 = self._reader(1) ++ self.assertTrue(g1.ok()) ++ with mock.patch("serve.telemetry.metal_working_set", lambda: WORKING_SET): ++ self.assertEqual(g1.read()["util"], 3) # the renderer, the higher of the two ++ self.assertEqual(g1.name(), "Apple M2 Ultra") ++ self.assertFalse(self._reader(5).ok()) ++ ++ def test_no_registry(self): ++ g = self._reader(text="") ++ self.assertFalse(g.ok()) ++ self.assertEqual(g.read(), {}) ++ self.assertIsNone(g.name()) ++ ++ def test_bad_process(self): ++ from serve import telemetry ++ g = telemetry._Metal(0, run=lambda: (_ for _ in ()).throw(OSError("no ioreg"))) ++ self.assertFalse(g.ok()) ++ self.assertEqual(g.read(), {}) ++ ++ def test_note_reaches_the_monitor(self): ++ from serve import telemetry ++ t = telemetry.Telemetry(extra=lambda: {"tok_s": 1.0}) ++ try: ++ with mock.patch.object(telemetry, "gpu_reader", lambda index=0, amd=False: self._reader()): ++ t2 = telemetry.Telemetry(extra=lambda: {"tok_s": 1.0}) ++ try: ++ self.assertEqual(t2.static["gpu_name"], "Apple M1 Max") ++ self.assertIn("powermetrics", t2.static["gpu_note"]) ++ with mock.patch("serve.telemetry.metal_working_set", lambda: WORKING_SET): ++ s = t2.sample() ++ self.assertEqual(s["gpu_util"], 76) ++ self.assertEqual(s["gpu_mem_used"], 13384105984) ++ self.assertEqual(s["gpu_mem_total"], WORKING_SET) ++ self.assertIsNone(s.get("gpu_temp")) # no reading without root; the cards show a dash ++ self.assertIsNone(s.get("gpu_pcie_gen")) ++ finally: ++ t2.close() ++ finally: ++ t.close() ++ ++ + class TelemetryNoDisks(unittest.TestCase): + """psutil.disk_io_counters() returns None where it finds no disk (a gVisor container; Windows with its disk counters + off): the sampler goes on without a disk reading instead of its thread ending at the first sample, which left +diff --git a/serve/test_slot_alloc.py b/serve/test_slot_alloc.py +index 474b4f78..b2b4d35b 100644 +--- a/serve/test_slot_alloc.py ++++ b/serve/test_slot_alloc.py +@@ -79,5 +79,30 @@ class OneGroup(unittest.TestCase): + self.assertEqual(e.pick_slot([1, 2, 3]), 3) # else the least recently used + + ++class OneSlotGroups(unittest.TestCase): ++ """--batch 2 on a two-GPU layer split: two pipelined groups of ONE slot. No slot sits below another in its group, ++ so no pad row writes an idle slot and each one keeps its chat (#2518): the next turn goes back to it.""" ++ ++ def test_the_slot_of_the_chat(self): ++ T = StrataEngine.TURN_ID ++ e = engine(2, 2) ++ self.assertEqual(e.slot_gs, 1) ++ hist_a, hist_b = list(range(1000, 1700)), list(range(2000, 2700)) ++ e.slot_held = [hist_a + [T, 7, 8] + [50, 51], hist_b + [T, 7, 8] + [60]] ++ e.slot_used = [10.0, 20.0] ++ again = lambda h: h + [T, 5, 6] + list(range(3000, 3050)) ++ self.assertEqual(e.pick_slot(again(hist_b)), 1) ++ self.assertEqual(e.pick_slot(again(hist_a)), 0) ++ self.assertEqual(e.pick_slot(list(range(5000, 5100))), 0) # a new chat: the least recently used one ++ ++ def test_groups_of_several_slots_still_ignore_held_tokens(self): ++ T = StrataEngine.TURN_ID ++ e = engine(4, 2) ++ self.assertEqual(e.slot_gs, 2) ++ hist = list(range(1000, 1700)) ++ e.slot_held[3] = hist + [T, 7, 8] + [50] ++ self.assertEqual(e.pick_slot(hist + [T, 5, 6] + list(range(3000, 3050))), 0) ++ ++ + if __name__ == "__main__": + unittest.main() +diff --git a/serve/test_web_sampling.py b/serve/test_web_sampling.py +new file mode 100644 +index 00000000..b502ace2 +--- /dev/null ++++ b/serve/test_web_sampling.py +@@ -0,0 +1,77 @@ ++"""serve/test_web_sampling.py - the web page sends no sampling of its own until the user changes it (0.1.45, #2685). ++ ++The page used to send temperature 0.6 / top_p 0.95 / top_k 20 with every request, so the server's model-card defaults ++(thinking 1.0, not thinking 0.7 + presence penalty) never applied to it. Pins, with node, the page's own functions: ++untouched (the shown defaults, or the old stored 0.6) sends nothing, a changed value is sent whole, 0 is greedy. ++Skipped without node. ++ ++ python -m unittest serve.test_web_sampling -v ++""" ++from __future__ import annotations ++ ++import json ++import re ++import shutil ++import subprocess ++import unittest ++from pathlib import Path ++ ++APP = Path(__file__).resolve().parent / "web" / "app.js" ++ ++ ++def slice_between(src: str, start: str, end: str) -> str: ++ a = src.index(start) ++ return src[a:src.index(end, a)] ++ ++ ++@unittest.skipUnless(shutil.which("node"), "node is not installed") ++class WebSampling(unittest.TestCase): ++ @classmethod ++ def setUpClass(cls): ++ src = APP.read_text(encoding="utf-8") ++ cls.defaults = slice_between(src, "const DEFAULTS", "let settings") ++ cls.shared = slice_between(src, "function sharedDefaults", "async function saveShared") ++ cls.body = slice_between(src, " if (!(settings.temperature > 0)) {", " if (settings.seed)") ++ ++ def run_js(self, settings: dict, pieces: str): ++ js = (self.defaults + "function projectionLoaded() { return false; }\n" + self.shared + ++ f"\nconst settings = {json.dumps(settings)};\nconst body = {{}};\n{pieces}\n" ++ "console.log(JSON.stringify({body, shared: sharedDefaults(settings)}));") ++ out = subprocess.run(["node", "-e", js], capture_output=True, text=True, timeout=30) ++ self.assertEqual(out.returncode, 0, out.stderr) ++ return json.loads(out.stdout) ++ ++ def sampling(self, **kw): ++ s = {"thinking": "high", "temperature": 1.0, "top_p": 0.95, "top_k": 20, "max": "", "seed": "", "esp": True} ++ s.update(kw) ++ r = self.run_js(s, self.body) ++ return r["body"], r["shared"] ++ ++ def test_untouched_sends_nothing(self): ++ for t in (1.0, 0.6): # the shown default, and the old stored default ++ body, shared = self.sampling(temperature=t) ++ self.assertEqual(body, {}) ++ self.assertEqual(shared, {"reasoning_effort": "high"}) ++ ++ def test_changed_values_are_sent_whole(self): ++ body, shared = self.sampling(temperature=0.3) ++ self.assertEqual(body, {"temperature": 0.3, "top_p": 0.95, "top_k": 20}) ++ self.assertEqual(shared["temperature"], 0.3) ++ body, _ = self.sampling(top_k=40) ++ self.assertEqual(body, {"temperature": 1.0, "top_p": 0.95, "top_k": 40}) ++ body, _ = self.sampling(temperature=0.6, top_p=0.9) ++ self.assertEqual(body["temperature"], 0.6) ++ ++ def test_zero_is_greedy(self): ++ body, shared = self.sampling(temperature=0) ++ self.assertEqual(body, {"temperature": 0}) ++ self.assertEqual(shared["temperature"], 0) ++ ++ def test_the_old_default_is_shown_as_the_new_one(self): ++ src = APP.read_text(encoding="utf-8") ++ self.assertRegex(src, r"if \(samplingUntouched\(settings\)\) settings\.temperature = DEFAULTS\.temperature;") ++ self.assertIsNone(re.search(r"temperature: 0\.6", src)) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/web/app.js b/serve/web/app.js +index d3ad404f..4bc2a46f 100644 +--- a/serve/web/app.js ++++ b/serve/web/app.js +@@ -560,8 +560,18 @@ function markdown(text) { + } + + // ------------------------------------------------------------------ Chat +-const DEFAULTS = {thinking: "high", temperature: 0.6, top_p: 0.95, top_k: 20, max: "", seed: "", show: true, esp: true, mcp: true}; ++// The sampling shown until the user changes it is the model card's thinking set. Untouched, the page sends NO sampling ++// fields: the server picks the card's set for the request's thinking state (thinking 1.0 / 0.95 / 20, not thinking ++// 0.7 / 0.8 / 20 with a presence penalty). Before 0.1.45 the page sent 0.6 itself; a stored 0.6 with the other two ++// values untouched counts as "not set" too. ++const DEFAULTS = {thinking: "high", temperature: 1.0, top_p: 0.95, top_k: 20, max: "", seed: "", show: true, esp: true, mcp: true}; ++const OLD_DEFAULT_TEMPERATURE = 0.6; ++function samplingUntouched(s) { ++ const t = +s.temperature; ++ return (t === DEFAULTS.temperature || t === OLD_DEFAULT_TEMPERATURE) && +s.top_p === DEFAULTS.top_p && +s.top_k === DEFAULTS.top_k; ++} + let settings = {...DEFAULTS, ...store.get("sampling", {})}; ++if (samplingUntouched(settings)) settings.temperature = DEFAULTS.temperature; + let messages = store.get("chat", []); + let attachments = []; // {name, url} + let busy = null; // {controller, msg} +@@ -790,10 +800,10 @@ async function send() { + + const body = {model: health.model, messages: apiMessages(), stream: true, + reasoning_effort: settings.thinking}; +- if (settings.temperature > 0) { +- Object.assign(body, {temperature: +settings.temperature, top_p: +settings.top_p, top_k: +settings.top_k}); +- } else { ++ if (!(settings.temperature > 0)) { + body.temperature = 0; ++ } else if (!samplingUntouched(settings)) { ++ Object.assign(body, {temperature: +settings.temperature, top_p: +settings.top_p, top_k: +settings.top_k}); + } + if (settings.seed) body.seed = +settings.seed; + if (settings.max) body.max_tokens = +settings.max; +@@ -1011,8 +1021,9 @@ async function loadShared() { + $("s-share").setAttribute("aria-checked", String(sharedOn)); + } + function sharedDefaults(s) { +- const d = {reasoning_effort: s.thinking, temperature: +s.temperature}; +- if (+s.temperature > 0) Object.assign(d, {top_p: +s.top_p, top_k: +s.top_k}); ++ const d = {reasoning_effort: s.thinking}; ++ if (!(+s.temperature > 0)) d.temperature = 0; ++ else if (!samplingUntouched(s)) Object.assign(d, {temperature: +s.temperature, top_p: +s.top_p, top_k: +s.top_k}); + if (s.seed) d.seed = +s.seed; + if (s.max) d.max_tokens = +s.max; + if (projectionLoaded()) d.experimental_speed_projection = s.esp !== false; +diff --git a/tools/strata_tokenizer.py b/tools/strata_tokenizer.py +index c9e84ee7..42a27fed 100644 +--- a/tools/strata_tokenizer.py ++++ b/tools/strata_tokenizer.py +@@ -15,9 +15,11 @@ emoji (4-byte), combining marks, whitespace runs, and C0 control bytes. + """ + from __future__ import annotations + ++import bisect + import json + import pathlib + import sys ++import threading + + import regex + +@@ -115,6 +117,8 @@ class Tokenizer: + # through __dict__.setdefault on each encoded segment or decoded token. + self._piece_ids: dict[str, list[int]] = {} + self._bytes_cache: dict[int, bytes] = {} ++ # How far a special-token match can look ahead of where it starts (PromptEncoder's margin). ++ self.max_special_len = max((len(t) for t in self.special_tokens), default=1) + + @staticmethod + def _alt(literals: list[str]): +@@ -234,13 +238,16 @@ class Tokenizer: + out.extend(got) + return out + +- def _encode_matching(self, text: str, pat, plain=()) -> list[int]: ++ def _encode_matching(self, text: str, pat, plain=(), marks: list | None = None) -> list[int]: + """Encode `text`, emitting any literal `pat` matches as single tokens and BPE-ing the rest. + + The split happens on the RAW text, before the byte mapping, because a special token's string is a + literal to match rather than bytes to decompose. Everything between the matches is tokenized + normally - which is why a near-miss like `<|im_star` still costs ordinary tokens. A match that starts + inside one of the `plain` (start, end) spans is left to the text around it (#537). ++ ++ `marks`, when given, receives (end of the match in `text`, ids so far) for every match: the points where ++ the encoding of a longer text with the same beginning can resume (PromptEncoder). + """ + if pat is None: + return self._encode_plain(text) +@@ -269,6 +276,8 @@ class Tokenizer: + out.extend(self._encode_plain(text[pos:start])) + out.append(self.special_tokens[m.group(0)]) + pos = m.end() ++ if marks is not None: ++ marks.append((pos, len(out))) + if pos < len(text): + out.extend(self._encode_plain(text[pos:])) + return out +@@ -283,6 +292,11 @@ class Tokenizer: + """ + return self._encode_matching(text, self._special_re if parse_special else self._always_re, plain) + ++ def encode_marked(self, text: str, parse_special: bool = False, plain=()) -> tuple[list[int], list[tuple[int, int]]]: ++ """encode() and its resume points: (end offset in `text`, ids so far) after every special-token match.""" ++ marks: list[tuple[int, int]] = [] ++ return self._encode_matching(text, self._special_re if parse_special else self._always_re, plain, marks), marks ++ + def token_bytes(self, i: int) -> bytes: + """The raw bytes of one token (a multi-byte character can be split across tokens).""" + cache = self._bytes_cache +@@ -303,6 +317,111 @@ class Tokenizer: + return b"".join(self.token_bytes(i) for i in ids).decode("utf-8", errors=errors) + + ++# ------------------------------------------------------------------ incremental prompts (#567) ++def common_prefix_len(a: str, b: str) -> int: ++ """Length of the longest common prefix, by slices (C speed) rather than a loop over characters.""" ++ n = min(len(a), len(b)) ++ lo, step = 0, 4096 ++ while lo < n: # whole chunks first, doubling ++ hi = min(n, lo + step) ++ if a[lo:hi] != b[lo:hi]: ++ break ++ lo, step = hi, min(step * 2, 1 << 20) ++ else: ++ return n ++ while hi - lo > 1: # the first difference is in [lo, hi): bisect it ++ mid = (lo + hi) // 2 ++ if a[lo:mid] == b[lo:mid]: ++ lo = mid ++ else: ++ hi = mid ++ return lo ++ ++ ++def plain_agree_until(a, b, limit: int) -> int: ++ """The first position below `limit` where the `plain` spans (#537) of two texts differ - a position inside a ++ span in one and not in the other - or `limit`. Spans are (start, end), sorted and disjoint.""" ++ def clip(spans): ++ return [(x, min(y, limit)) for x, y in spans if x < limit and y > x] ++ ca, cb = clip(a), clip(b) ++ for (xa, ya), (xb, yb) in zip(ca, cb): ++ if (xa, ya) != (xb, yb): ++ return min(xa, xb) if xa != xb else min(ya, yb) ++ if len(ca) != len(cb): ++ return (ca if len(ca) > len(cb) else cb)[min(len(ca), len(cb))][0] ++ return limit ++ ++ ++class PromptEncoder: ++ """`tok.encode(text, parse_special=True, plain=...)` for chat prompts, reusing the ids of an earlier prompt. ++ ++ A chat client sends the whole conversation every turn, and the rendered prompt of turn n+1 starts with most of ++ turn n's. Running BPE over all of it again costs hundreds of milliseconds at 100K tokens, even when the engine ++ reuses the whole prefix. So the ids up to a SPECIAL-TOKEN BOUNDARY both prompts share are taken as they were, ++ and only the rest is encoded. ++ ++ WHY THAT IS EXACT. Encoding splits the text at special-token matches and BPE-encodes each stretch between two ++ matches on its own (`_encode_matching`): nothing crosses a match. Whether a match is taken at position q ++ depends only on text[q : q + max_special_len] and on whether q lies inside a `plain` span (#537: a ++ quoted in a message is text). So if the two texts agree up to L and their plain spans agree below L, every ++ decision the scan makes at a position before c = the end of some match, with c + max_special_len - 1 <= L, is ++ the same for both - the same matches, the same stretches, the same ids up to c - and from c the scan starts ++ afresh, exactly as it does on text[c:] with the spans shifted by c. Hence encode(new) == ids_old[:ids at c] + ++ encode(new[c:]). The margin matters only for a vocabulary where a literal overlaps another's end; it costs ++ re-encoding one short stretch. ++ ++ A few recent prompts are kept (one per conversation: a prompt replaces the one it continues, that is, one whose ++ text it reuses for more than half; a prompt that shares only the template's head with every entry is a new ++ entry), so a second client does not evict the first. `encode(..., store=False)` reads the kept prompts without ++ writing one (a token count is not a conversation's turn). The tokenizer needs `encode_marked` and ++ `max_special_len`. ++ """ ++ ++ def __init__(self, tok, keep: int = 4): ++ self.tok, self.keep = tok, keep ++ # (text, plain spans, ids, mark ends, mark counts) ++ self.entries: list[tuple[str, tuple, list[int], list[int], list[int]]] = [] ++ self.lock = threading.Lock() ++ self.last_reused = 0 # characters taken from an earlier prompt (for tests) ++ ++ def encode(self, text: str, plain=(), store: bool = True) -> list[int]: ++ plain = tuple(sorted(plain)) ++ with self.lock: ++ entries = self.entries # replaced, never changed in place ++ margin = self.tok.max_special_len - 1 ++ src, cut_k = None, -1 ++ for e in entries: ++ same = plain_agree_until(e[1], plain, common_prefix_len(e[0], text)) ++ k = bisect.bisect_right(e[3], same - margin) - 1 # the last boundary c with c <= same - margin ++ if k >= 0 and (src is None or e[3][k] > src[3][cut_k]): ++ src, cut_k = e, k ++ if src is None: ++ ids, marks = self.tok.encode_marked(text, parse_special=True, plain=plain) ++ ends, counts = [m[0] for m in marks], [m[1] for m in marks] ++ self.last_reused = 0 ++ else: ++ c, n = src[3][cut_k], src[4][cut_k] ++ rest = tuple((max(x - c, 0), y - c) for x, y in plain if y > c) ++ tail, marks = self.tok.encode_marked(text[c:], parse_special=True, plain=rest) ++ ids = src[2][:n] + tail ++ ends = src[3][:cut_k + 1] + [c + m[0] for m in marks] ++ counts = src[4][:cut_k + 1] + [n + m[1] for m in marks] ++ self.last_reused = c ++ if not store: ++ return ids ++ # every rendered prompt starts with the template's header, so sharing a special boundary does not make two ++ # prompts one conversation: only the entry whose text is mostly reused is replaced ++ old = src if src is not None and 2 * src[3][cut_k] > len(src[0]) else None ++ with self.lock: ++ self.entries = [(text, plain, ids, ends, counts)] + [e for e in self.entries if e is not old][:self.keep - 1] ++ return list(ids) ++ ++ def forget(self, ids: list[int]) -> None: ++ """The prompt with these ids was refused: it is not kept.""" ++ with self.lock: ++ self.entries = [e for e in self.entries if e[2] != ids] ++ ++ + # ------------------------------------------------------------------ the pack's tokenizer/ directory + def extract(gguf_path, out_dir) -> dict: + """Write the tokenizer into `/tokenizer/` so the engine never opens the weight shards for it.""" +diff --git a/tools/test_strata_tokenizer.py b/tools/test_strata_tokenizer.py +index 17c3f698..e3cb498c 100644 +--- a/tools/test_strata_tokenizer.py ++++ b/tools/test_strata_tokenizer.py +@@ -1,9 +1,18 @@ +-"""The tokenizer's piece cache (0.1.41): the ids are the ones the merge loop gives, cache cold, warm or switched off. ++"""Tests for tools/strata_tokenizer.py: llama.cpp's reference vectors, the piece cache (0.1.41: the ids are the ones the ++merge loop gives, cache cold, warm or switched off), and the incremental prompt encoder, which must give exactly the ++ids of a plain encode. + +- python -m unittest tools.test_strata_tokenizer ++ python -m unittest tools.test_strata_tokenizer -v ++ ++The real vocabulary comes from $STRATA_TOKENIZER (a pack's tokenizer/ directory) when set, else from llama.cpp's ++`models/ggml-vocab-qwen35.gguf` next to the gguf-py the tools use ($STRATA_GGUF_PY, third_party/llama.cpp); the ++tests that need it are skipped without one. The chat template's USER_DEFINED tokens (``, ``, ++...) are added to the llama.cpp vocabulary, which does not have them, as the pack's has. + """ + from __future__ import annotations + ++import json ++import os + import random + import sys + import unittest +@@ -11,8 +20,295 @@ from pathlib import Path + + ROOT = Path(__file__).resolve().parents[1] + sys.path.insert(0, str(ROOT / "tools")) ++sys.path.insert(0, str(ROOT)) + import strata_tokenizer as ST # noqa: E402 + ++ADDED = ["", "", "", "", "", "", ++ "<|vision_start|>", "<|image_pad|>", "<|vision_end|>"] ++ ++ ++def vocab_gguf() -> Path | None: ++ cands = [os.environ.get("STRATA_GGUF_PY"), ROOT / "third_party" / "llama.cpp" / "gguf-py", "/opt/llama.cpp/gguf-py"] ++ for c in cands: ++ if c and (Path(c).parent / "models" / "ggml-vocab-qwen35.gguf").is_file(): ++ return Path(c).parent / "models" / "ggml-vocab-qwen35.gguf" ++ return None ++ ++ ++def load_tokenizer(): ++ pack = os.environ.get("STRATA_TOKENIZER") ++ if pack and (Path(pack) / "vocab.json").is_file(): ++ t = Path(pack) ++ vocab = json.loads((t / "vocab.json").read_text(encoding="utf-8")) ++ tokens = [None] * len(vocab) ++ for s, i in vocab.items(): ++ tokens[i] = s ++ return ST.Tokenizer(tokens, (t / "merges.txt").read_text(encoding="utf-8").split("\n"), ++ json.loads((t / "token_type.json").read_text())) ++ path = vocab_gguf() ++ return None if path is None else gguf_tokenizer(path) ++ ++ ++def gguf_tokenizer(path): ++ """llama.cpp's qwen35 vocabulary with the template's added tokens: the one its test vectors were made with.""" ++ from gguf_reader import GGUFFile ++ md = GGUFFile(path).metadata ++ tokens, types = list(md["tokenizer.ggml.tokens"]), list(md["tokenizer.ggml.token_type"]) ++ for s in ADDED: ++ if s not in tokens: ++ tokens.append(s) ++ types.append(4) ++ return ST.Tokenizer(tokens, list(md["tokenizer.ggml.merges"]), types) ++ ++ ++def byte_tokenizer(specials: list[str]) -> ST.Tokenizer: ++ """A vocabulary of the 256 byte tokens and `specials` (CONTROL), no merges: small enough to build per test.""" ++ tokens = [ST.BYTE_TO_UNICODE[b] for b in range(256)] + specials ++ return ST.Tokenizer(tokens, [], [1] * 256 + [3] * len(specials)) ++ ++ ++WORDS = ["the", "tokenizer", " ", " ", "\n", "\n\n", "\t", "def", "f(x):", "return", "x", "==", "->", "café", ++ "你好", "\U0001f600", "é", "1234", "don't", "I'm", "<|im_", "", ++ " ", "\r\n", "Hello,", "world!", "<", ">", "|", " ", "x" * 70] ++ ++ ++def random_text(rnd: random.Random, n: int, specials=()) -> str: ++ parts = [] ++ for _ in range(n): ++ parts.append(rnd.choice(WORDS) if not specials or rnd.random() > 0.05 else rnd.choice(specials)) ++ if rnd.random() < 0.5: ++ parts.append(" ") ++ return "".join(parts) ++ ++ ++def conversation_prompts(template, rnd: random.Random) -> list[tuple[str, str]]: ++ """Rendered prompts, in the order a server would see them, from clients that grow a conversation turn after turn, ++ use tools, send images, change the start of the prompt, edit an earlier turn, send a shorter conversation, and ++ interleave with each other: (what happens, prompt).""" ++ out: list[tuple[str, str]] = [] ++ specials = ["<|im_start|>", "<|im_end|>", "", "", "", ""] ++ weather = [{"type": "function", "function": {"name": "get_weather", "description": "Get the weather", ++ "parameters": {"type": "object", "properties": { ++ "city": {"type": "string"}}, "required": ["city"]}}}] ++ ++ def add(what, msgs, tools=None, **kw): ++ out.append((what, template.render(msgs, tools=tools, **kw))) ++ ++ # several turns, with unicode, a literal special token inside the text, and reasoning kept or not ++ msgs = [{"role": "system", "content": "Be brief. Réponds en français si on te parle en français."}] ++ for turn in range(6): ++ msgs.append({"role": "user", "content": random_text(rnd, 60, specials) + " é你\U0001f600 " * turn}) ++ add(f"turn {turn}", msgs) ++ answer = {"role": "assistant", "content": random_text(rnd, 30)} ++ if turn % 2: ++ answer["reasoning_content"] = random_text(rnd, 30) ++ msgs.append(answer) ++ grown = list(msgs) ++ ++ # tools: definitions, a call, its response, then a following turn ++ tmsgs = [{"role": "user", "content": "Weather in Oslo?"}] ++ add("tools 1", tmsgs, weather) ++ tmsgs += [{"role": "assistant", "content": "", "tool_calls": [{"type": "function", "function": { ++ "name": "get_weather", "arguments": {"city": "Oslo"}}}]}, ++ {"role": "tool", "content": "{\"temp\": 3}", "name": "get_weather"}] ++ add("tools 2", tmsgs, weather) ++ tmsgs += [{"role": "assistant", "content": "3 degrees."}, {"role": "user", "content": "And in Bergen? " * 20}] ++ add("tools 3", tmsgs, weather) ++ ++ # images: one per user turn, the placeholder tokens are specials ++ imsgs = [] ++ for turn in range(3): ++ imsgs.append({"role": "user", "content": [{"type": "image", "source": f"img{turn}"}, ++ {"type": "text", "text": f"What is on picture {turn}?"}]}) ++ add(f"images {turn}", imsgs) ++ imsgs.append({"role": "assistant", "content": "A cat."}) ++ ++ # the start of the prompt changes (a thinking level, a system message), then the old conversation grows again ++ add("effort changed", grown, reasoning_effort="low") ++ add("system changed", [dict(grown[0], content=grown[0]["content"] + " Be kind.")] + grown[1:]) ++ grown.append({"role": "user", "content": "one more question"}) ++ add("grown again", grown) ++ ++ # an earlier turn edited, then a conversation that got shorter, then the full one again ++ edited = list(grown) ++ edited[3] = dict(edited[3], content=edited[3]["content"] + " (edited)") ++ add("edited", edited) ++ add("shorter", grown[:5]) ++ add("empty user", [{"role": "user", "content": ""}]) ++ add("full again", grown) ++ return out ++ ++ ++class Oracle(unittest.TestCase): ++ """llama.cpp's own tokenizer test vectors for qwen35.""" ++ ++ def test_llama_cpp_vectors(self): ++ path = vocab_gguf() ++ if path is None or not path.with_suffix(".gguf.inp").is_file(): ++ self.skipTest("llama.cpp's ggml-vocab-qwen35.gguf test vectors are not here") ++ tk = gguf_tokenizer(path) # not $STRATA_TOKENIZER: the vectors belong to this vocabulary ++ inp = path.with_suffix(".gguf.inp").read_text(encoding="utf-8").split("\n__ggml_vocab_test__\n") ++ out = path.with_suffix(".gguf.out").read_text(encoding="utf-8").split("\n") ++ n = 0 ++ for text, want in zip(inp, out): ++ if text.endswith("\n__ggml_vocab_test__"): ++ text = text[:-len("\n__ggml_vocab_test__")] ++ self.assertEqual(tk.encode(text), [int(x) for x in want.split()], repr(text[:40])) ++ n += 1 ++ self.assertGreater(n, 40) ++ ++ ++class Marks(unittest.TestCase): ++ @classmethod ++ def setUpClass(cls): ++ cls.tk = load_tokenizer() ++ if cls.tk is None: ++ raise unittest.SkipTest("no tokenizer (set STRATA_TOKENIZER or STRATA_GGUF_PY)") ++ ++ def test_marks(self): ++ text = "a<|im_start|>user\nhi<|im_end|>\n\n" ++ ids, marks = self.tk.encode_marked(text, parse_special=True) ++ self.assertEqual(ids, self.tk.encode(text, parse_special=True)) ++ self.assertEqual([m[0] for m in marks], [text.index("user"), text.index("<|im_end|>") + len("<|im_end|>"), ++ len(text) - 1]) ++ for end, n in marks: # each mark: the ids of the text up to it ++ self.assertEqual(ids[:n], self.tk.encode(text[:end], parse_special=True)) ++ ++ ++class Prompts(unittest.TestCase): ++ """PromptEncoder: the ids of a plain encode, whatever the earlier prompts were.""" ++ ++ @classmethod ++ def setUpClass(cls): ++ cls.tk = load_tokenizer() ++ if cls.tk is None: ++ raise unittest.SkipTest("no tokenizer (set STRATA_TOKENIZER or STRATA_GGUF_PY)") ++ from serve.frontend import ChatTemplate ++ cls.template = ChatTemplate(ROOT / "serve" / "chat_template.jinja") ++ ++ def check(self, enc, text): ++ self.assertEqual(enc.encode(text), self.tk.encode(text, parse_special=True)) ++ ++ def test_chat_golden(self): ++ cases = json.loads((ROOT / "serve" / "chat_golden.json").read_text(encoding="utf-8")) ++ enc = ST.PromptEncoder(self.tk) ++ for c in cases + cases: # the second round finds its own earlier prompts ++ if c.get("rendered"): ++ with self.subTest(case=c["name"]): ++ self.check(enc, c["rendered"]) ++ ++ def test_conversations(self): ++ for seed in range(3): ++ prompts = conversation_prompts(self.template, random.Random(seed)) ++ wants = [self.tk.encode(prompt, parse_special=True) for _, prompt in prompts] ++ for keep in (1, 2, 4): ++ enc = ST.PromptEncoder(self.tk, keep=keep) ++ for (what, prompt), want in zip(prompts, wants): ++ with self.subTest(seed=seed, keep=keep, what=what): ++ self.assertEqual(enc.encode(prompt), want) ++ ++ def test_growing_turn_reuses_the_previous_prompt(self): ++ enc = ST.PromptEncoder(self.tk) ++ msgs = [{"role": "user", "content": "first question " * 50}] ++ prev = "" ++ for turn in range(4): ++ prompt = self.template.render(msgs) ++ self.check(enc, prompt) ++ if turn: ++ self.assertGreater(enc.last_reused, len(prev) - 100) ++ prev = prompt ++ msgs += [{"role": "assistant", "content": f"answer {turn}"}, {"role": "user", "content": f"question {turn}"}] ++ ++ def test_prefix_of_an_earlier_prompt_and_unrelated_prompts(self): ++ enc = ST.PromptEncoder(self.tk, keep=2) ++ a = self.template.render([{"role": "user", "content": "first question " * 50}]) ++ b = self.template.render([{"role": "user", "content": "another conversation"}]) ++ c = self.template.render([{"role": "user", "content": "a third one"}]) ++ for text in (a, a[:len(a) // 2], b, c, a, a + "tail", "", "<|im_end|>", "<|im_end|><|im_end|>x"): ++ with self.subTest(text=text[:30]): ++ self.check(enc, text) ++ ++ def test_a_read_only_encode_gives_the_ids_and_stores_nothing(self): ++ enc = ST.PromptEncoder(self.tk) ++ a = self.template.render([{"role": "user", "content": "first question " * 50}]) ++ enc.encode(a) ++ kept = list(enc.entries) ++ for text in (a + "<|im_start|>user\nand a follow-up<|im_end|>\n", a[:len(a) // 2], "<|im_end|>x", ++ self.template.render([{"role": "user", "content": "another conversation"}])): ++ with self.subTest(text=text[:30]): ++ self.assertEqual(enc.encode(text, store=False), self.tk.encode(text, parse_special=True)) ++ self.assertTrue(enc.entries == kept, "a read-only encode wrote an entry") ++ self.assertEqual(enc.encode(a + "tail", store=False), self.tk.encode(a + "tail", parse_special=True)) ++ self.assertGreater(enc.last_reused, len(a) // 2) ++ ++ def test_the_returned_ids_are_the_callers(self): ++ enc = ST.PromptEncoder(self.tk) ++ text = "<|im_start|>user\nhi<|im_end|>\n" + "and a tail longer than the look-ahead margin. " * 3 ++ want = self.tk.encode(text, parse_special=True) ++ got = enc.encode(text) ++ got[0] = -1 # inside the part the next call takes from this one ++ got.append(-1) ++ self.assertEqual(enc.encode(text), want) ++ self.assertGreater(enc.last_reused, 0) ++ ++ ++class Margin(unittest.TestCase): ++ """A special literal that is a prefix of a longer one: cutting at the shorter one's end without the look-ahead ++ margin would give different ids.""" ++ ++ def test_longer_literal_past_the_common_prefix(self): ++ tk = byte_tokenizer(["<|a|>", "<|a|>zz"]) ++ enc = ST.PromptEncoder(tk) ++ first = "x<|a|>y" + "<|a|>" * 3 + "q" ++ self.assertEqual(enc.encode(first), tk.encode(first, parse_special=True)) ++ second = "x<|a|>y" + "<|a|>" * 3 + "zz" # shares everything up to the last <|a|> ++ got = enc.encode(second) ++ self.assertEqual(got, tk.encode(second, parse_special=True)) ++ self.assertEqual(got[-1], tk.special_tokens["<|a|>zz"]) ++ self.assertGreater(enc.last_reused, 0) # the earlier boundaries were still reused ++ ++ def test_literal_ending_inside_another(self): ++ tk = byte_tokenizer(["<|a|>", "b<|a|>c"]) ++ enc = ST.PromptEncoder(tk) ++ for text in ("zzb<|a|>d", "zzb<|a|>c", "zzb<|a|>cz<|a|>", "zzb<|a|>cz<|a|>b<|a|>c"): ++ with self.subTest(text=text): ++ self.assertEqual(enc.encode(text), tk.encode(text, parse_special=True)) ++ ++ def test_random_specials(self): ++ """Literals that extend or end inside one another, texts that change right after one of them.""" ++ rnd = random.Random(3) ++ ++ def word(lo, hi): ++ return "".join(rnd.choice("<|ab>") for _ in range(rnd.randint(lo, hi))) ++ for trial in range(60): ++ a, b = word(2, 4), word(2, 4) ++ lits = sorted({a, a + word(1, 3), b, word(1, 2) + b + word(0, 2)}) ++ tk = byte_tokenizer(lits) ++ enc = ST.PromptEncoder(tk, keep=3) ++ pieces = lits + [x[:rnd.randint(1, len(x))] for x in lits] + ["x", "y", "a", "|"] ++ base = "".join(rnd.choice(pieces) for _ in range(60)) ++ for _ in range(30): ++ cut = rnd.randint(0, len(base)) ++ text = base[:cut] + "".join(rnd.choice(pieces) for _ in range(rnd.randint(0, 6))) ++ self.assertEqual(enc.encode(text), tk.encode(text, parse_special=True), (lits, text)) ++ if rnd.random() < 0.3: ++ base = text + base[cut:] ++ ++ ++class CommonPrefix(unittest.TestCase): ++ def test_against_a_loop(self): ++ rnd = random.Random(5) ++ for _ in range(500): ++ a = "".join(rnd.choice("ab") for _ in range(rnd.randint(0, 9000))) ++ k = rnd.randint(0, len(a)) ++ b = a[:k] + rnd.choice(["", "c", "cd"]) + a[k:k + rnd.randint(0, 50)] ++ want = 0 ++ while want < min(len(a), len(b)) and a[want] == b[want]: ++ want += 1 ++ self.assertEqual(ST.common_prefix_len(a, b), want) ++ self.assertEqual(ST.common_prefix_len(b, a), want) ++ ++ + + def toy() -> ST.Tokenizer: + base = [ST.BYTE_TO_UNICODE[b] for b in range(256)] diff --git a/results/round39/serve-post.log b/results/round39/serve-post.log new file mode 100644 index 0000000..54e4726 --- /dev/null +++ b/results/round39/serve-post.log @@ -0,0 +1,802 @@ +.........[strata] done: 12 tokens in 0 s (185725.6 tok/s) (stop, cancel=False) +..[strata] other apps now use the Chat settings: temperature=0.5 +.[strata] done: 12 tokens in 0 s (299593.1 tok/s) (stop, cancel=False) +.....[strata] done: 12 tokens in 0 s (324720.3 tok/s) (stop, cancel=False) +.[strata] other apps now use the Chat settings: temperature=0.5 +[strata] other apps now use the Chat settings: temperature=0.5 +[strata] other apps now use the Chat settings: temperature=0.5 +..[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] 400 invalid request: request exceeds the context window: prompt (150722 tokens) + max tokens (16) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] 400 invalid request: request exceeds the context window: prompt (1024 tokens) + max tokens (131072) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 130040 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +.[strata] done: 1 tokens in 0 s (47662.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) + max tokens (1) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] done: 1 tokens in 0 s (99864.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] done: 1 tokens in 0 s (91180.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) + max tokens (1) exceeds the context (131072); requests are never truncated. Send a smaller max_tokens (at most 0 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] done: 1 tokens in 0 s (91180.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] 400 invalid request: request exceeds the context window: prompt (131064 tokens) leaves no room to answer in the context (131072); requests are never truncated +[strata] done: 1 tokens in 0 s (43240.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 1 tokens in 0 s (58254.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +....sssss....[strata] 400 invalid request: unknown reasoning effort 'huge': use none, low, medium or high +.[strata] done: 24 tokens in 0 s (241979.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (333322.2 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (342392.2 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 500) +[strata] done: 24 tokens in 0 s (348315.9 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 2047) +[strata] done: 24 tokens in 0 s (375609.3 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 2048) +[strata] done: 24 tokens in 0 s (381300.4 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 8191) +[strata] done: 24 tokens in 0 s (375609.3 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 8192) +[strata] done: 24 tokens in 0 s (327893.5 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 31999) +[strata] done: 24 tokens in 0 s (370085.6 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (351969.6 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (343560.7 tok/s) (stop, cancel=False) +.[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (375609.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (387166.5 tok/s) (stop, cancel=False) +[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (276547.5 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (394758.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (374213.0 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (305040.3 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (405900.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (385683.1 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (311651.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (381300.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (379861.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (324720.3 tok/s) (stop, cancel=False) +.[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (348315.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (387166.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (387166.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (399457.5 tok/s) (stop, cancel=False) +[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (409200.4 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (405900.4 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (405900.4 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (388661.4 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (399457.5 tok/s) (stop, cancel=False) +[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (394758.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (388661.4 tok/s) (stop, cancel=False) +.[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (391686.0 tok/s) (stop, cancel=False) +[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (401049.0 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (407543.7 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (387166.5 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (387166.5 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (391686.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (387166.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (391686.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (391686.0 tok/s) (stop, cancel=False) +.[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (292625.9 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (311651.1 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (306900.3 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (296941.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (328965.0 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (285975.3 tok/s) (stop, cancel=False) +.[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (399457.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (381300.4 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (315558.9 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (401049.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (387166.5 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (333322.2 tok/s) (stop, cancel=False) +.[strata] thinking limited to 5964 of max_tokens 7951 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (371451.3 tok/s) (stop, cancel=False) +[strata] thinking limited to 6073 of max_tokens 8097 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (385683.1 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (351969.6 tok/s) (stop, cancel=False) +.[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (381300.4 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (347114.8 tok/s) (stop, cancel=False) +[strata] thinking capped at 25 of max_tokens 50 to leave room for the answer (budget 20000) +[strata] done: 24 tokens in 0 s (368730.0 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (303202.7 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (370085.6 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (381300.4 tok/s) (stop, cancel=False) +.[strata] done: 24 tokens in 0 s (113232.1 tok/s) (stop, cancel=False) +[strata] thinking limited to 6005 of max_tokens 8006 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (139422.8 tok/s) (stop, cancel=False) +[strata] thinking limited to 6114 of max_tokens 8152 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (133860.8 tok/s) (stop, cancel=False) +[strata] thinking limited to 5952 of max_tokens 7935 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (150919.5 tok/s) (stop, cancel=False) +.[strata] thinking limited to 5266 of max_tokens 7021 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (358232.4 tok/s) (stop, cancel=False) +[strata] thinking limited to 5213 of max_tokens 6950 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (393216.0 tok/s) (stop, cancel=False) +[strata] done: 24 tokens in 0 s (381300.4 tok/s) (stop, cancel=False) +.................[strata] the engine stopped unexpectedly (exit code -15). Unrecoverable native verification failure: verify: timed out at layer 18; its GPU waits were released and the GPU finished (#267) The next request starts the engine again. +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -15); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 1 tokens in 0 s (20068.4 tok/s) (length, cancel=False) +...[strata] the engine stopped unexpectedly (exit code -15). Unrecoverable native verification failure: verify: timed out at layer 18; its GPU waits were released and the GPU finished (#267) The next request starts the engine again. +[strata] done: 0 tokens in 0 s (0.0 tok/s) (error, cancel=False) +[strata] the engine had stopped (exit code -15); starting it again (a minute or two) ... +[strata] the engine is running again +[strata] done: 1 tokens in 0 s (27060.0 tok/s) (length, cancel=False) +................................[strata] image: 512x768, n=2, 8 steps, seed 11 +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5726: ResourceWarning: unclosed file <_io.TextIOWrapper name=11 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 0.2 s (keep) +.[strata] 400 invalid request: size 999x999: each side a multiple of 32 from 256 to 2048 +[strata] image: 1024x1024, n=1, 8 steps, seed 1390336375 +[strata] image job failed: sd-cli exited with code 1 +[INFO ] loading tensors completed, taking 0.10s +[ERROR] ggml_vulkan: ErrorDeviceLost +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5730: ResourceWarning: unclosed file <_io.TextIOWrapper name=11 encoding='UTF-8'> + return +ResourceWarning: Enable tracemalloc to get the object allocation traceback +.[strata] image: 1024x1024, n=1, 8 steps, seed 320404728 +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5726: ResourceWarning: unclosed file <_io.TextIOWrapper name=11 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 1.6 s (keep) +[strata] done: 8 tokens in 0 s (127100.1 tok/s) (length, cancel=False) +...[strata] image: 1024x1024, n=1, 8 steps, seed 0 +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5726: ResourceWarning: unclosed file <_io.TextIOWrapper name=11 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 0.3 s (keep) +[strata] image: 1024x1024, n=1, 8 steps, seed 1 +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5726: ResourceWarning: unclosed file <_io.TextIOWrapper name=9 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 0.3 s (keep) +[strata] image: 1024x1024, n=1, 8 steps, seed 2 +[strata] image done in 0.3 s (keep) +[strata] image: 1024x1024, n=1, 8 steps, seed 3 +/Users/dawn/proj/swift1.5-qwen3.8-next-flash-ps-scr-qat-persona/vendor/strata/serve/server.py:5726: ResourceWarning: unclosed file <_io.TextIOWrapper name=9 encoding='UTF-8'> + images, figures = gen.run(job, cancel) +ResourceWarning: Enable tracemalloc to get the object allocation traceback +[strata] image done in 0.3 s (keep) +.[strata] image job: the conversation (123 tokens) saved; unloading the model +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] the conversation from before the image job is back (123 tokens) +...[strata] image job: the engine's expert cache gave VRAM back (9000 MiB free) +.[strata] image job: the conversation (123 tokens) saved; unloading the model +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] the conversation from before the image job is back (123 tokens) +.........[strata] done: 3 tokens in 0 s (73584.3 tok/s) (stop, cancel=False) +F[strata] loading the model again (it was unloaded) ... +.[strata] model unloaded; the next request loads it again +.....[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 7 tokens in 0 s (132252.8 tok/s) (stop, cancel=False) +[strata] done: 7 tokens in 0 s (248814.6 tok/s) (stop, cancel=False) +[strata] model unloaded; the next request loads it again +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 7 tokens in 0 s (105993.2 tok/s) (stop, cancel=False) +.[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +....[strata] 400 invalid request: logit_bias: expected an object or a list of [token_id, bias] pairs +[strata] 400 invalid request: logit_bias: this engine does not support token biases; use a supporting CUDA/HIP build +[strata] 400 invalid request: logit_bias: biases must be finite numbers between -100 and 100; false is a pair-list ban +..........[strata] MCP server 'missing' did not start: could not start 'strata-no-such-program': [Errno 2] No such file or directory: 'strata-no-such-program'. Its tools are left out; the chat works without them. +[strata] MCP server 'paged': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +...[strata] MCP server 'solo': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'solo' had stopped; starting it again +[strata] MCP server 'solo': 6 tools (echo, fail, sleep, big, die, add) +........[strata] MCP server 'down' did not start: could not reach http://127.0.0.1:9/mcp: . Its tools are left out; the chat works without them. +[strata] MCP server 'web': 6 tools (echo, fail, sleep, big, die, add) +..[strata] thinking limited to 5955 of max_tokens 7939 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (117008.5 tok/s) (stop, cancel=False) +.[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 83 tokens in 0 s (113065.0 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +.[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 120 tokens in 0 s (142139.6 tok/s) (stop, cancel=False) +[strata] tool fake__echo: ok, 5 characters in 0.0 s +[strata] thinking limited to 4520 of max_tokens 6026 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 120 tokens in 0 s (170211.9 tok/s) (stop, cancel=False) +[strata] tool fake__echo: ok, 5 characters in 0.0 s +[strata] thinking limited to 4380 of max_tokens 5839 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 120 tokens in 0 s (183491.2 tok/s) (stop, cancel=False) +.[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 141 tokens in 0 s (122493.1 tok/s) (stop, cancel=False) +[strata] tool fake__add: ok, 1 characters in 0.0 s +[strata] thinking limited to 4507 of max_tokens 6009 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (288497.1 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] refused an API request from the web page 'http://evil.example' (no API key; add its host to "allowed_hosts" or its origin to "cors_origins" in the config) +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 13 tokens in 0 s (175324.6 tok/s) (stop, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4590 of max_tokens 6119 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 109 tokens in 0 s (144357.2 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 5949 of max_tokens 7932 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 116 tokens in 0 s (108699.6 tok/s) (stop, cancel=False) +[strata] thinking limited to 5949 of max_tokens 7932 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 16 tokens in 0 s (141579.9 tok/s) (stop, cancel=False) +[strata] thinking limited to 5949 of max_tokens 7932 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 16 tokens in 0 s (188508.0 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4665 of max_tokens 6220 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 120 tokens in 0 s (107408.6 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 134 tokens in 0 s (152272.2 tok/s) (stop, cancel=False) +[strata] tool fake__echo: ok, 19 characters in 0.0 s +[strata] thinking limited to 4499 of max_tokens 5998 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 31 tokens in 0 s (144150.1 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 84 tokens in 0 s (149225.6 tok/s) (stop, cancel=False) +[strata] tool fake__fail: error, 27 characters in 0.0 s +[strata] thinking limited to 4530 of max_tokens 6040 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 21 tokens in 0 s (164023.1 tok/s) (stop, cancel=False) +.[strata] MCP server 'fake': 6 tools (echo, fail, sleep, big, die, add) +[strata] MCP server 'broken' did not start: the server stopped (exit code 2): cannot start: missing configuration. Its tools are left out; the chat works without them. +[strata] thinking limited to 4660 of max_tokens 6213 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 115 tokens in 0 s (146922.0 tok/s) (stop, cancel=False) +[strata] tool fake__big: ok, 3,000 characters in 0.0 s (truncated for the model) +[strata] thinking limited to 4090 of max_tokens 5453 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 16 tokens in 0 s (211034.2 tok/s) (stop, cancel=False) +.[strata] done: 8 tokens in 0 s (147816.9 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 8 tokens in 0 s (209715.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 8 tokens in 0 s (229824.9 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] done: 8 tokens in 0 s (215092.5 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] done: 8 tokens in 0 s (268435.5 tok/s) (length, cancel=False) +.[strata] done: 8 tokens in 0 s (228261.4 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 8 tokens in 0 s (286790.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +..[strata] done: 8 tokens in 0 s (250406.2 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] thinking limited to 2882 of max_tokens 3842 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 24 tokens in 0 s (292625.9 tok/s) (stop, cancel=False) +....[strata] 400 invalid request: request exceeds the context window: prompt (183 tokens) + max tokens (20000) exceeds the context (16384); requests are never truncated. Send a smaller max_tokens (at most 16193 here), or add "fit_max_tokens": true to the model's strata-.json to shorten it to the room left (#545) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] thinking capped at 128 of max_tokens 256 to leave room for the answer (budget 1024) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +[strata] thinking capped at 128 of max_tokens 256 to leave room for the answer (budget 1024) +[strata] done: 0 tokens in 0 s (0.0 tok/s) (length, cancel=False) +[strata] the reply reached max tokens while still thinking, so it has no answer: a thinking budget (reasoning_budget_tokens, in the request or in strata-.json for every request) leaves room to answer +.[strata] thinking limited to 12098 of max_tokens 16130 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 50 tokens in 0 s (252364.9 tok/s) (stop, cancel=False) +.[strata] done: 50 tokens in 0 s (246144.6 tok/s) (stop, cancel=False) +[strata] done: 50 tokens in 0 s (112087.2 tok/s) (stop, cancel=False) +[strata] thinking capped at 128 of max_tokens 256 to leave room for the answer (budget 1024) +[strata] done: 50 tokens in 0 s (264458.0 tok/s) (stop, cancel=False) +[strata] thinking capped at 128 of max_tokens 256 to leave room for the answer (budget 1024) +[strata] done: 50 tokens in 0 s (111610.0 tok/s) (stop, cancel=False) +..[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] done: 50 tokens in 0 s (253585.5 tok/s) (stop, cancel=False) +[strata] done: 50 tokens in 0 s (170777.9 tok/s) (stop, cancel=False) +.[strata] thinking limited to 2879 of max_tokens 3838 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 7 tokens in 0 s (180123.5 tok/s) (stop, cancel=False) +..[strata] "parallel" must be a whole number of requests (2..8), not '4': ignored +[strata] "parallel" must be a whole number of requests (2..8), not True: ignored +[strata] "parallel": 12 - the engine runs at most 8 requests together; it will use 8 +.[strata] answering: 6 of max 8 tokens, 6.0 tok/s, 1 s +[strata] answering: 5 of max 8 tokens, 8.1 tok/s, 1 s +[strata] done: 8 tokens in 1 s (6.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 1 s so far +[strata] done: 8 tokens in 2 s (6.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 2 s so far +[strata] answering: 7 of max 8 tokens, 6.1 tok/s, 3 s +[strata] done: 8 tokens in 3 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 3 s so far +[strata] answering: 8 of max 8 tokens, 6.3 tok/s, 3 s +[strata] done: 8 tokens in 3 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 3 s so far +[strata] answering: 7 of max 8 tokens, 6.9 tok/s, 4 s +[strata] done: 8 tokens in 4 s (7.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 4 s so far +[strata] answering: 8 of max 8 tokens, 7.6 tok/s, 4 s +[strata] done: 8 tokens in 4 s (7.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 4 s so far +[strata] answering: 6 of max 8 tokens, 5.3 tok/s, 5 s +[strata] answering: 7 of max 8 tokens, 7.0 tok/s, 5 s +[strata] done: 8 tokens in 5 s (5.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 5 s so far +[strata] done: 8 tokens in 6 s (5.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 6 s so far +[strata] answering: 6 of max 8 tokens, 5.7 tok/s, 6 s +[strata] done: 8 tokens in 7 s (6.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 7 s so far +[strata] answering: 8 of max 8 tokens, 5.8 tok/s, 7 s +[strata] done: 8 tokens in 7 s (5.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 7 s so far +[strata] answering: 7 of max 8 tokens, 6.2 tok/s, 8 s +[strata] done: 8 tokens in 8 s (6.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 8 s so far +[strata] answering: 8 of max 8 tokens, 7.5 tok/s, 8 s +[strata] done: 8 tokens in 8 s (7.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 8 s so far +[strata] answering: 8 of max 8 tokens, 7.5 tok/s, 9 s +[strata] done: 8 tokens in 9 s (7.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 9 s so far +[strata] answering: 8 of max 8 tokens, 7.0 tok/s, 9 s +[strata] done: 8 tokens in 9 s (7.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 9 s so far +[strata] answering: 6 of max 8 tokens, 5.4 tok/s, 10 s +[strata] answering: 7 of max 8 tokens, 6.9 tok/s, 10 s +[strata] done: 8 tokens in 10 s (5.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 10 s so far +[strata] done: 8 tokens in 11 s (5.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 11 s so far +[strata] answering: 6 of max 8 tokens, 5.4 tok/s, 12 s +[strata] done: 8 tokens in 12 s (5.8 tok/s) (length, cancel=False) +[strata] answering: 7 of max 8 tokens, 6.9 tok/s, 12 s +[strata] reading the prompt: 32 of 46 tokens, 12 s so far +[strata] done: 8 tokens in 12 s (6.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 12 s so far +[strata] answering: 7 of max 8 tokens, 6.4 tok/s, 13 s +[strata] done: 8 tokens in 13 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 13 s so far +[strata] answering: 8 of max 8 tokens, 6.7 tok/s, 13 s +[strata] done: 8 tokens in 13 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 14 s so far +[strata] answering: 7 of max 8 tokens, 6.4 tok/s, 14 s +[strata] done: 8 tokens in 14 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 14 s so far +[strata] answering: 8 of max 8 tokens, 6.4 tok/s, 15 s +[strata] done: 8 tokens in 15 s (6.4 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 15 s so far +[strata] answering: 6 of max 8 tokens, 5.7 tok/s, 16 s +[strata] done: 8 tokens in 16 s (6.0 tok/s) (length, cancel=False) +[strata] answering: 7 of max 8 tokens, 6.7 tok/s, 16 s +[strata] reading the prompt: 32 of 46 tokens, 16 s so far +[strata] done: 8 tokens in 16 s (5.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 16 s so far +[strata] answering: 6 of max 8 tokens, 5.7 tok/s, 17 s +[strata] answering: 7 of max 8 tokens, 7.0 tok/s, 17 s +[strata] done: 8 tokens in 17 s (5.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 17 s so far +[strata] done: 8 tokens in 18 s (5.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 18 s so far +[strata] answering: 6 of max 8 tokens, 5.5 tok/s, 18 s +[strata] done: 8 tokens in 19 s (5.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 19 s so far +[strata] answering: 8 of max 8 tokens, 5.9 tok/s, 19 s +[strata] done: 8 tokens in 19 s (5.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 19 s so far +[strata] answering: 6 of max 8 tokens, 5.8 tok/s, 20 s +[strata] done: 8 tokens in 20 s (6.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 20 s so far +[strata] answering: 8 of max 8 tokens, 6.3 tok/s, 20 s +[strata] done: 8 tokens in 20 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 20 s so far +[strata] answering: 6 of max 8 tokens, 6.0 tok/s, 21 s +[strata] done: 8 tokens in 21 s (6.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 22 s so far +[strata] answering: 8 of max 8 tokens, 6.5 tok/s, 22 s +[strata] done: 8 tokens in 22 s (6.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 22 s so far +[strata] answering: 7 of max 8 tokens, 6.9 tok/s, 23 s +[strata] done: 8 tokens in 23 s (6.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 23 s so far +[strata] answering: 8 of max 8 tokens, 6.1 tok/s, 23 s +[strata] done: 8 tokens in 23 s (6.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 23 s so far +[strata] answering: 6 of max 8 tokens, 5.7 tok/s, 24 s +[strata] done: 8 tokens in 24 s (6.0 tok/s) (length, cancel=False) +[strata] answering: 7 of max 8 tokens, 6.9 tok/s, 24 s +[strata] reading the prompt: 32 of 46 tokens, 24 s so far +[strata] done: 8 tokens in 24 s (5.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 25 s so far +[strata] answering: 6 of max 8 tokens, 6.0 tok/s, 25 s +[strata] done: 8 tokens in 25 s (6.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 26 s so far +[strata] answering: 8 of max 8 tokens, 6.5 tok/s, 26 s +[strata] done: 8 tokens in 26 s (6.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 26 s so far +[strata] answering: 7 of max 8 tokens, 6.7 tok/s, 27 s +[strata] done: 8 tokens in 27 s (6.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 27 s so far +[strata] answering: 8 of max 8 tokens, 6.3 tok/s, 27 s +[strata] done: 8 tokens in 27 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 27 s so far +[strata] answering: 6 of max 8 tokens, 5.6 tok/s, 28 s +[strata] done: 8 tokens in 28 s (6.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 28 s so far +[strata] answering: 8 of max 8 tokens, 6.6 tok/s, 28 s +[strata] done: 8 tokens in 28 s (6.6 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 28 s so far +[strata] answering: 7 of max 8 tokens, 6.7 tok/s, 29 s +[strata] done: 8 tokens in 29 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 29 s so far +[strata] answering: 8 of max 8 tokens, 6.5 tok/s, 30 s +[strata] done: 8 tokens in 30 s (6.5 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 30 s so far +[strata] answering: 7 of max 8 tokens, 6.8 tok/s, 30 s +[strata] done: 8 tokens in 31 s (6.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 31 s so far +[strata] answering: 8 of max 8 tokens, 6.9 tok/s, 31 s +[strata] done: 8 tokens in 31 s (6.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 31 s so far +[strata] answering: 8 of max 8 tokens, 7.1 tok/s, 32 s +[strata] done: 8 tokens in 32 s (7.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 32 s so far +[strata] answering: 8 of max 8 tokens, 6.7 tok/s, 32 s +[strata] done: 8 tokens in 32 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 32 s so far +[strata] answering: 7 of max 8 tokens, 6.4 tok/s, 33 s +[strata] done: 8 tokens in 33 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 33 s so far +[strata] answering: 8 of max 8 tokens, 7.0 tok/s, 33 s +[strata] done: 8 tokens in 33 s (7.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 33 s so far +[strata] answering: 7 of max 8 tokens, 6.8 tok/s, 34 s +[strata] done: 8 tokens in 34 s (6.8 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 34 s so far +[strata] answering: 8 of max 8 tokens, 6.0 tok/s, 35 s +[strata] done: 8 tokens in 35 s (6.0 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 35 s so far +[strata] answering: 6 of max 8 tokens, 5.8 tok/s, 35 s +[strata] done: 8 tokens in 36 s (6.3 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 36 s so far +[strata] answering: 8 of max 8 tokens, 6.7 tok/s, 36 s +[strata] done: 8 tokens in 36 s (6.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 36 s so far +[strata] answering: 8 of max 8 tokens, 7.7 tok/s, 37 s +[strata] done: 8 tokens in 37 s (7.7 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 45 tokens, 37 s so far +[strata] answering: 8 of max 8 tokens, 7.9 tok/s, 37 s +[strata] done: 8 tokens in 37 s (7.9 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 37 s so far +[strata] answering: 7 of max 8 tokens, 7.0 tok/s, 38 s +[strata] done: 8 tokens in 38 s (7.2 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 38 s so far +[strata] answering: 8 of max 8 tokens, 7.1 tok/s, 38 s +[strata] done: 8 tokens in 38 s (7.1 tok/s) (length, cancel=False) +[strata] reading the prompt: 32 of 46 tokens, 38 s so far +[strata] answering: 6 of max 8 tokens, 5.8 tok/s, 39 s +[strata] done: 8 tokens in 39 s (6.1 tok/s) (length, cancel=False) +[strata] answering: 8 of max 8 tokens, 7.0 tok/s, 39 s +[strata] done: 8 tokens in 39 s (7.0 tok/s) (length, cancel=False) +.[strata] thinking: 1 of max 200 tokens, 1.6 tok/s, 1 s +[strata] reading the prompt: 32 of 249 tokens, 1 s so far +[strata] thinking: 5 of max 200 tokens, 3.0 tok/s, 2 s +[strata] thinking: 4 of max 200 tokens, 6.1 tok/s, 2 s +[strata] thinking: 12 of max 200 tokens, 4.4 tok/s, 3 s +[strata] done: 10 tokens in 3 s (6.5 tok/s) (stop, cancel=False) +[strata] thinking: 18 of max 200 tokens, 4.8 tok/s, 4 s +[strata] thinking: 26 of max 200 tokens, 5.4 tok/s, 5 s +[strata] thinking: 34 of max 200 tokens, 5.7 tok/s, 6 s +[strata] thinking: 41 of max 200 tokens, 5.8 tok/s, 7 s +[strata] thinking: 49 of max 200 tokens, 6.0 tok/s, 9 s +[strata] done: 55 tokens in 10 s (6.0 tok/s) (stop, cancel=False) +.[strata] the engine stopped unexpectedly (exit code 1). The engine exited (code 1) after it reported: verify batch: layer 34 never rang (graph finished) - please report it at github.com/Niko1221/Strata/issues with the log. The next request starts the engine again. +[strata] done: 3 tokens in 1 s (3.2 tok/s) (error, cancel=False) +[strata] the engine stopped unexpectedly (exit code 1). The engine exited (code 1) after it reported: verify batch: layer 34 never rang (graph finished) - please report it at github.com/Niko1221/Strata/issues with the log. The next request starts the engine again. +[strata] done: 1 tokens in 1 s (2.9 tok/s) (error, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +.[strata] parallel: slot(s) 0, 1 were marked busy with no request using them (issue #1603); freed, the waiting request goes on +[strata] reading the prompt: 32 of 52 tokens, 1 s so far +[strata] answering: 8 of max 16 tokens, 7.5 tok/s, 2 s +[strata] done: 10 tokens in 3 s (7.2 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 288 of 3,037 tokens, 1 s so far +[strata] answering: 5 of max 64 tokens, 5.5 tok/s, 1 s +[strata] reading the prompt: 928 of 3,037 tokens, 2 s so far +[strata] done: 10 tokens in 2 s (4.6 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,376 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 1,952 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 2,464 of 3,037 tokens, 5 s so far +[strata] reading the prompt: 2,912 of 3,037 tokens, 6 s so far +[strata] answering: 6 of max 64 tokens, 6.7 tok/s, 7 s +[strata] done: 10 tokens in 8 s (7.1 tok/s) (stop, cancel=False) +.[strata] done: 10 tokens in 1 s (10.0 tok/s) (stop, cancel=False) +[strata] answering: 7 of max 64 tokens, 7.6 tok/s, 1 s +[strata] done: 10 tokens in 2 s (6.4 tok/s) (stop, cancel=False) +[strata] loading the model again (it was unloaded) ... +[strata] the engine is running again +[strata] answering: 8 of max 16 tokens, 8.2 tok/s, 1 s +[strata] done: 10 tokens in 1 s (7.5 tok/s) (stop, cancel=False) +.[strata] answering: 7 of max 200 tokens, 7.2 tok/s, 1 s +[strata] answering: 6 of max 200 tokens, 8.0 tok/s, 1 s +[strata] done: 10 tokens in 2 s (7.8 tok/s) (stop, cancel=False) +[strata] answering: 15 of max 200 tokens, 7.1 tok/s, 2 s +[strata] answering: 24 of max 200 tokens, 7.6 tok/s, 3 s +[strata] answering: 31 of max 200 tokens, 7.3 tok/s, 4 s +[strata] answering: 43 of max 200 tokens, 8.1 tok/s, 5 s +[strata] answering: 51 of max 200 tokens, 8.0 tok/s, 6 s +[strata] done: 55 tokens in 7 s (7.8 tok/s) (stop, cancel=False) +.[strata] answering: 8 of max 200 tokens, 8.1 tok/s, 1 s +[strata] answering: 7 of max 200 tokens, 8.9 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.4 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 1 s (9.0 tok/s) (stop, cancel=False) +.[strata] answering: 5 of max 60 tokens, 4.2 tok/s, 1 s +[strata] reading the prompt: 128 of 3,037 tokens, 1 s so far +[strata] answering: 3 of max 64 tokens, 4.1 tok/s, 1 s +[strata] reading the prompt: 608 of 3,037 tokens, 2 s so far +[strata] answering: 9 of max 60 tokens, 3.9 tok/s, 2 s +[strata] done: 10 tokens in 3 s (3.9 tok/s) (stop, cancel=False) +[strata] answering: 8 of max 64 tokens, 4.1 tok/s, 2 s +[strata] reading the prompt: 1,120 of 3,037 tokens, 3 s so far +[strata] done: 10 tokens in 3 s (4.0 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,600 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 2,112 of 3,037 tokens, 5 s so far +[strata] reading the prompt: 2,624 of 3,037 tokens, 6 s so far +[strata] answering: 4 of max 64 tokens, 8.4 tok/s, 7 s +[strata] done: 10 tokens in 8 s (7.6 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 2,560 of 3,037 tokens, 6 s so far +[strata] answering: 1 of max 64 tokens, 135300.1 tok/s, 5 s +[strata] answering: 3 of max 60 tokens, 0.3 tok/s, 11 s +[strata] answering: 2 of max 64 tokens, 0.4 tok/s, 11 s +[strata] reading the prompt: 2,560 of 3,037 tokens, 11 s so far +[strata] answering: 3 of max 64 tokens, 0.4 tok/s, 12 s +[strata] answering: 2 of max 64 tokens, 6.4 tok/s, 13 s +[strata] answering: 5 of max 60 tokens, 0.4 tok/s, 13 s +[strata] done: 10 tokens in 13 s (1.3 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 13 s (9.6 tok/s) (stop, cancel=False) +[strata] answering: 15 of max 60 tokens, 1.1 tok/s, 14 s +[strata] answering: 24 of max 60 tokens, 1.6 tok/s, 15 s +[strata] answering: 32 of max 60 tokens, 2.0 tok/s, 16 s +[strata] answering: 39 of max 60 tokens, 2.3 tok/s, 17 s +[strata] answering: 50 of max 60 tokens, 2.8 tok/s, 18 s +[strata] done: 55 tokens in 19 s (3.0 tok/s) (stop, cancel=False) +.[strata] parallel: slot(s) 0 were marked busy with no request using them (issue #1603); freed, the waiting request goes on +..[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (1/6) (issue #1994) +[strata] answering: 6 of max 64 tokens, 6.5 tok/s, 1 s +[strata] answering: 2 of max 64 tokens, 6.0 tok/s, 1 s +[strata] answering: 13 of max 64 tokens, 6.4 tok/s, 2 s +[strata] done: 10 tokens in 2 s (7.6 tok/s) (stop, cancel=False) +[strata] answering: 23 of max 64 tokens, 7.2 tok/s, 3 s +[strata] answering: 31 of max 64 tokens, 7.2 tok/s, 4 s +[strata] answering: 40 of max 64 tokens, 7.6 tok/s, 5 s +[strata] answering: 47 of max 64 tokens, 7.5 tok/s, 6 s +[strata] answering: 54 of max 64 tokens, 7.3 tok/s, 8 s +[strata] done: 55 tokens in 8 s (7.3 tok/s) (stop, cancel=False) +[strata] answering: 8 of max 64 tokens, 8.1 tok/s, 1 s +[strata] done: 10 tokens in 2 s (6.8 tok/s) (stop, cancel=False) +.[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (1/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (2/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (3/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (4/6) (issue #1994) +[strata] answering: 8 of max 64 tokens, 6.7 tok/s, 1 s +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (5/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (6/6) (issue #1994) +[strata] parallel: the engine refused slot 1 (BGEN: no such free slot (--batch 2) or a bad max_new); trying another slot (7/6) (issue #1994) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (length, cancel=False) +/opt/homebrew/Cellar/python@3.14/3.14.7/Frameworks/Python.framework/Versions/3.14/lib/python3.14/tempfile.py:484: ResourceWarning: Implicitly cleaning up + _warnings.warn(self.warn_message, ResourceWarning) +[strata] answering: 14 of max 64 tokens, 6.3 tok/s, 2 s +[strata] answering: 23 of max 64 tokens, 6.9 tok/s, 3 s +[strata] answering: 31 of max 64 tokens, 7.0 tok/s, 4 s +[strata] answering: 40 of max 64 tokens, 7.4 tok/s, 5 s +[strata] answering: 48 of max 64 tokens, 7.4 tok/s, 6 s +[strata] done: 55 tokens in 7 s (7.4 tok/s) (stop, cancel=False) +[strata] answering: 8 of max 64 tokens, 8.4 tok/s, 1 s +[strata] done: 10 tokens in 2 s (6.9 tok/s) (stop, cancel=False) +.[strata] answering: 5 of max 64 tokens, 5.4 tok/s, 1 s +[strata] answering: 7 of max 64 tokens, 11.6 tok/s, 1 s +[strata] done: 10 tokens in 1 s (11.4 tok/s) (stop, cancel=False) +[strata] reading the prompt: 32 of 66 tokens, 1 s so far +[strata] answering: 15 of max 64 tokens, 7.3 tok/s, 2 s +[strata] answering: 8 of max 64 tokens, 7.8 tok/s, 2 s +[strata] answering: 22 of max 64 tokens, 7.2 tok/s, 3 s +[strata] done: 10 tokens in 3 s (7.3 tok/s) (stop, cancel=False) +[strata] answering: 29 of max 64 tokens, 7.2 tok/s, 4 s +[strata] answering: 37 of max 64 tokens, 7.3 tok/s, 5 s +[strata] answering: 47 of max 64 tokens, 7.6 tok/s, 6 s +[strata] done: 55 tokens in 7 s (7.8 tok/s) (stop, cancel=False) +.[strata] parallel requests: 4 asked, the engine runs 2 (its log says why) +.[strata] parallel requests: 4 asked, the engine runs one at a time (its log says why) +[strata] answering: 8 of max 64 tokens, 8.4 tok/s, 1 s +[strata] done: 10 tokens in 1 s (8.8 tok/s) (stop, cancel=False) +.[strata] answering: 5 of max 200 tokens, 4.6 tok/s, 1 s +[strata] answering: 3 of max 200 tokens, 6.0 tok/s, 1 s +[strata] done: 10 tokens in 2 s (8.4 tok/s) (stop, cancel=False) +[strata] answering: 15 of max 200 tokens, 7.0 tok/s, 2 s +[strata] answering: 22 of max 200 tokens, 6.9 tok/s, 3 s +[strata] answering: 29 of max 200 tokens, 6.8 tok/s, 4 s +[strata] answering: 37 of max 200 tokens, 6.9 tok/s, 5 s +[strata] answering: 44 of max 200 tokens, 6.8 tok/s, 6 s +[strata] answering: 51 of max 200 tokens, 6.8 tok/s, 8 s +[strata] done: 55 tokens in 8 s (6.8 tok/s) (stop, cancel=False) +.[strata] answering: 6 of max 200 tokens, 5.9 tok/s, 1 s +[strata] answering: 5 of max 64 tokens, 6.8 tok/s, 1 s +[strata] done: 10 tokens in 2 s (7.4 tok/s) (stop, cancel=False) +[strata] answering: 14 of max 200 tokens, 6.8 tok/s, 2 s +[strata] answering: 25 of max 200 tokens, 8.1 tok/s, 3 s +[strata] answering: 33 of max 200 tokens, 8.1 tok/s, 4 s +[strata] answering: 43 of max 200 tokens, 8.3 tok/s, 5 s +[strata] answering: 53 of max 200 tokens, 8.6 tok/s, 6 s +[strata] done: 55 tokens in 6 s (8.8 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 672 of 3,037 tokens, 1 s so far +[strata] reading the prompt: 1,312 of 3,037 tokens, 2 s so far +[strata] reading the prompt: 1,984 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 2,720 of 3,037 tokens, 4 s so far +[strata] answering: 1 of max 64 tokens, 1.8 tok/s, 5 s +[strata] answering: 1 of max 64 tokens, 0.6 tok/s, 6 s +[strata] answering: 1 of max 64 tokens, 0.4 tok/s, 7 s +[strata] answering: 1 of max 64 tokens, 0.3 tok/s, 8 s +[strata] answering: 1 of max 64 tokens, 0.2 tok/s, 9 s +[strata] reading the prompt: 32 of 3,037 tokens, 10 s so far +[strata] answering: 3 of max 64 tokens, 0.5 tok/s, 10 s +[strata] reading the prompt: 704 of 3,037 tokens, 11 s so far +[strata] answering: 9 of max 64 tokens, 1.3 tok/s, 11 s +[strata] done: 10 tokens in 12 s (1.4 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,440 of 3,037 tokens, 12 s so far +[strata] reading the prompt: 2,016 of 3,037 tokens, 13 s so far +[strata] reading the prompt: 2,624 of 3,037 tokens, 14 s so far +[strata] answering: 4 of max 64 tokens, 11.0 tok/s, 15 s +[strata] done: 10 tokens in 16 s (9.9 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 2,560 of 3,037 tokens, 4 s so far +[strata] answering: 1 of max 64 tokens, 0.2 tok/s, 9 s +[strata] answering: 3 of max 64 tokens, 0.3 tok/s, 14 s +[strata] reading the prompt: 2,560 of 3,037 tokens, 14 s so far +[strata] answering: 4 of max 64 tokens, 9.6 tok/s, 15 s +[strata] answering: 6 of max 64 tokens, 0.6 tok/s, 15 s +[strata] done: 10 tokens in 15 s (0.9 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 15 s (9.4 tok/s) (stop, cancel=False) +.[strata] the engine's logits became non-finite before the first token (#879); it dropped its cached state: sending the request again once +[strata] done: 10 tokens in 1 s (9.5 tok/s) (stop, cancel=False) +[strata] the engine's logits became non-finite before the first token (#879); it dropped its cached state: sending the request again once +[strata] answering: 6 of max 64 tokens, 6.1 tok/s, 1 s +[strata] answering: 5 of max 64 tokens, 9.3 tok/s, 1 s +[strata] done: 10 tokens in 2 s (9.2 tok/s) (stop, cancel=False) +[strata] answering: 16 of max 64 tokens, 7.8 tok/s, 2 s +[strata] answering: 25 of max 64 tokens, 8.0 tok/s, 3 s +[strata] answering: 33 of max 64 tokens, 8.0 tok/s, 4 s +[strata] answering: 42 of max 64 tokens, 8.0 tok/s, 5 s +[strata] answering: 52 of max 64 tokens, 8.2 tok/s, 6 s +[strata] done: 55 tokens in 7 s (8.0 tok/s) (stop, cancel=False) +[strata] the engine's logits became non-finite after 2 tokens (#879): the reply ends here as "length"; the engine dropped its cached state, so the next request reads its prompt from the start +[strata] done: 2 tokens in 1 s (7.4 tok/s) (length, cancel=False) +[strata] answering: 7 of max 64 tokens, 6.6 tok/s, 1 s +[strata] answering: 18 of max 64 tokens, 8.2 tok/s, 2 s +[strata] answering: 26 of max 64 tokens, 7.9 tok/s, 3 s +[strata] answering: 33 of max 64 tokens, 7.6 tok/s, 4 s +[strata] answering: 41 of max 64 tokens, 7.7 tok/s, 5 s +[strata] answering: 49 of max 64 tokens, 7.6 tok/s, 7 s +[strata] done: 55 tokens in 7 s (7.4 tok/s) (stop, cancel=False) +.[strata] done: 9 tokens in 1 s (8.6 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 1 s (10.5 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 2,560 of 3,037 tokens, 5 s so far +[strata] answering: 1 of max 64 tokens, 174762.7 tok/s, 5 s +[strata] reading the prompt: 2,560 of 3,037 tokens, 10 s so far +[strata] answering: 2 of max 64 tokens, 0.4 tok/s, 10 s +[strata] answering: 3 of max 64 tokens, 0.5 tok/s, 11 s +[strata] answering: 2 of max 64 tokens, 6.9 tok/s, 11 s +[strata] done: 10 tokens in 12 s (1.5 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 12 s (9.4 tok/s) (stop, cancel=False) +.[strata] reading the prompt: 32 of 3,037 tokens, 1 s so far +[strata] reading the prompt: 352 of 3,037 tokens, 2 s so far +[strata] reading the prompt: 864 of 3,037 tokens, 3 s so far +[strata] reading the prompt: 1,408 of 3,037 tokens, 4 s so far +[strata] reading the prompt: 1,952 of 3,037 tokens, 5 s so far +[strata] reading the prompt: 2,464 of 3,037 tokens, 6 s so far +[strata] reading the prompt: 32 of 3,037 tokens, 7 s so far +[strata] answering: 2 of max 64 tokens, 5.5 tok/s, 8 s +[strata] reading the prompt: 608 of 3,037 tokens, 8 s so far +[strata] answering: 7 of max 64 tokens, 4.1 tok/s, 9 s +[strata] reading the prompt: 1,088 of 3,037 tokens, 9 s so far +[strata] done: 10 tokens in 10 s (4.1 tok/s) (stop, cancel=False) +[strata] reading the prompt: 1,600 of 3,037 tokens, 10 s so far +[strata] reading the prompt: 1,952 of 3,037 tokens, 11 s so far +[strata] reading the prompt: 2,112 of 3,037 tokens, 11 s so far +[strata] reading the prompt: 2,688 of 3,037 tokens, 12 s so far +[strata] reading the prompt: 32 of 43 tokens, 13 s so far +[strata] answering: 2 of max 64 tokens, 3.6 tok/s, 13 s +[strata] answering: 7 of max 64 tokens, 6.9 tok/s, 14 s +[strata] done: 10 tokens in 14 s (7.3 tok/s) (stop, cancel=False) +[strata] done: 10 tokens in 14 s (8.7 tok/s) (stop, cancel=False) +...............[strata] reading the prompt: 8 of max 64 tokens, 8.1 tok/s, 1 s +[strata] reading the prompt: 7 of max 64 tokens, 7.5 tok/s, 1 s +[strata] reading the prompt: 10 of max 64 tokens, 9.0 tok/s, 1 s +[strata] answering: 17 of max 64 tokens, 8.4 tok/s, 2 s +[strata] answering: 15 of max 64 tokens, 7.4 tok/s, 2 s +[strata] answering: 19 of max 64 tokens, 8.6 tok/s, 2 s +[strata] answering: 27 of max 64 tokens, 8.7 tok/s, 3 s +[strata] answering: 27 of max 64 tokens, 8.5 tok/s, 3 s +[strata] answering: 28 of max 64 tokens, 8.5 tok/s, 3 s +[strata] answering: 36 of max 64 tokens, 8.6 tok/s, 4 s +[strata] answering: 35 of max 64 tokens, 8.1 tok/s, 4 s +[strata] answering: 36 of max 64 tokens, 8.2 tok/s, 4 s +[strata] answering: 44 of max 64 tokens, 8.5 tok/s, 5 s +[strata] answering: 45 of max 64 tokens, 8.4 tok/s, 5 s +[strata] answering: 44 of max 64 tokens, 8.0 tok/s, 6 s +[strata] done: 51 tokens in 6 s (8.3 tok/s) (stop, cancel=False) +[strata] done: 51 tokens in 6 s (8.3 tok/s) (stop, cancel=False) +[strata] done: 51 tokens in 6 s (7.9 tok/s) (stop, cancel=False) +.[strata] done: 22 tokens in 0 s (292935.5 tok/s) (stop, cancel=False) +............[strata] the config has a "vision" section but no --vision in its args: added --vision (run setup again with --vision to write them) +....[strata] thinking limited to 24386 of max_tokens 32514 so the answer has room ("answer_room": false in the config turns this off) +[strata] done: 3 tokens in 0 s (9970.6 tok/s) (stop, cancel=False) +...[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (16912.5 tok/s) (length, cancel=False) +.[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (5683.3 tok/s) (length, cancel=False) +.[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (11781.8 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (5849.8 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (7194.3 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (14266.3 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (5957.8 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 1 s (5289.2 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] answering: 1 of max 1 tokens, 167772.2 tok/s, 1 s +[strata] done: 1 tokens in 1 s (4810.0 tok/s) (length, cancel=False) +[strata] done: 0 tokens in 1 s (0.0 tok/s) (cancel, cancel=True) +[strata] done: 1 tokens in 0 s (41527.8 tok/s) (length, cancel=False) + +[deadline exceeded after 300s; process group 91074 terminated after SIGTERM] diff --git a/results/round39/serve-python-packages.txt b/results/round39/serve-python-packages.txt new file mode 100644 index 0000000..af24459 --- /dev/null +++ b/results/round39/serve-python-packages.txt @@ -0,0 +1,13 @@ +Using Python 3.14.7 environment at: /private/tmp/r39-serve-venv +certifi==2026.7.22 +charset-normalizer==3.5.2 +idna==3.20 +jinja2==3.1.6 +markupsafe==3.0.4 +numpy==2.5.4 +pillow==13.0.0 +psutil==7.2.2 +pyyaml==6.0.3 +regex==2026.9.29 +requests==2.34.2 +urllib3==2.8.0 diff --git a/results/round39/serve-source-hashes.json b/results/round39/serve-source-hashes.json new file mode 100644 index 0000000..748909d --- /dev/null +++ b/results/round39/serve-source-hashes.json @@ -0,0 +1,134 @@ +[ + { + "path": "serve/frontend.py", + "sha256": "1f40ac6fc9f4ad021ac0c202f35f184f5784009355dd7485ee0e44746bfd75e1", + "upstream_sha256": "1f40ac6fc9f4ad021ac0c202f35f184f5784009355dd7485ee0e44746bfd75e1", + "matches_upstream": true + }, + { + "path": "serve/logit_bias.py", + "sha256": "b6564c97e91c987e5787349b3df209fa0cf00fc8e80357ad93d40fd1e594da09", + "upstream_sha256": "1cff46f6ac1c58790adb76572670aecc35e4846eda561e995a6561aed0bc1292", + "matches_upstream": false + }, + { + "path": "serve/runconfig.py", + "sha256": "58aa118367a9d6dbb534cae5f35d6f418112f96985db85b048fa23eb042e403a", + "upstream_sha256": "58aa118367a9d6dbb534cae5f35d6f418112f96985db85b048fa23eb042e403a", + "matches_upstream": true + }, + { + "path": "serve/server.py", + "sha256": "7ee5536c9293583244abb3774851fb16aa2ac7e7e55e200fef13bb4be4aac19b", + "upstream_sha256": "9eb13d12dc3ed0f6c94d20ac1c0018efe1e00c331deae76bd0fb9f420e005232", + "matches_upstream": false + }, + { + "path": "serve/telemetry.py", + "sha256": "973772d0209167bf0dbaa8fd98fab8e64801f12dd241c72d5e06bbaaf96b646e", + "upstream_sha256": "d26913657b57442e2450501a4f7f21c68ee90fa9771fae0694f21c588711738a", + "matches_upstream": false + }, + { + "path": "serve/test_admit_order.py", + "sha256": "bbfe80b214f0dfcf2339bb4aab08e1549330951bcf5b7a607acea4f436cae47b", + "upstream_sha256": "bbfe80b214f0dfcf2339bb4aab08e1549330951bcf5b7a607acea4f436cae47b", + "matches_upstream": true + }, + { + "path": "serve/test_effort_paths.py", + "sha256": "9501d8787e396b95f5143eb298334973a8553b298feff0e93b8708dfcc4f6db9", + "upstream_sha256": "9501d8787e396b95f5143eb298334973a8553b298feff0e93b8708dfcc4f6db9", + "matches_upstream": true + }, + { + "path": "serve/test_logit_bias.py", + "sha256": "d4afa9e4f50bfa3595db67bc9e2487adabcaf7e902be2c3dca3bbccec434e92a", + "upstream_sha256": "f605101160feb919c3469002988de9c0f568c73ad74de1d538b81d20dae59b76", + "matches_upstream": false + }, + { + "path": "serve/test_model_card_sampling.py", + "sha256": "38a90fe3d2339d1bca61b073a5e0bafae936c36ab1b1953ffdc6f756a06bf5f3", + "upstream_sha256": "38a90fe3d2339d1bca61b073a5e0bafae936c36ab1b1953ffdc6f756a06bf5f3", + "matches_upstream": true + }, + { + "path": "serve/test_parallel.py", + "sha256": "0f6147f1e74badb61f974d0fd71d2b39e4c141eaeee6dc3cfbf5f5bbb585eeec", + "upstream_sha256": "0f6147f1e74badb61f974d0fd71d2b39e4c141eaeee6dc3cfbf5f5bbb585eeec", + "matches_upstream": true + }, + { + "path": "serve/test_parser_call_scan.py", + "sha256": "71aab3ee7ea39bf8810099e3d59d1c8580836923b7aa3446a72c281b6398232d", + "upstream_sha256": "71aab3ee7ea39bf8810099e3d59d1c8580836923b7aa3446a72c281b6398232d", + "matches_upstream": true + }, + { + "path": "serve/test_parser_hold.py", + "sha256": "d50e799c0bbf003318b51938f09e0b40fb11b8d572a7377ddd4617b9ddf9d7cf", + "upstream_sha256": "d50e799c0bbf003318b51938f09e0b40fb11b8d572a7377ddd4617b9ddf9d7cf", + "matches_upstream": true + }, + { + "path": "serve/test_prompt_encoder.py", + "sha256": "fad128ea329949fd38ce034a026b3f86d060c484fcfe1c0fea2a13f94524a73c", + "upstream_sha256": "fad128ea329949fd38ce034a026b3f86d060c484fcfe1c0fea2a13f94524a73c", + "matches_upstream": true + }, + { + "path": "serve/test_q36_config.py", + "sha256": "278dae2f526ad8cda6fe8a719e075456c71824bbd8b2362b87fe62139b6b0ea8", + "upstream_sha256": "278dae2f526ad8cda6fe8a719e075456c71824bbd8b2362b87fe62139b6b0ea8", + "matches_upstream": true + }, + { + "path": "serve/test_reasoning_loop_recovery.py", + "sha256": "c448df0757ae1bffb0a4fe8b48f45ee57ae0f06e7056d6bc96885dc8b024bf7b", + "upstream_sha256": "c448df0757ae1bffb0a4fe8b48f45ee57ae0f06e7056d6bc96885dc8b024bf7b", + "matches_upstream": true + }, + { + "path": "serve/test_server.py", + "sha256": "576f7e5eb3ea07cb8dc0939714d282b59e9f51d0d4c6e00549394f4cd9753d17", + "upstream_sha256": "64ae891b7ef81a2819a2a1d3076aa035208615318aeb8627a8d2455205f25042", + "matches_upstream": false + }, + { + "path": "serve/test_slot_alloc.py", + "sha256": "7ea0cb07ed01f95d720fb0ae1d945cee705a31a34213bd5ada22dfa35c68d40c", + "upstream_sha256": "7ea0cb07ed01f95d720fb0ae1d945cee705a31a34213bd5ada22dfa35c68d40c", + "matches_upstream": true + }, + { + "path": "serve/test_web_sampling.py", + "sha256": "8e4e717aaab9049b99a1a44beff2bae1d48e6f60d08de98a28e2d875d2eba18d", + "upstream_sha256": "8e4e717aaab9049b99a1a44beff2bae1d48e6f60d08de98a28e2d875d2eba18d", + "matches_upstream": true + }, + { + "path": "serve/web/app.js", + "sha256": "7d15c4efeccbc4b1f4944494dc76987e18e7a6cf40d12279978dfe93951bb8fd", + "upstream_sha256": "7d15c4efeccbc4b1f4944494dc76987e18e7a6cf40d12279978dfe93951bb8fd", + "matches_upstream": true + }, + { + "path": "tools/strata_tokenizer.py", + "sha256": "e0cd8d6ae38dda0a4ad7c14f812eb64d5fd6f4fed9c920d4427a14bf585dc9b3", + "upstream_sha256": "e0cd8d6ae38dda0a4ad7c14f812eb64d5fd6f4fed9c920d4427a14bf585dc9b3", + "matches_upstream": true + }, + { + "path": "tools/test_strata_tokenizer.py", + "sha256": "6a3cadf78c1972e83acd0ef90a8f5119868c7c3a4613fcfa000faae1195017b5", + "upstream_sha256": "6a3cadf78c1972e83acd0ef90a8f5119868c7c3a4613fcfa000faae1195017b5", + "matches_upstream": true + }, + { + "path": "serve/test_image_omission.py", + "sha256": "88c88b3a2f15d41958a2b60760734897c4656d275353b906b6c49d54591d4b39", + "upstream_sha256": null, + "matches_upstream": false + } +] diff --git a/results/round39/serve-upstream-review.diff b/results/round39/serve-upstream-review.diff new file mode 100644 index 0000000..e5de5fb --- /dev/null +++ b/results/round39/serve-upstream-review.diff @@ -0,0 +1,875 @@ +diff --git a/serve/server.py b/serve/server.py +index 5ca7f4e4..2e805773 100644 +--- a/serve/server.py ++++ b/serve/server.py +@@ -147,6 +147,7 @@ VISION_START = "<|vision_start|>" + # broken state that answers one token forever (an issue saw 36,689 tokens of "!"). The config's "repeat_stop_tokens" + # sets it; 0 turns it off. + REPEAT_STOP_TOKENS = 256 ++ANSWER_ROOM_MIN_CAP = 1024 # 0.1.45: from this max_tokens on, thinking without a budget is closed early enough to leave room to answer + ANSWER_RESERVE_MIN = 512 # #984: the tokens kept for the answer when a thinking budget is cut down to fit max_tokens + # #123: what closes the thinking when it reaches reasoning_budget_tokens (the model's own end-of-thinking tag after it) + # #1053 (opt-in "reasoning_close_retry": true): a reply that ends on its stop token still inside , with no answer +@@ -254,6 +255,15 @@ def _timeout_env(name: str, default: float) -> float | None: + # still working is never ended by this, only reported once (the recommendation: engine_silence_s, or this variable). + # STRATA_ENGINE_STALL_S overrides it (0 = off). Needs psutil (setup installs it); without it nothing is ended. + ENGINE_STALL_S = _timeout_env("STRATA_ENGINE_STALL_S", 90.0) ++# --batch: new requests take the control lines (one prompt read at a time) short prompts first: a prompt of at most ++# STRATA_ADMIT_SHORT_FIRST tokens goes before longer ones (opt-in; unset or 0: whoever gets the lock first), and a ++# longer one that waited STRATA_ADMIT_MAX_WAIT_S (default 30) counts as short. A chat question then waits for the one ++# read in progress, not for every agent read that came before it. ++try: ++ ADMIT_SHORT_FIRST = max(0, int(os.environ.get("STRATA_ADMIT_SHORT_FIRST") or 0)) ++except ValueError: ++ ADMIT_SHORT_FIRST = 0 ++ADMIT_MAX_WAIT_S = _timeout_env("STRATA_ADMIT_MAX_WAIT_S", 30.0) or 0.0 + STALL_CPU_EPS_S = 0.5 # CPU seconds the engine may use in the whole stall window and still count as frozen + STALL_IO_EPS_B = 1 << 20 # bytes read or written in that window + +@@ -705,6 +715,19 @@ class ConvCacheLog: + return dict(self.state) + + ++def conversation_slots(engine) -> int | None: ++ """How many conversations the engine's cache holds: its --conversation-cache-slots argument, which is known ++ before the engine has started (a lazy start has reported no INFO line yet), else the INFO line it reported.""" ++ args = engine.spawn[1] if hasattr(engine, "spawn") else [] ++ if "--conversation-cache-slots" in args[:-1]: ++ try: ++ return int(args[args.index("--conversation-cache-slots") + 1]) ++ except ValueError: ++ pass ++ slots = (getattr(engine, "info", None) or {}).get("conversation_cache_slots") ++ return slots if isinstance(slots, int) else None ++ ++ + def conversation_cache_view(info: dict, hist: list, totals: dict, parked: dict) -> dict: + """#596: the Monitor's Conversation cache card: the parked conversations (the engine's opt-in + --conversation-cache-mib: budget, slots, what its log says) and how much of the prompts the cache gave back.""" +@@ -1096,6 +1119,13 @@ class StrataEngine: + self.last.update(prompt_read=int(f[14])) + if len(f) >= 16: # #588 (engine 0.1.39+): routed experts read over PCIe + self.last.update(offloaded=int(f[15])) ++ for tok in f[16:]: # key=value fields after the positional ones (the prompt's ++ key, eq, value = tok.partition("=") # first chunk and its streamed expert blobs) ++ if eq and key in ("chunk", "experts_streamed"): ++ try: ++ self.last[key] = int(value) ++ except ValueError: ++ pass + + def vram(self, reserve_mib: int | None, timeout: float = 120.0) -> dict: + """#533: `VRAM ` between requests (the caller holds the service's FIFO): the engine shrinks its +@@ -1173,6 +1203,10 @@ class StrataEngine: + v = tune.get(k) + if isinstance(v, (int, float)) and not isinstance(v, bool) and 0.0 <= float(v) <= 1.0: + keys += f" {k}={float(v)!r}" ++ # a layer split's prompt chunks (--prefill-pipe): 0 off, 1 evened out, > 1 the rig's b/a in tokens ++ v = tune.get("prefill_pipe") ++ if isinstance(v, (int, float)) and not isinstance(v, bool) and 0.0 <= float(v) <= 262144.0: ++ keys += f" pipe={float(v)!r}" + # --aux-cpus for this request: 1 puts the engine's threads that are neither pool workers nor the host on the + # spare CPUs, 0 leaves them where they were created (true/false or 1/0; the engine started without a spare CPU + # ignores a 1). Anything else is not sent. +@@ -1416,10 +1450,11 @@ class StrataEngine: + + YIELDS_MAX = 2 # #656: how often one request's prompt read gives way to a shorter waiting one + +- def _take_control(self, cancel, plen: int, after_epoch: int | None = None): ++ def _take_control(self, cancel, plen: int, after_epoch: int | None = None, new: bool = False): + """Waits for the control lines (one prompt read at a time), yielding None heartbeats; False when cancelled. +- `after_epoch`: a request whose read gave way lets the requests waiting then go first.""" +- entry = [plen] ++ `after_epoch`: a request whose read gave way lets the requests waiting then go first. ++ `new`: the request's first admission (STRATA_ADMIT_SHORT_FIRST orders only these).""" ++ entry = [plen, time.monotonic(), new] # (length, arrival, ordered: STRATA_ADMIT_SHORT_FIRST) + with self.slot_cv: + self.waiting += 1 + self.wait_lens.append(entry) +@@ -1433,7 +1468,17 @@ class StrataEngine: + raise EngineDied("the engine stopped while this request waited; it was not sent") + with self.slot_cv: + turn = after_epoch is None or self.ctl_epoch > after_epoch or self.waiting <= 1 ++ if turn and new and ADMIT_SHORT_FIRST > 0: ++ turn = self._admit_next() is entry + if turn and self.ctl.acquire(timeout=0.5): ++ if new and ADMIT_SHORT_FIRST > 0: ++ with self.slot_cv: # a shorter one came while this one waited in acquire() ++ first = self._admit_next() is entry ++ if not first: ++ self.ctl.release() ++ with self.slot_cv: ++ self.slot_cv.notify_all() ++ continue + if self.alive(): + break + self.ctl.release() # it died (or is being restarted) while this one was getting +@@ -1456,6 +1501,18 @@ class StrataEngine: + self.slot_cv.notify_all() + return True + ++ def _admit_next(self): ++ """STRATA_ADMIT_SHORT_FIRST= (the caller holds slot_cv): the new request that takes the control lines ++ next - one whose prompt is at most that long (it reads little) before a longer one, in arrival order within ++ each; a longer one that has waited STRATA_ADMIT_MAX_WAIT_S counts as short (it is never held back for good). ++ Off (unset or 0): whoever gets the lock first, as before. Only the order changes, never what is read.""" ++ now = time.monotonic() ++ ordered = [e for e in self.wait_lens if len(e) > 2 and e[2]] ++ if not ordered: ++ return None ++ return min(ordered, key=lambda e: (0 if e[0] <= ADMIT_SHORT_FIRST or now - e[1] >= ADMIT_MAX_WAIT_S else 1, ++ e[1])) ++ + SOLO_AGAIN_MAX = 2 # how often a request left alone in a slot goes back to the solo path + SOLO_AGAIN_MIN_LEFT = 32 # ... only with at least this many tokens still allowed (max_tokens) + +@@ -1493,7 +1550,7 @@ class StrataEngine: + prompt, left = list(ids), int(max_new) + reused0 = None # the first read's reused tokens: a continued leg's count includes this request's own output + carry = {} # #1869: the figures of the legs that ended in a hand-back to the solo path (the final DONE adds to them) +- ok = yield from self._take_control(cancel, len(prompt)) ++ ok = yield from self._take_control(cancel, len(prompt), new=True) + if not ok: + return + holding, born = True, self.gen # (the engine it now has the control lines of) +@@ -1803,6 +1860,7 @@ class StrataEngine: + + slot_groups = 1 # --batch-groups (set when the engine starts) + SLOT_PREFIX_MIN = 512 # a held prefix shorter than this is not worth unbalancing the pipeline's groups ++ TURN_ID = 248045 # <|im_start|>: the engine's default --turn-token (its turn checkpoint sits at the last one) + + def pick_slot(self, prompt: list[int]) -> int | None: + """A free slot for `prompt` (the caller holds slot_cv): the one whose held tokens are the longest start of the +@@ -1812,18 +1870,30 @@ class StrataEngine: + if not free: + return None + gs = getattr(self, "slot_gs", 0) +- if gs: +- # pipelined groups: the engine keeps no pipelined slot as a conversation cache, so nothing is worth keeping +- # free for a next turn. A group's window spans its slots up to its last busy one, so first fill a hole +- # below a busy slot of the same group (it costs no extra row), then follow the spreading order ++ if gs > 1: ++ # pipelined groups of several slots: the engine keeps no idle slot below a busy one as a conversation cache ++ # (its pad row writes the slot), so nothing is worth keeping free for a next turn. A group's window spans ++ # its slots up to its last busy one, so first fill a hole below a busy slot of the same group (it costs no ++ # extra row), then follow the spreading order + # (reported on #793: least-recently-used picks left leading holes, [_ B], that cost a whole row each) + def fills_hole(b): + g0 = b - b % gs + return any(self.slot_busy[x] for x in range(b + 1, g0 + gs)) + return min(free, key=lambda b: (not fills_hole(b), self.slot_order.index(b))) + def held_prefix(b): ++ # what of `prompt` the slot's cache serves: the slot's whole held stream (the prompt and what it decoded) ++ # when this prompt continues it, else its turn checkpoint - the engine takes one at the last <|im_start|> of ++ # a prompt (the next request of the chat renders the history again, but not the reply it decoded or the ++ # assistant header), so a prompt that agrees with the held tokens up to there is that chat's next turn ++ # too (#2518: it was only found by the engine's own source search, after pick_slot had sent it to another ++ # slot and the request in that slot had overwritten the chat whose cache the next one needed) + h = self.slot_held[b] +- return len(h) if h and len(h) < len(prompt) and prompt[:len(h)] == h else 0 ++ if not h: ++ return 0 ++ if len(h) < len(prompt) and prompt[:len(h)] == h: ++ return len(h) ++ k = len(h) - 1 - h[::-1].index(self.TURN_ID) if self.TURN_ID in h else 0 ++ return k if 0 < k < len(prompt) and prompt[:k] == h[:k] else 0 + best = max(free, key=held_prefix) + # without --batch-groups a held prefix of any length is taken, as before + if held_prefix(best) >= (self.SLOT_PREFIX_MIN if self.slot_groups > 1 else 1): +@@ -2868,23 +2938,32 @@ class ByteTokenizer: + + ALWAYS = () # specials matched without parse_special (type 4, as ) + ++ @property ++ def max_special_len(self): ++ return max(len(s) for s in self.SPECIALS) ++ + @property + def control_tokens(self): + return [s for s in self.SPECIALS if s not in self.ALWAYS] + + def encode(self, text, parse_special=False, plain=()): +- out, i = [], 0 ++ return self.encode_marked(text, parse_special, plain)[0] ++ ++ def encode_marked(self, text, parse_special=False, plain=()): ++ """encode() and its resume points (after each special), as strata_tokenizer's for PromptEncoder.""" ++ out, marks, i = [], [], 0 + while i < len(text): + for k, s in enumerate(self.SPECIALS): + if (parse_special or s in self.ALWAYS) and text.startswith(s, i) and not any( + a <= i < b for a, b in plain): + out.append(256 + k) + i += len(s) ++ marks.append((i, len(out))) + break + else: + out.extend(text[i].encode("utf-8")) + i += 1 +- return out ++ return out, marks + + def decode(self, ids, errors="replace"): + raw = bytearray() +@@ -3031,10 +3110,12 @@ class Service: + self.literals = literal_tags(getattr(tokenizer, "control_tokens", ())) # texts that stay text inside a message + self.fit_max_tokens = fit_max_tokens # --fit-max-tokens: clamp the output cap instead of 400 + self.aliases: list[str] = [] # #297: other names of the model (the config's `aliases`) ++ self.answer_room = True # 0.1.45: see ANSWER_ROOM_MIN_CAP (the config's "answer_room": false turns it off) + self.reasoning_close_retry = False # #1053 (opt-in): close the thinking once when a reply ends inside it + self.codex_compaction_cache = False # #924 (opt-in): a Codex compaction is rendered with its conversation's tools + self.codex_thread_titles = False # #923 (opt-in): answer Codex's thread-title turns without the engine + self.sampling_defaults = dict(sampling_defaults or {}) # the run config's `sampling` block ++ self.model_arch = None # the config's family ("qwen4exp", "qwen35moe"): None = no model-card defaults + self.shared = {} # the web app's Chat settings for every client (POST /settings) + self.shared_path = None # where they are kept between starts (next to the config) + self.config_effort = None # #1641: the config's "reasoning_effort": the level for requests with none +@@ -3075,6 +3156,9 @@ class Service: + "prompt_ms": 0.0, "decode_ms": 0.0, + "drafts_offered": 0, "drafts_accepted": 0} # #457: the MTP drafts, summed where reported + self.last_timings = None # the last finished request's, llama.cpp's names (/v1/status) ++ # /metrics/prometheus' two throughput gauges need the totals as they stood at the last scrape, because ++ # llama.cpp averages them over the window between two scrapes (#2341). ++ self.scrape_bucket = (0, 0.0, 0, 0.0) + from serve.prometheus import Latencies + self.latencies = Latencies() # GET /metrics in Prometheus text: the latency histograms + self.last_request_at = None # when a request last started or finished +@@ -3104,6 +3188,17 @@ class Service: + self.restore_after_load = None # a session file an image job saved before it unloaded + self.stop_ids = set(tokenizer.encode(IM_END, parse_special=True) + + tokenizer.encode("<|endoftext|>", parse_special=True)) ++ # #567: a prompt re-encodes only what follows the last special token it shares with a recent prompt (the ++ # same ids as a full encode: tools/strata_tokenizer.py PromptEncoder). Tokenizers without resume points ++ # encode in full. ++ self.prompts = None ++ if hasattr(tokenizer, "encode_marked") and hasattr(tokenizer, "max_special_len"): ++ from strata_tokenizer import PromptEncoder ++ # one kept prompt per conversation the engine's cache can hold (--conversation-cache-slots); the ++ # encoder's own default without that setting ++ slots = conversation_slots(engine) ++ self.prompts = PromptEncoder(tokenizer, keep=slots) if slots is not None and slots > 0 \ ++ else PromptEncoder(tokenizer) + + def loaded(self) -> bool: + return not hasattr(self.engine, "alive") or self.engine.alive() +@@ -3296,20 +3391,17 @@ class Service: + threading.Thread(target=loop, daemon=True).start() + + def free_vram_mib(self) -> int | None: +- """Free VRAM on the engine's (first) GPU, from NVML (AMD backend: amdgpu's sysfs files, #301); None when it can't +- be read (then nothing is refused).""" ++ """Free VRAM on the engine's (first) GPU, from the backend's reader - NVML, the amdgpu sysfs files (#301) or ++ IOKit on macOS (Metal backend); None when it can't be read (then nothing is refused).""" + try: +- if getattr(self, "backend", None) == "hip": +- from serve.telemetry import free_vram_mib +- return free_vram_mib(int(getattr(self, "gpu_index", 0) or 0), amd=True) +- from serve.telemetry import _Nvml +- nv = _Nvml(int(getattr(self, "gpu_index", 0) or 0)) +- if not nv.ok(): ++ from serve.telemetry import gpu_reader ++ g = gpu_reader(int(getattr(self, "gpu_index", 0) or 0), amd=getattr(self, "backend", None) == "hip") ++ if not g.ok(): + return None +- m = nv.Mem() +- if nv.lib.nvmlDeviceGetMemoryInfo(nv.dev, ctypes.byref(m)) != 0: ++ r = g.read() ++ if r.get("mem_total") is None or r.get("mem_used") is None: + return None +- return int(m.free >> 20) ++ return int((r["mem_total"] - r["mem_used"]) >> 20) + except Exception: + return None + +@@ -3552,7 +3644,7 @@ class Service: + if effort: + if api == "openai": + ctk = req.get("chat_template_kwargs") if isinstance(req.get("chat_template_kwargs"), dict) else {} +- if not req.get("reasoning_effort") and not req.get("reasoning") and \ ++ if req.get("reasoning_effort") in (None, "") and not req.get("reasoning") and \ + "enable_thinking" not in ctk and "reasoning_effort" not in ctk: + req["reasoning_effort"] = effort + elif not req.get("thinking") and not req.get("output_config"): +@@ -3703,8 +3795,15 @@ class Service: + with self.status_lock: + s, t = dict(self.status), dict(self.totals) + running = len(self.live_reqs) or int(bool(s.get("busy"))) +- prompt_n = t["prompt_tokens"] - t["reused"] # what was read, as llama.cpp counts it +- prompt_s, decode_s = t["prompt_ms"] / 1000, t["decode_ms"] / 1000 ++ prompt_n = t["prompt_tokens"] - t["reused"] # what was read, as llama.cpp counts it ++ prompt_s, decode_s = t["prompt_ms"] / 1000, t["decode_ms"] / 1000 ++ # llama.cpp reads these two gauges from a bucket it resets on every scrape (server_metrics::bucket, ++ # "the gauges are averaged over the window between two scrapes"), so an idle server reports 0 rather ++ # than its lifetime average. Keep the same window here, or the names lie to llama.cpp dashboards. ++ was_prompt_n, was_prompt_s, was_predicted_n, was_decode_s = self.scrape_bucket ++ self.scrape_bucket = (prompt_n, prompt_s, t["output_tokens"], decode_s) ++ d_prompt_n, d_prompt_s = prompt_n - was_prompt_n, prompt_s - was_prompt_s ++ d_predicted_n, d_decode_s = t["output_tokens"] - was_predicted_n, decode_s - was_decode_s + rows = [ + ("llamacpp:prompt_tokens_total", "counter", "Prompt tokens read", prompt_n), + ("llamacpp:prompt_tokens_cached_total", "counter", "Prompt tokens the conversation cache already held", +@@ -3712,10 +3811,10 @@ class Service: + ("llamacpp:prompt_seconds_total", "counter", "Prompt read time", prompt_s), + ("llamacpp:tokens_predicted_total", "counter", "Tokens generated", t["output_tokens"]), + ("llamacpp:tokens_predicted_seconds_total", "counter", "Generation time", decode_s), +- ("llamacpp:prompt_tokens_seconds", "gauge", "Average prompt read throughput in tokens/s", +- prompt_n / prompt_s if prompt_s else 0), +- ("llamacpp:predicted_tokens_seconds", "gauge", "Average generation throughput in tokens/s", +- t["output_tokens"] / decode_s if decode_s else 0), ++ ("llamacpp:prompt_tokens_seconds", "gauge", "Average prompt read throughput in tokens/s, since the last scrape", ++ d_prompt_n / d_prompt_s if d_prompt_s > 0 else 0), ++ ("llamacpp:predicted_tokens_seconds", "gauge", "Average generation throughput in tokens/s, since the last scrape", ++ d_predicted_n / d_decode_s if d_decode_s > 0 else 0), + ("llamacpp:requests_processing", "gauge", "Requests running", running), + ("llamacpp:requests_deferred", "gauge", "Requests waiting their turn", s.get("queued") or 0), + ("strata:requests_total", "counter", "Requests finished", t["requests"]), +@@ -3787,17 +3886,21 @@ class Service: + tail = "" if off else EFFORT_TURN.format(EFFORT_TEXT.get(effort, EFFORT_TEXT["medium"])) + return history + tail + full[len(mine):] + +- def encode_prompt(self, messages, tools, kwargs) -> list[int]: ++ def encode_prompt(self, messages, tools, kwargs, store: bool = True) -> list[int]: + """The request's prompt: the template rendered and tokenized. #537: a / written inside a + message's text is encoded as the text it is, not as the model's reasoning markers (the template's own are). + A control token's text (<|im_start|>, <|im_end|>, <|endoftext|>, ...) inside a message is text as well: only +- the control tokens the template writes are control tokens.""" ++ the control tokens the template writes are control tokens. `store=False` (a token count) reuses the kept ++ prompts without adding or replacing one.""" + marked, marked_tools, changed = mark_think_literals(messages, tools, self.literals) + prompt = self.render_prompt(marked, marked_tools, kwargs) +- if not changed: +- return self.tok.encode(prompt, parse_special=True) +- prompt, plain = unmark_think_literals(prompt, self.literals) +- return self.tok.encode(prompt, parse_special=True, plain=plain) ++ plain = () ++ if changed: ++ prompt, plain = unmark_think_literals(prompt, self.literals) ++ if self.prompts is not None: # #567: only what follows the last special shared with a recent prompt ++ return self.prompts.encode(prompt, plain, store) ++ return self.tok.encode(prompt, parse_special=True, plain=plain) if plain else \ ++ self.tok.encode(prompt, parse_special=True) + + @contextlib.contextmanager + def _turn(self, cancel=None): +@@ -3845,6 +3948,11 @@ class Service: + content[n] = {"type": "text", "text": f"[image omitted: {why}]"} + return fetched + ++ def forget_prompt(self, ids): ++ if self.prompts is not None: ++ self.prompts.forget(ids) ++ ++ + def prepare(self, messages, tools, kwargs, max_new=None, force=None, req=None, cancel=None): + """-> (ids, thinking, max_new). An unset or non-positive max_new (some clients send -1) means "unlimited": + the rest of the context. `force` (forced_call): without thinking the reply starts with it, so it ends the +@@ -3852,7 +3960,7 @@ class Service: + # `cancel` (#1999): set when the client hangs up; the downloads, the wait for the turn and the encodes stop + # between pictures and raise ClientGone + fetched = self._note_unreadable_tool_images(messages, cancel) +- ids = self.encode_prompt(messages, tools, kwargs) ++ ids = encoded_ids = self.encode_prompt(messages, tools, kwargs) + if force and kwargs.get("enable_thinking", True) is False: + ids = ids + self.tok.encode(force, parse_special=True) + self.embeddings.path = None +@@ -3914,12 +4022,16 @@ class Service: + room = ctx - CTX_SLACK - len(ids) + if max_new is None or max_new <= 0 or (self.fit_max_tokens and room < 1): + if room < 1: ++ # A prompt refused for the context is not kept (#567): the client cannot resend it unchanged. The ++ # other refusals (engine starting, images) keep it: the same prompt comes back and its encode is reused. ++ self.forget_prompt(encoded_ids) + raise ValueError("request exceeds the context window: " + f"prompt ({len(ids)} tokens) leaves no room to answer in the context " + f"({ctx}); requests are never truncated") + max_new = room + elif max_new > room: + if not self.fit_max_tokens: ++ self.forget_prompt(encoded_ids) + raise ValueError("request exceeds the context window: " + f"prompt ({len(ids)} tokens) + max tokens ({max_new}) exceeds the context " + f"({ctx}); requests are never truncated. Send a smaller " +@@ -4076,12 +4188,23 @@ class Service: + def _run(self, ids, thinking, tools, max_new, sampling, cancel, force=None) -> Iterator[tuple[str, object]]: + """Yields ("event", Event) as text arrives, then ("done", {"finish": .., "completion_tokens": ..}). + `force` (forced_call): the opening of the call the reply must make - see prepare().""" +- defaults = {**self.sampling_defaults, **self.shared} # the config's, then the Chat settings shared with apps ++ # the model card's sampling for this thinking state, then the config's block, then the Chat settings shared with ++ # apps; the request's own fields win over all of them ++ defaults = {**model_card_sampling(self.model_arch, thinking), **self.sampling_defaults, **self.shared} + if defaults: # the request's own fields win (explicit 0 stays greedy) + req_values = {k: v for k, v in (sampling or {}).items() if v is not None} + sampling = {**defaults, **req_values} + # #123: read after the merge, so a budget shared through POST /settings is seen like the other keys + budget = self.reasoning_budget(sampling) if thinking else None ++ if thinking and not budget and self.answer_room and max_new and max_new >= ANSWER_ROOM_MIN_CAP: ++ # 0.1.45: max_tokens counts the thinking too. A client that caps the reply at 8192 (the desktop agent's ++ # harness sends max_completion_tokens: 8192 on every request) and a task the model thinks long about got ++ # the cut-off in the middle of the thinking and NO answer. So thinking is closed early enough to leave ++ # the answer its quarter of the cap (the budget's wrap-up: a short closing sentence, , the answer). ++ # A request with a budget of its own, or max_tokens below ANSWER_ROOM_MIN_CAP, is left as it was. ++ budget = max_new - min(max(ANSWER_RESERVE_MIN, max_new // 4), max_new // 2) ++ if max_new <= 32768: print(f"[strata] thinking limited to {budget} of max_tokens {max_new} so the answer has room " ++ "(\"answer_room\": false in the config turns this off)", flush=True) + if budget and max_new and max_new > 0: + # #984: a budget at or near max_tokens never fires before the reply is cut off while still thinking, so + # it has no answer. Close the thinking early enough to leave room for one. +@@ -4586,13 +4709,13 @@ def request_stats(segments: list[dict]) -> dict: + if not segments: + return {} + result = dict(segments[-1]) +- for key in ("reused", "prompt_read"): ++ for key in ("reused", "prompt_read", "chunk"): + if key in segments[0]: + result[key] = segments[0][key] + else: + result.pop(key, None) + for key in ("prompt_ms", "decode_ms", "generated", "drafts_offered", "drafts_accepted", +- "hits", "lookups", "offloaded", "ram_blobs", "file_blobs", "file_mb"): ++ "hits", "lookups", "offloaded", "ram_blobs", "file_blobs", "file_mb", "experts_streamed"): + if any(segment.get(key) is not None for segment in segments): + result[key] = sum(segment.get(key) or 0 for segment in segments) + return result +@@ -5612,7 +5735,7 @@ def make_handler(svc: Service): + if not svc.loaded() and not getattr(svc.engine, "unloaded", False): + self._json(503, {"error": {"message": "the engine is not running"}}) + return +- defaults = {**svc.sampling_defaults, **svc.shared} ++ defaults = {**model_card_sampling(svc.model_arch, True), **svc.sampling_defaults, **svc.shared} + names = {"repetition_penalty": "repeat_penalty", "penalty_last_n": "repeat_last_n"} + params = {names.get(k, k): v for k, v in defaults.items() + if k in ("temperature", "top_p", "top_k", "min_p", "seed", "repetition_penalty", +@@ -6045,7 +6168,7 @@ def make_handler(svc: Service): + messages, tools, kw = anthropic_to_messages(req, svc.anthropic_think_unasked) + stop_strings(req) # the same 400 as the request itself would get + check_request_sampling(req) # ... likewise a sampling field of the wrong type +- self._json(200, {"input_tokens": len(svc.encode_prompt(messages, tools, kw))}) ++ self._json(200, {"input_tokens": len(svc.encode_prompt(messages, tools, kw, store=False))}) + + def _anthropic(self, req): + svc.load() +@@ -6433,6 +6556,43 @@ def clean_shared_defaults(d) -> dict: + return out + + ++def split_joined_args(cfg: dict) -> None: ++ """#2267: an engine option written as ONE element of the config's "args" - "--ple-io mmap" or "--ple-io=mmap" - reaches ++ the engine as a single argument ("unknown argument: --ple-io mmap"; the usual form is two elements, "--ple-io", ++ "mmap"). Recommend, never force: such an element is split into the flag and its value, and said at startup. Only an ++ element that starts with a flag name (--letters-digits-dashes) followed by a space or "=" is touched, and never one ++ that sits in the value place of the flag before it.""" ++ args = cfg.get("args") ++ if not isinstance(args, list): ++ return ++ out, changed = [], [] ++ for a in args: ++ m = re.match(r"^(--[A-Za-z0-9][A-Za-z0-9-]*)(?:\s+|=)(\S.*)$", a) if isinstance(a, str) else None ++ prev = out[-1] if out else None ++ in_value_place = isinstance(prev, str) and re.match(r"^--[A-Za-z0-9][A-Za-z0-9-]*$", prev) is not None ++ if m is None or in_value_place: ++ out.append(a) ++ continue ++ out += [m.group(1), m.group(2).strip()] ++ changed.append(a) ++ if changed: ++ cfg["args"] = out ++ print("[strata] the config's args hold " + ", ".join(repr(c) for c in changed) + " as one element each: split into the " ++ 'flag and its value (write them as two elements, e.g. "--ple-io", "mmap")', flush=True) ++ ++ ++def read_config(path) -> dict: ++ """The run config's JSON. A mistake in it (a missing comma between two args, a trailing comma) is said with its ++ place in the file instead of a traceback (#2267: editing "args" by hand is where it happens).""" ++ text = Path(path).read_text(encoding="utf-8-sig") # Notepad adds a BOM ++ try: ++ return json.loads(text) ++ except json.JSONDecodeError as e: ++ raise SystemExit(f"[strata] {path} is not valid JSON: {e.msg} at line {e.lineno}, column {e.colno}. Check the commas " ++ 'between the entries of "args" (each option and each value is its own quoted entry, e.g. ' ++ '"--ple-io", "mmap",) and that no entry is followed by a comma at the end of a list.') from None ++ ++ + def lift_reasoning_effort_arg(cfg: dict) -> None: + """#1641: "--reasoning-effort high" in a config's "args" is meant for the thinking level, which is not an engine + option (the engine stopped at the unknown flag). It is the config's "reasoning_effort" (none, low, medium or high: +@@ -6615,6 +6775,25 @@ def check_request_sampling(req) -> None: + FAMILY_MODEL_NAMES = {"qwen4exp": "qwen3.8-flash-next", "qwen35moe": "qwen3.6-35b-a3b"} + + ++# The sampling each model card recommends ("Best Practices"), used when neither the request, the Chat settings nor the ++# run config's `sampling` block says otherwise. Greedy is only for a request or config that asks for temperature 0. ++# min_p 0, repetition_penalty 1 and presence_penalty 0 are the engine's own "off", so they are not listed. ++MODEL_CARD_SAMPLING = { ++ # Qwen3.8-Flash-Next (and Gyro, the same base model) ++ "qwen4exp": {"thinking": {"temperature": 1.0, "top_p": 0.95, "top_k": 20}, ++ "instruct": {"temperature": 0.7, "top_p": 0.8, "top_k": 20, "presence_penalty": 1.5}}, ++ # Qwen3.6-35B-A3B: the general-task sets (the card's precise-coding set is temperature 0.6, presence_penalty 0) ++ "qwen35moe": {"thinking": {"temperature": 1.0, "top_p": 0.95, "top_k": 20, "presence_penalty": 1.5}, ++ "instruct": {"temperature": 0.7, "top_p": 0.8, "top_k": 20, "presence_penalty": 1.5}}, ++} ++ ++ ++def model_card_sampling(arch, thinking) -> dict: ++ """The model card's sampling for the family `arch` and the thinking state ({} for an unknown or unset family).""" ++ sets = MODEL_CARD_SAMPLING.get(arch) if arch else None ++ return dict(sets["thinking" if thinking else "instruct"]) if sets else {} ++ ++ + def default_model_name(cfg: dict) -> str: + """The config's "model_name", else its family's name ("arch": qwen4exp, qwen35moe), else Flash-Next's.""" + return cfg.get("model_name") or FAMILY_MODEL_NAMES.get(cfg.get("arch") or "qwen4exp", "qwen3.8-flash-next") +@@ -6667,7 +6846,8 @@ def main() -> int: + "each, requests spread over them (also \"replicas\" in the config: a number, or " + "[{\"gpus\": [0, 1]}, {\"gpus\": [2, 3]}])") + a = ap.parse_args() +- cfg = json.loads(Path(a.config).read_text(encoding="utf-8-sig")) if a.config else {} # Notepad adds a BOM ++ cfg = read_config(a.config) if a.config else {} ++ split_joined_args(cfg) + lift_reasoning_effort_arg(cfg) + if a.replicas is not None: + cfg["replicas"] = a.replicas +@@ -6768,6 +6948,9 @@ def main() -> int: + model_name=default_model_name(cfg), vision=vision, + sampling_defaults=sampling_defaults, + fit_max_tokens=a.fit_max_tokens or cfg.get("fit_max_tokens") is True) ++ if a.engine == "strata": # the model card's sampling when nothing else is said (see MODEL_CARD_SAMPLING) ++ svc.model_arch = cfg.get("arch") or "qwen4exp" ++ svc.answer_room = cfg.get("answer_room") is not False + svc.reasoning_close_retry = cfg.get("reasoning_close_retry") is True # #1053: opt-in, off by default + if cfg.get("images"): # image generation (serve/images.py), set up by setup --images + icfg = {k: (os.path.abspath(os.path.join(cfg.get("cwd") or ".", v)) +diff --git a/serve/telemetry.py b/serve/telemetry.py +index 3232c63c..2e2a28ad 100644 +--- a/serve/telemetry.py ++++ b/serve/telemetry.py +@@ -3,7 +3,9 @@ + A background thread samples once a second and keeps the last 60 readings of each series for the sparklines: + - GPU: NVIDIA's own NVML library (nvml.dll / libnvidia-ml.so.1, installed with every driver) through ctypes, so no + pip package is needed: load, VRAM, temperature, power, PCIe link and throughput. With the AMD backend (#301): the +- amdgpu driver's Linux sysfs files - load, VRAM, temperature and power. ++ amdgpu driver's Linux sysfs files - load, VRAM, temperature and power. On macOS (Metal backend): IOKit's ++ IOAccelerator registry - load and VRAM; the temperature and power sensors need root (`powermetrics`), and an Apple ++ Silicon GPU sits on the die, so there is no PCIe link. + - CPU, RAM, disk: `psutil` when it is installed (setup installs it); without it the CPU and RAM readings fall back to + the OS (Windows GlobalMemoryStatusEx / GetSystemTimes, Linux /proc) and the disk rate is absent. + Anything that cannot be read is None; nothing here can stop the server. +@@ -14,6 +16,8 @@ import collections + import ctypes + import os + import platform ++import re ++import subprocess + import sys + import threading + import time +@@ -246,9 +250,108 @@ class _Amd: + return out + + ++def metal_working_set(): ++ """The Apple GPU's memory budget in bytes: MTLCreateSystemDefaultDevice's recommendedMaxWorkingSetSize, through ++ the Metal framework and objc_msgSend (no PyObjC) - the same property metal/src/mtl_rt.mm reads for the engine ++ (its "GPU working set" log line). None off macOS, or when it cannot be read: a selector the device does not ++ answer would raise an Objective-C exception no ctypes call can catch, so the class is asked first.""" ++ if sys.platform != "darwin": ++ return None ++ try: ++ metal = ctypes.CDLL("/System/Library/Frameworks/Metal.framework/Metal") ++ objc = ctypes.CDLL("/usr/lib/libobjc.dylib") ++ metal.MTLCreateSystemDefaultDevice.restype = ctypes.c_void_p ++ objc.sel_registerName.restype = ctypes.c_void_p ++ objc.object_getClass.argtypes = [ctypes.c_void_p] ++ objc.object_getClass.restype = ctypes.c_void_p ++ objc.class_respondsToSelector.argtypes = [ctypes.c_void_p, ctypes.c_void_p] ++ objc.class_respondsToSelector.restype = ctypes.c_ubyte ++ objc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] ++ objc.objc_msgSend.restype = ctypes.c_ulonglong ++ dev = metal.MTLCreateSystemDefaultDevice() ++ if not dev: ++ return None ++ sel = objc.sel_registerName(b"recommendedMaxWorkingSetSize") ++ if not objc.class_respondsToSelector(objc.object_getClass(dev), sel): ++ return None ++ return int(objc.objc_msgSend(ctypes.c_void_p(dev), ctypes.c_void_p(sel))) or None ++ except (OSError, AttributeError, ValueError): ++ return None ++ ++ ++class _Metal: ++ """macOS: an Apple GPU's readings from the IOKit registry (`ioreg -l -w0 -r -c IOAccelerator`, 22 ms), with ++ _Nvml's interface: the load and the memory are the driver's PerformanceStatistics of the index'th accelerator ++ node - Device Utilization % and the in-use system memory in bytes. The total is the device's Metal working set, ++ the budget one app gets of the unified memory. No temperature or power reading exists without root ++ (`powermetrics`), and an Apple Silicon GPU sits on the die, so the PCIe cards have nothing to show.""" ++ note = ("no GPU temperature or power reading without root (powermetrics), and an Apple Silicon GPU sits on the " ++ "die, so there is no PCIe link to show; the engine's own figures are in its log") ++ cmd = ("ioreg", "-l", "-w0", "-r", "-c", "IOAccelerator") ++ ++ def __init__(self, index=0, run=None): ++ self.index = index ++ self._run = run or self._ioreg ++ self.nodes = len(self._blocks()) ++ self._name = None ++ b = self._block() ++ if b: ++ m = re.search(r'"model" = "([^"]+)"', b) ++ self._name = m.group(1) if m else None ++ ++ @staticmethod ++ def _ioreg(): ++ try: ++ return subprocess.run(list(_Metal.cmd), capture_output=True, text=True, timeout=10).stdout or "" ++ except (OSError, subprocess.SubprocessError): ++ return "" ++ ++ def _blocks(self): ++ """The accelerator nodes, in the registry's order: every node that carries a PerformanceStatistics.""" ++ try: ++ out = self._run() or "" ++ except OSError: # ioreg missing, or a caller's runner that raises ++ return [] ++ return [b for b in re.split(r"^\s*\+-o ", out, flags=re.M) if '"PerformanceStatistics"' in b] ++ ++ def _block(self): ++ b = self._blocks() ++ return b[self.index] if 0 <= self.index < len(b) else None ++ ++ def ok(self): ++ return self.index < self.nodes ++ ++ def name(self): ++ return self._name ++ ++ def read(self): ++ b = self._block() ++ if b is None: ++ return {} ++ ++ def num(key): ++ m = re.search(r'"%s"=([0-9.]+)' % re.escape(key), b) ++ return float(m.group(1)) if m else None ++ ++ out = {} ++ u = num("Device Utilization %") ++ if u is None: # older chips report only the renderer and the tiler ++ r, t = num("Renderer Utilization %"), num("Tiler Utilization %") ++ u = max(v for v in (r, t) if v is not None) if (r is not None or t is not None) else None ++ out["util"] = int(u) if u is not None else None ++ used = num("In use system memory") ++ out["mem_used"] = int(used) if used is not None else None ++ out["mem_total"] = metal_working_set() ++ return out ++ ++ + def gpu_reader(index=0, amd=False): +- """The card's readings: NVML (NVIDIA), or the amdgpu sysfs files with the AMD backend (#301).""" +- return _Amd(index) if amd else _Nvml(index) ++ """The card's readings: NVML (NVIDIA), the amdgpu sysfs files with the AMD backend (#301), or IOKit's ++ IOAccelerator registry on macOS (Metal backend). `amd` comes first: no Mac runs the HIP backend, and the ++ tests' AMD trees are read on a Mac too.""" ++ if amd: ++ return _Amd(index) ++ return _Metal(index) if sys.platform == "darwin" else _Nvml(index) + + + _DISPLAY_CLASS = r"SYSTEM\CurrentControlSet\Control\Class\{4d36e968-e325-11ce-bfc1-08002be10318}" +@@ -413,15 +516,21 @@ class Telemetry: + cards = windows_amd_cards() + win_amd = [cards[i] for i in idx if 0 <= i < len(cards)] or cards[:1] + gpu_name = " + ".join(c["name"] for c in win_amd) if win_amd else None ++ # the macOS reader knows the device's working set (the VRAM budget of the unified memory); None elsewhere ++ vram_total = None ++ if win_amd is None and isinstance(self.gpu, _Metal) and self.gpu.ok(): ++ ws = metal_working_set() ++ vram_total = ws >> 20 if ws else None + self.static = { + "gpu_name": gpu_name, + "gpu_vram_total_mib": (sum(c["vram_bytes"] for c in win_amd) >> 20 +- if win_amd and all(c["vram_bytes"] for c in win_amd) else None), ++ if win_amd and all(c["vram_bytes"] for c in win_amd) else vram_total), + "gpu_count": len(self.gpus), + # #1380: the AMD readings are the amdgpu driver's Linux sysfs files; a Windows AMD card has none yet, and the + # dashboard said "not readable (NVML)" or showed empty tiles with no word why + "gpu_note": ("no GPU load or VRAM readings for AMD cards on Windows yet (Linux reads them from the amdgpu " +- "driver); the engine's own VRAM figures are in its log" if amd and not self.gpu.ok() else None), ++ "driver); the engine's own VRAM figures are in its log" if amd and not self.gpu.ok() ++ else getattr(self.gpu, "note", None)), + "cpu_name": _cpu_name(), + "cores": (self.ps.cpu_count(logical=False) if self.ps else None) or None, + "threads": os.cpu_count(), +diff --git a/tools/strata_tokenizer.py b/tools/strata_tokenizer.py +index c9e84ee7..42a27fed 100644 +--- a/tools/strata_tokenizer.py ++++ b/tools/strata_tokenizer.py +@@ -15,9 +15,11 @@ emoji (4-byte), combining marks, whitespace runs, and C0 control bytes. + """ + from __future__ import annotations + ++import bisect + import json + import pathlib + import sys ++import threading + + import regex + +@@ -115,6 +117,8 @@ class Tokenizer: + # through __dict__.setdefault on each encoded segment or decoded token. + self._piece_ids: dict[str, list[int]] = {} + self._bytes_cache: dict[int, bytes] = {} ++ # How far a special-token match can look ahead of where it starts (PromptEncoder's margin). ++ self.max_special_len = max((len(t) for t in self.special_tokens), default=1) + + @staticmethod + def _alt(literals: list[str]): +@@ -234,13 +238,16 @@ class Tokenizer: + out.extend(got) + return out + +- def _encode_matching(self, text: str, pat, plain=()) -> list[int]: ++ def _encode_matching(self, text: str, pat, plain=(), marks: list | None = None) -> list[int]: + """Encode `text`, emitting any literal `pat` matches as single tokens and BPE-ing the rest. + + The split happens on the RAW text, before the byte mapping, because a special token's string is a + literal to match rather than bytes to decompose. Everything between the matches is tokenized + normally - which is why a near-miss like `<|im_star` still costs ordinary tokens. A match that starts + inside one of the `plain` (start, end) spans is left to the text around it (#537). ++ ++ `marks`, when given, receives (end of the match in `text`, ids so far) for every match: the points where ++ the encoding of a longer text with the same beginning can resume (PromptEncoder). + """ + if pat is None: + return self._encode_plain(text) +@@ -269,6 +276,8 @@ class Tokenizer: + out.extend(self._encode_plain(text[pos:start])) + out.append(self.special_tokens[m.group(0)]) + pos = m.end() ++ if marks is not None: ++ marks.append((pos, len(out))) + if pos < len(text): + out.extend(self._encode_plain(text[pos:])) + return out +@@ -283,6 +292,11 @@ class Tokenizer: + """ + return self._encode_matching(text, self._special_re if parse_special else self._always_re, plain) + ++ def encode_marked(self, text: str, parse_special: bool = False, plain=()) -> tuple[list[int], list[tuple[int, int]]]: ++ """encode() and its resume points: (end offset in `text`, ids so far) after every special-token match.""" ++ marks: list[tuple[int, int]] = [] ++ return self._encode_matching(text, self._special_re if parse_special else self._always_re, plain, marks), marks ++ + def token_bytes(self, i: int) -> bytes: + """The raw bytes of one token (a multi-byte character can be split across tokens).""" + cache = self._bytes_cache +@@ -303,6 +317,111 @@ class Tokenizer: + return b"".join(self.token_bytes(i) for i in ids).decode("utf-8", errors=errors) + + ++# ------------------------------------------------------------------ incremental prompts (#567) ++def common_prefix_len(a: str, b: str) -> int: ++ """Length of the longest common prefix, by slices (C speed) rather than a loop over characters.""" ++ n = min(len(a), len(b)) ++ lo, step = 0, 4096 ++ while lo < n: # whole chunks first, doubling ++ hi = min(n, lo + step) ++ if a[lo:hi] != b[lo:hi]: ++ break ++ lo, step = hi, min(step * 2, 1 << 20) ++ else: ++ return n ++ while hi - lo > 1: # the first difference is in [lo, hi): bisect it ++ mid = (lo + hi) // 2 ++ if a[lo:mid] == b[lo:mid]: ++ lo = mid ++ else: ++ hi = mid ++ return lo ++ ++ ++def plain_agree_until(a, b, limit: int) -> int: ++ """The first position below `limit` where the `plain` spans (#537) of two texts differ - a position inside a ++ span in one and not in the other - or `limit`. Spans are (start, end), sorted and disjoint.""" ++ def clip(spans): ++ return [(x, min(y, limit)) for x, y in spans if x < limit and y > x] ++ ca, cb = clip(a), clip(b) ++ for (xa, ya), (xb, yb) in zip(ca, cb): ++ if (xa, ya) != (xb, yb): ++ return min(xa, xb) if xa != xb else min(ya, yb) ++ if len(ca) != len(cb): ++ return (ca if len(ca) > len(cb) else cb)[min(len(ca), len(cb))][0] ++ return limit ++ ++ ++class PromptEncoder: ++ """`tok.encode(text, parse_special=True, plain=...)` for chat prompts, reusing the ids of an earlier prompt. ++ ++ A chat client sends the whole conversation every turn, and the rendered prompt of turn n+1 starts with most of ++ turn n's. Running BPE over all of it again costs hundreds of milliseconds at 100K tokens, even when the engine ++ reuses the whole prefix. So the ids up to a SPECIAL-TOKEN BOUNDARY both prompts share are taken as they were, ++ and only the rest is encoded. ++ ++ WHY THAT IS EXACT. Encoding splits the text at special-token matches and BPE-encodes each stretch between two ++ matches on its own (`_encode_matching`): nothing crosses a match. Whether a match is taken at position q ++ depends only on text[q : q + max_special_len] and on whether q lies inside a `plain` span (#537: a ++ quoted in a message is text). So if the two texts agree up to L and their plain spans agree below L, every ++ decision the scan makes at a position before c = the end of some match, with c + max_special_len - 1 <= L, is ++ the same for both - the same matches, the same stretches, the same ids up to c - and from c the scan starts ++ afresh, exactly as it does on text[c:] with the spans shifted by c. Hence encode(new) == ids_old[:ids at c] + ++ encode(new[c:]). The margin matters only for a vocabulary where a literal overlaps another's end; it costs ++ re-encoding one short stretch. ++ ++ A few recent prompts are kept (one per conversation: a prompt replaces the one it continues, that is, one whose ++ text it reuses for more than half; a prompt that shares only the template's head with every entry is a new ++ entry), so a second client does not evict the first. `encode(..., store=False)` reads the kept prompts without ++ writing one (a token count is not a conversation's turn). The tokenizer needs `encode_marked` and ++ `max_special_len`. ++ """ ++ ++ def __init__(self, tok, keep: int = 4): ++ self.tok, self.keep = tok, keep ++ # (text, plain spans, ids, mark ends, mark counts) ++ self.entries: list[tuple[str, tuple, list[int], list[int], list[int]]] = [] ++ self.lock = threading.Lock() ++ self.last_reused = 0 # characters taken from an earlier prompt (for tests) ++ ++ def encode(self, text: str, plain=(), store: bool = True) -> list[int]: ++ plain = tuple(sorted(plain)) ++ with self.lock: ++ entries = self.entries # replaced, never changed in place ++ margin = self.tok.max_special_len - 1 ++ src, cut_k = None, -1 ++ for e in entries: ++ same = plain_agree_until(e[1], plain, common_prefix_len(e[0], text)) ++ k = bisect.bisect_right(e[3], same - margin) - 1 # the last boundary c with c <= same - margin ++ if k >= 0 and (src is None or e[3][k] > src[3][cut_k]): ++ src, cut_k = e, k ++ if src is None: ++ ids, marks = self.tok.encode_marked(text, parse_special=True, plain=plain) ++ ends, counts = [m[0] for m in marks], [m[1] for m in marks] ++ self.last_reused = 0 ++ else: ++ c, n = src[3][cut_k], src[4][cut_k] ++ rest = tuple((max(x - c, 0), y - c) for x, y in plain if y > c) ++ tail, marks = self.tok.encode_marked(text[c:], parse_special=True, plain=rest) ++ ids = src[2][:n] + tail ++ ends = src[3][:cut_k + 1] + [c + m[0] for m in marks] ++ counts = src[4][:cut_k + 1] + [n + m[1] for m in marks] ++ self.last_reused = c ++ if not store: ++ return ids ++ # every rendered prompt starts with the template's header, so sharing a special boundary does not make two ++ # prompts one conversation: only the entry whose text is mostly reused is replaced ++ old = src if src is not None and 2 * src[3][cut_k] > len(src[0]) else None ++ with self.lock: ++ self.entries = [(text, plain, ids, ends, counts)] + [e for e in self.entries if e is not old][:self.keep - 1] ++ return list(ids) ++ ++ def forget(self, ids: list[int]) -> None: ++ """The prompt with these ids was refused: it is not kept.""" ++ with self.lock: ++ self.entries = [e for e in self.entries if e[2] != ids] ++ ++ + # ------------------------------------------------------------------ the pack's tokenizer/ directory + def extract(gguf_path, out_dir) -> dict: + """Write the tokenizer into `/tokenizer/` so the engine never opens the weight shards for it.""" diff --git a/results/round39/service-health-after-ab.json b/results/round39/service-health-after-ab.json new file mode 100644 index 0000000..a55c1c1 --- /dev/null +++ b/results/round39/service-health-after-ab.json @@ -0,0 +1 @@ +{"status": "ok", "max_context": 32768, "model": "swift-1.5-k512-iq2xs", "images": false, "api_key": false, "image_generation": false, "loaded": true, "service": "strata"} \ No newline at end of file diff --git a/results/round39/service-health-final.json b/results/round39/service-health-final.json new file mode 100644 index 0000000..a55c1c1 --- /dev/null +++ b/results/round39/service-health-final.json @@ -0,0 +1 @@ +{"status": "ok", "max_context": 32768, "model": "swift-1.5-k512-iq2xs", "images": false, "api_key": false, "image_generation": false, "loaded": true, "service": "strata"} \ No newline at end of file diff --git a/results/round39/service-health.json b/results/round39/service-health.json new file mode 100644 index 0000000..a55c1c1 --- /dev/null +++ b/results/round39/service-health.json @@ -0,0 +1 @@ +{"status": "ok", "max_context": 32768, "model": "swift-1.5-k512-iq2xs", "images": false, "api_key": false, "image_generation": false, "loaded": true, "service": "strata"} \ No newline at end of file diff --git a/results/round39/tokenizer-baseline.log b/results/round39/tokenizer-baseline.log new file mode 100644 index 0000000..e047cf3 --- /dev/null +++ b/results/round39/tokenizer-baseline.log @@ -0,0 +1,5 @@ +...... +---------------------------------------------------------------------- +Ran 6 tests in 0.019s + +OK diff --git a/results/round39/tokenizer-post.log b/results/round39/tokenizer-post.log new file mode 100644 index 0000000..db6c8c5 --- /dev/null +++ b/results/round39/tokenizer-post.log @@ -0,0 +1,5 @@ +........ss..s +---------------------------------------------------------------------- +Ran 11 tests in 0.393s + +OK (skipped=3) diff --git a/results/round39/tokenizer-real.log b/results/round39/tokenizer-real.log new file mode 100644 index 0000000..eafe921 --- /dev/null +++ b/results/round39/tokenizer-real.log @@ -0,0 +1,24 @@ +test_decode_before_encode_cold_and_warm (tools.test_strata_tokenizer.ByteCache.test_decode_before_encode_cold_and_warm) ... ok +test_instances_with_different_vocabularies_stay_independent (tools.test_strata_tokenizer.ByteCache.test_instances_with_different_vocabularies_stay_independent) ... ok +test_invalid_ids_before_and_after_cache_use (tools.test_strata_tokenizer.ByteCache.test_invalid_ids_before_and_after_cache_use) ... ok +test_shared_tokenizer_concurrent_first_decode (tools.test_strata_tokenizer.ByteCache.test_shared_tokenizer_concurrent_first_decode) ... ok +test_against_a_loop (tools.test_strata_tokenizer.CommonPrefix.test_against_a_loop) ... ok +test_literal_ending_inside_another (tools.test_strata_tokenizer.Margin.test_literal_ending_inside_another) ... ok +test_longer_literal_past_the_common_prefix (tools.test_strata_tokenizer.Margin.test_longer_literal_past_the_common_prefix) ... ok +test_random_specials (tools.test_strata_tokenizer.Margin.test_random_specials) +Literals that extend or end inside one another, texts that change right after one of them. ... ok +test_marks (tools.test_strata_tokenizer.Marks.test_marks) ... ok +test_llama_cpp_vectors (tools.test_strata_tokenizer.Oracle.test_llama_cpp_vectors) ... skipped "llama.cpp's ggml-vocab-qwen35.gguf test vectors are not here" +test_same_ids_cold_warm_and_off (tools.test_strata_tokenizer.PieceCache.test_same_ids_cold_warm_and_off) ... ok +test_the_cache_is_bounded (tools.test_strata_tokenizer.PieceCache.test_the_cache_is_bounded) ... ok +test_a_read_only_encode_gives_the_ids_and_stores_nothing (tools.test_strata_tokenizer.Prompts.test_a_read_only_encode_gives_the_ids_and_stores_nothing) ... ok +test_chat_golden (tools.test_strata_tokenizer.Prompts.test_chat_golden) ... ok +test_conversations (tools.test_strata_tokenizer.Prompts.test_conversations) ... ok +test_growing_turn_reuses_the_previous_prompt (tools.test_strata_tokenizer.Prompts.test_growing_turn_reuses_the_previous_prompt) ... ok +test_prefix_of_an_earlier_prompt_and_unrelated_prompts (tools.test_strata_tokenizer.Prompts.test_prefix_of_an_earlier_prompt_and_unrelated_prompts) ... ok +test_the_returned_ids_are_the_callers (tools.test_strata_tokenizer.Prompts.test_the_returned_ids_are_the_callers) ... ok + +---------------------------------------------------------------------- +Ran 18 tests in 0.728s + +OK (skipped=1) diff --git a/results/round39/vendor-before.patch b/results/round39/vendor-before.patch new file mode 100644 index 0000000..20a36f7 --- /dev/null +++ b/results/round39/vendor-before.patch @@ -0,0 +1,17443 @@ +diff --git a/include/strata/core/scr.hpp b/include/strata/core/scr.hpp +new file mode 100644 +index 00000000..2a48f751 +--- /dev/null ++++ b/include/strata/core/scr.hpp +@@ -0,0 +1,227 @@ ++// Suffix cache reuse (SCR): after an edit in the middle of a prompt, the K/V of the text that survived the edit is ++// moved to its new positions instead of being read again. Ported from llama-cpp-bonsai's common/scr.cpp. ++// ++// The moved cells were computed over the text before the edit, so the result approximates a fresh read. The GDN ++// layers have no cells to move: their state continues from a conversation checkpoint at the end of the last moved ++// run (the plan's `adopt`), and everything after that run is read from it. Strata's indexer pools 4 cells into one ++// key and keeps only the pooled keys, so a run starts and ends on a block boundary of the OLD prompt and moves by a ++// multiple of 4: up to 3 pad tokens are read before a run to make its shift one (DeepSeek V4.1's recipe pads to an ++// even shift the same way). The moved blocks then keep their cells, and their pooled keys only turn by the shift. ++#pragma once ++ ++#include ++#include ++#include ++#include ++#include ++ ++namespace strata::core { ++ ++struct ScrParams { ++ int64_t min_span = 256; // shortest run worth moving ++ int64_t max_runs = 64; // most runs moved per request ++ int64_t seed = 32; // n-gram length that seeds a match ++ int64_t block = 4; // the indexer's cells per pooled key: read parts start and end on these ++}; ++ ++// the cells old[old_start, old_start + n) reappear at [new_start, new_start + n) of the prompt, after `pad` pad tokens ++// that the read inserts right before the run; `at` is where the run lands in the padded (physical) prompt ++struct ScrRun { ++ int64_t new_start = 0, old_start = 0, n = 0; ++ int64_t pad = 0, at = 0; ++ int64_t shift() const { return at - old_start; } ++}; ++ ++struct ScrPlan { ++ int64_t n_prefix = 0; // leading tokens equal in both prompts ++ int64_t restore = 0; // the checkpoint the read restarts from (at or before n_prefix) ++ int64_t adopt = -1; // old length of the checkpoint the GDN state continues from after the last run ++ std::vector runs; // in new order, at or after restore ++ ++ bool empty() const { return runs.empty(); } ++ // the prompt position the final read starts at, and the same in the padded prompt ++ int64_t resume() const { return runs.empty() ? restore : runs.back().new_start + runs.back().n; } ++ int64_t resume_at() const { return runs.empty() ? restore : runs.back().at + runs.back().n; } ++ int64_t pads() const { ++ int64_t n = 0; ++ for (const ScrRun& r : runs) n += r.pad; ++ return n; ++ } ++ int64_t reused() const { ++ int64_t n = 0; ++ for (const ScrRun& r : runs) n += r.n; ++ return n; ++ } ++ std::string str() const { ++ std::string s = "prefix " + std::to_string(n_prefix) + ", restore " + std::to_string(restore) + ", adopt " + ++ std::to_string(adopt); ++ for (const ScrRun& r : runs) ++ s += ", R " + std::to_string(r.new_start) + " <- " + std::to_string(r.old_start) + " x " + std::to_string(r.n) + ++ (r.pad ? " pad " + std::to_string(r.pad) : ""); ++ return s; ++ } ++}; ++ ++namespace scr_detail { ++// greedy maximal spans of cur[from, ...) found in old[from, ...), seeded by n-gram hashes, each old token used once ++template ++std::vector find_spans(const std::vector& old, const std::vector& cur, int64_t from, const ScrParams& p) { ++ std::vector spans; ++ const int64_t n_old = (int64_t) old.size(), n_cur = (int64_t) cur.size(), k = p.seed; ++ if (n_old - from < p.min_span || n_cur - from < p.min_span) return spans; ++ auto gram = [k](const auto& v, int64_t i) { ++ uint64_t h = 1469598103934665603ULL; ++ for (int64_t j = 0; j < k; ++j) h = (h ^ (uint32_t) v[(size_t) (i + j)]) * 1099511628211ULL; ++ return h; ++ }; ++ constexpr size_t max_bucket = 16; // the first few old positions of every seed n-gram ++ std::unordered_map> index; ++ index.reserve((size_t) (n_old - from)); ++ for (int64_t o = from; o + k <= n_old; ++o) { ++ auto& b = index[gram(old, o)]; ++ if (b.size() < max_bucket) b.push_back(o); ++ } ++ std::vector used((size_t) n_old, false); ++ for (int64_t q = from; q + k <= n_cur;) { ++ const auto it = index.find(gram(cur, q)); ++ int64_t best_o = -1, best_n = 0; ++ if (it != index.end()) ++ for (const int64_t o : it->second) { ++ int64_t n = 0; ++ while (q + n < n_cur && o + n < n_old && !used[(size_t) (o + n)] && ++ (int64_t) cur[(size_t) (q + n)] == (int64_t) old[(size_t) (o + n)]) ++ ++n; ++ if (n > best_n) { best_n = n; best_o = o; } ++ } ++ if (best_n >= p.min_span) { ++ spans.push_back({q, best_o, best_n}); ++ std::fill(used.begin() + best_o, used.begin() + best_o + best_n, true); ++ q += best_n; ++ } else { ++ ++q; ++ } ++ } ++ return spans; ++} ++inline int64_t round_up(int64_t x, int64_t b) { return (x + b - 1) / b * b; } ++inline int64_t round_down(int64_t x, int64_t b) { return x / b * b; } ++} // namespace scr_detail ++ ++// Plans the reuse of a session that holds old[0, n_old) for the prompt `cur`. `checkpoints`: the old lengths whose ++// GDN state can be restored (the live state counts as old.size()). A run starts on a block boundary of the new ++// prompt; every run but the last also ends on one, and the last one ends at a checkpoint. The plan always leaves at ++// least one token of `cur` to read. ++template ++ScrPlan scr_plan(const std::vector& old, const std::vector& cur, std::vector checkpoints, ++ const ScrParams& p) { ++ using namespace scr_detail; ++ ScrPlan plan; ++ const int64_t n_cur = (int64_t) cur.size(), bl = std::max(p.block, 1); ++ std::sort(checkpoints.begin(), checkpoints.end()); ++ int64_t a = 0; ++ while (a < (int64_t) old.size() && a < n_cur && (int64_t) old[(size_t) a] == (int64_t) cur[(size_t) a]) ++a; ++ plan.n_prefix = a; ++ for (const int64_t q : checkpoints) ++ if (q <= a) plan.restore = std::max(plan.restore, q); ++ auto spans = find_spans(old, cur, a, p); ++ // a run starts on a block boundary of the old prompt and (but the last) ends on one ++ for (ScrRun& s : spans) { ++ const int64_t cut = round_up(s.old_start, bl) - s.old_start; ++ s.new_start += cut; ++ s.old_start += cut; ++ s.n -= cut; ++ s.n = round_down(s.old_start + s.n, bl) - s.old_start; ++ } ++ spans.erase(std::remove_if(spans.begin(), spans.end(), [&](const ScrRun& s) { return s.n < p.min_span; }), ++ spans.end()); ++ // the largest runs, back in prompt order ++ std::sort(spans.begin(), spans.end(), [](const ScrRun& x, const ScrRun& y) { return x.n > y.n; }); ++ if ((int64_t) spans.size() > p.max_runs) spans.resize((size_t) p.max_runs); ++ std::sort(spans.begin(), spans.end(), [](const ScrRun& x, const ScrRun& y) { return x.new_start < y.new_start; }); ++ // the last run ends at a checkpoint, which also ends the reuse: everything after it is read from that state. ++ // It was cut back to a block boundary, so a checkpoint up to block - 1 cells past it may still match. ++ while (!spans.empty()) { ++ ScrRun& last = spans.back(); ++ int64_t end = -1; ++ for (const int64_t q : checkpoints) { ++ const int64_t n = q - last.old_start; ++ if (n < p.min_span || n >= last.n + bl || last.new_start + n >= n_cur || q <= end) continue; ++ if (n <= last.n || std::equal(old.begin() + last.old_start + last.n, old.begin() + q, ++ cur.begin() + last.new_start + last.n, ++ [](const A& x, const B& y) { return (int64_t) x == (int64_t) y; })) ++ end = q; ++ } ++ if (end >= 0) { ++ last.n = end - last.old_start; ++ plan.adopt = end; ++ break; ++ } ++ spans.pop_back(); ++ } ++ // the pads: walking the runs in prompt order, each gets what makes its padded position a multiple of the block ++ // away from its old one ++ int64_t pads = 0; ++ for (ScrRun& s : spans) { ++ s.pad = ((s.old_start - (s.new_start + pads)) % bl + bl) % bl; ++ pads += s.pad; ++ s.at = s.new_start + pads; ++ } ++ plan.runs = std::move(spans); ++ return plan; ++} ++ ++// The pads an earlier splice read (`pads`: their positions in `live`, sorted) that the prompt `cur` still shares with ++// the session: the ones before the last token of `cur` that `live` holds at the same place once the pads are ++// skipped. `out` is `cur` with those pads put back, which makes it start with `live` as far as they agree. ++template ++std::vector scr_carry(const std::vector& cur, const std::vector& live, const std::vector& pads, ++ int64_t pad_token, std::vector& out) { ++ size_t p = 0, i = 0, k = 0, kept = 0; ++ while (p < live.size()) { ++ if (k < pads.size() && pads[k] == (int64_t) p) { ++ if ((int64_t) live[p] != pad_token) break; ++ ++k; ++ ++p; ++ continue; ++ } ++ if (i + 1 >= cur.size() || (int64_t) live[p] != (int64_t) cur[i]) break; ++ ++i; ++ ++p; ++ kept = k; ++ } ++ std::vector carried(pads.begin(), pads.begin() + (long) kept); ++ out.clear(); ++ out.reserve(cur.size() + carried.size()); ++ size_t c = 0; ++ for (const int64_t at : carried) { ++ while ((int64_t) out.size() < at) out.push_back(cur[c++]); ++ out.push_back((B) pad_token); ++ } ++ out.insert(out.end(), cur.begin() + (long) c, cur.end()); ++ return carried; ++} ++ ++// The positions of the pads `plan` inserts, in the padded prompt. ++inline std::vector scr_pad_positions(const ScrPlan& plan) { ++ std::vector at; ++ for (const ScrRun& r : plan.runs) ++ for (int64_t j = r.pad; j > 0; --j) at.push_back(r.at - j); ++ return at; ++} ++ ++// The prompt the engine reads under `plan`: `cur` with each run's pads inserted right before it. ++template ++std::vector scr_padded(const std::vector& cur, const ScrPlan& plan, B pad_token) { ++ std::vector out; ++ out.reserve(cur.size() + (size_t) plan.pads()); ++ int64_t from = 0; ++ for (const ScrRun& r : plan.runs) { ++ out.insert(out.end(), cur.begin() + from, cur.begin() + r.new_start); ++ out.insert(out.end(), (size_t) r.pad, pad_token); ++ from = r.new_start; ++ } ++ out.insert(out.end(), cur.begin() + from, cur.end()); ++ return out; ++} ++ ++} // namespace strata::core +diff --git a/metal/CMakeLists.txt b/metal/CMakeLists.txt +index d23a26fc..92cd20c7 100644 +--- a/metal/CMakeLists.txt ++++ b/metal/CMakeLists.txt +@@ -120,4 +120,12 @@ target_compile_definitions(strata_metal_kernels PRIVATE + target_link_libraries(strata_metal_kernels PUBLIC strata_metal_host mtl_rt) + + add_executable(strata ${S}/program/generate.cpp) +-target_link_libraries(strata PRIVATE strata_metal_kernels strata_metal_host strata_metal_kernels) ++target_link_libraries(strata PRIVATE strata_metal_kernels) ++ ++add_executable(strata_window_cli src/strata_window_cli.cpp) ++target_link_libraries(strata_window_cli PRIVATE strata_metal_kernels) ++add_executable(convert_k256_q2_0 src/convert_k256_q2_0.cpp) ++target_link_libraries(convert_k256_q2_0 PRIVATE mtl_rt) ++add_executable(heal_k256_q2_0 src/heal_k256_q2_0.cpp) ++target_compile_definitions(heal_k256_q2_0 PRIVATE ACCELERATE_NEW_LAPACK) ++target_link_libraries(heal_k256_q2_0 PRIVATE mtl_rt "-framework Accelerate") +diff --git a/metal/kernels/dense_mm.metal b/metal/kernels/dense_mm.metal +index 45e29ace..ccf06f79 100644 +--- a/metal/kernels/dense_mm.metal ++++ b/metal/kernels/dense_mm.metal +@@ -8,7 +8,12 @@ using namespace mpp::tensor_ops; + + kernel void act_f16(device const float* x [[buffer(0)]], device half* xh [[buffer(1)]], constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { +- if (gid < n) xh[gid] = half(x[gid]); ++ const uint n4 = n >> 2; ++ if (gid < n4) { ++ ((device half4*)xh)[gid] = half4(((device const float4*)x)[gid]); ++ } else if (gid == n4) { ++ for (uint i = n4 << 2; i < n; ++i) xh[i] = half(x[i]); ++ } + } + + struct DqArgs { +@@ -19,27 +24,46 @@ struct DqArgs { + // streams of an output interleaved, for gemm_hc_up's epilogue) + }; + +-// one thread per unit: grid (units per row / 64, n_out), 64 threads ++// one thread per unit: grid ((units per row + nth.x - 1) / nth.x, n_out), nth.x threads + template + kernel void dq_f16(device const uchar* W [[buffer(0)]], device half* Wh [[buffer(1)]], constant DqArgs& a [[buffer(2)]], +- uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { ++ uint3 gid [[thread_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], ++ ushort3 nth [[threads_per_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, H = UNIT / 2; + threadgroup uint tab[fmt::TAB]; +- load_tab(tab, tid, 64); +- threadgroup_barrier(mem_flags::mem_threadgroup); +- const int u = int(tg.x) * 64 + tid, r = int(tg.y); ++ if (F == T_IQ4_XS || F == T_IQ4_NL) { ++ if (tid < 64) { ++ const uint r0 = tid & 15u, sh = (tid >> 4) * 8u; ++ const half h0 = half(as_type(uchar(((r0 & 8u) ? 0x7159453526190d01ul : 0xf6eaddcfbfad9881ul) >> ((r0 & 7u) * 8u)))); ++ tab[tid] = as_type(half2(h0, half(as_type(uchar(0xbfad9881u >> sh))))); ++ tab[tid + 64] = as_type(half2(h0, half(as_type(uchar(0xf6eaddcfu >> sh))))); ++ tab[tid + 128] = as_type(half2(h0, half(as_type(uchar(0x26190d01u >> sh))))); ++ tab[tid + 192] = as_type(half2(h0, half(as_type(uchar(0x71594535u >> sh))))); ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } else if (fmt::TAB > 1) { ++ load_tab(tab, tid, nth.x); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } ++ const int u = int(gid.x); + if (u * UNIT >= int(a.n_in)) return; +- // the unit's chunk origin: a 128-wide chunk holds whole units of every format +- const int k0 = (u * UNIT) / 128 * 128; ++ uint r = gid.y, ro = r; ++ if (a.perm) { ++ if (fmt::BB * 16 < fmt::QK * 17) { ++ r = a.perm == 4 ? (ro & 3u) * (a.n_out >> 2) + (ro >> 2) : (ro % a.perm) * (a.n_out / a.perm) + ro / a.perm; ++ } else { ++ const uint q = a.n_out / a.perm; ++ ro = (r % q) * a.perm + r / q; ++ } ++ } + float w[UNIT]; + int p0, p1; +- DW::run(W + ulong(r) * a.row_bytes, u, k0, w, p0, p1, tab); +- const int q = a.perm ? int(a.n_out / a.perm) : 1; +- const int ro = a.perm ? (r % q) * int(a.perm) + r / q : r; +- device half* o = Wh + ulong(ro) * a.n_in + k0; +- for (int j = 0; j < H; ++j) { +- o[p0 + j] = half(w[j]); +- o[p1 + j] = half(w[H + j]); ++ DW::run(W + ulong(r) * a.row_bytes, u, 0, w, p0, p1, tab); ++ device half* o = Wh + ulong(ro) * a.n_in; ++#pragma clang loop unroll(full) ++ for (int j = 0; j < H; j += 4) { ++ *(device half4*)(o + p0 + j) = half4(half(w[j]), half(w[j + 1]), half(w[j + 2]), half(w[j + 3])); ++ *(device half4*)(o + p1 + j) = half4(half(w[H + j]), half(w[H + j + 1]), half(w[H + j + 2]), half(w[H + j + 3])); + } + } + +@@ -61,9 +85,17 @@ kernel void gemm_f16(device half* A [[buffer(0)]], device half* B [[buffer(1)]], + + #define DQ(F) template [[host_name("dq_f16_" #F)]] kernel decltype(dq_f16) dq_f16; + #define DQ_(F) DQ(F) +-DQ_(T_F32) DQ_(T_BF16) DQ_(T_Q8_0) DQ_(T_Q4_0) DQ_(T_Q5_0) DQ_(T_IQ4_NL) DQ_(T_IQ4_XS) DQ_(T_Q2_0) DQ_(T_IQ3_S) DQ_(T_Q3_K) ++DQ_(T_F32) DQ_(T_BF16) DQ_(T_Q8_0) DQ_(T_Q4_0) DQ_(T_Q5_0) DQ_(T_Q5_1) DQ_(T_IQ4_NL) DQ_(T_IQ4_XS) DQ_(T_Q2_0) DQ_(T_IQ3_S) DQ_(T_Q3_K) + DQ_(T_Q4_K) DQ_(T_Q5_K) DQ_(T_Q6_K) + template [[host_name("gemm_f16_64_64_4")]] kernel decltype(gemm_f16<64, 64, 4>) gemm_f16<64, 64, 4>; + template [[host_name("gemm_f16_128_64_4")]] kernel decltype(gemm_f16<128, 64, 4>) gemm_f16<128, 64, 4>; + template [[host_name("gemm_f16_64_128_4")]] kernel decltype(gemm_f16<64, 128, 4>) gemm_f16<64, 128, 4>; + template [[host_name("gemm_f16_64_32_4")]] kernel decltype(gemm_f16<64, 32, 4>) gemm_f16<64, 32, 4>; ++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>; ++template [[host_name("gemm_f16_16_64_2")]] kernel decltype(gemm_f16<16, 64, 2>) gemm_f16<16, 64, 2>; ++template [[host_name("gemm_f16_32_64_2")]] kernel decltype(gemm_f16<32, 64, 2>) gemm_f16<32, 64, 2>; ++template [[host_name("gemm_f16_32_128_2")]] kernel decltype(gemm_f16<32, 128, 2>) gemm_f16<32, 128, 2>; ++template [[host_name("gemm_f16_64_32_2")]] kernel decltype(gemm_f16<64, 32, 2>) gemm_f16<64, 32, 2>; ++template [[host_name("gemm_f16_16_96_2")]] kernel decltype(gemm_f16<16, 96, 2>) gemm_f16<16, 96, 2>; +diff --git a/metal/kernels/gdn.metal b/metal/kernels/gdn.metal +index 516d89a4..c0570674 100644 +--- a/metal/kernels/gdn.metal ++++ b/metal/kernels/gdn.metal +@@ -27,6 +27,10 @@ struct GdnArgs { + // `plane` floats between planes) and the other planes and the per-16 sums are written too + }; + ++constant float gdn_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, ++ 0.00006103515625f}; ++constant float gdn_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; ++ + // The kernel is launched with IB * 128 threads (512 for gdn_step, 1024 for gdn_step8). Without the attribute the compiler + // is free to use as many registers as the 128-row state needs and a pipeline can end up with + // maxTotalThreadsPerThreadgroup below the launch size: an M2 Max gave 448 for the GDN_NTOK >= 2 specializations, and +@@ -45,10 +49,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 - 1][S]; ++ threadgroup float part[8], 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,109 +65,131 @@ 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) { ++ #pragma clang loop unroll(full) + 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 float4 cwv = *(device const float4*)(conv_w + ch * 4); ++ cw[0] = cwv.x; cw[1] = cwv.y; cw[2] = cwv.z; cw[3] = cwv.w; + } ++ const float gam = ib == 0 ? gamma[j] : 0.0f; ++ const float dth = tid == 384 ? dt[h] : 0.0f; ++ const float ssm_ah = tid == 384 ? ssm_a[h] : 0.0f; + float st[RPT]; +- for (int r = 0; r < RPT; ++r) st[r] = state[((i0 + r) * hv + h) * S + j]; ++ const int st_stride = hv * S; ++ const int st_off = (i0 * hv + h) * S + j; ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < RPT; ++r) st[r] = (state + st_off)[r * st_stride]; + + 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; ++ #pragma clang loop unroll(full) + for (int k = 0; k < 4; ++k) s = fma(vals[k], cw[k], s); + s = s + 0.0f; + cv = s / (1.0f + exp(-s)); + hst[0] = hst[1]; + hst[1] = hst[2]; + hst[2] = in; ++ if (t + 1 == commit_at && (tid >= 256 || h < int(a.h_k))) { ++ #pragma clang loop unroll(full) ++ for (int k = 0; k < 3; ++k) hist_out[ch * 3 + k] = hst[k]; ++ } + } +- // 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) { +- const float v = alpha[t * hv + h] + dt[h]; ++ in = in_next; ++ // raw q/k/v to shared memory + L2 norm partials of q (sg 0..3) and k (sg 4..7); gate and beta on idle lanes ++ if (tid < 256) { ++ const float ssq = simd_sum(cv * cv); ++ if (lane == 0) part[sg] = ssq; ++ if (tid < 128) qs[tid] = cv; ++ else ks[tid - 128] = cv; ++ } else if (tid < 384) { ++ vs[tid - 256] = cv; ++ } else if (tid == 384) { ++ const float v = alpha[t * hv + h] + dth; + // 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]); ++ gb[0] = exp(sp * ssm_ah); ++ } else if (tid == 416) { + gb[1] = 1.0f / (1.0f + exp(-beta[t * hv + h])); + } + threadgroup_barrier(mem_flags::mem_threadgroup); +- if (tid < 256) { +- const int base = tid < 128 ? 0 : 4; +- const float tot = (part[base] + part[base + 1]) + (part[base + 2] + part[base + 3]); +- const float rs = rsqrt(tot / float(S) + a.eps / float(S)); +- const float nv = (rs * cv) * scale; +- if (tid < 128) qs[tid] = nv; +- else ks[tid - 128] = nv; +- } else if (tid < 384) { +- vs[tid - 256] = cv; +- } +- threadgroup_barrier(mem_flags::mem_threadgroup); +- const float g = gb[0], be = gb[1]; ++ const float tot_k = (part[4] + part[5]) + (part[6] + part[7]); ++ const float rsk = rsqrt(tot_k / float(S) + a.eps / float(S)); ++ // the simdgroup's RPT rows of k and q one a lane, normalized on the fly and broadcast from the lane that holds them ++ const int qk_idx = i0 + (lane & (RPT - 1)); ++ const float kl = (rsk * ks[qk_idx]) * scale; + // kv_j + float kv = 0.0f; +- for (int r = 0; r < RPT; ++r) kv = fma(st[r], ks[i0 + r], kv); ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < RPT; ++r) kv = fma(st[r], simd_broadcast(kl, ushort(r)), kv); + red[ib][j] = kv; ++ const float tot_q = (part[0] + part[1]) + (part[2] + part[3]); ++ const float rsq = rsqrt(tot_q / float(S) + a.eps / float(S)); ++ const float ql = (rsq * qs[qk_idx]) * scale; ++ const float g = gb[0], be = gb[1], vj = vs[j]; + threadgroup_barrier(mem_flags::mem_threadgroup); + kv = 0.0f; ++ #pragma clang loop unroll(full) + for (int k = 0; k < IB; ++k) kv += red[k][j]; +- const float delta = (vs[j] - g * kv) * be; ++ const float delta = (vj - g * kv) * be; + float o = 0.0f; ++ #pragma clang loop unroll(full) + 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); ++ } ++ if (ib > 0) red2[ib - 1][j] = o; ++ if (t + 1 == commit_at) { ++ // the state and the conv history after token t (registers: st after this token's update, hst after its ++ // shift); the state is read only before the loop, so in place is safe. q/k history: the threadgroups ++ // h < h_k own their q/k head (h == hk); v history: every head its own - to hist_out, NOT in place: a q/k ++ // channel is read by the three threadgroups h % h_k at their start, and a threadgroup may start after its ++ // owner has finished (seen on the 10-core M4); the caller swaps the two ++ device float* st_base = state + st_off; ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < RPT; ++r) st_base[r * st_stride] = st[r]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); +- red[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]; +- o = o * scale; +- const float osq = ib == 0 ? simd_sum(o * o) : 0.0f; +- if (ib == 0 && lane == 0) part[cb] = osq; ++ // the readout and the output norm: column j's o (only ib == 0 computes readout and writes y) ++ if (ib == 0) { ++ #pragma clang loop unroll(full) ++ for (int k = 0; k < IB - 1; ++k) o += red2[k][j]; ++ o = o * scale; ++ const float osq = simd_sum(o * o); ++ if (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, +- 0.00006103515625f}; +- const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + device float* raw = y + t * hv * S + h * S + j; + const int jj = j & 15; +- raw[a.plane] = yv * q2s[jj & 7]; +- raw[2 * a.plane] = yv * p4s[jj & 3]; ++ raw[a.plane] = yv * gdn_q2s[jj & 7]; ++ raw[2 * a.plane] = yv * gdn_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; ++ #pragma clang loop unroll(full) ++ 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) { +- // the state and the conv history after token t (registers: st after this token's update, hst after its +- // shift); the state is read only before the loop, so in place is safe. q/k history: the threadgroups +- // h < h_k own their q/k head (h == hk); v history: every head its own - to hist_out, NOT in place: a q/k +- // channel is read by the three threadgroups h % h_k at their start, and a threadgroup may start after its +- // owner has finished (seen on the 10-core M4); the caller swaps the two +- for (int r = 0; r < RPT; ++r) state[((i0 + r) * hv + h) * S + j] = st[r]; +- 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 b2d530b1..c2c40200 100644 +--- a/metal/kernels/gdn_prompt.metal ++++ b/metal/kernels/gdn_prompt.metal +@@ -25,67 +25,161 @@ struct GdnPArgs { + + static inline float gp_conv(device const float* qkv, device const float* hist, device const float* cw, int ch, int t, + int channels) { ++ const float4 w = *(device const float4*)(cw + ch * 4); + float s = 0.0f; ++#pragma clang loop unroll(full) + for (int k = 0; k < 4; ++k) { + const int src = t - 3 + k; + const float x = src >= 0 ? qkv[src * channels + ch] : hist[ch * 3 + (src + 3)]; +- s = fma(x, cw[ch * 4 + k], s); ++ s = fma(x, w[k], s); + } + s = s + 0.0f; + return s / (1.0f + exp(-s)); + } + +-// grid (h_k + h_v / 2, T), 256 threads. x < h_k: q (threads 0..127) and k (128..255) of key head x; else value heads +-// 2 (x - h_k) and + 1. ++// grid ((h_k + h_v) / 8, T) = (8, T) or (h_k + h_v / 2, T) = (40, T), 256 threads: when grid_x == (h_k + h_v) / 8, ++// each simdgroup (32 lanes) owns one key head (x < h_k / 8: q, k and kq in registers) or one value head (else: vc and ++// gb), with zero threadgroup memory and zero barriers; grid_x == h_k + h_v / 2 keeps the legacy grid. + kernel void gdn_pre(device const float* qkv [[buffer(0)]], device const float* hist [[buffer(1)]], + device const float* conv_w [[buffer(2)]], device const float* alpha [[buffer(3)]], + device const float* beta [[buffer(4)]], device const float* dt [[buffer(5)]], + device const float* ssm_a [[buffer(6)]], device float* qn [[buffer(7)]], device float* kn [[buffer(8)]], + device float* vc [[buffer(9)]], device float2* gb [[buffer(10)]], device float* kq [[buffer(11)]], + device float* hist_out [[buffer(12)]], constant GdnPArgs& a [[buffer(13)]], +- uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], +- ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { ++ uint3 tg [[threadgroup_position_in_grid]], uint3 tgs [[threadgroups_per_grid]], ++ ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], ++ ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128; +- threadgroup float part[8]; +- threadgroup float prod[S]; + const int x = int(tg.x), t = int(tg.y), hk_n = int(a.h_k), hv = int(a.h_v), C = int(a.channels), T = int(a.n_tok); +- int ch; +- if (x < hk_n) ch = tid < 128 ? x * S + tid : hk_n * S + x * S + (tid - 128); +- else ch = 2 * hk_n * S + (2 * (x - hk_n) + (tid >> 7)) * S + (tid & 127); +- const float cv = gp_conv(qkv, hist, conv_w, ch, t, C); +- if (t == T - 1) +- for (int k = 0; k < 3; ++k) { +- const int src = T - 3 + k; +- hist_out[ch * 3 + k] = src >= 0 ? qkv[src * C + ch] : hist[ch * 3 + (src + 3)]; ++ const bool packed = int(tgs.x) * 8 == hk_n + hv; ++ // alpha and beta as the halves of one [T][2 h_v] GEMM (beta == alpha + h_v), else each [T][h_v] ++ const int ab = beta == alpha + hv ? 2 * hv : hv; ++ const int k_tgs = packed ? hk_n / 8 : hk_n; ++ if (x < k_tgs) { ++ const int kh = packed ? x * 8 + int(sg) : (sg == 0 ? x : hk_n); ++ if (kh >= hk_n) return; ++ float cv_q[4], cv_k[4]; ++ if (t >= 3) { ++ device const float* pq = qkv + (t - 3) * C + kh * S + int(lane); ++ device const float* pk = qkv + (t - 3) * C + (hk_n + kh) * S + int(lane); ++#pragma clang loop unroll(full) ++ for (int m = 0; m < 4; ++m) { ++ const int d = int(lane) + m * 32; ++ const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; ++ const float4 wq = *(device const float4*)(conv_w + ch_q * 4); ++ const float4 wk = *(device const float4*)(conv_w + ch_k * 4); ++ device const float* pqm = pq + m * 32; ++ device const float* pkm = pk + m * 32; ++ float sq = fma(pqm[0], wq.x, 0.0f); ++ sq = fma(pqm[C], wq.y, sq); ++ sq = fma(pqm[2 * C], wq.z, sq); ++ sq = fma(pqm[3 * C], wq.w, sq) + 0.0f; ++ cv_q[m] = sq / (1.0f + exp(-sq)); ++ float sk = fma(pkm[0], wk.x, 0.0f); ++ sk = fma(pkm[C], wk.y, sk); ++ sk = fma(pkm[2 * C], wk.z, sk); ++ sk = fma(pkm[3 * C], wk.w, sk) + 0.0f; ++ cv_k[m] = sk / (1.0f + exp(-sk)); ++ } ++ if (t == T - 1) ++ for (int m = 0; m < 4; ++m) { ++ const int d = int(lane) + m * 32; ++ const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; ++ for (int k = 0; k < 3; ++k) { ++ hist_out[ch_q * 3 + k] = pq[m * 32 + (k + 1) * C]; ++ hist_out[ch_k * 3 + k] = pk[m * 32 + (k + 1) * C]; ++ } ++ } ++ } else { ++#pragma clang loop unroll(full) ++ for (int m = 0; m < 4; ++m) { ++ const int d = int(lane) + m * 32; ++ const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; ++ cv_q[m] = gp_conv(qkv, hist, conv_w, ch_q, t, C); ++ cv_k[m] = gp_conv(qkv, hist, conv_w, ch_k, t, C); ++ } ++ if (t == T - 1) ++ for (int m = 0; m < 4; ++m) { ++ const int d = int(lane) + m * 32; ++ const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d; ++ for (int k = 0; k < 3; ++k) { ++ const int src = T - 3 + k; ++ hist_out[ch_q * 3 + k] = src >= 0 ? qkv[src * C + ch_q] : hist[ch_q * 3 + (src + 3)]; ++ hist_out[ch_k * 3 + k] = src >= 0 ? qkv[src * C + ch_k] : hist[ch_k * 3 + (src + 3)]; ++ } ++ } + } +- if (x < hk_n) { +- const float ssq = simd_sum(cv * cv); +- if (lane == 0) part[sg] = ssq; +- threadgroup_barrier(mem_flags::mem_threadgroup); +- const int base = tid < 128 ? 0 : 4; +- const float tot = (part[base] + part[base + 1]) + (part[base + 2] + part[base + 3]); +- const float rs = rsqrt(tot / float(S) + a.eps / float(S)); +- const float nv = (rs * cv) * (1.0f / sqrt(float(S))); +- if (tid < 128) { +- qn[(t * hk_n + x) * S + tid] = nv; +- prod[tid] = nv; ++ const float4 sq = simd_sum(float4(cv_q[0] * cv_q[0], cv_q[1] * cv_q[1], cv_q[2] * cv_q[2], cv_q[3] * cv_q[3])); ++ const float4 sk = simd_sum(float4(cv_k[0] * cv_k[0], cv_k[1] * cv_k[1], cv_k[2] * cv_k[2], cv_k[3] * cv_k[3])); ++ const float tot_q = (sq.x + sq.y) + (sq.z + sq.w), tot_k = (sk.x + sk.y) + (sk.z + sk.w); ++ const float rs_q = rsqrt(tot_q / float(S) + a.eps / float(S)); ++ const float rs_k = rsqrt(tot_k / float(S) + a.eps / float(S)); ++ const float scale = 1.0f / sqrt(float(S)); ++ const int base = (t * hk_n + kh) * S + int(lane); ++ float4 prod; ++#pragma clang loop unroll(full) ++ for (int m = 0; m < 4; ++m) { ++ const float nq = (rs_q * cv_q[m]) * scale, nk = (rs_k * cv_k[m]) * scale; ++ qn[base + m * 32] = nq; ++ kn[base + m * 32] = nk; ++ prod[m] = nk * nq; + } +- threadgroup_barrier(mem_flags::mem_threadgroup); +- if (tid >= 128) { +- kn[(t * hk_n + x) * S + tid - 128] = nv; +- const float p = simd_sum(nv * prod[tid - 128]); +- if (lane == 0) part[sg] = p; // simdgroups 4..7 (q's slots 0..3 are read already) ++ const float4 pv = simd_sum(prod); ++ if (lane == 0) kq[t * hk_n + kh] = (pv.x + pv.y) + (pv.z + pv.w); ++ } else if (packed) { ++ const int h = (x - k_tgs) * 8 + int(sg); ++ if (t >= 3) { ++ const int d = int(lane) * 4, ch = (2 * hk_n + h) * S + d; ++ device const float4* p = (device const float4*)(qkv + (t - 3) * C + ch); ++ device const float4* w = (device const float4*)(conv_w + ch * 4); ++ const float4 w0 = w[0], w1 = w[1], w2 = w[2], w3 = w[3]; ++ float4 cv = fma(p[0], float4(w0.x, w1.x, w2.x, w3.x), float4(0.0f)); ++ cv = fma(p[C / 4], float4(w0.y, w1.y, w2.y, w3.y), cv); ++ cv = fma(p[C / 2], float4(w0.z, w1.z, w2.z, w3.z), cv); ++ cv = fma(p[3 * C / 4], float4(w0.w, w1.w, w2.w, w3.w), cv); ++ cv = cv + 0.0f; ++ cv = cv / (1.0f + exp(-cv)); ++ *((device float4*)(vc + (t * hv + h) * S + d)) = cv; ++ if (t == T - 1) ++ for (int k = 0; k < 3; ++k) { ++ const float4 v = p[(k + 1) * C / 4]; ++ for (int m = 0; m < 4; ++m) hist_out[(ch + m) * 3 + k] = v[m]; ++ } ++ } else { ++#pragma clang loop unroll(full) ++ for (int m = 0; m < 4; ++m) { ++ const int d = int(lane) + m * 32; ++ const int ch = 2 * hk_n * S + h * S + d; ++ const float cv = gp_conv(qkv, hist, conv_w, ch, t, C); ++ if (t == T - 1) ++ for (int k = 0; k < 3; ++k) { ++ const int src = T - 3 + k; ++ hist_out[ch * 3 + k] = src >= 0 ? qkv[src * C + ch] : hist[ch * 3 + (src + 3)]; ++ } ++ vc[(t * hv + h) * S + d] = cv; ++ } ++ } ++ if (lane == 0) { ++ const float v = alpha[t * ab + h] + dt[h]; ++ 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[t * hv + h] = float2(exp(sp * ssm_a[h]), 1.0f / (1.0f + exp(-beta[t * ab + h]))); + } +- threadgroup_barrier(mem_flags::mem_threadgroup); +- if (tid == 0) kq[t * hk_n + x] = (part[4] + part[5]) + (part[6] + part[7]); + } else { +- const int h = 2 * (x - hk_n) + (tid >> 7); +- vc[(t * hv + h) * S + (tid & 127)] = cv; +- if ((tid & 127) == 0) { +- const float v = alpha[t * hv + h] + dt[h]; ++ const int h = 2 * (x - k_tgs) + (tid >> 7), d = tid & 127; ++ const int ch = 2 * hk_n * S + h * S + d; ++ const float cv = gp_conv(qkv, hist, conv_w, ch, t, C); ++ if (t == T - 1) ++ for (int k = 0; k < 3; ++k) { ++ const int src = T - 3 + k; ++ hist_out[ch * 3 + k] = src >= 0 ? qkv[src * C + ch] : hist[ch * 3 + (src + 3)]; ++ } ++ vc[(t * hv + h) * S + d] = cv; ++ if (d == 0) { ++ const float v = alpha[t * ab + h] + dt[h]; + 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[t * hv + h] = float2(exp(sp * ssm_a[h]), 1.0f / (1.0f + exp(-beta[t * hv + h]))); ++ gb[t * hv + h] = float2(exp(sp * ssm_a[h]), 1.0f / (1.0f + exp(-beta[t * ab + h]))); + } + } + } +@@ -156,55 +250,92 @@ constant int GDN_PROBE [[function_constant(6)]]; + constant int gdn_probe = is_function_constant_defined(GDN_PROBE) ? GDN_PROBE : 0; + + // the rows of a column split over IB simdgroups (128 / IB rows a lane, fewer registers): grid (h_v * 4), IB * 32 +-// threads; lane = column (32 columns a threadgroup); per token the IB partial dots are exchanged through threadgroup +-// memory (double-buffered: one barrier a token) ++// threads; lane = column (32 columns a threadgroup); when TB == 0 (single-pass), each simdgroup loads its R rows of q ++// and k across its lanes and broadcasts them with simd_broadcast (one initial dot pass at t = 0, one barrier a token, ++// only red[2][IB][32] in threadgroup memory); TB > 0 keeps the TB-token threadgroup-staged reference. + template +-kernel void gdn_scan_r(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]], +- device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]], ++kernel void gdn_scan_r(device float* state [[buffer(0)]], device const float* qn [[buffer(1)]], ++ device const float* kn [[buffer(2)]], device const float* vc [[buffer(3)]], + device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]], + device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]], + 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; +- threadgroup float4 qs[TB][32], ks[TB][32]; ++ constexpr int S = 128, R = S / IB; + threadgroup float2 red[2][IB][32]; +- threadgroup float vs[TB][32]; +- threadgroup float4 gk[TB]; // gate, beta, kq +- const int h = int(tg.x) / 4, hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k); ++ const int g_idx = int(tg.x) / 4, hv = int(a.h_v), hkn = int(a.h_k); ++ const int hk = (hv == 3 * hkn) ? (g_idx / 3) : (g_idx % hkn); ++ const int h = (hv == 3 * hkn) ? (hk + (g_idx % 3) * hkn) : g_idx; + const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * R; + const int T = int(a.n_tok); + const float scale = 1.0f / sqrt(float(S)); ++ const int rs = hv * S, s0 = (i0 * hv + h) * S + j; + float st[R]; + #pragma clang loop unroll(full) +- for (int i = 0; i < R; ++i) st[i] = state[((i0 + i) * hv + h) * S + j]; +- for (int t0 = 0; t0 < T; t0 += TB) { +- const int nb = min(TB, T - t0); +- threadgroup_barrier(mem_flags::mem_threadgroup); +- for (int idx = tid; idx < nb * 32; idx += IB * 32) { +- const int tt = idx >> 5, w = idx & 31; +- qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w]; +- ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w]; +- vs[tt][w] = vc[((t0 + tt) * hv + h) * S + (j - lane) + w]; +- if (w == 0) { +- const float2 g_b = gb[(t0 + tt) * hv + h]; +- gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f); +- } ++ for (int i = 0; i < R; ++i) st[i] = (state + s0)[i * rs]; ++ if (TB == 0) { ++ constexpr int C = (R + 31) / 32, R1 = R / C; ++ const ushort lr = lane & (R1 - 1); ++ const int ks = hkn * S, k0 = hk * S + i0 + int(lr), v0 = h * S + j; ++ float kl[C], ql[C]; ++#pragma clang loop unroll(full) ++ for (int c = 0; c < C; ++c) { ++ kl[c] = kn[k0 + c * 32]; ++ ql[c] = qn[k0 + c * 32]; + } +- threadgroup_barrier(mem_flags::mem_threadgroup); +- for (int tt = 0; tt < nb; ++tt) { +- const int t = t0 + tt, slot = t & 1; +- float ka = 0.0f, qa = 0.0f; ++ float2 g_b = gb[h]; ++ float v = vc[v0]; ++ 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); ++ for (int c = 0; c < C; ++c) ++#pragma clang loop unroll(full) ++ for (int r = 0; r < R1; ++r) { ++ ka = fma(st[c * 32 + r], simd_broadcast(kl[c], ushort(r)), ka); ++ qa = fma(st[c * 32 + r], simd_broadcast(ql[c], ushort(r)), qa); + } +- const float4 g_b = gk[tt]; +- const float kqv = g_b.z; +- const float v = vs[tt][lane]; ++ for (int t = 0; t < T - 1; ++t) { ++ const int slot = t & 1, tn = t + 1; ++ const int kn_off = tn * ks + k0, vn_off = tn * rs + v0; ++ float kn_l[C], qn_l[C]; ++#pragma clang loop unroll(full) ++ for (int c = 0; c < C; ++c) { ++ kn_l[c] = kn[kn_off + c * 32]; ++ qn_l[c] = qn[kn_off + c * 32]; ++ } ++ const float2 g_b_n = gb[tn * hv + h]; ++ const float v_n = vc[vn_off]; ++ const float kqv = sg == 0 ? kq[t * hkn + hk] : 0.0f; ++ float2 s = float2(ka, qa); ++ if (gdn_probe != 1) { ++ red[slot][sg][lane] = s; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ s = red[slot][0][lane]; ++#pragma clang loop unroll(full) ++ for (int b = 1; b < IB; ++b) s += red[slot][b][lane]; ++ } ++ const float g = g_b.x; ++ const float delta = (v - g * s.x) * g_b.y; ++ if (sg == 0) o[t * rs + v0] = fma(g, s.y, delta * kqv) * scale; ++ ka = 0.0f; ++ qa = 0.0f; ++ if (gdn_probe != 2) ++#pragma clang loop unroll(full) ++ for (int c = 0; c < C; ++c) ++#pragma clang loop unroll(full) ++ for (int r = 0; r < R1; ++r) { ++ const float kr = simd_broadcast(kl[c], ushort(r)); ++ const float knr = simd_broadcast(kn_l[c], ushort(r)), qnr = simd_broadcast(qn_l[c], ushort(r)); ++ st[c * 32 + r] = fma(g, st[c * 32 + r], kr * delta); ++ ka = fma(st[c * 32 + r], knr, ka); ++ qa = fma(st[c * 32 + r], qnr, qa); ++ } ++#pragma clang loop unroll(full) ++ for (int c = 0; c < C; ++c) kl[c] = kn_l[c]; ++ g_b = g_b_n; ++ v = v_n; ++ } ++ if (T > 0) { ++ const int t = T - 1, slot = t & 1; ++ const float kqv = sg == 0 ? kq[t * hkn + hk] : 0.0f; + float2 s = float2(ka, qa); + if (gdn_probe != 1) { + red[slot][sg][lane] = s; +@@ -215,20 +346,88 @@ 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; ++ if (sg == 0) o[t * rs + v0] = fma(g, s.y, delta * kqv) * scale; + if (gdn_probe != 2) ++#pragma clang loop unroll(full) ++ for (int c = 0; c < C; ++c) ++#pragma clang loop unroll(full) ++ for (int r = 0; r < R1; ++r) { ++ const float kr = simd_broadcast(kl[c], ushort(r)); ++ st[c * 32 + r] = fma(g, st[c * 32 + r], kr * delta); ++ } ++ } ++ } else { ++ constexpr int R4 = R / 4, TBS = TB > 0 ? TB : 1; ++ constexpr bool KC = R4 <= 8; ++ threadgroup float4 qs[TBS][32], ks[TBS][32]; ++ threadgroup float vs[TBS][32]; ++ threadgroup float4 gk[TBS]; ++ device const float4* qn4_p = (device const float4*)qn; ++ device const float4* kn4_p = (device const float4*)kn; ++ for (int t0 = 0; t0 < T; t0 += TBS) { ++ const int nb = min(TBS, T - t0); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ for (int idx = tid; idx < nb * 32; idx += IB * 32) { ++ const int tt = idx >> 5, w = idx & 31; ++ qs[tt][w] = qn4_p[((t0 + tt) * hkn + hk) * 32 + w]; ++ ks[tt][w] = kn4_p[((t0 + tt) * hkn + hk) * 32 + w]; ++ vs[tt][w] = vc[((t0 + tt) * hv + h) * S + (j - lane) + w]; ++ if (w == 0) { ++ const float2 g_b = gb[(t0 + tt) * hv + h]; ++ gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f); ++ } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ float ka = 0.0f, qa = 0.0f; ++ float4 kc[R4]; + #pragma clang loop unroll(full) + for (int w = 0; w < R4; ++w) { +- const float4 k4 = ks[tt][i0 / 4 + w]; +- 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); ++ 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; ++ const float4 g_b = gk[tt]; ++ const float kqv = g_b.z; ++ const float v = vs[tt][lane]; ++ float2 s = float2(ka, qa); ++ if (gdn_probe != 1) { ++ red[slot][sg][lane] = s; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ s = red[slot][0][lane]; ++#pragma clang loop unroll(full) ++ for (int b = 1; b < IB; ++b) s += red[slot][b][lane]; ++ } ++ 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 = 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); ++ } + } + } + } ++ device float* sp = state + s0; + #pragma clang loop unroll(full) +- for (int i = 0; i < R; ++i) state[((i0 + i) * hv + h) * S + j] = st[i]; ++ for (int i = 0; i < R; ++i) sp[i * rs] = st[i]; + } + + // gdn_scan_r<8, 16> with the 16 state rows of a lane in four float4 variables (no array): a register-allocation check +@@ -294,33 +493,52 @@ kernel void gdn_scan_v(device float* state [[buffer(0)]], device const float4* q + sp[12 * rs] = s3.x; sp[13 * rs] = s3.y; sp[14 * rs] = s3.z; sp[15 * rs] = s3.w; + } + +-// grid (h_v, T), 128 threads ++// grid (h_v / 4, T) or (h_v, T), 128 threads: one simdgroup (32 lanes) per head h computes the RMSNorm and gated SiLU ++// over the 128 channels (4 channels a lane) with zero threadgroup memory and zero barriers. + kernel void gdn_post(device const float* o [[buffer(0)]], device const float* z [[buffer(1)]], + device const float* gamma [[buffer(2)]], device half* yh [[buffer(3)]], device float* y [[buffer(4)]], + constant GdnPArgs& a [[buffer(5)]], +- uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], ++ uint3 tg [[threadgroup_position_in_grid]], uint3 tgs [[threadgroups_per_grid]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { +- threadgroup float part[4]; +- const int h = int(tg.x), t = int(tg.y), hv = int(a.h_v); +- const uint i = uint(t * hv + h) * 128 + tid; +- const float ov = o[i]; +- const float ss = simd_sum(ov * ov); +- if (lane == 0) part[sg] = ss; +- threadgroup_barrier(mem_flags::mem_threadgroup); +- const float tot = (part[0] + part[1]) + (part[2] + part[3]); ++ const int hv = int(a.h_v), t = int(tg.y); ++ const int h = int(tgs.x) == hv ? (sg == 0 ? int(tg.x) : hv) : int(tg.x) * 4 + int(sg); ++ if (h >= hv) return; ++ const uint base = uint(t * hv + h) * 128 + lane; ++ device const float* o_row = o + base; ++ device const float* z_row = z + base; ++ device const float* g_row = gamma + lane; ++ device half* yh_row = yh + base; ++ device float* y_row = y + base; ++ const float ov0 = o_row[0], ov1 = o_row[32], ov2 = o_row[64], ov3 = o_row[96]; ++ // the gate and gamma do not depend on the norm: computed while the simd_sums run ++ float sig[4]; ++#pragma clang loop unroll(full) ++ for (uint k = 0; k < 4; ++k) sig[k] = 1.0f / (1.0f + exp(-z_row[k * 32])); ++ const float gv[4] = {g_row[0], g_row[32], g_row[64], g_row[96]}; ++ const float4 pv = simd_sum(float4(ov0 * ov0, ov1 * ov1, ov2 * ov2, ov3 * ov3)); ++ const float tot = (pv.x + pv.y) + (pv.z + pv.w); + const float rs = rsqrt(tot / 128.0f + a.eps); +- const float yv = ((rs * ov) * gamma[tid]) * (1.0f / (1.0f + exp(-z[i]))); +- yh[i] = half(yv); +- if (a.f32) y[i] = yv; ++ const float ov[4] = {ov0, ov1, ov2, ov3}; ++ float yv[4]; ++#pragma clang loop unroll(full) ++ for (uint k = 0; k < 4; ++k) { ++ yv[k] = ((rs * ov[k]) * gv[k]) * sig[k]; ++ yh_row[k * 32] = half(yv[k]); ++ } ++ if (a.f32) { ++#pragma clang loop unroll(full) ++ for (uint k = 0; k < 4; ++k) y_row[k * 32] = yv[k]; ++ } + } + + template [[host_name("gdn_scan_1")]] kernel decltype(gdn_scan_t<1, 16>) gdn_scan_t<1, 16>; + template [[host_name("gdn_scan_2")]] kernel decltype(gdn_scan_t<2, 16>) gdn_scan_t<2, 16>; + template [[host_name("gdn_scan_4")]] kernel decltype(gdn_scan_t<4, 16>) gdn_scan_t<4, 16>; +-template [[host_name("gdn_scanr_2")]] kernel decltype(gdn_scan_r<2, 16>) gdn_scan_r<2, 16>; +-template [[host_name("gdn_scanr_4")]] kernel decltype(gdn_scan_r<4, 16>) gdn_scan_r<4, 16>; +-template [[host_name("gdn_scanr_8")]] kernel decltype(gdn_scan_r<8, 16>) gdn_scan_r<8, 16>; ++template [[host_name("gdn_scanr_2")]] kernel decltype(gdn_scan_r<2, 0>) gdn_scan_r<2, 0>; ++template [[host_name("gdn_scanr_4")]] kernel decltype(gdn_scan_r<4, 0>) gdn_scan_r<4, 0>; ++template [[host_name("gdn_scanr_8")]] kernel decltype(gdn_scan_r<8, 0>) gdn_scan_r<8, 0>; + template [[host_name("gdn_scanr_4_8")]] kernel decltype(gdn_scan_r<4, 8>) gdn_scan_r<4, 8>; ++template [[host_name("gdn_scanr_4_16")]] kernel decltype(gdn_scan_r<4, 16>) gdn_scan_r<4, 16>; + + // ---- the chunked form (chunks of GC = 32 tokens), the WY / UT form of the gated delta rule. With G_t the product + // of the chunk's gates up to t (log space: lg_t) and S0 the state entering the chunk: +diff --git a/metal/kernels/gemv.metal b/metal/kernels/gemv.metal +index 6d4ed3a5..49b7b621 100644 +--- a/metal/kernels/gemv.metal ++++ b/metal/kernels/gemv.metal +@@ -21,6 +21,85 @@ constant int NTOK [[function_constant(0)]]; // tokens of the window (1..8): th + constant int GEMV_PROBE [[function_constant(9)]]; + constant int gemv_probe = is_function_constant_defined(GEMV_PROBE) ? GEMV_PROBE : 0; + ++// Wider raw-dot windows and prepared Q2 at NR=1 did not improve on M5. ++template ++static inline bool gemv_direct_x() { ++ return F == T_Q8_0 && (NTOK == 1 || (NTOK == 2 && LPR == 16 && NR == 1)); ++} ++ ++// Keep unit_combine's four FMA chains without materialising the raw x slice. ++template ++static inline void gemv_raw_combine(device const float* xp, thread const float (&q)[NR][UNIT], ++ thread const float4* sc, thread float (&acc)[NR][8], int t) { ++ constexpr int H = UNIT / 2; ++ device const float4* x = (device const float4*)xp; ++ float s[NR][4]; ++#pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) s[r][0] = s[r][1] = s[r][2] = s[r][3] = 0.0f; ++#pragma clang loop unroll(full) ++ for (int j = 0; j < H; j += 4) { ++ const float4 a = x[j / 4], b = x[(H + j) / 4]; ++#pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) { ++ s[r][0] = fma(q[r][j], a.x, s[r][0]); ++ s[r][1] = fma(q[r][j + 1], a.y, s[r][1]); ++ s[r][2] = fma(q[r][H + j], b.x, s[r][2]); ++ s[r][3] = fma(q[r][H + j + 1], b.y, s[r][3]); ++ s[r][0] = fma(q[r][j + 2], a.z, s[r][0]); ++ s[r][1] = fma(q[r][j + 3], a.w, s[r][1]); ++ s[r][2] = fma(q[r][H + j + 2], b.z, s[r][2]); ++ s[r][3] = fma(q[r][H + j + 3], b.w, s[r][3]); ++ } ++ } ++#pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) ++ acc[r][t] += sc[r].x * (s[r][0] + s[r][1]) + sc[r].y * (s[r][2] + s[r][3]) - ++ (sc[r].z * 0.0f + sc[r].w * 0.0f); ++} ++ ++static inline float gemv_q2_raw_combine(ushort2 pk, float d, thread const float* xs) { ++ float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 8; j += 2) { ++ const uint2 q = uint2(pk) >> (2 * j); ++ const uint2 bits = ((q | (q << 14)) & 0x00030003u) | 0x64006400u; ++ const float4 v = float4(as_type(bits) - 1025.0h); ++ s0 = fma(v.x, xs[j], s0); ++ s1 = fma(v.y, xs[j + 1], s1); ++ s2 = fma(v.z, xs[8 + j], s2); ++ s3 = fma(v.w, xs[8 + j + 1], s3); ++ } ++ return d * (s0 + s1) + d * (s2 + s3) - 0.0f; ++} ++ ++// IQ4_XS kept as its 16 code bytes and scale (5 registers a row instead of 36), decoded inside the combine. M5, weights ++// from DRAM: 2560 -> 640 at 16,2,4,1 -1.5..-5% (1 token), -2..-3% (2..4); 3..4 tokens -1..-2% on 2560 -> 6144..12288 ++// (but 6144 -> 2560 at 16,2,4,1 3 tokens +1.4%); LPR 32 under 4 tokens +0.4..0.7%, so those keep the decoded unit. ++// IQ4_NL (18-byte blocks, ushort reads) measured no change and keeps the decoded unit. ++template ++static inline bool gemv_iq4_packed() { return LPR <= 16 || NTOK >= 4; } ++ ++static inline uint4 gemv_iq4xs_codes(device const uchar* row, int u, thread float& d) { ++ device const uchar* b = row + (u >> 3) * 136; ++ d = iq4xs_scale(b, u & 7); ++ return uint4(*(device const packed_uint4*)(b + 8 + 16 * (u & 7))); ++} ++ ++// unit_combine<32> of deq's q, decoding the codes as it goes: the same four FMA chains in the same order ++static inline float gemv_iq4_combine(uint4 pk, float d, thread const float* xs, threadgroup const uint* tab) { ++ float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 8; ++j) { ++ const uint v = pk[j >> 1] >> (16 * (j & 1)); ++ const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[(v >> 8) & 255]); ++ s0 = fma(float(w0.x), xs[2 * j], s0); ++ s1 = fma(float(w1.x), xs[2 * j + 1], s1); ++ s2 = fma(float(w0.y), xs[16 + 2 * j], s2); ++ s3 = fma(float(w1.y), xs[16 + 2 * j + 1], s3); ++ } ++ return d * (s0 + s1) + d * (s2 + s3) - 0.0f; ++} ++ + template + kernel void gemv(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], +@@ -29,8 +108,10 @@ kernel void gemv(device const uchar* W [[buffer(0)]], device const float* xp [[b + constexpr int UNIT = fmt::UNIT, RS = SG / KS, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + threadgroup float red[KS > 1 ? KS * RS * RG * NR * 8 : 1]; +- load_tab(tab, tid, SG * 32); +- threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (fmt::TAB > 1) { ++ load_tab(tab, tid, SG * 32); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } + const int ks = sg / RS, rs = sg % RS; + const int li = lane % LPR; + const int lrow0 = (rs * RG + lane / LPR) * NR; +@@ -38,18 +119,23 @@ kernel void gemv(device const uchar* W [[buffer(0)]], device const float* xp [[b + const int U = int(a.n_in) / UNIT; + const int u0 = ks * U / KS, u1 = (ks + 1) * U / KS; + float acc[NR][8]; ++#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) ++#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; ++#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) rp[r] = W + (ulong)min(row0 + r, int(a.n_out) - 1) * a.row_bytes; + const int s16n = int(a.n_in) / 16; + if (gemv_probe == 3) { + // rows outer, tokens inner: one row's decoded unit in registers at a time (x re-read per row from L1) + for (int u = u0 + li; u < u1; u += LPR) ++#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 sc; + Dec::run(rp[r], u, q, sc, tab); ++#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); +@@ -58,18 +144,58 @@ kernel void gemv(device const uchar* W [[buffer(0)]], device const float* xp [[b + } + } else + for (int u = u0 + li; u < u1; u += LPR) { +- float q[NR][UNIT]; +- float4 sc[NR]; +- for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); +- float xs[UNIT], xsum[2]; +- for (int t = 0; t < NTOK; ++t) { +- if (gemv_probe != 1 || t == 0) X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); +- if (gemv_probe != 2 || t == 0) +- for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); ++ if (NTOK == 1 && F == T_Q2_0R) { ++ ushort2 pk[NR]; ++ float d[NR]; ++#pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) { ++ device const uchar* b = rp[r] + (u >> 2) * 18; ++ device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); ++ pk[r] = ushort2(p[0], p[1]); ++ d[r] = rdh(b); ++ } ++ float xs[16]; ++ load16(xp + 16 * u, xs); ++#pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) acc[r][0] += gemv_q2_raw_combine(pk[r], d[r], xs); ++ } else if (F == T_IQ4_XS && gemv_iq4_packed()) { ++ uint4 pk[NR]; ++ float d[NR]; ++#pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) pk[r] = gemv_iq4xs_codes(rp[r], u, d[r]); ++#pragma clang loop unroll(full) ++ for (int t = 0; t < NTOK; ++t) { ++ float xs[32]; ++ load16(xp + t * int(a.x_stride) + 32 * u, xs); ++ load16(xp + t * int(a.x_stride) + 32 * u + 16, xs + 16); ++#pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) acc[r][t] += gemv_iq4_combine(pk[r], d[r], xs, tab); ++ } ++ } else { ++ float q[NR][UNIT]; ++ float4 sc[NR]; ++#pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); ++ float xs[UNIT], xsum[2]; ++#pragma clang loop unroll(full) ++ for (int t = 0; t < NTOK; ++t) { ++ if (gemv_direct_x()) { ++ if (gemv_probe != 2 || t == 0) ++ gemv_raw_combine(xp + (gemv_probe == 1 ? 0 : t) * int(a.x_stride) + u * UNIT, ++ q, sc, acc, t); ++ } else { ++ if (gemv_probe != 1 || t == 0) X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); ++ if (gemv_probe != 2 || t == 0) ++#pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); ++ } ++ } + } + } + constexpr int RPT = RS * RG * NR; ++#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) ++#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + if (KS == 1) { +@@ -81,12 +207,16 @@ kernel void gemv(device const uchar* W [[buffer(0)]], device const float* xp [[b + if (KS > 1) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ks == 0 && li == 0) ++#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) +- for (int t = 0; t < NTOK; ++t) { +- float v = 0.0f; +- for (int k = 0; k < KS; ++k) v += red[(k * RPT + lrow0 + r) * 8 + t]; +- if (row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = v; +- } ++ if (row0 + r < int(a.n_out)) ++#pragma clang loop unroll(full) ++ for (int t = 0; t < NTOK; ++t) { ++ float v = 0.0f; ++#pragma clang loop unroll(full) ++ for (int k = 0; k < KS; ++k) v += red[(k * RPT + lrow0 + r) * 8 + t]; ++ y[t * int(a.y_stride) + row0 + r] = v; ++ } + } + } + +@@ -114,10 +244,9 @@ kernel void gemv_pair(device const uchar* Wg [[buffer(0)]], device const float* + constexpr int TB = fmt::TAB > fmt::TAB ? fmt::TAB : fmt::TAB; + threadgroup uint tab1[fmt::TAB]; + threadgroup uint tab2[fmt::TAB]; +- (void)TB; +- load_tab(tab1, tid, SG * 32); +- load_tab(tab2, tid, SG * 32); +- threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (fmt::TAB > 1) load_tab(tab1, tid, SG * 32); ++ if (fmt::TAB > 1) load_tab(tab2, tid, SG * 32); ++ if (TB > 1) threadgroup_barrier(mem_flags::mem_threadgroup); + const int li = lane % LPR; + const int row = (int(tg.x) * SG + sg) * RG + lane / LPR; + const int rr = min(row, int(a.n_out) - 1); +@@ -125,11 +254,13 @@ kernel void gemv_pair(device const uchar* Wg [[buffer(0)]], device const float* + device const uchar* ru = Wu + ulong(rr) * rb_up; + const int s16n = int(a.n_in) / 16; + float ag[8], au[8]; +- for (int t = 0; t < 8; ++t) ag[t] = au[t] = 0.0f; ++#pragma clang loop unroll(full) ++ for (int t = 0; t < NTOK; ++t) ag[t] = au[t] = 0.0f; + for (int u = li; u < int(a.n_in) / U1; u += LPR) { + float q[U1]; + float4 sc; + Dec::run(rg, u, q, sc, tab1); ++#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[U1], xsum[2]; + X::load(xg + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); +@@ -140,12 +271,14 @@ kernel void gemv_pair(device const uchar* Wg [[buffer(0)]], device const float* + float q[U2]; + float4 sc; + Dec::run(ru, u, q, sc, tab2); ++#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[U2], xsum[2]; + X::load(xu + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + au[t] += unit_combine(q, sc, xs, xsum); + } + } ++#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float gv = lane_sum(ag[t]), uv = lane_sum(au[t]); + if (li == 0 && row < int(a.n_out)) h[t * int(a.y_stride) + row] = gv / (1.0f + exp(-gv)) * uv; +@@ -171,15 +304,20 @@ kernel void gemv_cat(device const uchar* W [[buffer(0)]], device const float* xp + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; +- load_tab(tab, tid, SG * 32); +- threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (fmt::TAB > 1) { ++ load_tab(tab, tid, SG * 32); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } + const int li = lane % LPR; + const int row0 = int(tg.x) * SG * RG * NR + (sg * RG + lane / LPR) * NR; + const int U = int(a.n_in) / UNIT, sp = int(n_split); + float acc[NR][8]; ++#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) ++#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; ++#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { + const int row = min(row0 + r, int(a.n_out) - 1); + rp[r] = row < sp ? W + (ulong)row * a.row_bytes : W2 + (ulong)(row - sp) * a.row_bytes; +@@ -188,14 +326,23 @@ kernel void gemv_cat(device const uchar* W [[buffer(0)]], device const float* xp + for (int u = li; u < U; u += LPR) { + float q[NR][UNIT]; + float4 sc[NR]; ++#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); ++#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { +- float xs[UNIT], xsum[2]; +- X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); +- for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); ++ if (gemv_direct_x()) { ++ gemv_raw_combine(xp + t * int(a.x_stride) + u * UNIT, q, sc, acc, t); ++ } else { ++ float xs[UNIT], xsum[2]; ++ X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); ++#pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); ++ } + } + } ++#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) ++#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + const int row = row0 + r; +@@ -209,6 +356,86 @@ kernel void gemv_cat(device const uchar* W [[buffer(0)]], device const float* xp + #define GVC(F, L, R, S) GVC_(F, L, R, S) + GVC(T_BF16, 32, 1, 4) GVC(T_BF16, 32, 1, 8) GVC(T_BF16, 16, 1, 4) GVC(T_BF16, 32, 2, 4) + GVC(T_F32, 32, 1, 4) GVC(T_F32, 32, 1, 8) GVC(T_F32, 16, 1, 4) GVC(T_F32, 32, 2, 4) ++GVC(T_Q8_0, 32, 1, 4) GVC(T_Q8_0, 32, 1, 8) GVC(T_Q8_0, 16, 1, 4) GVC(T_Q8_0, 32, 2, 4) GVC(T_Q8_0, 8, 1, 8) ++ ++// gemv_catk: gemv_cat with gemv's K split (KS simdgroups share a row, partial sums added in gemv's fixed order), so ++// the outputs are bitwise gemv_____'s. (The MTP draft step's k + v and shared gate + up: Q8_0 ++// 2560 -> 512 + 512 / 640 + 640 at 32,1,4,4 on the M5.) Buffers as gemv_cat. ++template ++kernel void gemv_catk(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], ++ device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], ++ device const uchar* W2 [[buffer(5)]], device float* y2 [[buffer(6)]], constant uint& n_split [[buffer(7)]], ++ uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], ++ ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { ++ constexpr int UNIT = fmt::UNIT, RS = SG / KS, RG = 32 / LPR; ++ threadgroup uint tab[fmt::TAB]; ++ threadgroup float red[KS * RS * RG * NR * 8]; ++ if (fmt::TAB > 1) { ++ load_tab(tab, tid, SG * 32); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } ++ const int ks = sg / RS, rs = sg % RS; ++ const int li = lane % LPR; ++ const int lrow0 = (rs * RG + lane / LPR) * NR; ++ const int row0 = int(tg.x) * RS * RG * NR + lrow0; ++ const int U = int(a.n_in) / UNIT, sp = int(n_split); ++ const int u0 = ks * U / KS, u1 = (ks + 1) * U / KS; ++ float acc[NR][8]; ++#pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) ++#pragma clang loop unroll(full) ++ for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; ++ device const uchar* rp[NR]; ++#pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) { ++ const int row = min(row0 + r, int(a.n_out) - 1); ++ rp[r] = row < sp ? W + (ulong)row * a.row_bytes : W2 + (ulong)(row - sp) * a.row_bytes; ++ } ++ const int s16n = int(a.n_in) / 16; ++ for (int u = u0 + li; u < u1; u += LPR) { ++ float q[NR][UNIT]; ++ float4 sc[NR]; ++#pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); ++#pragma clang loop unroll(full) ++ for (int t = 0; t < NTOK; ++t) { ++ if (gemv_direct_x()) { ++ gemv_raw_combine(xp + t * int(a.x_stride) + u * UNIT, q, sc, acc, t); ++ } else { ++ float xs[UNIT], xsum[2]; ++ X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); ++#pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); ++ } ++ } ++ } ++ constexpr int RPT = RS * RG * NR; ++#pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) ++#pragma clang loop unroll(full) ++ for (int t = 0; t < NTOK; ++t) { ++ const float v = lane_sum(acc[r][t]); ++ if (li == 0) red[(ks * RPT + lrow0 + r) * 8 + t] = v; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (ks == 0 && li == 0) ++#pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) { ++ const int row = row0 + r; ++ if (row < int(a.n_out)) ++#pragma clang loop unroll(full) ++ for (int t = 0; t < NTOK; ++t) { ++ float v = 0.0f; ++#pragma clang loop unroll(full) ++ for (int k = 0; k < KS; ++k) v += red[(k * RPT + lrow0 + r) * 8 + t]; ++ if (row < sp) y[t * int(a.y_stride) + row] = v; ++ else y2[t * int(a.y_stride) + row - sp] = v; ++ } ++ } ++} ++#define GVCK_(F, L, R, S, K) template [[host_name("gemvck_" #F "_" #L "_" #R "_" #S "_" #K)]] kernel decltype(gemv_catk) gemv_catk; ++#define GVCK(F, L, R, S, K) GVCK_(F, L, R, S, K) ++GVCK(T_Q8_0, 32, 1, 4, 4) GVCK(T_Q8_0, 32, 2, 4, 2) GVCK(T_Q8_0, 32, 1, 8, 8) + + // gemv_x: gemv (KS = 1) with the window's prepared x (the format's plane, NTOK rows, and its per-16 sums) staged in + // threadgroup memory once per threadgroup, so the per-unit x reads of every token come from threadgroup memory +@@ -223,9 +450,10 @@ kernel void gemv_x(device const uchar* W [[buffer(0)]], device const float* xp [ + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; +- load_tab(tab, tid, SG * 32); ++ if (fmt::TAB > 1) load_tab(tab, tid, SG * 32); + const int n = int(a.n_in), s16n = n / 16; + threadgroup float* st = xt + NTOK * n; ++#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + device const float4* src = (device const float4*)(xp + t * int(a.x_stride)); + threadgroup float4* dst = (threadgroup float4*)(xt + t * n); +@@ -237,21 +465,29 @@ kernel void gemv_x(device const uchar* W [[buffer(0)]], device const float* xp [ + const int row0 = int(tg.x) * SG * RG * NR + (sg * RG + lane / LPR) * NR; + const int U = n / UNIT; + float acc[NR][8]; ++#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) ++#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; ++#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) rp[r] = W + (ulong)min(row0 + r, int(a.n_out) - 1) * a.row_bytes; + for (int u = li; u < U; u += LPR) { + float q[NR][UNIT]; + float4 sc[NR]; ++#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); ++#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[UNIT], xsum[2]; + X::load((threadgroup const float*)(xt + t * n), (threadgroup const float*)(st + t * s16n), u, xs, xsum); ++#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } ++#pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) ++#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + if (li == 0 && row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = v; +diff --git a/metal/kernels/glue.metal b/metal/kernels/glue.metal +index b7bc8bcc..b34bb90d 100644 +--- a/metal/kernels/glue.metal ++++ b/metal/kernels/glue.metal +@@ -10,28 +10,29 @@ struct GlueArgs { + uint topk; + }; + ++constant float glue_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, ++ 0.00006103515625f}; ++constant float glue_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; ++ + // grid (n / 16, n_tok), 16 threads: one group of 16 values + kernel void swiglu_prep(device const float* g [[buffer(0)]], device const float* u [[buffer(1)]], + device float* xp [[buffer(2)]], constant GlueArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort j [[thread_index_in_threadgroup]]) { +- threadgroup float v[16]; +- const uint t = tg.y, d = tg.x * 16 + j, plane = a.n_tok * a.n; +- const float gv = g[t * a.n + d]; +- const float h = gv / (1.0f + exp(-gv)) * u[t * a.n + d]; +- const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; +- const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; +- device float* raw = xp + t * a.n + d; ++ const uint t = tg.y, d = tg.x * 16 + j, idx = t * a.n + d, plane = a.n_tok * a.n; ++ const float gv = g[idx]; ++ const float h = gv / (1.0f + exp(-gv)) * u[idx]; ++ device float* raw = xp + idx; + raw[0] = h; +- raw[plane] = h * q2s[j & 7]; +- raw[2 * plane] = h * p4s[j & 3]; ++ raw[plane] = h * glue_q2s[j & 7]; ++ raw[2 * plane] = h * glue_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? h * 0.00390625f : h; +- v[j] = h; +- threadgroup_barrier(mem_flags::mem_threadgroup); +- if (j == 0) { +- float s0 = 0.0f, s1 = 0.0f; +- for (int k = 0; k < 16; k += 2) { s0 += v[k]; s1 += v[k + 1]; } +- xp[4 * plane + t * (a.n / 16) + tg.x] = s0 + s1; ++ float s0 = 0.0f, s1 = 0.0f; ++#pragma clang loop unroll(full) ++ for (int k = 0; k < 16; k += 2) { ++ s0 += simd_shuffle_down(h, ushort(k)); ++ s1 += simd_shuffle_down(h, ushort(k + 1)); + } ++ if (j == 0) xp[4 * plane + t * (a.n / 16) + tg.x] = s0 + s1; + } + + kernel void moe_combine_sh(device const float* ye [[buffer(0)]], device const int* plan [[buffer(1)]], +@@ -40,8 +41,16 @@ kernel void moe_combine_sh(device const float* ye [[buffer(0)]], device const in + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.n_embd, k = gid % a.n_embd; + if (t >= a.n_tok) return; ++ const int row_base = int(t * a.topk); ++ device const int* te = plan + P_TOKENT + row_base; ++ device const float* yep = ye + k; + float s = 0.0f; +- for (uint j = 0; j < a.topk; ++j) s += ye[plan[P_TOKENT + t * a.topk + j] * int(a.n_embd) + k]; ++ if (a.topk == 10) { ++#pragma clang loop unroll(full) ++ for (uint j = 0; j < 10; ++j) s += yep[te[j] * int(a.n_embd)]; ++ } else { ++ for (uint j = 0; j < a.topk; ++j) s += yep[te[j] * int(a.n_embd)]; ++ } + const float gs = 1.0f / (1.0f + exp(-sg[t])); + bo[gid] = fma(gs, ysh[gid], s); + } +diff --git a/metal/kernels/gr.metal b/metal/kernels/gr.metal +index c1960fd5..c18d2ac7 100644 +--- a/metal/kernels/gr.metal ++++ b/metal/kernels/gr.metal +@@ -12,6 +12,11 @@ + + constant int GR_NTOK [[function_constant(1)]]; + ++// the prepared activations' per-position plane scales (units.metal's xprep): q2 planes 4^-(j%8), p4 planes 256^-(j%4) ++constant float gr_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, ++ 0.00006103515625f}; ++constant float gr_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; ++ + struct GrArgs { + uint n_embd; + uint hc; +@@ -31,80 +36,195 @@ kernel void gr_norm(device float* R [[buffer(0)]], device const float* bo [[buff + const uint c = tg.x, t = tg.y, n = a.n_embd; + device float* r = R + (t * a.hc + c) * n; + const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; ++ float rv[10]; + float ss = 0.0f; +- for (uint d = tid; d < n; d += 256) { +- float v = r[d]; ++ if (n == 2560) { ++ device float* rp = r + tid; + if (a.apply) { +- v = fma(bo[t * n + d], g, v); +- r[d] = v; ++ device const float* bp = bo + t * 2560u + tid; ++ #pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) { ++ const float v = fma(bp[i * 256], g, rp[i * 256]); ++ rp[i * 256] = v; ++ rv[i] = v; ++ ss = fma(v, v, ss); ++ } ++ } else { ++ #pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) { ++ rv[i] = rp[i * 256]; ++ ss = fma(rv[i], rv[i], ss); ++ } ++ } ++ } else { ++ for (uint d = tid; d < n; d += 256) { ++ float v = r[d]; ++ if (a.apply) { ++ v = fma(bo[t * n + d], g, v); ++ r[d] = v; ++ } ++ if (n <= 2560) rv[d / 256] = v; ++ ss = fma(v, v, ss); + } +- ss = fma(v, v, ss); + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; ++ #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device float* o = xn + (t * a.hc + c) * n; + device const float* w = w_norm + c * n; +- for (uint d = tid; d < n; d += 256) o[d] = r[d] * w[d] * rs; ++ if (n == 2560) { ++ o += tid; ++ w += tid; ++ #pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) o[i * 256] = rv[i] * w[i * 256] * rs; ++ } else { ++ for (uint d = tid; d < n; d += 256) ++ o[d] = (n <= 2560 ? rv[d / 256] : r[d]) * w[d] * rs; ++ } + } + + // rows: 0 .. hc_lr - 1 = w_down, hc_lr .. hc_lr + hc - 1 = w_inject (both bf16 rows of hc * n_embd). +-// SG simdgroups split the row's units KS ways; LPR = 32 (a simdgroup per row); RS = SG / KS rows per threadgroup. +-// one row's K-split partial over xn (units of format F) +-template +-static inline void grd_acc(device const uchar* w, int nin, int ks, int KS, ushort lane, device const float* xn, ++// SG simdgroups split the row's units KS ways; LPR = 32 (a simdgroup per row); RS = SG / KS simdgroups per ++// threadgroup, NR rows each. ++// one token's dot with a decoded unit (gr_up): x is the token's slice of lo ++template ++static inline float gr_dot(device const float* x, thread const float* q, float4 sc) { ++ float xs[UNIT], xsum[2] = {0.0f, 0.0f}; ++ #pragma clang loop unroll(full) ++ for (int k = 0; k < UNIT; k += 16) load16(x + k, xs + k); ++ return unit_combine(q, sc, xs, xsum); ++} ++ ++// a unit's decoded rows for grd_acc, read four codes at a time ++template ++struct GrdUnit { ++ float q[NR][fmt::UNIT]; ++ float4 sc[NR]; ++ void load(device const uchar* w, uint rb, int u, threadgroup const uint* notab) thread { ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) Dec::run(w + ulong(r) * rb, u, q[r], sc[r], notab); ++ } ++ float4 v(int r, int i) const thread { return float4(q[r][4 * i], q[r][4 * i + 1], q[r][4 * i + 2], q[r][4 * i + 3]); } ++}; ++ ++// Q8_0 two rows a simdgroup: the codes stay packed (8 registers a row, not 32 floats) and are widened per token. ++// Measured 10-16% faster past 6 tokens (where the token loop is rolled), no faster below, and slower for one row. ++template <> ++struct GrdUnit { ++ char4 q[2][8]; ++ float4 sc[2]; ++ void load(device const uchar* w, uint rb, int u, threadgroup const uint*) thread { ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < 2; ++r) { ++ device const uchar* b = w + ulong(r) * rb + 34 * u; ++ device const ushort* p = (device const ushort*)(b + 2); ++ #pragma clang loop unroll(full) ++ for (int k = 0; k < 8; ++k) q[r][k] = as_type(uint(p[2 * k]) | (uint(p[2 * k + 1]) << 16)); ++ const float d = rdh(b); ++ sc[r] = float4(d, d, 0.0f, 0.0f); ++ } ++ } ++ float4 v(int r, int i) const thread { return float4(q[r][i]); } ++}; ++ ++// NR consecutive rows' K-split partials over xn (units of format F; the rows rb bytes apart), acc[r * 8 + t]. ++// unit_combine's sums in its order, with each xn float4 read straight from memory once for all NR rows: xn is ++// 15-30x the weights' bytes at 4-8 tokens, so NR = 2 halves the kernel's traffic. GR_NTOK is a function constant, ++// which the compiler does not unroll over by itself: the pragma keeps acc in registers instead of the stack. Past 6 ++// tokens the unrolled loop measured slower than the rolled one, so it stays rolled there; ROLL keeps it rolled at ++// every count (the packed Q8_0 pair at KS = 4: 1-3% faster at 4-6 tokens, where KS = 8 measured 2-6% slower). ++template ++static inline void grd_acc(device const uchar* w, uint rb, int nin, int ks, int KS, ushort lane, device const float* xn, + thread float* acc, threadgroup const uint* notab) { +- constexpr int UNIT = fmt::UNIT; ++ constexpr int UNIT = fmt::UNIT, H = UNIT / 2; + const int U = nin / UNIT, u0 = ks * U / KS, u1 = (ks + 1) * U / KS; + for (int u = u0 + lane; u < u1; u += 32) { +- float q[UNIT]; +- float4 sc; +- Dec::run(w, u, q, sc, notab); +- for (int t = 0; t < GR_NTOK; ++t) { +- float xs[UNIT], xsum[2] = {0.0f, 0.0f}; +- for (int k = 0; k < UNIT; k += 16) load16(xn + t * nin + UNIT * u + k, xs + k); +- acc[t] += unit_combine(q, sc, xs, xsum); ++ GrdUnit d; ++ d.load(w, rb, u, notab); ++ auto tok = [&](int t) { ++ device const float4* x = (device const float4*)(xn + t * nin + UNIT * u); ++ float s[NR][4]; ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) s[r][0] = s[r][1] = s[r][2] = s[r][3] = 0.0f; ++ #pragma clang loop unroll(full) ++ for (int j = 0; j < H; j += 4) { ++ const float4 a = x[j / 4], b = x[(H + j) / 4]; ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) { ++ const float4 qa = d.v(r, j / 4), qb = d.v(r, (H + j) / 4); ++ s[r][0] = fma(qa.x, a.x, s[r][0]); ++ s[r][1] = fma(qa.y, a.y, s[r][1]); ++ s[r][2] = fma(qb.x, b.x, s[r][2]); ++ s[r][3] = fma(qb.y, b.y, s[r][3]); ++ s[r][0] = fma(qa.z, a.z, s[r][0]); ++ s[r][1] = fma(qa.w, a.w, s[r][1]); ++ s[r][2] = fma(qb.z, b.z, s[r][2]); ++ s[r][3] = fma(qb.w, b.w, s[r][3]); ++ } ++ } ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) ++ acc[r * 8 + t] += d.sc[r].x * (s[r][0] + s[r][1]) + d.sc[r].y * (s[r][2] + s[r][3]) - ++ (d.sc[r].z * 0.0f + d.sc[r].w * 0.0f); ++ }; ++ if (GR_NTOK <= 6 && !ROLL) { ++ #pragma clang loop unroll(full) ++ for (int t = 0; t < GR_NTOK; ++t) tok(t); ++ } else { ++ for (int t = 0; t < GR_NTOK; ++t) tok(t); + } + } + } + +-// FI: the inject rows' format (the Unsloth files: Q8_0 down, F32 inject); FI = F is the original kernel +-template ++// FI: the inject rows' format (the Unsloth files: Q8_0 down, F32 inject); FI = F is the original kernel. NR must ++// divide hc_lr and hc, so that a simdgroup's rows are all w_down or all w_inject. The sums do not depend on SG or NR: ++// every SG / NR with the same KS is bitwise the same kernel. ++template + kernel void gr_down_f(device const uchar* w_down [[buffer(0)]], device const uchar* w_inject [[buffer(1)]], + device const float* xn [[buffer(2)]], device float* lo [[buffer(3)]], + device float* inject [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + constexpr int RS = SG / KS; +- threadgroup float red[KS * RS * 8]; ++ threadgroup float red[KS * RS * NR * 8]; + threadgroup uint notab[1]; + const int ks = sg / RS, rs = sg % RS; +- const int row = int(tg.x) * RS + rs; ++ const int row0 = (int(tg.x) * RS + rs) * NR; + const int nrows = int(a.hc_lr + a.hc); +- const int rr = min(row, nrows - 1); ++ const int rr = min(row0, nrows - NR); + const int nin = int(a.hc * a.n_embd); + const uint rb = uint(nin / fmt::QK * fmt::BB), rbi = uint(nin / fmt::QK * fmt::BB); +- float acc[8]; +- for (int t = 0; t < 8; ++t) acc[t] = 0.0f; +- if (rr < int(a.hc_lr)) grd_acc(w_down + ulong(rr) * rb, nin, ks, KS, lane, xn, acc, notab); +- else grd_acc(w_inject + ulong(rr - int(a.hc_lr)) * rbi, nin, ks, KS, lane, xn, acc, notab); +- for (int t = 0; t < GR_NTOK; ++t) { +- const float v = simd_sum(acc[t]); +- if (lane == 0) red[(ks * RS + rs) * 8 + t] = v; ++ float acc[NR * 8]; ++ #pragma clang loop unroll(full) ++ for (int t = 0; t < NR * 8; ++t) acc[t] = 0.0f; ++ constexpr bool roll = F == T_Q8_0 && NR == 2 && KS == 4; ++ if (rr < int(a.hc_lr)) grd_acc(w_down + ulong(rr) * rb, rb, nin, ks, KS, lane, xn, acc, notab); ++ else grd_acc(w_inject + ulong(rr - int(a.hc_lr)) * rbi, rbi, nin, ks, KS, lane, xn, acc, notab); ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) { ++ #pragma clang loop unroll(full) ++ for (int t = 0; t < GR_NTOK; ++t) { ++ const float v = simd_sum(acc[r * 8 + t]); ++ if (lane == 0) red[((ks * RS + rs) * NR + r) * 8 + t] = v; ++ } + } + threadgroup_barrier(mem_flags::mem_threadgroup); +- if (ks == 0 && lane < GR_NTOK && row < nrows) { +- const int t = lane; +- float v = 0.0f; +- for (int k = 0; k < KS; ++k) v += red[(k * RS + rs) * 8 + t]; +- if (row < int(a.hc_lr)) { +- const float z = v / float(a.hc); +- lo[t * int(a.hc_lr) + row] = z / (1.0f + exp(-z)); +- } else { +- inject[t * int(a.hc) + row - int(a.hc_lr)] = v; ++ if (ks == 0 && lane < NR * 8) { ++ const int r = lane / 8, t = lane % 8, row = row0 + r; ++ if (t < GR_NTOK && row < nrows) { ++ float v = 0.0f; ++ #pragma clang loop unroll(full) ++ for (int k = 0; k < KS; ++k) v += red[((k * RS + rs) * NR + r) * 8 + t]; ++ if (row < int(a.hc_lr)) { ++ const float z = v / float(a.hc); ++ lo[t * int(a.hc_lr) + row] = z / (1.0f + exp(-z)); ++ } else { ++ inject[t * int(a.hc) + row - int(a.hc_lr)] = v; ++ } + } + } + } +@@ -119,7 +239,6 @@ kernel void gr_up_f(device const uchar* w_up [[buffer(0)]], device const float* + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float gate[8][4][16]; +- threadgroup float msum[8][16]; + threadgroup uint notab[1]; + const int n = int(a.n_embd), c = sg >> 1, dl = 8 * (sg & 1) + lane / 4, li = lane & 3; + const int d = int(tg.x) * 16 + dl; +@@ -127,17 +246,16 @@ kernel void gr_up_f(device const uchar* w_up [[buffer(0)]], device const float* + device const uchar* w = w_up + ulong(c * n + d) * (a.hc_lr / fmt::QK * fmt::BB); + const int U = int(a.hc_lr) / UNIT; + float acc[8]; ++ #pragma clang loop unroll(full) + for (int t = 0; t < 8; ++t) acc[t] = 0.0f; + for (int u = li; u < U; u += 4) { + float q[UNIT]; + float4 sc; + Dec::run(w, u, q, sc, notab); +- for (int t = 0; t < GR_NTOK; ++t) { +- float xs[UNIT], xsum[2] = {0.0f, 0.0f}; +- for (int k = 0; k < UNIT; k += 16) load16(lo + t * int(a.hc_lr) + UNIT * u + k, xs + k); +- acc[t] += unit_combine(q, sc, xs, xsum); +- } ++ #pragma clang loop unroll(full) ++ for (int t = 0; t < GR_NTOK; ++t) acc[t] += gr_dot(lo + t * int(a.hc_lr) + UNIT * u, q, sc); + } ++ #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) { + const float v = lane_sum<4>(acc[t]); + if (li == 0) gate[t][c][dl] = 1.0f / (1.0f + exp(-v)); +@@ -149,31 +267,41 @@ kernel void gr_up_f(device const uchar* w_up [[buffer(0)]], device const float* + const int dd = int(tg.x) * 16 + j; + const int hc = int(a.hc); + float m = 0.0f; +- for (int cc = 0; cc < hc; ++cc) m = fma(xn[(t * hc + cc) * n + dd], gate[t][cc][j], m); ++ if (hc == 4) { ++ #pragma clang loop unroll(full) ++ for (int cc = 0; cc < 4; ++cc) m = fma(xn[(t * 4 + cc) * n + dd], gate[t][cc][j], m); ++ } else { ++ for (int cc = 0; cc < hc; ++cc) m = fma(xn[(t * hc + cc) * n + dd], gate[t][cc][j], m); ++ } + m = m / float(hc); + const uint plane = a.n_tok * uint(n); + device float* raw = xp + t * n + dd; +- const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, +- 0.00006103515625f}; +- const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + raw[0] = m; +- raw[plane] = m * q2s[j & 7]; +- raw[2 * plane] = m * p4s[j & 3]; ++ raw[plane] = m * gr_q2s[j & 7]; ++ raw[2 * plane] = m * gr_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? m * 0.00390625f : m; +- msum[t][j] = m; +- } +- threadgroup_barrier(mem_flags::mem_threadgroup); +- if (t < GR_NTOK && j == 0) { +- // the per-16 sum in xprep's order (even and odd partial sums, then added) + float s0 = 0.0f, s1 = 0.0f; +- for (int k = 0; k < 16; k += 2) { s0 += msum[t][k]; s1 += msum[t][k + 1]; } +- xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1; ++ #pragma clang loop unroll(full) ++ for (int k = 0; k < 16; k += 2) { ++ s0 += simd_shuffle_down(m, ushort(k)); ++ s1 += simd_shuffle_down(m, ushort(k + 1)); ++ } ++ if (j == 0) xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1; + } + } + +-#define GRD(S, K) template [[host_name("gr_down_" #S "_" #K)]] kernel decltype(gr_down_f) gr_down_f; \ +- template [[host_name("gr_down_q8_" #S "_" #K)]] kernel decltype(gr_down_f) gr_down_f; \ +- template [[host_name("gr_down_q8f_" #S "_" #K)]] kernel decltype(gr_down_f) gr_down_f; ++#define GRD1(S, K, R, SFX) \ ++ template [[host_name("gr_down_" #S "_" #K SFX)]] kernel decltype(gr_down_f) gr_down_f; \ ++ template [[host_name("gr_down_q8_" #S "_" #K SFX)]] kernel decltype(gr_down_f) gr_down_f; \ ++ template [[host_name("gr_down_q8f_" #S "_" #K SFX)]] kernel decltype(gr_down_f) gr_down_f; ++#define GRD(S, K) GRD1(S, K, 1, "") + GRD(8, 8) GRD(8, 4) GRD(4, 4) GRD(4, 2) GRD(16, 16) GRD(16, 8) ++// two rows a simdgroup (gr_read past 4 tokens) ++GRD1(8, 8, 2, "_r2") GRD1(4, 4, 2, "_r2") + template [[host_name("gr_up")]] kernel decltype(gr_up_f) gr_up_f; + template [[host_name("gr_up_q8")]] kernel decltype(gr_up_f) gr_up_f; ++ ++template [[host_name("gr_down_q4_8_8")]] kernel decltype(gr_down_f<8, 8, 1, T_Q4_0, T_Q8_0>) gr_down_f<8, 8, 1, T_Q4_0, T_Q8_0>; ++template [[host_name("gr_up_q4")]] kernel decltype(gr_up_f) gr_up_f; ++template [[host_name("gr_down_q5_8_8")]] kernel decltype(gr_down_f<8, 8, 1, T_Q5_0, T_Q8_0>) gr_down_f<8, 8, 1, T_Q5_0, T_Q8_0>; ++template [[host_name("gr_up_q5")]] kernel decltype(gr_up_f) gr_up_f; +diff --git a/metal/kernels/gr2.metal b/metal/kernels/gr2.metal +index 1253a11d..58342f09 100644 +--- a/metal/kernels/gr2.metal ++++ b/metal/kernels/gr2.metal +@@ -87,6 +87,7 @@ kernel void gr_down3(device const uchar* w_down [[buffer(0)]], device const ucha + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < T) { + float tot = 0.0f; ++ #pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) tot += part[k][tid]; + const float rsv = rsqrt(tot / float(n) + a.eps); + rs_tg[tid] = rsv; +@@ -99,30 +100,44 @@ kernel void gr_down3(device const uchar* w_down [[buffer(0)]], device const ucha + // the rows: simdgroup sg holds RPS rows; lanes split the stream's units + int rows[RPS]; + device const uchar* wr[RPS]; ++ #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + rows[k] = (int(tg.x) * 8 + sg) * RPS + k; + const int rr = min(rows[k], nrows - 1); + wr[k] = rr < int(a.hc_lr) ? w_down + ulong(rr) * a.rb_down : w_inject + ulong(rr - int(a.hc_lr)) * a.rb_down; + } + float acc[RPS][TT]; +- for (int k = 0; k < RPS; ++k) ++ #pragma clang loop unroll(full) ++ for (int k = 0; k < RPS; ++k) { ++ #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) acc[k][t] = 0.0f; ++ } + const int U = n / UNIT, u0 = c * U; + for (int u = lane; u < U; u += 32) { + float q[RPS][UNIT]; + float4 sc[RPS]; ++ #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) Dec::run(wr[k], u0 + u, q[k], sc[k], tab); +- for (int t = 0; t < T; ++t) { ++ // TT-bounded with a T guard, so acc[][t] is indexed by a constant (registers, not the stack) ++ #pragma clang loop unroll(full) ++ for (int t = 0; t < TT; ++t) { ++ if (t >= T) break; + float xs[UNIT], xsum[2] = {0.0f, 0.0f}; ++ #pragma clang loop unroll(full) + for (int e = 0; e < UNIT; ++e) xs[e] = xs_tg[t][u * UNIT + e]; ++ #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) acc[k][t] += unit_combine(q[k], sc[k], xs, xsum); + } + } +- for (int k = 0; k < RPS; ++k) +- for (int t = 0; t < T; ++t) { ++ #pragma clang loop unroll(full) ++ for (int k = 0; k < RPS; ++k) { ++ #pragma clang loop unroll(full) ++ for (int t = 0; t < TT; ++t) { ++ if (t >= T) break; + const float v = simd_sum(acc[k][t]); + if (lane == 0 && rows[k] < nrows) D[((t0 + t) * nrows + rows[k]) * hc + c] = v; + } ++ } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } +@@ -138,7 +153,6 @@ kernel void gr_up3(device const uchar* w_up [[buffer(0)]], device const float* D + constexpr int UNIT = fmt::UNIT; + threadgroup float lo_tg[8][320]; + threadgroup float gate[8][4][16]; +- threadgroup float msum[8][16]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, 256); + const int n = int(a.n_embd), hc = int(a.hc), lr = int(a.hc_lr), nrows = lr + hc; +@@ -154,17 +168,21 @@ kernel void gr_up3(device const uchar* w_up [[buffer(0)]], device const float* D + device const uchar* w = w_up + ulong(c * n + d) * a.rb_up; + const int U = lr / UNIT; + float acc[8]; ++ #pragma clang loop unroll(full) + for (int t = 0; t < 8; ++t) acc[t] = 0.0f; + for (int u = li; u < U; u += 4) { + float q[UNIT]; + float4 sc; + Dec::run(w, u, q, sc, tab); ++ #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) { + float xs[UNIT], xsum[2] = {0.0f, 0.0f}; ++ #pragma clang loop unroll(full) + for (int e = 0; e < UNIT; ++e) xs[e] = lo_tg[t][UNIT * u + e]; + acc[t] += unit_combine(q, sc, xs, xsum); + } + } ++ #pragma clang loop unroll(full) + for (int t = 0; t < GR_NTOK; ++t) { + const float v = lane_sum<4>(acc[t]); + if (li == 0) gate[t][c][dl] = 1.0f / (1.0f + exp(-v)); +@@ -174,27 +192,32 @@ kernel void gr_up3(device const uchar* w_up [[buffer(0)]], device const float* D + if (t < GR_NTOK) { + const int dd = int(tg.x) * 16 + j; + float m = 0.0f; +- for (int cc = 0; cc < hc; ++cc) { +- const float xn = Rn[(t * hc + cc) * n + dd] * wn[cc * n + dd] * rs[t * hc + cc]; +- m = fma(xn, gate[t][cc][j], m); ++ if (hc == 4) { ++ #pragma clang loop unroll(full) ++ for (int cc = 0; cc < 4; ++cc) { ++ const float xn = Rn[(t * 4 + cc) * n + dd] * wn[cc * n + dd] * rs[t * 4 + cc]; ++ m = fma(xn, gate[t][cc][j], m); ++ } ++ } else { ++ for (int cc = 0; cc < hc; ++cc) { ++ const float xn = Rn[(t * hc + cc) * n + dd] * wn[cc * n + dd] * rs[t * hc + cc]; ++ m = fma(xn, gate[t][cc][j], m); ++ } + } + m = m / float(hc); + const uint plane = a.n_tok * uint(n); + device float* raw = xp + t * n + dd; +- const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, +- 0.00006103515625f}; +- const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + raw[0] = m; +- raw[plane] = m * q2s[j & 7]; +- raw[2 * plane] = m * p4s[j & 3]; ++ raw[plane] = m * gr_q2s[j & 7]; ++ raw[2 * plane] = m * gr_p4s[j & 3]; + raw[3 * plane] = (j & 1) ? m * 0.00390625f : m; +- msum[t][j] = m; +- } +- threadgroup_barrier(mem_flags::mem_threadgroup); +- if (t < GR_NTOK && j == 0) { + float s0 = 0.0f, s1 = 0.0f; +- for (int k = 0; k < 16; k += 2) { s0 += msum[t][k]; s1 += msum[t][k + 1]; } +- xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1; ++ #pragma clang loop unroll(full) ++ for (int k = 0; k < 16; k += 2) { ++ s0 += simd_shuffle_down(m, ushort(k)); ++ s1 += simd_shuffle_down(m, ushort(k + 1)); ++ } ++ if (j == 0) xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1; + } + } + +@@ -223,11 +246,13 @@ kernel void gr_down3b(device const uchar* w_down [[buffer(0)]], device const uch + const int nrows = int(a.hc_lr) + hc; + int rows[RPS]; + uint4 wr[RPS][5][2]; ++ #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + rows[k] = (int(tg.x) * 8 + sg) * RPS + k; + const int rr = min(rows[k], nrows - 1); + device const uchar* w = rr < int(a.hc_lr) ? w_down + ulong(rr) * a.rb_down : w_inject + ulong(rr - int(a.hc_lr)) * a.rb_down; + device const uint4* p = (device const uint4*)(w + 2 * (c * n)); ++ #pragma clang loop unroll(full) + for (int i = 0; i < 5; ++i) { + wr[k][i][0] = p[2 * (lane + 32 * i)]; + wr[k][i][1] = p[2 * (lane + 32 * i) + 1]; +@@ -261,6 +286,7 @@ kernel void gr_down3b(device const uchar* w_down [[buffer(0)]], device const uch + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < T) { + float tot = 0.0f; ++ #pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) tot += part[k][tid]; + const float rsv = rsqrt(tot / float(n) + a.eps); + rs_tg[tid] = rsv; +@@ -271,31 +297,44 @@ kernel void gr_down3b(device const uchar* w_down [[buffer(0)]], device const uch + for (int d = tid; d < n; d += 256) xs_tg[t][d] = xs_tg[t][d] * wn[c * n + d] * rs_tg[t]; + threadgroup_barrier(mem_flags::mem_threadgroup); + float acc[RPS][TT]; +- for (int k = 0; k < RPS; ++k) ++ #pragma clang loop unroll(full) ++ for (int k = 0; k < RPS; ++k) { ++ #pragma clang loop unroll(full) + for (int t = 0; t < TT; ++t) acc[k][t] = 0.0f; ++ } ++ #pragma clang loop unroll(full) + for (int i = 0; i < 5; ++i) { + const int u = lane + 32 * i; ++ #pragma clang loop unroll(full) + for (int k = 0; k < RPS; ++k) { + const uint v[8] = {wr[k][i][0].x, wr[k][i][0].y, wr[k][i][0].z, wr[k][i][0].w, + wr[k][i][1].x, wr[k][i][1].y, wr[k][i][1].z, wr[k][i][1].w}; + float q[16]; ++ #pragma clang loop unroll(full) + for (int e = 0; e < 8; ++e) { + q[2 * e] = as_type(v[e] << 16); + q[2 * e + 1] = as_type(v[e] & 0xFFFF0000u); + } + const float4 sc = float4(1.0f, 1.0f, 0.0f, 0.0f); +- for (int t = 0; t < T; ++t) { ++ #pragma clang loop unroll(full) ++ for (int t = 0; t < TT; ++t) { ++ if (t >= T) break; + float xs[16], xsum[2] = {0.0f, 0.0f}; ++ #pragma clang loop unroll(full) + for (int e = 0; e < 16; ++e) xs[e] = xs_tg[t][u * 16 + e]; + acc[k][t] += unit_combine<16>(q, sc, xs, xsum); + } + } + } +- for (int k = 0; k < RPS; ++k) +- for (int t = 0; t < T; ++t) { ++ #pragma clang loop unroll(full) ++ for (int k = 0; k < RPS; ++k) { ++ #pragma clang loop unroll(full) ++ for (int t = 0; t < TT; ++t) { ++ if (t >= T) break; + const float v = simd_sum(acc[k][t]); + if (lane == 0 && rows[k] < nrows) D[((t0 + t) * nrows + rows[k]) * hc + c] = v; + } ++ } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } +@@ -317,25 +356,129 @@ kernel void gr_norm_comb(device float* R [[buffer(0)]], device const float* ye [ + device float* r = R + (t * a.hc + c) * n; + const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); ++ float rv[10]; + float ss = 0.0f; +- for (uint d = tid; d < n; d += 256) { +- float v = r[d]; ++ if (n == 2560 && topk == 10) { ++ device const float* yep[10]; + if (a.apply) { +- float s = 0.0f; +- for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; +- const float bo = fma(gs, ysh[t * n + d], s); +- v = fma(bo, g, v); +- r[d] = v; ++ #pragma clang loop unroll(full) ++ for (uint j = 0; j < 10; ++j) ++ yep[j] = ye + uint(plan[P_TOKENT + t * 10 + j]) * 2560u + tid; ++ } ++ device const float* yshp = ysh + t * 2560u + tid; ++ device float* rp = r + tid; ++ #pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) { ++ float v = rp[i * 256]; ++ if (a.apply) { ++ float s = 0.0f; ++ #pragma clang loop unroll(full) ++ for (uint j = 0; j < 10; ++j) s += yep[j][i * 256]; ++ v = fma(fma(gs, yshp[i * 256], s), g, v); ++ rp[i * 256] = v; ++ } ++ rv[i] = v; ++ ss = fma(v, v, ss); ++ } ++ } else { ++ for (uint d = tid; d < n; d += 256) { ++ float v = r[d]; ++ if (a.apply) { ++ float s = 0.0f; ++ for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; ++ const float bo = fma(gs, ysh[t * n + d], s); ++ v = fma(bo, g, v); ++ r[d] = v; ++ } ++ if (n <= 2560) rv[d / 256] = v; ++ ss = fma(v, v, ss); + } +- ss = fma(v, v, ss); + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; ++ #pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device float* o = xn + (t * a.hc + c) * n; + device const float* w = w_norm + c * n; +- for (uint d = tid; d < n; d += 256) o[d] = r[d] * w[d] * rs; ++ if (n == 2560) { ++ o += tid; ++ w += tid; ++ #pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) o[i * 256] = rv[i] * w[i * 256] * rs; ++ } else { ++ for (uint d = tid; d < n; d += 256) ++ o[d] = (n <= 2560 ? rv[d / 256] : r[d]) * w[d] * rs; ++ } ++} ++ ++// Fixed n_embd = 2560, hc = 4; grid (1, n_tok), 256 threads. Form bo once for all streams. ++kernel void gr_norm_comb_4s(device float* R [[buffer(0)]], device const float* ye [[buffer(1)]], ++ device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], ++ device float* xn [[buffer(4)]], constant GrArgs& a [[buffer(5)]], ++ device const int* plan [[buffer(6)]], device const float* ysh [[buffer(7)]], ++ device const float* sgv [[buffer(8)]], constant uint& topk [[buffer(9)]], ++ uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], ++ ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { ++ threadgroup float part[4][8]; ++ const uint t = tg.y, n = 2560; ++ device float* rp = R + t * 4 * n + tid; ++ device float* op = xn + t * 4 * n + tid; ++ device const float* wp = w_norm + tid; ++ float rv[4][10], ss[4] = {0.0f, 0.0f, 0.0f, 0.0f}, g[4]; ++ #pragma clang loop unroll(full) ++ for (uint c = 0; c < 4; ++c) ++ g[c] = a.apply ? 2.0f / (1.0f + exp(-inj[t * 4 + c] / 4.0f)) : 0.0f; ++ const float gs = 1.0f / (1.0f + exp(-sgv[t])); ++ device const float* yep[10]; ++ if (a.apply && topk == 10) { ++ #pragma clang loop unroll(full) ++ for (uint j = 0; j < 10; ++j) ++ yep[j] = ye + uint(plan[P_TOKENT + t * 10 + j]) * 2560u + tid; ++ } ++ device const float* yshp = ysh + t * n + tid; ++ #pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) { ++ const uint d = tid + i * 256; ++ float bo = 0.0f; ++ if (a.apply) { ++ float s = 0.0f; ++ if (topk == 10) { ++ #pragma clang loop unroll(full) ++ for (uint j = 0; j < 10; ++j) s += yep[j][i * 256]; ++ } else { ++ for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; ++ } ++ bo = fma(gs, yshp[i * 256], s); ++ } ++ #pragma clang loop unroll(full) ++ for (uint c = 0; c < 4; ++c) { ++ device float* r = rp + c * n; ++ float v = r[i * 256]; ++ if (a.apply) { v = fma(bo, g[c], v); r[i * 256] = v; } ++ rv[c][i] = v; ++ ss[c] = fma(v, v, ss[c]); ++ } ++ } ++ #pragma clang loop unroll(full) ++ for (uint c = 0; c < 4; ++c) { ++ const float s = simd_sum(ss[c]); ++ if (lane == 0) part[c][sg] = s; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ #pragma clang loop unroll(full) ++ for (uint c = 0; c < 4; ++c) { ++ float tot = 0.0f; ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 8; ++i) tot += part[c][i]; ++ const float rs = rsqrt(tot / float(n) + a.eps); ++ device float* o = op + c * n; ++ device const float* w = wp + c * n; ++ #pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) { ++ o[i * 256] = rv[c][i] * w[i * 256] * rs; ++ } ++ } + } +diff --git a/metal/kernels/mma.metal b/metal/kernels/mma.metal +index 9bbb27e2..8fcc0abb 100644 +--- a/metal/kernels/mma.metal ++++ b/metal/kernels/mma.metal +@@ -17,11 +17,57 @@ template struct DW { + constexpr int UNIT = fmt::UNIT, UPB = fmt::QK / UNIT; + float q[UNIT]; + const float s = deq(row + (u / UPB) * fmt::BB, u % UPB, q, tab); ++#pragma clang loop unroll(full) + for (int j = 0; j < UNIT; ++j) w[j] = q[j] * s; + p0 = u * UNIT - k0; + p1 = p0 + UNIT / 2; + } + }; ++template <> struct DW { ++ static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, ++ threadgroup const uint* tab) { ++ device const uchar* b = row + u * 18; ++ device const ushort* p = (device const ushort*)(b + 2); ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 8; ++j) { ++ const ushort v = p[j]; ++ const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); ++ w[2 * j] = float(w0.x); ++ w[2 * j + 16] = float(w0.y); ++ w[2 * j + 1] = float(w1.x); ++ w[2 * j + 17] = float(w1.y); ++ } ++ const float d = rdh(b); ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 32; ++j) w[j] *= d; ++ p0 = 32 * u - k0; ++ p1 = p0 + 16; ++ } ++}; ++template <> struct DW { ++ static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, ++ threadgroup const uint* tab) { ++ device const uchar* b = row + (u >> 3) * 136; ++ const int ui = u & 7; ++ device const ushort* p = (device const ushort*)(b + 8 + 16 * ui); ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 8; ++j) { ++ const ushort v = p[j]; ++ const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); ++ w[2 * j] = float(w0.x); ++ w[2 * j + 16] = float(w0.y); ++ w[2 * j + 1] = float(w1.x); ++ w[2 * j + 17] = float(w1.y); ++ } ++ const uint sh = *(device const ushort*)(b + 2); ++ const int ls = int((b[4 + (ui >> 1)] >> (4 * (ui & 1))) & 15) | int(((sh >> (2 * ui)) & 3) << 4); ++ const float d = rdh(b) * float(ls - 32); ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 32; ++j) w[j] *= d; ++ p0 = 32 * u - k0; ++ p1 = p0 + 16; ++ } ++}; + // Q2_0: (code - 1) * d + template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, +@@ -30,6 +76,7 @@ template <> struct DW { + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + const ushort a = p[0], c = p[1]; + const float d = rdh(b); ++#pragma clang loop unroll(full) + for (int j = 0; j < 8; ++j) { + w[j] = float(int((a >> (2 * j)) & 3) - 1) * d; + w[8 + j] = float(int((c >> (2 * j)) & 3) - 1) * d; +@@ -44,6 +91,7 @@ template <> struct DW { + device const uchar* b = row + u * 24; + float q[32]; + const float d = deq(b, 0, q, tab), m = rdh(b + 2); ++#pragma clang loop unroll(full) + for (int j = 0; j < 32; ++j) w[j] = q[j] * d + m; + p0 = 32 * u - k0; + p1 = p0 + 16; +@@ -53,6 +101,7 @@ template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const float* p = (device const float*)(row + 64 * u); ++#pragma clang loop unroll(full) + for (int j = 0; j < 16; ++j) w[j] = p[j]; + p0 = 16 * u - k0; + p1 = p0 + 8; +@@ -64,6 +113,7 @@ template <> struct DW { + device const uint4* p = (device const uint4*)(row + 32 * u); + const uint4 a = p[0], c = p[1]; + const uint v[8] = {a.x, a.y, a.z, a.w, c.x, c.y, c.z, c.w}; ++#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) { + w[2 * k] = as_type(v[k] << 16); + w[2 * k + 1] = as_type(v[k] & 0xFFFF0000u); +@@ -84,6 +134,7 @@ template <> struct DW { + scale_min_k4(2 * c + 1, b + 4, s2, m2); + const float d1 = d * s1, mm1 = dm * m1, d2 = d * s2, mm2 = dm * m2; + device const uchar* q = b + 16 + 32 * c + 16 * m; ++#pragma clang loop unroll(full) + for (int j = 0; j < 16; ++j) { + w[j] = d1 * float(q[j] & 15) - mm1; + w[16 + j] = d2 * float(q[j] >> 4) - mm2; +@@ -105,6 +156,7 @@ template <> struct DW { + device const uchar* ql = b + 48 + 32 * c + 16 * m; + device const uchar* qh = b + 16 + 16 * m; + const uchar u1 = uchar(1 << (2 * c)), u2 = uchar(2 << (2 * c)); ++#pragma clang loop unroll(full) + for (int j = 0; j < 16; ++j) { + w[j] = d1 * float((ql[j] & 15) + ((qh[j] & u1) ? 16 : 0)) - mm1; + w[16 + j] = d2 * float((ql[j] >> 4) + ((qh[j] & u2) ? 16 : 0)) - mm2; +@@ -124,11 +176,20 @@ template <> struct DW { + device const char* sc = (device const char*)(b + 192); + const float d = rdh(b + 208); + const float ds1 = d * float(sc[8 * n + m + 2 * p]), ds2 = d * float(sc[8 * n + m + 2 * p + 4]); +- for (int j = 0; j < 16; ++j) { +- const int c1 = int((ql[j] & 15) | (((qh[j] >> (2 * p)) & 3) << 4)) - 32; +- const int c2 = int((ql[j] >> 4) | (((qh[j] >> (2 * p + 4)) & 3) << 4)) - 32; +- w[j] = ds1 * float(c1); +- w[16 + j] = ds2 * float(c2); ++ const uint sh = uint(2 * p), sh4 = sh + 4u; ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 16; j += 4) { ++ const uint la = rd32(ql + j), ha = rd32(qh + j); ++ const uchar4 c1 = as_type((la & 0x0F0F0F0Fu) | (((ha >> sh) & 0x03030303u) << 4)); ++ const uchar4 c2 = as_type(((la >> 4) & 0x0F0F0F0Fu) | (((ha >> sh4) & 0x03030303u) << 4)); ++ w[j] = ds1 * float(int(c1.x) - 32); ++ w[j + 1] = ds1 * float(int(c1.y) - 32); ++ w[j + 2] = ds1 * float(int(c1.z) - 32); ++ w[j + 3] = ds1 * float(int(c1.w) - 32); ++ w[16 + j] = ds2 * float(int(c2.x) - 32); ++ w[16 + j + 1] = ds2 * float(int(c2.y) - 32); ++ w[16 + j + 2] = ds2 * float(int(c2.z) - 32); ++ w[16 + j + 3] = ds2 * float(int(c2.w) - 32); + } + p0 = 256 * (u >> 3) + 128 * n + 32 * p + 16 * m - k0; + p1 = p0 + 64; +@@ -149,6 +210,7 @@ template <> struct DW { + const float d = rdh(b + 108); + const float dl1 = d * float(l1 - 32), dl2 = d * float(l2 - 32); + const uchar hb = uchar(1 << (4 * n + j)); ++#pragma clang loop unroll(full) + for (int l = 0; l < 16; ++l) { + w[l] = dl1 * float(int((q[l] >> (2 * j)) & 3) - ((hm[l] & hb) ? 0 : 4)); + w[16 + l] = dl2 * float(int((q[l + 16] >> (2 * j)) & 3) - ((hm[l + 16] & hb) ? 0 : 4)); +@@ -188,6 +250,7 @@ kernel void gemv_mma(device const uchar* W [[buffer(0)]], device const float* x + float w[UNIT]; + int p0, p1; + DW::run(W + ulong(rr) * a.row_bytes, k0 / UNIT + uu, k0, w, p0, p1, tab); ++#pragma clang loop unroll(full) + for (int j = 0; j < H; ++j) { + w8[r * 128 + p0 + j] = w[j]; + w8[r * 128 + p1 + j] = w[H + j]; +diff --git a/metal/kernels/moe.metal b/metal/kernels/moe.metal +index bb8fdbbc..878252e7 100644 +--- a/metal/kernels/moe.metal ++++ b/metal/kernels/moe.metal +@@ -20,7 +20,13 @@ + #define P_ETOK (P_GSTART + P_MAXE + 4) + #define P_EW (P_ETOK + P_MAXE) + #define P_TOKENT (P_EW + P_MAXE) +-#define P_INTS (P_TOKENT + P_MAXE) ++#define P_SPLIT (P_TOKENT + P_MAXE) // function constant 26 only: [0] resident groups, then the indirect grids ++#define P_IGU_R (P_SPLIT + 4) // ... of the resident groups (the GPU computes them during the host's turn) ++#define P_IDN_R (P_SPLIT + 8) ++#define P_IGU_M (P_SPLIT + 12) // ... and of the missing ones (after it) ++#define P_IDN_M (P_SPLIT + 16) ++#define P_GSLOT (P_SPLIT + 20) ++#define P_INTS (P_GSLOT + P_MAXE) + + // Expert e's matrices: gate at W + gate_off + e * gu_stride, up at W + up_off + e * gu_stride, down at + // W + down_off + e * d_stride (a pack's blobs: gate_off 0, strides = the blob; a GGUF: the three tensors). +@@ -57,43 +63,77 @@ static inline device const uchar* arena_slot(uint v, device const uchar* W, devi + } + #endif + ++// Function constant 26 (the SSD-streaming decode window, with smap): the groups whose expert is in the arena when the ++// routing is made come first, the missing ones after them, each part with its own indirect grids (P_SPLIT): the ++// resident part runs while the host reads the missing experts. The entries keep their (token, k) mapping, so the ++// combine (and every number) is the same; without it the plan is the old one, bit for bit. ++constant bool kSplitResFc [[function_constant(26)]]; ++constant bool kSplitRes = is_function_constant_defined(kSplitResFc) && kSplitResFc; ++ ++// s_grp, s_ent: 64 each (the upper halves for the missing groups) + static inline void plan_body(device const int* ids, device const float* w, device int* plan, constant MoeArgs& a, + ushort tid, ushort nth, ushort sg, ushort lane, threadgroup int* s_grp, +- threadgroup int* s_ent) { ++ threadgroup int* s_ent, device const uint* smap = nullptr) { + const int ne = int(a.n_tok * a.topk); + // one thread per expert id (nth == n_expert, at most 1024) + const int e = tid; + int cnt = 0; +- if (e < int(a.n_expert)) ++ if (e < int(a.n_expert)) { ++ #pragma clang loop unroll_count(10) + for (int i = 0; i < ne; ++i) cnt += ids[i] == e; ++ } + const int has = cnt > 0; +- const int pg = simd_prefix_exclusive_sum(has), pe = simd_prefix_exclusive_sum(cnt); +- if (lane == 31) { s_grp[sg] = pg + has; s_ent[sg] = pe + cnt; } ++ const bool split = kSplitRes && smap != nullptr; ++ const int miss = (split && has && smap[e] == 0xFFFFFFFFu) ? 1 : 0; ++ const int h0 = has - miss, c0 = miss ? 0 : cnt, c1 = miss ? cnt : 0; ++ const int pg = simd_prefix_exclusive_sum(h0), pe = simd_prefix_exclusive_sum(c0); ++ const int pg1 = split ? simd_prefix_exclusive_sum(miss) : 0; ++ const int pe1 = split ? simd_prefix_exclusive_sum(c1) : 0; ++ if (lane == 31) { ++ s_grp[sg] = pg + h0; s_ent[sg] = pe + c0; ++ if (split) { s_grp[32 + sg] = pg1 + miss; s_ent[32 + sg] = pe1 + c1; } ++ } + threadgroup_barrier(mem_flags::mem_threadgroup); +- if (tid == 0) { +- int og = 0, oe = 0; ++ if (sg == 0) { + const int nsg = (nth + 31) / 32; +- for (int i = 0; i < nsg; ++i) { +- const int g = s_grp[i], n = s_ent[i]; +- s_grp[i] = og; +- s_ent[i] = oe; +- og += g; +- oe += n; ++ const int g0 = lane < nsg ? s_grp[lane] : 0, n0 = lane < nsg ? s_ent[lane] : 0; ++ const int og0 = simd_prefix_exclusive_sum(g0), oe0 = simd_prefix_exclusive_sum(n0); ++ if (lane < nsg) { s_grp[lane] = og0; s_ent[lane] = oe0; } ++ const int ng0 = simd_broadcast(og0 + g0, 31), ne0 = simd_broadcast(oe0 + n0, 31); ++ int og = ng0, oe = ne0; ++ if (split) { ++ const int g1 = lane < nsg ? s_grp[32 + lane] : 0, n1 = lane < nsg ? s_ent[32 + lane] : 0; ++ const int og1 = ng0 + simd_prefix_exclusive_sum(g1), oe1 = ne0 + simd_prefix_exclusive_sum(n1); ++ if (lane < nsg) { s_grp[32 + lane] = og1; s_ent[32 + lane] = oe1; } ++ og = simd_broadcast(og1 + g1, 31); ++ oe = simd_broadcast(oe1 + n1, 31); ++ } ++ if (lane == 0) { ++ plan[P_NG] = og; ++ plan[P_IGU] = int(a.gx_gu); plan[P_IGU + 1] = og; plan[P_IGU + 2] = 1; ++ plan[P_IDN] = int(a.gx_dn); plan[P_IDN + 1] = og; plan[P_IDN + 2] = 1; ++ plan[P_GSTART + og] = oe; ++ if (split) { ++ plan[P_SPLIT] = ng0; ++ plan[P_IGU_R] = int(a.gx_gu); plan[P_IGU_R + 1] = ng0; plan[P_IGU_R + 2] = 1; ++ plan[P_IDN_R] = int(a.gx_dn); plan[P_IDN_R + 1] = ng0; plan[P_IDN_R + 2] = 1; ++ plan[P_IGU_M] = int(a.gx_gu); plan[P_IGU_M + 1] = og - ng0; plan[P_IGU_M + 2] = 1; ++ plan[P_IDN_M] = int(a.gx_dn); plan[P_IDN_M + 1] = og - ng0; plan[P_IDN_M + 2] = 1; ++ } + } +- plan[P_NG] = og; +- plan[P_IGU] = int(a.gx_gu); plan[P_IGU + 1] = og; plan[P_IGU + 2] = 1; +- plan[P_IDN] = int(a.gx_dn); plan[P_IDN + 1] = og; plan[P_IDN + 2] = 1; +- plan[P_GSTART + og] = oe; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (has) { +- const int gi = s_grp[sg] + pg, off = s_ent[sg] + pe; ++ const int gi = (miss ? s_grp[32 + sg] + pg1 : s_grp[sg] + pg), off = miss ? s_ent[32 + sg] + pe1 : s_ent[sg] + pe; + plan[P_GEXP + gi] = e; + plan[P_GSTART + gi] = off; ++ if (split && !miss) plan[P_GSLOT + gi] = int(smap[e]); ++ const bool topk10 = (a.topk == 10u); + int j = 0; ++ #pragma clang loop unroll_count(10) + for (int i = 0; i < ne; ++i) + if (ids[i] == e) { +- plan[P_ETOK + off + j] = i / int(a.topk); ++ plan[P_ETOK + off + j] = topk10 ? ((i * 205) >> 11) : (i / int(a.topk)); + plan[P_EW + off + j] = as_type(w[i]); + plan[P_TOKENT + i] = off + j; + ++j; +@@ -105,7 +145,7 @@ kernel void moe_plan(device const int* ids [[buffer(0)]], device const float* w + device int* plan [[buffer(2)]], constant MoeArgs& a [[buffer(3)]], + ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { +- threadgroup int s_grp[32], s_ent[32]; ++ threadgroup int s_grp[64], s_ent[64]; + plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); + } + +@@ -113,92 +153,432 @@ kernel void moe_plan(device const int* ids [[buffer(0)]], device const float* w + // token t exactly as native_router.cu (pinned llama.cpp topk-moe): softmax over the 512 logits, top 10 by + // probability (equal values: the lower expert index), weights = p / max(sum of the selected, 2^-14), the sum taken + // by winning lane in selection order and reduced with the same xor tree. ++#ifndef STRATA_MOCE_FC_METAL ++#define STRATA_MOCE_FC_METAL ++constant int kMoceProtectFc [[function_constant(24)]]; ++constant int kMoceProtect = is_function_constant_defined(kMoceProtectFc) ? kMoceProtectFc : 3; ++constant uint kMoceScaleR1BitsFc [[function_constant(25)]]; ++constant float kMoceScaleR1 = is_function_constant_defined(kMoceScaleR1BitsFc) ? as_type(kMoceScaleR1BitsFc) : 1.0f; ++#endif + kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]], + device float* w [[buffer(2)]], device int* plan [[buffer(3)]], constant MoeArgs& a [[buffer(4)]], ++ device const uint* smap [[buffer(5)]], + ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { +- threadgroup int s_grp[32], s_ent[32]; ++ threadgroup int s_grp[64], s_ent[64]; + if (sg < a.n_tok) { + device const float* l = logits + sg * 512; ++ const bool e512 = (a.n_expert == 512u); + float v[16]; +- for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i]; + float mx = -INFINITY; +- for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); +- for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + float sum = 0.0f; +- for (int i = 0; i < 16; ++i) { +- v[i] = exp(v[i] - mx); +- sum += v[i]; ++ if (e512) { ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) { ++ v[i] = l[lane + 32 * i]; ++ mx = max(mx, v[i]); ++ } ++ #pragma clang loop unroll(full) ++ for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) { ++ v[i] = exp(v[i] - mx); ++ sum += v[i]; ++ } ++ #pragma clang loop unroll(full) ++ for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); ++ const float rcp = 1.0f / sum; ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) { ++ v[i] *= rcp; ++ if (isnan(v[i])) v[i] = -FLT_MAX; ++ } ++ } else { ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) { ++ const int idx = lane + 32 * i; ++ v[i] = idx < int(a.n_expert) ? l[idx] : -INFINITY; ++ mx = max(mx, v[i]); ++ } ++ #pragma clang loop unroll(full) ++ for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) { ++ v[i] = (lane + 32 * i < int(a.n_expert)) ? exp(v[i] - mx) : 0.0f; ++ sum += v[i]; ++ } ++ #pragma clang loop unroll(full) ++ for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); ++ const float rcp = 1.0f / sum; ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) { ++ if (lane + 32 * i >= int(a.n_expert)) { ++ v[i] = -INFINITY; ++ } else { ++ v[i] *= rcp; ++ if (isnan(v[i])) v[i] = -FLT_MAX; ++ } ++ } + } +- for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); +- const float rcp = 1.0f / sum; +- for (int i = 0; i < 16; ++i) { +- v[i] *= rcp; +- if (isnan(v[i])) v[i] = -FLT_MAX; ++ const bool draft_moce = (a.pad & 0x80000000u) != 0u && sg > 0; ++ const float base_moce = (a.pad & 0x7FFFFFFFu) != 0u ? as_type(a.pad & 0x7FFFFFFFu) : 1.0f; ++ const float moce_scale = draft_moce ? (base_moce * base_moce) : base_moce; ++ const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f; ++ const int r_prot = draft_moce ? 1 : kMoceProtect; ++ const float r1_scale = kMoceScaleR1; ++ const float rcp_r1 = r1_scale > 1.0f ? (1.0f / r1_scale) : 1.0f; ++ const float r_step_scale = (r1_scale > 1.0f && moce_scale > 1.0f) ? (moce_scale * rcp_r1) : moce_scale; ++ const bool has_r1 = (r1_scale > 1.0f && r_prot > 1); ++ const bool has_moce = (moce_scale > 1.0f || r1_scale > 1.0f) && smap != nullptr; ++ const int min_scaled_rank = has_r1 ? 1 : r_prot; ++ ushort res_mask = 0; ++ if (has_moce) { ++ if (e512) { ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) ++ if (smap[lane + 32 * i] != 0xFFFFFFFFu) res_mask |= ushort(1u << i); ++ } else { ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) { ++ const int idx = lane + 32 * i; ++ if (idx < int(a.n_expert) && smap[idx] != 0xFFFFFFFFu) res_mask |= ushort(1u << i); ++ } ++ } + } ++ ushort scaled_mask = 0; + float selected = 0.0f, ssum = 0.0f; +- for (int rank = 0; rank < 10; ++rank) { +- float best = v[0]; +- int ex = lane; +- for (int i = 1; i < 16; ++i) +- if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } +- for (ushort m = 16; m > 0; m >>= 1) { +- const float ob = simd_shuffle_xor(best, m); +- const int oe = simd_shuffle_xor(ex, m); +- if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } ++ if (!has_moce) { ++ for (int rank = 0; rank < 10; ++rank) { ++ float best = v[0]; ++ int ex = lane; ++ #pragma clang loop unroll(full) ++ for (int i = 1; i < 16; ++i) ++ if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } ++ #pragma clang loop unroll(full) ++ for (ushort m = 16; m > 0; m >>= 1) { ++ const float ob = simd_shuffle_xor(best, m); ++ const int oe = simd_shuffle_xor(ex, m); ++ if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } ++ } ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) v[i] = (ex == lane + 32 * i) ? -INFINITY : v[i]; ++ if ((ex & 31) == lane) { ++ ids[sg * 10 + rank] = ex; ++ ssum += best; ++ } ++ if (rank == lane) selected = best; + } +- if ((ex & 31) == lane) { +- v[ex / 32] = -INFINITY; +- ids[sg * 10 + rank] = ex; +- ssum += best; ++ } else { ++ for (int rank = 0; rank < 10; ++rank) { ++ if (has_r1) { ++ if (rank == 1) { ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) { ++ if ((res_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) { ++ v[i] *= r1_scale; ++ scaled_mask |= ushort(1u << i); ++ } ++ } ++ } else if (rank == r_prot && moce_scale > 1.0f) { ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) { ++ if ((scaled_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) { ++ v[i] *= r_step_scale; ++ } ++ } ++ } ++ } else if (rank == r_prot && moce_scale > 1.0f) { ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) { ++ if ((res_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) { ++ v[i] *= moce_scale; ++ scaled_mask |= ushort(1u << i); ++ } ++ } ++ } ++ float best = v[0]; ++ int ex = lane; ++ #pragma clang loop unroll(full) ++ for (int i = 1; i < 16; ++i) ++ if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } ++ #pragma clang loop unroll(full) ++ for (ushort m = 16; m > 0; m >>= 1) { ++ const float ob = simd_shuffle_xor(best, m); ++ const int oe = simd_shuffle_xor(ex, m); ++ if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } ++ } ++ float true_prob = best; ++ if (rank >= min_scaled_rank && best > 0.0f) { ++ const ushort win_lane = ushort(ex & 31); ++ 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) { ++ // without a tail bonus, the rank-1 scale remains on unselected residents. ++ const bool r1_scaled = has_r1 && (rank < r_prot || moce_scale <= 1.0f); ++ true_prob = best * (r1_scaled ? rcp_r1 : rcp_moce); ++ } ++ } ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) v[i] = (ex == lane + 32 * i) ? -INFINITY : v[i]; ++ if ((ex & 31) == lane) { ++ ids[sg * 10 + rank] = ex; ++ ssum += true_prob; ++ } ++ if (rank == lane) selected = true_prob; + } +- if (rank == lane) selected = best; + } ++ #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + ssum = max(ssum, 6.103515625e-5f); + const float inv = 1.0f / ssum; + if (lane < 10) w[sg * 10 + lane] = selected * inv; + } + threadgroup_barrier(mem_flags::mem_device); +- plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); ++ plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent, smap); + } + + // ---- gate + up + SwiGLU ------------------------------------------------------------------------------------------- + // One pass over a group's rows for all its NE entries (1..8: an expert appears once per token); NE is uniform per + // threadgroup, so the kernel switches once into a body specialised for it. ++// ++// gu_fused: per row, gate and up decoded together and each entry's x slice loaded once for both (the 32-weight units, ++// whose x load is two runs and their sums: the M5 Air, 2560 x 640, NE 3-4: Q4_K / Q5_K / Q6_K / Q8_0 -8..-25%, ++// IQ4_XS -31..-36% from NE 2). The 16-weight units keep two x loads a row: the two decoded units held at once cost ++// them more (IQ2_XS / IQ3_S at NR 1, NE 3: +15..+21%). The same sums in the same order either way. ++template struct gu_fused { ++ enum { value = fmt::UNIT == 32 && (NE >= 3 || F == T_IQ4_XS) }; ++}; ++ ++template ++static inline float q2_combine(ushort2 pk, float d, thread const float* xs, float xsum0) { ++ float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; ++ #pragma clang loop unroll(full) ++ for (int j = 0; j < 8; j += 2) { ++ if (F == T_Q2_0R) { ++ const uint2 q = uint2(pk) >> (2 * j); ++ const uint2 bits = ((q | (q << 14)) & 0x00030003u) | 0x64006400u; ++ const float4 v = float4(as_type(bits) - 1025.0h); ++ s0 = fma(v.x, xs[j], s0); ++ s1 = fma(v.y, xs[j + 1], s1); ++ s2 = fma(v.z, xs[8 + j], s2); ++ s3 = fma(v.w, xs[8 + j + 1], s3); ++ } else { ++ s0 = fma(float(pk.x & ushort(3u << (2 * j))), xs[j], s0); ++ s1 = fma(float(pk.x & ushort(3u << (2 * j + 2))), xs[j + 1], s1); ++ s2 = fma(float(pk.y & ushort(3u << (2 * j))), xs[8 + j], s2); ++ s3 = fma(float(pk.y & ushort(3u << (2 * j + 2))), xs[8 + j + 1], s3); ++ } ++ } ++ const float dot = d * (s0 + s1) + d * (s2 + s3); ++ return F == T_Q2_0R ? (dot - 0.0f) : (dot - (d * xsum0 + 0.0f)); ++} ++ ++static inline float iq_combine(uint4 g, uint s16, float sc, thread const float* xs) { ++ float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; ++ #pragma clang loop unroll(full) ++ for (int k = 0; k < 2; ++k) { ++ const uchar4 v0 = as_type(g[k]), v1 = as_type(g[2 + k]); ++ const uint m0 = s16 >> (4 * k), m1 = s16 >> (8 + 4 * k); ++ #pragma clang loop unroll(full) ++ for (int j = 0; j < 4; j += 2) { ++ s0 = fma((m0 & (1u << j)) ? -float(v0[j]) : float(v0[j]), xs[4 * k + j], s0); ++ s1 = fma((m0 & (2u << j)) ? -float(v0[j + 1]) : float(v0[j + 1]), xs[4 * k + j + 1], s1); ++ s2 = fma((m1 & (1u << j)) ? -float(v1[j]) : float(v1[j]), xs[8 + 4 * k + j], s2); ++ s3 = fma((m1 & (2u << j)) ? -float(v1[j + 1]) : float(v1[j + 1]), xs[8 + 4 * k + j + 1], s3); ++ } ++ } ++ return sc * (s0 + s1) + sc * (s2 + s3) - 0.0f; ++} ++ ++template ++static inline float deq_iq_packed(device const uchar* b, int ui, thread uint4& g, thread uint& s16, threadgroup const uint* tab) { ++ const int ib = ui >> 1, h = ui & 1; ++ if (F == T_IQ2_S) { ++ const uint q2 = *(device const ushort*)(b + 2 + 2 * ui); ++ s16 = *(device const ushort*)(b + 34 + 2 * ui); ++ const uint qh = uint(b[66 + ib]) >> (4 * h); ++ const uint i0 = (q2 & 255) | ((qh << 8) & 0x300); ++ const uint i1 = (q2 >> 8) | ((qh << 6) & 0x300); ++ const uint2 g0 = ctab ? as_type(iq2s_grid[i0]) : ((threadgroup const uint2*)tab)[i0]; ++ const uint2 g1 = ctab ? as_type(iq2s_grid[i1]) : ((threadgroup const uint2*)tab)[i1]; ++ g = uint4(g0.x, g0.y, g1.x, g1.y); ++ const uint sc = b[74 + ib]; ++ return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); ++ } else if (F == T_IQ2_XXS) { ++ device const uchar* qs = b + 2 + 4 * ib + 2 * ui; ++ const uint a1 = rd32(b + 2 + 8 * ib + 4), s0 = a1 >> (14 * h); ++ const uint sa = s0 & 127, sb = (s0 >> 7) & 127; ++ s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12); ++ const uint2 g0 = ctab ? as_type(iq2xxs_grid[qs[0]]) : ((threadgroup const uint2*)tab)[qs[0]]; ++ const uint2 g1 = ctab ? as_type(iq2xxs_grid[qs[1]]) : ((threadgroup const uint2*)tab)[qs[1]]; ++ g = uint4(g0.x, g0.y, g1.x, g1.y); ++ return rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f); ++ } else if (F == T_IQ2_XS) { ++ device const ushort* qs = (device const ushort*)(b + 2) + 2 * ui; ++ const uint v0 = qs[0], v1 = qs[1]; ++ const uint sa = v0 >> 9, sb = v1 >> 9; ++ s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12); ++ const uint2 g0 = ctab ? as_type(iq2xs_grid[v0 & 511]) : ((threadgroup const uint2*)tab)[v0 & 511]; ++ const uint2 g1 = ctab ? as_type(iq2xs_grid[v1 & 511]) : ((threadgroup const uint2*)tab)[v1 & 511]; ++ g = uint4(g0.x, g0.y, g1.x, g1.y); ++ const uint sc = b[66 + ib]; ++ return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); ++ } else if (F == T_IQ3_S) { ++ const int l0 = 2 * h, l1 = 2 * h + 1; ++ const uint qh = b[66 + ib]; ++ s16 = *(device const ushort*)(b + 74 + 2 * ui); ++ const uint i0 = uint(b[2 + 8 * ib + 2 * l0]) | ((qh << (8 - 2 * l0)) & 256); ++ const uint i1 = uint(b[2 + 8 * ib + 2 * l0 + 1]) | ((qh << (7 - 2 * l0)) & 256); ++ const uint i2 = uint(b[2 + 8 * ib + 2 * l1]) | ((qh << (8 - 2 * l1)) & 256); ++ const uint i3 = uint(b[2 + 8 * ib + 2 * l1 + 1]) | ((qh << (7 - 2 * l1)) & 256); ++ g = uint4(g_iq3s(tab, i0), g_iq3s(tab, i1), g_iq3s(tab, i2), g_iq3s(tab, i3)); ++ const uint sc = b[106 + (ib >> 1)]; ++ return rdh(b) * fma(float((sc >> (4 * (ib & 1))) & 15), 2.0f, 1.0f); ++ } else { ++ const uint aux = rd32(b + 2 + 64 + 4 * ib), s0 = aux >> (14 * h); ++ const uint sa = s0 & 127, sb = (s0 >> 7) & 127; ++ s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12); ++ device const uchar* qs = b + 2 + 4 * ui; ++ g = uint4(g_iq3xxs(tab, qs[0]), g_iq3xxs(tab, qs[1]), g_iq3xxs(tab, qs[2]), g_iq3xxs(tab, qs[3])); ++ return rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f); ++ } ++} ++ + template + static inline void gu_body(device const uchar* g0, device const uchar* u0, device const float* x, + device const float* xs16, device float* h, device const int* etok, int c, int row0, int li, + constant MoeArgs& a, threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT; + float ag[NR][NE], au[NR][NE]; ++ #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) ++ #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) ag[r][t] = au[r][t] = 0.0f; + const int U = int(a.n_embd) / UNIT; ++ if (NE == 1) { ++ const int tok = etok[0]; ++ x += tok * int(a.n_embd); ++ xs16 += tok * int(a.n_embd / 16); ++ } + for (int u = li; u < U; u += LPR) { +- float qg[NR][UNIT], qu[NR][UNIT]; +- float4 cg[NR], cu[NR]; +- for (int r = 0; r < NR; ++r) { +- Dec::run(g0 + r * a.gu_row, u, qg[r], cg[r], tab); +- Dec::run(u0 + r * a.gu_row, u, qu[r], cu[r], tab); +- } +- for (int t = 0; t < NE; ++t) { +- const int tok = etok[t]; ++ if (NE == 1 && UNIT == 16 && F != T_Q2_0 && F != T_Q2_0R) { + float xs[UNIT], xsum[2]; +- X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); ++ X::load(x, xs16, u, xs, xsum); ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) { ++ float q[UNIT]; ++ float4 c; ++ Dec::run(g0 + r * a.gu_row, u, q, c, tab); ++ ag[r][0] += unit_combine(q, c, xs, xsum); ++ Dec::run(u0 + r * a.gu_row, u, q, c, tab); ++ au[r][0] += unit_combine(q, c, xs, xsum); ++ } ++ } else if (NE == 1) { ++ float xs[UNIT], xsum[2]; ++ X::load(x, xs16, u, xs, xsum); ++ for (int r = 0; r < NR; ++r) { ++ float q[UNIT]; ++ float4 c; ++ Dec::run(g0 + r * a.gu_row, u, q, c, tab); ++ ag[r][0] += unit_combine(q, c, xs, xsum); ++ Dec::run(u0 + r * a.gu_row, u, q, c, tab); ++ au[r][0] += unit_combine(q, c, xs, xsum); ++ } ++ } else if (F == T_Q2_0 || F == T_Q2_0R) { ++ const int boff = (u >> 2) * 18, qoff = 2 + 4 * (u & 3); ++ ushort2 pkg[NR], pku[NR]; ++ float dg[NR], du[NR]; ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) { ++ device const uchar* bg = g0 + r * a.gu_row + boff; ++ device const uchar* bu = u0 + r * a.gu_row + boff; ++ device const ushort* pg = (device const ushort*)(bg + qoff); ++ device const ushort* pu = (device const ushort*)(bu + qoff); ++ pkg[r] = ushort2(pg[0], pg[1]); ++ pku[r] = ushort2(pu[0], pu[1]); ++ dg[r] = rdh(bg); ++ du[r] = rdh(bu); ++ } ++ #pragma clang loop unroll(full) ++ for (int t = 0; t < NE; ++t) { ++ const int tok = etok[t]; ++ float xs[UNIT], xsum[2]; ++ X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) { ++ ag[r][t] += q2_combine(pkg[r], dg[r], xs, xsum[0]); ++ au[r][t] += q2_combine(pku[r], du[r], xs, xsum[0]); ++ } ++ } ++ } else if (gu_fused::value) { ++ for (int r = 0; r < NR; ++r) { ++ float qg[UNIT], qu[UNIT]; ++ float4 cg, cu; ++ Dec::run(g0 + r * a.gu_row, u, qg, cg, tab); ++ Dec::run(u0 + r * a.gu_row, u, qu, cu, tab); ++ for (int t = 0; t < NE; ++t) { ++ const int tok = etok[t]; ++ float xs[UNIT], xsum[2]; ++ X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); ++ ag[r][t] += unit_combine(qg, cg, xs, xsum); ++ au[r][t] += unit_combine(qu, cu, xs, xsum); ++ } ++ } ++ } else if (NR == 1 && NE == 4 && (F == T_IQ2_S || F == T_IQ2_XXS || F == T_IQ2_XS || F == T_IQ3_S || F == T_IQ3_XXS)) { ++ const int boff = (u >> 4) * fmt::BB, ui = (LPR % 16 == 0) ? (li & 15) : (u & 15); ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) { ++ uint4 gg, gu; ++ uint sg, su; ++ const float dg = deq_iq_packed(g0 + r * a.gu_row + boff, ui, gg, sg, tab); ++ const float du = deq_iq_packed(u0 + r * a.gu_row + boff, ui, gu, su, tab); ++ #pragma clang loop unroll(full) ++ for (int t = 0; t < NE; ++t) { ++ const int tok = etok[t]; ++ float xs[UNIT], xsum[2]; ++ X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); ++ ag[r][t] += iq_combine(gg, sg, dg, xs); ++ au[r][t] += iq_combine(gu, su, du, xs); ++ } ++ } ++ } else { ++ #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) { +- ag[r][t] += unit_combine(qg[r], cg[r], xs, xsum); +- au[r][t] += unit_combine(qu[r], cu[r], xs, xsum); ++ float q[UNIT]; ++ float4 c; ++ Dec::run(g0 + r * a.gu_row, u, q, c, tab); ++ #pragma clang loop unroll(full) ++ for (int t = 0; t < NE; ++t) { ++ const int tok = etok[t]; ++ float xs[UNIT], xsum[2]; ++ X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); ++ ag[r][t] += unit_combine(q, c, xs, xsum); ++ } ++ Dec::run(u0 + r * a.gu_row, u, q, c, tab); ++ #pragma clang loop unroll(full) ++ for (int t = 0; t < NE; ++t) { ++ const int tok = etok[t]; ++ float xs[UNIT], xsum[2]; ++ X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); ++ au[r][t] += unit_combine(q, c, xs, xsum); ++ } + } + } + } ++ #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) ++ #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const float gv = lane_sum(ag[r][t]), uv = lane_sum(au[r][t]); + if (li == 0) h[(c + t) * int(a.n_ff) + row0 + r] = gv / (1.0f + exp(-gv)) * uv; + } + } + ++// function constant 27, with 26's plan: 1 = only the groups resident at routing time (indirect grid P_IGU_R / P_IDN_R), ++// 2 = only the missing ones (P_IGU_M / P_IDN_M); 0 or unset = all (P_IGU / P_IDN) ++constant int kPhaseFc [[function_constant(27)]]; ++constant int kPhase = is_function_constant_defined(kPhaseFc) ? kPhaseFc : 0; ++ + // bench probe (function constant 22): 1 = no table load (wrong values for the i-quants) + constant int MOE_PROBE [[function_constant(22)]]; + constant int moe_probe = is_function_constant_defined(MOE_PROBE) ? MOE_PROBE : 0; +@@ -213,18 +593,19 @@ kernel void moe_gu(device const uchar* W [[buffer(0)]], device const float* x [[ + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; +- if (moe_probe != 1) { ++ if (fmt::TAB > 1 && moe_probe != 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } +- const int g = int(tg.y); +- if (g >= plan[P_NG]) return; ++ const int g = int(tg.y) + (kPhase == 2 ? plan[P_SPLIT] : 0); ++ if (g >= (kPhase == 1 ? plan[P_SPLIT] : plan[P_NG])) return; + const int li = lane % LPR; + const int row0 = ((int(tg.x) * SG + sg) * RG + lane / LPR) * NR; + const ulong ex = (ulong)plan[P_GEXP + g]; + const int c = plan[P_GSTART + g], ne = plan[P_GSTART + g + 1] - c; + device const int* etok = plan + P_ETOK + c; +- device const uchar* E = kArena ? arena_slot(smap[ex], W, W1, W2, W3) : W + ex * a.gu_stride; ++ const uint scode = (kArena && kPhase == 1) ? uint(plan[P_GSLOT + g]) : (kArena ? smap[ex] : 0u); ++ device const uchar* E = kArena ? arena_slot(scode, W, W1, W2, W3) : W + ex * a.gu_stride; + 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) { +@@ -250,20 +631,59 @@ static inline void dn_body(device const uchar* d0, device const float* h, device + threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT; + float ac[NR][NE]; ++ #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) ++ #pragma clang loop unroll(full) + 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) { +- float q[NR][UNIT]; +- float4 sc[NR]; +- for (int r = 0; r < NR; ++r) Dec::run(d0 + r * a.d_row, u, q[r], sc[r], tab); +- for (int t = 0; t < NE; ++t) { ++ if (F == T_Q2_0 || F == T_Q2_0R) { ++ device const uchar* bu = d0 + (u >> 2) * 18; ++ const int qoff = 2 + 4 * (u & 3); ++ ushort2 pk[NR]; ++ float d[NR]; ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) { ++ device const uchar* b = bu + r * a.d_row; ++ device const ushort* p = (device const ushort*)(b + qoff); ++ pk[r] = ushort2(p[0], p[1]); ++ d[r] = rdh(b); ++ } ++ #pragma clang loop unroll(full) ++ for (int t = 0; t < NE; ++t) { ++ float xs[UNIT], xsum[2]; ++ X::load(h + (c + t) * int(a.n_ff), hs16 + (c + t) * int(a.n_ff / 16), u, xs, xsum); ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) ++ ac[r][t] += q2_combine(pk[r], d[r], xs, xsum[0]); ++ } ++ } else if (NE == 1) { + float xs[UNIT], xsum[2]; +- X::load(h + (c + t) * int(a.n_ff), hs16 + (c + t) * int(a.n_ff / 16), u, xs, xsum); +- for (int r = 0; r < NR; ++r) ac[r][t] += unit_combine(q[r], sc[r], xs, xsum); ++ X::load(h + c * int(a.n_ff), hs16 + c * int(a.n_ff / 16), u, xs, xsum); ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) { ++ float q[UNIT]; ++ float4 sc; ++ Dec::run(d0 + r * a.d_row, u, q, sc, tab); ++ ac[r][0] += unit_combine(q, sc, xs, xsum); ++ } ++ } else { ++ float q[NR][UNIT]; ++ float4 sc[NR]; ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) Dec::run(d0 + r * a.d_row, u, q[r], sc[r], tab); ++ #pragma clang loop unroll(full) ++ for (int t = 0; t < NE; ++t) { ++ float xs[UNIT], xsum[2]; ++ X::load(h + (c + t) * int(a.n_ff), hs16 + (c + t) * int(a.n_ff / 16), u, xs, xsum); ++ #pragma clang loop unroll(full) ++ for (int r = 0; r < NR; ++r) ac[r][t] += unit_combine(q[r], sc[r], xs, xsum); ++ } + } + } ++ #pragma clang loop unroll(full) + for (int r = 0; r < NR; ++r) ++ #pragma clang loop unroll(full) + for (int t = 0; t < NE; ++t) { + const float v = lane_sum(ac[r][t]); + if (li == 0) ye[(c + t) * int(a.n_embd) + row0 + r] = as_type(plan[P_EW + c + t]) * v; +@@ -280,15 +700,18 @@ kernel void moe_down(device const uchar* W [[buffer(0)]], device const float* h + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; +- load_tab(tab, tid, SG * 32); +- threadgroup_barrier(mem_flags::mem_threadgroup); +- const int g = int(tg.y); +- if (g >= plan[P_NG]) return; ++ if (fmt::TAB > 1) { ++ load_tab(tab, tid, SG * 32); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } ++ const int g = int(tg.y) + (kPhase == 2 ? plan[P_SPLIT] : 0); ++ if (g >= (kPhase == 1 ? plan[P_SPLIT] : plan[P_NG])) return; + const int li = lane % LPR; + const int row0 = ((int(tg.x) * SG + sg) * RG + lane / LPR) * NR; + const ulong ex = (ulong)plan[P_GEXP + g]; + const int c = plan[P_GSTART + g], ne = plan[P_GSTART + g + 1] - c; +- device const uchar* d0 = (kArena ? arena_slot(smap[ex], W, W1, W2, W3) : W + ex * a.d_stride) + a.down_off + (ulong)row0 * a.d_row; ++ const uint scode = (kArena && kPhase == 1) ? uint(plan[P_GSLOT + g]) : (kArena ? smap[ex] : 0u); ++ device const uchar* d0 = (kArena ? arena_slot(scode, W, W1, W2, W3) : W + ex * a.d_stride) + a.down_off + (ulong)row0 * a.d_row; + switch (ne) { + case 1: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + case 2: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; +@@ -310,8 +733,14 @@ kernel void moe_combine(device const float* ye [[buffer(0)]], device float* y [[ + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.n_embd, k = gid % a.n_embd; + if (t >= a.n_tok) return; ++ device const int* te = plan + P_TOKENT + t * a.topk; + float s = 0.0f; +- for (uint j = 0; j < a.topk; ++j) s += ye[plan[P_TOKENT + t * a.topk + j] * int(a.n_embd) + k]; ++ if (a.topk == 10) { ++#pragma clang loop unroll(full) ++ for (uint j = 0; j < 10; ++j) s += ye[te[j] * int(a.n_embd) + k]; ++ } else { ++ for (uint j = 0; j < a.topk; ++j) s += ye[te[j] * int(a.n_embd) + k]; ++ } + y[gid] = s; + } + +diff --git a/metal/kernels/moe_mm.metal b/metal/kernels/moe_mm.metal +index ff28897a..03ac90a8 100644 +--- a/metal/kernels/moe_mm.metal ++++ b/metal/kernels/moe_mm.metal +@@ -32,28 +32,26 @@ kernel void mm_hist(device const int* ids [[buffer(0)]], device uint* bcnt [[buf + constant MmPlanArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup atomic_uint h[1024]; +- for (uint e = tid; e < a.n_expert; e += 1024) atomic_store_explicit(&h[e], 0u, memory_order_relaxed); ++ if (tid < a.n_expert) atomic_store_explicit(&h[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint i = tg.x * 1024 + tid; +- if (i < a.n_ent) atomic_fetch_add_explicit(&h[ids[i]], 1u, memory_order_relaxed); ++ if (i < a.n_ent && ids[i] >= 0) atomic_fetch_add_explicit(&h[ids[i]], 1u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); +- for (uint e = tid; e < a.n_expert; e += 1024) +- bcnt[tg.x * a.n_expert + e] = atomic_load_explicit(&h[e], memory_order_relaxed); ++ if (tid < a.n_expert) ++ bcnt[tg.x * a.n_expert + tid] = atomic_load_explicit(&h[tid], memory_order_relaxed); + } + + // one threadgroup of n_expert (<= 1024) threads. boff[block][e] = expert e's rows before this block (within e); + // eoff[e] = expert e's first padded row; tiles[i] = (expert, row0); meta[0] = n_tiles, meta[1] = padded rows; + // indirect grids: meta[2..4] = (gx_gu, n_tiles, 1) for the gate/up GEMM, meta[5..7] = (gx_down, n_tiles, 1) for the + // down GEMM (gx_gu in meta[2], gx_down in meta[5] written by the host once), meta[8..10] = (rows, 1, 1) for mm_gather. ++// Padding rows [r0 + tot, r0 + rows) of srt_tok (when bound at buffer 6) are cleared to -1 here so mm_clear is not needed. + kernel void mm_scan(device const uint* bcnt [[buffer(0)]], device uint* boff [[buffer(1)]], + device uint* eoff [[buffer(2)]], device int2* tiles [[buffer(3)]], device uint* meta [[buffer(4)]], +- constant MmPlanArgs& a [[buffer(5)]], ++ constant MmPlanArgs& a [[buffer(5)]], device int* srt_tok [[buffer(6)]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], +- ushort lane [[thread_index_in_simdgroup]]) { ++ ushort lane [[thread_index_in_simdgroup]], ushort nsg [[simdgroups_per_threadgroup]]) { + threadgroup uint s_rows[32], s_tiles[32]; +- // simdgroups past the launched ones must count 0 (512 threads fill only 16 of the 32 slots) +- if (tid < 32) { s_rows[tid] = 0; s_tiles[tid] = 0; } +- threadgroup_barrier(mem_flags::mem_threadgroup); + const uint e = tid; + uint tot = 0; + if (e < a.n_expert) +@@ -66,8 +64,8 @@ kernel void mm_scan(device const uint* bcnt [[buffer(0)]], device uint* boff [[b + if (lane == 31) { s_rows[sg] = pr + rows; s_tiles[sg] = pt + nt; } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint orow = 0, otile = 0, all_rows = 0, all_tiles = 0; +- for (int i = 0; i < 32; ++i) { +- if (i < int(sg)) { orow += s_rows[i]; otile += s_tiles[i]; } ++ for (uint i = 0; i < nsg; ++i) { ++ if (i < sg) { orow += s_rows[i]; otile += s_tiles[i]; } + all_rows += s_rows[i]; + all_tiles += s_tiles[i]; + } +@@ -75,6 +73,8 @@ kernel void mm_scan(device const uint* bcnt [[buffer(0)]], device uint* boff [[b + const uint r0 = orow + pr, t0 = otile + pt; + eoff[e] = r0; + for (uint k = 0; k < nt; ++k) tiles[t0 + k] = int2(int(e), int(r0 + k * a.tm)); ++ if (srt_tok) ++ for (uint r = tot; r < rows; ++r) srt_tok[r0 + r] = -1; + } + if (tid == 0) { + meta[0] = all_tiles; +@@ -89,21 +89,56 @@ kernel void mm_scan(device const uint* bcnt [[buffer(0)]], device uint* boff [[b + } + } + +-// grid (n_blocks), 1024 threads. srt_* must be cleared to the padding value (-1 / 0) beforehand by mm_clear. ++static inline uint mm_sum_nibbles(uint c) { ++ c = (c & 0x0F0F0F0Fu) + ((c >> 4) & 0x0F0F0F0Fu); ++ c = (c + (c >> 8)) & 0x00FF00FFu; ++ return (c + (c >> 16)) & 0xFFFFu; ++} ++ ++// grid (n_blocks), 1024 threads. Two-level deterministic intra-block rank: 32-lane sub-block nibble counts in p_cnt ++// plus simd_broadcast scan within the caller's 32-lane simdgroup (no s_id[1024] threadgroup array). + kernel void mm_place(device const int* ids [[buffer(0)]], device const float* w [[buffer(1)]], + device const uint* boff [[buffer(2)]], device const uint* eoff [[buffer(3)]], + device int* srt_tok [[buffer(4)]], device float* srt_w [[buffer(5)]], device int* pos [[buffer(6)]], + constant MmPlanArgs& a [[buffer(7)]], +- uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { +- threadgroup int s_id[1024]; +- const uint i = tg.x * 1024 + tid; ++ uint tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], ++ ushort lane [[thread_index_in_simdgroup]]) { ++ threadgroup atomic_uint p_cnt[4096]; ++ const uint n4e = 4 * a.n_expert; ++ if (a.n_expert == 512) { ++ atomic_store_explicit(&p_cnt[tid], 0u, memory_order_relaxed); ++ atomic_store_explicit(&p_cnt[tid + 1024], 0u, memory_order_relaxed); ++ } else { ++ for (uint k = tid; k < n4e; k += 1024) atomic_store_explicit(&p_cnt[k], 0u, memory_order_relaxed); ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ const uint i = tg * 1024 + tid; + const int e = i < a.n_ent ? ids[i] : -1; +- s_id[tid] = e; ++ const uint sub = tid >> 5; ++ if (e >= 0 && sub < 31) { ++ const uint idx = (sub >> 3) * a.n_expert + uint(e); ++ const uint sh = (sub & 7u) * 4u; ++ atomic_fetch_add_explicit(&p_cnt[idx], 1u << sh, memory_order_relaxed); ++ } ++ uint rank = 0; ++ #pragma clang loop unroll(full) ++ for (ushort l = 0; l < 32; ++l) { ++ const int el = simd_broadcast(e, l); ++ if (l < lane && el == e) ++rank; ++ } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (i >= a.n_ent) return; +- uint rank = 0; +- for (uint j = 0; j < tid; ++j) rank += s_id[j] == e; +- const uint row = eoff[e] + boff[tg.x * a.n_expert + e] + rank; ++ if (e < 0) { pos[i] = -1; return; } ++ for (uint g = 0; g < (sub >> 3); ++g) { ++ const uint c = atomic_load_explicit(&p_cnt[g * a.n_expert + uint(e)], memory_order_relaxed); ++ rank += mm_sum_nibbles(c); ++ } ++ if (sub & 7u) { ++ const uint m = (1u << ((sub & 7u) * 4u)) - 1u; ++ const uint c = atomic_load_explicit(&p_cnt[(sub >> 3) * a.n_expert + uint(e)], memory_order_relaxed) & m; ++ rank += mm_sum_nibbles(c); ++ } ++ const uint row = eoff[e] + boff[tg * a.n_expert + e] + rank; + srt_tok[row] = int(i / a.topk); + srt_w[row] = w[i]; + pos[i] = int(row); +@@ -113,6 +148,7 @@ kernel void mm_clear(device int* srt_tok [[buffer(0)]], constant uint& n [[buffe + if (gid < n) srt_tok[gid] = -1; + } + ++#ifndef MM_GATHER_ONLY + // ---- the GEMMs ----------------------------------------------------------------------------------------------------- + // weights of one expert: gate at W + gate_off + e * gu_stride (+ up_off for up), down at W + down_off + e * d_stride + struct MmArgs { +@@ -138,6 +174,7 @@ static inline void stage_w(device const uchar* w0, uint row_bytes, int k0, threa + float wv[UNIT]; + int p0, p1; + DW::run(w0 + ulong(r) * row_bytes, k0 / UNIT + uu, k0, wv, p0, p1, tab); ++#pragma clang loop unroll(full) + for (int j = 0; j < H; ++j) { + dst[r * TK + p0 + j] = half(wv[j]); + dst[r * TK + p1 + j] = half(wv[H + j]); +@@ -145,6 +182,7 @@ static inline void stage_w(device const uchar* w0, uint row_bytes, int k0, threa + } + } + ++#ifndef MM_V3_ONLY + template + kernel void mm_gu(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]], + device float* H [[buffer(2)]], device const int* srt_tok [[buffer(3)]], +@@ -236,6 +274,7 @@ kernel void mm_down(device const uchar* W [[buffer(0)]], device const float* H [ + } + + // grid (T * n_embd / 256), 256 threads ++#endif + kernel void mm_combine(device const float* YE [[buffer(0)]], device const int* pos [[buffer(1)]], + device const float* srt_w [[buffer(2)]], device float* y [[buffer(3)]], + constant uint4& a [[buffer(4)]], // (T, n_embd, topk, -) +@@ -245,11 +284,12 @@ kernel void mm_combine(device const float* YE [[buffer(0)]], device const int* p + float s = 0.0f; + for (uint k = 0; k < a.z; ++k) { + const int r = pos[t * a.z + k]; +- s = fma(srt_w[r], YE[ulong(r) * a.y + d], s); ++ if (r >= 0) s = fma(srt_w[r], YE[ulong(r) * a.y + d], s); + } + y[gid] = s; + } + ++#ifndef MM_V3_ONLY + #define MMGU_(F, M, N, K, S) template [[host_name("mm_gu_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu) mm_gu; + #define MMDN_(F, M, N, K, S) template [[host_name("mm_down_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down) mm_down; + #define MMGU(F, M, N, K, S) MMGU_(F, M, N, K, S) +@@ -258,7 +298,9 @@ kernel void mm_combine(device const float* YE [[buffer(0)]], device const int* p + #define MM_DN_SET(F) MMDN(F, 32, 64, 64, 4) MMDN(F, 64, 64, 64, 4) MMDN(F, 16, 64, 64, 2) MMDN(F, 32, 32, 64, 2) + MM_GU_SET(T_Q2_0) MM_GU_SET(T_IQ3_XXS) MM_GU_SET(T_IQ2_XXS) MM_GU_SET(T_IQ2_XS) MM_GU_SET(T_IQ2_S) MM_GU_SET(T_IQ3_S) + MM_DN_SET(T_Q2_0) MM_DN_SET(T_IQ4_NL) ++#endif + ++#endif + // ---- v2: activations gathered once per layer as fp16 rows, read by matmul2d from device memory --------------------- + // mm_gather: XS[row][k] = half(x[srt_tok[row]][k]) (0 on padding rows). grid (rows), 256 threads. + kernel void mm_gather(device const float* x [[buffer(0)]], device const int* srt_tok [[buffer(1)]], +@@ -266,9 +308,57 @@ kernel void mm_gather(device const float* x [[buffer(0)]], device const int* srt + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const int t = srt_tok[tg.x]; + device half* o = XS + ulong(tg.x) * n; +- for (uint k = tid; k < n; k += 256) o[k] = t >= 0 ? half(x[ulong(t) * n + k]) : half(0.0f); ++ if (n == 2560) { ++ #pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) { ++ const uint k = tid + i * 256; ++ o[k] = t >= 0 ? half(x[ulong(t) * n + k]) : half(0.0f); ++ } ++ } else { ++ for (uint k = tid; k < n; k += 256) o[k] = t >= 0 ? half(x[ulong(t) * n + k]) : half(0.0f); ++ } ++} ++ ++// mm_gather_h: gathers straight from fp16 rows (mixh) with 128-bit uint4 copies across nth threads per row. ++kernel void mm_gather_h(device const half* xh [[buffer(0)]], device const int* srt_tok [[buffer(1)]], ++ device half* XS [[buffer(2)]], constant uint& n [[buffer(3)]], ++ uint tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], ++ ushort nth [[threads_per_threadgroup]]) { ++ const int t = srt_tok[tg]; ++ device uint4* o8 = (device uint4*)(XS + ulong(tg) * n); ++ const uint n8 = n >> 3; ++ if (n == 2560) { ++ device const uint4* x8 = (device const uint4*)(xh + ulong(max(t, 0)) * n); ++ if (nth == 32) { ++ #pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) o8[tid + i * 32] = t >= 0 ? x8[tid + i * 32] : uint4(0u); ++ return; ++ } ++ if (nth == 64) { ++ #pragma clang loop unroll(full) ++ for (uint i = 0; i < 5; ++i) o8[tid + i * 64] = t >= 0 ? x8[tid + i * 64] : uint4(0u); ++ return; ++ } ++ if (nth == 160) { ++ o8[tid] = t >= 0 ? x8[tid] : uint4(0u); ++ o8[tid + 160] = t >= 0 ? x8[tid + 160] : uint4(0u); ++ return; ++ } ++ if (nth == 256) { ++ o8[tid] = t >= 0 ? x8[tid] : uint4(0u); ++ if (tid < 64) o8[tid + 256] = t >= 0 ? x8[tid + 256] : uint4(0u); ++ return; ++ } ++ } ++ if (t >= 0) { ++ device const uint4* x8 = (device const uint4*)(xh + ulong(t) * n); ++ for (uint k = tid; k < n8; k += nth) o8[k] = x8[k]; ++ } else { ++ for (uint k = tid; k < n8; k += nth) o8[k] = uint4(0u); ++ } + } + ++#ifndef MM_GATHER_ONLY + // The SSD-streaming arena (as in moe.metal): with function constant 20, expert e's slot is smap[e] in the chunks W, + // W1..W3 (buffers 5..8 of mm_gu2 / mm_down2 / mm_gu3 / mm_down3); gate_off / up_off / down_off are offsets in the slot. + #ifndef STRATA_ARENA_METAL +@@ -287,6 +377,7 @@ static inline device const uchar* arena_slot(uint v, device const uchar* W, devi + constant int MM_PROBE [[function_constant(5)]]; + constant int mm_probe = is_function_constant_defined(MM_PROBE) ? MM_PROBE : 0; + ++#ifndef MM_V3_ONLY + template + kernel void mm_gu2(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], +@@ -396,6 +487,8 @@ MM_GU2_SET(T_Q2_0) MM_GU2_SET(T_IQ3_XXS) MM_GU2_SET(T_IQ2_XXS) MM_GU2_SET(T_IQ2_ + MM_DN2_SET(T_Q2_0) + MMDN2(T_IQ4_NL, 16, 64, 32, 2) MMDN2(T_IQ4_NL, 16, 64, 64, 2) MMDN2(T_IQ4_NL, 16, 32, 32, 2) MMDN2(T_IQ4_NL, 32, 32, 32, 2) MMDN2(T_IQ4_NL, 16, 128, 32, 4) + ++#endif ++ + // ---- v3: the weights decoded straight into matmul2d's right-input cooperative tensor (single-simdgroup scope), + // no threadgroup staging and no barriers in the K loop. Each simdgroup owns TN output rows; a threadgroup SG * TN. + // The right input's element layout (tests/coop_layout): runs of 4 consecutive k at one n, so a lane decodes 4 values +@@ -422,14 +515,20 @@ static inline half4 w4_signed(uint g, uint s, float sc) { + (s & 4) ? -float(v[2]) : float(v[2]), (s & 8) ? -float(v[3]) : float(v[3])); + return half4(q * sc); + } ++ + template <> struct W4 { // 64 values a block: fp16 d, 16 bytes of 2-bit codes (value v at byte v / 4) +- // fp32 math then fp16 (an fp16-math decode was measured slower: 4.2 vs 3.5 ms a layer) ++ // the codes c land in fp16 mantissas (1024 + c), so c - 1 and its product with d are exact in fp16: bit for bit the ++ // fp32 decode for every d and code byte (checked exhaustively), and 8-11% off mm_down3 at T 1024-2048 ++ static inline half4 dec(uint q, half d) { ++ const uint2 c = (uint2(q | (q << 14), (q >> 4) | (q << 10)) & 0x00030003u) | 0x64006400u; ++ return (as_type(c) - 1025.0h) * d; ++ } ++ template ++ static inline half4 run_q(device const uchar* b, uint kl4) { ++ return dec((b + kl4)[OFF], *(device const half*)b); ++ } + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { +- device const uchar* b = row + (k >> 6) * 18; +- const float d = float(*(device const half*)b); +- const uint q = b[2 + ((k & 63) >> 2)]; +- return half4(float4(float(int(q & 3) - 1), float(int((q >> 2) & 3) - 1), float(int((q >> 4) & 3) - 1), +- float(int(q >> 6) - 1)) * d); ++ return run_q<2>(row + (k >> 6) * 18, uint(k & 63) >> 2); + } + }; + template <> struct W4 { +@@ -437,10 +536,8 @@ template <> struct W4 { + device const uchar* b = row + (k >> 8) * 98; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint aux = rd32(b + 2 + 64 + 4 * ib); +- threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 256); +- const uint s = ks[(aux >> (7 * l)) & 127]; +- const float sc = rdh(b) * (0.5f + float(aux >> 28)) * 0.5f; +- return w4_signed(tab[b[2 + 8 * ib + 2 * l + hh]], hh ? s >> 4 : s, sc); ++ const float sc = rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f); ++ return w4_signed(g_iq3xxs(tab, b[2 + ((k & 255) >> 2)]), sign4_iq2xs((aux >> (7 * l)) & 127, hh), sc); + } + }; + template <> struct W4 { +@@ -448,47 +545,43 @@ template <> struct W4 { + device const uchar* b = row + (k >> 8) * 66; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint a0 = rd32(b + 2 + 8 * ib), a1 = rd32(b + 2 + 8 * ib + 4); +- threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 512); +- const uint idx = (a0 >> (8 * l)) & 255, s = ks[(a1 >> (7 * l)) & 127]; +- const float sc = rdh(b) * (0.5f + float(a1 >> 28)) * 0.25f; +- return w4_signed(tab[2 * idx + hh], hh ? s >> 4 : s, sc); ++ const uint idx = (a0 >> (8 * l)) & 255; ++ const float sc = rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f); ++ return w4_signed(g_iq2xxs(tab, idx, hh), sign4_iq2xs((a1 >> (7 * l)) & 127, hh), sc); + } + }; + template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 74; +- const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; +- const uint v = ((device const ushort*)(b + 2))[4 * ib + l]; +- threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 1024); +- const uint idx = v & 511, s = ks[v >> 9]; ++ const int ib = (k & 255) >> 5, hh = (k & 7) >> 2; ++ const uint v = ((device const ushort*)(b + 2))[(k & 255) >> 3]; + const uint scb = b[66 + ib]; +- const float sc = rdh(b) * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f; +- return w4_signed(tab[2 * idx + hh], hh ? s >> 4 : s, sc); ++ const float sc = rdh(b) * fma(float((scb >> ((k & 16) >> 2)) & 15), 0.25f, 0.125f); ++ return w4_signed(g_iq2xs(tab, v & 511, hh), sign4_iq2xs(v >> 9, hh), sc); + } + }; + template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 82; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; +- const uint qh = b[66 + ib]; +- const uint idx = uint(b[2 + 4 * ib + l]) | ((qh << (8 - 2 * l)) & 0x300); +- const uint s = b[2 + 32 + 4 * ib + l]; ++ const uint qh = b[66 + ib], o = (k & 255) >> 3; ++ const uint idx = uint(b[2 + o]) | ((qh << (8 - 2 * l)) & 0x300); ++ const uint s = b[34 + o]; + const uint scb = b[74 + ib]; +- const float sc = rdh(b) * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f; +- return w4_signed(tab[2 * idx + hh], hh ? s >> 4 : s, sc); ++ const float sc = rdh(b) * fma(float((scb >> ((k & 16) >> 2)) & 15), 0.25f, 0.125f); ++ return w4_signed(g_iq2s(tab, idx, hh), hh ? s >> 4 : s, sc); + } + }; + template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 110; +- const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; ++ const int ib = (k & 255) >> 5, hh = (k & 7) >> 2; + const uint qh = b[66 + ib]; +- const uint i = hh ? (uint(b[2 + 8 * ib + 2 * l + 1]) | ((qh << (7 - 2 * l)) & 256)) +- : (uint(b[2 + 8 * ib + 2 * l]) | ((qh << (8 - 2 * l)) & 256)); +- const uint s = b[74 + 4 * ib + l]; ++ const uint i = uint(b[2 + ((k & 255) >> 2)]) | ((qh << (8 - ((k & 31) >> 2))) & 256u); ++ const uint s = b[74 + ((k & 255) >> 3)]; + const uint scb = b[106 + (ib >> 1)]; + const float sc = rdh(b) * float(1 + 2 * ((scb >> (4 * (ib & 1))) & 15)); +- return w4_signed(tab[i], hh ? s >> 4 : s, sc); ++ return w4_signed(g_iq3s(tab, i), hh ? s >> 4 : s, sc); + } + }; + template <> struct W4 { // 256 a block: d, scales_h, scales_l[4], qs[128]; 32-groups as IQ4_NL's +@@ -580,17 +673,12 @@ template <> struct W4 { // 256 a block: d, dmin, scales[12], qh[32], q + template <> struct W4 { // 256 a block: ql[128], qh[64], int8 scales[16], d; halves of 128 + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 8) * 210; +- const int j = k & 255, n = j >> 7, jj = j & 127, qd = jj >> 5, l = jj & 31; +- device const uchar* ql = b + 64 * n + l + ((qd & 1) ? 32 : 0); +- device const uchar* qh = b + 128 + 32 * n + l; +- const float d = rdh(b + 208); +- const float ds = d * float(((device const char*)(b + 192 + 8 * n))[l / 16 + 2 * qd]); +- float4 w; +- for (int i = 0; i < 4; ++i) { +- const int q = int(((qd >> 1) ? (ql[i] >> 4) : (ql[i] & 15)) | (((qh[i] >> (2 * qd)) & 3) << 4)) - 32; +- w[i] = ds * float(q); +- } +- return half4(w); ++ const uint j = uint(k & 255), j4 = j >> 4; ++ const uint la = rd32(b + (((j >> 1) & 64u) | (j & 63u))); ++ const uint ha = rd32(b + 128u + (((j >> 2) & 32u) | (j & 31u))); ++ const float ds = rdh(b + 208) * float(((device const char*)(b + 192))[j4]); ++ const uint qa = ((la >> (j4 & 4u)) & 0x0F0F0F0Fu) | (((ha >> (j4 & 6u)) & 0x03030303u) << 4); ++ return half4(ds * (float4(as_type(qa)) - 32.0f)); + } + }; + +@@ -607,88 +695,102 @@ template <> struct W8 { + const P bl = row + (k >> 8) * 98; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint aux = rd32(bl + 2 + 64 + 4 * ib); +- threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 256); +- const float sc = rdh(bl) * (0.5f + float(aux >> 28)) * 0.5f; +- const P qs = bl + 2 + 8 * ib + hh; ++ const float sc = rdh(bl) * fma(float(aux >> 28), 0.5f, 0.25f); ++ const P qs = bl + 2 + ((k & 255) >> 2); + if (mm_probe == 3) { // bench probe: no threadgroup-table lookups (wrong values) +- a = w4_signed(qs[2 * l] * 0x01010101u, aux, sc); +- b = w4_signed(qs[2 * l + 4] * 0x01010101u, aux >> 7, sc); ++ a = w4_signed(qs[0] * 0x01010101u, aux, sc); ++ b = w4_signed(qs[4] * 0x01010101u, aux >> 7, sc); + return; + } +- uint s0, s1; +- if (mm_probe == 4) { // the signs computed (ksigns_iq2xs[i] = i | parity(i) << 7) instead of looked up +- const uint i0 = (aux >> (7 * l)) & 127, i1 = (aux >> (7 * (l + 2))) & 127; +- s0 = i0 | ((popcount(i0) & 1u) << 7); +- s1 = i1 | ((popcount(i1) & 1u) << 7); +- } else { +- s0 = ksg(tab, 256, (aux >> (7 * l)) & 127); +- s1 = ksg(tab, 256, (aux >> (7 * (l + 2))) & 127); +- } +- a = w4_signed(g_iq3xxs(tab, qs[2 * l]), hh ? s0 >> 4 : s0, sc); +- b = w4_signed(g_iq3xxs(tab, qs[2 * l + 4]), hh ? s1 >> 4 : s1, sc); ++ const uint j0 = (aux >> (7 * l)) & 127, j1 = (aux >> (7 * (l + 2))) & 127; ++ a = w4_signed(g_iq3xxs(tab, qs[0]), sign4_iq2xs(j0, hh), sc); ++ b = w4_signed(g_iq3xxs(tab, qs[4]), sign4_iq2xs(j1, hh), sc); + } + }; + template <> struct W8 { +- static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { +- device const uchar* bl = row + (k >> 8) * 66; +- const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; +- const uint a0 = rd32(bl + 2 + 8 * ib), a1 = rd32(bl + 2 + 8 * ib + 4); +- threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 512); +- const float sc = rdh(bl) * (0.5f + float(a1 >> 28)) * 0.25f; +- const uint i0 = (a0 >> (8 * l)) & 255, i1 = (a0 >> (8 * (l + 2))) & 255; +- const uint s0 = ks[(a1 >> (7 * l)) & 127], s1 = ks[(a1 >> (7 * (l + 2))) & 127]; +- a = w4_signed(tab[2 * i0 + hh], hh ? s0 >> 4 : s0, sc); +- b = w4_signed(tab[2 * i1 + hh], hh ? s1 >> 4 : s1, sc); ++ template ++ static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { ++ const int ib = (k & 255) >> 5, l = (k & 31) >> 3; ++ const P p = bl + 2 + 8 * ib; ++ const uint a1 = rd32(p + 4); ++ const float sc = d * fma(float(a1 >> 28), 0.25f, 0.125f); ++ const uint i0 = p[l], i1 = p[l + 2]; ++ const uint j0 = (a1 >> (7 * l)) & 127, j1 = (a1 >> (7 * (l + 2))) & 127; ++ threadgroup const uint* th = tab + hh; ++ a = w4_signed(ctab ? uint(iq2xxs_grid[i0] >> (32 * hh)) : th[2 * i0], sign4_iq2xs(j0, hh), sc); ++ b = w4_signed(ctab ? uint(iq2xxs_grid[i1] >> (32 * hh)) : th[2 * i1], sign4_iq2xs(j1, hh), sc); ++ } ++ template ++ static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { ++ const P bl = row + (k >> 8) * 66; ++ block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b); + } + }; + template <> struct W8 { ++ static inline void block(device const uchar* bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { ++ const int ib = (k & 255) >> 5; ++ device const ushort* q = (device const ushort*)(bl + 2) + ((k & 255) >> 3); ++ const uint v0 = q[0], v1 = q[2]; ++ const uint scb = bl[66 + ib]; ++ const float sc0 = d * fma(float(scb & 15), 0.25f, 0.125f); // l < 2 (k % 32 < 16): runs l, l + 2 take nibbles 0, 1 ++ const float sc1 = d * fma(float(scb >> 4), 0.25f, 0.125f); ++ threadgroup const uint* th = tab + hh; ++ a = w4_signed(ctab ? uint(iq2xs_grid[v0 & 511] >> (32 * hh)) : th[2 * (v0 & 511)], sign4_iq2xs(v0 >> 9, hh), sc0); ++ b = w4_signed(ctab ? uint(iq2xs_grid[v1 & 511] >> (32 * hh)) : th[2 * (v1 & 511)], sign4_iq2xs(v1 >> 9, hh), sc1); ++ } + static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 74; +- const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; +- device const ushort* q = (device const ushort*)(bl + 2) + 4 * ib; +- const uint v0 = q[l], v1 = q[l + 2]; +- threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 1024); +- const uint scb = bl[66 + ib]; +- const float d = rdh(bl); +- const float sc0 = d * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f; +- const float sc1 = d * (0.5f + float((scb >> (4 * ((l + 2) >> 1))) & 15)) * 0.25f; +- const uint s0 = ks[v0 >> 9], s1 = ks[v1 >> 9]; +- a = w4_signed(tab[2 * (v0 & 511) + hh], hh ? s0 >> 4 : s0, sc0); +- b = w4_signed(tab[2 * (v1 & 511) + hh], hh ? s1 >> 4 : s1, sc1); ++ block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b); + } + }; + template <> struct W8 { ++ // grid words read through tab + hh and the sign byte as s >> 4 * hh (not a select): 8% off mm_gu3 at T 2048; ++ // the same for IQ2_XXS / IQ2_XS, whose sign4_iq2xs select stays (a shift form there measured 10% slower) ++ template ++ static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { ++ const int ib = (k & 255) >> 5, l = (k & 31) >> 3; ++ threadgroup const uint* th = tab + hh; ++ const uint sh = uint(hh) << 2; ++ const uint qh = bl[66 + ib], scb = bl[74 + ib]; ++ const P p = bl + 2 + ((k & 255) >> 3); ++ const uint i0 = uint(p[0]) | ((qh << (8 - 2 * l)) & 0x300); ++ const uint i1 = uint(p[2]) | ((qh << (4 - 2 * l)) & 0x300); ++ const uint s0 = p[32], s1 = p[34]; ++ const float sc0 = d * fma(float(scb & 15), 0.25f, 0.125f); ++ const float sc1 = d * fma(float(scb >> 4), 0.25f, 0.125f); ++ a = w4_signed(ctab ? uint(iq2s_grid[i0] >> (32 * hh)) : th[2 * i0], s0 >> sh, sc0); ++ b = w4_signed(ctab ? uint(iq2s_grid[i1] >> (32 * hh)) : th[2 * i1], s1 >> sh, sc1); ++ } + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + const P bl = row + (k >> 8) * 82; +- const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; +- const uint qh = bl[66 + ib], scb = bl[74 + ib]; +- const float d = rdh(bl); +- const uint i0 = uint(bl[2 + 4 * ib + l]) | ((qh << (8 - 2 * l)) & 0x300); +- const uint i1 = uint(bl[2 + 4 * ib + l + 2]) | ((qh << (8 - 2 * (l + 2))) & 0x300); +- const uint s0 = bl[2 + 32 + 4 * ib + l], s1 = bl[2 + 32 + 4 * ib + l + 2]; +- const float sc0 = d * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f; +- const float sc1 = d * (0.5f + float((scb >> (4 * ((l + 2) >> 1))) & 15)) * 0.25f; +- a = w4_signed(g_iq2s(tab, i0, hh), hh ? s0 >> 4 : s0, sc0); +- b = w4_signed(g_iq2s(tab, i1, hh), hh ? s1 >> 4 : s1, sc1); ++ block(bl, k, rdh(bl), tab, (k & 7) >> 2, a, b); + } + }; + template <> struct W8 { + template +- static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { +- const P bl = row + (k >> 8) * 110; +- const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; ++ static inline void block(P bl, int k, float d, threadgroup const uint* tab, int hh, thread half4& a, thread half4& b) { ++ const int ib = k >> 5; + const uint qh = bl[66 + ib], scb = bl[106 + (ib >> 1)]; +- const float sc = rdh(bl) * float(1 + 2 * ((scb >> (4 * (ib & 1))) & 15)); +- const int l1 = l + 2; +- const uint i0 = hh ? (uint(bl[2 + 8 * ib + 2 * l + 1]) | ((qh << (7 - 2 * l)) & 256)) +- : (uint(bl[2 + 8 * ib + 2 * l]) | ((qh << (8 - 2 * l)) & 256)); +- const uint i1 = hh ? (uint(bl[2 + 8 * ib + 2 * l1 + 1]) | ((qh << (7 - 2 * l1)) & 256)) +- : (uint(bl[2 + 8 * ib + 2 * l1]) | ((qh << (8 - 2 * l1)) & 256)); +- const uint s0 = bl[74 + 4 * ib + l], s1 = bl[74 + 4 * ib + l1]; ++ const float sc = d * fma(float((scb >> (4 * (ib & 1))) & 15), 2.0f, 1.0f); ++ const P qs = bl + 2 + (k >> 2); ++ const int sh0 = 8 - ((k & 31) >> 2); ++ const uint i0 = uint(qs[0]) | ((qh << sh0) & 256u); ++ const uint i1 = uint(qs[4]) | ((qh << (sh0 - 4)) & 256u); ++ const P sp = bl + 74 + (k >> 3); ++ const uint s0 = sp[0], s1 = sp[2]; ++ if (mm_probe == 5) { // bench probe: the loads and the scale, no grid lookup / sign math (wrong values) ++ a = half4(half(float(i0 ^ s0) * sc)); ++ b = half4(half(float(i1 ^ s1) * sc)); ++ return; ++ } ++ a = w4_signed(g_iq3s(tab, i0), s0 >> (4 * hh), sc); ++ b = w4_signed(g_iq3s(tab, i1), s1 >> (4 * hh), sc); ++ } ++ template ++ static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + if (mm_probe == 7) { +- // bench probe: a repacked layout's cost - one aligned 16-byte load a lane a row a 32-group holding every +- // field the lane needs (wrong values: the bytes come from an arbitrary aligned spot of the row) ++ const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint4 r = *(device const uint4*)((ulong(row) + ulong(k >> 5) * 16) & ~ulong(15)); + const uint qhv = r.z & 0xFF, scv = (r.z >> 8) & 0xFF; + const float scr = float(as_type(ushort(r.w & 0xFFFF))) * float(1 + 2 * ((scv >> (4 * (ib & 1))) & 15)); +@@ -698,13 +800,8 @@ template <> struct W8 { + b = w4_signed(tab[j1], (r.y >> (16 + 8 * hh)) & 0xFF, scr); + return; + } +- if (mm_probe == 5) { // bench probe: the loads and the scale, no grid lookup / sign math (wrong values) +- a = half4(half(float(i0 ^ s0) * sc)); +- b = half4(half(float(i1 ^ s1) * sc)); +- return; +- } +- a = w4_signed(g_iq3s(tab, i0), hh ? s0 >> 4 : s0, sc); +- b = w4_signed(g_iq3s(tab, i1), hh ? s1 >> 4 : s1, sc); ++ const P bl = row + (k >> 8) * 110; ++ block(bl, k & 255, rdh(bl), tab, (k & 7) >> 2, a, b); + } + }; + +@@ -761,22 +858,19 @@ template <> struct W8 { + template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 210; +- const int j = k & 255, n = j >> 7, jj = j & 127, qd = jj >> 5, l = jj & 31; // l < 16; the pair at l + 16 +- device const uchar* ql = bl + 64 * n + l + ((qd & 1) ? 32 : 0); +- device const uchar* qh = bl + 128 + 32 * n + l; ++ const uint j = uint(k & 255), j4 = j >> 4; // l = j & 31 < 16; the pair at l + 16 ++ device const uchar* ql = bl + (((j >> 1) & 64u) | (j & 63u)); ++ device const uchar* qh = bl + 128u + (((j >> 2) & 32u) | (j & 31u)); + const float d = rdh(bl + 208); +- device const char* sc = (device const char*)(bl + 192 + 8 * n); +- const float dsa = d * float(sc[2 * qd]), dsb = d * float(sc[1 + 2 * qd]); +- // the 4-byte runs as two ushorts each (210-byte blocks: 2-aligned) +- device const ushort* l2 = (device const ushort*)ql; +- device const ushort* h2 = (device const ushort*)qh; +- const uint la = uint(l2[0]) | (uint(l2[1]) << 16), lb = uint(l2[8]) | (uint(l2[9]) << 16); +- const uint ha = uint(h2[0]) | (uint(h2[1]) << 16), hb = uint(h2[8]) | (uint(h2[9]) << 16); +- const uint sl = (qd >> 1) ? 4u : 0u, shh = uint(2 * qd); +- const uint4 qa = (uint4(as_type((la >> sl) & 0x0F0F0F0Fu)) | (uint4(as_type((ha >> shh) & 0x03030303u)) << 4)); +- const uint4 qb = (uint4(as_type((lb >> sl) & 0x0F0F0F0Fu)) | (uint4(as_type((hb >> shh) & 0x03030303u)) << 4)); +- a = half4(dsa * (float4(qa) - 32.0f)); +- b = half4(dsb * (float4(qb) - 32.0f)); ++ device const char* sc = (device const char*)(bl + 192u + j4); ++ const float dsa = d * float(sc[0]), dsb = d * float(sc[1]); ++ const uint la = rd32(ql), lb = rd32(ql + 16); ++ const uint ha = rd32(qh), hb = rd32(qh + 16); ++ const uint sl = j4 & 4u, shh = j4 & 6u; ++ const uint qa = ((la >> sl) & 0x0F0F0F0Fu) | (((ha >> shh) & 0x03030303u) << 4); ++ const uint qb = ((lb >> sl) & 0x0F0F0F0Fu) | (((hb >> shh) & 0x03030303u) << 4); ++ a = half4(dsa * (float4(as_type(qa)) - 32.0f)); ++ b = half4(dsb * (float4(as_type(qb)) - 32.0f)); + } + }; + +@@ -788,8 +882,10 @@ kernel void mm_gu3(device const uchar* W [[buffer(0)]], device const half* XS [[ + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float ot[SG][TM * TN]; + threadgroup uint tab[fmt::TAB]; +- load_tab(tab, tid, SG * 32); +- threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (fmt::TAB > 1) { ++ load_tab(tab, tid, SG * 32); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN; + device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride; +@@ -798,54 +894,128 @@ kernel void mm_gu3(device const uchar* W [[buffer(0)]], device const half* XS [[ + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); +- auto rg = op.template get_right_input_cooperative_tensor(); +- auto ru = op.template get_right_input_cooperative_tensor(); ++ auto rw = op.template get_right_input_cooperative_tensor(); + auto tA = tX.template slice(0, row0); +- auto accg = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); +- auto accu = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); ++ auto accg = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); ++ auto accu = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } ++ if (coop_k4 && TK == 32 && mm_probe != 1) { ++ const uint16_t half_cap = rw.get_capacity() / 2; ++ const auto ix0 = rw.get_multidimensional_index(uint16_t(0)); ++ const int kl = ix0[0], hh = (kl >> 2) & 1; ++ const ulong roff0 = ulong(ix0[1]) * a.gu_row; ++ device const uchar* g0_l = g0 + roff0; ++ device const uchar* u0_l = u0 + roff0; ++ const uint step8 = 8u * a.gu_row; ++ // IQ2 formats: a 256 super-block at a time, its d read once; the lane's hh is fixed for the whole K loop ++ constexpr bool sb = F == T_IQ2_S || F == T_IQ2_XXS || F == T_IQ2_XS; ++ for (int kb = 0; kb < int(a.n_embd); kb += (sb ? 256 : TK)) { ++ const uint sboff = uint(kb >> 8) * fmt::BB; ++ float dg[TN * TK / 256], du[TN * TK / 256]; ++ if (sb) { ++#pragma clang loop unroll(full) ++ for (uint16_t i = 0; i < half_cap / 4; ++i) { ++ dg[i] = rdh(g0_l + i * step8 + sboff); ++ du[i] = rdh(u0_l + i * step8 + sboff); ++ } ++ } ++ for (int ki = 0; ki < (sb ? 256 : TK); ki += TK) { ++ const int k0 = kb + ki; ++ tA = tX.template slice(k0, row0); ++ const uint boff = F == T_IQ2_XS ? sboff : (fmt::QK == 256 ? uint(k0 >> 8) * fmt::BB : (fmt::QK == 64 ? uint(k0 >> 6) * fmt::BB : 0u)); ++ const int j_in = fmt::QK == 256 ? ((k0 + kl) & 255) : (fmt::QK == 64 ? ((k0 + kl) & 63) : (k0 + kl)); ++ device const uchar* rgw = fmt::QK >= 64 ? g0_l + boff : g0_l; ++#pragma clang loop unroll(full) ++ for (uint16_t i = 0; i < half_cap; i += 4, rgw += step8) { ++ half4 wg, wg2; ++ if (F == T_IQ2_S) W8::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2); ++ else if (F == T_IQ2_XXS) W8::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2); ++ else if (F == T_IQ2_XS) W8::block(rgw, j_in, dg[i / 4], tab, hh, wg, wg2); ++ else if (F == T_IQ3_S) W8::block(rgw, j_in, rdh(rgw), tab, hh, wg, wg2); ++ else W8::run(rgw, j_in, tab, wg, wg2); ++ rw[i] = wg.x; rw[i + 1] = wg.y; rw[i + 2] = wg.z; rw[i + 3] = wg.w; ++ const uint16_t j = i + half_cap; ++ rw[j] = wg2.x; rw[j + 1] = wg2.y; rw[j + 2] = wg2.z; rw[j + 3] = wg2.w; ++ } ++ op.run(tA, rw, accg); ++ device const uchar* ruw = fmt::QK >= 64 ? u0_l + boff : u0_l; ++#pragma clang loop unroll(full) ++ for (uint16_t i = 0; i < half_cap; i += 4, ruw += step8) { ++ half4 wu, wu2; ++ if (F == T_IQ2_S) W8::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2); ++ else if (F == T_IQ2_XXS) W8::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2); ++ else if (F == T_IQ2_XS) W8::block(ruw, j_in, du[i / 4], tab, hh, wu, wu2); ++ else if (F == T_IQ3_S) W8::block(ruw, j_in, rdh(ruw), tab, hh, wu, wu2); ++ else W8::run(ruw, j_in, tab, wu, wu2); ++ rw[i] = wu.x; rw[i + 1] = wu.y; rw[i + 2] = wu.z; rw[i + 3] = wu.w; ++ const uint16_t j = i + half_cap; ++ rw[j] = wu2.x; rw[j + 1] = wu2.y; rw[j + 2] = wu2.z; rw[j + 3] = wu2.w; ++ } ++ op.run(tA, rw, accu); ++ } ++ } ++ } else + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { +- const uint16_t cap = rg.get_capacity(), half_cap = cap / 2; ++ const uint16_t cap = rw.get_capacity(), half_cap = cap / 2; ++ tA = tX.template slice(k0, row0); + if (!coop_k4) { + for (uint16_t i = 0; i < cap; ++i) { +- const auto ix = rg.get_multidimensional_index(i); ++ const auto ix = rw.get_multidimensional_index(i); + const int k = k0 + ix[0], n = ix[1]; +- rg[i] = W4::run(g0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3]; +- ru[i] = W4::run(u0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3]; ++ rw[i] = W4::run(g0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3]; + } +- } else ++ op.run(tA, rw, accg); ++ for (uint16_t i = 0; i < cap; ++i) { ++ const auto ix = rw.get_multidimensional_index(i); ++ const int k = k0 + ix[0], n = ix[1]; ++ rw[i] = W4::run(u0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3]; ++ } ++ op.run(tA, rw, accu); ++ } else { + #pragma clang loop unroll(full) +- for (uint16_t i = 0; i < half_cap; i += 4) { +- const auto ix = rg.get_multidimensional_index(i); +- const auto iy = rg.get_multidimensional_index(i + half_cap); +- const int k = k0 + ix[0], n = ix[1]; +- device const uchar* rgw = g0 + ulong(n) * a.gu_row; +- device const uchar* ruw = u0 + ulong(n) * a.gu_row; +- half4 wg, wu, wg2, wu2; +- if (mm_probe == 1) { // the bench probe: no decode (constant weights) +- wg = half4(half(n & 7)); wu = wg; wg2 = half4(half(k & 7)); wu2 = wg2; +- } else if (TK == 32) { +- // the lane's two runs of this row sit in one 32-group: element i + cap / 2 is (k + 16, n) - part of +- // the COOP_K4 layout (coop_probe checks it); a runtime check of it here cost 30% +- W8::run(rgw, k, tab, wg, wg2); +- W8::run(ruw, k, tab, wu, wu2); +- } else { +- const int k2 = k0 + iy[0]; +- device const uchar* rgw2 = g0 + ulong(iy[1]) * a.gu_row; +- device const uchar* ruw2 = u0 + ulong(iy[1]) * a.gu_row; +- wg = W4::run(rgw, k, tab); wu = W4::run(ruw, k, tab); +- wg2 = W4::run(rgw2, k2, tab); wu2 = W4::run(ruw2, k2, tab); ++ for (uint16_t i = 0; i < half_cap; i += 4) { ++ const auto ix = rw.get_multidimensional_index(i); ++ const auto iy = rw.get_multidimensional_index(i + half_cap); ++ const int k = k0 + ix[0], n = ix[1]; ++ device const uchar* rgw = g0 + ulong(n) * a.gu_row; ++ half4 wg, wg2; ++ if (mm_probe == 1) { ++ wg = half4(half(n & 7)); wg2 = half4(half(k & 7)); ++ } else if (TK == 32) { ++ W8::run(rgw, k, tab, wg, wg2); ++ } else { ++ const int k2 = k0 + iy[0]; ++ wg = W4::run(rgw, k, tab); ++ wg2 = W4::run(g0 + ulong(iy[1]) * a.gu_row, k2, tab); ++ } ++ rw[i] = wg.x; rw[i + 1] = wg.y; rw[i + 2] = wg.z; rw[i + 3] = wg.w; ++ const uint16_t j = i + half_cap; ++ rw[j] = wg2.x; rw[j + 1] = wg2.y; rw[j + 2] = wg2.z; rw[j + 3] = wg2.w; + } +- rg[i] = wg.x; rg[i + 1] = wg.y; rg[i + 2] = wg.z; rg[i + 3] = wg.w; +- ru[i] = wu.x; ru[i + 1] = wu.y; ru[i + 2] = wu.z; ru[i + 3] = wu.w; +- const uint16_t j = i + half_cap; +- rg[j] = wg2.x; rg[j + 1] = wg2.y; rg[j + 2] = wg2.z; rg[j + 3] = wg2.w; +- ru[j] = wu2.x; ru[j + 1] = wu2.y; ru[j + 2] = wu2.z; ru[j + 3] = wu2.w; ++ op.run(tA, rw, accg); ++#pragma clang loop unroll(full) ++ for (uint16_t i = 0; i < half_cap; i += 4) { ++ const auto ix = rw.get_multidimensional_index(i); ++ const auto iy = rw.get_multidimensional_index(i + half_cap); ++ const int k = k0 + ix[0], n = ix[1]; ++ device const uchar* ruw = u0 + ulong(n) * a.gu_row; ++ half4 wu, wu2; ++ if (mm_probe == 1) { ++ wu = half4(half(n & 7)); wu2 = half4(half(k & 7)); ++ } else if (TK == 32) { ++ W8::run(ruw, k, tab, wu, wu2); ++ } else { ++ const int k2 = k0 + iy[0]; ++ wu = W4::run(ruw, k, tab); ++ wu2 = W4::run(u0 + ulong(iy[1]) * a.gu_row, k2, tab); ++ } ++ rw[i] = wu.x; rw[i + 1] = wu.y; rw[i + 2] = wu.z; rw[i + 3] = wu.w; ++ const uint16_t j = i + half_cap; ++ rw[j] = wu2.x; rw[j + 1] = wu2.y; rw[j + 2] = wu2.z; rw[j + 3] = wu2.w; ++ } ++ op.run(tA, rw, accu); + } +- tA = tX.template slice(k0, row0); +- op.run(tA, rg, accg); +- op.run(tA, ru, accu); + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { +@@ -855,6 +1025,17 @@ kernel void mm_gu3(device const uchar* W [[buffer(0)]], device const half* XS [[ + tensor, tensor_inline> tO(ot[sg], dextents(TN, TM)); + accg.store(tO); + simdgroup_barrier(mem_flags::mem_threadgroup); ++ if ((F == T_IQ2_XS || F == T_IQ3_XXS || F == T_IQ2_XXS || F == T_Q2_0 || F == T_IQ2_S || F == T_IQ3_S) && 128 % TN == 0 && (TM * TN) % 128 == 0) { ++ // a lane stores 4 adjacent outputs a pass, 128 a pass over the simdgroup ++ constexpr int RP = 128 / TN; // rows a pass ++ const int mr = lane / (TN / 4), f = (lane % (TN / 4)) * 4; ++ threadgroup const float4* ot_p = (threadgroup const float4*)(&ot[sg][mr * TN + f]); ++ device half* hs_p = HS + ulong(row0 + mr) * a.n_ff + n0 + f; ++ const uint hs_step = uint(RP) * a.n_ff; ++#pragma clang loop unroll(full) ++ for (int m = 0; m < TM / RP; ++m, hs_p += hs_step) ++ *(device half4*)hs_p = half4(ot_p[m * (RP * TN / 4)]); ++ } else + for (int it = lane; it < TM * TN; it += 32) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + m) * a.n_ff + n0 + f] = half(ot[sg][it]); +@@ -868,8 +1049,10 @@ kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint tab[fmt::TAB]; +- load_tab(tab, tid, SG * 32); +- threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (fmt::TAB > 1) { ++ load_tab(tab, tid, SG * 32); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN; + device const uchar* w0 = (kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.d_stride) + a.down_off + ulong(n0) * a.d_row; +@@ -881,6 +1064,38 @@ kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; ++ const auto ix0 = rw.get_multidimensional_index(uint16_t(0)); ++ const int kl = ix0[0]; ++ device const uchar* w0_l = w0 + ulong(ix0[1]) * a.d_row; ++ const uint step8 = 8u * a.d_row; ++ if (coop_k4 && F == T_Q2_0 && TK == 32) { ++ const uint16_t half_cap = rw.get_capacity() / 2; ++ const uint kl4 = uint(kl) >> 2; ++ for (int k0 = 0; k0 < int(a.n_ff); k0 += 64, w0_l += 18) { ++ device const uchar* rdw = w0_l; ++#pragma clang loop unroll(full) ++ for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) { ++ const half4 w = W4::run_q<2>(rdw, kl4); ++ const half4 w2 = W4::run_q<6>(rdw, kl4); ++ rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; ++ const uint16_t j = i + half_cap; ++ rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; ++ } ++ tA = tH.template slice(k0, row0); ++ op.run(tA, rw, acc); ++ rdw = w0_l; ++#pragma clang loop unroll(full) ++ for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) { ++ const half4 w = W4::run_q<10>(rdw, kl4); ++ const half4 w2 = W4::run_q<14>(rdw, kl4); ++ rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; ++ const uint16_t j = i + half_cap; ++ rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; ++ } ++ tA = tH.template slice(k0 + 32, row0); ++ op.run(tA, rw, acc); ++ } ++ } else + for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) { + if (!coop_k4) { + for (uint16_t i = 0; i < rw.get_capacity(); ++i) { +@@ -891,11 +1106,13 @@ kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS + } else if (TK == 32) { + // the paired runs (element i + cap / 2 is (k + 16, n), the COOP_K4 layout): one W8 a pair + const uint16_t half_cap = rw.get_capacity() / 2; ++ const uint boff = fmt::QK == 64 ? uint(k0 >> 6) * fmt::BB : (fmt::QK == 32 ? uint(k0 >> 5) * fmt::BB : 0u); ++ const int j_in = fmt::QK == 64 ? ((k0 + kl) & 63) : (fmt::QK == 32 ? kl : (k0 + kl)); ++ device const uchar* rdw = w0_l + boff; + #pragma clang loop unroll(full) +- for (uint16_t i = 0; i < half_cap; i += 4) { +- const auto ix = rw.get_multidimensional_index(i); ++ for (uint16_t i = 0; i < half_cap; i += 4, rdw += step8) { + half4 w, w2; +- W8::run(w0 + ulong(ix[1]) * a.d_row, k0 + ix[0], tab, w, w2); ++ W8::run(rdw, j_in, tab, w, w2); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + const uint16_t j = i + half_cap; + rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; +@@ -926,6 +1143,7 @@ kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS + MM3_GU(T_Q4_K) MM3_GU(T_Q5_K) MM3_GU(T_Q6_K) MM3_GU(T_Q8_0) MM3_GU(T_IQ4_XS) MM3_GU(T_Q2_0) MM3_GU(T_IQ3_XXS) MM3_GU(T_IQ2_XXS) MM3_GU(T_IQ2_XS) MM3_GU(T_IQ2_S) MM3_GU(T_IQ3_S) + MM3_DN(T_Q2_0) MM3_DN(T_IQ4_NL) MM3_DN(T_Q5_1) MM3_DN(T_Q8_0) + ++#ifndef MM_V3_ONLY + // mm_gu4: mm_gu3 with each decoded right operand (TN x TK weights of gate and of up) applied to MS 16-row token + // tiles (a tile of 16 * MS rows from the plan, tm = 16 * MS): the decode, which bounds mm_gu3, is shared by MS times + // as many token rows (more padding rows per expert, but the matrix ops are cheap next to the decode). The same +@@ -1016,6 +1234,8 @@ MM4_GU(T_Q2_0) MM4_GU(T_IQ3_XXS) MM4_GU(T_IQ2_XXS) MM4_GU(T_IQ2_XS) MM4_GU(T_IQ2 + template + kernel void mm_gu5(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], ++ device const uint* smap [[buffer(5)]], device const uchar* W1 [[buffer(6)]], ++ device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int TM = 16, TK = 32, BB = fmt::BB, ROWS = SG * TN, RU = BB / 2; +@@ -1025,8 +1245,9 @@ kernel void mm_gu5(device const uchar* W [[buffer(0)]], device const half* XS [[ + load_tab(tab, tid, SG * 32); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, nt0 = int(tg.x) * ROWS, n0 = nt0 + sg * TN; +- device const uchar* gt = W + a.gate_off + ulong(e) * a.gu_stride + ulong(nt0) * a.gu_row; +- device const uchar* ut = W + a.up_off + ulong(e) * a.gu_stride + ulong(nt0) * a.gu_row; ++ device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride; ++ device const uchar* gt = E + a.gate_off + ulong(nt0) * a.gu_row; ++ device const uchar* ut = E + a.up_off + ulong(nt0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); +@@ -1086,4 +1307,6 @@ kernel void mm_gu5(device const uchar* W [[buffer(0)]], device const half* XS [[ + #define MMGU5_(F, N, S) template [[host_name("mm_gu5_" #F "_" #N "_" #S)]] kernel decltype(mm_gu5) mm_gu5; + #define MMGU5(F, N, S) MMGU5_(F, N, S) + MMGU5(T_IQ3_S, 32, 2) MMGU5(T_IQ3_S, 32, 1) MMGU5(T_IQ3_S, 16, 4) MMGU5(T_IQ3_XXS, 32, 2) MMGU5(T_IQ3_XXS, 32, 1) +-MMGU5(T_IQ3_XXS, 16, 4) MMGU5(T_IQ2_S, 32, 2) MMGU5(T_IQ2_S, 32, 1) MMGU5(T_IQ2_S, 16, 4) ++MMGU5(T_IQ3_XXS, 16, 4) MMGU5(T_IQ2_S, 32, 2) MMGU5(T_IQ2_S, 32, 1) MMGU5(T_IQ2_S, 16, 4) MMGU5(T_IQ2_XXS, 32, 2) MMGU5(T_IQ2_XXS, 16, 4) ++#endif ++#endif +diff --git a/metal/kernels/prompt.metal b/metal/kernels/prompt.metal +index 64833b0a..a3fd96be 100644 +--- a/metal/kernels/prompt.metal ++++ b/metal/kernels/prompt.metal +@@ -7,33 +7,98 @@ + // gr_epi_up after the HC up GEMM (G = lo . w_up, T x hc*n): mixed = mean_c xn[c] * sigmoid(G[c]) -> float and + // half (the next GEMMs' operand) + // swiglu_h h = half(silu(g) * u) +-// moe_out bo = y_routed + sigmoid(sg) * y_shared ++// gemm_hc_down, gemm_swiglu_up (after gemm_hc_up) the HC down and shared-expert up GEMMs with gr_epi_down / ++// swiglu_h as register epilogues ++// moe_out bo = y_routed + sigmoid(sg) * y_shared (prompt_chain.inc; the prompt chunk folds it into the next ++// read: gr_norm_h_moe, gr_apply_moe, ple_moe_gate_norm; with mm_combine folded in too: gr_norm_h_moec, ++// gr_apply_moec) + + kernel void route_topk(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]], +- device float* w [[buffer(2)]], constant uint& n_tok [[buffer(3)]], ++ device float* w [[buffer(2)]], constant uint& n_tok_packed [[buffer(3)]], ++ device const uint* smap [[buffer(4)]], constant float& prompt_moce [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { ++ const uint n_tok = n_tok_packed & 0xFFFFu; ++ const int n_exp = (n_tok_packed >> 16) ? int(n_tok_packed >> 16) : 512; + const uint t = tg.x * 8 + sg; + if (t >= n_tok) return; +- device const float* l = logits + ulong(t) * 512; ++ device const float* l = logits + ulong(t) * ulong(n_exp); + float v[16]; +- for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i]; + float mx = -INFINITY; +- for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); +- for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + float sum = 0.0f; +- for (int i = 0; i < 16; ++i) { +- v[i] = exp(v[i] - mx); +- sum += v[i]; +- } +- for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); +- const float rcp = 1.0f / sum; +- for (int i = 0; i < 16; ++i) { +- v[i] *= rcp; +- if (isnan(v[i])) v[i] = -FLT_MAX; ++ if (n_exp == 512) { ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) { ++ v[i] = l[lane + 32 * i]; ++ mx = max(mx, v[i]); ++ } ++ for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) { ++ v[i] = exp(v[i] - mx); ++ sum += v[i]; ++ } ++ for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); ++ const float rcp = 1.0f / sum; ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) { ++ v[i] *= rcp; ++ if (isnan(v[i])) v[i] = -FLT_MAX; ++ } ++ } else { ++ for (int i = 0; i < 16; ++i) { ++ const int idx = lane + 32 * i; ++ v[i] = idx < n_exp ? l[idx] : -INFINITY; ++ } ++ for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); ++ for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); ++ for (int i = 0; i < 16; ++i) { ++ v[i] = (lane + 32 * i < n_exp) ? exp(v[i] - mx) : 0.0f; ++ sum += v[i]; ++ } ++ for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); ++ const float rcp = 1.0f / sum; ++ for (int i = 0; i < 16; ++i) { ++ if (lane + 32 * i >= n_exp) { ++ v[i] = -INFINITY; ++ } else { ++ v[i] *= rcp; ++ if (isnan(v[i])) v[i] = -FLT_MAX; ++ } ++ } + } ++ const int tail_exact = int(prompt_moce) / 10000; ++ const float rem_moce = prompt_moce - float(tail_exact * 10000); ++ const int body_prot = int(rem_moce) / 100; ++ const float body_bonus = rem_moce - float(body_prot * 100); ++ const bool in_tail = tail_exact > 0 && (t + uint(tail_exact) >= n_tok); ++ const int p_prot = in_tail ? 10 : body_prot; ++ const float p_bonus = in_tail ? 0.0f : body_bonus; ++ const float moce_scale = p_bonus > 0.0f ? exp(min(p_bonus, 40.0f)) : 1.0f; ++ const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f; ++ ushort scaled_mask = 0; + float selected = 0.0f, ssum = 0.0f; + for (int rank = 0; rank < 10; ++rank) { ++ if (rank == p_prot && moce_scale > 1.0f) { ++ if (n_exp == 512) { ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) { ++ const int idx = lane + 32 * i; ++ if (v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) { ++ v[i] *= moce_scale; ++ scaled_mask |= ushort(1u << i); ++ } ++ } ++ } else { ++ for (int i = 0; i < 16; ++i) { ++ const int idx = lane + 32 * i; ++ if (idx < n_exp && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) { ++ v[i] *= moce_scale; ++ scaled_mask |= ushort(1u << i); ++ } ++ } ++ } ++ } + float best = v[0]; + int ex = lane; + for (int i = 1; i < 16; ++i) +@@ -43,14 +108,41 @@ kernel void route_topk(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; } + } ++ float true_prob = best; ++ if (rank >= p_prot && moce_scale > 1.0f && best > 0.0f) { ++ const ushort win_lane = ushort(ex & 31); ++ 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 * rcp_moce; ++ } + if ((ex & 31) == lane) { +- v[ex / 32] = -INFINITY; ++ // v stays in registers: no runtime index into it ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) v[i] = i == (ex >> 5) ? -INFINITY : v[i]; + ids[t * 10 + rank] = ex; +- ssum += best; ++ ssum += true_prob; + } +- if (rank == lane) selected = best; ++ if (rank == lane) selected = true_prob; + } + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); ++ if (prompt_moce < 0.0f) { ++ const int skip_tail = int(-prompt_moce); ++ const float tau = (-prompt_moce) - float(skip_tail); ++ const bool skip_active = tau > 0.0f && !(skip_tail > 0 && (t + uint(skip_tail) >= n_tok)); ++ if (skip_active) { ++ float kept_prob = (lane < 10) ? selected : 0.0f; ++ if (lane >= 3 && lane < 10) { ++ const int ex_r = ids[t * 10 + lane]; ++ if (smap[ex_r] == 0xFFFFFFFFu && selected < tau * ssum) { ++ ids[t * 10 + lane] = -1; ++ kept_prob = 0.0f; ++ } ++ } ++ selected = kept_prob; ++ ssum = kept_prob; ++ for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); ++ } ++ } + ssum = max(ssum, 6.103515625e-5f); + if (lane < 10) w[t * 10 + lane] = selected * (1.0f / ssum); + } +@@ -121,57 +213,478 @@ kernel void gr_norm_h(device float* R [[buffer(0)]], device const float* bo [[bu + const uint c = tg.x, t = tg.y, n = a.n_embd; + device float* r = R + (t * a.hc + c) * n; + const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; +- float ss = 0.0f; +- for (uint d = tid; d < n; d += 256) { ++ float rv[10], ss = 0.0f; ++ for (uint d = tid, i = 0; d < n; d += 256, ++i) { + float v = r[d]; + if (a.apply) { + v = fma(bo[t * n + d], g, v); + r[d] = v; + } ++ rv[i] = v; + ss = fma(v, v, ss); + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; ++#pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device half* o = xnh + (t * a.hc + c) * n; + device const float* w = w_norm + c * n; +- for (uint d = tid; d < n; d += 256) o[d] = half(r[d] * w[d] * rs); ++ for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[i] * w[d] * rs); ++} ++ ++// gr_norm_h's tail for one token's hc streams: the sums of squares ss[c] (per thread) to the norm, xnh = R' * w * rs ++inline void gr_norm_h_out(thread const float (*rv)[10], device const float* w_norm, device half* xnh, uint t, uint hc, ++ uint n, float eps, thread const float* ss, threadgroup float (*part)[8], ushort tid, ++ ushort sg, ushort lane) { ++ for (uint c = 0; c < hc; ++c) { ++ const float s = simd_sum(ss[c]); ++ if (lane == 0) part[c][sg] = s; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ for (uint c = 0; c < hc; ++c) { ++ float tot = 0.0f; ++#pragma clang loop unroll(full) ++ for (int i = 0; i < 8; ++i) tot += part[c][i]; ++ const float rs = rsqrt(tot / float(n) + eps); ++ device half* o = xnh + (t * hc + c) * n; ++ device const float* w = w_norm + c * n; ++ for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[c][i] * w[d] * rs); ++ } ++} ++ ++// stream-wise reduction keeps one residual stream plus bo live, rather than all four residual streams. ++inline void gr_norm_h_stream_out(thread const float* rv, device const float* w, device half* o, uint n, float eps, ++ float ss, threadgroup float* part, ushort tid, ushort sg, ushort lane) { ++ ss = simd_sum(ss); ++ if (lane == 0) part[sg] = ss; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ float tot = 0.0f; ++#pragma clang loop unroll(full) ++ for (int i = 0; i < 8; ++i) tot += part[i]; ++ const float rs = rsqrt(tot / float(n) + eps); ++ if (n == 2560) { ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) { ++ const uint d = tid + i * 256; ++ o[d] = half(rv[i] * w[d] * rs); ++ } ++ } else { ++ for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[i] * w[d] * rs); ++ } ++} ++ ++// gr_norm_h4: gr_norm_h with one threadgroup per token over all hc streams, so bo is read once a token rather than ++// hc times; R', xnh bitwise gr_norm_h's. grid (1, T), 256 threads; hc <= 4 ++kernel void gr_norm_h4(device float* R [[buffer(0)]], device const float* bo [[buffer(1)]], ++ device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], ++ device half* xnh [[buffer(4)]], constant GrArgs& a [[buffer(5)]], ++ uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], ++ ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { ++ threadgroup float part[4][8]; ++ const uint t = tg.y, n = a.n_embd, hc = a.hc; ++ device float* r = R + t * hc * n; ++ if (!a.apply) { ++ for (uint c = 0; c < hc; ++c) { ++ float rv[10], ss = 0.0f; ++ if (n == 2560) { ++ device const float* rc = r + c * 2560 + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) { ++ const float v = rc[i * 256]; ++ rv[i] = v; ++ ss = fma(v, v, ss); ++ } ++ } else { ++ for (uint d = tid, i = 0; d < n; d += 256, ++i) { ++ const float v = r[c * n + d]; ++ rv[i] = v; ++ ss = fma(v, v, ss); ++ } ++ } ++ gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); ++ } ++ return; ++ } ++ float bv[10]; ++ if (n == 2560) { ++ device const float* brow = bo + t * 2560 + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) bv[i] = brow[i * 256]; ++ } else { ++ for (uint d = tid, i = 0; d < n; d += 256, ++i) bv[i] = bo[t * n + d]; ++ } ++ for (uint c = 0; c < hc; ++c) { ++ const float g = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); ++ float rv[10], ss = 0.0f; ++ if (n == 2560) { ++ device float* rc = r + c * 2560 + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) { ++ const float v = fma(bv[i], g, rc[i * 256]); ++ rc[i * 256] = v; ++ rv[i] = v; ++ ss = fma(v, v, ss); ++ } ++ } else { ++ for (uint d = tid, i = 0; d < n; d += 256, ++i) { ++ const float v = fma(bv[i], g, r[c * n + d]); ++ r[c * n + d] = v; ++ rv[i] = v; ++ ss = fma(v, v, ss); ++ } ++ } ++ gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); ++ } ++} ++ ++// gr_norm_h_moe: gr_norm_h with the sublayer output bo formed on the fly as moe_out - bo = fma(sigmoid(sg[t]), ys, yr) ++// - so the MoE's write never round-trips through bo; R', xnh bitwise moe_out + gr_norm_h. Always applies. ++// One threadgroup per token t (all hc streams: bo is formed once a token), 256 threads; hc <= 4. ++kernel void gr_norm_h_moe(device float* R [[buffer(0)]], device const float* yr [[buffer(1)]], ++ device const float* ys [[buffer(2)]], device const float* sgv [[buffer(3)]], ++ device const float* inj [[buffer(4)]], device const float* w_norm [[buffer(5)]], ++ device half* xnh [[buffer(6)]], constant GrArgs& a [[buffer(7)]], ++ uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], ++ ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { ++ threadgroup float part[4][8]; ++ const uint t = tg.y, n = a.n_embd, hc = a.hc; ++ device float* r = R + t * hc * n; ++ float rv[4][10], g[4], ss[4] = {0.0f, 0.0f, 0.0f, 0.0f}; ++ for (uint c = 0; c < hc; ++c) g[c] = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); ++ const float gs = 1.0f / (1.0f + exp(-sgv[t])); ++ for (uint d = tid, i = 0; d < n; d += 256, ++i) { ++ const float bo = fma(gs, ys[t * n + d], yr[t * n + d]); ++ for (uint c = 0; c < hc; ++c) { ++ const float v = fma(bo, g[c], r[c * n + d]); ++ r[c * n + d] = v; ++ rv[c][i] = v; ++ ss[c] = fma(v, v, ss[c]); ++ } ++ } ++ gr_norm_h_out(rv, w_norm, xnh, t, hc, n, a.eps, ss, part, tid, sg, lane); ++} ++ ++// token t's routed slots for moe_comb: compacted valid YE row offsets (uint(r) * n) and their weights in threadgroup memory ++inline uint moe_comb_slots(device const int* pos, device const float* srt_w, uint t, uint topk, uint n, ++ threadgroup uint* roff, threadgroup float* wk, ushort tid) { ++ int r = -1; ++ float w = 0.0f; ++ if (tid < topk) { ++ r = pos[t * topk + tid]; ++ if (r >= 0) w = srt_w[r]; ++ } ++ const bool valid = r >= 0; ++ const ushort prefix = simd_prefix_exclusive_sum(ushort(valid)); ++ const ushort cnt = simd_sum(ushort(valid)); ++ if (valid) { ++ roff[prefix] = uint(r) * n; ++ wk[prefix] = w; ++ } ++ if (tid == 0) roff[16] = cnt; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ return roff[16]; ++} ++ ++// the routed experts' write at (t, d) as mm_combine (moe_mm.metal) forms it: sum_k w(t, k) * YE[pos(t, k)][d], k in order ++inline float moe_comb(device const float* YE, threadgroup const uint* roff, threadgroup const float* wk, uint cnt, ++ uint d) { ++ float s = 0.0f; ++ for (uint k = 0; k < cnt; ++k) s = fma(wk[k], YE[roff[k] + d], s); ++ return s; ++} ++ ++// gr_norm_h_moec: gr_norm_h_moe with mm_combine folded in - yr formed from the expert rows YE (pos, srt_w: mm_place's) ++// as it is read, so it is never written; R', xnh bitwise mm_combine + gr_norm_h_moe. grid (1, T), 256 threads; ++// hc <= 4, topk <= 16 ++kernel void gr_norm_h_moec(device float* R [[buffer(0)]], device const float* YE [[buffer(1)]], ++ device const int* pos [[buffer(2)]], device const float* srt_w [[buffer(3)]], ++ device const float* ys [[buffer(4)]], device const float* sgv [[buffer(5)]], ++ device const float* inj [[buffer(6)]], device const float* w_norm [[buffer(7)]], ++ device half* xnh [[buffer(8)]], constant GrArgs& a [[buffer(9)]], constant uint& topk [[buffer(10)]], ++ uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], ++ ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { ++ threadgroup float part[4][8]; ++ threadgroup uint roff[17]; ++ threadgroup float wk[16]; ++ const uint t = tg.y, n = a.n_embd, hc = a.hc; ++ const uint cnt = moe_comb_slots(pos, srt_w, t, topk, n, roff, wk, tid); ++ device float* r = R + t * hc * n; ++ float bv[10]; ++ const float gs = 1.0f / (1.0f + exp(-sgv[t])); ++ if (n == 2560) { ++ // k outer: each slot's weight and row read once, its ten loads in flight together; s[i]'s fma order is moe_comb's ++ float s[10] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; ++ for (uint k = 0; k < cnt; ++k) { ++ const float w = wk[k]; ++ device const float* yk = YE + roff[k] + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) s[i] = fma(w, yk[i * 256], s[i]); ++ } ++ device const float* yrow = ys + t * 2560 + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) bv[i] = fma(gs, yrow[i * 256], s[i]); ++ } else { ++ for (uint d = tid, i = 0; d < n; d += 256, ++i) bv[i] = fma(gs, ys[t * n + d], moe_comb(YE, roff, wk, cnt, d)); ++ } ++ for (uint c = 0; c < hc; ++c) { ++ const float g = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))); ++ float rv[10], ss = 0.0f; ++ if (n == 2560) { ++ device float* rc = r + c * 2560 + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) { ++ const float v = fma(bv[i], g, rc[i * 256]); ++ rc[i * 256] = v; ++ rv[i] = v; ++ ss = fma(v, v, ss); ++ } ++ } else { ++ for (uint d = tid, i = 0; d < n; d += 256, ++i) { ++ const float v = fma(bv[i], g, r[c * n + d]); ++ r[c * n + d] = v; ++ rv[i] = v; ++ ss = fma(v, v, ss); ++ } ++ } ++ gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane); ++ } ++} ++ ++// gr_apply_moe: the last layer's MoE write applied to R alone (R' of moe_out + gr_norm; its xn has no reader). ++// grid (hc * n / 256, T), 256 threads ++kernel void gr_apply_moe(device float* R [[buffer(0)]], device const float* yr [[buffer(1)]], ++ device const float* ys [[buffer(2)]], device const float* sgv [[buffer(3)]], ++ device const float* inj [[buffer(4)]], constant GrArgs& a [[buffer(5)]], ++ uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { ++ const uint ch = tg.x * 256 + tid, t = tg.y, n = a.n_embd; ++ if (ch >= a.hc * n) return; ++ const uint c = ch / n, d = ch - c * n; ++ const float g = 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))); ++ const float gs = 1.0f / (1.0f + exp(-sgv[t])); ++ device float* r = R + t * a.hc * n + ch; ++ *r = fma(fma(gs, ys[t * n + d], yr[t * n + d]), g, *r); ++} ++ ++// gr_apply_moec: gr_apply_moe with mm_combine folded in (as gr_norm_h_moec); R' bitwise mm_combine + gr_apply_moe. ++// grid (ceil(n / 256), T), 256 threads: a thread's d over all hc streams (yr formed once); topk <= 16 ++kernel void gr_apply_moec(device float* R [[buffer(0)]], device const float* YE [[buffer(1)]], ++ device const int* pos [[buffer(2)]], device const float* srt_w [[buffer(3)]], ++ device const float* ys [[buffer(4)]], device const float* sgv [[buffer(5)]], ++ device const float* inj [[buffer(6)]], constant GrArgs& a [[buffer(7)]], ++ constant uint& topk [[buffer(8)]], uint3 tg [[threadgroup_position_in_grid]], ++ ushort tid [[thread_index_in_threadgroup]]) { ++ threadgroup uint roff[17]; ++ threadgroup float wk[16]; ++ const uint d = tg.x * 256 + tid, t = tg.y, n = a.n_embd, hc = a.hc; ++ const uint cnt = moe_comb_slots(pos, srt_w, t, topk, n, roff, wk, tid); ++ if (d >= n) return; ++ const float gs = 1.0f / (1.0f + exp(-sgv[t])); ++ const float bo = fma(gs, ys[t * n + d], moe_comb(YE, roff, wk, cnt, d)); ++ device float* r = R + t * hc * n + d; ++ for (uint c = 0; c < hc; ++c) r[c * n] = fma(bo, 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))), r[c * n]); + } + + // gemm_hc_up: G = lo . w_up^T with w_up's rows interleaved by dq_f16 (perm = hc: row d * hc + c), so a TN-column + // tile holds TN / hc outputs d with all their streams; the epilogue mixed[t][d] = mean_c xn[t][c][d] * +-// sigmoid(G[t][d * hc + c]) -> fp32 and fp16. G is never written. grid (hc * n / TN, ceil(T / TM)), 32 * SG threads. +-template +-kernel void gemm_hc_up(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device const half* xnh [[buffer(2)]], +- device float* mixed [[buffer(3)]], device half* mixh [[buffer(4)]], constant GemmArgs& g [[buffer(5)]], +- constant EpiArgs& e [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], +- ushort tid [[thread_index_in_threadgroup]]) { +- threadgroup float ot[TM * TN]; ++// sigmoid(G[t][d * hc + c]) -> fp16 and, if mixed is bound, fp32. G is never written. ++// grid (hc * n / TN, ceil(T / TM)), 32 * SG threads. ++template ++inline void hc_up_tile(device half* A, device half* B, device const half* xnh, device float* mixed, ++ device half* mixh, constant GemmArgs& g, constant EpiArgs& e, ++ int t0, int c0, ushort tid, threadgroup float* ot) { + tensor, tensor_inline> tA(A, dextents(g.K, g.M)); + tensor, tensor_inline> tB(B, dextents(g.K, g.N)); +- tensor, tensor_inline> tO(ot, dextents(TN, TM)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; +- auto mA = tA.slice(0, int(tg.y) * TM); +- auto mB = tB.slice(0, int(tg.x) * TN); +- op.run(mA, mB, tO); +- threadgroup_barrier(mem_flags::mem_threadgroup); +- const int hc = int(e.hc), n = int(e.n_embd), DPT = TN / hc; +- const int t0 = int(tg.y) * TM, d0 = int(tg.x) * DPT; +- for (int idx = tid; idx < TM * DPT; idx += SG * 32) { +- const int m = idx / DPT, dd = idx - m * DPT, t = t0 + m; +- if (t >= int(e.n_tok)) continue; +- float s = 0.0f; +- for (int c = 0; c < hc; ++c) +- s = fma(float(xnh[ulong(t) * hc * n + c * n + d0 + dd]), 1.0f / (1.0f + exp(-ot[m * TN + dd * hc + c])), s); +- s = s / float(hc); +- mixed[ulong(t) * n + d0 + dd] = s; +- mixh[ulong(t) * n + d0 + dd] = half(s); ++ auto mA = tA.slice(0, t0); ++ auto mB = tB.slice(0, c0); ++ if constexpr (REG) { ++ auto acc = op.template get_destination_cooperative_tensor(); ++ op.run(mA, mB, acc); ++ const int n = int(e.n_embd), d0 = c0 / 4; ++ // On M5, each lane owns four consecutive output columns: one output's HC streams. ++ #pragma clang loop unroll(full) ++ for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { ++ const auto ix = acc.get_multidimensional_index(i); ++ const int t = t0 + ix[1], d = d0 + ix[0] / 4; ++ if (t >= int(e.n_tok)) continue; ++ device const half* xnh_t = xnh + ulong(t) * 4 * n + d; ++ float s = 0.0f; ++ #pragma clang loop unroll(full) ++ for (int c = 0; c < 4; ++c) ++ s = fma(float(xnh_t[c * n]), 1.0f / (1.0f + exp(-acc[i + c])), s); ++ s *= 0.25f; ++ const ulong out_idx = ulong(t) * n + d; ++ if (mixed) mixed[out_idx] = s; ++ mixh[out_idx] = half(s); ++ } ++ } else { ++ tensor, tensor_inline> tO(ot, dextents(TN, TM)); ++ op.run(mA, mB, tO); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ const int hc = int(e.hc), n = int(e.n_embd), DPT = TN / hc; ++ const int d0 = c0 / hc; ++ for (int idx = tid; idx < TM * DPT; idx += SG * 32) { ++ const int m = idx / DPT, dd = idx - m * DPT, t = t0 + m; ++ if (t >= int(e.n_tok)) continue; ++ float s = 0.0f; ++ for (int c = 0; c < hc; ++c) ++ s = fma(float(xnh[ulong(t) * hc * n + c * n + d0 + dd]), 1.0f / (1.0f + exp(-ot[m * TN + dd * hc + c])), s); ++ s = s / float(hc); ++ if (mixed) mixed[ulong(t) * n + d0 + dd] = s; ++ mixh[ulong(t) * n + d0 + dd] = half(s); ++ } + } + } ++template ++kernel void gemm_hc_up(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device const half* xnh [[buffer(2)]], ++ device float* mixed [[buffer(3)]], device half* mixh [[buffer(4)]], constant GemmArgs& g [[buffer(5)]], ++ constant EpiArgs& e [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], ++ ushort tid [[thread_index_in_threadgroup]]) { ++ threadgroup float ot[REG ? 1 : TM * TN]; ++ const int t0 = int(tg.y) * TM, c0 = int(tg.x) * TN; ++ if constexpr (REG && TM > 16) { ++ if (t0 + 16 >= int(e.n_tok)) { ++ hc_up_tile<16, TN, SG, REG>(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot); ++ return; ++ } ++ } ++ if constexpr (REG && TM > 32) { ++ if (t0 + 32 >= int(e.n_tok)) { ++ hc_up_tile<32, TN, SG, REG>(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot); ++ return; ++ } ++ } ++ hc_up_tile(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot); ++} + template [[host_name("gemm_hc_up_64_64_4")]] kernel decltype(gemm_hc_up<64, 64, 4>) gemm_hc_up<64, 64, 4>; + template [[host_name("gemm_hc_up_64_32_4")]] kernel decltype(gemm_hc_up<64, 32, 4>) gemm_hc_up<64, 32, 4>; + template [[host_name("gemm_hc_up_32_64_4")]] kernel decltype(gemm_hc_up<32, 64, 4>) gemm_hc_up<32, 64, 4>; ++template [[host_name("gemm_hc_up_128_64_4")]] kernel decltype(gemm_hc_up<128, 64, 4>) gemm_hc_up<128, 64, 4>; ++template [[host_name("gemm_hc_up_64_128_4")]] kernel decltype(gemm_hc_up<64, 128, 4>) gemm_hc_up<64, 128, 4>; ++template [[host_name("gemm_hc_up_reg_64_64_4")]] kernel decltype(gemm_hc_up<64, 64, 4, true>) gemm_hc_up<64, 64, 4, true>; ++template [[host_name("gemm_hc_up_reg_128_64_4")]] kernel decltype(gemm_hc_up<128, 64, 4, true>) gemm_hc_up<128, 64, 4, true>; ++template [[host_name("gemm_hc_up_reg_64_128_4")]] kernel decltype(gemm_hc_up<64, 128, 4, true>) gemm_hc_up<64, 128, 4, true>; ++template [[host_name("gemm_hc_up_reg_64_32_4")]] kernel decltype(gemm_hc_up<64, 32, 4, true>) gemm_hc_up<64, 32, 4, true>; ++template [[host_name("gemm_hc_up_reg_32_128_4")]] kernel decltype(gemm_hc_up<32, 128, 4, true>) gemm_hc_up<32, 128, 4, true>; ++ ++// gemm_hc_down: the HC down GEMM V = xn . [w_down; w_inject]^T (T x (hc_lr + hc), K = hc * n_embd) with gr_epi_down ++// as its register epilogue: lo_h = half(silu(v / hc)), inject = v. V is never written; lo_h, inject bitwise gemm_f16 + ++// gr_epi_down. M5 layout only (coop_k4: a lane's elements 4j .. 4j + 3 are four consecutive columns of one row); ++// hc_lr % 4 == 0, hc == 4. grid (ceil((hc_lr + hc) / TN), ceil(T / TM)), 32 * SG threads ++template ++inline void hc_down_tile(device half* A, device half* B, device half* lo_h, device float* inject, ++ constant EpiArgs& e, int t0, int c0) { ++ const int K = int(e.hc * e.n_embd), nr = int(e.hc_lr + e.hc); ++ tensor, tensor_inline> tA(A, dextents(K, int(e.n_tok))); ++ tensor, tensor_inline> tB(B, dextents(K, nr)); ++ constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); ++ matmul2d> op; ++ auto mA = tA.slice(0, t0); ++ auto mB = tB.slice(0, c0); ++ auto acc = op.template get_destination_cooperative_tensor(); ++ op.run(mA, mB, acc); ++ #pragma clang loop unroll(full) ++ for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { ++ const auto ix = acc.get_multidimensional_index(i); ++ const int r = c0 + ix[0], t = t0 + ix[1]; ++ if (t >= int(e.n_tok) || r >= nr) continue; ++ if (r < int(e.hc_lr)) { ++ half4 o; ++ #pragma clang loop unroll(full) ++ for (int c = 0; c < 4; ++c) { ++ const float z = acc[i + c] * 0.25f; // / hc, exact: hc == 4 ++ o[c] = half(z / (1.0f + exp(-z))); ++ } ++ *(device half4*)(lo_h + ulong(t) * e.hc_lr + r) = o; ++ } else { ++ *(device float4*)(inject + ulong(t) * e.hc) = float4(acc[i], acc[i + 1], acc[i + 2], acc[i + 3]); ++ } ++ } ++} ++template ++inline void hc_down_rows(device half* A, device half* B, device half* lo_h, device float* inject, ++ constant EpiArgs& e, int t0, int c0) { ++ // Trimming 32-row parents regresses full tiles on M5. ++ if constexpr (TM > 32) { ++ if (t0 + 16 >= int(e.n_tok)) { ++ hc_down_tile<16, TN, SG>(A, B, lo_h, inject, e, t0, c0); ++ return; ++ } ++ if (t0 + 32 >= int(e.n_tok)) { ++ hc_down_tile<32, TN, SG>(A, B, lo_h, inject, e, t0, c0); ++ return; ++ } ++ } ++ hc_down_tile(A, B, lo_h, inject, e, t0, c0); ++} ++template ++kernel void gemm_hc_down(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device half* lo_h [[buffer(2)]], ++ device float* inject [[buffer(3)]], constant EpiArgs& e [[buffer(4)]], ++ uint3 tg [[threadgroup_position_in_grid]]) { ++ const int c0 = int(tg.x) * TN, t0 = int(tg.y) * TM; ++ if constexpr (TN == 128) { ++ // A single 96-column tail beats serial 64+16 on M5, including SG=2. ++ if (e.hc == 4 && c0 + 64 == int(e.hc_lr)) { ++ hc_down_rows(A, B, lo_h, inject, e, t0, c0); ++ return; ++ } ++ } ++ hc_down_rows(A, B, lo_h, inject, e, t0, c0); ++} ++template [[host_name("gemm_hc_down_16_64_2")]] kernel decltype(gemm_hc_down<16, 64, 2>) gemm_hc_down<16, 64, 2>; ++template [[host_name("gemm_hc_down_32_64_2")]] kernel decltype(gemm_hc_down<32, 64, 2>) gemm_hc_down<32, 64, 2>; ++template [[host_name("gemm_hc_down_32_128_4")]] kernel decltype(gemm_hc_down<32, 128, 4>) gemm_hc_down<32, 128, 4>; ++template [[host_name("gemm_hc_down_64_128_4")]] kernel decltype(gemm_hc_down<64, 128, 4>) gemm_hc_down<64, 128, 4>; ++template [[host_name("gemm_hc_down_32_128_2")]] kernel decltype(gemm_hc_down<32, 128, 2>) gemm_hc_down<32, 128, 2>; ++ ++// gemm_swiglu_up: the shared expert's up GEMM u = x . w_up^T (M x N, K) with swiglu_h as its register epilogue against ++// the gate GEMM's g: h = half(silu(g) * u). u is never written; h bitwise gemm_f16 + swiglu_h. One accumulator: the ++// gate stays a plain gemm_f16 before it. M5 layout only (as gemm_hc_down); N % 4 == 0. ++// grid (ceil(N / TN), ceil(M / TM)), 32 * SG threads ++template ++inline void swiglu_up_tile(device half* A, device half* B, device const float* g, device half* h, ++ constant GemmArgs& a, int t0, int c0) { ++ tensor, tensor_inline> tA(A, dextents(a.K, a.M)); ++ tensor, tensor_inline> tB(B, dextents(a.K, a.N)); ++ constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); ++ matmul2d> op; ++ auto mA = tA.slice(0, t0); ++ auto mB = tB.slice(0, c0); ++ auto acc = op.template get_destination_cooperative_tensor(); ++ op.run(mA, mB, acc); ++ #pragma clang loop unroll(full) ++ for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { ++ const auto ix = acc.get_multidimensional_index(i); ++ const int c = c0 + ix[0], t = t0 + ix[1]; ++ if (t >= a.M || c >= a.N) continue; ++ const ulong idx = ulong(t) * uint(a.N) + uint(c); ++ const float4 gv = *(device const float4*)(g + idx); ++ half4 o; ++ #pragma clang loop unroll(full) ++ for (int k = 0; k < 4; ++k) o[k] = half(gv[k] / (1.0f + exp(-gv[k])) * acc[i + k]); ++ *(device half4*)(h + idx) = o; ++ } ++} ++template ++kernel void gemm_swiglu_up(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device const float* g [[buffer(2)]], ++ device half* h [[buffer(3)]], constant GemmArgs& a [[buffer(4)]], ++ uint3 tg [[threadgroup_position_in_grid]]) { ++ const int t0 = int(tg.y) * TM, c0 = int(tg.x) * TN; ++ if constexpr (TM > 32) { ++ if (t0 + 16 >= a.M) { ++ swiglu_up_tile<16, TN, SG>(A, B, g, h, a, t0, c0); ++ return; ++ } ++ if (t0 + 32 >= a.M) { ++ swiglu_up_tile<32, TN, SG>(A, B, g, h, a, t0, c0); ++ return; ++ } ++ } ++ swiglu_up_tile(A, B, g, h, a, t0, c0); ++} ++template [[host_name("gemm_swiglu_up_32_64_2")]] kernel decltype(gemm_swiglu_up<32, 64, 2>) gemm_swiglu_up<32, 64, 2>; ++template [[host_name("gemm_swiglu_up_64_128_4")]] kernel decltype(gemm_swiglu_up<64, 128, 4>) gemm_swiglu_up<64, 128, 4>; +diff --git a/metal/kernels/qsa.metal b/metal/kernels/qsa.metal +index 9642e58e..63259b46 100644 +--- a/metal/kernels/qsa.metal ++++ b/metal/kernels/qsa.metal +@@ -35,6 +35,13 @@ constant int QSA_NTOK [[function_constant(3)]]; + // the pooled indexer keys too (position pos_base + 4 b of a block). Off: exactly the plain RoPE. + constant bool QSA_MROPE [[function_constant(24)]]; + constant bool qsa_mrope = is_function_constant_defined(QSA_MROPE) && QSA_MROPE; ++// QSA_Q_F16_FC (function constant 25): qsa_prep writes qcur directly as fp16 (device half*) for qsa_flash, skipping qsa_q_f16. ++constant bool QSA_Q_F16_FC [[function_constant(25)]]; ++constant bool qsa_qf16 = is_function_constant_defined(QSA_Q_F16_FC) ? QSA_Q_F16_FC : false; ++// QSA_ROW0 (function constant 28): qsa_prep's threadgroup x is row QSA_ROW0 + x (24: the K/V rows only, for rows whose ++// queries nobody reads). Undefined: 0. ++constant int QSA_ROW0 [[function_constant(28)]]; ++constant int qsa_row0 = is_function_constant_defined(QSA_ROW0) ? QSA_ROW0 : 0; + + struct QsaArgs { + uint pos0; // the window's first cell (relative to the sequence) +@@ -54,34 +61,43 @@ static inline float tg_sum256(float v, threadgroup float* part, ushort sg, ushor + threadgroup_barrier(mem_flags::mem_threadgroup); + float t = 0.0f; + for (int i = 0; i < nsg; ++i) t += part[i]; +- threadgroup_barrier(mem_flags::mem_threadgroup); + return t; + } + + // one row (width w <= 256 = threads) : out = rope((rs * x) * g) (RoPE on the first 64 dims, NEOX pairs i, i+32) + static inline float norm_rope(float x, float xhi_unused, device const float* g, threadgroup float* row, + threadgroup float* part, ushort tid, ushort sg, ushort lane, int w, float eps, +- device const float2* cs, uint absp, bool rope, device const int* mtab) { ++ device const float2* cs, uint absp, bool rope, device const int* mtab, ++ bool reuse = false) { + const float ss = tg_sum256(tid < w ? x * x : 0.0f, part, sg, lane, 8); + const float rs = rsqrt(ss / float(w) + eps); + const float y = tid < w ? (rs * x) * g[tid] : 0.0f; +- row[tid] = y; +- threadgroup_barrier(mem_flags::mem_threadgroup); + float o = y; +- if (rope && tid < QROT) { +- const int p = tid & 31; +- const uint pp = qsa_mrope ? uint(mtab[absp * 3 + uint(p % 3)]) : absp; +- const float2 c = cs[pp * 32 + p]; +- const float a = row[p], b = row[p + 32]; +- o = tid < 32 ? a * c.x - b * c.y : a * c.y + b * c.x; ++ if (rope) { ++ if (tid < QROT) row[tid] = y; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (tid < QROT) { ++ const int p = tid & 31; ++ const uint pp = qsa_mrope ? uint(mtab[absp * 3 + uint(p % 3)]) : absp; ++ const float2 c = cs[pp * 32 + p]; ++ const float a = row[p], b = row[p + 32]; ++ o = tid < 32 ? a * c.x - b * c.y : a * c.y + b * c.x; ++ } + } +- threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (reuse) threadgroup_barrier(mem_flags::mem_threadgroup); + (void)xhi_unused; + return o; + } + ++// QSA_KV_STRIDE (function constant 26): kcur, vcur, iqraw and ikraw are column slices of one [n_tok][QSA_KV_STRIDE] ++// projection (prompt_chunk's fused K | V | indexer K | indexer Q GEMM, 1152 or 1664 columns). Undefined: their own ++// [n_tok][rows] buffers. ++constant int QSA_KV_STRIDE [[function_constant(26)]]; ++constant int qsa_kv_s = is_function_constant_defined(QSA_KV_STRIDE) ? QSA_KV_STRIDE : 0; ++ + // grid (32 rows + 1, n_tok): rows 0..23 q heads, 24..25 k heads, 26..27 v heads, 28..31 indexer q heads; +-// threadgroup 32 (x = 32, y = 0) is the indexer-key threadgroup. 256 threads. ++// threadgroup 32 (x = 32, y = 0) is the indexer-key threadgroup. 256 threads. A grid of 28 (or of 4 from QSA_ROW0 = ++// 24) writes no indexer state: the MTP layer's attention is a dense window, never scored. + kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float* kcur [[buffer(1)]], + device const float* vcur [[buffer(2)]], device const float* iqraw [[buffer(3)]], + device const float* ikraw [[buffer(4)]], device const float* qn [[buffer(5)]], +@@ -95,30 +111,55 @@ kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; +- threadgroup float row[256]; ++ threadgroup float row[QROT]; + threadgroup float amax[8]; +- const int r = int(tg.x); ++ const int r = int(tg.x) + qsa_row0; ++ // the row strides of kcur / vcur, ikraw, iqraw ++ const int s_kv = qsa_kv_s ? qsa_kv_s : QKV * QD, s_ik = qsa_kv_s ? qsa_kv_s : ID; ++ const int s_iq = qsa_kv_s ? qsa_kv_s : IH * ID; + if (r == 32) { +- // the indexer keys, every window token in order (only threadgroup (32, 0) works) +- if (tg.y != 0) return; +- for (int t = 0; t < int(a.n_tok); ++t) { ++ // the indexer keys: threadgroup (32, 0) handles t < 3 (boundary with prior tail + cell 0 dead) and final tail; ++ // every threadgroup (32, t >= 3) with (cell & 3) == 3 pools its 4 in-window keys in parallel. ++ const int t_self = int(tg.y); ++ if (t_self >= 3) { ++ const uint cell = a.pos0 + uint(t_self); ++ if ((cell & 3) == 3) { ++ float m = 0.0f; ++ if (tid < ID) { ++ m = (((ikraw[(t_self - 3) * s_ik + tid] + ikraw[(t_self - 2) * s_ik + tid]) + ++ ikraw[(t_self - 1) * s_ik + tid]) + ikraw[t_self * s_ik + tid]) * 0.25f; ++ } ++ const uint b = cell >> 2; ++ const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab); ++ if (tid < ID) pooled[b * ID + tid] = o; ++ } ++ return; ++ } ++ if (t_self != 0) return; ++ const int n_pre = min(int(a.n_tok), 3); ++ for (int t = 0; t < n_pre; ++t) { + const uint cell = a.pos0 + uint(t); +- const float x = tid < ID ? ikraw[t * ID + tid] : 0.0f; ++ const float x = tid < ID ? ikraw[t * s_ik + tid] : 0.0f; + if (tid < ID) tail[(cell & 3) * ID + tid] = x; + threadgroup_barrier(mem_flags::mem_device); + if (cell == 0) { +- const float o = norm_rope(x, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, 0, false, mtab); ++ const float o = norm_rope(x, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, 0, false, mtab, true); + if (tid < ID) dead[tid] = o; + } + if ((cell & 3) == 3) { + float m = 0.0f; + if (tid < ID) m = (((tail[tid] + tail[ID + tid]) + tail[2 * ID + tid]) + tail[3 * ID + tid]) * 0.25f; + const uint b = cell >> 2; +- const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab); ++ const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab, true); + if (tid < ID) pooled[b * ID + tid] = o; + } + threadgroup_barrier(mem_flags::mem_device); + } ++ const int t_tail0 = max(n_pre, int(a.n_tok) - 4); ++ for (int t = t_tail0; t < int(a.n_tok); ++t) { ++ const uint cell = a.pos0 + uint(t); ++ if (tid < ID) tail[(cell & 3) * ID + tid] = ikraw[t * s_ik + tid]; ++ } + return; + } + const int t = int(tg.y); +@@ -126,11 +167,12 @@ kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float + if (r < QH) { + const float x = qfull[(t * QH + r) * 2 * QD + tid]; + const float o = norm_rope(x, 0.0f, qn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); +- qcur[(t * QH + r) * QD + tid] = o; ++ if (qsa_qf16) ((device half*)qcur)[(t * QH + r) * QD + tid] = half(o); ++ else qcur[(t * QH + r) * QD + tid] = o; + } else if (r < QH + 2 * QKV) { + const bool isv = r >= QH + QKV; + const int h = r - QH - (isv ? QKV : 0); +- float x = (isv ? vcur : kcur)[(t * QKV + h) * QD + tid]; ++ float x = (isv ? vcur : kcur)[t * s_kv + h * QD + tid]; + if (!isv) x = norm_rope(x, 0.0f, kn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); + // int8 with an fp16 scale per 64 values: groups = simdgroup pairs (64 threads) + float m = simd_max(fabs(x)); +@@ -146,8 +188,9 @@ kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float + (isv ? vq : kq)[rowi * QD + tid] = char(q); + if ((tid & 63) == 0) (isv ? vsc : ksc)[rowi * 4 + g] = sh; + } else { ++ if (a.pos0 + a.n_tok <= uint(QBUDGET)) return; + const int h = r - QH - 2 * QKV; +- const float x = tid < ID ? iqraw[(t * IH + h) * ID + tid] : 0.0f; ++ const float x = tid < ID ? iqraw[t * s_iq + 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); + if (tid < ID) qidx[(t * IH + h) * ID + tid] = o; + } +@@ -164,12 +207,21 @@ kernel void qsa_score(device const float* qidx [[buffer(0)]], device const float + const uint b = tg.x * 8 + sg; + if (b > n_bid) return; + device const float* key = b == n_bid ? dead : pooled + b * ID; +- float k[4]; ++ device const float* qp = qidx + (t * IH) * ID + lane; ++ float k[4], q[IH][4]; ++ #pragma clang loop unroll(full) + for (int i = 0; i < 4; ++i) k[i] = key[lane + 32 * i]; ++ #pragma clang loop unroll(full) ++ for (int h = 0; h < IH; ++h) { ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 4; ++i) q[h][i] = qp[h * ID + 32 * i]; ++ } + float s = 0.0f; ++ #pragma clang loop unroll(full) + for (int h = 0; h < IH; ++h) { + float d = 0.0f; +- for (int i = 0; i < 4; ++i) d = fma(qidx[(t * IH + h) * ID + lane + 32 * i], k[i], d); ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 4; ++i) d = fma(q[h][i], k[i], d); + d = simd_sum(d); + s += max(d, 0.0f); + } +@@ -190,7 +242,8 @@ kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [ + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint hist[256]; + threadgroup uint sel[4]; +- threadgroup uint pref[32]; ++ threadgroup uint2 pref[32]; ++ threadgroup uint2 totals; + const int t = int(tg.x); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2, nb = n_bid + 1; + device int* out = ids + t * QBUDGET; +@@ -199,6 +252,10 @@ kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [ + for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); + return; + } ++ const bool cached = nb <= 4096; ++ uint keys[4]; ++ if (cached) ++ for (uint b = tid, i = 0; b < nb; b += 1024, ++i) keys[i] = okey(sc[b]); + // radix select of the `need`-th largest cell key + uint prefix = 0, pmask = 0, need = QBUDGET; + for (int pass = 0; pass < 4; ++pass) { +@@ -206,28 +263,35 @@ kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [ + if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint b = tid; b < nb; b += 1024) { +- const uint k = okey(sc[b]); ++ const uint k = cached ? keys[b / 1024] : okey(sc[b]); + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); +- if (tid == 0) { +- uint acc = 0, d = 0; +- for (int v = 255; v >= 0; --v) { +- const uint h = atomic_load_explicit(&hist[v], memory_order_relaxed); +- if (acc + h >= need) { d = uint(v); break; } +- acc += h; ++ if (sg == 0) { ++ uint counts[8], sum = 0; ++ for (uint j = 0; j < 8; ++j) { ++ counts[j] = atomic_load_explicit(&hist[255 - (lane * 8 + j)], memory_order_relaxed); ++ sum += counts[j]; ++ } ++ uint acc = simd_prefix_exclusive_sum(sum); ++ if (acc < need && acc + sum >= need) { ++ for (uint j = 0; j < 8; ++j) { ++ if (acc + counts[j] >= need) { ++ sel[0] = 255 - (lane * 8 + j); ++ sel[1] = need - acc; ++ break; ++ } ++ acc += counts[j]; ++ } + } +- sel[0] = d; +- sel[1] = need - acc; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; +- threadgroup_barrier(mem_flags::mem_threadgroup); + } + // T = prefix; take every cell with key > T and the first `need` cells (lowest index) with key == T, ascending. + // Pass over blocks in order of index with a prefix scan of each block's selected cell count. +@@ -237,36 +301,35 @@ kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [ + const uint b = b0 + tid; + uint w = 0, weq = 0; + if (b < nb) { +- const uint k = okey(sc[b]); ++ const uint k = cached ? keys[b / 1024] : okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + if (k > T) w = cells; + else if (k == T) weq = cells; + } +- // equal-key cells: the lowest first, `need` of them in total + const uint eq_before = simd_prefix_exclusive_sum(weq); +- if (lane == 31) pref[sg] = eq_before + weq; ++ const uint gt_before = simd_prefix_exclusive_sum(w); ++ if (lane == 31) pref[sg] = uint2(eq_before + weq, gt_before + w); + threadgroup_barrier(mem_flags::mem_threadgroup); +- uint eq_sg = 0; +- for (int i = 0; i < int(sg); ++i) eq_sg += pref[i]; +- uint eq_tot = 0; +- for (int i = 0; i < 32; ++i) eq_tot += pref[i]; ++ if (sg == 0) { ++ const uint2 count = pref[lane]; ++ const uint2 before = uint2(simd_prefix_exclusive_sum(count.x), ++ simd_prefix_exclusive_sum(count.y)); ++ pref[lane] = before; ++ if (lane == 31) totals = before + count; ++ } + threadgroup_barrier(mem_flags::mem_threadgroup); +- const uint eq0 = eq_taken + eq_sg + eq_before; +- uint take_eq = 0; +- if (weq && eq0 < need) take_eq = min(weq, need - eq0); ++ const uint2 group_before = pref[sg], count = totals; ++ simdgroup_barrier(mem_flags::mem_threadgroup); ++ const uint eq0 = eq_taken + group_before.x + eq_before; ++ const uint take_eq = eq0 < need ? min(weq, need - eq0) : 0u; + const uint mine = w + take_eq; +- const uint before = simd_prefix_exclusive_sum(mine); +- if (lane == 31) pref[sg] = before + mine; +- threadgroup_barrier(mem_flags::mem_threadgroup); +- uint off = base; +- for (int i = 0; i < int(sg); ++i) off += pref[i]; +- uint tot = 0; +- for (int i = 0; i < 32; ++i) tot += pref[i]; +- threadgroup_barrier(mem_flags::mem_threadgroup); +- off += before; ++ // Clipping the equal-key prefix preserves the lowest-cell tie break. ++ const uint remaining = need - min(need, eq_taken); ++ const uint off = base + group_before.y + gt_before + ++ min(remaining, group_before.x + eq_before); + for (uint c = 0; c < mine; ++c) out[off + c] = int(4 * b + c); +- base += tot; +- eq_taken += eq_tot; ++ base += count.y + min(remaining, count.x); ++ eq_taken += count.x; + } + } + +@@ -293,19 +356,28 @@ kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* + return; + } + const int nc = min(QCH, int(width) - c0); ++ #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + g * HPG + h) * QD + tid]; + if (tid < QCH) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + // scores: 12 x 64 dots of 256; simdgroup sg takes cells sg, sg + 8, ...; lanes split the 256 dims (8 each) ++ threadgroup const float* q_lane = &qs[0][lane * 8]; + for (int c = sg; c < nc; c += 8) { + const uint rowi = uint(cid[c]) * QKV + uint(g); +- device const char* kr = kq + rowi * QD; +- float kv[8]; ++ device const char* kr = kq + rowi * QD + lane * 8; ++ const char4 k0 = *(device const char4*)(kr); ++ const char4 k1 = *(device const char4*)(kr + 4); + const float scl = float(ksc[rowi * 4 + (lane >> 3)]); +- for (int i = 0; i < 8; ++i) kv[i] = float(kr[lane * 8 + i]) * scl; ++ const float kv[8] = { ++ float(k0.x) * scl, float(k0.y) * scl, float(k0.z) * scl, float(k0.w) * scl, ++ float(k1.x) * scl, float(k1.y) * scl, float(k1.z) * scl, float(k1.w) * scl ++ }; ++ #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { ++ threadgroup const float* qh = q_lane + h * QD; + float d = 0.0f; +- for (int i = 0; i < 8; ++i) d = fma(qs[h][lane * 8 + i], kv[i], d); ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 8; ++i) d = fma(qh[i], kv[i], d); + d = simd_sum(d); + if (lane == 0) s[h][c] = d * 0.0625f; + } +@@ -313,28 +385,44 @@ kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* + threadgroup_barrier(mem_flags::mem_threadgroup); + // per head: max, exp, sum (simdgroup h % 8 handles heads h, h + 8) + threadgroup float mh[HPG], lh[HPG]; ++ // nc <= QCH = 64: a lane holds cells lane and lane + 32 in registers (the same max and sum as a strided loop) ++ const bool in0 = int(lane) < nc, in1 = int(lane) + 32 < nc; + for (int h = sg; h < HPG; h += 8) { +- float m = -INFINITY; +- for (int c = lane; c < nc; c += 32) m = max(m, s[h][c]); +- m = simd_max(m); +- float l = 0.0f; +- for (int c = lane; c < nc; c += 32) { +- const float p = exp(s[h][c] - m); +- s[h][c] = p; +- l += p; +- } +- l = simd_sum(l); ++ const float s0 = in0 ? s[h][lane] : -INFINITY, s1 = in1 ? s[h][lane + 32] : -INFINITY; ++ const float m = simd_max(max(s0, s1)); ++ const float p0 = in0 ? exp(s0 - m) : 0.0f, p1 = in1 ? exp(s1 - m) : 0.0f; ++ if (in0) s[h][lane] = p0; ++ if (in1) s[h][lane + 32] = p1; ++ const float l = simd_sum(p0 + p1); + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // acc[h][j] = sum_c p[h][c] v[c][j]: thread = dim j + float acc[HPG]; ++ #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; +- for (int c = 0; c < nc; ++c) { ++ // four cells' V loads in flight at once; the fma order per head stays c, c + 1, ... ++ auto vcell = [&](int c) { + const uint rowi = uint(cid[c]) * QKV + uint(g); +- const float v = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); ++ return float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); ++ }; ++ int c = 0; ++ for (; c + 4 <= nc; c += 4) { ++ const float v0 = vcell(c), v1 = vcell(c + 1), v2 = vcell(c + 2), v3 = vcell(c + 3); ++ #pragma clang loop unroll(full) ++ for (int h = 0; h < HPG; ++h) { ++ float x = fma(s[h][c], v0, acc[h]); ++ x = fma(s[h][c + 1], v1, x); ++ x = fma(s[h][c + 2], v2, x); ++ acc[h] = fma(s[h][c + 3], v3, x); ++ } ++ } ++ for (; c < nc; ++c) { ++ const float v = vcell(c); ++ #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][c], v, acc[h]); + } ++ #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (g * HPG + h) * (QD + 2); + o[tid] = acc[h]; +@@ -342,43 +430,63 @@ kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* + } + } + ++constant float qsa_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; ++constant float qsa_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; ++ ++static inline void qsa_write_prepared(float y, device float* xp, constant QsaArgs& a, ++ int t, int h, ushort tid) { ++ // prepared planes (xprep layout) of the n = QH * QD vector ++ const uint n = QH * QD, plane = a.n_tok * n, d = uint(h) * QD + tid; ++ const int j = tid & 15; ++ device float* raw = xp + uint(t) * n + d; ++ raw[0] = y; ++ raw[plane] = y * qsa_q2s[j & 7]; ++ raw[2 * plane] = y * qsa_p4s[j & 3]; ++ raw[3 * plane] = (j & 1) ? y * 0.00390625f : y; ++ float s0 = 0.0f, s1 = 0.0f; ++ #pragma clang loop unroll(full) ++ for (int k = 0; k < 16; k += 2) { ++ s0 += simd_shuffle_down(y, ushort(k)); ++ s1 += simd_shuffle_down(y, ushort(k + 1)); ++ } ++ if (j == 0) xp[4 * plane + uint(t) * (n / 16) + d / 16] = s0 + s1; ++} ++ + // grid (QH, n_tok), 256 threads: out = merge * sigmoid(gate), as prepared activations of n = QH * QD per token + kernel void qsa_merge(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], + device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { +- threadgroup float msum[256]; + const int h = int(tg.x), t = int(tg.y); + const int nchunks = (QBUDGET + QCH - 1) / QCH; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nc = int((width + QCH - 1) / QCH); ++ constexpr uint STRIDE = QH * (QD + 2); + device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2); ++ // the chunks' (m, l) and weights exp(m - M) once per threadgroup, not once per thread (nc <= 33 < 64) ++ threadgroup float mk[64], lk[64], ek[64], L_sh; ++ if (tid < uint(nc)) { ++ const float2 ml = *(device const float2*)(pp + tid * STRIDE + QD); ++ mk[tid] = ml.x; ++ lk[tid] = ml.y; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); + float M = -INFINITY; +- for (int k = 0; k < nc; ++k) M = max(M, pp[k * QH * (QD + 2) + QD]); +- float L = 0.0f, o = 0.0f; +- for (int k = 0; k < nc; ++k) { +- device const float* q = pp + k * QH * (QD + 2); +- const float e = exp(q[QD] - M); +- L = fma(q[QD + 1], e, L); +- o = fma(q[tid], e, o); ++ for (int k = tid & 31; k < nc; k += 32) M = max(M, mk[k]); ++ M = simd_max(M); ++ if (tid < uint(nc)) ek[tid] = exp(mk[tid] - M); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (tid == 0) { ++ float L = 0.0f; ++ for (int k = 0; k < nc; ++k) L = fma(lk[k], ek[k], L); ++ L_sh = L; + } +- o = o / L; ++ float o = 0.0f; ++ device const float* q = pp + tid; ++ for (int k = 0; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ o = o / L_sh; + const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; + const float y = o * (1.0f / (1.0f + exp(-gate))); +- // prepared planes (xprep layout) of the n = QH * QD vector +- const uint n = QH * QD, plane = a.n_tok * n, d = uint(h) * QD + tid; +- const int j = tid & 15; +- const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; +- const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; +- device float* raw = xp + uint(t) * n + d; +- raw[0] = y; +- raw[plane] = y * q2s[j & 7]; +- raw[2 * plane] = y * p4s[j & 3]; +- raw[3 * plane] = (j & 1) ? y * 0.00390625f : y; +- msum[tid] = y; +- threadgroup_barrier(mem_flags::mem_threadgroup); +- if (j == 0) { +- float s0 = 0.0f, s1 = 0.0f; +- for (int k = 0; k < 16; k += 2) { s0 += msum[tid + k]; s1 += msum[tid + k + 1]; } +- xp[4 * plane + uint(t) * (n / 16) + d / 16] = s0 + s1; +- } ++ qsa_write_prepared(y, xp, a, t, h, tid); + } ++ +diff --git a/metal/kernels/qsa2.metal b/metal/kernels/qsa2.metal +index 8ae189b6..d5b594c3 100644 +--- a/metal/kernels/qsa2.metal ++++ b/metal/kernels/qsa2.metal +@@ -29,9 +29,11 @@ kernel void qsa_attn2(device const float* qcur [[buffer(0)]], device const char* + if (uint(c0) >= width) return; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + const int nc = min(QCH2, int(width) - c0); ++ #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + h0 + h) * QD + tid]; + if (tid < QCH2) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); ++ threadgroup const float* q_lane = &qs[0][lane * 8]; + for (int c = sg; c < nc; c += 8) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char4* kr = (device const char4*)(kq + rowi * QD + lane * 8); +@@ -39,9 +41,11 @@ kernel void qsa_attn2(device const float* qcur [[buffer(0)]], device const char* + const float scl = float(ksc[rowi * 4 + (lane >> 3)]); + const float kv[8] = {float(k0.x) * scl, float(k0.y) * scl, float(k0.z) * scl, float(k0.w) * scl, + float(k1.x) * scl, float(k1.y) * scl, float(k1.z) * scl, float(k1.w) * scl}; ++ #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + float d = 0.0f; +- for (int i = 0; i < 8; ++i) d = fma(qs[h][lane * 8 + i], kv[i], d); ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 8; ++i) d = fma(q_lane[h * QD + i], kv[i], d); + d = simd_sum(d); + if (lane == 0) s[h][c] = d * 0.0625f; + } +@@ -58,20 +62,29 @@ kernel void qsa_attn2(device const float* qcur [[buffer(0)]], device const char* + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float acc[HPG]; ++ #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; +- for (int cb = 0; cb < nc; cb += 8) { ++ const uint v_scale_off = tid >> 6; ++ int cb = 0; ++ for (; cb + 8 <= nc; cb += 8) { + float v[8]; +- const int m = min(8, nc - cb); ++ #pragma clang loop unroll(full) + for (int c = 0; c < 8; ++c) { +- v[c] = 0.0f; +- if (c < m) { +- const uint rowi = uint(cid[cb + c]) * QKV + uint(g); +- v[c] = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); +- } ++ const uint rowi = uint(cid[cb + c]) * QKV + uint(g); ++ v[c] = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + v_scale_off]); + } +- for (int c = 0; c < m; ++c) ++ #pragma clang loop unroll(full) ++ for (int c = 0; c < 8; ++c) ++ #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb + c], v[c], acc[h]); + } ++ for (; cb < nc; ++cb) { ++ const uint rowi = uint(cid[cb]) * QKV + uint(g); ++ const float v = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + v_scale_off]); ++ #pragma clang loop unroll(full) ++ for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb], v, acc[h]); ++ } ++ #pragma clang loop unroll(full) + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (h0 + h) * (QD + 2); + o[tid] = acc[h]; +@@ -82,40 +95,37 @@ kernel void qsa_attn2(device const float* qcur [[buffer(0)]], device const char* + kernel void qsa_merge2(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], + device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { +- threadgroup float msum[256]; ++ threadgroup float mk[128], lk[128], ek[128], L_sh; + const int h = int(tg.x), t = int(tg.y); + const int nchunks = (QBUDGET + QCH2 - 1) / QCH2; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nc = int((width + QCH2 - 1) / QCH2); ++ constexpr uint STRIDE = QH * (QD + 2); + device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2); ++ if (tid < uint(nc)) { ++ const float2 ml = *(device const float2*)(pp + tid * STRIDE + QD); ++ mk[tid] = ml.x; ++ lk[tid] = ml.y; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); + float M = -INFINITY; +- for (int k = 0; k < nc; ++k) M = max(M, pp[k * QH * (QD + 2) + QD]); +- float L = 0.0f, o = 0.0f; +- for (int k = 0; k < nc; ++k) { +- device const float* q = pp + k * QH * (QD + 2); +- const float e = exp(q[QD] - M); +- L = fma(q[QD + 1], e, L); +- o = fma(q[tid], e, o); ++ for (int k = tid & 31; k < nc; k += 32) M = max(M, mk[k]); ++ M = simd_max(M); ++ if (tid < uint(nc)) ek[tid] = exp(mk[tid] - M); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (tid == 0) { ++ float L = 0.0f; ++ for (int k = 0; k < nc; ++k) L = fma(lk[k], ek[k], L); ++ L_sh = L; + } +- o = o / L; ++ float o = 0.0f; ++ device const float* q = pp + tid; ++ for (int k = 0; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ o = o / L_sh; + const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; + const float y = o * (1.0f / (1.0f + exp(-gate))); +- const uint n = QH * QD, plane = a.n_tok * n, d = uint(h) * QD + tid; +- const int j = tid & 15; +- const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; +- const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; +- device float* raw = xp + uint(t) * n + d; +- raw[0] = y; +- raw[plane] = y * q2s[j & 7]; +- raw[2 * plane] = y * p4s[j & 3]; +- raw[3 * plane] = (j & 1) ? y * 0.00390625f : y; +- msum[tid] = y; +- threadgroup_barrier(mem_flags::mem_threadgroup); +- if (j == 0) { +- float s0 = 0.0f, s1 = 0.0f; +- for (int k = 0; k < 16; k += 2) { s0 += msum[tid + k]; s1 += msum[tid + k + 1]; } +- xp[4 * plane + uint(t) * (n / 16) + d / 16] = s0 + s1; +- } ++ qsa_write_prepared(y, xp, a, t, h, tid); + } + + kernel void qsa_topk2(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]], +@@ -133,37 +143,76 @@ kernel void qsa_topk2(device const float* scores [[buffer(0)]], device int* ids + for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); + return; + } ++ if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); ++ threadgroup_barrier(mem_flags::mem_threadgroup); + uint prefix = 0, pmask = 0, need = QBUDGET; +- for (int pass = 0; pass < 4; ++pass) { +- const int shift = 24 - 8 * pass; +- if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); +- threadgroup_barrier(mem_flags::mem_threadgroup); +- for (uint b = tid; b < nb; b += 1024) { +- const uint k = okey(sc[b]); +- if ((k & pmask) == prefix) { +- const uint w = b < n_bid ? 4u : (n_kv & 3u); +- if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); +- } ++ if (nb <= 4096) { ++ uint keys[4]; ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 4; ++i) { ++ const uint b = tid + 1024u * uint(i); ++ keys[i] = b < nb ? okey(sc[b]) : 0u; + } +- threadgroup_barrier(mem_flags::mem_threadgroup); +- // thread i < 256 holds bin 255 - i: an exclusive prefix over i = the cells of the higher digits +- uint cnt = 0, ex = 0; +- if (tid < 256) { +- cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); +- ex = simd_prefix_exclusive_sum(cnt); +- if (lane == 31) sgs[sg] = ex + cnt; ++ for (int pass = 0; pass < 4; ++pass) { ++ const int shift = 24 - 8 * pass; ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 4; ++i) { ++ const uint b = tid + 1024u * uint(i); ++ if (b >= nb) break; ++ const uint k = keys[i]; ++ if ((k & pmask) == prefix) { ++ const uint w = b < n_bid ? 4u : (n_kv & 3u); ++ if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); ++ } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ uint cnt = 0, ex = 0; ++ if (tid < 256) { ++ cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); ++ if (pass < 3) atomic_store_explicit(&hist[255 - tid], 0u, memory_order_relaxed); ++ ex = simd_prefix_exclusive_sum(cnt); ++ if (lane == 31) sgs[sg] = ex + cnt; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (tid < 256) { ++ uint above = ex; ++ for (int i = 0; i < int(sg); ++i) above += sgs[i]; ++ if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ prefix |= sel[0] << shift; ++ pmask |= 255u << shift; ++ need = sel[1]; + } +- threadgroup_barrier(mem_flags::mem_threadgroup); +- if (tid < 256) { +- uint above = ex; +- for (int i = 0; i < int(sg); ++i) above += sgs[i]; +- if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } ++ } else { ++ for (int pass = 0; pass < 4; ++pass) { ++ const int shift = 24 - 8 * pass; ++ for (uint b = tid; b < nb; b += 1024) { ++ const uint k = okey(sc[b]); ++ if ((k & pmask) == prefix) { ++ const uint w = b < n_bid ? 4u : (n_kv & 3u); ++ if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); ++ } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ uint cnt = 0, ex = 0; ++ if (tid < 256) { ++ cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); ++ if (pass < 3) atomic_store_explicit(&hist[255 - tid], 0u, memory_order_relaxed); ++ ex = simd_prefix_exclusive_sum(cnt); ++ if (lane == 31) sgs[sg] = ex + cnt; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (tid < 256) { ++ uint above = ex; ++ for (int i = 0; i < int(sg); ++i) above += sgs[i]; ++ if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ prefix |= sel[0] << shift; ++ pmask |= 255u << shift; ++ need = sel[1]; + } +- threadgroup_barrier(mem_flags::mem_threadgroup); +- prefix |= sel[0] << shift; +- pmask |= 255u << shift; +- need = sel[1]; +- threadgroup_barrier(mem_flags::mem_threadgroup); + } + const uint T = prefix; + const uint per = (nb + 1023) / 1024, b0 = min(nb, uint(tid) * per), b1 = min(nb, b0 + per); +diff --git a/metal/kernels/qsa_prompt.metal b/metal/kernels/qsa_prompt.metal +index df7caf79..444e0c46 100644 +--- a/metal/kernels/qsa_prompt.metal ++++ b/metal/kernels/qsa_prompt.metal +@@ -30,8 +30,8 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* + constexpr int HPG = QH / QKV, NCH = (QBUDGET + QCH - 1) / QCH; + threadgroup float s[16 * QCH]; + threadgroup half p[16 * QCH]; +- threadgroup int cid[QCH]; +- threadgroup half vs[QCH][4]; ++ threadgroup int cid[QCH]; // coop_k4: byte offsets; otherwise cell ids ++ threadgroup half4 ks[QCH], vs[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; +@@ -42,9 +42,16 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* + } + const int nc = min(QCH, int(width) - c0); + for (int c = lane; c < QCH; c += 32) { +- const int id = c < nc ? ids[t * QBUDGET + c0 + c] : -1; +- cid[c] = id; +- for (int j = 0; j < 4; ++j) vs[c][j] = id >= 0 ? vsc[(uint(id) * QKV + g) * 4 + j] : half(0.0f); ++ const int id = c < nc ? (n_kv <= QBUDGET ? c0 + c : ids[t * QBUDGET + c0 + c]) : -1; ++ cid[c] = coop_k4 ? int((uint(max(id, 0)) * QKV + g) * QD) : id; ++ if (id >= 0) { ++ const uint row4 = (uint(id) * QKV + g) * 4; ++ if (coop_k4) ks[c] = *(device const half4*)(ksc + row4); ++ vs[c] = *(device const half4*)(vsc + row4); ++ } else { ++ ks[c] = half4(0.0h); ++ vs[c] = half4(0.0h); ++ } + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (qsa_stop == 1) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 1.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } +@@ -59,6 +66,17 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; ++ // the second half repeats these eight cells, 16 k later; empty cells read row 0 at scale 0. ++ uint ko[8]; ++ float ks8[8][4]; ++ if (coop_k4) ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 8; ++j) { ++ const int c = rk.get_multidimensional_index(uint16_t(4 * j))[1]; ++ ko[j] = uint(cid[c]); ++#pragma clang loop unroll(full) ++ for (int h = 0; h < 4; ++h) ks8[j][h] = float(ks[c][h]); ++ } + for (int k0 = 0; k0 < QD; k0 += 32) { + if (!coop_k4) { + for (uint16_t i = 0; i < rk.get_capacity(); ++i) { +@@ -70,14 +88,8 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* + } else + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < rk.get_capacity(); i += 4) { +- const auto ix = rk.get_multidimensional_index(i); +- const int k = k0 + ix[0], c = ix[1], id = cid[c]; +- half4 w = half4(0.0h); +- if (id >= 0) { +- const uint row = uint(id) * QKV + g; +- const char4 q4 = *(device const char4*)(kq + row * QD + k); +- w = half4(float4(q4) * float(ksc[row * 4 + (k >> 6)])); +- } ++ const int k = k0 + rk.get_multidimensional_index(i)[0], j = (i >> 2) & 7; ++ const half4 w = half4(float4(*(device const char4*)(kq + (ko[j] + uint(k)))) * ks8[j][k >> 6]); + rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w; + } + tA = tQ.template slice<32, 16>(k0, 0); +@@ -91,13 +103,16 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* + // ---- softmax pieces per head (lane = head for the reductions; 12 heads) + float mh = -INFINITY, lh = 0.0f; + if (lane < HPG) { ++#pragma clang loop unroll_count(4) + for (int c = 0; c < nc; ++c) mh = max(mh, s[lane * QCH + c] * 0.0625f); ++#pragma clang loop unroll_count(4) + for (int c = 0; c < QCH; ++c) { + const float e = c < nc ? exp(s[lane * QCH + c] * 0.0625f - mh) : 0.0f; + lh += e; + p[lane * QCH + c] = half(e); + } + } else if (lane < 16) { ++#pragma clang loop unroll_count(4) + for (int c = 0; c < QCH; ++c) p[lane * QCH + c] = half(0.0f); + } + simdgroup_barrier(mem_flags::mem_threadgroup); +@@ -115,12 +130,30 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; +- for (int c0 = 0; c0 < QCH; c0 += 32) { ++ for (int c0 = 0; c0 < nc; c0 += 32) { ++ if (!coop_k4) { + #pragma clang loop unroll(full) +- for (uint16_t i = 0; i < rv.get_capacity(); ++i) { +- const auto ix = rv.get_multidimensional_index(i); // (cell, dim) +- const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; +- rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vs[c][d >> 6])) : half(0.0f); ++ for (uint16_t i = 0; i < rv.get_capacity(); ++i) { ++ const auto ix = rv.get_multidimensional_index(i); // (cell, dim) ++ const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; ++ rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vs[c][d >> 6])) : half(0.0f); ++ } ++ } else { ++ uint vo[8]; ++ float vs8[8]; ++ const int cb = c0 + rv.get_multidimensional_index(uint16_t(0))[0]; ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 8; ++j) { ++ const int c = cb + (j & 3) + 16 * (j >> 2); ++ vo[j] = uint(cid[c]) + uint(d0); ++ vs8[j] = float(vs[c][d0 >> 6]); ++ } ++#pragma clang loop unroll(full) ++ for (uint16_t i = 0; i < rv.get_capacity(); i += 4) { ++ const int d = rv.get_multidimensional_index(i)[1], r = i < rv.get_capacity() / 2 ? 0 : 4; ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 4; ++j) rv[i + j] = half(float(vq[vo[r + j] + uint(d)]) * vs8[r + j]); ++ } + } + tPh = tP.template slice<32, 16>(c0, 0); + op.run(tPh, rv, acc); +@@ -136,7 +169,13 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* + simdgroup_barrier(mem_flags::mem_threadgroup); + acc.store(tT); + simdgroup_barrier(mem_flags::mem_threadgroup); +- for (int it = lane; it < HPG * 64; it += 32) pp[(it >> 6) * (QD + 2) + d0 + (it & 63)] = s[it]; ++ for (int it = lane; it < HPG * 16; it += 32) { ++ const int h = it >> 4, d = (it & 15) * 4; ++ const float4 v = *(threadgroup const float4*)(s + h * 64 + d); ++ // odd heads start on an 8-byte boundary, not a float4 boundary. ++ device float* dst = pp + h * (QD + 2) + d0 + d; ++ dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w; ++ } + } + } + if (lane < HPG) { pp[lane * (QD + 2) + QD] = mh; pp[lane * (QD + 2) + QD + 1] = lh; } +@@ -160,6 +199,10 @@ kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const fl + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float ot[64 * 64]; + const int b0 = int(tg.x) * 64, t0 = int(tg.y) * 16; ++ const uint max_t = min(uint(t0 + 15), a.n_tok - 1u); ++ const uint max_bid = (a.pos0 + max_t + 1u) >> 2; ++ // column zero still owns every token's tail score. ++ if (tg.x != 0 && uint(b0) >= max_bid) return; + const uint nb_all = (a.pos0 + a.n_tok) / 4 + 1; // blocks any token of the chunk can score + tensor, tensor_inline> tQ((device float*)qidx, dextents(ID, int(a.n_tok) * IH)); + tensor, tensor_inline> tK((device float*)pooled, dextents(ID, int(nb_all))); +@@ -171,9 +214,10 @@ kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const fl + op.run(mA, mB, tO); + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint stride = a.max_cells / 4 + 2; +- for (int idx = tid; idx < 16 * 64; idx += 128) { +- const int tt = idx >> 6, j = idx & 63, t = t0 + tt, b = b0 + j; +- if (t >= int(a.n_tok)) continue; ++ const int j = tid & 63, b = b0 + j; ++ for (int tt = tid >> 6; tt < 16; tt += 2) { ++ const int t = t0 + tt; ++ if (t >= int(a.n_tok)) break; + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + if (uint(b) >= n_bid) continue; + float s = 0.0f; +@@ -181,13 +225,13 @@ kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const fl + scores[uint(t) * stride + b] = s; + } + // the tail block of each token in this threadgroup's block range (the first column threadgroup does it) +- if (tg.x == 0) ++ if (tg.x == 0) { ++ float k[4]; ++ for (int i = 0; i < 4; ++i) k[i] = dead[lane + 32 * i]; + for (int tt = sg; tt < 16; tt += 4) { + const int t = t0 + tt; + if (t >= int(a.n_tok)) break; + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; +- float k[4]; +- for (int i = 0; i < 4; ++i) k[i] = dead[lane + 32 * i]; + float s = 0.0f; + for (int h = 0; h < IH; ++h) { + float d = 0.0f; +@@ -198,6 +242,7 @@ kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const fl + if ((n_kv & 3) != 0) s += 1e9f; + if (lane == 0) scores[uint(t) * stride + n_bid] = s; + } ++ } + } + + // qsa_flash: the whole attention of a (kv head g, token t) in one threadgroup (no split-K partials, no merge): the +@@ -211,10 +256,10 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq + 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 HPG = QH / QKV; +- threadgroup float sp[4][16 * QCH]; +- threadgroup half p[16 * QCH]; +- threadgroup int cid[QCH]; +- threadgroup float vs[QCH]; ++ threadgroup float sp[4][HPG * QCH]; ++ threadgroup half* p = (threadgroup half*)sp[1]; // P [16][QCH] in fp16 over a dead S partial ++ threadgroup int cid[QCH]; // coop_k4: byte offsets; otherwise cell ids ++ threadgroup half4 ks_tg[QCH], vs_tg[QCH]; + threadgroup float fac[16], lsum[16], mrun[16]; + const int g = int(tg.x), t = int(tg.y); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); +@@ -240,15 +285,26 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) tot[i] = 0.0f; + const int d0 = sg * 64; ++ if (tid < 16 - HPG) fac[HPG + tid] = 0.0f; // the padding rows' factor (P's rows are zeroed each chunk) + for (int ch = 0; ch < nch; ++ch) { + const int c0 = ch * QCH, nc = min(QCH, int(width) - c0); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < QCH) { +- const int id = tid < nc ? ids[t * QBUDGET + c0 + tid] : -1; +- cid[tid] = id; +- vs[tid] = id >= 0 ? float(vsc[(uint(id) * QKV + g) * 4 + sg * 0]) : 0.0f; // (per-simdgroup scale below) ++ const int id = tid < nc ? (n_kv <= QBUDGET ? c0 + tid : ids[t * QBUDGET + c0 + tid]) : -1; ++ cid[tid] = coop_k4 ? int((uint(max(id, 0)) * QKV + g) * QD) : id; ++ if (id >= 0) { ++ const uint row4 = (uint(id) * QKV + g) * 4; ++ ks_tg[tid] = *(device const half4*)(ksc + row4); ++ vs_tg[tid] = *(device const half4*)(vsc + row4); ++ } else { ++ ks_tg[tid] = half4(0.0h); ++ vs_tg[tid] = half4(0.0h); ++ } + } + threadgroup_barrier(mem_flags::mem_threadgroup); ++ // Eight lanes per head; keep the denominator's original sequential sum. ++ const int h = sg * 4 + (lane >> 3), hlane = lane & 7; ++ float sc[8]; + // S partial over this simdgroup's 64 dims + { + auto rk = opS.template get_right_input_cooperative_tensor(); +@@ -257,68 +313,111 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq + __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < s.get_capacity(); ++i) + if (s.is_valid_element(i)) s[i] = 0.0f; ++ // the lane's eight cells (one a run of four k; the second half repeats them 16 k on): row pointers and ++ // scales once a chunk, unconditional loads (an empty cell reads row 0 at scale 0, and a +-0 operand sums to ++ // the same value as the skipped one; the branchy loads cost half the kernel) ++ uint ko[8]; ++ float ks8[8]; ++ if (coop_k4) { ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 8; ++j) { ++ const int c = rk.get_multidimensional_index(uint16_t(4 * j))[1]; ++ ko[j] = uint(cid[c]) + uint(d0 + rk.get_multidimensional_index(uint16_t(0))[0]); ++ ks8[j] = float(ks_tg[c][sg]); ++ } ++ } ++#pragma clang loop unroll(full) + for (int k0 = d0; k0 < d0 + 64; k0 += 32) { + if (!coop_k4) { + for (uint16_t i = 0; i < rk.get_capacity(); ++i) { + const auto ix = rk.get_multidimensional_index(i); + const int k = k0 + ix[0], id = cid[ix[1]]; + const uint row = uint(id) * QKV + g; +- rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)])) : half(0.0f); ++ rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ks_tg[ix[1]][k >> 6])) : half(0.0f); + } + } else + #pragma clang loop unroll(full) + for (uint16_t i = 0; i < rk.get_capacity(); i += 4) { +- const auto ix = rk.get_multidimensional_index(i); +- const int k = k0 + ix[0], c = ix[1], id = cid[c]; +- half4 w = half4(0.0h); +- if (id >= 0) { +- const uint row = uint(id) * QKV + g; +- const char4 q4 = *(device const char4*)(kq + row * QD + k); +- w = half4(float4(q4) * float(ksc[row * 4 + (k >> 6)])); +- } ++ const int k = k0 - d0 + (i >= rk.get_capacity() / 2 ? 16 : 0), j = (i >> 2) & 7; ++ const half4 w = half4(float4(*(device const char4*)(kq + ko[j] + uint(k))) * ks8[j]); + rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w; + } + tA = tQ.template slice<32, 16>(k0, 0); + opS.run(tA, rk, s); + } +- tensor, tensor_inline> tS(sp[sg], dextents(QCH, 16)); ++ // Preserve ((s0 + s1) + (s2 + s3)) while storing all partials together. ++ tensor, tensor_inline> tS(sp[sg], dextents(QCH, HPG)); + s.store(tS); + } + threadgroup_barrier(mem_flags::mem_threadgroup); +- // online softmax per head row: simdgroup 0, lane = row +- if (sg == 0 && lane < 16) { +- if (lane < HPG) { +- float m = mrun[lane]; +- for (int c = 0; c < nc; ++c) { +- const int q = lane * QCH + c; +- m = max(m, (((sp[0][q] + sp[1][q]) + (sp[2][q] + sp[3][q])) * 0.0625f)); +- } +- const float f = exp(mrun[lane] - m); +- float l = lsum[lane] * f; +- for (int c = 0; c < QCH; ++c) { +- const int q = lane * QCH + c; +- const float e = c < nc ? exp(((sp[0][q] + sp[1][q]) + (sp[2][q] + sp[3][q])) * 0.0625f - m) : 0.0f; +- l += e; +- p[q] = half(e); ++ if (h < HPG) ++#pragma clang loop unroll(full) ++ for (int k = 0; k < 8; ++k) { ++ const int q = h * QCH + hlane + 8 * k; ++ sc[k] = ((sp[0][q] + sp[1][q]) + (sp[2][q] + sp[3][q])) * 0.0625f; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); // P overwrites sp[1] ++ if (h < HPG) { ++ float m = mrun[h]; ++#pragma clang loop unroll(full) ++ for (int k = 0; k < 8; ++k) ++ if (hlane + 8 * k < nc) m = max(m, sc[k]); ++ m = max(m, simd_shuffle_xor(m, 4)); ++ m = max(m, simd_shuffle_xor(m, 2)); ++ m = max(m, simd_shuffle_xor(m, 1)); ++#pragma clang loop unroll(full) ++ for (int k = 0; k < 8; ++k) { ++ const int c = hlane + 8 * k, q = h * QCH + c; ++ const float e = c < nc ? exp(sc[k] - m) : 0.0f; ++ p[q] = half(e); ++ sp[0][q] = e; ++ } ++ simdgroup_barrier(mem_flags::mem_threadgroup); ++ if (hlane == 0) { ++ const float f = exp(mrun[h] - m); ++ float l = lsum[h] * f; ++#pragma clang loop unroll(full) ++ for (int c = 0; c < QCH; c += 4) { // in order, four loads at a time ++ const float4 v = *(threadgroup const float4*)(sp[0] + h * QCH + c); ++ l += v.x; l += v.y; l += v.z; l += v.w; + } +- mrun[lane] = m; +- lsum[lane] = l; +- fac[lane] = f; +- } else { +- for (int c = 0; c < QCH; ++c) p[lane * QCH + c] = half(0.0f); +- fac[lane] = 0.0f; ++ mrun[h] = m; ++ lsum[h] = l; ++ fac[h] = f; + } ++ } else { ++ for (int c = hlane; c < QCH; c += 8) p[h * QCH + c] = half(0.0f); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // rescale this simdgroup's accumulator rows, then acc += P V over its 64 dims + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; +- for (int c0 = 0; c0 < QCH; c0 += 32) { // K = 32 cells a run (K = 64 is wrong on the M4) ++ for (int c0 = 0; c0 < nc; c0 += 32) { // K = 32 cells a run (K = 64 is wrong on the M4) ++ if (!coop_k4) { ++#pragma clang loop unroll(full) ++ for (uint16_t i = 0; i < rv.get_capacity(); ++i) { ++ const auto ix = rv.get_multidimensional_index(i); // (cell, dim) ++ const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; ++ rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vs_tg[c][sg])) : half(0.0f); ++ } ++ } else { ++ // a lane holds the same eight cells (two runs of four, 16 apart) at each of its dims: as for rk, row ++ // pointers and scales once a run, unconditional loads ++ uint vo[8]; ++ float vs8[8]; ++ const int cb = c0 + rv.get_multidimensional_index(uint16_t(0))[0]; ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 8; ++j) { ++ const int c = cb + (j & 3) + 16 * (j >> 2); ++ vo[j] = uint(cid[c]) + uint(d0); ++ vs8[j] = float(vs_tg[c][sg]); ++ } ++#pragma clang loop unroll(full) ++ for (uint16_t i = 0; i < rv.get_capacity(); i += 4) { ++ const int d = rv.get_multidimensional_index(i)[1], r = i < rv.get_capacity() / 2 ? 0 : 4; + #pragma clang loop unroll(full) +- for (uint16_t i = 0; i < rv.get_capacity(); ++i) { +- const auto ix = rv.get_multidimensional_index(i); // (cell, dim) +- const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; +- rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vsc[(uint(id) * QKV + g) * 4 + sg])) : half(0.0f); ++ for (int j = 0; j < 4; ++j) rv[i + j] = half(float(vq[vo[r + j] + uint(d)]) * vs8[r + j]); ++ } + } + tPh = tP.template slice<32, 16>(c0, 0); + opO.run(tPh, rv, acc); +@@ -327,14 +426,28 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq + if (acc.is_valid_element(i)) tot[i] = fma(tot[i], fac[acc.get_multidimensional_index(i)[1]], acc[i]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); +- for (uint16_t i = 0; i < acc.get_capacity(); ++i) +- if (acc.is_valid_element(i)) { +- const auto ix = acc.get_multidimensional_index(i); // (dim, row) ++ if (!coop_k4) { ++ for (uint16_t i = 0; i < acc.get_capacity(); ++i) ++ if (acc.is_valid_element(i)) { ++ const auto ix = acc.get_multidimensional_index(i); // (dim, row) ++ const int h = ix[1], d = d0 + ix[0]; ++ if (h < HPG) { ++ const int hh = g * HPG + h; ++ const float gate = qfull[(uint(t) * QH + hh) * 2 * QD + QD + d]; ++ xqh[(uint(t) * QH + hh) * QD + d] = half(tot[i] / lsum[h] * (1.0f / (1.0f + exp(-gate)))); ++ } ++ } ++ } else ++#pragma clang loop unroll(full) ++ for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { // runs of four dims of one row ++ const auto ix = acc.get_multidimensional_index(i); + const int h = ix[1], d = d0 + ix[0]; + if (h < HPG) { + const int hh = g * HPG + h; +- const float gate = qfull[(uint(t) * QH + hh) * 2 * QD + QD + d]; +- xqh[(uint(t) * QH + hh) * QD + d] = half(tot[i] / lsum[h] * (1.0f / (1.0f + exp(-gate)))); ++ const float l = lsum[h]; ++ const float4 gate = *(device const float4*)(qfull + (uint(t) * QH + hh) * 2 * QD + QD + d); ++ const float4 v = float4(tot[i], tot[i + 1], tot[i + 2], tot[i + 3]) / l * (1.0f / (1.0f + exp(-gate))); ++ *(device half4*)(xqh + (uint(t) * QH + hh) * QD + d) = half4(v); + } + } + } +diff --git a/metal/kernels/quant.metal b/metal/kernels/quant.metal +index 51fecca7..fe667be5 100644 +--- a/metal/kernels/quant.metal ++++ b/metal/kernels/quant.metal +@@ -34,11 +34,11 @@ template struct fmt; + // QK block elements, BB block bytes, UNIT weights per unit, TAB words of threadgroup table + template <> struct fmt { enum { QK = 64, BB = 18, UNIT = 16, TAB = 1 }; }; + template <> struct fmt { enum { QK = 64, BB = 18, UNIT = 16, TAB = 1 }; }; +-template <> struct fmt { enum { QK = 256, BB = 98, UNIT = 16, TAB = 256 + 32 }; }; +-template <> struct fmt { enum { QK = 256, BB = 66, UNIT = 16, TAB = 512 + 32 }; }; +-template <> struct fmt { enum { QK = 256, BB = 74, UNIT = 16, TAB = 1024 + 32 }; }; +-template <> struct fmt { enum { QK = 256, BB = 82, UNIT = 16, TAB = 2048 + 32 }; }; +-template <> struct fmt { enum { QK = 256, BB = 110, UNIT = 16, TAB = 512 + 32 }; }; ++template <> struct fmt { enum { QK = 256, BB = 98, UNIT = 16, TAB = 256 }; }; ++template <> struct fmt { enum { QK = 256, BB = 66, UNIT = 16, TAB = 512 }; }; ++template <> struct fmt { enum { QK = 256, BB = 74, UNIT = 16, TAB = 1024 }; }; ++template <> struct fmt { enum { QK = 256, BB = 82, UNIT = 16, TAB = 2048 }; }; ++template <> struct fmt { enum { QK = 256, BB = 110, UNIT = 16, TAB = 512 }; }; + template <> struct fmt { enum { QK = 32, BB = 18, UNIT = 32, TAB = 256 }; }; + template <> struct fmt { enum { QK = 256, BB = 136, UNIT = 32, TAB = 256 }; }; + template <> struct fmt { enum { QK = 256, BB = 110, UNIT = 32, TAB = 1 }; }; +@@ -59,15 +59,16 @@ static inline uint rd32(device const uchar* p) { + static inline float rdh(device const uchar* p) { return float(*(device const half*)p); } + + // ---- threadgroup tables ------------------------------------------------------------------------------------------- +-// layout in `tab` (uint words): [0, G) the grid, [G, G + 32) ksigns_iq2xs as bytes; IQ4_NL/XS: 16 floats. ++// layout in `tab` (uint words): [0, G) the grid; IQ4_NL/XS: 256 half2. + template static inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth); + +-// QUANT_CTAB (function constant 23): the i-quant grids and the sign table read from ggml-common's constant arrays ++// QUANT_CTAB (function constant 23): the i-quant grids read from ggml-common's constant arrays + // instead of a threadgroup copy (no table load and no barrier wait at the threadgroup's start); the same values. + constant bool QUANT_CTAB [[function_constant(23)]]; + constant bool ctab = is_function_constant_defined(QUANT_CTAB) && QUANT_CTAB; +-static inline uint ksg(threadgroup const uint* tab, int at, uint i) { +- return ctab ? uint(ksigns_iq2xs[i]) : uint(((threadgroup const uchar*)(tab + at))[i]); ++// ksigns_iq2xs[j] = j | parity(j) << 7 (an even count of minus signs): the low 4 bits of run hh (0 or 1). ++static inline uint sign4_iq2xs(uint j, int hh) { ++ return hh ? ((j >> 4) | ((popcount(j) & 1u) << 3)) : j; + } + static inline uint g_iq3xxs(threadgroup const uint* tab, uint i) { return ctab ? iq3xxs_grid[i] : tab[i]; } + static inline uint g_iq3s(threadgroup const uint* tab, uint i) { return ctab ? iq3s_grid[i] : tab[i]; } +@@ -88,34 +89,25 @@ template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} + template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} + template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} + +-static inline void load_signs(threadgroup uint* tab, int at, ushort tid, ushort nth) { +- threadgroup uchar* s = (threadgroup uchar*)(tab + at); +- for (int i = tid; i < 128; i += nth) s[i] = ksigns_iq2xs[i]; +-} + template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 256; i += nth) tab[i] = iq3xxs_grid[i]; +- load_signs(tab, 256, tid, nth); + } + template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 512; i += nth) tab[i] = iq3s_grid[i]; +- load_signs(tab, 512, tid, nth); + } + template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 256; i += nth) { tab[2 * i] = uint(iq2xxs_grid[i]); tab[2 * i + 1] = uint(iq2xxs_grid[i] >> 32); } +- load_signs(tab, 512, tid, nth); + } + template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 512; i += nth) { tab[2 * i] = uint(iq2xs_grid[i]); tab[2 * i + 1] = uint(iq2xs_grid[i] >> 32); } +- load_signs(tab, 1024, tid, nth); + } + template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 1024; i += nth) { tab[2 * i] = uint(iq2s_grid[i]); tab[2 * i + 1] = uint(iq2s_grid[i] >> 32); } +- load_signs(tab, 2048, tid, nth); + } + // IQ4: a byte -> (codebook[low nibble], codebook[high nibble]) as half2, so one lookup gives two weights + template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { +@@ -153,74 +145,82 @@ static inline void put_signed4(thread float* q, uint g, uint s) { + // IQ3_XXS: 256 = 8 x 32; unit = half of a 32 (two grid pairs). + template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; +- const uint aux = rd32(b + 2 + 64 + 4 * ib); +- device const uchar* qs = b + 2 + 8 * ib + 4 * h; ++ const uint aux = rd32(b + 2 + 64 + 4 * ib), s0 = aux >> (14 * h); ++ device const uchar* qs = b + 2 + 4 * ui; ++ #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { +- const int l = 2 * h + l2; +- const uint s = ksg(tab, 256, (aux >> (7 * l)) & 127); ++ const uint s = (s0 >> (7 * l2)) & 127; + put_signed4(q + 8 * l2, g_iq3xxs(tab, qs[2 * l2]), s); +- put_signed4(q + 8 * l2 + 4, g_iq3xxs(tab, qs[2 * l2 + 1]), s >> 4); ++ put_signed4(q + 8 * l2 + 4, g_iq3xxs(tab, qs[2 * l2 + 1]), sign4_iq2xs(s, 1)); + } +- return rdh(b) * (0.5f + float(aux >> 28)) * 0.5f; ++ return rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f); + } + + // IQ2_XXS: per 32: two u32 (4 grid indices, then 4 x 7 sign bits + 4-bit scale). + template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; +- const uint a0 = rd32(b + 2 + 8 * ib), a1 = rd32(b + 2 + 8 * ib + 4); ++ device const uchar* qs = b + 2 + 4 * ib + 2 * ui; ++ const uint a1 = rd32(b + 2 + 8 * ib + 4), s0 = a1 >> (14 * h); ++ #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { +- const int l = 2 * h + l2; +- const uint idx = (a0 >> (8 * l)) & 255; +- const uint s = ksg(tab, 512, (a1 >> (7 * l)) & 127); +- put_signed4(q + 8 * l2, g_iq2xxs(tab, idx, 0), s); +- put_signed4(q + 8 * l2 + 4, g_iq2xxs(tab, idx, 1), s >> 4); ++ const uint idx = qs[l2], s = (s0 >> (7 * l2)) & 127; ++ const uint2 g = ctab ? as_type(iq2xxs_grid[idx]) : ((threadgroup const uint2*)tab)[idx]; ++ put_signed4(q + 8 * l2, g.x, s); ++ put_signed4(q + 8 * l2 + 4, g.y, sign4_iq2xs(s, 1)); + } +- return rdh(b) * (0.5f + float(a1 >> 28)) * 0.25f; ++ return rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f); + } + + // IQ2_XS: uint16 qs[32] (9-bit grid index + 7-bit sign index), scales[8] nibbles. + template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; +- device const ushort* qs = (device const ushort*)(b + 2); ++ device const ushort* qs = (device const ushort*)(b + 2) + 2 * ui; ++ #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { +- const uint v = qs[4 * ib + 2 * h + l2]; +- const uint idx = v & 511, s = ksg(tab, 1024, v >> 9); +- put_signed4(q + 8 * l2, g_iq2xs(tab, idx, 0), s); +- put_signed4(q + 8 * l2 + 4, g_iq2xs(tab, idx, 1), s >> 4); ++ const uint v = qs[l2]; ++ const uint idx = v & 511, s = v >> 9; ++ const uint2 g = ctab ? as_type(iq2xs_grid[idx]) : ((threadgroup const uint2*)tab)[idx]; ++ put_signed4(q + 8 * l2, g.x, s); ++ put_signed4(q + 8 * l2 + 4, g.y, sign4_iq2xs(s, 1)); + } + const uint sc = b[66 + ib]; +- return rdh(b) * (0.5f + float((sc >> (4 * h)) & 15)) * 0.25f; ++ return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); + } + + // IQ2_S: qs[32] low 8 bits of the 10-bit index, signs[32], qh[8] (2 high bits per group of 8), scales[8]. + template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; +- const uint qh = b[66 + ib]; ++ const uint q2 = *(device const ushort*)(b + 2 + 2 * ui); ++ const uint s2 = *(device const ushort*)(b + 34 + 2 * ui); ++ const uint qh = uint(b[66 + ib]) >> (4 * h); ++ #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { +- const int l = 2 * h + l2; +- const uint idx = uint(b[2 + 4 * ib + l]) | ((qh << (8 - 2 * l)) & 0x300); +- const uint s = b[2 + 32 + 4 * ib + l]; +- put_signed4(q + 8 * l2, g_iq2s(tab, idx, 0), s); +- put_signed4(q + 8 * l2 + 4, g_iq2s(tab, idx, 1), s >> 4); ++ const uint idx = ((q2 >> (8 * l2)) & 255) | ((qh << (8 - 2 * l2)) & 0x300); ++ const uint s = (s2 >> (8 * l2)) & 255; ++ const uint2 g = ctab ? as_type(iq2s_grid[idx]) : ((threadgroup const uint2*)tab)[idx]; ++ put_signed4(q + 8 * l2, g.x, s); ++ put_signed4(q + 8 * l2 + 4, g.y, s >> 4); + } + const uint sc = b[74 + ib]; +- return rdh(b) * (0.5f + float((sc >> (4 * h)) & 15)) * 0.25f; ++ return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f); + } + + // IQ3_S: qs[64] low 8 bits of 9-bit indices, qh[8], signs[32], scales[4] (nibble per 32). + template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; ++ const uint s2 = *(device const ushort*)(b + 74 + 2 * ui); + const uint qh = b[66 + ib]; ++ #pragma clang loop unroll(full) + for (int l2 = 0; l2 < 2; ++l2) { + const int l = 2 * h + l2; + const uint i1 = uint(b[2 + 8 * ib + 2 * l]) | ((qh << (8 - 2 * l)) & 256); + const uint i2 = uint(b[2 + 8 * ib + 2 * l + 1]) | ((qh << (7 - 2 * l)) & 256); +- const uint s = b[74 + 4 * ib + l]; ++ const uint s = (s2 >> (8 * l2)) & 255; + put_signed4(q + 8 * l2, g_iq3s(tab, i1), s); + put_signed4(q + 8 * l2 + 4, g_iq3s(tab, i2), s >> 4); + } + const uint sc = b[106 + (ib >> 1)]; +- return rdh(b) * float(1 + 2 * ((sc >> (4 * (ib & 1))) & 15)); ++ return rdh(b) * fma(float((sc >> (4 * (ib & 1))) & 15), 2.0f, 1.0f); + } + + // IQ4_NL: 32 = 1 unit; low nibbles are elements 0..15, high 16..31 (one half2 lookup per byte). +@@ -238,6 +238,11 @@ template <> inline float deq(device const uchar* b, int, thread float* + } + + // IQ4_XS: 256 = 8 units of 32; 6-bit scale from scales_l nibble + 2 bits of scales_h, minus 32. ++static inline float iq4xs_scale(device const uchar* b, int ui) { ++ const uint sh = *(device const ushort*)(b + 2); ++ const int ls = int((b[4 + (ui >> 1)] >> (4 * (ui & 1))) & 15) | int(((sh >> (2 * ui)) & 3) << 4); ++ return rdh(b) * float(ls - 32); ++} + template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + device const ushort* p = (device const ushort*)(b + 8 + 16 * ui); + for (int j = 0; j < 8; ++j) { +@@ -248,9 +253,7 @@ template <> inline float deq(device const uchar* b, int ui, thread flo + q[2 * j + 1] = float(w1.x); + q[2 * j + 17] = float(w1.y); + } +- const uint sh = *(device const ushort*)(b + 2); +- const int ls = int((b[4 + (ui >> 1)] >> (4 * (ui & 1))) & 15) | int(((sh >> (2 * ui)) & 3) << 4); +- return rdh(b) * float(ls - 32); ++ return iq4xs_scale(b, ui); + } + + template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { +diff --git a/metal/kernels/router.metal b/metal/kernels/router.metal +index e5c75274..5c6bf590 100644 +--- a/metal/kernels/router.metal ++++ b/metal/kernels/router.metal +@@ -5,44 +5,108 @@ + // the counter. No threadgroup waits for another. ids / w / plan are bitwise moe_route's. + // Buffers: W, xp, s16, logits, args (GemvArgs: n_out = n_expert), ids, w, plan, MoeArgs, counter (one uint, zero). + ++#ifndef STRATA_MOCE_FC_METAL ++#define STRATA_MOCE_FC_METAL ++constant int kMoceProtectFc [[function_constant(24)]]; ++constant int kMoceProtect = is_function_constant_defined(kMoceProtectFc) ? kMoceProtectFc : 3; ++constant uint kMoceScaleR1BitsFc [[function_constant(25)]]; ++constant float kMoceScaleR1 = is_function_constant_defined(kMoceScaleR1BitsFc) ? as_type(kMoceScaleR1BitsFc) : 1.0f; ++#endif + static inline void route_body(device const float* logits, device int* ids, device float* w, device int* plan, +- constant MoeArgs& a, ushort tid, ushort nth, ushort sg, ushort lane, threadgroup int* s_grp, +- threadgroup int* s_ent) { ++ constant MoeArgs& a, device const uint* smap, ushort tid, ushort nth, ushort sg, ushort lane, ++ threadgroup int* s_grp, threadgroup int* s_ent) { + if (sg < a.n_tok) { + device const float* l = logits + sg * 512; + float v[16]; +- for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i]; ++ for (int i = 0; i < 16; ++i) { ++ const int idx = lane + 32 * i; ++ v[i] = idx < int(a.n_expert) ? l[idx] : -INFINITY; ++ } + float mx = -INFINITY; + for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + float sum = 0.0f; + for (int i = 0; i < 16; ++i) { +- v[i] = exp(v[i] - mx); ++ v[i] = (lane + 32 * i < int(a.n_expert)) ? exp(v[i] - mx) : 0.0f; + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + for (int i = 0; i < 16; ++i) { +- v[i] *= rcp; +- if (isnan(v[i])) v[i] = -FLT_MAX; ++ if (lane + 32 * i >= int(a.n_expert)) { ++ v[i] = -INFINITY; ++ } else { ++ v[i] *= rcp; ++ if (isnan(v[i])) v[i] = -FLT_MAX; ++ } + } ++ const bool draft_moce = (a.pad & 0x80000000u) != 0u && sg > 0; ++ const float base_moce = (a.pad & 0x7FFFFFFFu) != 0u ? as_type(a.pad & 0x7FFFFFFFu) : 1.0f; ++ const float moce_scale = draft_moce ? (base_moce * base_moce) : base_moce; ++ const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f; ++ const int r_prot = draft_moce ? 1 : kMoceProtect; ++ const float r1_scale = kMoceScaleR1; ++ const float rcp_r1 = r1_scale > 1.0f ? (1.0f / r1_scale) : 1.0f; ++ const float r_step_scale = (r1_scale > 1.0f && moce_scale > 1.0f) ? (moce_scale * rcp_r1) : moce_scale; ++ ushort scaled_mask = 0; + float selected = 0.0f, ssum = 0.0f; + for (int rank = 0; rank < 10; ++rank) { ++ if (r1_scale > 1.0f && r_prot > 1) { ++ if (rank == 1 && smap != nullptr) { ++ for (int i = 0; i < 16; ++i) { ++ const int idx = lane + 32 * i; ++ if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) { ++ v[i] *= r1_scale; ++ scaled_mask |= ushort(1u << i); ++ } ++ } ++ } else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) { ++ for (int i = 0; i < 16; ++i) { ++ if ((scaled_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) { ++ v[i] *= r_step_scale; ++ } ++ } ++ } ++ } else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) { ++ for (int i = 0; i < 16; ++i) { ++ const int idx = lane + 32 * i; ++ if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) { ++ v[i] *= moce_scale; ++ scaled_mask |= ushort(1u << i); ++ } ++ } ++ } + float best = v[0]; + int ex = lane; ++ #pragma clang loop unroll(full) + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } ++ #pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } ++ float true_prob = best; ++ const int min_scaled_rank = (r1_scale > 1.0f && r_prot > 1) ? 1 : r_prot; ++ if (rank >= min_scaled_rank && (moce_scale > 1.0f || r1_scale > 1.0f) && best > 0.0f) { ++ const ushort win_lane = ushort(ex & 31); ++ 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) { ++ // 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); ++ } ++ } ++ #pragma clang loop unroll(full) ++ for (int i = 0; i < 16; ++i) v[i] = (ex == lane + 32 * i) ? -INFINITY : v[i]; + if ((ex & 31) == lane) { +- v[ex / 32] = -INFINITY; + ids[sg * 10 + rank] = ex; +- ssum += best; ++ ssum += true_prob; + } +- if (rank == lane) selected = best; ++ if (rank == lane) selected = true_prob; + } + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + ssum = max(ssum, 6.103515625e-5f); +@@ -50,7 +114,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); +- plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); ++ plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent, smap); + } + + template +@@ -58,13 +122,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)]], +- device atomic_uint* counter [[buffer(9)]], ++ device atomic_uint* counter [[buffer(9)]], device const uint* smap [[buffer(10)]], + uint tgx [[threadgroup_position_in_grid]], uint ngx [[threadgroups_per_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort nth [[threads_per_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int UNIT = fmt::UNIT; + threadgroup uint tab[fmt::TAB]; +- threadgroup int s_grp[32], s_ent[32]; ++ threadgroup int s_grp[64], s_ent[64]; + threadgroup uint last; + load_tab(tab, tid, nth); + threadgroup_barrier(mem_flags::mem_threadgroup); +@@ -74,17 +138,20 @@ kernel void router_route(device const uchar* W [[buffer(0)]], device const float + device const uchar* rp = W + (ulong)rr * ga.row_bytes; + const int U = int(ga.n_in) / UNIT, s16n = int(ga.n_in) / 16; + float acc[8]; ++#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) acc[t] = 0.0f; + for (int u = lane; u < U; u += 32) { + float q[UNIT]; + float4 sc; + Dec::run(rp, u, q, sc, tab); ++#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(ga.x_stride), s16 + t * s16n, u, xs, xsum); + acc[t] += unit_combine(q, sc, xs, xsum); + } + } ++#pragma clang loop unroll(full) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum<32>(acc[t]); + if (lane == 0 && row < int(ga.n_out)) logits[t * int(ga.y_stride) + row] = v; +@@ -96,9 +163,11 @@ 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); +- route_body(logits, ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); ++ route_body(logits, ids, w, plan, a, smap, tid, nth, sg, lane, s_grp, s_ent); + if (tid == 0) atomic_store_explicit(counter, 0u, memory_order_relaxed); + } + + template [[host_name("router_route_30")]] kernel decltype(router_route) router_route; + template [[host_name("router_route_0")]] kernel decltype(router_route) router_route; ++ ++template [[host_name("router_route_8")]] kernel decltype(router_route) router_route; +diff --git a/metal/kernels/sample.metal b/metal/kernels/sample.metal +index 8443d360..4531ca2a 100644 +--- a/metal/kernels/sample.metal ++++ b/metal/kernels/sample.metal +@@ -13,6 +13,7 @@ static inline void better(thread float& bv, thread int& bi, float v, int i) { + + static inline void tg_argmax(thread float& bv, thread int& bi, threadgroup float* sv, threadgroup int* si, ushort sg, + ushort lane) { ++#pragma clang loop unroll(full) + for (ushort m = 16; m > 0; m >>= 1) { + const float ov = simd_shuffle_xor(bv, m); + const int oi = simd_shuffle_xor(bi, m); +@@ -22,6 +23,7 @@ static inline void tg_argmax(thread float& bv, thread int& bi, threadgroup float + threadgroup_barrier(mem_flags::mem_threadgroup); + bv = -INFINITY; + bi = 0x7fffffff; ++#pragma clang loop unroll(full) + for (int k = 0; k < 8; ++k) better(bv, bi, sv[k], si[k]); + } + +diff --git a/metal/kernels/units.metal b/metal/kernels/units.metal +index 15f584cb..25d0761a 100644 +--- a/metal/kernels/units.metal ++++ b/metal/kernels/units.metal +@@ -107,26 +107,33 @@ template <> struct X { + } + }; + +-// xprep: one thread per 16 consecutive values of one token row; XP = [4 kinds][n_tok][n] then S16 [n_tok][n / 16] ++// xprep: one thread per 16 consecutive values of one token row; XP = [4 kinds][n_tok][n] then S16 [n_tok][n / 16]. ++// raw and plane are 16-float aligned: the four planes are written as float4s (mtp.metal's mtp_prepare, bit for bit) + kernel void xprep(device const float* x [[buffer(0)]], device float* xp [[buffer(1)]], constant uint& n [[buffer(2)]], + constant uint& n_tok [[buffer(3)]], uint gid [[thread_position_in_grid]]) { + const uint per = n / 16; + if (gid >= per * n_tok) return; + const uint t = gid / per, i = gid % per; +- device const float* src = x + t * n + 16 * i; ++ device const float4* src = (device const float4*)(x + t * n + 16 * i); + const uint plane = n_tok * n; + device float* raw = xp + t * n + 16 * i; +- float v[16]; +- load16(src, v); ++ device float4* r0 = (device float4*)raw; ++ device float4* r1 = (device float4*)(raw + plane); ++ device float4* r2 = (device float4*)(raw + 2 * plane); ++ device float4* r3 = (device float4*)(raw + 3 * plane); ++ const float4 q2lo = float4(1.0f, 0.25f, 0.0625f, 0.015625f), ++ q2hi = float4(0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f), ++ p4 = float4(1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f); + float s0 = 0, s1 = 0; +- for (int j = 0; j < 16; j += 2) { s0 += v[j]; s1 += v[j + 1]; } +- const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; +- const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; +- for (int j = 0; j < 16; ++j) { +- raw[j] = v[j]; +- raw[plane + j] = v[j] * q2s[j & 7]; +- raw[2 * plane + j] = v[j] * p4s[j & 3]; +- raw[3 * plane + j] = (j & 1) ? v[j] * 0.00390625f : v[j]; ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 4; ++j) { ++ const float4 v = src[j]; ++ s0 += v.x; s0 += v.z; ++ s1 += v.y; s1 += v.w; ++ r0[j] = v; ++ r1[j] = v * ((j & 1) ? q2hi : q2lo); ++ r2[j] = v * p4; ++ r3[j] = float4(v.x, v.y * 0.00390625f, v.z, v.w * 0.00390625f); + } + xp[4 * plane + t * per + i] = s0 + s1; + } +diff --git a/metal/src/engine/arena_budget.h b/metal/src/engine/arena_budget.h +index 4cc024b7..513a944d 100644 +--- a/metal/src/engine/arena_budget.h ++++ b/metal/src/engine/arena_budget.h +@@ -44,6 +44,9 @@ struct ArenaMem { + double keep_gb = -1; // STRATA_METAL_KEEP_FREE_GB (< 0: unset) + double arena_gb = -1; // STRATA_METAL_ARENA_GB (< 0: unset) + double leave_share = -1; // STRATA_METAL_LEAVE_SHARE: the share of the RAM a big Mac leaves to macOS (< 0: kLeaveShare) ++ double dense_file = 0; // the dense weights the GPU reads straight from the GGUF's map and that are NOT wired: they are ++ // clean file cache, so `available` counts them as free - an arena that takes them makes ++ // every command buffer page them back in (STRATA_METAL_MLOCK_DENSE=0) + }; + + constexpr double kGB = 1e9; +@@ -112,7 +115,7 @@ inline double arena_ple_reserve(const ArenaMem& m) { return m.ple_mapped > 0 ? m + inline double arena_budget(const ArenaMem& m) { + if (m.arena_gb >= 0) return m.arena_gb * kGB; + const double gpu = m.working_set - footprint_base(m) - 1.0 * kGB; +- const double ram = m.available - arena_keep(m) - arena_ple_reserve(m); ++ const double ram = m.available - arena_keep(m) - arena_ple_reserve(m) - m.dense_file; + const double foot = footprint_cap(m) - footprint_base(m) - arena_ple_reserve(m); // the whole footprint under the cap + return std::max(1.0 * kGB, std::min(std::min(gpu, ram), foot)); + } +@@ -154,7 +157,7 @@ inline double arena_shrink(const ArenaMem& m, double planned, double available_n + /// (the explicit STRATA_METAL_ARENA_GB always fills) + inline bool arena_grow(const ArenaMem& m, double available_now, double step) { + if (m.arena_gb >= 0) return true; +- return available_now - step >= arena_keep(m) + arena_ple_reserve(m); ++ return available_now - step >= arena_keep(m) + arena_ple_reserve(m) + m.dense_file; + } + + /// The fill as a plan, for the test: the pairs' bytes in rank order, steps of `step` bytes, `avail(filled)` the free +diff --git a/metal/src/engine/msl/mtp.metal b/metal/src/engine/msl/mtp.metal +index 7159330f..051dcf72 100644 +--- a/metal/src/engine/msl/mtp.metal ++++ b/metal/src/engine/msl/mtp.metal +@@ -1,68 +1,179 @@ + // metal/src/engine/msl/mtp.metal - the small steps of the MTP draft layer's input branches (include/strata/core/mtp.hpp) + // around mac-kern's kernels (GEMVs, the hyper-connection read, QSA, MoE, the head): + // +-// mtp_rms y[r] = x[r] * rsqrt(mean(x[r]^2) + eps) * w rows of width n (w: (1 + w) folded by the host) +-// mtp_xprep xprep (units.metal) of rows read with a stride: the prepared activation of one stream of rows +-// mtp_combine R[t][c][d] = h[c][t][d] + e[t][d] the streams of fc_hidden plus fc_embedding ++// mtp_rms_xprep weighted RMSNorm into prepared embedding / hidden rows ++// mtp_xprep strided prepared rows for the hidden streams past the first MAXT ++// mtp_combine_gnorm R[t][c][d] = h[t][c][d] + e[t][d], then the attention hyper-connection norm + // + // Prepared activations: [4 kinds][n_tok][n] then S16 [n_tok][n / 16] (units.metal's xprep layout). + + #include + using namespace metal; + +-struct RmsArgs { uint n; float eps; }; ++struct XpArgs { uint n; uint n_tok; uint plane_tok; uint stride; }; + +-// grid (rows), 256 threads +-kernel void mtp_rms(device const float* x [[buffer(0)]], device const float* w [[buffer(1)]], device float* y [[buffer(2)]], +- constant RmsArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], +- ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], +- ushort lane [[thread_index_in_simdgroup]]) { +- threadgroup float part[8]; +- device const float* xr = x + tg.x * a.n; +- float ss = 0.0f; +- for (uint d = tid; d < a.n; d += 256) ss = fma(xr[d], xr[d], ss); +- ss = simd_sum(ss); +- if (lane == 0) part[sg] = ss; +- threadgroup_barrier(mem_flags::mem_threadgroup); +- float tot = 0.0f; +- for (int i = 0; i < 8; ++i) tot += part[i]; +- const float rs = rsqrt(tot / float(a.n) + a.eps); +- device float* yr = y + tg.x * a.n; +- for (uint d = tid; d < a.n; d += 256) yr[d] = xr[d] * rs * w[d]; ++// raw and plane are 16-float aligned: the four planes are written as float4s ++static inline void mtp_prepare(thread const float* v, device float* raw, device float* sum, uint plane) { ++ float s0 = 0, s1 = 0; ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 16; j += 2) { s0 += v[j]; s1 += v[j + 1]; } ++ const float4 q2lo = float4(1.0f, 0.25f, 0.0625f, 0.015625f), ++ q2hi = float4(0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f), ++ p4 = float4(1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f); ++ device float4* r0 = (device float4*) raw; ++ device float4* r1 = (device float4*) (raw + plane); ++ device float4* r2 = (device float4*) (raw + 2 * plane); ++ device float4* r3 = (device float4*) (raw + 3 * plane); ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 4; ++j) { ++ const float4 x = float4(v[4 * j], v[4 * j + 1], v[4 * j + 2], v[4 * j + 3]); ++ r0[j] = x; ++ r1[j] = x * ((j & 1) ? q2hi : q2lo); ++ r2[j] = x * p4; ++ r3[j] = float4(x.x, x.y * 0.00390625f, x.z, x.w * 0.00390625f); ++ } ++ *sum = s0 + s1; + } + +-struct XpArgs { uint n; uint n_tok; uint plane_tok; uint stride; }; +- + // one thread per 16 values of one row: row t at x + t * stride; out planes of plane_tok rows + kernel void mtp_xprep(device const float* x [[buffer(0)]], device float* xp [[buffer(1)]], constant XpArgs& a [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + const uint per = a.n / 16; + if (gid >= per * a.n_tok) return; +- const uint t = gid / per, i = gid % per; ++ const uint t = per == 160u ? gid / 160u : gid / per, i = gid - t * per; + device const float* src = x + t * a.stride + 16 * i; + const uint plane = a.plane_tok * a.n; + device float* raw = xp + t * a.n + 16 * i; + float v[16]; +- for (int j = 0; j < 16; ++j) v[j] = src[j]; +- float s0 = 0, s1 = 0; +- for (int j = 0; j < 16; j += 2) { s0 += v[j]; s1 += v[j + 1]; } +- const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; +- const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; +- for (int j = 0; j < 16; ++j) { +- raw[j] = v[j]; +- raw[plane + j] = v[j] * q2s[j & 7]; +- raw[2 * plane + j] = v[j] * p4s[j & 3]; +- raw[3 * plane + j] = (j & 1) ? v[j] * 0.00390625f : v[j]; ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 4; ++j) { ++ const float4 s = ((device const float4*) src)[j]; ++ v[4 * j] = s.x; v[4 * j + 1] = s.y; v[4 * j + 2] = s.z; v[4 * j + 3] = s.w; ++ } ++ mtp_prepare(v, raw, xp + 4 * plane + t * per + i, plane); ++} ++ ++struct RmsXpArgs { uint n; uint xp_n; uint plane_tok; uint split; uint keep; float eps; }; ++ ++// A wide hidden norm splits into its n / xp_n streams as prepared rows [t][c] (the first plane_tok of them); keep ++// retains every normalized row for mtp_xprep. Unsplit: the first xp_n values of each row. ++kernel void mtp_rms_xprep(device const float* x [[buffer(0)]], device const float* w [[buffer(1)]], ++ device float* y [[buffer(2)]], device float* xp [[buffer(3)]], ++ constant RmsXpArgs& a [[buffer(4)]], uint3 tg [[threadgroup_position_in_grid]], ++ ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], ++ ushort lane [[thread_index_in_simdgroup]]) { ++ threadgroup float part[8]; ++ device const float* xr = x + tg.x * a.n; ++ float ss = 0.0f; ++ if (a.n == 2560) { ++ device const float* xrp = xr + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < 10; ++i) ss = fma(xrp[256 * i], xrp[256 * i], ss); ++ } else if (a.n == 10240) { ++ device const float* xrp = xr + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < 40; ++i) ss = fma(xrp[256 * i], xrp[256 * i], ss); ++ } else { ++ for (uint d = tid; d < a.n; d += 256) ss = fma(xr[d], xr[d], ss); ++ } ++ ss = simd_sum(ss); ++ if (lane == 0) part[sg] = ss; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ float tot = 0.0f; ++#pragma clang loop unroll(full) ++ for (int i = 0; i < 8; ++i) tot += part[i]; ++ const float rs = rsqrt(tot / float(a.n) + a.eps); ++ const uint per = a.xp_n / 16, plane = a.plane_tok * a.xp_n; ++ auto scale16 = [&](uint i, uint n_row, thread float* v) { ++ device const float4* x4 = (device const float4*) (xr + 16 * i); ++ device const float4* w4 = (device const float4*) (w + 16 * i); ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 4; ++j) { ++ const float4 s = x4[j] * rs * w4[j]; ++ v[4 * j] = s.x; v[4 * j + 1] = s.y; v[4 * j + 2] = s.z; v[4 * j + 3] = s.w; ++ if (a.keep) ((device float4*) (y + tg.x * n_row + 16 * i))[j] = s; ++ } ++ }; ++ if (a.n == 2560 && a.xp_n == 2560) { ++ if (tid < 160) { ++ float v[16]; ++ scale16(tid, 2560, v); ++ if (!a.split || tg.x < a.plane_tok) ++ mtp_prepare(v, xp + tg.x * 2560 + 16 * tid, xp + 4 * plane + tg.x * 160 + tid, plane); ++ } ++ } else if (a.n == 10240 && a.xp_n == 2560) { ++ for (uint i = tid; i < 640; i += 256) { ++ float v[16]; ++ scale16(i, 10240, v); ++ const uint q = i / 160u, k = i - q * 160u; ++ const uint t = a.split ? tg.x * 4u + q : tg.x; ++ if (a.split ? t >= a.plane_tok : q > 0u) continue; ++ mtp_prepare(v, xp + t * 2560 + 16 * k, xp + 4 * plane + t * 160 + k, plane); ++ } ++ } else { ++ for (uint i = tid; i < a.n / 16; i += 256) { ++ float v[16]; ++ scale16(i, a.n, v); ++ const uint t = a.split ? tg.x * (a.n / a.xp_n) + i / per : tg.x; ++ if (a.split ? t >= a.plane_tok : i >= per) continue; ++ const uint k = i % per; ++ mtp_prepare(v, xp + t * a.xp_n + 16 * k, xp + 4 * plane + t * per + k, plane); ++ } + } +- xp[4 * plane + t * per + i] = s0 + s1; + } + +-struct CombArgs { uint n; uint hc; uint n_tok; uint h_stride; }; // h: stream c's rows at h + c * h_stride ++struct CombNormArgs { uint n; uint hc; float eps; }; + +-kernel void mtp_combine(device const float* h [[buffer(0)]], device const float* e [[buffer(1)]], device float* R [[buffer(2)]], +- constant CombArgs& a [[buffer(3)]], uint gid [[thread_position_in_grid]]) { +- const uint per = a.hc * a.n; +- if (gid >= a.n_tok * per) return; +- const uint t = gid / per, c = (gid % per) / a.n, d = gid % a.n; +- R[gid] = h[c * a.h_stride + t * a.n + d] + e[t * a.n + d]; ++// Same reduction and multiply order as gr_norm(apply=false); keep R for subsequent residual writes. h, R and xn are ++// [t][c][n]. ++kernel void mtp_combine_gnorm(device const float* h [[buffer(0)]], device const float* e [[buffer(1)]], ++ device float* R [[buffer(2)]], device const float* w_norm [[buffer(3)]], ++ device float* xn [[buffer(4)]], constant CombNormArgs& a [[buffer(5)]], ++ uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], ++ ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { ++ threadgroup float part[8]; ++ const uint c = tg.x, t = tg.y, n = a.n; ++ const uint row = (t * a.hc + c) * n; ++ device float* r = R + row; ++ float values[10]; // N = 2560: ten values per thread, retained across the reduction. ++ float ss = 0.0f; ++ if (n == 2560) { ++ device const float* hp = h + row + tid; ++ device const float* ep = e + t * 2560 + tid; ++ device float* rp = r + tid; ++#pragma clang loop unroll(full) ++ for (int i = 0; i < 10; ++i) { ++ const float v = hp[256 * i] + ep[256 * i]; ++ rp[256 * i] = v; ++ values[i] = v; ++ ss = fma(v, v, ss); ++ } ++ } else { ++ for (uint d = tid; d < n; d += 256) { ++ const float v = h[row + d] + e[t * n + d]; ++ r[d] = v; ++ values[d / 256] = v; ++ ss = fma(v, v, ss); ++ } ++ } ++ ss = simd_sum(ss); ++ if (lane == 0) part[sg] = ss; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ float tot = 0.0f; ++#pragma clang loop unroll(full) ++ for (int i = 0; i < 8; ++i) tot += part[i]; ++ const float rs = rsqrt(tot / float(n) + a.eps); ++ device float* o = xn + row; ++ device const float* w = w_norm + c * n; ++ if (n == 2560) { ++ o += tid; ++ w += tid; ++#pragma clang loop unroll(full) ++ for (int i = 0; i < 10; ++i) o[256 * i] = values[i] * w[256 * i] * rs; ++ } else { ++ for (uint d = tid; d < n; d += 256) { ++ const float v = values[d / 256]; ++ o[d] = v * w[d] * rs; ++ } ++ } + } +diff --git a/metal/src/engine/msl/mtp_chain.metal b/metal/src/engine/msl/mtp_chain.metal +index 62500df2..efee47ee 100644 +--- a/metal/src/engine/msl/mtp_chain.metal ++++ b/metal/src/engine/msl/mtp_chain.metal +@@ -13,6 +13,255 @@ + // a 16-wide unit has one sum: m1) + // and its position in the row is what X::load reads when x holds the positions 0 .. n-1. + ++struct DraftSlopArgs { ++ slop4::State state; ++ slop4::Knobs knobs; ++ uint nv, j, vocab; ++}; ++ ++constant int32_t kSlRather[] = {4598, 24944, 61745, 70714}, kSlAtherPre[] = {2515, 2697, 15656, 43487}; ++constant int32_t kSlRight[] = {1245, 1246, 5791, 9785, 26806, 51140}; ++constant int32_t kSlEndPre[] = {400, 809, 3601, 3839, 14637, 55293}; ++constant int32_t kSlShrt[] = {2716, 8412, 10382, 12115, 36639}; ++constant int32_t kSlExactly[] = {6681, 63151, 66172}, kSlExPre[] = {325, 5291}; ++constant int32_t kSlHere[] = {1532, 5514, 6527, 8160, 475}, kSlWhat[] = {1092, 3437, 3710, 12195, 29348}; ++constant int32_t kSlDet[] = {1719, 760, 279, 561, 276, 449, 1473, 2014, 2329, 821, 2921, 4888, 1323, 381, ++ 1602, 32985, 1379, 3444, 7261, 6183, 777, 4254, 678}; ++constant int32_t kSlClause[] = {13, 11, 0, 30, 25, 26, 1076, 2218, 1892, 12, 471, 8, ++ 321, 694, 748, 429, 7834, 684, 1541, 1921, 1987, 2939, 3850, 4272}; ++constant int32_t kSlH[] = {71, 39, 3101, 75833}; ++constant int32_t kSlSentence[] = {13, 0, 30, 25, 26, 1076, 2218, 1892, 8}; ++constant int32_t kSlCandor[] = {1865, 47045, 24969, 5507, 17290, 7435, 62179, 66842, 86775, ++ 12116, 6379, 6999, 13542, 11040, 73018}; ++constant int32_t kSlBePre[] = {310, 951, 3172, 728, 1066, 16032}; ++constant int32_t kSlIntens[] = {12116, 6379, 6999, 13542, 11040, 73018, 66842}; ++constant int32_t kSlFrameAdj[] = {10429, 1865, 47045, 24969, 5507, 7435, 2716, 4145, 3841, 26860, 13983, ++ 8412, 12115, 10382, 7700, 12402, 8535, 4880, 32254}; ++constant int32_t kSlFrameDet[] = {1719, 760, 279, 561, 2329, 821, 2921, 4888, 276, 449}; ++constant int32_t kSlFrame[] = {4087, 1831, 56004, 2243, 7793, 1301, 11782, 14996, 3296}; ++constant int32_t kSlThe[] = {1719, 760, 279, 561}, kSlTo[] = {1206, 1946, 951, 310}; ++constant int32_t kSlCandorBranch[] = {1865, 47045, 24969, 5507, 7435, 2716, 26860}; ++constant int32_t kSlLineFrame[] = {8412, 12115, 7700, 12402}; ++constant int32_t kSlMd[] = {12, 9, 6458, 471, 348, 7061, 332, 2972, 2, 550, 13962, 794, 653, 7460, 16112, ++ 4277, 11862, 4006, 2218, 10, 49300}; ++ ++template static inline bool sl_in(int32_t t, constant const int32_t (&s)[N]) { ++ bool found = false; ++#pragma clang loop unroll(full) ++ for (size_t i = 0; i < N; ++i) found |= (s[i] == t); ++ return found; ++} ++ ++static inline int mtp_slop_penalties(slop4::State s, uchar c_p0, uchar c_p1, slop4::Knobs k, ++ threadgroup int* ids, threadgroup float* vals) { ++ if ((s.f & (slop4::KNOWN | slop4::PROSE | slop4::TOOL | slop4::FENCE)) != (slop4::KNOWN | slop4::PROSE)) return 0; ++ int n = 0; ++ auto pen = [&](int32_t id, float v) { ++ if (v > 0.0f && n < 64) { ++ ids[n] = id; ++ vals[n] = v; ++ ++n; ++ } ++ }; ++ const int32_t p0 = s.t1, p1 = s.t2; ++ if (sl_in(p0, kSlRather) || (p0 == 1792 && sl_in(p1, kSlAtherPre))) pen(1056, k.lex); ++ if (sl_in(p0, kSlRight) || (p0 == 481 && p1 == 2697)) pen(3023, k.lex); ++ if (p0 == 310 && sl_in(p1, kSlEndPre)) pen(809, k.lex); ++ if (sl_in(p0, kSlShrt)) pen(2243, k.lex); ++ if (sl_in(p0, kSlExactly) || (p0 == 31654 && sl_in(p1, kSlExPre))) pen(1092, k.lex); ++ if (p0 == 579 && sl_in(p1, kSlHere)) pen(1092, k.lex); ++ if (p0 == 579 && sl_in(p1, kSlWhat)) pen(3404, k.lex); ++ const bool line_start = c_p0 == 1 || (s.f & slop4::BOL); ++ const bool sentence_start = sl_in(p0, kSlSentence) || line_start; ++ const bool clause_start = sl_in(p0, kSlClause) || line_start; ++ pen(92115, k.hon); ++ pen(68540, k.hon); ++ pen(51397, k.hon); ++ if (sl_in(p0, kSlDet) || sentence_start) pen(10429, k.hon); ++ if (clause_start) pen(25211, k.hon); ++ if (sl_in(p0, kSlH)) { pen(33961, 1.5f * k.hon); pen(77696, 1.5f * k.hon); } ++ if (p0 == 660) { pen(46253, k.hon); pen(25211, k.hon); } ++ if (p0 == 1602 || p0 == 32985) { ++#pragma clang loop unroll(full) ++ for (int32_t id : kSlCandor) pen(id, k.hon); ++ } ++ if ((p0 == 381 || p0 == 1323) && sl_in(p1, kSlBePre)) { ++#pragma clang loop unroll(full) ++ for (int32_t id : kSlCandor) pen(id, k.hon); ++ pen(2050, k.hon); ++ if (p1 == 310 || p1 == 951) { pen(6415, k.hon); pen(2708, k.hon); } ++ } ++ if (sl_in(p0, kSlIntens) && (p1 == 1602 || p1 == 32985 || p1 == 381 || p1 == 1323)) { ++ pen(10429, k.hon); ++#pragma clang loop unroll(full) ++ for (int32_t id : kSlCandor) pen(id, k.hon); ++ } ++ if (p0 == 2688 && (p1 == 353 || p1 == 585 || p1 == 40 || p1 == 72)) { pen(1602, 0.5f * k.hon); pen(10429, k.hon); } ++ if (clause_start) pen(40285, k.hon); ++ if (p0 == 660) { pen(49034, k.hon); pen(71853, k.hon); } ++ if (sl_in(p0, kSlFrameAdj) && (sl_in(p1, kSlFrameDet) || sl_in(p1, kSlSentence) || c_p1 == 1 || p1 == slop4::ETHINK)) { ++#pragma clang loop unroll(full) ++ for (int32_t id : kSlFrame) pen(id, k.hon); ++ } ++ const bool p1_root = sl_in(p1, kSlSentence) || c_p1 == 1 || p1 == slop4::ETHINK || p1 < 0; ++ if (sl_in(p0, kSlThe) && (p1_root || (c_p0 == 0 && (s.f & slop4::PBOL)))) { ++#pragma clang loop unroll(full) ++ for (int32_t id : kSlCandorBranch) pen(id, 0.5f * k.hon); ++ } ++ if (sl_in(p0, kSlTo) && (p1_root || (s.f & slop4::PBOL))) pen(381, 0.5f * k.hon); ++ if (s.f & slop4::BOL) { ++#pragma clang loop unroll(full) ++ for (int32_t id : kSlLineFrame) pen(id, 0.5f * k.hon); ++ const float v = k.md + ((s.f & slop4::COLON) ? k.colon : 0.0f); ++#pragma clang loop unroll(full) ++ for (int32_t id : kSlMd) pen(id, v); ++#pragma clang loop unroll(full) ++ for (int32_t id : slop4::kDigit) pen(id, (s.f & slop4::COLON) ? v : 0.25f * v); ++ } ++ if ((s.f & slop4::PBOL) && sl_in(p0, slop4::kDigit)) { pen(13, k.md + k.colon); pen(8, k.md + k.colon); } ++ if (sl_in(p0, slop4::kColon) && s.len <= 6) { pen(slop4::NL, k.colon); pen(slop4::NL2, k.colon); } ++ return n; ++} ++// one threadgroup; the state advances only along the draft branch, never the verifier's history. ++kernel void mtp_slop_pen(device float* lg [[buffer(0)]], device const int* inverse [[buffer(1)]], ++ device const int* dout [[buffer(2)]], device slop4::State* state [[buffer(3)]], ++ constant DraftSlopArgs& a [[buffer(4)]], ushort tid [[thread_index_in_threadgroup]], ++ ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { ++ threadgroup int ids[80]; ++ threadgroup float vals[80], rem[80], part[8]; ++ threadgroup int hold_ids[4]; ++ threadgroup int count; ++ threadgroup slop4::State st; ++ threadgroup float mx, total; ++ threadgroup bool hold; ++ threadgroup uchar tcls[2]; ++ const bool any_knob = a.knobs.lex > 0.0f || a.knobs.hon > 0.0f || a.knobs.md > 0.0f || a.knobs.colon > 0.0f; ++ if (a.j == 0 && !any_knob) { ++ if (tid == 0) state[0] = a.state; ++ return; ++ } ++ const slop4::State s0 = a.j == 0 ? a.state : state[0]; ++ const int32_t tok0 = a.j == 0 ? s0.t1 : dout[a.j - 1]; ++ const int32_t tok1 = a.j == 0 ? s0.t2 : s0.t1; ++ if (tid < 2) tcls[tid] = 0; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ constexpr uint kNL = sizeof(slop4::kNewline) / sizeof(int32_t); ++ constexpr uint kBL = sizeof(slop4::kBlank) / sizeof(int32_t); ++ if (tid < kNL) { ++ const int32_t x = slop4::kNewline[tid]; ++ if (x == tok0) tcls[0] = 1; ++ if (x == tok1) tcls[1] = 1; ++ } ++ if (tid < kBL) { ++ const int32_t x = slop4::kBlank[tid]; ++ if (x == tok0) tcls[0] = 2; ++ if (x == tok1) tcls[1] = 2; ++ } ++ if (tid < 4) { ++ const int hid = slop4::kHold[tid]; ++ hold_ids[tid] = uint(hid) < a.vocab ? inverse[hid] : -1; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (tid == 0) { ++ slop4::State s = s0; ++ const uchar c0 = tcls[0], c1 = tcls[1]; ++ if (a.j != 0) { ++ s.t2 = s.t1; ++ s.t1 = tok0; ++ s.f |= slop4::KNOWN; ++ const uint8_t bol = s.f & slop4::BOL; ++ s.f &= (uint8_t) ~slop4::PBOL; ++ if (tok0 == slop4::IM_START || tok0 == slop4::IM_END || tok0 == slop4::THINK) { ++ s.f = slop4::KNOWN | slop4::BOL; ++ s.len = 0; ++ } else if (tok0 == slop4::ETHINK) { ++ s.f = slop4::KNOWN | slop4::PROSE | slop4::BOL; ++ s.len = 0; ++ } else if (tok0 == slop4::TOOL_CALL) { ++ s.f |= slop4::TOOL; ++ } else if (tok0 == slop4::ETOOL_CALL) { ++ s.f = (uint8_t) ((s.f & ~(slop4::TOOL | slop4::COLON)) | slop4::BOL); ++ s.len = 0; ++ } else if (c0 == 1) { ++ if (!bol) s.f = (uint8_t) ((s.f & ~slop4::COLON) | (slop4::in(s.t2, slop4::kColon) ? slop4::COLON : 0)); ++ s.f |= slop4::BOL; ++ s.len = 0; ++ } else if (c0 == 2) { ++ } else { ++ if (slop4::in(tok0, slop4::kFence) && (s.f & slop4::PROSE) && !(s.f & slop4::TOOL)) s.f ^= slop4::FENCE; ++ s.f = (uint8_t) ((s.f & ~(slop4::BOL | slop4::COLON)) | (bol ? slop4::PBOL : 0)); ++ if (s.len < 255) ++s.len; ++ } ++ } ++ state[0] = s; ++ st = s; ++ const slop4::Knobs kn = a.knobs; ++ count = (kn.lex > 0.0f || kn.hon > 0.0f || kn.md > 0.0f || kn.colon > 0.0f) ++ ? mtp_slop_penalties(s, c0, c1, kn, ids, vals) ++ : 0; ++ hold = false; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (count > 0) { ++ if (int(tid) < count) { ++ const int raw_id = ids[tid]; ++ ids[tid] = raw_id >= 0 && uint(raw_id) < a.vocab ? inverse[raw_id] : -1; ++ } ++ float m = -INFINITY; ++#pragma clang loop unroll_count(4) ++ for (uint i = tid; i < a.nv; i += 256) if (!isnan(lg[i])) m = max(m, lg[i]); ++ m = simd_max(m); ++ if (lane == 0) part[sg] = m; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ float lmx = -INFINITY; ++#pragma clang loop unroll(full) ++ for (int k = 0; k < 8; ++k) lmx = max(lmx, part[k]); ++ if (tid == 0) mx = lmx; ++ if (int(tid) < count) { ++ const int id = ids[tid]; ++ if (id >= 0) { ++ const float lgi = lg[id]; ++ if (lgi > lmx - 12.0f) hold = true; ++ rem[tid] = exp(lgi - lmx) * (1.0f - exp(-vals[tid])); ++ } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } ++ if (hold) { ++ float sum = 0.0f; ++#pragma clang loop unroll_count(4) ++ for (uint i = tid; i < a.nv; i += 256) sum += exp(lg[i] - mx); ++ sum = simd_sum(sum); ++ if (lane == 0) part[sg] = sum; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (tid == 0) { ++ total = 0.0f; ++#pragma clang loop unroll(full) ++ for (int k = 0; k < 8; ++k) total += part[k]; ++ float removed = 0.0f; ++ for (int i = 0; i < count; ++i) ++ if (ids[i] >= 0) removed += rem[i]; ++ const float v = -log(1.0f - min(removed / total, 0.999999f)); ++#pragma clang loop unroll(full) ++ for (int k = 0; k < 4; ++k) { ++ ids[count] = hold_ids[k]; ++ vals[count++] = v; ++ } ++ } ++ } ++ if (tid == 0) { ++ if (a.knobs.lex > 0.0f && st.t2 == slop4::ETHINK && (st.t1 == slop4::NL || st.t1 == slop4::NL2)) { ++ const int empty[] = {slop4::IM_END, slop4::THINK}; ++ for (int id : empty) { ++ ids[count] = uint(id) < a.vocab ? inverse[id] : -1; ++ vals[count++] = a.knobs.lex; ++ } ++ } ++ // serial subtraction preserves additive penalties when two rules name the same token. ++ for (int i = 0; i < count; ++i) if (ids[i] >= 0) lg[ids[i]] -= vals[i]; ++ } ++} ++ + struct PickArgs { + uint nv; // the draft head's rows (the draft vocabulary, or the whole vocabulary) + uint j; // the draft step +@@ -28,12 +277,14 @@ kernel void mtp_pick(device const float* lg [[buffer(0)]], device const int* pic + const int di = pick[0]; + const float mx = lg[di]; + float s = 0.0f; ++#pragma clang loop unroll_count(4) + for (uint i = tid; i < a.nv; i += 1024) s += exp(lg[i] - mx); + s = simd_sum(s); + if (lane == 0) part[sg] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + float t = 0.0f; ++#pragma clang loop unroll(full) + for (int k = 0; k < 32; ++k) t += part[k]; + dout[a.j] = a.use_map ? vmap[di] : di; + dprob[a.j] = 1.0f / t; +@@ -60,28 +311,54 @@ template static inline float kind_scale(int pos) { + return 1.0f; + } + ++template static inline void mtp_embed_unit(device const uchar* row, int u, device float* y, ++ threadgroup const uint* tab) { ++ constexpr int UNIT = fmt::UNIT; ++ float q[UNIT]; ++ float4 sc; ++ Dec::run(row, u, q, sc, tab); ++ const int p0 = (F == T_Q4_K || F == T_Q5_K) ? (256 * (u >> 3) + 64 * ((u >> 1) & 3) + 16 * (u & 1)) ++ : (F == T_Q6_K) ? (256 * (u >> 3) + 128 * ((u >> 2) & 1) + 32 * (u & 1) + 16 * ((u >> 1) & 1)) ++ : (UNIT * u); ++ device float4* y0 = (device float4*)(y + p0); ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 4; ++j) { ++ y0[j] = float4(sc.x * (q[4 * j + 0] * kind_scale(4 * j + 0)) - sc.z, ++ sc.x * (q[4 * j + 1] * kind_scale(4 * j + 1)) - sc.z, ++ sc.x * (q[4 * j + 2] * kind_scale(4 * j + 2)) - sc.z, ++ sc.x * (q[4 * j + 3] * kind_scale(4 * j + 3)) - sc.z); ++ } ++ if (UNIT == 32) { ++ const int p1 = p0 + ((F == T_Q4_K || F == T_Q5_K) ? 32 : (F == T_Q6_K ? 64 : 16)); ++ device float4* y1 = (device float4*)(y + p1); ++#pragma clang loop unroll(full) ++ for (int j = 0; j < 4; ++j) { ++ y1[j] = float4(sc.y * (q[16 + 4 * j + 0] * kind_scale(4 * j + 0)) - sc.w, ++ sc.y * (q[16 + 4 * j + 1] * kind_scale(4 * j + 1)) - sc.w, ++ sc.y * (q[16 + 4 * j + 2] * kind_scale(4 * j + 2)) - sc.w, ++ sc.y * (q[16 + 4 * j + 3] * kind_scale(4 * j + 3)) - sc.w); ++ } ++ } ++} ++ + // one threadgroup of 256; posv = the floats 0 .. n-1 + template + kernel void mtp_embed(device const uchar* W [[buffer(0)]], device const int* dout [[buffer(1)]], + device const float* posv [[buffer(2)]], device float* y [[buffer(3)]], constant EmbArgs& a [[buffer(4)]], + ushort tid [[thread_index_in_threadgroup]]) { +- constexpr int UNIT = fmt::UNIT, H = UNIT / 2; ++ (void)posv; ++ constexpr int UNIT = fmt::UNIT; + threadgroup uint tab[fmt::TAB]; +- load_tab(tab, tid, 256); +- threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (fmt::TAB > 1) { ++ load_tab(tab, tid, 256); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } + device const uchar* row = W + (ulong) dout[a.j] * a.row_bytes; + const int U = int(a.n) / UNIT; +- for (int u = tid; u < U; u += 256) { +- float q[UNIT], ps[UNIT], xsum[2]; +- float4 sc; +- Dec::run(row, u, q, sc, tab); +- X::load(posv, posv, u, ps, xsum); +- for (int j = 0; j < UNIT; ++j) { +- const int pos = int(ps[j]); +- const float d = j < H ? sc.x : sc.y; +- const float m = (UNIT == 16 || j < 16) ? sc.z : sc.w; +- y[pos] = d * (q[j] * kind_scale(pos)) - m; +- } ++ if (U <= 256) { ++ if (int(tid) < U) mtp_embed_unit(row, int(tid), y, tab); ++ } else { ++ for (int u = tid; u < U; u += 256) mtp_embed_unit(row, u, y, tab); + } + } + +diff --git a/metal/src/engine/msl/ple.metal b/metal/src/engine/msl/ple.metal +index 7036eaeb..d323f0ea 100644 +--- a/metal/src/engine/msl/ple.metal ++++ b/metal/src/engine/msl/ple.metal +@@ -9,43 +9,32 @@ + // R += gated + conv + // + // History: hist[row + 9 * channel] (row-fastest, oldest first, the committed rows); the window's normalized rows +-// are written to norm_w[t][channel] (the host appends the kept ones at commit). ++// are written to norm_w[t][channel] (the window's host appends the kept ones at commit; a prompt chunk's ple_hist_adv). + + #include + using namespace metal; + ++#ifndef P_TOKENT ++#define P_TOKENT 532 // moe.metal's plan layout ++#endif ++ + struct PleArgs { + uint n_embd; // 2560 + uint hc; // 4 + float eps; + }; + ++// One barrier per reduction; callers with back-to-back reductions alternate part[0..7] and part[8..15]. + static inline float tg_sum(float v, threadgroup float* part, ushort sg, ushort lane) { + v = simd_sum(v); + if (lane == 0) part[sg] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + float t = 0.0f; ++#pragma clang loop unroll(full) + for (int i = 0; i < 8; ++i) t += part[i]; +- threadgroup_barrier(mem_flags::mem_threadgroup); + return t; + } + +-// grid (hc, T), 256 threads: y[t][c] = grouped_norm(x[t][c]) with w[c] +-kernel void ple_gnorm(device const float* x [[buffer(0)]], device const float* w [[buffer(1)]], device float* y [[buffer(2)]], +- constant PleArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], +- ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], +- ushort lane [[thread_index_in_simdgroup]]) { +- threadgroup float part[8]; +- const uint c = tg.x, t = tg.y, n = a.n_embd; +- device const float* xc = x + (t * a.hc + c) * n; +- float ss = 0.0f; +- for (uint d = tid; d < n; d += 256) ss += xc[d] * xc[d]; +- const float mean = tg_sum(ss, part, sg, lane) / float(n); +- const float scale = 1.0f / sqrt(mean + a.eps); +- device float* yc = y + (t * a.hc + c) * n; +- for (uint d = tid; d < n; d += 256) yc[d] = xc[d] * scale * w[c * n + d]; +-} +- + // grid (hc, T), 256 threads: the gate, then gated = value * gate and normalized = grouped_norm(gated, norm_conv) + kernel void ple_gate_norm(device const float* key [[buffer(0)]], device const float* query [[buffer(1)]], + device const float* value [[buffer(2)]], device const float* w_conv [[buffer(3)]], +@@ -53,7 +42,7 @@ kernel void ple_gate_norm(device const float* key [[buffer(0)]], device const fl + constant PleArgs& a [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { +- threadgroup float part[8]; ++ threadgroup float part[16]; + const uint c = tg.x, t = tg.y, n = a.n_embd, row = (t * a.hc + c) * n; + float acc = 0.0f; + for (uint d = tid; d < n; d += 256) acc += key[row + d] * query[row + d]; +@@ -67,25 +56,381 @@ kernel void ple_gate_norm(device const float* key [[buffer(0)]], device const fl + gated[row + d] = g; + ss += g * g; + } +- const float mean = tg_sum(ss, part, sg, lane) / float(n); ++ const float mean = tg_sum(ss, part + 8, sg, lane) / float(n); + const float scale = 1.0f / sqrt(mean + a.eps); + for (uint d = tid; d < n; d += 256) norm_w[row + d] = gated[row + d] * scale * w_conv[c * n + d]; + } + +-// grid (hc * n_embd / 256, T), 256 threads: the conv over [history | the window's rows], SiLU, and the residual ++// The fused forms keep a thread's NPT values of a row in registers (n_embd <= 256 * NPT), so the key and query norms ++// and gated never reach memory; each is bitwise the ple_gnorm / ple_gate_norm / ple_conv_add chain it replaces. ++constant constexpr uint NPT = 10; ++ ++// y = grouped_norm(y) in place, ple_gnorm's arithmetic ++static inline void gnorm_regs(thread float* y, device const float* w, uint n, uint tid, float eps, ++ threadgroup float* part, ushort sg, ushort lane) { ++ float ss = 0.0f; ++ if (n == 2560) { ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < NPT; ++i) ss += y[i] * y[i]; ++ } else { ++ for (uint i = 0, d = tid; d < n; ++i, d += 256) ss += y[i] * y[i]; ++ } ++ const float mean = tg_sum(ss, part, sg, lane) / float(n); ++ const float scale = 1.0f / sqrt(mean + eps); ++ if (n == 2560) { ++ device const float* wt = w + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < NPT; ++i) y[i] = y[i] * scale * wt[i * 256]; ++ } else { ++ for (uint i = 0, d = tid; d < n; ++i, d += 256) y[i] = y[i] * scale * w[d]; ++ } ++} ++ ++// grid (hc, T), 256 threads: y[t][c] = grouped_norm(x[t][c]) with w[c]; a row in registers when it fits (x read once) ++kernel void ple_gnorm(device const float* x [[buffer(0)]], device const float* w [[buffer(1)]], device float* y [[buffer(2)]], ++ constant PleArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], ++ ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], ++ ushort lane [[thread_index_in_simdgroup]]) { ++ threadgroup float part[8]; ++ const uint c = tg.x, t = tg.y, n = a.n_embd; ++ device const float* xc = x + (t * a.hc + c) * n; ++ device float* yc = y + (t * a.hc + c) * n; ++ if (n == 2560) { ++ float xv[NPT]; ++ device const float* xt = xc + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < NPT; ++i) xv[i] = xt[i * 256]; ++ gnorm_regs(xv, w + c * n, n, tid, a.eps, part, sg, lane); ++ device float* yt = yc + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < NPT; ++i) yt[i * 256] = xv[i]; ++ return; ++ } ++ float ss = 0.0f; ++ for (uint d = tid; d < n; d += 256) ss += xc[d] * xc[d]; ++ const float mean = tg_sum(ss, part, sg, lane) / float(n); ++ const float scale = 1.0f / sqrt(mean + a.eps); ++ for (uint d = tid; d < n; d += 256) yc[d] = xc[d] * scale * w[c * n + d]; ++} ++ ++// ple_gate_norm on register rows: the gate from key . query, then normalized written and the gated rows added to the ++// residual in the same pass: R = rb + gated, which ple_conv_res completes to ple_conv_add's (R + gated) + conv, so ++// gated never reaches memory. part has 16 floats: the prior gnorm_regs left part[0..7] live, so the first reduction ++// uses part + 8 and the second reuses part[0..7]. ++static inline void gate_norm_regs(thread const float* key, thread const float* query, thread const float* value, ++ device const float* w_conv, device float* R, thread const float* rb, ++ device float* norm_w, uint n, uint tid, float eps, threadgroup float* part, ushort sg, ++ ushort lane) { ++ float acc = 0.0f; ++ if (n == 2560) { ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < NPT; ++i) acc += key[i] * query[i]; ++ } else { ++ for (uint i = 0, d = tid; d < n; ++i, d += 256) acc += key[i] * query[i]; ++ } ++ const float s = tg_sum(acc, part + 8, sg, lane) * (1.0f / sqrt(float(n))); ++ const float mag = sqrt(max(fabs(s), 1e-6f)); ++ const float sgn = s > 0.0f ? 1.0f : (s < 0.0f ? -1.0f : 0.0f); ++ const float gate = 1.0f / (1.0f + exp(-(sgn * mag))); ++ float gv[NPT]; ++ float ss = 0.0f; ++ if (n == 2560) { ++ device float* rt = R + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < NPT; ++i) { ++ gv[i] = value[i] * gate; ++ rt[i * 256] = rb[i] + gv[i]; ++ ss += gv[i] * gv[i]; ++ } ++ } else { ++ for (uint i = 0, d = tid; d < n; ++i, d += 256) { ++ gv[i] = value[i] * gate; ++ R[d] = rb[i] + gv[i]; ++ ss += gv[i] * gv[i]; ++ } ++ } ++ const float mean = tg_sum(ss, part, sg, lane) / float(n); ++ const float scale = 1.0f / sqrt(mean + eps); ++ if (n == 2560) { ++ device float* nwt = norm_w + tid; ++ device const float* wct = w_conv + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < NPT; ++i) nwt[i * 256] = gv[i] * scale * wct[i * 256]; ++ } else { ++ for (uint i = 0, d = tid; d < n; ++i, d += 256) norm_w[d] = gv[i] * scale * w_conv[d]; ++ } ++} ++ ++// grid (hc, T), 256 threads: ple_gnorm(R, w_query) + ple_gate_norm with the key norm already in key (the window: ++// its key norm runs beside layer 0's first read); R += gated, ple_conv_res adds the conv ++kernel void ple_query_gate_norm(device const float* key [[buffer(0)]], device float* R [[buffer(1)]], ++ device const float* w_query [[buffer(2)]], device const float* value [[buffer(3)]], ++ device const float* w_conv [[buffer(4)]], device float* norm_w [[buffer(5)]], ++ constant PleArgs& a [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], ++ ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], ++ ushort lane [[thread_index_in_simdgroup]]) { ++ threadgroup float part[16]; ++ const uint c = tg.x, t = tg.y, n = a.n_embd, row = (t * a.hc + c) * n; ++ float k[NPT], q[NPT], rb[NPT], v[NPT]; ++ if (n == 2560) { ++ device const float* kt = key + row + tid; ++ device const float* rt = R + row + tid; ++ device const float* vt = value + t * 2560 + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < NPT; ++i) { ++ k[i] = kt[i * 256]; ++ q[i] = rb[i] = rt[i * 256]; ++ v[i] = vt[i * 256]; ++ } ++ } else { ++ for (uint i = 0, d = tid; d < n; ++i, d += 256) { ++ k[i] = key[row + d]; ++ q[i] = rb[i] = R[row + d]; ++ v[i] = value[t * n + d]; ++ } ++ } ++ gnorm_regs(q, w_query + c * n, n, tid, a.eps, part, sg, lane); ++ gate_norm_regs(k, q, v, w_conv + c * n, R + row, rb, norm_w + row, n, tid, a.eps, part, sg, lane); ++} ++ ++// grid (hc, T), 256 threads: ple_query_gate_norm with layer 0's MoE write formed and applied as R is read, gr_norm_comb's ++// R' = R + (sum_k ye[entry(t, k)] (k in order) + sigmoid(sg) ysh) * 2 sigmoid(inj / hc): bitwise gr_norm_comb + ++// ple_query_gate_norm, with R written once and none of gr_norm's xn (nothing reads it: layer 1's read norms R again ++// after the conv) ++kernel void ple_comb_query_gate_norm(device const float* key [[buffer(0)]], device float* R [[buffer(1)]], ++ device const float* w_query [[buffer(2)]], device const float* value [[buffer(3)]], ++ device const float* w_conv [[buffer(4)]], device float* norm_w [[buffer(5)]], ++ constant PleArgs& a [[buffer(6)]], device const float* ye [[buffer(7)]], ++ device const float* inj [[buffer(8)]], device const int* plan [[buffer(9)]], ++ device const float* ysh [[buffer(10)]], device const float* sgv [[buffer(11)]], ++ constant uint& topk [[buffer(12)]], uint3 tg [[threadgroup_position_in_grid]], ++ ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], ++ ushort lane [[thread_index_in_simdgroup]]) { ++ threadgroup float part[16]; ++ const uint c = tg.x, t = tg.y, n = a.n_embd, row = (t * a.hc + c) * n; ++ const float g = 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))); ++ const float gs = 1.0f / (1.0f + exp(-sgv[t])); ++ float k[NPT], q[NPT], rb[NPT], v[NPT]; ++ if (n == 2560 && topk == 10) { ++ device const float* yep[10]; ++#pragma clang loop unroll(full) ++ for (uint j = 0; j < 10; ++j) yep[j] = ye + uint(plan[P_TOKENT + t * 10 + j]) * 2560u + tid; ++ device const float* yshp = ysh + t * 2560u + tid; ++ device const float* kt = key + row + tid; ++ device const float* rt = R + row + tid; ++ device const float* vt = value + t * 2560 + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < NPT; ++i) { ++ float s = 0.0f; ++#pragma clang loop unroll(full) ++ for (uint j = 0; j < 10; ++j) s += yep[j][i * 256]; ++ q[i] = rb[i] = fma(fma(gs, yshp[i * 256], s), g, rt[i * 256]); ++ k[i] = kt[i * 256]; ++ v[i] = vt[i * 256]; ++ } ++ } else { ++ for (uint i = 0, d = tid; d < n; ++i, d += 256) { ++ float s = 0.0f; ++ for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; ++ q[i] = rb[i] = fma(fma(gs, ysh[t * n + d], s), g, R[row + d]); ++ k[i] = key[row + d]; ++ v[i] = value[t * n + d]; ++ } ++ } ++ gnorm_regs(q, w_query + c * n, n, tid, a.eps, part, sg, lane); ++ gate_norm_regs(k, q, v, w_conv + c * n, R + row, rb, norm_w + row, n, tid, a.eps, part, sg, lane); ++} ++ ++// stream c of token t from the residual rows rb (layer 0's MoE write applied): the key and query norms, the gate, ++// R = rb + gated and the normalized rows. part has 16 floats: gnorm_regs(k) uses part + 8, gnorm_regs(q) uses ++// part[0..7], and gate_norm_regs uses part + 8 then part[0..7]. ++static inline void moe_gate_norm_row(device float* R, thread const float* rb, device const float* pkey, ++ device const float* w_key, device const float* w_query, thread const float* v, ++ device const float* w_conv, device float* norm_w, uint row, uint c, uint n, ++ uint tid, float eps, threadgroup float* part, ushort sg, ushort lane) { ++ float k[NPT], q[NPT]; ++ if (n == 2560) { ++ device const float* pkt = pkey + row + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < NPT; ++i) { ++ q[i] = rb[i]; ++ k[i] = pkt[i * 256]; ++ } ++ } else { ++ for (uint i = 0, d = tid; d < n; ++i, d += 256) { ++ q[i] = rb[i]; ++ k[i] = pkey[row + d]; ++ } ++ } ++ gnorm_regs(k, w_key + c * n, n, tid, eps, part + 8, sg, lane); ++ gnorm_regs(q, w_query + c * n, n, tid, eps, part, sg, lane); ++ gate_norm_regs(k, q, v, w_conv + c * n, R + row, rb, norm_w + row, n, tid, eps, part, sg, lane); ++} ++ ++// grid (hc, T), 256 threads: a prompt chunk's whole PLE norm stage - layer 0's MoE write applied to R (moe_out + ++// gr_norm's R' = R + fma(sigmoid(sg), ys, yr) * 2 sigmoid(inj / hc)), the key and query norms, the gate, the ++// normalized rows and R' + gated written back; ple_conv_res adds the conv ++kernel void ple_moe_gate_norm(device float* R [[buffer(0)]], device const float* yr [[buffer(1)]], ++ device const float* ys [[buffer(2)]], device const float* sgv [[buffer(3)]], ++ device const float* inj [[buffer(4)]], device const float* pkey [[buffer(5)]], ++ device const float* w_key [[buffer(6)]], device const float* w_query [[buffer(7)]], ++ device const float* value [[buffer(8)]], device const float* w_conv [[buffer(9)]], ++ device float* norm_w [[buffer(10)]], constant PleArgs& a [[buffer(11)]], ++ uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], ++ ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { ++ threadgroup float part[16]; ++ const uint c = tg.x, t = tg.y, n = a.n_embd, row = (t * a.hc + c) * n; ++ const float g = 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))); ++ const float gs = 1.0f / (1.0f + exp(-sgv[t])); ++ float rb[NPT], v[NPT]; ++ if (n == 2560) { ++ device const float* yst = ys + t * 2560 + tid; ++ device const float* yrt = yr + t * 2560 + tid; ++ device const float* rt = R + row + tid; ++ device const float* vt = value + t * 2560 + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < NPT; ++i) { ++ rb[i] = fma(fma(gs, yst[i * 256], yrt[i * 256]), g, rt[i * 256]); ++ v[i] = vt[i * 256]; ++ } ++ } else { ++ for (uint i = 0, d = tid; d < n; ++i, d += 256) { ++ rb[i] = fma(fma(gs, ys[t * n + d], yr[t * n + d]), g, R[row + d]); ++ v[i] = value[t * n + d]; ++ } ++ } ++ moe_gate_norm_row(R, rb, pkey, w_key, w_query, v, w_conv, norm_w, row, c, n, tid, a.eps, part, sg, lane); ++} ++ ++// ple_moe_gate_norm with mm_combine folded in (prompt.metal's gr_norm_h_moec): yr formed from the expert rows YE ++// (pos, srt_w: mm_place's) as it is read, bitwise mm_combine + ple_moe_gate_norm. grid (1, T), 256 threads: a ++// token's streams in turn, so YE, ys and value are read once a token; hc <= 4, topk <= 16 ++kernel void ple_moe_gate_norm_c(device float* R [[buffer(0)]], device const float* YE [[buffer(1)]], ++ device const int* pos [[buffer(2)]], device const float* srt_w [[buffer(3)]], ++ device const float* ys [[buffer(4)]], device const float* sgv [[buffer(5)]], ++ device const float* inj [[buffer(6)]], device const float* pkey [[buffer(7)]], ++ device const float* w_key [[buffer(8)]], device const float* w_query [[buffer(9)]], ++ device const float* value [[buffer(10)]], device const float* w_conv [[buffer(11)]], ++ device float* norm_w [[buffer(12)]], constant PleArgs& a [[buffer(13)]], ++ constant uint& topk [[buffer(14)]], uint3 tg [[threadgroup_position_in_grid]], ++ ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], ++ ushort lane [[thread_index_in_simdgroup]]) { ++ threadgroup float part[16]; ++ threadgroup uint roff[17]; ++ threadgroup float wk[16]; ++ const uint t = tg.y, n = a.n_embd; ++ int r = -1; ++ float w = 0.0f; ++ if (tid < topk) { ++ r = pos[t * topk + tid]; ++ if (r >= 0) w = srt_w[r]; ++ } ++ const bool valid = r >= 0; ++ const ushort prefix = simd_prefix_exclusive_sum(ushort(valid)); ++ const ushort cnt_s = simd_sum(ushort(valid)); ++ if (valid) { ++ roff[prefix] = uint(r) * n; ++ wk[prefix] = w; ++ } ++ if (tid == 0) roff[16] = cnt_s; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ const uint cnt = roff[16]; ++ const float gs = 1.0f / (1.0f + exp(-sgv[t])); ++ float bo[NPT], v[NPT]; ++ if (n == 2560) { ++ device const float* yst = ys + t * 2560 + tid; ++ device const float* vt = value + t * 2560 + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < NPT; ++i) { ++ const uint d = tid + i * 256; ++ float s = 0.0f; // mm_combine's sum, k in order ++ for (uint k = 0; k < cnt; ++k) s = fma(wk[k], YE[roff[k] + d], s); ++ bo[i] = fma(gs, yst[i * 256], s); ++ v[i] = vt[i * 256]; ++ } ++ } else { ++ for (uint i = 0, d = tid; d < n; ++i, d += 256) { ++ float s = 0.0f; // mm_combine's sum, k in order ++ for (uint k = 0; k < cnt; ++k) s = fma(wk[k], YE[roff[k] + d], s); ++ bo[i] = fma(gs, ys[t * n + d], s); ++ v[i] = value[t * n + d]; ++ } ++ } ++ for (uint c = 0; c < a.hc; ++c) { ++ const uint row = (t * a.hc + c) * n; ++ const float g = 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))); ++ float rb[NPT]; ++ if (n == 2560) { ++ device const float* rt = R + row + tid; ++#pragma clang loop unroll(full) ++ for (uint i = 0; i < NPT; ++i) rb[i] = fma(bo[i], g, rt[i * 256]); ++ } else { ++ for (uint i = 0, d = tid; d < n; ++i, d += 256) rb[i] = fma(bo[i], g, R[row + d]); ++ } ++ moe_gate_norm_row(R, rb, pkey, w_key, w_query, v, w_conv, norm_w, row, c, n, tid, a.eps, part, sg, lane); ++ } ++} ++ ++// the conv over [history | the window's rows] at (t, ch), SiLU'd ++static inline float ple_conv(device const float* hist, device const float* norm_w, device const half* kw, uint ch, ++ uint t, uint hcd) { ++ const half4 w = *(device const half4*)(kw + 4 * ch); ++ float acc = 0.0f; ++ if (t >= 9) { ++ device const float* nw = norm_w + (t - 9) * hcd + ch; ++ const uint s3 = 3 * hcd; ++ acc += float(w.x) * nw[0]; ++ acc += float(w.y) * nw[s3]; ++ acc += float(w.z) * nw[2 * s3]; ++ acc += float(w.w) * nw[3 * s3]; ++ } else { ++#pragma clang loop unroll(full) ++ for (int k = 0; k < 4; ++k) { ++ const int p = int(t) - (3 - k) * 3; // the row this tap reads: < 0 = the committed history ++ const float v = p >= 0 ? norm_w[uint(p) * hcd + ch] : hist[uint(9 + p) + 9 * ch]; ++ acc += float(w[k]) * v; ++ } ++ } ++ return acc / (1.0f + exp(-acc)); ++} ++ ++// grid (hc * n_embd / 256, T), 256 threads: the conv and the residual, R += gated + conv + kernel void ple_conv_add(device const float* hist [[buffer(0)]], device const float* norm_w [[buffer(1)]], + device const half* kw [[buffer(2)]], device const float* gated [[buffer(3)]], + device float* R [[buffer(4)]], constant PleArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const uint ch = tg.x * 256 + tid, t = tg.y, hcd = a.hc * a.n_embd; + if (ch >= hcd) return; +- float acc = 0.0f; +- for (int k = 0; k < 4; ++k) { +- const int p = int(t) - (3 - k) * 3; // the row this tap reads: < 0 = the committed history +- const float v = p >= 0 ? norm_w[uint(p) * hcd + ch] : hist[uint(9 + p) + 9 * ch]; +- acc += float(kw[k + 4 * ch]) * v; +- } +- const float conv = acc / (1.0f + exp(-acc)); ++ const float conv = ple_conv(hist, norm_w, kw, ch, t, hcd); + device float* r = R + t * hcd; + r[ch] = r[ch] + gated[t * hcd + ch] + conv; + } ++ ++// ple_conv_add's grid after ple_query_gate_norm / ple_moe_gate_norm(_c), which already added gated to R ++kernel void ple_conv_res(device const float* hist [[buffer(0)]], device const float* norm_w [[buffer(1)]], ++ device const half* kw [[buffer(2)]], device float* R [[buffer(3)]], ++ constant PleArgs& a [[buffer(4)]], uint3 tg [[threadgroup_position_in_grid]], ++ ushort tid [[thread_index_in_threadgroup]]) { ++ const uint ch = tg.x * 256 + tid, t = tg.y, hcd = a.hc * a.n_embd; ++ if (ch >= hcd) return; ++ R[t * hcd + ch] += ple_conv(hist, norm_w, kw, ch, t, hcd); ++} ++ ++// grid (hc * n_embd / 256, 1), 256 threads, after ple_conv_res: the history advanced in place past the window's T rows, ++// row r of the new history = window row T - 9 + r, or the old history's row r + T before the window. A thread turns ++// its own channel's 9 rows, read before any is written ++kernel void ple_hist_adv(device float* hist [[buffer(0)]], device const float* norm_w [[buffer(1)]], ++ constant PleArgs& a [[buffer(2)]], constant uint& T [[buffer(3)]], ++ uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { ++ const uint ch = tg.x * 256 + tid, hcd = a.hc * a.n_embd; ++ if (ch >= hcd) return; ++ device float* h = hist + 9 * ch; ++ float col[9]; ++#pragma clang loop unroll(full) ++ for (int r = 0; r < 9; ++r) { ++ const int src = int(T) - 9 + r; ++ col[r] = src >= 0 ? norm_w[uint(src) * hcd + ch] : h[uint(r) + T]; ++ } ++#pragma clang loop unroll(full) ++ for (int r = 0; r < 9; ++r) h[r] = col[r]; ++} +diff --git a/metal/src/engine/prompt_chain.inc b/metal/src/engine/prompt_chain.inc +index 573b460e..a0951240 100644 +--- a/metal/src/engine/prompt_chain.inc ++++ b/metal/src/engine/prompt_chain.inc +@@ -272,16 +272,14 @@ bool Window::pc_chunk(const int64_t* tokens, int nT, int64_t pos0, std::string& + dq(T(bl + "ffn_gate_shexp.weight"), b.Wshg); + dq(T(bl + "ffn_up_shexp.weight"), b.Wshu); + dq(T(bl + "ffn_down_shexp.weight"), b.Wshd); +- gemm(b.mixh, b.Wr, b.logits, nT, NEXP, N); +- const uint32_t nt = (uint32_t) nT; ++ gemm(b.mixh, b.Wr, b.logits, nT, n_exp_, N); ++ const uint32_t nt = (uint32_t) nT | ((uint32_t) n_exp_ << 16); + D(P_(b.lib, "route_topk"), {mtl::Arg::B(b.logits), mtl::Arg::B(b.ids), mtl::Arg::B(b.wts), mtl::Arg::V(nt)}, (nT + 7) / 8, 1, 256); + const int nent = nT * TOPK, nblk = (nent + 1023) / 1024; +- PcMmPlanArgs pa{(uint32_t) nent, TOPK, NEXP, (uint32_t) PC_TM, (uint32_t) nblk}; +- const uint32_t nclr = (uint32_t) b.maxrows; +- D(P_(b.lib, "mm_clear"), {mtl::Arg::B(b.stok), mtl::Arg::V(nclr)}, (b.maxrows + 255) / 256, 1, 256); ++ PcMmPlanArgs pa{(uint32_t) nent, TOPK, (uint32_t) n_exp_, (uint32_t) PC_TM, (uint32_t) nblk}; + D(P_(b.lib, "mm_hist"), {mtl::Arg::B(b.ids), mtl::Arg::B(b.bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024); + D(P_(b.lib, "mm_scan"), {mtl::Arg::B(b.bcnt), mtl::Arg::B(b.boff), mtl::Arg::B(b.eoff), mtl::Arg::B(b.tiles), mtl::Arg::B(b.meta), +- mtl::Arg::V(pa)}, 1, 1, NEXP); ++ mtl::Arg::V(pa), mtl::Arg::B(b.stok)}, 1, 1, (uint32_t) n_exp_); + 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), +@@ -290,8 +288,8 @@ bool Window::pc_chunk(const int64_t* tokens, int nT, int64_t pos0, std::string& + ma.d_stride = (uint64_t) ma.d_row * N; + const uint32_t ne = N; + mtl::Dispatch g; +- g.p = P_(b.lib, "mm_gather"); +- g.args = {mtl::Arg::B(b.mixed), mtl::Arg::B(b.stok), mtl::Arg::B(b.XS), mtl::Arg::V(ne)}; ++ g.p = P_(b.lib, "mm_gather_h"); ++ g.args = {mtl::Arg::B(b.mixh), mtl::Arg::B(b.stok), mtl::Arg::B(b.XS), mtl::Arg::V(ne)}; + g.tg[0] = 256; + g.indirect = &b.meta; + g.ind_off = 8 * 4; +diff --git a/metal/src/engine/spec_min_p.hpp b/metal/src/engine/spec_min_p.hpp +new file mode 100644 +index 00000000..e2f25dc4 +--- /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 9db74e25..97ae303e 100644 +--- a/metal/src/engine/stream.cpp ++++ b/metal/src/engine/stream.cpp +@@ -62,6 +62,11 @@ bool ExpertStream::open(const gguf::File& f, std::string& err) { + if (fd_ < 0) { err = "metal stream: cannot open " + f.path; return false; } + fcntl(fd_, F_NOCACHE, 1); + fcntl(fd_, F_RDAHEAD, 0); ++ const gguf::Tensor* g0 = f.find("blk.0.ffn_gate_exps.weight"); ++ if (g0 && g0->shape.size() >= 3 && g0->shape[2] >= 1 && g0->shape[2] <= NE) { ++ ne_ = (int) g0->shape[2]; ++ } ++ if (const char* mb = std::getenv("STRATA_METAL_MOCE_BONUS")) moce_bonus_ = (float) std::atof(mb); + lay_.resize(NL); + for (int l = 0; l < NL; ++l) { + const std::string b = "blk." + std::to_string(l) + "."; +@@ -69,7 +74,7 @@ bool ExpertStream::open(const gguf::File& f, std::string& err) { + const gguf::Tensor* u = f.find(b + "ffn_up_exps.weight"); + const gguf::Tensor* d = f.find(b + "ffn_down_exps.weight"); + if (!g || !u || !d) { err = "metal stream: layer " + std::to_string(l) + " has no routed experts"; return false; } +- lay_[(size_t) l] = {g->offset, g->bytes / NE, u->offset, u->bytes / NE, d->offset, d->bytes / NE, -1}; ++ lay_[(size_t) l] = {g->offset, g->bytes / (uint64_t) ne_, u->offset, u->bytes / (uint64_t) ne_, d->offset, d->bytes / (uint64_t) ne_, -1}; + } + // the expert-contiguous copy beside the GGUF, when it is this GGUF's (sizes in its header) and its first and + // last experts match the GGUF's bytes +@@ -88,16 +93,18 @@ bool ExpertStream::open(const gguf::File& f, std::string& err) { + uint32_t ver = 0, nl = 0, ne = 0; + uint64_t src = 0; + std::memcpy(&ver, h + 4, 4); std::memcpy(&nl, h + 8, 4); std::memcpy(&ne, h + 12, 4); std::memcpy(&src, h + 16, 8); +- ok = ok && ver == 1 && nl == (uint32_t) NL && ne == (uint32_t) NE && src == (uint64_t) st.st_size; ++ ok = ok && ver == 1 && nl == (uint32_t) NL && ne == (uint32_t) ne_ && src == (uint64_t) st.st_size; + for (int l = 0; ok && l < NL; ++l) { + const Lay& x = lay_[(size_t) l]; + ok = t[(size_t) l * 4 + 1] == x.gs && t[(size_t) l * 4 + 2] == x.us && t[(size_t) l * 4 + 3] == x.ds; + off[(size_t) l] = t[(size_t) l * 4]; + } +- if (ok) { // spot check: layer 0's first and layer 47's last expert, every slice ++ const Lay& last_lay = lay_[(size_t) (NL - 1)]; ++ const bool has_gguf_exps = (uint64_t) st.st_size >= last_lay.d + (uint64_t) ne_ * last_lay.ds; ++ if (ok && has_gguf_exps) { // spot check: layer 0's first and layer 47's last expert, every slice + for (const int l : {0, NL - 1}) { + const Lay& x = lay_[(size_t) l]; +- const int e = l == 0 ? 0 : NE - 1; ++ const int e = l == 0 ? 0 : ne_ - 1; + std::vector a(x.gs + x.us + x.ds), b(a.size()); + ok = ok && ::pread(cfd, a.data(), a.size(), (off_t) (off[(size_t) l] + (uint64_t) e * a.size())) == (ssize_t) a.size() && + ::pread(fd_, b.data(), x.gs, (off_t) (x.g + (uint64_t) e * x.gs)) == (ssize_t) x.gs && +@@ -148,6 +155,7 @@ bool ExpertStream::open(const gguf::File& f, std::string& err) { + } + nio_ = std::max(1, env_int("STRATA_METAL_IO_THREADS", 6)); + split_ = std::max(1, std::min(8, env_int("STRATA_METAL_READ_SPLIT", 1))); ++ nio_urgent_ = std::max(0, std::min(nio_ - 1, env_int("STRATA_METAL_IO_URGENT", 2))); + // The read-ahead's depth and where its experts wait in the LRU, measured on the M4 (notes-mac-ssd.md, 3 interleaved + // pairs each): Q2_0 (1.38 MB experts) top 10 at the LRU front 12.2 tok/s vs top 6 at the tail 11.3; the IQ3_XXS + // mix (1.75 MB, the SSD is the limit) top 6 at the tail 8.06 vs top 10 at the front 7.43. +@@ -158,6 +166,11 @@ bool ExpertStream::open(const gguf::File& f, std::string& err) { + pre_k2_ = env_int("STRATA_METAL_PREFETCH2", 0); + prefetch_ = pre_k_ > 0; + pre_mru_ = env_int("STRATA_METAL_PREFETCH_MRU", big ? 0 : 1) != 0; ++ // a predicted expert that is already resident keeps its LRU place (touching it cost 4% more misses in the ++ // closed-loop replay, audits/opus_strata_metal_opts_v1.md); =1 the old behaviour ++ pre_touch_ = env_int("STRATA_METAL_PREFETCH_TOUCH", 0) != 0; ++ // the read-ahead bookkeeping after the GPU is let go (after_layer); =0 all of it inside the handshake ++ split_turn_ = env_int("STRATA_METAL_SPLIT_TURN", 1) != 0; + // timing only (WRONG answers): every missing expert is pointed at slot 0, nothing is read + fake_ = env_int("STRATA_METAL_STREAM_FAKE", 0) != 0; + return true; +@@ -178,6 +191,52 @@ bool ExpertStream::read(void* dst, uint64_t len, uint64_t off) const { + bool ExpertStream::load_dense(gguf::File& f, mtl::Context& cx, const std::vector& skip, mtl::Buffer& out, + std::string& err) { + const double t0 = now_ms(); ++ uint64_t max_dense_end = 0; ++ uint64_t min_exp_off = UINT64_MAX; ++ size_t dense_cnt = 0; ++ for (const auto& kv : f.tensors) { ++ const gguf::Tensor& t = kv.second; ++ if (t.name.find("_exps") != std::string::npos) { ++ min_exp_off = std::min(min_exp_off, t.offset); ++ continue; ++ } ++ ++dense_cnt; ++ max_dense_end = std::max(max_dense_end, (t.offset + t.bytes + PAGE - 1) / PAGE * PAGE); ++ } ++ if (f.map != nullptr && max_dense_end > 0 && max_dense_end <= f.map_bytes && max_dense_end <= min_exp_off && ++ env_int("STRATA_METAL_MMAP_DENSE", 1) != 0) { ++ const size_t alloc_before = cx.allocated(); ++ uint64_t locked_bytes = 0, skipped_hc_bytes = 0, unpinned_embd_bytes = 0; ++ for (const auto& kv : f.tensors) { ++ const gguf::Tensor& t = kv.second; ++ if (t.name.find("_exps") != std::string::npos) continue; ++ const uintptr_t p0 = ((uintptr_t) f.map + t.offset) / PAGE * PAGE; ++ const uintptr_t p1 = ((uintptr_t) f.map + t.offset + t.bytes + PAGE - 1) / PAGE * PAGE; ++ const bool is_skip = std::find(skip.begin(), skip.end(), t.name) != skip.end(); ++ const bool is_embd = t.name == "token_embd.weight"; ++ if (is_skip || is_embd) { ++ if (is_skip) skipped_hc_bytes += t.bytes; ++ if (is_embd) unpinned_embd_bytes += t.bytes; ++ madvise((void*) p0, p1 - p0, MADV_DONTNEED); ++ } else if (env_int("STRATA_METAL_MLOCK_DENSE", 1) != 0) { ++ madvise((void*) p0, p1 - p0, MADV_WILLNEED); ++ if (mlock((const void*) p0, p1 - p0) == 0) locked_bytes += (p1 - p0); ++ else dense_unwired_ += p1 - p0; ++ } else { ++ dense_unwired_ += p1 - p0; ++ } ++ } ++ if (f.map_bytes > max_dense_end) { ++ madvise((uint8_t*) f.map + max_dense_end, f.map_bytes - (size_t) max_dense_end, MADV_DONTNEED); ++ } ++ out = cx.wrap(f.map, (size_t) max_dense_end, /*count_allocated=*/false); ++ (void) alloc_before; ++ std::fprintf(stderr, "strata metal stream: locked %.2f GB of active GPU dense tensors (unpinned %.2f GB CPU token_embd + %.2f GB skipped BF16 HC in %.2f GB file map)\n", ++ locked_bytes / 1e9, unpinned_embd_bytes / 1e9, skipped_hc_bytes / 1e9, max_dense_end / 1e9); ++ std::fprintf(stderr, "strata metal stream: %zu dense tensors wrapped zero-copy from file (%.2f GB clean file-backed map, %.1f ms)\n", ++ dense_cnt, max_dense_end / 1e9, now_ms() - t0); ++ return true; ++ } + std::vector ts; + uint64_t tot = 0; + for (auto& kv : f.tensors) { +@@ -250,19 +309,30 @@ void ExpertStream::lru_front(int s) { + if (k.tail < 0) k.tail = s; + } + +-int ExpertStream::victim(int ci) { ++// avoid_layer: a read-ahead for that layer queued while the GPU may already route it (after_layer runs beside the ++// GPU): none of its resident experts may go, or the GPU's routing (MoCE) and its resident-first expert pass would race ++// the eviction ++int ExpertStream::victim(int ci, int avoid_layer, bool allow_hot) { + Cls& k = cls_[(size_t) ci]; +- for (int s = k.tail; s >= 0; s = slots_[(size_t) s].prev) +- if (slots_[(size_t) s].state.load() == 0 && slots_[(size_t) s].pin != stamp_) { ++ const int n_passes = (in_prompt_ && allow_hot) ? 2 : 1; ++ for (int pass = 0; pass < n_passes; ++pass) { ++ for (int s = k.tail; s >= 0; s = slots_[(size_t) s].prev) ++ if (slots_[(size_t) s].state.load() == 0 && slots_[(size_t) s].pin != stamp_ && !slots_[(size_t) s].limbo && ++ (avoid_layer < 0 || slots_[(size_t) s].layer != avoid_layer) && ++ (gpu_layer_ < 0 || slots_[(size_t) s].layer != gpu_layer_) && ++ (!after_turn_ || slots_[(size_t) s].layer != gpu_layer_ + 1) && ++ !(split_turn_ && !in_prompt_ && slots_[(size_t) s].pre) && (!in_prompt_ || pass == 1 || !slots_[(size_t) s].hot)) { + Slot& x = slots_[(size_t) s]; + if (x.layer >= 0) { + const int key = x.layer * NE + x.e; + where_[(size_t) key] = -1; + smap_.as()[key] = 0xFFFFFFFFu; + } ++ if (in_prompt_ && multi_chunk_) x.hot = false; + x.layer = x.e = -1; + return s; + } ++ } + return -1; + } + +@@ -272,9 +342,10 @@ int ExpertStream::victim(int ci) { + void ExpertStream::read_part(int s, int l, int e, int part) { + const Lay& x = lay_[(size_t) l]; + if (cfd_ >= 0) { // the contiguous copy: the whole expert, piece part of split_ (the slot keeps gate | up | down) ++ const int sp = in_prompt_ ? 1 : split_; + const uint64_t len = x.gs + x.us + x.ds, src = coff_[(size_t) l] + (uint64_t) e * len; +- const uint64_t a = len * (uint64_t) part / (uint64_t) split_ / 4096 * 4096; +- const uint64_t b = part + 1 == split_ ? len : len * (uint64_t) (part + 1) / (uint64_t) split_ / 4096 * 4096; ++ const uint64_t a = len * (uint64_t) part / (uint64_t) sp / 4096 * 4096; ++ const uint64_t b = part + 1 == sp ? len : len * (uint64_t) (part + 1) / (uint64_t) sp / 4096 * 4096; + uint8_t* p = host_[(size_t) s] + a; + uint64_t n = b - a, o = src + a; + while (n > 0) { +@@ -285,7 +356,7 @@ void ExpertStream::read_part(int s, int l, int e, int part) { + } + p += r; o += (uint64_t) r; n -= (uint64_t) r; + } +- bytes_read_ += b - a; ++ (in_prompt_ ? p_bytes_read_ : bytes_read_) += b - a; + return; + } + const int k = part / split_, p = part % split_; +@@ -298,15 +369,17 @@ void ExpertStream::read_part(int s, int l, int e, int part) { + std::fprintf(stderr, "strata metal stream: reading expert %d of layer %d failed\n", e, l); + std::exit(6); + } +- bytes_read_ += b - a; ++ (in_prompt_ ? p_bytes_read_ : bytes_read_) += b - a; + } + +-void ExpertStream::io_main() { ++void ExpertStream::io_main(int idx) { + for (;;) { + Job j; + { + std::unique_lock g(mu_); +- cv_.wait(g, [&] { return stop_ || !urgent_q_.empty() || !pre_q_.empty(); }); ++ // the first nio_urgent_ threads are kept for the misses the GPU waits on while decoding ++ const auto urgent_only = [&] { return idx < nio_urgent_ && !in_prompt_; }; ++ cv_.wait(g, [&] { return stop_ || !urgent_q_.empty() || (!urgent_only() && !pre_q_.empty()); }); + if (stop_) return; + if (!urgent_q_.empty()) { j = urgent_q_.front(); urgent_q_.pop_front(); } + else { j = pre_q_.front(); pre_q_.pop_front(); } +@@ -318,10 +391,10 @@ void ExpertStream::io_main() { + std::lock_guard g(mu_); + Slot& x = slots_[(size_t) j.slot]; + if (--x.left == 0) { +- smap_.as()[j.layer * NE + j.e] = code_[(size_t) j.slot]; ++ if (!x.pre) smap_.as()[j.layer * NE + j.e] = code_[(size_t) j.slot]; + x.state = 0; + --outstanding_; +- ++reads_; ++ ++(in_prompt_ ? p_reads_ : reads_); + done = true; + } + } +@@ -335,7 +408,7 @@ void ExpertStream::submit(int s, int l, int e, bool urgent) { + x.e = (int16_t) e; + x.state = 1; + x.pre = !urgent; +- const int parts = cfd_ >= 0 ? split_ : 3 * split_; ++ const int parts = cfd_ >= 0 ? (in_prompt_ ? 1 : split_) : 3 * split_; + x.left = parts; + ++outstanding_; + where_[(size_t) (l * NE + e)] = s; +@@ -377,11 +450,11 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri + char magic[4]; + uint32_t h[5]; + if (std::fread(magic, 1, 4, f) == 4 && std::memcmp(magic, "STRP", 4) == 0 && std::fread(h, 4, 5, f) == 5 && +- h[1] == (uint32_t) NL && h[2] == (uint32_t) NE) { ++ h[1] == (uint32_t) NL && h[2] == (uint32_t) ne_) { + for (uint32_t i = 0; i < h[4]; ++i) { + uint16_t p[2]; + if (std::fread(p, 2, 2, f) != 2) break; +- if (p[0] < NL && p[1] < NE && !seen[(size_t) p[0] * NE + p[1]]) { ++ if (p[0] < NL && p[1] < (uint16_t) ne_ && !seen[(size_t) p[0] * NE + p[1]]) { + seen[(size_t) p[0] * NE + p[1]] = 1; + rank.push_back({p[0], p[1]}); + } +@@ -392,9 +465,11 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri + std::fprintf(stderr, "strata metal stream: no expert profile at %s: the arena starts layer-interleaved\n", profile.c_str()); + } + const size_t from_profile = rank.size(); +- for (int e = 0; e < NE; ++e) ++ for (int e = 0; e < ne_; ++e) + for (int l = 0; l < NL; ++l) + if (!seen[(size_t) l * NE + e]) rank.push_back({l, e}); ++ for (int l = 0; l < NL; ++l) layer_rank_[l].clear(); ++ for (const auto& p : rank) if (p.first >= 0 && p.first < NL) layer_rank_[p.first].push_back(p.second); + // the pool of each size class: at least min_slots (a prompt chunk's layer needs up to 512 at once), then the + // top-ranked pairs that fit in the budget + uint64_t used = 0; +@@ -403,7 +478,7 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri + for (const auto& p : rank) { + Cls& k = cls_[(size_t) lay_[(size_t) p.first].cls]; + const uint64_t cost = k.n < min_slots ? 0 : k.pages * PAGE; +- if (used + cost > budget) break; ++ if (used + cost > budget) continue; + used += cost; + ++k.n; + fill.push_back(p); +@@ -475,7 +550,7 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri + k.head = k.n ? k.base : -1; + k.tail = k.n ? k.base + k.n - 1 : -1; + } +- for (int i = 0; i < nio_; ++i) io_.emplace_back([this] { io_main(); }); ++ for (int i = 0; i < nio_; ++i) io_.emplace_back([this, i] { io_main(i); }); + // the fill: the ranked pairs in steps (best first; each step in file order) on the I/O threads; before each step + // grow() asks whether macOS still has the memory for it (arena_grow); the best-ranked ends at the LRU front + uint64_t step = 1536ull << 20; +@@ -524,7 +599,17 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri + ++nretired; + } + } +- // the LRU order = the rank order (best at the front) ++ // the LRU order = the rank order (best at the front); mark top (k.n - (ne_ + 32)) STRP slots hot ++ std::vector cls_seen(cls_.size(), 0); ++ for (size_t i = 0; i < done; ++i) { ++ const int s = where_[(size_t) (fill[i].first * NE + fill[i].second)]; ++ if (s >= 0) { ++ const int ci = cls_of(s); ++ const int hot_cap = std::max(0, cls_[(size_t) ci].n - (ne_ + 32)); ++ slots_[(size_t) s].hot = (cls_seen[(size_t) ci] < hot_cap); ++ ++cls_seen[(size_t) ci]; ++ } ++ } + for (size_t i = done; i-- > 0;) lru_front(where_[(size_t) (fill[i].first * NE + fill[i].second)]); + } + if (stopped) +@@ -533,6 +618,8 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri + arena_live_ = arena_bytes_ - retired_bytes; + fill_s_ = (now_ms() - t0) / 1e3; + fill_bytes_ = bytes_read_.load(); ++ bytes_read_ = 0; ++ reads_ = 0; + uint64_t pairs = 0; + for (const auto& k : cls_) pairs += (uint64_t) k.n; + std::fprintf(stderr, "strata metal stream: expert arena %.2f GB, %llu slots in %zu size classes, %zu experts read from the " +@@ -571,23 +658,47 @@ bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::stri + // the experts a predicted routing names, most likely first (by the best logit over the window's tokens) + void ExpertStream::ahead(int ln, int T, const float* pred, int K) { + K = std::min(K, 32); ++ static const bool unpert_all = env_int("STRATA_METAL_PREFETCH_UNPERT", 1) != 0; ++ const uint32_t* sm = smap_.as() + (size_t) ln * NE; + std::pair pw[8 * 32]; + int np = 0; ++ static const int k_draft_max = env_int("STRATA_METAL_PREFETCH_DRAFT", 1); + for (int t = 0; t < T; ++t) { ++ const int Kt = (t == 0) ? K : std::min(K, k_draft_max); ++ if (Kt <= 0) continue; + const float* lg = pred + (size_t) t * NE; ++ auto eff_lg = [&](int e) -> float { ++ return lg[e] + (moce_bonus_ > 0.0f && sm[e] != 0xFFFFFFFFu ? moce_bonus_ : 0.0f); ++ }; ++ const int k_unpert = unpert_all ? Kt : std::min(3, Kt); + int top[32]; + int nt = 0; +- for (int e = 0; e < NE; ++e) { +- if (nt < K) { top[nt++] = e; continue; } +- int m = 0; +- for (int j = 1; j < K; ++j) if (lg[top[j]] < lg[top[m]]) m = j; +- if (lg[e] > lg[top[m]]) top[m] = e; ++ for (int r = 0; r < k_unpert; ++r) { ++ int best_e = -1; ++ float best_v = -INFINITY; ++ for (int e = 0; e < ne_; ++e) { ++ bool seen_e = false; ++ for (int j = 0; j < nt; ++j) if (top[j] == e) { seen_e = true; break; } ++ if (!seen_e && lg[e] > best_v) { best_v = lg[e]; best_e = e; } ++ } ++ if (best_e >= 0) top[nt++] = best_e; ++ } ++ for (int e = 0; e < ne_; ++e) { ++ if (k_unpert >= Kt) break; ++ bool in_top = false; ++ for (int j = 0; j < k_unpert; ++j) if (top[j] == e) { in_top = true; break; } ++ if (in_top) continue; ++ if (nt < Kt) { top[nt++] = e; continue; } ++ int m = k_unpert; ++ for (int j = k_unpert + 1; j < Kt; ++j) if (eff_lg(top[j]) < eff_lg(top[m])) m = j; ++ if (eff_lg(e) > eff_lg(top[m])) top[m] = e; + } + for (int j = 0; j < nt; ++j) { ++ const float prio = (j < k_unpert) ? (lg[top[j]] + moce_bonus_ + 1e3f) : eff_lg(top[j]); + int q = 0; + while (q < np && pw[q].second != top[j]) ++q; +- if (q == np) pw[np++] = {lg[top[j]], top[j]}; +- else pw[q].first = std::max(pw[q].first, lg[top[j]]); ++ if (q == np) pw[np++] = {prio, top[j]}; ++ else pw[q].first = std::max(pw[q].first, prio); + } + } + std::sort(pw, pw + np, [](const auto& a, const auto& b) { return a.first > b.first; }); +@@ -597,8 +708,8 @@ void ExpertStream::ahead(int ln, int T, const float* pred, int K) { + for (int j = 0; j < np; ++j) { + const int key = ln * NE + pw[j].second; + const int s0 = where_[(size_t) key]; +- if (s0 >= 0) { lru_front(s0); continue; } +- const int s = victim(lay_[(size_t) ln].cls); ++ if (s0 >= 0) { if (pre_touch_) lru_front(s0); continue; } ++ const int s = victim(lay_[(size_t) ln].cls, split_turn_ ? ln : -1); + if (s < 0) break; + submit(s, ln, pw[j].second, false); + slots_[(size_t) s].pin = stamp_; +@@ -611,17 +722,32 @@ void ExpertStream::ahead(int ln, int T, const float* pred, int K) { + if (any) cv_.notify_all(); + } + ++// Do not write future-layer residency while its GPU router can be running. The ++// existing shared-event handshake publishes completed preads before GPU resume. ++void ExpertStream::publish_ready(int lo, int hi) { ++ for (int l = lo; l <= hi; ++l) ++ for (int e = 0; e < ne_; ++e) { ++ const int key = l * NE + e, s = where_[(size_t) key]; ++ if (s < 0) continue; ++ const Slot& x = slots_[(size_t) s]; ++ if (!x.limbo && x.state.load(std::memory_order_acquire) == 0) ++ smap_.as()[key] = code_[(size_t) s]; ++ } ++} ++ + void ExpertStream::at_layer(int l, int T, int topk, const int32_t* ids, const float* pred, const float* pred2) { + const double t0 = now_ms(); + const bool prompt = T > 8; // a prompt chunk (the GEMM path): nearly every expert of the layer is used + if (l == 0) { if (prompt) ++p_chunks_; else ++windows_; } ++ gpu_layer_ = l; ++ after_turn_ = false; + ++stamp_; // the slots this layer uses are pinned until the next handshake + bool used[NE] = {}; + int want[NE]; + int nw = 0; + for (int i = 0; i < T * topk; ++i) { + const int e = ids[i]; +- if (e < 0 || e >= NE || used[e]) continue; ++ if (e < 0 || e >= ne_ || used[e]) continue; + used[e] = true; + want[nw++] = e; + } +@@ -632,16 +758,45 @@ void ExpertStream::at_layer(int l, int T, int topk, const int32_t* ids, const fl + return; + } + std::vector wait; ++ const bool split = split_turn_ && !prompt; + { + std::lock_guard g(mu_); + // read-aheads of this layer (or older) that were not started and are not wanted: dropped (the slot goes free) +- { ++ if (!split) { + std::vector drop; + for (const Job& j : pre_q_) { + if (j.layer > l || slots_[(size_t) j.slot].state.load() != 1) continue; + if (!(j.layer == l && used[j.e]) && std::find(drop.begin(), drop.end(), j.slot) == drop.end()) drop.push_back(j.slot); + } + for (int s : drop) cancel_slot(s); ++ } else { // the unwanted read-aheads of this layer: queued ones dropped now (the misses below get the SSD), all ++ // parked until the next turn's after_layer ++ retire_unused_pre(l, l, &limbo_new_, used); ++ } ++ if (!split) { ++ std::unique_lock g2(mu_, std::adopt_lock); ++ done_cv_.wait(g2, [&] { ++ for (int e = 0; e < ne_; ++e) { ++ if (used[e]) continue; ++ const int s = where_[(size_t) (l * NE + e)]; ++ if (s >= 0 && slots_[(size_t) s].pre && slots_[(size_t) s].state.load() != 0) return false; ++ } ++ return true; ++ }); ++ g2.release(); ++ // Evict any unchosen speculative read-ahead of this layer that finished so where_, smap_, and LRU order stay 100% deterministic ++ for (int e = 0; e < ne_; ++e) { ++ if (used[e]) continue; ++ const int s = where_[(size_t) (l * NE + e)]; ++ if (s >= 0 && slots_[(size_t) s].pre && slots_[(size_t) s].state.load() == 0) { ++ Slot& x = slots_[(size_t) s]; ++ where_[(size_t) (l * NE + e)] = -1; ++ smap_.as()[l * NE + e] = 0xFFFFFFFFu; ++ x.layer = x.e = -1; ++ x.pre = false; ++ lru_back(s); ++ } ++ } + } + for (int j = 0; j < nw; ++j) // pass 1: pin what is resident (or on its way), so no miss below evicts it + if (where_[(size_t) (l * NE + want[j])] >= 0) slots_[(size_t) where_[(size_t) (l * NE + want[j])]].pin = stamp_; +@@ -661,19 +816,54 @@ void ExpertStream::at_layer(int l, int T, int topk, const int32_t* ids, const fl + promote_slot(s); // a queued read-ahead becomes urgent + wait.push_back(s); + } +- if (!(prompt && x.pre)) lru_front(s); // a prompt's streamed experts stay at the LRU tail ++ if (!prompt || (!multi_chunk_ ? !x.pre : x.hot)) lru_front(s); // a prompt's streamed experts stay at the LRU tail ++ else if (multi_chunk_) lru_back(s); + x.pre = false; ++ if (st == 0) smap_.as()[key] = code_[(size_t) s]; + continue; + } + ++(prompt ? p_misses_ : misses_); +- s = victim(lay_[(size_t) l].cls); ++ s = victim(lay_[(size_t) l].cls, -1, true); + if (s < 0) { std::fprintf(stderr, "strata metal stream: no free slot (layer %d)\n", l); std::exit(6); } + submit(s, l, e, true); + slots_[(size_t) s].pin = stamp_; + if (!prompt) lru_front(s); ++ else if (multi_chunk_) lru_back(s); + wait.push_back(s); + } ++ // a prompt's lookahead: layer l + 1's experts as its router ranks them on this layer's FFN input, the most ++ // tokens first, queued behind this layer's misses now (not after the wait), so the SSD never runs dry ++ // (not with a MoCE prompt, 0 < STRATA_METAL_PROMPT_MOCE < 100: its routing reads residency, as the static read-ahead) ++ const float pm = std::getenv("STRATA_METAL_PROMPT_MOCE") ? (float) std::atof(std::getenv("STRATA_METAL_PROMPT_MOCE")) : 0.0f; ++ for (int distance = 1; prompt && distance <= 2; ++distance) { ++ const int32_t* predicted = distance == 1 ? p_pred_ : p_pred2_; ++ if (!predicted || l + distance >= NL || (pm > 0.0f && pm < 100.0f)) continue; ++ const int ln = l + distance; ++ int cnt[NE] = {}; ++ for (int i = 0; i < T * p_pred_k_; ++i) { ++ const int e = predicted[i]; ++ if (e >= 0 && e < ne_) ++cnt[e]; ++ } ++ int order[NE], no = 0; ++ static const int pmin = std::max(1, env_int("STRATA_METAL_PROMPT_PRED_MIN", 1)); ++ for (int e = 0; e < ne_; ++e) if (cnt[e] >= pmin) order[no++] = e; ++ std::stable_sort(order, order + no, [&](int a, int b) { return cnt[a] > cnt[b]; }); ++ int queued = 0; ++ const int cap = distance == 1 ? NE : std::clamp(env_int("STRATA_METAL_PROMPT_PRED2_MAX", 64), 0, NE); ++ for (int j = 0; j < no && queued < cap; ++j) { ++ const int e = order[j]; ++ if (where_[(size_t) (ln * NE + e)] >= 0) continue; ++ const int s = victim(lay_[(size_t) ln].cls, l + 1, !multi_chunk_ || distance == 1); ++ if (s < 0) break; ++ submit(s, ln, e, false); ++ ++queued; ++ slots_[(size_t) s].pin = stamp_; ++ lru_back(s); ++ ++p_streamed_; ++ } ++ } + } ++ if (prompt && (p_pred_ || p_pred2_)) cv_.notify_all(); + if (!wait.empty()) { + cv_.notify_all(); + const double tw = now_ms(); +@@ -684,20 +874,31 @@ void ExpertStream::at_layer(int l, int T, int topk, const int32_t* ids, const fl + }); + (prompt ? p_wait_ms_ : wait_ms_) += now_ms() - tw; + } ++ const float cur_pm = std::getenv("STRATA_METAL_PROMPT_MOCE") ? (float) std::atof(std::getenv("STRATA_METAL_PROMPT_MOCE")) : 0.0f; ++ if (prompt && cur_pm > 0.0f && env_int("STRATA_METAL_PUBLISH_PREFETCH", 1) != 0) { ++ std::lock_guard g(mu_); ++ publish_ready(l + 1, std::min(l + 2, NL - 1)); ++ } + if (prompt) { + // a chunk uses nearly all of the next layer's experts: every one not in the arena is streamed in behind this + // layer's compute, into the LRU tail (the decode-hot experts stay) +- if (l + 1 < NL) { ++ const float p_moce_env = std::getenv("STRATA_METAL_PROMPT_MOCE") ? (float) std::atof(std::getenv("STRATA_METAL_PROMPT_MOCE")) : 0.0f; ++ if (l + 1 < NL && !p_pred_ && (p_moce_env <= 0.0f || p_moce_env >= 100.0f)) { + bool any = false; + std::lock_guard g(mu_); + const int ln = l + 1; +- for (int e = 0; e < NE; ++e) { ++ const int max_pre = std::getenv("STRATA_METAL_PROMPT_MAX_PRE") ? std::atoi(std::getenv("STRATA_METAL_PROMPT_MAX_PRE")) : (p_moce_env < 0.0f ? 96 : (ln >= 40 ? 128 : 256)); ++ int queued = 0; ++ for (int e : layer_rank_[ln]) { ++ if (queued >= max_pre) break; + if (where_[(size_t) (ln * NE + e)] >= 0) continue; +- const int s = victim(lay_[(size_t) ln].cls); ++ const int s = victim(lay_[(size_t) ln].cls, -1, false); + if (s < 0) break; + submit(s, ln, e, false); + slots_[(size_t) s].pin = stamp_; ++ lru_back(s); + ++p_streamed_; ++ ++queued; + any = true; + } + if (any) cv_.notify_all(); +@@ -705,15 +906,87 @@ void ExpertStream::at_layer(int l, int T, int topk, const int32_t* ids, const fl + (prompt ? p_host_ms_ : host_ms_) += now_ms() - t0; + return; + } ++ if (split) { // the rest after the GPU is let go (after_layer) ++ host_ms_ += now_ms() - t0; ++ return; ++ } + // read-ahead: layer l + 1 (each token's top pre_k_ of the predicted router logits), then layer l + 2 (top pre_k2_) + if (prefetch_ && pred && l + 1 < NL) ahead(l + 1, T, pred, pre_k_); + if (pre_k2_ > 0 && pred2 && l + 2 < NL) ahead(l + 2, T, pred2, pre_k2_); + host_ms_ += now_ms() - t0; + } + ++// Read-aheads of layers lo..hi that the routing did not claim: queued pieces dropped, the slots parked in `into` (still ++// mapped, no victims) until release_limbo(), whatever their I/O state now, so which slot is ++// free when depends on the order of the turns only, not on the SSD's timing (MoCE routing reads residency). ++void ExpertStream::retire_unused_pre(int lo, int hi, std::vector* into, const bool* used) { ++ for (int key = lo * NE; key < (hi + 1) * NE; ++key) { ++ const int s = where_[(size_t) key]; ++ if (s < 0 || (used && used[key % NE])) continue; ++ Slot& x = slots_[(size_t) s]; ++ if (!x.pre || x.limbo) continue; ++ if (x.state.load() == 1) { // not started: its pieces dropped, nothing will write the slot ++ for (auto it = pre_q_.begin(); it != pre_q_.end();) it = it->slot == s ? pre_q_.erase(it) : it + 1; ++ x.state = 0; ++ x.left = 0; ++ --outstanding_; ++ ++pre_cancel_; ++ } ++ x.limbo = true; ++ into->push_back(s); ++ } ++} ++ ++void ExpertStream::release_limbo(std::unique_lock& g) { ++ if (limbo_.empty()) return; ++ done_cv_.wait(g, [&] { // a read that is still running (rare: it had a whole layer of GPU time) ++ for (int s : limbo_) if (slots_[(size_t) s].state.load() != 0) return false; ++ return true; ++ }); ++ std::sort(limbo_.begin(), limbo_.end()); ++ for (int s : limbo_) { ++ Slot& x = slots_[(size_t) s]; ++ if (x.layer >= 0 && where_[(size_t) (x.layer * NE + x.e)] == s) { ++ const int key = x.layer * NE + x.e; ++ where_[(size_t) key] = -1; ++ smap_.as()[key] = 0xFFFFFFFFu; ++ } ++ x.layer = x.e = -1; ++ x.pre = x.limbo = false; ++ lru_back(s); ++ } ++ limbo_.clear(); ++} ++ ++void ExpertStream::after_layer(int l, int T, const float* pred, const float* pred2) { ++ if (!split_turn_ || fake_) return; ++ const double t0 = now_ms(); ++ after_turn_ = true; ++ { ++ std::unique_lock g(mu_); ++ release_limbo(g); // the last turn's (they had a whole layer of GPU time to finish reading) ++ limbo_.swap(limbo_new_); // this turn's, retired in at_layer ++ } ++ if (prefetch_ && pred && l + 1 < NL) ahead(l + 1, T, pred, pre_k_); ++ if (pre_k2_ > 0 && pred2 && l + 2 < NL) ahead(l + 2, T, pred2, pre_k2_); ++ host_after_ms_ += now_ms() - t0; ++} ++ + // A region of `bytes` at the end of one arena buffer, for a prompt's scratch buffers: its slots are emptied and held + // (state 3) while a prompt runs (hold_region), and are ordinary slots again for decoding. + uint8_t* ExpertStream::region(uint64_t bytes) { ++ if (env_int("STRATA_METAL_SEPARATE_PROMPT_BUF", 1) != 0) { ++ const uint64_t sz = (bytes + PAGE - 1) / PAGE * PAGE; ++ void* m = mmap(nullptr, (size_t) sz, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0); ++ if (m == MAP_FAILED) return nullptr; ++ region_slots_.clear(); ++ region_ = (uint8_t*) m; ++ region_bytes_ = 0; // fresh zero pages from MAP_ANON; no arena slots to memset ++ sep_region_bytes_ = sz; ++ std::fprintf(stderr, "strata metal stream: prompt scratch %.2f GB mapped in separate anonymous mmap\n", ++ sz / 1e9); ++ return region_; ++ } + for (int c = (int) chunks_.size() - 1; c >= 0; --c) { + if (chunks_[(size_t) c].bytes < bytes + (64ull << 20)) continue; + uint8_t* end = chunks_[(size_t) c].as() + chunks_[(size_t) c].bytes; +@@ -730,14 +1003,14 @@ uint8_t* ExpertStream::region(uint64_t bytes) { + std::vector stolen(cls_.size(), 0); + for (int s : region_slots_) ++stolen[(size_t) cls_of(s)]; + bool tight = false; +- for (size_t i = 0; i < cls_.size(); ++i) tight = tight || (cls_[i].n > 0 && cls_[i].n - stolen[i] < NE); ++ for (size_t i = 0; i < cls_.size(); ++i) tight = tight || (cls_[i].n > 0 && cls_[i].n - stolen[i] < ne_); + if (tight) { + const uint64_t sz = (bytes + PAGE - 1) / PAGE * PAGE; + void* m = mmap(nullptr, (size_t) sz, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0); + if (m == MAP_FAILED) return nullptr; + region_slots_.clear(); + region_ = (uint8_t*) m; +- region_bytes_ = 0; // fresh zero pages; nothing in the arena to empty ++ region_bytes_ = 0; sep_region_bytes_ = sz; // fresh zero pages; unmapped in release_region() + std::fprintf(stderr, "strata metal stream: the arena is small: the prompt's %.2f GB of buffers are their own " + "allocation, not a region of its slots\n", bytes / 1e9); + return region_; +@@ -752,9 +1025,22 @@ uint8_t* ExpertStream::region(uint64_t bytes) { + return nullptr; + } + +-void ExpertStream::hold_region(bool hold) { ++void ExpertStream::hold_region(bool hold, bool multi_chunk) { + std::unique_lock g(mu_); ++ in_prompt_ = hold; ++ multi_chunk_ = hold && multi_chunk; ++ gpu_layer_ = -1; ++ after_turn_ = false; + if (hold) { ++ // Refresh hot flags from current LRU order (protecting top k.n - (ne_ + 32) slots in every prompt turn) ++ for (size_t ci = 0; ci < cls_.size(); ++ci) { ++ const int hot_cap = std::max(0, cls_[ci].n - (ci == 0 ? cls_[ci].n / 4 : (2 * ne_ + 32))); ++ int rank_idx = 0; ++ for (int s = cls_[ci].head; s >= 0; s = slots_[(size_t) s].next) { ++ slots_[(size_t) s].hot = (rank_idx < hot_cap); ++ ++rank_idx; ++ } ++ } + // nothing may be reading into those slots: the queued read-aheads are dropped, the running ones waited for + std::vector drop; + for (const Job& j : pre_q_) +@@ -762,6 +1048,10 @@ void ExpertStream::hold_region(bool hold) { + drop.push_back(j.slot); + for (int s : drop) cancel_slot(s); + done_cv_.wait(g, [&] { return outstanding_ == 0; }); ++ limbo_.insert(limbo_.end(), limbo_new_.begin(), limbo_new_.end()); ++ limbo_new_.clear(); ++ retire_unused_pre(0, NL - 1, &limbo_); // decode's last read-aheads: a prompt's victims may take them ++ release_limbo(g); + for (int s : region_slots_) { + Slot& x = slots_[(size_t) s]; + if (x.layer >= 0) { +@@ -773,43 +1063,99 @@ void ExpertStream::hold_region(bool hold) { + x.pre = false; + x.state = 3; + } +- if (region_) std::memset(region_, 0, (size_t) region_bytes_); ++ if (region_ && region_bytes_ > 0) std::memset(region_, 0, (size_t) region_bytes_); ++ // Kick off layer 0's hottest non-resident experts immediately so they overlap with prompt setup + layer 0 GDN ++ const float p_moce_env = std::getenv("STRATA_METAL_PROMPT_MOCE") ? (float) std::atof(std::getenv("STRATA_METAL_PROMPT_MOCE")) : 0.0f; ++ if (!fake_ && (p_moce_env <= 0.0f || p_moce_env >= 100.0f)) { ++ bool any = false; ++ // with the chunk's lookahead (STRATA_METAL_PROMPT_PRED) the whole profile rank: what layer 0 does not use and ++ // the SSD has not started on by its handshake is dropped there ++ const int max_pre0 = std::getenv("STRATA_METAL_PROMPT_MAX_PRE") ? std::atoi(std::getenv("STRATA_METAL_PROMPT_MAX_PRE")) ++ : (p_moce_env < 0.0f ? 96 : (env_int("STRATA_METAL_PROMPT_PRED", 1) != 0 ? NE : 256)); ++ const int layers = std::clamp(env_int("STRATA_METAL_PROMPT_START_LAYERS", 1), 1, 4); ++ for (int l = 0; l < layers; ++l) { ++ int queued = 0; ++ const int cap = l == 0 ? max_pre0 : std::min(max_pre0, 32); ++ for (int e : layer_rank_[l]) { ++ if (queued >= cap) break; ++ if (where_[(size_t) (l * NE + e)] >= 0) continue; ++ const int s = victim(lay_[(size_t) l].cls, 0, false); ++ if (s < 0) break; ++ submit(s, l, e, false); ++ lru_back(s); ++ ++p_streamed_; ++ ++queued; ++ any = true; ++ } ++ } ++ if (any) cv_.notify_all(); ++ } + } else { ++ std::vector drop; ++ for (const Job& j : pre_q_) ++ if (slots_[(size_t) j.slot].state.load() == 1 && std::find(drop.begin(), drop.end(), j.slot) == drop.end()) ++ drop.push_back(j.slot); ++ for (int s : drop) cancel_slot(s); ++ done_cv_.wait(g, [&] { return outstanding_ == 0; }); + for (int s : region_slots_) { + slots_[(size_t) s].state = 0; + if (!retired_[(size_t) s]) lru_back(s); // a retired slot's memory serves the prompt only + } ++ if (split_turn_) { // the prompt's unclaimed stream-ins: decode's victim() passes read-aheads by ++ limbo_.insert(limbo_.end(), limbo_new_.begin(), limbo_new_.end()); ++ limbo_new_.clear(); ++ retire_unused_pre(0, NL - 1, &limbo_); ++ release_limbo(g); ++ } + } + } + ++bool ExpertStream::reuse_region(int mode) { ++ if (mode != 1 || sep_region_bytes_ == 0 || region_ == nullptr) return false; ++ return true; ++} ++ ++uint64_t ExpertStream::release_region() { ++ std::unique_lock g(mu_); ++ if (sep_region_bytes_ > 0 && region_ != nullptr) { ++ const uint64_t freed = sep_region_bytes_; ++ munmap(region_, (size_t) sep_region_bytes_); ++ region_ = nullptr; ++ sep_region_bytes_ = 0; ++ return freed; ++ } ++ return 0; ++} ++ + void ExpertStream::reset_stats() { + lookups_ = hits_ = pre_hits_ = late_ = misses_ = pre_issued_ = pre_cancel_ = windows_ = 0; + p_lookups_ = p_hits_ = p_pre_ = p_misses_ = p_streamed_ = p_chunks_ = 0; + p_wait_ms_ = p_host_ms_ = 0; +- bytes_read_ = 0; +- reads_ = 0; +- wait_ms_ = host_ms_ = 0; ++ bytes_read_ = p_bytes_read_ = 0; ++ reads_ = p_reads_ = 0; ++ wait_ms_ = host_ms_ = host_after_ms_ = 0; + } + + void ExpertStream::report(FILE* f) const { + if (!ready()) return; + if (p_chunks_) + std::fprintf(f, "expert stream prompts %llu chunks: %llu distinct lookups (%.1f per layer), in the arena %.4f, streamed " +- "ahead %.4f, missed %.4f; %llu streamed in; host %.1f ms per chunk, of it waiting %.1f ms\n", ++ "ahead %.4f, missed %.4f; %llu streamed in (%.1f MB/chunk, %llu reads); host %.1f ms per chunk, of it waiting %.1f ms\n", + (unsigned long long) p_chunks_, (unsigned long long) p_lookups_, p_lookups_ / 48.0 / p_chunks_, + p_hits_ / std::max(1.0, (double) p_lookups_), p_pre_ / std::max(1.0, (double) p_lookups_), +- p_misses_ / std::max(1.0, (double) p_lookups_), (unsigned long long) p_streamed_, p_host_ms_ / p_chunks_, +- p_wait_ms_ / p_chunks_); ++ p_misses_ / std::max(1.0, (double) p_lookups_), (unsigned long long) p_streamed_, ++ p_bytes_read_.load() / 1e6 / p_chunks_, (unsigned long long) p_reads_.load(), ++ p_host_ms_ / p_chunks_, p_wait_ms_ / p_chunks_); + const double w = std::max(1, windows_); + std::fprintf(f, + "expert stream %llu windows: %llu distinct lookups (%.1f per window), hits %.4f, read-ahead hits " + "%.4f (%llu of them late), misses %.4f (%.1f per window); read-aheads %llu (%.1f per window, %llu dropped); %.1f MB " +- "read per window (%llu reads); host %.2f ms per window, of it waiting %.2f ms\n", ++ "read per window (%llu reads); host %.2f ms per window, of it waiting %.2f ms (+ %.2f ms after the GPU went on)\n", + (unsigned long long) windows_, (unsigned long long) lookups_, lookups_ / w, + hits_ / std::max(1.0, (double) lookups_), pre_hits_ / std::max(1.0, (double) lookups_), + (unsigned long long) late_, misses_ / std::max(1.0, (double) lookups_), misses_ / w, + (unsigned long long) pre_issued_, pre_issued_ / w, (unsigned long long) pre_cancel_, +- (bytes_read_.load() - 0) / 1e6 / w, (unsigned long long) reads_.load(), host_ms_ / w, wait_ms_ / w); ++ (bytes_read_.load() - 0) / 1e6 / w, (unsigned long long) reads_.load(), host_ms_ / w, wait_ms_ / w, host_after_ms_ / w); + } + + } // namespace strata::metal_engine +diff --git a/metal/src/engine/stream.h b/metal/src/engine/stream.h +index 2806938a..6cbac41a 100644 +--- a/metal/src/engine/stream.h ++++ b/metal/src/engine/stream.h +@@ -53,6 +53,7 @@ public: + const std::function& grow = {}); + uint64_t shrink_to() const { return shrink_to_; } + uint64_t arena_used_bytes() const { return arena_live_; } // the arena's bytes in use (its retired slots left out) ++ uint64_t dense_unwired_bytes() const { return dense_unwired_; } // load_dense's mapped dense bytes left unwired + bool ready() const { return !chunks_.empty(); } + // the arena's buffers (at most 4, each within the device's maxBufferLength): chunk(i) for i >= count is chunk 0 + const mtl::Buffer& chunk(int i) const { return chunks_[(size_t) i < chunks_.size() ? (size_t) i : 0]; } +@@ -63,32 +64,52 @@ public: + // computed from this layer's FFN input (nullptr: no read-ahead). + // pred2: layer l + 2's, from the same input (STRATA_METAL_PREFETCH2 = its top-k, 0 = off) + void at_layer(int l, int T, int topk, const int32_t* ids, const float* pred, const float* pred2 = nullptr); ++ // a prompt chunk's lookahead: T x topk ids of layer l + 1's router on layer l's FFN input (nullptr: the static ++ // expert-profile rank instead); ids2 optionally names layer l + 2; read by at_layer ++ void set_prompt_pred(const int32_t* ids, int topk, const int32_t* ids2 = nullptr) { ++ p_pred_ = ids; p_pred_k_ = topk; p_pred2_ = ids2; ++ } ++ // the rest of a decode layer's turn, after the GPU has been let go (run_graph_seg's after_cut): the read-aheads ++ // the routing did not use are retired, the next layer's are queued. Only when split_turn() (else at_layer did it). ++ void after_layer(int l, int T, const float* pred, const float* pred2 = nullptr); ++ bool split_turn() const { return split_turn_; } + int ahead2() const { return pre_k2_; } + // a prompt's scratch buffers inside the arena: region() places them (once), hold_region() empties and holds its + // slots while a prompt runs and gives them back after + uint8_t* region(uint64_t bytes); +- void hold_region(bool hold); ++ void hold_region(bool hold, bool multi_chunk = false); ++ uint64_t release_region(); ++ bool reuse_region(int mode); ++ size_t region_stolen_slots() const { return region_slots_.size(); } + void report(FILE* f) const; + void reset_stats(); + + private: +- uint64_t shrink_to_ = 0; ++ uint64_t shrink_to_ = 0, dense_unwired_ = 0; + struct Lay { uint64_t g, gs, u, us, d, ds; int cls; }; +- struct Slot { int16_t layer = -1; int16_t e = -1; int prev = -1, next = -1; bool pre = false; int left = 0; uint64_t pin = 0; std::atomic state{0}; }; // 0 idle, 1 queued, 2 reading ++ struct Slot { int16_t layer = -1; int16_t e = -1; int prev = -1, next = -1; bool pre = false, hot = false, limbo = false; int left = 0; uint64_t pin = 0; std::atomic state{0}; }; // 0 idle, 1 queued, 2 reading + struct Cls { uint64_t pages = 0; int n = 0; int base = 0; int head = -1, tail = -1; }; // slot range [base, base+n) + struct Job { int slot; int layer, e, part; }; // part: slice * split_ + piece + void lru_unlink(int s); + void lru_front(int s); + void lru_back(int s); + int cls_of(int s) const; +- int victim(int cls); ++ int victim(int cls, int avoid_layer = -1, bool allow_hot = true); + void read_part(int slot, int l, int e, int part); + void cancel_slot(int s); + void promote_slot(int s); +- void io_main(); ++ + void ahead(int ln, int T, const float* pred, int K); + void submit(int slot, int l, int e, bool urgent); + void wait_slot(int s); ++ // the unclaimed read-aheads of layers lo..hi (not in `used`) parked in `into` (mu_ held) ++ void retire_unused_pre(int lo, int hi, std::vector* into, const bool* used = nullptr); ++ void release_limbo(std::unique_lock& g); ++ std::vector limbo_; // retired read-aheads: not victims until released one turn later (deterministic LRU) ++ std::vector limbo_new_; // ... retired in this turn's at_layer (released at the next turn's after_layer) ++ int nio_urgent_ = 0; // I/O threads that take urgent reads only (STRATA_METAL_IO_URGENT) ++ void io_main(int idx); ++ void publish_ready(int lo, int hi); // mu_ held, GPU at the layer handshake + + int fd_ = -1; + // .experts (metal/tools/stream_experts.py): each expert's gate, up and down rows together, read with one +@@ -100,6 +121,12 @@ private: + std::vector slots_; + std::vector retired_; // slots left unused when the stepped fill stopped at macOS's floor (never in an LRU) + std::vector where_; // (l * 512 + e) -> slot, -1 ++ std::vector layer_rank_[48]; // per-layer experts ordered hottest-first from expert-profile.bin ++ const int32_t* p_pred_ = nullptr; ++ const int32_t* p_pred2_ = nullptr; ++ int gpu_layer_ = -1; ++ bool after_turn_ = false; ++ int p_pred_k_ = 0; + std::vector chunks_; + std::vector region_slots_; + uint8_t* region_ = nullptr; +@@ -116,15 +143,18 @@ private: + std::deque urgent_q_, pre_q_; + bool stop_ = false; + int outstanding_ = 0; +- int nio_ = 6, pre_k_ = 10, pre_k2_ = 0, split_ = 2; +- bool prefetch_ = true, fake_ = false, pre_mru_ = false; ++ int ne_ = 512, nio_ = 6, pre_k_ = 10, pre_k2_ = 0, split_ = 2; ++ float moce_bonus_ = 0.0f; ++ bool in_prompt_ = false, multi_chunk_ = false; ++ uint64_t sep_region_bytes_ = 0; ++ bool prefetch_ = true, fake_ = false, pre_mru_ = false, pre_touch_ = false, split_turn_ = true; + // stats + uint64_t stamp_ = 1; // the handshake count: slots pinned with the current one are in use on the GPU + uint64_t p_lookups_ = 0, p_hits_ = 0, p_pre_ = 0, p_misses_ = 0, p_streamed_ = 0, p_chunks_ = 0; + double p_wait_ms_ = 0, p_host_ms_ = 0; + uint64_t lookups_ = 0, hits_ = 0, pre_hits_ = 0, late_ = 0, misses_ = 0, pre_issued_ = 0, pre_cancel_ = 0, windows_ = 0; +- std::atomic bytes_read_{0}, reads_{0}; +- double wait_ms_ = 0, host_ms_ = 0, fill_s_ = 0; ++ std::atomic bytes_read_{0}, reads_{0}, p_bytes_read_{0}, p_reads_{0}; ++ double wait_ms_ = 0, host_ms_ = 0, host_after_ms_ = 0, fill_s_ = 0; + uint64_t fill_bytes_ = 0; + }; + +diff --git a/metal/src/engine/suffix_window.hpp b/metal/src/engine/suffix_window.hpp +new file mode 100644 +index 00000000..edc7cd61 +--- /dev/null ++++ b/metal/src/engine/suffix_window.hpp +@@ -0,0 +1,30 @@ ++#pragma once ++ ++#include "strata/spec/suffix_drafter.hpp" ++ ++#include ++#include ++ ++namespace strata::spec { ++ ++inline bool suffix_stop(int32_t token) { ++ return token == 248046 || token == 248044; ++} ++ ++inline int suffix_limit(int requested, int remaining, int64_t pos, int64_t max_cells, bool unrestricted_qsa = true) { ++ const int block_drafts = unrestricted_qsa ? 3 : (3 - int(pos & 3)); ++ const int64_t room = std::max(0, max_cells - pos - 3); ++ const int64_t drafts = std::min({requested, remaining - 1, block_drafts, room}); ++ return int(std::max(0, drafts)); ++} ++ ++inline int suffix_keep(const int32_t* input, const int32_t* output, int nd, bool stop_on_eos = true) { ++ int keep = 1; ++ for (int i = 0; i < nd; ++i) { ++ if ((stop_on_eos && suffix_stop(output[i])) || output[i] != input[i + 1]) break; ++ ++keep; ++ } ++ return keep; ++} ++ ++} // namespace strata::spec +diff --git a/metal/src/engine/window.cpp b/metal/src/engine/window.cpp +index 9a6d1d32..074b64fb 100644 +--- a/metal/src/engine/window.cpp ++++ b/metal/src/engine/window.cpp +@@ -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 + #include + #include + #include + #include + #include + #include ++#include ++#include + #include ++#include + #include + #include + #include +@@ -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, +- P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_INTS = P_TOKENT + P_MAXE, P_IGU = 4, P_IDN = 8; ++ P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_SPLIT = P_TOKENT + P_MAXE, P_GSLOT = P_SPLIT + 20, P_INTS = P_GSLOT + P_MAXE, ++ P_IGU = 4, P_IDN = 8, P_IGU_R = P_SPLIT + 4, P_IDN_R = P_SPLIT + 8, P_IGU_M = P_SPLIT + 12, P_IDN_M = P_SPLIT + 16; + 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 +79,21 @@ bool metal_ctab() { + return on; + } + ++// the drafter's attention sink (STRATA_METAL_MTP_SINK, default 64; 0: the sliding window alone) ++int64_t mtp_sink_cap() { ++ static const int64_t cap = [] { ++ const char* e = std::getenv("STRATA_METAL_MTP_SINK"); ++ return e ? std::max(0, std::min(1024, std::atoll(e))) : 64; ++ }(); ++ return cap; ++} ++ ++std::array tile_env(const char* name) { ++ std::array t{0, 0, 0}; ++ if (const char* e = std::getenv(name)) std::sscanf(e, "%d,%d,%d", &t[0], &t[1], &t[2]); ++ return t; ++} ++ + int xkind(int t) { + switch (t) { + case ref::Q2_0: return 1; +@@ -81,6 +103,78 @@ int xkind(int t) { + } + } + ++// Q4_0 with a signed-d 2-pass Lloyd-Max scale per block (codes -8..7) ++void quant_q4_0_lloyd(const float* x, size_t n, uint8_t* out) { ++ const size_t nb = n / 32; ++ for (size_t b = 0; b < nb; ++b) { ++ const float* xb = x + b * 32; ++ float amax = 0.0f; ++ for (int j = 0; j < 32; ++j) amax = std::max(amax, std::fabs(xb[j])); ++ uint16_t d16 = 0; ++ uint8_t qs[16] = {}; ++ if (amax > 0.0f) { ++ // Signed-d 2-pass Lloyd-Max for Q4_0: codes in [-8..+7] ++ float best_score = -1.0f, best_qw = 0.0f, best_qq = 1.0f; ++ int best_q[32] = {}; ++ const float seeds[6] = {amax / 7.5f, amax / 8.0f, amax / 8.5f, -amax / 7.5f, -amax / 8.0f, -amax / 8.5f}; ++ for (int s = 0; s < 6; ++s) { ++ float d = seeds[s]; ++ for (int it = 0; it < 2; ++it) { ++ if (std::fabs(d) <= 1e-12f) break; ++ const float inv = 1.0f / d; ++ float qw = 0.0f, qq = 0.0f; ++ int qc[32]; ++ for (int j = 0; j < 32; ++j) { ++ int q = (int) std::lrintf(xb[j] * inv); ++ q = std::max(-8, std::min(7, q)); ++ qc[j] = q; ++ qw += (float) q * xb[j]; ++ qq += (float) (q * q); ++ } ++ if (qq <= 0.0f) break; ++ const float sc = (qw * qw) / qq; ++ if (sc > best_score) { ++ best_score = sc; best_qw = qw; best_qq = qq; ++ std::memcpy(best_q, qc, sizeof(qc)); ++ } ++ d = qw / qq; ++ } ++ } ++ d16 = ref::f32_to_f16(best_qw / best_qq); ++ for (int j = 0; j < 16; ++j) { ++ const uint8_t lo = (uint8_t) (best_q[j] + 8); ++ const uint8_t hi = (uint8_t) (best_q[j + 16] + 8); ++ qs[j] = (lo & 0x0Fu) | (uint8_t) ((hi & 0x0Fu) << 4); ++ } ++ } ++ std::memcpy(out + b * 18, &d16, 2); ++ std::memcpy(out + b * 18 + 2, qs, 16); ++ } ++} ++ ++// ggml's block_q5_0: d, qh, qs ++void quant_q5_0(const float* x, size_t n, uint8_t* out) { ++ for (size_t b = 0; b < n / 32; ++b) { ++ const float* xb = x + b * 32; ++ float amax = 0.0f, mx = 0.0f; ++ for (int j = 0; j < 32; ++j) ++ if (std::fabs(xb[j]) > amax) { amax = std::fabs(xb[j]); mx = xb[j]; } ++ const float d = mx / -16.0f, id = d != 0.0f ? 1.0f / d : 0.0f; ++ const uint16_t d16 = ref::f32_to_f16(d); ++ uint32_t qh = 0; ++ uint8_t qs[16]; ++ for (int j = 0; j < 16; ++j) { ++ const int x0 = std::min(31, (int) (xb[j] * id + 16.5f)), x1 = std::min(31, (int) (xb[j + 16] * id + 16.5f)); ++ qs[j] = (uint8_t) ((x0 & 0xF) | ((x1 & 0xF) << 4)); ++ qh |= (uint32_t) ((x0 & 0x10) >> 4) << j; ++ qh |= (uint32_t) ((x1 & 0x10) >> 4) << (j + 16); ++ } ++ std::memcpy(out + b * 22, &d16, 2); ++ std::memcpy(out + b * 22 + 2, &qh, 4); ++ std::memcpy(out + b * 22 + 6, qs, 16); ++ } ++} ++ + // a prepared activation: 4 planes of MAXT rows of n, then the per-16 sums + struct Act { + mtl::Buffer b; +@@ -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; + ++// Anti-slop v4: a sampler shaper over the head logits, outside , and ``` fences only. ++// Each position's state (the last two tokens, the region, where the line stands) follows from the one before it, so ++// it is kept per position: a draft row, a rejected suffix, a prompt read from a cached prefix all start from the ++// state of the position before them. Tables: the Qwen3.8-Flash-Next vocab (models/*/tokenizer/vocab.json). ++#include "slop4.h" ++ ++namespace slop4 { ++ ++// The ends of a turn or a thought must not gain from a penalty: with a likely token pushed down, the rest of the row ++// scales up by 1 / (1 - b), b the share taken away, and <|im_end|> is often next in line (a reply cut off mid-sentence). ++// Their logits go down by log(1 / (1 - b)) so their probability stays what it was. lse is the row's log-sum-exp; a ++// row whose penalized tokens all sit 12+ below its max loses under 6e-6 each and is skipped (no lse needed). ++ ++inline bool needs_hold(const Pen* p, int n, const float* lg, float row_max) { ++ for (int i = 0; i < n; ++i) if (lg[p[i].id] > row_max - 12.0f) return true; ++ return false; ++} ++ ++inline int hold(const Pen* p, int n, const float* lg, double lse, Pen* out) { ++ double b = 0.0; ++ for (int i = 0; i < n; ++i) b += std::exp((double) lg[p[i].id] - lse) * (1.0 - std::exp(-(double) p[i].v)); ++ if (b <= 0.0) return 0; ++ const float v = (float) -std::log1p(-std::min(b, 0.999999)); ++ int m = 0; ++ for (int32_t id : kHold) out[m++] = {id, v}; ++ return m; ++} ++ ++inline Knobs knobs() { ++ auto f = [](const char* e, float d) { const char* v = std::getenv(e); return v ? (float) std::atof(v) : d; }; ++ return {f("STRATA_ANTI_SLOP_PENALTY", 9.0f), f("STRATA_ANTI_SLOP_HONEST", 6.0f), f("STRATA_ANTI_SLOP_MD", 4.0f), ++ f("STRATA_ANTI_SLOP_COLON", 3.0f)}; ++} ++ ++} // namespace slop4 ++ + class Window { + public: + bool init(const strata::core::ModelGeometry& g, strata::core::SessionState& ss, const std::string& shard1, int max_t, +@@ -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; ++ // anti-slop v4: the prompt paths do not pass finish, so the prompt's tokens advance the per-position states here ++ template void slop_scan(const I* tokens, int64_t n, int64_t pos0) { ++ slop4::State st = slop_at(pos0 - 1); ++ for (int64_t i = 0; i < n; ++i) slop_put(pos0 + i) = st = slop4::step(st, (int32_t) tokens[i]); ++ } + int64_t vocab() const { return n_vocab_; } + bool ready() const { return cx_ != nullptr; } + bool streaming() const { return stream_; } +@@ -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); ++ // a disk session: the running state of checkpoint_save, ss.ple_prev and the positional cells [0, n_cells) of every ++ // QSA layer (and the MTP layer's), read and written straight into the buffers (scr_state.hpp's file layout) ++ bool session_save(const std::string& path, const int32_t* ids, int64_t n_cells, std::string& err); ++ bool session_load(const std::string& path, std::vector& ids, int64_t& n_cells, std::string& err); ++ // SCR: every QSA layer's cells [c_src, c_src + n) moved to c_dst, keys and pooled keys re-rotated by c_dst - c_src ++ bool scr_relocate(int64_t c_src, int64_t c_dst, int64_t n, std::string& err); + 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) +@@ -176,7 +317,18 @@ 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); ++ int ntok, size_t y_off = 0, const mtl::Buffer* wb = nullptr, int tile_ntok = 0); ++ std::array gemv_tile(int ty, int n_in, int n_out, int ntok); ++ bool gemv_cat(std::vector& ds, const gguf::Tensor& ta, const gguf::Tensor& tb, const Act& x, ++ const mtl::Buffer& ya, const mtl::Buffer& yb, size_t y_stride, int ntok, size_t y_off); ++ // a deferred MoE combine: the buffer it writes (nullptr: none) and its token count ++ const mtl::Buffer* comb_bo_ = nullptr; ++ int comb_ntok_ = 0; ++ void comb_push(std::vector& ds, const mtl::Buffer& bo, int ntok); ++ void comb_defer(std::vector& ds, const mtl::Buffer& bo, int ntok); ++ void comb_flush(std::vector& ds); ++ void gr_norm(std::vector& ds, mtl::Buffer& R, const mtl::Buffer& bo, const mtl::Buffer& inj, ++ const mtl::Buffer& wb, size_t w_off, bool apply, int ntok); + 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); +@@ -191,7 +343,8 @@ private: + // token's embedding into memb_ row 0 (the next step's input) + void mtp_pick(std::vector& ds, int j, bool embed); + void gr_read_m(std::vector& ds, const std::string& pre, mtl::Buffer& Rb, const mtl::Buffer& bo, +- const mtl::Buffer& inj_in, const mtl::Buffer& inj_out, bool apply, const Act& out, bool final_, int ntok); ++ const mtl::Buffer& inj_in, const mtl::Buffer& inj_out, bool apply, const Act& out, bool final_, int ntok, ++ bool already_normed = false); + const gguf::Tensor& TM(const std::string& name) const { + const gguf::Tensor* x = fm_.find(name); + if (!x) { std::fprintf(stderr, "strata metal: the MTP model has no tensor %s\n", name.c_str()); std::exit(2); } +@@ -214,10 +367,11 @@ private: + int64_t max_cells_ = 0; + int64_t n_vocab_ = 0; + float eps_ = 1e-6f; ++ double rope_base_ = 1e7; + // 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_, +- emb_, pkey_, pval_, qkvL_, alL_, beL_, zL_, ydump_, keyn_, qn_, gated_, normW_, pleHist_; ++ emb_, pkey_, pval_, qkvL_, alL_, beL_, zL_, ydump_, keyn_, normW_, pleHist_; + Act XA_, XF_, XY_, XQ_, XH_, XE_; + int cur_inj_ = 0; + // state +@@ -254,6 +408,8 @@ private: + struct Cached { + std::vector ds; + std::vector cuts; // SSD streaming: the segment ends (one per layer, before its experts) ++ std::vector sigs; // where each layer's handshake starts (after its router): the shared ++ // expert runs between sigs[s] and cuts[s], beside the host's turn + std::vector> qsa; // (dispatch, kind): 0 QsaArgs n_tok = T, 1 the merge (MAXT), 2 the scores + }; + std::map wcache_; +@@ -262,8 +418,15 @@ private: + std::unique_ptr st_; + bool stream_ = false, st_pred_ = false; + std::vector* cuts_ = nullptr; ++ std::vector* sigs_ = nullptr; + mtl::Buffer predL_, predL2_; // 48 x MAXT x 512: layer l + 1's (l + 2's) router logits from layer l's FFN input +- int seg_T_ = 0; ++ mtl::Buffer predW_; // the routers in Q8_0 for those guesses (empty: the GGUF's own) ++ mtl::Buffer routW_; // the routers as BF16 when the F32 file holds BF16 values (exact: half the bytes) ++ std::vector routT_; ++ std::vector predT_; ++ int seg_T_ = 0, n_exp_ = NEXP; ++ uint32_t moce_pad_ = 0; ++ int pg_lib_ = -1, pg_lgdn_ = -1, pg_lqsa_ = -1; + bool stream_arena(std::string& err); + double run_ds(Cached& c); + // prompts on the matrix units (prompt_chunk below) +@@ -276,9 +439,12 @@ private: + public: + bool prompt_gemm_on(int64_t n) const; + bool prompt_gemm_init(std::string& err); ++ void free_prompt_gemm(); + bool prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::string& err); + int prompt_chunk_rows() const; +- void prompt_region(bool hold) { if (stream_ && st_ && st_->ready()) st_->hold_region(hold); } ++ void prompt_region(bool hold, bool multi_chunk = false); ++ uint64_t prompt_gemm_need_bytes() const; ++ void reset_stream_stats() { if (st_) st_->reset_stats(); for (int t = 0; t <= MAXT; ++t) { win_n_[t] = 0; win_ms_[t] = win_gpu_[t] = 0; } } + void stream_report(FILE* f) const { + if (!st_) return; + st_->report(f); +@@ -289,17 +455,136 @@ public: + } + int win_n_[MAXT + 1] = {}; + double win_ms_[MAXT + 1] = {}, win_gpu_[MAXT + 1] = {}; ++ void dump_moe_layer(const char* pfx, int l, bool dump_in, bool dump_out); ++ std::thread* ple_join_ = nullptr; // a prompt chunk's PLE gather, joined at layer 0's handshake + void at_cut(int s) { ++ if (s == 0 && ple_join_) { ++ if (ple_join_->joinable()) ple_join_->join(); ++ ple_join_ = nullptr; ++ } ++ if (cut_ids_ && pg_) { ++ if (const char* pfx = std::getenv("STRATA_METAL_DUMP_MOE")) { ++ if (s > 0) dump_moe_layer(pfx, s - 1, false, true); ++ dump_moe_layer(pfx, s, true, false); ++ } ++ } + st_->at_layer(s, seg_T_, TOPK, cut_ids_ ? cut_ids_ : ids_.as(), st_pred_ && !cut_ids_ ? predL_.as() + (size_t) s * MAXT * 512 : nullptr, + predL2_.ptr && !cut_ids_ ? predL2_.as() + (size_t) s * MAXT * 512 : nullptr); ++ if (cut_ids_) st_->after_layer(s, seg_T_, nullptr); // the prompt chain has no after_cut: its bookkeeping here ++ } ++ // the pipelines of every window size up to max_t_ and of the MTP layer, made at load (window_init) rather than when ++ // a window first needs them (a T=3 window's first build compiled its #3 specializations in the middle of a decode: ++ // 282 ms; the first request's T=1 window and MTP prefill / draft stalled the same way). Past 2051, so the QSA ++ // selection's pipelines too (a superset of a short context's: layer() makes qsa_score / qsa_topk either way). ++ int prewarmed_t_ = 0; // window sizes 1..prewarmed_t_ made ++ bool mtp_prewarmed_ = false; ++ void prewarm_pipelines(bool overlap) { ++ auto* rec = rec_; ++ rec_ = nullptr; // the QSA patch list belongs to the window being built ++ std::vector c, s; ++ for (int t = prewarmed_t_ + 1; t <= max_t_; ++t) { ++ std::vector ds; ++ c.clear(); s.clear(); ++ cuts_ = stream_ ? &c : nullptr; ++ sigs_ = stream_ && overlap ? &s : nullptr; ++ nb_ = false; ++ cur_inj_ = 0; ++ fold(ds, t, 0); ++ if (ple_on_) { ++ xprep(ds, emb_, XE_, N); ++ gemv(ds, T("blk.1.ple_key.weight"), XE_, pkey_, 10240, t); ++ gemv(ds, T("blk.1.ple_value.weight"), XE_, pval_, N, t); ++ ple_key_norm(ds, t); ++ } ++ layer(ds, 0, t, 4096); ++ if (ple_on_) ple_after_layer0(ds, t); ++ for (int l = 1; l < 48; ++l) layer(ds, l, t, 4096); ++ gr_read(ds, "output_hc", true, XA_, true, t); ++ gemv(ds, T("output.weight"), XA_, lm_, (size_t) n_vocab_, t); ++ } ++ prewarmed_t_ = std::max(prewarmed_t_, max_t_); ++ P(lsmp_, "argmax_part", {}, "argmax_part"); ++ P(lsmp_, "argmax_final", {}, "argmax_final"); ++ cuts_ = sigs_ = nullptr; ++ if (mtp_on_ && !mtp_prewarmed_) { ++ std::vector ds; ++ nb_ = false; ++ for (int k = 1; k <= MAXT; ++k) mtp_encode(ds, mR_, 0, k, nullptr, 0, false); // mtp_prefill's chunks ++ mtp_encode(ds, mR_, 0, 1, nullptr, 0, true); // a draft step ++ if (chain_) mtp_pick(ds, 0, true); ++ mtp_prewarmed_ = true; ++ } ++ rec_ = rec; ++ cur_inj_ = 0; ++ nb_ = false; ++ comb_bo_ = nullptr; ++ comb_ntok_ = 0; ++ } ++ // window_init's, else the first window's (STRATA_METAL_LAZY_ARENA: the MoE dispatches bind the arena, made then) ++ void prewarm() { ++ static const bool on = !(std::getenv("STRATA_METAL_PREWARM") && std::atoi(std::getenv("STRATA_METAL_PREWARM")) == 0); ++ static const bool overlap = !(std::getenv("STRATA_METAL_SHEXP_OVERLAP") && std::atoi(std::getenv("STRATA_METAL_SHEXP_OVERLAP")) == 0); ++ if (!on || (stream_ && !st_->ready())) return; ++ if (prewarmed_t_ < max_t_ || (mtp_on_ && !mtp_prewarmed_)) prewarm_pipelines(overlap); ++ } ++ // the PLE block on the GPU (msl/ple.metal) around layer 0: the key norm reads only the key projection, so it runs ++ // beside layer 0's first read, not behind layer 0 ++ struct PleArgs { uint32_t n_embd, hc; float eps; }; ++ void ple_key_norm(std::vector& ds, int Tn) { ++ const PleArgs pa{N, HC, strata::kernels::NG_RMS_EPS}; ++ mtl::Dispatch d; ++ d.p = P(lple_, "ple_gnorm", {}, "ple_gnorm"); ++ d.args = {mtl::Arg::B(pkey_), mtl::Arg::B(W_, T("blk.1.ple_norm_key.weight").offset), mtl::Arg::B(keyn_), mtl::Arg::V(pa)}; ++ d.grid[0] = HC; d.grid[1] = Tn; d.tg[0] = 256; ++ push(ds, d); ++ nb_ = true; ++ } ++ // after layer 0: its FFN write applied to R, the query norm, the gate, the normalized rows, the conv. Layer 0's MoE ++ // combine deferred: ple_comb_query_gate_norm forms it as it reads R (STRATA_METAL_PLE_COMB=0: gr_norm_comb, whose ++ // xn nothing reads, then ple_query_gate_norm) ++ void ple_after_layer0(std::vector& ds, int Tn) { ++ static const bool fuse = !(std::getenv("STRATA_METAL_PLE_COMB") && std::atoi(std::getenv("STRATA_METAL_PLE_COMB")) == 0); ++ const PleArgs pa{N, HC, strata::kernels::NG_RMS_EPS}; ++ mtl::Dispatch g, cv; ++ g.args = {mtl::Arg::B(keyn_), mtl::Arg::B(R_), mtl::Arg::B(W_, T("blk.1.ple_norm_query.weight").offset), ++ mtl::Arg::B(pval_), mtl::Arg::B(W_, T("blk.1.ple_norm_conv.weight").offset), mtl::Arg::B(normW_), ++ mtl::Arg::V(pa)}; ++ if (fuse && comb_bo_ == &bo_ && comb_ntok_ == Tn) { ++ comb_bo_ = nullptr; ++ const uint32_t tk = TOPK; ++ g.p = P(lple_, "ple_comb_query_gate_norm", {}, "ple_comb_query_gate_norm"); ++ g.args.insert(g.args.end(), {mtl::Arg::B(ye_), mtl::Arg::B(inj_[cur_inj_]), mtl::Arg::B(plan_), mtl::Arg::B(ysh_), ++ mtl::Arg::B(sg_), mtl::Arg::V(tk)}); ++ } else { ++ gr_norm(ds, R_, bo_, inj_[cur_inj_], W_, T("blk.1.hc_attn_norm.weight").offset, true, Tn); ++ g.p = P(lple_, "ple_query_gate_norm", {}, "ple_query_gate_norm"); ++ } ++ g.grid[0] = HC; g.grid[1] = Tn; g.tg[0] = 256; ++ push(ds, g); ++ cv.p = P(lple_, "ple_conv_res", {}, "ple_conv_res"); ++ cv.args = {mtl::Arg::B(pleHist_), mtl::Arg::B(normW_), kw_arg(), mtl::Arg::B(R_), mtl::Arg::V(pa)}; ++ cv.grid[0] = HC * N / 256; cv.grid[1] = Tn; cv.tg[0] = 256; ++ push(ds, cv); ++ } ++ void after_cut(int s) { // decode only (the prompt chain runs without it) ++ st_->after_layer(s, seg_T_, st_pred_ ? predL_.as() + (size_t) s * MAXT * 512 : nullptr, ++ predL2_.ptr ? predL2_.as() + (size_t) s * MAXT * 512 : nullptr); + } + private: + std::vector ple_hist_; // NG_HIST x hc_dim normalized rows, row-fastest + bool ple_on_ = false; + int last_T_ = 0; ++ int64_t last_pos0_ = 0; ++ std::vector tail_bak_; + double last_gpu_ms_ = 0; + int last_n_disp_ = 0; + int32_t last_tokens_[MAXT] = {}; ++ std::vector slop_; // anti-slop v4: the state after the token at each position (State{} = unknown) ++ slop4::State slop_at(int64_t p) const { return p >= 0 && (size_t) p < slop_.size() ? slop_[(size_t) p] : slop4::State{}; } ++ slop4::State& slop_put(int64_t p) { ++ if ((size_t) p >= slop_.size()) slop_.resize(((size_t) p + 4096) & ~(size_t) 4095); ++ return slop_[(size_t) p]; ++ } + 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 +606,26 @@ public: + ~Window(); + void prof_add(bool draft, double ms) { Prof& x = draft ? dprof_ : cprof_; ++x.n; x.all += ms; } + private: +- // the Q8_0 copies of the hyper-connection matrices (the default; STRATA_HC_Q8=0: the BF16 originals) +- bool hc_q8_ = false; ++ // the Q8_0 (or Q4_0 when STRATA_HC_Q4=1) copies of the hyper-connection matrices ++ bool hc_q8_ = false, hc_q4_ = false; // hc_q4_: some HC matrices below 8 bits (q8type_) ++ int hc_q4g_ = 0, hc_q5g_ = 0; // their groups: bit 0 hc_attn, bit 1 hc_ffn + mtl::Buffer Q8_; + std::map q8off_; ++ std::map q8type_; + // MTP + bool mtp_on_ = false; + int lmtp_ = -1; ++ // the drafter's cells of committed tokens: the run mtp_cell0_..mtp_cell1_-1 (the sliding window's) and the attention ++ // sink 0..mtp_sink_-1 (at most mtp_sink_cap() cells, read beside the window once it has moved past them) ++ int64_t mtp_cell0_ = 0, mtp_cell1_ = 0, mtp_sink_ = 0; ++ void mtp_wrote(int64_t c0, int64_t n) { // cells c0..c0+n-1 now hold committed tokens; the ones past them are stale ++ mtp_cell1_ = std::min(mtp_cell1_, c0); ++ mtp_sink_ = std::min(mtp_sink_, c0); ++ if (n <= 0) return; ++ if (c0 != mtp_cell1_ || c0 < mtp_cell0_) mtp_cell0_ = c0; ++ mtp_cell1_ = c0 + n; ++ if (c0 == mtp_sink_) mtp_sink_ = std::min(mtp_sink_cap(), c0 + n); ++ } + 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 +679,16 @@ private: + // per step with the host's pick and embedding (as before) + int lchain_ = -1; + bool chain_ = false; +- mtl::Buffer dout_, dprob_, posv_, dmap_; ++ mtl::Buffer dout_, dprob_, posv_, dmap_, dslop_, dinverse_; ++ bool draft_slop_ = false; ++ struct DraftSlopArgs { ++ slop4::State state; ++ slop4::Knobs knobs; ++ uint32_t nv, j, vocab; ++ }; ++ DraftSlopArgs draft_slop_args_{}; ++ void mtp_slop(std::vector& ds, int j); ++ void mtp_slop_host(slop4::State state); + }; + + namespace { +@@ -411,10 +718,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; +- if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_ff != NFF || g.n_expert != NEXP || g.n_layers != 48) { ++ if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_ff != NFF || g.n_expert < 1 || g.n_expert > NEXP || g.n_layers != 48) { + err = "metal window: the geometry differs from Qwen3.8-Flash-Next's"; + return false; + } ++ n_exp_ = (int) g.n_expert; ++ if (const char* mb = std::getenv("STRATA_METAL_MOCE_BONUS")) { ++ const float bonus = (float) std::atof(mb); ++ if (bonus > 0.0f) { ++ const float sc = std::exp(bonus); ++ std::memcpy(&moce_pad_, &sc, 4); ++ } ++ } ++ if (!(std::getenv("STRATA_METAL_KEEP_HOST_WEIGHTS") && std::string(std::getenv("STRATA_METAL_KEEP_HOST_WEIGHTS")) == "1")) { ++ strata::gpu::purge_device_allocations(1u << 20); ++ } + 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 +741,9 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta + return false; + } + if (!f_.open(shard1, err)) return false; ++ if (const gguf::Tensor* ge0 = f_.find("blk.0.ffn_gate_exps.weight")) { ++ if (ge0->shape.size() >= 3 && ge0->shape[2] >= 1 && ge0->shape[2] <= NEXP) n_exp_ = (int) ge0->shape[2]; ++ } + cx_ = std::make_unique(); + stream_ = ExpertStream::wanted(f_, *cx_); + if (stream_) { +@@ -460,13 +781,14 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta + // the GDN layers' window inputs, kept per layer so commit() can fold the accepted prefix into the state + qkvL_ = A(4 * (size_t) 36 * MAXT * 10240); alL_ = A(4 * (size_t) 36 * MAXT * 48); beL_ = A(4 * (size_t) 36 * MAXT * 48); + zL_ = A(4 * (size_t) 36 * MAXT * 6144); ydump_ = A(4 * MAXT * 6144); +- keyn_ = A(4 * MAXT * 10240); qn_ = A(4 * MAXT * 10240); gated_ = A(4 * MAXT * 10240); normW_ = A(4 * MAXT * 10240); ++ keyn_ = A(4 * MAXT * 10240); normW_ = A(4 * MAXT * 10240); + mrope_ = A(12 * (size_t) (MAXC + 16)); + pleHist_ = A(4 * 9 * 10240); + auto act = [&](int n) { Act a; a.n = n; a.b = A(4 * ((size_t) 4 * MAXT * n + (size_t) MAXT * n / 16)); return a; }; + 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; ++ rope_base_ = base; + 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 +817,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"); ++ // STRATA_HC_Q4 / STRATA_HC_Q5 = 1 (every hc_*_down/_up) | attn | ffn: those in Q4_0 / Q5_0 (Q4 wins where both) ++ auto groups = [](const char* name) { ++ const char* e = std::getenv(name); ++ const std::string v = e ? e : ""; ++ return v == "attn" ? 1 : v == "ffn" ? 2 : (!v.empty() && std::atoi(v.c_str()) != 0) ? 3 : 0; ++ }; ++ hc_q4g_ = hc_q8_ ? groups("STRATA_HC_Q4") : 0; ++ hc_q5g_ = hc_q8_ ? groups("STRATA_HC_Q5") : 0; ++ hc_q4_ = hc_q4g_ || hc_q5g_; + } + if (hc_q8_) { + const double tq = now_ms(); + size_t tot = 0; + std::vector names; ++ auto low_type = [&](const std::string& nm) { // Q4_0, Q5_0 or Q8_0 for this HC matrix ++ if (!hc_q4_) return (int) ref::Q8_0; ++ if (nm.find("output_hc") != std::string::npos || nm.find("_inject") != std::string::npos) return (int) ref::Q8_0; ++ if (nm.find("_down.weight") == std::string::npos && nm.find("_up.weight") == std::string::npos) return (int) ref::Q8_0; ++ const int g = nm.find("hc_attn") != std::string::npos ? 1 : nm.find("hc_ffn") != std::string::npos ? 2 : 0; ++ return (hc_q4g_ & g) ? (int) ref::Q4_0 : (hc_q5g_ & g) ? (int) ref::Q5_0 : (int) ref::Q8_0; ++ }; + for (const auto& kv : f_.tensors) { + const gguf::Tensor& x = kv.second; + if (x.type != ref::BF16 || x.name.find("hc_") == std::string::npos || x.name.find("_norm") != std::string::npos) continue; + uint64_t ne = 1; + for (uint64_t d : x.shape) ne *= d; ++ const int qt = low_type(x.name); + q8off_[x.name] = tot; ++ q8type_[x.name] = qt; + names.push_back(x.name); +- tot += ne / 32 * 34; ++ tot += ne / 32 * (qt == ref::Q4_0 ? 18 : qt == ref::Q5_0 ? 22 : 34); + } + 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 +868,18 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta + src = raw.data(); + } + ref::dequant(ref::BF16, src, tmp.data(), ne); +- ref::quantize_q8_0(tmp.data(), ne, Q8_.as() + q8off_[names[i]]); ++ if (q8type_.at(names[i]) == ref::Q4_0) ++ quant_q4_0_lloyd(tmp.data(), ne, Q8_.as() + q8off_[names[i]]); ++ else if (q8type_.at(names[i]) == ref::Q5_0) ++ quant_q5_0(tmp.data(), ne, Q8_.as() + q8off_[names[i]]); ++ else ++ ref::quantize_q8_0(tmp.data(), ne, Q8_.as() + q8off_[names[i]]); + } + }); + for (auto& x : th) x.join(); +- std::fprintf(stderr, "strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): %zu matrices, %.0f MB, made in %.0f ms\n", +- names.size(), tot / 1e6, now_ms() - tq); ++ std::fprintf(stderr, "strata metal: HC read in %s: %zu matrices, %.0f MB, made in %.0f ms\n", ++ hc_q4_ ? "Q8_0 with Q4_0/Q5_0 (STRATA_HC_Q4/_Q5)" : "Q8_0 (STRATA_HC_Q8=0: BF16)", names.size(), tot / 1e6, ++ now_ms() - tq); + } + if (stream_) { + std::vector skip; // the BF16 originals of the Q8_0 HC copies are not read again +@@ -543,9 +889,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); ++ // the read-ahead's router copies in Q8_0 (STRATA_METAL_PRED_Q8=0: the F32 originals): the guess only picks what is ++ // read ahead, never the routing, and an F32 router is 2.6 MB a layer read twice a token (92 MB a token saved) ++ const char* pq = std::getenv("STRATA_METAL_PRED_Q8"); ++ if (st_pred_ && !(pq && std::atoi(pq) == 0) && T("blk.1.ffn_gate_inp.weight").type == ref::F32) { ++ const size_t rb = ref::row_bytes(ref::Q8_0, N), per = rb * (size_t) n_exp_; ++ predW_ = A(per * 48); ++ predT_.resize(48); ++ for (int l = 0; l < 48; ++l) { ++ const gguf::Tensor& src = T("blk." + std::to_string(l) + ".ffn_gate_inp.weight"); ++ ref::quantize_q8_0((const float*) ((const uint8_t*) W_.ptr + src.offset), (size_t) N * n_exp_, ++ predW_.as() + per * (size_t) l); ++ predT_[(size_t) l] = src; ++ predT_[(size_t) l].type = ref::Q8_0; ++ predT_[(size_t) l].offset = per * (size_t) l; ++ predT_[(size_t) l].bytes = per; ++ } ++ } + if (st_pred_ && st_->ahead2() > 0) predL2_ = A(4 * (size_t) 48 * MAXT * 512); + } +- ple_on_ = ss.ple.ready() && f_.find("blk.1.ple_key.weight") != nullptr; ++ { // the routers: an F32 router whose values are all BF16 (the upcast originals) is read as BF16, bit for bit the same ++ // weights at half the bytes (63 MB a token). STRATA_METAL_ROUTER_BF16=0: the F32 tensors. ++ const char* rb = std::getenv("STRATA_METAL_ROUTER_BF16"); ++ bool exact = !(rb && std::atoi(rb) == 0) && T("blk.0.ffn_gate_inp.weight").type == ref::F32; ++ for (int l = 0; exact && l < 48; ++l) { ++ const gguf::Tensor& t = T("blk." + std::to_string(l) + ".ffn_gate_inp.weight"); ++ const uint32_t* p = (const uint32_t*) ((const uint8_t*) W_.ptr + t.offset); ++ for (size_t i = 0; exact && i < t.bytes / 4; ++i) exact = (p[i] & 0xFFFFu) == 0; ++ } ++ if (exact) { ++ const size_t per = (size_t) N * n_exp_ * 2; ++ routW_ = cx_->alloc(per * 48); ++ routT_.resize(48); ++ for (int l = 0; l < 48; ++l) { ++ const gguf::Tensor& t = T("blk." + std::to_string(l) + ".ffn_gate_inp.weight"); ++ const uint32_t* p = (const uint32_t*) ((const uint8_t*) W_.ptr + t.offset); ++ uint16_t* q = (uint16_t*) (routW_.as() + per * (size_t) l); ++ for (size_t i = 0; i < t.bytes / 4; ++i) q[i] = (uint16_t) (p[i] >> 16); ++ routT_[(size_t) l] = t; ++ routT_[(size_t) l].type = ref::BF16; ++ routT_[(size_t) l].offset = per * (size_t) l; ++ routT_[(size_t) l].bytes = per; ++ } ++ } ++ } ++ ple_on_ = (ss.ple.ready() || (ss.ple.table != nullptr && ss.ple.table->is_open())) && f_.find("blk.1.ple_key.weight") != nullptr; + 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 +951,64 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta + return true; + } + +-void Window::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, const mtl::Buffer* wb) { +- const int n_in = (int) t.shape[0], n_out = t.shape.size() > 1 ? (int) t.shape[1] : 1; +- const int ty = t.type; ++// gemv()'s tile for a matrix: {LPR, NR, SG, KS} ++std::array Window::gemv_tile(int ty, int n_in, int n_out, int ntok) { + int L = 16, Rr = 1, S = 4, K = 1; +- if (n_in >= 6144) { L = 32; Rr = 1; S = 8; K = 2; } +- else if (n_in == 640) { L = 8; Rr = 1; S = 4; K = 1; } +- else if (n_out <= 640 || ty == ref::Q2_0 || ty == ref::BF16) { L = 32; Rr = 1; S = 4; K = 1; } ++ const bool m5 = m5_layout(); ++ if (n_in >= 6144) { ++ if (m5 && ty == ref::IQ4_XS) { L = (ntok >= 2 && ntok <= 3) ? 16 : 32; Rr = 2; S = 4; K = 1; } ++ else if (m5 && ty == ref::IQ3_S) { ++ if (ntok >= 2) { L = 8; Rr = 1; S = 8; K = 1; } ++ else { L = 32; Rr = 2; S = 4; K = 1; } ++ } else if (m5 && ty == ref::Q6_K) { L = ntok >= 2 ? 32 : 16; Rr = ntok >= 4 ? 2 : 1; S = 4; K = 1; } ++ else { L = 32; Rr = 1; S = 8; K = 2; } ++ } else if (n_in == 640) { ++ // M5, weights from DRAM, 640 -> 2560: Q2_0 at 4 tokens 8,2,4,1 -4..-6% (the same sums as 8,1,4,1) ++ if (m5 && ty == ref::IQ4_NL && ntok >= 2) { L = 4; Rr = 2; S = 4; K = 1; } ++ else if (m5 && ty == ref::Q8_0) { L = ntok >= 3 ? 4 : 32; Rr = 1; S = 4; K = 1; } ++ else { L = 8; Rr = m5 && ty == ref::Q2_0 && ntok >= 4 ? 2 : 1; S = 4; K = 1; } ++ } else if (n_out <= 640) { ++ // M5, weights from DRAM, 2560 -> 512 (QSA k, v) against 32,1,4,1: 32,2,4,1 keeps LPR=32 (100% bitwise identical) ++ // and cuts IQ4_XS by -3.2%..-25.5%, IQ3_S by -9.0%..-24.4%, and Q6_K by -7.6%..-13.9%. ++ if (m5 && ty == ref::IQ4_XS && n_out == 640) { L = 16; Rr = 2; S = 4; K = 1; } ++ else if (m5 && (ty == ref::IQ4_XS || ty == ref::IQ3_S || ty == ref::Q6_K) && n_out == 512) { L = 32; Rr = 2; S = 4; K = 1; } ++ else if (m5 && ty == ref::IQ3_S && n_out == 640) { ++ if (ntok >= 2) { L = 8; Rr = 1; S = 8; K = 1; } ++ else { L = 32; Rr = 2; S = 4; K = 1; } ++ } else if (m5 && ty == ref::Q6_K && n_out == 640) { ++ if (ntok >= 2) { L = 8; Rr = 1; S = 4; K = 1; } ++ else { L = 32; Rr = 2; S = 4; K = 1; } ++ } else if (m5 && ty == ref::Q2_0 && n_out == 640) { ++ if (ntok >= 2) { L = 16; Rr = 1; S = 4; K = 1; } ++ else { L = 32; Rr = 2; S = 4; K = 1; } ++ } else if (m5 && ty == ref::Q8_0 && n_out == 512) { ++ if (ntok >= 2) { L = 16; Rr = 1; S = 4; K = 1; } ++ else { L = 32; Rr = 1; S = 4; K = 4; } ++ } else if (m5 && ty == ref::BF16 && n_out >= 128 && n_out <= 512) { ++ if (ntok >= 3 && n_out == 128) { L = 16; Rr = 1; S = 4; K = 1; } ++ else { L = 32; Rr = 1; S = 8; K = 8; } ++ } else { L = 32; Rr = 1; S = 4; K = 1; } ++ } else if (ty == ref::Q2_0 || ty == ref::BF16) { ++ L = 32; Rr = 1; S = 4; K = 1; ++ } else if (m5 && ty == ref::Q8_0) { ++ if (ntok == 1) { L = 32; Rr = 1; S = 4; K = 4; } ++ else if (ntok <= 3) { L = 32; Rr = 2; S = 4; K = 1; } ++ else { L = 16; Rr = 1; S = 4; K = 1; } ++ } else if (m5 && ty == ref::IQ4_XS) { ++ if (ntok >= 2) { L = 16; Rr = 2; S = 4; K = 1; } ++ else { L = 32; Rr = 2; S = 4; K = 1; } ++ } else if (m5 && ty == ref::IQ3_S && n_out >= 6144) { ++ if (ntok >= 2) { L = 8; Rr = 1; S = 4; K = 1; } ++ else { L = 32; Rr = 2; S = 4; K = 1; } ++ } else if (m5 && ty == ref::Q6_K && n_out >= 6144) { ++ if (ntok == 1) { L = 32; Rr = 1; S = 8; K = 2; } ++ else if (ntok == 2) { L = 32; Rr = 2; S = 4; K = 1; } ++ else if (n_out == 6144) { L = 16; Rr = 1; S = 4; K = 1; } ++ else { L = 8; Rr = 1; S = 4; K = 1; } ++ } else if (m5 && (ty == ref::Q4_K || ty == ref::Q5_K || ty == ref::Q3_K) && n_out >= 6144) { ++ if (ntok >= 2) { L = 8; Rr = 1; S = 8; K = 1; } ++ else if (ty != ref::Q3_K) { L = 32; Rr = 1; S = 4; K = 1; } ++ } + { + // 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 +1019,15 @@ void Window::gemv(std::vector& ds, const gguf::Tensor& t, const A + if (K == 1 || (L == 32 && K == 2)) Rr = 2; + } + } ++ return {L, Rr, S, K}; ++} ++ ++void Window::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, const mtl::Buffer* wb, int tile_ntok) { ++ // tile_ntok > 0: the tile of a window that wide, so these ntok tokens' answers are bitwise its first tokens' ++ const int n_in = (int) t.shape[0], n_out = t.shape.size() > 1 ? (int) t.shape[1] : 1; ++ const int ty = t.type; ++ const auto [L, Rr, S, K] = gemv_tile(ty, n_in, n_out, tile_ntok > 0 ? tile_ntok : ntok); + 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 +1043,37 @@ void Window::gemv(std::vector& ds, const gguf::Tensor& t, const A + push(ds, d); + } + ++// two GEMVs of the same format and shape over the same x in one dispatch (gemv_cat): bitwise the two gemv() calls when ++// gemv() would run them KS = 1 (a row's sums depend only on LPR). False when not instantiated / not applicable. ++// STRATA_METAL_GEMV_CAT=0: off. ++bool Window::gemv_cat(std::vector& ds, const gguf::Tensor& ta, const gguf::Tensor& tb, const Act& x, ++ const mtl::Buffer& ya, const mtl::Buffer& yb, size_t y_stride, int ntok, size_t y_off) { ++ static const bool on = !(std::getenv("STRATA_METAL_GEMV_CAT") && std::atoi(std::getenv("STRATA_METAL_GEMV_CAT")) == 0); ++ if (!on || ta.type != tb.type || (ta.type != ref::BF16 && ta.type != ref::F32 && ta.type != ref::Q8_0) || ta.shape != tb.shape) ++ return false; ++ const int n_in = (int) ta.shape[0], n_out = ta.shape.size() > 1 ? (int) ta.shape[1] : 1; ++ const auto [L, Rr, S, K] = gemv_tile(ta.type, n_in, n_out, ntok); ++ // gemv.metal's GVC instances ++ const bool inst = K == 1 && ((L == 32 && Rr == 1 && (S == 4 || S == 8)) || (L == 16 && Rr == 1 && S == 4) || ++ (L == 32 && Rr == 2 && S == 4) || (ta.type == ref::Q8_0 && L == 8 && Rr == 1 && S == 8)); ++ if (!inst) return false; ++ const std::string fn = "gemvc_" + std::to_string(ta.type) + "_" + std::to_string(L) + "_" + std::to_string(Rr) + "_" + ++ std::to_string(S); ++ GemvArgs a{(uint32_t) n_in, (uint32_t) (2 * n_out), (uint32_t) ntok, (uint32_t) ref::row_bytes(ta.type, n_in), ++ (uint32_t) n_in, (uint32_t) y_stride}; ++ const uint32_t sp = (uint32_t) n_out; ++ mtl::Dispatch d; ++ d.p = P(lib_, fn, mtl::Constants().i(0, ntok), fn + "#" + std::to_string(ntok)); ++ d.args = {mtl::Arg::B(W_, ta.offset), mtl::Arg::B(x.b, x.plane_bytes() * xkind(ta.type)), ++ mtl::Arg::B(x.b, x.plane_bytes() * 4), mtl::Arg::B(ya, y_off), mtl::Arg::V(a), mtl::Arg::B(W_, tb.offset), ++ mtl::Arg::B(yb, y_off), mtl::Arg::V(sp)}; ++ const int rpt = S * (32 / L) * Rr; ++ d.grid[0] = (2 * n_out + rpt - 1) / rpt; ++ d.tg[0] = S * 32; ++ push(ds, d); ++ return true; ++} ++ + 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 +1084,58 @@ void Window::xprep(std::vector& ds, const mtl::Buffer& x, const A + push(ds, d); + } + ++// the MoE combine (moe_combine_sh) of the last layer into bo, pushed now ++void Window::comb_push(std::vector& ds, const mtl::Buffer& bo, int ntok) { ++ GlueArgs a{N, (uint32_t) ntok, N, TOPK}; ++ mtl::Dispatch d; ++ d.p = P(lib_, "moe_combine_sh", {}, "moe_combine_sh"); ++ d.args = {mtl::Arg::B(ye_), mtl::Arg::B(plan_), mtl::Arg::B(ysh_), mtl::Arg::B(sg_), mtl::Arg::B(bo), mtl::Arg::V(a)}; ++ d.grid[0] = (ntok * N + 255) / 256; ++ d.tg[0] = 256; ++ push(ds, d); ++} ++ ++// the MoE combine into bo, left to the next gr_norm that reads bo (gr_norm_comb: bitwise moe_combine_sh + gr_norm; ++// one dispatch, one barrier and bo's round trip less). STRATA_METAL_GR_COMB=0: the combine's own dispatch. ++void Window::comb_defer(std::vector& ds, const mtl::Buffer& bo, int ntok) { ++ static const bool on = !(std::getenv("STRATA_METAL_GR_COMB") && std::atoi(std::getenv("STRATA_METAL_GR_COMB")) == 0); ++ comb_flush(ds); ++ if (!on) { comb_push(ds, bo, ntok); return; } ++ comb_bo_ = &bo; ++ comb_ntok_ = ntok; ++} ++ ++void Window::comb_flush(std::vector& ds) { ++ if (!comb_bo_) return; ++ const mtl::Buffer* bo = comb_bo_; ++ comb_bo_ = nullptr; ++ comb_push(ds, *bo, comb_ntok_); ++} ++ ++// gr_norm (R' = R + bo * 2 sigmoid(inj / hc) when `apply`, xn), or gr_norm_comb when bo is a deferred MoE combine ++void Window::gr_norm(std::vector& ds, mtl::Buffer& R, const mtl::Buffer& bo, const mtl::Buffer& inj, ++ const mtl::Buffer& wb, size_t w_off, bool apply, int ntok) { ++ const bool comb = apply && comb_bo_ == &bo && comb_ntok_ == ntok; ++ if (comb) comb_bo_ = nullptr; ++ else comb_flush(ds); ++ GrArgs a{N, HC, LR, (uint32_t) ntok, apply ? 1u : 0u, eps_}; ++ mtl::Dispatch d; ++ if (comb) { ++ static_assert(N == 2560 && HC == 4, "gr_norm_comb shape"); ++ const uint32_t tk = TOPK; ++ d.p = P(lib_, "gr_norm_comb", {}, "gr_norm_comb"); ++ d.args = {mtl::Arg::B(R), mtl::Arg::B(ye_), mtl::Arg::B(inj), mtl::Arg::B(wb, w_off), mtl::Arg::B(xn_), ++ mtl::Arg::V(a), mtl::Arg::B(plan_), mtl::Arg::B(ysh_), mtl::Arg::B(sg_), mtl::Arg::V(tk)}; ++ } else { ++ d.p = P(lib_, "gr_norm", {}, "gr_norm"); ++ d.args = {mtl::Arg::B(R), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wb, w_off), mtl::Arg::B(xn_), mtl::Arg::V(a)}; ++ } ++ d.grid[0] = HC; ++ d.grid[1] = ntok; ++ d.tg[0] = 256; ++ push(ds, d); ++} ++ + // 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 +1144,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_}; +- mtl::Dispatch d1; +- d1.p = P(lib_, "gr_norm", {}, "gr_norm"); +- d1.args = {mtl::Arg::B(R_), mtl::Arg::B(bo_), mtl::Arg::B(inj_[cur_inj_]), mtl::Arg::B(W_, tn.offset), mtl::Arg::B(xn_), +- mtl::Arg::V(a)}; +- d1.grid[0] = HC; +- d1.grid[1] = ntok; +- d1.tg[0] = 256; +- push(ds, d1); ++ gr_norm(ds, R_, bo_, inj_[cur_inj_], W_, tn.offset, apply, ntok); + // 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 +1156,25 @@ 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); ++ const int lt = q8 && hc_q4_ && !final_ && q8type_.count(tdn.name) ? q8type_.at(tdn.name) : (int) ref::Q8_0; ++ const bool q4 = lt == ref::Q4_0, q5 = lt == ref::Q5_0; + const bool injf = q8 && !final_ && !q8off_.count(tin.name) && tin.type == ref::F32; + mtl::Dispatch d2; +- const char* dn = injf ? "gr_down_q8f_8_8" : 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)); ++ // M5: KS = 4 (4 simdgroups a threadgroup, bitwise the old 8_4), else KS = 8. Past 4 tokens two rows a simdgroup ++ // (_r2), which reads each xn value once for both ++ static_assert(LR % 2 == 0 && HC % 2 == 0, "gr_down _r2 rows"); ++ const bool m5_gr = m5_layout() && !q4 && !q5, r2 = !q4 && !q5 && ntok > 4; ++ std::string dn = q4 ? "gr_down_q4_8_8" : q5 ? "gr_down_q5_8_8" ++ : std::string(injf ? "gr_down_q8f_" : q8 ? "gr_down_q8_" : "gr_down_") + (m5_gr ? "4_4" : "8_8"); ++ if (r2) dn += "_r2"; ++ d2.p = P(lib_, dn, mtl::Constants().i(1, ntok), dn + "#" + std::to_string(ntok)); + d2.args = {mtl::Arg::B(*bd, od), mtl::Arg::B(*bi, oi), mtl::Arg::B(xn_), mtl::Arg::B(lo_), + mtl::Arg::B(inj_[1 - cur_inj_]), mtl::Arg::V(a)}; +- d2.grid[0] = LR + HC; +- d2.tg[0] = 256; ++ d2.grid[0] = r2 ? (LR + HC) / 2 : (LR + HC); ++ d2.tg[0] = m5_gr ? 128 : 256; + push(ds, d2); + mtl::Dispatch d3; +- const char* up = q8 ? "gr_up_q8" : "gr_up"; ++ const char* up = q4 ? "gr_up_q4" : q5 ? "gr_up_q5" : q8 ? "gr_up_q8" : "gr_up"; + 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 +1200,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; +- gemv(ds, T(b + "ssm_alpha.weight"), XA_, alL_, 48, ntok, o_ab); +- nb_ = true; +- gemv(ds, T(b + "ssm_beta.weight"), XA_, beL_, 48, ntok, o_ab); ++ // ssm_alpha + ssm_beta: one dispatch (gemv_cat), else two ++ if (!gemv_cat(ds, T(b + "ssm_alpha.weight"), T(b + "ssm_beta.weight"), XA_, alL_, beL_, 48, ntok, o_ab)) { ++ gemv(ds, T(b + "ssm_alpha.weight"), XA_, alL_, 48, ntok, o_ab); ++ nb_ = true; ++ gemv(ds, T(b + "ssm_beta.weight"), XA_, beL_, 48, ntok, o_ab); ++ } + 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 +1224,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; ++ const bool sel = !(fuse_ && pos0 + ntok <= 2051); + gemv(ds, T(b + "attn_q.weight"), XA_, qfull_, 12288, ntok); + nb_ = true; + gemv(ds, T(b + "attn_k.weight"), XA_, kcur_, 512, ntok); + nb_ = true; + gemv(ds, T(b + "attn_v.weight"), XA_, vcur_, 512, ntok); +- nb_ = true; +- gemv(ds, T(b + "indexer.q_proj.weight"), XA_, iqr_, 512, ntok); ++ if (sel) { ++ nb_ = true; ++ gemv(ds, T(b + "indexer.q_proj.weight"), XA_, iqr_, 512, ntok); ++ } + 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 +1258,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; +- const bool sel = !(fuse_ && pos0 + ntok <= 2051); + if (sel) { + push(ds, s); + if (rec_) rec_->push_back({(int) ds.size() - 1, 2}); +@@ -739,12 +1275,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)}; +- at.grid[0] = 33; ++ at.grid[0] = (uint32_t) std::min(33, std::max(1, (pos0 + ntok + 63) / 64)); + at.grid[1] = 2; + at.grid[2] = ntok; + at.tg[0] = 256; + push(ds, at); +- if (rec_) rec_->push_back({(int) ds.size() - 1, 0}); ++ if (rec_) rec_->push_back({(int) ds.size() - 1, 3}); + mtl::Dispatch m; + m.p = P(lqsa_, "qsa_merge", {}, "qsa_merge"); + QsaArgs am = a; +@@ -759,28 +1295,51 @@ 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); +- const gguf::Tensor& trt = T(b + "ffn_gate_inp.weight"); ++ const gguf::Tensor& trt = routT_.empty() ? T(b + "ffn_gate_inp.weight") : routT_[(size_t) l]; ++ const mtl::Buffer& rtw = routT_.empty() ? W_ : routW_; + const bool rf = routef_ && (trt.type == ref::BF16 || trt.type == ref::F32); + if (!rf) { +- gemv(ds, trt, XF_, logits_, 512, ntok); ++ gemv(ds, trt, XF_, logits_, 512, ntok, 0, &rtw); + nb_ = true; + } + if (stream_ && st_pred_ && l + 1 < 48) { // the next layer's routing guessed from this layer's FFN input ++ // STRATA_METAL_PREFETCH_DRAFT <= 0: the read-ahead reads only token 0's guess (stream.cpp ahead), so only it is ++ // computed, on the window's tile (bitwise the window's token 0). M5, Q8_0 2560 -> 512 + 512 at 4 tokens: -16% ++ // hot, -6% from DRAM (gemv_parity --pred; gemv_cat of the two at 1 token is slower than the pair) ++ static const int draft_k = [] { const char* e = std::getenv("STRATA_METAL_PREFETCH_DRAFT"); return e && *e ? std::atoi(e) : 1; }(); ++ const int pt = draft_k <= 0 ? 1 : ntok; + nb_ = !rf; // beside the router GEMV; with router_route it is the first reader of XF_ here: a barrier +- gemv(ds, T("blk." + std::to_string(l + 1) + ".ffn_gate_inp.weight"), XF_, predL_, 512, ntok, +- (size_t) l * MAXT * 512 * 4); ++ if (predW_.ptr) ++ gemv(ds, predT_[(size_t) l + 1], XF_, predL_, 512, pt, (size_t) l * MAXT * 512 * 4, &predW_, ntok); ++ else ++ gemv(ds, T("blk." + std::to_string(l + 1) + ".ffn_gate_inp.weight"), XF_, predL_, 512, pt, ++ (size_t) l * MAXT * 512 * 4, nullptr, ntok); + if (predL2_.ptr && l + 2 < 48) { + nb_ = true; +- gemv(ds, T("blk." + std::to_string(l + 2) + ".ffn_gate_inp.weight"), XF_, predL2_, 512, ntok, +- (size_t) l * MAXT * 512 * 4); ++ if (predW_.ptr) ++ gemv(ds, predT_[(size_t) l + 2], XF_, predL2_, 512, pt, (size_t) l * MAXT * 512 * 4, &predW_, ntok); ++ else ++ gemv(ds, T("blk." + std::to_string(l + 2) + ".ffn_gate_inp.weight"), XF_, predL2_, 512, pt, ++ (size_t) l * MAXT * 512 * 4, nullptr, ntok); + } + nb_ = true; + } +- gemv(ds, T(b + "ffn_gate_shexp.weight"), XF_, gsh_, NFF, ntok); +- nb_ = true; +- gemv(ds, T(b + "ffn_up_shexp.weight"), XF_, ush_, NFF, ntok); +- nb_ = true; +- gemv(ds, T(b + "ffn_gate_inp_shexp.weight"), XF_, sg_, 1, ntok); ++ // the shared expert: its own list when streaming, placed between the handshake's signal and its wait (nothing in it ++ // reads what the host writes), else where it always was ++ std::vector shx; ++ std::vector& dsh = sigs_ ? shx : ds; ++ static const bool xsplit = !(std::getenv("STRATA_METAL_EXPERT_SPLIT") && std::atoi(std::getenv("STRATA_METAL_EXPERT_SPLIT")) == 0); ++ const bool esplit = sigs_ && xsplit; ++ auto shared_gu = [&] { ++ gemv(dsh, T(b + "ffn_gate_shexp.weight"), XF_, gsh_, NFF, ntok); ++ nb_ = true; ++ gemv(dsh, T(b + "ffn_up_shexp.weight"), XF_, ush_, NFF, ntok); ++ nb_ = true; ++ gemv(dsh, T(b + "ffn_gate_inp_shexp.weight"), XF_, sg_, 1, ntok); ++ }; ++ if (sigs_) nb_ = false; ++ if (!esplit) shared_gu(); ++ if (sigs_) nb_ = false; // moe_route below waits for the router GEMV (router_route sets its own) + { + MoeArgs a{}; + const gguf::Tensor &tg_ = T(b + "ffn_gate_exps.weight"), &tu = T(b + "ffn_up_exps.weight"), +@@ -801,34 +1360,63 @@ void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0 + a.n_ff = NFF; + a.n_tok = ntok; + a.topk = TOPK; +- a.n_expert = NEXP; ++ a.n_expert = (uint32_t) n_exp_; ++ static const bool dmoce = !(std::getenv("STRATA_METAL_DRAFT_MOCE") && std::atoi(std::getenv("STRATA_METAL_DRAFT_MOCE")) == 0); ++ a.pad = stream_ ? (moce_pad_ | (dmoce && win_head_ && ntok > 1 && moce_pad_ ? 0x80000000u : 0u)) : 0u; + const bool dq2 = td.type == ref::Q2_0; +- const int gl = 32, gr = 1, gs = 4, ds_ = 4; +- // 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) ++ int gl = 32, gr = 1, gs = 4, ds_ = 4; ++ // the down tile: Q2_0 8,2 (the M5 Air's KPROF sweep: 2.185 ms a token vs 2.431 at 8,1), 8,1 on the M4; 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(); +- const int dl = dq2 || (m4 && td.type == ref::Q8_0) ? 8 : 4, +- dr = dq2 || (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1)) ? 1 : 2; ++ if (!m4 && (tg_.type == ref::IQ2_XXS || tg_.type == ref::IQ3_XXS || ++ (ntok >= 3 && (tg_.type == ref::IQ2_S || tg_.type == ref::IQ2_XS || tg_.type == ref::IQ3_S)))) ++ gl = 16; ++ int dl = dq2 || (!m4 && td.type == ref::IQ4_NL) || (m4 && td.type == ref::Q8_0) ? 8 : 4, ++ dr = dq2 ? (m4 ? 1 : 2) : (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1)) ? 1 : 2; ++ // for sweeping (KPROF): STRATA_METAL_MOE_GU_TILE / _DN_TILE = "LPR,NR,SG" of an instantiated moe_gu_ / moe_down_ ++ static const std::array gu_t = tile_env("STRATA_METAL_MOE_GU_TILE"), dn_t = tile_env("STRATA_METAL_MOE_DN_TILE"); ++ if (gu_t[0] > 0) { gl = gu_t[0]; gr = gu_t[1]; gs = gu_t[2]; } ++ if (dn_t[0] > 0) { dl = dn_t[0]; dr = dn_t[1]; ds_ = dn_t[2]; } + a.gx_gu = NFF / (gs * (32 / gl) * gr); + a.gx_dn = N / (ds_ * (32 / dl) * dr); + mtl::Dispatch r; ++ const int mp_fc = std::getenv("STRATA_METAL_MOCE_PROTECT") ? std::max(0, std::min(10, std::atoi(std::getenv("STRATA_METAL_MOCE_PROTECT")))) : 3; ++ uint32_t mr1_bits = 0x3F800000u; // 1.0f ++ if (const char* mr1_env = std::getenv("STRATA_METAL_MOCE_BONUS_R1")) { ++ const float mr1_b = (float) std::atof(mr1_env); ++ if (mr1_b > 0.0f) { ++ const float mr1_s = std::exp(mr1_b); ++ std::memcpy(&mr1_bits, &mr1_s, 4); ++ } ++ } ++ mtl::Constants r_fc; ++ std::string r_fc_tag; ++ if (mp_fc != 3) { r_fc.i(24, mp_fc); r_fc_tag += "#mp" + std::to_string(mp_fc); } ++ if (mr1_bits != 0x3F800000u) { r_fc.u(25, mr1_bits); r_fc_tag += "#mr1_" + std::to_string(mr1_bits); } ++ // the resident experts first (moe.metal function constant 26): they run beside the host's turn ++ if (esplit) { r_fc.b(26, true); r_fc_tag += "#xs"; } + if (rf) { + // the router GEMV + moe_route (bitwise both), beside the shared expert's GEMVs +- r.p = P(lib_, "router_route_" + std::to_string(trt.type), mtl::Constants().i(0, ntok), "router_route_" + std::to_string(trt.type) + "#" + std::to_string(ntok)); +- GemvArgs ga{(uint32_t) N, 512, (uint32_t) ntok, (uint32_t) ref::row_bytes(trt.type, N), (uint32_t) N, 512}; +- r.args = {mtl::Arg::B(W_, trt.offset), mtl::Arg::B(XF_.b, 0), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4), ++ mtl::Constants rfc2 = r_fc; ++ rfc2.i(0, ntok); ++ r.p = P(lib_, "router_route_" + std::to_string(trt.type), rfc2, "router_route_" + std::to_string(trt.type) + "#" + std::to_string(ntok) + r_fc_tag); ++ GemvArgs ga{(uint32_t) N, (uint32_t) n_exp_, (uint32_t) ntok, (uint32_t) ref::row_bytes(trt.type, N), (uint32_t) N, 512}; ++ r.args = {mtl::Arg::B(rtw, trt.offset), mtl::Arg::B(XF_.b, 0), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4), + mtl::Arg::B(logits_), mtl::Arg::V(ga), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), +- mtl::Arg::V(a), mtl::Arg::B(rcounter_)}; +- r.grid[0] = 512 / 16; ++ mtl::Arg::V(a), mtl::Arg::B(rcounter_), ++ stream_ ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(rcounter_)}; ++ r.grid[0] = (uint32_t) ((n_exp_ + 15) / 16); + r.tg[0] = 512; +- nb_ = true; ++ nb_ = !sigs_ || (stream_ && st_pred_ && l + 1 < 48); + } else { +- r.p = P(lib_, "moe_route", {}, "moe_route"); +- r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a)}; ++ r.p = P(lib_, "moe_route", r_fc, "moe_route" + r_fc_tag); ++ r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a), ++ stream_ ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(rcounter_)}; + r.tg[0] = 512; + } + push(ds, r); +- { ++ if (sigs_) sigs_->push_back((int) ds.size()); // the handshake starts: the routing is read ++ auto shared_prep = [&] { + GlueArgs ga{NFF, MAXT, N, TOPK}; + mtl::Dispatch d; + d.p = P(lib_, "swiglu_prep", {}, "swiglu_prep"); +@@ -836,53 +1424,67 @@ 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; +- nb_ = !rf; // with router_route the shared GEMVs are not behind a barrier yet: this one waits for them +- push(ds, d); +- } ++ nb_ = esplit || (!sigs_ && !rf); ++ push(dsh, d); ++ }; ++ if (!esplit) shared_prep(); ++ if (sigs_ && !esplit) gemv(shx, T(b + "ffn_down_shexp.weight"), XH_, ysh_, N, ntok); // the shared expert's down too + const std::string gun = "moe_gu_" + std::to_string(tg_.type) + "_" + std::to_string(gl) + "_" + std::to_string(gr) + + "_" + std::to_string(gs); + const std::string dnn = "moe_down_" + std::to_string(dq2 ? 142 : td.type) + "_" + std::to_string(dl) + "_" + + std::to_string(dr) + "_" + std::to_string(ds_); +- mtl::Dispatch g; + const mtl::Buffer& EW = stream_ ? st_->chunk(0) : W_; +- // function constant 23 (QUANT_CTAB): the i-quant grids from constant memory, no threadgroup copy (same values) +- const bool ct = metal_ctab(); +- mtl::Constants ac = stream_ ? mtl::Constants().b(20, true) : mtl::Constants(); +- ac.b(23, ct); +- const std::string ak = std::string(stream_ ? "#arena" : "") + (ct ? "#ct" : ""); +- g.p = P(lib_, gun, ac, gun + ak); +- g.args = {mtl::Arg::B(EW), mtl::Arg::B(XF_.b, XF_.plane_bytes() * xkind(tg_.type)), mtl::Arg::B(h_), mtl::Arg::B(plan_), +- mtl::Arg::V(a), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4)}; +- if (stream_) +- for (int c = -1; c < 3; ++c) +- g.args.push_back(c < 0 ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(st_->chunk(c + 1))); +- g.tg[0] = 32 * gs; +- g.indirect = &plan_; +- g.ind_off = P_IGU * 4; ++ // function constant 23 (QUANT_CTAB): the i-quant grids from constant memory, no threadgroup copy (same values); ++ // IQ3_S uses threadgroup tab on M5 (118.4 vs 105.8 GB/s: its 2 indices per loop diverge across lanes) ++ const bool ct = metal_ctab() && tg_.type != ref::IQ3_S; ++ // phase (function constant 27): 0 every group, 1 the resident ones, 2 the missing ones ++ auto experts = [&](std::vector& out, int phase) { ++ mtl::Constants ac = stream_ ? mtl::Constants().b(20, true) : mtl::Constants(); ++ ac.b(23, ct); ++ std::string ak = std::string(stream_ ? "#arena" : "") + (ct ? "#ct" : ""); ++ if (phase) { ac.i(27, phase); ak += "#ph" + std::to_string(phase); } ++ mtl::Dispatch g; ++ g.p = P(lib_, gun, ac, gun + ak); ++ g.args = {mtl::Arg::B(EW), mtl::Arg::B(XF_.b, XF_.plane_bytes() * xkind(tg_.type)), mtl::Arg::B(h_), mtl::Arg::B(plan_), ++ mtl::Arg::V(a), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4)}; ++ if (stream_) ++ for (int c = -1; c < 3; ++c) ++ g.args.push_back(c < 0 ? (phase == 1 ? mtl::Arg::B(plan_) : mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4)) : mtl::Arg::B(st_->chunk(c + 1))); ++ g.tg[0] = 32 * gs; ++ g.indirect = &plan_; ++ g.ind_off = (phase == 1 ? P_IGU_R : phase == 2 ? P_IGU_M : P_IGU) * 4; ++ if (phase == 1) nb_ = false; ++ push(out, g); ++ if (phase == 1) { ++ nb_ = true; ++ shared_gu(); ++ } ++ if (!sigs_) { ++ nb_ = true; ++ gemv(out, T(b + "ffn_down_shexp.weight"), XH_, ysh_, N, ntok); ++ } ++ mtl::Dispatch dd; ++ dd.p = P(lib_, dnn, ac, dnn + ak); ++ dd.args = {mtl::Arg::B(EW), mtl::Arg::B(h_), mtl::Arg::B(ye_), mtl::Arg::B(plan_), mtl::Arg::V(a), mtl::Arg::B(h_)}; ++ if (stream_) ++ for (int c = -1; c < 3; ++c) ++ dd.args.push_back(c < 0 ? (phase == 1 ? mtl::Arg::B(plan_) : mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4)) : mtl::Arg::B(st_->chunk(c + 1))); ++ dd.tg[0] = 32 * ds_; ++ dd.indirect = &plan_; ++ dd.ind_off = (phase == 1 ? P_IDN_R : phase == 2 ? P_IDN_M : P_IDN) * 4; ++ push(out, dd); ++ if (phase == 1) { ++ // zip independent shared work beside resident experts, then wait for the prepared activations ++ shared_prep(); ++ gemv(out, T(b + "ffn_down_shexp.weight"), XH_, ysh_, N, ntok); ++ } ++ }; ++ if (esplit) experts(shx, 1); // the experts the routing found in the arena: they cannot be evicted (pinned) ++ if (sigs_) ds.insert(ds.end(), shx.begin(), shx.end()); + if (cuts_) cuts_->push_back((int) ds.size()); // the host's turn: this layer's experts put in place +- push(ds, g); +- nb_ = true; +- gemv(ds, T(b + "ffn_down_shexp.weight"), XH_, ysh_, N, ntok); +- mtl::Dispatch dd; +- dd.p = P(lib_, dnn, ac, dnn + ak); +- dd.args = {mtl::Arg::B(EW), mtl::Arg::B(h_), mtl::Arg::B(ye_), mtl::Arg::B(plan_), mtl::Arg::V(a), mtl::Arg::B(h_)}; +- if (stream_) +- for (int c = -1; c < 3; ++c) +- dd.args.push_back(c < 0 ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(st_->chunk(c + 1))); +- dd.tg[0] = 32 * ds_; +- dd.indirect = &plan_; +- dd.ind_off = P_IDN * 4; +- push(ds, dd); +- } +- { +- GlueArgs a{N, (uint32_t) ntok, N, TOPK}; +- mtl::Dispatch d; +- d.p = P(lib_, "moe_combine_sh", {}, "moe_combine_sh"); +- d.args = {mtl::Arg::B(ye_), mtl::Arg::B(plan_), mtl::Arg::B(ysh_), mtl::Arg::B(sg_), mtl::Arg::B(bo_), mtl::Arg::V(a)}; +- d.grid[0] = (ntok * N + 255) / 256; +- d.tg[0] = 256; +- push(ds, d); ++ experts(ds, esplit ? 2 : 0); + } ++ comb_defer(ds, bo_, ntok); // into the next layer's (or the PLE's, the final mixer's) gr_norm + } + + namespace { +@@ -952,6 +1554,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(); ++ last_pos0_ = pos0; ++ if (Tn > 1 && Tn <= 4) { ++ tail_bak_.resize(tail_.size() * 4 * 128); ++ for (size_t j = 0; j < tail_.size(); ++j) ++ std::memcpy(tail_bak_.data() + j * 4 * 128, tail_[j].ptr, 4 * 4 * 128); ++ } + // ---- 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 +1608,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); ++ if (kind == 3) d.grid[0] = (uint32_t) std::min(33, std::max(1, (pos0 + Tn + 63) / 64)); + } + seg_T_ = Tn; + const double r0 = now_ms(); +@@ -1007,7 +1616,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; +- if (kt == Tn && head && ++seen == 20 && !stream_) kprof(c.ds, Tn); ++ if (kt == Tn && head && ++seen == 20) kprof(c.ds, Tn); + } + tp_rg_ = now_ms() - r0; + last_n_disp_ = (int) c.ds.size(); +@@ -1017,6 +1626,11 @@ 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; ++ // STRATA_METAL_SHEXP_OVERLAP=0: the shared expert after the handshake (as before) ++ static const bool overlap = !(std::getenv("STRATA_METAL_SHEXP_OVERLAP") && std::atoi(std::getenv("STRATA_METAL_SHEXP_OVERLAP")) == 0); ++ prewarm(); // window_init's already, unless the arena came late ++ cuts_ = stream_ ? &cw.cuts : nullptr; ++ sigs_ = stream_ && overlap ? &cw.sigs : nullptr; + 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) +@@ -1025,45 +1639,11 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t + gemv(ds, T("blk.1.ple_key.weight"), XE_, pkey_, 10240, Tn); + nb_ = true; + gemv(ds, T("blk.1.ple_value.weight"), XE_, pval_, N, Tn); ++ ple_key_norm(ds, Tn); + } + cur_inj_ = 0; + layer(ds, 0, Tn, pos0); +- if (ple_on_) { +- // layer 0's FFN write applied to R (gr_norm with apply writes R back; its xn is not used) +- const gguf::Tensor& tn = T("blk.1.hc_attn_norm.weight"); +- GrArgs a{N, HC, LR, (uint32_t) Tn, 1u, eps_}; +- mtl::Dispatch d1; +- d1.p = P(lib_, "gr_norm", {}, "gr_norm"); +- d1.args = {mtl::Arg::B(R_), mtl::Arg::B(bo_), mtl::Arg::B(inj_[cur_inj_]), mtl::Arg::B(W_, tn.offset), mtl::Arg::B(xn_), +- mtl::Arg::V(a)}; +- d1.grid[0] = HC; +- d1.grid[1] = Tn; +- d1.tg[0] = 256; +- push(ds, d1); +- // the PLE block on the GPU (msl/ple.metal): norms, gate, conv, the residual update +- struct PleArgs { uint32_t n_embd, hc; float eps; } pa{N, HC, strata::kernels::NG_RMS_EPS}; +- auto off = [&](const char* n) { return T(n).offset; }; +- mtl::Dispatch g1, g2, g3, g4; +- g1.p = P(lple_, "ple_gnorm", {}, "ple_gnorm"); +- g1.args = {mtl::Arg::B(pkey_), mtl::Arg::B(W_, off("blk.1.ple_norm_key.weight")), mtl::Arg::B(keyn_), mtl::Arg::V(pa)}; +- g1.grid[0] = HC; g1.grid[1] = Tn; g1.tg[0] = 256; +- push(ds, g1); +- g2.p = g1.p; +- g2.args = {mtl::Arg::B(R_), mtl::Arg::B(W_, off("blk.1.ple_norm_query.weight")), mtl::Arg::B(qn_), mtl::Arg::V(pa)}; +- g2.grid[0] = HC; g2.grid[1] = Tn; g2.tg[0] = 256; +- nb_ = true; // beside the key norm +- 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")), +- mtl::Arg::B(gated_), mtl::Arg::B(normW_), mtl::Arg::V(pa)}; +- g3.grid[0] = HC; g3.grid[1] = Tn; g3.tg[0] = 256; +- push(ds, g3); +- g4.p = P(lple_, "ple_conv_add", {}, "ple_conv_add"); +- g4.args = {mtl::Arg::B(pleHist_), mtl::Arg::B(normW_), kw_arg(), +- mtl::Arg::B(gated_), mtl::Arg::B(R_), mtl::Arg::V(pa)}; +- g4.grid[0] = HC * N / 256; g4.grid[1] = Tn; g4.tg[0] = 256; +- push(ds, g4); +- } ++ if (ple_on_) ple_after_layer0(ds, Tn); + for (int l = 1; l < 48; ++l) layer(ds, l, Tn, pos0); + // ---- the final mixer and the head + gr_read(ds, "output_hc", true, XA_, true, Tn); +@@ -1085,6 +1665,7 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t + } + rec_ = nullptr; + cuts_ = nullptr; ++ sigs_ = nullptr; + seg_T_ = Tn; + const double r0 = now_ms(); + last_gpu_ms_ = 1000.0 * run_ds(cw); +@@ -1095,7 +1676,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); +- return cx_->run_graph_seg(c.ds, c.cuts, [](void* u, int s) { ((Window*) u)->at_cut(s); }, this); ++ const std::vector& sigs = c.sigs.size() == c.cuts.size() ? c.sigs : c.cuts; // SHEXP_OVERLAP=0: none between ++ return cx_->run_graph_seg(c.ds, c.cuts, sigs, [](void* u, int s) { ((Window*) u)->at_cut(s); }, ++ [](void* u, int s) { ((Window*) u)->after_cut(s); }, this); + } + + namespace { +@@ -1164,6 +1747,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 ++ am.dense_file = (double) st_->dense_unwired_bytes(); + const double keep = arena_keep(am); + const double avail = am.available; + double budget = arena_budget(am); +@@ -1179,25 +1763,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); ++ const char* epg = std::getenv("STRATA_METAL_PROMPT_GEMM"); ++ const bool pg_on = !(epg && std::string(epg) == "0"); ++ const char* esep = std::getenv("STRATA_METAL_SEPARATE_PROMPT_BUF"); ++ const bool sep_pg = !(esep && std::string(esep) == "0"); ++ const double pg_scratch_need = (pg_on && sep_pg) ? (double) prompt_gemm_need_bytes() : 0.0; ++ if (am.arena_gb < 0 && pg_scratch_need > 0.0) { ++ budget = std::max(1.0 * gb, budget - pg_scratch_need); ++ } + { // the planned total: arena + the Metal buffers (or the process's own count when it is more) + the n-gram table's + // reserve, against the cap (arena_budget.h: footprint_cap) + const FootprintPlan fp = footprint_plan(am, budget); +- std::fprintf(stderr, "strata metal stream: planned resident footprint %.1f GB = arena %.1f + %.1f held before it + %.1f " +- "reserved for the n-gram table; cap %.1f GB = the lesser of the GPU working set %.1f GB and the RAM " +- "%.1f GB less %.1f GB left to macOS and its file cache%s\n", +- fp.total / gb, fp.arena / gb, fp.base / gb, fp.ple / gb, fp.cap / gb, cx_->recommended_working_set() / gb, ++ std::fprintf(stderr, "strata metal stream: planned resident footprint %.1f GB decode (%.1f GB virtual peak during prefill with %.2f GB transient prompt scratch) " ++ "= arena %.1f + %.1f held before it + %.1f reserved for the n-gram table; cap %.1f GB = the lesser of the GPU working set %.1f GB " ++ "and the RAM %.1f GB less %.1f GB left to macOS and its file cache%s\n", ++ fp.total / gb, (fp.total + pg_scratch_need) / gb, pg_scratch_need / gb, ++ fp.arena / gb, fp.base / gb, fp.ple / gb, fp.cap / gb, cx_->recommended_working_set() / gb, + phys / gb, (am.phys - fp.cap) / gb, +- am.arena_gb >= 0 && fp.over() ? " - STRATA_METAL_ARENA_GB is above the cap: kept as set (your choice), but " +- "a Mac that runs short of memory can freeze (watchdog panics seen at 90+ GB " +- "resident of 128)" : ""); ++ am.arena_gb >= 0 && (fp.total + pg_scratch_need > fp.cap) ? " - peak prefill above cap" : ""); + } +- const char* epg = std::getenv("STRATA_METAL_PROMPT_GEMM"); +- const bool pg_on = !(epg && std::string(epg) == "0"); + bool minimal = false; // the arena's least layout (the min slots) is not shrunk further: it goes ahead, filled in steps + for (int tries = 0;; ++tries) { + const double planned = budget; + std::function recheck = [&]() -> uint64_t { return (uint64_t) arena_shrink(am, planned, avail_now()); }; +- if (st_->make_arena(*cx_, (uint64_t) budget, prof, pg_on ? 576 : 64, err, minimal ? std::function() : recheck, ++ const int min_sl = std::getenv("STRATA_METAL_MIN_SLOTS") ? std::max(64, std::atoi(std::getenv("STRATA_METAL_MIN_SLOTS"))) : (pg_on ? (n_exp_ + 64) : 64); ++ if (st_->make_arena(*cx_, (uint64_t) budget, prof, min_sl, err, minimal ? std::function() : recheck, + [&](uint64_t step) { return arena_grow(am, avail_now(), (double) step); })) + break; + if (st_->shrink_to() == 0 || tries == 4) return false; +@@ -1206,7 +1796,7 @@ bool Window::stream_arena(std::string& err) { + budget / gb, st_->shrink_to() / gb); + budget = (double) st_->shrink_to(); + } +- if (pg_on && !prompt_gemm_init(err)) return false; ++ if (pg_on && !sep_pg && !prompt_gemm_init(err)) return false; + { // 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 +1890,8 @@ void Window::kprof(const std::vector& ds, int Tn) { + } + + Window::~Window() { ++ free_prompt_gemm(); ++ if (st_) st_->release_region(); + if (logpos_) std::fclose(logpos_); + if (std::getenv("STRATA_METAL_PROFILE") == nullptr) return; + auto line = [](const char* what, const Prof& x) { +@@ -1327,6 +1919,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; + } ++ { // anti-slop v4 (slop4 above): every row's state is stored at its position; the head rows get the penalties ++ static const slop4::Knobs kn = slop4::knobs(); ++ const bool on = head && (kn.lex > 0.0f || kn.hon > 0.0f || kn.md > 0.0f || kn.colon > 0.0f); ++ const int per = (n_vocab_ + 119) / 120; ++ slop4::State st = slop_at(pos0 - 1); ++ for (int t = 0; t < Tn; ++t) { ++ st = slop4::step(st, tokens[t]); ++ slop_put(pos0 + t) = st; ++ if (!on) continue; ++ slop4::Pen pen[80]; ++ int np = slop4::penalties(st, kn, pen); ++ float* lg = lm_.as() + (size_t) t * n_vocab_; ++ float* pv = apv_.as() + (size_t) t * 120; ++ int32_t* pi = api_.as() + (size_t) t * 120; ++ if (np > 0) { ++ float mx = -INFINITY; ++ for (int k = 0; k < 120; ++k) mx = std::max(mx, pv[k]); ++ if (slop4::needs_hold(pen, np, lg, mx)) { ++ double s = 0.0; ++ for (int64_t i = 0; i < n_vocab_; ++i) s += std::exp((double) (lg[i] - mx)); ++ np += slop4::hold(pen, np, lg, (double) mx + std::log(s), pen + np); ++ } ++ } ++ // no empty turn: <|im_end|> or right after \n\n ++ if (kn.lex > 0.0f && st.t2 == slop4::ETHINK && (st.t1 == slop4::NL2 || st.t1 == slop4::NL)) { ++ pen[np++] = {slop4::IM_END, kn.lex}; ++ pen[np++] = {slop4::THINK, kn.lex}; ++ } ++ if (np == 0) continue; ++ for (int i = 0; i < np; ++i) lg[pen[i].id] -= pen[i].v; ++ // the argmax (tokout_) from the GPU's 120 partition maxima: a partition whose max was pushed down is rescanned ++ bool moved = false; ++ for (int i = 0; i < np; ++i) { ++ const int part = pen[i].id / per; ++ if (pi[part] != pen[i].id) continue; ++ const int b0 = part * per, b1 = std::min(b0 + per, (int) n_vocab_); ++ float bv = -INFINITY; int bi = 0x7fffffff; ++ for (int j = b0; j < b1; ++j) if (!std::isnan(lg[j]) && (lg[j] > bv || (lg[j] == bv && j < bi))) { bv = lg[j]; bi = j; } ++ pv[part] = bv; pi[part] = bi; ++ moved = true; ++ } ++ if (!moved) continue; ++ float bv = -INFINITY; int bi = 0x7fffffff; ++ for (int k = 0; k < 120; ++k) if (pv[k] > bv || (pv[k] == bv && pi[k] < bi)) { bv = pv[k]; bi = pi[k]; } ++ tokout_.as()[t] = bi; ++ } ++ } + 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 +2027,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); +- c.tails.resize(tail_.size() * ts); +- c.dead.resize(dead_.size() * dsz); ++ const size_t nq = tail_.size() + (mtp_on_ ? mtail_.size() : 0); ++ c.tails.resize(nq * ts); ++ c.dead.resize(nq * dsz); + for (size_t j = 0; j < tail_.size(); ++j) { + std::memcpy(c.tails.data() + j * ts, tail_[j].ptr, ts); + std::memcpy(c.dead.data() + j * dsz, dead_[j].ptr, dsz); + } ++ if (mtp_on_ && !mtail_.empty()) { ++ std::memcpy(c.tails.data() + tail_.size() * ts, mtail_[0].ptr, ts); ++ std::memcpy(c.dead.data() + dead_.size() * dsz, mdead_[0].ptr, dsz); ++ } + c.block_pos.clear(); + return true; + } +@@ -1412,6 +2056,7 @@ void Window::reset() { + zero(tail_[j], 4 * 4 * 128); + zero(dead_[j], 4 * 128); + } ++ mtp_cell0_ = mtp_cell1_ = mtp_sink_ = 0; + for (size_t j = 0; j < mtail_.size(); ++j) { + zero(mtail_[j], 4 * 4 * 128); + zero(mdead_[j], 4 * 128); +@@ -1420,8 +2065,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; ++ const size_t nq = tail_.size() + (mtp_on_ ? mtail_.size() : 0); + if (c.gdn.size() != gst_.size() * (gs + hs) || c.ple.size() != (size_t) 4 * 9 * 10240 || +- c.tails.size() != tail_.size() * ts || c.dead.size() != dead_.size() * dsz) { ++ c.tails.size() != nq * ts || c.dead.size() != nq * dsz) { + err = "metal window: the checkpoint is not this engine's (sizes differ)"; + return false; + } +@@ -1436,12 +2082,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); + } ++ if (mtp_on_ && !mtail_.empty()) { ++ std::memcpy(mtail_[0].ptr, c.tails.data() + tail_.size() * ts, ts); ++ std::memcpy(mdead_[0].ptr, c.dead.data() + dead_.size() * dsz, dsz); ++ } ++ return true; ++} ++ ++namespace { ++// one contiguous range of a buffer in the session file ++struct SessionSec { ++ void* p; ++ size_t bytes; ++ size_t off = 0; ++}; ++// the sections in file order from `first` on, each on a 16 KiB boundary; returns the end ++size_t session_place(std::vector& secs, size_t first) { ++ size_t o = scr::align_up(first); ++ for (auto& s : secs) { s.off = o; o = scr::align_up(o + s.bytes); } ++ return o; ++} ++// pread / pwrite of every section, 8 MiB pieces over 4 threads (one stream does not fill the SSD's queue) ++bool session_io(int fd, const std::vector& secs, bool write, std::string& err) { ++ struct Piece { uint8_t* p; size_t n; off_t off; }; ++ std::vector ps; ++ constexpr size_t PIECE = (size_t) 8 << 20; ++ for (const auto& s : secs) ++ for (size_t o = 0; o < s.bytes; o += PIECE) ++ ps.push_back({(uint8_t*) s.p + o, std::min(PIECE, s.bytes - o), (off_t) (s.off + o)}); ++ std::atomic next{0}; ++ std::atomic ok{true}; ++ auto work = [&] { ++ for (size_t i; ok && (i = next++) < ps.size();) { ++ const Piece& q = ps[i]; ++ for (size_t d = 0; d < q.n;) { ++ const ssize_t r = write ? ::pwrite(fd, q.p + d, q.n - d, q.off + (off_t) d) : ::pread(fd, q.p + d, q.n - d, q.off + (off_t) d); ++ if (r <= 0) { if (r < 0 && errno == EINTR) continue; ok = false; break; } ++ d += (size_t) r; ++ } ++ } ++ }; ++ std::vector th; ++ for (int t = 0; t < 4; ++t) th.emplace_back(work); ++ for (auto& t : th) t.join(); ++ if (!ok) err = std::string("metal window: session ") + (write ? "write" : "read") + " failed: " + std::strerror(errno); ++ return ok; ++} ++} // namespace ++ ++// the session's buffers in file order; the cells [0, n) of the positional caches ++static void session_sections(std::vector& secs, const std::vector& gst, const std::vector& ghist, ++ const mtl::Buffer& ple, const std::vector* const* q, int64_t n) { ++ for (size_t i = 0; i < gst.size(); ++i) { ++ secs.push_back({gst[i].ptr, scr::GST_BYTES}); ++ secs.push_back({ghist[i].ptr, scr::GHIST_BYTES}); ++ } ++ secs.push_back({ple.ptr, scr::PLE_BYTES}); ++ const size_t un = (size_t) n; ++ for (int set = 0; set < 2; ++set) { // the 12 QSA layers, then the MTP layer's (when loaded) ++ const std::vector* const* b = q + (set ? 7 : 0); ++ if (set && b[0] == nullptr) break; ++ for (size_t j = 0; j < b[0]->size(); ++j) { ++ secs.push_back({(*b[0])[j].ptr, un * scr::KQ_CELL}); ++ secs.push_back({(*b[1])[j].ptr, un * scr::KQ_CELL}); ++ secs.push_back({(*b[2])[j].ptr, un * scr::SC_CELL}); ++ secs.push_back({(*b[3])[j].ptr, un * scr::SC_CELL}); ++ secs.push_back({(*b[4])[j].ptr, scr::TAIL_BYTES}); ++ secs.push_back({(*b[5])[j].ptr, scr::DEAD_BYTES}); ++ secs.push_back({(*b[6])[j].ptr, un / 4 * scr::POOL_BLOCK}); ++ } ++ } ++} ++ ++bool Window::session_save(const std::string& path, const int32_t* ids, int64_t n_cells, std::string& err) { ++ if (n_cells < 0 || n_cells >= max_cells_) { err = "metal window: session cell count out of range"; return false; } ++ if (mrope_on()) { err = "metal window: sessions of M-RoPE (vision) prompts are not supported"; return false; } ++ if (pend_fold_ > 0) { // as checkpoint_save: the state after the last commit's fold ++ std::vector ds; ++ nb_ = false; ++ fold(ds, pend_fold_, 1 - hp_); ++ cx_->run_graph(ds); ++ pend_fold_ = 0; ++ } ++ const std::vector* q[14] = {&kq_, &vq_, &ksc_, &vsc_, &tail_, &dead_, &pooled_}; ++ if (mtp_on_) { ++ const std::vector* m[7] = {&mkq_, &mvq_, &mksc_, &mvsc_, &mtail_, &mdead_, &mpooled_}; ++ std::copy(m, m + 7, q + 7); ++ } ++ std::vector secs; ++ session_sections(secs, gst_, ghist_[hp_], pleHist_, q, n_cells); ++ scr::SessionHeader h; ++ h.flags = (ple_on_ ? 1u : 0u) | (mtp_on_ ? 2u : 0u); ++ h.n_cells = n_cells; ++ h.n_ids = ids ? n_cells : 0; ++ h.model_bytes = f_.file_bytes; ++ h.model_tensors = f_.tensors.size(); ++ h.rope_base = rope_base_; ++ h.ple_prev[0] = ss_->ple_prev[0]; ++ h.ple_prev[1] = ss_->ple_prev[1]; ++ const size_t head = sizeof h + (size_t) h.n_ids * 4, end = session_place(secs, head); ++ h.payload_bytes = end - head; ++ const std::string tmp = path + ".partial"; ++ const int fd = ::open(tmp.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644); ++ if (fd < 0) { err = "metal window: cannot create " + tmp + ": " + std::strerror(errno); return false; } ++ ::fcntl(fd, F_NOCACHE, 1); ++ bool ok = ::ftruncate(fd, (off_t) end) == 0 && ::pwrite(fd, &h, sizeof h, 0) == (ssize_t) sizeof h && ++ (h.n_ids == 0 || ::pwrite(fd, ids, (size_t) h.n_ids * 4, sizeof h) == (ssize_t) (h.n_ids * 4)); ++ if (!ok) err = "metal window: session header write failed: " + std::string(std::strerror(errno)); ++ ok = ok && session_io(fd, secs, true, err); ++ if (ok) ::fsync(fd); ++ ok = (::close(fd) == 0) && ok; ++ if (ok && std::rename(tmp.c_str(), path.c_str()) != 0) { err = "metal window: rename to " + path + " failed"; ok = false; } ++ if (!ok) ::unlink(tmp.c_str()); ++ return ok; ++} ++ ++bool Window::session_load(const std::string& path, std::vector& ids, int64_t& n_cells, std::string& err) { ++ const int fd = ::open(path.c_str(), O_RDONLY); ++ if (fd < 0) { err = "metal window: cannot open " + path + ": " + std::strerror(errno); return false; } ++ ::fcntl(fd, F_NOCACHE, 1); ++ scr::SessionHeader h, want; ++ struct stat st{}; ++ bool ok = ::fstat(fd, &st) == 0 && ::pread(fd, &h, sizeof h, 0) == (ssize_t) sizeof h; ++ if (!ok || std::memcmp(h.magic, want.magic, 8) != 0 || h.version != want.version) err = "metal window: " + path + " is not a session file"; ++ else if (h.model_bytes != f_.file_bytes || h.model_tensors != f_.tensors.size() || h.rope_base != rope_base_) ++ err = "metal window: the session is another model's"; ++ else if (h.n_cells < 0 || h.n_cells >= max_cells_) err = "metal window: the session is longer than this context (--ctx)"; ++ else if ((h.flags & 1u) != (ple_on_ ? 1u : 0u) || ((h.flags & 2u) && !mtp_on_)) err = "metal window: the session's PLE / MTP setup differs"; ++ if (!err.empty()) { ::close(fd); return false; } ++ const std::vector* q[14] = {&kq_, &vq_, &ksc_, &vsc_, &tail_, &dead_, &pooled_}; ++ if (h.flags & 2u) { ++ const std::vector* m[7] = {&mkq_, &mvq_, &mksc_, &mvsc_, &mtail_, &mdead_, &mpooled_}; ++ std::copy(m, m + 7, q + 7); ++ } ++ std::vector secs; ++ session_sections(secs, gst_, ghist_[hp_], pleHist_, q, h.n_cells); ++ const size_t head = sizeof h + (size_t) h.n_ids * 4, end = session_place(secs, head); ++ if (end - head != h.payload_bytes || (size_t) st.st_size != end) { ::close(fd); err = "metal window: the session file is truncated or of another layout"; return false; } ++ ids.resize((size_t) h.n_ids); ++ ok = h.n_ids == 0 || ::pread(fd, ids.data(), (size_t) h.n_ids * 4, sizeof h) == (ssize_t) (h.n_ids * 4); ++ ok = ok && session_io(fd, secs, false, err); ++ ::close(fd); ++ if (!ok) { if (err.empty()) err = "metal window: session read failed"; reset(); return false; } ++ pend_fold_ = 0; ++ mtp_cell0_ = mtp_cell1_ = mtp_sink_ = 0; // which of the drafter's cells the file holds is not recorded ++ last_T_ = ids.empty() ? 0 : 1; ++ if (!ids.empty()) last_tokens_[0] = ids.back(); ++ slop_scan(ids.data(), (int64_t) ids.size(), 0); ++ ss_->ple_prev[0] = h.ple_prev[0]; ++ ss_->ple_prev[1] = h.ple_prev[1]; ++ n_cells = h.n_cells; + return true; + } + ++bool Window::scr_relocate(int64_t c_src, int64_t c_dst, int64_t n, std::string& err) { ++ if (mrope_on()) { err = "metal window: SCR of M-RoPE (vision) cells is not supported"; return false; } ++ if (n < 0 || c_src < 0 || c_dst < 0 || c_src + n >= max_cells_ || c_dst + n >= max_cells_) { err = "metal window: SCR span out of range"; return false; } ++ if (((c_dst - c_src) & 3) != 0) { err = "metal window: SCR needs a shift by a multiple of 4 cells (pad the new span)"; return false; } ++ const scr::Rot rot(rope_base_, c_dst - c_src); ++ std::vector L[5]; ++ auto add = [&](const std::vector& k, const std::vector& v, const std::vector& ks, ++ const std::vector& vs, const std::vector& p) { ++ for (size_t j = 0; j < k.size(); ++j) { L[0].push_back(&k[j]); L[1].push_back(&v[j]); L[2].push_back(&ks[j]); L[3].push_back(&vs[j]); L[4].push_back(&p[j]); } ++ }; ++ add(kq_, vq_, ksc_, vsc_, pooled_); ++ if (mtp_on_) add(mkq_, mvq_, mksc_, mvsc_, mpooled_); ++ std::atomic ok{true}; ++ std::vector th; // the layers are independent: one thread each ++ for (size_t j = 0; j < L[0].size(); ++j) ++ th.emplace_back([&, j] { ++ if (!scr::relocate_layer(L[0][j]->as(), L[1][j]->as(), L[2][j]->as(), L[3][j]->as(), ++ L[4][j]->as(), c_src, c_dst, n, rot)) ++ ok = false; ++ }); ++ for (auto& t : th) t.join(); ++ if (!ok) err = "metal window: SCR relocation refused"; ++ if (ok && n > 0 && (size_t) (c_src + n) <= slop_.size()) { // the anti-slop states move with their cells ++ slop_put(c_dst + n - 1); ++ std::memmove(slop_.data() + c_dst, slop_.data() + c_src, (size_t) n * sizeof(slop4::State)); ++ } ++ return ok; ++} ++ + bool Window::commit(int n_keep, std::string& err) { + if (n_keep < 1 || n_keep > last_T_) { err = "metal window: commit count out of range"; return false; } + using namespace strata::kernels; ++ if (n_keep < last_T_ && last_T_ <= 4 && tail_bak_.size() == tail_.size() * 4 * 128) { ++ for (int t = n_keep; t < last_T_; ++t) { ++ const size_t s = (size_t) ((last_pos0_ + t) & 3) * 128; ++ for (size_t j = 0; j < tail_.size(); ++j) ++ std::memcpy(tail_[j].as() + s, tail_bak_.data() + j * 4 * 128 + s, 128 * 4); ++ } ++ } + 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 +2339,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")) { ++ // STRATA_METAL_MTP_Q4 / STRATA_METAL_MTP_Q5 = 1: the GEMVs a draft step reads most (q/o_proj, fc_embedding, ++ // fc_hidden, the shared expert) in Q4_0 / Q5_0 (Q4 wins where both). K/V, the router and the ++ // hyper-connections (gr_read_m reads Q8_0 or BF16 only) stay Q8_0. ++ auto on = [](const char* name) { const char* v = std::getenv(name); return v && std::atoi(v) != 0; }; ++ const int low = on("STRATA_METAL_MTP_Q4") ? (int) ref::Q4_0 : on("STRATA_METAL_MTP_Q5") ? (int) ref::Q5_0 : (int) ref::Q8_0; ++ auto low_type = [&](const std::string& nm) { ++ for (const char* s : {"self_attn.q_proj.", "self_attn.o_proj.", "mtp.fc_embedding.", "mtp.fc_hidden.", ++ "mlp.shared_expert."}) ++ if (nm.find(s) != std::string::npos) return low; ++ return (int) ref::Q8_0; ++ }; + size_t tot = 0; + std::vector names; + for (const auto& kv : fm_.tensors) { +@@ -1514,10 +2357,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; ++ if (x.name.find("self_attn.indexer.") != std::string::npos) continue; // the window is never scored + gguf::Tensor q = x; +- q.type = ref::Q8_0; ++ q.type = low_type(x.name); + q.offset = tot; +- q.bytes = ne / 32 * 34; ++ q.bytes = ne / 32 * (q.type == ref::Q4_0 ? 18 : q.type == ref::Q5_0 ? 22 : 34); + mq8t_[x.name] = q; + names.push_back(x.name); + tot += (q.bytes + 255) / 256 * 256; +@@ -1533,12 +2377,43 @@ 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); +- ref::quantize_q8_0(tmp.data(), ne, mq8_.as() + mq8t_[names[i]].offset); ++ const gguf::Tensor& q = mq8t_.at(names[i]); ++ uint8_t* dst = mq8_.as() + q.offset; ++ if (q.type == ref::Q4_0) quant_q4_0_lloyd(tmp.data(), ne, dst); ++ else if (q.type == ref::Q5_0) quant_q5_0(tmp.data(), ne, dst); ++ else ref::quantize_q8_0(tmp.data(), ne, dst); + } + }); + for (auto& x : th) x.join(); + } + } ++ { ++ // What the draft reads from the map (the routed experts, 0.71 GB of Q2_0) is wired, as stream.cpp wires the ++ // dense weights: clean file pages are what macOS drops first under the arena and a prompt's scratch, and the ++ // draft then faults them back from the SSD (mtp_ms 212 -> 432-467 on 2K prompts). What it never reads (the ++ // BF16 now in mq8_, the F32 norms now in mnorm_, the indexer) is let go, interior pages only. ++ const char* ml = std::getenv("STRATA_METAL_MLOCK_DENSE"); ++ const bool mlock_on = !(ml && std::atoi(ml) == 0); ++ const uintptr_t base = (uintptr_t) fm_.map, end = base + fm_.map_bytes, PG = 16384; ++ size_t locked = 0, unlocked = 0, dropped = 0; ++ for (const auto& kv : fm_.tensors) { ++ const gguf::Tensor& x = kv.second; ++ const bool skip = mq8t_.count(x.name) || mnorm_off_.count(x.name) || ++ x.name.find("self_attn.indexer.") != std::string::npos; ++ if (skip) { ++ const uintptr_t p0 = (base + x.offset + PG - 1) & ~(PG - 1), p1 = (base + x.offset + x.bytes) & ~(PG - 1); ++ if (p1 > p0) { madvise((void*) p0, p1 - p0, MADV_DONTNEED); dropped += p1 - p0; } ++ } else if (mlock_on) { ++ const uintptr_t p0 = (base + x.offset) & ~(PG - 1), ++ p1 = std::min(end, (base + x.offset + x.bytes + PG - 1) & ~(PG - 1)); ++ madvise((void*) p0, p1 - p0, MADV_WILLNEED); ++ if (mlock((const void*) p0, p1 - p0) == 0) locked += p1 - p0; ++ else unlocked += p1 - p0; ++ } ++ } ++ std::fprintf(stderr, "strata metal: MTP map: %.2f GB wired (mlock; %.2f GB failed), %.2f GB unused released\n", ++ locked / 1e9, unlocked / 1e9, dropped / 1e9); ++ } + auto A = [&](size_t bytes) { return cx_->alloc(bytes); }; + const int MAXC = (int) max_cells_, MB = MAXC / 4 + 2; + mR_ = A(4 * MAXT * HC * N); memb_ = A(4 * MAXT * N); men_ = A(4 * MAXT * N); me2_ = A(4 * MAXT * N); +@@ -1575,6 +2450,25 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) + mhead_t_.bytes = dvocab_.size() * rb; + } + } ++ const char* align = std::getenv("STRATA_METAL_MTP_SLOP"); ++ draft_slop_ = !(align && std::string(align) == "0"); ++ if (draft_slop_) { ++ dinverse_ = A(4 * n_vocab_); ++ int32_t* inv = dinverse_.as(); ++ std::fill(inv, inv + n_vocab_, -1); ++ if (dvocab_.empty()) { ++ for (int i = 0; i < n_vocab_; ++i) inv[i] = i; ++ } else { ++ for (size_t i = 0; i < dvocab_.size(); ++i) { ++ if (inv[dvocab_[i]] >= 0) { ++ err = "metal mtp: slop alignment requires unique draft vocabulary ids"; ++ return false; ++ } ++ inv[dvocab_[i]] = (int32_t) i; ++ } ++ } ++ dslop_ = A(sizeof(slop4::State)); ++ } + mtp_on_ = true; + { + const int ty = T("token_embd.weight").type; +@@ -1582,7 +2476,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_) { +- lchain_ = cx_->compile({"quant.metal", "units.metal", std::string(STRATA_METAL_ENGINE_KERNELS_DEFAULT) + "/mtp_chain.metal"}); ++ lchain_ = cx_->compile({"quant.metal", "units.metal", std::string(STRATA_METAL_ENGINE_KERNELS_DEFAULT) + "/../slop4.h", ++ std::string(STRATA_METAL_ENGINE_KERNELS_DEFAULT) + "/mtp_chain.metal"}); + 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())); +@@ -1623,25 +2518,23 @@ bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) + + void Window::gr_read_m(std::vector& ds, const std::string& pre, mtl::Buffer& Rb, const mtl::Buffer& bo, + const mtl::Buffer& inj_in, const mtl::Buffer& inj_out, bool apply, const Act& out, bool final_, +- int ntok) { ++ int ntok, bool already_normed) { + const gguf::Tensor& tdn = TQ(pre + ".input_mix_weight_down.weight"); + const gguf::Tensor& tup = TQ(pre + ".input_mix_weight_up.weight"); + const gguf::Tensor& tin = final_ ? tdn : TQ(pre + ".block_inject_weight.weight"); + 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_}; +- mtl::Dispatch d1; +- d1.p = P(lib_, "gr_norm", {}, "gr_norm"); +- d1.args = {mtl::Arg::B(Rb), mtl::Arg::B(bo), mtl::Arg::B(inj_in), mtl::Arg::B(mnorm_, NM(pre + ".hc_norm.weight")), +- mtl::Arg::B(xn_), mtl::Arg::V(a)}; +- d1.grid[0] = HC; d1.grid[1] = ntok; d1.tg[0] = 256; +- push(ds, d1); ++ if (apply || !already_normed) ++ gr_norm(ds, Rb, bo, inj_in, mnorm_, NM(pre + ".hc_norm.weight"), apply, ntok); + 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)); ++ const bool r2 = ntok > 4; ++ std::string dn = q8 ? "gr_down_q8_8_8" : "gr_down_8_8"; ++ if (r2) dn += "_r2"; ++ d2.p = P(lib_, dn, mtl::Constants().i(1, ntok), dn + "#" + std::to_string(ntok)); + d2.args = {mtl::Arg::B(WB, tdn.offset), mtl::Arg::B(WB, tin.offset), mtl::Arg::B(xn_), mtl::Arg::B(lo_), + mtl::Arg::B(inj_out), mtl::Arg::V(a)}; +- d2.grid[0] = LR + HC; d2.tg[0] = 256; ++ d2.grid[0] = r2 ? (LR + HC) / 2 : (LR + HC); d2.tg[0] = 256; + push(ds, d2); + mtl::Dispatch d3; + const char* up = q8 ? "gr_up_q8" : "gr_up"; +@@ -1655,9 +2548,10 @@ void Window::gr_read_m(std::vector& ds, const std::string& pre, m + + void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, size_t src_off, int n, const int32_t* tokens, + int64_t cell0, bool full, int erow, int slot) { +- struct RmsArgs { uint32_t n; float eps; }; ++ static_assert(N == 2560, "mtp_combine_gnorm retains ten values per thread"); ++ struct RmsXpArgs { uint32_t n, xp_n, plane_tok, split, keep; float eps; }; + struct XpArgs { uint32_t n, n_tok, plane_tok, stride; }; +- struct CombArgs { uint32_t n, hc, n_tok, h_stride; }; ++ struct CombNormArgs { uint32_t n, hc; float eps; }; + // ---- the embedding rows (host tokens: dequantized on the CPU into the shared buffer) + const gguf::Tensor& te = T("token_embd.weight"); + const size_t rb = ref::row_bytes(te.type, N); +@@ -1666,12 +2560,13 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, + ref::dequant(te.type, (const uint8_t*) W_.ptr + te.offset + (size_t) tokens[i] * rb, + (membp_ ? *membp_ : memb_).as() + (size_t) (erow + i) * N, N); + const std::string L = "mtp.layers.0."; +- auto rms = [&](const mtl::Buffer& x, size_t xo, const std::string& w, mtl::Buffer& y, int width, int rows) { ++ auto rms_xp = [&](const mtl::Buffer& x, size_t xo, const std::string& w, mtl::Buffer& y, const Act& out, ++ int width, bool split, bool keep) { + mtl::Dispatch d; +- d.p = P(lmtp_, "mtp_rms", {}, "mtp_rms"); +- RmsArgs a{(uint32_t) width, eps_}; +- d.args = {mtl::Arg::B(x, xo), mtl::Arg::B(mnorm_, NM(w)), mtl::Arg::B(y), mtl::Arg::V(a)}; +- d.grid[0] = rows; d.tg[0] = 256; ++ d.p = P(lmtp_, "mtp_rms_xprep", {}, "mtp_rms_xprep"); ++ RmsXpArgs a{(uint32_t) width, N, MAXT, split ? 1u : 0u, keep ? 1u : 0u, eps_}; ++ d.args = {mtl::Arg::B(x, xo), mtl::Arg::B(mnorm_, NM(w)), mtl::Arg::B(y), mtl::Arg::B(out.b), mtl::Arg::V(a)}; ++ d.grid[0] = n; d.tg[0] = 256; + push(ds, d); + }; + auto xp = [&](const mtl::Buffer& x, size_t xo, const Act& out, int width, int rows, int stride) { +@@ -1682,14 +2577,44 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, + d.grid[0] = (width / 16 * rows + 255) / 256; d.tg[0] = 256; + push(ds, d); + }; ++ // the GEMV tile (L, R, S, K) of an MTP tensor ++ auto tile = [&](int ty, int n_in, int n_out, int rows) { ++ // STRATA_METAL_MTP_TILE=L,R,S,K: the tile of the wide ones (n_out > 640, n_in < 6144), one of gemv.metal's GV_SET ++ static const std::array wide = [] { ++ std::array t{0, 0, 0, 0}; ++ if (const char* e = std::getenv("STRATA_METAL_MTP_TILE")) std::sscanf(e, "%d,%d,%d,%d", &t[0], &t[1], &t[2], &t[3]); ++ return t; ++ }(); ++ int Lq = 16, Rr = 1, S = 4, K = 1; ++ const bool m5 = m5_layout(); ++ if (n_in >= 6144) { ++ if (m5 && ty == ref::Q8_0 && rows == 1) { Lq = 32; Rr = 1; S = 8; K = 8; } ++ else { Lq = 32; Rr = 1; S = 8; K = 2; } ++ } else if (n_in == 640) { ++ if (m5 && ty == ref::Q8_0) { Lq = rows >= 3 ? 4 : 32; Rr = 1; S = 4; K = 1; } ++ else { Lq = 8; Rr = 1; S = 4; K = 1; } ++ } else if (n_out <= 640) { ++ if (m5 && ty == ref::Q8_0 && rows == 1 && n_out >= 512) { Lq = 32; Rr = 1; S = 4; K = 4; } ++ else if (m5 && ty == ref::Q8_0 && rows >= 2) { Lq = 16; Rr = 1; S = 4; K = 1; } ++ else { Lq = 32; Rr = 1; S = 4; K = 1; } ++ } else if (wide[0] > 0) { ++ Lq = wide[0]; Rr = wide[1]; S = wide[2]; K = wide[3]; ++ } else if (m5 && ty == ref::Q8_0) { ++ if (rows == 1) { Lq = 32; Rr = 1; S = 4; K = 4; } ++ else if (rows <= 3) { Lq = 32; Rr = 2; S = 4; K = 1; } ++ else { Lq = 16; Rr = 1; S = 4; K = 1; } ++ } else if (m5 && (ty == ref::Q4_0 || ty == ref::Q5_0)) { ++ Lq = 32; Rr = 2; S = 4; K = 1; ++ } else if (ty == ref::Q2_0 || ty == ref::BF16) { ++ Lq = 32; Rr = 1; S = 4; K = 1; ++ } ++ return std::array{Lq, Rr, S, K}; ++ }; + auto gm = [&](const std::string& name, const Act& x, const mtl::Buffer& y, size_t y_stride, int rows, size_t y_off = 0) { + // 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; +- int Lq = 16, Rr = 1, S = 4, K = 1; +- if (n_in >= 6144) { Lq = 32; Rr = 1; S = 8; K = 2; } +- else if (n_in == 640) { Lq = 8; Rr = 1; S = 4; K = 1; } +- else if (n_out <= 640 || ty == ref::Q2_0 || ty == ref::BF16) { Lq = 32; Rr = 1; S = 4; K = 1; } ++ const auto [Lq, Rr, S, K] = tile(ty, n_in, n_out, rows); + 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, +@@ -1702,38 +2627,82 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, + d.grid[0] = (n_out + rpt - 1) / rpt; d.tg[0] = S * 32; + push(ds, d); + }; ++ // gm(a) and gm(b) side by side (same format and shape) as one gemv_cat / gemv_catk at gm's tile: bitwise the two ++ // GEMVs. A draft step only (rows = 1, GPU time even): for a prefill window (16,1,4,1) it measured 1.4-4.3% slower. ++ // Else (STRATA_METAL_GEMV_CAT=0, or no instance) the two GEMVs in one wave. ++ auto gm_cat = [&](const std::string& na, const std::string& nb, const Act& x, const mtl::Buffer& ya, ++ const mtl::Buffer& yb, size_t y_stride, int rows) { ++ static const bool on = !(std::getenv("STRATA_METAL_GEMV_CAT") && std::atoi(std::getenv("STRATA_METAL_GEMV_CAT")) == 0); ++ const gguf::Tensor &ta = TQ(na), &tb = TQ(nb); ++ const int n_in = (int) ta.shape[0], n_out = ta.shape.size() > 1 ? (int) ta.shape[1] : 1, ty = ta.type; ++ const auto [Lq, Rr, S, K] = tile(ty, n_in, n_out, rows); ++ // gemv.metal's GVC (K = 1) and GVCK instances ++ const bool inst = K == 1 ? (Lq == 32 && Rr == 1 && (S == 4 || S == 8)) || (Lq == 16 && Rr == 1 && S == 4) || ++ (Lq == 32 && Rr == 2 && S == 4) || (ty == ref::Q8_0 && Lq == 8 && Rr == 1 && S == 8) ++ : ty == ref::Q8_0 && Lq == 32 && ((Rr == 1 && S == 4 && K == 4) || ++ (Rr == 2 && S == 4 && K == 2) || (Rr == 1 && S == 8 && K == 8)); ++ if (!on || rows != 1 || !inst || tb.type != ty || tb.shape != ta.shape || ++ (ty != ref::Q8_0 && ty != ref::BF16 && ty != ref::F32)) { ++ gm(na, x, ya, y_stride, rows); ++ nb_ = true; ++ gm(nb, x, yb, y_stride, rows); ++ return; ++ } ++ const std::string fn = (K == 1 ? "gemvc_" : "gemvck_") + std::to_string(ty) + "_" + std::to_string(Lq) + "_" + ++ std::to_string(Rr) + "_" + std::to_string(S) + (K == 1 ? "" : "_" + std::to_string(K)); ++ GemvArgs a{(uint32_t) n_in, (uint32_t) (2 * n_out), (uint32_t) rows, (uint32_t) ref::row_bytes(ty, n_in), ++ (uint32_t) n_in, (uint32_t) y_stride}; ++ const uint32_t sp = (uint32_t) n_out; ++ mtl::Dispatch d; ++ d.p = P(lib_, fn, mtl::Constants().i(0, rows), fn + "#" + std::to_string(rows)); ++ d.args = {mtl::Arg::B(MB(na), ta.offset), mtl::Arg::B(x.b, x.plane_bytes() * xkind(ty)), ++ mtl::Arg::B(x.b, x.plane_bytes() * 4), mtl::Arg::B(ya), mtl::Arg::V(a), mtl::Arg::B(MB(nb), tb.offset), ++ mtl::Arg::B(yb), mtl::Arg::V(sp)}; ++ const int rpt = (S / K) * (32 / Lq) * Rr; ++ d.grid[0] = (2 * n_out + rpt - 1) / rpt; d.tg[0] = S * 32; ++ push(ds, d); ++ }; + // ---- e = fc_embedding(rms(embed(tok)) * (1 + w)), h = fc_hidden per stream (rms_10240(R) * (1 + w)), R = h + e +- rms(membp_ ? *membp_ : memb_, (size_t) erow * N * 4, "mtp.pre_fc_norm_embedding.weight", men_, N, n); +- xp(men_, 0, XEn_, N, n, N); ++ // the embedding and the hidden chains side by side (disjoint buffers): a barrier per step, not per dispatch ++ rms_xp(membp_ ? *membp_ : memb_, (size_t) erow * N * 4, "mtp.pre_fc_norm_embedding.weight", ++ men_, XEn_, N, false, false); ++ nb_ = true; ++ // fc_hidden over the n * HC streams as rows [t][c], MAXT per GEMV (one weight read for a draft or two accepted ++ // tokens); the norm prepares the first MAXT, mtp_xprep each further slice from the retained normalized rows. ++ // Every GEMV has >= 4 rows (tile() varies with rows only below 4): a stream's h is bitwise the same for any n. ++ static_assert(MAXT % HC == 0, "a slice holds whole tokens"); ++ const int hrows = n * HC; ++ rms_xp(srcR, src_off, "mtp.pre_fc_norm_hidden.weight", mhn_, XHn_, HC * N, true, hrows > MAXT); + gm("mtp.fc_embedding.weight", XEn_, me2_, N, n); +- rms(srcR, src_off, "mtp.pre_fc_norm_hidden.weight", mhn_, HC * N, n); +- for (int c = 0; c < HC; ++c) { +- 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); ++ nb_ = true; ++ gm("mtp.fc_hidden.weight", XHn_, mh2_, N, std::min(hrows, MAXT)); ++ for (int r = MAXT; r < hrows; r += MAXT) { ++ const int k = std::min(hrows - r, MAXT); ++ xp(mhn_, (size_t) r * N * 4, XHn_, N, k, N); ++ gm("mtp.fc_hidden.weight", XHn_, mh2_, N, k, (size_t) r * N * 4); + } + { + mtl::Dispatch d; +- d.p = P(lmtp_, "mtp_combine", {}, "mtp_combine"); +- CombArgs a{N, HC, (uint32_t) n, (uint32_t) (MAXT * N)}; +- d.args = {mtl::Arg::B(mh2_), mtl::Arg::B(me2_), mtl::Arg::B(mR_), mtl::Arg::V(a)}; +- d.grid[0] = (n * HC * N + 255) / 256; d.tg[0] = 256; ++ d.p = P(lmtp_, "mtp_combine_gnorm", {}, "mtp_combine_gnorm"); ++ CombNormArgs a{N, HC, eps_}; ++ d.args = {mtl::Arg::B(mh2_), mtl::Arg::B(me2_), mtl::Arg::B(mR_), ++ mtl::Arg::B(mnorm_, NM(L + "attn_hyper_connection.hc_norm.weight")), mtl::Arg::B(xn_), mtl::Arg::V(a)}; ++ d.grid[0] = HC; d.grid[1] = n; d.tg[0] = 256; + push(ds, d); + } +- // ---- the attention hyper-connection, K/V (and the indexer keys) into the layer's own cache +- gr_read_m(ds, L + "attn_hyper_connection", mR_, mbo_, minj_[1], minj_[0], false, XA_, false, n); +- gm(L + "self_attn.k_proj.weight", XA_, kcur_, 512, n); +- nb_ = true; +- gm(L + "self_attn.v_proj.weight", XA_, vcur_, 512, n); +- nb_ = true; +- gm(L + "self_attn.indexer.index_qk_proj.weight", XA_, midx_, 640, n); ++ // ---- the attention hyper-connection, K/V into the layer's own cache. No indexer: the attention below is a dense ++ // window (never qsa_score / qsa_topk), so its projection and keys would be read by nobody. ++ gr_read_m(ds, L + "attn_hyper_connection", mR_, mbo_, minj_[1], minj_[0], false, XA_, false, n, true); ++ gm_cat(L + "self_attn.k_proj.weight", L + "self_attn.v_proj.weight", XA_, kcur_, vcur_, 512, n); + if (full) { + nb_ = true; + gm(L + "self_attn.q_proj.weight", XA_, qfull_, 12288, n); + } + QsaArgs qa{(uint32_t) cell0, 0, (uint32_t) max_cells_, (uint32_t) n, eps_}; + { ++ // rows 0..27 (q, K, V) with the queries, else rows 24..27 (QSA_ROW0 = 24) + mtl::Dispatch d; +- d.p = P(lqsa_, "qsa_prep", qsa_c(), qsa_k()); ++ d.p = full ? P(lqsa_, "qsa_prep", qsa_c(), qsa_k()) : P(lqsa_, "qsa_prep", qsa_c().i(28, 24), qsa_k() + "#kv"); + 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 +2710,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_)}; +- d.grid[0] = 33; d.grid[1] = n; d.tg[0] = 256; ++ d.grid[0] = full ? 28 : 4; d.grid[1] = n; d.tg[0] = 256; + push(ds, d); + } + if (!full) return; +- // ---- dense attention over the window of cells up to this one (the draft layer's rule), the output gate, o_proj ++ // ---- dense attention over the attention sink and the local sliding window of cells up to this one (default 64 + ++ // 256 cells), the output gate, o_proj + { ++ static const int64_t mtp_win = std::getenv("STRATA_METAL_MTP_WIN") ? std::max(64, std::min(2050, std::atoll(std::getenv("STRATA_METAL_MTP_WIN")))) : 256; + int* ids = mqids_.as() + (size_t) slot * 2051; +- const int64_t first = std::max(0, cell0 - 2050); +- for (int64_t c = first; c <= cell0; ++c) ids[c - first] = (int) c; ++ // the window's cells past mtp_cell1_ are this round's drafts ++ const int64_t first = std::min(cell0, std::max(mtp_cell0_, cell0 - mtp_win)); ++ const int64_t ns = std::min({mtp_sink_, first, 2050 - (cell0 - first)}); ++ for (int64_t c = 0; c < ns; ++c) ids[c] = (int) c; ++ for (int64_t c = first; c <= cell0; ++c) ids[ns + c - first] = (int) c; ++ QsaArgs qa_at = qa; ++ qa_at.pos0 = (uint32_t) (ns + cell0 - first); ++ const uint32_t nch = (uint32_t) std::max(1, (qa_at.pos0 + 1 + 63) / 64); + mtl::Dispatch at; + at.p = P(lqsa_, "qsa_attn", {}, "qsa_attn"); + at.args = {mtl::Arg::B(qcur_), mtl::Arg::B(mkq_[0]), mtl::Arg::B(mvq_[0]), mtl::Arg::B(mksc_[0]), mtl::Arg::B(mvsc_[0]), +- mtl::Arg::B(mqids_, (size_t) slot * 2051 * 4), mtl::Arg::B(part_), mtl::Arg::V(qa)}; +- at.grid[0] = 33; at.grid[1] = 2; at.grid[2] = 1; at.tg[0] = 256; ++ mtl::Arg::B(mqids_, (size_t) slot * 2051 * 4), mtl::Arg::B(part_), mtl::Arg::V(qa_at)}; ++ at.grid[0] = nch; at.grid[1] = 2; at.grid[2] = 1; at.tg[0] = 256; + push(ds, at); + mtl::Dispatch m; + m.p = P(lqsa_, "qsa_merge", {}, "qsa_merge"); +- QsaArgs am = qa; ++ QsaArgs am = qa_at; + 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; +@@ -1769,9 +2746,7 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, + gr_read_m(ds, L + "mlp_hyper_connection", mR_, mbo_, minj_[0], minj_[1], true, XF_, false, 1); + gm(L + "mlp.gate.weight", XF_, logits_, 512, 1); + nb_ = true; +- gm(L + "mlp.shared_expert.gate_proj.weight", XF_, gsh_, NFF, 1); +- nb_ = true; +- gm(L + "mlp.shared_expert.up_proj.weight", XF_, ush_, NFF, 1); ++ gm_cat(L + "mlp.shared_expert.gate_proj.weight", L + "mlp.shared_expert.up_proj.weight", XF_, gsh_, ush_, NFF, 1); + nb_ = true; + gm(L + "mlp.shared_expert_gate.weight", XF_, sg_, 1, 1); + { +@@ -1790,13 +2765,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(); +- const int dl = dq2 || (m4 && td.type == ref::Q8_0) ? 8 : 4, +- dr = dq2 || (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1)) ? 1 : 2; ++ const int dl = dq2 || (!m4 && td.type == ref::IQ4_NL) || (m4 && td.type == ref::Q8_0) ? 8 : 4, ++ dr = dq2 ? (m4 ? 1 : 2) : (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1)) ? 1 : 2; + a.gx_gu = NFF / (gs * (32 / gl) * gr); + a.gx_dn = N / (ds_ * (32 / dl) * dr); + mtl::Dispatch r; + r.p = P(lib_, "moe_route", {}, "moe_route"); +- r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a)}; ++ r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a), mtl::Arg::B(rcounter_)}; + r.tg[0] = 512; + push(ds, r); + { +@@ -1829,18 +2804,14 @@ void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, + dd.indirect = &plan_; + dd.ind_off = P_IDN * 4; + push(ds, dd); +- GlueArgs ca{N, 1, N, TOPK}; +- mtl::Dispatch cb; +- cb.p = P(lib_, "moe_combine_sh", {}, "moe_combine_sh"); +- cb.args = {mtl::Arg::B(ye_), mtl::Arg::B(plan_), mtl::Arg::B(ysh_), mtl::Arg::B(sg_), mtl::Arg::B(mbo_), mtl::Arg::V(ca)}; +- cb.grid[0] = (N + 255) / 256; cb.tg[0] = 256; +- push(ds, cb); ++ comb_defer(ds, mbo_, 1); // into the final mixer's gr_norm + } + // the layer's write (R updated in place: the next draft step's input residual) and the final mixer + gr_read_m(ds, "mtp.hyper_connection_mixer", mR_, mbo_, minj_[1], minj_[0], true, XA_, true, 1); + const int64_t nv = dvocab_.empty() ? n_vocab_ : (int64_t) dvocab_.size(); + if (dvocab_.empty()) gemv(ds, T("output.weight"), XA_, mlm_, (size_t) n_vocab_, 1); + else gemv(ds, mhead_t_, XA_, mlm_, (size_t) nv, 1, 0, &mhead_); ++ if (draft_slop_ && chain_) mtp_slop(ds, slot); + { + ArgArgs a{(uint32_t) nv, 120}; + mtl::Dispatch d1, d2; +@@ -1858,26 +2829,91 @@ 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; } +- mrope_sync(cell0, n); +- // every group of MAXT rows in one command buffer: the embedding rows of all of them in mpemb_ (row c + i) +- if ((int64_t) mpemb_rows_ < n) { ++ static const int64_t tail_cap = [] { ++ if (const char* e = std::getenv("STRATA_METAL_MTP_PREFILL_TAIL")) ++ return std::max(64, std::atoll(e)); ++ const char* e = std::getenv("STRATA_METAL_MTP_WIN"); ++ return e ? std::clamp(std::atoll(e), 64, 2050) : int64_t{256}; ++ }(); ++ const int64_t skip = (n > tail_cap) ? (n - tail_cap) : 0; ++ const int64_t rem = n - skip; ++ if (rem <= 0) return true; ++ // the skipped rows' head is the attention sink's (cells below mtp_sink_cap()): those too, before the tail ++ const int64_t nsk = std::max(0, std::min(skip, mtp_sink_cap() - cell0)); ++ mtp_wrote(cell0, nsk); ++ mtp_wrote(cell0 + skip, rem); ++ mrope_sync(cell0, nsk); ++ mrope_sync(cell0 + skip, rem); ++ if ((int64_t) mpemb_rows_ < nsk + rem) { + if (mpemb_rows_ > 0) cx_->free(mpemb_); +- mpemb_rows_ = (size_t) std::max(n, 64); ++ mpemb_rows_ = (size_t) std::max(nsk + rem, 64); + mpemb_ = cx_->alloc(4 * mpemb_rows_ * N); + } + membp_ = &mpemb_; + std::vector ds; + nb_ = false; +- for (int64_t c = 0; c < n; c += MAXT) { +- const int k = (int) std::min(MAXT, n - c); +- mtp_encode(ds, mpsrc_ ? *mpsrc_ : mtp_src_ ? *mtp_src_ : R_, (size_t) c * HC * N * 4, k, next_tokens + c, +- cell0 + c, false, (int) c, 0); ++ const mtl::Buffer& src = mpsrc_ ? *mpsrc_ : mtp_src_ ? *mtp_src_ : R_; ++ for (int64_t c = 0; c < nsk; c += MAXT) { ++ const int k = (int) std::min(MAXT, nsk - c); ++ mtp_encode(ds, src, (size_t) c * HC * N * 4, k, next_tokens + c, cell0 + c, false, (int) c, 0); ++ } ++ for (int64_t c = 0; c < rem; c += MAXT) { ++ const int k = (int) std::min(MAXT, rem - c); ++ mtp_encode(ds, src, (size_t) (skip + c) * HC * N * 4, k, next_tokens + skip + c, cell0 + skip + c, false, ++ (int) (nsk + c), 0); + } + membp_ = nullptr; + cx_->run_graph(ds); + return true; + } + ++void Window::mtp_slop(std::vector& ds, int j) { ++ DraftSlopArgs a = draft_slop_args_; ++ a.j = (uint32_t) j; ++ mtl::Dispatch d; ++ d.p = P(lchain_, "mtp_slop_pen", {}, "mtp_slop_pen"); ++ d.args = {mtl::Arg::B(mlm_), mtl::Arg::B(dinverse_), mtl::Arg::B(dout_), ++ mtl::Arg::B(dslop_), mtl::Arg::V(a)}; ++ d.tg[0] = 256; ++ push(ds, d); ++} ++ ++void Window::mtp_slop_host(slop4::State state) { ++ const slop4::Knobs kn = slop4::knobs(); ++ const int nv = dvocab_.empty() ? (int) n_vocab_ : (int) dvocab_.size(); ++ float* lg = mlm_.as(); ++ const int32_t* inv = dinverse_.as(); ++ slop4::Pen pen[80]; ++ int np = slop4::penalties(state, kn, pen); ++ // hold is calibrated on the draft head, before the empty-turn guard, just as in finish. ++ const float mx = lg[tokout_.as()[0]]; ++ bool hold = false; ++ for (int i = 0; i < np; ++i) ++ if (inv[pen[i].id] >= 0 && lg[inv[pen[i].id]] > mx - 12.0f) hold = true; ++ if (hold) { ++ double sum = 0.0, removed = 0.0; ++ for (int i = 0; i < nv; ++i) sum += std::exp((double) (lg[i] - mx)); ++ for (int i = 0; i < np; ++i) { ++ const int di = inv[pen[i].id]; ++ if (di >= 0) removed += std::exp((double) (lg[di] - mx)) * (1.0 - std::exp(-(double) pen[i].v)); ++ } ++ const float v = (float) -std::log1p(-std::min(removed / sum, 0.999999)); ++ for (int32_t id : slop4::kHold) pen[np++] = {id, v}; ++ } ++ if (kn.lex > 0.0f && state.t2 == slop4::ETHINK && (state.t1 == slop4::NL || state.t1 == slop4::NL2)) { ++ pen[np++] = {slop4::IM_END, kn.lex}; ++ pen[np++] = {slop4::THINK, kn.lex}; ++ } ++ for (int i = 0; i < np; ++i) if (inv[pen[i].id] >= 0) lg[inv[pen[i].id]] -= pen[i].v; ++ int best = 0x7fffffff; ++ float bv = -INFINITY; ++ for (int i = 0; i < nv; ++i) ++ if (!std::isnan(lg[i]) && (lg[i] > bv || (lg[i] == bv && i < best))) { ++ bv = lg[i]; best = i; ++ } ++ tokout_.as()[0] = best; ++} ++ + 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,9 +2943,42 @@ 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; ++ slop4::State draft_state{}; ++ if (draft_slop_) { ++ // tokens are verifier outputs: row a predicts the token following the committed input at p+a. ++ draft_state = slop4::step(slop_at(p + a), tokens[a]); ++ draft_slop_args_ = {draft_state, slop4::knobs(), ++ (uint32_t) (dvocab_.empty() ? n_vocab_ : dvocab_.size()), 0, (uint32_t) n_vocab_}; ++ } + 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]; ++ const bool prefetch_ple = ple_on_ && ss_ && ss_->ple.table && ss_->ple.table->is_open(); ++ int32_t ple_prev[2] = {}; ++ auto prefetch = [&](int32_t token) { ++ uint32_t rows[strata::kernels::PLE_N_HEADS]; ++ strata::kernels::ngram_rows(&token, ple_prev, 1, ss_->ple.consts, rows); ++ ss_->ple.table->prefetch_rows(rows); ++ ple_prev[0] = ple_prev[1]; ++ ple_prev[1] = token; ++ }; ++ if (prefetch_ple) { ++ // commit advanced through the kept inputs, not the verifier's next token. ++ ple_prev[0] = ss_->ple_prev[0]; ++ ple_prev[1] = ss_->ple_prev[1]; ++ prefetch(tok); ++ } + if (chain_ && max_steps > 0) { + std::vector ds; + nb_ = false; +@@ -1923,10 +2992,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 (!policy.keep(j, pj, min_p, cumulative)) break; + drafts[j] = dout_.as()[j]; + if (probs) probs[j] = pj; + ++n; +- if (min_p > 0.0f && pj < min_p) break; + } + if (n_drafts) *n_drafts = n; + return true; +@@ -1941,6 +3010,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); ++ if (draft_slop_) mtp_slop_host(draft_state); + 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 +@@ -1950,11 +3020,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 (!policy.enabled() && min_p > 0.0f && pj < min_p) break; + } + if (n_drafts) *n_drafts = n; + return true; +@@ -1990,6 +3062,10 @@ std::string gu2_tile(int t) { + return "_16_64_16_2"; + } + std::string dn2_tile(int t) { return t == ref::IQ4_NL || t == ref::Q5_1 || t == ref::Q8_0 ? "_16_64_32_2" : "_16_64_16_2"; } ++// a prompt chunk's bo in yr's rows (prompt_gemm_init): off while STRATA_METAL_DUMP_MOE reads yr after the next attention ++bool prompt_bo_alias() { return std::getenv("STRATA_METAL_DUMP_MOE") == nullptr; } ++// an output left unbound (nil): the kernel skips it (gemm_hc_up's fp32 mixed) ++const mtl::Buffer kUnbound{}; + } // namespace + + struct Window::PromptGemm { +@@ -1998,11 +3074,91 @@ struct Window::PromptGemm { + bool k4 = false, gu_v3 = false, dn_v3 = false, flash = false; + 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; ++ XS, HS, YE, yr, gsh, ush, hh, ys, sgv, emb, embh, pkey, pval, normW; ++ mtl::Buffer ids3; ++ mtl::Buffer logits2, ids2, wts2, Wr2; // SSD streaming: layer l + 1's router on layer l's FFN input (the read-ahead) + mtl::Buffer WD, WU, Wqkv, Wz, Wal, Wbe, Wout, Wq, Wk, Wv, Wiq, Wik, Wr, Wsg, Wshg, Wshu, Wshd, Wpk, Wpv; ++ std::vector all_bufs; + }; + ++ ++void Window::free_prompt_gemm() { ++ if (!pg_) return; ++ mtp_src_ = nullptr; ++ for (mtl::Buffer& b : pg_->all_bufs) cx_->free(b); ++ delete pg_; ++ pg_ = nullptr; ++} ++ + int Window::prompt_chunk_rows() const { return pg_ ? pg_->MT : 0; } ++void Window::dump_moe_layer(const char* pfx, int l, bool dump_in, bool dump_out) { ++ if (!pfx || !pg_) return; ++ static std::string s_pfx; ++ static size_t s_prev_T = 0; ++ const size_t T = (size_t) seg_T_; ++ const bool first_chunk = (last_T_ == 0 && s_pfx != pfx); ++ if (l == 0 && dump_in && first_chunk) { ++ s_pfx = pfx; ++ s_prev_T = 0; ++ } ++ auto wr = [&](const char* ext, const void* ptr, size_t row_bytes) { ++ std::string fn = std::string(pfx) + ext; ++ const size_t chunk_bytes = T * row_bytes; ++ if (s_prev_T == 0) { ++ FILE* fp = std::fopen(fn.c_str(), l == 0 ? "wb" : "r+b"); ++ if (!fp) fp = std::fopen(fn.c_str(), "wb"); ++ if (!fp) { std::perror(("dump_moe fopen " + fn).c_str()); std::exit(2); } ++ if (std::fseek(fp, (long)((size_t) l * chunk_bytes), SEEK_SET) != 0 || ++ std::fwrite(ptr, 1, chunk_bytes, fp) != chunk_bytes) { ++ std::perror(("dump_moe fwrite " + fn).c_str()); std::exit(2); ++ } ++ std::fclose(fp); ++ } else { ++ // Multi-chunk prompt: splice new T rows into [NL, s_prev_T + T, row_bytes] ++ const size_t old_l_bytes = s_prev_T * row_bytes; ++ const size_t new_l_bytes = (s_prev_T + T) * row_bytes; ++ static std::map> prev_files; ++ std::vector& prev_file = prev_files[ext]; ++ if (l == 0) { ++ FILE* rf = std::fopen(fn.c_str(), "rb"); ++ if (!rf) { std::perror(("dump_moe multi-chunk r " + fn).c_str()); std::exit(2); } ++ prev_file.resize((size_t) 48 * old_l_bytes); ++ if (std::fread(prev_file.data(), 1, prev_file.size(), rf) != prev_file.size()) { ++ std::perror(("dump_moe multi-chunk fread " + fn).c_str()); std::exit(2); ++ } ++ std::fclose(rf); ++ } ++ FILE* wf = std::fopen(fn.c_str(), l == 0 ? "wb" : "r+b"); ++ if (!wf) { std::perror(("dump_moe multi-chunk w " + fn).c_str()); std::exit(2); } ++ std::fseek(wf, (long)((size_t) l * new_l_bytes), SEEK_SET); ++ if (std::fwrite(prev_file.data() + (size_t) l * old_l_bytes, 1, old_l_bytes, wf) != old_l_bytes || ++ std::fwrite(ptr, 1, chunk_bytes, wf) != chunk_bytes) { ++ std::perror(("dump_moe multi-chunk fwrite " + fn).c_str()); std::exit(2); ++ } ++ std::fclose(wf); ++ if (l == 47) prev_file.clear(); ++ } ++ }; ++ if (dump_in) { ++ wr(".mixed", pg_->mixed.ptr, N * sizeof(float)); ++ wr(".ids", pg_->ids.ptr, TOPK * sizeof(int32_t)); ++ wr(".wts", pg_->wts.ptr, TOPK * sizeof(float)); ++ } ++ if (dump_out) { ++ wr(".yr", pg_->yr.ptr, N * sizeof(float)); ++ wr(".hh", pg_->hh.ptr, NFF * sizeof(uint16_t)); ++ wr(".sgv", pg_->sgv.ptr, sizeof(float)); ++ if (l == 47) { ++ std::string fn_r = std::string(pfx) + ".R_post47"; ++ FILE* fpr = std::fopen(fn_r.c_str(), s_prev_T == 0 ? "wb" : "ab"); ++ if (fpr) { std::fwrite(pg_->R.ptr, 1, T * 4 * N * sizeof(float), fpr); std::fclose(fpr); } ++ std::string fn_i = std::string(pfx) + ".inj47"; ++ FILE* fpi = std::fopen(fn_i.c_str(), s_prev_T == 0 ? "wb" : "ab"); ++ if (fpi) { std::fwrite(pg_->inj[0].ptr, 1, T * 4 * sizeof(float), fpi); std::fclose(fpi); } ++ s_prev_T += T; ++ } ++ } ++} + + bool Window::prompt_gemm_on(int64_t n) const { + static const bool on = [] { +@@ -2018,30 +3174,41 @@ bool Window::prompt_gemm_on(int64_t n) const { + + bool Window::prompt_gemm_init(std::string& err) { + if (pg_) return true; +- if (stream_ && !st_->ready()) return stream_arena(err); // it makes the prompt buffers after the arena ++ if (stream_ && !st_->ready() && !stream_arena(err)) return false; // it makes the prompt buffers after the arena + 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 + // same). The longer chunk's buffers (4.3 GB) come out of the arena only while a prompt runs, so it is the default on a +- // big Mac; a small one keeps 1024 (the buffers would cost it cache). ++ // big Mac; a 24 GB one takes 2048 (1.19 GB at a 64K context, 1024's 0.74), a smaller one keeps 1024. + uint64_t phys = 0; + size_t plen = sizeof phys; + sysctlbyname("hw.memsize", &phys, &plen, nullptr, 0); +- const int stream_chunk = phys >= (96ull << 30) ? 4096 : 1024; ++ const int stream_chunk = phys >= (96ull << 30) ? 4096 : phys >= (24ull << 30) ? 2048 : 1024; + const int C = std::max(64, ec ? std::atoi(ec) : (stream_ ? stream_chunk : 2048)) / 64 * 64; + auto pg = std::make_unique(); + PromptGemm& p = *pg; + p.MT = C; + p.QS = std::min(C, 256); +- p.lib = cx_->compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "moe_mm.metal", "gr.metal", "gemv.metal", +- "prompt.metal"}); +- p.lgdn = cx_->compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); +- p.lqsa = cx_->compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); ++ if (pg_lib_ < 0) { ++ pg_lib_ = cx_->compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "moe_mm.metal", "gr.metal", "gemv.metal", ++ "prompt.metal"}); ++ pg_lgdn_ = cx_->compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); ++ pg_lqsa_ = cx_->compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); ++ } ++ p.lib = pg_lib_; ++ p.lgdn = pg_lgdn_; ++ p.lqsa = pg_lqsa_; + 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; + const char* eq = std::getenv("STRATA_METAL_PROMPT_QSA"); // v1 | flash (default: flash with the M5's layout) + p.flash = eq ? std::string(eq) == "flash" : p.k4; ++ // the long-context QSA selection first runs at chunk 1 (pos0 + T > 2051): compiled there, it stalls the host ~200 ms ++ P(p.lqsa, "qsa_score_mm", {}, "pg:qsa_score_mm"); ++ P(p.lqsa, "qsa_topk", {}, "pg:qsa_topk"); ++ for (int kvn : {1152, 1664}) ++ P(p.lqsa, "qsa_prep", (p.flash ? qsa_c().b(25, true) : qsa_c()).i(26, kvn), ++ "pg:" + (p.flash ? qsa_k() + "#qf16" : qsa_k()) + "#kv" + std::to_string(kvn)); + const size_t MT = (size_t) C; + p.maxrows = C * TOPK + NEXP * p.TM; + p.nblk_max = (C * TOPK + 1023) / 1024; +@@ -2054,8 +3221,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{}; } +- if (region == nullptr) return cx_->alloc(bytes); +- mtl::Buffer r = cx_->wrap(region, b); ++ if (region == nullptr) { ++ mtl::Buffer r = cx_->alloc(bytes); ++ p.all_bufs.push_back(r); ++ return r; ++ } ++ mtl::Buffer r = cx_->wrap(region, b, /*count_allocated=*/false); ++ p.all_bufs.push_back(r); + region += b; + return r; + }; +@@ -2063,31 +3235,53 @@ bool Window::prompt_gemm_init(std::string& err) { + auto F16 = [&](size_t n) { return AB(2 * n); }; + auto make = [&] { + const size_t mr = (size_t) p.maxrows; +- p.R = F32(MT * HC * N); p.bo = F32(MT * N); p.inj[0] = F32(MT * HC); p.inj[1] = F32(MT * HC); p.xnh = F16(MT * HC * N); +- p.V = F32(MT * (LR + HC)); p.loh = F16(MT * LR); p.mixed = F32(MT * N); p.mixh = F16(MT * N); +- p.qkv = F32(MT * 10240); p.z = F32(MT * 6144); p.al = F32(MT * 48); p.be = F32(MT * 48); p.yg = F32(MT * 6144); +- p.ygh = F16(MT * 6144); p.gqn = F32(MT * 16 * 128); p.gkn = F32(MT * 16 * 128); p.gvc = F32(MT * 6144); ++ p.R = F32(MT * HC * N); p.inj[0] = F32(MT * HC); p.inj[1] = F32(MT * HC); p.xnh = F16(MT * HC * N); ++ p.V = p.k4 ? F32(16) : F32(MT * (LR + HC)); p.loh = F16(MT * LR); p.mixed = prompt_bo_alias() ? F32(16) : F32(MT * N); p.mixh = F16(MT * N); ++ // QSA (l % 4 == 3) and GDN (l % 4 != 3) layers are mutually exclusive and consume their scratch within the layer, ++ // so QSA's layer-local activations and dequantized weight buffers share GDN's allocations. ++ // qkv also holds the expert rows' HS (gate/up -> down, between the attention and layer 0's PLE projections) ++ p.qkv = AB(std::max(4 * MT * 12288, 2 * mr * NFF)); p.z = F32(MT * 6144); p.al = F32(MT * 96); p.yg = F32(16); ++ p.ygh = F16(MT * 6144); p.gqn = F32(MT * 16 * 128); p.gkn = F32(MT * 16 * 128); ++ p.gvc = F32(MT * std::max(6144, (size_t) MB)); + p.ggb = F32(MT * 96); p.gkq = F32(MT * 16); p.go = F32(MT * 6144); +- p.qfull = F32(MT * 12288); p.kcur = F32(MT * 512); p.vcur = F32(MT * 512); p.iqr = F32(MT * 512); p.ikr = F32(MT * 128); +- p.qcur = F32(MT * 24 * 256); p.qcurh = F16(MT * 24 * 256); p.qidx = F32(MT * 4 * 128); p.scores = F32(MT * (size_t) MB); +- p.qids = AB(4 * MT * 2051); ++ // kcur is QSA's fused K | V | indexer-K | indexer-Q projection, [T][1152 or 1664] in gqn's MT * 2048 floats: vcur, ++ // ikr and iqr are its column slices (prompt_chunk), so they have no buffers of their own ++ p.qfull = p.qkv; p.kcur = p.gqn; ++ p.qcur = p.flash ? F32(16) : p.z; p.qcurh = p.ygh; p.qidx = F32(MT * 4 * 128); p.scores = p.gvc; ++ p.qids = p.go; + 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.XQ = p.flash ? F32(16) : F32(4 * MT * 6144 + 3 * (size_t) p.QS * 6144 + MT * 384); p.xqh = p.flash ? p.z : F16(MT * 6144); // flash's QSA leaves z idle + p.logits = F32(MT * NEXP); p.ids = AB(4 * MT * TOPK); p.wts = F32(MT * TOPK); ++ if (stream_) { ++ p.logits2 = F32(MT * NEXP); p.ids2 = AB(4 * MT * TOPK); p.wts2 = F32(MT * TOPK); ++ p.Wr2 = F16((size_t) NEXP * N); ++ if (std::getenv("STRATA_METAL_PROMPT_PRED2") && std::atoi(std::getenv("STRATA_METAL_PROMPT_PRED2")) != 0) ++ p.ids3 = AB(4 * MT * TOPK); ++ } + 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); +- p.YE = F32(mr * N); p.yr = F32(MT * N); p.gsh = F32(MT * NFF); p.ush = F32(MT * NFF); p.hh = F16(MT * NFF); ++ p.stok = AB(4 * mr); p.sw = F32(mr); p.pos = AB(4 * MT * TOPK); p.XS = F16(mr * N); p.HS = p.qkv; ++ // with mm_combine folded into every reader yr is never written; bo (the attention write, read by the FFN's HC ++ // read) shares ys's rows: ys lives from the shared expert to the next layer's HC read, before its attention ++ // GEMM writes bo ++ p.YE = F32(mr * N); p.yr = prompt_bo_alias() ? F32(16) : F32(MT * N); ++ // the shared expert runs between the MoE and the next attention (or layer 0's PLE), when the GDN's q/k rows and ++ // ygh are idle; hh keeps its own while STRATA_METAL_DUMP_MOE reads it after the next attention ++ p.gsh = p.gqn; p.ush = p.gkn; p.hh = prompt_bo_alias() ? p.ygh : F16(MT * NFF); + p.ys = F32(MT * N); p.sgv = F32(MT); +- p.WD = F16((size_t) (LR + HC) * HC * N); p.WU = F16((size_t) HC * N * LR); p.Wqkv = F16((size_t) 10240 * N); +- p.Wz = F16((size_t) 6144 * N); p.Wal = F16((size_t) 48 * N); p.Wbe = F16((size_t) 48 * N); p.Wout = F16((size_t) N * 6144); +- p.Wq = F16((size_t) 12288 * N); p.Wk = F16((size_t) 512 * N); p.Wv = F16((size_t) 512 * N); p.Wiq = F16((size_t) 512 * N); +- p.Wik = F16((size_t) 128 * N); p.Wr = F16((size_t) NEXP * N); p.Wsg = F16((size_t) N); p.Wshg = F16((size_t) NFF * N); ++ p.bo = prompt_bo_alias() ? p.ys : F32(MT * N); ++ p.WD = F16((size_t) (LR + HC) * HC * N); p.WU = F16((size_t) HC * N * LR); p.Wqkv = F16((size_t) 12288 * N); ++ p.Wz = F16((size_t) 6144 * N); p.Wal = F16((size_t) 512 * N); p.Wbe = F16((size_t) 512 * N); p.Wout = F16((size_t) N * 6144); ++ // Wk holds the fused projection's rows (attn_k, attn_v, indexer k, indexer q): 1664 of Wz's 6144 ++ p.Wq = p.Wqkv; p.Wk = p.Wz; ++ p.Wr = F16((size_t) NEXP * N); p.Wsg = F16((size_t) N); p.Wshg = F16((size_t) NFF * N); + p.Wshu = F16((size_t) NFF * N); p.Wshd = F16((size_t) N * NFF); + if (ple_on_) { +- p.emb = F32(MT * N); p.embh = F16(MT * N); p.pkey = F32(MT * 10240); p.pval = F32(MT * N); p.keyn = F32(MT * 10240); +- p.qn = F32(MT * 10240); p.gated = F32(MT * 10240); p.normW = F32(MT * 10240); +- p.Wpk = F16((size_t) 10240 * N); p.Wpv = F16((size_t) N * N); ++ // the PLE projections run after layer 0's MoE, when layer 0's attention scratch is idle until layer 1; emb ++ // (the gather's, on the CPU beside layer 0) keeps its own. normW lives from the PLE norms to ple_hist_adv, ++ // after layer 0's expert gate/up read XS and before layer 1's gather writes it ++ p.emb = F32(MT * N); p.normW = 2 * mr * N >= 4 * MT * 10240 ? p.XS : F32(MT * 10240); ++ p.embh = p.ygh; p.pkey = p.qkv; p.pval = p.z; p.Wpk = p.Wqkv; p.Wpv = p.Wz; + } + }; + make(); +@@ -2097,14 +3291,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(); ++ if (need != prompt_gemm_need_bytes()) { err = "metal window: prompt_gemm_need_bytes() != need"; return false; } + pg_ = pg.release(); ++ const char* esep_log = std::getenv("STRATA_METAL_SEPARATE_PROMPT_BUF"); ++ const bool sep_pg_log = !(esep_log && std::string(esep_log) == "0"); + std::fprintf(stderr, "strata metal: prompts on the matrix units in chunks of %d (experts %s, QSA %s, matmul2d layout %s; " + "%.2f GB of buffers%s)\n", C, p.gu_v3 ? "v3" : "v2", p.flash ? "flash" : "v1", p.k4 ? "K4" : "other", +- need / 1e9, stream_ ? ", in the expert arena while a prompt runs" : ""); ++ need / 1e9, stream_ ? (sep_pg_log ? ", in a separate transient mmap while a prompt runs" : ", in the expert arena while a prompt runs") : ""); + (void) err; + return true; + } + ++uint64_t Window::prompt_gemm_need_bytes() const { ++ const char* ec = std::getenv("STRATA_METAL_PROMPT_CHUNK"); ++ uint64_t phys = 0; ++ size_t plen = sizeof phys; ++ sysctlbyname("hw.memsize", &phys, &plen, nullptr, 0); ++ const int stream_chunk = phys >= (96ull << 30) ? 4096 : phys >= (24ull << 30) ? 2048 : 1024; ++ const int C = std::max(64, ec ? std::atoi(ec) : (stream_ ? stream_chunk : 2048)) / 64 * 64; ++ const int QS = std::min(C, 256); ++ const bool k4 = ::coop_k4(*cx_); ++ const char* eq = std::getenv("STRATA_METAL_PROMPT_QSA"); ++ const bool flash = eq ? std::string(eq) == "flash" : k4; ++ PromptGemm p_tm; ++ const size_t MT = (size_t) C; ++ const size_t mr = (size_t) (C * TOPK + NEXP * p_tm.TM); ++ const size_t nblk_max = (size_t) ((C * TOPK + 1023) / 1024); ++ const size_t MB = (size_t) (max_cells_ / 4 + 2); ++ uint64_t need = 0; ++ auto AB = [&](size_t bytes) { need += (bytes + 16383) / 16384 * 16384; }; ++ auto F32 = [&](size_t n) { AB(4 * n); }; ++ auto F16 = [&](size_t n) { AB(2 * n); }; ++ F32(MT * HC * N); F32(MT * HC); F32(MT * HC); F16(MT * HC * N); ++ F32(k4 ? 16 : MT * (LR + HC)); F16(MT * LR); F32(prompt_bo_alias() ? 16 : MT * N); F16(MT * N); ++ AB(std::max(4 * MT * 12288, 2 * mr * NFF)); F32(MT * 6144); F32(MT * 96); F32(16); ++ F16(MT * 6144); F32(MT * 16 * 128); F32(MT * 16 * 128); F32(MT * std::max(6144, MB)); ++ F32(MT * 96); F32(MT * 16); F32(MT * 6144); ++ if (flash) F32(16); ++ F32(MT * 4 * 128); ++ if (flash) F32(16); else F32((size_t) QS * 33 * 24 * 258); ++ if (flash) F32(16); else F32(4 * MT * 6144 + 3 * (size_t) QS * 6144 + MT * 384); ++ if (!flash) F16(MT * 6144); ++ F32(MT * NEXP); AB(4 * MT * TOPK); F32(MT * TOPK); ++ if (stream_) { ++ F32(MT * NEXP); AB(4 * MT * TOPK); F32(MT * TOPK); F16((size_t) NEXP * N); ++ if (std::getenv("STRATA_METAL_PROMPT_PRED2") && std::atoi(std::getenv("STRATA_METAL_PROMPT_PRED2")) != 0) ++ AB(4 * MT * TOPK); ++ } ++ AB(4 * nblk_max * NEXP); AB(4 * nblk_max * NEXP); ++ AB(4 * NEXP); AB(8 * (mr / p_tm.TM + 1)); AB(64); ++ AB(4 * mr); F32(mr); AB(4 * MT * TOPK); F16(mr * N); ++ F32(mr * N); F32(prompt_bo_alias() ? 16 : MT * N); ++ if (!prompt_bo_alias()) F16(MT * NFF); ++ F32(MT * N); F32(MT); ++ if (!prompt_bo_alias()) F32(MT * N); ++ F16((size_t) (LR + HC) * HC * N); F16((size_t) HC * N * LR); F16((size_t) 12288 * N); ++ F16((size_t) 6144 * N); F16((size_t) 512 * N); F16((size_t) 512 * N); F16((size_t) N * 6144); ++ F16((size_t) NEXP * N); F16((size_t) N); F16((size_t) NFF * N); ++ F16((size_t) NFF * N); F16((size_t) N * NFF); ++ if (ple_on_) { ++ F32(MT * N); ++ if (2 * mr * N < 4 * MT * 10240) F32(MT * 10240); ++ } ++ return need; ++} ++ ++void Window::prompt_region(bool hold, bool multi_chunk) { ++ if (!stream_ || !st_ || !st_->ready()) return; ++ if (hold) { ++ st_->hold_region(true, multi_chunk); ++ } else { ++ st_->hold_region(false); ++ const char* esep = std::getenv("STRATA_METAL_SEPARATE_PROMPT_BUF"); ++ const bool sep_pg = !(esep && std::string(esep) == "0"); ++ if (sep_pg && pg_ != nullptr) { ++ const char* reuse = std::getenv("STRATA_METAL_PROMPT_BUF_REUSE"); ++ const int mode = reuse ? std::atoi(reuse) : 0; ++ if (st_->reuse_region(mode)) { ++ mtp_src_ = nullptr; ++ return; ++ } ++ const ProcMem pm_before = proc_mem(); ++ free_prompt_gemm(); ++ const uint64_t unmapped = st_->release_region(); ++ strata::gpu::purge_device_allocations(1u << 20); ++ const ProcMem pm_after = proc_mem(); ++ if (unmapped > 0) { ++ std::fprintf(stderr, "strata metal stream: prompt scratch %.2f GB unmapped after prefill (phys_footprint %.2f -> %.2f GB, resident RSS %.2f -> %.2f GB)\n", ++ unmapped / 1e9, pm_before.footprint / 1e9, pm_after.footprint / 1e9, pm_before.rss / 1e9, pm_after.rss / 1e9); ++ } ++ } ++ } ++} ++ + 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 +3400,7 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri + pend_fold_ = 0; + } + const int T = Tn; ++ const double tb0 = now_ms(); + // ---- the embeddings into every stream, the PLE rows + { + const gguf::Tensor& te = TN("token_embd.weight"); +@@ -2133,19 +3413,36 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri + for (int c = 1; c < HC; ++c) std::memcpy(r + (size_t) c * N, r, (size_t) N * 4); + } + } ++ const double tbe = now_ms(); + strata::core::SessionState& ss = *ss_; ++ // Under streaming, the PLE reads run on a thread beside layer 0's attention and are joined at layer 0's handshake, ++ // before the PLE projections read p.emb (layer 0 has no PLE). ++ static const bool ple_async = !(std::getenv("STRATA_METAL_PLE_ASYNC") && std::atoi(std::getenv("STRATA_METAL_PLE_ASYNC")) == 0); ++ const bool ple_late = ple_on_ && stream_ && ple_async; ++ std::vector ple_rows; ++ std::thread ple_th; ++ std::string ple_err; ++ bool ple_ok = true; + if (ple_on_) { +- std::vector rows((size_t) T * strata::kernels::PLE_N_HEADS); ++ ple_rows.resize((size_t) T * strata::kernels::PLE_N_HEADS); + int32_t prev[2] = {ss.ple_prev[0], ss.ple_prev[1]}; + for (int t = 0; t < T; ++t) { +- strata::kernels::ngram_rows(&tokens[t], prev, 1, ss.ple.consts, rows.data() + (size_t) t * strata::kernels::PLE_N_HEADS); ++ strata::kernels::ngram_rows(&tokens[t], prev, 1, ss.ple.consts, ple_rows.data() + (size_t) t * strata::kernels::PLE_N_HEADS); + prev[0] = prev[1]; + prev[1] = tokens[t]; + } +- if (!ss.ple.table->gather_batch(rows.data(), (size_t) T, p.emb.as(), err)) return false; ++ auto gather = [&] { ple_ok = ss.ple.table->gather_batch(ple_rows.data(), (size_t) T, p.emb.as(), ple_err); }; ++ if (ple_late) { ++ ple_th = std::thread(gather); ++ ple_join_ = &ple_th; ++ } else { ++ gather(); ++ if (!ple_ok) { err = ple_err; return false; } ++ } + } ++ const double tb1 = now_ms(); + std::vector ds; +- std::vector cuts; ++ std::vector cuts, sigs; + nb_ = false; + auto D = [&](const mtl::Pipeline* pp, std::vector a, uint32_t gx, uint32_t gy, uint32_t tg, uint32_t gz = 1) { + mtl::Dispatch d; +@@ -2159,110 +3456,196 @@ 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) { +- const int tm = 64, tn = Nn <= 64 ? 32 : 64; +- const std::string fn = "gemm_f16_" + std::to_string(tm) + "_" + std::to_string(tn) + "_4"; ++ // every tile gives the same bits (dense_mm_parity --tiles); the fastest on M5 per shape, from its sweep at ++ // M = 256..2048 (audits/opus_r11_dense.md) ++ int tm = 64, tn = 128, sg = 4; ++ auto set = [&](int m, int n, int s) { tm = m; tn = n; sg = s; }; ++ if (M <= 128) { ++ if (Nn <= 64) set(64, 32, 4); ++ else if (Nn <= 640 || K > 2560) set(64, Nn <= 128 ? 32 : 64, 4); ++ else set(M >= 128 ? 128 : 64, 64, 4); ++ } else if (Nn <= 64) { ++ M < 1024 ? set(16, 64, 2) : set(64, 32, 2); ++ } else if (Nn == 96) { // GDN's alpha | beta: from 640 rows one 96-wide column tile, p.mixh read once ++ M < 640 ? set(16, 64, 2) : set(16, 96, 2); ++ } else if (Nn <= 128) { ++ M < 768 ? set(16, 64, 2) : set(32, 64, 2); ++ } else if (Nn <= 640 && K > 2560) { // the HC read ++ M <= 512 ? set(64, 64, 4) : M < 1536 ? set(128, 64, 4) : set(32, 128, 2); ++ } else if (Nn <= 512 || (Nn <= 640 && M < 384)) { ++ M < 384 ? set(32, 64, 2) : set(32, 128, 2); ++ } else if (Nn <= 1664) { // QSA's fused K | V | indexer rows (1152, 1664): 128_64_4 is 10% slower at M = 1024 ++ M < 768 ? set(32, 128, 2) : set(64, 128, 4); ++ } else { ++ // 128-row tiles only where they pad no more rows than 64-row ones: M % 128 in 1..64 is 7-21% slower ++ M < 1536 && (M % 128 == 0 || M % 128 > 64) ? set(128, 64, 4) : set(64, 128, 4); ++ } ++ if (Nn > 64) ++ if (const char* ge = std::getenv("STRATA_PG_GEMM")) std::sscanf(ge, "%d,%d,%d", &tm, &tn, &sg); ++ const std::string fn = "gemm_f16_" + std::to_string(tm) + "_" + std::to_string(tn) + "_" + std::to_string(sg); + PgGemm g{M, Nn, K, 0}; +- D(PP(p.lib, fn), {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(Y), mtl::Arg::V(g)}, (Nn + tn - 1) / tn, (M + tm - 1) / tm, 128); ++ D(PP(p.lib, fn), {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(Y), mtl::Arg::V(g)}, (Nn + tn - 1) / tn, (M + tm - 1) / tm, sg * 32); + }; + 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) { ++ const int u = n_in / pg_unit(type), tg = u <= 64 ? 64 : u <= 128 ? 128 : u % 192 == 0 ? 192 : u % 160 == 0 ? 160 : 256; + PgDq a{(uint32_t) n_in, (uint32_t) n_out, (uint32_t) ref::row_bytes(type, n_in), perm}; + D(PP(p.lib, "dq_f16_" + std::to_string(type)), {mtl::Arg::B(src, off), mtl::Arg::B(out, out_off), mtl::Arg::V(a)}, +- (n_in / pg_unit(type) + 63) / 64, n_out, 64); ++ (u + tg - 1) / tg, n_out, tg); + }; + 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; +- if (hc_q8_ && q8off_.count(t.name)) dqb(Q8_, q8off_.at(t.name), ref::Q8_0, n_in, n_out, out, out_off, perm); ++ if (hc_q8_ && q8off_.count(t.name)) dqb(Q8_, q8off_.at(t.name), q8type_.count(t.name) ? q8type_.at(t.name) : ref::Q8_0, n_in, n_out, out, out_off, perm); + 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) { + const uint32_t nn = (uint32_t) n; +- D(PP(p.lib, "act_f16"), {mtl::Arg::B(x), mtl::Arg::B(xh), mtl::Arg::V(nn)}, (nn + 255) / 256, 1, 256); ++ D(PP(p.lib, "act_f16"), {mtl::Arg::B(x), mtl::Arg::B(xh), mtl::Arg::V(nn)}, (nn + 1023) / 1024, 1, 256); + }; + int ci = 0; +- auto hc_read = [&](const std::string& pre, bool apply) { ++ // mm_combine folded into the next read (gr_norm_h_moec, ple_moe_gate_norm_c, gr_apply_moec): yr is formed from YE ++ // as it is read and never written - except while STRATA_METAL_DUMP_MOE reads yr ++ const bool comb_fold = prompt_bo_alias(); ++ const uint32_t topk = TOPK; ++ // moe: the write to apply is the last MoE's, formed from yr, ys, sgv (moe_out folded in) rather than read from bo ++ auto hc_read = [&](const std::string& pre, bool apply, bool moe) { + 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); ++ nb_ = true; + dq(tin, p.WD, (size_t) 2 * LR * HC * N); ++ nb_ = true; + dq(tup, p.WU, 0, HC); + GrArgs a{N, HC, LR, (uint32_t) T, apply ? 1u : 0u, eps_}; + PgEpi ea{(uint32_t) T, N, HC, LR}; +- 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); +- D(PP(p.lib, "gr_epi_down"), {mtl::Arg::B(p.V), mtl::Arg::B(p.loh), mtl::Arg::B(p.inj[1 - ci]), mtl::Arg::V(ea)}, +- (T * (LR + HC) + 255) / 256, 1, 256); ++ nb_ = true; ++ if (apply && moe && comb_fold) ++ D(PP(p.lib, "gr_norm_h_moec"), {mtl::Arg::B(p.R), mtl::Arg::B(p.YE), mtl::Arg::B(p.pos), mtl::Arg::B(p.sw), ++ mtl::Arg::B(p.ys), mtl::Arg::B(p.sgv), mtl::Arg::B(p.inj[ci]), mtl::Arg::B(W_, tn.offset), ++ mtl::Arg::B(p.xnh), mtl::Arg::V(a), mtl::Arg::V(topk)}, 1, T, 256); ++ else if (apply && moe) ++ D(PP(p.lib, "gr_norm_h_moe"), {mtl::Arg::B(p.R), mtl::Arg::B(p.yr), mtl::Arg::B(p.ys), mtl::Arg::B(p.sgv), ++ mtl::Arg::B(p.inj[ci]), mtl::Arg::B(W_, tn.offset), mtl::Arg::B(p.xnh), mtl::Arg::V(a)}, ++ 1, T, 256); ++ else ++ // applying: a threadgroup a token (bo read once, not hc times); else the (stream, token) grid is the faster ++ D(PP(p.lib, apply ? "gr_norm_h4" : "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)}, ++ apply ? 1 : HC, T, 256); ++ if (p.k4) { ++ // M5: gr_epi_down as the GEMM's register epilogue, V never written; the tiles from comb_fold_parity T epi bench ++ const auto [tm, tn, sg] = T <= 192 ? std::array{16, 64, 2} : T <= 384 ? std::array{32, 64, 2} ++ : T <= 512 ? std::array{32, 128, 4} : T < 1536 ? std::array{64, 128, 4} : std::array{32, 128, 2}; ++ const std::string fn = "gemm_hc_down_" + std::to_string(tm) + "_" + std::to_string(tn) + "_" + std::to_string(sg); ++ D(PP(p.lib, fn), {mtl::Arg::B(p.xnh), mtl::Arg::B(p.WD), mtl::Arg::B(p.loh), mtl::Arg::B(p.inj[1 - ci]), mtl::Arg::V(ea)}, ++ (LR + HC + tn - 1) / tn, (T + tm - 1) / tm, sg * 32); ++ } else { ++ gemm(p.xnh, p.WD, p.V, T, LR + HC, HC * N); ++ D(PP(p.lib, "gr_epi_down"), {mtl::Arg::B(p.V), mtl::Arg::B(p.loh), mtl::Arg::B(p.inj[1 - ci]), mtl::Arg::V(ea)}, ++ (T * (LR + HC) + 255) / 256, 1, 256); ++ } + PgGemm g{T, HC * N, LR, 0}; +- D(PP(p.lib, "gemm_hc_up_64_64_4"), {mtl::Arg::B(p.loh), mtl::Arg::B(p.WU), mtl::Arg::B(p.xnh), mtl::Arg::B(p.mixed), +- mtl::Arg::B(p.mixh), mtl::Arg::V(g), mtl::Arg::V(ea)}, HC * N / 64, (T + 63) / 64, 128); ++ // M5 keeps each output's four HC streams in one lane; other layouts use the shared-memory epilogue. ++ const int up_tm = p.k4 ? 32 : 64, up_tn = p.k4 ? 128 : 64; ++ const char* up = p.k4 ? "gemm_hc_up_reg_32_128_4" : "gemm_hc_up_64_64_4"; ++ // the fp32 mixed has no reader but STRATA_METAL_DUMP_MOE ++ D(PP(p.lib, up), {mtl::Arg::B(p.loh), mtl::Arg::B(p.WU), mtl::Arg::B(p.xnh), ++ mtl::Arg::B(prompt_bo_alias() ? kUnbound : p.mixed), ++ mtl::Arg::B(p.mixh), mtl::Arg::V(g), mtl::Arg::V(ea)}, HC * N / up_tn, (T + up_tm - 1) / up_tm, 128); + ci = 1 - ci; + }; + // the PLE projections (the window's gemv pair as GEMMs) +- if (ple_on_) { ++ auto ple_proj = [&] { + act16(p.emb, p.embh, (size_t) T * N); ++ nb_ = true; + dq(TN("blk.1.ple_key.weight"), p.Wpk); ++ nb_ = true; + dq(TN("blk.1.ple_value.weight"), p.Wpv); + gemm(p.embh, p.Wpk, p.pkey, T, 10240, N); ++ nb_ = true; + gemm(p.embh, p.Wpv, p.pval, T, N, N); +- } ++ }; + const int MB = (int) max_cells_ / 4 + 2; + (void) MB; ++ // STRATA_METAL_PROMPT_PRED=0: the read-ahead of a chunk's next layer from the static expert profile instead ++ static const bool pred_on = !(std::getenv("STRATA_METAL_PROMPT_PRED") && std::atoi(std::getenv("STRATA_METAL_PROMPT_PRED")) == 0); + for (int l = 0; l < 48; ++l) { + const std::string b = "blk." + std::to_string(l) + "."; +- hc_read(b + "hc_attn", l > 0 && !(l == 1 && ple_on_)); ++ hc_read(b + "hc_attn", l > 0 && !(l == 1 && ple_on_), true); + if (l % 4 != 3) { + const int gi = l - l / 4; ++ nb_ = true; + dq(TN(b + "attn_qkv.weight"), p.Wqkv); ++ nb_ = true; + dq(TN(b + "attn_gate.weight"), p.Wz); ++ // alpha's and beta's 48 rows as one 96-row GEMM (one dispatch, one p.mixh read): al is [T][96], alpha then beta, ++ // and gdn_pre reads the interleave from the two pointers ++ nb_ = true; + dq(TN(b + "ssm_alpha.weight"), p.Wal); +- dq(TN(b + "ssm_beta.weight"), p.Wbe); ++ nb_ = true; ++ dq(TN(b + "ssm_beta.weight"), p.Wal, (size_t) 2 * 48 * N); ++ nb_ = true; + dq(TN(b + "ssm_out.weight"), p.Wout); + gemm(p.mixh, p.Wqkv, p.qkv, T, 10240, N); ++ nb_ = true; + gemm(p.mixh, p.Wz, p.z, T, 6144, N); +- gemm(p.mixh, p.Wal, p.al, T, 48, N); +- gemm(p.mixh, p.Wbe, p.be, T, 48, N); ++ nb_ = true; ++ gemm(p.mixh, p.Wal, p.al, T, 96, N); + PgGdnP ga{(uint32_t) T, 16, 48, 10240, eps_, 0}; + D(PP(p.lgdn, "gdn_pre"), + {mtl::Arg::B(p.qkv), mtl::Arg::B(ghist_[hp_][(size_t) gi]), mtl::Arg::B(W_, TN(b + "ssm_conv1d.weight").offset), +- mtl::Arg::B(p.al), mtl::Arg::B(p.be), mtl::Arg::B(W_, TN(b + "ssm_dt.bias").offset), mtl::Arg::B(W_, TN(b + "ssm_a").offset), ++ mtl::Arg::B(p.al), mtl::Arg::B(p.al, 4 * 48), mtl::Arg::B(W_, TN(b + "ssm_dt.bias").offset), mtl::Arg::B(W_, TN(b + "ssm_a").offset), + mtl::Arg::B(p.gqn), mtl::Arg::B(p.gkn), mtl::Arg::B(p.gvc), mtl::Arg::B(p.ggb), mtl::Arg::B(p.gkq), +- mtl::Arg::B(ghist_[1 - hp_][(size_t) gi]), mtl::Arg::V(ga)}, 16 + 24, T, 256); ++ mtl::Arg::B(ghist_[1 - hp_][(size_t) gi]), mtl::Arg::V(ga)}, (16 + 48) / 8, T, 256); + D(PP(p.lgdn, "gdn_scanr_4", mtl::Constants().i(6, 0), "gdn_scanr_4"), + {mtl::Arg::B(gst_[(size_t) gi]), mtl::Arg::B(p.gqn), mtl::Arg::B(p.gkn), mtl::Arg::B(p.gvc), mtl::Arg::B(p.ggb), + mtl::Arg::B(p.gkq), mtl::Arg::B(p.go), mtl::Arg::V(ga)}, 48 * 4, 1, 32 * 4); + D(PP(p.lgdn, "gdn_post"), {mtl::Arg::B(p.go), mtl::Arg::B(p.z), mtl::Arg::B(W_, TN(b + "ssm_norm.weight").offset), +- mtl::Arg::B(p.ygh), mtl::Arg::B(p.yg), mtl::Arg::V(ga)}, 48, T, 128); ++ mtl::Arg::B(p.ygh), mtl::Arg::B(p.yg), mtl::Arg::V(ga)}, 48 / 4, T, 128); + gemm(p.ygh, p.Wout, p.bo, T, N, 6144); + } else { + const int qi = l / 4; ++ // K, V, indexer K (and past the budget indexer Q) as one GEMM in Wk / kcur: rows 0..511 K, 512..1023 V, ++ // 1024..1151 indexer K, 1152..1663 indexer Q. The same bits as the separate GEMMs, one dispatch and one ++ // p.mixh read instead of 3-4 (audits/opus_r16_kvik.md); qsa_prep reads the slices at the row stride ++ // QSA_KV_STRIDE (function constant 26) ++ const bool sel = pos0 + T > 2051; ++ const int kvn = sel ? 1664 : 1152; ++ nb_ = true; + dq(TN(b + "attn_q.weight"), p.Wq); ++ nb_ = true; + dq(TN(b + "attn_k.weight"), p.Wk); +- dq(TN(b + "attn_v.weight"), p.Wv); +- dq(TN(b + "indexer.q_proj.weight"), p.Wiq); +- dq(TN(b + "indexer.k_proj.weight"), p.Wik); ++ nb_ = true; ++ dq(TN(b + "attn_v.weight"), p.Wk, (size_t) 2 * 512 * N); ++ nb_ = true; ++ dq(TN(b + "indexer.k_proj.weight"), p.Wk, (size_t) 2 * 1024 * N); ++ if (sel) { ++ nb_ = true; ++ dq(TN(b + "indexer.q_proj.weight"), p.Wk, (size_t) 2 * 1152 * N); ++ } ++ nb_ = true; + dq(TN(b + "attn_output.weight"), p.Wout); + gemm(p.mixh, p.Wq, p.qfull, T, 12288, N); +- gemm(p.mixh, p.Wk, p.kcur, T, 512, N); +- gemm(p.mixh, p.Wv, p.vcur, T, 512, N); +- gemm(p.mixh, p.Wiq, p.iqr, T, 512, N); +- gemm(p.mixh, p.Wik, p.ikr, T, 128, N); ++ nb_ = true; ++ gemm(p.mixh, p.Wk, p.kcur, T, kvn, 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()), +- {mtl::Arg::B(p.qfull), mtl::Arg::B(p.kcur), mtl::Arg::B(p.vcur), mtl::Arg::B(p.iqr), mtl::Arg::B(p.ikr), ++ const mtl::Constants qc = (p.flash ? qsa_c().b(25, true) : qsa_c()).i(26, kvn); ++ D(PP(p.lqsa, "qsa_prep", qc, (p.flash ? qsa_k() + "#qf16" : qsa_k()) + "#kv" + std::to_string(kvn)), ++ {mtl::Arg::B(p.qfull), mtl::Arg::B(p.kcur), mtl::Arg::B(p.kcur, 4 * 512), mtl::Arg::B(p.kcur, 4 * 1152), ++ mtl::Arg::B(p.kcur, 4 * 1024), + mtl::Arg::B(W_, TN(b + "attn_q_norm.weight").offset), mtl::Arg::B(W_, TN(b + "attn_k_norm.weight").offset), + mtl::Arg::B(W_, TN(b + "indexer.q_norm.weight").offset), mtl::Arg::B(W_, TN(b + "indexer.k_norm.weight").offset), +- mtl::Arg::B(cs_), mtl::Arg::B(p.qcur), mtl::Arg::B(p.qidx), 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(tail_[(size_t) qi]), +- mtl::Arg::B(dead_[(size_t) qi]), mtl::Arg::B(pooled_[(size_t) qi]), mtl::Arg::V(a), mtl::Arg::B(mrope_)}, 33, T, 256); +- if (pos0 + T > 2051) ++ mtl::Arg::B(cs_), mtl::Arg::B(p.flash ? p.qcurh : p.qcur), mtl::Arg::B(p.qidx), 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(tail_[(size_t) qi]), mtl::Arg::B(dead_[(size_t) qi]), mtl::Arg::B(pooled_[(size_t) qi]), ++ mtl::Arg::V(a), mtl::Arg::B(mrope_)}, 33, T, 256); ++ if (sel) + D(PP(p.lqsa, "qsa_score_mm"), {mtl::Arg::B(p.qidx), mtl::Arg::B(pooled_[(size_t) qi]), mtl::Arg::B(dead_[(size_t) qi]), + mtl::Arg::B(p.scores), mtl::Arg::V(a)}, + (uint32_t) (((pos0 + T) / 4 + 1 + 63) / 64), (T + 15) / 16, 128); +- D(PP(p.lqsa, "qsa_topk"), {mtl::Arg::B(p.scores), mtl::Arg::B(p.qids), mtl::Arg::V(a)}, T, 1, 1024); ++ if (sel || !p.flash) ++ 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; +- 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]), + mtl::Arg::B(vsc_[(size_t) qi]), mtl::Arg::B(p.qids), mtl::Arg::B(p.qfull), mtl::Arg::B(p.xqh), mtl::Arg::V(a)}, +@@ -2287,28 +3670,95 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri + // layer 0's attention write applied to R happens in the FFN read below; the PLE block comes after the FFN + } + // ---- the MoE +- hc_read(b + "hc_ffn", true); ++ hc_read(b + "hc_ffn", true, false); + const gguf::Tensor &tr = TN(b + "ffn_gate_inp.weight"), &tsg = TN(b + "ffn_gate_inp_shexp.weight"), + &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"); +- dq(tr, p.Wr); ++ // under streaming the next layer's router is read ahead beside this one's: its dequant, GEMM and top-k in this ++ // one's waves (disjoint buffers; the two GEMMs read p.mixh together) ++ const bool ahead = stream_ && pred_on && l + 1 < 48; ++ // ... into Wr and Wr2 in turn: layer l + 1's router is the one read ahead (unless PRED2 reuses Wr2 for l + 2) ++ const bool reuse_wr = stream_ && pred_on && !p.ids3.ptr; ++ const mtl::Buffer& w_r = reuse_wr && (l & 1) ? p.Wr2 : p.Wr; ++ const mtl::Buffer& w_r2 = reuse_wr && (l & 1) ? p.Wr : p.Wr2; ++ if (l == 0 || !reuse_wr) { ++ nb_ = true; ++ dq(tr, w_r); ++ } ++ nb_ = true; + dq(tsg, p.Wsg); ++ nb_ = true; + dq(tg_, p.Wshg); ++ nb_ = true; + dq(tu_, p.Wshu); ++ nb_ = true; + dq(td_, p.Wshd); +- gemm(p.mixh, p.Wr, p.logits, T, NEXP, N); +- const uint32_t nt = (uint32_t) T; +- D(PP(p.lib, "route_topk"), {mtl::Arg::B(p.logits), mtl::Arg::B(p.ids), mtl::Arg::B(p.wts), mtl::Arg::V(nt)}, (T + 7) / 8, 1, 256); ++ if (ahead) { ++ nb_ = true; ++ dq(TN("blk." + std::to_string(l + 1) + ".ffn_gate_inp.weight"), w_r2); ++ } ++ gemm(p.mixh, w_r, p.logits, T, n_exp_, N); ++ if (ahead) { ++ nb_ = true; ++ gemm(p.mixh, w_r2, p.logits2, T, n_exp_, N); ++ } ++ const uint32_t nt = (uint32_t) T | ((uint32_t) n_exp_ << 16); ++ const float p_moce = stream_ ? (std::getenv("STRATA_METAL_PROMPT_MOCE") ? (float) std::atof(std::getenv("STRATA_METAL_PROMPT_MOCE")) ++ : ((pos0 > 0) ? 320335.0f : 0.0f)) : 0.0f; ++ D(PP(p.lib, "route_topk"), {mtl::Arg::B(p.logits), mtl::Arg::B(p.ids), mtl::Arg::B(p.wts), mtl::Arg::V(nt), stream_ ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(p.ids), mtl::Arg::V(p_moce)}, (T + 7) / 8, 1, 256); ++ if (ahead) { ++ nb_ = true; ++ D(PP(p.lib, "route_topk"), {mtl::Arg::B(p.logits2), mtl::Arg::B(p.ids2), mtl::Arg::B(p.wts2), mtl::Arg::V(nt), ++ mtl::Arg::B(st_->smap(), (size_t) (l + 1) * 512 * 4), mtl::Arg::V(p_moce)}, (T + 7) / 8, 1, 256); ++ if (p.ids3.ptr && l + 2 < 48) { // STRATA_METAL_PROMPT_PRED2: two layers ahead, after the first ++ const int ln = l + 2; ++ dq(TN("blk." + std::to_string(ln) + ".ffn_gate_inp.weight"), p.Wr2); ++ gemm(p.mixh, p.Wr2, p.logits2, T, n_exp_, N); ++ D(PP(p.lib, "route_topk"), {mtl::Arg::B(p.logits2), mtl::Arg::B(p.ids3), mtl::Arg::B(p.wts2), mtl::Arg::V(nt), ++ mtl::Arg::B(st_->smap(), (size_t) ln * 512 * 4), mtl::Arg::V(p_moce)}, (T + 7) / 8, 1, 256); ++ } ++ } ++ if (stream_ && comb_fold) sigs.push_back((int) ds.size()); + const int nent = T * TOPK, nblk = (nent + 1023) / 1024; +- PgPlan pa{(uint32_t) nent, TOPK, NEXP, (uint32_t) p.TM, (uint32_t) nblk}; +- const uint32_t nclr = (uint32_t) p.maxrows; +- D(PP(p.lib, "mm_clear"), {mtl::Arg::B(p.stok), mtl::Arg::V(nclr)}, (p.maxrows + 255) / 256, 1, 256); ++ PgPlan pa{(uint32_t) nent, TOPK, (uint32_t) n_exp_, (uint32_t) p.TM, (uint32_t) nblk}; ++ // with comb_fold the shared expert runs beside the plan, wave for wave (disjoint buffers): mm_hist | gate (and up) ++ // GEMMs, mm_scan | swiglu, mm_place | down GEMM; mm_scan's one threadgroup no longer holds the GPU alone. ++ // M5 past 160 rows: swiglu_h as the up GEMM's register epilogue, ush never written (the up GEMM waits for the ++ // gate's; at <= 160 rows the two overlapped are the faster) ++ const bool swu = p.k4 && T > 160; + D(PP(p.lib, "mm_hist"), {mtl::Arg::B(p.ids), mtl::Arg::B(p.bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024); ++ if (comb_fold) { ++ nb_ = true; ++ gemm(p.mixh, p.Wshg, p.gsh, T, NFF, N); ++ if (!swu) { ++ nb_ = true; ++ gemm(p.mixh, p.Wshu, p.ush, T, NFF, N); ++ } ++ nb_ = true; ++ gemm(p.mixh, p.Wsg, p.sgv, T, 1, N); ++ } + D(PP(p.lib, "mm_scan"), {mtl::Arg::B(p.bcnt), mtl::Arg::B(p.boff), mtl::Arg::B(p.eoff), mtl::Arg::B(p.tiles), +- mtl::Arg::B(p.meta), mtl::Arg::V(pa)}, 1, 1, NEXP); ++ mtl::Arg::B(p.meta), mtl::Arg::V(pa), mtl::Arg::B(p.stok)}, 1, 1, (uint32_t) n_exp_); ++ if (comb_fold) { ++ nb_ = true; ++ if (swu) { ++ const int tm = T < 384 ? 32 : 64, tn = T < 384 ? 64 : 128, sg = T < 384 ? 2 : 4; ++ const PgGemm g{T, NFF, N, 0}; ++ D(PP(p.lib, T < 384 ? "gemm_swiglu_up_32_64_2" : "gemm_swiglu_up_64_128_4"), ++ {mtl::Arg::B(p.mixh), mtl::Arg::B(p.Wshu), mtl::Arg::B(p.gsh), mtl::Arg::B(p.hh), mtl::Arg::V(g)}, ++ (NFF + tn - 1) / tn, (T + tm - 1) / tm, sg * 32); ++ } else { ++ const uint32_t nh = (uint32_t) T * NFF; ++ D(PP(p.lib, "swiglu_h"), {mtl::Arg::B(p.gsh), mtl::Arg::B(p.ush), mtl::Arg::B(p.hh), mtl::Arg::V(nh)}, (nh + 255) / 256, 1, 256); ++ } ++ } + 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); ++ if (comb_fold) { ++ nb_ = true; ++ gemm(p.hh, p.Wshd, p.ys, T, N, NFF); ++ } + PgMm ma{eg.offset, TN(b + "ffn_up_exps.weight").offset, ed.offset, 0, 0, (uint32_t) ref::row_bytes(eg.type, N), + (uint32_t) ref::row_bytes(ed.type, NFF), N, NFF}; + ma.gu_stride = (uint64_t) ma.gu_row * NFF; +@@ -2317,8 +3767,8 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri + const uint32_t ne_ = N; + { + mtl::Dispatch g; +- g.p = PP(p.lib, "mm_gather"); +- g.args = {mtl::Arg::B(p.mixed), mtl::Arg::B(p.stok), mtl::Arg::B(p.XS), mtl::Arg::V(ne_)}; ++ g.p = PP(p.lib, "mm_gather_h"); ++ g.args = {mtl::Arg::B(p.mixh), mtl::Arg::B(p.stok), mtl::Arg::B(p.XS), mtl::Arg::V(ne_)}; + g.tg[0] = 256; + g.indirect = &p.meta; + g.ind_off = 8 * 4; +@@ -2349,7 +3799,10 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri + arena_args(gu); + gu.indirect = &p.meta; + gu.ind_off = 2 * 4; +- if (stream_) cuts.push_back((int) ds.size()); // the host's turn: this layer's experts put in place ++ if (stream_) { ++ if (!comb_fold) sigs.push_back((int) ds.size()); ++ cuts.push_back((int) ds.size()); // the host's turn: this layer's experts put in place ++ } + push(ds, gu); + mtl::Dispatch dn; + if (p.dn_v3) { +@@ -2371,75 +3824,87 @@ bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::stri + dn.ind_off = 5 * 4; + push(ds, dn); + const uint32_t ca[4] = {(uint32_t) T, N, TOPK, 0}; +- D(PP(p.lib, "mm_combine"), {mtl::Arg::B(p.YE), mtl::Arg::B(p.pos), mtl::Arg::B(p.sw), mtl::Arg::B(p.yr), mtl::Arg::V(ca)}, +- (T * N + 255) / 256, 1, 256); +- gemm(p.mixh, p.Wshg, p.gsh, T, NFF, N); +- gemm(p.mixh, p.Wshu, p.ush, T, NFF, N); +- const uint32_t nh = (uint32_t) T * NFF; +- D(PP(p.lib, "swiglu_h"), {mtl::Arg::B(p.gsh), mtl::Arg::B(p.ush), mtl::Arg::B(p.hh), mtl::Arg::V(nh)}, (nh + 255) / 256, 1, 256); +- gemm(p.hh, p.Wshd, p.ys, T, N, NFF); +- gemm(p.mixh, p.Wsg, p.sgv, T, 1, N); +- const uint32_t oa[2] = {(uint32_t) T, N}; +- D(PP(p.lib, "moe_out"), {mtl::Arg::B(p.yr), mtl::Arg::B(p.ys), mtl::Arg::B(p.sgv), mtl::Arg::B(p.bo), mtl::Arg::V(oa)}, +- (T * N + 255) / 256, 1, 256); ++ if (!comb_fold) { ++ D(PP(p.lib, "mm_combine"), {mtl::Arg::B(p.YE), mtl::Arg::B(p.pos), mtl::Arg::B(p.sw), mtl::Arg::B(p.yr), mtl::Arg::V(ca)}, ++ (T * N + 255) / 256, 1, 256); ++ gemm(p.mixh, p.Wshg, p.gsh, T, NFF, N); ++ nb_ = true; ++ gemm(p.mixh, p.Wshu, p.ush, T, NFF, N); ++ nb_ = true; ++ gemm(p.mixh, p.Wsg, p.sgv, T, 1, N); ++ const uint32_t nh = (uint32_t) T * NFF; ++ D(PP(p.lib, "swiglu_h"), {mtl::Arg::B(p.gsh), mtl::Arg::B(p.ush), mtl::Arg::B(p.hh), mtl::Arg::V(nh)}, (nh + 255) / 256, 1, 256); ++ gemm(p.hh, p.Wshd, p.ys, T, N, NFF); ++ } ++ // no moe_out: the next read forms the MoE write from yr (or YE), ys, sgv (gr_norm_h_moe(c), ple_moe_gate_norm(_c), ++ // gr_apply_moe(c)) + if (l == 0 && ple_on_) { +- // 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"); +- GrArgs a{N, HC, LR, (uint32_t) T, 1u, eps_}; +- D(P(lib_, "gr_norm", {}, "gr_norm"), {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.YE), mtl::Arg::V(a)}, HC, T, 256); ++ // the projections into layer 0's idle attention scratch (prompt_gemm_init), after layer 0's handshake: the ++ // gather has been joined ++ ple_proj(); ++ // layer 0's FFN write applied to R (with comb_fold its yr formed from YE as it is read), the PLE norms and ++ // R + gated in one dispatch, then the conv (as the window, for T rows) + struct PleArgs { uint32_t n_embd, hc; float eps; } pl{N, HC, strata::kernels::NG_RMS_EPS}; + auto off = [&](const char* n) { return TN(n).offset; }; +- D(P(lple_, "ple_gnorm", {}, "ple_gnorm"), {mtl::Arg::B(p.pkey), mtl::Arg::B(W_, off("blk.1.ple_norm_key.weight")), +- mtl::Arg::B(p.keyn), mtl::Arg::V(pl)}, HC, T, 256); +- nb_ = true; +- D(P(lple_, "ple_gnorm", {}, "ple_gnorm"), {mtl::Arg::B(p.R), mtl::Arg::B(W_, off("blk.1.ple_norm_query.weight")), +- mtl::Arg::B(p.qn), mtl::Arg::V(pl)}, HC, T, 256); +- D(P(lple_, "ple_gate_norm", {}, "ple_gate_norm"), +- {mtl::Arg::B(p.keyn), mtl::Arg::B(p.qn), mtl::Arg::B(p.pval), mtl::Arg::B(W_, off("blk.1.ple_norm_conv.weight")), +- mtl::Arg::B(p.gated), mtl::Arg::B(p.normW), mtl::Arg::V(pl)}, HC, T, 256); +- D(P(lple_, "ple_conv_add", {}, "ple_conv_add"), +- {mtl::Arg::B(pleHist_), mtl::Arg::B(p.normW), kw_arg(), mtl::Arg::B(p.gated), +- mtl::Arg::B(p.R), mtl::Arg::V(pl)}, HC * N / 256, T, 256); ++ const char* fn = comb_fold ? "ple_moe_gate_norm_c" : "ple_moe_gate_norm"; ++ std::vector a = {mtl::Arg::B(p.R)}; ++ if (comb_fold) a.insert(a.end(), {mtl::Arg::B(p.YE), mtl::Arg::B(p.pos), mtl::Arg::B(p.sw)}); ++ else a.push_back(mtl::Arg::B(p.yr)); ++ a.insert(a.end(), {mtl::Arg::B(p.ys), mtl::Arg::B(p.sgv), mtl::Arg::B(p.inj[ci]), mtl::Arg::B(p.pkey), ++ mtl::Arg::B(W_, off("blk.1.ple_norm_key.weight")), ++ mtl::Arg::B(W_, off("blk.1.ple_norm_query.weight")), mtl::Arg::B(p.pval), ++ mtl::Arg::B(W_, off("blk.1.ple_norm_conv.weight")), mtl::Arg::B(p.normW), mtl::Arg::V(pl)}); ++ if (comb_fold) a.push_back(mtl::Arg::V(topk)); ++ D(P(lple_, fn, {}, fn), std::move(a), comb_fold ? 1 : HC, T, 256); ++ D(P(lple_, "ple_conv_res", {}, "ple_conv_res"), ++ {mtl::Arg::B(pleHist_), mtl::Arg::B(p.normW), kw_arg(), mtl::Arg::B(p.R), mtl::Arg::V(pl)}, HC * N / 256, T, 256); ++ // the PLE history advanced past the chunk once the conv has read it, so normW is dead after layer 0 ++ const uint32_t tu = (uint32_t) T; ++ D(P(lple_, "ple_hist_adv", {}, "ple_hist_adv"), {mtl::Arg::B(pleHist_), mtl::Arg::B(p.normW), mtl::Arg::V(pl), mtl::Arg::V(tu)}, ++ HC * N / 256, 1, 256); + } + } + // the last layer's write applied to R (the draft layer's K/V read the final residual rows) + { +- const gguf::Tensor& tn = TN("output_hc_norm.weight"); + GrArgs a{N, HC, LR, (uint32_t) T, 1u, eps_}; +- D(P(lib_, "gr_norm", {}, "gr_norm"), {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.YE), mtl::Arg::V(a)}, HC, T, 256); ++ if (comb_fold) ++ D(PP(p.lib, "gr_apply_moec"), {mtl::Arg::B(p.R), mtl::Arg::B(p.YE), mtl::Arg::B(p.pos), mtl::Arg::B(p.sw), ++ mtl::Arg::B(p.ys), mtl::Arg::B(p.sgv), mtl::Arg::B(p.inj[ci]), mtl::Arg::V(a), ++ mtl::Arg::V(topk)}, (N + 255) / 256, T, 256); ++ else ++ D(PP(p.lib, "gr_apply_moe"), {mtl::Arg::B(p.R), mtl::Arg::B(p.yr), mtl::Arg::B(p.ys), mtl::Arg::B(p.sgv), ++ mtl::Arg::B(p.inj[ci]), mtl::Arg::V(a)}, HC * N / 256, T, 256); + } + const double t0 = now_ms(); + if (stream_) { + seg_T_ = T; + cut_ids_ = p.ids.as(); +- cx_->run_graph_seg(ds, cuts, [](void* u, int s) { ((Window*) u)->at_cut(s); }, this); ++ st_->set_prompt_pred(pred_on ? p.ids2.as() : nullptr, TOPK, ++ pred_on && p.ids3.ptr ? p.ids3.as() : nullptr); ++ cx_->run_graph_seg(ds, cuts, sigs, [](void* u, int s) { ((Window*) u)->at_cut(s); }, nullptr, this); ++ st_->set_prompt_pred(nullptr, 0); ++ if (ple_th.joinable()) ple_th.join(); // only if the graph stopped before layer 0's handshake ++ ple_join_ = nullptr; ++ if (!ple_ok) { err = ple_err; return false; } ++ if (const char* pfx = std::getenv("STRATA_METAL_DUMP_MOE")) { ++ dump_moe_layer(pfx, 47, false, true); ++ } + cut_ids_ = nullptr; + } else { + cx_->run_graph(ds); + } + pg_ms_ += now_ms() - t0; + pg_tok_ += T; ++ if (std::getenv("STRATA_METAL_PKPROF")) kprof(ds, T); + static const bool trace = std::getenv("STRATA_METAL_TRACE") != nullptr; +- if (trace) std::fprintf(stderr, "strata metal: prompt chunk pos %lld, %d rows: %.0f ms\n", (long long) pos0, T, now_ms() - t0); +- // the window's state after the chunk: the conv history turned, the PLE history and token window advanced ++ if (trace) std::fprintf(stderr, "strata metal: prompt chunk pos %lld, %d rows: %.0f ms (embd %.0f ms, ple %.0f ms, graph build %.0f ms)\n", (long long) pos0, T, now_ms() - t0, tbe - tb0, tb1 - tbe, t0 - tb1); ++ // the window's state after the chunk: the conv history turned (the PLE history on the GPU: ple_hist_adv), the token ++ // window advanced + hp_ = 1 - hp_; +- if (ple_on_) { +- using namespace strata::kernels; +- float* hist = pleHist_.as(); +- const float* nw = p.normW.as(); +- for (int c = 0; c < NG_HC_DIM; ++c) { +- float col[NG_HIST]; +- for (int r = 0; r < NG_HIST; ++r) { +- const int src = T - NG_HIST + r; // row r of the new history: chunk row src, or the old history +- col[r] = src >= 0 ? nw[(size_t) src * NG_HC_DIM + c] : hist[(size_t) c * NG_HIST + (size_t) (r + T)]; +- } +- std::memcpy(hist + (size_t) c * NG_HIST, col, sizeof col); +- } +- } + 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]; } ++ last_T_ = 1; ++ last_tokens_[0] = tokens[T - 1]; + mtp_src_ = &p.R; // the draft layer's prompt K/V read this chunk's final rows + return true; + } +@@ -2479,8 +3944,15 @@ bool window_init(void* w, const strata::core::ModelGeometry& g, strata::core::Se + if (!x->init(g, ss, g_shard1, MAXT, max_cells, err)) return false; + window_register(w); + if (!mtp_load_pending(x, err)) return false; ++ if (x->streaming() && !(std::getenv("STRATA_METAL_LAZY_ARENA") && std::atoi(std::getenv("STRATA_METAL_LAZY_ARENA")) != 0)) { ++ const char* esep = std::getenv("STRATA_METAL_SEPARATE_PROMPT_BUF"); ++ const bool sep_pg = !(esep && std::string(esep) == "0"); ++ if (!x->prompt_gemm_init(err)) return false; ++ if (!sep_pg) x->prompt_region(false); ++ } + } + x->set_max_t(max_t); ++ x->prewarm(); + *n_vocab = x->vocab(); + return true; + } +@@ -2494,6 +3966,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); + } ++void window_reset_stats(void* w) { if (w) { if (std::getenv("STRATA_METAL_PKPROF") || std::getenv("STRATA_METAL_TRACE")) ((Window*) w)->stream_report(stderr); ((Window*) w)->reset_stream_stats(); } } + void window_reset() { + if (g_the_window != nullptr && g_the_window->ready()) g_the_window->reset(); + } +@@ -2505,6 +3978,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); + } ++bool window_session_save(const char* path, const int32_t* ids, int64_t n_cells, std::string& err) { ++ if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; } ++ return g_the_window->session_save(path, ids, n_cells, err); ++} ++bool window_session_load(const char* path, std::vector& ids, int64_t* n_cells, std::string& err) { ++ if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; } ++ return g_the_window->session_load(path, ids, *n_cells, err); ++} ++bool window_scr_relocate(int64_t c_src, int64_t c_dst, int64_t n, std::string& err) { ++ if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; } ++ return g_the_window->scr_relocate(c_src, c_dst, n, err); ++} + 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,23 +4040,26 @@ 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; } ++ g_window->slop_scan(tokens, n, pos0); + static const bool windows = std::getenv("STRATA_METAL_PROMPT") && std::string(std::getenv("STRATA_METAL_PROMPT")) == "window"; + static const int64_t min_rows = std::getenv("STRATA_METAL_PROMPT_MIN") ? std::atoll(std::getenv("STRATA_METAL_PROMPT_MIN")) : 32; + if (!windows && !g_window->streaming() && n >= min_rows) return g_window->prompt(tokens, n, pos0, on_chunk, err); + if (!windows && g_window->streaming()) { + const char* e = std::getenv("STRATA_METAL_PROMPT_GEMM"); + const char* em = std::getenv("STRATA_METAL_PROMPT_GEMM_MIN"); +- const int64_t mn = em ? std::atoll(em) : 64; ++ const int64_t mn = em ? std::atoll(em) : 128; + 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(); + std::vector tk32((size_t) C); +- g_window->prompt_region(true); ++ g_window->prompt_region(true, n > C); + bool ok = true; + 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]; +- ok = g_window->prompt_chunk(tk32.data(), k, pos0 + c, err) && (!on_chunk || on_chunk(nullptr, k, pos0 + c, err)); ++ // the drafter reads the last chunk's rows, and the first's below the attention sink ++ const int64_t sink = c + k < n ? std::min(k, mtp_sink_cap() - (pos0 + c)) : k; ++ ok = g_window->prompt_chunk(tk32.data(), k, pos0 + c, err) && (!on_chunk || sink <= 0 || on_chunk(nullptr, sink, pos0 + c, err)); + } + g_window->prompt_region(false); + return ok; +diff --git a/metal/src/mtl_rt.h b/metal/src/mtl_rt.h +index e5d72592..9e12c80a 100644 +--- a/metal/src/mtl_rt.h ++++ b/metal/src/mtl_rt.h +@@ -88,7 +88,7 @@ public: + Pipeline pipeline(int lib, const std::string& fn, const Constants& c = {}); + + Buffer alloc(size_t bytes); // shared storage, zeroed +- Buffer wrap(void* page_aligned, size_t bytes); // no-copy wrap (page-aligned, size multiple of page) ++ Buffer wrap(void* page_aligned, size_t bytes, bool count_allocated = true); // no-copy wrap (page-aligned, size multiple of page) + void free(Buffer& b); + + // Encodes `ds` into ONE compute encoder (serial or concurrent dispatch type) and ONE command buffer, `reps` +@@ -103,6 +103,11 @@ public: + // missing experts are put in place). Returns the GPU time in seconds. + double run_graph_seg(const std::vector& ds, const std::vector& cuts, void (*at_cut)(void*, int), + void* user); ++ // sigs (same length as cuts, sigs[s] <= cuts[s]): the GPU signals at sigs[s] and waits at cuts[s]; the dispatches ++ // between run while the host takes its turn (they must not read what the host writes). after_cut(s) runs on the ++ // host once the GPU has been let go at cut s (work the GPU does not wait for). ++ double run_graph_seg(const std::vector& ds, const std::vector& cuts, const std::vector& sigs, ++ void (*at_cut)(void*, int), void (*after_cut)(void*, int), void* user); + // One command buffer per `run_cb` call, empty or not: measures commit+wait round trips. + double run_wall(const std::vector& ds, int reps, double* gpu_seconds = nullptr); + +diff --git a/metal/src/mtl_rt.mm b/metal/src/mtl_rt.mm +index 4c6045b2..106657d5 100644 +--- a/metal/src/mtl_rt.mm ++++ b/metal/src/mtl_rt.mm +@@ -183,13 +183,13 @@ Buffer Context::alloc(size_t bytes) { + return r; + } + +-Buffer Context::wrap(void* p, size_t bytes) { ++Buffer Context::wrap(void* p, size_t bytes, bool count_allocated) { + id b = [impl_->dev newBufferWithBytesNoCopy:p length:bytes options:MTLResourceStorageModeShared deallocator:nil]; + if (!b) { + std::fprintf(stderr, "mtl: wrap %zu failed\n", bytes); + std::exit(4); + } +- allocated_ += bytes; ++ if (count_allocated) allocated_ += bytes; + Buffer r; + r.handle = (__bridge_retained void*)b; + r.ptr = p; +@@ -267,14 +267,19 @@ double Context::run_graph(const std::vector& ds) { + + double Context::run_graph_seg(const std::vector& ds, const std::vector& cuts, void (*at_cut)(void*, int), + void* user) { ++ return run_graph_seg(ds, cuts, cuts, at_cut, nullptr, user); ++} ++ ++double Context::run_graph_seg(const std::vector& ds, const std::vector& cuts, const std::vector& sigs, ++ void (*at_cut)(void*, int), void (*after_cut)(void*, int), void* user) { + @autoreleasepool { + if (impl_->ev == nil) impl_->ev = [impl_->dev newSharedEvent]; + id ev = impl_->ev; + const uint64_t base = impl_->ev_v; + id cb = [impl_->q commandBufferWithUnretainedReferences]; + size_t i = 0; +- for (size_t s = 0; s <= cuts.size(); ++s) { +- const size_t end = s < cuts.size() ? (size_t) cuts[s] : ds.size(); ++ auto run_to = [&](size_t end) { // [i, end) as one concurrent encoder ++ if (i >= end) return; + id enc = [cb computeCommandEncoderWithDispatchType:MTLDispatchTypeConcurrent]; + const size_t first = i; + for (; i < end; ++i) { +@@ -282,11 +287,14 @@ double Context::run_graph_seg(const std::vector& ds, const std::vector + encode(enc, ds[i]); + } + [enc endEncoding]; +- if (s < cuts.size()) { +- [cb encodeSignalEvent:ev value:base + 2 * s + 1]; +- [cb encodeWaitForEvent:ev value:base + 2 * s + 2]; +- } ++ }; ++ for (size_t s = 0; s < cuts.size(); ++s) { ++ run_to((size_t) sigs[s]); ++ [cb encodeSignalEvent:ev value:base + 2 * s + 1]; ++ run_to((size_t) cuts[s]); // beside the host's turn ++ [cb encodeWaitForEvent:ev value:base + 2 * s + 2]; + } ++ run_to(ds.size()); + [cb commit]; + for (size_t s = 0; s < cuts.size(); ++s) { + const uint64_t want = base + 2 * s + 1; +@@ -295,6 +303,7 @@ double Context::run_graph_seg(const std::vector& ds, const std::vector + } + at_cut(user, (int) s); + ev.signaledValue = base + 2 * s + 2; ++ if (after_cut) after_cut(user, (int) s); + } + impl_->ev_v = base + 2 * cuts.size() + 2; + [cb waitUntilCompleted]; +diff --git a/metal/tests/comb_fold_parity.cpp b/metal/tests/comb_fold_parity.cpp +new file mode 100644 +index 00000000..8606e130 +--- /dev/null ++++ b/metal/tests/comb_fold_parity.cpp +@@ -0,0 +1,1291 @@ ++// metal/tests/comb_fold_parity.cpp - Round 5 Track A: the prompt chunk's folded mm_combine and the per-token HC read ++// against the dispatch chains they replace, BITWISE; `comb_fold_parity [T] bench` adds their GPU times at T (default ++// 1024), 48 reps = a chunk's layers: ++// gr_norm_h_moec == mm_combine + gr_norm_h_moe (R', xnh) ++// gr_apply_moec == mm_combine + gr_apply_moe (R') ++// gr_norm_h4 == gr_norm_h (R', xnh; apply 0 and 1) ++// gemm_hc_up, mixed unbound == gemm_hc_up's mixh ++// `comb_fold_parity T grvec [baseline-prompt.metal]` checks/times GR only in a short process. ++// `comb_fold_parity T hcup [bench|confirm]` checks HC-up tiles alone (both outputs, null mixed too). ++// `comb_fold_parity T epi [bench [tile,tile...]]` checks gemm_hc_down / gemm_swiglu_up against gemm_f16 + gr_epi_down / swiglu_h. ++// `comb_fold_parity T denseab baseline-dense_mm.metal [shape]` checks/times identical dense tiles across source revisions. ++// `comb_fold_parity T gu [bench [tile,tile...]]` checks/times the fused shared-expert gate + up against gate + up + swiglu_h. ++// `comb_fold_parity T topk baseline-prompt.metal` checks/times route_topk against a baseline copy's. ++// `comb_fold_parity T ab baseline-prompt.metal [gr|gemm]` checks/times this prompt.metal's kernels against a baseline's. ++// `comb_fold_parity 0 mtp [bench] [baseline-mtp.metal]` checks/times MTP's batched fc_hidden (mtp_parity). ++// `comb_fold_parity 0 mtpcat` checks/times MTP's gm_cat against its two GEMVs (mtp_cat_parity). ++#include "mtl_rt.h" ++#include "ref_quant.h" ++#include "coop_probe.h" ++ ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++ ++namespace { ++struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; ++struct GemmArgs { int M, N, K, pad; }; ++struct EpiArgs { uint32_t n_tok, n_embd, hc, hc_lr; }; ++bool same(const mtl::Buffer& a, const mtl::Buffer& b, size_t bytes) { return std::memcmp(a.ptr, b.ptr, bytes) == 0; } ++ ++int hc_up_parity(uint32_t T, bool bench, bool confirm = false) { ++ mtl::Context cx; ++ // Compile just this kernel for short standalone runs, without the quant/MoE libraries. ++ const std::string source = mtl::read_file(cx.kernel_dir + "/prompt.metal"); ++ const std::string prelude = "#define NO_GGML_COMMON 1\n" ++ "#include \n" ++ "using namespace metal; using namespace mpp::tensor_ops;\n" ++ "struct GemmArgs { int M, N, K, pad; };\n" ++ "struct EpiArgs { uint n_tok, n_embd, hc, hc_lr; };\n"; ++ const auto begin = source.find("// gemm_hc_up:"); ++ if (begin == std::string::npos) { ++ std::fprintf(stderr, "gemm_hc_up source marker missing\n"); ++ return 1; ++ } ++ std::string extra; ++ for (const auto shape : {std::pair{32, 128}, {64, 128}, {32, 64}, {64, 64}, {16, 64}}) { ++ const auto [tm, tn] = shape; ++ const std::string args = std::to_string(tm) + ", " + std::to_string(tn) + ", 2, true"; ++ extra += "template [[host_name(\"gemm_hc_up_reg_" + std::to_string(tm) + "_" + std::to_string(tn) + ++ "_2\")]] kernel decltype(gemm_hc_up<" + args + ">) gemm_hc_up<" + args + ">;\n"; ++ } ++ const int lib = cx.compile({}, prelude + source.substr(begin) + extra); ++ const bool k4 = coop_k4(cx); ++ constexpr uint32_t N = 2560, HC = 4, LR = 320; ++ const size_t RN = size_t(N) * HC; ++ auto A = cx.alloc(2 * T * LR), B = cx.alloc(2 * RN * LR), X = cx.alloc(2 * T * RN); ++ auto mr = cx.alloc(4 * T * N), hr = cx.alloc(2 * T * N); ++ auto mc = cx.alloc(4 * T * N), hc = cx.alloc(2 * T * N); ++ const mtl::Buffer nil{}; ++ ref::Rng rng(23 + T); ++ for (const auto* b : {&A, &B, &X}) ++ for (size_t i = 0; i < b->bytes / 2; ++i) ++ b->as()[i] = ref::f32_to_f16(rng.sym() * (b == &B ? 0.1f : 1.0f)); ++ const GemmArgs g{int(T), int(RN), int(LR), 0}; ++ const EpiArgs e{T, N, HC, LR}; ++ auto dispatch = [&](const mtl::Pipeline& p, int tm, int tn, const mtl::Buffer& m, const mtl::Buffer& h, int sg = 4) { ++ mtl::Dispatch d; ++ d.p = &p; ++ d.args = {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(X), mtl::Arg::B(m), mtl::Arg::B(h), ++ mtl::Arg::V(g), mtl::Arg::V(e)}; ++ d.grid[0] = RN / tn; ++ d.grid[1] = (T + tm - 1) / tm; ++ d.tg[0] = sg * 32; ++ return d; ++ }; ++ auto baseline = cx.pipeline(lib, "gemm_hc_up_64_64_4"); ++ cx.run({dispatch(baseline, 64, 64, mr, hr)}); ++ int fails = 0; ++ std::printf("device %s, T %u, register layout %s\n", cx.device_name().c_str(), T, k4 ? "K4" : "SKIP"); ++ for (bool reg : {false, true}) { ++ if (reg && !k4) continue; ++ for (const auto shape : {std::array{64, 64, 4}, {128, 64, 4}, {64, 128, 4}, {64, 32, 4}, ++ {32, 128, 4}, {32, 128, 2}, {64, 128, 2}, {32, 64, 2}, {64, 64, 2}, {16, 64, 2}}) { ++ const auto [tm, tn, sg] = shape; ++ if (!reg && (sg == 2 || (tm == 32 && tn == 128))) continue; ++ const std::string name = std::string("gemm_hc_up_") + (reg ? "reg_" : "") + ++ std::to_string(tm) + "_" + std::to_string(tn) + "_" + std::to_string(sg); ++ auto p = cx.pipeline(lib, name); ++ for (bool bound : {true, false}) { ++ std::memset(mc.ptr, 0xCD, mc.bytes); ++ std::memset(hc.ptr, 0xCD, hc.bytes); ++ auto d = dispatch(p, tm, tn, bound ? mc : nil, hc, sg); ++ cx.run({d}); ++ const bool ok = same(hr, hc, hr.bytes) && (!bound || same(mr, mc, mr.bytes)); ++ fails += !ok; ++ std::printf("%s mixed %s: bitwise %s tgmem %d", name.c_str(), bound ? "bound" : "unbound", ++ ok ? "PASS" : "FAIL", p.static_tg_mem); ++ if (bench) { ++ cx.warm({d}, 0.01); ++ double best = 1e9; ++ for (int i = 0; i < 3; ++i) best = std::min(best, cx.run({d}, 48)); ++ std::printf(" ms48 %.3f", best * 1e3); ++ } ++ std::printf("\n"); ++ } ++ } ++ } ++ if (fails) return 1; ++ if (confirm && k4) { ++ auto r64 = cx.pipeline(lib, "gemm_hc_up_reg_64_64_4"); ++ auto r128 = cx.pipeline(lib, "gemm_hc_up_reg_64_128_4"); ++ for (bool bound : {false, true}) { ++ const auto& m = bound ? mc : nil; ++ const std::vector ds = {dispatch(baseline, 64, 64, m, hc), ++ dispatch(r64, 64, 64, m, hc), dispatch(r128, 64, 128, m, hc)}; ++ for (const auto& d : ds) cx.warm({d}, 0.01); ++ for (int rep = 0; rep < 5; ++rep) { ++ double ms[3]; ++ for (int j = 0; j < 3; ++j) { ++ const int i = (j + rep) % 3; ++ ms[i] = cx.run({ds[i]}, 48) * 1e3; ++ } ++ std::printf("confirm mixed %s rep %d baseline %.3f reg64 %.3f reg128 %.3f\n", ++ bound ? "bound" : "unbound", rep, ms[0], ms[1], ms[2]); ++ } ++ } ++ } ++ return 0; ++} ++ ++// the window's gemm tile for (M, N, K) (window.cpp prompt_chunk: gemm), as "TM_TN_SG" ++std::string window_tile(int M, int Nn, int K) { ++ int tm = 64, tn = 128, sg = 4; ++ auto set = [&](int m, int n, int s) { tm = m; tn = n; sg = s; }; ++ if (M <= 128) { ++ if (Nn <= 64) set(64, 32, 4); ++ else if (Nn <= 640 || K > 2560) set(64, Nn <= 128 ? 32 : 64, 4); ++ else set(M >= 128 ? 128 : 64, 64, 4); ++ } else if (Nn <= 64) { ++ M < 1024 ? set(16, 64, 2) : set(64, 32, 2); ++ } else if (Nn == 96) { ++ M < 640 ? set(16, 64, 2) : set(16, 96, 2); ++ } else if (Nn <= 128) { ++ M < 768 ? set(16, 64, 2) : set(32, 64, 2); ++ } else if (Nn <= 640 && K > 2560) { ++ M <= 512 ? set(64, 64, 4) : M < 1536 ? set(128, 64, 4) : set(32, 128, 2); ++ } else if (Nn <= 512 || (Nn <= 640 && M < 384)) { ++ M < 384 ? set(32, 64, 2) : set(32, 128, 2); ++ } else if (Nn <= 1664) { ++ M < 768 ? set(32, 128, 2) : set(64, 128, 4); ++ } else { ++ M < 1536 && (M % 128 == 0 || M % 128 > 64) ? set(128, 64, 4) : set(64, 128, 4); ++ } ++ return std::to_string(tm) + "_" + std::to_string(tn) + "_" + std::to_string(sg); ++} ++ ++// `comb_fold_parity T epi [bench]`: gemm_hc_down == gemm_f16 + gr_epi_down (lo_h, inject) and gemm_swiglu_up == ++// gemm_f16 (up) + swiglu_h (h), BITWISE, against the window's gemm_f16 tile, every epilogue tile; T 0 = the sweep ++// 1..2048. bench: paired GPU times at T, 96 HC reads / 48 shared experts (a chunk's), alternating order, median of 5 ++int epi_parity(uint32_t TB, bool bench, const char* only) { ++ mtl::Context cx; ++ const std::string dm = mtl::read_file(cx.kernel_dir + "/dense_mm.metal"), pr = mtl::read_file(cx.kernel_dir + "/prompt.metal"); ++ const auto gb = dm.find("template "), ge = dm.find("#define DQ(F)"); ++ const auto ib = dm.find("template [[host_name(\"gemm_f16_"); ++ const auto eb = pr.find("struct EpiArgs {"), ee = pr.find("// ---- the HC read"), hb = pr.find("// gemm_hc_up:"); ++ if (gb == std::string::npos || ge == std::string::npos || ib == std::string::npos || eb == std::string::npos || ++ ee == std::string::npos || hb == std::string::npos) { ++ std::fprintf(stderr, "epilogue source markers missing\n"); ++ return 1; ++ } ++ const std::string prelude = "#define NO_GGML_COMMON 1\n" ++ "#include \n" ++ "using namespace metal; using namespace mpp::tensor_ops;\n" ++ "struct GemmArgs { int M, N, K; int pad; };\n"; ++ if (!coop_k4(cx)) { ++ std::printf("device %s: not the M5 layout, the register epilogues are not used: SKIP\n", cx.device_name().c_str()); ++ return 0; ++ } ++ // the sweep's other tiles (and gemm_f16_32_128_4, so each epilogue tile has its own plain GEMM), only here ++ std::string extra = "template [[host_name(\"gemm_f16_32_128_4\")]] kernel decltype(gemm_f16<32, 128, 4>) gemm_f16<32, 128, 4>;\n"; ++ for (const char* t : {"64_64_4", "128_64_4", "64_128_4", "32_128_4", "32_128_2", "16_64_2", "32_64_2"}) { ++ int tm, tn, sg; ++ std::sscanf(t, "%d_%d_%d", &tm, &tn, &sg); ++ const std::string targs = std::to_string(tm) + ", " + std::to_string(tn) + ", " + std::to_string(sg); ++ for (const char* k : {"gemm_hc_down", "gemm_swiglu_up"}) { ++ const std::string name = std::string(k) + "_" + t; ++ if (pr.find("\"" + name + "\"") == std::string::npos) ++ extra += "template [[host_name(\"" + name + "\")]] kernel decltype(" + k + "<" + targs + ">) " + k + "<" + targs + ">;\n"; ++ } ++ } ++ const int lib = cx.compile({}, prelude + dm.substr(gb, ge - gb) + dm.substr(ib) + pr.substr(eb, ee - eb) + pr.substr(hb) + extra); ++ constexpr uint32_t N = 2560, HC = 4, LR = 320, NFF = 640, NR = LR + HC, KD = HC * N; ++ std::printf("device %s\n", cx.device_name().c_str()); ++ std::map ps; ++ auto P = [&](const std::string& n) -> const mtl::Pipeline& { ++ auto it = ps.find(n); ++ if (it == ps.end()) it = ps.emplace(n, cx.pipeline(lib, n)).first; ++ return it->second; ++ }; ++ auto dims = [](const std::string& tile, int& tm, int& tn, int& sg) { std::sscanf(tile.c_str(), "%d_%d_%d", &tm, &tn, &sg); }; ++ auto D = [](const mtl::Pipeline& p, std::vector a, uint32_t gx, uint32_t gy, uint32_t th, bool barrier = true) { ++ mtl::Dispatch d; ++ d.p = &p; ++ d.args = std::move(a); ++ d.grid[0] = gx; ++ d.grid[1] = gy; ++ d.tg[0] = th; ++ d.barrier = barrier; ++ return d; ++ }; ++ std::vector tiles = {"32_128_2", "64_128_4", "128_64_4", "32_128_4", "16_64_2", "64_64_4", "32_64_2"}; ++ if (only) { ++ tiles.clear(); ++ for (const char *s = only, *e; *s; s = *e ? e + 1 : e) { ++ e = std::strchr(s, ','); ++ if (!e) e = s + std::strlen(s); ++ tiles.emplace_back(s, e); ++ } ++ } ++ int fails = 0; ++ const std::vector Ts = TB ? std::vector{TB} : ++ std::vector{1u, 7u, 16u, 17u, 32u, 33u, 64u, 256u, 289u, 297u, 306u, 512u, ++ 1024u, 1025u, 1041u, 1983u, 1987u, 2048u}; ++ for (uint32_t T : Ts) { ++ mtl::Buffer xnh = cx.alloc(2 * size_t(T) * KD), WD = cx.alloc(2 * size_t(NR) * KD), V = cx.alloc(4 * size_t(T) * NR), ++ la = cx.alloc(2 * T * LR), lb = cx.alloc(2 * T * LR), ia = cx.alloc(4 * T * HC), ib_ = cx.alloc(4 * T * HC), ++ mixh = cx.alloc(2 * size_t(T) * N), Wg = cx.alloc(2 * size_t(NFF) * N), Wu = cx.alloc(2 * size_t(NFF) * N), ++ gsh = cx.alloc(4 * size_t(T) * NFF), ush = cx.alloc(4 * size_t(T) * NFF), ha = cx.alloc(2 * size_t(T) * NFF), ++ hb_ = cx.alloc(2 * size_t(T) * NFF); ++ ref::Rng r(31 + T); ++ for (const auto& [b, s] : {std::pair{&xnh, 1.0f}, {&WD, 0.02f}, {&mixh, 1.0f}, {&Wg, 0.05f}, {&Wu, 0.05f}}) ++ for (size_t i = 0; i < b->bytes / 2; ++i) b->as()[i] = ref::f32_to_f16(s * r.sym()); ++ const EpiArgs ea{T, N, HC, LR}; ++ const GemmArgs gd{int(T), int(NR), int(KD), 0}, gs{int(T), int(NFF), int(N), 0}; ++ auto gemm = [&](const std::string& tile, const mtl::Buffer& A, const mtl::Buffer& B, const mtl::Buffer& Y, ++ const GemmArgs& g, bool barrier = true) { ++ int tm, tn, sg; ++ dims(tile, tm, tn, sg); ++ return D(P("gemm_f16_" + tile), {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(Y), mtl::Arg::V(g)}, ++ (g.N + tn - 1) / tn, (g.M + tm - 1) / tm, sg * 32, barrier); ++ }; ++ auto down = [&](const std::string& tile) { ++ int tm, tn, sg; ++ dims(tile, tm, tn, sg); ++ return D(P("gemm_hc_down_" + tile), {mtl::Arg::B(xnh), mtl::Arg::B(WD), mtl::Arg::B(lb), mtl::Arg::B(ib_), ++ mtl::Arg::V(ea)}, (NR + tn - 1) / tn, (T + tm - 1) / tm, sg * 32); ++ }; ++ auto swup = [&](const std::string& tile) { ++ int tm, tn, sg; ++ dims(tile, tm, tn, sg); ++ return D(P("gemm_swiglu_up_" + tile), {mtl::Arg::B(mixh), mtl::Arg::B(Wu), mtl::Arg::B(gsh), mtl::Arg::B(hb_), ++ mtl::Arg::V(gs)}, (NFF + tn - 1) / tn, (T + tm - 1) / tm, sg * 32); ++ }; ++ const std::string wd = window_tile(int(T), int(NR), int(KD)), ws = window_tile(int(T), int(NFF), int(N)); ++ const mtl::Dispatch epi = D(P("gr_epi_down"), {mtl::Arg::B(V), mtl::Arg::B(la), mtl::Arg::B(ia), mtl::Arg::V(ea)}, ++ (T * NR + 255) / 256, 1, 256); ++ const uint32_t nh = T * NFF; ++ const mtl::Dispatch swi = D(P("swiglu_h"), {mtl::Arg::B(gsh), mtl::Arg::B(ush), mtl::Arg::B(ha), mtl::Arg::V(nh)}, ++ (nh + 255) / 256, 1, 256); ++ cx.run({gemm(wd, xnh, WD, V, gd), epi}); ++ cx.run({gemm(ws, mixh, Wg, gsh, gs), gemm(ws, mixh, Wu, ush, gs), swi}); ++ for (const auto& tile : tiles) { ++ std::memset(lb.ptr, 0xCD, lb.bytes); ++ std::memset(ib_.ptr, 0xCD, ib_.bytes); ++ std::memset(hb_.ptr, 0xCD, hb_.bytes); ++ cx.run({down(tile)}); ++ cx.run({swup(tile)}); ++ const bool okd = same(la, lb, la.bytes) && same(ia, ib_, ia.bytes), oks = same(ha, hb_, ha.bytes); ++ fails += !okd + !oks; ++ std::printf("T %4u %-11s gemm_hc_down vs %s + gr_epi_down: bitwise %s; gemm_swiglu_up vs %s + swiglu_h: bitwise %s\n", ++ T, tile.c_str(), wd.c_str(), okd ? "PASS" : "FAIL", ws.c_str(), oks ? "PASS" : "FAIL"); ++ } ++ if (!bench) continue; ++ // 96 HC reads: [gemm_f16, gr_epi_down] vs [gemm_hc_down]; 48 shared experts: [gate, up (overlapped), swiglu_h] ++ // vs [gate, gemm_swiglu_up] - each a dependency graph (run_graph) with the window's barriers; every tile both ++ // ways, so a row's change is the fusion's alone (vs its own tile) and the window's (vs its current tile) ++ auto rep = [](std::vector one, int n) { ++ std::vector all; ++ for (int i = 0; i < n; ++i) all.insert(all.end(), one.begin(), one.end()); ++ return all; ++ }; ++ struct Cand { std::string name, tile; bool fused; std::vector ds; }; ++ std::vector cand; ++ for (const auto& tile : tiles) { ++ cand.push_back({"down", tile, false, rep({gemm(tile, xnh, WD, V, gd), epi}, 96)}); ++ cand.push_back({"down", tile, true, rep({down(tile)}, 96)}); ++ } ++ for (const auto& tile : tiles) { ++ cand.push_back({"swiglu", tile, false, rep({gemm(ws, mixh, Wg, gsh, gs), gemm(tile, mixh, Wu, ush, gs, false), swi}, 48)}); ++ cand.push_back({"swiglu", tile, true, rep({gemm(ws, mixh, Wg, gsh, gs), swup(tile)}, 48)}); ++ } ++ cand.push_back({"down", wd, false, rep({gemm(wd, xnh, WD, V, gd)}, 96)}); // the GEMM alone: the epilogue's ceiling ++ cand.push_back({"down", wd, false, rep({gemm(wd, xnh, WD, V, gd)}, 96)}); ++ for (const auto& c : cand) cx.warm(c.ds, 0.01); ++ constexpr int NS = 5; ++ std::vector> ms(cand.size()); ++ for (int k = 0; k < NS; ++k) ++ for (size_t j = 0; j < cand.size(); ++j) { ++ const size_t i = (k % 2) ? cand.size() - 1 - j : j; ++ ms[i].push_back(cx.run_graph(cand[i].ds) * 1e3); ++ } ++ std::printf("T %u, ms per chunk (median of %d, alternating order); window tiles: down %s, swiglu gate/up %s\n", T, NS, ++ wd.c_str(), ws.c_str()); ++ std::vector med(cand.size()); ++ for (size_t i = 0; i < cand.size(); ++i) { ++ std::sort(ms[i].begin(), ms[i].end()); ++ med[i] = ms[i][NS / 2]; ++ } ++ for (size_t i = 0; i + 2 < cand.size(); i += 2) { ++ double wbase = 0; ++ for (size_t j = 0; j + 2 < cand.size(); j += 2) ++ if (cand[j].name == cand[i].name && cand[j].tile == (cand[i].name == "down" ? wd : ws)) wbase = med[j]; ++ std::printf(" %-6s %-8s plain %8.3f fused %8.3f (%+5.1f%% vs own tile", cand[i].name.c_str(), cand[i].tile.c_str(), ++ med[i], med[i + 1], 100.0 * (med[i + 1] / med[i] - 1)); ++ if (wbase > 0) std::printf(", %+5.1f%% vs window", 100.0 * (med[i + 1] / wbase - 1)); ++ std::printf(")\n"); ++ } ++ std::printf(" down %-8s GEMM alone, no gr_epi_down: %8.3f / %8.3f\n", wd.c_str(), med[cand.size() - 2], med.back()); ++ } ++ std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); ++ return fails ? 1 : 0; ++} ++ ++// Same tile and buffers, different dense_mm source revisions; --ab in dense_mm_parity compares tile policies. ++int dense_ab_parity(uint32_t T, const char* base_path, const char* only) { ++ if (!T || T > 2048) return 1; ++ mtl::Context cx; ++ int lib[2]; ++ for (int v = 0; v < 2; ++v) { ++ const auto src = mtl::read_file(v ? base_path : cx.kernel_dir + "/dense_mm.metal"); ++ const auto b = src.find("struct GemmArgs"), e = src.find("#define DQ(F)"); ++ const auto i = src.find("template [[host_name(\"gemm_f16_"); ++ if (b == std::string::npos || e == std::string::npos || i == std::string::npos || e <= b) { ++ std::fprintf(stderr, "dense source markers missing\n"); ++ return 1; ++ } ++ lib[v] = cx.compile({}, "#define NO_GGML_COMMON 1\n" ++ "#include \n" ++ "using namespace metal; using namespace mpp::tensor_ops;\n" + src.substr(b, e - b) + src.substr(i)); ++ } ++ struct Shape { const char* name; int n, k; }; ++ const Shape shapes[] = {{"hc_read", 324, 10240}, {"Wqkv", 10240, 2560}, {"Wz", 6144, 2560}, ++ {"Wout", 2560, 6144}, {"Wq", 12288, 2560}, {"Wpv", 2560, 2560}, {"Wshgu", 640, 2560}, ++ {"Wshd", 2560, 640}, {"W512", 512, 2560}, {"Wk", 1152, 2560}, {"Wk_sel", 1664, 2560}, ++ {"Walbe", 96, 2560}, {"Wsg", 1, 2560}}; ++ auto A = cx.alloc(2 * size_t(T) * 10240), B = cx.alloc(2 * size_t(12288) * 2560); ++ mtl::Buffer Y[2] = {cx.alloc(4 * size_t(T) * 12288), cx.alloc(4 * size_t(T) * 12288)}; ++ ref::Rng rng(61 + T); ++ for (const auto* b : {&A, &B}) ++ for (size_t i = 0; i < b->bytes / 2; ++i) b->as()[i] = ref::f32_to_f16(rng.sym() * 0.1f); ++ int fails = 0, checked = 0; ++ for (const auto& s : shapes) { ++ if (only && std::strcmp(only, s.name) != 0) continue; ++ ++checked; ++ const auto tile = window_tile(T, s.n, s.k); ++ int tm, tn, sg; ++ std::sscanf(tile.c_str(), "%d_%d_%d", &tm, &tn, &sg); ++ mtl::Pipeline p[2] = {cx.pipeline(lib[0], "gemm_f16_" + tile), cx.pipeline(lib[1], "gemm_f16_" + tile)}; ++ const GemmArgs g{int(T), s.n, s.k, 0}; ++ mtl::Dispatch d[2][2]; ++ for (int v = 0; v < 2; ++v) ++ for (int b = 0; b < 2; ++b) { ++ auto& x = d[v][b]; ++ x.p = &p[v]; ++ x.args = {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(Y[b]), mtl::Arg::V(g)}; ++ x.grid[0] = (s.n + tn - 1) / tn; ++ x.grid[1] = (T + tm - 1) / tm; ++ x.tg[0] = sg * 32; ++ } ++ for (int v = 0; v < 2; ++v) { ++ std::memset(Y[v].ptr, 0xCD, Y[v].bytes); ++ cx.run({d[v][v]}); ++ } ++ const bool ok = same(Y[0], Y[1], 4 * size_t(T) * s.n); ++ fails += !ok; ++ for (int v = 0; v < 2; ++v) cx.warm({d[v][v]}, 0.005); ++ std::vector ratios; ++ for (int r = 0; r < 12; ++r) { ++ double ms[2]; ++ for (int j = 0; j < 2; ++j) { ++ const int v = (r / 2 + j) % 2; ++ ms[v] = cx.run({d[v][(r + v) % 2]}, 8); ++ } ++ ratios.push_back(ms[0] / ms[1]); ++ } ++ std::sort(ratios.begin(), ratios.end()); ++ std::printf("M %u %-8s %s bitwise %s ratio %.4f [%.4f, %.4f]\n", T, s.name, tile.c_str(), ++ ok ? "PASS" : "FAIL", (ratios[5] + ratios[6]) * 0.5, ratios[2], ratios[9]); ++ } ++ std::printf("%s\n", fails || !checked ? "FAIL" : "ALL PASS"); ++ return fails || !checked ? 1 : 0; ++} ++ ++// Round 29's shared-expert gate + up fusions, measured here and not shipped (audits/opus_r29_prompt.md): compute-bound, ++// they tie gate + gemm_swiglu_up past 160 rows outside a few tile-aligned T ++const char* kSwigluFused = R"MSL( ++// gemm_swiglu_gu: the shared expert's gate and up GEMMs in one kernel, h = half(silu(x . w_gate^T) * (x . w_up^T)), ++// both over the same A tile; g and u are never written. h bitwise gemm_f16 (gate, up) + swiglu_h. Two accumulators ++// (silu(g) waits in the first). M5 layout only (as gemm_hc_down); N % 4 == 0. grid (ceil(N / TN), ceil(M / TM)), ++// 32 * SG threads ++template ++kernel void gemm_swiglu_gu(device half* A [[buffer(0)]], device half* Bg [[buffer(1)]], device half* Bu [[buffer(2)]], ++ device half* h [[buffer(3)]], constant GemmArgs& a [[buffer(4)]], ++ uint3 tg [[threadgroup_position_in_grid]]) { ++ tensor, tensor_inline> tA(A, dextents(a.K, a.M)); ++ tensor, tensor_inline> tBg(Bg, dextents(a.K, a.N)); ++ tensor, tensor_inline> tBu(Bu, dextents(a.K, a.N)); ++ constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); ++ matmul2d> op; ++ auto mA = tA.slice(0, int(tg.y) * TM); ++ auto mBg = tBg.slice(0, int(tg.x) * TN); ++ auto mBu = tBu.slice(0, int(tg.x) * TN); ++ auto sg = op.template get_destination_cooperative_tensor(); ++ auto acc = op.template get_destination_cooperative_tensor(); ++ op.run(mA, mBg, sg); ++ #pragma clang loop unroll(full) ++ for (uint16_t i = 0; i < sg.get_capacity(); ++i) sg[i] = sg[i] / (1.0f + exp(-sg[i])); ++ op.run(mA, mBu, acc); ++ #pragma clang loop unroll(full) ++ for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { ++ const auto ix = acc.get_multidimensional_index(i); ++ const int c = int(tg.x) * TN + ix[0], t = int(tg.y) * TM + ix[1]; ++ if (t >= a.M || c >= a.N) continue; ++ half4 o; ++ #pragma clang loop unroll(full) ++ for (int k = 0; k < 4; ++k) o[k] = half(sg[i + k] * acc[i + k]); ++ *(device half4*)(h + ulong(t) * uint(a.N) + uint(c)) = o; ++ } ++} ++ ++// gemm_swiglu_il: gate and up as one GEMM over B = w_gate and w_up rows interleaved (row 2j = gate j, 2j + 1 = up j; ++// dq_f16 perm 2 over the pair), one accumulator; on M5 a lane's four columns are (g, u) of two outputs. ++// grid (ceil(2N / TN), ceil(M / TM)), a.N = N ++template ++kernel void gemm_swiglu_il(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device half* h [[buffer(2)]], ++ constant GemmArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]]) { ++ tensor, tensor_inline> tA(A, dextents(a.K, a.M)); ++ tensor, tensor_inline> tB(B, dextents(a.K, 2 * a.N)); ++ constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); ++ matmul2d> op; ++ auto mA = tA.slice(0, int(tg.y) * TM); ++ auto mB = tB.slice(0, int(tg.x) * TN); ++ auto acc = op.template get_destination_cooperative_tensor(); ++ op.run(mA, mB, acc); ++ #pragma clang loop unroll(full) ++ for (uint16_t i = 0; i < acc.get_capacity(); i += 4) { ++ const auto ix = acc.get_multidimensional_index(i); ++ const int c = (int(tg.x) * TN + ix[0]) / 2, t = int(tg.y) * TM + ix[1]; ++ if (t >= a.M || c >= a.N) continue; ++ const half2 o(half(acc[i] / (1.0f + exp(-acc[i])) * acc[i + 1]), half(acc[i + 2] / (1.0f + exp(-acc[i + 2])) * acc[i + 3])); ++ *(device half2*)(h + ulong(t) * uint(a.N) + uint(c)) = o; ++ } ++} ++)MSL"; ++ ++// `comb_fold_parity T gu [bench [tile,tile...]]`: gemm_swiglu_gu / gemm_swiglu_il (gate + up + swiglu_h in one kernel) ++// against gemm_f16 (gate) + gemm_f16 (up) + swiglu_h and the window's gate + gemm_swiglu_up, BITWISE; T 0 = 1, 7, 161, ++// 256, 512, 1024, 1983, 2048; every tile (TM_TN_SG, or ilTM_TN_SG). bench: paired GPU times of 48 shared experts (a chunk's): the window's old ++// paths (gate ∥ up + swiglu_h, gate + gemm_swiglu_up) vs each tile, alternating order, median of 9 ++int gu_parity(uint32_t TB, bool bench, const char* only) { ++ mtl::Context cx; ++ const std::string dm = mtl::read_file(cx.kernel_dir + "/dense_mm.metal"), pr = mtl::read_file(cx.kernel_dir + "/prompt.metal"); ++ const auto gb = dm.find("template "), ge = dm.find("#define DQ(F)"); ++ const auto ib = dm.find("template [[host_name(\"gemm_f16_"); ++ const auto eb = pr.find("struct EpiArgs {"), ee = pr.find("// ---- the HC read"), hb = pr.find("// gemm_hc_up:"); ++ if (gb == std::string::npos || ge == std::string::npos || ib == std::string::npos || eb == std::string::npos || ++ ee == std::string::npos || hb == std::string::npos) { ++ std::fprintf(stderr, "gu source markers missing\n"); ++ return 1; ++ } ++ if (!coop_k4(cx)) { ++ std::printf("device %s: not the M5 layout, the fused gate + up is M5 only: SKIP\n", cx.device_name().c_str()); ++ return 0; ++ } ++ std::vector tiles = {"32_64_2", "64_64_4", "32_128_4", "64_128_4", "128_64_4", "64_128_8", "128_128_8"}; ++ if (only) { ++ tiles.clear(); ++ for (const char *s = only, *e; *s; s = *e ? e + 1 : e) { ++ e = std::strchr(s, ','); ++ if (!e) e = s + std::strlen(s); ++ tiles.emplace_back(s, e); ++ } ++ } ++ // a tile TM_TN_SG runs gemm_swiglu_gu, ilTM_TN_SG gemm_swiglu_il (its B the interleaved gate and up) ++ auto parse = [](const std::string& v, int& tm, int& tn, int& sg) { ++ const bool il = v.rfind("il", 0) == 0; ++ std::sscanf(v.c_str() + (il ? 2 : 0), "%d_%d_%d", &tm, &tn, &sg); ++ return il; ++ }; ++ auto kname = [](const std::string& v) { return "gemm_swiglu_" + (v.rfind("il", 0) == 0 ? v : "gu_" + v); }; ++ std::string extra; ++ for (const auto& v : tiles) { ++ int tm = 0, tn = 0, sg = 0; ++ const std::string k = parse(v, tm, tn, sg) ? "gemm_swiglu_il" : "gemm_swiglu_gu"; ++ const std::string targs = std::to_string(tm) + ", " + std::to_string(tn) + ", " + std::to_string(sg); ++ extra += "template [[host_name(\"" + kname(v) + "\")]] kernel decltype(" + k + "<" + targs + ">) " + k + "<" + targs + ">;\n"; ++ } ++ const std::string prelude = "#define NO_GGML_COMMON 1\n" ++ "#include \n" ++ "using namespace metal; using namespace mpp::tensor_ops;\n" ++ "struct GemmArgs { int M, N, K; int pad; };\n"; ++ const int lib = cx.compile({}, prelude + dm.substr(gb, ge - gb) + dm.substr(ib) + pr.substr(eb, ee - eb) + pr.substr(hb) + ++ kSwigluFused + extra); ++ constexpr uint32_t N = 2560, NFF = 640; ++ std::printf("device %s\n", cx.device_name().c_str()); ++ std::map ps; ++ auto P = [&](const std::string& n) -> const mtl::Pipeline& { ++ auto it = ps.find(n); ++ if (it == ps.end()) it = ps.emplace(n, cx.pipeline(lib, n)).first; ++ return it->second; ++ }; ++ auto D = [](const mtl::Pipeline& p, std::vector a, uint32_t gx, uint32_t gy, uint32_t th, bool barrier = true) { ++ mtl::Dispatch d; ++ d.p = &p; ++ d.args = std::move(a); ++ d.grid[0] = gx; ++ d.grid[1] = gy; ++ d.tg[0] = th; ++ d.barrier = barrier; ++ return d; ++ }; ++ int fails = 0; ++ const std::vector Ts = TB ? std::vector{TB} : ++ std::vector{1u, 7u, 161u, 256u, 512u, 1024u, 1983u, 2048u}; ++ for (uint32_t T : Ts) { ++ mtl::Buffer mixh = cx.alloc(2 * size_t(T) * N), Wg = cx.alloc(2 * size_t(NFF) * N), Wu = cx.alloc(2 * size_t(NFF) * N), ++ gsh = cx.alloc(4 * size_t(T) * NFF), ush = cx.alloc(4 * size_t(T) * NFF), ha = cx.alloc(2 * size_t(T) * NFF), ++ hw = cx.alloc(2 * size_t(T) * NFF), hb_ = cx.alloc(2 * size_t(T) * NFF); ++ ref::Rng r(37 + T); ++ for (const auto& [b, s] : {std::pair{&mixh, 1.0f}, {&Wg, 0.05f}, {&Wu, 0.05f}}) ++ for (size_t i = 0; i < b->bytes / 2; ++i) b->as()[i] = ref::f32_to_f16(s * r.sym()); ++ const GemmArgs gs{int(T), int(NFF), int(N), 0}; ++ mtl::Buffer Wgu = cx.alloc(4 * size_t(NFF) * N); // gate and up rows interleaved (row 2j gate j, 2j + 1 up j) ++ for (uint32_t j = 0; j < NFF; ++j) { ++ std::memcpy(Wgu.as() + size_t(2 * j) * N, Wg.as() + size_t(j) * N, 2 * N); ++ std::memcpy(Wgu.as() + size_t(2 * j + 1) * N, Wu.as() + size_t(j) * N, 2 * N); ++ } ++ const std::string ws = window_tile(int(T), int(NFF), int(N)); ++ int wm, wn, wsg; ++ std::sscanf(ws.c_str(), "%d_%d_%d", &wm, &wn, &wsg); ++ auto gemm = [&](const mtl::Buffer& B, const mtl::Buffer& Y, bool barrier = true) { ++ return D(P("gemm_f16_" + ws), {mtl::Arg::B(mixh), mtl::Arg::B(B), mtl::Arg::B(Y), mtl::Arg::V(gs)}, ++ (NFF + wn - 1) / wn, (T + wm - 1) / wm, wsg * 32, barrier); ++ }; ++ // the window past 160 rows: gate, then gemm_swiglu_up at its tile ++ const bool small = T < 384; ++ const mtl::Dispatch swup = D(P(small ? "gemm_swiglu_up_32_64_2" : "gemm_swiglu_up_64_128_4"), ++ {mtl::Arg::B(mixh), mtl::Arg::B(Wu), mtl::Arg::B(gsh), mtl::Arg::B(hw), mtl::Arg::V(gs)}, ++ (NFF + (small ? 64 : 128) - 1) / (small ? 64 : 128), (T + (small ? 32 : 64) - 1) / (small ? 32 : 64), small ? 64 : 128); ++ const uint32_t nh = T * NFF; ++ const mtl::Dispatch swi = D(P("swiglu_h"), {mtl::Arg::B(gsh), mtl::Arg::B(ush), mtl::Arg::B(ha), mtl::Arg::V(nh)}, ++ (nh + 255) / 256, 1, 256); ++ cx.run({gemm(Wg, gsh), gemm(Wu, ush), swi}); ++ std::memset(hw.ptr, 0xCD, hw.bytes); ++ cx.run({swup}); ++ const bool okw = same(ha, hw, ha.bytes); ++ fails += !okw; ++ std::printf("T %4u window gate %s + gemm_swiglu_up vs gate + up + swiglu_h: bitwise %s\n", T, ws.c_str(), okw ? "PASS" : "FAIL"); ++ auto gu = [&](const std::string& v, const mtl::Buffer& h) { ++ int tm = 0, tn = 0, sg = 0; ++ if (parse(v, tm, tn, sg)) ++ return D(P(kname(v)), {mtl::Arg::B(mixh), mtl::Arg::B(Wgu), mtl::Arg::B(h), mtl::Arg::V(gs)}, ++ (2 * NFF + tn - 1) / tn, (T + tm - 1) / tm, sg * 32); ++ return D(P(kname(v)), {mtl::Arg::B(mixh), mtl::Arg::B(Wg), mtl::Arg::B(Wu), mtl::Arg::B(h), mtl::Arg::V(gs)}, ++ (NFF + tn - 1) / tn, (T + tm - 1) / tm, sg * 32); ++ }; ++ for (const auto& v : tiles) { ++ std::memset(hb_.ptr, 0xCD, hb_.bytes); ++ cx.run({gu(v, hb_)}); ++ const bool ok = same(ha, hb_, ha.bytes); ++ fails += !ok; ++ std::printf("T %4u %-26s vs gate + up + swiglu_h: bitwise %s\n", T, kname(v).c_str(), ok ? "PASS" : "FAIL"); ++ } ++ if (!bench) continue; ++ auto rep = [](std::vector one, int n) { ++ std::vector all; ++ for (int i = 0; i < n; ++i) all.insert(all.end(), one.begin(), one.end()); ++ return all; ++ }; ++ // [gate, up overlapped, swiglu_h] (the window at <= 160 rows), [gate, gemm_swiglu_up] (past 160), each tile ++ std::vector> cand = {rep({gemm(Wg, gsh), gemm(Wu, ush, false), swi}, 48), ++ rep({gemm(Wg, gsh), swup}, 48)}; ++ for (const auto& v : tiles) cand.push_back(rep({gu(v, hb_)}, 48)); ++ for (const auto& c : cand) cx.warm(c, 0.01); ++ constexpr int NS = 9; ++ std::vector> ms(cand.size()); ++ for (int k = 0; k < NS; ++k) ++ for (size_t j = 0; j < cand.size(); ++j) { ++ const size_t i = (k % 2) ? cand.size() - 1 - j : j; ++ ms[i].push_back(cx.run_graph(cand[i]) * 1e3); ++ } ++ std::vector med(cand.size()); ++ for (size_t i = 0; i < cand.size(); ++i) { ++ std::sort(ms[i].begin(), ms[i].end()); ++ med[i] = ms[i][NS / 2]; ++ } ++ std::printf("T %u, ms per chunk (48 shared experts, median of %d, alternating order)\n" ++ " gate ∥ up %s + swiglu_h %8.3f\n gate %s + gemm_swiglu_up_%-8s %8.3f\n", ++ T, NS, ws.c_str(), med[0], ws.c_str(), small ? "32_64_2" : "64_128_4", med[1]); ++ for (size_t i = 0; i < tiles.size(); ++i) ++ std::printf(" %-26s %8.3f (%+5.1f%% vs gate ∥ up, %+5.1f%% vs gemm_swiglu_up)\n", kname(tiles[i]).c_str(), ++ med[i + 2], 100.0 * (med[i + 2] / med[0] - 1), 100.0 * (med[i + 2] / med[1] - 1)); ++ } ++ std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); ++ return fails ? 1 : 0; ++} ++ ++// `comb_fold_parity T topk baseline-prompt.metal`: this prompt.metal's route_topk against a baseline copy's, BITWISE ++// (ids and w) for prompt_moce 0, 320335 (the streamed default: tail 32, protect 3, bonus 35) and -16.3 / -0.05 (skip), ++// with paired GPU times: 48 reps a sample, 11 samples in alternating order, median of the per-sample ratios. T 0 = ++// 512, 1024, 2048. Logits are rounded to 1/16 so ties are common; half the experts are not resident in smap. ++int topk_parity(uint32_t TB, const char* base_path) { ++ mtl::Context cx; ++ int lib[2]; ++ for (int v = 0; v < 2; ++v) { ++ const std::string src = mtl::read_file(v ? std::string(base_path) : cx.kernel_dir + "/prompt.metal"); ++ const auto e = src.find("struct EpiArgs {"); ++ if (e == std::string::npos) { ++ std::fprintf(stderr, "route_topk source marker missing\n"); ++ return 1; ++ } ++ lib[v] = cx.compile({}, "using namespace metal;\n" + src.substr(0, e)); ++ } ++ const mtl::Pipeline pn = cx.pipeline(lib[0], "route_topk"), pb = cx.pipeline(lib[1], "route_topk"); ++ std::printf("device %s\n", cx.device_name().c_str()); ++ constexpr uint32_t K = 10; ++ int fails = 0; ++ std::vector> cases; // (T, n_exp) ++ for (uint32_t T : TB ? std::vector{TB} : std::vector{512u, 1024u, 2048u}) ++ for (uint32_t NE : {512u, 256u}) cases.emplace_back(T, NE); ++ for (const auto& [T, NE] : cases) { ++ auto lg = cx.alloc(4 * size_t(T) * NE), sm = cx.alloc(4 * NE); ++ mtl::Buffer ids[2] = {cx.alloc(4 * T * K), cx.alloc(4 * T * K)}, w[2] = {cx.alloc(4 * T * K), cx.alloc(4 * T * K)}; ++ ref::Rng r(41 + T); ++ for (size_t i = 0; i < size_t(T) * NE; ++i) lg.as()[i] = std::round(16.0f * 4.0f * r.sym()) / 16.0f; ++ for (uint32_t i = 0; i < NE; ++i) sm.as()[i] = r.sym() > 0.0f ? i : 0xFFFFFFFFu; ++ const uint32_t nt = T | (NE << 16); ++ auto D = [&](const mtl::Pipeline& p, int v, float moce) { ++ mtl::Dispatch d; ++ d.p = &p; ++ d.args = {mtl::Arg::B(lg), mtl::Arg::B(ids[v]), mtl::Arg::B(w[v]), mtl::Arg::V(nt), mtl::Arg::B(sm), mtl::Arg::V(moce)}; ++ d.grid[0] = (T + 7) / 8; ++ d.tg[0] = 256; ++ return d; ++ }; ++ for (const float moce : {0.0f, 320335.0f, -16.3f, -0.05f}) { ++ for (int v = 0; v < 2; ++v) { ++ std::memset(ids[v].ptr, 0xCD, ids[v].bytes); ++ std::memset(w[v].ptr, 0xCD, w[v].bytes); ++ cx.run({D(v ? pb : pn, v, moce)}); ++ } ++ const bool ok = same(ids[0], ids[1], ids[0].bytes) && same(w[0], w[1], w[0].bytes); ++ fails += !ok; ++ std::printf("T %4u n_exp %u route_topk moce %9.2f: ids + w bitwise %s\n", T, NE, moce, ok ? "PASS" : "FAIL"); ++ } ++ for (const float moce : {0.0f, 320335.0f}) { ++ if (NE != 512) break; // the times: the model's 512 experts only ++ std::vector dn(48, D(pn, 0, moce)), db(48, D(pb, 1, moce)); ++ cx.warm(dn, 0.01); ++ cx.warm(db, 0.01); ++ std::vector ratio, tn, tb; ++ for (int k = 0; k < 11; ++k) { ++ double a, b; ++ if (k % 2) { b = cx.run(db); a = cx.run(dn); } else { a = cx.run(dn); b = cx.run(db); } ++ tn.push_back(a * 1e3); ++ tb.push_back(b * 1e3); ++ ratio.push_back(a / b); ++ } ++ for (auto* v : {&ratio, &tn, &tb}) std::sort(v->begin(), v->end()); ++ std::printf("T %4u route_topk moce %9.2f x48: new %.3f ms, baseline %.3f ms, ratio %.3f (median of 11)\n", T, moce, ++ tn[5], tb[5], ratio[5]); ++ } ++ } ++ std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); ++ return fails ? 1 : 0; ++} ++ ++// `comb_fold_parity T ab baseline-prompt.metal [gr|gemm]`: this prompt.metal's GR reads and register-epilogue GEMMs ++// against a baseline copy's, BITWISE (outputs of the two builds), with paired GPU times: 8 reps a sample, 11 samples in ++// alternating order, median of the per-sample ratios. T 0 = 512, 1024, 2048; the window's tiles at each T ++int ab_parity(uint32_t TB, const char* base_path, const char* only) { ++ mtl::Context cx; ++ const std::string gr_pre = "#define NO_GGML_COMMON 1\nusing namespace metal;\n" ++ "struct GrArgs { uint n_embd, hc, hc_lr, n_tok, apply; float eps; };\n"; ++ const std::string mm_pre = "#define NO_GGML_COMMON 1\n" ++ "#include \n" ++ "using namespace metal; using namespace mpp::tensor_ops;\n" ++ "struct GemmArgs { int M, N, K, pad; };\n" ++ "struct EpiArgs { uint n_tok, n_embd, hc, hc_lr; };\n"; ++ const bool do_gr = !only || std::strcmp(only, "gr") == 0, do_mm = !only || std::strcmp(only, "gemm") == 0; ++ int gl[2], ml[2]; ++ for (int v = 0; v < 2; ++v) { ++ const std::string src = mtl::read_file(v ? std::string(base_path) : cx.kernel_dir + "/prompt.metal"); ++ const auto b = src.find("// ---- the HC read"), e = src.find("// gemm_hc_up:"); ++ if (b == std::string::npos || e == std::string::npos || e <= b) { ++ std::fprintf(stderr, "source markers missing\n"); ++ return 1; ++ } ++ gl[v] = do_gr ? cx.compile({}, gr_pre + src.substr(b, e - b)) : -1; ++ ml[v] = do_mm ? cx.compile({}, mm_pre + src.substr(e)) : -1; ++ } ++ if (do_mm && !coop_k4(cx)) { ++ std::printf("device %s: not the M5 layout: SKIP\n", cx.device_name().c_str()); ++ return 0; ++ } ++ std::printf("device %s, baseline %s\n", cx.device_name().c_str(), base_path); ++ constexpr uint32_t N = 2560, HC = 4, LR = 320, NR = LR + HC, NFF = 640, TOPK = 10, NEXP = 512, TMR = 16; ++ constexpr size_t RN = size_t(HC) * N; ++ std::map ps; ++ auto P = [&](int lib, const std::string& n) -> const mtl::Pipeline& { ++ const std::string key = std::to_string(lib) + n; ++ auto it = ps.find(key); ++ if (it == ps.end()) it = ps.emplace(key, cx.pipeline(lib, n)).first; ++ return it->second; ++ }; ++ int fails = 0; ++ for (uint32_t T : TB ? std::vector{TB} : std::vector{512u, 1024u, 2048u}) { ++ const size_t mr = size_t(T) * TOPK + NEXP * TMR; ++ // per build v (0 this, 1 baseline): its own outputs (and R, updated in place) ++ mtl::Buffer R0 = cx.alloc(4 * T * RN), R[2] = {cx.alloc(4 * T * RN), cx.alloc(4 * T * RN)}, ++ xn[2] = {cx.alloc(2 * T * RN), cx.alloc(2 * T * RN)}, mh[2] = {cx.alloc(2 * T * N), cx.alloc(2 * T * N)}, ++ lo[2] = {cx.alloc(2 * T * LR), cx.alloc(2 * T * LR)}, ij[2] = {cx.alloc(4 * T * HC), cx.alloc(4 * T * HC)}, ++ hh[2] = {cx.alloc(2 * T * NFF), cx.alloc(2 * T * NFF)}; ++ mtl::Buffer YE = cx.alloc(do_gr ? 4 * mr * N : 4), pos = cx.alloc(4 * T * TOPK), sw = cx.alloc(4 * mr), ++ ys = cx.alloc(4 * T * N), sgv = cx.alloc(4 * T), inj = cx.alloc(4 * T * HC), wn = cx.alloc(4 * RN), ++ xin = cx.alloc(2 * T * RN), loh = cx.alloc(2 * T * LR), WU = cx.alloc(2 * RN * LR), ++ WD = cx.alloc(2 * NR * RN), mixh = cx.alloc(2 * T * N), Wu = cx.alloc(2 * NFF * N), gsh = cx.alloc(4 * T * NFF), ++ nil{}; ++ ref::Rng r(23 + T); ++ auto fill = [&](const mtl::Buffer& b, float s, float o) { ++ for (size_t i = 0; i < b.bytes / 4; ++i) b.as()[i] = o + s * r.sym(); ++ }; ++ auto fillh = [&](const mtl::Buffer& b, float s) { ++ for (size_t i = 0; i < b.bytes / 2; ++i) b.as()[i] = ref::f32_to_f16(s * r.sym()); ++ }; ++ fill(R0, 1.0f, 0.0f); fill(YE, 0.5f, 0.0f); fill(sw, 0.1f, 0.1f); fill(ys, 0.5f, 0.0f); fill(sgv, 3.0f, 0.0f); ++ fill(inj, 4.0f, 0.0f); fill(wn, 0.2f, 1.0f); fill(gsh, 4.0f, 0.0f); fillh(xin, 1.0f); fillh(loh, 1.0f); ++ fillh(WU, 0.1f); fillh(WD, 0.02f); fillh(mixh, 1.0f); fillh(Wu, 0.05f); ++ std::vector perm(mr); ++ std::iota(perm.begin(), perm.end(), 0); ++ std::shuffle(perm.begin(), perm.end(), std::mt19937(T)); ++ for (size_t i = 0; i < size_t(T) * TOPK; ++i) pos.as()[i] = (i % 37 == 5) ? -1 : perm[i]; ++ const GrArgs a{N, HC, LR, T, 1u, 1e-6f}, a0{N, HC, LR, T, 0u, 1e-6f}; ++ const EpiArgs ea{T, N, HC, LR}; ++ const GemmArgs gu{int(T), int(RN), int(LR), 0}, gs{int(T), int(NFF), int(N), 0}; ++ auto D = [](const mtl::Pipeline& p, std::vector args, uint32_t gx, uint32_t gy, uint32_t th) { ++ mtl::Dispatch d; ++ d.p = &p; ++ d.args = std::move(args); ++ d.grid[0] = gx; ++ d.grid[1] = gy; ++ d.tg[0] = th; ++ return d; ++ }; ++ // the window's tiles (window.cpp prompt_chunk) ++ const auto [dm, dn, dsg] = T <= 192 ? std::array{16, 64, 2} : T <= 384 ? std::array{32, 64, 2} ++ : T <= 512 ? std::array{32, 128, 4} : T < 1536 ? std::array{64, 128, 4} : std::array{32, 128, 2}; ++ const std::string down = "gemm_hc_down_" + std::to_string(dm) + "_" + std::to_string(dn) + "_" + std::to_string(dsg); ++ const std::string swup = T < 384 ? "gemm_swiglu_up_32_64_2" : "gemm_swiglu_up_64_128_4"; ++ const int sm = T < 384 ? 32 : 64, sn = T < 384 ? 64 : 128, ssg = T < 384 ? 2 : 4; ++ // build v's dispatches on buffer set b; sample s times build v on set (s + v) % 2, so placement is not the build's ++ struct Case { std::string name; mtl::Dispatch d[2][2]; std::vector out[2]; }; ++ std::vector cs; ++ for (int v = 0; v < 2; ++v) ++ for (int b = 0; b < 2; ++b) { ++ std::vector> ds; ++ std::vector> outs; ++ if (do_gr) { ++ const int l = gl[v]; ++ auto nh4 = [&](const GrArgs& g) { ++ return D(P(l, "gr_norm_h4"), {mtl::Arg::B(R[b]), mtl::Arg::B(ys), mtl::Arg::B(inj), mtl::Arg::B(wn), ++ mtl::Arg::B(xn[b]), mtl::Arg::V(g)}, 1, T, 256); ++ }; ++ ds.push_back({"gr_norm_h_moec", D(P(l, "gr_norm_h_moec"), {mtl::Arg::B(R[b]), mtl::Arg::B(YE), mtl::Arg::B(pos), ++ mtl::Arg::B(sw), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(xn[b]), ++ mtl::Arg::V(a), mtl::Arg::V(TOPK)}, 1, T, 256)}); ++ ds.push_back({"gr_apply_moec", D(P(l, "gr_apply_moec"), {mtl::Arg::B(R[b]), mtl::Arg::B(YE), mtl::Arg::B(pos), ++ mtl::Arg::B(sw), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::V(a), mtl::Arg::V(TOPK)}, ++ (N + 255) / 256, T, 256)}); ++ ds.push_back({"gr_norm_h4 apply 1", nh4(a)}); ++ ds.push_back({"gr_norm_h4 apply 0", nh4(a0)}); ++ outs.insert(outs.end(), {{&R[b], &xn[b]}, {&R[b]}, {&R[b], &xn[b]}, {&xn[b]}}); ++ } ++ if (do_mm) { ++ const int l = ml[v]; ++ ds.push_back({down, D(P(l, down), {mtl::Arg::B(xin), mtl::Arg::B(WD), mtl::Arg::B(lo[b]), mtl::Arg::B(ij[b]), ++ mtl::Arg::V(ea)}, (NR + dn - 1) / dn, (T + dm - 1) / dm, dsg * 32)}); ++ ds.push_back({"gemm_hc_up_reg_32_128_4", D(P(l, "gemm_hc_up_reg_32_128_4"), {mtl::Arg::B(loh), mtl::Arg::B(WU), ++ mtl::Arg::B(xin), mtl::Arg::B(nil), mtl::Arg::B(mh[b]), mtl::Arg::V(gu), mtl::Arg::V(ea)}, ++ RN / 128, (T + 31) / 32, 128)}); ++ ds.push_back({swup, D(P(l, swup), {mtl::Arg::B(mixh), mtl::Arg::B(Wu), mtl::Arg::B(gsh), mtl::Arg::B(hh[b]), ++ mtl::Arg::V(gs)}, (NFF + sn - 1) / sn, (T + sm - 1) / sm, ssg * 32)}); ++ outs.insert(outs.end(), {{&lo[b], &ij[b]}, {&mh[b]}, {&hh[b]}}); ++ } ++ if (cs.empty()) cs.resize(ds.size()); ++ for (size_t i = 0; i < ds.size(); ++i) { ++ cs[i].name = ds[i].first; ++ cs[i].d[v][b] = ds[i].second; ++ if (b == v) cs[i].out[v] = outs[i]; ++ } ++ } ++ for (auto& c : cs) { ++ for (int v = 0; v < 2; ++v) { ++ std::memcpy(R[v].ptr, R0.ptr, R0.bytes); ++ for (const auto* b : c.out[v]) ++ if (b != &R[v]) std::memset(b->ptr, 0xCD, b->bytes); ++ cx.run({c.d[v][v]}); ++ } ++ bool ok = true; ++ for (size_t i = 0; i < c.out[0].size(); ++i) ok = ok && same(*c.out[0][i], *c.out[1][i], c.out[0][i]->bytes); ++ fails += !ok; ++ for (int v = 0; v < 2; ++v) cx.warm({c.d[v][v]}, 0.02); ++ constexpr int NS = 12, REPS = 8; ++ std::vector t[2], ratio; ++ for (int s = 0; s < NS; ++s) { ++ double ms[2]; ++ for (int j = 0; j < 2; ++j) { ++ const int v = (s / 2 + j) % 2; ++ ms[v] = cx.run({c.d[v][(s + v) % 2]}, REPS) * 1e3 / REPS; ++ t[v].push_back(ms[v]); ++ } ++ ratio.push_back(ms[0] / ms[1]); ++ } ++ for (auto* x : {&t[0], &t[1], &ratio}) std::sort(x->begin(), x->end()); ++ const auto med = [](const std::vector& x) { return 0.5 * (x[x.size() / 2 - 1] + x[x.size() / 2]); }; ++ std::printf("T %4u %-24s bitwise %s this %7.4f ms baseline %7.4f ms ratio %.4f [%.4f, %.4f]\n", T, ++ c.name.c_str(), ok ? "PASS" : "FAIL", med(t[0]), med(t[1]), med(ratio), ratio[2], ratio[NS - 3]); ++ } ++ } ++ std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); ++ return fails ? 1 : 0; ++} ++// `comb_fold_parity 0 mtp [bench] [baseline-mtp.metal]`: window.cpp mtp_encode's hidden branch (mtp_rms_xprep, the ++// Q8_0 2560 -> 2560 fc_hidden GEMVs, mtp_combine_gnorm) with the n * HC streams batched as rows [t][c], MAXT per GEMV, ++// against the per-stream chain it replaces (one GEMV of n rows per stream at the window's tile for n rows, M5 Q8_0): ++// R and xn BITWISE wherever both pick one tile (n = 1, n >= 4), and every n BITWISE against n single-token chains. ++// baseline: the per-stream chain on that mtp.metal (its [c][t] h layout) == the per-stream chain on this one. ++// bench: paired GPU times per n, the weights walking 16 copies (from DRAM, as in a decode), median of 9. ++int mtp_parity(bool bench, const char* baseline) { ++ struct RmsXpArgs { uint32_t n, xp_n, plane_tok, split, keep; float eps; }; ++ struct XpArgs { uint32_t n, n_tok, plane_tok, stride; }; ++ struct CombNormArgs { uint32_t n, hc; float eps; }; ++ struct OldCombNormArgs { uint32_t n, hc, h_stride; float eps; }; ++ struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; ++ constexpr int N = 2560, HC = 4, MAXT = 8; ++ constexpr float eps = 1e-6f; ++ mtl::Context cx; ++ const std::string msl = cx.kernel_dir + "/../src/engine/msl/mtp.metal"; ++ const int lg = cx.compile({"quant.metal", "units.metal", "gemv.metal"}); ++ const int lm = cx.compile({msl}, "#define NO_GGML_COMMON 1\n"); ++ const int lb = baseline ? cx.compile({baseline}, "#define NO_GGML_COMMON 1\n") : -1; ++ std::printf("device %s\n", cx.device_name().c_str()); ++ const size_t rb = ref::row_bytes(ref::Q8_0, N), mb = rb * N; ++ const int copies = bench ? 16 : 1; ++ auto W = cx.alloc(mb * copies); ++ ref::random_blocks(ref::Q8_0, W.as(), N, N, 41); ++ for (int k = 1; k < copies; ++k) std::memcpy(W.as() + k * mb, W.ptr, mb); ++ const size_t RN = (size_t) HC * N, plane = 4ull * MAXT * N; ++ auto src = cx.alloc(4 * MAXT * RN), wn = cx.alloc(4 * RN), whc = cx.alloc(4 * RN), e = cx.alloc(4ull * MAXT * N); ++ auto mhn = cx.alloc(4 * MAXT * RN), xp = cx.alloc(4 * plane + 4ull * MAXT * N / 16), h = cx.alloc(4 * MAXT * RN); ++ mtl::Buffer R[3], xn[3]; // per-stream chain, batched, n single-token chains ++ for (int i = 0; i < 3; ++i) { R[i] = cx.alloc(4 * MAXT * RN); xn[i] = cx.alloc(4 * MAXT * RN); } ++ ref::Rng rng(77); ++ for (size_t i = 0; i < MAXT * RN; ++i) src.as()[i] = rng.sym() * 3.0f; ++ for (size_t i = 0; i < RN; ++i) { wn.as()[i] = 1.0f + 0.2f * rng.sym(); whc.as()[i] = 1.0f + 0.2f * rng.sym(); } ++ for (size_t i = 0; i < (size_t) MAXT * N; ++i) e.as()[i] = rng.sym(); ++ mtl::Pipeline prms = cx.pipeline(lm, "mtp_rms_xprep"), pxp = cx.pipeline(lm, "mtp_xprep"), ++ pcomb = cx.pipeline(lm, "mtp_combine_gnorm"); ++ mtl::Pipeline brms, bxp, bcomb; ++ if (lb >= 0) { ++ brms = cx.pipeline(lb, "mtp_rms_xprep"); bxp = cx.pipeline(lb, "mtp_xprep"); bcomb = cx.pipeline(lb, "mtp_combine_gnorm"); ++ } ++ // window.cpp tile() for a wide Q8_0 tensor on the M5 ++ auto tile = [](int rows) { ++ return rows == 1 ? std::array{32, 1, 4, 4} : rows <= 3 ? std::array{32, 2, 4, 1} ++ : std::array{16, 1, 4, 1}; ++ }; ++ std::map gp; ++ auto gemv = [&](int rows, uint32_t y_stride, size_t y_off, int copy) { ++ const auto [L, Rr, S, K] = tile(rows); ++ const std::string fn = "gemv_" + std::to_string(ref::Q8_0) + "_" + std::to_string(L) + "_" + std::to_string(Rr) + ++ "_" + std::to_string(S) + "_" + std::to_string(K); ++ const std::string key = fn + "#" + std::to_string(rows); ++ if (!gp.count(key)) gp.emplace(key, cx.pipeline(lg, fn, mtl::Constants().i(0, rows))); ++ mtl::Dispatch d; ++ d.p = &gp.at(key); ++ const GemvArgs a{N, N, (uint32_t) rows, (uint32_t) rb, N, y_stride}; ++ d.args = {mtl::Arg::B(W, copy * mb), mtl::Arg::B(xp), mtl::Arg::B(xp, plane * 4), mtl::Arg::B(h, y_off), mtl::Arg::V(a)}; ++ const int rpt = (S / K) * (32 / L) * Rr; ++ d.grid[0] = (N + rpt - 1) / rpt; d.tg[0] = S * 32; ++ return d; ++ }; ++ auto rms = [&](const mtl::Pipeline& p, int t0, int n, bool split, bool keep) { ++ mtl::Dispatch d; ++ d.p = &p; ++ const RmsXpArgs a{(uint32_t) RN, N, MAXT, split, keep, eps}; ++ d.args = {mtl::Arg::B(src, t0 * RN * 4), mtl::Arg::B(wn), mtl::Arg::B(mhn), mtl::Arg::B(xp), mtl::Arg::V(a)}; ++ d.grid[0] = n; d.tg[0] = 256; ++ return d; ++ }; ++ auto xprep = [&](const mtl::Pipeline& p, size_t off, int rows, uint32_t stride) { ++ mtl::Dispatch d; ++ d.p = &p; ++ const XpArgs a{N, (uint32_t) rows, MAXT, stride}; ++ d.args = {mtl::Arg::B(mhn, off), mtl::Arg::B(xp), mtl::Arg::V(a)}; ++ d.grid[0] = (N / 16 * rows + 255) / 256; d.tg[0] = 256; ++ return d; ++ }; ++ auto comb = [&](bool old, int t0, int n, int out) { ++ mtl::Dispatch d; ++ d.p = old ? &bcomb : &pcomb; ++ const size_t o = t0 * RN * 4; ++ d.args = {mtl::Arg::B(h), mtl::Arg::B(e, t0 * N * 4), mtl::Arg::B(R[out], o), mtl::Arg::B(whc), mtl::Arg::B(xn[out], o)}; ++ if (old) d.args.push_back(mtl::Arg::V(OldCombNormArgs{N, HC, (uint32_t) (MAXT * N), eps})); ++ else d.args.push_back(mtl::Arg::V(CombNormArgs{N, HC, eps})); ++ d.grid[0] = HC; d.grid[1] = n; d.tg[0] = 256; ++ return d; ++ }; ++ // the replaced chain: a GEMV of n rows per stream; old: the baseline's kernels and [c][t] h ++ auto per_stream = [&](int n, int copy, bool old) { ++ const mtl::Pipeline &r = old ? brms : prms, &x = old ? bxp : pxp; ++ std::vector ds{rms(r, 0, n, n == 1, n > 1)}; ++ if (n == 1) ds.push_back(gemv(HC, old ? MAXT * N : N, 0, copy)); ++ for (int c = 0; c < HC && n > 1; ++c) { ++ if (c > 0) ds.push_back(xprep(x, (size_t) c * N * 4, n, RN)); ++ ds.push_back(old ? gemv(n, N, (size_t) c * MAXT * N * 4, copy) : gemv(n, RN, (size_t) c * N * 4, copy)); ++ } ++ ds.push_back(comb(old, 0, n, 0)); ++ return ds; ++ }; ++ auto batched = [&](int t0, int n, int copy, int out) { ++ const int rows = n * HC; ++ std::vector ds{rms(prms, t0, n, true, rows > MAXT)}; ++ ds.push_back(gemv(std::min(rows, MAXT), N, 0, copy)); ++ for (int r = MAXT; r < rows; r += MAXT) { ++ const int k = std::min(rows - r, MAXT); ++ ds.push_back(xprep(pxp, (size_t) r * N * 4, k, N)); ++ ds.push_back(gemv(k, N, (size_t) r * N * 4, copy)); ++ } ++ ds.push_back(comb(false, t0, n, out)); ++ return ds; ++ }; ++ int fails = 0; ++ for (int n = 1; n <= MAXT; ++n) { ++ const size_t bytes = 4 * n * RN; ++ for (int i = 0; i < 3; ++i) { std::memset(R[i].ptr, 0xCD, R[i].bytes); std::memset(xn[i].ptr, 0xCD, xn[i].bytes); } ++ cx.run_graph(per_stream(n, 0, false)); ++ cx.run_graph(batched(0, n, 0, 1)); ++ for (int t = 0; t < n; ++t) cx.run_graph(batched(t, 1, 0, 2)); ++ const bool one = tile(n == 1 ? HC : n) == tile(std::min(n * HC, MAXT)); ++ const bool vs_old = same(R[0], R[1], bytes) && same(xn[0], xn[1], bytes); ++ const bool vs_one = same(R[1], R[2], bytes) && same(xn[1], xn[2], bytes); ++ double md = 0; ++ for (size_t i = 0; i < n * RN; ++i) md = std::max(md, (double) std::fabs(xn[0].as()[i] - xn[1].as()[i])); ++ const bool ok = vs_one && (!one || vs_old); ++ fails += !ok; ++ std::printf("mtp n %d: batched == per-stream %s (%s tile), == n x (n = 1) %s, max |dxn| %.3g: %s\n", n, ++ vs_old ? "bitwise" : "differs", one ? "same" : "other", vs_one ? "bitwise" : "FAIL", md, ok ? "PASS" : "FAIL"); ++ if (lb >= 0) { ++ std::memset(R[2].ptr, 0xCD, R[2].bytes); std::memset(xn[2].ptr, 0xCD, xn[2].bytes); ++ auto ds = per_stream(n, 0, true); ++ ds.back() = comb(true, 0, n, 2); ++ cx.run_graph(ds); ++ const bool b = same(R[0], R[2], bytes) && same(xn[0], xn[2], bytes); ++ fails += !b; ++ std::printf("mtp n %d: per-stream on the baseline's kernels bitwise %s\n", n, b ? "PASS" : "FAIL"); ++ } ++ } ++ if (bench && !fails) { ++ for (int n : {1, 2, 3, 4, 8}) { ++ std::vector a, b; ++ for (int k = 0; k < copies; ++k) { ++ auto x = per_stream(n, k, false), y = batched(0, n, k, 1); ++ a.insert(a.end(), x.begin(), x.end()); ++ b.insert(b.end(), y.begin(), y.end()); ++ } ++ cx.warm(a, 0.1); ++ cx.warm(b, 0.1); ++ std::vector ta, tb; ++ for (int p = 0; p < 9; ++p) { ++ if (p % 2) { tb.push_back(cx.run_graph(b)); ta.push_back(cx.run_graph(a)); } ++ else { ta.push_back(cx.run_graph(a)); tb.push_back(cx.run_graph(b)); } ++ } ++ std::sort(ta.begin(), ta.end()); std::sort(tb.begin(), tb.end()); ++ const double ma = ta[4] * 1e6 / copies, mb2 = tb[4] * 1e6 / copies; ++ std::printf("mtp n %d cold: per-stream %zu dispatches %.1f us, batched %zu dispatches %.1f us (%+.1f%%)\n", n, ++ a.size() / copies, ma, b.size() / copies, mb2, (mb2 / ma - 1) * 100); ++ } ++ } ++ std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); ++ return fails ? 1 : 0; ++} ++// `comb_fold_parity 0 mtpcat`: window.cpp mtp_encode's gm_cat (k + v 2560 -> 512 + 512, shared gate + up 2560 -> 640 + ++// 640, Q8_0, one token) as gemv_cat / gemv_catk against its two GEMVs beside each other: bitwise, then paired GPU times ++// behind a barrier per pair, hot and with the weights walking copies (cold), median of 15. ++int mtp_cat_parity() { ++ struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; ++ constexpr int NI = 2560; ++ mtl::Context cx; ++ const int lg = cx.compile({"quant.metal", "units.metal", "gemv.metal"}); ++ std::printf("device %s\n", cx.device_name().c_str()); ++ auto x = cx.alloc(4 * NI), xp = cx.alloc(4 * (4 * NI + NI / 16)); ++ ref::Rng rng(9); ++ for (int i = 0; i < NI; ++i) x.as()[i] = rng.sym(); ++ { ++ auto prep = cx.pipeline(lg, "xprep"); ++ mtl::Dispatch d; ++ d.p = &prep; ++ d.args = {mtl::Arg::B(x), mtl::Arg::B(xp), mtl::Arg::V((uint32_t) NI), mtl::Arg::V(1u)}; ++ d.grid[0] = (NI / 16 + 255) / 256; d.tg[0] = 256; ++ cx.run({d}); ++ } ++ int fails = 0; ++ for (int no : {512, 640}) ++ for (const auto t : {std::array{32, 1, 4, 4}, std::array{32, 1, 4, 1}}) { ++ const auto [L, Rr, S, K] = t; ++ const size_t rb = ref::row_bytes(ref::Q8_0, NI), mb = rb * no; ++ const int copies = 32; ++ auto wa = cx.alloc(mb * copies), wb = cx.alloc(mb * copies); ++ ref::random_blocks(ref::Q8_0, wa.as(), no, NI, 3); ++ ref::random_blocks(ref::Q8_0, wb.as(), no, NI, 4); ++ for (int k = 1; k < copies; ++k) { ++ std::memcpy(wa.as() + k * mb, wa.ptr, mb); ++ std::memcpy(wb.as() + k * mb, wb.ptr, mb); ++ } ++ auto ya = cx.alloc(4 * no), yb = cx.alloc(4 * no), za = cx.alloc(4 * no), zb = cx.alloc(4 * no); ++ const std::string tl = std::to_string(L) + "_" + std::to_string(Rr) + "_" + std::to_string(S); ++ const std::string gn = "gemv_" + std::to_string(ref::Q8_0) + "_" + tl + "_" + std::to_string(K), ++ cn = K == 1 ? "gemvc_" + std::to_string(ref::Q8_0) + "_" + tl ++ : "gemvck_" + std::to_string(ref::Q8_0) + "_" + tl + "_" + std::to_string(K); ++ auto pg = cx.pipeline(lg, gn, mtl::Constants().i(0, 1)), pc = cx.pipeline(lg, cn, mtl::Constants().i(0, 1)); ++ const int rpt = (S / K) * (32 / L) * Rr; ++ const GemvArgs a{NI, (uint32_t) no, 1, (uint32_t) rb, NI, (uint32_t) no}; ++ GemvArgs a2 = a; ++ a2.n_out = 2 * no; ++ auto gemv = [&](const mtl::Buffer& w, int k, const mtl::Buffer& y, bool barrier) { ++ mtl::Dispatch d; ++ d.p = &pg; ++ d.args = {mtl::Arg::B(w, k * mb), mtl::Arg::B(xp), mtl::Arg::B(xp, 4ull * 4 * NI), mtl::Arg::B(y), mtl::Arg::V(a)}; ++ d.grid[0] = (no + rpt - 1) / rpt; d.tg[0] = S * 32; ++ d.barrier = barrier; ++ return d; ++ }; ++ auto cat = [&](int k) { ++ mtl::Dispatch d; ++ d.p = &pc; ++ d.args = {mtl::Arg::B(wa, k * mb), mtl::Arg::B(xp), mtl::Arg::B(xp, 4ull * 4 * NI), mtl::Arg::B(yb), ++ mtl::Arg::V(a2), mtl::Arg::B(wb, k * mb), mtl::Arg::B(zb), mtl::Arg::V((uint32_t) no)}; ++ d.grid[0] = (2 * no + rpt - 1) / rpt; d.tg[0] = S * 32; ++ return d; ++ }; ++ cx.run_graph({gemv(wa, 0, ya, true), gemv(wb, 0, za, false)}); ++ cx.run_graph({cat(0)}); ++ const bool ok = same(ya, yb, 4 * no) && same(za, zb, 4 * no); ++ fails += !ok; ++ std::printf("gm_cat Q8_0 2560 -> %d+%d tile %d,%d,%d,%d: bitwise %s", no, no, L, Rr, S, K, ok ? "PASS" : "FAIL"); ++ for (bool cold : {false, true}) { ++ std::vector two, one; ++ for (int r = 0; r < 256 / (cold ? copies : 1); ++r) ++ for (int k = 0; k < (cold ? copies : 1); ++k) { ++ two.push_back(gemv(wa, k, ya, true)); ++ two.push_back(gemv(wb, k, za, false)); ++ one.push_back(cat(k)); ++ } ++ cx.warm(two, 0.05); ++ cx.warm(one, 0.05); ++ std::vector tt, to; ++ for (int p = 0; p < 15; ++p) { ++ if (p % 2) { to.push_back(cx.run_graph(one)); tt.push_back(cx.run_graph(two)); } ++ else { tt.push_back(cx.run_graph(two)); to.push_back(cx.run_graph(one)); } ++ } ++ std::sort(tt.begin(), tt.end()); std::sort(to.begin(), to.end()); ++ const double m2 = tt[7] * 1e6 / one.size(), m1 = to[7] * 1e6 / one.size(); ++ std::printf(" %s: 2 gemv %.2f us, cat %.2f us (%+.1f%%)", cold ? "cold" : "hot", m2, m1, (m1 / m2 - 1) * 100); ++ } ++ std::printf("\n"); ++ } ++ std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); ++ return fails ? 1 : 0; ++} ++} // namespace ++ ++int main(int argc, char** argv) { ++ setvbuf(stdout, nullptr, _IONBF, 0); ++ if (argc > 2 && std::strcmp(argv[2], "mtpcat") == 0) return mtp_cat_parity(); ++ if (argc > 2 && std::strcmp(argv[2], "mtp") == 0) { ++ const bool bench = argc > 3 && std::strcmp(argv[3], "bench") == 0; ++ return mtp_parity(bench, argc > 3 + bench ? argv[3 + bench] : nullptr); ++ } ++ if (argc > 2 && std::strcmp(argv[2], "epi") == 0) ++ return epi_parity(uint32_t(std::atoi(argv[1])), argc > 3 && std::strcmp(argv[3], "bench") == 0, argc > 4 ? argv[4] : nullptr); ++ if (argc > 3 && std::strcmp(argv[2], "topk") == 0) return topk_parity(uint32_t(std::atoi(argv[1])), argv[3]); ++ if (argc > 2 && std::strcmp(argv[2], "gu") == 0) ++ return gu_parity(uint32_t(std::atoi(argv[1])), argc > 3 && std::strcmp(argv[3], "bench") == 0, argc > 4 ? argv[4] : nullptr); ++ if (argc > 2 && std::strcmp(argv[2], "hcup") == 0) ++ return hc_up_parity(uint32_t(std::atoi(argv[1])), argc > 3 && std::strcmp(argv[3], "bench") == 0, ++ argc > 3 && std::strcmp(argv[3], "confirm") == 0); ++ if (argc > 3 && std::strcmp(argv[2], "ab") == 0) ++ return ab_parity(uint32_t(std::atoi(argv[1])), argv[3], argc > 4 ? argv[4] : nullptr); ++ if (argc > 3 && std::strcmp(argv[2], "denseab") == 0) ++ return dense_ab_parity(uint32_t(std::atoi(argv[1])), argv[3], argc > 4 ? argv[4] : nullptr); ++ const bool gr_only = argc > 2 && std::strcmp(argv[2], "grvec") == 0; ++ mtl::Context cx; ++ const std::string source = gr_only ? mtl::read_file(cx.kernel_dir + "/prompt.metal") : ""; ++ const std::string prelude = "#define NO_GGML_COMMON 1\nusing namespace metal;\n" ++ "struct GrArgs { uint n_embd, hc, hc_lr, n_tok, apply; float eps; };\n"; ++ const auto begin = source.find("// ---- the HC read"); ++ const auto end = source.find("// gemm_hc_up:"); ++ if (gr_only && (begin == std::string::npos || end == std::string::npos || end <= begin)) { ++ std::fprintf(stderr, "GR source markers missing\n"); ++ return 1; ++ } ++ const int lib = gr_only ? cx.compile({}, prelude + source.substr(begin, end - begin)) : ++ cx.compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "gr.metal", "prompt.metal"}); ++ const std::string mm = gr_only ? mtl::read_file(cx.kernel_dir + "/moe_mm.metal") : ""; ++ const auto mb = mm.find("kernel void mm_combine("); ++ const auto me = mm.find("\n}", mb == std::string::npos ? 0 : mb); ++ if (gr_only && (mb == std::string::npos || me == std::string::npos)) { ++ std::fprintf(stderr, "mm_combine source marker missing\n"); ++ return 1; ++ } ++ const int lmm = gr_only ? cx.compile({}, prelude + mm.substr(mb, me + 2 - mb)) : ++ cx.compile({"quant.metal", "units.metal", "mma.metal", "moe_mm.metal"}); ++ std::printf("device %s\n", cx.device_name().c_str()); ++ const uint32_t N = 2560, HC = 4, LR = 320, TOPK = 10, NEXP = 512, TMR = 16; ++ const size_t RN = (size_t) HC * N; ++ const uint32_t TB = argc > 1 ? (uint32_t) std::atoi(argv[1]) : 1024; ++ const bool bench = argc > 2 && std::strcmp(argv[2], "bench") == 0; ++ mtl::Pipeline pcomb = cx.pipeline(lmm, "mm_combine"), pnhm = cx.pipeline(lib, "gr_norm_h_moe"), ++ pnhc = cx.pipeline(lib, "gr_norm_h_moec"), pam = cx.pipeline(lib, "gr_apply_moe"), ++ pamc = cx.pipeline(lib, "gr_apply_moec"), pnh = cx.pipeline(lib, "gr_norm_h"), ++ pnh4 = cx.pipeline(lib, "gr_norm_h4"), pup = gr_only ? mtl::Pipeline{} : cx.pipeline(lib, "gemm_hc_up_64_64_4"); ++ std::vector baseline; ++ if (gr_only && argc > 3) { ++ const std::string ref = mtl::read_file(argv[3]); ++ const auto rb = ref.find("// ---- the HC read"), re = ref.find("// gemm_hc_up:"); ++ if (rb == std::string::npos || re == std::string::npos || re <= rb) return 1; ++ const int rl = cx.compile({}, prelude + ref.substr(rb, re - rb)); ++ for (const char* name : {"gr_norm_h_moec", "gr_apply_moec", "gr_norm_h4", "gr_norm_h4"}) ++ baseline.push_back(cx.pipeline(rl, name)); ++ } ++ auto D = [](const mtl::Pipeline& p, std::vector a, uint32_t gx, uint32_t gy, uint32_t th = 256) { ++ mtl::Dispatch d; ++ d.p = &p; ++ d.args = std::move(a); ++ d.grid[0] = gx; ++ d.grid[1] = gy; ++ d.tg[0] = th; ++ return d; ++ }; ++ int fails = 0; ++ for (uint32_t T : gr_only ? std::vector{TB} : std::vector{1u, 7u, 64u, TB}) { ++ const size_t mr = (size_t) T * TOPK + NEXP * TMR; ++ mtl::Buffer R0 = cx.alloc(4 * T * RN), Ra = cx.alloc(4 * T * RN), Rb = cx.alloc(4 * T * RN), YE = cx.alloc(4 * mr * N), ++ pos = cx.alloc(4 * T * TOPK), sw = cx.alloc(4 * mr), yr = cx.alloc(4 * T * N), ys = cx.alloc(4 * T * N), ++ sgv = cx.alloc(4 * T), inj = cx.alloc(4 * T * HC), wn = cx.alloc(4 * RN), xa = cx.alloc(2 * T * RN), ++ xb = cx.alloc(2 * T * RN), loh = cx.alloc(2 * T * LR), WU = cx.alloc(2 * RN * LR), ++ mixed = cx.alloc(4 * T * N), mha = cx.alloc(2 * T * N), mhb = cx.alloc(2 * T * N), nil{}; ++ ref::Rng r(23 + T); ++ auto fill = [&](const mtl::Buffer& b, float s, float o) { ++ for (size_t i = 0; i < b.bytes / 4; ++i) b.as()[i] = o + s * r.sym(); ++ }; ++ auto fillh = [&](const mtl::Buffer& b, float s) { ++ for (size_t i = 0; i < b.bytes / 2; ++i) b.as()[i] = ref::f32_to_f16(s * r.sym()); ++ }; ++ fill(R0, 1.0f, 0.0f); fill(YE, 0.5f, 0.0f); fill(sw, 0.1f, 0.1f); fill(ys, 0.5f, 0.0f); fill(sgv, 3.0f, 0.0f); ++ fill(inj, 4.0f, 0.0f); fill(wn, 0.2f, 1.0f); fillh(loh, 1.0f); fillh(WU, 0.1f); fillh(xa, 1.0f); ++ std::memcpy(xb.ptr, xa.ptr, xb.bytes); ++ // mm_place's slots: distinct rows of YE (expert-sorted, so scattered), a few dropped (-1) ++ std::vector perm(mr); ++ std::iota(perm.begin(), perm.end(), 0); ++ std::shuffle(perm.begin(), perm.end(), std::mt19937(T)); ++ for (size_t i = 0; i < (size_t) T * TOPK; ++i) pos.as()[i] = (i % 37 == 5) ? -1 : perm[i]; ++ const GrArgs a{N, HC, LR, T, 1u, 1e-6f}, a0{N, HC, LR, T, 0u, 1e-6f}; ++ const uint32_t ca[4] = {T, N, TOPK, 0}; ++ const mtl::Dispatch comb = D(pcomb, {mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(yr), mtl::Arg::V(ca)}, ++ (T * N + 255) / 256, 1); ++ auto nhm = [&](const mtl::Buffer& R, const mtl::Buffer& x) { ++ return D(pnhm, {mtl::Arg::B(R), mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::B(wn), ++ mtl::Arg::B(x), mtl::Arg::V(a)}, 1, T); ++ }; ++ auto nhc = [&](const mtl::Buffer& R, const mtl::Buffer& x) { ++ return D(pnhc, {mtl::Arg::B(R), mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(ys), mtl::Arg::B(sgv), ++ mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(x), mtl::Arg::V(a), mtl::Arg::V(TOPK)}, 1, T); ++ }; ++ auto am = [&](const mtl::Buffer& R) { ++ return D(pam, {mtl::Arg::B(R), mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::V(a)}, ++ RN / 256, T); ++ }; ++ auto amc = [&](const mtl::Buffer& R) { ++ return D(pamc, {mtl::Arg::B(R), mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(ys), mtl::Arg::B(sgv), ++ mtl::Arg::B(inj), mtl::Arg::V(a), mtl::Arg::V(TOPK)}, (N + 255) / 256, T); ++ }; ++ auto nh = [&](const mtl::Pipeline& p, const mtl::Buffer& R, const mtl::Buffer& x, const GrArgs& g, uint32_t gx) { ++ return D(p, {mtl::Arg::B(R), mtl::Arg::B(ys), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(x), mtl::Arg::V(g)}, gx, T); ++ }; ++ const GemmArgs gg{(int) T, (int) RN, (int) LR, 0}; ++ const EpiArgs ea{T, N, HC, LR}; ++ auto up = [&](const mtl::Buffer& m, const mtl::Buffer& mh) { ++ return D(pup, {mtl::Arg::B(loh), mtl::Arg::B(WU), mtl::Arg::B(xa), mtl::Arg::B(m), mtl::Arg::B(mh), mtl::Arg::V(gg), ++ mtl::Arg::V(ea)}, RN / 64, (T + 63) / 64, 128); ++ }; ++ auto reset = [&] { ++ std::memcpy(Ra.ptr, R0.ptr, Ra.bytes); ++ std::memcpy(Rb.ptr, R0.ptr, Rb.bytes); ++ std::memset(yr.ptr, 0, yr.bytes); ++ }; ++ auto check = [&](const char* what, bool ok) { ++ fails += !ok; ++ std::printf("%-26s T %4u: bitwise %s\n", what, T, ok ? "PASS" : "FAIL"); ++ }; ++ reset(); ++ cx.run({comb, nhm(Ra, xa)}); ++ cx.run({nhc(Rb, xb)}); ++ check("gr_norm_h_moec", same(Ra, Rb, Ra.bytes) && same(xa, xb, xa.bytes)); ++ reset(); ++ cx.run({comb, am(Ra)}); ++ cx.run({amc(Rb)}); ++ check("gr_apply_moec", same(Ra, Rb, Ra.bytes)); ++ for (const GrArgs* g : {&a0, &a}) { ++ reset(); ++ cx.run({nh(pnh, Ra, xa, *g, HC)}); ++ cx.run({nh(pnh4, Rb, xb, *g, 1)}); ++ check(g->apply ? "gr_norm_h4 (apply 1)" : "gr_norm_h4 (apply 0)", same(Ra, Rb, Ra.bytes) && same(xa, xb, xa.bytes)); ++ } ++ if (gr_only) { ++ const std::vector ds = { ++ nhc(Rb, xb), amc(Rb), nh(pnh4, Rb, xb, a, 1), nh(pnh4, Rb, xb, a0, 1)}; ++ for (size_t i = 0; i < ds.size(); ++i) { ++ const auto& d = ds[i]; ++ auto ref = d; ++ if (!baseline.empty()) ref.p = &baseline[i]; ++ cx.run({d}); ++ if (!baseline.empty()) cx.run({ref}); ++ std::printf("grvec %s apply %u ms4", d.p->name.c_str(), i == 3 ? 0u : 1u); ++ std::vector ct, bt, ratios; ++ for (int sample = 0; sample < 5; ++sample) { ++ double ms, rms = 0.0; ++ if (!baseline.empty() && sample % 2 == 0) rms = cx.run({ref}, 4) * 1e3; ++ ms = cx.run({d}, 4) * 1e3; ++ if (!baseline.empty() && sample % 2 != 0) rms = cx.run({ref}, 4) * 1e3; ++ std::printf(" %.3f", ms); ++ if (!baseline.empty()) { ++ std::printf("/%.3f", rms); ++ ct.push_back(ms); bt.push_back(rms); ratios.push_back(ms / rms); ++ } ++ } ++ if (!baseline.empty()) { ++ std::sort(ct.begin(), ct.end()); std::sort(bt.begin(), bt.end()); std::sort(ratios.begin(), ratios.end()); ++ std::printf(" med %.3f/%.3f ms ratio %.4f", ct[2], bt[2], ratios[2]); ++ } ++ std::printf("\n"); ++ } ++ continue; ++ } ++ cx.run({up(mixed, mha)}); ++ cx.run({up(nil, mhb)}); ++ check("gemm_hc_up, mixed unbound", same(mha, mhb, mha.bytes)); ++ if (T != TB || !bench) continue; ++ // ---- GPU times: 48 reps in one command buffer (a chunk's layers), best of 3 ++ auto time = [&](const std::vector& ds) { ++ cx.warm(ds, 0.03); ++ double best = 1e9; ++ for (int i = 0; i < 3; ++i) best = std::min(best, cx.run(ds, 48)); ++ return best * 1e3; ++ }; ++ const double t_comb = time({comb}), t_nhm = time({nhm(Ra, xa)}), t_pair = time({comb, nhm(Ra, xa)}), ++ t_nhc = time({nhc(Rb, xb)}), t_am = time({comb, am(Ra)}), t_amc = time({amc(Rb)}), ++ t_nh = time({nh(pnh, Ra, xa, a, HC)}), t_nh4 = time({nh(pnh4, Rb, xb, a, 1)}), ++ t_nh0 = time({nh(pnh, Ra, xa, a0, HC)}), t_nh40 = time({nh(pnh4, Rb, xb, a0, 1)}), ++ t_up = time({up(mixed, mha)}), t_upn = time({up(nil, mhb)}); ++ std::printf("T %u, ms per 48 reps (best of 3):\n", T); ++ std::printf(" mm_combine %.2f, gr_norm_h_moe %.2f, mm_combine + gr_norm_h_moe %.2f -> gr_norm_h_moec %.2f (%+.1f%%)\n", ++ t_comb, t_nhm, t_pair, t_nhc, 100.0 * (t_nhc / t_pair - 1)); ++ std::printf(" mm_combine + gr_apply_moe %.2f -> gr_apply_moec %.2f (%+.1f%%)\n", t_am, t_amc, 100.0 * (t_amc / t_am - 1)); ++ std::printf(" gr_norm_h apply 1 (hc, T) %.2f -> gr_norm_h4 (1, T) %.2f (%+.1f%%)\n", t_nh, t_nh4, 100.0 * (t_nh4 / t_nh - 1)); ++ std::printf(" gr_norm_h apply 0 (hc, T) %.2f -> gr_norm_h4 (1, T) %.2f (%+.1f%%)\n", t_nh0, t_nh40, 100.0 * (t_nh40 / t_nh0 - 1)); ++ std::printf(" gemm_hc_up with mixed %.2f -> unbound %.2f (%+.1f%%)\n", t_up, t_upn, 100.0 * (t_upn / t_up - 1)); ++ } ++ std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); ++ return fails ? 1 : 0; ++} +diff --git a/metal/tests/dense_mm_parity.cpp b/metal/tests/dense_mm_parity.cpp +index 33ce4500..969cb790 100644 +--- a/metal/tests/dense_mm_parity.cpp ++++ b/metal/tests/dense_mm_parity.cpp +@@ -3,13 +3,22 @@ + // awkward row counts; then the speed on the packs' shapes for a 2048-token chunk. + // dense_mm_parity parity + // dense_mm_parity --bench speed ++// dense_mm_parity --tiles M [shape] every gemm_f16 tile on the prompt chunk's 13 (N, K) shapes at M rows: ++// memcmp vs 128_64_4 and the fastest of interleaved rounds; [shape] runs one (its name, ++// as printed); DENSE_TILES=a,b,... only those tiles, DENSE_ROUNDS=n rounds ++// dense_mm_parity --ab M [shape] the same for the tile prompt_chunk picked before round 35 vs now ++// dense_mm_parity --dq-ab [type] [shape] dq_f16 at 64 threads vs dq_tg's; with DQ_AB_KERNELS=, that dir's ++// dq_f16 vs this build's at dq_tg's (memcmp + paired timing) + #include "mtl_rt.h" + #include "ref_quant.h" + + #include ++#include + #include + #include ++#include + #include ++#include + #include + #include + +@@ -24,13 +33,226 @@ int unit_of(int t) { + default: return 32; + } + } ++int dq_tg(int u) { return u <= 64 ? 64 : u <= 128 ? 128 : u % 192 == 0 ? 192 : u % 160 == 0 ? 160 : 256; } ++ ++using Tile = std::array; // TM, TN, SG ++ ++// Window::prompt_chunk's gemm tile before round 11 ++Tile old_tile(int M, int Nn, int K) { ++ int tm = (M >= 128 && Nn > 64) ? 128 : 64, tn = Nn <= 64 ? 32 : 64, sg = 4; ++ if (Nn > 64) { ++ if (M <= 128 && (Nn <= 640 || K > 2560)) { tm = 64; tn = Nn <= 128 ? 32 : 64; } ++ else if (M <= 512 && Nn <= 324) { tm = 64; tn = 64; } ++ else if (M >= 2048 && Nn >= 2560 && K >= 2560) { tm = 64; tn = 128; } ++ } ++ return {tm, tn, sg}; ++} ++ ++// ... before round 35 ++Tile r34_tile(int M, int Nn, int K) { ++ if (M <= 128) return old_tile(M, Nn, K); ++ if (Nn <= 64) return M < 1024 ? Tile{16, 64, 2} : Tile{64, 32, 2}; ++ if (Nn == 96) return M < 768 ? Tile{16, 64, 2} : Tile{64, 32, 2}; ++ if (Nn <= 128) return M < 768 ? Tile{16, 64, 2} : Tile{32, 64, 2}; ++ if (Nn <= 640 && K > 2560) return M <= 512 ? Tile{64, 64, 4} : M < 1536 ? Tile{128, 64, 4} : Tile{32, 128, 2}; ++ if (Nn <= 512 || (Nn <= 640 && M < 384)) return M < 384 ? Tile{32, 64, 2} : Tile{32, 128, 2}; ++ if (Nn <= 1664) return M < 768 ? Tile{32, 128, 2} : Tile{64, 128, 4}; ++ return M < 1536 ? Tile{128, 64, 4} : Tile{64, 128, 4}; ++} ++ ++// ... and now (window.cpp, kept in step by hand) ++Tile new_tile(int M, int Nn, int K) { ++ if (M <= 128) return old_tile(M, Nn, K); ++ if (Nn == 96) return M < 640 ? Tile{16, 64, 2} : Tile{16, 96, 2}; ++ if (Nn > 1664 && M < 1536) return M % 128 && M % 128 <= 64 ? Tile{64, 128, 4} : Tile{128, 64, 4}; ++ return r34_tile(M, Nn, K); ++} ++ ++std::string tile_name(const Tile& t) { ++ return std::to_string(t[0]) + "_" + std::to_string(t[1]) + "_" + std::to_string(t[2]); ++} ++ ++void fill_h(const mtl::Buffer& b, size_t n, uint64_t seed, float scale) { ++ ref::Rng r(seed); ++ for (size_t i = 0; i < n; ++i) b.as()[i] = ref::f32_to_f16(r.sym() * scale); ++} ++ ++// --tiles: every tile; --ab: r34_tile vs new_tile. Each candidate's Y is compared with 128_64_4's (memcmp), then ++// timed over 5 (--ab: 9) rounds that take the candidates in turn ++int tile_bench(mtl::Context& cx, int lib, int M, bool ab, const char* only) { ++ struct Sh { const char* name; int N, K, calls; }; // calls: a 48-layer chunk's (streaming: the read-ahead router) ++ const std::vector shapes = {{"hc_read", 324, 10240, 96}, {"Wqkv", 10240, 2560, 36}, {"Wz", 6144, 2560, 36}, ++ {"Wout", 2560, 6144, 48}, {"Wq", 12288, 2560, 12}, {"Wpv", 2560, 2560, 1}, ++ {"Wshgu", 640, 2560, 96}, {"Wshd", 2560, 640, 48}, {"W512", 512, 2560, 95}, ++ {"Wk", 1152, 2560, 12}, {"Wk_sel", 1664, 2560, 12}, {"Walbe", 96, 2560, 36}, ++ {"Wsg", 1, 2560, 48}}; ++ const std::vector all = {{128, 64, 4}, {64, 64, 4}, {64, 128, 4}, {128, 128, 4}, {64, 32, 4}, {64, 64, 2}, ++ {64, 128, 2}, {16, 64, 2}, {32, 64, 2}, {32, 128, 2}, {64, 32, 2}, {16, 96, 2}}; ++ mtl::Buffer A = cx.alloc(2 * (size_t)M * 10240), B = cx.alloc(2 * (size_t)12288 * 2560), ++ Y = cx.alloc(4 * (size_t)M * 12288), Yref = cx.alloc(4 * (size_t)M * 12288); ++ fill_h(A, (size_t)M * 10240, 11, 1.0f); ++ fill_h(B, (size_t)12288 * 2560, 12, 0.05f); ++ std::map pl; ++ auto disp = [&](const Tile& tl, const Sh& s, const mtl::Buffer& y) { ++ const std::string fn = "gemm_f16_" + tile_name(tl); ++ auto it = pl.find(fn); ++ if (it == pl.end()) it = pl.emplace(fn, cx.pipeline(lib, fn)).first; ++ GemmArgs g{M, s.N, s.K, 0}; ++ mtl::Dispatch d; d.p = &it->second; ++ d.args = {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(y), mtl::Arg::V(g)}; ++ d.grid[0] = (s.N + tl[1] - 1) / tl[1]; d.grid[1] = (M + tl[0] - 1) / tl[0]; d.tg[0] = 32 * tl[2]; ++ return d; ++ }; ++ int fails = 0; ++ double tot_old = 0, tot_new = 0; ++ std::printf("M %d: shape N x K calls | tile us/call (DIFF: not bitwise 128_64_4) | old -> new tile\n", M); ++ for (const Sh& s : shapes) { ++ if (only && std::string(only) != s.name) continue; ++ const Tile to = r34_tile(M, s.N, s.K), tn = new_tile(M, s.N, s.K); ++ std::vector tiles = ab ? std::vector{to, tn} : all; ++ if (const char* only_t = std::getenv("DENSE_TILES")) { // a comma list of tile names, e.g. 64_32_2,16_96_2 ++ const std::string want = std::string(",") + only_t + ","; ++ std::vector keep; ++ for (const Tile& tl : tiles) ++ if (want.find("," + tile_name(tl) + ",") != std::string::npos) keep.push_back(tl); ++ tiles = keep; ++ } ++ const size_t ny = (size_t)M * s.N * 4; ++ std::memset(Yref.ptr, 0xFF, ny); ++ cx.run({disp({128, 64, 4}, s, Yref)}); ++ const int reps = std::clamp((int)(1e9 / (2.0 * M * s.N * s.K)), 1, 32); ++ cx.warm({disp({128, 64, 4}, s, Y)}, &s == &shapes[0] || only ? 0.3 : 0.05); // the first: the clocks up ++ std::vector cmp(tiles.size()); ++ for (size_t i = 0; i < tiles.size(); ++i) { ++ std::memset(Y.ptr, 0xFF, ny); ++ cx.run({disp(tiles[i], s, Y)}); ++ cmp[i] = std::memcmp(Y.ptr, Yref.ptr, ny); ++ fails += cmp[i] != 0; ++ } ++ std::vector> ts(tiles.size()); ++ // the minimum over the rounds: a busy GPU (another process) only adds time ++ const int rounds = std::getenv("DENSE_ROUNDS") ? std::atoi(std::getenv("DENSE_ROUNDS")) : ab ? 9 : 5; ++ for (int round = 0; round < rounds; ++round) ++ for (size_t i = 0; i < tiles.size(); ++i) { ++ std::vector ds(reps, disp(tiles[i], s, Y)); ++ ts[i].push_back(cx.run(ds) / reps * 1e6); ++ } ++ double uo = 0, un = 0; ++ std::printf("%-8s %5d x %5d x%3d |", s.name, s.N, s.K, s.calls); ++ for (size_t i = 0; i < tiles.size(); ++i) { ++ const double us = *std::min_element(ts[i].begin(), ts[i].end()); ++ if (tiles[i] == to) uo = us; ++ if (tiles[i] == tn) un = us; ++ std::printf(" %s %.1f%s", tile_name(tiles[i]).c_str(), us, cmp[i] ? " DIFF" : ""); ++ } ++ std::printf(" | %s -> %s %.2fx, %.2f -> %.2f ms/chunk\n", tile_name(to).c_str(), tile_name(tn).c_str(), uo / un, ++ uo * s.calls / 1e3, un * s.calls / 1e3); ++ tot_old += uo * s.calls; ++ tot_new += un * s.calls; ++ } ++ std::printf("M %d: old %.2f ms, new %.2f ms per chunk (%.2fx); %s\n", M, tot_old / 1e3, tot_new / 1e3, ++ tot_old / tot_new, fails ? "SOME TILES DIFFER" : "ALL BITWISE"); ++ return fails ? 1 : 0; ++} ++ ++int dq_ab_bench(mtl::Context& cx, int lib, int lib_a, const char* only_ty, const char* only_sh) { ++ struct DqSh { const char* name; int n_in, n_out; uint32_t perm; int calls; }; ++ const std::vector shapes = { ++ {"2560x12288", 2560, 12288, 0, 12}, {"2560x10240", 2560, 10240, 0, 37}, {"320x10240p4", 320, 10240, 4, 96}, ++ {"2560x6144", 2560, 6144, 0, 36}, {"6144x2560", 6144, 2560, 0, 48}, {"2560x640", 2560, 640, 0, 96}, ++ {"640x2560", 640, 2560, 0, 48}, {"2560x512", 2560, 512, 0, 119}, {"2560x128", 2560, 128, 0, 12}, ++ {"2560x48", 2560, 48, 0, 72}, {"10240x320", 10240, 320, 0, 96}, {"10240x4", 10240, 4, 0, 96}, ++ {"2560x1", 2560, 1, 0, 48}, ++ }; ++ const std::vector types = {ref::F32, ref::BF16, ref::Q8_0, ref::Q4_0, ref::Q5_0, ref::Q5_1, ref::IQ4_NL, ++ ref::IQ4_XS, ref::Q2_0, ref::IQ3_S, ref::Q3_K, ref::Q4_K, ref::Q5_K, ref::Q6_K}; ++ const size_t pool_w = (size_t)256 << 20, max_wh = (size_t)2 * 2560 * 12288; ++ mtl::Buffer W = cx.alloc(pool_w), wh = cx.alloc(max_wh), wh_ref = cx.alloc(max_wh); ++ int fails = 0; ++ for (int ty : types) { ++ if (only_ty && std::string(only_ty) != ref::type_name(ty)) continue; ++ mtl::Pipeline pdq = cx.pipeline(lib, "dq_f16_" + std::to_string(ty)); ++ mtl::Pipeline pdq_a = cx.pipeline(lib_a, "dq_f16_" + std::to_string(ty)); ++ const int unit = unit_of(ty); ++ double tot64 = 0, tot_sel = 0; ++ for (const auto& s : shapes) { ++ if (only_sh && std::string(only_sh) != s.name) continue; ++ if (s.n_in % 256 && ref::block_elems(ty) == 256) continue; ++ const size_t rb = ref::row_bytes(ty, s.n_in), mat_w = rb * (size_t)s.n_out; ++ const size_t stride_w = (mat_w + 16383) / 16384 * 16384, nbytes = (size_t)2 * s.n_in * s.n_out; ++ const int nslots = std::max(1, std::min(16, (int)(pool_w / stride_w))), rinit = std::min(s.n_out, 64); ++ ref::random_blocks(ty, W.as(), rinit, s.n_in, 17 + ty + s.n_in, 0.02f); ++ for (int row = rinit; row < s.n_out; row += rinit) ++ std::memcpy(W.as() + (size_t)row * rb, W.as(), (size_t)std::min(rinit, s.n_out - row) * rb); ++ for (int sl = 1; sl < nslots; ++sl) std::memcpy(W.as() + (size_t)sl * stride_w, W.as(), mat_w); ++ const int u = s.n_in / unit, tgs[2] = {lib_a != lib ? dq_tg(u) : 64, dq_tg(u)}; ++ DqArgs a{(uint32_t)s.n_in, (uint32_t)s.n_out, (uint32_t)rb, s.perm}; ++ auto disp = [&](int i, size_t woff, const mtl::Buffer& out) { ++ const int tg = tgs[i]; ++ mtl::Dispatch d; d.p = i ? &pdq : &pdq_a; d.args = {mtl::Arg::B(W, woff), mtl::Arg::B(out), mtl::Arg::V(a)}; ++ d.grid[0] = (u + tg - 1) / tg; d.grid[1] = s.n_out; d.tg[0] = tg; d.barrier = true; return d; ++ }; ++ std::memset(wh_ref.ptr, 0xAA, nbytes); ++ std::memset(wh.ptr, 0x55, nbytes); ++ cx.run({disp(0, 0, wh_ref)}); ++ cx.run({disp(1, 0, wh)}); ++ const int cmp = std::memcmp(wh.ptr, wh_ref.ptr, nbytes); ++ fails += (cmp != 0); ++ const int reps = std::clamp((int)(4e7 / ((double)s.n_in * s.n_out)), 4, 16); ++ cx.warm({disp(0, 0, wh), disp(1, 0, wh)}, (&s == &shapes[0] || only_sh) ? 0.3 : 0.02); ++ std::vector ts[2]; ++ double best_pair = 1e30; ++ for (int round = 0; round < 13; ++round) { ++ double v[2]; ++ for (int step = 0; step < 2; ++step) { ++ const int i = (round & 1) ? (1 - step) : step; ++ std::vector ds; ++ for (int r = 0; r < reps; ++r) ds.push_back(disp(i, (size_t)(r % nslots) * stride_w, wh)); ++ v[i] = cx.run_graph(ds) / reps * 1e6; ++ } ++ best_pair = std::min(best_pair, std::min(v[0], v[1])); ++ ts[0].push_back(v[0]); ++ ts[1].push_back(v[1]); ++ } ++ std::vector f0, f1; ++ for (size_t i = 0; i < ts[0].size(); ++i) { ++ if (std::max(ts[0][i], ts[1][i]) <= 1.15 * best_pair) { ++ f0.push_back(ts[0][i]); ++ f1.push_back(ts[1][i]); ++ } ++ } ++ if (!f0.empty()) { ts[0] = f0; ts[1] = f1; } ++ const double u0 = median(ts[0]), u1 = median(ts[1]); ++ tot64 += u0 * s.calls / 1e3; tot_sel += u1 * s.calls / 1e3; ++ std::printf("%-7s %-13s x%3d (u=%3d, tg %3d->%3d): %6.1f -> %6.1f us (%.2fx) memcmp=%d\n", ++ ref::type_name(ty), s.name, s.calls, u, tgs[0], tgs[1], u0, u1, u0 / u1, cmp); ++ } ++ std::printf("CHUNK %-7s: %.2f -> %.2f ms (%.2fx)\n", ref::type_name(ty), tot64, tot_sel, tot64 / tot_sel); ++ } ++ return fails ? 1 : 0; ++} ++ + } // namespace + + int main(int argc, char** argv) { ++ if (argc > 1 && std::string(argv[1]) == "--dq-ab") { ++ mtl::Context cx; ++ const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal"}); ++ const char* kd = std::getenv("DQ_AB_KERNELS"); ++ if (!kd) kd = std::getenv("STRATA_DQ_AB_KERNELS"); ++ int argi = 2; ++ if (argc > argi && (std::strchr(argv[argi], '/') || std::strchr(argv[argi], '.'))) kd = argv[argi++]; ++ const std::string d = kd ? kd : ""; ++ const int lib_a = kd ? cx.compile({d + "/quant.metal", d + "/units.metal", d + "/mma.metal", d + "/dense_mm.metal"}) : lib; ++ return dq_ab_bench(cx, lib, lib_a, argc > argi ? argv[argi] : nullptr, argc > argi + 1 ? argv[argi + 1] : nullptr); ++ } + const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal"}); + std::printf("device %s\n", cx.device_name().c_str()); ++ const std::string mode = argc > 1 ? argv[1] : ""; ++ if (mode == "--tiles" || mode == "--ab") ++ return tile_bench(cx, lib, argc > 2 ? std::atoi(argv[2]) : 1983, mode == "--ab", argc > 3 ? argv[3] : nullptr); + mtl::Pipeline pact = cx.pipeline(lib, "act_f16"); + auto run_all = [&](int ty, const mtl::Buffer& W, size_t woff, int n_in, int n_out, int T, const mtl::Buffer& x, + const mtl::Buffer& xh, const mtl::Buffer& wh, const mtl::Buffer& y, const std::string& gk, +@@ -45,12 +267,13 @@ int main(int argc, char** argv) { + if (which & 1) { + const uint32_t n = (uint32_t)T * n_in; + mtl::Dispatch d; d.p = &pact; d.args = {mtl::Arg::B(x), mtl::Arg::B(xh), mtl::Arg::V(n)}; +- d.grid[0] = (n + 255) / 256; d.tg[0] = 256; ds.push_back(d); ++ d.grid[0] = (n + 1023) / 1024; d.tg[0] = 256; ds.push_back(d); + } + if (which & 2) { ++ const int u = n_in / unit_of(ty), tg = dq_tg(u); + DqArgs a{(uint32_t)n_in, (uint32_t)n_out, (uint32_t)ref::row_bytes(ty, n_in)}; + mtl::Dispatch d; d.p = pdq; d.args = {mtl::Arg::B(W, woff), mtl::Arg::B(wh), mtl::Arg::V(a)}; +- d.grid[0] = (n_in / unit_of(ty) + 63) / 64; d.grid[1] = n_out; d.tg[0] = 64; ds.push_back(d); ++ d.grid[0] = (u + tg - 1) / tg; d.grid[1] = n_out; d.tg[0] = tg; ds.push_back(d); + } + if (which & 4) { + GemmArgs g{T, n_out, n_in, 0}; +@@ -58,8 +281,8 @@ int main(int argc, char** argv) { + d.grid[0] = (n_out + tn - 1) / tn; d.grid[1] = (T + tm - 1) / tm; d.tg[0] = 128; ds.push_back(d); + } + }; +- const std::vector types = {ref::F32, ref::BF16, ref::Q8_0, ref::Q4_0, ref::Q5_0, ref::IQ4_NL, ref::IQ4_XS, ref::Q2_0, +- ref::IQ3_S, ref::Q3_K, ref::Q4_K, ref::Q5_K, ref::Q6_K}; ++ const std::vector types = {ref::F32, ref::BF16, ref::Q8_0, ref::Q4_0, ref::Q5_0, ref::Q5_1, ref::IQ4_NL, ref::IQ4_XS, ++ ref::Q2_0, ref::IQ3_S, ref::Q3_K, ref::Q4_K, ref::Q5_K, ref::Q6_K}; + if (!bench) { + int fails = 0; + for (int ty : types) +@@ -78,6 +301,11 @@ int main(int argc, char** argv) { + cx.run(ds); + std::vector wf((size_t)n_in * n_out); + ref::dequant(ty, W.ptr, wf.data(), wf.size()); ++ std::vector wh_ref(wf.size()), xh_ref((size_t)T * n_in); ++ for (size_t i = 0; i < wf.size(); ++i) wh_ref[i] = ref::f32_to_f16(wf[i]); ++ for (size_t i = 0; i < xh_ref.size(); ++i) xh_ref[i] = ref::f32_to_f16(x.as()[i]); ++ const int wh_cmp = std::memcmp(wh.ptr, wh_ref.data(), wh_ref.size() * 2); ++ const int xh_cmp = std::memcmp(xh.ptr, xh_ref.data(), xh_ref.size() * 2); + double se16 = 0, se = 0, sr = 0; + for (int t = 0; t < T; ++t) + for (int o = 0; o < n_out; ++o) { +@@ -93,47 +321,50 @@ int main(int argc, char** argv) { + sr += ae * ae; + } + const double r16 = std::sqrt(se16 / sr), re = std::sqrt(se / sr); +- const bool ok = r16 < 1e-5; ++ const bool ok = r16 < 1e-5 && wh_cmp == 0 && xh_cmp == 0; + fails += !ok; +- std::printf("dense mm %-7s %5d -> %4d, T %d: vs fp16 operands %.2e, vs exact %.2e %s\n", +- ref::type_name(ty), n_in, n_out, T, r16, re, ok ? "PASS" : "FAIL"); ++ std::printf("dense mm %-7s %5d -> %4d, T %d: vs fp16 operands %.2e, vs exact %.2e, memcmp(wh)=%d, memcmp(xh)=%d %s\n", ++ ref::type_name(ty), n_in, n_out, T, r16, re, wh_cmp, xh_cmp, ok ? "PASS" : "FAIL"); + } + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; + } +- struct Shape { const char* name; int type, n_in, n_out; }; +- const std::vector shapes = {{"gdn qkv", ref::IQ4_XS, 2560, 10240}, {"gdn qkv", ref::Q3_K, 2560, 10240}, +- {"gdn gate", ref::Q4_K, 2560, 6144}, {"gdn out", ref::Q6_K, 6144, 2560}, +- {"qsa q", ref::Q2_0, 2560, 12288}, {"hc down", ref::BF16, 10240, 324}, +- {"hc up", ref::BF16, 320, 10240}, {"router", ref::BF16, 2560, 512}, +- {"alpha", ref::BF16, 2560, 48}}; +- const int T = 2048; +- for (const Shape& s : shapes) { +- if (s.n_in % 128) { std::printf("%-9s %-6s: n_in %d not a multiple of 128 (needs the 32-chunk path)\n", s.name, ref::type_name(s.type), s.n_in); continue; } +- const size_t rb = ref::row_bytes(s.type, s.n_in); +- mtl::Buffer W = cx.alloc(rb * s.n_out), x = cx.alloc(4 * (size_t)T * s.n_in), xh = cx.alloc(2 * (size_t)T * s.n_in), +- wh = cx.alloc(2 * (size_t)s.n_out * s.n_in), y = cx.alloc(4 * (size_t)T * s.n_out); +- ref::random_blocks(s.type, W.as(), s.n_out, s.n_in, 3, 0.02f); +- auto timeit = [&](int which, const std::string& gk) { +- std::vector ds; +- for (int i = 0; i < 4; ++i) run_all(s.type, W, 0, s.n_in, s.n_out, T, x, xh, wh, y, gk, ds, which); +- cx.warm(ds, 0.2); +- std::vector t; +- for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds) / 4); +- return median(t); +- }; +- const double tdq = timeit(2, "gemm_f16_64_64_4"), tact = timeit(1, "gemm_f16_64_64_4"); +- double best = 1e9; +- std::string bk; +- for (const char* gk : {"gemm_f16_64_64_4", "gemm_f16_128_64_4", "gemm_f16_64_128_4", "gemm_f16_64_32_4"}) { +- const double tg = timeit(4, gk); +- if (tg < best) { best = tg; bk = gk; } ++ const size_t max_w = (size_t)4 * 2560 * 12288; ++ const size_t max_wh = (size_t)2 * 2560 * 12288; ++ mtl::Buffer W = cx.alloc(max_w), wh = cx.alloc(max_wh); ++ std::vector wf(2560 * 12288); ++ std::vector wh_ref(2560 * 12288); ++ struct DqShape { const char* name; int n_in, n_out; }; ++ const std::vector dq_shapes = {{"2560x10240", 2560, 10240}, {"2560x12288", 2560, 12288}, ++ {"2560x6144", 2560, 6144}, {"10240x324", 10240, 324}}; ++ int dq_fails = 0; ++ for (const auto& ds_sh : dq_shapes) { ++ std::printf("--- dq_f16 all 13 formats on %s (%d x %d): bitwise memcmp + speed ---\n", ds_sh.name, ds_sh.n_in, ds_sh.n_out); ++ for (int ty : types) { ++ const size_t rb = ref::row_bytes(ty, ds_sh.n_in); ++ ref::random_blocks(ty, W.as(), std::min(ds_sh.n_out, 64), ds_sh.n_in, 17 + ty, 0.02f); ++ for (int row = 64; row < ds_sh.n_out; row += 64) { ++ const int ncopy = std::min(64, ds_sh.n_out - row); ++ std::memcpy(W.as() + (size_t)row * rb, W.as(), (size_t)ncopy * rb); ++ } ++ mtl::Pipeline pdq = cx.pipeline(lib, "dq_f16_" + std::to_string(ty)); ++ const int u = ds_sh.n_in / unit_of(ty), tg = dq_tg(u); ++ DqArgs a{(uint32_t)ds_sh.n_in, (uint32_t)ds_sh.n_out, (uint32_t)rb, 0}; ++ mtl::Dispatch d; d.p = &pdq; d.args = {mtl::Arg::B(W), mtl::Arg::B(wh), mtl::Arg::V(a)}; ++ d.grid[0] = (u + tg - 1) / tg; d.grid[1] = ds_sh.n_out; d.tg[0] = tg; d.barrier = true; ++ std::vector ds(4, d); ++ cx.run_graph(ds); ++ const size_t nelem = (size_t)ds_sh.n_in * ds_sh.n_out; ++ ref::dequant(ty, W.ptr, wf.data(), nelem); ++ for (size_t i = 0; i < nelem; ++i) wh_ref[i] = ref::f32_to_f16(wf[i]); ++ const int cmp = std::memcmp(wh.ptr, wh_ref.data(), nelem * 2); ++ dq_fails += (cmp != 0); ++ std::vector ts; ++ for (int i = 0; i < 5; ++i) ts.push_back(cx.run_graph(ds) / 4); ++ std::printf("dq_f16 %-7s %-10s: %6.1f us, memcmp=%d %s\n", ++ ref::type_name(ty), ds_sh.name, median(ts) * 1e6, cmp, cmp == 0 ? "PASS" : "FAIL"); + } +- const double fl = 2.0 * T * s.n_in * s.n_out; +- std::printf("%-9s %-6s %5d -> %5d, T %d: dequant %7.1f us, act %6.1f us, GEMM %7.1f us (%5.1f TFLOP/s, %s) -> %.2f us a token\n", +- s.name, ref::type_name(s.type), s.n_in, s.n_out, T, tdq * 1e6, tact * 1e6, best * 1e6, fl / best / 1e12, +- bk.c_str(), (tdq + best) / T * 1e6); + } +- return 0; ++ return dq_fails ? 1 : 0; + } +diff --git a/metal/tests/expert_parity.cpp b/metal/tests/expert_parity.cpp +index 3bd15ba3..10d082b7 100644 +--- a/metal/tests/expert_parity.cpp ++++ b/metal/tests/expert_parity.cpp +@@ -2,6 +2,7 @@ + // reference, and their speed. + // + // expert_parity [--gu Q2_0] [--dn Q2_0] [--ntok 1..8] [--gucfg L,R,S] [--dncfg L,R,S] parity (synthetic) ++// [--same] every token routed to the same 10 experts (each group has ntok entries) + // expert_parity --bench [--gu ..] [--dn ..] [--ntok N] [--layers 8] [--reps 96] [--sweep] speed + // + // Parity: random valid blocks for 64 experts, random activations x ~ U(-1,1), random distinct top-10 per token, +@@ -69,7 +70,7 @@ Layout layout(int gu, int dn) { + } + + void random_routing(ref::Rng& r, int ntok, int n_expert, std::vector& ids, std::vector& w, +- double share = 0.0, const std::vector* prev = nullptr) { ++ double share = 0.0, const std::vector* prev = nullptr, bool same = false) { + ids.assign(ntok * TOPK, 0); + w.assign(ntok * TOPK, 0.0f); + for (int t = 0; t < ntok; ++t) { +@@ -78,6 +79,7 @@ void random_routing(ref::Rng& r, int ntok, int n_expert, std::vector& ids, + for (int k = 0; k < TOPK; ++k) { + int e; + do { ++ if (same && t > 0) { e = ids[(k + t) % TOPK]; break; } // token 0's experts, rotated + // with probability `share`, reuse an expert of the previous token (verify windows share experts) + if (t > 0 && share > 0 && r.uni() < share) e = ids[(t - 1) * TOPK + (r.next() % TOPK)]; + else e = (int)(r.next() % n_expert); +@@ -192,7 +194,9 @@ int main(int argc, char** argv) { + Cfg cgu{32, 2, 4}, cdn{8, 2, 4}; + bool bench = false, sweep = false; + double share = 0.0; ++ bool same = false; + bool cfg_gu_set = false, cfg_dn_set = false; ++ std::string ab_ref; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + auto nx = [&] { return i + 1 < argc ? argv[++i] : (std::fprintf(stderr, "%s needs a value\n", a.c_str()), std::exit(2), (char*)nullptr); }; +@@ -207,6 +211,8 @@ int main(int argc, char** argv) { + else if (a == "--reps") reps = std::atoi(nx()); + else if (a == "--share") share = std::atof(nx()); + else if (a == "--dn-prep") g_dn_prep = true; ++ else if (a == "--same") same = true; ++ else if (a == "--ab") ab_ref = nx(); + else { std::fprintf(stderr, "unknown argument %s\n", a.c_str()); return 2; } + } + (void)cfg_gu_set; +@@ -214,21 +220,33 @@ int main(int argc, char** argv) { + if (dn == ref::IQ4_NL && cdn.L == 8 && !cfg_dn_set) cdn = {4, 2, 4}; + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal"}); ++ int ref_lib = -1; ++ if (!ab_ref.empty()) { ++ if (ab_ref.size() >= 6 && ab_ref.substr(ab_ref.size() - 6) == ".metal") { ++ ref_lib = cx.compile({"quant.metal", "units.metal", ab_ref}); ++ } else { ++ const std::string cur = cx.kernel_dir; ++ cx.kernel_dir = ab_ref; ++ ref_lib = cx.compile({"quant.metal", "units.metal", "moe.metal"}); ++ cx.kernel_dir = cur; ++ } ++ } + const Layout L = layout(gu, dn); + std::printf("device %s; gate/up %s, down %s, blob %zu bytes; %d token(s)\n", cx.device_name().c_str(), + ref::type_name(gu), ref::type_name(dn), L.blob, ntok); + +- auto make = [&](Cfg a, Cfg b) { ++ auto make_from = [&](int l, Cfg a, Cfg b) { + Kernels K; +- K.plan = cx.pipeline(lib, "moe_plan"); +- K.prep = cx.pipeline(lib, "xprep"); +- K.comb = cx.pipeline(lib, "moe_combine"); +- K.gu = cx.pipeline(lib, kname("moe_gu", gu, a), mtl::Constants().i(22, std::getenv("MOE_PROBE") ? std::atoi(std::getenv("MOE_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)); +- K.dn = cx.pipeline(lib, kname("moe_down", dn == ref::Q2_0 && !g_dn_prep ? 142 : dn, b)); ++ K.plan = cx.pipeline(l, "moe_plan"); ++ K.prep = cx.pipeline(l, "xprep"); ++ K.comb = cx.pipeline(l, "moe_combine"); ++ K.gu = cx.pipeline(l, kname("moe_gu", gu, a), mtl::Constants().i(22, std::getenv("MOE_PROBE") ? std::atoi(std::getenv("MOE_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)); ++ K.dn = cx.pipeline(l, kname("moe_down", dn == ref::Q2_0 && !g_dn_prep ? 142 : dn, b)); + K.cgu = a; + K.cdn = b; + return K; + }; ++ auto make = [&](Cfg a, Cfg b) { return make_from(lib, a, b); }; + auto args_for = [&](const Kernels& K, int nt, int n_expert) { + MoeArgs a{}; + a.gate_off = 0; +@@ -262,6 +280,117 @@ int main(int argc, char** argv) { + cx.run({prep_dispatch(pp, x, XPX, N_EMBD, 8)}); + } + ++ if (!ab_ref.empty()) { ++ const int NE_AB = 64; ++ mtl::Buffer W = cx.alloc(L.blob * NE_AB); ++ for (int e = 0; e < NE_AB; ++e) { ++ uint8_t* b = W.as() + e * L.blob; ++ ref::random_blocks(gu, b, 2 * N_FF, N_EMBD, 1000 + e); ++ ref::random_blocks(dn, b + L.down_off, N_EMBD, N_FF, 5000 + e); ++ } ++ mtl::Buffer h_r = cx.alloc(h.bytes), ye_r = cx.alloc(ye.bytes), y_r = cx.alloc(y.bytes), plan_r = cx.alloc(4 * P_INTS); ++ mtl::Buffer bid = cx.alloc(4 * 8 * TOPK), bw = cx.alloc(4 * 8 * TOPK), plan = cx.alloc(4 * P_INTS); ++ Kernels Kc = make_from(lib, cgu, cdn), Kr = make_from(ref_lib, cgu, cdn); ++ int ab_fails = 0; ++ for (bool sm : {true, false}) { ++ for (int nt : {1, 2, 3, 4, 8}) { ++ std::vector ids; ++ std::vector w; ++ random_routing(rng, nt, NE_AB, ids, w, sm ? 0.0 : 0.35, nullptr, sm); ++ std::memcpy(bid.ptr, ids.data(), 4 * ids.size()); ++ std::memcpy(bw.ptr, w.data(), 4 * w.size()); ++ MoeArgs a = args_for(Kc, nt, 512); ++ std::memset(h.ptr, 0xa5, h.bytes); std::memset(ye.ptr, 0xa5, ye.bytes); ++ std::memset(y.ptr, 0xa5, y.bytes); std::memset(plan.ptr, 0xa5, plan.bytes); ++ std::memset(h_r.ptr, 0xa5, h_r.bytes); std::memset(ye_r.ptr, 0xa5, ye_r.bytes); ++ std::memset(y_r.ptr, 0xa5, y_r.bytes); std::memset(plan_r.ptr, 0xa5, plan_r.bytes); ++ std::vector dsc, dsr; ++ encode_layer(dsc, Kc, W, 0, x, h, ye, y, bid, bw, plan, 0, a, true, true, true, true); ++ encode_layer(dsr, Kr, W, 0, x, h_r, ye_r, y_r, bid, bw, plan_r, 0, a, true, true, true, true); ++ cx.run(dsc); ++ cx.run(dsr); ++ const int ne_tot = nt * TOPK; ++ if (std::memcmp(plan.ptr, plan_r.ptr, plan.bytes) != 0 || ++ std::memcmp(h.ptr, h_r.ptr, 4 * ne_tot * N_FF) != 0 || ++ std::memcmp(ye.ptr, ye_r.ptr, 4 * ne_tot * N_EMBD) != 0 || ++ std::memcmp(y.ptr, y_r.ptr, 4 * nt * N_EMBD) != 0) { ++ ++ab_fails; ++ std::printf("AB MISMATCH same=%d ntok=%d\n", (int)sm, nt); ++ } ++ } ++ } ++ std::printf("expert_parity --ab bitwise (%s/%s%s gu %d,%d,%d dn %d,%d,%d, NE=1..4,8 + random): %d fail(s) %s\n", ++ ref::type_name(gu), ref::type_name(dn), g_dn_prep ? "(prep)" : "(142)", ++ cgu.L, cgu.R, cgu.S, cdn.L, cdn.R, cdn.S, ab_fails, ab_fails == 0 ? "ALL PASS" : "FAIL"); ++ if (ab_fails || !bench) return ab_fails ? 1 : 0; ++ ++ // Paired A/B benchmark across isolated NE=1,2,3,4 (same=true) and verify windows ntok=1,2,4,8 (share=0.3) ++ const int NE_B = 512; ++ const size_t layer_bytes = L.blob * NE_B; ++ mtl::Buffer Wb = cx.alloc(layer_bytes * layers); ++ for (size_t e = 0; e < (size_t)NE_B * layers; ++e) ++ std::memcpy(Wb.as() + e * L.blob, W.as() + (e % NE_AB) * L.blob, L.blob); ++ mtl::Buffer bid_b = cx.alloc(4 * 8 * TOPK * reps), bw_b = cx.alloc(4 * 8 * TOPK * reps), plan_b = cx.alloc(4 * P_INTS * reps); ++ auto run_ab_case = [&](int nt, bool sm, double sh, const char* tag) { ++ double groups = 0; ++ for (int r = 0; r < reps; ++r) { ++ std::vector ids; std::vector w; ++ random_routing(rng, nt, NE_B, ids, w, sh, nullptr, sm); ++ std::memcpy(bid_b.as() + r * 8 * TOPK, ids.data(), 4 * ids.size()); ++ std::memcpy(bw_b.as() + r * 8 * TOPK, w.data(), 4 * w.size()); ++ groups += std::set(ids.begin(), ids.end()).size(); ++ } ++ groups /= reps; ++ const MoeArgs a = args_for(Kc, nt, NE_B); ++ std::vector dp, dgc, dgr, ddc, ddr; ++ for (int r = 0; r < reps; ++r) { ++ const size_t woff = (size_t)(r % layers) * layer_bytes; ++ const size_t poff = (size_t)r * P_INTS * 4; ++ std::vector p1, gc, gr, dc, dr; ++ encode_layer(p1, Kc, Wb, woff, x, h, ye, y, bid_b, bw_b, plan_b, poff, a, true, false, false, false); ++ p1[0].args[0] = mtl::Arg::B(bid_b, (size_t)r * 8 * TOPK * 4); ++ p1[0].args[1] = mtl::Arg::B(bw_b, (size_t)r * 8 * TOPK * 4); ++ dp.push_back(p1[0]); ++ encode_layer(gc, Kc, Wb, woff, x, h, ye, y, bid_b, bw_b, plan_b, poff, a, false, true, false, false); ++ encode_layer(gr, Kr, Wb, woff, x, h, ye, y, bid_b, bw_b, plan_b, poff, a, false, true, false, false); ++ encode_layer(dc, Kc, Wb, woff, x, h, ye, y, bid_b, bw_b, plan_b, poff, a, false, false, true, false); ++ encode_layer(dr, Kr, Wb, woff, x, h, ye, y, bid_b, bw_b, plan_b, poff, a, false, false, true, false); ++ dgc.push_back(gc[0]); dgr.push_back(gr[0]); ++ for (auto& z : dc) ddc.push_back(z); ++ for (auto& z : dr) ddr.push_back(z); ++ } ++ cx.run(dp); ++ cx.run(dgc); ++ cx.warm(ddr, 0.04); cx.warm(ddc, 0.04); ++ std::vector b_dn, c_dn, r_dn, b_gu, c_gu, r_gu; ++ for (int it = 0; it < 7; ++it) { ++ double bd, cd, bg, cg; ++ if (it & 1) { ++ cd = cx.run(ddc) / reps; bd = cx.run(ddr) / reps; ++ cg = cx.run(dgc) / reps; bg = cx.run(dgr) / reps; ++ } else { ++ bd = cx.run(ddr) / reps; cd = cx.run(ddc) / reps; ++ bg = cx.run(dgr) / reps; cg = cx.run(dgc) / reps; ++ } ++ b_dn.push_back(bd * 1e6); c_dn.push_back(cd * 1e6); r_dn.push_back(cd / bd); ++ b_gu.push_back(bg * 1e6); c_gu.push_back(cg * 1e6); r_gu.push_back(cg / bg); ++ } ++ const double bd_bytes = groups * N_EMBD * L.d_row; ++ std::printf("ab %-14s groups %4.1f | dn base %6.1f us cand %6.1f us (%5.1f GB/s) ratio %.4f | gu base %6.1f us cand %6.1f us ratio %.4f\n", ++ tag, groups, median(b_dn), median(c_dn), bd_bytes / (median(c_dn) * 1e-6) / 1e9, median(r_dn), ++ median(b_gu), median(c_gu), median(r_gu)); ++ }; ++ for (int ne : {1, 2, 3, 4}) { ++ char buf[32]; std::snprintf(buf, sizeof(buf), "NE=%d (same)", ne); ++ run_ab_case(ne, true, 0.0, buf); ++ } ++ for (int nt : {1, 2, 4, 8}) { ++ char buf[32]; std::snprintf(buf, sizeof(buf), "ntok=%d sh=0.3", nt); ++ run_ab_case(nt, false, 0.3, buf); ++ } ++ return 0; ++ } ++ + if (!bench) { + // ---------------- parity ---------------- + const int NE = 64; +@@ -273,7 +402,7 @@ int main(int argc, char** argv) { + } + std::vector ids; + std::vector w; +- random_routing(rng, ntok, NE, ids, w, 0.3); ++ random_routing(rng, ntok, NE, ids, w, 0.3, nullptr, same); + mtl::Buffer bid = cx.alloc(4 * ids.size()), bw = cx.alloc(4 * w.size()), plan = cx.alloc(4 * P_INTS); + std::memcpy(bid.ptr, ids.data(), 4 * ids.size()); + std::memcpy(bw.ptr, w.data(), 4 * w.size()); +@@ -364,7 +493,7 @@ int main(int argc, char** argv) { + auto set_routing = [&](int nt) { + double groups = 0; + for (int r = 0; r < reps; ++r) { +- random_routing(rng, nt, NE, rids[r], rw[r], share); ++ random_routing(rng, nt, NE, rids[r], rw[r], share, nullptr, same); + std::memcpy(bid.as() + r * 8 * TOPK, rids[r].data(), 4 * rids[r].size()); + std::memcpy(bw.as() + r * 8 * TOPK, rw[r].data(), 4 * rw[r].size()); + groups += std::set(rids[r].begin(), rids[r].end()).size(); +diff --git a/metal/tests/fuse_parity.cpp b/metal/tests/fuse_parity.cpp +index e67d9e58..6c6fe1e6 100644 +--- a/metal/tests/fuse_parity.cpp ++++ b/metal/tests/fuse_parity.cpp +@@ -1,14 +1,18 @@ +-// metal/tests/fuse_parity.cpp - the decode latency fusions against the kernels they replace, BITWISE, for 1..4 ++// metal/tests/fuse_parity.cpp - the decode latency fusions against the kernels they replace, BITWISE, for 1..8 + // tokens: + // gemv_cat (gemv.metal) ssm_alpha + ssm_beta (BF16 2560 -> 48 + 48) in one dispatch == two gemv dispatches + // router_route the router GEMV + moe_route in one dispatch == gemv + moe_route (logits, ids, w, plan; + // twice in a row: the counter resets) + // gr_norm_comb (gr2.metal) the MoE combine folded into the next gr_norm == moe_combine_sh + gr_norm (R', xn) ++// --gr-unroll BASE_DIR NTOK [bench]: gr.metal / gr2.metal's kernels against the copies in BASE_DIR (the same entry ++// points, or gr_down's variants against the entry point they replace, compiled from the old source), bitwise, and ++// paired timings of the decode window's gr_down / gr_up + #include "mtl_rt.h" + #include "ref_quant.h" + + #include + #include ++#include + #include + #include + +@@ -18,12 +22,292 @@ struct MoeArgs { + uint64_t gate_off, up_off, down_off, gu_stride, d_stride; + uint32_t gu_row, d_row, n_embd, n_ff, n_tok, topk, n_expert, gx_gu, gx_dn, pad; + }; +-constexpr int P_INTS = 16 + 128 + 128 + 4 + 128 * 3, P_TOKENT = 16 + 128 + 128 + 4 + 128 * 2; ++constexpr int P_INTS = 16 + 128 + 128 + 4 + 128 * 3 + 20 + 128, P_TOKENT = 16 + 128 + 128 + 4 + 128 * 2; + struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; + struct GlueArgs { uint32_t n, n_tok, n_embd, topk; }; ++struct GrArgs3 { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; uint32_t rb_down, rb_up, comb, topk; }; ++ ++// The A/B source holds baseline and candidate entry points in one library. ++int gr_ab(const char* source, int nt, bool down) { ++ if (nt < 1 || nt > 8) return 2; ++ mtl::Context cx; ++ std::printf("device %s\n", cx.device_name().c_str()); ++ const int lib = down ? cx.compile({"quant.metal", "units.metal", source}) ++ : cx.compile({source}, "// NO_GGML_COMMON"); ++ constexpr int N = 2560, HC = 4, LR = 320, TK = 10; ++ ref::Rng rng(1802); ++ auto floats = [&](size_t n) { ++ auto b = cx.alloc(4 * n); ++ for (size_t i = 0; i < n; ++i) b.as()[i] = rng.sym(); ++ return b; ++ }; ++ auto R0 = floats(nt * HC * N), R1 = floats(nt * HC * N), R2 = floats(nt * HC * N); ++ auto x1 = floats(nt * HC * N), x2 = floats(nt * HC * N); ++ GrArgs a{N, HC, LR, (uint32_t)nt, 1, 1e-6f}; ++ int fails = 0; ++ auto bench = [&](mtl::Dispatch b, mtl::Dispatch v, const char* name) { ++ cx.warm({b}, 0.3); cx.warm({v}, 0.1); ++ for (int pair = 0; pair < 5; ++pair) { ++ double tb, tv; ++ if (pair % 2) { tv = cx.run({v}, 64); tb = cx.run({b}, 64); } ++ else { tb = cx.run({b}, 64); tv = cx.run({v}, 64); } ++ std::printf("bench %s ntok %d pair %d base_us %.6f variant_us %.6f\n", name, nt, pair, ++ tb * 1e6 / 64, tv * 1e6 / 64); ++ } ++ }; ++ if (down) { ++ const size_t rb = ref::row_bytes(ref::Q8_0, HC * N); ++ auto wd = cx.alloc(rb * LR), wi = cx.alloc(rb * HC); ++ ref::random_blocks(ref::Q8_0, wd.as(), LR, HC * N, 1803, 0.05f); ++ ref::random_blocks(ref::Q8_0, wi.as(), HC, HC * N, 1804, 0.05f); ++ auto lo1 = floats(nt * LR), lo2 = floats(nt * LR), in1 = floats(nt * HC), in2 = floats(nt * HC); ++ auto pb = cx.pipeline(lib, "baseline_down", mtl::Constants().i(1, nt)); ++ auto pv = cx.pipeline(lib, "optimized_down", mtl::Constants().i(1, nt)); ++ mtl::Dispatch b, v; ++ b.p = &pb; v.p = &pv; ++ b.args = {mtl::Arg::B(wd), mtl::Arg::B(wi), mtl::Arg::B(R0), mtl::Arg::B(lo1), mtl::Arg::B(in1), mtl::Arg::V(a)}; ++ v.args = {mtl::Arg::B(wd), mtl::Arg::B(wi), mtl::Arg::B(R0), mtl::Arg::B(lo2), mtl::Arg::B(in2), mtl::Arg::V(a)}; ++ b.grid[0] = v.grid[0] = (LR + HC) / 2; ++ b.tg[0] = v.tg[0] = 256; ++ cx.run({b}); cx.run({v}); ++ const bool ok = std::memcmp(lo1.ptr, lo2.ptr, lo1.bytes) == 0 && ++ std::memcmp(in1.ptr, in2.ptr, in1.bytes) == 0; ++ fails += !ok; ++ std::printf("down ntok %d bitwise %s\n", nt, ok ? "PASS" : "FAIL"); ++ bench(b, v, "down"); ++ } else { ++ auto ye = floats(nt * TK * N), ysh = floats(nt * N), sgv = floats(nt), inj = floats(nt * HC), wn = floats(HC * N); ++ auto plan = cx.alloc(4 * P_INTS); ++ for (int i = 0; i < nt * TK; ++i) plan.as()[P_TOKENT + i] = (i * 7) % (nt * TK); ++ auto pb = cx.pipeline(lib, "baseline_4s"); ++ for (const char* name : {"gr_norm_comb_4s", "gr_norm_comb", "gr_norm_comb_staged"}) { ++ auto pv = cx.pipeline(lib, name); ++ for (uint32_t tk : {1u, 10u}) for (uint32_t apply : {0u, 1u}) { ++ a.apply = apply; ++ mtl::Dispatch b, v; ++ auto args = [&](const mtl::Buffer& R, const mtl::Buffer& x) { ++ return std::vector{mtl::Arg::B(R), mtl::Arg::B(ye), mtl::Arg::B(inj), mtl::Arg::B(wn), ++ mtl::Arg::B(x), mtl::Arg::V(a), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::V(tk)}; ++ }; ++ b.p = &pb; v.p = &pv; b.args = args(R1, x1); v.args = args(R2, x2); ++ b.grid[0] = 1; v.grid[0] = std::strcmp(name, "gr_norm_comb") == 0 ? HC : 1; ++ b.grid[1] = v.grid[1] = nt; b.tg[0] = v.tg[0] = 256; ++ std::memcpy(R1.ptr, R0.ptr, R0.bytes); std::memcpy(R2.ptr, R0.ptr, R0.bytes); ++ cx.run({b}); cx.run({v}); ++ const bool ok = std::memcmp(R1.ptr, R2.ptr, R0.bytes) == 0 && ++ std::memcmp(x1.ptr, x2.ptr, x1.bytes) == 0; ++ fails += !ok; ++ std::printf("%s ntok %d topk %u apply %u R/xn bitwise %s\n", ++ name, nt, tk, apply, ok ? "PASS" : "FAIL"); ++ if (apply && tk == TK) bench(b, v, name); ++ } ++ } ++ } ++ return fails ? 1 : 0; ++} ++ ++// Each kernel runs from both libraries on the same inputs into separate outputs, which must match byte for byte. ++int gr_unroll(const std::string& base, int nt, bool timed) { ++ if (nt < 1 || nt > 8) return 2; ++ mtl::Context cx; ++ std::printf("device %s\n", cx.device_name().c_str()); ++ const std::vector srcs = {"quant.metal", "units.metal", "gr.metal", "gr2.metal"}; ++ std::vector old; ++ for (const auto& s : srcs) old.push_back(base + "/" + s); ++ const int lb = cx.compile(old), lv = cx.compile(srcs); ++ constexpr int N = 2560, HC = 4, LR = 320, NR = LR + HC, TK = 10, MAXT = 8; ++ ref::Rng rng(2101); ++ auto floats = [&](size_t n, float scale = 1.0f, float bias = 0.0f) { ++ auto b = cx.alloc(4 * n); ++ for (size_t i = 0; i < n; ++i) b.as()[i] = bias + scale * rng.sym(); ++ return b; ++ }; ++ auto weights = [&](int type, int rows, int cols, uint64_t seed) { ++ auto b = cx.alloc(ref::row_bytes(type, cols) * rows); ++ ref::random_blocks(type, b.as(), rows, cols, seed, 0.05f); ++ return b; ++ }; ++ auto wd8 = weights(ref::Q8_0, LR, HC * N, 11), wi8 = weights(ref::Q8_0, HC, HC * N, 12), ++ wif = weights(ref::F32, HC, HC * N, 13), wu8 = weights(ref::Q8_0, HC * N, LR, 14), ++ wd16 = weights(ref::BF16, LR, HC * N, 15), wi16 = weights(ref::BF16, HC, HC * N, 16), ++ wu16 = weights(ref::BF16, HC * N, LR, 17), wd4 = weights(ref::Q4_0, LR, HC * N, 18), ++ wd5 = weights(ref::Q5_0, LR, HC * N, 19); ++ auto R0 = floats((size_t)nt * HC * N), bo = floats((size_t)nt * N), inj = floats(nt * HC, 3.0f), ++ wn = floats(HC * N, 0.5f, 1.0f), xn = floats((size_t)nt * HC * N), lo = floats(nt * LR), ++ ye = floats((size_t)MAXT * TK * N), ysh = floats(nt * N), sgv = floats(nt, 3.0f), ++ Dprev = floats((size_t)nt * NR * HC, 0.3f), rs = floats(nt * HC, 0.1f, 1.0f); ++ auto plan = cx.alloc(4 * P_INTS); ++ for (int i = 0; i < nt * TK; ++i) plan.as()[P_TOKENT + i] = (i * 7) % (MAXT * TK); ++ int fails = 0; ++ // outs[i]: the baseline's and the variant's copy of output i; `fresh` restores in-place inputs before a run ++ struct Out { mtl::Buffer b, v; }; ++ auto outs = [&](std::vector sizes) { ++ std::vector o; ++ for (size_t s : sizes) o.push_back({cx.alloc(s), cx.alloc(s)}); ++ return o; ++ }; ++ auto check = [&](const std::string& name, mtl::Dispatch b, mtl::Dispatch v, std::vector& o, ++ const std::function& fresh, bool bench) { ++ for (auto& x : o) { std::memset(x.b.ptr, 0, x.b.bytes); std::memset(x.v.ptr, 0, x.v.bytes); } ++ if (fresh) fresh(); ++ cx.run({b}); ++ cx.run({v}); ++ bool ok = true; ++ for (auto& x : o) ok = ok && std::memcmp(x.b.ptr, x.v.ptr, x.b.bytes) == 0; ++ fails += !ok; ++ std::printf("%s ntok %d bitwise %s\n", name.c_str(), nt, ok ? "PASS" : "FAIL"); ++ if (!bench) return; ++ cx.warm({b, v}, 0.2); ++ // the median pair ratio: a lone mean drifts with the GPU clock and whatever else shares the GPU ++ constexpr int PAIRS = 16, IT = 128; ++ double sb = 0, sv = 0; ++ std::vector ratio; ++ for (int pair = 0; pair < PAIRS; ++pair) { ++ double tb, tv; ++ if (pair % 2) { tv = cx.run({v}, IT); tb = cx.run({b}, IT); } ++ else { tb = cx.run({b}, IT); tv = cx.run({v}, IT); } ++ sb += tb; sv += tv; ++ ratio.push_back(tv / tb); ++ } ++ std::sort(ratio.begin(), ratio.end()); ++ std::printf("bench %s ntok %d mean base_us %.3f variant_us %.3f (%+.1f%%) median %+.1f%%\n", name.c_str(), nt, ++ sb * 1e6 / (PAIRS * IT), sv * 1e6 / (PAIRS * IT), 100.0 * (sv / sb - 1), ++ 100.0 * ((ratio[PAIRS / 2 - 1] + ratio[PAIRS / 2]) / 2 - 1)); ++ }; ++ const mtl::Constants fc = mtl::Constants().i(1, nt); ++ // ---- gr_norm (apply 0, 1); no function constant, R written in place ++ { ++ auto pb = cx.pipeline(lb, "gr_norm"), pv = cx.pipeline(lv, "gr_norm"); ++ for (uint32_t apply : {0u, 1u}) { ++ GrArgs a{N, HC, LR, (uint32_t)nt, apply, 1e-6f}; ++ auto o = outs({R0.bytes, xn.bytes}); ++ mtl::Dispatch b, v; ++ b.p = &pb; v.p = &pv; ++ b.args = {mtl::Arg::B(o[0].b), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(o[1].b), mtl::Arg::V(a)}; ++ v.args = {mtl::Arg::B(o[0].v), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(o[1].v), mtl::Arg::V(a)}; ++ b.grid[0] = v.grid[0] = HC; b.grid[1] = v.grid[1] = nt; b.tg[0] = v.tg[0] = 256; ++ check("gr_norm apply " + std::to_string(apply), b, v, o, [&] { ++ std::memcpy(o[0].b.ptr, R0.ptr, R0.bytes); std::memcpy(o[0].v.ptr, R0.ptr, R0.bytes); }, timed); ++ } ++ } ++ // ---- gr_norm_comb (gr2.metal) ++ { ++ auto pb = cx.pipeline(lb, "gr_norm_comb"), pv = cx.pipeline(lv, "gr_norm_comb"); ++ GrArgs a{N, HC, LR, (uint32_t)nt, 1, 1e-6f}; ++ const uint32_t tk = TK; ++ auto o = outs({R0.bytes, xn.bytes}); ++ auto args = [&](const mtl::Buffer& R, const mtl::Buffer& x) { ++ return std::vector{mtl::Arg::B(R), mtl::Arg::B(ye), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(x), ++ mtl::Arg::V(a), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::V(tk)}; ++ }; ++ mtl::Dispatch b, v; ++ b.p = &pb; v.p = &pv; b.args = args(o[0].b, o[1].b); v.args = args(o[0].v, o[1].v); ++ b.grid[0] = v.grid[0] = HC; b.grid[1] = v.grid[1] = nt; b.tg[0] = v.tg[0] = 256; ++ check("gr_norm_comb", b, v, o, [&] { ++ std::memcpy(o[0].b.ptr, R0.ptr, R0.bytes); std::memcpy(o[0].v.ptr, R0.ptr, R0.bytes); }, timed); ++ } ++ // ---- gr_down_f: rows = w_down then w_inject, RS = SG / KS simdgroups a threadgroup, NR rows a simdgroup; a ++ // variant entry point (vname) runs against the baseline's (name) with its own rows per threadgroup ++ struct Down { const char* name; const char* vname; int rows, vrows, btg, vtg; const mtl::Buffer* wd; const mtl::Buffer* wi; bool bench; }; ++ for (const Down& dn : {Down{"gr_down_q8_8_4", "gr_down_q8_4_4", 2, 1, 256, 128, &wd8, &wi8, false}, ++ Down{"gr_down_q8_8_4", "gr_down_q8_4_4_r2", 2, 2, 256, 128, &wd8, &wi8, false}, ++ Down{"gr_down_q8f_8_4", "gr_down_q8f_4_4", 2, 1, 256, 128, &wd8, &wif, false}, ++ Down{"gr_down_q8f_8_4", "gr_down_q8f_4_4_r2", 2, 2, 256, 128, &wd8, &wif, false}, ++ // the decode window's (M5) and the MTP head's entry points, each against its own copy ++ Down{"gr_down_q8_4_4", "gr_down_q8_4_4", 1, 1, 128, 128, &wd8, &wi8, true}, ++ Down{"gr_down_q8_4_4_r2", "gr_down_q8_4_4_r2", 2, 2, 128, 128, &wd8, &wi8, true}, ++ Down{"gr_down_q8f_4_4", "gr_down_q8f_4_4", 1, 1, 128, 128, &wd8, &wif, true}, ++ Down{"gr_down_q8f_4_4_r2", "gr_down_q8f_4_4_r2", 2, 2, 128, 128, &wd8, &wif, true}, ++ Down{"gr_down_q8_8_8", "gr_down_q8_8_8", 1, 1, 256, 256, &wd8, &wi8, true}, ++ Down{"gr_down_q8_8_8_r2", "gr_down_q8_8_8_r2", 2, 2, 256, 256, &wd8, &wi8, true}, ++ // the dispatch choice, one row a simdgroup against two (gr_read switches from 4 tokens, gr_read_m from 5) ++ Down{"gr_down_q8_4_4", "gr_down_q8_4_4_r2", 1, 2, 128, 128, &wd8, &wi8, true}, ++ Down{"gr_down_q8f_4_4", "gr_down_q8f_4_4_r2", 1, 2, 128, 128, &wd8, &wif, true}, ++ Down{"gr_down_q8_8_8", "gr_down_q8_8_8_r2", 1, 2, 256, 256, &wd8, &wi8, true}, ++ Down{"gr_down_q8f_8_8", "gr_down_q8f_8_8", 1, 1, 256, 256, &wd8, &wif, false}, ++ Down{"gr_down_q8f_8_8", "gr_down_q8f_8_8_r2", 1, 2, 256, 256, &wd8, &wif, false}, ++ Down{"gr_down_8_4", "gr_down_4_4", 2, 1, 256, 128, &wd16, &wi16, false}, ++ Down{"gr_down_8_4", "gr_down_4_4_r2", 2, 2, 256, 128, &wd16, &wi16, false}, ++ Down{"gr_down_8_8", "gr_down_8_8", 1, 1, 256, 256, &wd16, &wi16, false}, ++ Down{"gr_down_8_8", "gr_down_8_8_r2", 1, 2, 256, 256, &wd16, &wi16, false}, ++ Down{"gr_down_q4_8_8", "gr_down_q4_8_8", 1, 1, 256, 256, &wd4, &wi8, false}, ++ Down{"gr_down_q5_8_8", "gr_down_q5_8_8", 1, 1, 256, 256, &wd5, &wi8, false}}) { ++ auto pb = cx.pipeline(lb, dn.name, fc), pv = cx.pipeline(lv, dn.vname, fc); ++ GrArgs a{N, HC, LR, (uint32_t)nt, 1, 1e-6f}; ++ auto o = outs({lo.bytes, (size_t)4 * nt * HC}); ++ mtl::Dispatch b, v; ++ b.p = &pb; v.p = &pv; ++ b.args = {mtl::Arg::B(*dn.wd), mtl::Arg::B(*dn.wi), mtl::Arg::B(xn), mtl::Arg::B(o[0].b), mtl::Arg::B(o[1].b), mtl::Arg::V(a)}; ++ v.args = {mtl::Arg::B(*dn.wd), mtl::Arg::B(*dn.wi), mtl::Arg::B(xn), mtl::Arg::B(o[0].v), mtl::Arg::B(o[1].v), mtl::Arg::V(a)}; ++ b.grid[0] = (NR + dn.rows - 1) / dn.rows; ++ v.grid[0] = (NR + dn.vrows - 1) / dn.vrows; ++ b.tg[0] = dn.btg; ++ v.tg[0] = dn.vtg; ++ check(dn.vname, b, v, o, {}, timed && dn.bench); ++ } ++ // ---- gr_up_f: the prepared planes have MAXT rows ++ for (const char* up : {"gr_up_q8", "gr_up"}) { ++ const bool q8 = std::strcmp(up, "gr_up_q8") == 0; ++ auto pb = cx.pipeline(lb, up, fc), pv = cx.pipeline(lv, up, fc); ++ GrArgs a{N, HC, LR, MAXT, 1, 1e-6f}; ++ auto o = outs({4 * ((size_t)4 * MAXT * N + MAXT * N / 16)}); ++ const mtl::Buffer& w = q8 ? wu8 : wu16; ++ mtl::Dispatch b, v; ++ b.p = &pb; v.p = &pv; ++ b.args = {mtl::Arg::B(w), mtl::Arg::B(lo), mtl::Arg::B(xn), mtl::Arg::B(o[0].b), mtl::Arg::V(a)}; ++ v.args = {mtl::Arg::B(w), mtl::Arg::B(lo), mtl::Arg::B(xn), mtl::Arg::B(o[0].v), mtl::Arg::V(a)}; ++ b.grid[0] = v.grid[0] = N / 16; ++ b.tg[0] = v.tg[0] = 256; ++ check(up, b, v, o, {}, timed && q8); ++ } ++ // ---- gr_down3 / gr_down3b (apply and comb 0 / 1), gr_up3 ++ for (int f : {30, 8}) for (int tt = 1; tt <= 3; ++tt) for (int rps = 1; rps <= 2; ++rps) for (int vb = 0; vb < 1 + (f == 30); ++vb) { ++ const std::string name = (vb ? "gr_down3b_" : "gr_down3_" + std::to_string(f) + "_") + std::to_string(tt) + "_" + std::to_string(rps); ++ auto pb = cx.pipeline(lb, name, fc), pv = cx.pipeline(lv, name, fc); ++ for (uint32_t apply : {0u, 1u}) for (uint32_t comb : {0u, 1u}) { ++ if (comb && !apply) continue; ++ GrArgs3 a{N, HC, LR, MAXT, apply, 1e-6f, (uint32_t)(f == 8 ? HC * N / 32 * 34 : HC * N * 2), 0, comb, TK}; ++ auto o = outs({R0.bytes, Dprev.bytes, rs.bytes}); ++ const mtl::Buffer& wd = f == 8 ? wd8 : wd16; ++ const mtl::Buffer& wi = f == 8 ? wi8 : wi16; ++ auto args = [&](const mtl::Buffer& Ro, const mtl::Buffer& D, const mtl::Buffer& r) { ++ return std::vector{mtl::Arg::B(wd), mtl::Arg::B(wi), mtl::Arg::B(R0), mtl::Arg::B(Ro), mtl::Arg::B(bo), ++ mtl::Arg::B(Dprev), mtl::Arg::B(wn), mtl::Arg::B(D), mtl::Arg::B(r), mtl::Arg::V(a), mtl::Arg::B(plan), ++ mtl::Arg::B(ye), mtl::Arg::B(ysh), mtl::Arg::B(sgv)}; ++ }; ++ mtl::Dispatch b, v; ++ b.p = &pb; v.p = &pv; b.args = args(o[0].b, o[1].b, o[2].b); v.args = args(o[0].v, o[1].v, o[2].v); ++ b.grid[0] = v.grid[0] = (NR + 8 * rps - 1) / (8 * rps); ++ b.grid[1] = v.grid[1] = HC; ++ b.tg[0] = v.tg[0] = 256; ++ check(name + " apply " + std::to_string(apply) + " comb " + std::to_string(comb), b, v, o, {}, false); ++ } ++ } ++ for (int f : {30, 8}) { ++ const std::string name = "gr_up3_" + std::to_string(f); ++ auto pb = cx.pipeline(lb, name, fc), pv = cx.pipeline(lv, name, fc); ++ GrArgs3 a{N, HC, LR, MAXT, 1, 1e-6f, 0, (uint32_t)(f == 8 ? LR / 32 * 34 : LR * 2), 0, TK}; ++ auto o = outs({4 * ((size_t)4 * MAXT * N + MAXT * N / 16)}); ++ const mtl::Buffer& w = f == 8 ? wu8 : wu16; ++ mtl::Dispatch b, v; ++ b.p = &pb; v.p = &pv; ++ b.args = {mtl::Arg::B(w), mtl::Arg::B(Dprev), mtl::Arg::B(R0), mtl::Arg::B(o[0].b), mtl::Arg::V(a), mtl::Arg::B(wn), mtl::Arg::B(rs)}; ++ v.args = {mtl::Arg::B(w), mtl::Arg::B(Dprev), mtl::Arg::B(R0), mtl::Arg::B(o[0].v), mtl::Arg::V(a), mtl::Arg::B(wn), mtl::Arg::B(rs)}; ++ b.grid[0] = v.grid[0] = N / 16; ++ b.tg[0] = v.tg[0] = 256; ++ check(name, b, v, o, {}, false); ++ } ++ std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); ++ return fails ? 1 : 0; ++} ++ + } // namespace + +-int main() { ++int main(int argc, char** argv) { ++ if (argc == 5 && std::string(argv[1]) == "--gr-ab") ++ return gr_ab(argv[2], std::stoi(argv[3]), std::string(argv[4]) == "down"); ++ if (argc >= 4 && std::string(argv[1]) == "--gr-unroll") ++ return gr_unroll(argv[2], std::stoi(argv[3]), argc > 4 && std::string(argv[4]) == "bench"); + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal", "gemv.metal", "router.metal", "gr.metal", + "gr2.metal", "glue.metal"}); +@@ -37,7 +321,7 @@ int main() { + ref::random_blocks(ref::BF16, Wb.as(), NO, N, 2, 0.05f); + ref::Rng r(3); + for (size_t i = 0; i < xp.bytes / 4; ++i) xp.as()[i] = r.sym(); +- for (int nt = 1; nt <= 4; ++nt) { ++ for (int nt = 1; nt <= MAXT; ++nt) { + mtl::Pipeline pg = cx.pipeline(lib, "gemv_30_32_1_4_1", mtl::Constants().i(0, nt)), + pc = cx.pipeline(lib, "gemvc_30_32_1_4", mtl::Constants().i(0, nt)); + const size_t s16 = (size_t)4 * 4 * MAXT * N; +@@ -75,7 +359,7 @@ int main() { + w2 = cx.alloc(4 * MAXT * 10), pl1 = cx.alloc(4 * P_INTS), pl2 = cx.alloc(4 * P_INTS), cnt = cx.alloc(16); + ref::random_blocks(rty, Wr.as(), NE, N, 9, 0.05f); + mtl::Pipeline pr = cx.pipeline(lib, "moe_route"); +- for (int nt = 1; nt <= 4; ++nt) { ++ for (int nt = 1; nt <= MAXT; ++nt) { + mtl::Pipeline pg = cx.pipeline(lib, "gemv_" + std::to_string(rty) + "_32_1_4_1", mtl::Constants().i(0, nt)), + pf = cx.pipeline(lib, "router_route_" + std::to_string(rty), mtl::Constants().i(0, nt)); + const size_t s16 = (size_t)4 * 4 * MAXT * N; +@@ -131,12 +415,14 @@ int main() { + for (int i = 0; i < MAXT * HC; ++i) inj.as()[i] = r.sym() * 3; + for (int i = 0; i < HC * N; ++i) wn.as()[i] = 0.5f + r.uni(); + for (int i = 0; i < MAXT * TOPK; ++i) plan.as()[P_TOKENT + i] = (i * 7) % (MAXT * TOPK); +- mtl::Pipeline pcm = cx.pipeline(lib, "moe_combine_sh"), pn = cx.pipeline(lib, "gr_norm"), pnc = cx.pipeline(lib, "gr_norm_comb"); +- for (int nt = 1; nt <= 4; ++nt) { ++ mtl::Pipeline pcm = cx.pipeline(lib, "moe_combine_sh"), pn = cx.pipeline(lib, "gr_norm"), pnc = cx.pipeline(lib, "gr_norm_comb"), ++ pnc4 = cx.pipeline(lib, "gr_norm_comb_4s"); ++ for (int nt = 1; nt <= MAXT; ++nt) ++ for (uint32_t apply : {0u, 1u}) for (bool four : {false, true}) { + std::memcpy(R1.ptr, R0.ptr, R0.bytes); + std::memcpy(R2.ptr, R0.ptr, R0.bytes); + GlueArgs ga{(uint32_t)N, (uint32_t)nt, (uint32_t)N, TOPK}; +- GrArgs gr{(uint32_t)N, HC, 320, (uint32_t)nt, 1, 1e-6f}; ++ GrArgs gr{(uint32_t)N, HC, 320, (uint32_t)nt, apply, 1e-6f}; + mtl::Dispatch c; c.p = &pcm; + c.args = {mtl::Arg::B(ye), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::B(bo), mtl::Arg::V(ga)}; + c.grid[0] = (nt * N + 255) / 256; c.tg[0] = 256; +@@ -145,14 +431,15 @@ int main() { + n1.grid[0] = HC; n1.grid[1] = nt; n1.tg[0] = 256; + cx.run({c, n1}); + const uint32_t tk = TOPK; +- mtl::Dispatch n2; n2.p = &pnc; ++ mtl::Dispatch n2; n2.p = four ? &pnc4 : &pnc; + n2.args = {mtl::Arg::B(R2), mtl::Arg::B(ye), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(x2), mtl::Arg::V(gr), + mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::V(tk)}; +- n2.grid[0] = HC; n2.grid[1] = nt; n2.tg[0] = 256; ++ n2.grid[0] = four ? 1 : HC; n2.grid[1] = nt; n2.tg[0] = 256; + cx.run({n2}); + const bool ok = std::memcmp(R1.ptr, R2.ptr, 4 * (size_t)nt * HC * N) == 0 && std::memcmp(x1.ptr, x2.ptr, 4 * (size_t)nt * HC * N) == 0; + fails += !ok; +- std::printf("gr_norm_comb ntok %d: R', xn bitwise %s\n", nt, ok ? "PASS" : "FAIL"); ++ std::printf("gr_norm_comb%s ntok %d apply %u: R', xn bitwise %s\n", ++ four ? "_4s" : "", nt, apply, ok ? "PASS" : "FAIL"); + } + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); +diff --git a/metal/tests/gdn_parity.cpp b/metal/tests/gdn_parity.cpp +index 7288eb1d..09e74497 100644 +--- a/metal/tests/gdn_parity.cpp ++++ b/metal/tests/gdn_parity.cpp +@@ -4,6 +4,7 @@ + // gdn_parity parity for windows of 1..4 tokens: y, the committed state and conv history; a verify + // window leaves the state untouched; its first m outputs are bitwise a fold of m + // gdn_parity --bench GPU time per layer: 36 distinct layer states rotated (113 MB, beyond the SLC) ++// gdn_parity --ab ref.metal|ref_dir [--bench] bitwise parity (ntok 1..8, commit 0/1/partial, plane 0/49152) + paired A/B + #include "mtl_rt.h" + #include "ref_quant.h" + +@@ -12,12 +13,17 @@ + #include + #include + #include ++#include + #include + #include + + namespace { + constexpr int S = 128, HK = 16, HV = 48, C = (2 * HK + HV) * S, MAXT = 8; ++constexpr uint32_t PLANE = MAXT * HV * S; // 49152 ++constexpr size_t Y_ELEMS = 4 * (size_t)PLANE + PLANE / 16; + struct GdnArgs { uint32_t h_k, h_v, channels, commit; float eps; uint32_t plane; }; ++struct GlueArgs { uint32_t n, n_tok, n_embd, topk; }; ++constexpr int P_INTS = 16 + 128 + 128 + 4 + 128 * 3 + 20 + 128, P_TOKENT = 16 + 128 + 128 + 4 + 128 * 2; + double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } + struct Err { + double se = 0, sr = 0; +@@ -27,17 +33,30 @@ struct Err { + } // namespace + + int main(int argc, char** argv) { +- const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; ++ const bool ab = argc > 1 && std::string(argv[1]) == "--ab"; ++ if (ab && argc < 3) { ++ std::fprintf(stderr, "usage: gdn_parity --ab ref.metal|ref_dir [--bench|--no-bench]\n"); ++ return 2; ++ } ++ const bool no_bench = ab && argc > 3 && std::string(argv[3]) == "--no-bench"; ++ const bool bench = (argc > 1 && std::string(argv[1]) == "--bench") || (ab && !no_bench); ++ std::string ref_gdn = ab ? std::string(argv[2]) : ""; ++ std::string ref_glue = ""; ++ if (ab && std::filesystem::is_directory(ref_gdn)) { ++ if (std::filesystem::exists(ref_gdn + "/glue.metal")) ref_glue = ref_gdn + "/glue.metal"; ++ ref_gdn += "/gdn.metal"; ++ } + mtl::Context cx; + const int lib = cx.compile({"gdn.metal"}, "#define NO_GGML_COMMON 1\n"); ++ const int ref_lib = ab ? cx.compile({ref_gdn}, "#define NO_GGML_COMMON 1\n") : -1; + std::printf("device %s\n", cx.device_name().c_str()); + const int LAYERS = bench ? 36 : 1; + const size_t st_elems = (size_t)S * HV * S; + mtl::Buffer state = cx.alloc(4 * st_elems * LAYERS), st0 = cx.alloc(4 * st_elems), hist = cx.alloc(4 * C * 3 * LAYERS), + h0 = cx.alloc(4 * C * 3), qkv = cx.alloc(4 * C * MAXT), al = cx.alloc(4 * HV * MAXT), + be = cx.alloc(4 * HV * MAXT), z = cx.alloc(4 * HV * S * MAXT), cw = cx.alloc(4 * C * 4), +- dt = cx.alloc(4 * HV), sa = cx.alloc(4 * HV), gm = cx.alloc(4 * S), y = cx.alloc(4 * HV * S * MAXT), +- y2 = cx.alloc(4 * HV * S * MAXT), hist2 = cx.alloc(4 * C * 3 * LAYERS); ++ dt = cx.alloc(4 * HV), sa = cx.alloc(4 * HV), gm = cx.alloc(4 * S), y = cx.alloc(4 * Y_ELEMS), ++ y2 = cx.alloc(4 * Y_ELEMS), hist2 = cx.alloc(4 * C * 3 * LAYERS); + ref::Rng r(17); + for (size_t i = 0; i < st_elems; ++i) st0.as()[i] = r.sym() * 0.1f; + for (int i = 0; i < C * 3; ++i) h0.as()[i] = r.sym(); +@@ -51,8 +70,8 @@ int main(int argc, char** argv) { + std::memcpy(state.as() + l * st_elems, st0.ptr, 4 * st_elems); + std::memcpy(hist.as() + l * C * 3, h0.ptr, 4 * C * 3); + } +- auto disp = [&](const mtl::Pipeline& p, int layer, int commit, const mtl::Buffer& out) { +- GdnArgs a{HK, HV, C, (uint32_t)commit, 1e-6f, 0}; ++ auto disp = [&](const mtl::Pipeline& p, int layer, int commit, const mtl::Buffer& out, uint32_t plane = 0) { ++ GdnArgs a{HK, HV, C, (uint32_t)commit, 1e-6f, plane}; + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(state, 4 * st_elems * layer), mtl::Arg::B(hist, 4 * (size_t)C * 3 * layer), mtl::Arg::B(qkv), +@@ -78,6 +97,149 @@ int main(int argc, char** argv) { + return 1; + } + } ++ if (ab) { ++ const char* kn = std::getenv("GDN_IB8") ? "gdn_step8" : "gdn_step"; ++ int ab_fails = 0; ++ mtl::Buffer st_ref = cx.alloc(4 * st_elems), h_ref = cx.alloc(4 * C * 3); ++ for (int nt = 1; nt <= MAXT; ++nt) { ++ mtl::Pipeline pc = cx.pipeline(lib, kn, mtl::Constants().i(2, nt)); ++ mtl::Pipeline pr = cx.pipeline(ref_lib, kn, mtl::Constants().i(2, nt)); ++ for (uint32_t plane : {0u, PLANE}) { ++ std::vector commits = {0, 1}; ++ for (int k = 1; k < nt; ++k) commits.push_back(0x100 | k); ++ for (int commit : commits) { ++ const size_t y_bytes = 4 * (plane ? Y_ELEMS : (size_t)HV * S * nt); ++ std::memcpy(state.ptr, st0.ptr, 4 * st_elems); ++ std::memcpy(hist.ptr, h0.ptr, 4 * C * 3); ++ std::memset(hist2.ptr, 0xa5, 4 * C * 3); ++ std::memset(y.ptr, 0xa5, y_bytes); ++ cx.run({disp(pr, 0, commit, y, plane)}); ++ std::memcpy(st_ref.ptr, state.ptr, 4 * st_elems); ++ std::memcpy(h_ref.ptr, hist2.ptr, 4 * C * 3); ++ ++ std::memcpy(state.ptr, st0.ptr, 4 * st_elems); ++ std::memcpy(hist.ptr, h0.ptr, 4 * C * 3); ++ std::memset(hist2.ptr, 0xa5, 4 * C * 3); ++ std::memset(y2.ptr, 0xa5, y_bytes); ++ cx.run({disp(pc, 0, commit, y2, plane)}); ++ ++ const bool y_ok = std::memcmp(y.ptr, y2.ptr, y_bytes) == 0; ++ const bool s_ok = std::memcmp(state.ptr, st_ref.ptr, 4 * st_elems) == 0; ++ const bool h_ok = std::memcmp(hist2.ptr, h_ref.ptr, 4 * C * 3) == 0; ++ const bool ok = y_ok && s_ok && h_ok; ++ ab_fails += !ok; ++ std::printf("ab ntok %d commit 0x%x plane %u: y %s state %s hist %s %s\n", ++ nt, commit, plane, y_ok ? "ok" : "DIFF", s_ok ? "ok" : "DIFF", h_ok ? "ok" : "DIFF", ++ ok ? "PASS" : "FAIL"); ++ } ++ } ++ } ++ std::printf("ab bitwise summary: %d fail(s) %s\n", ab_fails, ab_fails == 0 ? "ALL PASS" : "FAIL"); ++ if (ab_fails || !bench) return ab_fails ? 1 : 0; ++ for (uint32_t plane : {0u, PLANE}) { ++ for (int nt = 1; nt <= MAXT; ++nt) { ++ mtl::Pipeline pc = cx.pipeline(lib, kn, mtl::Constants().i(2, nt)); ++ mtl::Pipeline pr = cx.pipeline(ref_lib, kn, mtl::Constants().i(2, nt)); ++ for (int commit : {1, 0}) { ++ std::vector dsc, dsr; ++ for (int k = 0; k < 72; ++k) { ++ dsc.push_back(disp(pc, k % LAYERS, commit, y, plane)); ++ dsr.push_back(disp(pr, k % LAYERS, commit, y, plane)); ++ } ++ cx.warm(dsr, 0.015); ++ cx.warm(dsc, 0.015); ++ std::vector bt, ct, rt; ++ for (int i = 0; i < 7; ++i) { ++ double b, c; ++ if (i & 1) { c = cx.run(dsc) / 72; b = cx.run(dsr) / 72; } ++ else { b = cx.run(dsr) / 72; c = cx.run(dsc) / 72; } ++ bt.push_back(b * 1e6); ++ ct.push_back(c * 1e6); ++ rt.push_back(c / b); ++ } ++ std::printf("ab bench plane %5u ntok %d %-7s base %7.2f us cand %7.2f us ratio %.4f (x36: %.3f -> %.3f ms)\n", ++ plane, nt, commit ? "commit" : "verify", median(bt), median(ct), median(rt), ++ 36 * median(bt) * 1e-3, 36 * median(ct) * 1e-3); ++ } ++ } ++ } ++ if (!ref_glue.empty()) { ++ const int glib_c = cx.compile({"glue.metal"}, "#define NO_GGML_COMMON 1\n#define P_TOKENT 532\n"); ++ const int glib_r = cx.compile({ref_glue}, "#define NO_GGML_COMMON 1\n#define P_TOKENT 532\n"); ++ const mtl::Pipeline sw_c = cx.pipeline(glib_c, "swiglu_prep"), sw_r = cx.pipeline(glib_r, "swiglu_prep"); ++ const mtl::Pipeline mc_c = cx.pipeline(glib_c, "moe_combine_sh"), mc_r = cx.pipeline(glib_r, "moe_combine_sh"); ++ constexpr int N = 2560, NFF = 1024, TOPK = 10; ++ const size_t xp_elems = 4 * (size_t)MAXT * NFF + MAXT * NFF / 16; ++ mtl::Buffer gsh = cx.alloc(4 * MAXT * NFF), ush = cx.alloc(4 * MAXT * NFF), ++ xp1 = cx.alloc(4 * xp_elems), xp2 = cx.alloc(4 * xp_elems), ++ ye = cx.alloc(4 * (size_t)MAXT * TOPK * N), ysh = cx.alloc(4 * MAXT * N), ++ sgv = cx.alloc(4 * MAXT), plan = cx.alloc(4 * P_INTS), ++ bo1 = cx.alloc(4 * MAXT * N), bo2 = cx.alloc(4 * MAXT * N); ++ for (int i = 0; i < MAXT * NFF; ++i) { gsh.as()[i] = r.sym() * 2.0f; ush.as()[i] = r.sym() * 2.0f; } ++ for (size_t i = 0; i < (size_t)MAXT * TOPK * N; ++i) ye.as()[i] = r.sym(); ++ for (int i = 0; i < MAXT * N; ++i) ysh.as()[i] = r.sym(); ++ for (int i = 0; i < MAXT; ++i) sgv.as()[i] = r.sym() * 3.0f; ++ for (int i = 0; i < MAXT * TOPK; ++i) plan.as()[P_TOKENT + i] = (i * 7) % (MAXT * TOPK); ++ int glue_fails = 0; ++ for (int nt = 1; nt <= MAXT; ++nt) { ++ GlueArgs ga_sw{NFF, MAXT, N, TOPK}; ++ mtl::Dispatch dsw_r{&sw_r, {mtl::Arg::B(gsh), mtl::Arg::B(ush), mtl::Arg::B(xp1), mtl::Arg::V(ga_sw)}, {NFF / 16, (uint32_t)nt, 1}, {16, 1, 1}}; ++ mtl::Dispatch dsw_c{&sw_c, {mtl::Arg::B(gsh), mtl::Arg::B(ush), mtl::Arg::B(xp2), mtl::Arg::V(ga_sw)}, {NFF / 16, (uint32_t)nt, 1}, {16, 1, 1}}; ++ std::memset(xp1.ptr, 0xa5, xp1.bytes); std::memset(xp2.ptr, 0xa5, xp2.bytes); ++ cx.run({dsw_r, dsw_c}); ++ const bool sw_ok = std::memcmp(xp1.ptr, xp2.ptr, xp1.bytes) == 0; ++ ++ for (uint32_t tk : {10u, 8u}) { ++ GlueArgs ga_mc{N, (uint32_t)nt, N, tk}; ++ const uint32_t gx = (nt * N + 255) / 256; ++ mtl::Dispatch dmc_r{&mc_r, {mtl::Arg::B(ye), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::B(bo1), mtl::Arg::V(ga_mc)}, {gx, 1, 1}, {256, 1, 1}}; ++ mtl::Dispatch dmc_c{&mc_c, {mtl::Arg::B(ye), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::B(bo2), mtl::Arg::V(ga_mc)}, {gx, 1, 1}, {256, 1, 1}}; ++ std::memset(bo1.ptr, 0xa5, bo1.bytes); std::memset(bo2.ptr, 0xa5, bo2.bytes); ++ cx.run({dmc_r, dmc_c}); ++ const bool mc_ok = std::memcmp(bo1.ptr, bo2.ptr, 4 * (size_t)nt * N) == 0; ++ glue_fails += (!sw_ok || !mc_ok); ++ if (tk == 10u) ++ std::printf("glue ab ntok %d: swiglu_prep %s moe_combine_sh(topk=10) %s %s\n", ++ nt, sw_ok ? "ok" : "DIFF", mc_ok ? "ok" : "DIFF", (sw_ok && mc_ok) ? "PASS" : "FAIL"); ++ else if (!mc_ok) ++ std::printf("glue ab ntok %d: moe_combine_sh(topk=8) DIFF FAIL\n", nt); ++ } ++ } ++ std::printf("glue ab bitwise summary: %d fail(s) %s\n", glue_fails, glue_fails == 0 ? "ALL PASS" : "FAIL"); ++ if (glue_fails) return 1; ++ for (int nt : {1, 2, 4, 8}) { ++ GlueArgs ga_sw{NFF, MAXT, N, TOPK}; ++ mtl::Dispatch dsw_r{&sw_r, {mtl::Arg::B(gsh), mtl::Arg::B(ush), mtl::Arg::B(xp1), mtl::Arg::V(ga_sw)}, {NFF / 16, (uint32_t)nt, 1}, {16, 1, 1}}; ++ mtl::Dispatch dsw_c{&sw_c, {mtl::Arg::B(gsh), mtl::Arg::B(ush), mtl::Arg::B(xp2), mtl::Arg::V(ga_sw)}, {NFF / 16, (uint32_t)nt, 1}, {16, 1, 1}}; ++ cx.warm({dsw_r}, 0.01); cx.warm({dsw_c}, 0.01); ++ std::vector bt, ct, rt; ++ for (int i = 0; i < 7; ++i) { ++ double b, c; ++ if (i & 1) { c = cx.run({dsw_c}, 256, false, true) / 256; b = cx.run({dsw_r}, 256, false, true) / 256; } ++ else { b = cx.run({dsw_r}, 256, false, true) / 256; c = cx.run({dsw_c}, 256, false, true) / 256; } ++ bt.push_back(b * 1e6); ct.push_back(c * 1e6); rt.push_back(c / b); ++ } ++ std::printf("glue bench swiglu_prep ntok %d base %6.2f us cand %6.2f us ratio %.4f\n", ++ nt, median(bt), median(ct), median(rt)); ++ ++ GlueArgs ga_mc{N, (uint32_t)nt, N, TOPK}; ++ const uint32_t gx = (nt * N + 255) / 256; ++ mtl::Dispatch dmc_r{&mc_r, {mtl::Arg::B(ye), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::B(bo1), mtl::Arg::V(ga_mc)}, {gx, 1, 1}, {256, 1, 1}}; ++ mtl::Dispatch dmc_c{&mc_c, {mtl::Arg::B(ye), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::B(bo2), mtl::Arg::V(ga_mc)}, {gx, 1, 1}, {256, 1, 1}}; ++ cx.warm({dmc_r}, 0.01); cx.warm({dmc_c}, 0.01); ++ bt.clear(); ct.clear(); rt.clear(); ++ for (int i = 0; i < 7; ++i) { ++ double b, c; ++ if (i & 1) { c = cx.run({dmc_c}, 256, false, true) / 256; b = cx.run({dmc_r}, 256, false, true) / 256; } ++ else { b = cx.run({dmc_r}, 256, false, true) / 256; c = cx.run({dmc_c}, 256, false, true) / 256; } ++ bt.push_back(b * 1e6); ct.push_back(c * 1e6); rt.push_back(c / b); ++ } ++ std::printf("glue bench moe_combine_sh ntok %d base %6.2f us cand %6.2f us ratio %.4f\n", ++ nt, median(bt), median(ct), median(rt)); ++ } ++ } ++ return 0; ++ } + if (bench) { + for (int nt : {1, 2, 3, 4}) { + const char* kn = std::getenv("GDN_IB8") ? "gdn_step8" : "gdn_step"; +@@ -201,12 +363,36 @@ int main(int argc, char** argv) { + partial = partial && std::memcmp(state.ptr, sk.data(), 4 * st_elems) == 0 && + std::memcmp(hist2.ptr, hk.data(), 4 * C * 3) == 0 && std::memcmp(y2.ptr, yv.data(), 4 * HV * S * nt) == 0; + } +- const bool ok = untouched && same_y && prefix && partial && ey.l2() < 1e-5 && es.l2() < 1e-5 && eh.l2() == 0; ++ // plane = PLANE (49152): verify prepared planes and per-16 sums against raw yv ++ std::memcpy(state.ptr, st0.ptr, 4 * st_elems); ++ std::memcpy(hist.ptr, h0.ptr, 4 * C * 3); ++ std::memset(y2.ptr, 0, 4 * Y_ELEMS); ++ cx.run({disp(p, 0, 0, y2, PLANE)}); ++ bool plane_ok = std::memcmp(y.ptr, y2.ptr, 4 * HV * S * nt) == 0; ++ const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; ++ const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; ++ const float* yp = y2.as(); ++ for (int idx = 0; idx < HV * S * nt && plane_ok; ++idx) { ++ const float yv = yp[idx]; ++ const int jj = idx & 15; ++ if (yp[PLANE + idx] != yv * q2s[jj & 7]) plane_ok = false; ++ if (yp[2 * PLANE + idx] != yv * p4s[jj & 3]) plane_ok = false; ++ if (yp[3 * PLANE + idx] != ((jj & 1) ? yv * 0.00390625f : yv)) plane_ok = false; ++ if (jj == 0) { ++ float s0 = 0.0f, s1 = 0.0f; ++ for (int k = 0; k < 16; k += 2) { ++ s0 += yp[idx + k]; ++ s1 += yp[idx + k + 1]; ++ } ++ if (yp[4 * PLANE + idx / 16] != s0 + s1) plane_ok = false; ++ } ++ } ++ const bool ok = untouched && same_y && prefix && partial && plane_ok && ey.l2() < 1e-5 && es.l2() < 1e-5 && eh.l2() == 0; + fails += !ok; + std::printf("gdn parity ntok %d: y rel L2 %.2e, state %.2e, history %.1e; verify leaves state: %s, verify y == " +- "commit y: %s, first %d of a longer window bitwise: %s, partial commits bitwise: %s %s\n", nt, ey.l2(), es.l2(), +- eh.l2(), untouched ? "yes" : "NO", same_y ? "yes" : "NO", nt, prefix ? "yes" : "NO", partial ? "yes" : "NO", +- ok ? "PASS" : "FAIL"); ++ "commit y: %s, first %d of a longer window bitwise: %s, partial commits bitwise: %s, plane bitwise: %s %s\n", ++ nt, ey.l2(), es.l2(), eh.l2(), untouched ? "yes" : "NO", same_y ? "yes" : "NO", nt, ++ prefix ? "yes" : "NO", partial ? "yes" : "NO", plane_ok ? "yes" : "NO", ok ? "PASS" : "FAIL"); + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; +diff --git a/metal/tests/gdn_prompt_parity.cpp b/metal/tests/gdn_prompt_parity.cpp +index 24cb85bf..beaee258 100644 +--- a/metal/tests/gdn_prompt_parity.cpp ++++ b/metal/tests/gdn_prompt_parity.cpp +@@ -3,6 +3,7 @@ + // against the CPU reference by gdn_parity): y, the final state and the conv history. The scan reassociates + // o = S_new^T q as g S^T q + delta (k . q), so the match is to rounding, not bitwise. Then the speed for T = 2048. + // gdn_prompt_parity [--bench] ++// gdn_prompt_parity --ab baseline.metal T f32 [bench] + #include "mtl_rt.h" + #include "ref_quant.h" + +@@ -10,6 +11,7 @@ + #include + #include + #include ++#include + #include + #include + +@@ -18,14 +20,22 @@ constexpr int HK = 16, HV = 48, CH = 10240, S = 128; + struct GdnArgs { uint32_t h_k, h_v, channels, commit; float eps; uint32_t plane; }; + struct GdnPArgs { uint32_t n_tok, h_k, h_v, channels; float eps; uint32_t f32; }; + double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } ++bool same(const mtl::Buffer& a, const mtl::Buffer& b, size_t bytes) { return std::memcmp(a.ptr, b.ptr, bytes) == 0; } + } // namespace + + int main(int argc, char** argv) { ++ const bool ab = argc > 1 && std::string(argv[1]) == "--ab"; ++ if (ab && (argc < 5 || std::atoi(argv[3]) < 1 || ++ (std::string(argv[4]) != "0" && std::string(argv[4]) != "1"))) { ++ std::fprintf(stderr, "usage: gdn_prompt_parity --ab baseline.metal T f32 [bench]\n"); ++ return 2; ++ } + const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; + mtl::Context cx; +- const int lib = cx.compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); ++ const int lib = cx.compile(ab ? std::vector{"gdn_prompt.metal"} : ++ std::vector{"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + std::printf("device %s\n", cx.device_name().c_str()); +- const int TMAX = 2048; ++ const int TMAX = ab ? std::atoi(argv[3]) : 2048; + auto F = [&](size_t n) { return cx.alloc(4 * n); }; + mtl::Buffer qkv = F((size_t)TMAX * CH), al = F(TMAX * HV), be = F(TMAX * HV), z = F((size_t)TMAX * HV * S), cw = F(CH * 4), + dt = F(HV), sa = F(HV), gam = F(S), st0 = F(S * HV * S), stA = F(S * HV * S), stB = F(S * HV * S), +@@ -53,7 +63,7 @@ int main(int argc, char** argv) { + d1.args = {mtl::Arg::B(qkv), mtl::Arg::B(hB), mtl::Arg::B(cw), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(dt), + mtl::Arg::B(sa), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), + mtl::Arg::B(hA2), mtl::Arg::V(a)}; +- d1.grid[0] = HK + HV / 2; ++ d1.grid[0] = (HK + HV) / 8; + d1.grid[1] = T; + d1.tg[0] = 256; + ds.push_back(d1); +@@ -85,14 +95,150 @@ int main(int argc, char** argv) { + mtl::Dispatch d3; + d3.p = &post; + d3.args = {mtl::Arg::B(o), mtl::Arg::B(z), mtl::Arg::B(gam), mtl::Arg::B(yh), mtl::Arg::B(yB), mtl::Arg::V(a)}; +- d3.grid[0] = HV; ++ d3.grid[0] = HV / 4; + d3.grid[1] = T; + d3.tg[0] = 128; + ds.push_back(d3); + }; + mtl::Pipeline pre = cx.pipeline(lib, "gdn_pre"), post = cx.pipeline(lib, "gdn_post"); + int fails = 0; +- if (!bench) ++ if (ab) { ++ const int T = TMAX; ++ const uint32_t f32 = std::atoi(argv[4]); ++ const bool timing = argc > 5 && std::string(argv[5]) == "bench"; ++ const int ref_lib = cx.compile({argv[2]}, "#define NO_GGML_COMMON 1\n"); ++ mtl::Pipeline pre_ref = cx.pipeline(ref_lib, "gdn_pre"), post_ref = cx.pipeline(ref_lib, "gdn_post"), ++ scan_ref = cx.pipeline(ref_lib, "gdn_scanr_4", mtl::Constants().i(6, 0)), ++ scan = cx.pipeline(lib, "gdn_scanr_4", mtl::Constants().i(6, 0)); ++ std::copy(h0.as(), h0.as() + CH * 3, hB.as()); ++ std::copy(st0.as(), st0.as() + S * HV * S, stB.as()); ++ std::vector candidate, reference; ++ prompt_ds(T, pre, scan, post, -4, candidate); ++ GdnPArgs a{(uint32_t)T, HK, HV, CH, 1e-6f, f32}; ++ candidate[0].args.back() = mtl::Arg::V(a); ++ candidate[1].args.back() = mtl::Arg::V(a); ++ candidate[2].args.back() = mtl::Arg::V(a); ++ reference = candidate; ++ reference[0].p = &pre_ref; reference[1].p = &scan_ref; reference[2].p = &post_ref; ++ const mtl::Buffer* outputs[] = {&qn, &kn, &vc, &gb, &kq, &hA2, &o, &stB, &yh, &yB}; ++ const size_t sizes[] = {4 * (size_t)T * HK * S, 4 * (size_t)T * HK * S, 4 * (size_t)T * HV * S, ++ 8 * (size_t)T * HV, 4 * (size_t)T * HK, 4 * CH * 3, 4 * (size_t)T * HV * S, ++ 4 * S * HV * S, 2 * (size_t)T * HV * S, 4 * (size_t)T * HV * S}; ++ const char* names[] = {"qn", "kn", "vc", "gb", "kq", "history", "o", "state", "yh", "y"}; ++ std::vector> expected; ++ // Each stage uses identical inputs even if an earlier candidate stage fails. ++ for (int stage = 0, out = 0; stage < 3; ++stage) { ++ const int end = stage == 0 ? 6 : stage == 1 ? 8 : 10; ++ if (stage == 2) std::memset(yB.ptr, 0xa5, yB.bytes); ++ cx.run({reference[stage]}); ++ for (int j = out; j < end; ++j) { ++ const auto* bytes = outputs[j]->as(); ++ expected.emplace_back(bytes, bytes + sizes[j]); ++ } ++ if (stage == 1) std::copy(st0.as(), st0.as() + S * HV * S, stB.as()); ++ if (stage == 2) std::memset(yB.ptr, 0xa5, yB.bytes); ++ cx.run({candidate[stage]}); ++ for (int j = out; j < end; ++j) { ++ const bool ok = std::memcmp(outputs[j]->ptr, expected[j].data(), sizes[j]) == 0; ++ std::printf("T %d f32 %u %s %s\n", T, f32, names[j], ok ? "PASS" : "FAIL"); ++ fails += !ok; ++ std::memcpy(outputs[j]->ptr, expected[j].data(), sizes[j]); ++ } ++ out = end; ++ } ++ if (fails || !timing) return fails ? 1 : 0; ++ for (int stage = 0; stage < 3; ++stage) { ++ cx.warm({reference[stage], candidate[stage]}, 0.05); ++ cx.run({reference[stage]}, 36); ++ cx.run({candidate[stage]}, 36); ++ std::vector ratios, bt, ct; ++ for (int pair = 0; pair < 3; ++pair) { ++ double b, c; ++ if (pair & 1) { c = cx.run({candidate[stage]}, 36); b = cx.run({reference[stage]}, 36); } ++ else { b = cx.run({reference[stage]}, 36); c = cx.run({candidate[stage]}, 36); } ++ bt.push_back(b * 1e3); ct.push_back(c * 1e3); ratios.push_back(c / b); ++ std::printf("pair %d stage %d baseline %.3f candidate %.3f ms ratio %.4f\n", pair, stage, b * 1e3, c * 1e3, c / b); ++ } ++ std::printf("T %d f32 %u stage %d 36 layers baseline %.3f candidate %.3f ms ratio %.4f\n", ++ T, f32, stage, median(bt), median(ct), median(ratios)); ++ } ++ { ++ cx.warm({reference[0], reference[1], reference[2], candidate[0], candidate[1], candidate[2]}, 0.05); ++ cx.run(reference, 36); ++ cx.run(candidate, 36); ++ std::vector ratios, bt, ct; ++ for (int pair = 0; pair < 3; ++pair) { ++ double b, c; ++ if (pair & 1) { c = cx.run(candidate, 36); b = cx.run(reference, 36); } ++ else { b = cx.run(reference, 36); c = cx.run(candidate, 36); } ++ bt.push_back(b * 1e3); ct.push_back(c * 1e3); ratios.push_back(c / b); ++ std::printf("pair %d total baseline %.3f candidate %.3f ms ratio %.4f\n", pair, b * 1e3, c * 1e3, c / b); ++ } ++ std::printf("T %d f32 %u total 36 layers baseline %.3f candidate %.3f ms ratio %.4f\n", ++ T, f32, median(bt), median(ct), median(ratios)); ++ } ++ return 0; ++ } ++ if (!bench) { ++ // Bitwise parity: 1-simdgroup-per-head packed grid ((8, T) pre, (12, T) post) vs legacy grid ((40, T) pre, (48, T) post) ++ mtl::Buffer qn2 = F((size_t)TMAX * HK * S), kn2 = F((size_t)TMAX * HK * S), vc2 = F((size_t)TMAX * HV * S), ++ gb2 = F(2 * TMAX * HV), kq2 = F(TMAX * HK), h2 = F(CH * 3), yh2 = cx.alloc(2 * (size_t)TMAX * HV * S), ++ y2 = F((size_t)TMAX * HV * S); ++ for (int T : {1, 3, 37, 300}) { ++ GdnPArgs a{(uint32_t)T, HK, HV, CH, 1e-6f, 1u}; ++ std::copy(h0.as(), h0.as() + CH * 3, hB.as()); ++ mtl::Dispatch pa, pb, qa, qb; ++ pa.p = pb.p = ⪯ ++ pa.args = {mtl::Arg::B(qkv), mtl::Arg::B(hB), mtl::Arg::B(cw), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(dt), ++ mtl::Arg::B(sa), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), ++ mtl::Arg::B(hA2), mtl::Arg::V(a)}; ++ pb.args = {mtl::Arg::B(qkv), mtl::Arg::B(hB), mtl::Arg::B(cw), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(dt), ++ mtl::Arg::B(sa), mtl::Arg::B(qn2), mtl::Arg::B(kn2), mtl::Arg::B(vc2), mtl::Arg::B(gb2), mtl::Arg::B(kq2), ++ mtl::Arg::B(h2), mtl::Arg::V(a)}; ++ pa.grid[0] = (HK + HV) / 8; pa.grid[1] = T; pa.tg[0] = 256; ++ pb.grid[0] = HK + HV / 2; pb.grid[1] = T; pb.tg[0] = 256; ++ cx.run({pa, pb}); ++ const bool pre_ok = same(qn, qn2, 4 * (size_t)T * HK * S) && same(kn, kn2, 4 * (size_t)T * HK * S) && ++ same(vc, vc2, 4 * (size_t)T * HV * S) && same(gb, gb2, 8 * (size_t)T * HV) && ++ same(kq, kq2, 4 * (size_t)T * HK) && same(hA2, h2, 4 * CH * 3); ++ qa.p = qb.p = &post; ++ qa.args = {mtl::Arg::B(vc), mtl::Arg::B(z), mtl::Arg::B(gam), mtl::Arg::B(yh), mtl::Arg::B(yB), mtl::Arg::V(a)}; ++ qb.args = {mtl::Arg::B(vc), mtl::Arg::B(z), mtl::Arg::B(gam), mtl::Arg::B(yh2), mtl::Arg::B(y2), mtl::Arg::V(a)}; ++ qa.grid[0] = HV / 4; qa.grid[1] = T; qa.tg[0] = 128; ++ qb.grid[0] = HV; qb.grid[1] = T; qb.tg[0] = 128; ++ cx.run({qa, qb}); ++ const bool post_ok = same(yh, yh2, 2 * (size_t)T * HV * S) && same(yB, y2, 4 * (size_t)T * HV * S); ++ fails += !pre_ok + !post_ok; ++ std::printf("T %3d, bitwise 1-sg vs legacy: gdn_pre %s, gdn_post %s\n", T, pre_ok ? "PASS" : "FAIL", ++ post_ok ? "PASS" : "FAIL"); ++ } ++ mtl::Pipeline scanr4 = cx.pipeline(lib, "gdn_scanr_4", mtl::Constants().i(6, 0)), ++ scanr4_16 = cx.pipeline(lib, "gdn_scanr_4_16", mtl::Constants().i(6, 0)); ++ for (int T : {1, 3, 16, 17, 32, 37, 300, 2048}) { ++ GdnPArgs a{(uint32_t)T, HK, HV, CH, 1e-6f, 1u}; ++ std::copy(h0.as(), h0.as() + CH * 3, hB.as()); ++ std::copy(st0.as(), st0.as() + S * HV * S, stA.as()); ++ std::copy(st0.as(), st0.as() + S * HV * S, stB.as()); ++ mtl::Dispatch dp, da_s, db_s; ++ dp.p = ⪯ ++ dp.args = {mtl::Arg::B(qkv), mtl::Arg::B(hB), mtl::Arg::B(cw), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(dt), ++ mtl::Arg::B(sa), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), ++ mtl::Arg::B(hA2), mtl::Arg::V(a)}; ++ dp.grid[0] = (HK + HV) / 8; dp.grid[1] = T; dp.tg[0] = 256; ++ da_s.p = &scanr4; ++ da_s.args = {mtl::Arg::B(stA), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), ++ mtl::Arg::B(o), mtl::Arg::V(a)}; ++ da_s.grid[0] = HV * 4; da_s.tg[0] = 128; ++ db_s.p = &scanr4_16; ++ db_s.args = {mtl::Arg::B(stB), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), ++ mtl::Arg::B(y2), mtl::Arg::V(a)}; ++ db_s.grid[0] = HV * 4; db_s.tg[0] = 128; ++ cx.run({dp, da_s, db_s}); ++ const bool o_ok = same(o, y2, 4 * (size_t)T * HV * S), st_ok = same(stA, stB, 4 * (size_t)S * HV * S); ++ fails += !o_ok + !st_ok; ++ std::printf("T %4d, bitwise single-pass vs TB=16 gdn_scanr_4: o %s, state %s\n", T, o_ok ? "PASS" : "FAIL", ++ st_ok ? "PASS" : "FAIL"); ++ } + for (int T : {1, 3, 37, 300}) + for (int cb : {1, 4, -2, -4, -8, -9, -100}) { + mtl::Constants pc; +@@ -135,6 +281,7 @@ int main(int argc, char** argv) { + std::printf("T %3d, scan CB %d: y rel %.2e (fp16 %.2e), state rel %.2e, history %s %s\n", T, cb, ry, + std::sqrt(ryh / syh), rs, rh == 0 ? "exact" : "DIFFERS", ok ? "PASS" : "FAIL"); + } ++ } + if (fails || !bench) { + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; +@@ -149,14 +296,14 @@ int main(int argc, char** argv) { + mtl::Arg::B(cw), mtl::Arg::B(dt), mtl::Arg::B(sa), mtl::Arg::B(gam), mtl::Arg::B(yA), mtl::Arg::V(da), mtl::Arg::B(hA2)}; + d.grid[0] = HV; + d.tg[0] = 512; +- auto timeit = [&](const std::vector& ds) { +- cx.warm(ds, 0.3); ++ auto timeit = [&](const std::vector& ds, int reps = 1) { ++ cx.warm(ds, 0.08); + std::vector t; +- for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds)); ++ for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds, reps)); + return median(t); + }; + std::printf("T %d: decode step over the chunk %.2f ms\n", T, timeit({d}) * 1e3); +- for (const char* nm : {"gdn_scanr_4", "chunked", "gdn_scanr_4", "chunked"}) { ++ for (const char* nm : {"gdn_scanr_4", "chunked"}) { + const bool ch = std::string(nm) == "chunked"; + const int cb = ch ? -100 : nm[9] == 'v' ? -8 : nm[8] == 'r' ? -(nm[10] - '0') : nm[9] - '0'; + mtl::Constants pc; +@@ -171,5 +318,30 @@ int main(int argc, char** argv) { + std::printf("T %d, %-13s: pre %.3f + scan %.3f + post %.3f ms = %.3f ms together\n", T, nm, tp * 1e3, ts * 1e3, + to * 1e3, all * 1e3); + } ++ mtl::Pipeline scanr4 = cx.pipeline(lib, "gdn_scanr_4", mtl::Constants().i(6, 0)), ++ scanr4_16 = cx.pipeline(lib, "gdn_scanr_4_16", mtl::Constants().i(6, 0)); ++ for (int Tb : {512, 1024, 2048}) { ++ GdnPArgs a0{(uint32_t)Tb, HK, HV, CH, 1e-6f, 0u}; ++ mtl::Dispatch dpre, dscan, dscan16, dpost; ++ dpre.p = ⪯ ++ dpre.args = {mtl::Arg::B(qkv), mtl::Arg::B(hB), mtl::Arg::B(cw), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(dt), ++ mtl::Arg::B(sa), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), ++ mtl::Arg::B(hA2), mtl::Arg::V(a0)}; ++ dpre.grid[0] = (HK + HV) / 8; dpre.grid[1] = Tb; dpre.tg[0] = 256; ++ dscan.p = &scanr4; ++ dscan.args = {mtl::Arg::B(stB), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), ++ mtl::Arg::B(o), mtl::Arg::V(a0)}; ++ dscan.grid[0] = HV * 4; dscan.tg[0] = 128; ++ dscan16.p = &scanr4_16; ++ dscan16.args = dscan.args; ++ dscan16.grid[0] = HV * 4; dscan16.tg[0] = 128; ++ dpost.p = &post; ++ dpost.args = {mtl::Arg::B(o), mtl::Arg::B(z), mtl::Arg::B(gam), mtl::Arg::B(yh), mtl::Arg::B(yB), mtl::Arg::V(a0)}; ++ dpost.grid[0] = HV / 4; dpost.grid[1] = Tb; dpost.tg[0] = 128; ++ const double tpre = timeit({dpre}, 36) * 1e3, tscan = timeit({dscan}, 36) * 1e3, ++ tscan16 = timeit({dscan16}, 36) * 1e3, tpost = timeit({dpost}, 36) * 1e3; ++ std::printf("T %4d, 36 GDN layers (f32=0): gdn_pre %.2f ms (%.3f), gdn_scanr_4 %.2f ms (%.3f vs TB=16 %.2f ms / %.3f), gdn_post %.2f ms (%.3f)\n", ++ Tb, tpre, tpre / 36.0, tscan, tscan / 36.0, tscan16, tscan16 / 36.0, tpost, tpost / 36.0); ++ } + return 0; + } +diff --git a/metal/tests/gemv_parity.cpp b/metal/tests/gemv_parity.cpp +index 5687527f..93d57603 100644 +--- a/metal/tests/gemv_parity.cpp ++++ b/metal/tests/gemv_parity.cpp +@@ -4,6 +4,13 @@ + // gemv_parity parity: every format, n_in 2560/640/10240, 1..8 tokens + // gemv_parity --bench [--pack Q2_0|IQ3_XXS] [--ntok N] [--all] speed on the pack's shapes (best tile per shape) + // gemv_parity --one TYPE NIN NOUT [--ntok N] sweep one shape ++// gemv_parity --ab BASE_DIR TYPE NIN NOUT NTOK MODE [L R S K] bitwise A/B, paired hot-weight timings ++// MODE: gemv | pair | cat | catk; GEMV_AB_PARITY_ONLY=1 skips timing; GEMV_AB_COLD=1 (gemv): the timed ++// dispatches walk copies of W (up to 512 MB), so the weights come from DRAM as in a decode ++// gemv_parity --tiles TYPE NIN NOUT NTOK L,R,S,K ... paired timings of several tiles (GEMV_AB_COLD=1; ++// GEMV_BESIDE=TYPE,NOUT,L,R,S,K: each tile as k + v beside a big GEMV) ++// gemv_parity --cat TYPE NIN NOUT NTOK L,R,S gemv_cat vs its two gemvs: bitwise, paired timings ++// gemv_parity --pred TYPE NIN NOUT NTOK L,R,S token 0 of two gemvs #NTOK vs #1 vs gemv_cat #1 + #include "gguf_lite.h" + #include "mtl_rt.h" + #include "ref_quant.h" +@@ -13,6 +20,7 @@ + #include + #include + #include ++#include + #include + #include + +@@ -62,7 +70,7 @@ mtl::Dispatch xprep_dispatch(const mtl::Pipeline& p, const mtl::Buffer& x, const + + int type_of(const std::string& s) { + for (int t : {ref::Q2_0, ref::Q4_0, ref::Q5_0, ref::Q8_0, ref::Q3_K, ref::Q4_K, ref::Q5_K, ref::Q6_K, ref::IQ4_NL, +- ref::IQ4_XS, ref::IQ3_S, ref::BF16}) ++ ref::IQ4_XS, ref::IQ3_S, ref::BF16, ref::F32}) + if (s == ref::type_name(t)) return t; + std::fprintf(stderr, "unknown type %s\n", s.c_str()); + std::exit(2); +@@ -70,9 +78,409 @@ int type_of(const std::string& s) { + + struct Shape { const char* name; int type, n_in, n_out, per_token; }; // per_token: calls per decoded token + ++// A short, paired comparison with a source snapshot, including byte-for-byte outputs. ++int ab(int argc, char** argv) { ++ if (argc != 8 && argc != 12) return 2; ++ const int type = std::string(argv[3]) == "Q2_0R" ? ref::Q2_0 : type_of(argv[3]); ++ const bool raw_q2 = std::string(argv[3]) == "Q2_0R"; ++ const int kernel_type = raw_q2 ? 142 : type; ++ const int ni = std::atoi(argv[4]), no = std::atoi(argv[5]), nt = std::atoi(argv[6]); ++ const std::string mode = argv[7]; ++ if (ni <= 0 || no <= 0 || nt < 1 || nt > 8 || ni % ref::block_elems(type) || ++ (mode != "gemv" && mode != "pair" && mode != "cat" && mode != "catk")) return 2; ++ Cfg c{32, 2, 4, 1}; ++ if (mode == "pair") c = {32, 1, 4, 1}; ++ if (mode == "catk") c = {32, 1, 4, 4}; ++ if (argc == 12) c = {std::atoi(argv[8]), std::atoi(argv[9]), std::atoi(argv[10]), std::atoi(argv[11])}; ++ const bool known_cfg = std::any_of(kCfgs.begin(), kCfgs.end(), [&](Cfg k) { ++ return k.L == c.L && k.R == c.R && k.S == c.S && k.K == c.K; ++ }); ++ if (!known_cfg || !cfg_ok(c, type, ni) || ++ (mode == "pair" && (raw_q2 || type != ref::Q2_0 || c.L != 32 || c.R != 1 || c.S != 4 || c.K != 1)) || ++ (mode == "cat" && ((type != ref::Q8_0 && type != ref::BF16 && type != ref::F32) || c.K != 1 || ++ !((c.L == 32 && c.R == 1 && (c.S == 4 || c.S == 8)) || ++ (c.L == 16 && c.R == 1 && c.S == 4) || ++ (c.L == 32 && c.R == 2 && c.S == 4) || ++ (type == ref::Q8_0 && c.L == 8 && c.R == 1 && c.S == 8)))) || ++ (mode == "catk" && (type != ref::Q8_0 || c.L != 32 || ++ !((c.R == 1 && c.S == 4 && c.K == 4) || ++ (c.R == 2 && c.S == 4 && c.K == 2) || ++ (c.R == 1 && c.S == 8 && c.K == 8))))) return 2; ++ mtl::Context cx; ++ const std::string current = cx.kernel_dir; ++ cx.kernel_dir = argv[2]; ++ const int lb = cx.compile({"quant.metal", "units.metal", "gemv.metal"}); ++ cx.kernel_dir = current; ++ const int lv = cx.compile({"quant.metal", "units.metal", "gemv.metal"}); ++ const size_t rb = ref::row_bytes(type, ni), mb = rb * no; ++ const bool cold = std::getenv("GEMV_AB_COLD") && mode == "gemv"; ++ const int copies = cold ? (int)std::max(4, std::min(64, (512ull << 20) / mb)) : 1; ++ auto w = cx.alloc(mb * copies), w2 = cx.alloc(mb), x = cx.alloc(4ull * ni * nt); ++ auto xp = cx.alloc(4ull * (4 * ni * nt + ni / 16 * nt)); ++ auto yb = cx.alloc(4ull * no * nt), yv = cx.alloc(4ull * no * nt); ++ auto zb = cx.alloc(4ull * no * nt), zv = cx.alloc(4ull * no * nt); ++ ref::random_blocks(type, w.as(), no, ni, 99); ++ ref::random_blocks(type, w2.as(), no, ni, 100); ++ for (int k = 1; k < copies; ++k) std::memcpy(w.as() + k * mb, w.ptr, mb); ++ ref::Rng rng(5); ++ for (int i = 0; i < ni * nt; ++i) x.as()[i] = rng.sym(); ++ auto prep = cx.pipeline(lb, "xprep"); ++ cx.run({xprep_dispatch(prep, x, xp, ni, nt)}); ++ const size_t plane = 4ull * ni * nt; ++ std::string name = kname(kernel_type, c); ++ if (mode == "pair") name = "gemv_pair_" + std::to_string(kernel_type) + "_" + std::to_string(kernel_type); ++ if (mode == "cat" || mode == "catk") { ++ name = (mode == "cat" ? "gemvc_" : "gemvck_") + std::to_string(kernel_type) + "_" + ++ std::to_string(c.L) + "_" + std::to_string(c.R) + "_" + std::to_string(c.S); ++ if (mode == "catk") name += "_" + std::to_string(c.K); ++ } ++ auto pb = cx.pipeline(lb, name, mtl::Constants().i(0, nt)); ++ auto pv = cx.pipeline(lv, name, mtl::Constants().i(0, nt)); ++ GemvArgs a{(uint32_t)ni, (uint32_t)(mode == "cat" || mode == "catk" ? no * 2 : no), ++ (uint32_t)nt, (uint32_t)rb, (uint32_t)ni, (uint32_t)no}; ++ auto dispatch = [&](const mtl::Pipeline& p, const mtl::Buffer& y, const mtl::Buffer& z, int copy = 0) { ++ mtl::Dispatch d; ++ d.p = &p; ++ const size_t off = plane * (raw_q2 ? 0 : xkind(type)); ++ d.args = {mtl::Arg::B(w, copy * mb), mtl::Arg::B(xp, off), mtl::Arg::B(xp, plane * 4), ++ mtl::Arg::B(y), mtl::Arg::V(a)}; ++ if (mode == "pair") ++ d.args = {mtl::Arg::B(w), mtl::Arg::B(xp, off), mtl::Arg::B(w2), mtl::Arg::B(xp, off), ++ mtl::Arg::B(xp, plane * 4), mtl::Arg::B(y), mtl::Arg::V(a), mtl::Arg::V((uint32_t)rb)}; ++ if (mode == "cat" || mode == "catk") { ++ d.args.push_back(mtl::Arg::B(w2)); d.args.push_back(mtl::Arg::B(z)); ++ d.args.push_back(mtl::Arg::V((uint32_t)no)); ++ } ++ const int rows = mode == "pair" ? 4 : rows_per_tg(c); ++ d.grid[0] = (a.n_out + rows - 1) / rows; ++ d.tg[0] = c.S * 32; ++ return d; ++ }; ++ auto b = dispatch(pb, yb, zb), v = dispatch(pv, yv, zv); ++ cx.run({b}); cx.run({v}); ++ const bool exact = !std::memcmp(yb.ptr, yv.ptr, yb.bytes) && !std::memcmp(zb.ptr, zv.ptr, zb.bytes); ++ std::printf("parity %s %s %d %d nt %d tile %d,%d,%d,%d %s\n", mode.c_str(), argv[3], ni, no, nt, ++ c.L, c.R, c.S, c.K, exact ? "bitwise" : "FAIL"); ++ if (!exact) return 1; ++ if (std::getenv("GEMV_AB_PARITY_ONLY")) return 0; ++ cx.warm({b, v}, 0.15); ++ std::vector bs, vs; ++ for (int k = 0; k < copies; ++k) { bs.push_back(dispatch(pb, yb, zb, k)); vs.push_back(dispatch(pv, yv, zv, k)); } ++ const int reps = cold ? std::max(1, 64 / copies) : 64, per = reps * copies; ++ for (int pair = 0; pair < 15; ++pair) { ++ double tb, tv; ++ if (pair % 2) { tv = cx.run(vs, reps); tb = cx.run(bs, reps); } ++ else { tb = cx.run(bs, reps); tv = cx.run(vs, reps); } ++ std::printf("ab %d %.6f %.6f\n", pair, tb * 1e6 / per, tv * 1e6 / per); ++ } ++ return 0; ++} ++ ++// Paired tile timings on one shape: every listed tile in each round (order rotated), the weights from DRAM ++// (GEMV_AB_COLD=1: the dispatches walk copies of W, up to 512 MB) or hot. Each dispatch behind a barrier (a decode's ++// GEMV waits for its x). Prints each tile's median us and its rel L2 against double. ++int tiles(int argc, char** argv) { ++ if (argc < 7) return 2; ++ const bool raw_q2 = std::string(argv[2]) == "Q2_0R"; ++ const int type = raw_q2 ? ref::Q2_0 : type_of(argv[2]), kernel_type = raw_q2 ? 142 : type; ++ const int ni = std::atoi(argv[3]), no = std::atoi(argv[4]), nt = std::atoi(argv[5]); ++ if (ni <= 0 || no <= 0 || nt < 1 || nt > 8 || ni % ref::block_elems(type)) return 2; ++ std::vector cs; ++ for (int i = 6; i < argc; ++i) { ++ Cfg c{}; ++ if (std::sscanf(argv[i], "%d,%d,%d,%d", &c.L, &c.R, &c.S, &c.K) != 4 || !cfg_ok(c, type, ni)) return 2; ++ cs.push_back(c); ++ } ++ mtl::Context cx; ++ const int lib = cx.compile({"quant.metal", "units.metal", "gemv.metal"}); ++ const size_t rb = ref::row_bytes(type, ni), mb = rb * no; ++ const bool cold = std::getenv("GEMV_AB_COLD") != nullptr; ++ const int copies = cold ? (int)std::max(4, std::min(256, (512ull << 20) / mb)) : 1; ++ auto w = cx.alloc(mb * copies), x = cx.alloc(4ull * ni * nt), y = cx.alloc(4ull * no * nt); ++ auto xp = cx.alloc(4ull * (4 * ni * nt + ni / 16 * nt)); ++ ref::random_blocks(type, w.as(), no, ni, 99); ++ for (int k = 1; k < copies; ++k) std::memcpy(w.as() + k * mb, w.ptr, mb); ++ ref::Rng rng(5); ++ for (int i = 0; i < ni * nt; ++i) x.as()[i] = rng.sym(); ++ auto prep = cx.pipeline(lib, "xprep"); ++ cx.run({xprep_dispatch(prep, x, xp, ni, nt)}); ++ if (raw_q2) std::memcpy(xp.ptr, x.ptr, 4ull * ni * nt); // Q2_0R reads raw x ++ const size_t plane = 4ull * ni * nt; ++ std::vector wf((size_t)ni * no); ++ ref::dequant(type, w.ptr, wf.data(), wf.size()); ++ GemvArgs a{(uint32_t)ni, (uint32_t)no, (uint32_t)nt, (uint32_t)rb, (uint32_t)ni, (uint32_t)no}; ++ std::vector ps; ++ for (Cfg c : cs) ps.push_back(cx.pipeline(lib, kname(kernel_type, c), mtl::Constants().i(0, nt))); ++ std::vector> runs(cs.size()); ++ for (size_t i = 0; i < cs.size(); ++i) ++ for (int k = 0; k < copies; ++k) { ++ mtl::Dispatch d; ++ d.p = &ps[i]; ++ d.args = {mtl::Arg::B(w, k * mb), mtl::Arg::B(xp, raw_q2 ? 0 : plane * xkind(type)), mtl::Arg::B(xp, plane * 4), ++ mtl::Arg::B(y), mtl::Arg::V(a)}; ++ d.grid[0] = (no + rows_per_tg(cs[i]) - 1) / rows_per_tg(cs[i]); ++ d.tg[0] = cs[i].S * 32; ++ runs[i].push_back(d); ++ } ++ std::vector rel(cs.size()); ++ for (size_t i = 0; i < cs.size(); ++i) { ++ cx.run({runs[i][0]}); ++ double se = 0, sr = 0; ++ for (int t = 0; t < nt; ++t) ++ for (int o = 0; o < no; ++o) { ++ double r = 0; ++ for (int k = 0; k < ni; ++k) r += (double)wf[(size_t)o * ni + k] * x.as()[t * ni + k]; ++ const double e = y.as()[t * no + o] - r; ++ se += e * e; ++ sr += r * r; ++ } ++ rel[i] = std::sqrt(se / sr); ++ } ++ // GEMV_BESIDE=TYPE,NOUT,L,R,S,K: each timed dispatch runs twice (as the k and v beside each other) beside one ++ // GEMV of that shape over the same n_in (as attn_q), all behind one barrier: the tile's cost inside the group ++ mtl::Buffer bw, by; ++ mtl::Pipeline bp; ++ GemvArgs ba{}; ++ if (const char* bs = std::getenv("GEMV_BESIDE")) { ++ char bt[16] = {}; ++ int bno = 0; ++ Cfg bc{}; ++ if (std::sscanf(bs, "%15[^,],%d,%d,%d,%d,%d", bt, &bno, &bc.L, &bc.R, &bc.S, &bc.K) != 6) return 2; ++ const int btype = type_of(bt); ++ const size_t bmb = ref::row_bytes(btype, ni) * bno; ++ const int bcopies = cold ? (int)std::max(2, std::min(64, (512ull << 20) / bmb)) : 1; ++ bw = cx.alloc(bmb * bcopies); ++ by = cx.alloc(4ull * bno * nt); ++ ref::random_blocks(btype, bw.as(), bno, ni, 7); ++ for (int k = 1; k < bcopies; ++k) std::memcpy(bw.as() + k * bmb, bw.ptr, bmb); ++ bp = cx.pipeline(lib, kname(btype, bc), mtl::Constants().i(0, nt)); ++ ba = {(uint32_t)ni, (uint32_t)bno, (uint32_t)nt, (uint32_t)ref::row_bytes(btype, ni), (uint32_t)ni, (uint32_t)bno}; ++ for (auto& r : runs) { ++ std::vector g; ++ for (int k = 0; k < copies; ++k) { ++ mtl::Dispatch d; ++ d.p = &bp; ++ d.args = {mtl::Arg::B(bw, (size_t)(k % bcopies) * bmb), mtl::Arg::B(xp, plane * xkind(btype)), ++ mtl::Arg::B(xp, plane * 4), mtl::Arg::B(by), mtl::Arg::V(ba)}; ++ d.grid[0] = (bno + rows_per_tg(bc) - 1) / rows_per_tg(bc); ++ d.tg[0] = bc.S * 32; ++ g.push_back(d); ++ mtl::Dispatch t1 = r[k], t2 = r[(k + copies / 2) % copies]; ++ t1.barrier = t2.barrier = false; ++ g.push_back(t1); ++ g.push_back(t2); ++ } ++ r = g; ++ } ++ } ++ const int reps = cold ? std::max(1, 256 / copies) : 256; ++ std::vector all; ++ for (auto& r : runs) all.insert(all.end(), r.begin(), r.end()); ++ cx.warm(all, 0.2); ++ std::vector> t(cs.size()); ++ for (int round = 0; round < 15; ++round) ++ for (size_t j = 0; j < cs.size(); ++j) { ++ const size_t i = (j + round) % cs.size(); ++ std::vector ds; ++ for (int r = 0; r < reps; ++r) ds.insert(ds.end(), runs[i].begin(), runs[i].end()); ++ t[i].push_back(cx.run_graph(ds) * 1e6 / (reps * copies)); ++ } ++ for (size_t i = 0; i < cs.size(); ++i) ++ std::printf("tile %s %d %d nt %d %2d,%d,%d,%d %s: %8.3f us rel L2 %.2e\n", argv[2], ni, no, nt, cs[i].L, cs[i].R, ++ cs[i].S, cs[i].K, cold ? "cold" : "hot", median(t[i]), rel[i]); ++ return 0; ++} ++ ++// gemv_cat against the two gemv dispatches it replaces, at one tile (KS = 1): bitwise, then paired timings with each ++// pair behind a barrier as in a decode (the two gemvs beside each other, no barrier between). GEMV_AB_COLD=1: the ++// dispatches walk copies of both matrices (up to 512 MB). ++int cat2(int argc, char** argv) { ++ if (argc != 7) return 2; ++ const int type = type_of(argv[2]), ni = std::atoi(argv[3]), no = std::atoi(argv[4]), nt = std::atoi(argv[5]); ++ Cfg c{}; ++ if (std::sscanf(argv[6], "%d,%d,%d", &c.L, &c.R, &c.S) != 3) return 2; ++ c.K = 1; ++ if (ni <= 0 || no <= 0 || nt < 1 || nt > 8 || ni % ref::block_elems(type) || !cfg_ok(c, type, ni)) return 2; ++ mtl::Context cx; ++ const int lib = cx.compile({"quant.metal", "units.metal", "gemv.metal"}); ++ const size_t rb = ref::row_bytes(type, ni), mb = rb * no; ++ const bool cold = std::getenv("GEMV_AB_COLD") != nullptr; ++ const int copies = cold ? (int)std::max(4, std::min(128, (256ull << 20) / mb)) : 1; ++ auto w = cx.alloc(mb * copies), w2 = cx.alloc(mb * copies), x = cx.alloc(4ull * ni * nt); ++ auto xp = cx.alloc(4ull * (4 * ni * nt + ni / 16 * nt)); ++ auto ya = cx.alloc(4ull * no * nt), yb = cx.alloc(4ull * no * nt), za = cx.alloc(4ull * no * nt), ++ zb = cx.alloc(4ull * no * nt); ++ ref::random_blocks(type, w.as(), no, ni, 99); ++ ref::random_blocks(type, w2.as(), no, ni, 100); ++ for (int k = 1; k < copies; ++k) { ++ std::memcpy(w.as() + k * mb, w.ptr, mb); ++ std::memcpy(w2.as() + k * mb, w2.ptr, mb); ++ } ++ ref::Rng rng(5); ++ for (int i = 0; i < ni * nt; ++i) x.as()[i] = rng.sym(); ++ auto prep = cx.pipeline(lib, "xprep"); ++ cx.run({xprep_dispatch(prep, x, xp, ni, nt)}); ++ const size_t plane = 4ull * ni * nt; ++ const std::string cn = "gemvc_" + std::to_string(type) + "_" + std::to_string(c.L) + "_" + std::to_string(c.R) + "_" + ++ std::to_string(c.S); ++ auto pg = cx.pipeline(lib, kname(type, c), mtl::Constants().i(0, nt)); ++ auto pc = cx.pipeline(lib, cn, mtl::Constants().i(0, nt)); ++ const int rows = rows_per_tg(c); ++ GemvArgs a{(uint32_t)ni, (uint32_t)no, (uint32_t)nt, (uint32_t)rb, (uint32_t)ni, (uint32_t)no}; ++ GemvArgs a2 = a; ++ a2.n_out = 2 * no; ++ auto gemv = [&](const mtl::Buffer& W, int k, const mtl::Buffer& y, bool barrier) { ++ mtl::Dispatch d; ++ d.p = &pg; ++ d.args = {mtl::Arg::B(W, k * mb), mtl::Arg::B(xp, plane * xkind(type)), mtl::Arg::B(xp, plane * 4), ++ mtl::Arg::B(y), mtl::Arg::V(a)}; ++ d.grid[0] = (no + rows - 1) / rows; ++ d.tg[0] = c.S * 32; ++ d.barrier = barrier; ++ return d; ++ }; ++ auto cat = [&](int k) { ++ mtl::Dispatch d; ++ d.p = &pc; ++ d.args = {mtl::Arg::B(w, k * mb), mtl::Arg::B(xp, plane * xkind(type)), mtl::Arg::B(xp, plane * 4), ++ mtl::Arg::B(yb), mtl::Arg::V(a2), mtl::Arg::B(w2, k * mb), mtl::Arg::B(zb), ++ mtl::Arg::V((uint32_t)no)}; ++ d.grid[0] = (2 * no + rows - 1) / rows; ++ d.tg[0] = c.S * 32; ++ return d; ++ }; ++ std::vector two, one; ++ for (int k = 0; k < copies; ++k) { ++ two.push_back(gemv(w, k, ya, true)); ++ two.push_back(gemv(w2, k, za, false)); ++ one.push_back(cat(k)); ++ } ++ cx.run_graph({two[0], two[1]}); ++ cx.run_graph({one[0]}); ++ const bool exact = !std::memcmp(ya.ptr, yb.ptr, ya.bytes) && !std::memcmp(za.ptr, zb.ptr, za.bytes); ++ std::printf("parity cat %s %d %d+%d nt %d tile %d,%d,%d: %s\n", argv[2], ni, no, no, nt, c.L, c.R, c.S, ++ exact ? "bitwise" : "FAIL"); ++ if (!exact) return 1; ++ if (std::getenv("GEMV_AB_PARITY_ONLY")) return 0; ++ const int reps = cold ? std::max(1, 256 / copies) : 256; ++ std::vector ts, os; ++ for (int r = 0; r < reps; ++r) { ++ ts.insert(ts.end(), two.begin(), two.end()); ++ os.insert(os.end(), one.begin(), one.end()); ++ } ++ cx.warm(ts, 0.1); ++ cx.warm(os, 0.1); ++ std::vector tt, to; ++ for (int pair = 0; pair < 15; ++pair) { ++ if (pair % 2) { to.push_back(cx.run_graph(os)); tt.push_back(cx.run_graph(ts)); } ++ else { tt.push_back(cx.run_graph(ts)); to.push_back(cx.run_graph(os)); } ++ } ++ const double per = 1e6 / (reps * copies), m2 = median(tt) * per, m1 = median(to) * per; ++ std::printf("cat %s %d %d+%d nt %d tile %d,%d,%d %s: 2 gemv %.3f us, gemv_cat %.3f us (%+.1f%%)\n", argv[2], ni, no, ++ no, nt, c.L, c.R, c.S, cold ? "cold" : "hot", m2, m1, (m1 / m2 - 1) * 100); ++ return 0; ++} ++ ++// the read-ahead's router guesses (window.cpp predL_ / predL2_, Q8_0 2560 -> 512 for layers l + 1 and l + 2) when only ++// token 0's guess is read (STRATA_METAL_PREFETCH_DRAFT=0): the two window-wide gemvs at NTOK = nt against the same two ++// at NTOK = 1 on the window's tile, and gemv_cat of both at NTOK = 1. Token 0 bitwise, then paired timings (each ++// variant behind one barrier, its two gemvs beside each other). GEMV_AB_COLD=1: the weights from DRAM. ++int pred(int argc, char** argv) { ++ if (argc != 7) return 2; ++ const int type = type_of(argv[2]), ni = std::atoi(argv[3]), no = std::atoi(argv[4]), nt = std::atoi(argv[5]); ++ Cfg c{}; ++ if (std::sscanf(argv[6], "%d,%d,%d", &c.L, &c.R, &c.S) != 3) return 2; ++ c.K = 1; ++ if (ni <= 0 || no <= 0 || nt < 1 || nt > 8 || ni % ref::block_elems(type) || !cfg_ok(c, type, ni)) return 2; ++ mtl::Context cx; ++ const int lib = cx.compile({"quant.metal", "units.metal", "gemv.metal"}); ++ const size_t rb = ref::row_bytes(type, ni), mb = rb * no; ++ const bool cold = std::getenv("GEMV_AB_COLD") != nullptr; ++ const int copies = cold ? (int)std::max(4, std::min(128, (256ull << 20) / mb)) : 1; ++ auto w = cx.alloc(mb * copies), w2 = cx.alloc(mb * copies), x = cx.alloc(4ull * ni * nt); ++ auto xp = cx.alloc(4ull * (4 * ni * nt + ni / 16 * nt)); ++ std::vector ys, zs; ++ for (int v = 0; v < 3; ++v) { ys.push_back(cx.alloc(4ull * no * nt)); zs.push_back(cx.alloc(4ull * no * nt)); } ++ ref::random_blocks(type, w.as(), no, ni, 99); ++ ref::random_blocks(type, w2.as(), no, ni, 100); ++ for (int k = 1; k < copies; ++k) { ++ std::memcpy(w.as() + k * mb, w.ptr, mb); ++ std::memcpy(w2.as() + k * mb, w2.ptr, mb); ++ } ++ ref::Rng rng(5); ++ for (int i = 0; i < ni * nt; ++i) x.as()[i] = rng.sym(); ++ auto prep = cx.pipeline(lib, "xprep"); ++ cx.run({xprep_dispatch(prep, x, xp, ni, nt)}); ++ const size_t plane = 4ull * ni * nt; ++ auto pw = cx.pipeline(lib, kname(type, c), mtl::Constants().i(0, nt)); ++ auto p1 = cx.pipeline(lib, kname(type, c), mtl::Constants().i(0, 1)); ++ auto pc = cx.pipeline(lib, "gemvc_" + std::to_string(type) + "_" + std::to_string(c.L) + "_" + std::to_string(c.R) + ++ "_" + std::to_string(c.S), mtl::Constants().i(0, 1)); ++ const int rows = rows_per_tg(c); ++ auto gemv = [&](const mtl::Pipeline& p, int n, const mtl::Buffer& W, int k, const mtl::Buffer& y, bool barrier) { ++ mtl::Dispatch d; ++ d.p = &p; ++ const GemvArgs a{(uint32_t)ni, (uint32_t)no, (uint32_t)n, (uint32_t)rb, (uint32_t)ni, (uint32_t)no}; ++ d.args = {mtl::Arg::B(W, k * mb), mtl::Arg::B(xp, plane * xkind(type)), mtl::Arg::B(xp, plane * 4), ++ mtl::Arg::B(y), mtl::Arg::V(a)}; ++ d.grid[0] = (no + rows - 1) / rows; ++ d.tg[0] = c.S * 32; ++ d.barrier = barrier; ++ return d; ++ }; ++ auto cat = [&](int k) { ++ mtl::Dispatch d; ++ d.p = &pc; ++ const GemvArgs a{(uint32_t)ni, (uint32_t)(2 * no), 1, (uint32_t)rb, (uint32_t)ni, (uint32_t)no}; ++ d.args = {mtl::Arg::B(w, k * mb), mtl::Arg::B(xp, plane * xkind(type)), mtl::Arg::B(xp, plane * 4), ++ mtl::Arg::B(ys[2]), mtl::Arg::V(a), mtl::Arg::B(w2, k * mb), mtl::Arg::B(zs[2]), ++ mtl::Arg::V((uint32_t)no)}; ++ d.grid[0] = (2 * no + rows - 1) / rows; ++ d.tg[0] = c.S * 32; ++ return d; ++ }; ++ std::vector> runs(3); ++ for (int k = 0; k < copies; ++k) { ++ runs[0].push_back(gemv(pw, nt, w, k, ys[0], true)); ++ runs[0].push_back(gemv(pw, nt, w2, k, zs[0], false)); ++ runs[1].push_back(gemv(p1, 1, w, k, ys[1], true)); ++ runs[1].push_back(gemv(p1, 1, w2, k, zs[1], false)); ++ runs[2].push_back(cat(k)); ++ } ++ for (int v = 0; v < 3; ++v) cx.run_graph(std::vector(runs[v].begin(), runs[v].begin() + (v < 2 ? 2 : 1))); ++ bool exact = true; ++ for (int v = 1; v < 3; ++v) ++ exact = exact && !std::memcmp(ys[0].ptr, ys[v].ptr, 4ull * no) && !std::memcmp(zs[0].ptr, zs[v].ptr, 4ull * no); ++ std::printf("parity pred %s %d %d+%d tile %d,%d,%d: token 0 of #%d == #1 == gemvc #1: %s\n", argv[2], ni, no, no, ++ c.L, c.R, c.S, nt, exact ? "bitwise" : "FAIL"); ++ if (!exact) return 1; ++ if (std::getenv("GEMV_AB_PARITY_ONLY")) return 0; ++ const int reps = cold ? std::max(1, 256 / copies) : 256; ++ std::vector> rep(3); ++ for (int v = 0; v < 3; ++v) { ++ for (int r = 0; r < reps; ++r) rep[v].insert(rep[v].end(), runs[v].begin(), runs[v].end()); ++ cx.warm(rep[v], 0.1); ++ } ++ std::vector> t(3); ++ for (int round = 0; round < 15; ++round) ++ for (int j = 0; j < 3; ++j) { ++ const int v = (j + round) % 3; ++ t[v].push_back(cx.run_graph(rep[v])); ++ } ++ const double per = 1e6 / (reps * copies), m0 = median(t[0]) * per, m1 = median(t[1]) * per, m2 = median(t[2]) * per; ++ std::printf("pred %s %d %d+%d tile %d,%d,%d %s: 2 gemv #%d %.3f us, 2 gemv #1 %.3f us (%+.1f%%), gemvc #1 %.3f us " ++ "(%+.1f%%)\n", argv[2], ni, no, no, c.L, c.R, c.S, cold ? "cold" : "hot", nt, m0, m1, (m1 / m0 - 1) * 100, ++ m2, (m2 / m0 - 1) * 100); ++ return 0; ++} ++ + } // namespace + + int main(int argc, char** argv) { ++ if (argc > 1 && std::string(argv[1]) == "--ab") return ab(argc, argv); ++ if (argc > 1 && std::string(argv[1]) == "--tiles") return tiles(argc, argv); ++ if (argc > 1 && std::string(argv[1]) == "--cat") return cat2(argc, argv); ++ if (argc > 1 && std::string(argv[1]) == "--pred") return pred(argc, argv); + bool bench = false, all = false; + int ntok = 1; + std::string gguf_path; +@@ -180,8 +588,9 @@ int main(int argc, char** argv) { + ref::dequant(t, W.ptr, wf.data(), wf.size()); + for (int nt : {1, 3, 8}) { + double worst = 0; +- int not_exact = 0; ++ int not_exact = 0, not_same = 0; + std::string worst_cfg; ++ std::map> by_lpr; // KS = 1: a row's sums depend only on LPR + for (Cfg c : kCfgs) { + if (!cfg_ok(c, t, n_in)) continue; + mtl::Pipeline p = cx.pipeline(lib, kname(t, c), mtl::Constants().i(0, nt)); +@@ -206,6 +615,11 @@ int main(int argc, char** argv) { + } + const double rel = std::sqrt(se / sr); + if (rel > worst) { worst = rel; worst_cfg = kname(t, c); } ++ if (c.K == 1) { ++ const float* yv = (const float*)y.ptr; ++ auto [it, first] = by_lpr.try_emplace(c.L, yv, yv + nt * n_out); ++ if (!first && std::memcmp(it->second.data(), yv, 4ull * nt * n_out) != 0) ++not_same; ++ } + if (nt > 1) { + // every token of a window must be bitwise what a one-token call gives + std::vector ywin((float*)y.ptr, (float*)y.ptr + nt * n_out); +@@ -223,12 +637,13 @@ int main(int argc, char** argv) { + } + } + } +- const bool ok = worst < 1e-5 && not_exact == 0; ++ const bool ok = worst < 1e-5 && not_exact == 0 && not_same == 0; + fails += !ok; + std::printf("parity %-7s n_in %5d ntok %d: worst rel L2 over %zu tiles %.2e (%s), window == one " +- "token: %s %s\n", ++ "token: %s, same LPR: %s %s\n", + ref::type_name(t), n_in, nt, kCfgs.size(), worst, worst_cfg.c_str(), +- nt == 1 ? "-" : not_exact ? "NO" : "bitwise", ok ? "PASS" : "FAIL"); ++ nt == 1 ? "-" : not_exact ? "NO" : "bitwise", not_same ? "NO" : "bitwise", ++ ok ? "PASS" : "FAIL"); + } + cx.free(W); + cx.free(x); +diff --git a/metal/tests/moe_mm_parity.cpp b/metal/tests/moe_mm_parity.cpp +index 6f51f1e7..8e06a774 100644 +--- a/metal/tests/moe_mm_parity.cpp ++++ b/metal/tests/moe_mm_parity.cpp +@@ -3,6 +3,8 @@ + // units' operands) - and against the exact float reference for the size of that rounding; then the speed. + // + // moe_mm_parity [--gu Q2_0] [--dn Q2_0] [--T 256] [--tile 32,64,64,4] parity (64 experts) ++// moe_mm_parity --gather-bench [T=2048] [experts=512] [width=2560] short plan/copy sweep ++// MM_BASELINE_KERNEL=/path/to/moe_mm.metal compares every plan buffer byte-for-byte + // moe_mm_parity --bench [--gu ..] [--dn ..] [--T 2048] [--tile ..] per-layer time, 512 experts + #include "coop_probe.h" + #include "mtl_rt.h" +@@ -28,30 +30,190 @@ int type_of(const std::string& s) { + std::exit(2); + } + float h16(float v) { return ref::f16_to_f32(ref::f32_to_f16(v)); } ++ ++int gather_bench(int T, int NE, int n) { ++ if (T <= 0 || NE < TOPK || NE > 1024 || NE % 32 || n <= 0 || n % 8) { ++ std::fprintf(stderr, "gather-bench needs T > 0, experts in [32, 1024] divisible by 32, and width divisible by 8\n"); ++ return 2; ++ } ++ mtl::Context cx; ++ const int lib = cx.compile({"moe_mm.metal"}, "#define NO_GGML_COMMON 1\n#define MM_GATHER_ONLY 1\n"); ++ mtl::Pipeline phist = cx.pipeline(lib, "mm_hist"), pscan = cx.pipeline(lib, "mm_scan"), pplace = cx.pipeline(lib, "mm_place"), ++ pgat = cx.pipeline(lib, "mm_gather_h"), pgat_f = cx.pipeline(lib, "mm_gather"); ++ const char* baseline = std::getenv("MM_BASELINE_KERNEL"); ++ const int bl = baseline ? cx.compile({baseline}, "#define NO_GGML_COMMON 1\n#define MM_GATHER_ONLY 1\n") : lib; ++ mtl::Pipeline bh = cx.pipeline(bl, "mm_hist"), bs = cx.pipeline(bl, "mm_scan"), bp = cx.pipeline(bl, "mm_place"), ++ bg = cx.pipeline(bl, "mm_gather_h"), bf = cx.pipeline(bl, "mm_gather"); ++ std::printf("device %s; mm_plan + mm_gather_h bitwise parity & benchmark\n", cx.device_name().c_str()); ++ const int TM = 16; ++ const int nent = T * TOPK, nblk = (nent + 1023) / 1024, maxrows = nent + NE * TM; ++ mtl::Buffer x = cx.alloc(4 * (size_t)T * n), xh = cx.alloc(2 * (size_t)T * n), ++ ids = cx.alloc(4 * nent), w = cx.alloc(4 * nent), ++ bcnt = cx.alloc(4 * (size_t)nblk * NE), boff = cx.alloc(4 * (size_t)nblk * NE), eoff = cx.alloc(4 * NE), ++ tiles = cx.alloc(8 * (size_t)(maxrows / TM + 1)), meta = cx.alloc(64), stok = cx.alloc(4 * (size_t)maxrows), ++ sw = cx.alloc(4 * (size_t)maxrows), pos = cx.alloc(4 * nent), ++ XS = cx.alloc(2 * (size_t)maxrows * n), XS_ref = cx.alloc(2 * (size_t)maxrows * n); ++ ref::Rng r(7); ++ for (size_t i = 0; i < (size_t)T * n; ++i) { ++ const float v = r.sym(); ++ x.as()[i] = v; ++ xh.as()[i] = ref::f32_to_f16(v); ++ } ++ for (int t = 0; t < T; ++t) { ++ std::set used; ++ float s = 0; ++ for (int k = 0; k < TOPK; ++k) { ++ int e; ++ do e = (int)(r.next() % NE); while (used.count(e)); ++ used.insert(e); ++ ids.as()[t * TOPK + k] = e; ++ s += (w.as()[t * TOPK + k] = 0.2f + r.uni()); ++ } ++ for (int k = 0; k < TOPK; ++k) w.as()[t * TOPK + k] /= s; ++ } ++ MmPlanArgs pa{(uint32_t)nent, TOPK, (uint32_t)NE, (uint32_t)TM, (uint32_t)nblk}; ++ std::vector ds_plan; ++ mtl::Dispatch h; h.p = &phist; h.args = {mtl::Arg::B(ids), mtl::Arg::B(bcnt), mtl::Arg::V(pa)}; h.grid[0] = nblk; h.tg[0] = 1024; ds_plan.push_back(h); ++ mtl::Dispatch sc; sc.p = &pscan; sc.args = {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta), mtl::Arg::V(pa), mtl::Arg::B(stok)}; sc.tg[0] = NE; ds_plan.push_back(sc); ++ mtl::Dispatch pl; pl.p = &pplace; pl.args = {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok), mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}; pl.grid[0] = nblk; pl.tg[0] = 1024; ds_plan.push_back(pl); ++ bool ok_bits = true; ++ if (baseline) { ++ auto bd = ds_plan; ++ bd[0].p = &bh; bd[1].p = &bs; bd[2].p = &bp; ++ cx.run(bd); ++ std::vector outputs{&bcnt, &boff, &eoff, &tiles, &meta, &stok, &sw, &pos}; ++ std::vector sizes{4u * nblk * NE, 4u * nblk * NE, 4u * NE, ++ 8u * (maxrows / TM + 1), 64, 4u * maxrows, 4u * maxrows, 4u * nent}; ++ std::vector> expected; ++ for (size_t i = 0; i < outputs.size(); ++i) { ++ const auto* bytes = outputs[i]->as(); ++ expected.emplace_back(bytes, bytes + sizes[i]); ++ std::memset(outputs[i]->ptr, 0, sizes[i]); ++ } ++ cx.run(ds_plan); ++ const char* names[]{"bcnt", "boff", "eoff", "tiles", "meta", "srt_tok", "srt_w", "pos"}; ++ for (size_t i = 0; i < outputs.size(); ++i) { ++ if (std::memcmp(outputs[i]->ptr, expected[i].data(), sizes[i]) != 0) { ++ std::fprintf(stderr, "%s memcmp failed\n", names[i]); ++ ok_bits = false; ++ } ++ } ++ cx.warm(bd, 0.02); ++ cx.warm(ds_plan, 0.02); ++ std::vector bt, ct; ++ for (int i = 0; i < 7; ++i) { ++ if (i & 1) { ++ ct.push_back(cx.run(ds_plan, 8, false, true)); ++ bt.push_back(cx.run(bd, 8, false, true)); ++ } else { ++ bt.push_back(cx.run(bd, 8, false, true)); ++ ct.push_back(cx.run(ds_plan, 8, false, true)); ++ } ++ } ++ std::printf(" plan baseline %.3f candidate %.3f us/L; all 8 buffers memcmp=%s\n", ++ median(bt) * 1e6 / 8, median(ct) * 1e6 / 8, ok_bits ? "PASS" : "FAIL"); ++ } else cx.run(ds_plan); ++ const uint32_t tot_rows = meta.as()[1], ne = n; ++ const size_t xs_bytes = 2 * (size_t)tot_rows * n; ++ bool ok_plan = true; ++ std::vector seen(tot_rows, 0); ++ for (int i = 0; i < nent; ++i) { ++ const int row = pos.as()[i]; ++ if (row < 0 || row >= (int)tot_rows || stok.as()[row] != i / TOPK) ok_plan = false; ++ else seen[row]++; ++ } ++ for (uint32_t r0 = 0; r0 < tot_rows; ++r0) ++ if (seen[r0] > 1 || (seen[r0] == 0 && stok.as()[r0] != -1)) ok_plan = false; ++ mtl::Dispatch d_ref; d_ref.p = &pgat_f; d_ref.args = {mtl::Arg::B(x), mtl::Arg::B(stok), mtl::Arg::B(XS_ref), mtl::Arg::V(ne)}; d_ref.grid[0] = tot_rows; d_ref.tg[0] = 256; ++ cx.run({d_ref}); ++ auto b_ref = d_ref; b_ref.p = &bf; b_ref.args[2] = mtl::Arg::B(XS); ++ cx.run({b_ref}); ++ bool ok_xs = std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) == 0; ++ for (uint32_t tg_sz : {32u, 64u, 80u, 128u, 160u, 256u}) { ++ std::memset(XS.ptr, 0x55, xs_bytes); ++ mtl::Dispatch d; d.p = &pgat; d.args = {mtl::Arg::B(xh), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; d.grid[0] = tot_rows; d.tg[0] = tg_sz; ++ cx.run({d}); ++ if (std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) != 0) ok_xs = false; ++ d.grid[0] = 1; d.indirect = &meta; d.ind_off = 8 * 4; ++ cx.run({d}); ++ if (std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) != 0) ok_xs = false; ++ auto b_gat = d; b_gat.p = &bg; ++ cx.warm({b_gat, d}, 0.02); ++ std::vector bt, ct; ++ for (int i = 0; i < 7; ++i) { ++ if (i & 1) { ++ ct.push_back(cx.run({d}, 8, false, true)); ++ bt.push_back(cx.run({b_gat}, 8, false, true)); ++ } else { ++ bt.push_back(cx.run({b_gat}, 8, false, true)); ++ ct.push_back(cx.run({d}, 8, false, true)); ++ } ++ } ++ std::printf(" T=%d NE=%d rows=%u tg=%u baseline %.3f candidate %.3f us/L XS_memcmp=%s\n", ++ T, NE, tot_rows, tg_sz, median(bt) * 1e6 / 8, median(ct) * 1e6 / 8, ok_xs ? "PASS" : "FAIL"); ++ } ++ auto combined = ds_plan; ++ mtl::Dispatch gather = d_ref; ++ gather.p = &pgat; ++ gather.args[0] = mtl::Arg::B(xh); ++ gather.args[2] = mtl::Arg::B(XS); ++ gather.indirect = &meta; gather.ind_off = 8 * 4; ++ combined.push_back(gather); ++ std::vector times; ++ for (int i = 0; i < 3; ++i) times.push_back(cx.run(combined, 8, false, true)); ++ ok_xs &= std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) == 0; ++ std::printf(" combined plan+gather tg256 %.3f us/L plan=%s XS_memcmp=%s\n", ++ median(times) * 1e6 / 8, ok_plan ? "PASS" : "FAIL", ok_xs ? "PASS" : "FAIL"); ++ if (!ok_plan || !ok_xs || !ok_bits) return 1; ++ return 0; ++} ++ + } // namespace + + int main(int argc, char** argv) { ++ if (argc > 1 && std::string(argv[1]) == "--gather-bench") { ++ return gather_bench(argc > 2 ? std::atoi(argv[2]) : 2048, ++ argc > 3 ? std::atoi(argv[3]) : 512, ++ argc > 4 ? std::atoi(argv[4]) : N_EMBD); ++ } + int gu = ref::Q2_0, dn = ref::Q2_0, T = 256, TM = 16, TN = 64, TK = 16, SG = 2, DTN = -1, DTK = -1, DSG = -1; + bool bench = false, v2 = false, v3 = false, v4 = false, v5 = false; + int MS = 1; ++ bool gu_only = false, dn_only = false; ++ const char* dump = nullptr, *ab_dir = nullptr; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + if (a == "--gu") gu = type_of(argv[++i]); + else if (a == "--dn") dn = type_of(argv[++i]); +- else if (a == "--T") T = std::atoi(argv[++i]); ++ else if (a == "--T" || a == "-T") T = std::atoi(argv[++i]); + else if (a == "--tile") std::sscanf(argv[++i], "%d,%d,%d,%d", &TM, &TN, &TK, &SG); +- else if (a == "--dtile") std::sscanf(argv[++i], "%d,%d,%d", &DTN, &DTK, &DSG); // the down kernel's TN,TK,SG ++ else if (a == "--dtile") { ++ int m, n, k, s; ++ const char* value = argv[++i]; ++ if (std::sscanf(value, "%d,%d,%d,%d", &m, &n, &k, &s) == 4) { ++ if (m != TM) { ++ std::fprintf(stderr, "down tile TM must match gate/up TM\n"); ++ return 2; ++ } ++ DTN = n; DTK = k; DSG = s; ++ } else std::sscanf(value, "%d,%d,%d", &DTN, &DTK, &DSG); ++ } + else if (a == "--bench") bench = true; ++ else if (a == "--gu-bench") bench = gu_only = true; // --bench timing the gate/up GEMM alone ++ else if (a == "--dn-bench") bench = dn_only = true; + else if (a == "--v2") v2 = true; + else if (a == "--v3") v2 = v3 = true; // v2's gather + the register-operand GEMMs (TN per simdgroup) + else if (a == "--v4") { v2 = v3 = v4 = true; MS = std::atoi(argv[++i]); } // mm_gu4 with MS 16-row subtiles + else if (a == "--v5") { v2 = v3 = true; v5 = true; } // mm_gu5 (weight blocks staged in threadgroup memory) ++ else if (a == "--dump") dump = argv[++i]; // HS and y to a file: bitwise comparison of two kernel versions ++ else if (a == "--ab") ab_dir = argv[++i]; + } ++ const bool isolated_bench = gu_only || dn_only; + if (bench && T == 256) T = 2048; + if (v4) TM = 16 * MS; // the plan's tile height; mm_gu4's name carries MS instead + mtl::Context cx; +- const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "moe_mm.metal"}); +- const int NE = bench ? 512 : 64, LAYERS = bench ? 4 : 1; ++ const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "moe_mm.metal"}, (v3 && !v4 && !v5) ? "#define MM_V3_ONLY 1\n" : ""); ++ const int NE = bench ? 512 : 64, LAYERS = bench && !isolated_bench ? 4 : 1; + const size_t gu_row = ref::row_bytes(gu, N_EMBD), d_row = ref::row_bytes(dn, N_FF); + const size_t up_off = (size_t)N_FF * gu_row, down_off = 2 * up_off, blob = down_off + (size_t)N_EMBD * d_row; + std::printf("device %s; %s/%s, T %d, %d experts, tile TM %d TN %d TK %d SG %d\n", cx.device_name().c_str(), +@@ -67,14 +229,18 @@ int main(int argc, char** argv) { + } + const int nent = T * TOPK, nblk = (nent + 1023) / 1024; + const int maxrows = nent + NE * TM; +- mtl::Buffer x = cx.alloc(4 * (size_t)T * N_EMBD), ids = cx.alloc(4 * nent), w = cx.alloc(4 * nent), ++ mtl::Buffer x = cx.alloc(4 * (size_t)T * N_EMBD), xh = cx.alloc(2 * (size_t)T * N_EMBD), ids = cx.alloc(4 * nent), w = cx.alloc(4 * nent), + bcnt = cx.alloc(4 * (size_t)nblk * NE), boff = cx.alloc(4 * (size_t)nblk * NE), eoff = cx.alloc(4 * NE), + tiles = cx.alloc(8 * (size_t)(maxrows / TM + 1)), meta = cx.alloc(64), stok = cx.alloc(4 * (size_t)maxrows), + sw = cx.alloc(4 * (size_t)maxrows), pos = cx.alloc(4 * nent), H = cx.alloc(4 * (size_t)maxrows * N_FF), + YE = cx.alloc(4 * (size_t)maxrows * N_EMBD), y = cx.alloc(4 * (size_t)T * N_EMBD), + XS = cx.alloc(2 * (size_t)maxrows * N_EMBD), HS = cx.alloc(2 * (size_t)maxrows * N_FF); + ref::Rng r(7); +- for (size_t i = 0; i < (size_t)T * N_EMBD; ++i) x.as()[i] = r.sym(); ++ for (size_t i = 0; i < (size_t)T * N_EMBD; ++i) { ++ const float v = r.sym(); ++ x.as()[i] = v; ++ xh.as()[i] = ref::f32_to_f16(v); ++ } + for (int t = 0; t < T; ++t) { + std::set used; + float s = 0; +@@ -88,7 +254,7 @@ int main(int argc, char** argv) { + for (int k = 0; k < TOPK; ++k) w.as()[t * TOPK + k] /= s; + } + mtl::Pipeline phist = cx.pipeline(lib, "mm_hist"), pscan = cx.pipeline(lib, "mm_scan"), pplace = cx.pipeline(lib, "mm_place"), +- pclear = cx.pipeline(lib, "mm_clear"), pcomb = cx.pipeline(lib, "mm_combine"); ++ pcomb = cx.pipeline(lib, "mm_combine"); + if (DTN < 0) { DTN = TN; DTK = TK; DSG = SG; } + if (dn == ref::IQ4_NL && DTK < 32) DTK = 32; + const std::string tl = std::to_string(TM) + "_" + std::to_string(TN) + "_" + std::to_string(TK) + "_" + std::to_string(SG); +@@ -105,15 +271,13 @@ int main(int argc, char** argv) { + : cx.pipeline(lib, (v2 ? "mm_gu2_" : "mm_gu_") + std::to_string(gu) + "_" + tl, probe), + pdn = v3 ? cx.pipeline(lib, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4)) + : cx.pipeline(lib, (v2 ? "mm_down2_" : "mm_down_") + std::to_string(dn) + "_" + (v2 ? dtl : tl)), +- pgat = cx.pipeline(lib, "mm_gather"); ++ pgat = cx.pipeline(lib, "mm_gather_h"); + MmPlanArgs pa{(uint32_t)nent, TOPK, (uint32_t)NE, (uint32_t)TM, (uint32_t)nblk}; + auto encode = [&](int layer, std::vector& ds, bool plan, bool gemm) { + if (plan) { +- const uint32_t nclr = maxrows; +- mtl::Dispatch c; c.p = &pclear; c.args = {mtl::Arg::B(stok), mtl::Arg::V(nclr)}; c.grid[0] = (maxrows + 255) / 256; c.tg[0] = 256; ds.push_back(c); + mtl::Dispatch h; h.p = &phist; h.args = {mtl::Arg::B(ids), mtl::Arg::B(bcnt), mtl::Arg::V(pa)}; h.grid[0] = nblk; h.tg[0] = 1024; ds.push_back(h); + meta.as()[2] = N_FF / TN; +- mtl::Dispatch sc; sc.p = &pscan; sc.args = {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta), mtl::Arg::V(pa)}; sc.tg[0] = NE; ds.push_back(sc); ++ mtl::Dispatch sc; sc.p = &pscan; sc.args = {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta), mtl::Arg::V(pa), mtl::Arg::B(stok)}; sc.tg[0] = NE; ds.push_back(sc); + mtl::Dispatch pl; pl.p = &pplace; pl.args = {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok), mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}; pl.grid[0] = nblk; pl.tg[0] = 1024; ds.push_back(pl); + } + if (gemm) { +@@ -121,7 +285,7 @@ int main(int argc, char** argv) { + blob, blob, (uint32_t)gu_row, (uint32_t)d_row, N_EMBD, N_FF}; + if (v2) { + const uint32_t ne = N_EMBD; +- mtl::Dispatch ga; ga.p = &pgat; ga.args = {mtl::Arg::B(x), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; ++ mtl::Dispatch ga; ga.p = &pgat; ga.args = {mtl::Arg::B(xh), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; + ga.grid[0] = meta.as()[1]; ga.tg[0] = 256; ds.push_back(ga); + } + mtl::Dispatch g; g.p = &pgu; +@@ -142,6 +306,24 @@ int main(int argc, char** argv) { + } + const uint32_t ntiles = meta.as()[0]; + std::printf("plan: %u tiles of %d, %u padded rows for %d entries\n", ntiles, TM, meta.as()[1], nent); ++ auto compare_ab = [&](const std::vector& ds, mtl::Pipeline& pgu_b, mtl::Pipeline& pdn_b, size_t pr) { ++ std::vector hs_a(HS.as(), HS.as() + 2 * pr * N_FF); ++ std::vector ye_a(YE.as(), YE.as() + 4 * pr * N_EMBD); ++ std::vector y_a(y.as(), y.as() + 4 * (size_t)T * N_EMBD); ++ std::memset(HS.ptr, 0, hs_a.size()); ++ std::memset(YE.ptr, 0, ye_a.size()); ++ std::memset(y.ptr, 0, y_a.size()); ++ auto ds_b = ds; ++ ds_b[v2 ? 1 : 0].p = &pgu_b; ++ ds_b[v2 ? 2 : 1].p = &pdn_b; ++ cx.run(ds_b); ++ const bool ok_hs = std::memcmp(HS.ptr, hs_a.data(), hs_a.size()) == 0; ++ const bool ok_ye = std::memcmp(YE.ptr, ye_a.data(), ye_a.size()) == 0; ++ const bool ok_y = std::memcmp(y.ptr, y_a.data(), y_a.size()) == 0; ++ std::printf("ab bitwise (T=%d): HS_memcmp=%s YE_memcmp=%s Y_memcmp=%s\n", ++ T, ok_hs ? "PASS" : "FAIL", ok_ye ? "PASS" : "FAIL", ok_y ? "PASS" : "FAIL"); ++ return ok_hs && ok_ye && ok_y; ++ }; + if (!bench) { + // the plan: every entry placed once, rows of an expert contiguous and in entry order + bool ok_plan = true; +@@ -155,6 +337,22 @@ int main(int argc, char** argv) { + std::vector ds; + encode(0, ds, false, true); + cx.run(ds); ++ const size_t pr = meta.as()[1]; ++ if (dump) { ++ FILE* f = std::fopen(dump, "wb"); ++ if (v2) std::fwrite(HS.as(), 2, pr * N_FF, f); ++ std::fwrite(YE.as(), 4, pr * N_EMBD, f); ++ std::fwrite(y.as(), 4, (size_t)T * N_EMBD, f); ++ std::fclose(f); ++ return 0; ++ } ++ if (ab_dir) { ++ const std::string kd(ab_dir); ++ const int lib_b = cx.compile({kd + "/quant.metal", kd + "/units.metal", kd + "/mma.metal", kd + "/moe_mm.metal"}, (v3 && !v4 && !v5) ? "#define MM_V3_ONLY 1\n" : ""); ++ mtl::Pipeline pgu_b = cx.pipeline(lib_b, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)); ++ mtl::Pipeline pdn_b = cx.pipeline(lib_b, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4)); ++ return (ok_plan && compare_ab(ds, pgu_b, pdn_b, pr)) ? 0 : 1; ++ } + // reference: fp16 operands (x, weights, h), fp32+ accumulation; and the exact one + double se16 = 0, se = 0, sr = 0; + std::vector wg((size_t)N_FF * N_EMBD), wu(wg.size()), wd((size_t)N_EMBD * N_FF); +@@ -217,18 +415,85 @@ int main(int argc, char** argv) { + } + auto timeit = [&](bool plan, bool gemm, int part) { + std::vector ds; +- for (int l = 0; l < 8; ++l) { ++ const int repeats = isolated_bench ? 4 : 8; ++ for (int l = 0; l < repeats; ++l) { + std::vector one; + encode(l % LAYERS, one, plan, gemm); + if (part >= 0 && gemm) { ds.push_back(one[part]); continue; } + for (auto& d : one) ds.push_back(d); + } +- cx.warm(ds, 0.3); ++ cx.warm(ds, isolated_bench ? 0.1 : 0.3); + std::vector t; +- for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds) / 8); ++ for (int i = 0; i < (isolated_bench ? 3 : 5); ++i) t.push_back(cx.run(ds) / repeats); + return median(t); + }; + const int o = v2 ? 1 : 0; ++ if (isolated_bench) { ++ std::vector ds; ++ encode(0, ds, false, true); ++ cx.run(ds); ++ } ++ if (ab_dir && isolated_bench) { ++ const std::string kd(ab_dir); ++ const int lib_b = cx.compile({kd + "/quant.metal", kd + "/units.metal", kd + "/mma.metal", kd + "/moe_mm.metal"}, (v3 && !v4 && !v5) ? "#define MM_V3_ONLY 1\n" : ""); ++ mtl::Pipeline pgu_b = cx.pipeline(lib_b, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)); ++ mtl::Pipeline pdn_b = cx.pipeline(lib_b, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4)); ++ std::vector parity; ++ encode(0, parity, false, true); ++ const size_t pr = meta.as()[1]; ++ std::memset(HS.ptr, 0, 2 * pr * N_FF); ++ std::memset(YE.ptr, 0, 4 * pr * N_EMBD); ++ cx.run(parity); ++ if (!compare_ab(parity, pgu_b, pdn_b, pr)) return 1; ++ const int part = dn_only ? (o + 1) : o; ++ std::vector one; ++ encode(0, one, false, true); ++ std::vector ds_a{one[part]}; ++ mtl::Dispatch db = one[part]; ++ db.p = dn_only ? &pdn_b : &pgu_b; ++ std::vector ds_b{db}; ++ cx.warm(ds_a, 0.05); ++ cx.warm(ds_b, 0.05); ++ std::vector ta, tb; ++ double best_pair = 1e30; ++ for (int i = 0; i < 19; ++i) { ++ double va, vb; ++ if (i & 1) { ++ vb = cx.run(ds_b, 2) / 2; ++ va = cx.run(ds_a, 2) / 2; ++ } else { ++ va = cx.run(ds_a, 2) / 2; ++ vb = cx.run(ds_b, 2) / 2; ++ } ++ best_pair = std::min(best_pair, std::min(va, vb)); ++ ta.push_back(va); ++ tb.push_back(vb); ++ } ++ std::vector fa, fb; ++ for (size_t i = 0; i < ta.size(); ++i) { ++ if (std::max(ta[i], tb[i]) <= 1.10 * best_pair) { ++ fa.push_back(ta[i]); ++ fb.push_back(tb[i]); ++ } ++ } ++ if (fa.size() >= 3) { ta = fa; tb = fb; } ++ const double ma = median(ta) * 1e6, mb = median(tb) * 1e6; ++ const double mina = *std::min_element(ta.begin(), ta.end()) * 1e6; ++ const double minb = *std::min_element(tb.begin(), tb.end()) * 1e6; ++ std::printf("%s A/B (T=%d): A med=%.1f (min=%.1f) | B med=%.1f (min=%.1f) | diff med=%+.2f%% min=%+.2f%%\n", ++ dn_only ? "down" : "gate/up", T, ma, mina, mb, minb, (mb - ma) / ma * 100.0, (minb - mina) / mina * 100.0); ++ return 0; ++ } ++ if (dn_only) { ++ const double td = timeit(false, true, o + 1); ++ std::printf("down %.1f us (%.2f TFLOP/s useful)\n", td * 1e6, 2.0 * nent * N_FF * N_EMBD / td / 1e12); ++ return 0; ++ } ++ if (gu_only) { ++ const double tg = timeit(false, true, o); ++ std::printf("gate/up %.1f us (%.2f TFLOP/s useful)\n", tg * 1e6, 2.0 * nent * 2.0 * N_FF * N_EMBD / tg / 1e12); ++ return 0; ++ } + const double tp = timeit(true, false, -1), tg = timeit(false, true, o), td = timeit(false, true, o + 1), + tc = timeit(false, true, o + 2), ta = timeit(false, true, -1); + if (v2) std::printf("gather %.1f us\n", timeit(false, true, 0) * 1e6); +diff --git a/metal/tests/ple_moe_parity.cpp b/metal/tests/ple_moe_parity.cpp +new file mode 100644 +index 00000000..fb40e33e +--- /dev/null ++++ b/metal/tests/ple_moe_parity.cpp +@@ -0,0 +1,340 @@ ++// metal/tests/ple_moe_parity.cpp - the prompt chunk's moe_out folds and the fused PLE norms against the dispatch ++// chains they replace, BITWISE: ++// gr_norm_h_moe (prompt.metal) == moe_out + gr_norm_h (R', xnh) ++// gr_apply_moe (prompt.metal) == moe_out + gr_norm (R') ++// ple_moe_gate_norm + ple_conv_res == moe_out + gr_norm + ple_gnorm (key) + ple_gnorm (query) + ple_gate_norm ++// (msl/ple.metal) + ple_conv_add (R', normalized) ++// ple_moe_gate_norm_c + ple_conv_res == mm_combine + the ple_moe_gate_norm chain above (R', normalized) ++// ple_query_gate_norm + ple_conv_res == ple_gnorm (query) + ple_gate_norm + ple_conv_add (R', normalized; key ++// from ple_gnorm) ++// ple_comb_query_gate_norm == gr_norm_comb + ple_query_gate_norm (R', normalized) ++// ple_conv_res + ple_hist_adv == ple_conv_res + the host's history advance (R', the history), T = 1 .. 2048 ++// `ple_moe_parity [T] bench` (or `--bench`) adds the PLE stage's GPU times at T (default 2048). ++// `ple_moe_parity --ab baseline.metal [T]` runs short paired A/B bitwise and timing checks, including history, at any T. ++#include "mtl_rt.h" ++#include "ref_quant.h" ++ ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++ ++namespace { ++struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; ++struct PleArgs { uint32_t n_embd, hc; float eps; }; ++constexpr uint32_t P_TOKENT = 532; ++bool same(const mtl::Buffer& a, const mtl::Buffer& b, size_t bytes) { return std::memcmp(a.ptr, b.ptr, bytes) == 0; } ++double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } ++} // namespace ++ ++int main(int argc, char** argv) { ++ const bool ab = argc > 1 && std::strcmp(argv[1], "--ab") == 0; ++ if (ab && argc < 3) { ++ std::fprintf(stderr, "usage: ple_moe_parity --ab baseline.metal [T]\n"); ++ return 2; ++ } ++ const bool flag_bench = argc > 1 && std::strcmp(argv[1], "--bench") == 0; ++ const uint32_t TB = ab ? (argc > 3 ? (uint32_t) std::atoi(argv[3]) : 2048u) : ++ (argc > 1 && !flag_bench ? (uint32_t) std::atoi(argv[1]) : 2048u); ++ const bool bench = ab || flag_bench || (argc > 2 && (std::strcmp(argv[2], "bench") == 0 || std::strcmp(argv[2], "--bench") == 0)); ++ mtl::Context cx; ++ const int lple = cx.compile({"../src/engine/msl/ple.metal"}, "#define NO_GGML_COMMON 1\n"); ++ std::printf("device %s\n", cx.device_name().c_str()); ++ const uint32_t N = 2560, HC = 4, LR = 320, TOPK = 10, NEXP = 512, TMR = 16; ++ const size_t RN = (size_t) HC * N; ++ mtl::Pipeline pg = cx.pipeline(lple, "ple_gnorm"), ++ pgn = cx.pipeline(lple, "ple_gate_norm"), pca = cx.pipeline(lple, "ple_conv_add"), ++ pcr = cx.pipeline(lple, "ple_conv_res"), pf = cx.pipeline(lple, "ple_moe_gate_norm"), ++ pfc = cx.pipeline(lple, "ple_moe_gate_norm_c"), pq = cx.pipeline(lple, "ple_query_gate_norm"), ++ pcq = cx.pipeline(lple, "ple_comb_query_gate_norm"), pha = cx.pipeline(lple, "ple_hist_adv"); ++ auto D = [](const mtl::Pipeline& p, std::vector a, uint32_t gx, uint32_t gy) { ++ mtl::Dispatch d; ++ d.p = &p; ++ d.args = std::move(a); ++ d.grid[0] = gx; ++ d.grid[1] = gy; ++ d.tg[0] = 256; ++ return d; ++ }; ++ int fails = 0; ++ if (ab) { ++ const int lref = cx.compile({argv[2]}, "#define NO_GGML_COMMON 1\n"); ++ mtl::Pipeline pg_r = cx.pipeline(lref, "ple_gnorm"), pgn_r = cx.pipeline(lref, "ple_gate_norm"), ++ pcr_r = cx.pipeline(lref, "ple_conv_res"), pf_r = cx.pipeline(lref, "ple_moe_gate_norm"), ++ pfc_r = cx.pipeline(lref, "ple_moe_gate_norm_c"), pq_r = cx.pipeline(lref, "ple_query_gate_norm"), ++ pcq_r = cx.pipeline(lref, "ple_comb_query_gate_norm"), pha_r = cx.pipeline(lref, "ple_hist_adv"); ++ const std::vector ts = argc > 3 ? std::vector{TB} : ++ std::vector{1u, 3u, 7u, 9u, 10u, 64u, 512u, 1024u, 2048u}; ++ for (uint32_t T : ts) { ++ const size_t mr = (size_t) T * TOPK + NEXP * TMR; ++ mtl::Buffer R0 = cx.alloc(4 * T * RN), Ra = cx.alloc(4 * T * RN), Rb = cx.alloc(4 * T * RN), yr = cx.alloc(4 * T * N), ++ ys = cx.alloc(4 * T * N), sgv = cx.alloc(4 * T), inj = cx.alloc(4 * T * HC), ++ wk = cx.alloc(4 * RN), wq = cx.alloc(4 * RN), wc = cx.alloc(4 * RN), ++ pkey = cx.alloc(4 * T * RN), pval = cx.alloc(4 * T * N), keyn = cx.alloc(4 * T * RN), ++ ga = cx.alloc(4 * T * RN), gb = cx.alloc(4 * T * RN), ++ na = cx.alloc(4 * T * RN), nb = cx.alloc(4 * T * RN), hist = cx.alloc(4 * 9 * RN), kw = cx.alloc(2 * 4 * RN), ++ YE = cx.alloc(4 * mr * N), pos = cx.alloc(4 * T * TOPK), sw = cx.alloc(4 * mr), ++ plan = cx.alloc(4 * (P_TOKENT + T * TOPK)), ++ ha = cx.alloc(4 * 9 * RN), hb = cx.alloc(4 * 9 * RN); ++ ref::Rng r(11 + T); ++ auto fill = [&](const mtl::Buffer& b, float s, float o) { ++ for (size_t i = 0; i < b.bytes / 4; ++i) b.as()[i] = o + s * r.sym(); ++ }; ++ fill(R0, 1.0f, 0.0f); fill(yr, 0.5f, 0.0f); fill(ys, 0.5f, 0.0f); fill(sgv, 3.0f, 0.0f); fill(inj, 4.0f, 0.0f); ++ fill(wk, 0.2f, 1.0f); fill(wq, 0.2f, 1.0f); fill(wc, 0.2f, 1.0f); ++ fill(pkey, 1.0f, 0.0f); fill(pval, 1.0f, 0.0f); fill(keyn, 1.0f, 0.0f); fill(hist, 1.0f, 0.0f); ++ fill(YE, 0.5f, 0.0f); fill(sw, 0.1f, 0.1f); ++ for (size_t i = 0; i < kw.bytes / 2; ++i) kw.as()[i] = ref::f32_to_f16(0.5f * r.sym()); ++ std::vector perm(mr); ++ std::iota(perm.begin(), perm.end(), 0); ++ std::shuffle(perm.begin(), perm.end(), std::mt19937(T)); ++ for (size_t i = 0; i < (size_t) T * TOPK; ++i) { ++ pos.as()[i] = (i % 37 == 5) ? -1 : perm[i]; ++ plan.as()[P_TOKENT + i] = perm[i]; ++ } ++ const PleArgs pl{N, HC, 1e-6f}; ++ auto reset = [&] { ++ std::memcpy(Ra.ptr, R0.ptr, 4 * T * RN); ++ std::memcpy(Rb.ptr, R0.ptr, 4 * T * RN); ++ std::memset(na.ptr, 0, 4 * T * RN); ++ std::memset(nb.ptr, 0, 4 * T * RN); ++ }; ++ const mtl::Dispatch g_r = D(pg_r, {mtl::Arg::B(pkey), mtl::Arg::B(wk), mtl::Arg::B(na), mtl::Arg::V(pl)}, HC, T), ++ g_c = D(pg, {mtl::Arg::B(pkey), mtl::Arg::B(wk), mtl::Arg::B(nb), mtl::Arg::V(pl)}, HC, T), ++ gn_r = D(pgn_r, {mtl::Arg::B(keyn), mtl::Arg::B(R0), mtl::Arg::B(pval), mtl::Arg::B(wc), ++ mtl::Arg::B(ga), mtl::Arg::B(na), mtl::Arg::V(pl)}, HC, T), ++ gn_c = D(pgn, {mtl::Arg::B(keyn), mtl::Arg::B(R0), mtl::Arg::B(pval), mtl::Arg::B(wc), ++ mtl::Arg::B(gb), mtl::Arg::B(nb), mtl::Arg::V(pl)}, HC, T), ++ res_r = D(pcr_r, {mtl::Arg::B(hist), mtl::Arg::B(na), mtl::Arg::B(kw), mtl::Arg::B(Ra), mtl::Arg::V(pl)}, HC * N / 256, T), ++ res_c = D(pcr, {mtl::Arg::B(hist), mtl::Arg::B(nb), mtl::Arg::B(kw), mtl::Arg::B(Rb), mtl::Arg::V(pl)}, HC * N / 256, T), ++ f_r = D(pf_r, {mtl::Arg::B(Ra), mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), ++ mtl::Arg::B(pkey), mtl::Arg::B(wk), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(na), mtl::Arg::V(pl)}, HC, T), ++ f_c = D(pf, {mtl::Arg::B(Rb), mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), ++ mtl::Arg::B(pkey), mtl::Arg::B(wk), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(nb), mtl::Arg::V(pl)}, HC, T), ++ fc_r = D(pfc_r, {mtl::Arg::B(Ra), mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(ys), ++ mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::B(pkey), mtl::Arg::B(wk), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(na), mtl::Arg::V(pl), mtl::Arg::V(TOPK)}, 1, T), ++ fc_c = D(pfc, {mtl::Arg::B(Rb), mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(ys), ++ mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::B(pkey), mtl::Arg::B(wk), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(nb), mtl::Arg::V(pl), mtl::Arg::V(TOPK)}, 1, T), ++ q_r = D(pq_r, {mtl::Arg::B(keyn), mtl::Arg::B(Ra), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(na), mtl::Arg::V(pl)}, HC, T), ++ q_c = D(pq, {mtl::Arg::B(keyn), mtl::Arg::B(Rb), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(nb), mtl::Arg::V(pl)}, HC, T), ++ cq_r = D(pcq_r, {mtl::Arg::B(keyn), mtl::Arg::B(Ra), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(na), mtl::Arg::V(pl), ++ mtl::Arg::B(YE), mtl::Arg::B(inj), mtl::Arg::B(plan), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::V(TOPK)}, HC, T), ++ cq_c = D(pcq, {mtl::Arg::B(keyn), mtl::Arg::B(Rb), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(nb), mtl::Arg::V(pl), ++ mtl::Arg::B(YE), mtl::Arg::B(inj), mtl::Arg::B(plan), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::V(TOPK)}, HC, T); ++ reset(); cx.run({g_r}); cx.run({g_c}); ++ const bool ok_g = same(na, nb, 4 * T * RN); ++ reset(); std::memset(ga.ptr, 0, 4 * T * RN); std::memset(gb.ptr, 0, 4 * T * RN); ++ cx.run({gn_r}); cx.run({gn_c}); ++ const bool ok_gn = same(ga, gb, 4 * T * RN) && same(na, nb, 4 * T * RN); ++ reset(); cx.run({f_r, res_r}); cx.run({f_c, res_c}); ++ const bool ok_f = same(Ra, Rb, 4 * T * RN) && same(na, nb, 4 * T * RN); ++ reset(); cx.run({fc_r, res_r}); cx.run({fc_c, res_c}); ++ const bool ok_fc = same(Ra, Rb, 4 * T * RN) && same(na, nb, 4 * T * RN); ++ reset(); cx.run({q_r, res_r}); cx.run({q_c, res_c}); ++ const bool ok_q = same(Ra, Rb, 4 * T * RN) && same(na, nb, 4 * T * RN); ++ reset(); cx.run({cq_r, res_r}); cx.run({cq_c, res_c}); ++ const bool ok_cq = same(Ra, Rb, 4 * T * RN) && same(na, nb, 4 * T * RN); ++ const mtl::Dispatch h_r = D(pha_r, {mtl::Arg::B(ha), mtl::Arg::B(na), mtl::Arg::V(pl), mtl::Arg::V(T)}, HC * N / 256, 1), ++ h_c = D(pha, {mtl::Arg::B(hb), mtl::Arg::B(nb), mtl::Arg::V(pl), mtl::Arg::V(T)}, HC * N / 256, 1); ++ std::memcpy(ha.ptr, hist.ptr, hist.bytes); ++ std::memcpy(hb.ptr, hist.ptr, hist.bytes); ++ cx.run({h_r}); cx.run({h_c}); ++ const bool ok_h = same(ha, hb, hist.bytes); ++ fails += !ok_g + !ok_gn + !ok_f + !ok_fc + !ok_q + !ok_cq + !ok_h; ++ std::printf("A/B T %4u: ple_gnorm %s, ple_gate_norm %s, ple_hist_adv %s\n", ++ T, ok_g ? "PASS" : "FAIL", ok_gn ? "PASS" : "FAIL", ok_h ? "PASS" : "FAIL"); ++ std::printf("A/B T %4u: ple_moe_gate_norm %s, ple_moe_gate_norm_c %s, ple_query_gate_norm %s, ple_comb_query_gate_norm %s\n", ++ T, ok_f ? "PASS" : "FAIL", ok_fc ? "PASS" : "FAIL", ok_q ? "PASS" : "FAIL", ok_cq ? "PASS" : "FAIL"); ++ const bool single_t = argc > 3; ++ const int reps = single_t ? (T <= 10 ? 64 : (T <= 64 ? 16 : 4)) : 2; ++ const int pairs = single_t ? 7 : 5; ++ auto ab_time = [&](const char* name, const std::vector& ref_ds, const std::vector& cand_ds, int r = 0) { ++ if (!r) r = reps; ++ cx.warm(ref_ds, 0.003); cx.warm(cand_ds, 0.003); ++ std::vector bt, ct, rt; ++ for (int p = 0; p < pairs; ++p) { ++ double b, c; ++ if (p & 1) { c = cx.run(cand_ds, r) / r * 1e3; b = cx.run(ref_ds, r) / r * 1e3; } ++ else { b = cx.run(ref_ds, r) / r * 1e3; c = cx.run(cand_ds, r) / r * 1e3; } ++ bt.push_back(b); ct.push_back(c); rt.push_back(c / b); ++ } ++ std::printf(" T %4u %-34s base %.4f -> cand %.4f ms (ratio %.4f, %+.1f%%)\n", ++ T, name, median(bt), median(ct), median(rt), 100.0 * (median(rt) - 1.0)); ++ }; ++ ab_time("ple_gnorm", {g_r}, {g_c}); ++ ab_time("ple_gate_norm", {gn_r}, {gn_c}); ++ ab_time("ple_query_gate_norm", {q_r}, {q_c}); ++ ab_time("ple_comb_query_gate_norm", {cq_r}, {cq_c}); ++ ab_time("ple_moe_gate_norm", {f_r}, {f_c}); ++ ab_time("ple_moe_gate_norm_c", {fc_r}, {fc_c}); ++ ab_time("ple_hist_adv", {h_r}, {h_c}, 128); ++ ab_time("ple_conv_res", {res_r}, {res_c}); ++ ab_time("ple_comb_query_gate_norm + conv_res", {cq_r, res_r}, {cq_c, res_c}); ++ ab_time("ple_moe_gate_norm_c + conv_res", {fc_r, res_r}, {fc_c, res_c}); ++ ab_time("ple_query_gate_norm + conv_res", {q_r, res_r}, {q_c, res_c}); ++ } ++ std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); ++ return fails ? 1 : 0; ++ } ++ const std::string prelude = "#define NO_GGML_COMMON 1\nusing namespace metal;\n"; ++ const std::string gr_src = mtl::read_file(cx.kernel_dir + "/gr.metal"); ++ const std::string pr_src = mtl::read_file(cx.kernel_dir + "/prompt.metal"); ++ const std::string mm_src = mtl::read_file(cx.kernel_dir + "/moe_mm.metal"); ++ const auto gn_end = gr_src.find("// rows: 0 .. hc_lr - 1"); ++ const auto mo_beg = pr_src.find("// grid (T * n_embd / 256); sg: T values"); ++ const auto pr_end = pr_src.find("// gemm_hc_up:"); ++ const auto mc_beg = mm_src.find("kernel void mm_combine("); ++ const auto mc_end = mm_src.find("\n}", mc_beg == std::string::npos ? 0 : mc_beg); ++ const int lib = cx.compile({}, prelude + gr_src.substr(0, gn_end) + pr_src.substr(mo_beg, pr_end - mo_beg)); ++ const int lmm = cx.compile({}, prelude + mm_src.substr(mc_beg, mc_end + 2 - mc_beg)); ++ mtl::Pipeline pmo = cx.pipeline(lib, "moe_out"), pnh = cx.pipeline(lib, "gr_norm_h"), pn = cx.pipeline(lib, "gr_norm"), ++ pnhm = cx.pipeline(lib, "gr_norm_h_moe"), pam = cx.pipeline(lib, "gr_apply_moe"), ++ pcomb = cx.pipeline(lmm, "mm_combine"); ++ for (uint32_t T : {1u, 3u, 7u, 9u, 10u, 64u, TB}) { ++ if (T == TB && !bench) continue; ++ const size_t mr = (size_t) T * TOPK + NEXP * TMR; ++ mtl::Buffer R0 = cx.alloc(4 * T * RN), Ra = cx.alloc(4 * T * RN), Rb = cx.alloc(4 * T * RN), yr = cx.alloc(4 * T * N), ++ ys = cx.alloc(4 * T * N), sgv = cx.alloc(4 * T), inj = cx.alloc(4 * T * HC), bo = cx.alloc(4 * T * N), ++ wn = cx.alloc(4 * RN), wk = cx.alloc(4 * RN), wq = cx.alloc(4 * RN), wc = cx.alloc(4 * RN), ++ pkey = cx.alloc(4 * T * RN), pval = cx.alloc(4 * T * N), xa = cx.alloc(4 * T * RN), xb = cx.alloc(4 * T * RN), ++ keyn = cx.alloc(4 * T * RN), qn = cx.alloc(4 * T * RN), ga = cx.alloc(4 * T * RN), ++ na = cx.alloc(4 * T * RN), nb = cx.alloc(4 * T * RN), hist = cx.alloc(4 * 9 * RN), kw = cx.alloc(2 * 4 * RN), ++ YE = cx.alloc(4 * mr * N), pos = cx.alloc(4 * T * TOPK), sw = cx.alloc(4 * mr); ++ ref::Rng r(11 + T); ++ auto fill = [&](const mtl::Buffer& b, float s, float o) { ++ for (size_t i = 0; i < b.bytes / 4; ++i) b.as()[i] = o + s * r.sym(); ++ }; ++ fill(R0, 1.0f, 0.0f); fill(yr, 0.5f, 0.0f); fill(ys, 0.5f, 0.0f); fill(sgv, 3.0f, 0.0f); fill(inj, 4.0f, 0.0f); ++ fill(wn, 0.2f, 1.0f); fill(wk, 0.2f, 1.0f); fill(wq, 0.2f, 1.0f); fill(wc, 0.2f, 1.0f); ++ fill(pkey, 1.0f, 0.0f); fill(pval, 1.0f, 0.0f); fill(hist, 1.0f, 0.0f); fill(YE, 0.5f, 0.0f); fill(sw, 0.1f, 0.1f); ++ for (size_t i = 0; i < kw.bytes / 2; ++i) kw.as()[i] = ref::f32_to_f16(0.5f * r.sym()); ++ // mm_place's slots: distinct rows of YE (expert-sorted, so scattered), a few dropped (-1) ++ std::vector perm(mr); ++ std::iota(perm.begin(), perm.end(), 0); ++ std::shuffle(perm.begin(), perm.end(), std::mt19937(T)); ++ for (size_t i = 0; i < (size_t) T * TOPK; ++i) pos.as()[i] = (i % 37 == 5) ? -1 : perm[i]; ++ const GrArgs a{N, HC, LR, T, 1u, 1e-6f}; ++ const PleArgs pl{N, HC, 1e-6f}; ++ const uint32_t oa[2] = {T, N}, ca[4] = {T, N, TOPK, 0}; ++ const mtl::Dispatch mo = D(pmo, {mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(bo), mtl::Arg::V(oa)}, ++ (T * N + 255) / 256, 1), ++ comb = D(pcomb, {mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(yr), mtl::Arg::V(ca)}, ++ (T * N + 255) / 256, 1); ++ auto reset = [&] { ++ std::memcpy(Ra.ptr, R0.ptr, 4 * T * RN); ++ std::memcpy(Rb.ptr, R0.ptr, 4 * T * RN); ++ std::memset(bo.ptr, 0, 4 * T * N); ++ }; ++ auto check = [&](const char* what, bool ok) { ++ fails += !ok; ++ std::printf("%-34s T %4u: bitwise %s\n", what, T, ok ? "PASS" : "FAIL"); ++ }; ++ // the unfused PLE stage on R (layer 0's MoE write from yr applied first if moe) ++ auto chain = [&](const mtl::Buffer& R, bool moe) { ++ std::vector ds; ++ if (moe) { ++ ds.push_back(mo); ++ ds.push_back(D(pn, {mtl::Arg::B(R), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(xa), ++ mtl::Arg::V(a)}, HC, T)); ++ } ++ ds.push_back(D(pg, {mtl::Arg::B(pkey), mtl::Arg::B(wk), mtl::Arg::B(keyn), mtl::Arg::V(pl)}, HC, T)); ++ ds.push_back(D(pg, {mtl::Arg::B(R), mtl::Arg::B(wq), mtl::Arg::B(qn), mtl::Arg::V(pl)}, HC, T)); ++ ds.push_back(D(pgn, {mtl::Arg::B(keyn), mtl::Arg::B(qn), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(ga), ++ mtl::Arg::B(na), mtl::Arg::V(pl)}, HC, T)); ++ ds.push_back(D(pca, {mtl::Arg::B(hist), mtl::Arg::B(na), mtl::Arg::B(kw), mtl::Arg::B(ga), mtl::Arg::B(R), ++ mtl::Arg::V(pl)}, HC * N / 256, T)); ++ return ds; ++ }; ++ const mtl::Dispatch res = D(pcr, {mtl::Arg::B(hist), mtl::Arg::B(nb), mtl::Arg::B(kw), mtl::Arg::B(Rb), mtl::Arg::V(pl)}, ++ HC * N / 256, T), ++ fused = D(pf, {mtl::Arg::B(Rb), mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), ++ mtl::Arg::B(pkey), mtl::Arg::B(wk), mtl::Arg::B(wq), mtl::Arg::B(pval), ++ mtl::Arg::B(wc), mtl::Arg::B(nb), mtl::Arg::V(pl)}, HC, T), ++ fused_c = D(pfc, {mtl::Arg::B(Rb), mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(ys), ++ mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::B(pkey), mtl::Arg::B(wk), ++ mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), mtl::Arg::B(nb), ++ mtl::Arg::V(pl), mtl::Arg::V(TOPK)}, 1, T), ++ query = D(pq, {mtl::Arg::B(keyn), mtl::Arg::B(Rb), mtl::Arg::B(wq), mtl::Arg::B(pval), mtl::Arg::B(wc), ++ mtl::Arg::B(nb), mtl::Arg::V(pl)}, HC, T); ++ const bool small = T != TB; ++ if (small && (T == 1u || T == 7u || T == 64u)) { ++ // ---- gr_norm_h_moe ++ reset(); ++ cx.run({mo, D(pnh, {mtl::Arg::B(Ra), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(xa), mtl::Arg::V(a)}, HC, T)}); ++ cx.run({D(pnhm, {mtl::Arg::B(Rb), mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::B(wn), ++ mtl::Arg::B(xb), mtl::Arg::V(a)}, 1, T)}); ++ check("gr_norm_h_moe", same(Ra, Rb, 4 * T * RN) && same(xa, xb, 2 * T * RN)); ++ // ---- gr_apply_moe ++ reset(); ++ cx.run({mo, D(pn, {mtl::Arg::B(Ra), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(xa), mtl::Arg::V(a)}, HC, T)}); ++ cx.run({D(pam, {mtl::Arg::B(Rb), mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(inj), mtl::Arg::V(a)}, ++ HC * N / 256, T)}); ++ check("gr_apply_moe", same(Ra, Rb, 4 * T * RN)); ++ } ++ // ---- ple_moe_gate_norm + ple_conv_res ++ reset(); ++ cx.run(chain(Ra, true)); ++ cx.run({fused, res}); ++ check("ple_moe_gate_norm + ple_conv_res", same(Ra, Rb, 4 * T * RN) && same(na, nb, 4 * T * RN)); ++ // ---- ple_moe_gate_norm_c + ple_conv_res (yr from mm_combine in the chain) ++ reset(); ++ std::vector cc = chain(Ra, true); ++ cc.insert(cc.begin(), comb); ++ cx.run(cc); ++ std::memset(nb.ptr, 0, 4 * T * RN); ++ cx.run({fused_c, res}); ++ check("ple_moe_gate_norm_c + ple_conv_res", same(Ra, Rb, 4 * T * RN) && same(na, nb, 4 * T * RN)); ++ // ---- ple_query_gate_norm + ple_conv_res (the window: keyn from the chain's ple_gnorm) ++ reset(); ++ cx.run(chain(Ra, false)); ++ std::memset(nb.ptr, 0, 4 * T * RN); ++ cx.run({query, res}); ++ check("ple_query_gate_norm + ple_conv_res", same(Ra, Rb, 4 * T * RN) && same(na, nb, 4 * T * RN)); ++ if (small) continue; ++ // ---- GPU times: one PLE stage a chunk, ms, best of 5 x 8 reps ++ auto time = [&](const std::vector& ds) { ++ cx.warm(ds, 0.1); ++ double best = 1e9; ++ for (int i = 0; i < 5; ++i) best = std::min(best, cx.run(ds, 8) / 8); ++ return best * 1e3; ++ }; ++ const double t_f = time({comb, fused, res}), t_fc = time({fused_c, res}), t_q = time({query, res}); ++ std::printf("T %u, ms per PLE stage (best of 5):\n", T); ++ std::printf(" mm_combine + ple_moe_gate_norm + ple_conv_res %.3f -> ple_moe_gate_norm_c + ple_conv_res %.3f (%+.1f%%)\n", ++ t_f, t_fc, 100.0 * (t_fc / t_f - 1)); ++ std::printf(" ple_query_gate_norm + ple_conv_res %.3f\n", t_q); ++ } ++ // ---- ple_hist_adv after ple_conv_res: the conv reads the old history, then the history moves past the T rows (T < 9 ++ // keeps 9 - T old rows); against window.cpp's former host loop ++ for (uint32_t T : {1u, 3u, 7u, 8u, 9u, 10u, 64u, 2048u}) { ++ mtl::Buffer nw = cx.alloc(4 * T * RN), hist = cx.alloc(4 * 9 * RN), kw = cx.alloc(2 * 4 * RN), Ra = cx.alloc(4 * T * RN), ++ Rb = cx.alloc(4 * T * RN); ++ ref::Rng r(97 + T); ++ for (size_t i = 0; i < nw.bytes / 4; ++i) nw.as()[i] = r.sym(); ++ for (size_t i = 0; i < hist.bytes / 4; ++i) hist.as()[i] = r.sym(); ++ for (size_t i = 0; i < Ra.bytes / 4; ++i) Ra.as()[i] = Rb.as()[i] = r.sym(); ++ for (size_t i = 0; i < kw.bytes / 2; ++i) kw.as()[i] = ref::f32_to_f16(0.5f * r.sym()); ++ const PleArgs pl{N, HC, 1e-6f}; ++ cx.run({D(pcr, {mtl::Arg::B(hist), mtl::Arg::B(nw), mtl::Arg::B(kw), mtl::Arg::B(Ra), mtl::Arg::V(pl)}, HC * N / 256, T)}); ++ std::vector want(hist.as(), hist.as() + 9 * RN); ++ for (size_t c = 0; c < RN; ++c) { ++ float col[9]; ++ for (int k = 0; k < 9; ++k) { ++ const int src = (int) T - 9 + k; ++ col[k] = src >= 0 ? nw.as()[(size_t) src * RN + c] : want[c * 9 + (size_t) k + T]; ++ } ++ std::memcpy(want.data() + c * 9, col, sizeof col); ++ } ++ cx.run({D(pcr, {mtl::Arg::B(hist), mtl::Arg::B(nw), mtl::Arg::B(kw), mtl::Arg::B(Rb), mtl::Arg::V(pl)}, HC * N / 256, T), ++ D(pha, {mtl::Arg::B(hist), mtl::Arg::B(nw), mtl::Arg::V(pl), mtl::Arg::V(T)}, HC * N / 256, 1)}); ++ const bool ok = same(Ra, Rb, 4 * T * RN) && std::memcmp(hist.ptr, want.data(), 4 * 9 * RN) == 0; ++ fails += !ok; ++ std::printf("%-34s T %4u: bitwise %s\n", "ple_conv_res + ple_hist_adv", T, ok ? "PASS" : "FAIL"); ++ } ++ std::printf("%s\n", fails ? "FAIL" : "ALL PASS"); ++ return fails ? 1 : 0; ++} +diff --git a/metal/tests/prompt_chain.cpp b/metal/tests/prompt_chain.cpp +index 876afeca..88759e98 100644 +--- a/metal/tests/prompt_chain.cpp ++++ b/metal/tests/prompt_chain.cpp +@@ -413,11 +413,9 @@ int main(int argc, char** argv) { + 256, "router"); + const int nent = T * TOPK, nblk = (nent + 1023) / 1024; + MmPlanArgs pa{(uint32_t)nent, TOPK, NEXP, (uint32_t)TM, (uint32_t)nblk}; +- const uint32_t nclr = maxrows; +- D(P(lib, "mm_clear"), {mtl::Arg::B(stok), mtl::Arg::V(nclr)}, (maxrows + 255) / 256, 1, 256, "experts"); + D(P(lib, "mm_hist"), {mtl::Arg::B(ids), mtl::Arg::B(bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024, "experts"); + D(P(lib, "mm_scan"), {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta), +- mtl::Arg::V(pa)}, 1, 1, NEXP, "experts"); ++ mtl::Arg::V(pa), mtl::Arg::B(stok)}, 1, 1, NEXP, "experts"); + D(P(lib, "mm_place"), {mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok), + mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}, nblk, 1, 1024, "experts"); + MmArgs ma{eg.offset, eu.offset, ed.offset, 0, 0, (uint32_t)ref::row_bytes(eg.type, N), +@@ -426,8 +424,8 @@ int main(int argc, char** argv) { + ma.d_stride = (uint64_t)ma.d_row * N; + const uint32_t ne = N; + mtl::Dispatch g; +- g.p = P(lib, "mm_gather"); +- g.args = {mtl::Arg::B(mixed), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; ++ g.p = P(lib, "mm_gather_h"); ++ g.args = {mtl::Arg::B(mixh), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; + g.tg[0] = 256; + g.indirect = &meta; + g.ind_off = 8 * 4; +diff --git a/metal/tests/qsa_parity.cpp b/metal/tests/qsa_parity.cpp +index 8162ec2e..2897b5ca 100644 +--- a/metal/tests/qsa_parity.cpp ++++ b/metal/tests/qsa_parity.cpp +@@ -4,6 +4,15 @@ + // 2051), then a window of 3 tokens: prep, score, top-k, attention, merge. Each stage is + // checked against the CPU fed with the GPU's inputs to that stage; the ids exactly. + // qsa_parity --bench [P] GPU time of the window's kernels at context P, 1 and 3 tokens ++// qsa_parity --bitwise REF.metal ++// qsa_attn + qsa_merge against the same kernels compiled from REF.metal (e.g. the file ++// before an edit): part and xp memcmp-equal over pos0 0..30000 and 1, 2, 4, 8 tokens on ++// random KV codes and ids, then paired GPU times of both ++// QSA_ATTN=2 tests qsa_attn2 + qsa_merge2; REF is then qsa2.metal ++// qsa_parity --paired REF_DIR [P] [--tokens N] [--prep-only] ++// full-buffer memcmp and paired per-stage timings; prep-only avoids attention allocation ++// qsa_parity --ab REF_DIR [P] ++// paired A/B across decode T = 1, 2, 3, 4, 8 and prompt T = 1984, 2048 (with warmup + 12x reps) + #include "mtl_rt.h" + #include "ref_quant.h" + +@@ -25,17 +34,120 @@ struct Err { + double l2() const { return sr > 0 ? std::sqrt(se / sr) : std::sqrt(se); } + }; + float h2f(uint16_t h) { return ref::f16_to_f32(h); } ++// --bitwise: the attention and merge of this tree against REF's, bit for bit on the same inputs, and their times ++int bitwise(const char* ref_path) { ++ constexpr int MAXT = 8, CELLS = 32768; ++ mtl::Context cx; ++ const bool v2 = std::getenv("QSA_ATTN") && std::atoi(std::getenv("QSA_ATTN")) == 2; ++ const int nch = v2 ? (BUDGET + 31) / 32 : NCH; ++ std::vector src_new{"qsa.metal"}, src_ref; ++ if (v2) { src_new.push_back("qsa2.metal"); src_ref.push_back("qsa.metal"); } ++ src_ref.push_back(ref_path); ++ const int lnew = cx.compile(src_new, "#define NO_GGML_COMMON 1\n"); ++ const int lref = cx.compile(src_ref, "#define NO_GGML_COMMON 1\n"); ++ const char* attn = v2 ? "qsa_attn2" : "qsa_attn"; ++ const char* merge = v2 ? "qsa_merge2" : "qsa_merge"; ++ mtl::Pipeline an = cx.pipeline(lnew, attn), mn = cx.pipeline(lnew, merge), ++ ar = cx.pipeline(lref, attn), mr = cx.pipeline(lref, merge); ++ const size_t part_bytes = 4 * (size_t)MAXT * nch * QH * (QD + 2), xp_bytes = 4 * (4 * MAXT * QH * QD + MAXT * QH * QD / 16); ++ mtl::Buffer qcur = cx.alloc(4 * MAXT * QH * QD), qfull = cx.alloc(4 * MAXT * QH * 2 * QD), ++ kq = cx.alloc((size_t)CELLS * QKV * QD), vq = cx.alloc((size_t)CELLS * QKV * QD), ++ ksc = cx.alloc(2 * (size_t)CELLS * QKV * 4), vsc = cx.alloc(2 * (size_t)CELLS * QKV * 4), ++ ids = cx.alloc(4 * MAXT * BUDGET), part[2] = {cx.alloc(part_bytes), cx.alloc(part_bytes)}, ++ xp[2] = {cx.alloc(xp_bytes), cx.alloc(xp_bytes)}; ++ ref::Rng r(41); ++ for (size_t i = 0; i < qcur.bytes / 4; ++i) qcur.as()[i] = r.sym() * 1.5f; ++ for (size_t i = 0; i < qfull.bytes / 4; ++i) qfull.as()[i] = r.sym() * 2.0f; ++ for (size_t i = 0; i < kq.bytes; ++i) { kq.as()[i] = (int8_t)(r.next() % 255 - 127); vq.as()[i] = (int8_t)(r.next() % 255 - 127); } ++ for (size_t i = 0; i < ksc.bytes / 2; ++i) { ++ ksc.as()[i] = ref::f32_to_f16(0.004f + 0.012f * r.uni()); ++ vsc.as()[i] = ref::f32_to_f16(0.004f + 0.012f * r.uni()); ++ } ++ auto fill_ids = [&](uint32_t pos0, int nt) { ++ for (int t = 0; t < nt; ++t) { ++ const int n_kv = (int)pos0 + t + 1, width = std::min(n_kv, BUDGET); ++ int* o = ids.as() + t * BUDGET; ++ if (n_kv <= BUDGET) { for (int c = 0; c < width; ++c) o[c] = c; continue; } ++ // width distinct cells of [0, n_kv) ascending, as the top-k emits them (folded into the KV buffer) ++ std::vector pick; ++ for (int c = 0, left = width; c < n_kv && left; ++c) ++ if (r.next() % (uint64_t)(n_kv - c) < (uint64_t)left) { pick.push_back(c % CELLS); --left; } ++ std::copy(pick.begin(), pick.end(), o); ++ } ++ }; ++ auto dispatches = [&](int v, uint32_t pos0, int nt) { ++ QsaArgs a{pos0, 0, (uint32_t)CELLS, (uint32_t)nt, 1e-6f}; ++ mtl::Dispatch at, m; ++ at.p = v ? &an : &ar; ++ at.args = {mtl::Arg::B(qcur), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), ++ mtl::Arg::B(ids), mtl::Arg::B(part[v]), mtl::Arg::V(a)}; ++ at.grid[0] = nch; at.grid[1] = v2 ? QKV * 2 : QKV; at.grid[2] = nt; at.tg[0] = 256; ++ m.p = v ? &mn : &mr; ++ m.args = {mtl::Arg::B(part[v]), mtl::Arg::B(qfull), mtl::Arg::B(xp[v]), mtl::Arg::V(a)}; ++ m.grid[0] = QH; m.grid[1] = nt; m.tg[0] = 256; ++ return std::vector{at, m}; ++ }; ++ int fails = 0; ++ for (uint32_t pos0 : {0u, 1u, 63u, 64u, 512u, 2047u, 2048u, 2050u, 8192u, 30000u}) ++ for (int nt : {1, 2, 4, 8}) { ++ fill_ids(pos0, nt); ++ for (int v = 0; v < 2; ++v) { ++ std::memset(part[v].ptr, 0, part_bytes); ++ std::memset(xp[v].ptr, 0, xp_bytes); ++ cx.run(dispatches(v, pos0, nt)); ++ } ++ const bool part_ok = !std::memcmp(part[0].ptr, part[1].ptr, part_bytes); ++ const bool xp_ok = !std::memcmp(xp[0].ptr, xp[1].ptr, xp_bytes); ++ const bool ok = part_ok && xp_ok; ++ fails += !ok; ++ if (!ok) std::printf("pos0 %5u ntok %d: part %s xp %s FAIL\n", pos0, nt, ++ part_ok ? "PASS" : "FAIL", xp_ok ? "PASS" : "FAIL"); ++ } ++ std::printf("bitwise: %d of 40 (pos0, ntok) cases differ %s\n", fails, fails ? "FAIL" : "PASS"); ++ std::printf("device %s: GPU us per dispatch, median of 9 paired rounds of 32 (ref / this)\n", cx.device_name().c_str()); ++ for (uint32_t pos0 : {512u, 2048u, 8192u}) ++ for (int nt : {1, 4}) { ++ fill_ids(pos0, nt); ++ std::vector d[2] = {dispatches(0, pos0, nt), dispatches(1, pos0, nt)}; ++ double us[2][2]; ++ for (int i = 0; i < 2; ++i) { ++ std::vector one[2] = {std::vector(32, d[0][i]), std::vector(32, d[1][i])}; ++ cx.warm(one[0], 0.05); ++ std::vector t[2]; ++ for (int k = 0; k < 9; ++k) ++ for (int turn = 0; turn < 2; ++turn) { ++ const int v = (k + turn) & 1; ++ t[v].push_back(cx.run(one[v]) / 32); ++ } ++ for (int v = 0; v < 2; ++v) us[i][v] = median(t[v]) * 1e6; ++ } ++ std::printf(" pos0 %5u ntok %d attn %6.2f / %6.2f merge %5.2f / %5.2f\n", pos0, nt, us[0][0], us[0][1], ++ us[1][0], us[1][1]); ++ } ++ return fails ? 1 : 0; ++} + } // namespace + + int main(int argc, char** argv) { + bool bench = false; ++ bool ab_suite = false; ++ const char* paired_ref = nullptr; ++ bool prep_only = false; ++ int paired_nt = 3; + int P = 3000; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; ++ if (a == "--prep-only") { prep_only = true; continue; } ++ if (a == "--ab" && i + 1 < argc) { paired_ref = argv[++i]; ab_suite = true; continue; } ++ if (a == "--paired" && i + 1 < argc) { paired_ref = argv[++i]; continue; } ++ if (a == "--tokens" && i + 1 < argc) { paired_nt = std::atoi(argv[++i]); continue; } ++ if (a == "--bitwise" && i + 1 < argc) return bitwise(argv[i + 1]); + if (a == "--bench") bench = true; + else P = std::atoi(argv[i]); + } +- const int MAXT = 8, MAXC = std::max(4096, (P + 64) / 4 * 4 + 64); ++ if ((prep_only && !paired_ref) || (paired_nt > 128 && !prep_only)) return 2; ++ if (paired_ref && (paired_nt < 1 || paired_nt > 2048 || P < 0)) return 2; ++ const int MAXT = ab_suite ? 2048 : std::max(8, paired_nt), MAXC = std::max(4096, (P + MAXT + 64) / 4 * 4 + 64); + mtl::Context cx; + const int lib = cx.compile({"qsa.metal", "qsa2.metal"}, "#define NO_GGML_COMMON 1\n"); + // QSA_TOPK=1: the v1 top-k; QSA_ATTN=1|2|3 (default 3): which attention (2 has its own chunking and merge) +@@ -46,6 +158,7 @@ int main(int argc, char** argv) { + pattn = cx.pipeline(lib, av == 1 ? "qsa_attn" : av == 2 ? "qsa_attn2" : "qsa_attn3"), + pmerge = cx.pipeline(lib, v2 ? "qsa_merge2" : "qsa_merge"); + const int MB = MAXC / 4 + 2; ++ const int AT = prep_only ? 1 : MAXT; + // per-cell inputs for every cell (context + window) + const int NC = P + MAXT; + mtl::Buffer qfull = cx.alloc(4 * (size_t)NC * QH * 2 * QD), kc = cx.alloc(4 * (size_t)NC * QKV * QD), +@@ -55,8 +168,8 @@ int main(int argc, char** argv) { + qidx = cx.alloc(4 * MAXT * IH * ID), kq = cx.alloc((size_t)MAXC * QKV * QD), vq = cx.alloc((size_t)MAXC * QKV * QD), + ksc = cx.alloc(2 * (size_t)MAXC * QKV * 4), vsc = cx.alloc(2 * (size_t)MAXC * QKV * 4), tail = cx.alloc(4 * 4 * ID), + dead = cx.alloc(4 * ID), pooled = cx.alloc(4 * (size_t)MB * ID), scores = cx.alloc(4 * (size_t)MAXT * MB), +- ids = cx.alloc(4 * MAXT * BUDGET), part = cx.alloc(4 * (size_t)MAXT * 2 * NCH * QH * (QD + 2)), +- xp = cx.alloc(4 * (4 * MAXT * QH * QD + MAXT * QH * QD / 16)); ++ ids = cx.alloc(4 * AT * BUDGET), part = cx.alloc(4 * (size_t)AT * 2 * NCH * QH * (QD + 2)), ++ xp = cx.alloc(4 * (4 * AT * QH * QD + AT * QH * QD / 16)); + ref::Rng r(23); + for (size_t i = 0; i < qfull.bytes / 4; ++i) qfull.as()[i] = r.sym() * 2.0f; + for (size_t i = 0; i < kc.bytes / 4; ++i) kc.as()[i] = r.sym() * 2.0f; +@@ -131,6 +244,74 @@ int main(int argc, char** argv) { + m.tg[0] = 256; + ds.push_back(m); + }; ++ if (paired_ref) { ++ const std::string dir = paired_ref; ++ const int lr = cx.compile({dir + "/qsa.metal", dir + "/qsa2.metal"}, "#define NO_GGML_COMMON 1\n"); ++ mtl::Pipeline pr[] = { ++ cx.pipeline(lr, "qsa_prep", mtl::Constants().b(24, mrope)), cx.pipeline(lr, "qsa_score"), ++ cx.pipeline(lr, std::getenv("QSA_TOPK") ? "qsa_topk" : "qsa_topk2"), ++ cx.pipeline(lr, av == 1 ? "qsa_attn" : av == 2 ? "qsa_attn2" : "qsa_attn3"), ++ cx.pipeline(lr, v2 ? "qsa_merge2" : "qsa_merge")}; ++ std::vector outputs{&qcur, &qidx, &kq, &vq, &ksc, &vsc, &tail, &dead, ++ &pooled, &scores, &ids, &part, &xp}; ++ const char* names[] = {"qcur", "qidx", "kq", "vq", "ksc", "vsc", "tail", "dead", ++ "pooled", "scores", "ids", "part", "xp"}; ++ auto run_case = [&](int nt, bool p_only) -> int { ++ std::vector ds[2]; ++ window(P, nt, ds[1]); ++ if (p_only) ds[1].resize(1); ++ ds[0] = ds[1]; ++ for (size_t i = 0; i < ds[0].size(); ++i) ds[0][i].p = &pr[i]; ++ std::vector> expected; ++ int fails = 0; ++ for (int v = 0; v < 2; ++v) { ++ for (auto* b : outputs) std::memset(b->ptr, 0, b->bytes); ++ ref::Rng history(103); ++ for (auto* b : {&tail, &dead, &pooled}) ++ for (size_t i = 0; i < b->bytes / 4; ++i) b->as()[i] = history.sym(); ++ cx.run(ds[v]); ++ for (size_t i = 0; i < outputs.size(); ++i) { ++ auto* b = outputs[i]; ++ const auto* bytes = static_cast(b->ptr); ++ if (!v) expected.emplace_back(bytes, bytes + b->bytes); ++ else { ++ const bool ok = !std::memcmp(bytes, expected[i].data(), b->bytes); ++ if (!ab_suite || !ok) std::printf("%s %s\n", names[i], ok ? "PASS" : "FAIL"); ++ fails += !ok; ++ } ++ } ++ } ++ if (fails) return 1; ++ std::printf("paired P=%d nt=%d mrope=%d attn=%d%s, ref / candidate GPU us\n", P, nt, mrope, av, ++ ab_suite ? " (memcmp PASS)" : ""); ++ const char* stages[] = {"prep", "score", "topk", "attn", "merge"}; ++ for (size_t i = 0; i < ds[0].size(); ++i) { ++ const int repeats = nt > 128 ? 4 : 16, pairs = ab_suite ? 5 : 9; ++ cx.warm({ds[0][i], ds[1][i]}, 0.01); ++ cx.run({ds[0][i]}, repeats); ++ cx.run({ds[1][i]}, repeats); ++ std::vector times[2], ratios; ++ for (int k = 0; k < pairs; ++k) { ++ double t0 = 0, t1 = 0; ++ for (int turn = 0; turn < 2; ++turn) { ++ const int v = (k + turn) & 1; ++ const double dt = cx.run({ds[v][i]}, repeats) * 1e6 / repeats; ++ if (v == 0) t0 = dt; else t1 = dt; ++ times[v].push_back(dt); ++ } ++ ratios.push_back(t1 / t0); ++ } ++ std::printf("%s %.3f / %.3f (ratio %.4f)\n", stages[i], median(times[0]), median(times[1]), median(ratios)); ++ } ++ return 0; ++ }; ++ if (ab_suite) { ++ for (int nt : {1, 2, 3, 4, 8, 1984, 2048}) ++ if (run_case(nt, prep_only || nt > 8)) return 1; ++ return 0; ++ } ++ return run_case(paired_nt, prep_only); ++ } + // the context, one cell at a time + { + std::vector ds; +diff --git a/metal/tests/qsa_prompt_parity.cpp b/metal/tests/qsa_prompt_parity.cpp +index 18e1b3af..02dde834 100644 +--- a/metal/tests/qsa_prompt_parity.cpp ++++ b/metal/tests/qsa_prompt_parity.cpp +@@ -1,7 +1,12 @@ + // metal/tests/qsa_prompt_parity.cpp - qsa_attn_mm (qsa_prompt.metal, the matrix units, fp16 operands) against qsa_attn + // (fp32, itself checked by qsa_parity), both followed by qsa_merge: the attention output of T query tokens on random + // int8 caches, within the budget (all cells, causal) and past it (2051 selected cells per token); then the speed. +-// qsa_prompt_parity [--bench] ++// qsa_prompt_parity [--bench | --flash-compare T pos0 baseline_qsa_prompt.metal [layers pairs]] ++// qsa_prompt_parity --score-compare T pos0 baseline_qsa_prompt.metal baseline_qsa.metal pairs ++// qsa_prompt_parity --topk-compare T pos0 baseline_qsa.metal pairs ++// QSA_TOPK_EDGE=1: ties/signed zero; QSA_FLASH_EDGE=1|2: signed zero/tiny cancellation ++// qsa_prompt_parity --mm-compare attn|score T pos0 baseline_qsa_prompt.metal pairs ++// qsa_prompt_parity --ab T pos0 baseline_qsa_prompt.metal baseline_qsa.metal pairs + #include "coop_probe.h" + #include "mtl_rt.h" + #include "ref_quant.h" +@@ -18,9 +23,391 @@ namespace { + constexpr int QH = 24, QKV = 2, QD = 256, QB = 2051, NCH = 33; + struct QsaArgs { uint32_t pos0, pos_base, max_cells, n_tok; float eps; }; + double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } ++ ++int topk_compare(int T, uint32_t pos0, const std::string& baseline, int pairs) { ++ if (T < 1 || T > 1984 || pos0 > 30000 || pairs < 1 || pairs > 15) return 2; ++ mtl::Context cx; ++ const int lib = cx.compile({"qsa.metal"}, "#define NO_GGML_COMMON 1\n"); ++ const int old = cx.compile({baseline}, "#define NO_GGML_COMMON 1\n"); ++ const uint32_t cells = pos0 + T, stride = cells / 4 + 2; ++ auto scores = cx.alloc(4 * size_t(T) * stride), out = cx.alloc(4 * size_t(T) * QB), ++ ref = cx.alloc(out.bytes); ++ ref::Rng r(5); ++ const bool edge = std::getenv("QSA_TOPK_EDGE") != nullptr; ++ for (int t = 0; t < T; ++t) ++ for (uint32_t b = 0; b < stride; ++b) { ++ float v = std::fabs(r.sym()) * 32.0f; ++ if (edge) { ++ if (t % 4 == 0) v = 1.0f; ++ if (t % 4 == 1) v = float(int(r.next() % 7) - 3); ++ if (t % 4 == 2) v = b & 1 ? -0.0f : 0.0f; ++ } ++ if (!edge && b == (pos0 + t + 1) / 4 && ((pos0 + t + 1) & 3)) v += 1e9f; ++ scores.as()[size_t(t) * stride + b] = v; ++ } ++ std::memset(out.ptr, 0xa5, out.bytes); std::memset(ref.ptr, 0xa5, ref.bytes); ++ auto p = cx.pipeline(lib, "qsa_topk"), pold = cx.pipeline(old, "qsa_topk"); ++ QsaArgs a{pos0, 0, cells, uint32_t(T), 1e-6f}; ++ mtl::Dispatch d; ++ d.p = &pold; ++ d.args = {mtl::Arg::B(scores), mtl::Arg::B(ref), mtl::Arg::V(a)}; ++ d.grid[0] = T; d.tg[0] = 1024; ++ auto candidate = d; ++ candidate.p = &p; candidate.args[1] = mtl::Arg::B(out); ++ cx.run({d}); cx.run({candidate}); ++ const bool equal = std::memcmp(out.ptr, ref.ptr, out.bytes) == 0; ++ std::vector before, after; ++ for (int i = 0; i < pairs; ++i) { ++ if (i & 1) { after.push_back(cx.run({candidate})); before.push_back(cx.run({d})); } ++ else { before.push_back(cx.run({d})); after.push_back(cx.run({candidate})); } ++ } ++ std::printf("device %s; topk T %d pos0 %u memcmp %s before %.6f after %.6f ms\n", ++ cx.device_name().c_str(), T, pos0, equal ? "PASS" : "FAIL", median(before) * 1e3, median(after) * 1e3); ++ for (int i = 0; i < pairs; ++i) ++ std::printf("pair %d before %.6f after %.6f ms\n", i, before[i] * 1e3, after[i] * 1e3); ++ return equal ? 0 : 1; ++} ++ ++int score_compare(int T, uint32_t pos0, const std::string& baseline_prompt, ++ const std::string& baseline_qsa, int pairs) { ++ if (T < 1 || T > 1984 || pos0 > 30000 || pairs < 1 || pairs > 15) return 2; ++ mtl::Context cx; ++ const int lib = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); ++ const int old = cx.compile({baseline_qsa, baseline_prompt}, "#define NO_GGML_COMMON 1\n"); ++ const size_t cells = pos0 + T, nb = cells / 4 + 1, stride = cells / 4 + 2; ++ auto q = cx.alloc(4 * size_t(T) * 4 * 128), k = cx.alloc(4 * nb * 128), dead = cx.alloc(4 * 128); ++ auto scores = cx.alloc(4 * size_t(T) * stride), ref_scores = cx.alloc(scores.bytes), ++ ids = cx.alloc(4 * size_t(T) * QB), ref_ids = cx.alloc(ids.bytes); ++ ref::Rng r(5); ++ for (auto* b : {&q, &k, &dead}) ++ for (size_t i = 0; i < b->bytes / 4; ++i) b->as()[i] = r.sym(); ++ for (auto* b : {&scores, &ref_scores, &ids, &ref_ids}) std::memset(b->ptr, 0xa5, b->bytes); ++ const auto constants = mtl::Constants().b(7, coop_k4(cx)); ++ auto score = cx.pipeline(lib, "qsa_score_mm", constants), ++ old_score = cx.pipeline(old, "qsa_score_mm", constants), ++ topk = cx.pipeline(lib, "qsa_topk"), old_topk = cx.pipeline(old, "qsa_topk"); ++ QsaArgs a{pos0, 0, uint32_t(cells), uint32_t(T), 1e-6f}; ++ mtl::Dispatch ds, dt; ++ ds.p = &old_score; ++ ds.args = {mtl::Arg::B(q), mtl::Arg::B(k), mtl::Arg::B(dead), mtl::Arg::B(ref_scores), mtl::Arg::V(a)}; ++ ds.grid[0] = (nb + 63) / 64; ds.grid[1] = (T + 15) / 16; ds.tg[0] = 128; ++ dt.p = &old_topk; ++ dt.args = {mtl::Arg::B(ref_scores), mtl::Arg::B(ref_ids), mtl::Arg::V(a)}; ++ dt.grid[0] = T; dt.tg[0] = 1024; ++ auto cs = ds, ct = dt; ++ cs.p = &score; cs.args[3] = mtl::Arg::B(scores); ++ ct.p = &topk; ct.args[0] = mtl::Arg::B(scores); ct.args[1] = mtl::Arg::B(ids); ++ cx.run({ds, dt}, 1, false, true); cx.run({cs, ct}, 1, false, true); ++ const bool same_scores = std::memcmp(scores.ptr, ref_scores.ptr, scores.bytes) == 0; ++ const bool same_ids = std::memcmp(ids.ptr, ref_ids.ptr, ids.bytes) == 0; ++ std::printf("device %s; score+topk T %d pos0 %u scores %s qids %s tg %d -> %d bytes\n", ++ cx.device_name().c_str(), T, pos0, same_scores ? "PASS" : "FAIL", same_ids ? "PASS" : "FAIL", ++ old_score.static_tg_mem, score.static_tg_mem); ++ const std::vector before_dispatches[] = {{ds}, {dt}, {ds, dt}}; ++ const std::vector after_dispatches[] = {{cs}, {ct}, {cs, ct}}; ++ for (int stage = 0; stage < 3; ++stage) { ++ const auto& before_dispatch = before_dispatches[stage]; ++ const auto& after_dispatch = after_dispatches[stage]; ++ cx.run(before_dispatch, 1, false, true); cx.run(after_dispatch, 1, false, true); ++ std::vector before, after; ++ for (int i = 0; i < pairs; ++i) { ++ if (i & 1) { ++ after.push_back(cx.run(after_dispatch, 1, false, true)); ++ before.push_back(cx.run(before_dispatch, 1, false, true)); ++ } else { ++ before.push_back(cx.run(before_dispatch, 1, false, true)); ++ after.push_back(cx.run(after_dispatch, 1, false, true)); ++ } ++ } ++ std::printf("%s before %.6f after %.6f ms\n", stage == 0 ? "score" : stage == 1 ? "topk" : "combined", ++ median(before) * 1e3, median(after) * 1e3); ++ for (int i = 0; i < pairs; ++i) ++ std::printf("pair %d before %.6f after %.6f ms\n", i, before[i] * 1e3, after[i] * 1e3); ++ } ++ const bool stable = std::memcmp(scores.ptr, ref_scores.ptr, scores.bytes) == 0 && ++ std::memcmp(ids.ptr, ref_ids.ptr, ids.bytes) == 0; ++ std::printf("post-timing memcmp %s\n", stable ? "PASS" : "FAIL"); ++ return same_scores && same_ids && stable ? 0 : 1; ++} ++ ++int flash_compare(int T, uint32_t pos0, const std::string& baseline, int layers, int pairs) { ++ if (T < 1 || T > 1984 || pos0 > 30000 || layers < 1 || pairs < 1) return 2; ++ mtl::Context cx; ++ const int lib = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); ++ const int old = cx.compile({"qsa.metal", baseline}, "#define NO_GGML_COMMON 1\n"); ++ const size_t cells = pos0 + T, nq = (size_t)T * QH * QD; ++ mtl::Buffer kq = cx.alloc(cells * QKV * QD), vq = cx.alloc(kq.bytes), ++ ks = cx.alloc(2 * cells * QKV * 4), vs = cx.alloc(ks.bytes), ++ qh = cx.alloc(2 * nq), gate = cx.alloc(8 * nq), ids = cx.alloc(4 * (size_t)T * QB), ++ out = cx.alloc(2 * nq), ref = cx.alloc(out.bytes); ++ ref::Rng r(5); ++ for (size_t i = 0; i < kq.bytes; ++i) { ++ kq.as()[i] = int8_t(int(r.next() % 255) - 127); ++ vq.as()[i] = int8_t(int(r.next() % 255) - 127); ++ } ++ for (size_t i = 0; i < ks.bytes / 2; ++i) { ++ ks.as()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); ++ vs.as()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); ++ } ++ for (size_t i = 0; i < nq; ++i) qh.as()[i] = ref::f32_to_f16(r.sym() * 3.0f); ++ for (size_t i = 0; i < 2 * nq; ++i) gate.as()[i] = r.sym(); ++ for (int t = 0; t < T; ++t) { ++ const uint32_t n = pos0 + t + 1; ++ int* o = ids.as() + (size_t)t * QB; ++ if (n <= QB) { ++ for (uint32_t c = 0; c < n; ++c) o[c] = int(c); ++ } else { ++ std::set selected; ++ while (selected.size() < QB) selected.insert(int(r.next() % n)); ++ std::copy(selected.begin(), selected.end(), o); ++ } ++ } ++ if (const char* edge = std::getenv("QSA_FLASH_EDGE")) { ++ const int mode = std::atoi(edge); ++ for (size_t i = 0; i < kq.bytes; ++i) { ++ kq.as()[i] = 0; ++ vq.as()[i] = mode == 2 ? int8_t((i / (QKV * QD)) & 1 ? -1 : 1) : 0; ++ } ++ for (size_t i = 0; i < ks.bytes / 2; ++i) { ++ ks.as()[i] = i & 1 ? 0x8000 : 0; ++ vs.as()[i] = mode == 2 ? 0x0400 : (i & 1 ? 0x8000 : 0); ++ } ++ for (size_t i = 0; i < nq; ++i) qh.as()[i] = i & 1 ? 0x8000 : 0; ++ } ++ const auto constants = mtl::Constants().b(7, coop_k4(cx)); ++ auto p = cx.pipeline(lib, "qsa_flash", constants), pold = cx.pipeline(old, "qsa_flash", constants); ++ QsaArgs a{pos0, 0, uint32_t(cells), uint32_t(T), 1e-6f}; ++ mtl::Dispatch d; ++ d.p = &pold; ++ d.args = {mtl::Arg::B(qh), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ks), mtl::Arg::B(vs), ++ mtl::Arg::B(ids), mtl::Arg::B(gate), mtl::Arg::B(ref), mtl::Arg::V(a)}; ++ d.grid[0] = QKV; d.grid[1] = T; d.tg[0] = 128; ++ cx.run({d}); ++ // Poison only implicit rows, including the early rows of a chunk crossing the budget. ++ for (int t = 0; t < T && pos0 + t + 1 <= QB; ++t) ++ std::fill_n(ids.as() + (size_t)t * QB, QB, -1); ++ auto candidate = d; ++ candidate.p = &p; candidate.args[7] = mtl::Arg::B(out); ++ cx.run({candidate}); ++ const bool equal = std::memcmp(out.ptr, ref.ptr, out.bytes) == 0; ++ // Baseline needs explicit ids for the paired timings below. ++ for (int t = 0; t < T && pos0 + t + 1 <= QB; ++t) ++ for (uint32_t c = 0; c < pos0 + t + 1; ++c) ids.as()[(size_t)t * QB + c] = int(c); ++ // paired, alternating order: each sample is one command buffer of `layers` back-to-back dispatches (the 12 QSA ++ // layers of a prompt chunk) ++ std::vector before, after; ++ for (int i = 0; i < pairs; ++i) { ++ if (i & 1) { after.push_back(cx.run({candidate}, layers)); before.push_back(cx.run({d}, layers)); } ++ else { before.push_back(cx.run({d}, layers)); after.push_back(cx.run({candidate}, layers)); } ++ } ++ std::printf("device %s; T %d pos0 %u: memcmp %s; flash x%d before median %.3f min %.3f ms, after median %.3f min %.3f " ++ "ms; tg %d -> %d bytes\n", cx.device_name().c_str(), T, pos0, equal ? "PASS" : "FAIL", layers, ++ median(before) * 1e3, *std::min_element(before.begin(), before.end()) * 1e3, median(after) * 1e3, ++ *std::min_element(after.begin(), after.end()) * 1e3, pold.static_tg_mem, p.static_tg_mem); ++ for (int i = 0; i < pairs; ++i) ++ std::printf("pair %d before %.6f after %.6f ms\n", i, before[i] * 1e3, after[i] * 1e3); ++ return equal ? 0 : 1; ++} ++ ++int mm_compare(bool score, int T, uint32_t pos0, const std::string& baseline, int pairs) { ++ if (T < 1 || T > 1984 || pos0 > 30000 || pairs < 1 || pairs > 15) return 2; ++ mtl::Context cx; ++ const int lib = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); ++ const int old = cx.compile({"qsa.metal", baseline}, "#define NO_GGML_COMMON 1\n"); ++ const size_t cells = pos0 + T; ++ QsaArgs a{pos0, 0, uint32_t(cells), uint32_t(T), 1e-6f}; ++ const bool k4 = coop_k4(cx); ++ const auto constants = mtl::Constants().b(7, k4); ++ std::printf("device %s coop_k4 %d\n", cx.device_name().c_str(), int(k4)); ++ const std::string name = score ? "qsa_score_mm" : "qsa_attn_mm"; ++ auto p = cx.pipeline(lib, name.c_str(), constants), pold = cx.pipeline(old, name.c_str(), constants); ++ mtl::Dispatch d; ++ d.p = &pold; ++ mtl::Buffer out, ref; ++ auto compare = [&](const mtl::Dispatch& before_dispatch, const mtl::Dispatch& after_dispatch) { ++ cx.run({before_dispatch}); cx.run({after_dispatch}); ++ const bool equal = std::memcmp(out.ptr, ref.ptr, out.bytes) == 0; ++ cx.warm({before_dispatch, after_dispatch}, 0.05); ++ std::vector before, after; ++ for (int i = 0; i < pairs; ++i) { ++ if (i & 1) { ++ after.push_back(cx.run({after_dispatch})); before.push_back(cx.run({before_dispatch})); ++ } else { ++ before.push_back(cx.run({before_dispatch})); after.push_back(cx.run({after_dispatch})); ++ } ++ } ++ std::printf("%s T %d pos0 %u memcmp %s before %.6f after %.6f ms tg %d -> %d\n", name.c_str(), T, pos0, ++ equal ? "PASS" : "FAIL", median(before) * 1e3, median(after) * 1e3, pold.static_tg_mem, p.static_tg_mem); ++ for (int i = 0; i < pairs; ++i) ++ std::printf("pair %d before %.6f after %.6f ms\n", i, before[i] * 1e3, after[i] * 1e3); ++ const bool stable = std::memcmp(out.ptr, ref.ptr, out.bytes) == 0; ++ std::printf("post-timing memcmp %s\n", stable ? "PASS" : "FAIL"); ++ return equal && stable ? 0 : 1; ++ }; ++ ref::Rng r(5); ++ if (score) { ++ const size_t nb = cells / 4 + 1, stride = cells / 4 + 2; ++ auto q = cx.alloc(4 * size_t(T) * 4 * 128), k = cx.alloc(4 * nb * 128), dead = cx.alloc(4 * 128); ++ for (auto* b : {&q, &k, &dead}) ++ for (size_t i = 0; i < b->bytes / 4; ++i) b->as()[i] = r.sym(); ++ out = cx.alloc(4 * size_t(T) * stride); ref = cx.alloc(out.bytes); ++ std::memset(out.ptr, 0xa5, out.bytes); std::memset(ref.ptr, 0xa5, ref.bytes); ++ d.args = {mtl::Arg::B(q), mtl::Arg::B(k), mtl::Arg::B(dead), mtl::Arg::B(ref), mtl::Arg::V(a)}; ++ d.grid[0] = (nb + 63) / 64; d.grid[1] = (T + 15) / 16; d.tg[0] = 128; ++ auto candidate = d; ++ candidate.p = &p; candidate.args[3] = mtl::Arg::B(out); ++ return compare(d, candidate); ++ } ++ auto k = cx.alloc(cells * QKV * QD), v = cx.alloc(k.bytes), ks = cx.alloc(2 * cells * QKV * 4), vs = cx.alloc(ks.bytes), ++ q = cx.alloc(2 * size_t(T) * QH * QD), ids = cx.alloc(4 * size_t(T) * QB); ++ for (size_t i = 0; i < k.bytes; ++i) { ++ k.as()[i] = int8_t(int(r.next() % 255) - 127); v.as()[i] = int8_t(int(r.next() % 255) - 127); ++ } ++ for (auto* b : {&ks, &vs}) ++ for (size_t i = 0; i < b->bytes / 2; ++i) b->as()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); ++ for (size_t i = 0; i < q.bytes / 2; ++i) q.as()[i] = ref::f32_to_f16(r.sym() * 3.0f); ++ for (int t = 0; t < T; ++t) { ++ const uint32_t n = pos0 + t + 1; ++ for (int c = 0; c < QB; ++c) ids.as()[size_t(t) * QB + c] = n <= QB ? -1 : int(uint64_t(c) * n / QB); ++ } ++ out = cx.alloc(4 * size_t(T) * NCH * QH * (QD + 2)); ref = cx.alloc(out.bytes); ++ std::memset(out.ptr, 0xa5, out.bytes); std::memset(ref.ptr, 0xa5, ref.bytes); ++ d.args = {mtl::Arg::B(q), mtl::Arg::B(k), mtl::Arg::B(v), mtl::Arg::B(ks), mtl::Arg::B(vs), ++ mtl::Arg::B(ids), mtl::Arg::B(ref), mtl::Arg::V(a)}; ++ d.grid[0] = NCH; d.grid[1] = QKV; d.grid[2] = T; d.tg[0] = 32; ++ auto candidate = d; ++ candidate.p = &p; candidate.args[6] = mtl::Arg::B(out); ++ const int result = compare(d, candidate); ++ const size_t nq = size_t(T) * QH * QD; ++ auto gate = cx.alloc(8 * nq), merged = cx.alloc(4 * (4 * nq + nq / 16)), merged_ref = cx.alloc(merged.bytes); ++ for (size_t i = 0; i < gate.bytes / 4; ++i) gate.as()[i] = r.sym(); ++ auto merge = cx.pipeline(lib, "qsa_merge"); ++ mtl::Dispatch dm; ++ dm.p = &merge; ++ dm.args = {mtl::Arg::B(ref), mtl::Arg::B(gate), mtl::Arg::B(merged_ref), mtl::Arg::V(a)}; ++ dm.grid[0] = QH; dm.grid[1] = T; dm.tg[0] = 256; ++ auto cm = dm; ++ cm.args[0] = mtl::Arg::B(out); cm.args[2] = mtl::Arg::B(merged); ++ cx.run({dm}); cx.run({cm}); ++ const bool same_merge = std::memcmp(merged.ptr, merged_ref.ptr, merged.bytes) == 0; ++ std::printf("merge+prepared memcmp %s\n", same_merge ? "PASS" : "FAIL"); ++ return result || !same_merge; ++} + } // namespace + ++int ab_compare(int T, uint32_t pos0, const std::string& baseline_prompt, const std::string& baseline_qsa, int pairs) { ++ if (T < 1 || T > 2048 || pos0 > 30000 || pairs < 1 || pairs > 15) return 2; ++ mtl::Context cx; ++ const int lib = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); ++ const int old = cx.compile({baseline_qsa, baseline_prompt}, "#define NO_GGML_COMMON 1\n"); ++ const size_t cells = pos0 + T, nb = cells / 4 + 1, stride = cells / 4 + 2; ++ QsaArgs a{pos0, 0, uint32_t(cells), uint32_t(T), 1e-6f}; ++ const bool mrope = std::getenv("QSA_MROPE") != nullptr; ++ const bool qf16 = std::getenv("QSA_QF16") != nullptr; ++ const auto prep_c = mtl::Constants().b(24, mrope).b(25, qf16); ++ const auto score_c = mtl::Constants().b(7, coop_k4(cx)); ++ auto prep = cx.pipeline(lib, "qsa_prep", prep_c), old_prep = cx.pipeline(old, "qsa_prep", prep_c); ++ auto score = cx.pipeline(lib, "qsa_score_mm", score_c), old_score = cx.pipeline(old, "qsa_score_mm", score_c); ++ mtl::Buffer qfull = cx.alloc(4 * size_t(T) * QH * 2 * QD), kc = cx.alloc(4 * size_t(T) * QKV * QD), ++ vc = cx.alloc(4 * size_t(T) * QKV * QD), iq = cx.alloc(4 * size_t(T) * 4 * 128), ++ ik = cx.alloc(4 * size_t(T) * 128), qn = cx.alloc(4 * QD), kn = cx.alloc(4 * QD), ++ iqn = cx.alloc(4 * 128), ikn = cx.alloc(4 * 128), cs = cx.alloc(8 * (cells + 16) * 32), ++ mtab = cx.alloc(4 * 3 * (cells + 16)); ++ mtl::Buffer qcur[2] = {cx.alloc((qf16 ? 2 : 4) * size_t(T) * QH * QD), cx.alloc((qf16 ? 2 : 4) * size_t(T) * QH * QD)}, ++ qidx[2] = {cx.alloc(4 * size_t(T) * 4 * 128), cx.alloc(4 * size_t(T) * 4 * 128)}, ++ kq[2] = {cx.alloc(cells * QKV * QD), cx.alloc(cells * QKV * QD)}, ++ vq[2] = {cx.alloc(cells * QKV * QD), cx.alloc(cells * QKV * QD)}, ++ ksc[2] = {cx.alloc(2 * cells * QKV * 4), cx.alloc(2 * cells * QKV * 4)}, ++ vsc[2] = {cx.alloc(2 * cells * QKV * 4), cx.alloc(2 * cells * QKV * 4)}, ++ tail[2] = {cx.alloc(4 * 4 * 128), cx.alloc(4 * 4 * 128)}, ++ dead[2] = {cx.alloc(4 * 128), cx.alloc(4 * 128)}, ++ pooled[2] = {cx.alloc(4 * stride * 128), cx.alloc(4 * stride * 128)}, ++ scores[2] = {cx.alloc(4 * size_t(T) * stride), cx.alloc(4 * size_t(T) * stride)}; ++ ref::Rng r(5); ++ for (auto* b : {&qfull, &kc, &vc, &iq, &ik}) ++ for (size_t i = 0; i < b->bytes / 4; ++i) b->as()[i] = r.sym(); ++ for (int i = 0; i < QD; ++i) { qn.as()[i] = 1.0f + 0.2f * r.sym(); kn.as()[i] = 1.0f + 0.2f * r.sym(); } ++ for (int i = 0; i < 128; ++i) { iqn.as()[i] = 1.0f + 0.2f * r.sym(); ikn.as()[i] = 1.0f + 0.2f * r.sym(); } ++ for (size_t p = 0; p < cells + 16; ++p) { ++ for (int k = 0; k < 3; ++k) mtab.as()[p * 3 + k] = int(r.next() % (cells + 16)); ++ for (int i = 0; i < 32; ++i) { ++ const double inv = std::pow(1e7, -2.0 * i / 64.0), th = double(p) * inv; ++ cs.as()[(p * 32 + i) * 2] = float(std::cos(th)); ++ cs.as()[(p * 32 + i) * 2 + 1] = float(std::sin(th)); ++ } ++ } ++ auto reset_state = [&](int v) { ++ ref::Rng h(103); ++ for (auto* b : {&qcur[v], &qidx[v], &kq[v], &vq[v], &ksc[v], &vsc[v], &scores[v]}) ++ std::memset(b->ptr, 0xa5, b->bytes); ++ for (auto* b : {&tail[v], &dead[v], &pooled[v]}) ++ for (size_t i = 0; i < b->bytes / 4; ++i) b->as()[i] = h.sym(); ++ }; ++ mtl::Dispatch dp[2], ds[2]; ++ for (int v = 0; v < 2; ++v) { ++ dp[v].p = v ? &prep : &old_prep; ++ dp[v].args = {mtl::Arg::B(qfull), mtl::Arg::B(kc), mtl::Arg::B(vc), mtl::Arg::B(iq), mtl::Arg::B(ik), ++ mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(iqn), mtl::Arg::B(ikn), mtl::Arg::B(cs), ++ mtl::Arg::B(qcur[v]), mtl::Arg::B(qidx[v]), mtl::Arg::B(kq[v]), mtl::Arg::B(vq[v]), ++ mtl::Arg::B(ksc[v]), mtl::Arg::B(vsc[v]), mtl::Arg::B(tail[v]), mtl::Arg::B(dead[v]), ++ mtl::Arg::B(pooled[v]), mtl::Arg::V(a), mtl::Arg::B(mtab)}; ++ dp[v].grid[0] = 33; dp[v].grid[1] = T; dp[v].tg[0] = 256; ++ ds[v].p = v ? &score : &old_score; ++ ds[v].args = {mtl::Arg::B(qidx[v]), mtl::Arg::B(pooled[v]), mtl::Arg::B(dead[v]), mtl::Arg::B(scores[v]), mtl::Arg::V(a)}; ++ ds[v].grid[0] = (nb + 63) / 64; ds[v].grid[1] = (T + 15) / 16; ds[v].tg[0] = 128; ++ reset_state(v); ++ cx.run({dp[v], ds[v]}, 1, false, true); ++ } ++ const mtl::Buffer* b0[] = {&qcur[0], &qidx[0], &kq[0], &vq[0], &ksc[0], &vsc[0], &tail[0], &dead[0], &pooled[0], &scores[0]}; ++ const mtl::Buffer* b1[] = {&qcur[1], &qidx[1], &kq[1], &vq[1], &ksc[1], &vsc[1], &tail[1], &dead[1], &pooled[1], &scores[1]}; ++ bool ok = true; ++ for (int i = 0; i < 10; ++i) ok = ok && (std::memcmp(b0[i]->ptr, b1[i]->ptr, b0[i]->bytes) == 0); ++ std::printf("device %s; ab T %d pos0 %u mrope %d qf16 %d memcmp %s\n", ++ cx.device_name().c_str(), T, pos0, int(mrope), int(qf16), ok ? "PASS" : "FAIL"); ++ if (!ok) return 1; ++ const char* sname[] = {"qsa_prep", "qsa_score_mm", "combined"}; ++ const std::vector d0[] = {{dp[0]}, {ds[0]}, {dp[0], ds[0]}}; ++ const std::vector d1[] = {{dp[1]}, {ds[1]}, {dp[1], ds[1]}}; ++ for (int s = 0; s < 3; ++s) { ++ cx.warm(d0[s], 0.05); ++ cx.run(d0[s], 12, false, true); ++ cx.run(d1[s], 12, false, true); ++ std::vector before, after, ratios; ++ for (int i = 0; i < pairs; ++i) { ++ double t0 = 0, t1 = 0; ++ if (i & 1) { ++ t1 = cx.run(d1[s], 12, false, true) * 1e3 / 12.0; ++ t0 = cx.run(d0[s], 12, false, true) * 1e3 / 12.0; ++ } else { ++ t0 = cx.run(d0[s], 12, false, true) * 1e3 / 12.0; ++ t1 = cx.run(d1[s], 12, false, true) * 1e3 / 12.0; ++ } ++ before.push_back(t0); after.push_back(t1); ratios.push_back(t1 / t0); ++ } ++ std::printf(" %-12s before %.4f after %.4f ms (ratio %.4f)\n", sname[s], median(before), median(after), median(ratios)); ++ } ++ return 0; ++} ++ + int main(int argc, char** argv) { ++ if (argc == 7 && std::string(argv[1]) == "--ab") ++ return ab_compare(std::atoi(argv[2]), uint32_t(std::strtoul(argv[3], nullptr, 10)), ++ argv[4], argv[5], std::atoi(argv[6])); ++ if (argc == 7 && std::string(argv[1]) == "--score-compare") ++ return score_compare(std::atoi(argv[2]), uint32_t(std::strtoul(argv[3], nullptr, 10)), ++ argv[4], argv[5], std::atoi(argv[6])); ++ if (argc == 6 && std::string(argv[1]) == "--topk-compare") ++ return topk_compare(std::atoi(argv[2]), uint32_t(std::strtoul(argv[3], nullptr, 10)), ++ argv[4], std::atoi(argv[5])); ++ if (argc == 7 && std::string(argv[1]) == "--mm-compare") { ++ const std::string target = argv[2]; ++ if (target != "attn" && target != "score") return 2; ++ return mm_compare(target == "score", std::atoi(argv[3]), ++ uint32_t(std::strtoul(argv[4], nullptr, 10)), argv[5], std::atoi(argv[6])); ++ } ++ if ((argc == 5 || argc == 7) && std::string(argv[1]) == "--flash-compare") ++ return flash_compare(std::atoi(argv[2]), uint32_t(std::strtoul(argv[3], nullptr, 10)), argv[4], ++ argc == 7 ? std::atoi(argv[5]) : 1, argc == 7 ? std::atoi(argv[6]) : 3); + const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; + std::setvbuf(stdout, nullptr, _IOLBF, 0); + const bool noflash = std::getenv("QSA_NOFLASH") != nullptr; +@@ -44,7 +431,7 @@ int main(int argc, char** argv) { + std::printf("pipelines: qsa_attn_mm max threads %d, tg mem %d; qsa_flash max threads %d, tg mem %d\n", pmm.max_threads, + pmm.static_tg_mem, pf.max_threads, pf.static_tg_mem); + int fails = 0; +- for (uint32_t pos0 : {0u, 30000u}) { ++ for (uint32_t pos0 : {0u, 1984u, 2050u, 30000u}) { + // ids: within the budget the identity; past it 2051 distinct ascending cells < n_kv + for (int t = 0; t < T; ++t) { + const uint32_t n_kv = pos0 + t + 1; +@@ -82,6 +469,8 @@ int main(int argc, char** argv) { + } + cx.run(ref_ds); + std::fill(part.as(), part.as() + part.bytes / 4, std::nanf("")); // a kernel that writes nothing shows ++ for (int t = 0; t < T && pos0 + t + 1 <= QB; ++t) ++ std::fill_n(ids.as() + (size_t)t * QB, QB, -1); + cx.run(mm_ds); + if (std::getenv("QSA_DEBUG")) { + // how much of the partials the matrix-unit kernel wrote, and token 5's first chunk (m, l, acc[0..3]) +@@ -155,6 +544,77 @@ int main(int argc, char** argv) { + std::printf("score pos0 %u, T %d: qsa_score_mm vs qsa_score worst rel %.2e %s\n", pos0, T, worst, ok ? "PASS" : "FAIL"); + } + } ++ // qsa_prep with QSA_Q_F16_FC=true vs qsa_prep (fp32) + qsa_q_f16: 100% bitwise parity on qcurh and qsa_flash output ++ { ++ const uint32_t pos0 = 0, MB = MAXC / 4 + 2; ++ mtl::Buffer kc = cx.alloc(4 * (size_t)T * QKV * QD), vc = cx.alloc(4 * (size_t)T * QKV * QD), ++ iq = cx.alloc(4 * (size_t)T * 4 * 128), ik = cx.alloc(4 * (size_t)T * 128), ++ qn = cx.alloc(4 * QD), kn = cx.alloc(4 * QD), iqn = cx.alloc(4 * 128), ikn = cx.alloc(4 * 128), ++ cs = cx.alloc(8 * (size_t)(T + 16) * 32), qidx = cx.alloc(4 * (size_t)T * 4 * 128), ++ tail = cx.alloc(4 * 4 * 128), dead = cx.alloc(4 * 128), pooled = cx.alloc(4 * (size_t)MB * 128), ++ mtab = cx.alloc(4 * 3 * (size_t)(T + 16)), ++ qh_ref = cx.alloc(2 * (size_t)T * QH * QD), qh_fused = cx.alloc(2 * (size_t)T * QH * QD), ++ xf_ref = cx.alloc(2 * (size_t)T * QH * QD), xf_fused = cx.alloc(2 * (size_t)T * QH * QD); ++ for (size_t i = 0; i < kc.bytes / 4; ++i) { kc.as()[i] = r.sym(); vc.as()[i] = r.sym(); } ++ for (size_t i = 0; i < iq.bytes / 4; ++i) iq.as()[i] = r.sym(); ++ for (size_t i = 0; i < ik.bytes / 4; ++i) ik.as()[i] = r.sym(); ++ for (int i = 0; i < QD; ++i) { qn.as()[i] = 1.0f + 0.2f * r.sym(); kn.as()[i] = 1.0f + 0.2f * r.sym(); } ++ for (int i = 0; i < 128; ++i) { iqn.as()[i] = 1.0f + 0.2f * r.sym(); ikn.as()[i] = 1.0f + 0.2f * r.sym(); } ++ for (int p = 0; p < T + 16; ++p) ++ for (int i = 0; i < 32; ++i) { ++ const double inv = std::pow(1e7, -2.0 * i / 64.0), th = (double)p * inv; ++ cs.as()[(p * 32 + i) * 2] = (float)std::cos(th); ++ cs.as()[(p * 32 + i) * 2 + 1] = (float)std::sin(th); ++ } ++ QsaArgs a{pos0, 0, (uint32_t)MAXC, (uint32_t)T, 1e-6f}; ++ const uint32_t nq = T * QH * QD; ++ mtl::Pipeline pprep_f32 = cx.pipeline(lib, "qsa_prep"), ++ pprep_f16 = cx.pipeline(lib, "qsa_prep", mtl::Constants().b(25, true)); ++ auto make_prep = [&](const mtl::Pipeline& pp, const mtl::Buffer& qout) { ++ mtl::Dispatch d; ++ d.p = &pp; ++ d.args = {mtl::Arg::B(qfull), mtl::Arg::B(kc), mtl::Arg::B(vc), mtl::Arg::B(iq), mtl::Arg::B(ik), ++ mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(iqn), mtl::Arg::B(ikn), mtl::Arg::B(cs), ++ mtl::Arg::B(qout), mtl::Arg::B(qidx), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), ++ mtl::Arg::B(vsc), mtl::Arg::B(tail), mtl::Arg::B(dead), mtl::Arg::B(pooled), mtl::Arg::V(a), ++ mtl::Arg::B(mtab)}; ++ d.grid[0] = 33; d.grid[1] = T; d.tg[0] = 256; ++ return d; ++ }; ++ mtl::Dispatch d_prep_f32 = make_prep(pprep_f32, qcur); ++ mtl::Dispatch d_qf16; d_qf16.p = &pq; d_qf16.args = {mtl::Arg::B(qcur), mtl::Arg::B(qh_ref), mtl::Arg::V(nq)}; ++ d_qf16.grid[0] = (nq + 255) / 256; d_qf16.tg[0] = 256; ++ mtl::Dispatch d_prep_f16 = make_prep(pprep_f16, qh_fused); ++ for (int t = 0; t < T; ++t) { ++ int* o = ids.as() + (size_t)t * QB; ++ for (int c = 0; c <= t && c < QB; ++c) o[c] = c; ++ } ++ mtl::Dispatch d_fl_ref; d_fl_ref.p = &pf; ++ d_fl_ref.args = {mtl::Arg::B(qh_ref), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), ++ mtl::Arg::B(ids), mtl::Arg::B(qfull), mtl::Arg::B(xf_ref), mtl::Arg::V(a)}; ++ d_fl_ref.grid[0] = QKV; d_fl_ref.grid[1] = T; d_fl_ref.tg[0] = 128; ++ mtl::Dispatch d_fl_fused = d_fl_ref; ++ d_fl_fused.args[0] = mtl::Arg::B(qh_fused); ++ d_fl_fused.args[7] = mtl::Arg::B(xf_fused); ++ if (bench) { ++ auto timeit = [&](const std::vector& ds) { ++ cx.warm(ds, 0.3); ++ std::vector t; ++ for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds, 1, false, true)); ++ return median(t); ++ }; ++ std::printf("prep T %d: qsa_prep(f32)+qsa_q_f16 %.3f ms (prep %.3f ms, q_f16 %.3f ms), qsa_prep(qf16) %.3f ms\n", ++ T, timeit({d_prep_f32, d_qf16}) * 1e3, timeit({d_prep_f32}) * 1e3, timeit({d_qf16}) * 1e3, timeit({d_prep_f16}) * 1e3); ++ } else { ++ cx.run({d_prep_f32, d_qf16, d_fl_ref}, 1, false, true); ++ cx.run({d_prep_f16, d_fl_fused}, 1, false, true); ++ const bool ok_qh = std::memcmp(qh_ref.ptr, qh_fused.ptr, qh_ref.bytes) == 0; ++ const bool ok_xf = std::memcmp(xf_ref.ptr, xf_fused.ptr, xf_ref.bytes) == 0; ++ fails += (!ok_qh || !ok_xf); ++ std::printf("qsa_prep(QSA_Q_F16_FC) vs qsa_prep+qsa_q_f16: qcurh memcmp %s, qsa_flash memcmp %s\n", ++ ok_qh ? "PASS" : "FAIL", ok_xf ? "PASS" : "FAIL"); ++ } ++ } + if (!bench) std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; + } +diff --git a/metal/tests/route_parity.cpp b/metal/tests/route_parity.cpp +index 2f72329a..f5513183 100644 +--- a/metal/tests/route_parity.cpp ++++ b/metal/tests/route_parity.cpp +@@ -1,5 +1,10 @@ + // metal/tests/route_parity.cpp - moe_route (softmax + top-10 + the expert plan) against a CPU replica of +-// native_router.cu's rule, for 1..8 tokens; logits with deliberate ties. ++// native_router.cu's rule, for 1..8 tokens: ++// Phase 1: 200 windows with n_expert = 512, pad = 0 (deliberate exact ties) ++// Phase 2: 200 windows with n_expert = 256, pad = bit_cast(exp(0.85f)), random smap residency, and huge traps in 256..511 ++// route_parity --router-ab BASE_KERNEL_DIR TYPE NTOK SPLIT bitwise fused-router A/B + short paired timing ++// route_parity --mtp-ab BASE_MTP_METAL bitwise mtp.metal (mtp_rms_xprep, mtp_xprep, mtp_combine_gnorm) A/B + timing ++// route_parity --ab ref.metal [--bench] bitwise parity (moe_route + moe_plan, kSplitRes 0/1) + paired A/B + #include "mtl_rt.h" + #include "ref_quant.h" + +@@ -7,27 +12,341 @@ + #include + #include + #include ++#include + #include + #include ++#include ++#include + #include + + namespace { + constexpr int P_NG = 0, P_MAXE = 128, P_GEXP = 16, P_GSTART = P_GEXP + P_MAXE, P_ETOK = P_GSTART + P_MAXE + 4, +- P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_INTS = P_TOKENT + P_MAXE; ++ P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_INTS = P_TOKENT + P_MAXE + 20 + P_MAXE; + struct MoeArgs { + uint64_t gate_off, up_off, down_off, gu_stride, d_stride; + uint32_t gu_row, d_row, n_embd, n_ff, n_tok, topk, n_expert, gx_gu, gx_dn, pad; + }; ++double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } ++struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; ++ ++// One specialization per invocation keeps the GPU checks short. ++int router_ab(int argc, char** argv) { ++ if (argc != 6) return 2; ++ const int type = std::atoi(argv[3]), nt = std::atoi(argv[4]), split = std::atoi(argv[5]); ++ if ((type != ref::BF16 && type != ref::F32 && type != ref::Q8_0) || ++ nt < 1 || nt > 8 || split < 0 || split > 1) return 2; ++ mtl::Context cx; ++ auto compile = [&](const std::string& dir) { ++ return cx.compile({dir + "/quant.metal", dir + "/units.metal", dir + "/moe.metal", ++ dir + "/router.metal"}, ++ "#include \nusing namespace metal;\n" ++ "struct GemvArgs { uint n_in, n_out, n_tok, row_bytes, x_stride, y_stride; };\n" ++ "constant int NTOK [[function_constant(0)]];\n"); ++ }; ++ const int baseline = compile(argv[2]), candidate = compile(cx.kernel_dir); ++ mtl::Constants fc; ++ fc.i(0, nt).b(26, split != 0); ++ const std::string name = "router_route_" + std::to_string(type); ++ auto pb = cx.pipeline(baseline, name, fc), pc = cx.pipeline(candidate, name, fc); ++ constexpr int ni = 2560; ++ const int rb = int(ref::row_bytes(type, ni)); ++ auto W = cx.alloc(512 * rb), x = cx.alloc(4 * nt * ni), sums = cx.alloc(4 * nt * ni / 16); ++ auto smap = cx.alloc(4 * 512); ++ mtl::Buffer lg[2], ids[2], w[2], plan[2], counter[2]; ++ for (int j = 0; j < 2; ++j) { ++ lg[j] = cx.alloc(4 * nt * 512); ids[j] = cx.alloc(4 * nt * 10); ++ w[j] = cx.alloc(4 * nt * 10); plan[j] = cx.alloc(4 * P_INTS); counter[j] = cx.alloc(4); ++ } ++ ref::Rng rng(928); ++ mtl::Dispatch d[2]; ++ for (int trial = 0; trial < 4; ++trial) { ++ ref::random_blocks(type, W.as(), 512, ni, 928 + trial); ++ for (int i = 0; i < nt * ni; ++i) x.as()[i] = rng.sym(); ++ for (int e = 0; e < 512; ++e) smap.as()[e] = rng.sym() > 0 ? e : 0xFFFFFFFFu; ++ if (trial == 2) for (int e = 1; e < 512; ++e) ++ std::memcpy(W.as() + e * rb, W.ptr, rb); ++ MoeArgs a{}; ++ a.n_tok = nt; a.topk = 10; a.n_expert = trial == 1 ? 256 : 512; ++ a.gx_gu = 5; a.gx_dn = 20; ++ const float scale = std::exp(0.85f); ++ if (trial & 1) std::memcpy(&a.pad, &scale, 4); ++ if (trial == 3) a.pad |= 0x80000000u; ++ GemvArgs ga{ni, a.n_expert, uint32_t(nt), uint32_t(rb), ni, 512}; ++ for (int j = 0; j < 2; ++j) { ++ for (auto b : {lg[j], ids[j], w[j], plan[j]}) std::memset(b.ptr, 0xa5, b.bytes); ++ d[j].p = j ? &pc : &pb; ++ d[j].args = {mtl::Arg::B(W), mtl::Arg::B(x), mtl::Arg::B(sums), mtl::Arg::B(lg[j]), ++ mtl::Arg::V(ga), mtl::Arg::B(ids[j]), mtl::Arg::B(w[j]), mtl::Arg::B(plan[j]), ++ mtl::Arg::V(a), mtl::Arg::B(counter[j]), mtl::Arg::B(smap)}; ++ d[j].tg[0] = 512; d[j].grid[0] = a.n_expert / 16; ++ cx.run({d[j]}); ++ } ++ for (auto pair : {std::pair{lg[0], lg[1]}, std::pair{ids[0], ids[1]}, ++ std::pair{w[0], w[1]}, std::pair{plan[0], plan[1]}}) ++ if (std::memcmp(pair.first.ptr, pair.second.ptr, pair.first.bytes)) { ++ std::printf("router parity FAIL type %d nt %d split %d trial %d\n", type, nt, split, trial); ++ return 1; ++ } ++ if (*counter[0].as() || *counter[1].as()) return 1; ++ } ++ MoeArgs timing{}; ++ timing.n_tok = nt; timing.topk = 10; timing.n_expert = 512; ++ timing.gx_gu = 5; timing.gx_dn = 20; ++ for (auto& dispatch : d) dispatch.args[8] = mtl::Arg::V(timing); ++ d[1].args = d[0].args; ++ cx.warm({d[0], d[1]}, 0.05); ++ std::vector bt, ct, ratios; ++ constexpr int reps = 16; ++ for (int i = 0; i < 7; ++i) { ++ double t[2]; ++ for (int j = 0; j < 2; ++j) { ++ int k = (i + j) & 1; ++ t[k] = cx.run({d[k]}, reps) / reps; ++ } ++ bt.push_back(t[0] * 1e6); ct.push_back(t[1] * 1e6); ratios.push_back(t[1] / t[0]); ++ } ++ std::printf("router type %d nt %d split %d parity PASS base %.3f us cand %.3f us ratio %.5f\n", ++ type, nt, split, median(bt), median(ct), median(ratios)); ++ return 0; ++} ++int mtp_ab(int argc, char** argv) { ++ if (argc != 3) return 2; ++ struct RmsXpArgs { uint32_t n, xp_n, plane_tok, split, keep; float eps; }; ++ struct XpArgs { uint32_t n, n_tok, plane_tok, stride; }; ++ struct CombNormArgs { uint32_t n, hc; float eps; }; ++ constexpr uint32_t N = 2560, HC = 4, RN = HC * N, MAXT = 8; ++ constexpr float eps = 1e-6f; ++ mtl::Context cx; ++ const int lb = cx.compile({argv[2]}, "#define NO_GGML_COMMON 1\n"); ++ const int lc = cx.compile({cx.kernel_dir + "/../src/engine/msl/mtp.metal"}, "#define NO_GGML_COMMON 1\n"); ++ mtl::Pipeline prms[2] = {cx.pipeline(lb, "mtp_rms_xprep"), cx.pipeline(lc, "mtp_rms_xprep")}; ++ mtl::Pipeline pxp[2] = {cx.pipeline(lb, "mtp_xprep"), cx.pipeline(lc, "mtp_xprep")}; ++ mtl::Pipeline pcomb[2] = {cx.pipeline(lb, "mtp_combine_gnorm"), cx.pipeline(lc, "mtp_combine_gnorm")}; ++ const size_t plane = 4ull * MAXT * N; ++ auto x = cx.alloc(4 * MAXT * RN), w = cx.alloc(4 * RN), e = cx.alloc(4 * MAXT * N), h = cx.alloc(4 * MAXT * RN); ++ mtl::Buffer y[2], xp[2], R[2], xn[2]; ++ for (int k = 0; k < 2; ++k) { ++ y[k] = cx.alloc(4 * MAXT * RN); ++ xp[k] = cx.alloc(4 * plane + 4ull * MAXT * N / 16); ++ R[k] = cx.alloc(4 * MAXT * RN); ++ xn[k] = cx.alloc(4 * MAXT * RN); ++ } ++ ref::Rng rng(101); ++ for (size_t i = 0; i < MAXT * RN; ++i) { x.as()[i] = rng.sym() * 2.0f; h.as()[i] = rng.sym(); } ++ for (size_t i = 0; i < RN; ++i) w.as()[i] = 1.0f + 0.2f * rng.sym(); ++ for (size_t i = 0; i < MAXT * N; ++i) e.as()[i] = rng.sym(); ++ auto bench_pair = [&](const char* label, mtl::Dispatch db, mtl::Dispatch dc, int reps = 128) { ++ dc.args = db.args; ++ cx.warm({db, dc}, 0.01); ++ std::vector bt, ct, rt; ++ for (int i = 0; i < 11; ++i) { ++ double b, c; ++ if (i & 1) { c = cx.run({dc}, reps) / reps; b = cx.run({db}, reps) / reps; } ++ else { b = cx.run({db}, reps) / reps; c = cx.run({dc}, reps) / reps; } ++ bt.push_back(b * 1e6); ct.push_back(c * 1e6); rt.push_back(c / b); ++ } ++ std::printf("%-32s parity PASS base %6.3f us cand %6.3f us ratio %.4f (%+.2f%%)\n", ++ label, median(bt), median(ct), median(rt), 100.0 * (median(rt) - 1.0)); ++ }; ++ for (int nt : {1, 4, 8}) { ++ // 1. mtp_rms_xprep embedding (n=2560, split=0, keep=0) ++ { ++ RmsXpArgs a{N, N, MAXT, 0, 0, eps}; ++ mtl::Dispatch d[2]; ++ for (int k = 0; k < 2; ++k) { ++ std::memset(xp[k].ptr, 0xa5, xp[k].bytes); ++ d[k] = {&prms[k], {mtl::Arg::B(x), mtl::Arg::B(w), mtl::Arg::B(y[k]), mtl::Arg::B(xp[k]), mtl::Arg::V(a)}}; ++ d[k].grid[0] = nt; d[k].tg[0] = 256; ++ cx.run({d[k]}); ++ } ++ if (std::memcmp(xp[0].ptr, xp[1].ptr, xp[0].bytes)) return 1; ++ char lbl[64]; std::snprintf(lbl, sizeof(lbl), "mtp_rms_xprep n=2560 nt=%d", nt); ++ bench_pair(lbl, d[0], d[1]); ++ } ++ // 2. mtp_rms_xprep hidden (n=10240, split=1, keep=(nt*4 > 8)) ++ { ++ RmsXpArgs a{RN, N, MAXT, 1, uint32_t(nt * HC > MAXT), eps}; ++ mtl::Dispatch d[2]; ++ for (int k = 0; k < 2; ++k) { ++ std::memset(y[k].ptr, 0xa5, y[k].bytes); ++ std::memset(xp[k].ptr, 0xa5, xp[k].bytes); ++ d[k] = {&prms[k], {mtl::Arg::B(x), mtl::Arg::B(w), mtl::Arg::B(y[k]), mtl::Arg::B(xp[k]), mtl::Arg::V(a)}}; ++ d[k].grid[0] = nt; d[k].tg[0] = 256; ++ cx.run({d[k]}); ++ } ++ if (std::memcmp(y[0].ptr, y[1].ptr, y[0].bytes) || std::memcmp(xp[0].ptr, xp[1].ptr, xp[0].bytes)) return 1; ++ char lbl[64]; std::snprintf(lbl, sizeof(lbl), "mtp_rms_xprep n=10240 nt=%d", nt); ++ bench_pair(lbl, d[0], d[1]); ++ } ++ // 3. mtp_combine_gnorm (hc=4, nt) ++ { ++ CombNormArgs a{N, HC, eps}; ++ mtl::Dispatch d[2]; ++ for (int k = 0; k < 2; ++k) { ++ std::memset(R[k].ptr, 0xa5, R[k].bytes); ++ std::memset(xn[k].ptr, 0xa5, xn[k].bytes); ++ d[k] = {&pcomb[k], {mtl::Arg::B(h), mtl::Arg::B(e), mtl::Arg::B(R[k]), mtl::Arg::B(w), mtl::Arg::B(xn[k]), mtl::Arg::V(a)}}; ++ d[k].grid[0] = HC; d[k].grid[1] = nt; d[k].tg[0] = 256; ++ cx.run({d[k]}); ++ } ++ if (std::memcmp(R[0].ptr, R[1].ptr, R[0].bytes) || std::memcmp(xn[0].ptr, xn[1].ptr, xn[0].bytes)) return 1; ++ char lbl[64]; std::snprintf(lbl, sizeof(lbl), "mtp_combine_gnorm nt=%d", nt); ++ bench_pair(lbl, d[0], d[1]); ++ } ++ } ++ // 4. mtp_xprep (rows=8, n=2560) ++ { ++ XpArgs a{N, MAXT, MAXT, N}; ++ mtl::Dispatch d[2]; ++ for (int k = 0; k < 2; ++k) { ++ std::memset(xp[k].ptr, 0xa5, xp[k].bytes); ++ d[k] = {&pxp[k], {mtl::Arg::B(x), mtl::Arg::B(xp[k]), mtl::Arg::V(a)}}; ++ d[k].grid[0] = (N / 16 * MAXT + 255) / 256; d[k].tg[0] = 256; ++ cx.run({d[k]}); ++ } ++ if (std::memcmp(xp[0].ptr, xp[1].ptr, xp[0].bytes)) return 1; ++ bench_pair("mtp_xprep rows=8", d[0], d[1]); ++ } ++ return 0; ++} + } // namespace + +-int main() { ++int main(int argc, char** argv) { ++ if (argc > 1 && std::string(argv[1]) == "--router-ab") return router_ab(argc, argv); ++ if (argc > 1 && std::string(argv[1]) == "--mtp-ab") return mtp_ab(argc, argv); ++ const bool ab = argc > 1 && std::string(argv[1]) == "--ab"; ++ if (ab && argc < 3) { ++ std::fprintf(stderr, "usage: route_parity --ab ref.metal [--bench]\n"); ++ return 2; ++ } ++ const bool bench = (argc > 1 && std::string(argv[1]) == "--bench") || ++ (ab && argc > 3 && std::string(argv[3]) == "--bench"); + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal"}); ++ const int ref_lib = ab ? cx.compile({"quant.metal", "units.metal", argv[2]}) : -1; + mtl::Pipeline p = cx.pipeline(lib, "moe_route"); +- mtl::Buffer lg = cx.alloc(4 * 8 * 512), ids = cx.alloc(4 * 80), w = cx.alloc(4 * 80), plan = cx.alloc(4 * P_INTS); ++ mtl::Buffer lg = cx.alloc(4 * 8 * 512), ids = cx.alloc(4 * 80), w = cx.alloc(4 * 80), plan = cx.alloc(4 * P_INTS), ++ smap = cx.alloc(4 * 512); + ref::Rng r(9); +- int fails = 0; +- double worst_w = 0; ++ int fails_512 = 0, fails_256_moce = 0; ++ double worst_w_512 = 0, worst_w_256_moce = 0; ++ ++ if (ab) { ++ mtl::Buffer ids_r = cx.alloc(4 * 80), w_r = cx.alloc(4 * 80), plan_r = cx.alloc(4 * P_INTS); ++ int ab_fails = 0; ++ const float moce_scale = std::exp(0.85f); ++ uint32_t moce_pad = 0; ++ std::memcpy(&moce_pad, &moce_scale, 4); ++ for (int split : {0, 1}) { ++ mtl::Constants fc; ++ fc.b(26, split != 0); ++ mtl::Pipeline pr_c = cx.pipeline(lib, "moe_route", fc); ++ mtl::Pipeline pr_r = cx.pipeline(ref_lib, "moe_route", fc); ++ mtl::Pipeline pp_c = cx.pipeline(lib, "moe_plan", fc); ++ mtl::Pipeline pp_r = cx.pipeline(ref_lib, "moe_plan", fc); ++ for (int trial = 0; trial < 100; ++trial) { ++ const int nt = 1 + trial % 8; ++ const int ne_exp = (trial & 1) ? 256 : 512; ++ for (int e = 0; e < 512; ++e) smap.as()[e] = (r.sym() > 0.0f) ? (uint32_t)e : 0xFFFFFFFFu; ++ for (int i = 0; i < nt * 512; ++i) lg.as()[i] = r.sym() * 4.0f; ++ MoeArgs a{}; ++ a.n_tok = nt; ++ a.topk = 10; ++ a.n_expert = ne_exp; ++ a.gx_gu = 5; ++ a.gx_dn = 20; ++ a.pad = (trial & 1) ? moce_pad : 0; ++ ++ std::memset(ids_r.ptr, 0xa5, ids_r.bytes); ++ std::memset(w_r.ptr, 0xa5, w_r.bytes); ++ std::memset(plan_r.ptr, 0xa5, plan_r.bytes); ++ mtl::Dispatch dr{&pr_r, {mtl::Arg::B(lg), mtl::Arg::B(ids_r), mtl::Arg::B(w_r), mtl::Arg::B(plan_r), mtl::Arg::V(a), mtl::Arg::B(smap)}}; ++ dr.tg[0] = 512; ++ cx.run({dr}); ++ ++ std::memset(ids.ptr, 0xa5, ids.bytes); ++ std::memset(w.ptr, 0xa5, w.bytes); ++ std::memset(plan.ptr, 0xa5, plan.bytes); ++ mtl::Dispatch dc{&pr_c, {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(smap)}}; ++ dc.tg[0] = 512; ++ cx.run({dc}); ++ ++ if (std::memcmp(ids.ptr, ids_r.ptr, 4 * nt * 10) != 0 || ++ std::memcmp(w.ptr, w_r.ptr, 4 * nt * 10) != 0 || ++ std::memcmp(plan.ptr, plan_r.ptr, plan.bytes) != 0) { ++ ++ab_fails; ++ } ++ ++ // Also verify standalone moe_plan bitwise parity ++ std::memset(plan_r.ptr, 0x5a, plan_r.bytes); ++ std::memset(plan.ptr, 0x5a, plan.bytes); ++ mtl::Dispatch dpr{&pp_r, {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan_r), mtl::Arg::V(a)}}; ++ mtl::Dispatch dpc{&pp_c, {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a)}}; ++ dpr.tg[0] = dpc.tg[0] = 512; ++ cx.run({dpr, dpc}); ++ if (std::memcmp(plan.ptr, plan_r.ptr, plan.bytes) != 0) ++ab_fails; ++ } ++ } ++ std::printf("route_parity --ab bitwise (moe_route + moe_plan, split 0/1, 200 trials): %d fail(s) %s\n", ++ ab_fails, ab_fails == 0 ? "ALL PASS" : "FAIL"); ++ if (ab_fails || !bench) return ab_fails ? 1 : 0; ++ ++ for (int split : {0, 1}) { ++ mtl::Constants fc; ++ fc.b(26, split != 0); ++ mtl::Pipeline pr_c = cx.pipeline(lib, "moe_route", fc); ++ mtl::Pipeline pr_r = cx.pipeline(ref_lib, "moe_route", fc); ++ mtl::Pipeline pp_c = cx.pipeline(lib, "moe_plan", fc); ++ mtl::Pipeline pp_r = cx.pipeline(ref_lib, "moe_plan", fc); ++ for (int nt : {1, 2, 4, 8}) { ++ for (int e = 0; e < 512; ++e) smap.as()[e] = (r.sym() > 0.0f) ? (uint32_t)e : 0xFFFFFFFFu; ++ for (int i = 0; i < nt * 512; ++i) lg.as()[i] = r.sym() * 4.0f; ++ MoeArgs a{}; ++ a.n_tok = nt; ++ a.topk = 10; ++ a.n_expert = 512; ++ a.gx_gu = 5; ++ a.gx_dn = 20; ++ a.pad = split ? moce_pad : 0; ++ mtl::Dispatch dr{&pr_r, {mtl::Arg::B(lg), mtl::Arg::B(ids_r), mtl::Arg::B(w_r), mtl::Arg::B(plan_r), mtl::Arg::V(a), mtl::Arg::B(smap)}}; ++ mtl::Dispatch dc{&pr_c, {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(smap)}}; ++ dr.tg[0] = dc.tg[0] = 512; ++ cx.run({dc}); ++ cx.warm({dr}, 0.03); ++ cx.warm({dc}, 0.03); ++ std::vector bt, ct, rt; ++ const int reps = 192; // 4x48 layers ++ for (int i = 0; i < 9; ++i) { ++ double b, c; ++ if (i & 1) { c = cx.run({dc}, reps) / reps; b = cx.run({dr}, reps) / reps; } ++ else { b = cx.run({dr}, reps) / reps; c = cx.run({dc}, reps) / reps; } ++ bt.push_back(b * 1e6); ct.push_back(c * 1e6); rt.push_back(c / b); ++ } ++ std::printf("ab bench moe_route split %d ntok %d: base %6.2f us cand %6.2f us ratio %.4f (x48: %.3f -> %.3f ms)\n", ++ split, nt, median(bt), median(ct), median(rt), 48 * median(bt) * 1e-3, 48 * median(ct) * 1e-3); ++ ++ if (split == 0) { ++ mtl::Dispatch dpr{&pp_r, {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan_r), mtl::Arg::V(a)}}; ++ mtl::Dispatch dpc{&pp_c, {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a)}}; ++ dpr.tg[0] = dpc.tg[0] = 512; ++ cx.warm({dpr}, 0.03); ++ cx.warm({dpc}, 0.03); ++ bt.clear(); ct.clear(); rt.clear(); ++ for (int i = 0; i < 9; ++i) { ++ double b, c; ++ if (i & 1) { c = cx.run({dpc}, reps) / reps; b = cx.run({dpr}, reps) / reps; } ++ else { b = cx.run({dpr}, reps) / reps; c = cx.run({dpc}, reps) / reps; } ++ bt.push_back(b * 1e6); ct.push_back(c * 1e6); rt.push_back(c / b); ++ } ++ std::printf("ab bench moe_plan split %d ntok %d: base %6.2f us cand %6.2f us ratio %.4f (x48: %.3f -> %.3f ms)\n", ++ split, nt, median(bt), median(ct), median(rt), 48 * median(bt) * 1e-3, 48 * median(ct) * 1e-3); ++ } ++ } ++ } ++ return 0; ++ } ++ ++ // Phase 1: K=512, MoCE off (pad = 0) + for (int trial = 0; trial < 200; ++trial) { + const int nt = 1 + trial % 8; + for (int i = 0; i < nt * 512; ++i) { +@@ -41,9 +360,10 @@ int main() { + a.n_expert = 512; + a.gx_gu = 1; + a.gx_dn = 1; ++ a.pad = 0; + mtl::Dispatch d; + d.p = &p; +- d.args = {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a)}; ++ d.args = {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(smap)}; + d.tg[0] = 512; + cx.run({d}); + for (int t = 0; t < nt; ++t) { +@@ -54,19 +374,17 @@ int main() { + double s = 0; + for (int e = 0; e < 512; ++e) s += (pr[e] = std::exp((double)l[e] - mx)); + for (int e = 0; e < 512; ++e) pr[e] /= s; +- // top 10: largest probability, ties -> lower index (probabilities of equal logits are equal) + std::vector order(512); + for (int e = 0; e < 512; ++e) order[e] = e; + std::stable_sort(order.begin(), order.end(), [&](int x, int y) { return l[x] > l[y]; }); + double ssum = 0; + for (int k = 0; k < 10; ++k) ssum += pr[order[k]]; + for (int k = 0; k < 10; ++k) { +- if (ids.as()[t * 10 + k] != order[k]) { ++fails; break; } ++ if (ids.as()[t * 10 + k] != order[k]) { ++fails_512; break; } + const double ref = pr[order[k]] / std::max(ssum, 6.103515625e-5); +- worst_w = std::max(worst_w, std::fabs(w.as()[t * 10 + k] - ref) / ref); ++ worst_w_512 = std::max(worst_w_512, std::fabs(w.as()[t * 10 + k] - ref) / ref); + } + } +- // the plan + const int* P = plan.as(); + std::set distinct(ids.as(), ids.as() + nt * 10); + bool ok = P[P_NG] == (int)distinct.size(); +@@ -82,9 +400,100 @@ int main() { + if (j != P[P_GSTART + gi + 1]) ok = false; + ++gi; + } +- fails += !ok; ++ fails_512 += !ok; ++ } ++ ++ // Phase 2: K=256, MoCE on (pad = bit_cast(exp(0.85f))), random smap residency, and huge trap logits in 256..511 ++ const float moce_scale = std::exp(0.85f); ++ uint32_t moce_pad = 0; ++ std::memcpy(&moce_pad, &moce_scale, 4); ++ for (int trial = 0; trial < 200; ++trial) { ++ const int nt = 1 + trial % 8; ++ for (int e = 0; e < 512; ++e) { ++ smap.as()[e] = (r.sym() > 0.0f) ? (uint32_t) e : 0xFFFFFFFFu; ++ } ++ for (int t = 0; t < nt; ++t) { ++ for (int e = 0; e < 256; ++e) { ++ float v = r.sym() * 4.0f; ++ if (trial % 3 == 0) v = std::round(v * 4.0f) / 4.0f; ++ lg.as()[t * 512 + e] = v; ++ } ++ for (int e = 256; e < 512; ++e) { ++ lg.as()[t * 512 + e] = 50.0f + (float) e; // trap logits past n_expert = 256 ++ } ++ } ++ MoeArgs a{}; ++ a.n_tok = nt; ++ a.topk = 10; ++ a.n_expert = 256; ++ a.gx_gu = 1; ++ a.gx_dn = 1; ++ a.pad = moce_pad; ++ mtl::Dispatch d; ++ d.p = &p; ++ d.args = {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(smap)}; ++ d.tg[0] = 512; ++ cx.run({d}); ++ ++ const uint32_t* sm = smap.as(); ++ for (int t = 0; t < nt; ++t) { ++ const float* l = lg.as() + t * 512; ++ float mx = -INFINITY; ++ for (int e = 0; e < 256; ++e) mx = std::max(mx, l[e]); ++ std::vector pr(256); ++ float s = 0.0f; ++ for (int e = 0; e < 256; ++e) s += (pr[e] = std::exp(l[e] - mx)); ++ const float rcp = 1.0f / s; ++ for (int e = 0; e < 256; ++e) pr[e] *= rcp; ++ ++ std::vector work = pr; ++ std::vector picked(10, -1); ++ double ssum = 0.0; ++ for (int rank = 0; rank < 10; ++rank) { ++ if (rank == 3) { ++ for (int e = 0; e < 256; ++e) { ++ if (work[e] > 0.0f && sm[e] != 0xFFFFFFFFu) work[e] *= moce_scale; ++ } ++ } ++ int best_e = 0; ++ float best_v = work[0]; ++ for (int e = 1; e < 256; ++e) { ++ if (work[e] > best_v) { best_v = work[e]; best_e = e; } ++ } ++ picked[rank] = best_e; ++ work[best_e] = -INFINITY; ++ ssum += pr[best_e]; ++ } ++ for (int k = 0; k < 10; ++k) { ++ const int got_id = ids.as()[t * 10 + k]; ++ if (got_id != picked[k] || got_id >= 256) { ++fails_256_moce; break; } ++ const double ref = (double) pr[picked[k]] / std::max(ssum, 6.103515625e-5); ++ worst_w_256_moce = std::max(worst_w_256_moce, std::fabs(w.as()[t * 10 + k] - ref) / ref); ++ } ++ } ++ const int* P = plan.as(); ++ std::set distinct(ids.as(), ids.as() + nt * 10); ++ bool ok = P[P_NG] == (int)distinct.size(); ++ int gi = 0; ++ for (int e : distinct) { ++ if (e >= 256 || P[P_GEXP + gi] != e) ok = false; ++ int j = P[P_GSTART + gi]; ++ for (int i = 0; i < nt * 10; ++i) ++ if (ids.as()[i] == e) { ++ if (P[P_ETOK + j] != i / 10 || P[P_TOKENT + i] != j) ok = false; ++ ++j; ++ } ++ if (j != P[P_GSTART + gi + 1]) ok = false; ++ ++gi; ++ } ++ fails_256_moce += !ok; + } +- std::printf("route parity, 200 windows of 1..8 tokens (a third with exact ties): ids/plan mismatches %d, worst weight " +- "rel err %.2e %s\n", fails, worst_w, fails == 0 && worst_w < 2e-6 ? "PASS" : "FAIL"); +- return fails ? 1 : 0; ++ ++ const int total_fails = fails_512 + fails_256_moce; ++ const double worst_w = std::max(worst_w_512, worst_w_256_moce); ++ std::printf("route parity K=512 (pad=0, 200 windows): mismatches %d, worst weight rel err %.2e\n", ++ fails_512, worst_w_512); ++ std::printf("route parity K=256 + MoCE=0.85 (random smap + 256..511 traps, 200 windows): mismatches %d, worst weight rel err %.2e %s\n", ++ 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/sample_parity.cpp b/metal/tests/sample_parity.cpp +index 049439af..2cb5b4a4 100644 +--- a/metal/tests/sample_parity.cpp ++++ b/metal/tests/sample_parity.cpp +@@ -1,61 +1,260 @@ +-// metal/tests/sample_parity.cpp - greedy argmax (sample.metal) vs the CPU over 248,320 logits with planted ties, +-// and its time. ++// metal/tests/sample_parity.cpp - exact sampling parity and short paired GPU timings. ++// sample_parity [--mtp] [baseline-source.metal]; without a baseline, compare the kernels to themselves. + #include "mtl_rt.h" + #include "ref_quant.h" ++#include "../src/engine/slop4.h" + + #include + #include + #include ++#include ++#include + #include + +-int main() { +- constexpr uint32_t V = 248320, PARTS = 120; +- mtl::Context cx; +- const int lib = cx.compile({"sample.metal"}, "#define NO_GGML_COMMON 1\n"); +- mtl::Pipeline pp = cx.pipeline(lib, "argmax_part"), pf = cx.pipeline(lib, "argmax_final"); +- const int MAXT = 8; +- mtl::Buffer lg = cx.alloc(4 * (size_t)V * MAXT), pv = cx.alloc(4 * PARTS * MAXT), pi = cx.alloc(4 * PARTS * MAXT), +- out = cx.alloc(4 * MAXT); ++struct ArgArgs { uint32_t n_vocab, parts; }; ++struct PickArgs { uint32_t nv, j, use_map; }; ++struct DraftArgs { slop4::State state; slop4::Knobs knobs; uint32_t nv, j, vocab; }; ++struct EmbArgs { uint32_t n, row_bytes, j; }; ++ ++static void timing(mtl::Context& cx, const char* name, const std::vector& base, ++ const std::vector& cand, int reps = 64) { ++ cx.warm(base, 0.005); ++ cx.warm(cand, 0.005); ++ std::vector b, c, r; ++ for (int k = 0; k < 9; ++k) { ++ double bt, ct; ++ if (k & 1) { ct = cx.run(cand, reps); bt = cx.run(base, reps); } ++ else { bt = cx.run(base, reps); ct = cx.run(cand, reps); } ++ b.push_back(bt / reps * 1e6); ++ c.push_back(ct / reps * 1e6); ++ r.push_back(ct / bt); ++ } ++ std::sort(b.begin(), b.end()); ++ std::sort(c.begin(), c.end()); ++ std::sort(r.begin(), r.end()); ++ std::printf("%s baseline %.3f us candidate %.3f us ratio %.4f (%+.2f%%)\n", ++ name, b[4], c[4], r[4], 100.0 * (r[4] - 1.0)); ++} ++ ++static int argmax(mtl::Context& cx, const char* baseline) { ++ const int lb = cx.compile({baseline}, "#define NO_GGML_COMMON 1\n"); ++ const int lc = cx.compile({"sample.metal"}, "#define NO_GGML_COMMON 1\n"); ++ mtl::Pipeline pp[2] = {cx.pipeline(lb, "argmax_part"), cx.pipeline(lc, "argmax_part")}; ++ mtl::Pipeline pf[2] = {cx.pipeline(lb, "argmax_final"), cx.pipeline(lc, "argmax_final")}; ++ constexpr uint32_t V = 248320, PARTS = 120, MAXT = 8; ++ auto lg = cx.alloc(4 * size_t(V) * MAXT); ++ mtl::Buffer pv[2], pi[2], out[2]; ++ for (int k = 0; k < 2; ++k) { ++ pv[k] = cx.alloc(4 * PARTS * MAXT); ++ pi[k] = cx.alloc(4 * PARTS * MAXT); ++ out[k] = cx.alloc(4 * MAXT); ++ } + ref::Rng r(5); + int fails = 0; +- struct A { uint32_t n_vocab, parts; } a{V, PARTS}; +- for (int trial = 0; trial < 20; ++trial) { +- const int nt = 1 + trial % MAXT; +- for (size_t i = 0; i < (size_t)V * nt; ++i) lg.as()[i] = std::round(r.sym() * 1000.0f) / 100.0f; +- for (int t = 0; t < nt; ++t) { +- // plant a tie at the maximum +- const size_t i1 = r.next() % V, i2 = r.next() % V; +- lg.as()[t * V + i1] = 20.0f; +- lg.as()[t * V + i2] = 20.0f; ++ struct Case { uint32_t nv, nt; }; ++ for (auto shape : {Case{V, 1}, Case{V, 2}, Case{V, 3}, Case{V, 4}, Case{V, 5}, ++ Case{V, 6}, Case{V, 7}, Case{V, 8}, ++ Case{248064, 1}, Case{248064, 2}, Case{248064, 3}, Case{248064, 4}, ++ Case{248064, 5}, Case{248064, 6}, Case{248064, 7}, Case{248064, 8}, ++ Case{40525, 1}, Case{40525, 2}, Case{40525, 4}, Case{40525, 8}, Case{7, 1}}) { ++ const ArgArgs a{shape.nv, PARTS}; ++ const uint32_t per = (a.n_vocab + PARTS - 1) / PARTS; ++ std::vector ds[2]; ++ for (int k = 0; k < 2; ++k) { ++ mtl::Dispatch d1, d2; ++ d1.p = &pp[k]; ++ d1.args = {mtl::Arg::B(lg), mtl::Arg::B(pv[k]), mtl::Arg::B(pi[k]), mtl::Arg::V(a)}; ++ d1.grid[0] = PARTS; ++ d1.grid[1] = shape.nt; ++ d1.tg[0] = 256; ++ d2.p = &pf[k]; ++ d2.args = {mtl::Arg::B(pv[k]), mtl::Arg::B(pi[k]), mtl::Arg::B(out[k]), mtl::Arg::V(a)}; ++ d2.grid[0] = shape.nt; ++ d2.tg[0] = 256; ++ ds[k] = {d1, d2}; + } +- mtl::Dispatch d1, d2; +- d1.p = &pp; +- d1.args = {mtl::Arg::B(lg), mtl::Arg::B(pv), mtl::Arg::B(pi), mtl::Arg::V(a)}; +- d1.grid[0] = PARTS; +- d1.grid[1] = nt; +- d1.tg[0] = 256; +- d2.p = &pf; +- d2.args = {mtl::Arg::B(pv), mtl::Arg::B(pi), mtl::Arg::B(out), mtl::Arg::V(a)}; +- d2.grid[0] = nt; +- d2.tg[0] = 256; +- cx.run({d1, d2}); +- for (int t = 0; t < nt; ++t) { +- const float* l = lg.as() + (size_t)t * V; +- int best = 0; +- for (uint32_t i = 1; i < V; ++i) +- if (l[i] > l[best]) best = (int)i; +- fails += out.as()[t] != best; ++ for (int trial = 0; trial < 7; ++trial) { ++ for (size_t i = 0; i < size_t(a.n_vocab) * shape.nt; ++i) { ++ float v = std::round(r.sym() * 1000.0f) / 100.0f; ++ if (trial == 1) v = -INFINITY; ++ if (trial == 2) v = std::numeric_limits::quiet_NaN(); ++ if (trial == 3) v = (i & 1) ? -0.0f : 0.0f; ++ if (trial == 4 && i % 17 == 0) v = std::numeric_limits::quiet_NaN(); ++ lg.as()[i] = v; ++ } ++ if (trial == 0 || trial >= 5) { ++ for (uint32_t t = 0; t < shape.nt; ++t) { ++ float* l = lg.as() + size_t(t) * a.n_vocab; ++ l[1] = l[a.n_vocab - 1] = trial == 5 ? INFINITY : 20.0f; ++ } ++ } ++ for (int k = 0; k < 2; ++k) cx.run(ds[k]); ++ fails += std::memcmp(pv[0].ptr, pv[1].ptr, 4 * PARTS * shape.nt) != 0; ++ fails += std::memcmp(pi[0].ptr, pi[1].ptr, 4 * PARTS * shape.nt) != 0; ++ fails += std::memcmp(out[0].ptr, out[1].ptr, 4 * shape.nt) != 0; ++ for (uint32_t t = 0; t < shape.nt; ++t) { ++ float best = -INFINITY; ++ int id = 0x7fffffff; ++ for (uint32_t p = 0; p < PARTS; ++p) { ++ float v = -INFINITY; ++ int ix = 0x7fffffff; ++ for (uint32_t i = p * per; i < std::min((p + 1) * per, a.n_vocab); ++i) { ++ const float x = lg.as()[size_t(t) * a.n_vocab + i]; ++ if (x > v || (x == v && int(i) < ix)) { v = x; ix = int(i); } ++ } ++ const size_t off = t * PARTS + p; ++ fails += std::memcmp(&v, pv[1].as() + off, 4) != 0; ++ fails += ix != pi[1].as()[off]; ++ if (v > best || (v == best && ix < id)) { best = v; id = ix; } ++ } ++ fails += out[1].as()[t] != id; ++ } + } +- if (trial == 19) { +- std::vector ds; +- for (int k = 0; k < 64; ++k) { ds.push_back(d1); ds.push_back(d2); } +- cx.warm(ds, 0.05); +- std::vector ts; +- for (int k = 0; k < 7; ++k) ts.push_back(cx.run(ds) / 64); +- std::sort(ts.begin(), ts.end()); +- std::printf("argmax %d token(s) over %u logits: %.2f us (2 dispatches)\n", nt, V, ts[3] * 1e6); ++ for (size_t i = 0; i < size_t(a.n_vocab) * shape.nt; ++i) ++ lg.as()[i] = std::round(r.sym() * 1000.0f) / 100.0f; ++ char label[80]; ++ std::snprintf(label, sizeof(label), "argmax nv=%u nt=%u", a.n_vocab, shape.nt); ++ timing(cx, label, ds[0], ds[1]); ++ } ++ std::printf("argmax pv/pi/out bitwise + cpu: %d wrong %s\n", fails, fails ? "FAIL" : "PASS"); ++ return fails; ++} ++ ++static int mtp(mtl::Context& cx, const char* baseline) { ++ const std::string engine = cx.kernel_dir + "/../src/engine/"; ++ const int lb = cx.compile({"quant.metal", "units.metal", engine + "slop4.h", baseline}); ++ const int lc = cx.compile({"quant.metal", "units.metal", engine + "slop4.h", engine + "msl/mtp_chain.metal"}); ++ mtl::Pipeline pick[2] = {cx.pipeline(lb, "mtp_pick"), cx.pipeline(lc, "mtp_pick")}; ++ mtl::Pipeline pen[2] = {cx.pipeline(lb, "mtp_slop_pen"), cx.pipeline(lc, "mtp_slop_pen")}; ++ constexpr uint32_t V = 248320; ++ auto lg = cx.alloc(4 * V), ix = cx.alloc(4), map = cx.alloc(4 * V); ++ mtl::Buffer dout[2], prob[2], state[2], penalized[2]; ++ for (int k = 0; k < 2; ++k) { ++ dout[k] = cx.alloc(16); ++ prob[k] = cx.alloc(16); ++ state[k] = cx.alloc(sizeof(slop4::State)); ++ penalized[k] = cx.alloc(4 * V); ++ } ++ auto pen_dispatch = [&](int k, const DraftArgs& a) { ++ mtl::Dispatch d; ++ d.p = &pen[k]; ++ d.args = {mtl::Arg::B(penalized[k]), mtl::Arg::B(map), mtl::Arg::B(dout[k]), ++ mtl::Arg::B(state[k]), mtl::Arg::V(a)}; ++ d.tg[0] = 256; ++ return d; ++ }; ++ ref::Rng r(19); ++ int fails = 0; ++ for (uint32_t nv : {40525u, V}) { ++ for (uint32_t i = 0; i < V; ++i) map.as()[i] = V - i - 1; ++ PickArgs a{nv, 2, 1}; ++ std::vector ds[2]; ++ for (int k = 0; k < 2; ++k) { ++ mtl::Dispatch d; ++ d.p = &pick[k]; ++ d.args = {mtl::Arg::B(lg), mtl::Arg::B(ix), mtl::Arg::B(map), mtl::Arg::B(dout[k]), ++ mtl::Arg::B(prob[k]), mtl::Arg::V(a)}; ++ d.tg[0] = 1024; ++ ds[k] = {d}; ++ } ++ for (int trial = 0; trial < 8; ++trial) { ++ ix.as()[0] = trial; ++ for (uint32_t i = 0; i < nv; ++i) lg.as()[i] = r.sym() * 10.0f; ++ if (trial == 6) lg.as()[nv - 1] = -INFINITY; ++ if (trial == 7) lg.as()[nv - 1] = std::numeric_limits::quiet_NaN(); ++ a.j = trial % 4; ++ a.use_map = trial & 1; ++ for (int k = 0; k < 2; ++k) { ++ ds[k][0].args[5] = mtl::Arg::V(a); ++ cx.run(ds[k]); ++ } ++ const int expected = a.use_map ? map.as()[trial] : trial; ++ fails += dout[1].as()[a.j] != expected; ++ fails += std::memcmp(dout[0].ptr, dout[1].ptr, 16) != 0; ++ fails += std::memcmp(prob[0].ptr, prob[1].ptr, 16) != 0; ++ } ++ for (uint32_t i = 0; i < nv; ++i) lg.as()[i] = r.sym() * 10.0f; ++ timing(cx, nv == V ? "mtp_pick nv=248320" : "mtp_pick nv=40525", ds[0], ds[1]); ++ for (int trial = 0; trial < 10; ++trial) { ++ DraftArgs a{{}, {1.0f, 1.0f, 1.0f, 1.0f}, nv, 0, V}; ++ a.state.f = slop4::KNOWN | slop4::PROSE | slop4::BOL; ++ a.state.t1 = 1602; ++ if (trial == 0) a.knobs = {}; ++ if (trial >= 4) { a.j = 1; a.state.t1 = slop4::NL; } ++ for (uint32_t i = 0; i < V; ++i) map.as()[i] = i < nv ? int(i) : -1; ++ for (uint32_t i = 0; i < nv; ++i) lg.as()[i] = r.sym() * 10.0f; ++ if (trial == 6) std::fill_n(lg.as(), nv, -INFINITY); ++ if (trial == 7) std::fill_n(lg.as(), nv, std::numeric_limits::quiet_NaN()); ++ if (trial == 8) lg.as()[100] = INFINITY; ++ if (trial == 9) std::fill_n(lg.as(), nv, 0.0f); ++ if (trial == 2) lg.as()[100] = 100.0f; ++ if (trial == 3) lg.as()[100] = std::numeric_limits::quiet_NaN(); ++ for (int k = 0; k < 2; ++k) { ++ *state[k].as() = a.state; ++ dout[k].as()[0] = slop4::NL; ++ std::memcpy(penalized[k].ptr, lg.ptr, 4 * nv); ++ ds[k] = {pen_dispatch(k, a)}; ++ cx.run(ds[k]); ++ } ++ fails += std::memcmp(penalized[0].ptr, penalized[1].ptr, 4 * nv) != 0; ++ fails += std::memcmp(state[0].ptr, state[1].ptr, sizeof(slop4::State)) != 0; ++ if (trial == 1 || trial == 2) ++ fails += std::memcmp(lg.ptr, penalized[1].ptr, 4 * nv) == 0; ++ } ++ // Small penalties keep the hold branch stable across timed repetitions. ++ for (int mode = 0; mode < 3; ++mode) { ++ DraftArgs a{{}, {0.0f, mode ? 1e-6f : 0.0f, 0.0f, 0.0f}, nv, 0, V}; ++ a.state.f = slop4::KNOWN | slop4::PROSE; ++ a.state.t1 = 1602; ++ for (int k = 0; k < 2; ++k) { ++ std::fill_n(penalized[k].as(), nv, 0.0f); ++ if (mode == 1) penalized[k].as()[100] = 100.0f; ++ ds[k] = {pen_dispatch(k, a)}; ++ } ++ char label[80]; ++ std::snprintf(label, sizeof(label), "mtp_slop_pen nv=%u mode=%d", nv, mode); ++ timing(cx, label, ds[0], ds[1]); + } + } +- std::printf("argmax parity, 20 windows with planted ties: %d wrong %s\n", fails, fails ? "FAIL" : "PASS"); +- return fails ? 1 : 0; ++ for (int format : {ref::BF16, ref::Q8_0, ref::Q4_0, ref::Q5_0, ref::IQ4_NL, ref::IQ4_XS, ref::Q2_0, ++ ref::IQ3_S, ref::Q3_K, ref::Q4_K, ref::Q5_K, ref::Q6_K, ref::IQ3_XXS}) { ++ const int n = 2560, rows = 3; ++ const int rb = int(ref::row_bytes(format, n)); ++ auto w = cx.alloc(size_t(rb) * rows), pos = cx.alloc(4 * n); ++ auto y0 = cx.alloc(4 * n), y1 = cx.alloc(4 * n); ++ for (int i = 0; i < n; ++i) pos.as()[i] = float(i); ++ ref::random_blocks(format, w.as(), rows, n, 31 + format); ++ mtl::Pipeline emb[2] = {cx.pipeline(lb, "mtp_embed_" + std::to_string(format)), ++ cx.pipeline(lc, "mtp_embed_" + std::to_string(format))}; ++ std::vector ds[2]; ++ EmbArgs a{uint32_t(n), uint32_t(rb), 1}; ++ for (int trial = 0; trial < rows; ++trial) { ++ for (int k = 0; k < 2; ++k) { ++ dout[k].as()[1] = trial; ++ mtl::Dispatch d; ++ d.p = &emb[k]; ++ d.args = {mtl::Arg::B(w), mtl::Arg::B(dout[k]), mtl::Arg::B(pos), ++ mtl::Arg::B(k ? y1 : y0), mtl::Arg::V(a)}; ++ d.tg[0] = 256; ++ ds[k] = {d}; ++ cx.run(ds[k]); ++ } ++ fails += std::memcmp(y0.ptr, y1.ptr, 4 * n) != 0; ++ } ++ char label[40]; ++ std::snprintf(label, sizeof(label), "mtp_embed format=%d", format); ++ timing(cx, label, ds[0], ds[1], 1024); ++ } ++ std::printf("mtp pick/pen/embed bitwise: %d wrong %s\n", fails, fails ? "FAIL" : "PASS"); ++ return fails; ++} ++ ++int main(int argc, char** argv) { ++ mtl::Context cx; ++ const bool chain = argc > 1 && std::strcmp(argv[1], "--mtp") == 0; ++ const char* baseline = argc > (chain ? 2 : 1) ? argv[chain ? 2 : 1] : nullptr; ++ const std::string self = cx.kernel_dir + "/../src/engine/msl/mtp_chain.metal"; ++ std::printf("device %s; baseline %s\n", cx.device_name().c_str(), baseline ? baseline : "self"); ++ return (chain ? mtp(cx, baseline ? baseline : self.c_str()) ++ : argmax(cx, baseline ? baseline : "sample.metal")) ? 1 : 0; + } +diff --git a/metal/tests/spec_min_p.cpp b/metal/tests/spec_min_p.cpp +new file mode 100644 +index 00000000..ec5ef0c6 +--- /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 a47c7630..a0add9b2 100644 +--- a/metal/tests/token_chain.cpp ++++ b/metal/tests/token_chain.cpp +@@ -24,7 +24,7 @@ namespace { + + constexpr int N = 2560, HC = 4, LR = 320, NFF = 640, TOPK = 10, NEXP = 512, MAXT = 8; + constexpr int P_MAXE = 128, P_GEXP = 16, P_GSTART = P_GEXP + P_MAXE, P_ETOK = P_GSTART + P_MAXE + 4, +- P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_INTS = P_TOKENT + P_MAXE, P_IGU = 4, P_IDN = 8; ++ P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_INTS = P_TOKENT + P_MAXE + 20 + P_MAXE, P_IGU = 4, P_IDN = 8; + 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 GrArgs3 { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; uint32_t rb_down, rb_up, comb, topk; }; +@@ -552,7 +552,7 @@ int main(int argc, char** argv) { + GemvArgs ga{(uint32_t)N, 512, (uint32_t)ntok, (uint32_t)ref::row_bytes(tr.type, N), (uint32_t)N, 512}; + r.args = {mtl::Arg::B(W, tr.offset), mtl::Arg::B(XF.b, 0), mtl::Arg::B(XF.b, XF.plane_bytes() * 4), + mtl::Arg::B(logits), mtl::Arg::V(ga), mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::B(plan), +- mtl::Arg::V(a), mtl::Arg::B(rcounter)}; ++ mtl::Arg::V(a), mtl::Arg::B(rcounter), mtl::Arg::B(rcounter)}; + r.grid[0] = 512 / 16; + r.tg[0] = 512; + nb = true; +@@ -560,7 +560,7 @@ int main(int argc, char** argv) { + } else { + mtl::Dispatch r; + r.p = proute; +- r.args = {mtl::Arg::B(logits), mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::B(plan), mtl::Arg::V(a)}; ++ r.args = {mtl::Arg::B(logits), mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(plan)}; + r.tg[0] = 512; + push(r, "router", 0); + } +diff --git a/serve/server.py b/serve/server.py +index 5ca7f4e4..26a68232 100644 +--- a/serve/server.py ++++ b/serve/server.py +@@ -3851,6 +3851,13 @@ class Service: + prompt; with thinking, Service.run writes it once the thinking is over.""" + # `cancel` (#1999): set when the client hangs up; the downloads, the wait for the turn and the encodes stop + # between pictures and raise ClientGone ++ if self.vision is None and os.environ.get("STRATA_OMIT_IMAGES_WHEN_NO_VISION") == "1": ++ for m in messages: ++ if isinstance(m.get("content"), list): ++ m["content"] = [ ++ {"type": "text", "text": "[image omitted on text-only engine]"} if isinstance(it, dict) and it.get("type") == "image" else it ++ for it in m["content"] ++ ] + fetched = self._note_unreadable_tool_images(messages, cancel) + ids = self.encode_prompt(messages, tools, kwargs) + if force and kwargs.get("enable_thinking", True) is False: diff --git a/scripts/strata_metal_complete.diff b/scripts/strata_metal_complete.diff index 4f9b715..fee4544 100644 --- a/scripts/strata_metal_complete.diff +++ b/scripts/strata_metal_complete.diff @@ -1,6 +1,6 @@ diff --git a/include/strata/core/scr.hpp b/include/strata/core/scr.hpp new file mode 100644 -index 0000000..2a48f75 +index 00000000..2a48f751 --- /dev/null +++ b/include/strata/core/scr.hpp @@ -0,0 +1,227 @@ @@ -232,7 +232,7 @@ index 0000000..2a48f75 + +} // namespace strata::core diff --git a/metal/CMakeLists.txt b/metal/CMakeLists.txt -index d23a26f..92cd20c 100644 +index d23a26fc..92cd20c7 100644 --- a/metal/CMakeLists.txt +++ b/metal/CMakeLists.txt @@ -120,4 +120,12 @@ target_compile_definitions(strata_metal_kernels PRIVATE @@ -250,7 +250,7 @@ index d23a26f..92cd20c 100644 +target_compile_definitions(heal_k256_q2_0 PRIVATE ACCELERATE_NEW_LAPACK) +target_link_libraries(heal_k256_q2_0 PRIVATE mtl_rt "-framework Accelerate") diff --git a/metal/kernels/dense_mm.metal b/metal/kernels/dense_mm.metal -index 45e29ac..ccf06f7 100644 +index 45e29ace..ccf06f79 100644 --- a/metal/kernels/dense_mm.metal +++ b/metal/kernels/dense_mm.metal @@ -8,7 +8,12 @@ using namespace mpp::tensor_ops; @@ -348,7 +348,7 @@ index 45e29ac..ccf06f7 100644 +template [[host_name("gemm_f16_64_32_2")]] kernel decltype(gemm_f16<64, 32, 2>) gemm_f16<64, 32, 2>; +template [[host_name("gemm_f16_16_96_2")]] kernel decltype(gemm_f16<16, 96, 2>) gemm_f16<16, 96, 2>; diff --git a/metal/kernels/gdn.metal b/metal/kernels/gdn.metal -index 516d89a..c057067 100644 +index 516d89a4..c0570674 100644 --- a/metal/kernels/gdn.metal +++ b/metal/kernels/gdn.metal @@ -27,6 +27,10 @@ struct GdnArgs { @@ -568,7 +568,7 @@ index 516d89a..c057067 100644 } diff --git a/metal/kernels/gdn_prompt.metal b/metal/kernels/gdn_prompt.metal -index b2d530b..c2c4020 100644 +index b2d530b1..c2c40200 100644 --- a/metal/kernels/gdn_prompt.metal +++ b/metal/kernels/gdn_prompt.metal @@ -25,67 +25,161 @@ struct GdnPArgs { @@ -1065,7 +1065,7 @@ index b2d530b..c2c4020 100644 // ---- the chunked form (chunks of GC = 32 tokens), the WY / UT form of the gated delta rule. With G_t the product // of the chunk's gates up to t (log space: lg_t) and S0 the state entering the chunk: diff --git a/metal/kernels/gemv.metal b/metal/kernels/gemv.metal -index 6d4ed3a..49b7b62 100644 +index 6d4ed3a5..49b7b621 100644 --- a/metal/kernels/gemv.metal +++ b/metal/kernels/gemv.metal @@ -21,6 +21,85 @@ constant int NTOK [[function_constant(0)]]; // tokens of the window (1..8): th @@ -1504,7 +1504,7 @@ index 6d4ed3a..49b7b62 100644 const float v = lane_sum(acc[r][t]); if (li == 0 && row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = v; diff --git a/metal/kernels/glue.metal b/metal/kernels/glue.metal -index b7bc8bc..b34bb90 100644 +index b7bc8bcc..b34bb90d 100644 --- a/metal/kernels/glue.metal +++ b/metal/kernels/glue.metal @@ -10,28 +10,29 @@ struct GlueArgs { @@ -1571,7 +1571,7 @@ index b7bc8bc..b34bb90 100644 bo[gid] = fma(gs, ysh[gid], s); } diff --git a/metal/kernels/gr.metal b/metal/kernels/gr.metal -index c1960fd..c18d2ac 100644 +index c1960fd5..c18d2ac7 100644 --- a/metal/kernels/gr.metal +++ b/metal/kernels/gr.metal @@ -12,6 +12,11 @@ @@ -1910,7 +1910,7 @@ index c1960fd..c18d2ac 100644 +template [[host_name("gr_down_q5_8_8")]] kernel decltype(gr_down_f<8, 8, 1, T_Q5_0, T_Q8_0>) gr_down_f<8, 8, 1, T_Q5_0, T_Q8_0>; +template [[host_name("gr_up_q5")]] kernel decltype(gr_up_f) gr_up_f; diff --git a/metal/kernels/gr2.metal b/metal/kernels/gr2.metal -index 1253a11..58342f0 100644 +index 1253a11d..58342f09 100644 --- a/metal/kernels/gr2.metal +++ b/metal/kernels/gr2.metal @@ -87,6 +87,7 @@ kernel void gr_down3(device const uchar* w_down [[buffer(0)]], device const ucha @@ -2258,7 +2258,7 @@ index 1253a11..58342f0 100644 + } } diff --git a/metal/kernels/mma.metal b/metal/kernels/mma.metal -index 9bbb27e..8fcc0ab 100644 +index 9bbb27e2..8fcc0abb 100644 --- a/metal/kernels/mma.metal +++ b/metal/kernels/mma.metal @@ -17,11 +17,57 @@ template struct DW { @@ -2410,7 +2410,7 @@ index 9bbb27e..8fcc0ab 100644 w8[r * 128 + p0 + j] = w[j]; w8[r * 128 + p1 + j] = w[H + j]; diff --git a/metal/kernels/moe.metal b/metal/kernels/moe.metal -index bb8fdbb..878252e 100644 +index bb8fdbbc..878252e7 100644 --- a/metal/kernels/moe.metal +++ b/metal/kernels/moe.metal @@ -20,7 +20,13 @@ @@ -3135,7 +3135,7 @@ index bb8fdbb..878252e 100644 } diff --git a/metal/kernels/moe_mm.metal b/metal/kernels/moe_mm.metal -index ff28897..03ac90a 100644 +index ff28897a..03ac90a8 100644 --- a/metal/kernels/moe_mm.metal +++ b/metal/kernels/moe_mm.metal @@ -32,28 +32,26 @@ kernel void mm_hist(device const int* ids [[buffer(0)]], device uint* bcnt [[buf @@ -4036,7 +4036,7 @@ index ff28897..03ac90a 100644 +#endif +#endif diff --git a/metal/kernels/prompt.metal b/metal/kernels/prompt.metal -index 64833b0..a3fd96b 100644 +index 64833b0a..a3fd96be 100644 --- a/metal/kernels/prompt.metal +++ b/metal/kernels/prompt.metal @@ -7,33 +7,98 @@ @@ -4704,7 +4704,7 @@ index 64833b0..a3fd96b 100644 +template [[host_name("gemm_swiglu_up_32_64_2")]] kernel decltype(gemm_swiglu_up<32, 64, 2>) gemm_swiglu_up<32, 64, 2>; +template [[host_name("gemm_swiglu_up_64_128_4")]] kernel decltype(gemm_swiglu_up<64, 128, 4>) gemm_swiglu_up<64, 128, 4>; diff --git a/metal/kernels/qsa.metal b/metal/kernels/qsa.metal -index 9642e58..63259b4 100644 +index 9642e58e..63259b46 100644 --- a/metal/kernels/qsa.metal +++ b/metal/kernels/qsa.metal @@ -35,6 +35,13 @@ constant int QSA_NTOK [[function_constant(3)]]; @@ -5196,7 +5196,7 @@ index 9642e58..63259b4 100644 } + diff --git a/metal/kernels/qsa2.metal b/metal/kernels/qsa2.metal -index 8ae189b..d5b594c 100644 +index 8ae189b6..30439a16 100644 --- a/metal/kernels/qsa2.metal +++ b/metal/kernels/qsa2.metal @@ -29,9 +29,11 @@ kernel void qsa_attn2(device const float* qcur [[buffer(0)]], device const char* @@ -5262,7 +5262,7 @@ index 8ae189b..d5b594c 100644 for (int h = 0; h < HPG; ++h) { device float* o = pp + (h0 + h) * (QD + 2); o[tid] = acc[h]; -@@ -82,40 +95,37 @@ kernel void qsa_attn2(device const float* qcur [[buffer(0)]], device const char* +@@ -82,40 +95,45 @@ kernel void qsa_attn2(device const float* qcur [[buffer(0)]], device const char* kernel void qsa_merge2(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { @@ -5296,11 +5296,19 @@ index 8ae189b..d5b594c 100644 + float L = 0.0f; + for (int k = 0; k < nc; ++k) L = fma(lk[k], ek[k], L); + L_sh = L; - } -- o = o / L; ++ } + float o = 0.0f; + device const float* q = pp + tid; -+ for (int k = 0; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); ++ int k = 0; ++ for (; k + 3 < nc; k += 4, q += 4 * STRIDE) { ++ const float p0 = q[0], p1 = q[STRIDE], p2 = q[2 * STRIDE], p3 = q[3 * STRIDE]; ++ o = fma(p0, ek[k], o); ++ o = fma(p1, ek[k + 1], o); ++ o = fma(p2, ek[k + 2], o); ++ o = fma(p3, ek[k + 3], o); + } +- o = o / L; ++ for (; k < nc; ++k, q += STRIDE) o = fma(*q, ek[k], o); + threadgroup_barrier(mem_flags::mem_threadgroup); + o = o / L_sh; const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; @@ -5325,7 +5333,7 @@ index 8ae189b..d5b594c 100644 } kernel void qsa_topk2(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]], -@@ -133,37 +143,76 @@ kernel void qsa_topk2(device const float* scores [[buffer(0)]], device int* ids +@@ -133,37 +151,76 @@ kernel void qsa_topk2(device const float* scores [[buffer(0)]], device int* ids for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); return; } @@ -5430,7 +5438,7 @@ index 8ae189b..d5b594c 100644 const uint T = prefix; const uint per = (nb + 1023) / 1024, b0 = min(nb, uint(tid) * per), b1 = min(nb, b0 + per); diff --git a/metal/kernels/qsa_prompt.metal b/metal/kernels/qsa_prompt.metal -index df7caf7..444e0c4 100644 +index df7caf79..2ee4a4b6 100644 --- a/metal/kernels/qsa_prompt.metal +++ b/metal/kernels/qsa_prompt.metal @@ -30,8 +30,8 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* @@ -5464,7 +5472,7 @@ index df7caf7..444e0c4 100644 } simdgroup_barrier(mem_flags::mem_threadgroup); if (qsa_stop == 1) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 1.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } -@@ -59,6 +66,17 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* +@@ -59,25 +66,33 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); for (uint16_t i = 0; i < acc.get_capacity(); ++i) if (acc.is_valid_element(i)) acc[i] = 0.0f; @@ -5482,7 +5490,16 @@ index df7caf7..444e0c4 100644 for (int k0 = 0; k0 < QD; k0 += 32) { if (!coop_k4) { for (uint16_t i = 0; i < rk.get_capacity(); ++i) { -@@ -70,14 +88,8 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* + const auto ix = rk.get_multidimensional_index(i); + const int k = k0 + ix[0], id = cid[ix[1]]; +- const uint row = uint(id) * QKV + g; +- rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)])) : half(0.0f); ++ // the row is clamped and the select is on the value: a select with the device load in the arm ++ // loses the threadgroup on the M1 (macOS 26, measured); the discarded value cannot matter ++ const uint row = uint(max(id, 0)) * QKV + g; ++ const float v = float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)]); ++ rk[i] = id >= 0 ? half(v) : half(0.0f); + } } else #pragma clang loop unroll(full) for (uint16_t i = 0; i < rk.get_capacity(); i += 4) { @@ -5499,7 +5516,7 @@ index df7caf7..444e0c4 100644 rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w; } tA = tQ.template slice<32, 16>(k0, 0); -@@ -91,13 +103,16 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* +@@ -91,13 +106,16 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* // ---- softmax pieces per head (lane = head for the reductions; 12 heads) float mh = -INFINITY, lh = 0.0f; if (lane < HPG) { @@ -5516,7 +5533,7 @@ index df7caf7..444e0c4 100644 for (int c = 0; c < QCH; ++c) p[lane * QCH + c] = half(0.0f); } simdgroup_barrier(mem_flags::mem_threadgroup); -@@ -115,12 +130,30 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* +@@ -115,12 +133,32 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); for (uint16_t i = 0; i < acc.get_capacity(); ++i) if (acc.is_valid_element(i)) acc[i] = 0.0f; @@ -5531,7 +5548,9 @@ index df7caf7..444e0c4 100644 + for (uint16_t i = 0; i < rv.get_capacity(); ++i) { + const auto ix = rv.get_multidimensional_index(i); // (cell, dim) + const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; -+ rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vs[c][d >> 6])) : half(0.0f); ++ const uint row = uint(max(id, 0)) * QKV + g; ++ const float v = float(vq[row * QD + d]) * float(vs[c][d >> 6]); ++ rv[i] = id >= 0 ? half(v) : half(0.0f); + } + } else { + uint vo[8]; @@ -5552,7 +5571,7 @@ index df7caf7..444e0c4 100644 } tPh = tP.template slice<32, 16>(c0, 0); op.run(tPh, rv, acc); -@@ -136,7 +169,13 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* +@@ -136,7 +174,13 @@ kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* simdgroup_barrier(mem_flags::mem_threadgroup); acc.store(tT); simdgroup_barrier(mem_flags::mem_threadgroup); @@ -5567,7 +5586,7 @@ index df7caf7..444e0c4 100644 } } if (lane < HPG) { pp[lane * (QD + 2) + QD] = mh; pp[lane * (QD + 2) + QD + 1] = lh; } -@@ -160,6 +199,10 @@ kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const fl +@@ -160,6 +204,10 @@ kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const fl ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { threadgroup float ot[64 * 64]; const int b0 = int(tg.x) * 64, t0 = int(tg.y) * 16; @@ -5578,7 +5597,7 @@ index df7caf7..444e0c4 100644 const uint nb_all = (a.pos0 + a.n_tok) / 4 + 1; // blocks any token of the chunk can score tensor, tensor_inline> tQ((device float*)qidx, dextents(ID, int(a.n_tok) * IH)); tensor, tensor_inline> tK((device float*)pooled, dextents(ID, int(nb_all))); -@@ -171,9 +214,10 @@ kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const fl +@@ -171,9 +219,10 @@ kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const fl op.run(mA, mB, tO); threadgroup_barrier(mem_flags::mem_threadgroup); const uint stride = a.max_cells / 4 + 2; @@ -5592,7 +5611,7 @@ index df7caf7..444e0c4 100644 const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; if (uint(b) >= n_bid) continue; float s = 0.0f; -@@ -181,13 +225,13 @@ kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const fl +@@ -181,13 +230,13 @@ kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const fl scores[uint(t) * stride + b] = s; } // the tail block of each token in this threadgroup's block range (the first column threadgroup does it) @@ -5609,7 +5628,7 @@ index df7caf7..444e0c4 100644 float s = 0.0f; for (int h = 0; h < IH; ++h) { float d = 0.0f; -@@ -198,6 +242,7 @@ kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const fl +@@ -198,6 +247,7 @@ kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const fl if ((n_kv & 3) != 0) s += 1e9f; if (lane == 0) scores[uint(t) * stride + n_bid] = s; } @@ -5617,7 +5636,7 @@ index df7caf7..444e0c4 100644 } // qsa_flash: the whole attention of a (kv head g, token t) in one threadgroup (no split-K partials, no merge): the -@@ -211,10 +256,10 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq +@@ -211,10 +261,10 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq 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 HPG = QH / QKV; @@ -5632,7 +5651,7 @@ index df7caf7..444e0c4 100644 threadgroup float fac[16], lsum[16], mrun[16]; const int g = int(tg.x), t = int(tg.y); const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); -@@ -240,15 +285,26 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq +@@ -240,15 +290,26 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq for (uint16_t i = 0; i < acc.get_capacity(); ++i) if (acc.is_valid_element(i)) tot[i] = 0.0f; const int d0 = sg * 64; @@ -5662,7 +5681,7 @@ index df7caf7..444e0c4 100644 // S partial over this simdgroup's 64 dims { auto rk = opS.template get_right_input_cooperative_tensor(); -@@ -257,68 +313,111 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq +@@ -257,68 +318,114 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq __tensor_ops_detail::__remove_addrspace_t, float>(); for (uint16_t i = 0; i < s.get_capacity(); ++i) if (s.is_valid_element(i)) s[i] = 0.0f; @@ -5685,9 +5704,11 @@ index df7caf7..444e0c4 100644 for (uint16_t i = 0; i < rk.get_capacity(); ++i) { const auto ix = rk.get_multidimensional_index(i); const int k = k0 + ix[0], id = cid[ix[1]]; - const uint row = uint(id) * QKV + g; +- const uint row = uint(id) * QKV + g; - rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)])) : half(0.0f); -+ rk[i] = id >= 0 ? half(float(kq[row * QD + k]) * float(ks_tg[ix[1]][k >> 6])) : half(0.0f); ++ const uint row = uint(max(id, 0)) * QKV + g; ++ const float v = float(kq[row * QD + k]) * float(ks_tg[ix[1]][k >> 6]); ++ rk[i] = id >= 0 ? half(v) : half(0.0f); } } else #pragma clang loop unroll(full) @@ -5779,11 +5800,17 @@ index df7caf7..444e0c4 100644 - for (int c0 = 0; c0 < QCH; c0 += 32) { // K = 32 cells a run (K = 64 is wrong on the M4) + for (int c0 = 0; c0 < nc; c0 += 32) { // K = 32 cells a run (K = 64 is wrong on the M4) + if (!coop_k4) { -+#pragma clang loop unroll(full) + #pragma clang loop unroll(full) +- for (uint16_t i = 0; i < rv.get_capacity(); ++i) { +- const auto ix = rv.get_multidimensional_index(i); // (cell, dim) +- const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; +- rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vsc[(uint(id) * QKV + g) * 4 + sg])) : half(0.0f); + for (uint16_t i = 0; i < rv.get_capacity(); ++i) { + const auto ix = rv.get_multidimensional_index(i); // (cell, dim) + const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; -+ rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vs_tg[c][sg])) : half(0.0f); ++ const uint row = uint(max(id, 0)) * QKV + g; ++ const float v = float(vq[row * QD + d]) * float(vs_tg[c][sg]); ++ rv[i] = id >= 0 ? half(v) : half(0.0f); + } + } else { + // a lane holds the same eight cells (two runs of four, 16 apart) at each of its dims: as for rk, row @@ -5800,17 +5827,13 @@ index df7caf7..444e0c4 100644 +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); i += 4) { + const int d = rv.get_multidimensional_index(i)[1], r = i < rv.get_capacity() / 2 ? 0 : 4; - #pragma clang loop unroll(full) -- for (uint16_t i = 0; i < rv.get_capacity(); ++i) { -- const auto ix = rv.get_multidimensional_index(i); // (cell, dim) -- const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; -- rv[i] = id >= 0 ? half(float(vq[(uint(id) * QKV + g) * QD + d]) * float(vsc[(uint(id) * QKV + g) * 4 + sg])) : half(0.0f); ++#pragma clang loop unroll(full) + for (int j = 0; j < 4; ++j) rv[i + j] = half(float(vq[vo[r + j] + uint(d)]) * vs8[r + j]); + } } tPh = tP.template slice<32, 16>(c0, 0); opO.run(tPh, rv, acc); -@@ -327,14 +426,28 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq +@@ -327,14 +434,28 @@ kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq if (acc.is_valid_element(i)) tot[i] = fma(tot[i], fac[acc.get_multidimensional_index(i)[1]], acc[i]); } threadgroup_barrier(mem_flags::mem_threadgroup); @@ -5845,7 +5868,7 @@ index df7caf7..444e0c4 100644 } } diff --git a/metal/kernels/quant.metal b/metal/kernels/quant.metal -index 51fecca..fe667be 100644 +index 51fecca7..fe667be5 100644 --- a/metal/kernels/quant.metal +++ b/metal/kernels/quant.metal @@ -34,11 +34,11 @@ template struct fmt; @@ -6056,7 +6079,7 @@ index 51fecca..fe667be 100644 template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { diff --git a/metal/kernels/router.metal b/metal/kernels/router.metal -index e5c7527..5c6bf59 100644 +index e5c75274..5c6bf590 100644 --- a/metal/kernels/router.metal +++ b/metal/kernels/router.metal @@ -5,44 +5,108 @@ @@ -6237,7 +6260,7 @@ index e5c7527..5c6bf59 100644 + +template [[host_name("router_route_8")]] kernel decltype(router_route) router_route; diff --git a/metal/kernels/sample.metal b/metal/kernels/sample.metal -index 8443d36..4531ca2 100644 +index 8443d360..4531ca2a 100644 --- a/metal/kernels/sample.metal +++ b/metal/kernels/sample.metal @@ -13,6 +13,7 @@ static inline void better(thread float& bv, thread int& bi, float v, int i) { @@ -6257,7 +6280,7 @@ index 8443d36..4531ca2 100644 } diff --git a/metal/kernels/units.metal b/metal/kernels/units.metal -index 15f584c..25d0761 100644 +index 15f584cb..25d0761a 100644 --- a/metal/kernels/units.metal +++ b/metal/kernels/units.metal @@ -107,26 +107,33 @@ template <> struct X { @@ -6307,7 +6330,7 @@ index 15f584c..25d0761 100644 xp[4 * plane + t * per + i] = s0 + s1; } diff --git a/metal/src/engine/arena_budget.h b/metal/src/engine/arena_budget.h -index 4cc024b..513a944 100644 +index 4cc024b7..513a944d 100644 --- a/metal/src/engine/arena_budget.h +++ b/metal/src/engine/arena_budget.h @@ -44,6 +44,9 @@ struct ArenaMem { @@ -6339,7 +6362,7 @@ index 4cc024b..513a944 100644 /// The fill as a plan, for the test: the pairs' bytes in rank order, steps of `step` bytes, `avail(filled)` the free diff --git a/metal/src/engine/msl/mtp.metal b/metal/src/engine/msl/mtp.metal -index 7159330..051dcf7 100644 +index 7159330f..051dcf72 100644 --- a/metal/src/engine/msl/mtp.metal +++ b/metal/src/engine/msl/mtp.metal @@ -1,68 +1,179 @@ @@ -6565,7 +6588,7 @@ index 7159330..051dcf7 100644 + } } diff --git a/metal/src/engine/msl/mtp_chain.metal b/metal/src/engine/msl/mtp_chain.metal -index 62500df..efee47e 100644 +index 62500df2..efee47ee 100644 --- a/metal/src/engine/msl/mtp_chain.metal +++ b/metal/src/engine/msl/mtp_chain.metal @@ -13,6 +13,255 @@ @@ -6909,7 +6932,7 @@ index 62500df..efee47e 100644 } diff --git a/metal/src/engine/msl/ple.metal b/metal/src/engine/msl/ple.metal -index 7036eae..d323f0e 100644 +index 7036eaeb..d323f0ea 100644 --- a/metal/src/engine/msl/ple.metal +++ b/metal/src/engine/msl/ple.metal @@ -9,43 +9,32 @@ @@ -7364,7 +7387,7 @@ index 7036eae..d323f0e 100644 + for (int r = 0; r < 9; ++r) h[r] = col[r]; +} diff --git a/metal/src/engine/prompt_chain.inc b/metal/src/engine/prompt_chain.inc -index 573b460..a095124 100644 +index 573b460e..a0951240 100644 --- a/metal/src/engine/prompt_chain.inc +++ b/metal/src/engine/prompt_chain.inc @@ -272,16 +272,14 @@ bool Window::pc_chunk(const int64_t* tokens, int nT, int64_t pos0, std::string& @@ -7401,7 +7424,7 @@ index 573b460..a095124 100644 g.ind_off = 8 * 4; 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 +index 00000000..e2f25dc4 --- /dev/null +++ b/metal/src/engine/spec_min_p.hpp @@ -0,0 +1,62 @@ @@ -7468,7 +7491,7 @@ index 0000000..e2f25dc + +} // namespace strata::metal_engine diff --git a/metal/src/engine/stream.cpp b/metal/src/engine/stream.cpp -index 9db74e2..97ae303 100644 +index 9db74e25..97ae303e 100644 --- a/metal/src/engine/stream.cpp +++ b/metal/src/engine/stream.cpp @@ -62,6 +62,11 @@ bool ExpertStream::open(const gguf::File& f, std::string& err) { @@ -8248,7 +8271,7 @@ index 9db74e2..97ae303 100644 } // namespace strata::metal_engine diff --git a/metal/src/engine/stream.h b/metal/src/engine/stream.h -index 2806938..6cbac41 100644 +index 2806938a..6cbac41a 100644 --- a/metal/src/engine/stream.h +++ b/metal/src/engine/stream.h @@ -53,6 +53,7 @@ public: @@ -8355,7 +8378,7 @@ index 2806938..6cbac41 100644 diff --git a/metal/src/engine/suffix_window.hpp b/metal/src/engine/suffix_window.hpp new file mode 100644 -index 0000000..edc7cd6 +index 00000000..edc7cd61 --- /dev/null +++ b/metal/src/engine/suffix_window.hpp @@ -0,0 +1,30 @@ @@ -8390,7 +8413,7 @@ index 0000000..edc7cd6 + +} // namespace strata::spec diff --git a/metal/src/engine/window.cpp b/metal/src/engine/window.cpp -index 9a6d1d3..074b64f 100644 +index 9a6d1d32..074b64fb 100644 --- a/metal/src/engine/window.cpp +++ b/metal/src/engine/window.cpp @@ -22,15 +22,21 @@ @@ -11591,7 +11614,7 @@ index 9a6d1d3..074b64f 100644 g_window->prompt_region(false); return ok; diff --git a/metal/src/mtl_rt.h b/metal/src/mtl_rt.h -index e5d7259..9e12c80 100644 +index e5d72592..9e12c80a 100644 --- a/metal/src/mtl_rt.h +++ b/metal/src/mtl_rt.h @@ -88,7 +88,7 @@ public: @@ -11616,7 +11639,7 @@ index e5d7259..9e12c80 100644 double run_wall(const std::vector& ds, int reps, double* gpu_seconds = nullptr); diff --git a/metal/src/mtl_rt.mm b/metal/src/mtl_rt.mm -index 4c6045b..106657d 100644 +index 4c6045b2..106657d5 100644 --- a/metal/src/mtl_rt.mm +++ b/metal/src/mtl_rt.mm @@ -183,13 +183,13 @@ Buffer Context::alloc(size_t bytes) { @@ -11686,7 +11709,7 @@ index 4c6045b..106657d 100644 [cb waitUntilCompleted]; diff --git a/metal/tests/comb_fold_parity.cpp b/metal/tests/comb_fold_parity.cpp new file mode 100644 -index 0000000..8606e13 +index 00000000..8606e130 --- /dev/null +++ b/metal/tests/comb_fold_parity.cpp @@ -0,0 +1,1291 @@ @@ -12982,7 +13005,7 @@ index 0000000..8606e13 + return fails ? 1 : 0; +} diff --git a/metal/tests/dense_mm_parity.cpp b/metal/tests/dense_mm_parity.cpp -index 33ce450..969cb79 100644 +index 33ce4500..969cb790 100644 --- a/metal/tests/dense_mm_parity.cpp +++ b/metal/tests/dense_mm_parity.cpp @@ -3,13 +3,22 @@ @@ -13361,7 +13384,7 @@ index 33ce450..969cb79 100644 + return dq_fails ? 1 : 0; } diff --git a/metal/tests/expert_parity.cpp b/metal/tests/expert_parity.cpp -index 3bd15ba..10d082b 100644 +index 3bd15ba3..10d082b7 100644 --- a/metal/tests/expert_parity.cpp +++ b/metal/tests/expert_parity.cpp @@ -2,6 +2,7 @@ @@ -13585,7 +13608,7 @@ index 3bd15ba..10d082b 100644 std::memcpy(bw.as() + r * 8 * TOPK, rw[r].data(), 4 * rw[r].size()); groups += std::set(rids[r].begin(), rids[r].end()).size(); diff --git a/metal/tests/fuse_parity.cpp b/metal/tests/fuse_parity.cpp -index e67d9e5..6c6fe1e 100644 +index e67d9e58..6c6fe1e6 100644 --- a/metal/tests/fuse_parity.cpp +++ b/metal/tests/fuse_parity.cpp @@ -1,14 +1,18 @@ @@ -13959,7 +13982,7 @@ index e67d9e5..6c6fe1e 100644 } std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); diff --git a/metal/tests/gdn_parity.cpp b/metal/tests/gdn_parity.cpp -index 7288eb1..09e7449 100644 +index 7288eb1d..09e74497 100644 --- a/metal/tests/gdn_parity.cpp +++ b/metal/tests/gdn_parity.cpp @@ -4,6 +4,7 @@ @@ -14225,7 +14248,7 @@ index 7288eb1..09e7449 100644 std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); return fails ? 1 : 0; diff --git a/metal/tests/gdn_prompt_parity.cpp b/metal/tests/gdn_prompt_parity.cpp -index 24cb85b..beaee25 100644 +index 24cb85bf..beaee258 100644 --- a/metal/tests/gdn_prompt_parity.cpp +++ b/metal/tests/gdn_prompt_parity.cpp @@ -3,6 +3,7 @@ @@ -14490,7 +14513,7 @@ index 24cb85b..beaee25 100644 return 0; } diff --git a/metal/tests/gemv_parity.cpp b/metal/tests/gemv_parity.cpp -index 5687527..93d5760 100644 +index 5687527f..93d57603 100644 --- a/metal/tests/gemv_parity.cpp +++ b/metal/tests/gemv_parity.cpp @@ -4,6 +4,13 @@ @@ -14975,7 +14998,7 @@ index 5687527..93d5760 100644 cx.free(W); cx.free(x); diff --git a/metal/tests/moe_mm_parity.cpp b/metal/tests/moe_mm_parity.cpp -index 6f51f1e..8e06a77 100644 +index 6f51f1e7..a44397cf 100644 --- a/metal/tests/moe_mm_parity.cpp +++ b/metal/tests/moe_mm_parity.cpp @@ -3,6 +3,8 @@ @@ -15132,12 +15155,13 @@ index 6f51f1e..8e06a77 100644 } // namespace int main(int argc, char** argv) { +- int gu = ref::Q2_0, dn = ref::Q2_0, T = 256, TM = 16, TN = 64, TK = 16, SG = 2, DTN = -1, DTK = -1, DSG = -1; + if (argc > 1 && std::string(argv[1]) == "--gather-bench") { + return gather_bench(argc > 2 ? std::atoi(argv[2]) : 2048, + argc > 3 ? std::atoi(argv[3]) : 512, + argc > 4 ? std::atoi(argv[4]) : N_EMBD); + } - int gu = ref::Q2_0, dn = ref::Q2_0, T = 256, TM = 16, TN = 64, TK = 16, SG = 2, DTN = -1, DTK = -1, DSG = -1; ++ int gu = ref::Q2_0, dn = ref::Q2_0, T = 256, TM = 16, TN = 64, TK = 64, SG = 2, DTN = -1, DTK = -1, DSG = -1; bool bench = false, v2 = false, v3 = false, v4 = false, v5 = false; int MS = 1; + bool gu_only = false, dn_only = false; @@ -15212,7 +15236,7 @@ index 6f51f1e..8e06a77 100644 if (DTN < 0) { DTN = TN; DTK = TK; DSG = SG; } if (dn == ref::IQ4_NL && DTK < 32) DTK = 32; const std::string tl = std::to_string(TM) + "_" + std::to_string(TN) + "_" + std::to_string(TK) + "_" + std::to_string(SG); -@@ -105,15 +271,13 @@ int main(int argc, char** argv) { +@@ -105,31 +271,29 @@ int main(int argc, char** argv) { : cx.pipeline(lib, (v2 ? "mm_gu2_" : "mm_gu_") + std::to_string(gu) + "_" + tl, probe), pdn = v3 ? cx.pipeline(lib, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4)) : cx.pipeline(lib, (v2 ? "mm_down2_" : "mm_down_") + std::to_string(dn) + "_" + (v2 ? dtl : tl)), @@ -15230,15 +15254,27 @@ index 6f51f1e..8e06a77 100644 mtl::Dispatch pl; pl.p = &pplace; pl.args = {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok), mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}; pl.grid[0] = nblk; pl.tg[0] = 1024; ds.push_back(pl); } if (gemm) { -@@ -121,7 +285,7 @@ int main(int argc, char** argv) { + MmArgs ma{(uint64_t)layer * NE * blob, (uint64_t)layer * NE * blob + up_off, (uint64_t)layer * NE * blob + down_off, blob, blob, (uint32_t)gu_row, (uint32_t)d_row, N_EMBD, N_FF}; - if (v2) { +- if (v2) { ++ if (v2 || v3) { const uint32_t ne = N_EMBD; - mtl::Dispatch ga; ga.p = &pgat; ga.args = {mtl::Arg::B(x), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; + mtl::Dispatch ga; ga.p = &pgat; ga.args = {mtl::Arg::B(xh), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; ga.grid[0] = meta.as()[1]; ga.tg[0] = 256; ds.push_back(ga); } mtl::Dispatch g; g.p = &pgu; +- if (v2) g.args = {mtl::Arg::B(W), mtl::Arg::B(XS), mtl::Arg::B(HS), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; ++ if (v2 || v3) g.args = {mtl::Arg::B(W), mtl::Arg::B(XS), mtl::Arg::B(HS), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + else g.args = {mtl::Arg::B(W), mtl::Arg::B(x), mtl::Arg::B(H), mtl::Arg::B(stok), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + g.grid[0] = N_FF / (v3 ? TN * SG : TN); g.grid[1] = meta.as()[0]; g.tg[0] = 32 * SG; ds.push_back(g); +- mtl::Dispatch d; d.p = &pdn; d.args = {mtl::Arg::B(W), mtl::Arg::B(v2 ? HS : H), mtl::Arg::B(YE), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; +- d.grid[0] = N_EMBD / (v3 ? DTN * DSG : v2 ? DTN : TN); d.grid[1] = meta.as()[0]; d.tg[0] = 32 * (v2 ? DSG : SG); ds.push_back(d); ++ mtl::Dispatch d; d.p = &pdn; d.args = {mtl::Arg::B(W), mtl::Arg::B(v2 || v3 ? HS : H), mtl::Arg::B(YE), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; ++ d.grid[0] = N_EMBD / (v3 ? DTN * DSG : v2 ? DTN : TN); d.grid[1] = meta.as()[0]; d.tg[0] = 32 * ((v2 || v3) ? DSG : SG); ds.push_back(d); + const uint32_t ca[4] = {(uint32_t)T, N_EMBD, TOPK, 0}; + mtl::Dispatch cb; cb.p = &pcomb; cb.args = {mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(y), mtl::Arg::V(ca)}; + cb.grid[0] = (T * N_EMBD + 255) / 256; cb.tg[0] = 256; ds.push_back(cb); @@ -142,6 +306,24 @@ int main(int argc, char** argv) { } const uint32_t ntiles = meta.as()[0]; @@ -15378,7 +15414,7 @@ index 6f51f1e..8e06a77 100644 if (v2) std::printf("gather %.1f us\n", timeit(false, true, 0) * 1e6); diff --git a/metal/tests/ple_moe_parity.cpp b/metal/tests/ple_moe_parity.cpp new file mode 100644 -index 0000000..fb40e33 +index 00000000..fb40e33e --- /dev/null +++ b/metal/tests/ple_moe_parity.cpp @@ -0,0 +1,340 @@ @@ -15723,7 +15759,7 @@ index 0000000..fb40e33 + return fails ? 1 : 0; +} diff --git a/metal/tests/prompt_chain.cpp b/metal/tests/prompt_chain.cpp -index 876afec..88759e9 100644 +index 876afeca..88759e98 100644 --- a/metal/tests/prompt_chain.cpp +++ b/metal/tests/prompt_chain.cpp @@ -413,11 +413,9 @@ int main(int argc, char** argv) { @@ -15751,7 +15787,7 @@ index 876afec..88759e9 100644 g.indirect = &meta; g.ind_off = 8 * 4; diff --git a/metal/tests/qsa_parity.cpp b/metal/tests/qsa_parity.cpp -index 8162ec2..2897b5c 100644 +index 8162ec2e..2897b5ca 100644 --- a/metal/tests/qsa_parity.cpp +++ b/metal/tests/qsa_parity.cpp @@ -4,6 +4,15 @@ @@ -15987,7 +16023,7 @@ index 8162ec2..2897b5c 100644 { std::vector ds; diff --git a/metal/tests/qsa_prompt_parity.cpp b/metal/tests/qsa_prompt_parity.cpp -index 18e1b3a..02dde83 100644 +index 18e1b3af..02dde834 100644 --- a/metal/tests/qsa_prompt_parity.cpp +++ b/metal/tests/qsa_prompt_parity.cpp @@ -1,7 +1,12 @@ @@ -16493,7 +16529,7 @@ index 18e1b3a..02dde83 100644 return fails ? 1 : 0; } diff --git a/metal/tests/route_parity.cpp b/metal/tests/route_parity.cpp -index 2f72329..f551318 100644 +index 2f72329a..f5513183 100644 --- a/metal/tests/route_parity.cpp +++ b/metal/tests/route_parity.cpp @@ -1,5 +1,10 @@ @@ -16995,7 +17031,7 @@ index 2f72329..f551318 100644 + return total_fails == 0 && worst_w < 2e-6 ? 0 : 1; } diff --git a/metal/tests/sample_parity.cpp b/metal/tests/sample_parity.cpp -index 049439a..2cb5b4a 100644 +index 049439af..2cb5b4a4 100644 --- a/metal/tests/sample_parity.cpp +++ b/metal/tests/sample_parity.cpp @@ -1,61 +1,260 @@ @@ -17307,7 +17343,7 @@ index 049439a..2cb5b4a 100644 } diff --git a/metal/tests/spec_min_p.cpp b/metal/tests/spec_min_p.cpp new file mode 100644 -index 0000000..ec5ef0c +index 00000000..ec5ef0c6 --- /dev/null +++ b/metal/tests/spec_min_p.cpp @@ -0,0 +1,81 @@ @@ -17393,7 +17429,7 @@ index 0000000..ec5ef0c + 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 +index a47c7630..a0add9b2 100644 --- a/metal/tests/token_chain.cpp +++ b/metal/tests/token_chain.cpp @@ -24,7 +24,7 @@ namespace { @@ -17423,11 +17459,536 @@ index a47c763..a0add9b 100644 r.tg[0] = 512; push(r, "router", 0); } +diff --git a/serve/frontend.py b/serve/frontend.py +index c36a44da..3fda9c56 100644 +--- a/serve/frontend.py ++++ b/serve/frontend.py +@@ -637,6 +637,11 @@ def anthropic_to_messages(req: dict, think_unasked: bool = True) -> tuple[list[d + # "output_config": {"effort": "low" | "medium" | "high"} + thinking = req.get("thinking") + effort = (req.get("output_config") or {}).get("effort") if isinstance(req.get("output_config"), dict) else None ++ # 0.1.45: Anthropic's budget_tokens is a real budget, not only a level: the thinking is closed (wrap-up, , ++ # the answer) after that many tokens. A reasoning_budget_tokens of the request's own wins. ++ bt = thinking.get("budget_tokens") if isinstance(thinking, dict) and thinking.get("type") != "disabled" else None ++ if isinstance(bt, int) and not isinstance(bt, bool) and bt > 0 and req.get("reasoning_budget_tokens") is None: ++ req["reasoning_budget_tokens"] = bt + if isinstance(thinking, dict) and thinking.get("type") == "disabled": + kwargs["enable_thinking"] = False + elif effort: +@@ -1039,6 +1044,7 @@ class OutputParser: + return out + + def _reset_scan(self): ++ self._call_end_checked = 0 # closing-tag search progress for the current wrapped call + self.sp = 0 # how much of self.buf (the call body) the scanner has consumed + self.ss = "name" # name -> between -> str|raw -> ... -> done + self.scall = None # the ToolCall being streamed (its id is reused by the final event) +@@ -1046,6 +1052,20 @@ class OutputParser: + self.sval_started = False + self.sdeclared = {} + ++ def _wrapped_call_end(self) -> int: ++ # Without a closing tag the structural walk cannot finish. Search only ++ # newly appended text plus enough overlap for a tag split across feeds. ++ # A literal closing tag inside a value needs structural disambiguation: ++ # once any candidate appears, keep the original full walk for this call. ++ # Recovery can rewrite the buffer, so it always uses the full walk. ++ if not self.recover and self._call_end_checked >= 0: ++ start = max(0, self._call_end_checked - len(CALL_END) + 1) ++ if self.buf.find(CALL_END, start) < 0: ++ self._call_end_checked = len(self.buf) ++ return -1 ++ self._call_end_checked = -1 ++ return call_end(self.buf) ++ + def _scan(self) -> list[Event]: + """Advance the streaming view of the call body in self.buf (see stream_tools).""" + out = [] +@@ -1152,10 +1172,19 @@ class OutputParser: + def _hold(self, text: str, tags: tuple[str, ...]) -> int: + """Length of the longest suffix of `text` that is a proper prefix of one of `tags`.""" + best = 0 ++ size = len(text) + for tag in tags: +- for n in range(1, len(tag)): +- if text.endswith(tag[:n]): +- best = max(best, n) ++ if len(tag) <= 1: ++ continue ++ # Only a suffix beginning with the tag's first character can match. ++ # Ordinary token deltas contain no '<': one C-level search replaces ++ # checking every possible prefix. Earliest candidate is the longest. ++ start = text.find(tag[0], max(0, size - len(tag) + 1)) ++ while start >= 0 and size - start > best: ++ if tag.startswith(text[start:]): ++ best = size - start ++ break ++ start = text.find(tag[0], start + 1) + return best + + def feed(self, delta: str) -> list[Event]: +@@ -1351,7 +1380,7 @@ class OutputParser: + elif after and not CALL_END.startswith(after): + drop = 0 + else: +- i = call_end(self.buf) ++ i = self._wrapped_call_end() + if i >= 0: + drop = len(CALL_END) + if self.recover: # a batch in one wrapper: the next call follows this one's +diff --git a/serve/logit_bias.py b/serve/logit_bias.py +index 59873a77..700511c0 100644 +--- a/serve/logit_bias.py ++++ b/serve/logit_bias.py +@@ -1,9 +1,12 @@ + """Request-scoped token biases. Normalize before HTTP streaming starts.""" ++import json + import math + import re ++import threading ++from collections import OrderedDict + + +-def normalize(value, vocab_size=None): ++def _normalize(value, vocab_size=None): + if value is None: + return {} + if isinstance(value, dict): +@@ -36,9 +39,52 @@ def normalize(value, vocab_size=None): + return out + + ++# A long list (a standing ban of a whole script is tens of thousands of IDs) is sent with every request and used to be ++# normalized five times per request: in the OpenAI route, twice in validate_request and twice in engine_key (sampling_keys ++# runs for the cache key and again for the GEN line), about 170 ms for 65,929 entries. JSON text is an exact cache key ++# that tells False from 0 and "1" from 1, and catches in-place edits to the same request object. ++# Only a list that normalized without an error is stored, so a bad list still raises on every call. Short lists skip the ++# cache: they cost less than the key. ++_CACHE_MIN = 64 ++_CACHE_MAX = 4 ++_cache = OrderedDict() # (vocab_size, JSON text) -> [normalized dict, engine key or None] ++_lock = threading.Lock() ++ ++ ++def _entry(value, vocab_size): ++ if not isinstance(value, (dict, list)) or len(value) < _CACHE_MIN: ++ return [_normalize(value, vocab_size), None] ++ try: ++ key = (vocab_size, json.dumps(value, separators=(",", ":"))) ++ except (TypeError, ValueError): ++ key = None # something JSON cannot write: _normalize raises for it ++ entry = None ++ if key is not None: ++ with _lock: ++ entry = _cache.get(key) ++ if entry is not None: ++ _cache.move_to_end(key) ++ if entry is None: ++ entry = [_normalize(value, vocab_size), None] ++ if key is not None: ++ with _lock: ++ _cache[key] = entry ++ while len(_cache) > _CACHE_MAX: ++ _cache.popitem(last=False) ++ return entry ++ ++ ++def normalize(value, vocab_size=None): ++ return dict(_entry(value, vocab_size)[0]) ++ ++ + def engine_key(value): +- entries = normalize(value) +- return " logit_bias=" + ",".join(f"{i}:{b:g}" for i, b in sorted(entries.items())) if entries else "" ++ entry = _entry(value, None) ++ if not entry[0]: ++ return "" ++ if entry[1] is None: ++ entry[1] = " logit_bias=" + ",".join(f"{i}:{b:g}" for i, b in sorted(entry[0].items())) ++ return entry[1] + + + def validate_request(req, engine, vocab_size): +diff --git a/serve/runconfig.py b/serve/runconfig.py +index 472b0021..842554ab 100644 +--- a/serve/runconfig.py ++++ b/serve/runconfig.py +@@ -18,7 +18,7 @@ from pathlib import Path + # of the "sampling" block, ("arg", flag) for an engine option kept in "args". + EDITABLE = [ + ("sampling.temperature", ("sampling", "num>=0"), +- "Default temperature for requests that send none (0 = greedy, the default without a sampling block)"), ++ "Default temperature for requests that send none (0 = greedy; without it the model card's: 1.0 thinking, 0.7 not)"), + ("sampling.top_p", ("sampling", "0, with no answer +@@ -254,6 +256,15 @@ def _timeout_env(name: str, default: float) -> float | None: + # still working is never ended by this, only reported once (the recommendation: engine_silence_s, or this variable). + # STRATA_ENGINE_STALL_S overrides it (0 = off). Needs psutil (setup installs it); without it nothing is ended. + ENGINE_STALL_S = _timeout_env("STRATA_ENGINE_STALL_S", 90.0) ++# --batch: new requests take the control lines (one prompt read at a time) short prompts first: a prompt of at most ++# STRATA_ADMIT_SHORT_FIRST tokens goes before longer ones (opt-in; unset or 0: whoever gets the lock first), and a ++# longer one that waited STRATA_ADMIT_MAX_WAIT_S (default 30) counts as short. A chat question then waits for the one ++# read in progress, not for every agent read that came before it. ++try: ++ ADMIT_SHORT_FIRST = max(0, int(os.environ.get("STRATA_ADMIT_SHORT_FIRST") or 0)) ++except ValueError: ++ ADMIT_SHORT_FIRST = 0 ++ADMIT_MAX_WAIT_S = _timeout_env("STRATA_ADMIT_MAX_WAIT_S", 30.0) or 0.0 + STALL_CPU_EPS_S = 0.5 # CPU seconds the engine may use in the whole stall window and still count as frozen + STALL_IO_EPS_B = 1 << 20 # bytes read or written in that window + +@@ -705,6 +716,19 @@ class ConvCacheLog: + return dict(self.state) + + ++def conversation_slots(engine) -> int | None: ++ """How many conversations the engine's cache holds: its --conversation-cache-slots argument, which is known ++ before the engine has started (a lazy start has reported no INFO line yet), else the INFO line it reported.""" ++ args = engine.spawn[1] if hasattr(engine, "spawn") else [] ++ if "--conversation-cache-slots" in args[:-1]: ++ try: ++ return int(args[args.index("--conversation-cache-slots") + 1]) ++ except ValueError: ++ pass ++ slots = (getattr(engine, "info", None) or {}).get("conversation_cache_slots") ++ return slots if isinstance(slots, int) else None ++ ++ + def conversation_cache_view(info: dict, hist: list, totals: dict, parked: dict) -> dict: + """#596: the Monitor's Conversation cache card: the parked conversations (the engine's opt-in + --conversation-cache-mib: budget, slots, what its log says) and how much of the prompts the cache gave back.""" +@@ -1096,6 +1120,13 @@ class StrataEngine: + self.last.update(prompt_read=int(f[14])) + if len(f) >= 16: # #588 (engine 0.1.39+): routed experts read over PCIe + self.last.update(offloaded=int(f[15])) ++ for tok in f[16:]: # key=value fields after the positional ones (the prompt's ++ key, eq, value = tok.partition("=") # first chunk and its streamed expert blobs) ++ if eq and key in ("chunk", "experts_streamed"): ++ try: ++ self.last[key] = int(value) ++ except ValueError: ++ pass + + def vram(self, reserve_mib: int | None, timeout: float = 120.0) -> dict: + """#533: `VRAM ` between requests (the caller holds the service's FIFO): the engine shrinks its +@@ -1173,6 +1204,10 @@ class StrataEngine: + v = tune.get(k) + if isinstance(v, (int, float)) and not isinstance(v, bool) and 0.0 <= float(v) <= 1.0: + keys += f" {k}={float(v)!r}" ++ # a layer split's prompt chunks (--prefill-pipe): 0 off, 1 evened out, > 1 the rig's b/a in tokens ++ v = tune.get("prefill_pipe") ++ if isinstance(v, (int, float)) and not isinstance(v, bool) and 0.0 <= float(v) <= 262144.0: ++ keys += f" pipe={float(v)!r}" + # --aux-cpus for this request: 1 puts the engine's threads that are neither pool workers nor the host on the + # spare CPUs, 0 leaves them where they were created (true/false or 1/0; the engine started without a spare CPU + # ignores a 1). Anything else is not sent. +@@ -1416,10 +1451,11 @@ class StrataEngine: + + YIELDS_MAX = 2 # #656: how often one request's prompt read gives way to a shorter waiting one + +- def _take_control(self, cancel, plen: int, after_epoch: int | None = None): ++ def _take_control(self, cancel, plen: int, after_epoch: int | None = None, new: bool = False): + """Waits for the control lines (one prompt read at a time), yielding None heartbeats; False when cancelled. +- `after_epoch`: a request whose read gave way lets the requests waiting then go first.""" +- entry = [plen] ++ `after_epoch`: a request whose read gave way lets the requests waiting then go first. ++ `new`: the request's first admission (STRATA_ADMIT_SHORT_FIRST orders only these).""" ++ entry = [plen, time.monotonic(), new] # (length, arrival, ordered: STRATA_ADMIT_SHORT_FIRST) + with self.slot_cv: + self.waiting += 1 + self.wait_lens.append(entry) +@@ -1433,7 +1469,17 @@ class StrataEngine: + raise EngineDied("the engine stopped while this request waited; it was not sent") + with self.slot_cv: + turn = after_epoch is None or self.ctl_epoch > after_epoch or self.waiting <= 1 ++ if turn and new and ADMIT_SHORT_FIRST > 0: ++ turn = self._admit_next() is entry + if turn and self.ctl.acquire(timeout=0.5): ++ if new and ADMIT_SHORT_FIRST > 0: ++ with self.slot_cv: # a shorter one came while this one waited in acquire() ++ first = self._admit_next() is entry ++ if not first: ++ self.ctl.release() ++ with self.slot_cv: ++ self.slot_cv.notify_all() ++ continue + if self.alive(): + break + self.ctl.release() # it died (or is being restarted) while this one was getting +@@ -1456,6 +1502,18 @@ class StrataEngine: + self.slot_cv.notify_all() + return True + ++ def _admit_next(self): ++ """STRATA_ADMIT_SHORT_FIRST= (the caller holds slot_cv): the new request that takes the control lines ++ next - one whose prompt is at most that long (it reads little) before a longer one, in arrival order within ++ each; a longer one that has waited STRATA_ADMIT_MAX_WAIT_S counts as short (it is never held back for good). ++ Off (unset or 0): whoever gets the lock first, as before. Only the order changes, never what is read.""" ++ now = time.monotonic() ++ ordered = [e for e in self.wait_lens if len(e) > 2 and e[2]] ++ if not ordered: ++ return None ++ return min(ordered, key=lambda e: (0 if e[0] <= ADMIT_SHORT_FIRST or now - e[1] >= ADMIT_MAX_WAIT_S else 1, ++ e[1])) ++ + SOLO_AGAIN_MAX = 2 # how often a request left alone in a slot goes back to the solo path + SOLO_AGAIN_MIN_LEFT = 32 # ... only with at least this many tokens still allowed (max_tokens) + +@@ -1493,7 +1551,7 @@ class StrataEngine: + prompt, left = list(ids), int(max_new) + reused0 = None # the first read's reused tokens: a continued leg's count includes this request's own output + carry = {} # #1869: the figures of the legs that ended in a hand-back to the solo path (the final DONE adds to them) +- ok = yield from self._take_control(cancel, len(prompt)) ++ ok = yield from self._take_control(cancel, len(prompt), new=True) + if not ok: + return + holding, born = True, self.gen # (the engine it now has the control lines of) +@@ -1803,6 +1861,7 @@ class StrataEngine: + + slot_groups = 1 # --batch-groups (set when the engine starts) + SLOT_PREFIX_MIN = 512 # a held prefix shorter than this is not worth unbalancing the pipeline's groups ++ TURN_ID = 248045 # <|im_start|>: the engine's default --turn-token (its turn checkpoint sits at the last one) + + def pick_slot(self, prompt: list[int]) -> int | None: + """A free slot for `prompt` (the caller holds slot_cv): the one whose held tokens are the longest start of the +@@ -1812,18 +1871,30 @@ class StrataEngine: + if not free: + return None + gs = getattr(self, "slot_gs", 0) +- if gs: +- # pipelined groups: the engine keeps no pipelined slot as a conversation cache, so nothing is worth keeping +- # free for a next turn. A group's window spans its slots up to its last busy one, so first fill a hole +- # below a busy slot of the same group (it costs no extra row), then follow the spreading order ++ if gs > 1: ++ # pipelined groups of several slots: the engine keeps no idle slot below a busy one as a conversation cache ++ # (its pad row writes the slot), so nothing is worth keeping free for a next turn. A group's window spans ++ # its slots up to its last busy one, so first fill a hole below a busy slot of the same group (it costs no ++ # extra row), then follow the spreading order + # (reported on #793: least-recently-used picks left leading holes, [_ B], that cost a whole row each) + def fills_hole(b): + g0 = b - b % gs + return any(self.slot_busy[x] for x in range(b + 1, g0 + gs)) + return min(free, key=lambda b: (not fills_hole(b), self.slot_order.index(b))) + def held_prefix(b): ++ # what of `prompt` the slot's cache serves: the slot's whole held stream (the prompt and what it decoded) ++ # when this prompt continues it, else its turn checkpoint - the engine takes one at the last <|im_start|> of ++ # a prompt (the next request of the chat renders the history again, but not the reply it decoded or the ++ # assistant header), so a prompt that agrees with the held tokens up to there is that chat's next turn ++ # too (#2518: it was only found by the engine's own source search, after pick_slot had sent it to another ++ # slot and the request in that slot had overwritten the chat whose cache the next one needed) + h = self.slot_held[b] +- return len(h) if h and len(h) < len(prompt) and prompt[:len(h)] == h else 0 ++ if not h: ++ return 0 ++ if len(h) < len(prompt) and prompt[:len(h)] == h: ++ return len(h) ++ k = len(h) - 1 - h[::-1].index(self.TURN_ID) if self.TURN_ID in h else 0 ++ return k if 0 < k < len(prompt) and prompt[:k] == h[:k] else 0 + best = max(free, key=held_prefix) + # without --batch-groups a held prefix of any length is taken, as before + if held_prefix(best) >= (self.SLOT_PREFIX_MIN if self.slot_groups > 1 else 1): +@@ -2868,23 +2939,32 @@ class ByteTokenizer: + + ALWAYS = () # specials matched without parse_special (type 4, as ) + ++ @property ++ def max_special_len(self): ++ return max(len(s) for s in self.SPECIALS) ++ + @property + def control_tokens(self): + return [s for s in self.SPECIALS if s not in self.ALWAYS] + + def encode(self, text, parse_special=False, plain=()): +- out, i = [], 0 ++ return self.encode_marked(text, parse_special, plain)[0] ++ ++ def encode_marked(self, text, parse_special=False, plain=()): ++ """encode() and its resume points (after each special), as strata_tokenizer's for PromptEncoder.""" ++ out, marks, i = [], [], 0 + while i < len(text): + for k, s in enumerate(self.SPECIALS): + if (parse_special or s in self.ALWAYS) and text.startswith(s, i) and not any( + a <= i < b for a, b in plain): + out.append(256 + k) + i += len(s) ++ marks.append((i, len(out))) + break + else: + out.extend(text[i].encode("utf-8")) + i += 1 +- return out ++ return out, marks + + def decode(self, ids, errors="replace"): + raw = bytearray() +@@ -3031,10 +3111,12 @@ class Service: + self.literals = literal_tags(getattr(tokenizer, "control_tokens", ())) # texts that stay text inside a message + self.fit_max_tokens = fit_max_tokens # --fit-max-tokens: clamp the output cap instead of 400 + self.aliases: list[str] = [] # #297: other names of the model (the config's `aliases`) ++ self.answer_room = True # 0.1.45: see ANSWER_ROOM_MIN_CAP (the config's "answer_room": false turns it off) + self.reasoning_close_retry = False # #1053 (opt-in): close the thinking once when a reply ends inside it + self.codex_compaction_cache = False # #924 (opt-in): a Codex compaction is rendered with its conversation's tools + self.codex_thread_titles = False # #923 (opt-in): answer Codex's thread-title turns without the engine + self.sampling_defaults = dict(sampling_defaults or {}) # the run config's `sampling` block ++ self.model_arch = None # the config's family ("qwen4exp", "qwen35moe"): None = no model-card defaults + self.shared = {} # the web app's Chat settings for every client (POST /settings) + self.shared_path = None # where they are kept between starts (next to the config) + self.config_effort = None # #1641: the config's "reasoning_effort": the level for requests with none +@@ -3075,6 +3157,9 @@ class Service: + "prompt_ms": 0.0, "decode_ms": 0.0, + "drafts_offered": 0, "drafts_accepted": 0} # #457: the MTP drafts, summed where reported + self.last_timings = None # the last finished request's, llama.cpp's names (/v1/status) ++ # /metrics/prometheus' two throughput gauges need the totals as they stood at the last scrape, because ++ # llama.cpp averages them over the window between two scrapes (#2341). ++ self.scrape_bucket = (0, 0.0, 0, 0.0) + from serve.prometheus import Latencies + self.latencies = Latencies() # GET /metrics in Prometheus text: the latency histograms + self.last_request_at = None # when a request last started or finished +@@ -3104,6 +3189,18 @@ class Service: + self.restore_after_load = None # a session file an image job saved before it unloaded + self.stop_ids = set(tokenizer.encode(IM_END, parse_special=True) + + tokenizer.encode("<|endoftext|>", parse_special=True)) ++ # #567: a prompt re-encodes only what follows the last special token it shares with a recent prompt (the ++ # same ids as a full encode: tools/strata_tokenizer.py PromptEncoder). Tokenizers without resume points ++ # encode in full. ++ self.prompts = None ++ if hasattr(tokenizer, "encode_marked") and hasattr(tokenizer, "max_special_len"): ++ from strata_tokenizer import PromptEncoder ++ # Keep at most one prompt per engine conversation slot, bounded independently of engine config. ++ slots = conversation_slots(engine) ++ if slots is not None and slots > 0: ++ self.prompts = PromptEncoder(tokenizer, keep=min(slots, PROMPT_ENCODER_KEEP_MAX)) ++ else: ++ self.prompts = PromptEncoder(tokenizer) + + def loaded(self) -> bool: + return not hasattr(self.engine, "alive") or self.engine.alive() +@@ -3296,20 +3393,17 @@ class Service: + threading.Thread(target=loop, daemon=True).start() + + def free_vram_mib(self) -> int | None: +- """Free VRAM on the engine's (first) GPU, from NVML (AMD backend: amdgpu's sysfs files, #301); None when it can't +- be read (then nothing is refused).""" ++ """Free VRAM on the engine's (first) GPU, from the backend's reader - NVML, the amdgpu sysfs files (#301) or ++ IOKit on macOS (Metal backend); None when it can't be read (then nothing is refused).""" + try: +- if getattr(self, "backend", None) == "hip": +- from serve.telemetry import free_vram_mib +- return free_vram_mib(int(getattr(self, "gpu_index", 0) or 0), amd=True) +- from serve.telemetry import _Nvml +- nv = _Nvml(int(getattr(self, "gpu_index", 0) or 0)) +- if not nv.ok(): ++ from serve.telemetry import gpu_reader ++ g = gpu_reader(int(getattr(self, "gpu_index", 0) or 0), amd=getattr(self, "backend", None) == "hip") ++ if not g.ok(): + return None +- m = nv.Mem() +- if nv.lib.nvmlDeviceGetMemoryInfo(nv.dev, ctypes.byref(m)) != 0: ++ r = g.read() ++ if r.get("mem_total") is None or r.get("mem_used") is None: + return None +- return int(m.free >> 20) ++ return int((r["mem_total"] - r["mem_used"]) >> 20) + except Exception: + return None + +@@ -3552,7 +3646,7 @@ class Service: + if effort: + if api == "openai": + ctk = req.get("chat_template_kwargs") if isinstance(req.get("chat_template_kwargs"), dict) else {} +- if not req.get("reasoning_effort") and not req.get("reasoning") and \ ++ if req.get("reasoning_effort") in (None, "") and not req.get("reasoning") and \ + "enable_thinking" not in ctk and "reasoning_effort" not in ctk: + req["reasoning_effort"] = effort + elif not req.get("thinking") and not req.get("output_config"): +@@ -3703,8 +3797,15 @@ class Service: + with self.status_lock: + s, t = dict(self.status), dict(self.totals) + running = len(self.live_reqs) or int(bool(s.get("busy"))) +- prompt_n = t["prompt_tokens"] - t["reused"] # what was read, as llama.cpp counts it +- prompt_s, decode_s = t["prompt_ms"] / 1000, t["decode_ms"] / 1000 ++ prompt_n = t["prompt_tokens"] - t["reused"] # what was read, as llama.cpp counts it ++ prompt_s, decode_s = t["prompt_ms"] / 1000, t["decode_ms"] / 1000 ++ # llama.cpp reads these two gauges from a bucket it resets on every scrape (server_metrics::bucket, ++ # "the gauges are averaged over the window between two scrapes"), so an idle server reports 0 rather ++ # than its lifetime average. Keep the same window here, or the names lie to llama.cpp dashboards. ++ was_prompt_n, was_prompt_s, was_predicted_n, was_decode_s = self.scrape_bucket ++ self.scrape_bucket = (prompt_n, prompt_s, t["output_tokens"], decode_s) ++ d_prompt_n, d_prompt_s = prompt_n - was_prompt_n, prompt_s - was_prompt_s ++ d_predicted_n, d_decode_s = t["output_tokens"] - was_predicted_n, decode_s - was_decode_s + rows = [ + ("llamacpp:prompt_tokens_total", "counter", "Prompt tokens read", prompt_n), + ("llamacpp:prompt_tokens_cached_total", "counter", "Prompt tokens the conversation cache already held", +@@ -3712,10 +3813,10 @@ class Service: + ("llamacpp:prompt_seconds_total", "counter", "Prompt read time", prompt_s), + ("llamacpp:tokens_predicted_total", "counter", "Tokens generated", t["output_tokens"]), + ("llamacpp:tokens_predicted_seconds_total", "counter", "Generation time", decode_s), +- ("llamacpp:prompt_tokens_seconds", "gauge", "Average prompt read throughput in tokens/s", +- prompt_n / prompt_s if prompt_s else 0), +- ("llamacpp:predicted_tokens_seconds", "gauge", "Average generation throughput in tokens/s", +- t["output_tokens"] / decode_s if decode_s else 0), ++ ("llamacpp:prompt_tokens_seconds", "gauge", "Average prompt read throughput in tokens/s, since the last scrape", ++ d_prompt_n / d_prompt_s if d_prompt_s > 0 else 0), ++ ("llamacpp:predicted_tokens_seconds", "gauge", "Average generation throughput in tokens/s, since the last scrape", ++ d_predicted_n / d_decode_s if d_decode_s > 0 else 0), + ("llamacpp:requests_processing", "gauge", "Requests running", running), + ("llamacpp:requests_deferred", "gauge", "Requests waiting their turn", s.get("queued") or 0), + ("strata:requests_total", "counter", "Requests finished", t["requests"]), +@@ -3787,17 +3888,21 @@ class Service: + tail = "" if off else EFFORT_TURN.format(EFFORT_TEXT.get(effort, EFFORT_TEXT["medium"])) + return history + tail + full[len(mine):] + +- def encode_prompt(self, messages, tools, kwargs) -> list[int]: ++ def encode_prompt(self, messages, tools, kwargs, store: bool = True) -> list[int]: + """The request's prompt: the template rendered and tokenized. #537: a / written inside a + message's text is encoded as the text it is, not as the model's reasoning markers (the template's own are). + A control token's text (<|im_start|>, <|im_end|>, <|endoftext|>, ...) inside a message is text as well: only +- the control tokens the template writes are control tokens.""" ++ the control tokens the template writes are control tokens. `store=False` (a token count) reuses the kept ++ prompts without adding or replacing one.""" + marked, marked_tools, changed = mark_think_literals(messages, tools, self.literals) + prompt = self.render_prompt(marked, marked_tools, kwargs) +- if not changed: +- return self.tok.encode(prompt, parse_special=True) +- prompt, plain = unmark_think_literals(prompt, self.literals) +- return self.tok.encode(prompt, parse_special=True, plain=plain) ++ plain = () ++ if changed: ++ prompt, plain = unmark_think_literals(prompt, self.literals) ++ if self.prompts is not None: # #567: only what follows the last special shared with a recent prompt ++ return self.prompts.encode(prompt, plain, store) ++ return self.tok.encode(prompt, parse_special=True, plain=plain) if plain else \ ++ self.tok.encode(prompt, parse_special=True) + + @contextlib.contextmanager + def _turn(self, cancel=None): +@@ -3845,14 +3950,26 @@ class Service: + content[n] = {"type": "text", "text": f"[image omitted: {why}]"} + return fetched + ++ def forget_prompt(self, ids): ++ if self.prompts is not None: ++ self.prompts.forget(ids) ++ ++ + def prepare(self, messages, tools, kwargs, max_new=None, force=None, req=None, cancel=None): + """-> (ids, thinking, max_new). An unset or non-positive max_new (some clients send -1) means "unlimited": + the rest of the context. `force` (forced_call): without thinking the reply starts with it, so it ends the prompt; with thinking, Service.run writes it once the thinking is over.""" # `cancel` (#1999): set when the client hangs up; the downloads, the wait for the turn and the encodes stop # between pictures and raise ClientGone @@ -17439,5 +18000,2543 @@ index 5ca7f4e..26a6823 100644 + for it in m["content"] + ] fetched = self._note_unreadable_tool_images(messages, cancel) - ids = self.encode_prompt(messages, tools, kwargs) +- ids = self.encode_prompt(messages, tools, kwargs) ++ ids = encoded_ids = self.encode_prompt(messages, tools, kwargs) if force and kwargs.get("enable_thinking", True) is False: + ids = ids + self.tok.encode(force, parse_special=True) + self.embeddings.path = None +@@ -3914,12 +4031,16 @@ class Service: + room = ctx - CTX_SLACK - len(ids) + if max_new is None or max_new <= 0 or (self.fit_max_tokens and room < 1): + if room < 1: ++ # A prompt refused for the context is not kept (#567): the client cannot resend it unchanged. The ++ # other refusals (engine starting, images) keep it: the same prompt comes back and its encode is reused. ++ self.forget_prompt(encoded_ids) + raise ValueError("request exceeds the context window: " + f"prompt ({len(ids)} tokens) leaves no room to answer in the context " + f"({ctx}); requests are never truncated") + max_new = room + elif max_new > room: + if not self.fit_max_tokens: ++ self.forget_prompt(encoded_ids) + raise ValueError("request exceeds the context window: " + f"prompt ({len(ids)} tokens) + max tokens ({max_new}) exceeds the context " + f"({ctx}); requests are never truncated. Send a smaller " +@@ -4076,12 +4197,23 @@ class Service: + def _run(self, ids, thinking, tools, max_new, sampling, cancel, force=None) -> Iterator[tuple[str, object]]: + """Yields ("event", Event) as text arrives, then ("done", {"finish": .., "completion_tokens": ..}). + `force` (forced_call): the opening of the call the reply must make - see prepare().""" +- defaults = {**self.sampling_defaults, **self.shared} # the config's, then the Chat settings shared with apps ++ # the model card's sampling for this thinking state, then the config's block, then the Chat settings shared with ++ # apps; the request's own fields win over all of them ++ defaults = {**model_card_sampling(self.model_arch, thinking), **self.sampling_defaults, **self.shared} + if defaults: # the request's own fields win (explicit 0 stays greedy) + req_values = {k: v for k, v in (sampling or {}).items() if v is not None} + sampling = {**defaults, **req_values} + # #123: read after the merge, so a budget shared through POST /settings is seen like the other keys + budget = self.reasoning_budget(sampling) if thinking else None ++ if thinking and not budget and self.answer_room and max_new and max_new >= ANSWER_ROOM_MIN_CAP: ++ # 0.1.45: max_tokens counts the thinking too. A client that caps the reply at 8192 (the desktop agent's ++ # harness sends max_completion_tokens: 8192 on every request) and a task the model thinks long about got ++ # the cut-off in the middle of the thinking and NO answer. So thinking is closed early enough to leave ++ # the answer its quarter of the cap (the budget's wrap-up: a short closing sentence, , the answer). ++ # A request with a budget of its own, or max_tokens below ANSWER_ROOM_MIN_CAP, is left as it was. ++ budget = max_new - min(max(ANSWER_RESERVE_MIN, max_new // 4), max_new // 2) ++ if max_new <= 32768: print(f"[strata] thinking limited to {budget} of max_tokens {max_new} so the answer has room " ++ "(\"answer_room\": false in the config turns this off)", flush=True) + if budget and max_new and max_new > 0: + # #984: a budget at or near max_tokens never fires before the reply is cut off while still thinking, so + # it has no answer. Close the thinking early enough to leave room for one. +@@ -4586,13 +4718,13 @@ def request_stats(segments: list[dict]) -> dict: + if not segments: + return {} + result = dict(segments[-1]) +- for key in ("reused", "prompt_read"): ++ for key in ("reused", "prompt_read", "chunk"): + if key in segments[0]: + result[key] = segments[0][key] + else: + result.pop(key, None) + for key in ("prompt_ms", "decode_ms", "generated", "drafts_offered", "drafts_accepted", +- "hits", "lookups", "offloaded", "ram_blobs", "file_blobs", "file_mb"): ++ "hits", "lookups", "offloaded", "ram_blobs", "file_blobs", "file_mb", "experts_streamed"): + if any(segment.get(key) is not None for segment in segments): + result[key] = sum(segment.get(key) or 0 for segment in segments) + return result +@@ -5612,7 +5744,7 @@ def make_handler(svc: Service): + if not svc.loaded() and not getattr(svc.engine, "unloaded", False): + self._json(503, {"error": {"message": "the engine is not running"}}) + return +- defaults = {**svc.sampling_defaults, **svc.shared} ++ defaults = {**model_card_sampling(svc.model_arch, True), **svc.sampling_defaults, **svc.shared} + names = {"repetition_penalty": "repeat_penalty", "penalty_last_n": "repeat_last_n"} + params = {names.get(k, k): v for k, v in defaults.items() + if k in ("temperature", "top_p", "top_k", "min_p", "seed", "repetition_penalty", +@@ -6045,7 +6177,7 @@ def make_handler(svc: Service): + messages, tools, kw = anthropic_to_messages(req, svc.anthropic_think_unasked) + stop_strings(req) # the same 400 as the request itself would get + check_request_sampling(req) # ... likewise a sampling field of the wrong type +- self._json(200, {"input_tokens": len(svc.encode_prompt(messages, tools, kw))}) ++ self._json(200, {"input_tokens": len(svc.encode_prompt(messages, tools, kw, store=False))}) + + def _anthropic(self, req): + svc.load() +@@ -6433,6 +6565,43 @@ def clean_shared_defaults(d) -> dict: + return out + + ++def split_joined_args(cfg: dict) -> None: ++ """#2267: an engine option written as ONE element of the config's "args" - "--ple-io mmap" or "--ple-io=mmap" - reaches ++ the engine as a single argument ("unknown argument: --ple-io mmap"; the usual form is two elements, "--ple-io", ++ "mmap"). Recommend, never force: such an element is split into the flag and its value, and said at startup. Only an ++ element that starts with a flag name (--letters-digits-dashes) followed by a space or "=" is touched, and never one ++ that sits in the value place of the flag before it.""" ++ args = cfg.get("args") ++ if not isinstance(args, list): ++ return ++ out, changed = [], [] ++ for a in args: ++ m = re.match(r"^(--[A-Za-z0-9][A-Za-z0-9-]*)(?:\s+|=)(\S.*)$", a) if isinstance(a, str) else None ++ prev = out[-1] if out else None ++ in_value_place = isinstance(prev, str) and re.match(r"^--[A-Za-z0-9][A-Za-z0-9-]*$", prev) is not None ++ if m is None or in_value_place: ++ out.append(a) ++ continue ++ out += [m.group(1), m.group(2).strip()] ++ changed.append(a) ++ if changed: ++ cfg["args"] = out ++ print("[strata] the config's args hold " + ", ".join(repr(c) for c in changed) + " as one element each: split into the " ++ 'flag and its value (write them as two elements, e.g. "--ple-io", "mmap")', flush=True) ++ ++ ++def read_config(path) -> dict: ++ """The run config's JSON. A mistake in it (a missing comma between two args, a trailing comma) is said with its ++ place in the file instead of a traceback (#2267: editing "args" by hand is where it happens).""" ++ text = Path(path).read_text(encoding="utf-8-sig") # Notepad adds a BOM ++ try: ++ return json.loads(text) ++ except json.JSONDecodeError as e: ++ raise SystemExit(f"[strata] {path} is not valid JSON: {e.msg} at line {e.lineno}, column {e.colno}. Check the commas " ++ 'between the entries of "args" (each option and each value is its own quoted entry, e.g. ' ++ '"--ple-io", "mmap",) and that no entry is followed by a comma at the end of a list.') from None ++ ++ + def lift_reasoning_effort_arg(cfg: dict) -> None: + """#1641: "--reasoning-effort high" in a config's "args" is meant for the thinking level, which is not an engine + option (the engine stopped at the unknown flag). It is the config's "reasoning_effort" (none, low, medium or high: +@@ -6615,6 +6784,25 @@ def check_request_sampling(req) -> None: + FAMILY_MODEL_NAMES = {"qwen4exp": "qwen3.8-flash-next", "qwen35moe": "qwen3.6-35b-a3b"} + + ++# The sampling each model card recommends ("Best Practices"), used when neither the request, the Chat settings nor the ++# run config's `sampling` block says otherwise. Greedy is only for a request or config that asks for temperature 0. ++# min_p 0, repetition_penalty 1 and presence_penalty 0 are the engine's own "off", so they are not listed. ++MODEL_CARD_SAMPLING = { ++ # Qwen3.8-Flash-Next (and Gyro, the same base model) ++ "qwen4exp": {"thinking": {"temperature": 1.0, "top_p": 0.95, "top_k": 20}, ++ "instruct": {"temperature": 0.7, "top_p": 0.8, "top_k": 20, "presence_penalty": 1.5}}, ++ # Qwen3.6-35B-A3B: the general-task sets (the card's precise-coding set is temperature 0.6, presence_penalty 0) ++ "qwen35moe": {"thinking": {"temperature": 1.0, "top_p": 0.95, "top_k": 20, "presence_penalty": 1.5}, ++ "instruct": {"temperature": 0.7, "top_p": 0.8, "top_k": 20, "presence_penalty": 1.5}}, ++} ++ ++ ++def model_card_sampling(arch, thinking) -> dict: ++ """The model card's sampling for the family `arch` and the thinking state ({} for an unknown or unset family).""" ++ sets = MODEL_CARD_SAMPLING.get(arch) if arch else None ++ return dict(sets["thinking" if thinking else "instruct"]) if sets else {} ++ ++ + def default_model_name(cfg: dict) -> str: + """The config's "model_name", else its family's name ("arch": qwen4exp, qwen35moe), else Flash-Next's.""" + return cfg.get("model_name") or FAMILY_MODEL_NAMES.get(cfg.get("arch") or "qwen4exp", "qwen3.8-flash-next") +@@ -6667,7 +6855,8 @@ def main() -> int: + "each, requests spread over them (also \"replicas\" in the config: a number, or " + "[{\"gpus\": [0, 1]}, {\"gpus\": [2, 3]}])") + a = ap.parse_args() +- cfg = json.loads(Path(a.config).read_text(encoding="utf-8-sig")) if a.config else {} # Notepad adds a BOM ++ cfg = read_config(a.config) if a.config else {} ++ split_joined_args(cfg) + lift_reasoning_effort_arg(cfg) + if a.replicas is not None: + cfg["replicas"] = a.replicas +@@ -6768,6 +6957,9 @@ def main() -> int: + model_name=default_model_name(cfg), vision=vision, + sampling_defaults=sampling_defaults, + fit_max_tokens=a.fit_max_tokens or cfg.get("fit_max_tokens") is True) ++ if a.engine == "strata": # the model card's sampling when nothing else is said (see MODEL_CARD_SAMPLING) ++ svc.model_arch = cfg.get("arch") or "qwen4exp" ++ svc.answer_room = cfg.get("answer_room") is not False + svc.reasoning_close_retry = cfg.get("reasoning_close_retry") is True # #1053: opt-in, off by default + if cfg.get("images"): # image generation (serve/images.py), set up by setup --images + icfg = {k: (os.path.abspath(os.path.join(cfg.get("cwd") or ".", v)) +diff --git a/serve/telemetry.py b/serve/telemetry.py +index 3232c63c..697eeec7 100644 +--- a/serve/telemetry.py ++++ b/serve/telemetry.py +@@ -3,7 +3,9 @@ + A background thread samples once a second and keeps the last 60 readings of each series for the sparklines: + - GPU: NVIDIA's own NVML library (nvml.dll / libnvidia-ml.so.1, installed with every driver) through ctypes, so no + pip package is needed: load, VRAM, temperature, power, PCIe link and throughput. With the AMD backend (#301): the +- amdgpu driver's Linux sysfs files - load, VRAM, temperature and power. ++ amdgpu driver's Linux sysfs files - load, VRAM, temperature and power. On macOS (Metal backend): IOKit's ++ IOAccelerator registry - load and VRAM; the temperature and power sensors need root (`powermetrics`), and an Apple ++ Silicon GPU sits on the die, so there is no PCIe link. + - CPU, RAM, disk: `psutil` when it is installed (setup installs it); without it the CPU and RAM readings fall back to + the OS (Windows GlobalMemoryStatusEx / GetSystemTimes, Linux /proc) and the disk rate is absent. + Anything that cannot be read is None; nothing here can stop the server. +@@ -14,6 +16,8 @@ import collections + import ctypes + import os + import platform ++import re ++import subprocess + import sys + import threading + import time +@@ -246,9 +250,114 @@ class _Amd: + return out + + ++def metal_working_set(): ++ """The Apple GPU's memory budget in bytes: MTLCreateSystemDefaultDevice's recommendedMaxWorkingSetSize, through ++ the Metal framework and objc_msgSend (no PyObjC) - the same property metal/src/mtl_rt.mm reads for the engine ++ (its "GPU working set" log line). None off macOS, or when it cannot be read: a selector the device does not ++ answer would raise an Objective-C exception no ctypes call can catch, so the class is asked first.""" ++ if sys.platform != "darwin": ++ return None ++ try: ++ metal = ctypes.CDLL("/System/Library/Frameworks/Metal.framework/Metal") ++ objc = ctypes.CDLL("/usr/lib/libobjc.dylib") ++ metal.MTLCreateSystemDefaultDevice.restype = ctypes.c_void_p ++ objc.sel_registerName.restype = ctypes.c_void_p ++ objc.object_getClass.argtypes = [ctypes.c_void_p] ++ objc.object_getClass.restype = ctypes.c_void_p ++ objc.class_respondsToSelector.argtypes = [ctypes.c_void_p, ctypes.c_void_p] ++ objc.class_respondsToSelector.restype = ctypes.c_ubyte ++ objc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] ++ objc.objc_msgSend.restype = ctypes.c_ulonglong ++ objc.objc_release.argtypes = [ctypes.c_void_p] ++ objc.objc_release.restype = None ++ dev = metal.MTLCreateSystemDefaultDevice() ++ if not dev: ++ return None ++ try: ++ sel = objc.sel_registerName(b"recommendedMaxWorkingSetSize") ++ if not objc.class_respondsToSelector(objc.object_getClass(dev), sel): ++ return None ++ return int(objc.objc_msgSend(ctypes.c_void_p(dev), ctypes.c_void_p(sel))) or None ++ finally: ++ objc.objc_release(dev) # MTLDevice.h: MTLCreateSystemDefaultDevice is NS_RETURNS_RETAINED ++ except (OSError, AttributeError, ValueError): ++ return None ++ ++ ++class _Metal: ++ """macOS: an Apple GPU's readings from the IOKit registry (`ioreg -l -w0 -r -c IOAccelerator`, 22 ms), with ++ _Nvml's interface: the load and the memory are the driver's PerformanceStatistics of the index'th accelerator ++ node - Device Utilization % and the in-use system memory in bytes. The total is the device's Metal working set, ++ the budget one app gets of the unified memory. No temperature or power reading exists without root ++ (`powermetrics`), and an Apple Silicon GPU sits on the die, so the PCIe cards have nothing to show.""" ++ note = ("no GPU temperature or power reading without root (powermetrics), and an Apple Silicon GPU sits on the " ++ "die, so there is no PCIe link to show; the engine's own figures are in its log") ++ cmd = ("ioreg", "-l", "-w0", "-r", "-c", "IOAccelerator") ++ ++ def __init__(self, index=0, run=None): ++ self.index = index ++ self._run = run or self._ioreg ++ self.nodes = len(self._blocks()) ++ self.working_set = metal_working_set() ++ self._name = None ++ b = self._block() ++ if b: ++ m = re.search(r'"model" = "([^"]+)"', b) ++ self._name = m.group(1) if m else None ++ ++ @staticmethod ++ def _ioreg(): ++ try: ++ return subprocess.run(list(_Metal.cmd), capture_output=True, text=True, timeout=10).stdout or "" ++ except (OSError, subprocess.SubprocessError): ++ return "" ++ ++ def _blocks(self): ++ """The accelerator nodes, in the registry's order: every node that carries a PerformanceStatistics.""" ++ try: ++ out = self._run() or "" ++ except OSError: # ioreg missing, or a caller's runner that raises ++ return [] ++ return [b for b in re.split(r"^\s*\+-o ", out, flags=re.M) if '"PerformanceStatistics"' in b] ++ ++ def _block(self): ++ b = self._blocks() ++ return b[self.index] if 0 <= self.index < len(b) else None ++ ++ def ok(self): ++ return self.index < self.nodes ++ ++ def name(self): ++ return self._name ++ ++ def read(self): ++ b = self._block() ++ if b is None: ++ return {} ++ ++ def num(key): ++ m = re.search(r'"%s"=([0-9.]+)' % re.escape(key), b) ++ return float(m.group(1)) if m else None ++ ++ out = {} ++ u = num("Device Utilization %") ++ if u is None: # older chips report only the renderer and the tiler ++ r, t = num("Renderer Utilization %"), num("Tiler Utilization %") ++ u = max(v for v in (r, t) if v is not None) if (r is not None or t is not None) else None ++ out["util"] = int(u) if u is not None else None ++ used = num("In use system memory") ++ out["mem_used"] = int(used) if used is not None else None ++ out["mem_total"] = self.working_set ++ return out ++ ++ + def gpu_reader(index=0, amd=False): +- """The card's readings: NVML (NVIDIA), or the amdgpu sysfs files with the AMD backend (#301).""" +- return _Amd(index) if amd else _Nvml(index) ++ """The card's readings: NVML (NVIDIA), the amdgpu sysfs files with the AMD backend (#301), or IOKit's ++ IOAccelerator registry on macOS (Metal backend). `amd` comes first: no Mac runs the HIP backend, and the ++ tests' AMD trees are read on a Mac too.""" ++ if amd: ++ return _Amd(index) ++ return _Metal(index) if sys.platform == "darwin" else _Nvml(index) + + + _DISPLAY_CLASS = r"SYSTEM\CurrentControlSet\Control\Class\{4d36e968-e325-11ce-bfc1-08002be10318}" +@@ -413,15 +522,21 @@ class Telemetry: + cards = windows_amd_cards() + win_amd = [cards[i] for i in idx if 0 <= i < len(cards)] or cards[:1] + gpu_name = " + ".join(c["name"] for c in win_amd) if win_amd else None ++ # the macOS reader knows the device's working set (the VRAM budget of the unified memory); None elsewhere ++ vram_total = None ++ if win_amd is None and isinstance(self.gpu, _Metal) and self.gpu.ok(): ++ ws = self.gpu.working_set ++ vram_total = ws >> 20 if ws else None + self.static = { + "gpu_name": gpu_name, + "gpu_vram_total_mib": (sum(c["vram_bytes"] for c in win_amd) >> 20 +- if win_amd and all(c["vram_bytes"] for c in win_amd) else None), ++ if win_amd and all(c["vram_bytes"] for c in win_amd) else vram_total), + "gpu_count": len(self.gpus), + # #1380: the AMD readings are the amdgpu driver's Linux sysfs files; a Windows AMD card has none yet, and the + # dashboard said "not readable (NVML)" or showed empty tiles with no word why + "gpu_note": ("no GPU load or VRAM readings for AMD cards on Windows yet (Linux reads them from the amdgpu " +- "driver); the engine's own VRAM figures are in its log" if amd and not self.gpu.ok() else None), ++ "driver); the engine's own VRAM figures are in its log" if amd and not self.gpu.ok() ++ else getattr(self.gpu, "note", None)), + "cpu_name": _cpu_name(), + "cores": (self.ps.cpu_count(logical=False) if self.ps else None) or None, + "threads": os.cpu_count(), +diff --git a/serve/test_effort_paths.py b/serve/test_effort_paths.py +index dbc1ae6f..15525f6b 100644 +--- a/serve/test_effort_paths.py ++++ b/serve/test_effort_paths.py +@@ -130,6 +130,7 @@ class EffortPaths(unittest.TestCase): + "thinking": {"type": "enabled", "budget_tokens": 20000}}), "xhigh") + self.assertEqual(self.level("/v1/responses", {"input": "hello", "reasoning": {"effort": "high"}}), "xhigh") + self.assertEqual(self.level("/v1/chat/completions", {"messages": MSG, "reasoning_effort": "none"}), "none") ++ self.assertEqual(self.level("/v1/chat/completions", {"messages": MSG, "reasoning_effort": False}), "none") + + def test_the_chat_settings_shared_with_apps_beat_the_config(self): + self.svc.config_effort = "low" +diff --git a/serve/test_logit_bias.py b/serve/test_logit_bias.py +index ee3751b1..8543a9c2 100644 +--- a/serve/test_logit_bias.py ++++ b/serve/test_logit_bias.py +@@ -59,3 +59,116 @@ class LogitBias(unittest.TestCase): + finally: + httpd.shutdown() + httpd.server_close() ++ ++ ++class LogitBiasCache(unittest.TestCase): ++ """A long list is normalized once; the cache must not change what a call returns or raises.""" ++ ++ def setUp(self): ++ from serve import logit_bias ++ logit_bias._cache.clear() ++ ++ def test_same_answers_as_the_plain_function(self): ++ from serve.logit_bias import _normalize ++ forms = [{str(i): -100 for i in range(200)}, ++ [[i, False] for i in range(200)], ++ {str(i): (i % 7) - 3.5 for i in range(200)}, ++ [[str(i), 2] for i in range(200)]] ++ for vocab in (None, 1000): ++ for form in forms: ++ want = _normalize(form, vocab) ++ for _ in range(3): # the first call fills the cache, the others read it ++ self.assertEqual(normalize(form, vocab), want) ++ self.assertEqual(normalize(json.loads(json.dumps(form)), vocab), want) # an equal list, another object ++ for form in forms: ++ ref = " logit_bias=" + ",".join(f"{i}:{b:g}" for i, b in sorted(_normalize(form).items())) ++ self.assertEqual(engine_key(form), ref) ++ self.assertEqual(engine_key(form), ref) ++ ++ def test_false_is_not_zero(self): ++ ban = [[i, False] for i in range(100)] ++ zero = [[i, 0] for i in range(100)] ++ self.assertTrue(ban == zero) # Python calls them equal, so an == key would mix them up ++ for first, second in [(ban, zero), (zero, ban)]: ++ self.setUp() ++ self.assertEqual(set(normalize(first).values()), {-100.0 if first is ban else 0.0}) ++ self.assertEqual(set(normalize(second).values()), {-100.0 if second is ban else 0.0}) ++ ++ def test_a_bad_list_is_never_stored(self): ++ good = {str(i): -1 for i in range(100)} ++ normalize(good, 1000) ++ for bad in [{**good, "5": True}, {**good, "5": 101}, {**good, "5": float("nan")}, {**good, "5": "x"}]: ++ for _ in range(2): ++ with self.assertRaises(ValueError): ++ normalize(bad, 1000) ++ pairs = [[i, 1] for i in range(100)] ++ normalize(pairs) ++ with self.assertRaises(ValueError): # True == 1, but a bool is not a bias ++ normalize([[i, 1] for i in range(99)] + [[99, True]]) ++ with self.assertRaises(ValueError): # a duplicate ID ++ normalize([[i, 1] for i in range(100)] + [[0, 1]]) ++ ++ def test_vocabulary_size_is_part_of_the_key(self): ++ ids = {str(i): -1 for i in range(100)} ++ self.assertEqual(len(normalize(ids, 1000)), 100) ++ with self.assertRaises(ValueError): ++ normalize(ids, 50) ++ self.assertEqual(len(normalize(ids)), 100) ++ with self.assertRaises(ValueError): ++ normalize(ids, 50) ++ ++ def test_the_result_is_a_copy(self): ++ ids = {str(i): -1 for i in range(100)} ++ first = normalize(ids) ++ first[0] = 99.0 ++ self.assertEqual(normalize(ids)[0], -1.0) ++ self.assertEqual(validate_request({"logit_bias": ids}, SimpleNamespace(info={"logit_bias": 1}), 1000)[0], -1.0) ++ ++ def test_an_equal_list_in_a_later_request_is_a_hit(self): ++ from serve import logit_bias ++ a = {str(i): -100 for i in range(100)} ++ b = json.loads(json.dumps(a)) ++ self.assertIsNot(a, b) ++ key_a = engine_key(a) ++ key_b = engine_key(b) ++ self.assertIs(key_a, key_b) # one entry, one key string ++ self.assertEqual(len(logit_bias._cache), 1) ++ ++ def test_json_equivalent_pairs_and_token_keys(self): ++ pairs = [[i, 0] for i in range(64)] ++ tuples = [(i, 0) for i in range(64)] ++ numeric = {i: 0 for i in range(64)} ++ strings = {str(i): 0 for i in range(64)} ++ for first, second in ((pairs, tuples), (numeric, strings)): ++ self.assertEqual(normalize(first, 1000), normalize(second, 1000)) ++ self.assertEqual(engine_key(first), engine_key(second)) ++ duplicates = {**numeric, "0": 0} ++ for _ in range(2): ++ with self.assertRaisesRegex(ValueError, "duplicate"): ++ normalize(duplicates, 1000) ++ ++ def test_same_length_mutations_are_validated(self): ++ for value in ([[i, 0] for i in range(64)], {str(i): 0 for i in range(64)}): ++ with self.subTest(form=type(value).__name__): ++ self.assertEqual(normalize(value, 1000)[0], 0) ++ engine_key(value) ++ if isinstance(value, list): ++ value[0][1] = -100 ++ else: ++ value["0"] = -100 ++ self.assertEqual(normalize(value, 1000)[0], -100) ++ self.assertIn("0:-100", engine_key(value)) ++ if isinstance(value, list): ++ value[0][1] = True ++ else: ++ value["0"] = True ++ with self.assertRaises(ValueError): ++ normalize(value, 1000) ++ with self.assertRaises(ValueError): ++ engine_key(value) ++ ++ def test_the_cache_is_bounded(self): ++ from serve import logit_bias ++ for n in range(logit_bias._CACHE_MAX + 3): ++ normalize({str(i): -(n + 1) for i in range(100)}) ++ self.assertEqual(len(logit_bias._cache), logit_bias._CACHE_MAX) +diff --git a/serve/test_parallel.py b/serve/test_parallel.py +index 1a638919..bde52084 100644 +--- a/serve/test_parallel.py ++++ b/serve/test_parallel.py +@@ -248,6 +248,34 @@ class PickSlot(unittest.TestCase): + e.slot_busy = [False] * 3 + self.assertEqual(e.pick_slot([1, 2, 3]), 1) # a held prompt is never the WHOLE prompt + ++ def test_a_rewritten_next_turn_goes_to_the_slot_of_its_chat(self): ++ # #2518: two chats in two slots; each next turn renders the history again but not the reply it decoded or the ++ # assistant header, so the held stream (prompt + reply) is no start of it. The engine's turn checkpoint sits at ++ # the last <|im_start|> of the old prompt: the prompt that agrees up to there belongs to that slot. Before, both ++ # went to the least recently used slot in turn, and the first one overwrote the other chat's cache. ++ T = StrataEngine.TURN_ID ++ hist_a, hist_b = list(range(1000, 1600)), list(range(2000, 2600)) ++ held_a = hist_a + [T, 7, 8, 9] + [50, 51, 52] # prompt (ends with the assistant header) + reply ++ held_b = hist_b + [T, 7, 8, 9] + [60, 61] ++ e = self.engine(2) ++ e.slot_held = [held_a, held_b] ++ e.slot_used = [10.0, 20.0] # slot 0 was used longer ago ++ next_a = hist_a + [T, 5, 6, 7, 8, 9] + list(range(3000, 3100)) # the history again, then new turns ++ next_b = hist_b + [T, 5, 6, 7, 8, 9] + list(range(3000, 3100)) ++ self.assertEqual(e.pick_slot(next_b), 1) # the LRU slot 0 holds the other chat ++ self.assertEqual(e.pick_slot(next_a), 0) ++ e.slot_busy[0] = True # a busy slot is never picked ++ self.assertEqual(e.pick_slot(next_a), 1) ++ e.slot_busy[0] = False ++ # a chat that agrees with neither (a new one, or one that rewrote its history) keeps the old order ++ self.assertEqual(e.pick_slot(list(range(5000, 5200))), 0) ++ self.assertEqual(e.pick_slot(hist_a[:300] + [T, 5, 6] + list(range(3000, 3100))), 0) ++ e.slot_held[0] = [] # an empty slot still comes before evicting one ++ self.assertEqual(e.pick_slot(list(range(5000, 5200))), 0) ++ # the whole held stream continued is still the longer match ++ e.slot_held = [held_a, held_b + [T, 1, 2]] ++ self.assertEqual(e.pick_slot(held_b + [T, 1, 2, 3, 4]), 1) ++ + def test_batch_groups_spread_requests_over_the_groups(self): + # #1249: with --batch-groups 2 over 4 slots (groups {0,1} and {2,3}) a second request goes to the other group, + # and a short held prefix does not pull it back into the busy one +diff --git a/serve/test_reasoning_loop_recovery.py b/serve/test_reasoning_loop_recovery.py +index b53ee56a..c00fc1be 100644 +--- a/serve/test_reasoning_loop_recovery.py ++++ b/serve/test_reasoning_loop_recovery.py +@@ -58,6 +58,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + engine = RecordingEngine(self.tok, [self.loop, + 'I have verified the solution.\n\nComplete solution.']) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = enabled + events = list(service.run(self.ids, True, [], 40000, {}, threading.Event())) + return engine, events +@@ -66,6 +67,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + engine = RecordingEngine(self.tok, [self.loop, + 'Verified.\n\nComplete solution.']) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + request = {'temperature': .2, 'presence_penalty': 0, 'top_p': .95, 'top_k': 64, 'seed': 2} + list(service.run(self.ids, True, [], 40000, request, threading.Event())) +@@ -81,6 +83,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_frozen_cohort_sampler_is_unchanged_during_recovery(self): + engine = RecordingEngine(self.tok, [self.loop, 'Done.\n\nComplete solution.']) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + sampler = {'temperature': 1.0, 'top_p': .95, 'top_k': 20, 'presence_penalty': 1.5, 'seed': 2} + list(service.run(self.ids, True, [], 40000, sampler, threading.Event())) +@@ -91,6 +94,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_a_lower_top_p_and_top_k_are_never_raised_or_lowered(self): + engine = RecordingEngine(self.tok, [self.loop, 'Verified.\n\nComplete solution.']) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = "recover" + list(service.run(self.ids, True, [], 40000, {'temperature': 0, 'top_p': .5, 'top_k': 5}, threading.Event())) + self.assertEqual((engine.sampling_calls[1]['top_p'], engine.sampling_calls[1]['top_k']), (.5, 5)) +@@ -110,6 +114,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_stop_mode_ends_the_reply_at_the_loop(self): + engine = RecordingEngine(self.tok, [self.loop]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = "stop" + events = list(service.run(self.ids, True, [], 40000, {}, threading.Event())) + self.assertEqual(len(engine.calls), 1) # no second pass +@@ -121,6 +126,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + """#1753: the word count never sees the digits, so only the tail's period ends the reply.""" + engine = RecordingEngine(self.tok, ['9^104 = ' + '58747693387' * 3000]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = "stop" + events = list(service.run(self.ids, True, [], 40000, {}, threading.Event())) + self.assertEqual(len(engine.calls), 1) # no second pass +@@ -136,6 +142,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_default_single_token_guard_does_not_detect_passage_loop(self): + engine = RecordingEngine(self.tok, [self.loop]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + events = list(service.run(self.ids, True, [], 16384, {}, threading.Event())) + self.assertEqual(service.repeat_stop_tokens, 256) + self.assertEqual(events[-1][1]['finish'], 'length') +@@ -145,6 +152,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_recovery_is_once_and_never_extends_output_limit(self): + engine = RecordingEngine(self.tok, [self.loop, self.loop]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + events = list(service.run(self.ids, True, [], 16384, {}, threading.Event())) + self.assertEqual(len(engine.calls), 2) +@@ -155,6 +163,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_repeated_final_text_does_not_change_sampler_or_prompt(self): + engine = RecordingEngine(self.tok, [self.loop]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + events = list(service.run(self.ids, False, [], 40000, {}, threading.Event())) + self.assertEqual(len(engine.calls), 1) +@@ -163,6 +172,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_missing_native_high_hint_skips_recovery(self): + engine = RecordingEngine(self.tok, [self.loop]) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + ids = self.tok.encode('<|im_start|>user\nTask<|im_end|>\n', parse_special=True) + events = list(service.run(ids, True, [], 40000, {}, threading.Event())) +@@ -172,6 +182,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + def test_cancel_during_reasoning_does_not_start_recovery(self): + engine = RecordingEngine(self.tok, [self.loop, 'Unused']) + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + cancel = threading.Event() + events = [] +@@ -187,6 +198,7 @@ class ReasoningLoopRecovery(unittest.TestCase): + engine = RecordingEngine(self.tok, [self.loop, 'Done.\n\nAnswer.']) + engine.batch = 2 + service = Service(engine, self.tok, None) ++ service.answer_room = False # these cases test the loop guards at the cap itself + service.reasoning_loop_recovery = True + events = list(service.run(self.ids, True, [], 40000, {}, threading.Event())) + self.assertEqual(events[-1][1]['reasoning_recoveries'], 1) +diff --git a/serve/test_server.py b/serve/test_server.py +index 3a49ff25..5f58d09d 100644 +--- a/serve/test_server.py ++++ b/serve/test_server.py +@@ -26,7 +26,7 @@ from serve.server import (CTX_SLACK, ByteTokenizer, EngineDied, EngineNonFinite, + PP_DONE_TAIL, Service, + StrataEngine, api_key_of, engine_args, key_matches, layer_split_value, prompt_progress, + prompt_tokens_seen, +- request_timings, serve, start_failure_hint) ++ request_stats, request_timings, serve, start_failure_hint) + from types import SimpleNamespace # noqa: E402 + + ROOT = Path(__file__).resolve().parents[1] +@@ -1548,6 +1548,13 @@ class SamplingKeys(unittest.TestCase): + bad = self.keys(strata_tune={"pcie_frac": 3, "spec_min_p": True, "pool_workers": 2}) + self.assertFalse([x for x in bad if x.split("=")[0] in ("pcie_frac", "spec_min_p", "pool_workers")]) + ++ def test_tune_prefill_pipe(self): ++ self.assertIn("pipe=1536.0", self.keys(temperature=0, strata_tune={"prefill_pipe": 1536})) ++ self.assertIn("pipe=0.0", self.keys(temperature=0, strata_tune={"prefill_pipe": 0})) ++ for bad in (-1, True, "800", 1e9): ++ self.assertFalse([x for x in self.keys(temperature=0, strata_tune={"prefill_pipe": bad}) ++ if x.startswith("pipe=")], bad) ++ + def test_tune_aux_cpus(self): + for val, want in ((True, "aux_cpus=1"), (1, "aux_cpus=1"), (False, "aux_cpus=0"), (0, "aux_cpus=0"), + (1.0, "aux_cpus=1"), (0.0, "aux_cpus=0")): +@@ -1771,6 +1778,150 @@ class RecordingPrompt(MockEngine): + yield from super().generate(ids, max_new, sampling, cancel, embeddings) + + ++class IncrementalPrompts(unittest.TestCase): ++ """The prompt encoder: every request's ids are those of a full encode, and a turn reuses the previous one's.""" ++ ++ @classmethod ++ def setUpClass(cls): ++ tok = ByteTokenizer() ++ cls.engine = RecordingPrompt(tok, "Thinking.\n\n\nThe answer.", max_context=CTX) ++ cls.svc = Service(cls.engine, tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) ++ cls.httpd = serve(cls.svc, port=0) ++ cls.base = f"http://127.0.0.1:{cls.httpd.server_address[1]}" ++ ++ @classmethod ++ def tearDownClass(cls): ++ cls.httpd.shutdown() ++ cls.httpd.server_close() ++ ++ post = ClientShapes.post ++ ++ def test_turns(self): ++ self.assertIsNotNone(self.svc.prompts) ++ full = Service(self.engine, self.svc.tok, self.svc.template) ++ full.prompts = None # the reference: the same prompt, encoded in full ++ msgs = [{"role": "system", "content": "Be brief."}] ++ for turn in range(6): ++ msgs.append({"role": "user", "content": f"question {turn} <|im_end|> é你 " * (turn + 1)}) ++ status, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 64, "messages": msgs}) ++ self.assertEqual(status, 200, b) ++ prompt = self.svc.template.render(msgs) ++ self.assertEqual(self.engine.last_ids, full.encode_prompt(msgs, None, {})) ++ if turn: ++ self.assertGreater(self.svc.prompts.last_reused, len(prompt) // 3) ++ msgs.append({"role": "assistant", "content": b["choices"][0]["message"]["content"]}) ++ ++ def conversation(self, n): ++ """Messages of the n-th conversation: distinct from every other one from the number on.""" ++ return [{"role": "user", "content": f"conversation {n} says: " + f"topic {n} sentence {n} " * 20}] ++ ++ def service(self, **info): ++ engine = RecordingPrompt(self.svc.tok, "The answer.", max_context=CTX) ++ engine.info = info ++ return Service(engine, self.svc.tok, self.svc.template) ++ ++ def test_the_kept_prompts_follow_the_engines_conversation_slots(self): ++ svc = self.service(conversation_cache_slots=8) ++ texts = [] ++ for n in range(8): ++ msgs = self.conversation(n) ++ svc.encode_prompt(msgs, None, {}) ++ texts.append(svc.render_prompt(msgs, None, {})) ++ self.assertEqual(len(svc.prompts.entries), 8) ++ self.assertEqual(sorted(e[0] for e in svc.prompts.entries), sorted(texts)) ++ msgs = self.conversation(8) ++ svc.encode_prompt(msgs, None, {}) ++ kept = sorted(e[0] for e in svc.prompts.entries) ++ self.assertEqual(len(kept), 8) ++ self.assertEqual(kept, sorted(texts[1:] + [svc.render_prompt(msgs, None, {})])) ++ ++ def lazy_service(self, args, **info): ++ """A Service over the real StrataEngine class started lazily: no process yet, so no INFO line either.""" ++ engine = StrataEngine("strata", args, lazy=True) ++ engine.info.update(info) ++ return Service(engine, self.svc.tok, self.svc.template) ++ ++ def test_a_lazy_engine_keeps_the_prompts_its_arguments_ask_for(self): ++ svc = self.lazy_service(["--conversation-cache-slots", "8"]) ++ for n in range(10): ++ svc.encode_prompt(self.conversation(n), None, {}) ++ self.assertEqual(len(svc.prompts.entries), 8) ++ ++ def test_the_slot_argument_wins_over_the_info_line_and_a_bad_one_falls_back(self): ++ self.assertEqual(self.lazy_service(["--conversation-cache-slots", "6"], conversation_cache_slots=3) ++ .prompts.keep, 6) ++ self.assertEqual(self.lazy_service(["--conversation-cache-slots", "x"], conversation_cache_slots=3) ++ .prompts.keep, 3) ++ self.assertEqual(self.lazy_service(["--conversation-cache-slots"]).prompts.keep, 4) ++ self.assertEqual(self.lazy_service([]).prompts.keep, 4) ++ ++ def test_prompt_encoder_slot_count_is_capped(self): ++ from serve.server import PROMPT_ENCODER_KEEP_MAX ++ self.assertEqual(PROMPT_ENCODER_KEEP_MAX, 16) ++ self.assertEqual(self.service(conversation_cache_slots=10**9).prompts.keep, 16) ++ self.assertEqual(self.lazy_service(["--conversation-cache-slots", str(10**9)]).prompts.keep, 16) ++ ++ def test_without_a_slot_setting_four_prompts_are_kept(self): ++ svc = self.service() ++ for n in range(6): ++ svc.encode_prompt(self.conversation(n), None, {}) ++ self.assertEqual(len(svc.prompts.entries), 4) ++ ++ def test_counting_tokens_leaves_the_kept_prompts_alone(self): ++ svc = self.service(conversation_cache_slots=2) ++ httpd = serve(svc, port=0) ++ self.addCleanup(httpd.server_close) ++ self.addCleanup(httpd.shutdown) ++ self.base = f"http://127.0.0.1:{httpd.server_address[1]}" ++ main = self.conversation(0) ++ status, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 8, "messages": main}) ++ self.assertEqual(status, 200, b) ++ before = list(svc.prompts.entries) ++ full = self.service() ++ full.prompts = None ++ for n in range(1, 6): ++ other = self.conversation(n) ++ status, b = self.post("/v1/messages/count_tokens", {"model": "m", "messages": other}) ++ self.assertEqual(status, 200, b) ++ self.assertEqual(b["input_tokens"], len(full.encode_prompt(other, None, {}))) ++ self.assertTrue(svc.prompts.entries == before, "the count wrote into the kept prompts") ++ main.append({"role": "assistant", "content": "ok"}) ++ main.append({"role": "user", "content": "and then?"}) ++ svc.encode_prompt(main, None, {}) ++ self.assertGreater(svc.prompts.last_reused, len(svc.render_prompt(main, None, {})) // 2) ++ ++ def test_same_ids_as_a_full_encode_on_varied_conversations(self): ++ import random ++ sys.path.insert(0, str(ROOT / "tools")) ++ from test_strata_tokenizer import conversation_prompts, load_tokenizer ++ toks = [("byte", ByteTokenizer())] ++ qwen = load_tokenizer() ++ if qwen is not None: ++ toks.append(("qwen35", qwen)) ++ for name, tok in toks: ++ svc = Service(self.engine, tok, self.svc.template) ++ for seed in range(3): ++ for what, prompt in conversation_prompts(self.svc.template, random.Random(seed)): ++ with self.subTest(tokenizer=name, seed=seed, what=what): ++ self.assertEqual(svc.prompts.encode(prompt), tok.encode(prompt, parse_special=True)) ++ ++ def test_a_tokenizer_without_resume_points_encodes_in_full(self): ++ class Plain: ++ """ByteTokenizer without the resume points.""" ++ def __init__(self): ++ self.base = ByteTokenizer() ++ ++ def __getattr__(self, name): ++ if name in ("encode_marked", "max_special_len"): ++ raise AttributeError(name) ++ return getattr(self.base, name) ++ svc = Service(self.engine, Plain(), self.svc.template) ++ self.assertIsNone(svc.prompts) ++ msgs = [{"role": "user", "content": "hi"}] ++ self.assertEqual(svc.encode_prompt(msgs, None, {}), ++ ByteTokenizer().encode(svc.render_prompt(msgs, None, {}), parse_special=True)) ++ ++ + class DyingEngine(MockEngine): + """Issue #27: an engine that dies after a few tokens of its first answer, and comes back when restarted.""" + +@@ -1989,6 +2140,39 @@ class PrometheusMetrics(unittest.TestCase): + self.assertEqual((m["strata:requests_total"], m["strata:drafts_offered_total"], + m["strata:drafts_accepted_total"]), (2, 17, 10)) + ++ def test_throughput_gauges_cover_the_window_between_scrapes(self): ++ """#2341: llama.cpp reads these two from a bucket it resets on every scrape, so a scrape that saw no ++ traffic reports 0 - not the lifetime average, which froze at the last request's speed.""" ++ tok = ByteTokenizer() ++ engine = DoneLineEngine(tok, "\n\nok", max_context=CTX, ++ done_lines=["DONE 4 20 400.0 200.0 stop 7 12 5"]) ++ svc = Service(engine, tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) ++ httpd = serve(svc, port=0) ++ base = f"http://127.0.0.1:{httpd.server_address[1]}" ++ ++ def scrape(): ++ with urllib.request.urlopen(base + "/metrics/prometheus", timeout=10) as r: ++ return {line.split()[0]: float(line.split()[1]) ++ for line in r.read().decode().splitlines() if not line.startswith("#")} ++ ++ try: ++ body = json.dumps({"model": "m", "max_tokens": 10, "messages": [{"role": "user", "content": "hi"}]}) ++ with urllib.request.urlopen(urllib.request.Request(base + "/v1/chat/completions", data=body.encode(), ++ headers={"Content-Type": "application/json"}), ++ timeout=30) as r: ++ self.assertEqual(r.status, 200) ++ first = scrape() ++ self.assertGreater(first["llamacpp:prompt_tokens_seconds"], 0) ++ self.assertGreater(first["llamacpp:predicted_tokens_seconds"], 0) ++ idle = scrape() # nothing ran since the scrape above, so the window is empty ++ self.assertEqual((idle["llamacpp:prompt_tokens_seconds"], ++ idle["llamacpp:predicted_tokens_seconds"]), (0, 0)) ++ self.assertEqual(idle["llamacpp:tokens_predicted_total"], # counters stay cumulative ++ first["llamacpp:tokens_predicted_total"]) ++ finally: ++ httpd.shutdown() ++ httpd.server_close() ++ + + class PcieShare(unittest.TestCase): + """#588: the hit rate stays the VRAM share of the lookups; the routed experts the GPU read over PCIe (the DONE +@@ -2507,6 +2691,30 @@ class CancelledRead(unittest.TestCase): + StrataEngine._parse_done(e, "DONE 0 98179 17565.0 0.0 cancel 0 0 0 0 0 0 0 0.0") + self.assertNotIn("prompt_read", e.last) + ++ def test_parse_done_chunk_fields(self): ++ """The prompt's first chunk and its streamed expert blobs come as key=value fields after the positional ones, ++ with or without the five --lookup-chain fields; an older line has neither.""" ++ e = SimpleNamespace() ++ StrataEngine._parse_done(e, "DONE 5 12000 900.0 80.0 stop 3 4 0 7 9 0 0 0.0 12000 2 chunk=1536 " ++ "experts_streamed=24576") ++ self.assertEqual((e.last["chunk"], e.last["experts_streamed"], e.last["offloaded"]), (1536, 24576, 2)) ++ StrataEngine._parse_done(e, "DONE 5 12000 900.0 80.0 stop 3 4 0 7 9 0 0 0.0 12000 2 1 2 3 4 5 chunk=8192 " ++ "experts_streamed=0") ++ self.assertEqual((e.last["chunk"], e.last["experts_streamed"]), (8192, 0)) ++ StrataEngine._parse_done(e, "DONE 5 12000 900.0 80.0 stop 3 4 0 7 9 0 0 0.0 12000 2") ++ self.assertNotIn("chunk", e.last) ++ self.assertNotIn("experts_streamed", e.last) ++ StrataEngine._parse_done(e, "DONE 5 12000 900.0 80.0 stop 3 4 0 7 9 0 0 0.0 12000 2 chunk=x other=1") ++ self.assertNotIn("chunk", e.last) ++ self.assertNotIn("other", e.last) ++ ++ def test_request_stats_chunk_fields(self): ++ """Across a request's continuations the streamed blobs add up; the chunk is the request's own read's.""" ++ stats = request_stats([{"prompt_ms": 10.0, "chunk": 1536, "experts_streamed": 100}, ++ {"prompt_ms": 2.0, "chunk": 64, "experts_streamed": 5}]) ++ self.assertEqual((stats["chunk"], stats["experts_streamed"], stats["prompt_ms"]), (1536, 105, 12.0)) ++ self.assertNotIn("chunk", request_stats([{"prompt_ms": 1.0}, {"prompt_ms": 1.0, "chunk": 64}])) ++ + def test_prompt_tokens_seen(self): + last = {"finish": "cancel", "reused": 1000, "prompt_read": 2000} + self.assertEqual(prompt_tokens_seen(98179, last), 3000) +@@ -2822,6 +3030,46 @@ class ConfigReasoningEffort(unittest.TestCase): + self.assertEqual(cfg, {"args": ["--ctx", "4096"]}) + + ++class JoinedEngineArgs(unittest.TestCase): ++ """#2267: "--ple-io mmap" written as one element of the config's args reached the engine as one argument.""" ++ ++ def split(self, args): ++ from serve.server import split_joined_args ++ cfg = {"args": list(args)} ++ with mock.patch("builtins.print") as pr: ++ split_joined_args(cfg) ++ return cfg["args"], pr.called ++ ++ def test_one_element_becomes_flag_and_value(self): ++ self.assertEqual(self.split(["--spec", "4", "--ple-io mmap"]), (["--spec", "4", "--ple-io", "mmap"], True)) ++ self.assertEqual(self.split(["--ple-io=mmap", "--kv int8"]), (["--ple-io", "mmap", "--kv", "int8"], True)) ++ ++ def test_proper_args_are_untouched(self): ++ args = ["--ple-io", "mmap", "--native", "/data/My Models/a b.gguf", "--max-context", "131072", "--vision"] ++ self.assertEqual(self.split(args), (args, False)) ++ # a value in a flag's value place stays whole, whatever it looks like ++ args = ["--system-prompt", "--be brief please"] ++ self.assertEqual(self.split(args), (args, False)) ++ ++ def test_no_args(self): ++ from serve.server import split_joined_args ++ cfg = {} ++ split_joined_args(cfg) ++ self.assertEqual(cfg, {}) ++ ++ def test_bad_json_names_the_place(self): ++ import tempfile ++ from serve.server import read_config ++ with tempfile.TemporaryDirectory() as d: ++ p = Path(d) / "c.json" ++ p.write_text('{"args": ["--ple-io" "mmap"]}', encoding="utf-8") ++ with self.assertRaises(SystemExit) as cm: ++ read_config(p) ++ self.assertIn("line 1", str(cm.exception)) ++ p.write_text('{"args": ["--ple-io", "mmap"]}', encoding="utf-8") ++ self.assertEqual(read_config(p), {"args": ["--ple-io", "mmap"]}) ++ ++ + class SharedSettings(unittest.TestCase): + """The web app's "Use for other apps too": POST /settings makes its Chat settings every client's defaults.""" + +@@ -3776,6 +4024,45 @@ class ThinkingBudget(unittest.TestCase): + self.assertEqual(text, ThinkingEngine.ANSWER) + self.assertEqual(evs[-2]["delta"]["stop_reason"], "end_turn") + ++ def test_a_cap_that_the_thinking_would_eat_leaves_room_to_answer(self): ++ """0.1.45: max_tokens counts the thinking; a thinking that would fill it gets closed with a quarter left (the ++ agent harness sends max_completion_tokens 8192; the pagoda's thinking took all of it and no answer came).""" ++ self.engine.THOUGHT = "Let me think step by step about two plus two. " * 40 # 1,800 reasoning tokens ++ code, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 1024, ++ "messages": [{"role": "user", "content": "2+2?"}]}) ++ self.assertEqual(code, 200, b) ++ msg = b["choices"][0]["message"] ++ self.assertEqual(msg["content"], ThinkingEngine.ANSWER) ++ self.assertEqual(b["choices"][0]["finish_reason"], "stop") ++ self.assertLess(len(msg["reasoning_content"]), 1024 - 256) ++ self.assertEqual(len(self.engine.prompts), 2) # the wrap-up continuation ++ ++ def test_the_answer_room_has_its_limits(self): ++ self.engine.THOUGHT = "Let me think step by step about two plus two. " * 40 ++ code, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 300, # below the floor: as before ++ "messages": [{"role": "user", "content": "2+2?"}]}) ++ self.assertEqual(b["choices"][0]["finish_reason"], "length") ++ self.assertFalse(b["choices"][0]["message"]["content"]) ++ self.svc.answer_room = False # the config's "answer_room": false ++ code, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 1024, ++ "messages": [{"role": "user", "content": "2+2?"}]}) ++ self.assertEqual(b["choices"][0]["finish_reason"], "length") ++ self.svc.answer_room = True ++ code, b = self.post("/v1/chat/completions", {"model": "m", "max_tokens": 1024, "reasoning_effort": "none", ++ "messages": [{"role": "user", "content": "2+2?"}]}) ++ self.assertEqual(len(self.engine.prompts), 3) # thinking off: no budget, one prompt each ++ ++ def test_anthropic_thinking_budget_tokens_is_a_budget(self): ++ """0.1.45: thinking {type: enabled, budget_tokens: N} closes the thinking after N tokens (it was only a level).""" ++ from serve.server import REASONING_WRAP_UP ++ code, b = self.post("/v1/messages", {"model": "m", "max_tokens": 400, "thinking": {"type": "enabled", "budget_tokens": 30}, ++ "messages": [{"role": "user", "content": "2+2?"}]}) ++ self.assertEqual(code, 200, b) ++ thinking = "".join(c.get("thinking", "") for c in b["content"] if c["type"] == "thinking") ++ self.assertEqual(thinking, ThinkingEngine.THOUGHT[:30] + REASONING_WRAP_UP.split("")[0]) ++ self.assertEqual("".join(c.get("text", "") for c in b["content"] if c["type"] == "text"), ThinkingEngine.ANSWER) ++ self.assertEqual(len(self.engine.prompts), 2) ++ + def test_a_budget_the_thinking_stays_under(self): + code, b = self.openai(reasoning_budget_tokens=10_000) + self.assertEqual(b["choices"][0]["message"]["reasoning_content"], ThinkingEngine.THOUGHT) +@@ -4481,6 +4768,114 @@ class AmdTelemetry(unittest.TestCase): + self.assertEqual(svc.free_vram_mib(), 26 << 10) + + ++IOREG_FIXTURE = ''' ++ +-o AGXAcceleratorG13X ++ { ++ "IOClass" = "AGXAcceleratorG13X" ++ "model" = "Apple M1 Max" ++ "PerformanceStatistics" = {"In use system memory (driver)"=0,"Alloc system memory"=22991831040,"Tiler Utilization %"=36,"recoveryCount"=44,"lastRecoveryTime"=216970984728,"Renderer Utilization %"=35,"TiledSceneBytes"=1671168,"Device Utilization %"=76,"SplitSceneCount"=0,"Allocated PB Size"=76152832,"In use system memory"=13384105984} ++ +-o IOUserClientCreator ++ { ++ "IOUserClientCreator" = "pid 402, WindowServer" ++ } ++ } ++ ++ +-o AGXAccelerator ++ { ++ "IOClass" = "AGXAccelerator" ++ "model" = "Apple M2 Ultra" ++ "PerformanceStatistics" = {"Renderer Utilization %"=3,"Tiler Utilization %"=2,"In use system memory"=1234567890} ++ } ++''' ++ ++WORKING_SET = 26800603136 # bytes; the M1 Max's device in the fixture above ++ ++ ++class MacTelemetry(unittest.TestCase): ++ """The macOS (Metal backend) reader: IOKit's IOAccelerator nodes and the device's Metal working set.""" ++ ++ def _reader(self, index=0, text=IOREG_FIXTURE): ++ from serve import telemetry ++ with mock.patch.object(telemetry, "metal_working_set", return_value=WORKING_SET): ++ return telemetry._Metal(index, run=lambda: text) ++ ++ def test_readings(self): ++ g = self._reader() ++ self.assertTrue(g.ok()) ++ self.assertEqual(g.name(), "Apple M1 Max") ++ with mock.patch("serve.telemetry.metal_working_set", lambda: WORKING_SET): ++ self.assertEqual(g.read(), {"util": 76, "mem_used": 13384105984, "mem_total": WORKING_SET}) ++ # the second accelerator is index 1: the older chips' renderer/tiler-only utilization is the fallback ++ g1 = self._reader(1) ++ self.assertTrue(g1.ok()) ++ with mock.patch("serve.telemetry.metal_working_set", lambda: WORKING_SET): ++ self.assertEqual(g1.read()["util"], 3) # the renderer, the higher of the two ++ self.assertEqual(g1.name(), "Apple M2 Ultra") ++ self.assertFalse(self._reader(5).ok()) ++ ++ def test_working_set_is_read_once_for_samples_and_static(self): ++ from serve import telemetry ++ with mock.patch.object(telemetry, "metal_working_set", return_value=WORKING_SET) as budget: ++ g = telemetry._Metal(run=lambda: IOREG_FIXTURE) ++ for _ in range(3): ++ self.assertEqual(g.read()["mem_total"], WORKING_SET) ++ with mock.patch.object(telemetry, "gpu_reader", return_value=g): ++ t = telemetry.Telemetry() ++ try: ++ self.assertEqual(t.static["gpu_vram_total_mib"], WORKING_SET >> 20) ++ finally: ++ t.close() ++ budget.assert_called_once_with() ++ ++ def test_device_is_released_on_success_and_unsupported_selector(self): ++ from serve import telemetry ++ metal, objc = mock.MagicMock(), mock.MagicMock() ++ metal.MTLCreateSystemDefaultDevice.return_value = 1234 ++ objc.sel_registerName.return_value = 5678 ++ objc.objc_msgSend.return_value = WORKING_SET ++ for supported in (True, False): ++ with self.subTest(supported=supported): ++ objc.objc_release.reset_mock() ++ objc.class_respondsToSelector.return_value = supported ++ with mock.patch.object(telemetry.sys, "platform", "darwin"), \ ++ mock.patch.object(telemetry.ctypes, "CDLL", side_effect=[metal, objc]): ++ self.assertEqual(telemetry.metal_working_set(), WORKING_SET if supported else None) ++ objc.objc_release.assert_called_once_with(1234) ++ ++ def test_no_registry(self): ++ g = self._reader(text="") ++ self.assertFalse(g.ok()) ++ self.assertEqual(g.read(), {}) ++ self.assertIsNone(g.name()) ++ ++ def test_bad_process(self): ++ from serve import telemetry ++ g = telemetry._Metal(0, run=lambda: (_ for _ in ()).throw(OSError("no ioreg"))) ++ self.assertFalse(g.ok()) ++ self.assertEqual(g.read(), {}) ++ ++ def test_note_reaches_the_monitor(self): ++ from serve import telemetry ++ t = telemetry.Telemetry(extra=lambda: {"tok_s": 1.0}) ++ try: ++ with mock.patch.object(telemetry, "gpu_reader", lambda index=0, amd=False: self._reader()): ++ t2 = telemetry.Telemetry(extra=lambda: {"tok_s": 1.0}) ++ try: ++ self.assertEqual(t2.static["gpu_name"], "Apple M1 Max") ++ self.assertIn("powermetrics", t2.static["gpu_note"]) ++ with mock.patch("serve.telemetry.metal_working_set", lambda: WORKING_SET): ++ s = t2.sample() ++ self.assertEqual(s["gpu_util"], 76) ++ self.assertEqual(s["gpu_mem_used"], 13384105984) ++ self.assertEqual(s["gpu_mem_total"], WORKING_SET) ++ self.assertIsNone(s.get("gpu_temp")) # no reading without root; the cards show a dash ++ self.assertIsNone(s.get("gpu_pcie_gen")) ++ finally: ++ t2.close() ++ finally: ++ t.close() ++ ++ + class TelemetryNoDisks(unittest.TestCase): + """psutil.disk_io_counters() returns None where it finds no disk (a gVisor container; Windows with its disk counters + off): the sampler goes on without a disk reading instead of its thread ending at the first sample, which left +diff --git a/serve/test_slot_alloc.py b/serve/test_slot_alloc.py +index 474b4f78..b2b4d35b 100644 +--- a/serve/test_slot_alloc.py ++++ b/serve/test_slot_alloc.py +@@ -79,5 +79,30 @@ class OneGroup(unittest.TestCase): + self.assertEqual(e.pick_slot([1, 2, 3]), 3) # else the least recently used + + ++class OneSlotGroups(unittest.TestCase): ++ """--batch 2 on a two-GPU layer split: two pipelined groups of ONE slot. No slot sits below another in its group, ++ so no pad row writes an idle slot and each one keeps its chat (#2518): the next turn goes back to it.""" ++ ++ def test_the_slot_of_the_chat(self): ++ T = StrataEngine.TURN_ID ++ e = engine(2, 2) ++ self.assertEqual(e.slot_gs, 1) ++ hist_a, hist_b = list(range(1000, 1700)), list(range(2000, 2700)) ++ e.slot_held = [hist_a + [T, 7, 8] + [50, 51], hist_b + [T, 7, 8] + [60]] ++ e.slot_used = [10.0, 20.0] ++ again = lambda h: h + [T, 5, 6] + list(range(3000, 3050)) ++ self.assertEqual(e.pick_slot(again(hist_b)), 1) ++ self.assertEqual(e.pick_slot(again(hist_a)), 0) ++ self.assertEqual(e.pick_slot(list(range(5000, 5100))), 0) # a new chat: the least recently used one ++ ++ def test_groups_of_several_slots_still_ignore_held_tokens(self): ++ T = StrataEngine.TURN_ID ++ e = engine(4, 2) ++ self.assertEqual(e.slot_gs, 2) ++ hist = list(range(1000, 1700)) ++ e.slot_held[3] = hist + [T, 7, 8] + [50] ++ self.assertEqual(e.pick_slot(hist + [T, 5, 6] + list(range(3000, 3050))), 0) ++ ++ + if __name__ == "__main__": + unittest.main() +diff --git a/serve/web/app.js b/serve/web/app.js +index d3ad404f..4bc2a46f 100644 +--- a/serve/web/app.js ++++ b/serve/web/app.js +@@ -560,8 +560,18 @@ function markdown(text) { + } + + // ------------------------------------------------------------------ Chat +-const DEFAULTS = {thinking: "high", temperature: 0.6, top_p: 0.95, top_k: 20, max: "", seed: "", show: true, esp: true, mcp: true}; ++// The sampling shown until the user changes it is the model card's thinking set. Untouched, the page sends NO sampling ++// fields: the server picks the card's set for the request's thinking state (thinking 1.0 / 0.95 / 20, not thinking ++// 0.7 / 0.8 / 20 with a presence penalty). Before 0.1.45 the page sent 0.6 itself; a stored 0.6 with the other two ++// values untouched counts as "not set" too. ++const DEFAULTS = {thinking: "high", temperature: 1.0, top_p: 0.95, top_k: 20, max: "", seed: "", show: true, esp: true, mcp: true}; ++const OLD_DEFAULT_TEMPERATURE = 0.6; ++function samplingUntouched(s) { ++ const t = +s.temperature; ++ return (t === DEFAULTS.temperature || t === OLD_DEFAULT_TEMPERATURE) && +s.top_p === DEFAULTS.top_p && +s.top_k === DEFAULTS.top_k; ++} + let settings = {...DEFAULTS, ...store.get("sampling", {})}; ++if (samplingUntouched(settings)) settings.temperature = DEFAULTS.temperature; + let messages = store.get("chat", []); + let attachments = []; // {name, url} + let busy = null; // {controller, msg} +@@ -790,10 +800,10 @@ async function send() { + + const body = {model: health.model, messages: apiMessages(), stream: true, + reasoning_effort: settings.thinking}; +- if (settings.temperature > 0) { +- Object.assign(body, {temperature: +settings.temperature, top_p: +settings.top_p, top_k: +settings.top_k}); +- } else { ++ if (!(settings.temperature > 0)) { + body.temperature = 0; ++ } else if (!samplingUntouched(settings)) { ++ Object.assign(body, {temperature: +settings.temperature, top_p: +settings.top_p, top_k: +settings.top_k}); + } + if (settings.seed) body.seed = +settings.seed; + if (settings.max) body.max_tokens = +settings.max; +@@ -1011,8 +1021,9 @@ async function loadShared() { + $("s-share").setAttribute("aria-checked", String(sharedOn)); + } + function sharedDefaults(s) { +- const d = {reasoning_effort: s.thinking, temperature: +s.temperature}; +- if (+s.temperature > 0) Object.assign(d, {top_p: +s.top_p, top_k: +s.top_k}); ++ const d = {reasoning_effort: s.thinking}; ++ if (!(+s.temperature > 0)) d.temperature = 0; ++ else if (!samplingUntouched(s)) Object.assign(d, {temperature: +s.temperature, top_p: +s.top_p, top_k: +s.top_k}); + if (s.seed) d.seed = +s.seed; + if (s.max) d.max_tokens = +s.max; + if (projectionLoaded()) d.experimental_speed_projection = s.esp !== false; +diff --git a/src/kernels/ngram.cpp b/src/kernels/ngram.cpp +index 91de3705..1fd7a3a5 100644 +--- a/src/kernels/ngram.cpp ++++ b/src/kernels/ngram.cpp +@@ -508,27 +508,31 @@ bool PleTable::gather_batch(const uint32_t* rows, size_t n_tokens, float* out, s + if (impl_->pending) { err = "PleTable::gather_batch while a token is in flight"; return false; } + const size_t n = n_tokens * (size_t) PLE_N_HEADS; + if (impl_->mode == PleIo::Direct) { +- // Fast path: if all requested tokens match our in-flight/completed prefetch slots, collect them directly ++ // Fast path: match each requested token to any in-flight/completed prefetch slot, then collect directly + // without issuing a second PleReader ticket or copying through RowCache. +- if (impl_->n_prefetch >= n_tokens && n_tokens <= Impl::kMaxPrefetch) { ++ if (n_tokens <= Impl::kMaxPrefetch) { ++ size_t matched[Impl::kMaxPrefetch]; + bool exact_match = true; + for (size_t t = 0; t < n_tokens; ++t) { +- if (std::memcmp(impl_->prefetch_keys[t], rows + t * PLE_N_HEADS, sizeof(uint32_t) * PLE_N_HEADS) != 0) { +- exact_match = false; +- break; +- } ++ size_t i = 0; ++ while (i < impl_->n_prefetch && ++ std::memcmp(impl_->prefetch_keys[i], rows + t * PLE_N_HEADS, ++ sizeof(uint32_t) * PLE_N_HEADS) != 0) ++i; ++ if (i == impl_->n_prefetch) { exact_match = false; break; } ++ matched[t] = i; // duplicate requests may share one immutable prefetch slot + } + if (exact_match) { ++ bool collect_ok = true; + for (size_t i = 0; i < impl_->n_prefetch; ++i) { +- if (!impl_->reader.collect(impl_->prefetch_tickets[i], err)) { +- impl_->n_prefetch = 0; +- return false; +- } ++ std::string ignored; ++ if (!impl_->reader.collect(impl_->prefetch_tickets[i], collect_ok ? err : ignored)) ++ collect_ok = false; // retain the first error, but drain every output buffer + } + impl_->n_prefetch = 0; ++ if (!collect_ok) return false; + for (size_t t = 0; t < n_tokens; ++t) { + for (int h = 0; h < PLE_N_HEADS; ++h) { +- impl_->decode(impl_->prefetch_raw[t] + (size_t) h * impl_->rb, ++ impl_->decode(impl_->prefetch_raw[matched[t]] + (size_t) h * impl_->rb, + out + (t * PLE_N_HEADS + (size_t) h) * PLE_HEAD_DIM); + } + } +diff --git a/src/ngram/ple_reader_test.cpp b/src/ngram/ple_reader_test.cpp +index 4663455f..670b29cc 100644 +--- a/src/ngram/ple_reader_test.cpp ++++ b/src/ngram/ple_reader_test.cpp +@@ -242,9 +242,89 @@ void bad_sizes_are_refused(const std::string& dir) { + } + } + ++// #2622 (reported by @hotschmoe): gather_batch matched the prefetched tokens by position, so a prefetch of A..F ++// gathered as D,E,F,A,B,C (or A gathered as A,A,A) issued a second ticket for rows that were already in flight. ++// Every case: the bytes equal the mapped reader's, and the requested-row counter says how many rows were re-read. ++void prefetched_keys_are_reused(const std::string& dir) { ++ const k::PleFormatInfo* fmt = nullptr; ++ for (int i = 0; i < k::ple_format_count(); ++i) ++ if (std::string(k::ple_formats()[i].name) == "IQ4_NL") fmt = &k::ple_formats()[i]; ++ if (fmt == nullptr) { CHECK(false, "no IQ4_NL format"); return; } ++ const std::string path = write_format_table(dir, *fmt, 4096); ++ if (path.empty()) { CHECK(false, "cannot write the prefetch table"); return; } ++ struct Case { const char* name; std::vector pre, want; uint64_t extra_rows; }; ++ std::vector seventeen; ++ for (int i = 0; i < 17; ++i) seventeen.push_back(i); ++ const std::vector cases = { ++ {"ordered", {0, 1}, {0, 1}, 0}, ++ {"permutation", {0, 1, 2, 3, 4, 5}, {3, 4, 5, 0, 1, 2}, 0}, ++ {"subset", {0, 1, 2, 3, 4, 5}, {4, 1}, 0}, ++ {"duplicates", {0}, {0, 0, 0}, 0}, ++ {"duplicate prefetch", {0, 0, 1}, {0, 1}, 0}, ++ {"missing key", {0, 1}, {0, 2}, 32}, ++ {"extra prefetch", {0, 1, 2}, {1}, 0}, ++ {"none", {0}, {}, 0}, ++ {"over the window", {0, 1}, seventeen, 17u * k::PLE_N_HEADS}, ++ {"window overflow", seventeen, {16}, 0}, ++ }; ++ const auto keys = [](const std::vector& ids) { ++ std::vector r; ++ for (int id : ids) for (int h = 0; h < k::PLE_N_HEADS; ++h) r.push_back((uint32_t) (id * 32 + h)); ++ return r; ++ }; ++ const auto requests = [](const k::PleTable& t) { ++ unsigned long long n = 0; ++ const std::string s = t.io_report(); ++ return std::sscanf(s.c_str(), "ple io: %llu rows", &n) == 1 ? (uint64_t) n : ~0ull; ++ }; ++ const size_t tok = (size_t) k::PLE_N_HEADS * k::PLE_HEAD_DIM; ++ std::vector> refs(cases.size()); ++ { ++ k::PleTable mapped; ++ k::PleIoOptions mio; ++ mio.mode = k::PleIo::Mmap; ++ std::string err; ++ if (!mapped.open(path, err, mio)) { CHECK(false, "mapped open: %s", err.c_str()); return; } ++ for (size_t c = 0; c < cases.size(); ++c) { ++ const auto r = keys(cases[c].want); ++ refs[c].resize(cases[c].want.size() * tok); ++ CHECK(mapped.gather_batch(r.data(), cases[c].want.size(), refs[c].data(), err), "mapped gather: %s", err.c_str()); ++ } ++ } ++ for (uint64_t cache : {uint64_t(0), uint64_t(256)}) for (bool worker : {false, true}) { ++ for (size_t c = 0; c < cases.size(); ++c) { ++ k::PleTable t; ++ k::PleIoOptions io; ++ io.mode = k::PleIo::Direct; ++ io.cache_rows = cache; ++ io.io_thread = worker; ++ io.max_inflight = 8; ++ io.keepalive_ms = 0; ++ std::string err; ++ if (!t.open(path, err, io)) { CHECK(false, "direct open: %s", err.c_str()); return; } ++ const auto pr = keys(cases[c].pre), wr = keys(cases[c].want); ++ for (size_t i = 0; i < cases[c].pre.size(); ++i) t.prefetch_rows(pr.data() + i * k::PLE_N_HEADS); ++ const uint64_t rq = requests(t), by = t.bytes_read(); ++ const float sentinel = -12345.5f; ++ std::vector out(refs[c].size() + 2, sentinel); ++ const bool ok = t.gather_batch(wr.data(), cases[c].want.size(), out.data() + 1, err); ++ CHECK(ok, "%s (cache %llu worker %d): %s", cases[c].name, (unsigned long long) cache, (int) worker, err.c_str()); ++ CHECK(requests(t) - rq == cases[c].extra_rows, "%s (cache %llu worker %d): %llu rows re-requested, want %llu", ++ cases[c].name, (unsigned long long) cache, (int) worker, (unsigned long long) (requests(t) - rq), ++ (unsigned long long) cases[c].extra_rows); ++ CHECK(t.bytes_read() - by == cases[c].want.size() * k::PLE_N_HEADS * fmt->row_bytes, "%s: bytes_read", cases[c].name); ++ CHECK(refs[c].empty() || std::memcmp(out.data() + 1, refs[c].data(), refs[c].size() * sizeof(float)) == 0, ++ "%s (cache %llu worker %d): bytes differ from the mapped reader", cases[c].name, (unsigned long long) cache, (int) worker); ++ CHECK(out.front() == sentinel && out.back() == sentinel, "%s: wrote outside the output", cases[c].name); ++ } ++ } ++ if (g_fail == 0) std::printf("ple prefetched-key reuse (reordered, duplicated, subset), Direct vs Mmap bytes: OK\n"); ++} ++ + void all_formats_round_trip(const std::string& dir) { + bf16_widening_is_exact(); // BF16's own claim: a shift, exact for every pattern + bad_sizes_are_refused(dir); ++ prefetched_keys_are_reused(dir); + CHECK(k::ple_format_count() > 0, "no formats"); + for (int i = 0; i < k::ple_format_count(); ++i) { + const k::PleFormatInfo& f = k::ple_formats()[i]; +diff --git a/tools/strata_tokenizer.py b/tools/strata_tokenizer.py +index c9e84ee7..42a27fed 100644 +--- a/tools/strata_tokenizer.py ++++ b/tools/strata_tokenizer.py +@@ -15,9 +15,11 @@ emoji (4-byte), combining marks, whitespace runs, and C0 control bytes. + """ + from __future__ import annotations + ++import bisect + import json + import pathlib + import sys ++import threading + + import regex + +@@ -115,6 +117,8 @@ class Tokenizer: + # through __dict__.setdefault on each encoded segment or decoded token. + self._piece_ids: dict[str, list[int]] = {} + self._bytes_cache: dict[int, bytes] = {} ++ # How far a special-token match can look ahead of where it starts (PromptEncoder's margin). ++ self.max_special_len = max((len(t) for t in self.special_tokens), default=1) + + @staticmethod + def _alt(literals: list[str]): +@@ -234,13 +238,16 @@ class Tokenizer: + out.extend(got) + return out + +- def _encode_matching(self, text: str, pat, plain=()) -> list[int]: ++ def _encode_matching(self, text: str, pat, plain=(), marks: list | None = None) -> list[int]: + """Encode `text`, emitting any literal `pat` matches as single tokens and BPE-ing the rest. + + The split happens on the RAW text, before the byte mapping, because a special token's string is a + literal to match rather than bytes to decompose. Everything between the matches is tokenized + normally - which is why a near-miss like `<|im_star` still costs ordinary tokens. A match that starts + inside one of the `plain` (start, end) spans is left to the text around it (#537). ++ ++ `marks`, when given, receives (end of the match in `text`, ids so far) for every match: the points where ++ the encoding of a longer text with the same beginning can resume (PromptEncoder). + """ + if pat is None: + return self._encode_plain(text) +@@ -269,6 +276,8 @@ class Tokenizer: + out.extend(self._encode_plain(text[pos:start])) + out.append(self.special_tokens[m.group(0)]) + pos = m.end() ++ if marks is not None: ++ marks.append((pos, len(out))) + if pos < len(text): + out.extend(self._encode_plain(text[pos:])) + return out +@@ -283,6 +292,11 @@ class Tokenizer: + """ + return self._encode_matching(text, self._special_re if parse_special else self._always_re, plain) + ++ def encode_marked(self, text: str, parse_special: bool = False, plain=()) -> tuple[list[int], list[tuple[int, int]]]: ++ """encode() and its resume points: (end offset in `text`, ids so far) after every special-token match.""" ++ marks: list[tuple[int, int]] = [] ++ return self._encode_matching(text, self._special_re if parse_special else self._always_re, plain, marks), marks ++ + def token_bytes(self, i: int) -> bytes: + """The raw bytes of one token (a multi-byte character can be split across tokens).""" + cache = self._bytes_cache +@@ -303,6 +317,111 @@ class Tokenizer: + return b"".join(self.token_bytes(i) for i in ids).decode("utf-8", errors=errors) + + ++# ------------------------------------------------------------------ incremental prompts (#567) ++def common_prefix_len(a: str, b: str) -> int: ++ """Length of the longest common prefix, by slices (C speed) rather than a loop over characters.""" ++ n = min(len(a), len(b)) ++ lo, step = 0, 4096 ++ while lo < n: # whole chunks first, doubling ++ hi = min(n, lo + step) ++ if a[lo:hi] != b[lo:hi]: ++ break ++ lo, step = hi, min(step * 2, 1 << 20) ++ else: ++ return n ++ while hi - lo > 1: # the first difference is in [lo, hi): bisect it ++ mid = (lo + hi) // 2 ++ if a[lo:mid] == b[lo:mid]: ++ lo = mid ++ else: ++ hi = mid ++ return lo ++ ++ ++def plain_agree_until(a, b, limit: int) -> int: ++ """The first position below `limit` where the `plain` spans (#537) of two texts differ - a position inside a ++ span in one and not in the other - or `limit`. Spans are (start, end), sorted and disjoint.""" ++ def clip(spans): ++ return [(x, min(y, limit)) for x, y in spans if x < limit and y > x] ++ ca, cb = clip(a), clip(b) ++ for (xa, ya), (xb, yb) in zip(ca, cb): ++ if (xa, ya) != (xb, yb): ++ return min(xa, xb) if xa != xb else min(ya, yb) ++ if len(ca) != len(cb): ++ return (ca if len(ca) > len(cb) else cb)[min(len(ca), len(cb))][0] ++ return limit ++ ++ ++class PromptEncoder: ++ """`tok.encode(text, parse_special=True, plain=...)` for chat prompts, reusing the ids of an earlier prompt. ++ ++ A chat client sends the whole conversation every turn, and the rendered prompt of turn n+1 starts with most of ++ turn n's. Running BPE over all of it again costs hundreds of milliseconds at 100K tokens, even when the engine ++ reuses the whole prefix. So the ids up to a SPECIAL-TOKEN BOUNDARY both prompts share are taken as they were, ++ and only the rest is encoded. ++ ++ WHY THAT IS EXACT. Encoding splits the text at special-token matches and BPE-encodes each stretch between two ++ matches on its own (`_encode_matching`): nothing crosses a match. Whether a match is taken at position q ++ depends only on text[q : q + max_special_len] and on whether q lies inside a `plain` span (#537: a ++ quoted in a message is text). So if the two texts agree up to L and their plain spans agree below L, every ++ decision the scan makes at a position before c = the end of some match, with c + max_special_len - 1 <= L, is ++ the same for both - the same matches, the same stretches, the same ids up to c - and from c the scan starts ++ afresh, exactly as it does on text[c:] with the spans shifted by c. Hence encode(new) == ids_old[:ids at c] + ++ encode(new[c:]). The margin matters only for a vocabulary where a literal overlaps another's end; it costs ++ re-encoding one short stretch. ++ ++ A few recent prompts are kept (one per conversation: a prompt replaces the one it continues, that is, one whose ++ text it reuses for more than half; a prompt that shares only the template's head with every entry is a new ++ entry), so a second client does not evict the first. `encode(..., store=False)` reads the kept prompts without ++ writing one (a token count is not a conversation's turn). The tokenizer needs `encode_marked` and ++ `max_special_len`. ++ """ ++ ++ def __init__(self, tok, keep: int = 4): ++ self.tok, self.keep = tok, keep ++ # (text, plain spans, ids, mark ends, mark counts) ++ self.entries: list[tuple[str, tuple, list[int], list[int], list[int]]] = [] ++ self.lock = threading.Lock() ++ self.last_reused = 0 # characters taken from an earlier prompt (for tests) ++ ++ def encode(self, text: str, plain=(), store: bool = True) -> list[int]: ++ plain = tuple(sorted(plain)) ++ with self.lock: ++ entries = self.entries # replaced, never changed in place ++ margin = self.tok.max_special_len - 1 ++ src, cut_k = None, -1 ++ for e in entries: ++ same = plain_agree_until(e[1], plain, common_prefix_len(e[0], text)) ++ k = bisect.bisect_right(e[3], same - margin) - 1 # the last boundary c with c <= same - margin ++ if k >= 0 and (src is None or e[3][k] > src[3][cut_k]): ++ src, cut_k = e, k ++ if src is None: ++ ids, marks = self.tok.encode_marked(text, parse_special=True, plain=plain) ++ ends, counts = [m[0] for m in marks], [m[1] for m in marks] ++ self.last_reused = 0 ++ else: ++ c, n = src[3][cut_k], src[4][cut_k] ++ rest = tuple((max(x - c, 0), y - c) for x, y in plain if y > c) ++ tail, marks = self.tok.encode_marked(text[c:], parse_special=True, plain=rest) ++ ids = src[2][:n] + tail ++ ends = src[3][:cut_k + 1] + [c + m[0] for m in marks] ++ counts = src[4][:cut_k + 1] + [n + m[1] for m in marks] ++ self.last_reused = c ++ if not store: ++ return ids ++ # every rendered prompt starts with the template's header, so sharing a special boundary does not make two ++ # prompts one conversation: only the entry whose text is mostly reused is replaced ++ old = src if src is not None and 2 * src[3][cut_k] > len(src[0]) else None ++ with self.lock: ++ self.entries = [(text, plain, ids, ends, counts)] + [e for e in self.entries if e is not old][:self.keep - 1] ++ return list(ids) ++ ++ def forget(self, ids: list[int]) -> None: ++ """The prompt with these ids was refused: it is not kept.""" ++ with self.lock: ++ self.entries = [e for e in self.entries if e[2] != ids] ++ ++ + # ------------------------------------------------------------------ the pack's tokenizer/ directory + def extract(gguf_path, out_dir) -> dict: + """Write the tokenizer into `/tokenizer/` so the engine never opens the weight shards for it.""" +diff --git a/tools/test_strata_tokenizer.py b/tools/test_strata_tokenizer.py +index 17c3f698..e3cb498c 100644 +--- a/tools/test_strata_tokenizer.py ++++ b/tools/test_strata_tokenizer.py +@@ -1,9 +1,18 @@ +-"""The tokenizer's piece cache (0.1.41): the ids are the ones the merge loop gives, cache cold, warm or switched off. ++"""Tests for tools/strata_tokenizer.py: llama.cpp's reference vectors, the piece cache (0.1.41: the ids are the ones the ++merge loop gives, cache cold, warm or switched off), and the incremental prompt encoder, which must give exactly the ++ids of a plain encode. + +- python -m unittest tools.test_strata_tokenizer ++ python -m unittest tools.test_strata_tokenizer -v ++ ++The real vocabulary comes from $STRATA_TOKENIZER (a pack's tokenizer/ directory) when set, else from llama.cpp's ++`models/ggml-vocab-qwen35.gguf` next to the gguf-py the tools use ($STRATA_GGUF_PY, third_party/llama.cpp); the ++tests that need it are skipped without one. The chat template's USER_DEFINED tokens (``, ``, ++...) are added to the llama.cpp vocabulary, which does not have them, as the pack's has. + """ + from __future__ import annotations + ++import json ++import os + import random + import sys + import unittest +@@ -11,8 +20,295 @@ from pathlib import Path + + ROOT = Path(__file__).resolve().parents[1] + sys.path.insert(0, str(ROOT / "tools")) ++sys.path.insert(0, str(ROOT)) + import strata_tokenizer as ST # noqa: E402 + ++ADDED = ["", "", "", "", "", "", ++ "<|vision_start|>", "<|image_pad|>", "<|vision_end|>"] ++ ++ ++def vocab_gguf() -> Path | None: ++ cands = [os.environ.get("STRATA_GGUF_PY"), ROOT / "third_party" / "llama.cpp" / "gguf-py", "/opt/llama.cpp/gguf-py"] ++ for c in cands: ++ if c and (Path(c).parent / "models" / "ggml-vocab-qwen35.gguf").is_file(): ++ return Path(c).parent / "models" / "ggml-vocab-qwen35.gguf" ++ return None ++ ++ ++def load_tokenizer(): ++ pack = os.environ.get("STRATA_TOKENIZER") ++ if pack and (Path(pack) / "vocab.json").is_file(): ++ t = Path(pack) ++ vocab = json.loads((t / "vocab.json").read_text(encoding="utf-8")) ++ tokens = [None] * len(vocab) ++ for s, i in vocab.items(): ++ tokens[i] = s ++ return ST.Tokenizer(tokens, (t / "merges.txt").read_text(encoding="utf-8").split("\n"), ++ json.loads((t / "token_type.json").read_text())) ++ path = vocab_gguf() ++ return None if path is None else gguf_tokenizer(path) ++ ++ ++def gguf_tokenizer(path): ++ """llama.cpp's qwen35 vocabulary with the template's added tokens: the one its test vectors were made with.""" ++ from gguf_reader import GGUFFile ++ md = GGUFFile(path).metadata ++ tokens, types = list(md["tokenizer.ggml.tokens"]), list(md["tokenizer.ggml.token_type"]) ++ for s in ADDED: ++ if s not in tokens: ++ tokens.append(s) ++ types.append(4) ++ return ST.Tokenizer(tokens, list(md["tokenizer.ggml.merges"]), types) ++ ++ ++def byte_tokenizer(specials: list[str]) -> ST.Tokenizer: ++ """A vocabulary of the 256 byte tokens and `specials` (CONTROL), no merges: small enough to build per test.""" ++ tokens = [ST.BYTE_TO_UNICODE[b] for b in range(256)] + specials ++ return ST.Tokenizer(tokens, [], [1] * 256 + [3] * len(specials)) ++ ++ ++WORDS = ["the", "tokenizer", " ", " ", "\n", "\n\n", "\t", "def", "f(x):", "return", "x", "==", "->", "café", ++ "你好", "\U0001f600", "é", "1234", "don't", "I'm", "<|im_", "", ++ " ", "\r\n", "Hello,", "world!", "<", ">", "|", " ", "x" * 70] ++ ++ ++def random_text(rnd: random.Random, n: int, specials=()) -> str: ++ parts = [] ++ for _ in range(n): ++ parts.append(rnd.choice(WORDS) if not specials or rnd.random() > 0.05 else rnd.choice(specials)) ++ if rnd.random() < 0.5: ++ parts.append(" ") ++ return "".join(parts) ++ ++ ++def conversation_prompts(template, rnd: random.Random) -> list[tuple[str, str]]: ++ """Rendered prompts, in the order a server would see them, from clients that grow a conversation turn after turn, ++ use tools, send images, change the start of the prompt, edit an earlier turn, send a shorter conversation, and ++ interleave with each other: (what happens, prompt).""" ++ out: list[tuple[str, str]] = [] ++ specials = ["<|im_start|>", "<|im_end|>", "", "", "", ""] ++ weather = [{"type": "function", "function": {"name": "get_weather", "description": "Get the weather", ++ "parameters": {"type": "object", "properties": { ++ "city": {"type": "string"}}, "required": ["city"]}}}] ++ ++ def add(what, msgs, tools=None, **kw): ++ out.append((what, template.render(msgs, tools=tools, **kw))) ++ ++ # several turns, with unicode, a literal special token inside the text, and reasoning kept or not ++ msgs = [{"role": "system", "content": "Be brief. Réponds en français si on te parle en français."}] ++ for turn in range(6): ++ msgs.append({"role": "user", "content": random_text(rnd, 60, specials) + " é你\U0001f600 " * turn}) ++ add(f"turn {turn}", msgs) ++ answer = {"role": "assistant", "content": random_text(rnd, 30)} ++ if turn % 2: ++ answer["reasoning_content"] = random_text(rnd, 30) ++ msgs.append(answer) ++ grown = list(msgs) ++ ++ # tools: definitions, a call, its response, then a following turn ++ tmsgs = [{"role": "user", "content": "Weather in Oslo?"}] ++ add("tools 1", tmsgs, weather) ++ tmsgs += [{"role": "assistant", "content": "", "tool_calls": [{"type": "function", "function": { ++ "name": "get_weather", "arguments": {"city": "Oslo"}}}]}, ++ {"role": "tool", "content": "{\"temp\": 3}", "name": "get_weather"}] ++ add("tools 2", tmsgs, weather) ++ tmsgs += [{"role": "assistant", "content": "3 degrees."}, {"role": "user", "content": "And in Bergen? " * 20}] ++ add("tools 3", tmsgs, weather) ++ ++ # images: one per user turn, the placeholder tokens are specials ++ imsgs = [] ++ for turn in range(3): ++ imsgs.append({"role": "user", "content": [{"type": "image", "source": f"img{turn}"}, ++ {"type": "text", "text": f"What is on picture {turn}?"}]}) ++ add(f"images {turn}", imsgs) ++ imsgs.append({"role": "assistant", "content": "A cat."}) ++ ++ # the start of the prompt changes (a thinking level, a system message), then the old conversation grows again ++ add("effort changed", grown, reasoning_effort="low") ++ add("system changed", [dict(grown[0], content=grown[0]["content"] + " Be kind.")] + grown[1:]) ++ grown.append({"role": "user", "content": "one more question"}) ++ add("grown again", grown) ++ ++ # an earlier turn edited, then a conversation that got shorter, then the full one again ++ edited = list(grown) ++ edited[3] = dict(edited[3], content=edited[3]["content"] + " (edited)") ++ add("edited", edited) ++ add("shorter", grown[:5]) ++ add("empty user", [{"role": "user", "content": ""}]) ++ add("full again", grown) ++ return out ++ ++ ++class Oracle(unittest.TestCase): ++ """llama.cpp's own tokenizer test vectors for qwen35.""" ++ ++ def test_llama_cpp_vectors(self): ++ path = vocab_gguf() ++ if path is None or not path.with_suffix(".gguf.inp").is_file(): ++ self.skipTest("llama.cpp's ggml-vocab-qwen35.gguf test vectors are not here") ++ tk = gguf_tokenizer(path) # not $STRATA_TOKENIZER: the vectors belong to this vocabulary ++ inp = path.with_suffix(".gguf.inp").read_text(encoding="utf-8").split("\n__ggml_vocab_test__\n") ++ out = path.with_suffix(".gguf.out").read_text(encoding="utf-8").split("\n") ++ n = 0 ++ for text, want in zip(inp, out): ++ if text.endswith("\n__ggml_vocab_test__"): ++ text = text[:-len("\n__ggml_vocab_test__")] ++ self.assertEqual(tk.encode(text), [int(x) for x in want.split()], repr(text[:40])) ++ n += 1 ++ self.assertGreater(n, 40) ++ ++ ++class Marks(unittest.TestCase): ++ @classmethod ++ def setUpClass(cls): ++ cls.tk = load_tokenizer() ++ if cls.tk is None: ++ raise unittest.SkipTest("no tokenizer (set STRATA_TOKENIZER or STRATA_GGUF_PY)") ++ ++ def test_marks(self): ++ text = "a<|im_start|>user\nhi<|im_end|>\n\n" ++ ids, marks = self.tk.encode_marked(text, parse_special=True) ++ self.assertEqual(ids, self.tk.encode(text, parse_special=True)) ++ self.assertEqual([m[0] for m in marks], [text.index("user"), text.index("<|im_end|>") + len("<|im_end|>"), ++ len(text) - 1]) ++ for end, n in marks: # each mark: the ids of the text up to it ++ self.assertEqual(ids[:n], self.tk.encode(text[:end], parse_special=True)) ++ ++ ++class Prompts(unittest.TestCase): ++ """PromptEncoder: the ids of a plain encode, whatever the earlier prompts were.""" ++ ++ @classmethod ++ def setUpClass(cls): ++ cls.tk = load_tokenizer() ++ if cls.tk is None: ++ raise unittest.SkipTest("no tokenizer (set STRATA_TOKENIZER or STRATA_GGUF_PY)") ++ from serve.frontend import ChatTemplate ++ cls.template = ChatTemplate(ROOT / "serve" / "chat_template.jinja") ++ ++ def check(self, enc, text): ++ self.assertEqual(enc.encode(text), self.tk.encode(text, parse_special=True)) ++ ++ def test_chat_golden(self): ++ cases = json.loads((ROOT / "serve" / "chat_golden.json").read_text(encoding="utf-8")) ++ enc = ST.PromptEncoder(self.tk) ++ for c in cases + cases: # the second round finds its own earlier prompts ++ if c.get("rendered"): ++ with self.subTest(case=c["name"]): ++ self.check(enc, c["rendered"]) ++ ++ def test_conversations(self): ++ for seed in range(3): ++ prompts = conversation_prompts(self.template, random.Random(seed)) ++ wants = [self.tk.encode(prompt, parse_special=True) for _, prompt in prompts] ++ for keep in (1, 2, 4): ++ enc = ST.PromptEncoder(self.tk, keep=keep) ++ for (what, prompt), want in zip(prompts, wants): ++ with self.subTest(seed=seed, keep=keep, what=what): ++ self.assertEqual(enc.encode(prompt), want) ++ ++ def test_growing_turn_reuses_the_previous_prompt(self): ++ enc = ST.PromptEncoder(self.tk) ++ msgs = [{"role": "user", "content": "first question " * 50}] ++ prev = "" ++ for turn in range(4): ++ prompt = self.template.render(msgs) ++ self.check(enc, prompt) ++ if turn: ++ self.assertGreater(enc.last_reused, len(prev) - 100) ++ prev = prompt ++ msgs += [{"role": "assistant", "content": f"answer {turn}"}, {"role": "user", "content": f"question {turn}"}] ++ ++ def test_prefix_of_an_earlier_prompt_and_unrelated_prompts(self): ++ enc = ST.PromptEncoder(self.tk, keep=2) ++ a = self.template.render([{"role": "user", "content": "first question " * 50}]) ++ b = self.template.render([{"role": "user", "content": "another conversation"}]) ++ c = self.template.render([{"role": "user", "content": "a third one"}]) ++ for text in (a, a[:len(a) // 2], b, c, a, a + "tail", "", "<|im_end|>", "<|im_end|><|im_end|>x"): ++ with self.subTest(text=text[:30]): ++ self.check(enc, text) ++ ++ def test_a_read_only_encode_gives_the_ids_and_stores_nothing(self): ++ enc = ST.PromptEncoder(self.tk) ++ a = self.template.render([{"role": "user", "content": "first question " * 50}]) ++ enc.encode(a) ++ kept = list(enc.entries) ++ for text in (a + "<|im_start|>user\nand a follow-up<|im_end|>\n", a[:len(a) // 2], "<|im_end|>x", ++ self.template.render([{"role": "user", "content": "another conversation"}])): ++ with self.subTest(text=text[:30]): ++ self.assertEqual(enc.encode(text, store=False), self.tk.encode(text, parse_special=True)) ++ self.assertTrue(enc.entries == kept, "a read-only encode wrote an entry") ++ self.assertEqual(enc.encode(a + "tail", store=False), self.tk.encode(a + "tail", parse_special=True)) ++ self.assertGreater(enc.last_reused, len(a) // 2) ++ ++ def test_the_returned_ids_are_the_callers(self): ++ enc = ST.PromptEncoder(self.tk) ++ text = "<|im_start|>user\nhi<|im_end|>\n" + "and a tail longer than the look-ahead margin. " * 3 ++ want = self.tk.encode(text, parse_special=True) ++ got = enc.encode(text) ++ got[0] = -1 # inside the part the next call takes from this one ++ got.append(-1) ++ self.assertEqual(enc.encode(text), want) ++ self.assertGreater(enc.last_reused, 0) ++ ++ ++class Margin(unittest.TestCase): ++ """A special literal that is a prefix of a longer one: cutting at the shorter one's end without the look-ahead ++ margin would give different ids.""" ++ ++ def test_longer_literal_past_the_common_prefix(self): ++ tk = byte_tokenizer(["<|a|>", "<|a|>zz"]) ++ enc = ST.PromptEncoder(tk) ++ first = "x<|a|>y" + "<|a|>" * 3 + "q" ++ self.assertEqual(enc.encode(first), tk.encode(first, parse_special=True)) ++ second = "x<|a|>y" + "<|a|>" * 3 + "zz" # shares everything up to the last <|a|> ++ got = enc.encode(second) ++ self.assertEqual(got, tk.encode(second, parse_special=True)) ++ self.assertEqual(got[-1], tk.special_tokens["<|a|>zz"]) ++ self.assertGreater(enc.last_reused, 0) # the earlier boundaries were still reused ++ ++ def test_literal_ending_inside_another(self): ++ tk = byte_tokenizer(["<|a|>", "b<|a|>c"]) ++ enc = ST.PromptEncoder(tk) ++ for text in ("zzb<|a|>d", "zzb<|a|>c", "zzb<|a|>cz<|a|>", "zzb<|a|>cz<|a|>b<|a|>c"): ++ with self.subTest(text=text): ++ self.assertEqual(enc.encode(text), tk.encode(text, parse_special=True)) ++ ++ def test_random_specials(self): ++ """Literals that extend or end inside one another, texts that change right after one of them.""" ++ rnd = random.Random(3) ++ ++ def word(lo, hi): ++ return "".join(rnd.choice("<|ab>") for _ in range(rnd.randint(lo, hi))) ++ for trial in range(60): ++ a, b = word(2, 4), word(2, 4) ++ lits = sorted({a, a + word(1, 3), b, word(1, 2) + b + word(0, 2)}) ++ tk = byte_tokenizer(lits) ++ enc = ST.PromptEncoder(tk, keep=3) ++ pieces = lits + [x[:rnd.randint(1, len(x))] for x in lits] + ["x", "y", "a", "|"] ++ base = "".join(rnd.choice(pieces) for _ in range(60)) ++ for _ in range(30): ++ cut = rnd.randint(0, len(base)) ++ text = base[:cut] + "".join(rnd.choice(pieces) for _ in range(rnd.randint(0, 6))) ++ self.assertEqual(enc.encode(text), tk.encode(text, parse_special=True), (lits, text)) ++ if rnd.random() < 0.3: ++ base = text + base[cut:] ++ ++ ++class CommonPrefix(unittest.TestCase): ++ def test_against_a_loop(self): ++ rnd = random.Random(5) ++ for _ in range(500): ++ a = "".join(rnd.choice("ab") for _ in range(rnd.randint(0, 9000))) ++ k = rnd.randint(0, len(a)) ++ b = a[:k] + rnd.choice(["", "c", "cd"]) + a[k:k + rnd.randint(0, 50)] ++ want = 0 ++ while want < min(len(a), len(b)) and a[want] == b[want]: ++ want += 1 ++ self.assertEqual(ST.common_prefix_len(a, b), want) ++ self.assertEqual(ST.common_prefix_len(b, a), want) ++ ++ + + def toy() -> ST.Tokenizer: + base = [ST.BYTE_TO_UNICODE[b] for b in range(256)] +diff --git a/serve/test_admit_order.py b/serve/test_admit_order.py +new file mode 100644 +index 00000000..67766e55 +--- /dev/null ++++ b/serve/test_admit_order.py +@@ -0,0 +1,119 @@ ++"""serve/test_admit_order.py - STRATA_ADMIT_SHORT_FIRST: new requests take the control lines short prompts first ++(StrataEngine._take_control / _admit_next; no process, no GPU). ++ ++ python -m unittest serve.test_admit_order -v ++""" ++from __future__ import annotations ++ ++import sys ++import threading ++import time ++import unittest ++from pathlib import Path ++from unittest import mock ++ ++sys.path.insert(0, str(Path(__file__).resolve().parents[1])) ++import serve.server as server # noqa: E402 ++from serve.server import StrataEngine # noqa: E402 ++ ++ ++def engine_without_process(n=2): ++ """A StrataEngine with the admission state of a started one and no process: alive() is `e.up`.""" ++ e = StrataEngine("missing-executable", [], lazy=True) ++ e.batch = n ++ e.slot_busy = [False] * n ++ e.slot_claims, e.slot_releasing, e._orphan_seen = set(), set(), {} ++ e.slot_cv = threading.Condition() ++ e.waiting, e.admitting, e.wait_lens, e.ctl_epoch = 0, 0, [], 0 ++ e.ctl = threading.Lock() ++ e.gen = 1 ++ e.up = True ++ e.alive = lambda: e.up ++ return e ++ ++ ++class AdmitOrderTest(unittest.TestCase): ++ def setUp(self): ++ self.e = engine_without_process(16) ++ self.order, self.cancels = [], [] ++ self.order_lock = threading.Lock() ++ ++ def tearDown(self): ++ for c in self.cancels: ++ c.set() ++ ++ def take(self, name, plen, **kw): ++ """A request waiting for the control lines; once it has them it notes its name and gives them back.""" ++ c = threading.Event() ++ self.cancels.append(c) ++ ++ def go(): ++ gen = self.e._take_control(c, plen, **kw) ++ try: ++ while True: ++ next(gen) ++ except StopIteration as s: ++ if s.value: ++ with self.order_lock: ++ self.order.append(name) ++ time.sleep(0.05) # its read ++ self.e.ctl.release() ++ threading.Thread(target=go, daemon=True).start() ++ ++ def until(self, cond, s=5.0): ++ end = time.time() + s ++ while time.time() < end and not cond(): ++ time.sleep(0.01) ++ self.assertTrue(cond()) ++ ++ def queue(self, *reqs): ++ """While a long read holds the control lines, the requests arrive in this order; then the read ends.""" ++ self.assertTrue(self.e.ctl.acquire(timeout=1)) ++ for name, plen in reqs: ++ self.take(name, plen, new=True) ++ self.until(lambda n=len(self.order) + sum(1 for _ in self.e.wait_lens): self.e.waiting >= 1) ++ time.sleep(0.02) # (distinct arrival times) ++ self.until(lambda: self.e.waiting == len(reqs)) ++ self.e.ctl.release() ++ self.until(lambda: len(self.order) == len(reqs), s=15.0) ++ return self.order ++ ++ def test_short_first_puts_the_chat_before_the_agent_reads(self): ++ with mock.patch.object(server, "ADMIT_SHORT_FIRST", 8192), mock.patch.object(server, "ADMIT_MAX_WAIT_S", 60.0): ++ order = self.queue(("agent1", 120000), ("agent2", 190000), ("chat", 500), ("agent3", 150000)) ++ self.assertEqual(order[0], "chat") ++ self.assertEqual(order[1:], ["agent1", "agent2", "agent3"]) # arrival order among the long ones ++ ++ def test_a_long_request_that_waited_too_long_goes_first(self): ++ with mock.patch.object(server, "ADMIT_SHORT_FIRST", 8192), mock.patch.object(server, "ADMIT_MAX_WAIT_S", 0.2): ++ self.assertTrue(self.e.ctl.acquire(timeout=1)) ++ self.take("agent", 120000, new=True) ++ self.until(lambda: self.e.waiting == 1) ++ time.sleep(0.4) # past STRATA_ADMIT_MAX_WAIT_S ++ self.take("chat", 500, new=True) ++ self.until(lambda: self.e.waiting == 2) ++ self.e.ctl.release() ++ self.until(lambda: len(self.order) == 2, s=15.0) ++ self.assertEqual(self.order, ["agent", "chat"]) ++ ++ def test_off_by_default(self): ++ self.assertEqual(server.ADMIT_SHORT_FIRST if "STRATA_ADMIT_SHORT_FIRST" not in __import__("os").environ ++ else 0, 0) ++ with mock.patch.object(server, "ADMIT_SHORT_FIRST", 0): ++ order = self.queue(("agent1", 120000), ("chat", 500)) ++ self.assertEqual(sorted(order), ["agent1", "chat"]) # both get their turn, in any order ++ ++ def test_a_continuation_is_not_held_back(self): ++ """A request already admitted once (new=False: promoted from the solo path) is not part of the order.""" ++ with mock.patch.object(server, "ADMIT_SHORT_FIRST", 8192), mock.patch.object(server, "ADMIT_MAX_WAIT_S", 60.0): ++ self.assertTrue(self.e.ctl.acquire(timeout=1)) ++ self.take("chat", 500, new=True) ++ self.take("continuation", 190000) ++ self.until(lambda: self.e.waiting == 2) ++ self.e.ctl.release() ++ self.until(lambda: len(self.order) == 2, s=15.0) ++ self.assertEqual(sorted(self.order), ["chat", "continuation"]) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_model_card_sampling.py b/serve/test_model_card_sampling.py +new file mode 100644 +index 00000000..ee491506 +--- /dev/null ++++ b/serve/test_model_card_sampling.py +@@ -0,0 +1,108 @@ ++"""serve/test_model_card_sampling.py - a request that sends no sampling gets the model card's, not greedy (mock engine). ++ ++Qwen's "Best Practices": thinking mode temperature 1.0 / top_p 0.95 / top_k 20, non-thinking 0.7 / 0.8 / 20 with ++presence_penalty 1.5. The request's own fields win, then the Chat settings shared with apps, then the config's ++`sampling` block, then the card; temperature 0 asked for is still greedy. ++ ++ python -m unittest serve.test_model_card_sampling -v ++""" ++from __future__ import annotations ++ ++import json ++import sys ++import unittest ++import urllib.request ++from pathlib import Path ++ ++sys.path.insert(0, str(Path(__file__).resolve().parents[1])) ++from serve.frontend import ChatTemplate # noqa: E402 ++from serve.server import ByteTokenizer, MockEngine, Service, StrataEngine, model_card_sampling, serve # noqa: E402 ++ ++ROOT = Path(__file__).resolve().parents[1] ++ ++ ++class Recording(MockEngine): ++ def generate(self, ids, max_new, sampling, cancel, embeddings=None): ++ self.sampling = dict(sampling or {}) ++ yield from super().generate(ids, max_new, sampling, cancel, embeddings) ++ ++ ++class ModelCardSampling(unittest.TestCase): ++ @classmethod ++ def setUpClass(cls): ++ cls.tok = ByteTokenizer() ++ cls.engine = Recording(cls.tok, "\nhm\n\n\nok", max_context=8192) ++ cls.svc = Service(cls.engine, cls.tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) ++ cls.httpd = serve(cls.svc, port=0) ++ cls.base = f"http://127.0.0.1:{cls.httpd.server_address[1]}" ++ ++ @classmethod ++ def tearDownClass(cls): ++ cls.httpd.shutdown() ++ cls.httpd.server_close() ++ ++ def setUp(self): ++ self.svc.model_arch = "qwen4exp" ++ self.svc.sampling_defaults = {} ++ self.svc.shared = {} ++ self.svc.config_effort = None ++ ++ def tearDown(self): ++ self.svc.model_arch = None ++ ++ def keys(self, **body): ++ body = {"model": "m", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 8, **body} ++ self.engine.sampling = {} ++ req = urllib.request.Request(self.base + "/v1/chat/completions", data=json.dumps(body).encode(), ++ headers={"Content-Type": "application/json"}) ++ urllib.request.urlopen(req, timeout=10).read() ++ return StrataEngine.sampling_keys(self.engine.sampling) ++ ++ def test_tables(self): ++ self.assertEqual(model_card_sampling("qwen4exp", True), {"temperature": 1.0, "top_p": 0.95, "top_k": 20}) ++ self.assertEqual(model_card_sampling("qwen4exp", False), ++ {"temperature": 0.7, "top_p": 0.8, "top_k": 20, "presence_penalty": 1.5}) ++ self.assertEqual(model_card_sampling("qwen35moe", True)["presence_penalty"], 1.5) ++ self.assertEqual(model_card_sampling(None, True), {}) ++ self.assertEqual(model_card_sampling("unknown", False), {}) ++ ++ def test_thinking_default(self): ++ k = self.keys() ++ self.assertIn("temperature=1.0", k) ++ self.assertIn("top_p=0.95", k) ++ self.assertIn("top_k=20", k) ++ self.assertNotIn("penalty_present", k) ++ ++ def test_non_thinking_default(self): ++ k = self.keys(reasoning_effort="none") ++ self.assertIn("temperature=0.7", k) ++ self.assertIn("top_p=0.8", k) ++ self.assertIn("penalty_present=1.5", k) ++ ++ def test_request_wins_and_zero_is_greedy(self): ++ k = self.keys(temperature=0) ++ self.assertNotIn("temperature", k) ++ self.assertIn("top_k=20", k) ++ k = self.keys(temperature=0.3, top_p=0.5) ++ self.assertIn("temperature=0.3", k) ++ self.assertIn("top_p=0.5", k) ++ ++ def test_config_block_and_shared_settings_win_over_the_card(self): ++ self.svc.sampling_defaults = {"temperature": 0.6, "top_k": 40} ++ k = self.keys() ++ self.assertIn("temperature=0.6", k) ++ self.assertIn("top_k=40", k) ++ self.assertIn("top_p=0.95", k) ++ self.svc.shared = {"temperature": 0.2} ++ self.assertIn("temperature=0.2", self.keys()) ++ self.svc.sampling_defaults = {"temperature": 0} ++ self.svc.shared = {} ++ self.assertNotIn("temperature", self.keys()) # a config that asks for greedy gets it ++ ++ def test_no_family_no_defaults(self): ++ self.svc.model_arch = None ++ self.assertEqual(self.keys(), "") ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_parser_call_scan.py b/serve/test_parser_call_scan.py +new file mode 100644 +index 00000000..a4484514 +--- /dev/null ++++ b/serve/test_parser_call_scan.py +@@ -0,0 +1,61 @@ ++"""Wrapped-call precheck parity against the original structural walk.""" ++import itertools ++import random ++import unittest ++from unittest.mock import patch ++ ++from serve import frontend ++ ++ ++class ReferenceParser(frontend.OutputParser): ++ def _wrapped_call_end(self): ++ return frontend.call_end(self.buf) ++ ++ ++def events_without_ids(events): ++ return [(e.kind, e.text, (e.call.name, e.call.arguments) if e.call else None) for e in events] ++ ++ ++class CallScanTests(unittest.TestCase): ++ def test_chunk_boundaries_literal_tags_multiple_calls_and_recovery(self): ++ def call(value): ++ return "" + value + "" ++ bodies = [call("ordinary"), call("literal text here"), ++ call("literal prose"), call("a") + call("b"), ++ call("unfinished")[:-5], call("incomplete prose")[:-30], ++ "", ++ "bare", ++ "a" ++ "b"] ++ tools = [{"name": "write", "parameters": {"properties": {"text": {"type": "string"}}}}] ++ rng = random.Random(1030) ++ for text, stream, recover in itertools.product(bodies, (False, True), (False, True)): ++ # Every fixed chunk width, plus random partitions; compare after EACH feed. ++ for width in range(1, 16): ++ a = frontend.OutputParser(thinking=False, tools=tools, stream_tools=stream, recover=recover) ++ b = ReferenceParser(thinking=False, tools=tools, stream_tools=stream, recover=recover) ++ pos = 0 ++ while pos < len(text): ++ end = pos + (rng.randint(1, 20) if width == 15 else width) ++ delta = text[pos:end] ++ self.assertEqual(events_without_ids(a.feed(delta)), events_without_ids(b.feed(delta))) ++ pos = end ++ self.assertEqual(events_without_ids(a.finish()), events_without_ids(b.finish())) ++ ++ def test_no_structural_rescans_until_closer_and_reset_for_next_call(self): ++ parser = frontend.OutputParser(thinking=False) ++ with patch.object(frontend, "call_end", wraps=frontend.call_end) as walk: ++ for _ in range(2): ++ parser.feed("") ++ before = walk.call_count ++ for _ in range(100): ++ parser.feed("abcdefgh") ++ self.assertEqual(walk.call_count, before) ++ for char in "": ++ parser.feed(char) ++ self.assertEqual(walk.call_count, before + 1) ++ self.assertEqual(parser._call_end_checked, 0) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_parser_hold.py b/serve/test_parser_hold.py +new file mode 100644 +index 00000000..519f8192 +--- /dev/null ++++ b/serve/test_parser_hold.py +@@ -0,0 +1,60 @@ ++"""Partial-tag parity, including overlapping prefixes and arbitrary chunk boundaries.""" ++import itertools ++import random ++import unittest ++ ++from serve.frontend import OutputParser ++ ++ ++def reference_hold(text, tags): ++ return max((n for tag in tags for n in range(1, len(tag)) if text.endswith(tag[:n])), default=0) ++ ++ ++class ReferenceParser(OutputParser): ++ def _hold(self, text, tags): ++ return reference_hold(text, tags) ++ ++ ++class HoldTests(unittest.TestCase): ++ def test_overlapping_and_empty_tags(self): ++ parser = OutputParser() ++ tags = ("", "a", "aaa", "abab", "baab", "", "") ++ for size in range(8): ++ for chars in itertools.product("ab<", repeat=size): ++ text = "".join(chars) ++ self.assertEqual(parser._hold(text, tags), reference_hold(text, tags), text) ++ ++ def test_every_real_tag_split(self): ++ tags = ("", "", "") ++ parser = OutputParser() ++ for tag in tags: ++ for n in range(len(tag) + 1): ++ for prefix in ("", "prose\n", "<", "<<<", "🙂日本語"): ++ text = prefix + tag[:n] ++ self.assertEqual(parser._hold(text, tags), reference_hold(text, tags), text) ++ ++ def test_streamed_events_match_reference(self): ++ texts = ["Reasoning.\n\nAnswer café 🙂\n", ++ "Quoted `` text and prose.", ++ "\n```xml\n\n```\nAnswer\nabc\nDone", ++ "abcDone"] ++ rng = random.Random(1028) ++ for text, thinking, recover, stream_tools in itertools.product(texts, (False, True), (False, True), (False, True)): ++ for _ in range(5): ++ args = dict(thinking=thinking, tools=[{"name": "lookup"}], recover=recover, stream_tools=stream_tools) ++ actual, expected = OutputParser(**args), ReferenceParser(**args) ++ def normalized(events): ++ # Tool call IDs are deliberately random; compare their semantic contents. ++ return [(e.kind, e.text, (e.call.name, e.call.arguments) if e.call else None) for e in events] ++ pos = 0 ++ while pos < len(text): ++ end = pos + rng.randint(1, 9) ++ delta = text[pos:end] ++ self.assertEqual(normalized(actual.feed(delta)), normalized(expected.feed(delta))) ++ pos = end ++ self.assertEqual(normalized(actual.finish()), normalized(expected.finish())) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_prompt_encoder.py b/serve/test_prompt_encoder.py +new file mode 100644 +index 00000000..36cc06f2 +--- /dev/null ++++ b/serve/test_prompt_encoder.py +@@ -0,0 +1,162 @@ ++"""serve/test_prompt_encoder.py - #567's incremental prompt encoding with #537's plain spans: every prompt's ids ++are those of a full encode, and a turn reuses the previous one's. ++ ++ python -m unittest serve.test_prompt_encoder -v ++""" ++from __future__ import annotations ++ ++import random ++import re ++import sys ++import unittest ++from pathlib import Path ++ ++ROOT = Path(__file__).resolve().parents[1] ++sys.path.insert(0, str(ROOT)) ++sys.path.insert(0, str(ROOT / "tools")) ++from serve.frontend import ChatTemplate # noqa: E402 ++from serve.server import ByteTokenizer, MockEngine, Service # noqa: E402 ++from strata_tokenizer import PromptEncoder, plain_agree_until # noqa: E402 ++ ++ ++class ThinkTokenizer(ByteTokenizer): ++ """ByteTokenizer whose / are specials matched even without parse_special (type 4), as in the ++ Qwen vocabulary: the case #537's plain spans exist for.""" ++ SPECIALS = ByteTokenizer.SPECIALS + ["", ""] ++ ALWAYS = ("", "") ++ ++ ++PIECES = ["<|im_start|>user\n", "<|im_end|>\n", "", "", "hello ", "é你 ", "<|im_sta", "x" * 7, ++ "<|im_start|>assistant\n", "", "<|vision_start|>", "<|image_pad|>"] ++ ++ ++def spans_of(text, rng): ++ """Random plain spans over the specials' texts (, , <|im_start|>, ...), as unmark_think_literals ++ marks the ones quoted inside a message.""" ++ matches = re.finditer("|".join(map(re.escape, ThinkTokenizer.SPECIALS)), text) ++ return tuple((m.start(), m.end()) for m in matches if rng.random() < 0.5) # finditer: left to right, no overlap ++ ++ ++class Units(unittest.TestCase): ++ def test_plain_agree_until(self): ++ self.assertEqual(plain_agree_until((), (), 50), 50) ++ self.assertEqual(plain_agree_until(((5, 13),), ((5, 13),), 50), 50) ++ self.assertEqual(plain_agree_until(((5, 13),), (), 50), 5) ++ self.assertEqual(plain_agree_until(((5, 13),), ((5, 9),), 50), 9) ++ self.assertEqual(plain_agree_until(((60, 70),), (), 50), 50) # beyond the shared text: no matter ++ self.assertEqual(plain_agree_until(((5, 13), (20, 28)), ((5, 13),), 50), 20) ++ ++ def test_same_ids_as_a_full_encode(self): ++ tok = ThinkTokenizer() ++ rng = random.Random(7) ++ for conv in range(300): ++ enc = PromptEncoder(tok) ++ text, plain = "", () ++ for turn in range(8): ++ text += "".join(rng.choice(PIECES) for _ in range(rng.randint(1, 12))) ++ if rng.random() < 0.6: ++ plain = spans_of(text, rng) # the spans may change in the shared part too ++ with self.subTest(conv=conv, turn=turn): ++ self.assertEqual(enc.encode(text, plain), tok.encode(text, True, plain)) ++ if rng.random() < 0.2 and len(text) > 10: # an edit: a client re-rendering its history ++ cut = rng.randrange(len(text)) ++ text = text[:cut] ++ plain = tuple((a, b) for a, b in plain if b <= cut) ++ ++ def test_a_turn_reuses_the_previous_one(self): ++ tok = ThinkTokenizer() ++ enc = PromptEncoder(tok) ++ # as a chat turn: the earlier prompt goes on past its last special (the answer being written), and the ++ # next one shares it up to there - the resume point needs max_special_len characters after it ++ base = "<|im_start|>user\n" + "quoted here " * 200 + "<|im_end|>\n<|im_start|>assistant\nan answer" ++ plain = tuple((i, i + 8) for i in range(len(base)) if base.startswith("", i)) ++ enc.encode(base, plain) ++ more = base + " that goes on<|im_end|>\n<|im_start|>user\nnext<|im_end|>\n" ++ self.assertEqual(enc.encode(more, plain), tok.encode(more, True, plain)) ++ self.assertGreater(enc.last_reused, len(base) // 2) ++ # the same text with the quoted no longer plain: nothing before it may be reused past its start ++ self.assertEqual(enc.encode(more, ()), tok.encode(more, True, ())) ++ self.assertLessEqual(enc.last_reused, plain[0][0]) ++ ++ ++class Served(unittest.TestCase): ++ def test_quoted_think_across_turns(self): ++ tok = ThinkTokenizer() ++ template = ChatTemplate(ROOT / "serve/chat_template.jinja") ++ svc = Service(MockEngine(tok, "\n\nok", max_context=1 << 20), tok, template) ++ full = Service(MockEngine(tok, "\n\nok", max_context=1 << 20), tok, template) ++ full.prompts = None # the reference: a full encode every time ++ self.assertIsNotNone(svc.prompts) ++ msgs = [{"role": "system", "content": "Explain tags."}] ++ for turn in range(5): ++ msgs.append({"role": "user", "content": f"turn {turn}: what does mean? too " * 3}) ++ got, want = svc.encode_prompt(msgs, None, {}), full.encode_prompt(msgs, None, {}) ++ self.assertEqual(got, want) ++ if turn: ++ self.assertGreater(svc.prompts.last_reused, 0) ++ msgs.append({"role": "assistant", "content": "It closes the reasoning: is a marker."}) ++ ++ def service(self, max_context=1 << 22): ++ tok = ByteTokenizer() ++ return Service(MockEngine(tok, "ok", max_context=max_context), tok, ChatTemplate(ROOT / "serve/chat_template.jinja")) ++ ++ def test_another_conversation_does_not_evict_the_first(self): ++ """A subagent beside the main conversation: every rendered prompt starts with the template's header, which ++ must not make the second client's prompt a continuation of the first's entry.""" ++ svc = self.service() ++ turn2 = [{"role": "user", "content": "a long question " * 500}, {"role": "assistant", "content": "answer"}, ++ {"role": "user", "content": "next"}] ++ svc.encode_prompt(turn2, None, {}) ++ svc.encode_prompt([{"role": "user", "content": "one unrelated prompt " * 300}], None, {}) ++ self.assertEqual(len(svc.prompts.entries), 2) ++ turn3 = turn2 + [{"role": "assistant", "content": "x"}, {"role": "user", "content": "more"}] ++ svc.encode_prompt(turn3, None, {}) ++ self.assertGreater(svc.prompts.last_reused, len(turn2[0]["content"])) ++ self.assertEqual(len(svc.prompts.entries), 2) # turn 3 replaced turn 2, the other one stayed ++ ++ def test_a_refused_prompt_is_not_kept(self): ++ svc = self.service(max_context=4096) ++ msgs = [{"role": "user", "content": "kept " * 20}] ++ svc.prepare(msgs, None, {}, 16) ++ refused = msgs + [{"role": "assistant", "content": "a"}, {"role": "user", "content": "more " * 20}] ++ full = self.service(max_context=4096) ++ full.prompts = None ++ refused_ids = full.encode_prompt(refused, None, {}) ++ with self.assertRaisesRegex(ValueError, "exceeds the context"): ++ svc.prepare(refused, None, {}, 1 << 20) ++ self.assertNotIn(refused_ids, [e[2] for e in svc.prompts.entries]) ++ self.assertNotIn(full.render_prompt(refused, None, {}), [e[0] for e in svc.prompts.entries]) ++ ++ def test_random_messages_equal_a_full_encode(self): ++ """Messages quoting control-token texts and tags at random, growing, edited and cut back between ++ requests: the incremental encoder's ids are the full encode's, 2400 prompts.""" ++ tok = ThinkTokenizer() ++ template = ChatTemplate(ROOT / "serve/chat_template.jinja") ++ svc = Service(MockEngine(tok, "ok", max_context=1 << 20), tok, template) ++ full = Service(MockEngine(tok, "ok", max_context=1 << 20), tok, template) ++ full.prompts = None ++ rng = random.Random(567) ++ quoted = ThinkTokenizer.SPECIALS + ["<|im_sta", " 2: # a shorter conversation ++ msgs = msgs[:rng.randrange(1, len(msgs))] + msgs[-1:] ++ with self.subTest(conv=conv, turn=turn): ++ self.assertEqual(svc.encode_prompt(msgs, None, {}), full.encode_prompt(msgs, None, {})) ++ reused += svc.prompts.last_reused > 0 ++ msgs.append({"role": "assistant", "content": text()}) ++ self.assertGreater(reused, 1000) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_q36_config.py b/serve/test_q36_config.py +new file mode 100644 +index 00000000..981ae5f3 +--- /dev/null ++++ b/serve/test_q36_config.py +@@ -0,0 +1,48 @@ ++"""serve/test_q36_config.py - what the server hands the Qwen3.6-35B-A3B engine (arch qwen35moe): the config's pictures ++("vision" section -> --vision) and parallel requests ("parallel": N -> --batch N), as for Flash-Next's engine. ++ ++ python -m unittest serve.test_q36_config ++""" ++from __future__ import annotations ++ ++import unittest ++from unittest import mock ++ ++from serve.server import engine_args, parallel_args ++ ++ ++def cfg(**kw): ++ base = {"arch": "qwen35moe", "exe": "strata", "args": ["--native", "/m/Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf", ++ "--max-context", "32768", "--prefill", "auto", "--mtp", "2"]} ++ base.update(kw) ++ return base ++ ++ ++class Q36Config(unittest.TestCase): ++ def test_parallel_becomes_batch(self): ++ a = engine_args(cfg(parallel=4)) ++ self.assertEqual(a[a.index("--batch") + 1], "4") ++ self.assertNotIn("--batch", engine_args(cfg())) ++ self.assertNotIn("--batch", engine_args(cfg(parallel=1))) ++ ++ def test_more_than_the_window_is_said(self): ++ with mock.patch("builtins.print") as pr: ++ a = engine_args(cfg(parallel=12)) ++ self.assertEqual(a[a.index("--batch") + 1], "12") # the engine caps it at 8 and says so ++ self.assertTrue(any("at most 8" in str(c) for c in pr.call_args_list)) ++ ++ def test_a_vision_section_starts_the_engine_with_pictures(self): ++ v = {"exe": "strata-vision", "mmproj": "/m/mmproj-F16.gguf", "model": "/m/Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf", ++ "gpu": False, "max_tokens": 300} ++ self.assertIn("--vision", engine_args(cfg(vision=v))) ++ self.assertIn("--vision", engine_args(cfg(vision=v, args=cfg()["args"] + ["--vision"]))) ++ self.assertEqual(engine_args(cfg(vision=v, args=cfg()["args"] + ["--vision"])).count("--vision"), 1) ++ self.assertNotIn("--vision", engine_args(cfg())) ++ ++ def test_an_explicit_batch_in_the_args_wins(self): ++ c = cfg(parallel=4, args=cfg()["args"] + ["--batch", "2"]) ++ self.assertEqual(parallel_args(c, c["args"]), []) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_web_sampling.py b/serve/test_web_sampling.py +new file mode 100644 +index 00000000..b502ace2 +--- /dev/null ++++ b/serve/test_web_sampling.py +@@ -0,0 +1,77 @@ ++"""serve/test_web_sampling.py - the web page sends no sampling of its own until the user changes it (0.1.45, #2685). ++ ++The page used to send temperature 0.6 / top_p 0.95 / top_k 20 with every request, so the server's model-card defaults ++(thinking 1.0, not thinking 0.7 + presence penalty) never applied to it. Pins, with node, the page's own functions: ++untouched (the shown defaults, or the old stored 0.6) sends nothing, a changed value is sent whole, 0 is greedy. ++Skipped without node. ++ ++ python -m unittest serve.test_web_sampling -v ++""" ++from __future__ import annotations ++ ++import json ++import re ++import shutil ++import subprocess ++import unittest ++from pathlib import Path ++ ++APP = Path(__file__).resolve().parent / "web" / "app.js" ++ ++ ++def slice_between(src: str, start: str, end: str) -> str: ++ a = src.index(start) ++ return src[a:src.index(end, a)] ++ ++ ++@unittest.skipUnless(shutil.which("node"), "node is not installed") ++class WebSampling(unittest.TestCase): ++ @classmethod ++ def setUpClass(cls): ++ src = APP.read_text(encoding="utf-8") ++ cls.defaults = slice_between(src, "const DEFAULTS", "let settings") ++ cls.shared = slice_between(src, "function sharedDefaults", "async function saveShared") ++ cls.body = slice_between(src, " if (!(settings.temperature > 0)) {", " if (settings.seed)") ++ ++ def run_js(self, settings: dict, pieces: str): ++ js = (self.defaults + "function projectionLoaded() { return false; }\n" + self.shared + ++ f"\nconst settings = {json.dumps(settings)};\nconst body = {{}};\n{pieces}\n" ++ "console.log(JSON.stringify({body, shared: sharedDefaults(settings)}));") ++ out = subprocess.run(["node", "-e", js], capture_output=True, text=True, timeout=30) ++ self.assertEqual(out.returncode, 0, out.stderr) ++ return json.loads(out.stdout) ++ ++ def sampling(self, **kw): ++ s = {"thinking": "high", "temperature": 1.0, "top_p": 0.95, "top_k": 20, "max": "", "seed": "", "esp": True} ++ s.update(kw) ++ r = self.run_js(s, self.body) ++ return r["body"], r["shared"] ++ ++ def test_untouched_sends_nothing(self): ++ for t in (1.0, 0.6): # the shown default, and the old stored default ++ body, shared = self.sampling(temperature=t) ++ self.assertEqual(body, {}) ++ self.assertEqual(shared, {"reasoning_effort": "high"}) ++ ++ def test_changed_values_are_sent_whole(self): ++ body, shared = self.sampling(temperature=0.3) ++ self.assertEqual(body, {"temperature": 0.3, "top_p": 0.95, "top_k": 20}) ++ self.assertEqual(shared["temperature"], 0.3) ++ body, _ = self.sampling(top_k=40) ++ self.assertEqual(body, {"temperature": 1.0, "top_p": 0.95, "top_k": 40}) ++ body, _ = self.sampling(temperature=0.6, top_p=0.9) ++ self.assertEqual(body["temperature"], 0.6) ++ ++ def test_zero_is_greedy(self): ++ body, shared = self.sampling(temperature=0) ++ self.assertEqual(body, {"temperature": 0}) ++ self.assertEqual(shared["temperature"], 0) ++ ++ def test_the_old_default_is_shown_as_the_new_one(self): ++ src = APP.read_text(encoding="utf-8") ++ self.assertRegex(src, r"if \(samplingUntouched\(settings\)\) settings\.temperature = DEFAULTS\.temperature;") ++ self.assertIsNone(re.search(r"temperature: 0\.6", src)) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/serve/test_image_omission.py b/serve/test_image_omission.py +new file mode 100644 +index 00000000..677229aa +--- /dev/null ++++ b/serve/test_image_omission.py +@@ -0,0 +1,69 @@ ++"""The local text-only image switch must run before incremental prompt encoding.""" ++import copy ++import os ++import unittest ++from pathlib import Path ++from unittest import mock ++ ++from serve.frontend import ChatTemplate ++from serve.server import ByteTokenizer, MockEngine, Service ++ ++ ++ROOT = Path(__file__).resolve().parents[1] ++ ++ ++class ImageOmission(unittest.TestCase): ++ def service(self): ++ tok = ByteTokenizer() ++ return Service(MockEngine(tok, "ok", max_context=1 << 20), tok, ++ ChatTemplate(ROOT / "serve/chat_template.jinja")) ++ ++ def messages(self): ++ return [{"role": "user", "content": [ ++ {"type": "text", "text": "describe literally"}, ++ {"type": "image", "source": "unused.png"}, ++ ]}] ++ ++ @mock.patch.dict(os.environ, {"STRATA_OMIT_IMAGES_WHEN_NO_VISION": "1"}) ++ def test_omitted_images_are_cached_as_text_and_retries_are_stable(self): ++ svc = self.service() ++ original = self.messages() ++ messages = copy.deepcopy(original) ++ text = messages[0]["content"][0] ++ first, _, _ = svc.prepare(messages, None, {}, 16) ++ self.assertIs(messages[0]["content"][0], text) ++ self.assertEqual(messages[0]["content"][1], { ++ "type": "text", "text": "[image omitted on text-only engine]"}) ++ self.assertEqual(svc.prompts.entries[0][2], first) ++ full = self.service() ++ full.prompts = None ++ self.assertEqual(first, full.encode_prompt(messages, None, {})) ++ again, _, _ = svc.prepare(copy.deepcopy(original), None, {}, 16) ++ self.assertEqual(first, again) ++ self.assertGreater(svc.prompts.last_reused, 0) ++ same_object, _, _ = svc.prepare(messages, None, {}, 16) ++ self.assertEqual(first, same_object) ++ ++ @mock.patch.dict(os.environ, {"STRATA_OMIT_IMAGES_WHEN_NO_VISION": "0"}) ++ def test_disabled_switch_keeps_the_normal_vision_error(self): ++ svc = self.service() ++ messages = self.messages() ++ before = copy.deepcopy(messages) ++ with self.assertRaisesRegex(ValueError, "without the vision encoder"): ++ svc.prepare(messages, None, {}, 16) ++ self.assertEqual(messages, before) ++ ++ @mock.patch.dict(os.environ, {"STRATA_OMIT_IMAGES_WHEN_NO_VISION": "1"}) ++ def test_configured_vision_is_not_omitted(self): ++ svc = self.service() ++ svc.vision = object() ++ messages = self.messages() ++ before = copy.deepcopy(messages) ++ with mock.patch.object(svc, "_note_unreadable_tool_images", side_effect=RuntimeError("before encoding")): ++ with self.assertRaisesRegex(RuntimeError, "before encoding"): ++ svc.prepare(messages, None, {}, 16) ++ self.assertEqual(messages, before) ++ ++ ++if __name__ == "__main__": ++ unittest.main()