# Experiments — What We Tried, What Shipped, What Didn't

This is the decision log for speed / quality experiments layered on top of the baseline SWA 4-chunk pipeline. Anything *not* in `conversion/convert.py`'s shipping path lives here.

Legend:

- **Shipping** — in `convert.py` or a default path, used by pre-converted HuggingFace models.
- **Prototype** — runs end-to-end, numerically validated, not yet plumbed into the default pipeline or the Swift runtime.
- **Shelved** — implemented far enough to judge, then rejected. Reason recorded.
- **Planned** — designed but not implemented.

See `docs/SPEED_8K.md` for the overall speed roadmap and tier assignments.

---

## Attention-variant experiments (8K context)

### WFA — Windowed Full Attention  —  Shelved

- Files: `conversion/build_wfa.py`, `conversion/models/gemma4_swa_wfa.py`, `conversion/benchmark_wfa.py`
- Idea: Cap the 7 full-attention layers' KV at a fixed window `FW` (e.g. 2048) instead of the full context. Same shift-based update as the 28 sliding layers, so every layer is O(W) per step.
- Build speed: at `FW=2048`, decode throughput matches the 2K baseline (~31 tok/s) even at 8K context. Bandwidth cost is the bottleneck the knob hits.
- **Reason shelved**: recall drops on prompts that need attention to tokens beyond `FW`. The quality regression on long-context tasks is the exact class of failure this project refuses to absorb silently. See `docs/SPEED_8K.md §0` ("naive WFA quality NG past FW"). Kept around because DuoAttention-style per-head windowing (Tier A) reuses the same shift-cache machinery.

### Flash Decoding  —  Prototype

- Files: `conversion/build_flash.py`, `conversion/models/gemma4_swa_flash.py`
- Idea: Split `Q @ K^T` along the K dimension into fixed-size chunks (default 1024) and recombine with an online softmax. Mathematically identical to standard attention within fp16 rounding (cosine > 0.999999 vs standard attention on test tensors).
- Status: builder works, outputs validated. Not integrated into the Swift runtime because the expected win on ANE is small — ANE is not SRAM-bound for 8K × 512 K/V (1 MB per chunk fits easily in 32 MB SRAM), so tiling buys less than it does on GPU. Keep for future much-longer-context variants.

### Merged chunk2+chunk3  —  Prototype

- File: `conversion/models/gemma4_swa_merged.py`
- Idea: Merge chunks 2 and 3 (L8–24) so `kv13` / `kv14` are produced and consumed on the ANE side without a round-trip to the Swift runtime. Saves 32 MB of fp16 KV I/O per decode step.
- Status: builder works, not yet benched on-device. Trade-off vs current 4-chunk split: the merged chunk is larger, which may push the ANE compiler past the ~15-layers-per-chunk stability line we hit elsewhere. Measure before shipping.

### Lite 2-chunk variant  —  Prototype

- Files: `conversion/models/gemma4_lite_chunks.py`, `conversion/models/gemma4_lite_wrapper.py`
- Idea: Two chunks (embedding + L0–14) and (L15–34 + LM head). Fewer chunks = less per-step overhead, but risks the ANE compiler stability ceiling.
- Status: kept as a fallback in case the 4-chunk split has problems on a new OS release. Not the default. Note: predates SWA, uses ctx-sized KV cache for every layer — not suitable for shipping.

### SWA 2-chunk / 1-chunk consolidation  —  Prototype (2026-04-15)

- Files: `conversion/models/gemma4_swa_merged2.py` (2-chunk), `conversion/models/gemma4_swa_merged1.py` (1-chunk), `conversion/build_merged_chunks.py`, `conversion/test_merged_parity.py`, `Sources/CoreMLLLM/ChunkedEngine.swift` (auto-detect + dispatch), `docs/CHUNK_CONSOLIDATION_BENCH.md`
- Motivation: `docs/BASELINE_SPEED_AUDIT.md` shows 4×2.3 ms dispatch overhead per step. Halving dispatches ≈ +14 tok/s on the 2K decode path.
- Idea:
  - **2-chunk**: `MergedChunk12` (L0-14 + PLE, owns KV) -> `MergedChunk34` (L15-34 + norm + LM head). Reuses `_run_layer_swa` from the shipping builder, so layer math is byte-identical — composition only.
  - **1-chunk**: `MergedChunk1` (all 35 layers, PLE, norm, LM head) in a single graph. kv13/kv14 stay internal so the Swift runtime never materialises them.
