mlx-bun 0.0.10 → 0.0.12

package/docs/investigations/dspark-handoff.md

@@ -1,11 +1,176 @@
 # DSpark — handoff
 
-**State:** faithful DSpark (DFlash KV-injection + Markov + confidence) is **built
-and proven architecturally correct** (overfit τ=3.24, paper-range). It does not
-yet net a speedup on e4b — that's a data-scale + target-speed problem, not an
-architecture problem. Full design + measured numbers:
+**State (2026-07-06): the DSpark paper is CODE-COMPLETE and serve-integrated.**
+Architecture proven correct (overfit τ=3.24, paper-range); serve-integrated
+behind `--draft-model` (Phase 1 — the `DraftSource` seam is KV-borrowing-ready;
+sibling AssistantSource verified bit-lossless on real weights,
+`tests/spec-serve-assistant.test.ts`); and every remaining paper component
+landed in Phase 2 the same day: Alg-1 confidence-scheduled draft-length pruning
+(checkpoint-driven, zero change for uncalibrated checkpoints), STS calibration
+(§3.2.1, `calibration.ts` + `scripts/dspark-calibrate.ts`), the RNN sequential
+head (Eq 6, `--seq-head rnn`, init-equivalent to Markov), the tightened draft
+loop (on-device token chaining + deferred conf reads + `collectLogits:false`
+on serve — bit-identity pinned), and the `dspark` variant rename + central
+loader. Suites: smoke 21/21 · dspark tests 17/17 · adjacent server suites
+green. **What remains is exactly the GPU work below** — data scale + 12B
+retarget + train + calibrate + measure; e4b not netting a speedup is a
+data/target-speed problem, not architecture. Full design:
 [docs/design/dspark-speculative-decoding.md](../design/dspark-speculative-decoding.md).
 
