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:

(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:

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:

[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:

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:

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.