- Runtime: `ChunkedEngine` auto-detects layout by file presence; falls back to the 4-chunk path if merged files are missing. Merged layouts allocate their own 12×W + 3×ctx KV buffers. Prefill / speculative-verify paths keep the 4-chunk split (dispatch savings there are negligible and their shapes differ).
- Risk (not yet exercised on device): the ANE compiler stability ceiling. 15-layer merged chunk1 is near the line we hit in the WFA experiment; 35-layer merged_full is almost certainly past it. `ComputePlanAudit` gates shipping. See the bench doc for pass/ship criteria.
- Status as of 2026-04-15: builder, parity test and runtime wiring landed. Numerical parity on CPU-torch unverified (requires HF weights) but the merged forward paths re-use the reference `_run_layer_swa`, so failure would indicate a composition bug, not a math bug. Device bench pending.

### Stateless 4-chunk (no MLState)  —  Shipping

- File: `conversion/models/gemma4_stateless_chunks.py`
- Idea: Explicit KV input/output tensors instead of Apple's `MLState` API. Chosen over the Monolithic/`MLState` path because `MLState` introduces int64 state indices that break ANE placement on the model sizes we ship (see `docs/CONVERSION.md` on "Explicit KV I/O").
- This is what the default Gemma 4 E2B conversion actually produces.

---

## Quantization experiments

### INT4 palettization, group_size=32  —  Shipping

- File: `conversion/exporter.py :: _quantize_model`
- Chosen for weight-only compression. See `docs/CONVERSION.md` "Quantization" section for the size/quality/latency trade-off and why not INT8 or FP16.

### INT8 weight-only  —  Prototype

- File: `conversion/build_w8a8.py` (mode `w8`)
- Idea: symmetric per-channel INT8. Smaller than FP16, larger than INT4. Zero latency gain on ANE (ANE is FP16 internally, so weight-only INT8 gives size only — same as INT4 but at double the storage).
- Status: shelved as a *shipping* option for size reasons, kept as a baseline for measuring INT4's quality cost. Not useful standalone.

### W8A8 (naive calibration)  —  Shelved

- File: `conversion/build_w8a8.py` (mode `w8a8`)
- Idea: activation + weight INT8. This is the only quantization that unlocks ANE's INT8×INT8 compute path (~1.3–1.6× per Apple's ResNet-50 docs).
- Calibration used 5 random samples. Quality regressed visibly on chat outputs.
- **Reason shelved**: insufficient calibration. Superseded by `build_w8a8_proper.py`.

### W8A8 (realistic calibration)  —  Rejected (2026-04-13)

- File: `conversion/build_w8a8_proper.py`
- Idea: collect real activation traces by running the INT4 model on 32+ prompts at positions 0..31, then quantize. Also provides a W4A8 fallback (INT4 palette weights + INT8 activations) which is more stable than full W8A8 in practice.
- Outcome: Mac Studio M4 Max compiles and runs but **0% speedup** (ANE still runs FP16 internally); iPhone 17 Pro **`ANECCompile() FAILED`** — the `quantize`/`dequantize` MIL ops that `coremltools.optimize.coreml.linear_quantize_activations` inserts are not compilable by the iPhone ANE compiler. See `docs/SPEED_8K.md §1 A2 / Tier D`. No INT8 path reaches ANE on iOS 26 / coremltools 9.0.
- Side finding: `linear_quantize_activations` leaks ~1 temp `.mlpackage` per calibration op-group (tied to `atexit` rather than `__del__`), ~38 GB per Gemma 4 E2B chunk2 calibration run. Fillable disk. Worth filing upstream if anyone revives this path.

### W2A16 palettization (2-bit, post-training)  —  Rejected (2026-04-13)