+## PATH A (preferred, no training) — DeepSeek's published 12B drafter
+
+DeepSeek released a TRAINED drafter for our exact target:
+`deepseek-ai/dspark_gemma4_12b_block7` (6.86 GB bf16, γ=7; DeepSpec repo is
+the MIT reference impl + the oracle — see the "DeepSpec ground truth" section
+of the design doc). Once the deepspec-variant port lands:
+
+```sh
+# 1. download the trained drafter (~6.9 GB)
+mlx-bun get deepseek-ai/dspark_gemma4_12b_block7
+# 2. oracle gate: dump their reference trace (temp 0 = RNG-free deterministic)
+#    in a torch venv, then compare ours round-for-round (scripts staged)
+# 3. measure + serve (auto-detected via config.json architectures field)
+mlx-bun serve --model gemma-4-12B-it-OptiQ-4bit --draft-model <snapshot-of-dspark_gemma4_12b_block7>
+```
+
+### Oracle protocol (step 2, expanded)
+
+Two staged scripts implement the temp-0 deterministic trace protocol against
+DeepSpec's own reference implementation (github.com/deepseek-ai/DeepSpec, MIT).
+At temperature 0 their leaky-rejection verify degenerates to exact argmax
+token-match — the whole round trace is RNG-free and reproducible, so a
+plain round-for-round diff is a real bit-exact gate (same discipline as the
+D2 diffusion oracle, minus RNG parity entirely — see the design doc's
+"DeepSpec ground truth" section).
+
+```sh
+# a. dump the reference trace — needs its OWN torch venv (NOT the mlx-lm
+#    oracle venv at /Users/joshrossi/Code/mlx-lm/.venv, which has no
+#    torch/deepspec and shouldn't); see the script header for the exact
+#    venv-setup + pip install steps.
+.venv-deepspec/bin/python scripts/oracle-dspark-deepspec.py \
+    --target <snapshot-dir-of-google/gemma-4-12B-it> \
+    --drafter <snapshot-of-dspark_gemma4_12b_block7> \
+    --data prompts.jsonl --n 8 --max-new-tokens 128 \
+    --temperature 0 --confidence-threshold 0 \
+    --out goldens/dspark-deepspec/trace-thr0.jsonl
+
+# b. compare our port round-for-round
+bun scripts/dspark-deepspec-compare.ts \
+    --fixture goldens/dspark-deepspec/trace-thr0.jsonl \
+    --target gemma-4-12B-it-bf16 --drafter <snapshot-of-dspark_gemma4_12b_block7>
+```
+
+Run the pair twice — once at `--confidence-threshold 0` (no pruning), once at
+`0.5` — to gate both the base draft loop and the confidence-truncation path.
+
+Fixture location: `goldens/dspark-deepspec/*.jsonl` follows the house
+manifest convention (goldens/README.md) — these are machine-independent
+**JSON manifests** (prompts + token-id traces, no tensor blobs), so they are
+the **tracked** kind, not the `.bin` machine-specific kind. Regenerate by
+rerunning step (a); there is no `scripts/regen-*` wrapper for this family yet
+since the oracle run needs its own GPU-hosting torch venv, not the mlx-lm
+oracle venv this repo otherwise standardizes on.
+
+**The OptiQ-4bit caveat is load-bearing, not a footnote**: the fixture above
+is generated against the bf16 HF target. A `dspark-deepspec-compare.ts` run
+with `--target` pointing at our bf16 12B is the TRUE bit-exact gate. A run
+with `--target` pointing at the OptiQ-4bit-quantized 12B we actually serve is
+a **separate acceptance-rate measurement** — tapped hiddens differ
+numerically at 4-bit, so round-for-round divergence there is expected, not a
+bug; pass `--acceptance-only` on that arm so the script reports an aggregate
+accept-length reading instead of a PASS/FAIL assert.
+
+Caveat: trained against the bf16 HF target; our serving 12B is OptiQ-4bit —
+tapped hiddens differ numerically, so measure acceptance rather than assuming
+their 60–85%.
+
+**FIRST LIVE RUN (2026-07-07, loaded machine — directional only):** τ ≈ 2.8
+committed/target-forward (200 tok / 72 verify rounds; per-token acceptance
+26–33% — the predicted OptiQ-4bit degradation), output coherent, losslessness
+holding. **But wall-clock LOST ~3.4×** (spec 14.6 vs serial 49.8 agg tok/s;
+the 20 s bench TTFT is conc-4 queueing on the serial-forced spec lane, not
+prefill). The tax is the drafter itself. **The full step-by-step program for
+the next pass is
+[docs/design/dspark-serving-program.md](../design/dspark-serving-program.md)**
+(phases, gates, dependency graph; pickup at TurboQuant merge, only its
+Phase 5 waits for it). Lever summary:
+1. **Quantize the drafter** — the big one, and the ideal quantization target
+   PERIOD: drafter numerics only move ACCEPTANCE, never correctness
+   (losslessness is structural), so aggressive bpw is safe and the gate is a
+   single acceptance A/B. 6.9 GB bf16 → ~1.8 GB @4-bit kills both the
+   ~3 GB/round bandwidth tax (2 GB lm_head + 7×134 MB markov_w2 reads) and
+   the 14 GB two-model memory pressure on 24 GB. Uniform 4-bit (mlx-native)
+   first as the baseline; **TurboQuant (Phase 13), when it lands, competes
+   on acceptance-per-byte against that baseline** — the drafter is its
+   lowest-risk first customer. Enabler either way: quantized-weight support
+   in deepspec-module's forward (quantized_matmul, the house pattern).
+2. **Tighten deepspec draftBlock** — ~14 host syncs/round (7 argmax + 7
+   confidence reads); replay the on-device chaining + deferred-conf
+   tightening already done for our own module (threshold=0 needs NO conf
+   reads).
+3. **captureLayer in the GENERATED gemma forwards** — tapped verifies
+   currently fall back to the monolith (correct-but-unspecialized); teach
+   scripts/gen-model.ts to emit the tap so spec verifies keep the fast path.
+4. **γ sweep** — γ=7 at ~30% acceptance wastes drafts;
+   `--num-draft-tokens 2..3` may net better wall-clock free of any code.
+
+## PATH B — train our own module (custom targets / research)
+
+The payoff target is **12B** (27B is memory-infeasible to train on 24 GB; the
+pipeline stays dim-generic so 27B remains a future recipe). All three scripts
+take `--model` and (new, 2026-07-06) `--tap-layers` — the retarget is now a
+pure parameterization, no code edits. **`--tap-layers` MUST be identical across
+regen and train** (the shard feature dim is `m*H`, `m = tapLayers.length`).
+
+12B tap set: use **`5,17,29,41,46`** — DeepSeek's actual trained
+`target_layer_ids` for this model (their layer-output convention matches our
+tapLayers indexing; audited 2026-07-06). Our earlier `24,37,47,48` guess is
+superseded.
+
+```sh
+# 0. one-time: cache the 12B target + a prompt/topic corpus (thousands of
+#    on-distribution topics — the data-scale lever; 160 only reached per-pos ~0.17).
+mlx-bun get mlx-community/gemma-4-12B-it-OptiQ-4bit
+# (12B target may need the serve process stopped to free memory for regen.)
+
+# 1. regen multi-layer training shards from the 12B target's own generations
+bun scripts/dspark-regen-dflash.ts \
+  --model gemma-4-12B-it-OptiQ-4bit --tap-layers 5,17,29,41,46 \
+  --topics <thousands-of-topics.txt> --out data/dspark-12b --max-resp 320
+
+# 2. train the drafter against the 12B (warm-start survivable via --resume;
+#    GPU runs kept getting killed ~step 5000 pre-resume). Optional A/B:
+#    --seq-head rnn trains the Eq-6 RNN head instead of Markov (Eq 5).
+bun scripts/dspark-train-dflash.ts \
+  --model gemma-4-12B-it-OptiQ-4bit --tap-layers 5,17,29,41,46 \
+  --data data/dspark-12b --out ckpt/dspark-12b \
+  --iters 8000 --batch 8 [--seq-head rnn] [--ddraft 2560] [--resume ckpt/dspark-12b]
+
+# 3. STS-calibrate (§3.2.1): fit per-position confidence thresholds on real
+#    verify outcomes and stamp them into dspark.json — this ACTIVATES the
+#    Alg-1 confidence scheduler at decode (uncalibrated = fixed-γ, no change).
+bun scripts/dspark-calibrate.ts \
+  --model gemma-4-12B-it-OptiQ-4bit \
+  --drafter ckpt/dspark-12b --data <prompts.jsonl> [--target 0.5]
+
+# 4. measure live τ + tok/s vs vanilla 12B decode (clean machine — dirty
+#    numbers are garbage). Run BOTH pre- and post-calibration for the
+#    scheduler's own A/B (MLX_BUN_DSPARK_MINCONF=<c> = manual uniform knob).
+bun scripts/dspark-measure-dflash.ts \
+  --model gemma-4-12B-it-OptiQ-4bit \
+  --drafter ckpt/dspark-12b --data <prompts.jsonl>
+
+# 5. serve it (auto-detected as kind=dspark via dspark.json; γ pinned to cfg.gamma)
+mlx-bun serve --model gemma-4-12B-it-OptiQ-4bit --draft-model ckpt/dspark-12b
+```
+
+Decision gate (write it down, RESULTS.md only if clean-machine): does 12B live
+τ net a wall-clock win? If the draft loop dominates, do the loop-tightening
+(per-position host syncs + double 262K LM-head) first — see the open items in
+the design doc. Memory guard (`server.ts`) caps context to fit; a client
+sending huge `max_tokens` trips admission — lower it.
+
 ## What happened (short version)
 - Built v1 first as a single-vector shortcut — it dropped the paper's core KV
   injection and was never really DSpark. Chunk (τ≈1.26) and article (τ≈1.16)
