# n1.py Harvest Plan

Goal: move proven, no-quality-loss trainer improvements from
`C:\Users\Scott\Downloads\n1.py` into AGILLM-4 without replacing the AGILLM-4
long-context/model-scale branch.

## Ported

### 1. Exact M-Fold Expansion Attention

Status: done.

For ranks where `rank > d_k`, AGILLM-4 now computes:

```text
(q @ U) @ (k @ U).T == q @ (U @ U.T) @ k.T
```

This keeps attention scores and KV-cache keys in `d_k` width instead of
`rank` width while preserving the exact expanded-attention function. The
training path recomputes `U @ U.T` with gradients, and inference/no-grad caches
the metric until `U` changes.

Verification:

```bash
python agillm-4/verify_m_fold_agillm4.py \
  --presets pico_1x,micro_3x \
  --backends manual,sdpa \
  --cached_len 8 \
  --new_len 4
```

The verifier checks forward output, loss, input gradients, parameter gradients,
cached append equivalence, cache key width, and metric-cache invalidation.

### 2. Fused QKV Projection

Status: done.

n1 fuses separate `q/k/v` linear layers into one `qkv` linear while keeping
checkpoint compatibility by folding old state-dict keys on load. AGILLM-4 now
does the same. The parameter count and function are unchanged:

```text
[x Wq.T, x Wk.T, x Wv.T] == split(x [Wq; Wk; Wv].T)
```

Checkpoint compatibility:

- legacy `*.q.weight`, `*.k.weight`, `*.v.weight` triples load into
  `*.qkv.weight`
- warm-start shape filtering fuses legacy triples before filtering
- legacy AdamW q/k/v moment tensors are concatenated into qkv optimizer state
  when a full resume can be proven to match the old parameter layout
- if optimizer remap cannot be proven, model weights still load and optimizer
  state is reset with a warning

Verification:

```bash
python agillm-4/verify_qkv_agillm4.py \
  --presets pico_1x,micro_3x \
  --backends manual,sdpa,sublinear \
  --cached_len 8 \
  --new_len 4
```

The verifier checks fused-vs-unfused forward output, loss, input gradients,
parameter gradients, strict legacy state-dict loading, `_safe_load_any`
warm-start loading, and optimizer-state remap.

### 3. KV Cache Buffer

Status: done.

`KVBuffer` is a preallocated head-major `[B, H, T, d_k]` slab for decode-time
K/V. Callers size it once and call `append(k_new, v_new)` per step instead of
re-running `torch.cat` and reallocating the cache tensor every token.
`TuneableAttentionMHA.forward` accepts either a legacy `(k, v)` tuple cache
(no behaviour change) or a `KVBuffer` instance, and returns the same cache
object it received so the caller can reuse it across steps. Overflow raises;
callers must size the buffer to the maximum decode length.

Verification:

```bash
python /workspace/agillm-4/verify_anchor_memory_and_kv_buffer_agillm4.py
```

The verifier exercises buffer fill, slice view, overflow handling, and confirms
the legacy tuple-cache code path still works (regression test).

## Next Candidates

### 3. Combined ALiBi + Mask Cache

n1 pre-folds ALiBi into the mask once per encoder forward instead of rebuilding
the same layer-independent bias in every block.

Risk: cache semantics differ for KV decode where the ALiBi slice changes.

### 4. SAT Speculative Inference

n1 has proof-covered SAT-draft / AR-verify speculative decoding. This belongs
in AGILLM-4 after the SFT result tells us whether chat turns are sane.

Risk: inference control flow and cache rollback complexity.

### 5. Compact Checkpoint

n1 can compact `U` spectra post-training and save compatible checkpoints.

Risk: optimizer state must be dropped or remapped carefully; do only as a
separate command, never during a live training run.

## Rule

Every harvested feature needs its own AGILLM-4 verifier or profile artifact.
Do not rely on n1's proof suite alone after adapting the implementation.