- File: `conversion/build_flash.py --nbits 2`
- Idea: Apple's own 3.18B ships W2 palettized weights (per Foundation Models 2025 tech report). Reduce weight bandwidth 8× vs FP16 (2× vs INT4). Post-training palettization via `OpPalettizerConfig(nbits=2, granularity="per_grouped_channel", group_size=32)`. Expected ×1.4–2.0 decode speedup from bandwidth reduction.
- Size: chunk total **546 MB** (vs 1,091 MB INT4 = exactly 50%). Conversion succeeds, no compile errors.
- Quality: **complete gibberish**. Multilingual garbage tokens, zero coherent output on 3 test prompts (France capital, photosynthesis, haiku). Tested via `conversion/smoke_w2_quality.py` with real INT8 embeddings + RoPE, autoregressive generation on Mac CPU.
- **Root cause**: 4 codewords per group (2-bit) has insufficient representational capacity for post-training palettization. Apple sustains quality at W2 via **QAT (Quality-Aware Training)** — their shipping recipe includes quantization-aware fine-tuning, not post-training compression. Post-training W2 is not viable without QAT (days of GPU training).

### W3A16 palettization (3-bit, post-training)  —  Rejected (2026-04-13)

- File: `conversion/build_flash.py --nbits 3`
- Idea: 8 codewords per group might survive post-training where 4 did not. Weight size = 75% of INT4.
- Size: chunk total **818 MB** (vs 1,091 MB INT4 = 75%).
- Quality: **still gibberish**, though pattern differs from W2 — repetitive underscore-separated fragments instead of multilingual noise. Not usable.
- **Conclusion**: post-training palettization quality cliff is between 3-bit and 4-bit for Gemma 4 E2B. Only ≥4-bit palettization (already shipping) works without QAT. Sub-4-bit requires QAT or knowledge distillation, which is a multi-day GPU effort outside current scope.

---

## Speculative decoding experiments

### EAGLE-3  —  In training

- Files: `conversion/collect_eagle_hidden_states.py`, `conversion/download_eagle_corpus.py`, `conversion/train_eagle_draft.ipynb`, `conversion/train_eagle3_draft.ipynb` (notebooks are untracked while actively iterated)
- Idea: train a small decoder-layer draft model on Gemma 4 E2B hidden states; verify in-graph against the target.
- Training corpus: WikiText + C4 + Alpaca + Dolly + CodeAlpaca + UltraChat, formatted with Gemma 4's chat template. ~50 k samples.
- Current acceptance metric (acc0) tracked in notebook; see `docs/SPEED_8K.md §3 P1` for the latest snapshot (dated, not live).
- Integration path: once trained, draft model → `build_speculative.py` → paired with verify chunks.

### Medusa (3 heads)  —  Shelved

- Files: `conversion/train_medusa_heads.py`, parts of `conversion/build_speculative.py`
- Idea: 3 lightweight ResBlocks predict the next 3 tokens from the final hidden state; verify with target model.
- **Reason shelved**: published Medusa acceptance on Gemma-class models is ~1.3 %, far below EAGLE-3's 50–70 %. Confirmed on a small internal run. Kept because `build_speculative.py` reuses the same verify-chunk plumbing for EAGLE-3.

### MLState stateful KV cache  —  Rejected (2026-04-13)

- Files: `conversion/models/gemma4_stateful_chunks.py`, `conversion/build_stateful.py`
- Idea: use `coremltools.StateType` (iOS 18+) to declare KV cache as internal model state, eliminating per-dispatch IOSurface round-trip. Per FUNDAMENTAL_UNTRIED.md §2 / PRIORITY_ROADMAP.md Phase 1 item 2. Expected ×1.3–2.0 if dispatch overhead is the bottleneck.
- Implementation: StatefulChunk2 (L8-14) with `kv_sliding` (10,1,512,512) and `kv_full` (4,1,8192,512) as `ct.StateType`. CoreML conversion succeeds, INT4 palettized, 134 MB.
- **Result**: `error code: -14` ("Failed to build the model execution plan") on both Mac ANE and iPhone 17 Pro ANE. The `coreml_update_state` op does not generate a valid ANE execution plan.
- **Root cause**: MLState is GPU-only on current hardware/OS. HuggingFace's WWDC24 Mistral CoreML reference explicitly states stateful KV is "excellent for GPUs on Mac computers" and ANE requires "additional adaptations." Apple's own on-device Llama 3.1 and smpanaro/coreml-llm-cli both use stateless explicit-I/O KV, not MLState. The `coreml_update_state` MIL op is not supported by the ANE compiler as of iOS 26 / coremltools 9.0.
- **Conclusion**: dispatch-overhead hypothesis (FUNDAMENTAL_UNTRIED.md §0) remains valid as a bottleneck description, but MLState cannot address it on ANE. Alternative paths to reduce dispatch overhead: (1) chunk consolidation (4→2 chunks), (2) speculative decoding (amortize dispatch across multiple tokens per burst).