@@ -31,20 +196,21 @@ architecture problem. Full design + measured numbers:
 - Training `--resume` (warm-start from checkpoint) — the GPU runs kept getting killed ~step 5000; resume survives that.
 
 ## Next (in priority order)
-1. **Retarget to a slow model (27B or 12b).** The drafter reads dims from the
-   model + tapLayers is config, so it ports — but needs regen+train on that
-   model. 27B on 32GB is memory-tight (17.75GB weights); may need the serve
-   process stopped to free the ~14GB for regen. This is where a real speedup lives.
-2. **Scale training data** — thousands of on-distribution generations (not 160)
-   to lift generalizing τ toward the overfit ~0.75.
-3. **Tighten the draft inference loop** — per-position host syncs + double 262K
-   LM-head make the draft heavy; needed for τ to translate to wall-clock.
-4. **Finish the paper:** RNN head (Eq 6), STS calibration (§3.2.1), hardware-aware
-   prefix scheduler (Alg 1 — single-user = confidence-scheduled draft-length pruning).
-5. **Rename** `dflash`→`dspark` (faithful module IS DSpark; v1 = legacy single-vector).
+1. **Retarget to 12B + scale data + train + calibrate + measure** — the GPU
+   recipe above is the complete remaining work. (27B stays a future recipe:
+   memory-infeasible to train on the 24 GB M4 Pro.)
+2. ~~Tighten the draft inference loop~~ **DONE 2026-07-06** (on-device token
+   chaining, deferred conf reads, `collectLogits:false`; bit-identity pinned).
+3. ~~Finish the paper~~ **DONE 2026-07-06**: RNN head (Eq 6, `--seq-head rnn`),
+   STS calibration (§3.2.1), Alg-1 confidence-scheduled draft-length pruning
+   (single-user form). ⚠ Eq 6 + §3.2.1 shapes are design-doc-faithful readings
+   (paper PDF absent) — verify against the paper when available.
+4. ~~Rename `dflash`→`dspark`~~ **DONE 2026-07-06** (canonical variant +
+   central loader `src/spec/dspark/loader.ts`; legacy stamps still load).
 
 ## Key files