### MLState + KV heads padded to 32  —  Rejected (2026-04-15)

- Files: `conversion/models/gemma4_stateful_padded.py`, `conversion/build_stateful_padded.py` (on branch `worktree-agent-ad21e314`, not in main). Findings in `docs/SPLIT_ROTATE_FINDINGS.md` and `docs/SPLIT_ROTATE_BENCH.md`.
- Idea: hypothesis that error -14 was caused by the non-mod-32 `num_kv_heads=1` dim being rejected by ANE's 32-wide tile scheduler. Probe: pad KV heads 1→32 in both stateless (`PaddedKVChunk2`) and stateful (`StatefulPaddedChunk2`) variants.
- Part 1 result: PyTorch parity between `StatelessChunk2` and `PaddedKVChunk2` is bit-exact on all finite outputs; padded heads stay zero as expected.
- Part 2 result: `build_stateful_padded.py --ctx 512 --nbits 0 --smoke-test` → conversion emits the **same** `error code: -14` warning at save time; Mac CPU_ONLY predict returns finite outputs in 49 ms (graph is runtime-valid on CPU). Device retry would reproduce -14 on ANE.
- **Conclusion**: 32-alignment is **not** the cause of error -14. The ANE compiler does not schedule `coreml_update_state` regardless of the padded tensor widths. ANEMLL's `--split-rotate` is a separate multi-function-loading workaround, not an alignment fix. No further MLState-on-ANE work is warranted; any stateful path must target GPU (`CPU_AND_GPU`) as the WWDC24 Mistral demo does.

### SuffixDecoding (CPU-only draft)  —  Measured, demoted to auxiliary (2026-04-13)

- Files: `Sources/CoreMLLLM/SuffixTree.swift`, `Sources/CoreMLLLM/SuffixDecoding.swift` (on branch `claude/suffix-decoding-impl`, not merged to main)
- Idea: build suffix tree from all prior model outputs, draft K tokens via CPU trie lookup (~20µs), verify with Q=K ANE verifier. Paper reports 1.9-5.3× on chat/agentic workloads (NeurIPS 2025 Spotlight, arXiv 2411.04975).
- **T=1 instrumentation results on iPhone 17 Pro (4 multi-turn generations)**:
  - hit rate: 2% → 29% → 48% (climbs as tree grows, as expected)
  - **T1 accuracy: 18.4%** (of hits, how often the top-1 draft matches model output)
  - tree: 10k nodes after 4 generations
- **Performance overhead**: insert at end-of-generation blocks next generation start (~1-2s for 500 token sequence). Draft lookup adds ~2 tok/s overhead even at 4th-token sampling. Async insert + NSLock fixes the blocking but doesn't eliminate lookup cost.
- **Assessment**: T1=18% is too low to be the primary speculative method. EAGLE-3 (acc0=75%, workload-independent) is strictly better as the main draft source. SuffixDecoding is **workload-dependent** — the paper's high numbers come from production workloads with stable system prompts, RAG, and code editing patterns. Random diverse chat is the worst case.
- **Decision**: demote to auxiliary draft source. Use suffix tree when it has a high-confidence match (e.g., count > threshold), fall back to EAGLE-3 otherwise. Don't build Q=K verifier specifically for SuffixDecoding — build it for EAGLE-3, and SuffixDecoding can reuse it later.

### TriForce / Quest (sparse KV retrieval)  —  Planned

- See `docs/SPEED_8K.md §1 A3 / A4`. Requires block-static top-k redesign to stay on ANE. Not started.

### DuoAttention (retrieval vs streaming heads)  —  Planned

- See `docs/SPEED_8K.md §1 A1 / §3 P3`. Offline head classification + two KV banks per layer. High ROI, training-free at inference time. Next candidate after EAGLE-3 lands.

---

## MLComputePlan silent-fallback audit  —  Measured (2026-04-13)

- File: `Sources/CoreMLLLM/ComputePlanAudit.swift`
- Idea: use `MLComputePlan.deviceUsage(for:)` (iOS 17+) to walk every MIL op in chunk1-4 and identify any op whose preferred device is not Neural Engine. Per UNEXPLORED_APPROACHES_V2.md §G2 / PRIORITY_ROADMAP.md Phase 0a.
- **Result on iPhone 17 Pro (8K chunks)**:
  - **chunk1 (L0-7)**: 0 compute ops on CPU/GPU. All matmul/conv/softmax/attention on ANE.
  - **chunk2 (L8-14)**: 0 compute ops on CPU/GPU. Same.
  - **chunk3 (L15-24)**: 0 compute ops on CPU/GPU. 1 `identity` on unknown (no-op).
  - **chunk4 (L25-34 + LM head)**: **8 compute ops on CPU** — the entire `InModelArgmax` tail: `mul` (×2), `tanh`, `squeeze` (×2), `reduce_argmax`, `expand_dims`, `gather_along_axis`. Total estimated cost = 0.0028. This is the tanh-based softargmax + gather pipeline from `ane_ops.py::InModelArgmax`.
  - All `constexpr_lut_to_dense` (INT4 depalettization) and `const` ops reported as "unknown" device — these are weight-loading ops that run once at model load, not per decode step.
- **Interpretation**: the only per-step CPU fallback is chunk4's argmax tail (~8 ops, ~0.5-2 ms ANE↔CPU round-trip). chunk1-3 are 100% ANE. This confirms the dispatch-overhead hypothesis from FUNDAMENTAL_UNTRIED.md §0: the bottleneck is not individual ops falling to CPU, but the 4× per-step IOSurface round-trip between chunks.
- **Actionable**: (1) InModelArgmax CPU fallback is small but fixable — rewrite to ANE-compatible ops or accept the ~1-3% cost. (2) No other compute ops need fixing. (3) MLState (stateful KV to eliminate IOSurface round-trips) remains the highest-leverage next step.

---

## Vocabulary pruning  —  Abandoned

- File: `conversion/prune_vocab.py`
- Idea: the Gemma 4 vocab is 262 K tokens. Embedding + LM-head weights dominate on-device size. Analysis shows large blocks (rare scripts, emoji variants) are near-never used for English chat.
- **Reason abandoned**: (1) Gemma's tokenizer is sentencepiece — dropping tokens changes BPE merges and breaks round-trip tokenization; (2) the `gemma4_lite_wrapper.py` route (external per-layer embedding in Swift) already reclaimed the main memory win (~40 %). Keep the analysis as reference.

---

## 8K chunk4 rebuild  —  Shipping (maintenance fix)

- File: `conversion/rebuild_chunk4_8k.py`
- Purpose: regenerate `chunk4` with `causal_mask_full` at `(1,1,1,8192)` when the earlier build shipped with a 2048-sized mask. This is the kind of silent mismatch the `ChunkedEngine` auto-detection (see commit `4311991`) now guards against at load time.

---

## How to add a new experiment

1. Pick a short name (e.g. `wfa`, `flash`, `merged`).
2. Put the model variant in `conversion/models/gemma4_<name>.py`.
3. Put the builder in `conversion/build_<name>.py`. Keep it runnable standalone (`python build_<name>.py --output ./output/<name>`).
4. Put the A/B benchmark in `conversion/benchmark_<name>.py` comparing against the shipping baseline on the same prompts.
5. Write one row in this file: what you tried, status after the first real measurement, reason if shelved.

Rule: **every experiment that gets shelved gets a one-paragraph obituary here with the numeric reason**. Future-us will otherwise rebuild the same thing.