-`src/spec/dspark/{module,data,generate}-dflash.ts`, `{loss,sample}.ts`;
+`src/spec/dspark/{module,data,generate}-dflash.ts`, `{loss,sample,calibration,loader}.ts`;
+`src/spec/{source,serve-loop,dflash-source,assistant-source}.ts` (the serve seam);
 `src/model/gemma4.ts` (hiddenTap), `gemma4-base.ts` (trim bypass);
 `scripts/dspark-{regen,train,measure}-dflash.ts`, `dspark-dflash-smoke.ts`.
 

package/docs/investigations/indexshare-performance-spike.md

@@ -0,0 +1,357 @@
+# IndexShare / IndexCache performance spike
+
+**Date:** 2026-08-16
+**Status:** accepted research spike; paper/code audit complete; implementation and
+full-model measurement not started
+**Target:** GLM-5.2 DSA on the M1 Max 32 GB Colibri path
+
+## Recommendation
+
+Run the spike, but do not treat it as a new model feature. The GLM-5.2 path
+already implements the checkpoint's cross-layer IndexShare schedule: retained
+`full` layers run the DSA indexer and intervening `shared` layers reuse the most
+recent selection. The useful work is to make that correctness path production
+shaped, prove long-context parity, and measure whether the saved indexer work is
+visible behind expert I/O on Apple Silicon.
+
+The first prerequisite is the exact stock indexer overlay. The pinned public
+Colibri artifact has no `out-idx-*` files, so every full-model result recorded
+to date is DSA-off. The second prerequisite is a scalable indexer path: current
+code materializes a large `[heads, context, head_dim]` product, copies every
+score to the host, sorts the full score vector in JavaScript, and supports
+sparse DSA only for single-token decode. Those choices are good correctness
+scaffolding but cannot test the paper's long-prefill result.
+
+This spike must not change the trained `full`/`shared` pattern or expose it as a
+normal user knob. GLM-5.2 ships IndexShare as checkpoint architecture, not as an
+optional quality/performance preset.
+
+## Naming and source claims
+
+The linked paper calls the method **IndexCache**. The GLM-5.2 model card calls
+the deployed architecture **IndexShare**. They use the same core mechanism:
+cache the top-k positions produced by a retained DSA indexer and reuse them in
+subsequent shared layers.
+
+Primary sources:
+
+- [IndexCache paper, arXiv:2603.12201v1](https://arxiv.org/abs/2603.12201v1)
+- [authors' SGLang/vLLM reference patches](https://github.com/THUDM/IndexCache)
+- [pinned GLM-5.2 model card](https://huggingface.co/zai-org/GLM-5.2/blob/f6142f127a14b58dc602592e996cd7d8ff139351/README.md)
+- [pinned GLM-5.2 configuration](https://huggingface.co/zai-org/GLM-5.2/blob/f6142f127a14b58dc602592e996cd7d8ff139351/config.json)
+
+The sources establish three different claims that must not be blended:
+
+| Claim | Published evidence | What mlx-bun may say before local measurement |
+|---|---|---|
+| Cross-layer reuse | F layers compute top-k; S layers reuse the nearest preceding F layer | Implemented semantics, once parity is gated |
+| IndexCache speed | On a 30B DSA model on an 8x H100 node, 1/4 retention reached up to 1.82x prefill and 1.48x per-request decode at 200K context | Paper result only; not a Mac prediction |
+| GLM-5.2 IndexShare | The model card says one retained indexer per four sparse-attention layers reduces per-token FLOPs by 2.9x at 1M context | Vendor FLOP claim only; not end-to-end speed |
+
+The paper also reports that uniform 1/4 retention hurts a frozen 30B model's
+long-context average (50.2 to 43.0), while a loss-searched pattern recovers it
+to 49.9. Training-aware IndexCache makes uniform interleaving viable. Therefore
+mlx-bun must preserve GLM-5.2's trained schedule rather than search a new one
+during this spike.
+
+## GLM-5.2 schedule
+
+The pinned GLM-5.2 config has 78 layers, `index_topk=2048`,
+`index_topk_freq=4`, and `index_skip_topk_offset=3`. Its explicit schedule has
+21 full and 57 shared layers:
+
+```text
+FFFSSSFSSSFSSSF...FSSS
+```
+
+Relative to an all-full counterfactual, the checkpoint omits 57/78 = 73.1% of
+layer indexer invocations. It does not ship the missing shared-layer indexer
+weights, so a true same-checkpoint `IndexShare off` A/B is impossible. Any
+estimated all-full number must be labeled a component counterfactual, not an
+end-to-end model benchmark.
+
+`index_share_for_mtp_iteration` is a second, orthogonal optimization: reuse
+indices across speculative draft iterations. mlx-bun parses that field but does
+not currently consume it in the MTP runtime. Keep its measurement separate so
+cross-layer IndexShare and MTP-iteration sharing each receive an attributable
+result.
+
+## Current mlx-bun state
+
+Already present:
+
+- [`src/model/glm52-config.ts`](../../src/model/glm52-config.ts) parses explicit
+  `indexer_types` and the frequency/offset fallback.
+- [`src/model/glm52.ts`](../../src/model/glm52.ts) constructs DSA indexers only
+  for full layers and passes one request-local selection state through the
+  decoder stack.
+- [`src/model/glm52-dsa.ts`](../../src/model/glm52-dsa.ts) records a full
+  layer's exact dense/sparse selection and serves it to shared layers.
+- [`tests/glm52-config-container.test.ts`](../../tests/glm52-config-container.test.ts)
+  and [`tests/glm52-dsa.test.ts`](../../tests/glm52-dsa.test.ts) cover schedule
+  parsing, missing-owner refusal, context consistency, and exact reuse.
+
+Known blockers and likely bottlenecks:
+
+1. **No production DSA overlay.** The public Colibri artifact deliberately
+   disables DSA because every full-layer indexer tensor is absent. Producing the
+   exact overlay requires selective reads from 20 pinned source shards (about
+   99.90 GiB input, about 197 MB output). This is a Josh-run acquisition step,
+   not an agent-session download.
+2. **No sparse long prefill.** `Glm52DsaIndexer.projectAndSelect()` rejects a
+   multi-token prefill once context exceeds `index_topk`. The paper's prefill
+   headline is therefore out of scope until chunked/sparse DSA prefill exists.
+3. **Quadratic transient shape.** `glm52DsaScoresMlx()` forms the broadcast
+   product before reduction. At production geometry this scales as
+   `32 * context * 128` float elements rather than streaming/tiled score
+   accumulation.
+4. **GPU-to-host top-k.** The full score vector crosses to a `Float32Array`,
+   then `selectDsaThresholdTiesF32()` clones and sorts all positions. This is
+   `O(context log context)` host work in each full layer, unlike the production
+   GPU top-k assumed by the paper.
+5. **Shared-index copies.** Shared layers clone the 2,048-element JavaScript
+   positions list, and absorbed MLA uploads it again before each gather. The
+   paper's temporary index tensor is reused in place. Measure this separately;
+   remove copies only after ownership and asynchronous lifetime are gated.
+6. **The current short-context bottleneck is elsewhere.** At 128 tokens the
+   full model is limited by streamed expert delivery, and DSA cannot prune a
+   prefix smaller than 2,048. No short-context speedup should be expected.
+
+## Hypotheses
+
+- **H1 - component cost:** at the longest context that fits the 32 GB contract,
+  retained-indexer time is material and grows roughly linearly per decode token
+  with context length.
+- **H2 - sharing value:** avoiding 57 of 78 indexer invocations removes a
+  measurable amount of counterfactual work even after a tiled score/top-k
+  kernel is used.
+- **H3 - product value:** DSA plus the shipped IndexShare schedule improves
+  end-to-end long-context decode over DSA-off full attention, without token
+  drift from the official DSA runtime. This comparison measures the combined
+  sparse-attention product, not IndexShare in isolation.
+- **H4 - MTP interaction:** any win remains a win with native MTP enabled. A
+  decode-only win that disappears under the product MTP path is not a default
+  performance claim.
+
+## Bounded experiment
+
+### Stage 0 - artifact and parity gate
+
+1. Generate the immutable `out-idx-*` overlay from the pinned
+   `zai-org/GLM-5.2-FP8` revision. Record source revision, selected shard list,
+   tensor inventory, byte count, and SHA-256 values. Do not mutate the 357 GiB
+   serving snapshot.
+2. Header-validate exactly the 21 full-layer indexers required by the shipped
+   schedule and prove that shared-layer indexers are neither required nor
+   silently loaded.
+3. Add an official-runtime oracle at contexts above 2,048. Gate selected
+   positions, logits on tie-free rows, and greedy token IDs before taking any
+   performance number.
+
+#### Stage 0 result — pass (2026-08-17)
+
+The stock overlay was generated outside mlx-bun from
+`zai-org/GLM-5.2-FP8@ba978f7d347eaf65d22f1a86833408afdb953541`
+with `colibri@ecade075cfc2eae684097ea7de5570c3786ce199`. It contains 20
+`out-idx-*` files totaling 197,202,400 bytes. The existing header validator
+finds 118,646 tensors, 59,066 quantized tensors, 19,456 routed experts,
+complete DSA and MTP families, all 21 required full-layer indexers, and no
+shared-layer indexer requirement.
+
+The first-sparse gate prefills exactly 2,048 deterministic tokens, then sends
+one greedy token at context 2,049 with MTP off. The mlx-bun probe observes the
+21 configured full layers and exactly 2,048 selected positions per layer. A
+detached pinned Colibri build captures each official score row, selection,
+float32 logit vector, and token trajectory. Replaying the captured scores with
+`Glm52DsaSelectionState` gives:
+
+- 21/21 exact ordered position vectors;
+- 21/21 exact float32 thresholds;
+- 21/21 tie-free boundaries, minimum boundary gap 0.025123596;
+- exact greedy trajectory `[264, 264]` and the same sparse-step top-1 token;
+- 154,880-logit cosine 0.997645, RMSE 0.200109, maximum absolute drift
+  0.947218, with wide top-1 margins in both runtimes (9.571914 / 9.567580).
+
+The direct full-runtime selections are 10/21 ordered-exact and 14/21
+set-exact. Each of the seven set differences swaps only the single excluded
+position at this top-2,048-of-2,049 boundary. This is upstream numerical drift
+from the runtimes' different quantized-matmul accumulation before DSA, not a
+selector mismatch; replaying identical official scores proves the isolated
+DSA threshold/tie contract exactly. No performance claim is made from either
+correctness run.
+
+Reproduction is split between
+`scripts/probe-colibri-glm52-dsa-long.ts` and
+`scripts/check-colibri-glm52-dsa-long.ts`. The raw capture, external Colibri
+instrumentation patch, and SHA-256 inventory are machine-local at
+`~/.cache/mlx-bun/evidence/glm52-dsa-stage0-2026-08-17/`; manifest SHA-256 is
+`6ecb88d944b055146dd700a1714af97b61845d4218527282fcbea009663fe3b4`.
+
+### Stage 1 - production-shaped indexer kernel
+
+Build a model-free benchmark at exact geometry (`32` heads, head dimension
+`128`, top-k `2048`) for contexts `2,048`, `8,192`, and `32,768`:
+
+- tiled score accumulation that never materializes `[H,L,D]`;
+- on-device/native deterministic top-k with the current threshold and
+  lower-position tie contract;
+- one compact index buffer retained across F -> S layers;
+- direct gather consumption without a score-vector host copy or per-shared-
+  layer index re-upload.
+
+Compare the candidate against the existing reference composition for exact
+selected positions on random, tied, all-equal, and production-captured rows.
+Record kernel time, synchronization time, peak MLX allocation, process
+footprint, and bytes copied to/from the host. Negative results count.
+
+#### Stage 1 result — pass (2026-08-17)
+
+The production path now computes `[H,D] @ [D,L] -> [H,L]` and never forms the
+old `[H,L,D]` broadcast product. Deterministic device selection encodes the
+lexicographic `(score descending, position ascending)` rank into `uint64`, uses
+MLX `argpartition`, and then restores Colibri's observable two-scan ordering
+(strictly above threshold first, threshold ties second). A FULL layer retains
+only the resulting 2,048-entry, 8 KiB index buffer; SHARED layers borrow that
+same device array and MLA gathers from it directly. Normal inference performs
+no score/index host readback or per-SHARED-layer upload. The probe observer is
+the explicit diagnostic exception.
+
+Correctness gates pass for deterministic random rows, threshold ties,
+all-equal rows, the exact 2,049/2,048 boundary, borrowed MLA consumption, and
+all 21 production-captured Colibri score rows. The latter reproduce exact
+ordered positions and float32 thresholds on device. A fresh full-model run
+also reproduced all 21 prior mlx-bun position vectors and both greedy tokens;
+its complete 154,880-float decode-logit payload is byte-identical (SHA-256
+`d1f62be2de4877d2fc8ae8feab88e522958210252b50ef96eb786cb22581ba18`).
+Tiled accumulation moved nine diagnostic thresholds by at most
+`3.0517578125e-5`, but no selection boundary moved.
+
+The exact-geometry, warm synchronized component benchmark (seven repeats) is:
+
+| Context | Production behavior | Score median | Score + top-k median | Timed host copies | MLX peak |
+|---:|---|---:|---:|---:|---:|
+| 2,048 | dense boundary; production skips score/top-k | 0.448 ms (diagnostic only) | — | 0 B | 1.27 MiB |
+| 8,192 | sparse | 0.473 ms | 0.929 ms | 0 B | 5.41 MiB |
+| 32,768 | sparse | 0.748 ms | 1.269 ms | 0 B | 21.35 MiB |
+
+At those contexts the removed broadcast product would have been 32, 128, and
+512 MiB respectively; the tiled dot output is 0.25, 1, and 4 MiB. These are
+model-free component measurements, not end-to-end inference claims. The
+benchmark is `scripts/bench-glm52-dsa-indexer.ts`. Machine-local evidence is
+under `~/.cache/mlx-bun/evidence/glm52-dsa-stage1-2026-08-17/`; manifest
+SHA-256 is
+`600ec92c41d4c66735e1834a678a09f2ae55ecdddfab31c3dd1bf31ee9709fef`.
+
+### Stage 2 - decode matrix
+
+Use direct library inference, not a server. Run separate fresh processes and
+the standard quiet-machine preflight.
+
+| Axis | Cells |
+|---|---|
+| Context | 2,048; 8,192; 32,768; 64K only if the exact planner remains inside the 25 GiB process ceiling |
+| DSA | off/full attention; on with the checkpoint's 21F/57S schedule |
+| MTP | off to isolate DSA; on for the product decision |
+| Residency | the same fixed expert policy and seed profile in every paired arm |
+| Repeats | three fresh processes per default-eligible cell; cold and warm turn in each process |
+
+For each turn record TTFT, prefill wall time, decode tok/s, total wall time,
+DSA projection/score/top-k/gather time, MLA time, expert disk service/wait,
+target forwards, MTP acceptance, physical footprint, compression, swap, and
+exact token IDs.
+
+The DSA-off versus DSA-on result combines sparse core attention with
+IndexShare. To estimate IndexShare's isolated contribution, report the measured
+retained-indexer total and a clearly labeled all-full component counterfactual
+derived from the same per-layer exact-geometry benchmark. Do not present that
+counterfactual as a runnable checkpoint.
+
+#### Stage 2 result — pass, negative product-performance decision (2026-08-17)
+
+The fresh-process runner planned 36 cells. Twenty-four were eligible under the
+G5 25 GiB process ceiling and completed: contexts 2,048 and 8,192, four DSA/MTP
+arms, and three repeats. Every process ran cold and warm turns. Cold/warm,
+cross-repeat, and MTP-off/on greedy token vectors were exact; all emitted
+sixteen token IDs equal to `264`. The twelve 32K combinations were not run:
+the exact planner reports 29,334,150,784 bytes (27.320 GiB) MTP-off and
+30,644,281,216 bytes (28.540 GiB) MTP-on, both above the unchanged 25 GiB
+process limit. The initial refusal is retained in the evidence as an attempt
+record. This is contract ineligibility, not an allocation crash.
+
+Long prompts use an explicitly bounded benchmark seam: dense batched prompt
+construction ends at the exact prompt boundary, and every subsequent decode or
+MTP verification row uses exact DSA. These numbers characterize decode and
+whole-turn wall time only. They are not sparse-prefill measurements.
+
+Warm medians across three fresh processes are:
+
+| Context | DSA | MTP | Decode tok/s | Total wall |
+|---:|---|---|---:|---:|
+| 2,048 | off | off | 0.1363 | 260.9 s |
+| 2,048 | on | off | 0.1318 | 271.5 s |
+| 2,048 | off | on | 0.1312 | 281.7 s |
+| 2,048 | on | on | 0.0880 | 338.7 s |
+| 8,192 | off | off | 0.1292 | 722.1 s |
+| 8,192 | on | off | 0.1420 | 711.8 s |
+| 8,192 | off | on | 0.1587 | 711.8 s |
+| 8,192 | on | on | 0.1042 | 770.1 s |
+
+At 8K without MTP, the paired median DSA delta is +12.38% decode throughput
+and -1.89% total wall time. Expert I/O and prompt work hide most of the decode
+gain, leaving the whole-turn result below the required 5%. With MTP enabled,
+the paired median delta reverses to -34.33% decode throughput and +8.19% total
+wall time. The exact MTP trace explains part of the product result: DSA-off
+accepted ten drafts in five rounds (3.2 tokens/target forward, six target
+calls), whereas DSA-on accepted eight drafts in seven rounds (2.286
+tokens/target forward, eight target calls). Small verify batches now keep
+residual/MLP/expert work batched, but exact causal sparse attention still runs
+per verification row. Two of the three 8K DSA+MTP cells observed net swap
+growth, so the memory side of the win gate also does not close.
+
+The immutable external overlay retains 197,202,400 bytes of indexers. As a
+clearly labeled non-runnable component counterfactual, Stage 1's 8K
+score+top-k median implies 19.51 ms for the shipped 21 FULL layers versus
+72.48 ms for hypothetical 78 FULL layers, so IndexShare avoids about 52.97 ms
+of selection work per target row. At 32K the corresponding component-only
+figures are 26.64 ms, 98.95 ms, and 72.31 ms avoided. The checkpoint lacks the
+57 missing indexers, and these sums exclude the rest of model execution.
+
+The product decision is negative: retain the exact 21F/57S implementation as
+model semantics and land its correctness/scalability work, but make no local
+end-to-end speed claim. Decode is conclusive, so Stage 3 remains a separate
+deferred prefill experiment rather than a G6R exit requirement. Evidence is
+under `~/.cache/mlx-bun/evidence/glm52-dsa-stage2-2026-08-17/`; manifest
+SHA-256 is
+`90b3fe4ed53714604b7a747991b3bb1b87aedbf57a139915065f5b4be42cda38`.
+
+### Stage 3 - deferred prefill follow-up
+
+Only after decode is characterized, implement chunked sparse prefill with exact
+causal selection semantics and run the same context ladder. This is the stage
+that can test whether the paper's growing prefill benefit translates to MLX;
+the existing single-token path cannot support that claim.
+
+## Decision rules
+
+- Preserve the shipped 21F/57S schedule regardless of outcome; it is model
+  semantics.
+- Call the work a local performance win only when three fresh MTP-on processes
+  are token-identical and the paired median end-to-end wall time improves by at
+  least 5% at a context that fits the 32 GB contract, without new swap or a
+  footprint regression outside the existing G5 bounds.
+- If the kernel saves DSA time but expert I/O hides the wall win, land only the
+  correctness/scalability work and record the negative end-to-end result.
+- Do not quote the paper's H100 speedups or GLM-5.2's 1M-token FLOP reduction as
+  mlx-bun speed. Curated claims require this machine's `benchmark.sh`-quality
+  evidence and belong in `benchmarks/RESULTS.md` only after the gate passes.
+- Evaluate MTP-iteration index sharing as a separate follow-up with its own
+  on/off control and acceptance trace.
+
+## Exit
+
+The spike closes when the DSA overlay is provenance-pinned, long-context
+IndexShare matches an official runtime, the production-shaped score/top-k path
+has exact model-free gates, and the paired MTP-off/on decode matrix has a
+written decision. Prefill may remain a separately tracked follow-up only if the
+decode result is already conclusive; otherwise sparse long prefill is required
+before closure.

package/docs/investigations/lab-build-journal.md

@@ -1,7 +1,7 @@
 # Build journal — mlx-bun Lab (web UI + native quantize/train/dataset + pi web chat)
 
 *Branch `web-ui-and-native-lab`, 2026-06-13/14. Visual summary:
-[archive/mlx-bun-lab-report.html](../../archive/mlx-bun-lab-report.html). Evidence
+[archive/mlx-bun-lab-report.html](../../docs/archive/mlx-bun-lab-report.html). Evidence
 artifact: `benchmarks/lab-verification.json` (regenerable, gitignored — produced by
 `scripts/verify-lab.ts`).*
 
@@ -22,7 +22,7 @@ page. Compute is **native Bun/MLX (no Python)** per Josh's decision.
 
 2. **Gate the riskiest unknown first.** Native LoRA needs autograd across the
    FFI boundary. Rather than discover mid-build that it's infeasible, the first
-   thing built was `spikes/phase-train-vag.ts`: a tiny LoRA-shaped loss
+   thing built was `lab/spikes/phase-train-vag.ts`: a tiny LoRA-shaped loss
    (`mean((x + (x@A)@B − y)²)`) differentiated by `mlx_value_and_grad` through a
    `bun:ffi` JSCallback closure, checked against central finite differences.
    **PASS (<1e-2), exactly 2 grads returned (base frozen).** Only then was the