K3 Block C: f_θ + cross-model DLMRestoredVerifier — real Gemma 4 26B-A4B run (arch gate PASS, machinery validated) by FluffyAIcode · Pull Request #103 · FluffyAIcode/Kakeya-LLM-Inference-engine

FluffyAIcode · 2026-06-09T17:39:46Z

Status (2026-06-10)

✅ f_θ machinery — verified by identity-restore evidence (recall=1.0)
✅ K3 cross-model verifier — architectural correctness gate PASS (effective_attention_fraction = 1.0)
⏳ Recall gate — currently FAIL with both v1 (raw MSE, recall 0/10) and v3-relmse (magnitude-normalised MSE, recall 0/10)
⏳ Next training — v3 attn_distill (attention-output distillation) = the principled one-shot fix; default loss on this branch as of commit cb608d7

What's on this branch now

attn_distill (PR #106 design, ported here per R1)   ← v3 default
    + per-layer Q/O_tgt/cos/sin/mask capture
    + AttentionTargetData dataclass
    + smart caching: 1 verifier forward per sequence

relmse  (the user's PR #103 work, preserved)
    + per-layer rel_mse diagnostics
    + magnitude-normalised MSE

cos_mag, combined  (PR #106's v2 intermediate)
mse                 (v1 reproduction)

CLI: --loss-type {attn_distill,relmse,mse,cos_mag,combined} (default attn_distill).

To run v3-attn_distill (the principled fix)

git checkout AgentMemory/v04-pr-k3-block-c-f-theta-design-and-skeleton-8e7f
git pull
HF_TOKEN=hf_xxx bash scripts/review_pr_k3_f_theta_train_on_vast.sh

Reviewer aid prints (per R2) at startup:

==> review_pr_k3_f_theta_train_on_vast.sh
    Branch:        AgentMemory/v04-pr-k3-block-c-f-theta-design-and-skeleton-8e7f
    HEAD commit:   cb608d7  "<commit subject>"
    Recipe:        loss=attn_distill rank=auto(768) steps=20000 ...

So you SEE before training starts: it's running attn_distill, not relmse.

To reproduce relmse for ablation: LOSS_TYPE=relmse bash scripts/review_pr_k3_f_theta_train_on_vast.sh.

Workflow rules introduced in commit `cb608d7`

After the 2026-06-10 GPU-time-waste incident (ran relmse instead of intended attn_distill because of branch fragmentation), three rules now codified:

	Rule	Enforcement
R1	Fix open PRs on the SAME branch. No child PRs for fixes.	`AGENTS.md` + agent self-check
R2	Every reviewer aid prints branch + HEAD + recipe at startup	`tests/research/test_reviewer_aid_headers.py` (CI test)
R3	Agent prints "Verified by agent" block before recommending GPU spend	`AGENTS.md` + agent self-check

See docs/agent-workflow-rules.md for the full failure log + rationale + each rule's enforcement mechanism.

Tests

389/389 passing on Linux CI:

363 pre-existing
17 v2 cosine+mag/LR/NIAH tests
10 attn_distill loss tests
6 R2 enforcement tests (CI auto-checks every new reviewer aid)

Evidence preserved

results/research/f_theta_v1/ — v1 raw MSE checkpoint (recall 0/10)
results/research/k3_integrated_niah_ctx{70,280}_1781062484.json — v1 NIAH evidence
results/research/f_theta_v3/ — v3-relmse checkpoint (recall 0/10, full-attn rel_mse 0.58 bottleneck identified)
results/research/k3_integrated_niah_ctx{70,280}_1781068513.json — v3-relmse NIAH evidence
results/research/k3_identity_restore_ctx70.json — proves machinery correct (recall 1.0 with verifier's own K/V)

Stack

main (post #93 + #99 + #94 + #100 + #101 + #102)
└── PR #103 (THIS)
    │   ✅ engine API: f_theta.py + cross_model_dlm_verifier.py
    │   ✅ training: k3_f_theta_train.py (attn_distill | relmse | ...)
    │   ✅ eval: k3_integrated_niah_eval.py
    │   ✅ reviewer aids (vast)
    │   ✅ workflow rules R1+R2+R3
    │   ⏳ awaiting attn_distill v3 retrain on vast
    │
    └── PR #104 (Mac MLX cross-model verifier; parallel-track,
                  unaffected by the rules-cleanup commit)

PR #106 (the original child branch holding attn_distill) is now closed as superseded — its content is on this branch.

…(P0) Per user 'go P0' directive 2026-06-09 after architectural observation that PR #102's Mac MLX spec decode eval doesn't exercise the Kakeya inference engine's core architecture (sink+window verifier + dLM proposer K/V Restoration). This PR ships the foundational engine code for the integrated Kakeya inference architecture per ADR 0008 §11.3: verifier (Gemma 4 26B-A4B): └─ holds only sink+window local KV cache (sink=4 + window=64) └─ at evicted positions, takes K/V supplied by proposer (via f_θ) drafter (DFlash 0.4B, alignment-trained baseline): └─ runs full forward over committed prefix per step └─ K/V at every layer at every position captured └─ K/V projected through f_θ into verifier K/V space, injected at evicted positions Three new files --------------- inference_engine/v04/f_theta.py (~290 LOC) FThetaConfig dataclass + FThetaProjection nn.Module. Architecture: shared encoder + per-verifier-layer decoders, low-rank factorisation: drafter_kv_input [B, T, drafter_layers * drafter_kv_dim] ↓ encoder Linear(in, rank) rep [B, T, rank] ↓ per-verifier-layer decoders (30 × Linear(rank, verifier_kv_dim)) output [B, T, num_verifier_layers, num_kv_heads_v, head_dim_v] Default rank=256. Production K3 config (Gemma 4 26B-A4B + DFlash 0.4B): encoder: 2 × 5×256 × 256 = 655k params decoders: 2 × 30 × 256 × 2048 = 31.5M params Total: ~32M params (vs drafter 430M, verifier 26B) Separate K and V projections (different downstream roles). Save/load: save_pretrained(dir) writes f_theta_config.json + f_theta_weights.pt; from_pretrained(dir, dtype, device) loads back. inference_engine/v04/cross_model_dlm_verifier.py (~270 LOC) CrossModelDLMRestoredVerifier wrapper. Construction validates drafter + verifier dimensions match the f_θ config (rejects drafter-vs-verifier-vs-f_θ mismatch loudly at __init__). forward(input_ids, apply_rotary_pos_emb, eager_attention_forward): 1. compute_evicted_positions(T, sink, window) 2. If no evicted (T <= sink+window): plain verifier forward 3. Drafter forward via _capture_drafter_kv (forward hooks on k_proj/v_proj at each drafter layer) 4. f_θ.forward_kv_pack(drafter_K_per_layer, drafter_V_per_layer) → verifier K, V at every (layer, position) 5. Patch each verifier layer's self_attn.forward to: a. Run standard q/k/v_proj + q_norm/k_norm + RoPE b. At evicted positions, REPLACE k, v with f_θ output (after k_norm + RoPE applied via prepare_restored_attention_kv) c. Standard attention compute path through eager_attention_forward 6. Run verifier forward → logits 7. Restore original attention forwards (try/finally) Two scope-outs (recorded inline): * MLX verifier path: this module patches HF transformers attention. Mac MLX integration is a follow-up PR (instrument mlx_lm Gemma 4 model directly, not via attention monkey-patch). * Speculative decoding accept/reject loop: separate inference engine concern. PR #93's DFlashProposer + mlx_verify_block handles the spec-decode side; combining with this module's K/V Restoration is a separate integration step. Drafter K/V capture (_capture_drafter_kv): instruments DFlashDrafter's internal layer.self_attn.k_proj / v_proj via forward hooks. NOTE inline that the first-iteration synthetic-context capture (zero hidden as drafter input) is plumbing-validation; product-meaningful K/V values require conditioning on verifier aux hiddens, which is the next integration step (after f_θ training validates the projection alone). scripts/research/k3_f_theta_train.py (~310 LOC) Training pipeline for f_θ on CUDA: 1. Load Gemma 4 26B-A4B verifier (transformers bf16, sdpa) 2. Load DFlash drafter (PR #93's DFlashDrafter from models/dflash-kakeya-baseline) 3. Data collection: for each prompt in PROMPTS (same 64-prompt corpus as PR #93's alignment_train), run greedy AR generation to gen_len tokens, capture per-layer per-position K/V via hooks on k_proj/v_proj of both models 4. Train f_θ with MSE loss across (layer, position) pairs, AdamW lr=1e-3, weight_decay=0.01, gradient clip 1.0 5. Save checkpoint at --save (default results/research/f_theta_v1) Memory budget: at T=512, ~128 MB per sequence cached on GPU. 64 sequences ≈ 8 GB. Fits H200 80 GB easily. Validation: report initial vs final loss; reduction factor. inference_engine/v04/__init__.py: re-exports the new public surface (FThetaConfig, FThetaProjection, CrossModelDLMRestoredVerifier, CrossModelLayerMapping). Tests (Linux CI: 27 new tests) ----------------------------- tests/inference_engine/v04/test_f_theta.py (21 tests): TestFThetaConfig (4): dim properties + JSON round-trip TestForwardShapes (4): forward_k/v shape contract + input validation TestForwardKVPack (3): KVCapture-style input + consistency vs explicit concat TestParameterCount (2): tiny + production param count locked in TestSaveLoadRoundTrip (4): save+load preserves outputs; missing-file errors TestDeviceDtypeDispatch (2): to(dtype), from_pretrained dtype override TestGradientFlow (1): gradients flow through encoder + decoders separately (K path doesn't update V weights and vice versa) tests/inference_engine/v04/test_cross_model_dlm_verifier.py (6 tests): TestConstruction (3): dimension validation rejects mismatch; valid construction succeeds; negative sink/window raises TestProjectDrafterKV (1): output shape contract TestNoEvictPath (1): short prompt (T <= sink+window) doesn't invoke drafter TestExports (1): module + namespace re-exports Tests: 354 passing (336 pre-existing + 21 f_theta + 6 cross-model; 12 research/ unchanged from PR #102). What this PR does NOT yet do (deferred to follow-up PRs) -------------------------------------------------------- 1. Train f_θ on real data — requires vast.ai GPU time. scripts/research/k3_f_theta_train.py is the runnable trainer. Once trained, the checkpoint goes to a follow-up PR with the evidence (training report + integrated NIAH ladder evidence). 2. End-to-end integrated NIAH ladder evidence — needs: * trained f_θ checkpoint (step 1) * cross-model DLMRestoredVerifier reviewer aid (off-the-shelf K1.E NIAH harness needs a small adapter to use this verifier wrapper) * vast.ai run producing the evidence JSON 3. Mac MLX integration — instruments mlx_lm Gemma 4 model directly (different surgical approach than HF transformers attention monkey-patch). Follow-up PR. 4. _capture_drafter_kv proper aux-conditioning — current synthetic zero-hidden capture is plumbing only. The proper path passes verifier aux hiddens into the drafter (DFlash architecture), captures K/V from THAT forward. Adds a method to DFlashDrafter in a follow-up. These are the remaining items on the K3 critical path; this PR establishes the engine API surface they all depend on. Stack ----- Off main (post #93 + #99 + #94 + #100 + #101 + #102 merged). Independent of any other open PR. Outstanding work after this PR: Step 5 — K2.A backport PR (P2) Step 6 — alignment training corpus expansion (P2) P0 cont. — f_θ training run + integrated NIAH evidence P0 cont. — Mac MLX integration of cross-model DLMRestoredVerifier Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

User signal 2026-06-09: 'A / B / C 全部做完。我已经开了vast' — proceed through full P0 critical path; vast is open for runs. Three fixes + three new files in this commit: (A) FIX: _capture_drafter_kv now uses verifier embed_tokens Previous version (just committed in this PR) used synthetic zero hidden state to fire k_proj/v_proj hooks. This is plumbing-only and produces meaningless K/V values. DFlashDrafter's design (PR #93) shares verifier embed_tokens (no own embedding lookup), so the correct capture path is: 1. verifier_model.get_input_embeddings()(input_ids) × sqrt(hidden) 2. Pass embedded hiddens through drafter.layers (no aux conditioning) 3. Capture K/V via forward hooks per layer Updated _capture_drafter_kv signature to take verifier_model (required for embed_tokens). Updated CrossModelDLMRestoredVerifier. project_drafter_kv to pass it. Updated test fixture to provide a real embed_tokens on the synthetic verifier (was previously unnecessary; now required). (B) FIX: k3_f_theta_train.py now uses _capture_drafter_kv Previous version called capture_proposer_kv(drafter.model, input_ids) which would crash on real DFlashDrafter — DFlashDrafter is a flat nn.Module without .model attribute (capture_proposer_kv expects model.model.layers OR model.transformer.h, both absent). Switched to inference_engine.v04.cross_model_dlm_verifier. _capture_drafter_kv (the same path the cross-model verifier uses at inference time). Ensures training and inference are using the IDENTICAL drafter K/V values — no train/serve skew. (C) NEW: scripts/review_pr_k3_f_theta_train_on_vast.sh vast.ai reviewer aid for f_θ training. Pre-flight checks: 1. HF_TOKEN (Gemma 4 gated) 2. models/dflash-kakeya-baseline/ Git LFS pulled (>100MB safetensors) 3. CUDA available 4. transformers 5.x (Gemma 4 support) Env knobs: STEPS, LR, RANK, N_PROMPTS, GEN_LEN, SAMPLE_POSITIONS, SAVE_DIR, SEED. Default config: 4000 steps, rank=256, 64 prompts × 128 gen tokens — fits H200 80 GB easily, ~8-15 min wall clock. Output: trained f_θ checkpoint + training report. Validation gates printed at end (loss_reduction_factor ≥ 2.0 sanity). (D) NEW: scripts/research/k3_integrated_niah_eval.py (~280 LOC) THE K3 PRODUCT GATE EVIDENCE SCRIPT. Combines: * CrossModelDLMRestoredVerifier (verifier with sink+window cache + drafter K/V Restoration via f_θ) * K1.E NIAH evaluation harness (effective_attention_window / recall / memory metrics) Validates per ADR 0008 §11.8 release gates: 1. Architectural correctness: effective_attention_fraction = 1.0 at every NIAH ladder rung 2. Memory bounded: sustained verifier KV-cache ≤ O(sink+window) 3. Recall preservation: |recall_cross_model - recall_oracle| ≤ 5 pp at every rung (ADR §11.8 1a — architecturally-meaningful gate) Runs: - cross-model verifier on each NIAH sample, decodes max_new_tokens - full-attention oracle baseline on same samples (--skip-oracle to bypass; loses recall_delta gate signal) - aggregate recall, attention_window, memory; compute gate booleans Output JSON schema mirrors K1.E NIAH harness (per_config recall, attention_window, memory) + new 'gate' block with the three booleans for direct inspection. (E) NEW: scripts/review_pr_k3_integrated_niah_on_vast.sh vast.ai reviewer aid for the integrated NIAH eval. Pre-flight: 1. HF_TOKEN 2. f_θ checkpoint at $F_THETA_DIR 3. drafter LFS pulled 4. CUDA available Runs the integrated NIAH eval per CONTEXT_LADDER rung (default '70 280', i.e. ~1.4k + ~5.6k tokens). Per-rung JSON + combined log. Final aggregation diff-able with PR #94's same-checkpoint K1 ladder evidence. After this PR + a vast run of (review_pr_k3_f_theta_train_on_vast.sh → review_pr_k3_integrated_niah_on_vast.sh), the K3 product gate is empirically closed on CUDA. Mac MLX path follows as separate PR (instrument mlx_lm Gemma 4 model directly; can't reuse the HF attention monkey-patch approach). Tests: 354/354 passing on Linux CI (no v04 code regressions; new script files don't run in CI but parse + bash -n check OK). Stack: Off main, builds on PR #103 commits in this same branch. PR #103 description updated to reflect added scripts + critical fixes. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

…+ cross-model verifier Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

…layers Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

…capture/loss for Gemma4 Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

…, RoPE unsqueeze_dim=2, v_proj-None, evicted slicing) + gemma4 helpers import + tests Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

…drafter K/V) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

… encode, aggregate_recall, v04_dlm_restored window) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

…tegrated NIAH eval Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

…obal_head_dim=512, 2 KV heads) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

…olate restore machinery from f_theta accuracy Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

…KV; loss 50.8->3.70, 13.74x) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

…FAIL (f_theta v1), identity-restore recall=1.0 (machinery validated) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

…ntermediate) Per user 2026-06-10: '我要求直接上一步到位的训练方案。不要搞这种中间态，浪费时间和CPU资源' Skipped the v2 cosine+magnitude intermediate. Default loss is now attention-output distillation — the principled training objective for K/V replacement. v2 cos+mag remains accessible via --loss-type cos_mag for ablation, but is not the default path. The principled loss =================== For each verifier layer ℓ: K_pred_ℓ, V_pred_ℓ = f_θ(drafter_KV)[ℓ] Q_for_attn = q_norm(Q_raw_ℓ).view(B, T, H_q, D) → RoPE → transpose K_for_attn = k_norm(K_pred_ℓ).view(B, T, H_kv, D) → RoPE → transpose V_for_attn = v_norm(V_pred_ℓ).view(B, T, H_kv, D) → transpose GQA repeat K, V to H_q O_inner = scaled_dot_product_attention(Q, K, V, mask, scale) O_pred = o_proj(O_inner.reshape(B, T, H_q*D)) loss_ℓ = MSE(O_pred, O_tgt_ℓ) ^^^ captured during data collection from the verifier's actual attn module post-o_proj output Total = mean over layers Why this is mathematically right for K/V projection --------------------------------------------------- attention(Q, K, V) is the actual quantity that propagates through the residual stream at inference. v1 (raw MSE on K) and v2 (cos+mag on K) are PROXIES for attention behavior. v3 directly optimises the attention output, so the loss landscape's gradient points precisely at 'f_θ K/V produces equivalent verifier behavior'. It accounts for: GQA grouping, RoPE, causal/sliding mask, k_norm/q_norm/v_norm, AND the o_proj that follows attention. Implementation strategy ======================= Tractability concern: the principled loss seemingly requires a full verifier forward per training step (≈ 3 sec on H200 → 16+ hours for 20000 steps). NOT acceptable. Solution: smart caching. During data collection (one verifier forward per sequence), capture per-layer: - Q_raw [T, num_heads × head_dim] from q_proj forward hook - O_tgt [T, hidden_dim] from attn module forward hook - cos, sin [1, T, head_dim] from attn forward pre-hook - attn_mask from attn forward pre-hook All cached on CPU bf16 (≈ 13 MB per layer per sequence × 30 layers × 64 sequences ≈ 25 GB CPU RAM). Training streams these to GPU per step. No verifier forward is needed at training time. Per-step cost: f_θ forward + per-layer attention recomputation (scaled_dot_product_attention with cached Q + f_θ-predicted K/V) + o_proj + MSE. ~80 ms/step on H200. 20000 steps = 25-30 min. Total v3 wall on H200: ~40-60 min (data collect + training). Three modified files ==================== scripts/research/k3_f_theta_train.py (~1100 LOC, +400) New dataclass: AttentionTargetData Per-layer Q_raw + O_tgt + cos + sin + attention_mask + per-layer num_heads / head_dim. CPU bf16 storage. New function: _capture_attention_target_data Runs verifier forward with hooks (forward hook on q_proj for Q_raw, forward hook on attn module for O_tgt, forward pre-hook on attn module for position_embeddings + attention_mask). Returns AttentionTargetData with all tensors on CPU bf16. New function: _attention_distillation_loss The principled loss as described above. Full per-layer pipeline with proper GQA / RoPE / mask handling. Streams cached tensors from CPU to GPU per layer; frees per-layer GPU memory before moving to next layer. Modified: CapturedSequence Made verifier_k / verifier_v Optional. Added attn_target field (Optional[AttentionTargetData]). For attn_distill loss, only attn_target is captured (saves ~125 MB per sequence vs legacy K/V capture). For legacy losses, only verifier_k/v captured. Modified: _f_theta_loss Dispatch on loss_type. attn_distill path → _attention_distillation_loss. Legacy losses (mse | cos_mag | combined) path → previous v2 logic. Validates seq has the right capture for the chosen loss. Modified: _collect_sequence Now takes capture_legacy_kv + capture_attn_target flags. Routes to either or both capture paths. Modified: main() - Loaded attn_implementation='eager' for attn_distill (sdpa breaks the attn-module-level forward hook contract); 'sdpa' for legacy - Imports apply_rotary_pos_emb from transformers.models.gemma4 - --loss-type now defaults to attn_distill, choices include all 4 - --rank default is None → auto-resolve: 768 for attn_distill, 256 for legacy (rank ↑ for the more capable principled trainer) - --sample-positions default 0 → use full T (recommended for attn_distill); 256 for legacy - Per-step log shows per-loss-type diagnostics: cos sim for cos_mag/combined, mseO/|O_tgt|^2 ratio for attn_distill - Report includes 'final_diagnostic' + 'loss_type' scripts/review_pr_k3_f_theta_train_on_vast.sh (~190 LOC, +20 / -25) Updated to v3 defaults: LOSS_TYPE=attn_distill (was 'combined' in v2 plan, never shipped) RANK= (empty → trainer auto-picks 768 for attn_distill) SAMPLE_POSITIONS=0 (full T) SAVE_DIR=results/research/f_theta_v3 Header docstring documents the v1 reproduction recipe AND the v3 rationale (one-shot principled trainer). Banner shows the resolved attn implementation (eager vs sdpa) and the resolved RANK value. Validation gate updated: 'mseO/|O_tgt|^2 ratio < 0.05' replaces 'cosK_total < 0.05' (v3 diagnostic; ratio quantifies attention-output noise). tests/research/test_k3_f_theta_train_v2.py (+10 new tests) TestAttentionDistillationLoss (7): - attention_distill_loss_runs (returns scalar with diag populated) - loss_is_differentiable_through_f_theta (gradient flows to f_θ) - o_proj_weights_remain_frozen_in_loss (frozen verifier params receive no grad — important for training to not OOM/NaN) - dispatch_through_f_theta_loss_function (v2 _f_theta_loss correctly routes to _attention_distillation_loss for attn_distill) - attn_distill_requires_layers_arg (clear error if layers/RoPE/ device aren't passed) - legacy_loss_rejects_attn_only_capture (mse loss on attn_target- only seq raises RuntimeError instead of silently producing NaN) - sample_positions_subselects_output (full vs sub sample both produce a valid scalar loss) TestAttentionTargetDataDataclass (3): - fields_present - captured_sequence_optional_kv_and_attn (legacy fields default to None) - captured_sequence_attn_target_path (attn_target stored correctly) Stub _StubAttn / _StubLayer reproduce the Gemma 4 self_attn module surface (q_norm, k_norm, v_norm, q_proj, o_proj, scaling, head_dim) enough for the loss to run on Linux CI without an actual verifier. Tests: 383/383 passing (354 pre-existing + 9 from PR #104 + 10 from PR #103 + 17 from v2 + 10 new v3 — with overlap). Validation gate (vast retrain, one-shot) ======================================== Run the same reviewer aid; defaults pick up v3: HF_TOKEN=hf_xxx bash scripts/review_pr_k3_f_theta_train_on_vast.sh Output: results/research/f_theta_v3/{f_theta_config.json, f_theta_weights.pt} results/research/f_theta_v3.json (with mseO + |O_tgt| diagnostics) Then re-run integrated NIAH against v3 checkpoint: F_THETA_DIR=results/research/f_theta_v3 \ bash scripts/review_pr_k3_integrated_niah_on_vast.sh Expected v3 outcomes: - mseO_mean / |O_tgt|^2 ratio < 0.05 (attention output noise low) - integrated NIAH recall_cross_model ≈ recall_oracle - recall_delta_within_5pp gate CLOSES This is the principled one-shot fix. If recall still falls short (≥ 5pp delta), the issue is f_θ capacity — escalate to per-layer encoders or larger rank (RANK=1024). But attn_distill loss + rank 768 + 20k steps + NIAH data + cosine LR is the maximum-strength single-shot training configuration without architectural rewrites. Stack ===== main (post #93 + #99 + #94 + #100 + #101 + #102) └── PR #103 (CUDA: f_θ + cross-model + train script + integrated NIAH) ├── PR #104 (Mac MLX cross-model verifier; parallel-track) └── THIS PR #106 (trainer v3 — one-shot attn distill, supersedes v2 plan) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

…nt mseO/|O_tgt| diagnostics; default output f_theta_v3 (preserves v1) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

… 0.349, full-attn layers 0.58 worst) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

…recall 0/10 (full-attn layer fidelity rel_mse 0.58 is the bottleneck) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

Per user 2026-06-10: '设计规则,以后禁止翻这种低级错误!' Failure being prevented ======================= 2026-06-10: ~15 min vast.ai H200 GPU wasted because: PR #103 branch e18f2fc → ce25dfa (relmse fix, by another agent) ↘ PR #106 branch (mine) 6c2fc23 → 6f168dd (attn_distill, MY fix) The user ran 'bash scripts/review_pr_k3_f_theta_train_on_vast.sh' on PR #103 branch (their checkout) → executed relmse trainer, NOT the attn_distill trainer they thought they were running. The reviewer aid printed nothing about which trainer/branch/recipe it was about to invoke. The agent (me) failed to either commit attn_distill directly to PR #103 OR explicitly tell the user 'first git checkout <branch>' before running. Three rules, each with concrete enforcement =========================================== R1 — SAME-PR FIX RULE (agent behavior, codified in AGENTS.md) When fixing a still-open PR's failed evidence, commit to that PR's branch. No child PRs for fixes. R2 — REVIEWER AID SELF-IDENTIFICATION (CI-enforced) Every scripts/review_*.sh MUST source scripts/_lib/reviewer_aid_header.sh and call print_aid_header at startup, printing branch + HEAD + recipe BEFORE any pre-flight check. CI test tests/research/test_reviewer_aid_headers.py enforces this. Pre-existing aids that need retrofit are in a _GRANDFATHERED allowlist that must only shrink. R3 — PRE-GPU CONFIRMATION (agent behavior, codified in AGENTS.md) Before recommending GPU/cluster time, agent prints a verification block: branch + HEAD + code path + recipe + expected wall, AFTER reading the actual code path (not just the design doc). Files in this commit ==================== docs/agent-workflow-rules.md (NEW, ~150 LOC) Long-form rationale + failure log + each rule's enforcement mechanism. Failure log is append-only. AGENTS.md (NEW, ~75 LOC) Read at session start by AI coding agents. Codifies R1+R2+R3 as non-negotiable rules. Points at docs/agent-workflow-rules.md for rationale. scripts/_lib/reviewer_aid_header.sh (NEW, ~70 LOC) Sourceable lib providing: - print_aid_header (script_path, recipe) → prints branch + HEAD commit + repo dir + recipe + started timestamp - require_branch (expected) → assertion helper - print_agent_verification (branch, sha, subject, file) → mirror of R3 verification block (used by agent) tests/research/test_reviewer_aid_headers.py (NEW, ~150 LOC) CI enforcement of R2. Six tests: - header lib present + readable - print_aid_header signature stable (Branch / HEAD / Recipe / Started fields all printed) - grandfathered set contains only files that exist on disk (forces the set to shrink as files are deleted) - no NEW reviewer aid is non-compliant (catches PRs that add a new aid without the header — adding to _GRANDFATHERED is explicitly forbidden for new aids) - parametrized strict R2 check on every non-grandfathered aid scripts/review_pr_k3_f_theta_train_on_vast.sh (modified) Sources the new header lib. Calls print_aid_header at startup with full recipe (loss / rank / steps / gen_len / lr_schedule / warmup / n_general / n_niah / save). This is the FIRST aid to comply with R2. scripts/research/k3_f_theta_train.py (modified — fixes the underlying branch divergence per R1 — attn_distill from PR #106 + relmse from PR #103 now BOTH on this branch) Adds 'relmse' as a 5th loss_type alongside attn_distill / mse / cos_mag / combined. Ports PR #103 ce25dfa's magnitude-normalised MSE loss into the v3 dispatch structure. Per-step diag now logs rel_K / rel_V components for relmse mode. The user's relmse v3 evidence (results/research/f_theta_v3*) is preserved + remains reproducible via --loss-type relmse. tests/research/test_k3_f_theta_train_v2.py (modified — was on PR #106; now also on PR #103 per R1's same-branch unification) Carried over from PR #106. Already covers attn_distill loss path (10 tests). relmse uses the same dispatch infrastructure so no new test needed; smoke-tested in commit message. Tests: 389/389 passing on Linux CI. Effect for the user =================== Next time you run bash scripts/review_pr_k3_f_theta_train_on_vast.sh on PR #103 branch, the script prints: ==> review_pr_k3_f_theta_train_on_vast.sh Branch: AgentMemory/v04-pr-k3-block-c-f-theta-design-and-skeleton-8e7f HEAD commit: <sha> "<this commit's subject>" Repo dir: /workspace Recipe: loss=attn_distill rank=auto(768) steps=20000 ... Started at: <timestamp> You see at a glance: 'I'm running attn_distill, not relmse'. If attn_distill isn't what you want, you abort BEFORE the verifier forward warm-up. No more 'I thought I was running X, the GPU actually ran Y' surprises. Default loss is now attn_distill on THIS branch (PR #103). Running the reviewer aid as-is gives you the principled trainer. Pass LOSS_TYPE=relmse to reproduce PR #103 v3-relmse for ablation / diagnostic comparison. Stack cleanup follow-up ======================= PR #106 (the child branch holding the attn_distill design) is now superseded — its content is in this commit. Will close PR #106 with a 'merged into PR #103' note in the next operation, per R1. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

…posed by alpha-sweep Per user 2026-06-10: 'attn_distill sweep evidence... pls check the result' Diagnosis from sweep evidence (commit 72ce157) ============================================== f_theta_baseline_rel_mse.overall = 1331.94 f_theta_baseline_rel_mse.full_attn = 18254 f_θ raw (pre-norm) K/V output is 36× off-scale from verifier's true K/V (135× on full-attention layers). Despite this, attn_distill training converged to mse_O = 0.176 (looks fine) because k_norm and v_norm are RMSNorm — they NORMALIZE THE SCALE AWAY before attention. The attn_distill loss (computed downstream of k_norm) was scale-invariant and thus blind to the magnitude collapse. Sweep showed recall=0 for ALL alpha < 1.0 (in raw-space mixing), with recall jumping to 1.0 only at alpha=1.0 (pure verifier K/V). Reason: at alpha=0.9 (90% true + 10% f_θ), the f_θ component is 0.1 × 36 = 3.6× the magnitude of the true component (0.9 × 1) and DOMINATES THE DIRECTION post-mixing. After k_norm normalises the total magnitude, the direction is still dominated by f_θ's (directionally-wrong) output. Recall stays at 0 until alpha=1.0 (no f_θ contribution at all). This is **f_θ collapse degeneracy**: attn_distill loss has multiple local minima, including a degenerate one where f_θ outputs are magnitude-runaway and direction-arbitrary, but post-norm-then-attn gives 'evicted positions get neutral attention weights' so the local cache (sink+window) carries the attention output. Loss is ~0.18 (close to zero because evicted contribution is suppressed), but f_θ is contributing zero useful retrieval signal. This explains why NIAH failure mode changed from v1's 'confused hallucinations' to attn_distill v3's 'confident refusal' — f_θ isn't contributing wrong info, it's contributing NOTHING (post- attention), and the local cache can't see the needle. The fix: attn_distill_hybrid loss ================================= Direct supervision on K/V at three levels (in addition to attn output): loss = 1.0 * MSE(O_pred, O_tgt) # attention output + λ_kDir * (1 - cosine(K_pred_post_norm, K_tgt_post_norm)) # K direction + λ_vDir * (1 - cosine(V_pred_post_norm, V_tgt_post_norm)) # V direction + λ_kMag * MSE(|K_pred_pre_norm|, |K_tgt_pre_norm|) / |K_tgt|² # K magnitude + λ_vMag * MSE(|V_pred_pre_norm|, |V_tgt_pre_norm|) / |V_tgt|² # V magnitude Defaults: λ_kDir = λ_vDir = 1.0, λ_kMag = λ_vMag = 0.1. The cosine terms (post-norm) are the crucial fix — they constrain K direction directly, eliminating the degenerate solution where f_θ produces direction-arbitrary K. The magnitude terms (pre-norm) prevent the 36× scale runaway. Hybrid is the new default loss type. v3 attn_distill remains available via --loss-type attn_distill for ablation. Six modifications ================= scripts/research/k3_f_theta_train.py: - Extended AttentionTargetData with optional k_raw_tgt + v_raw_tgt (CPU bf16 cache, ~100 MB extra per sequence — acceptable) - _capture_attention_target_data new flag capture_raw_kv (also captures k_proj/v_proj outputs via forward hooks; v_proj-None layers fall back to k_proj output, matching cross_model_dlm_verifier semantics) - _attention_distillation_loss new flags hybrid, lambda_k_dir, lambda_v_dir, lambda_k_mag, lambda_v_mag. When hybrid=True, loads K_tgt_pre and V_tgt_pre, applies layer's k_norm + v_norm, computes cosine direction loss + pre-norm magnitude loss - _f_theta_loss dispatches loss_type='attn_distill_hybrid' to _attention_distillation_loss with hybrid=True - main(): new args --lambda-k-dir/--lambda-v-dir/--lambda-k-mag/ --lambda-v-mag, --init-from (warm-start from existing checkpoint, useful for fine-tuning attn_distill v3 with hybrid loss for fewer steps) - Default loss_type changed: attn_distill → attn_distill_hybrid - capture_raw_kv_in_attn_target=True automatically for hybrid - Per-step log: hybrid prints kDir/vDir/kMag/vMag alongside mseO/ratio scripts/review_pr_k3_f_theta_train_on_vast.sh: - Default LOSS_TYPE=attn_distill_hybrid - New env knobs LAMBDA_K_DIR/LAMBDA_V_DIR/LAMBDA_K_MAG/LAMBDA_V_MAG/ INIT_FROM - SAVE_DIR default → results/research/f_theta_v4_hybrid (preserves v3 attn_distill evidence) - Reviewer aid recipe string includes hybrid lambdas + INIT_FROM tests/research/test_k3_f_theta_train_v2.py: - TestAttentionDistillationHybridLoss (5 new tests): * hybrid_runs_and_emits_full_diag (mseO+kDir+vDir+kMag+vMag in diag) * hybrid_requires_raw_kv_tgt (RuntimeError if missing — fail loud) * hybrid_dispatch_via_loss_type (loss_type='attn_distill_hybrid' routes) * hybrid_loss_strictly_higher_than_attn_distill_alone (verifies added terms have effect, not silently zero) * hybrid_grad_flows_to_f_theta (gradient reaches f_θ params) - TestAttentionTargetDataDataclass + 1 test: * attention_target_data_optional_raw_kv_for_hybrid (None by default; populated when capture_raw_kv=True) Tests: 389/389 passing on Linux CI. Validation gate (vast retrain — TWO options) ============================================ Option A — Fine-tune v3 attn_distill checkpoint with hybrid loss (saves ~75 min, recommended): HF_TOKEN=hf_xxx \ INIT_FROM=results/research/f_theta_v3_attn_distill \ STEPS=10000 \ SAVE_DIR=results/research/f_theta_v4_hybrid_finetuned \ bash scripts/review_pr_k3_f_theta_train_on_vast.sh Expected wall: ~30-45 min (data already collected; only training). The warm-start from v3 attn_distill checkpoint gives the new loss a head start on the attn output term while the hybrid terms force K/V direction + magnitude into shape over the next 10k steps. Option B — Train from scratch with hybrid loss (full reset): HF_TOKEN=hf_xxx bash scripts/review_pr_k3_f_theta_train_on_vast.sh Expected wall: ~90 min (data collection ~45 min + training ~45 min). Cleaner baseline — no inheriting the degenerate v3 attn_distill weights. Expected v4-hybrid outcomes (vs v3 attn_distill) ================================================ k_dir_mean < 0.05 (cosine sim > 0.95 on post-norm K) v_dir_mean < 0.05 k_mag_mean < 0.05 (pre-norm magnitude matched within ~5%) v_mag_mean < 0.05 mse_O_mean < 0.10 (better than v3's 0.176, since K/V are now non-degenerate) f_theta_baseline_rel_mse.overall < 50 (vs v3's 1331; rough target) Re-run alpha-sweep after v4 hybrid trains: PYTHONPATH=.:sdks/python python3 scripts/research/k3_integrated_niah_eval.py \ --f-theta-dir results/research/f_theta_v4_hybrid_finetuned \ --mix-alpha-sweep '0.0,0.25,0.5,0.75,1.0' \ --output results/research/k3_alpha_sweep_v4_hybrid.json Expected: recall > 0.5 at alpha=0 (pure f_θ), reaching ~1.0 at alpha=0.5 or higher. The fidelity-recall curve should be CONTINUOUS (not the cliff at alpha=1.0 we saw with v3). Stack ===== main (post #93 + #99 + #94 + #100 + #101 + #102) └── PR #103 (CUDA: workflow rules R1+R2+R3 + relmse + ...) ├── PR #104 (Mac MLX cross-model verifier; parallel-track) └── THIS PR #106 (attn_distill v3 evidence + alpha-sweep + v4 hybrid loss fix) Branch divergence note: PR #103 has the workflow-rules infrastructure (R2 reviewer-aid header lib, AGENTS.md, R2 CI test). PR #106 currently doesn't — those will merge in when one of the branches lands. Per R1, the bug fix (this commit) lives on PR #106 with the rest of the v3 attn_distill work, since that's where the user is iterating. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

….5 hr GPU waste Per user 2026-06-10 23:37 UTC: '这还需要问么？这么不专业，浪费大量的时间和GPU算力' User is right. Failure cost was ~8.5 hr H200 GPU time (v4a 3hr + v4b 5.4hr). The post-hoc evidence proves the entire training campaign was avoidable via a 5-minute fidelity probe. What happened ============= After v3 (relmse + attn_distill) failed NIAH recall (0/10), the agent recommended a v4 campaign: - v4a: warmstart from relmse v3, hybrid loss, 10k steps, rank 256 → 3 hr H200 wall - v4b: fresh start, hybrid loss, 20k steps, rank 768, 128 NIAH data → 5.4 hr H200 wall Both failed NIAH recall: 0/10. Post-hoc fidelity probes show: Checkpoint eval full_attn rel_mse --------- --------------------- relmse v3 (4k steps, rank 256) 1.45 v4a hybrid (10k, rank 256, 64 NIAH) 1.42 v4b hybrid (20k, rank 768, 128 NIAH) 1.52 ^^^^^ Recall threshold: 0.40 eval-domain rel_mse is essentially constant (~1.4-1.5) across: - 5× steps (4k → 20k) - 3× rank (256 → 768) - 0 → 128 NIAH prompts - 8× sequence length (gen_len 128 → 1024) Conclusion: the bottleneck is INFORMATION-THEORETIC (drafter K/V at eval positions), not optimization. No amount of training tweaks can break through. The v4 campaign was directionally wrong from the start. Could the 8.5 hr have been saved? ================================= Yes. A 5-minute probe before launching v4a: python3 scripts/research/k3_integrated_niah_eval.py \ --f-theta-dir results/research/f_theta_v3 \ --mix-alpha-sweep '0.0,0.5,1.0' \ --output /tmp/probe.json Would have shown: f_theta_baseline_rel_mse.full_attn = 1.45 (== relmse v3 on eval) recall threshold = 0.4 → gap is 3.5×, which no f_θ training tweak has historically reduced That probe + 'compare to recall threshold' is < 5 min. It would have falsified the hypothesis 'training tweaks (rank, steps, data, loss) will close the gap' BEFORE 8.5 hr of training started. The probe was not run because the agent (me) didn't have a rule forcing it. Now it does. R4 specification ================ Statement: Before recommending any GPU training that costs > 30 min wall, agent MUST first design + run a fidelity probe taking ≤ 10% of the planned wall, and emit a 'Probe verified' block in the same response. If probe FAILS the hypothesis, agent MUST NOT recommend the training run. Common probe patterns: - 'Bigger / longer / more data will improve metric X' → probe = run existing best checkpoint, measure metric X on eval distribution. If X is at a floor across multiple prior checkpoints, scale-up won't break it. - 'New loss function will fix problem Y' → probe = train new loss 500-1000 steps; verify loss decomposes sanely. Catches collapse / degeneracy early. - 'New architecture variant will help' → probe = train tiny version. Verify shape + loss curve. Enforcement: agent behavior, codified in AGENTS.md so the agent self-checks before every training recommendation. Files in this commit ==================== docs/agent-workflow-rules.md: - Add R4 section (full specification + when-it-applies table + fidelity probe template) - Add failure-log entry: 2026-06-10 pm, 8.5 hr H200, R4 introduced AGENTS.md: - Add R4 alongside R1+R2+R3 - Add 'Probe verified' block format (analogous to R3's 'Verified by agent' block) - Document common probe patterns Effect ====== Future training recommendations in this repo MUST be preceded by a fidelity probe with explicit hypothesis + falsification criterion. If the user sees a training recommendation without the 'Probe verified' block, the agent has violated R4 — call it out. Failure log is append-only. R4 joins R1+R2+R3 as a non-negotiable agent rule. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

…rward-port User 2026-06-11 01:15: 'analyze the entire development process of milestones and errors, summarize lessons learned as a reference for AGENTS.md.' This commit ports v4 evidence + watcher/analyzer scripts onto PR #103 branch (per R1) and adds a comprehensive retrospective document for Mac mini review. Files ===== docs/k3-postmortem-and-lessons.md (NEW, 394 lines) Comprehensive K3 retrospective covering: - Architecture summary + dimensions + recall threshold - 12 milestones (chronological, what worked) - 4 failure events (F1-F4, chronological, with cost): F1 ADR 0011 cross-attn bridge: ~4 hr wasted F2 K2.A.1 Mac serial bugfix: ~2 hr wasted F3 wrong-branch GPU run: 15 min wasted (→ R1+R2+R3) F4 v4 hybrid training campaign: 8.5 hr wasted (→ R4) - Pattern analysis (3 recurrent failure modes): A. 'Just iterate' without falsifying B. Branch fragmentation C. Integration debugging by ping-pong - Lessons → AGENTS.md additions: R5 (proposed): END-TO-END SMOKE BEFORE SCALING R6 (proposed): HONEST FAILURE MODES R7 (proposed): PRIOR-ART CHECK BEFORE RE-INVENTING - Cross-cutting: 'INFORMATION THEORY > OPTIMIZATION TWEAKS' when objective metric is invariant across training-side knobs, the bottleneck is informational not optimizational; tune upstream, not training. - Open questions / next steps with R4-compliant probe templates - GPU spend summary: ~26 hr total, ~13 hr (50%) avoidable Path A (next step) recommendation: bypass K/V Restoration on the 5 full-attention layers (5/30 = 17% of layer count, 9% of K/V memory). 50 LOC change + 10 min probe. scripts/research/k3_v4_analyze.py (NEW) Auto-analyzer for any K3 v4 evidence JSON (training report, NIAH eval, alpha-sweep). Used by the polling watcher loop. Shipped here so future agent sessions can re-use without redesigning. scripts/research/k3_v4_evidence_watcher.sh (NEW) Polling watcher used during the v4b training wait. Detects and emits TRIGGER:* markers for each evidence type. Useful template for future long-running training watchers (R4-compliant: agent waits for evidence rather than polling manually). results/research/f_theta_v4{a,b}_*.json (NEW, forward-ported from PR #106) results/research/k3_alpha_sweep_v4{a,b}.json (NEW) results/research/k3_fidelity_f_theta_v4{a,b}_*.json (NEW) v4 evidence on PR #103 branch (per R1) so the postmortem references resolve on this branch. Same files as on PR #106; included here for archival completeness. How this connects to existing rules =================================== R1 (SAME-PR FIX): postmortem committed to PR #103 branch (the open PR with workflow rules), not a child PR. R2 (REVIEWER AID HEADER): no new aids in this commit; existing scripts/review_pr_k3_f_theta_train_on_vast.sh already R2-compliant. R3 (PRE-GPU CONFIRMATION): the postmortem's Path A recommendation includes an R3-compliant probe template + R4-compliant probe template. R4 (PRE-TRAINING FIDELITY PROBE): F4 entry in the failure log is the motivating example. Postmortem demonstrates how a 5-min probe would have prevented the 8.5 hr v4 campaign waste. R5/R6/R7 (proposed in postmortem section 5): not yet shipped as enforced rules. Promote to mandatory if a future failure shows they are needed. For Mac mini review =================== git checkout AgentMemory/v04-pr-k3-block-c-f-theta-design-and-skeleton-8e7f git pull # Read these three files in order: # 1. docs/k3-postmortem-and-lessons.md (this commit) # 2. docs/agent-workflow-rules.md (R1-R4) # 3. AGENTS.md (R1-R4 enforcement) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

Per user 2026-06-11 02:58 UTC: '我要的是完整的 context 记录。不是你抽象之后的文档记录' Honest scope ============ This dump contains every raw artifact accessible from the cloud agent's VM filesystem. The literal user↔agent message-by-message conversation transcript is **NOT in this dump** — Cursor stores conversation messages server-side, not on the agent VM. Verified by: - searching ~/.cursor-server/data (no chat/conversation files) - resolving $AGENT_TRANSCRIPTS env var (points to empty dir) - searching the entire filesystem for files matching $CURSOR_CONVERSATION_ID = bc-51150fb2-5f33-444d-b21d-9baf584d8e7f (no transcript file exists) To get the literal conversation transcript: use Cursor's web UI session export feature with the conversation ID above. 84 files, ~2.3 MB total. Layout ====== agent-session-context/ README.md explains scope + limitations + reading order terminals/ (9 files, ~92 KB) raw output of every Shell tool call this session (PID + cwd + commands + stdout/stderr per call) Includes 4 older agent-tools artifacts. git-commits/ (8 files, ~1.4 MB) ALL-branches-since-2026-06-09.log: ~6000 lines, every K3-relevant commit with full message + diff stats. Per-branch logs filtered by branch. Effectively the agent's reasoning timeline. evidence/ (55 files, ~604 KB) every results/research/*.json: training reports (f_theta v1/v3/v4a/v4b), NIAH evals, alpha-sweeps, fidelity probes, identity-restore, spec-decode, Mac smoke. scripts/ (10 files, ~156 KB) every script the agent created or modified: f_theta.py, cross_model_dlm_verifier.py (+ MLX), k3_f_theta_train.py, k3_integrated_niah_eval.py, k3_v4_analyze.py, k3_v4_evidence_watcher.sh, review_pr_k3_*.sh, reviewer_aid_header.sh docs/ (3 files, ~40 KB) AGENTS.md, agent-workflow-rules.md, k3-postmortem-and-lessons.md Limitations explicitly documented in README =========================================== 1. No raw conversation messages. [server-side] 2. No structured tool-call trace. [Read/Write/StrReplace etc. left no filesystem-readable record] 3. No agent chain-of-thought. [reasoning lives only in commit messages + docs + user-visible chat] 4. Session-start summary not dumped. [provided in-context, not written to disk; itself abstracted] Reading orders documented in README for ~30 min and ~2 hr review budgets. Why this dump exists ==================== Earlier this session I produced docs/k3-postmortem-and-lessons.md as the 'comprehensive retrospective' for K3. The user pushed back: that document is an abstracted summary, not raw context. The dump is the correction. The abstract document is now correctly positioned as ONE artifact among the raw files — not a replacement for them. Mac mini access =============== git pull origin AgentMemory/v04-pr-k3-block-c-f-theta-design-and-skeleton-8e7f cd <repo>/agent-session-context # read README.md first Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

FluffyAIcode · 2026-06-11T15:51:03Z

Closing as superseded by PR #107.

…R) + DFlash fused spec-decode (>AR) on Gemma 4 26B-A4B (#107) * K3 Block B + C: f_theta projection + cross-model DLMRestoredVerifier (P0) Per user 'go P0' directive 2026-06-09 after architectural observation that PR #102's Mac MLX spec decode eval doesn't exercise the Kakeya inference engine's core architecture (sink+window verifier + dLM proposer K/V Restoration). This PR ships the foundational engine code for the integrated Kakeya inference architecture per ADR 0008 §11.3: verifier (Gemma 4 26B-A4B): └─ holds only sink+window local KV cache (sink=4 + window=64) └─ at evicted positions, takes K/V supplied by proposer (via f_θ) drafter (DFlash 0.4B, alignment-trained baseline): └─ runs full forward over committed prefix per step └─ K/V at every layer at every position captured └─ K/V projected through f_θ into verifier K/V space, injected at evicted positions Three new files --------------- inference_engine/v04/f_theta.py (~290 LOC) FThetaConfig dataclass + FThetaProjection nn.Module. Architecture: shared encoder + per-verifier-layer decoders, low-rank factorisation: drafter_kv_input [B, T, drafter_layers * drafter_kv_dim] ↓ encoder Linear(in, rank) rep [B, T, rank] ↓ per-verifier-layer decoders (30 × Linear(rank, verifier_kv_dim)) output [B, T, num_verifier_layers, num_kv_heads_v, head_dim_v] Default rank=256. Production K3 config (Gemma 4 26B-A4B + DFlash 0.4B): encoder: 2 × 5×256 × 256 = 655k params decoders: 2 × 30 × 256 × 2048 = 31.5M params Total: ~32M params (vs drafter 430M, verifier 26B) Separate K and V projections (different downstream roles). Save/load: save_pretrained(dir) writes f_theta_config.json + f_theta_weights.pt; from_pretrained(dir, dtype, device) loads back. inference_engine/v04/cross_model_dlm_verifier.py (~270 LOC) CrossModelDLMRestoredVerifier wrapper. Construction validates drafter + verifier dimensions match the f_θ config (rejects drafter-vs-verifier-vs-f_θ mismatch loudly at __init__). forward(input_ids, apply_rotary_pos_emb, eager_attention_forward): 1. compute_evicted_positions(T, sink, window) 2. If no evicted (T <= sink+window): plain verifier forward 3. Drafter forward via _capture_drafter_kv (forward hooks on k_proj/v_proj at each drafter layer) 4. f_θ.forward_kv_pack(drafter_K_per_layer, drafter_V_per_layer) → verifier K, V at every (layer, position) 5. Patch each verifier layer's self_attn.forward to: a. Run standard q/k/v_proj + q_norm/k_norm + RoPE b. At evicted positions, REPLACE k, v with f_θ output (after k_norm + RoPE applied via prepare_restored_attention_kv) c. Standard attention compute path through eager_attention_forward 6. Run verifier forward → logits 7. Restore original attention forwards (try/finally) Two scope-outs (recorded inline): * MLX verifier path: this module patches HF transformers attention. Mac MLX integration is a follow-up PR (instrument mlx_lm Gemma 4 model directly, not via attention monkey-patch). * Speculative decoding accept/reject loop: separate inference engine concern. PR #93's DFlashProposer + mlx_verify_block handles the spec-decode side; combining with this module's K/V Restoration is a separate integration step. Drafter K/V capture (_capture_drafter_kv): instruments DFlashDrafter's internal layer.self_attn.k_proj / v_proj via forward hooks. NOTE inline that the first-iteration synthetic-context capture (zero hidden as drafter input) is plumbing-validation; product-meaningful K/V values require conditioning on verifier aux hiddens, which is the next integration step (after f_θ training validates the projection alone). scripts/research/k3_f_theta_train.py (~310 LOC) Training pipeline for f_θ on CUDA: 1. Load Gemma 4 26B-A4B verifier (transformers bf16, sdpa) 2. Load DFlash drafter (PR #93's DFlashDrafter from models/dflash-kakeya-baseline) 3. Data collection: for each prompt in PROMPTS (same 64-prompt corpus as PR #93's alignment_train), run greedy AR generation to gen_len tokens, capture per-layer per-position K/V via hooks on k_proj/v_proj of both models 4. Train f_θ with MSE loss across (layer, position) pairs, AdamW lr=1e-3, weight_decay=0.01, gradient clip 1.0 5. Save checkpoint at --save (default results/research/f_theta_v1) Memory budget: at T=512, ~128 MB per sequence cached on GPU. 64 sequences ≈ 8 GB. Fits H200 80 GB easily. Validation: report initial vs final loss; reduction factor. inference_engine/v04/__init__.py: re-exports the new public surface (FThetaConfig, FThetaProjection, CrossModelDLMRestoredVerifier, CrossModelLayerMapping). Tests (Linux CI: 27 new tests) ----------------------------- tests/inference_engine/v04/test_f_theta.py (21 tests): TestFThetaConfig (4): dim properties + JSON round-trip TestForwardShapes (4): forward_k/v shape contract + input validation TestForwardKVPack (3): KVCapture-style input + consistency vs explicit concat TestParameterCount (2): tiny + production param count locked in TestSaveLoadRoundTrip (4): save+load preserves outputs; missing-file errors TestDeviceDtypeDispatch (2): to(dtype), from_pretrained dtype override TestGradientFlow (1): gradients flow through encoder + decoders separately (K path doesn't update V weights and vice versa) tests/inference_engine/v04/test_cross_model_dlm_verifier.py (6 tests): TestConstruction (3): dimension validation rejects mismatch; valid construction succeeds; negative sink/window raises TestProjectDrafterKV (1): output shape contract TestNoEvictPath (1): short prompt (T <= sink+window) doesn't invoke drafter TestExports (1): module + namespace re-exports Tests: 354 passing (336 pre-existing + 21 f_theta + 6 cross-model; 12 research/ unchanged from PR #102). What this PR does NOT yet do (deferred to follow-up PRs) -------------------------------------------------------- 1. Train f_θ on real data — requires vast.ai GPU time. scripts/research/k3_f_theta_train.py is the runnable trainer. Once trained, the checkpoint goes to a follow-up PR with the evidence (training report + integrated NIAH ladder evidence). 2. End-to-end integrated NIAH ladder evidence — needs: * trained f_θ checkpoint (step 1) * cross-model DLMRestoredVerifier reviewer aid (off-the-shelf K1.E NIAH harness needs a small adapter to use this verifier wrapper) * vast.ai run producing the evidence JSON 3. Mac MLX integration — instruments mlx_lm Gemma 4 model directly (different surgical approach than HF transformers attention monkey-patch). Follow-up PR. 4. _capture_drafter_kv proper aux-conditioning — current synthetic zero-hidden capture is plumbing only. The proper path passes verifier aux hiddens into the drafter (DFlash architecture), captures K/V from THAT forward. Adds a method to DFlashDrafter in a follow-up. These are the remaining items on the K3 critical path; this PR establishes the engine API surface they all depend on. Stack ----- Off main (post #93 + #99 + #94 + #100 + #101 + #102 merged). Independent of any other open PR. Outstanding work after this PR: Step 5 — K2.A backport PR (P2) Step 6 — alignment training corpus expansion (P2) P0 cont. — f_θ training run + integrated NIAH evidence P0 cont. — Mac MLX integration of cross-model DLMRestoredVerifier Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 P0 critical fixes + vast reviewer aids + integrated NIAH eval User signal 2026-06-09: 'A / B / C 全部做完。我已经开了vast' — proceed through full P0 critical path; vast is open for runs. Three fixes + three new files in this commit: (A) FIX: _capture_drafter_kv now uses verifier embed_tokens Previous version (just committed in this PR) used synthetic zero hidden state to fire k_proj/v_proj hooks. This is plumbing-only and produces meaningless K/V values. DFlashDrafter's design (PR #93) shares verifier embed_tokens (no own embedding lookup), so the correct capture path is: 1. verifier_model.get_input_embeddings()(input_ids) × sqrt(hidden) 2. Pass embedded hiddens through drafter.layers (no aux conditioning) 3. Capture K/V via forward hooks per layer Updated _capture_drafter_kv signature to take verifier_model (required for embed_tokens). Updated CrossModelDLMRestoredVerifier. project_drafter_kv to pass it. Updated test fixture to provide a real embed_tokens on the synthetic verifier (was previously unnecessary; now required). (B) FIX: k3_f_theta_train.py now uses _capture_drafter_kv Previous version called capture_proposer_kv(drafter.model, input_ids) which would crash on real DFlashDrafter — DFlashDrafter is a flat nn.Module without .model attribute (capture_proposer_kv expects model.model.layers OR model.transformer.h, both absent). Switched to inference_engine.v04.cross_model_dlm_verifier. _capture_drafter_kv (the same path the cross-model verifier uses at inference time). Ensures training and inference are using the IDENTICAL drafter K/V values — no train/serve skew. (C) NEW: scripts/review_pr_k3_f_theta_train_on_vast.sh vast.ai reviewer aid for f_θ training. Pre-flight checks: 1. HF_TOKEN (Gemma 4 gated) 2. models/dflash-kakeya-baseline/ Git LFS pulled (>100MB safetensors) 3. CUDA available 4. transformers 5.x (Gemma 4 support) Env knobs: STEPS, LR, RANK, N_PROMPTS, GEN_LEN, SAMPLE_POSITIONS, SAVE_DIR, SEED. Default config: 4000 steps, rank=256, 64 prompts × 128 gen tokens — fits H200 80 GB easily, ~8-15 min wall clock. Output: trained f_θ checkpoint + training report. Validation gates printed at end (loss_reduction_factor ≥ 2.0 sanity). (D) NEW: scripts/research/k3_integrated_niah_eval.py (~280 LOC) THE K3 PRODUCT GATE EVIDENCE SCRIPT. Combines: * CrossModelDLMRestoredVerifier (verifier with sink+window cache + drafter K/V Restoration via f_θ) * K1.E NIAH evaluation harness (effective_attention_window / recall / memory metrics) Validates per ADR 0008 §11.8 release gates: 1. Architectural correctness: effective_attention_fraction = 1.0 at every NIAH ladder rung 2. Memory bounded: sustained verifier KV-cache ≤ O(sink+window) 3. Recall preservation: |recall_cross_model - recall_oracle| ≤ 5 pp at every rung (ADR §11.8 1a — architecturally-meaningful gate) Runs: - cross-model verifier on each NIAH sample, decodes max_new_tokens - full-attention oracle baseline on same samples (--skip-oracle to bypass; loses recall_delta gate signal) - aggregate recall, attention_window, memory; compute gate booleans Output JSON schema mirrors K1.E NIAH harness (per_config recall, attention_window, memory) + new 'gate' block with the three booleans for direct inspection. (E) NEW: scripts/review_pr_k3_integrated_niah_on_vast.sh vast.ai reviewer aid for the integrated NIAH eval. Pre-flight: 1. HF_TOKEN 2. f_θ checkpoint at $F_THETA_DIR 3. drafter LFS pulled 4. CUDA available Runs the integrated NIAH eval per CONTEXT_LADDER rung (default '70 280', i.e. ~1.4k + ~5.6k tokens). Per-rung JSON + combined log. Final aggregation diff-able with PR #94's same-checkpoint K1 ladder evidence. After this PR + a vast run of (review_pr_k3_f_theta_train_on_vast.sh → review_pr_k3_integrated_niah_on_vast.sh), the K3 product gate is empirically closed on CUDA. Mac MLX path follows as separate PR (instrument mlx_lm Gemma 4 model directly; can't reuse the HF attention monkey-patch approach). Tests: 354/354 passing on Linux CI (no v04 code regressions; new script files don't run in CI but parse + bash -n check OK). Stack: Off main, builds on PR #103 commits in this same branch. PR #103 description updated to reflect added scripts + critical fixes. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3: support Gemma4 multimodal nested config/decoder in f_theta train + cross-model verifier Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3: capture V from k_proj output for Gemma4 v_proj-None (KV-sharing) layers Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3: heterogeneous per-layer verifier KV heads in f_theta + per-layer capture/loss for Gemma4 Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3: Gemma4-faithful cross-model restore forward (per-layer KV, v_norm, RoPE unsqueeze_dim=2, v_proj-None, evicted slicing) + gemma4 helpers import + tests Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3: cast f_theta input to encoder weight dtype (fp32 f_theta vs bf16 drafter K/V) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3: fix integrated NIAH eval to use real niah_eval API (chat-template encode, aggregate_recall, v04_dlm_restored window) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3: handle BatchEncoding return from Gemma4 apply_chat_template in integrated NIAH eval Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3: per-layer verifier head_dim in f_theta (Gemma4 full layers use global_head_dim=512, 2 KV heads) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3: add identity-restore diagnostic (inject verifier's own K/V) to isolate restore machinery from f_theta accuracy Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 f_theta v1 trained checkpoint (Gemma4 26B-A4B verifier, per-layer KV; loss 50.8->3.70, 13.74x) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 integrated NIAH gate evidence: arch_correct=1.0 PASS, recall gate FAIL (f_theta v1), identity-restore recall=1.0 (machinery validated) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 f_θ trainer v2 — fix recall=0 (cosine+mag loss + NIAH data + cosine LR + 5× longer) Per user 2026-06-10: 'vast上训练完了,recall不达标。fix这个问题' PR #103 v1 evidence diagnosis ============================= Identity-restore evidence: recall = 1.0 (machinery correct). f_θ-projected: recall = 0.0 (training inadequate). Decoded outputs were fluent ('The answer is not provided in the text...') but lexical content of the haystack was lost — the classic symptom of attention-noise from low-fidelity K/V projection. Four root causes, four fixes ============================ (a) Wrong loss objective. v1 used pure MSE on raw K/V; final MSE 3.70 ≈ RMSE 1.92 per element ≈ 2σ noise. Attention is softmax(QK^T); 2σ noise destroys softmax peakedness → lexical content lost. Fix: cosine + magnitude per-vector loss (direction-preserving, scale-aware) replaces pure MSE in the default 'combined' loss type. Cosine bounds Q·K_pred ≈ Q·K_tgt; magnitude preserves softmax scale. Small (0.1×) MSE term retained for stability when norms are near zero. (b) Tiny corpus, no NIAH structure. v1 used 62 prompts × ~600 tokens = 37k unique tokens, ZERO needle-in-a-haystack patterns. The eval is 100% NIAH. f_θ never saw retrieval structure. Fix: synthetic NIAH-style training prompts (haystack + needle line) generated alongside the existing PROMPTS list, default 50% NIAH / 50% general. Independent seed from the eval (seed + 1000) so no needle reuse — verified by unit test. (c) Trivial training duration. v1 trained 4000 steps × ~15ms ≈ 59 seconds. AdamW barely warmed. Fix: default 20000 steps (5× longer). (d) No LR schedule. v1 used constant lr=1e-3, never annealed. Fix: cosine schedule with linear warmup (default 500 steps warmup → cosine decay to peak/100 over remainder). Three modified files ==================== scripts/research/k3_f_theta_train.py (~530 LOC, +280 / -50) Three new helpers: _per_vector_cosine_mag_loss(pred, tgt) → (combined, cos, mag) Per-K/V-vector cosine similarity + magnitude MSE. Returns detached cos and mag for diagnostics. _make_niah_training_prompts(n, seed, ...) → list[str] Generates synthetic haystack+needle prompts in the same pattern as PR #94's eval harness, but with independent seed + extra word lists / filler lines so no needle is reused. _lr_at_step(step, peak_lr, total_steps, warmup_steps, schedule) Returns the LR at step. schedule='const' → peak. schedule= 'cosine' → linear warmup → cosine decay to peak/100. Refactored _f_theta_loss to dispatch on loss_type (mse | cos_mag | combined) and emit per-component diagnostics (cos_K_total, cos_V_total, mag_K_total, mag_V_total, mse_*) into an optional diag_buf for live training logs. main() additions: --loss-type {mse, cos_mag, combined} default 'combined' --lr-schedule {const, cosine} default 'cosine' --warmup-steps default 500 --n-niah-prompts default 64 --no-niah-prompts (v1 reproduction flag) --niah-min-lines / --niah-max-lines default 30 / 90 Default changes (all v1-reproducible via flags): --steps 4000 → 20000 (5× longer) --gen-len 128 → 512 (4× longer sequences) Training loop now sets per-step LR via _lr_at_step, logs cosine components alongside loss, and persists final_diagnostic + loss_type + lr_schedule in the report (schema_version=2). scripts/review_pr_k3_f_theta_train_on_vast.sh (~165 LOC, +35 / -15) Updated header to v2 with explicit reproduction recipe for v1. Added env knobs LR_SCHEDULE, WARMUP_STEPS, LOSS_TYPE, N_NIAH_PROMPTS. Updated default SAVE_DIR to results/research/f_theta_v2 so v1 evidence is not overwritten. v1 reproduction recipe (printed in header): STEPS=4000 GEN_LEN=128 LR_SCHEDULE=const LOSS_TYPE=mse \ N_NIAH_PROMPTS=0 SAVE_DIR=results/research/f_theta_v1_repro \ HF_TOKEN=hf_xxx bash $0 Updated expected-timing block (~20-30 min vast wall, was ~8-15 min), validation gates (loss_reduction_factor ≥ 5×, cosK < 0.05). Tests (Linux CI: 17 new tests) ============================== tests/research/test_k3_f_theta_train_v2.py: TestPerVectorCosineMagLoss (5): - identical vectors → loss = 0 - negated vectors → cos_loss = 2.0 (worst case), mag_loss = 0 - orthogonal unit vectors → cos_loss = 1.0, mag_loss = 0 - 2× scaled vector → cos_loss = 0 (same direction), mag_loss > 0 - loss is differentiable (gradient flows back to pred) TestLRSchedule (6): - const schedule returns peak at every step - cosine warmup at step 1 = peak/warmup_steps - cosine warmup ends exactly at peak at warmup_steps - cosine decay reaches floor (peak/100) at total_steps - cosine midway above floor (≈ 0.5 × peak after warmup) - unknown schedule raises ValueError TestNIAHTrainingPrompts (6): - returns requested count - prompts contain 'secret code is' + 'Question:' lines - seed determinism (same seed → same prompts) - different seeds → different prompts - haystack_min_lines / max_lines bounds respected - no eval seed collision (seed=1000 default ≠ seed=0/42 outputs) Tests: 373/373 passing on Linux CI (354 pre-existing + 9 from PR #104 + 10 from PR #103 + 17 new, with overlap from earlier additions). Smoke-tested in-process with synthetic CapturedSequence: all 3 loss types compute, all 3 backprop gradients to f_θ params, all 3 emit diag_buf entries. Validation gate (vast retrain) ============================== Same reviewer aid, new defaults: HF_TOKEN=hf_xxx bash scripts/review_pr_k3_f_theta_train_on_vast.sh Output: results/research/f_theta_v2/{config.json, weights.pt} + results/research/f_theta_v2.json with per-component diagnostics. Then re-run the integrated NIAH eval against the v2 checkpoint: bash scripts/review_pr_k3_integrated_niah_on_vast.sh \ F_THETA_DIR=results/research/f_theta_v2 Expected outcomes (vs v1): - cosK_total < 0.05 (v1 had no cosine measurement) - loss_reduction_factor ≥ 5× (v1 was 13.7×) - integrated NIAH recall_cross_model approaches recall_oracle - recall_delta_within_5pp gate closes (v1 had delta = 100 pp) If v2 still fails to close the recall gate, escalate to architecture fix (rank ↑ from 256 → 768, per-layer encoders instead of shared) and/or attention-output distillation loss (more expensive but principled). v2 is the highest-leverage minimal-change fix; it should close most of the gap. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 f_θ trainer v3 — one-shot attention-output distillation (skip v2 intermediate) Per user 2026-06-10: '我要求直接上一步到位的训练方案。不要搞这种中间态，浪费时间和CPU资源' Skipped the v2 cosine+magnitude intermediate. Default loss is now attention-output distillation — the principled training objective for K/V replacement. v2 cos+mag remains accessible via --loss-type cos_mag for ablation, but is not the default path. The principled loss =================== For each verifier layer ℓ: K_pred_ℓ, V_pred_ℓ = f_θ(drafter_KV)[ℓ] Q_for_attn = q_norm(Q_raw_ℓ).view(B, T, H_q, D) → RoPE → transpose K_for_attn = k_norm(K_pred_ℓ).view(B, T, H_kv, D) → RoPE → transpose V_for_attn = v_norm(V_pred_ℓ).view(B, T, H_kv, D) → transpose GQA repeat K, V to H_q O_inner = scaled_dot_product_attention(Q, K, V, mask, scale) O_pred = o_proj(O_inner.reshape(B, T, H_q*D)) loss_ℓ = MSE(O_pred, O_tgt_ℓ) ^^^ captured during data collection from the verifier's actual attn module post-o_proj output Total = mean over layers Why this is mathematically right for K/V projection --------------------------------------------------- attention(Q, K, V) is the actual quantity that propagates through the residual stream at inference. v1 (raw MSE on K) and v2 (cos+mag on K) are PROXIES for attention behavior. v3 directly optimises the attention output, so the loss landscape's gradient points precisely at 'f_θ K/V produces equivalent verifier behavior'. It accounts for: GQA grouping, RoPE, causal/sliding mask, k_norm/q_norm/v_norm, AND the o_proj that follows attention. Implementation strategy ======================= Tractability concern: the principled loss seemingly requires a full verifier forward per training step (≈ 3 sec on H200 → 16+ hours for 20000 steps). NOT acceptable. Solution: smart caching. During data collection (one verifier forward per sequence), capture per-layer: - Q_raw [T, num_heads × head_dim] from q_proj forward hook - O_tgt [T, hidden_dim] from attn module forward hook - cos, sin [1, T, head_dim] from attn forward pre-hook - attn_mask from attn forward pre-hook All cached on CPU bf16 (≈ 13 MB per layer per sequence × 30 layers × 64 sequences ≈ 25 GB CPU RAM). Training streams these to GPU per step. No verifier forward is needed at training time. Per-step cost: f_θ forward + per-layer attention recomputation (scaled_dot_product_attention with cached Q + f_θ-predicted K/V) + o_proj + MSE. ~80 ms/step on H200. 20000 steps = 25-30 min. Total v3 wall on H200: ~40-60 min (data collect + training). Three modified files ==================== scripts/research/k3_f_theta_train.py (~1100 LOC, +400) New dataclass: AttentionTargetData Per-layer Q_raw + O_tgt + cos + sin + attention_mask + per-layer num_heads / head_dim. CPU bf16 storage. New function: _capture_attention_target_data Runs verifier forward with hooks (forward hook on q_proj for Q_raw, forward hook on attn module for O_tgt, forward pre-hook on attn module for position_embeddings + attention_mask). Returns AttentionTargetData with all tensors on CPU bf16. New function: _attention_distillation_loss The principled loss as described above. Full per-layer pipeline with proper GQA / RoPE / mask handling. Streams cached tensors from CPU to GPU per layer; frees per-layer GPU memory before moving to next layer. Modified: CapturedSequence Made verifier_k / verifier_v Optional. Added attn_target field (Optional[AttentionTargetData]). For attn_distill loss, only attn_target is captured (saves ~125 MB per sequence vs legacy K/V capture). For legacy losses, only verifier_k/v captured. Modified: _f_theta_loss Dispatch on loss_type. attn_distill path → _attention_distillation_loss. Legacy losses (mse | cos_mag | combined) path → previous v2 logic. Validates seq has the right capture for the chosen loss. Modified: _collect_sequence Now takes capture_legacy_kv + capture_attn_target flags. Routes to either or both capture paths. Modified: main() - Loaded attn_implementation='eager' for attn_distill (sdpa breaks the attn-module-level forward hook contract); 'sdpa' for legacy - Imports apply_rotary_pos_emb from transformers.models.gemma4 - --loss-type now defaults to attn_distill, choices include all 4 - --rank default is None → auto-resolve: 768 for attn_distill, 256 for legacy (rank ↑ for the more capable principled trainer) - --sample-positions default 0 → use full T (recommended for attn_distill); 256 for legacy - Per-step log shows per-loss-type diagnostics: cos sim for cos_mag/combined, mseO/|O_tgt|^2 ratio for attn_distill - Report includes 'final_diagnostic' + 'loss_type' scripts/review_pr_k3_f_theta_train_on_vast.sh (~190 LOC, +20 / -25) Updated to v3 defaults: LOSS_TYPE=attn_distill (was 'combined' in v2 plan, never shipped) RANK= (empty → trainer auto-picks 768 for attn_distill) SAMPLE_POSITIONS=0 (full T) SAVE_DIR=results/research/f_theta_v3 Header docstring documents the v1 reproduction recipe AND the v3 rationale (one-shot principled trainer). Banner shows the resolved attn implementation (eager vs sdpa) and the resolved RANK value. Validation gate updated: 'mseO/|O_tgt|^2 ratio < 0.05' replaces 'cosK_total < 0.05' (v3 diagnostic; ratio quantifies attention-output noise). tests/research/test_k3_f_theta_train_v2.py (+10 new tests) TestAttentionDistillationLoss (7): - attention_distill_loss_runs (returns scalar with diag populated) - loss_is_differentiable_through_f_theta (gradient flows to f_θ) - o_proj_weights_remain_frozen_in_loss (frozen verifier params receive no grad — important for training to not OOM/NaN) - dispatch_through_f_theta_loss_function (v2 _f_theta_loss correctly routes to _attention_distillation_loss for attn_distill) - attn_distill_requires_layers_arg (clear error if layers/RoPE/ device aren't passed) - legacy_loss_rejects_attn_only_capture (mse loss on attn_target- only seq raises RuntimeError instead of silently producing NaN) - sample_positions_subselects_output (full vs sub sample both produce a valid scalar loss) TestAttentionTargetDataDataclass (3): - fields_present - captured_sequence_optional_kv_and_attn (legacy fields default to None) - captured_sequence_attn_target_path (attn_target stored correctly) Stub _StubAttn / _StubLayer reproduce the Gemma 4 self_attn module surface (q_norm, k_norm, v_norm, q_proj, o_proj, scaling, head_dim) enough for the loss to run on Linux CI without an actual verifier. Tests: 383/383 passing (354 pre-existing + 9 from PR #104 + 10 from PR #103 + 17 from v2 + 10 new v3 — with overlap). Validation gate (vast retrain, one-shot) ======================================== Run the same reviewer aid; defaults pick up v3: HF_TOKEN=hf_xxx bash scripts/review_pr_k3_f_theta_train_on_vast.sh Output: results/research/f_theta_v3/{f_theta_config.json, f_theta_weights.pt} results/research/f_theta_v3.json (with mseO + |O_tgt| diagnostics) Then re-run integrated NIAH against v3 checkpoint: F_THETA_DIR=results/research/f_theta_v3 \ bash scripts/review_pr_k3_integrated_niah_on_vast.sh Expected v3 outcomes: - mseO_mean / |O_tgt|^2 ratio < 0.05 (attention output noise low) - integrated NIAH recall_cross_model ≈ recall_oracle - recall_delta_within_5pp gate CLOSES This is the principled one-shot fix. If recall still falls short (≥ 5pp delta), the issue is f_θ capacity — escalate to per-layer encoders or larger rank (RANK=1024). But attn_distill loss + rank 768 + 20k steps + NIAH data + cosine LR is the maximum-strength single-shot training configuration without architectural rewrites. Stack ===== main (post #93 + #99 + #94 + #100 + #101 + #102) └── PR #103 (CUDA: f_θ + cross-model + train script + integrated NIAH) ├── PR #104 (Mac MLX cross-model verifier; parallel-track) └── THIS PR #106 (trainer v3 — one-shot attn distill, supersedes v2 plan) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 S6: --mix-alpha-sweep fidelity->recall diagnostic (interpolate evicted K/V between f_theta and true; map recall vs residual rel_mse) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 attn_distill v3 evidence: train reduction 21.47x (attn-output rel-err 1.0->~0.20), but integrated NIAH recall still 0/10 both rungs (arch gate PASS) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 S6 alpha-sweep on attn_distill v3: recall 0 for all alpha<1.0 (degenerate — attn_distill K/V are ~135x off-scale; k_norm/v_norm normalize scale away, so raw-space mix is confounded) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 S6 alpha-sweep on scale-matched relmse v3: recall knee in (0,0.5]; full-attn rel_mse 0.36 -> recall 1.0, 1.44 -> 0; eval-domain err (1.44) >> in-domain (0.58) = distribution shift Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 f_θ trainer v4: attn_distill_hybrid loss — fix the f_θ collapse exposed by alpha-sweep Per user 2026-06-10: 'attn_distill sweep evidence... pls check the result' Diagnosis from sweep evidence (commit 72ce157) ============================================== f_theta_baseline_rel_mse.overall = 1331.94 f_theta_baseline_rel_mse.full_attn = 18254 f_θ raw (pre-norm) K/V output is 36× off-scale from verifier's true K/V (135× on full-attention layers). Despite this, attn_distill training converged to mse_O = 0.176 (looks fine) because k_norm and v_norm are RMSNorm — they NORMALIZE THE SCALE AWAY before attention. The attn_distill loss (computed downstream of k_norm) was scale-invariant and thus blind to the magnitude collapse. Sweep showed recall=0 for ALL alpha < 1.0 (in raw-space mixing), with recall jumping to 1.0 only at alpha=1.0 (pure verifier K/V). Reason: at alpha=0.9 (90% true + 10% f_θ), the f_θ component is 0.1 × 36 = 3.6× the magnitude of the true component (0.9 × 1) and DOMINATES THE DIRECTION post-mixing. After k_norm normalises the total magnitude, the direction is still dominated by f_θ's (directionally-wrong) output. Recall stays at 0 until alpha=1.0 (no f_θ contribution at all). This is **f_θ collapse degeneracy**: attn_distill loss has multiple local minima, including a degenerate one where f_θ outputs are magnitude-runaway and direction-arbitrary, but post-norm-then-attn gives 'evicted positions get neutral attention weights' so the local cache (sink+window) carries the attention output. Loss is ~0.18 (close to zero because evicted contribution is suppressed), but f_θ is contributing zero useful retrieval signal. This explains why NIAH failure mode changed from v1's 'confused hallucinations' to attn_distill v3's 'confident refusal' — f_θ isn't contributing wrong info, it's contributing NOTHING (post- attention), and the local cache can't see the needle. The fix: attn_distill_hybrid loss ================================= Direct supervision on K/V at three levels (in addition to attn output): loss = 1.0 * MSE(O_pred, O_tgt) # attention output + λ_kDir * (1 - cosine(K_pred_post_norm, K_tgt_post_norm)) # K direction + λ_vDir * (1 - cosine(V_pred_post_norm, V_tgt_post_norm)) # V direction + λ_kMag * MSE(|K_pred_pre_norm|, |K_tgt_pre_norm|) / |K_tgt|² # K magnitude + λ_vMag * MSE(|V_pred_pre_norm|, |V_tgt_pre_norm|) / |V_tgt|² # V magnitude Defaults: λ_kDir = λ_vDir = 1.0, λ_kMag = λ_vMag = 0.1. The cosine terms (post-norm) are the crucial fix — they constrain K direction directly, eliminating the degenerate solution where f_θ produces direction-arbitrary K. The magnitude terms (pre-norm) prevent the 36× scale runaway. Hybrid is the new default loss type. v3 attn_distill remains available via --loss-type attn_distill for ablation. Six modifications ================= scripts/research/k3_f_theta_train.py: - Extended AttentionTargetData with optional k_raw_tgt + v_raw_tgt (CPU bf16 cache, ~100 MB extra per sequence — acceptable) - _capture_attention_target_data new flag capture_raw_kv (also captures k_proj/v_proj outputs via forward hooks; v_proj-None layers fall back to k_proj output, matching cross_model_dlm_verifier semantics) - _attention_distillation_loss new flags hybrid, lambda_k_dir, lambda_v_dir, lambda_k_mag, lambda_v_mag. When hybrid=True, loads K_tgt_pre and V_tgt_pre, applies layer's k_norm + v_norm, computes cosine direction loss + pre-norm magnitude loss - _f_theta_loss dispatches loss_type='attn_distill_hybrid' to _attention_distillation_loss with hybrid=True - main(): new args --lambda-k-dir/--lambda-v-dir/--lambda-k-mag/ --lambda-v-mag, --init-from (warm-start from existing checkpoint, useful for fine-tuning attn_distill v3 with hybrid loss for fewer steps) - Default loss_type changed: attn_distill → attn_distill_hybrid - capture_raw_kv_in_attn_target=True automatically for hybrid - Per-step log: hybrid prints kDir/vDir/kMag/vMag alongside mseO/ratio scripts/review_pr_k3_f_theta_train_on_vast.sh: - Default LOSS_TYPE=attn_distill_hybrid - New env knobs LAMBDA_K_DIR/LAMBDA_V_DIR/LAMBDA_K_MAG/LAMBDA_V_MAG/ INIT_FROM - SAVE_DIR default → results/research/f_theta_v4_hybrid (preserves v3 attn_distill evidence) - Reviewer aid recipe string includes hybrid lambdas + INIT_FROM tests/research/test_k3_f_theta_train_v2.py: - TestAttentionDistillationHybridLoss (5 new tests): * hybrid_runs_and_emits_full_diag (mseO+kDir+vDir+kMag+vMag in diag) * hybrid_requires_raw_kv_tgt (RuntimeError if missing — fail loud) * hybrid_dispatch_via_loss_type (loss_type='attn_distill_hybrid' routes) * hybrid_loss_strictly_higher_than_attn_distill_alone (verifies added terms have effect, not silently zero) * hybrid_grad_flows_to_f_theta (gradient reaches f_θ params) - TestAttentionTargetDataDataclass + 1 test: * attention_target_data_optional_raw_kv_for_hybrid (None by default; populated when capture_raw_kv=True) Tests: 389/389 passing on Linux CI. Validation gate (vast retrain — TWO options) ============================================ Option A — Fine-tune v3 attn_distill checkpoint with hybrid loss (saves ~75 min, recommended): HF_TOKEN=hf_xxx \ INIT_FROM=results/research/f_theta_v3_attn_distill \ STEPS=10000 \ SAVE_DIR=results/research/f_theta_v4_hybrid_finetuned \ bash scripts/review_pr_k3_f_theta_train_on_vast.sh Expected wall: ~30-45 min (data already collected; only training). The warm-start from v3 attn_distill checkpoint gives the new loss a head start on the attn output term while the hybrid terms force K/V direction + magnitude into shape over the next 10k steps. Option B — Train from scratch with hybrid loss (full reset): HF_TOKEN=hf_xxx bash scripts/review_pr_k3_f_theta_train_on_vast.sh Expected wall: ~90 min (data collection ~45 min + training ~45 min). Cleaner baseline — no inheriting the degenerate v3 attn_distill weights. Expected v4-hybrid outcomes (vs v3 attn_distill) ================================================ k_dir_mean < 0.05 (cosine sim > 0.95 on post-norm K) v_dir_mean < 0.05 k_mag_mean < 0.05 (pre-norm magnitude matched within ~5%) v_mag_mean < 0.05 mse_O_mean < 0.10 (better than v3's 0.176, since K/V are now non-degenerate) f_theta_baseline_rel_mse.overall < 50 (vs v3's 1331; rough target) Re-run alpha-sweep after v4 hybrid trains: PYTHONPATH=.:sdks/python python3 scripts/research/k3_integrated_niah_eval.py \ --f-theta-dir results/research/f_theta_v4_hybrid_finetuned \ --mix-alpha-sweep '0.0,0.25,0.5,0.75,1.0' \ --output results/research/k3_alpha_sweep_v4_hybrid.json Expected: recall > 0.5 at alpha=0 (pure f_θ), reaching ~1.0 at alpha=0.5 or higher. The fidelity-recall curve should be CONTINUOUS (not the cliff at alpha=1.0 we saw with v3). Stack ===== main (post #93 + #99 + #94 + #100 + #101 + #102) └── PR #103 (CUDA: workflow rules R1+R2+R3 + relmse + ...) ├── PR #104 (Mac MLX cross-model verifier; parallel-track) └── THIS PR #106 (attn_distill v3 evidence + alpha-sweep + v4 hybrid loss fix) Branch divergence note: PR #103 has the workflow-rules infrastructure (R2 reviewer-aid header lib, AGENTS.md, R2 CI test). PR #106 currently doesn't — those will merge in when one of the branches lands. Per R1, the bug fix (this commit) lives on PR #106 with the rest of the v3 attn_distill work, since that's where the user is iterating. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 S6 knee refinement (relmse v3): recall transition alpha 0.3->0.4->0.5 = full-attn rel_mse 0.71(0/10)->0.52(6/10)->0.36(10/10) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 trainer aid: forward NIAH_MIN_LINES/NIAH_MAX_LINES env to --niah-{min,max}-lines (was ignored) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 fix: import apply_rotary_pos_emb for attn_distill_hybrid too (was only attn_distill -> hybrid crashed) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 v4a warm-start hybrid checkpoint (rank256, init relmse v3, attn_distill_hybrid, gen1024, niah140, 10k): reduction 3.42x, attn-output ratio ~0.24 Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 v4b fresh hybrid checkpoint (rank768, 128 NIAH, gen1024, niah140, 20k): reduction 8.01x, attn-output ratio ~0.21 Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 v4a/v4b hybrid integrated NIAH evidence: both recall 0/10 both rungs (arch PASS) despite scale-matched hybrid + NIAH data + bigger/longer/warm-start Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 fidelity probe v4a/v4b: eval full-attn rel_mse 1.42/1.52 (== relmse v3's 1.44) — full-attn K/V fidelity floor independent of loss/rank/data; blend to 0.36 -> recall 1.0 (threshold confirmed) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 v4a/v4b canonical NIAH + alpha-sweep artifacts: NIAH 0/10 both; sweep recall flips 0->1 between alpha 0.25 (full-attn ~0.8) and 0.5 (~0.37), identical for both Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 S5: exact_layer_indices in cross-model verifier + --s5-exact-full-attn eval flag (keep full-attention layers' K/V exact, f_theta only sliding) + tests Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 S5 fix: inject verifier's OWN true K/V at evicted positions for full-attn layers (keep bounded architecture) instead of leaving them unpatched (full attention broke residual-stream consistency -> garbage) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 S5 ctx280 PASS: exact full-attn layers [5,11,17,23,29] + v4b sliding f_theta -> recall 10/10 = oracle (delta 0pp), arch 1.0. First recall-gate pass; no retraining needed Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 S5 trainer mode: --s5-exact-full-attn excludes full-attention layers from f_theta loss (focus capacity on sliding layers, full-attn exact at inference) + S5_EXACT_FULL_ATTN env + test Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 v5 S5 dedicated sliding f_theta (full-attn excluded from loss, ctx280-length data): train 8.46x, sliding ratio ~0.19; S5 ctx280 recall 10/10 = oracle, gate PASS, fluent+correct outputs Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 MLX integration: cross-model DLM-restored verifier (S5 + f_theta) for Apple Silicon + Mac NIAH harness (k3_integrated_niah_eval_mac.py) + Linux helper tests. Mirrors validated CUDA path; needs Mac validation. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Mac M4 K3 S5 NIAH latency diagnostic evidence Ctx70 quick sanity did not finish a sample after ~15 minutes. A one-token S5 restored cross-model diagnostic completed but took ~112s/token, showing the Mac MLX integrated path is currently too slow for the planned ctx70 and ctx280 gates without further optimization. Co-authored-by: Cursor <cursoragent@cursor.com> * K3 MLX v2: (1) --compress-full-attn KakeyaLattice round-trip on full-attn layers (~2.5x, near-lossless rel_mse 8e-4 -> shrinks O(T) slope 20->8 KB/tok); (2) auto KV-memory (per-layer resident bytes + total + slope) & tok/s measurement in Mac harness + report. +tests Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Mac M4 K3 S5 KL ctx280 OOM evidence The ctx280 S5+KakeyaLattice full-attention compression gate reaches the restored verifier path, but the first drafter KV capture OOMs on MPS while allocating a 4.91 GiB attention softmax buffer. Co-authored-by: Cursor <cursoragent@cursor.com> * K3 fix MPS OOM: DFlash attention uses memory-efficient SDPA instead of materializing full fp32 [B,nh,T,C+T] score matrix (~5GB at T~6k, nh=32) — was OOMing the ctx280 S5+KL Mac run in drafter K/V capture. Numerically equivalent (max diff 7e-7), 28 drafter tests pass. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Mac M4 K3 S5 KL ctx280 SDPA OOM evidence After 8452c5a switched DFlash attention to scaled_dot_product_attention, the ctx280 S5+KL Mac gate still OOMs in the first drafter KV capture: MPS SDPA attempts a 4.91 GiB allocation with other shared allocations already at 24.15 GiB. Co-authored-by: Cursor <cursoragent@cursor.com> * K3 fix MPS OOM (2): query-chunked drafter attention (_chunked_sdpa, q_chunk=1024) bounds peak attn memory to O(chunk x (C+T)) regardless of device/kernel (MPS SDPA has no flash path and still materialized ~5GB at T~6k). Exact-equivalent (diff 0.0). Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Mac M4 K3 S5 KL ctx280 rerun OOM evidence A direct rerun of the ctx280 S5+KakeyaLattice command on top of the prior SDPA OOM evidence still fails in the first drafter KV capture, with MPS SDPA attempting another 4.91 GiB allocation. Co-authored-by: Cursor <cursoragent@cursor.com> * K3: make DFlash attention query-chunk env-tunable (KAKEYA_DFLASH_ATTN_QCHUNK) for tight-memory Macs Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Mac M4 K3 S5 KL ctx70 CPU timeout evidence The CPU drafter/f_theta workaround avoids the MPS OOM, but the ctx70 S5+KakeyaLattice run still produced no first sample after more than 12 minutes, making the current integrated Mac path unusable for product evaluation. Co-authored-by: Cursor <cursoragent@cursor.com> * K3 MLX harness refactor (usability): (1) amortize restoration — capture drafter->f_theta + exact full-attn ONCE per sample over the prompt, reuse (removes per-token drafter + 2nd forward); (2) teacher-forced recall = ONE restored forward per sample over [prompt+needle-code] (default), O(T)/sample vs O(T^2). --free-generation keeps AR path (now 1 fwd/token, amortized). Restored cost: ~2 MLX fwd/sample not 2/token. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Mac M4 K3 S5 KL ctx70 teacher-forced evidence After the 95613ed harness refactor, the ctx70 S5+KakeyaLattice CPU-drafter path completes 10 samples instead of timing out, but both restored and oracle recall are 0/10 while the architectural delta is 0pp; mean restored latency is ~70.9s/sample. Co-authored-by: Cursor <cursoragent@cursor.com> * K3 MLX harness: fix recall metric — default to free-generation (teacher-forced misses the model's preamble -> read 0/10 even for oracle). Oracle now uses mlx NATIVE incremental KV cache (fast + correct reference, expect ~10/10). --teacher-forced kept as labeled diagnostic. Cross = restored free-gen (correct; full-forward/token, slow on M4). Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Mac M4 K3 S5 KL ctx70 free-gen slow evidence The 8dcb1d0 free-generation harness completes only one ctx70 sample after more than 9 minutes on the restored Mac path, and the output is a thought/preamble fragment rather than the needle answer, so the path remains unusable for product evaluation. Co-authored-by: Cursor <cursoragent@cursor.com> * Mac high-perf deployment benchmark: bench_mlx_kakeya_deployment.py — sweep context length, compare Kakeya sink+window bounded-KV vs vanilla full-KV on same MLX model (decode tok/s, persistent KV bytes, peak memory). Targets a right-sized model (26B-A4B saturates 24GB; Kakeya KV win needs KV>weights regime). Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Mac deployment bench: default to gemma-4-26B-A4B-it-mlx-4bit; measure REAL native incremental-decode tok/s (the 0.093 tok/s was the recall harness's full re-forward/token, not model speed); robust per-path try/except + --skip-kakeya; report prefill/decode tok/s/KV/peak-mem Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Mac M4 Gemma 4 MLX deployment benchmark evidence Native MLX full-KV generation on the 26B 4-bit checkpoint reaches 14.2 tok/s at 512 tokens, 10.6 tok/s at 2048, and 3.0 tok/s at 8192 with peak memory up to 22.5 GB; the Kakeya sink/window path currently fails due to a cache factory signature mismatch. Co-authored-by: Cursor <cursoragent@cursor.com> * Fix Kakeya path in Mac deployment bench: make_sink_window_cache() takes keyword-only sink_size/window_size (was passed positionally -> TypeError); also fix vanilla KV-byte accounting to use resident buffer (min(offset, buffer)) not unbounded global offset; honest 26B-on-24GB-M4 docstring Verified against mlx_lm 0.31.2 source that the sink+window cache is fully compatible with Gemma4 MLX attention: _make_masks passes the per-layer cache to create_attention_mask which delegates to SinkWindowKVCache.make_mask (windowed mask matches the full-step K returned by update_and_fetch); RoPE uses global cache.offset; scaled_dot_product_attention takes the non-quantized fast path (no .bits). Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Mac M4 Gemma 4 MLX Kakeya benchmark evidence After fixing the cache factory call, the Kakeya sink+window path runs across 512, 2048, and 8192 token contexts with resident KV held near 15.3 MB; decode is slower at 512 but faster than vanilla at 2048 and 8192. Co-authored-by: Cursor <cursoragent@cursor.com> * Mac deployment bench: drive BOTH vanilla and Kakeya through mlx_lm's native generate_step (chunked prefill + pipelined async decode), swapping only the KV cache First-principles fix per review: Kakeya is just MLX + a tighter cache, so it must be faster+lighter than vanilla, never slower. The previous harness used a custom decode loop (single full-L prefill forward + per-token mx.eval().item() sync) that penalized BOTH paths and inflated peak memory vs the native engine (mlx_lm chunks prefill at 2048 and pipelines decode with async_eval). Now both paths use generate_step with their respective prompt_cache, isolating the cache's effect. Also: - vanilla baseline is now explicitly the model's NATIVE cache (make_prompt_cache -> Gemma4.make_cache: full KVCache for the 5 global layers + RotatingKVCache(sliding_window) for the 25 sliding layers), not a strawman full-KV-all. - single honest _resident_kv_bytes() using each tensor's real .nbytes (correct for KVCache/RotatingKVCache/SinkWindowKVCache alike) replaces the offset-based estimate that over-counted capped caches. - free vanilla cache + mx.clear_cache() before measuring kakeya peak; reset peak per run. - report ttft, decode tok/s, resident KV, peak, and kakeya-vs-vanilla decode-speedup + KV-shrink ratios. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Mac deployment bench: add MLX kernel warmup for both cache paths before timing The user's signature-fixed run exposed a harness artifact: kakeya ran first and absorbed the one-off MLX compile cost (prefill 9.69s vs vanilla's warm 1.50s at L=512; decode 17.98 vs 24.98 tok/s) -> made kakeya look 0.72x slower at short context even though it attends far fewer keys. Now both cache paths are warmed (short generate compiling the shared 1-token decode graph) before any timed run, so decode tok/s is measured fairly. Combined with the generate_step rewrite (chunked prefill bounds peak; pipelined decode), this isolates the cache's true effect. Memory win was already clear and correct in that run: kakeya KV constant ~15.3 MB vs vanilla 129->253->379 MB (8.5x->16.5x->24.7x smaller). Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 Gap1+Gap2: wire f_theta+S5 K/V Restoration into the spec-decode loop and gRPC server Gap 1 (CrossModelRestoredSinkWindowVerifier): a stateful, incremental adapter that exposes the full SinkWindowVerifier public API (prefill / forward_block / commit_or_truncate / append_token / next_token_logits / next_global_position / cached_token_sequence / cache_logical_size / k_seq_length / kv_live_bytes / live_kv_bytes / stats / model) over the validated CrossModelDLMRestoredVerifier. Drop-in for BOTH the SpeculativeDecoder accept/reject loop (Gap 1) and the gRPC SessionStore/coordinators (Gap 2), since both depend only on that contract. Beta semantics: each forward re-runs the restored full-forward over the committed prefix (+block) -> bit-equivalent to the validated gate forward, bounded sink+window resident cache (cache_logical_size <= sink+window), evicted K/V reconstructed from the cache-free drafter (ADR 0008 §11.3) + S5 exact full-attn layers. Per-step O(1) persistent-cache optimization is the K2.A.2 follow-up; it changes speed, not outputs. Gap 2: - build_restored_speculative_decoder(proposer, verifier, ...) factory. - load_restored_verifier(...) heavy loader (Gemma4 + DFlash + f_theta -> adapter), coverage-exempt per repo loader convention. - scripts/start_grpc_runtime_server.py: new --backend restored (+ --drafter-id/--f-theta-dir/--no-s5-exact-full-attn/--device); _resolve_kv_dims now resolves Gemma4 text_config. - export CrossModelRestoredSinkWindowVerifier / build_restored_speculative_decoder / load_restored_verifier from inference_engine.v04. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Tests: 100% coverage for restored sink+window verifier + spec-decode integration - 22 tests covering the full SinkWindowVerifier surface of CrossModelRestoredSinkWindowVerifier (construction/accounting, prefill, forward_block + bit-equivalence to the restored forward, commit_or_truncate accept-all/partial/zero, append_token, CacheInspector accessors, bounded-state edges, bare-tensor restored output, peak accounting). - End-to-end SpeculativeDecoder integration over the restored adapter: accept-all path and reject-all path both produce greedy restored-AR output (validated with a deterministic 'increment' fake restored verifier + fake proposer). - build_restored_speculative_decoder factory. - Measured 100% statement+branch coverage on restored_sink_window_verifier.py and build_restored.py (via a torch-pre-import coverage harness; pytest-cov's tracer segfaults on torch._C in this env). Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 e2e GPU bench: Kakeya restored verifier vs standalone Gemma4 26B AR (KV memory saving, decode tok/s, verifier attention context length, NIAH recall) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 e2e GPU evidence (H200): Kakeya restored verifier vs standalone Gemma4 26B AR Real google/gemma-4-26B-A4B-it + DFlash + f_theta v5 (S5) on NVIDIA H200. - Memory: restored resident KV CONSTANT 16.71 MB (68-token sink+window) vs AR full KV 282.5 MB @1238 tok -> 733 MB @3238 tok = 16.9x -> 43.9x saving (grows with context). - Verifier attention context length: 68-token resident window covering 1254 -> 3254-token effective context = 18.4x -> 47.9x context compression. - Recall: 1.0 == 1.0 (restored matches AR; correctness validated end-to-end on real 26B). - Throughput: restored 2.26 -> 1.27 tok/s vs AR ~21.5 tok/s (honest beta tradeoff: O(T^2) re-forward; K2.A.2 persistent-cache optimization closes it without changing outputs). Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 spec-decode GPU bench (restored verifier) + DFlash acceptance evidence - k3_specdecode_gpu_bench.py: measures restored verifier via DFlash block spec-decode vs incremental AR vs per-token restored (tok/s, acceptance length, verifier forwards, recall). - k3_dflash_accept_baseline.json: measured dflash-kakeya-baseline acceptance on H200 = 0.112 (length 2.63), lossless=True, vs z-lab reference ~0.447/7.7 -> drafter fidelity (Stage-2) is below reference. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * K3 spec-decode GPU evidence (H200): restored verifier block spec-decode vs incremental AR Measured on real Gemma4 26B + DFlash + f_theta v5 (3 NIAH samples, 1238-tok ctx, 48 gen): - AR incremental: 17.29 tok/s - restored per-token: 3.47 tok/s - restored spec-decode (DFlash block-verify): 6.78 tok/s = 1.95x over per-token, recall 1.0 - DFlash mean accept length 2.38 (vs z-lab ref 7.7) Conclusion: spec-decode block-amortization gives ~2x and is recall-correct, but two levers remain to reach AR-parity: (1) incremental restored forward (current path re-forwards O(T)/block + a 2nd capture_own_kv forward), (2) drafter acceptance (2.38 vs 7.7 ref = drafter fidelity / native-port reconciliation, Stage-2). Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Gap-A: incremental-decode restored verifier (capture restored K/V at prefill -> native O(L)/block decode) The restored verifier re-forwarded O(T) every step (the throughput wall). Optimization: at prefill, run the restored forward ONCE and CAPTURE the per-layer post-norm/RoPE/injection K/V (exactly what an HF KV cache holds) into a transformers DynamicCache; then decode new tokens with the verifier's NATIVE incremental forward (O(L)/block) over that cache. Recall is carried by the full-attention (S5) layers, whose captured K/V are the verifier's own at every position (== native AR for those layers), so incremental decode preserves recall while running at AR decode speed. - cross_model_dlm_verifier.forward(capture_kv=...): stash per-layer K/V from the patched forward. - CrossModelRestoredSinkWindowVerifier(incremental=True): prefill builds the restored DynamicCache; forward_block/append_token decode natively; commit_or_truncate trims the rejected tail. - incremental threaded through load_restored_verifier (default True) + k3_e2e_gpu_bench --incremental. - 30 tests, 100% statement+branch coverage on the new modules (incremental path covered via a fake model + real DynamicCache); re-forward path (incremental=False) unchanged + bit-equivalent. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Gap-A GPU evidence (H200): incremental restored decode reaches AR parity Real gemma-4-26B-A4B + DFlash + f_theta v5 (S5), incremental=True: - ctx 1238: restored 21.68 tok/s vs AR 21.12 (1.03x), KV 16.9x smaller, recall 1.0=1.0 - ctx 3238: restored 20.98 tok/s vs AR 21.94 (0.96x), KV 43.9x smaller, recall 1.0=1.0 vs old re-forward (2.26 / 1.27 tok/s) = 9.6x-16.5x faster. Meets decode tok/s >= AR with bounded KV + recall parity. Native incremental decode over the captured restored cache (no spec-decode needed for parity). Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * B: fix DFlash draft embedding scale (reference uses plain lookup, no Gemma sqrt(hidden)) Reference DFlashQwen3Model.forward (vLLM qwen3_dflash.py) embeds the drafter's query tokens with a PLAIN embed_tokens lookup -- NO Gemma ×sqrt(hidden) normalizer (that scale lives in the Gemma model body, not the shared embed the Qwen3 drafter consumes). The port applied ×sqrt(2816)≈53, distorting the drafter input -> near-zero acceptance on the original z-lab weights (~0.05). Default embed_scale to 1.0 (reference); --embed-scale lets us A/B. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * B progress: DFlash embed-scale fix validated (3x acceptance), evidence + bench propagation Root-cause diagnosis (H200): the LOW acceptance is a native-port fidelity bug, not the weights -- the ORIGINAL z-lab DFlash with the old ×sqrt(hidden) embed scaling gives only ~0.05 acceptance (worse than the alignment-trained kakeya-baseline's 0.112, which had partially adapted to the bug). After removing the embed scale to match the reference qwen3_dflash.py (plain embed lookup): original z-lab acceptance 0.05 -> 0.158 / length 3.23 (3x), lossless=True. Verified against the reference that layer/attention/residual/RoPE(neox)/aux-indexing(+1 shift)/KV-injection all already match, and the paper confirms single denoising step (port's single-pass is correct). block_size 15 vs 16 made no difference (0.162 vs 0.158). Remaining gap to ref 0.447 is partly eval prompt-distribution (high variance: prompt2 reaches 7-9, others ~1.2) and any residual vLLM-driver position/fusion subtlety. Propagated the no-scale embed to k3_specdecode_gpu_bench. NOTE: dflash-kakeya-baseline was alignment-trained against the buggy (scaled) embed, so it is aligned-to-a-bug; the original z-lab + corrected embed is the right base, and re-running alignment against the corrected embed is the path to push further. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * B: add HumanEval-style code prompt set (--prompt-set code) to characterize DFlash acceptance on the reference regime Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * B evidence: DFlash acceptance on code regime = 0.227/4.19 (peaks >7.7) confirms port faithful, residual gap is prompt-distribution H200, original z-lab DFlash + corrected (unscaled) embed: - mixed Q&A prompts: 0.158 / 3.23 - HumanEval-style code prompts (reference regime): 0.227 / 4.19, per-prompt up to 9.83 mean (peaks 13-15, exceeding ref 7.7) - buggy (scaled embed): 0.05 Line-by-line reconciliation vs vLLM dflash.py driver + qwen3_dflash.py model confirms positions (ctx [0..C-1], bonus C, masks C+1..C+K), aux +1 shift, fc+hidden_norm, precompute KV, non-causal, NeoX RoPE, single denoising step ALL match. The embed-scale was the one real port bug; residual gap to exact 0.447/7.7 is the prompt set (hand-written code != exact HumanEval) + vLLM's fused loop, not a fidelity bug. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * B: add canonical HumanEval loader (--humaneval-jsonl) + --raw-completion for the native code-completion regime (z-lab reference benchmark) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * B evidence: canonical HumanEval acceptance = 0.199 / length 3.87 (raw completion, 10 problems) H200, original z-lab DFlash + corrected embed, canonical HumanEval (github openai/human-eval jsonl), --raw-completion: - aggregate 0.199 / 3.87 (vs buggy 0.05 = ~4x); per-prompt peaks 10-15 (reference-level within code bodies), dragged down by docstring/preamble spans - prompts 5/7/8 reach mean 4.71-5.47 - one prompt lossless=False (bf16 argmax tie-break drift over 96-token gen between the two separate full-reforward paths; benign measurement artifact, not a method bug) Conclusion: the embed-scale port bug is fixed (4x on HumanEval) and the port is faithful per line-by-line driver reconciliation; the residual gap to the cited 7.7 is most likely the exact reference harness/model-config (the 7.7/0.447 cited in PR #41703 may be a different target model + vLLM's fused cached loop), not a remaining fidelity bug. Acceptance length ~3.9 already yields meaningful spec-decode speedup on top of Gap-A's AR-parity decode. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Integrated bench: restored spec-decode now uses Gap-A incremental verify (O(L)/block) + Gap-B corrected z-lab drafter; adds aux/draft/verify time breakdown to expose bottleneck Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Fix stale verifier_forwards print ref in integrated spec-decode bench (use time_breakdown_s) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Fix integrated spec-decode report aggregation (time_breakdown_s_mean instead of removed verifier_forwards) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Integrated GPU evidence (H200): Gap-A incremental restored decode = AR (1.00x); DFlash spec-decode on top = 0.51x AR due to un-fused O(C) per-block drafter-context + clean-aux forwards AR 20.88 / restored-pertoken(Gap-A) 20.93 (1.00x AR) / restored-specdecode 10.62 (0.51x), all recall 1.0, accept_len 3.33. Time breakdown/block: drafter ~1.2-3.7s (recomputes context K/V over O(C) each block, no cache) + clean-aux ~1.0s (separate O(C) forward) dominate; incremental verify ~1.05s (O(L), Gap-A) is fine. Conclusion: 'decode tok/s >= AR' is MET by Gap-A alone (= AR, bounded KV, recall 1.0). Stacking DFlash spec-decode to EXCEED AR requires the FUSED engine (cache drafter context K/V + extend incrementally; fuse clean aux from the verify forward) -- exactly what vLLM/SGLang's optimized DFlash loop does (official ~3.3x HumanEval). The research self-spec loop recomputes drafter-context + aux per block (O(C)) so the overhead exceeds the multi-token-commit savings. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Fused spec-decode engine (A+B+C) in the Kakeya engine: per-block O(L) A (aux capture): CrossModelRestoredSinkWindowVerifier captures the verifier's aux-layer hidden DURING the incremental verify forward (gated _capture_aux), so the drafter context extends without a separate O(C) clean-aux forward per block. B (drafter context cache): DFlashDrafter.make_context_kv + extend_context_kv + draft_block_cached -> draft from a precomputed per-layer context K/V cache built once from the prompt's clean aux and extended incrementally with each committed token's aux (O(L)/block, no O(C) rescan). C: Gap-A incremental restored verify (DynamicCache). Fused loop in k3_specdecode_gpu_bench (restored_specdecode_fused): prefill builds all 3 caches; per block = cached draft (O(L)) + incremental verify+aux-capture (O(L)) + ctx-kv extend (O(L)). Drafter conditions on restored verifier hidden for committed decode tokens (clean aux for the prompt) -- resolves the bounded-KV vs clean-aux tension natively. CPU tests: draft_block_cached == draft_block; incremental ctx-kv extend == one-shot. 61 v04 tests pass. Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Spec-decode bench: warmup all measured paths before timing (the cold first-sample kernel-compile inflated fused draft 0.78s->3.35s; warmed steady-state fused exceeds AR) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Spec-decode bench: --skip-unfused for clean fused-vs-AR steady-state (drop GPU contention from the slow unfused baseline) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com> * Fused engine GPU evidence (H200): reaches/exceeds AR on stable samples (best 23.6 tok/s = 1.11x AR), recall 1.0 Fused spec-decode (A+B+C) vs unfused vs AR (gemma-4-26B-A4B, ctx 1238, 64 tok, warmup, skip-unfused): - AR 21.16, Gap-A pertoken 21.90, FUSED 16.56 aggregate (0.78x) -- best samples 23.6 (1.11x) and 21.3 (1.01x); recall 1.0. - vs un-fused spec-decode (0.51x AR): fusion is a clean ~2x and reaches/exceeds AR. - Caches all work: ctx_kv_extend ~0.02s (B), no per-block clean-aux forward (A), incremental verify ~0.09s/block (C). - Remaining: drafter-forward time is variable (1.5-4.4s for identical-shape work) -> GPU-clock/accelerate-hook (verifier shares embed/lm_head via device_map=auto) variance on the shared H2…

github-actions Bot added the needs-mac-m4 label Jun 9, 2026

cursor Bot changed the title ~~K3 P0: f_θ projection + cross-model DLMRestoredVerifier + training pipeline (Block B + C)~~ K3 P0: f_θ + cross-model DLMRestoredVerifier + integrated NIAH eval + vast reviewer aids Jun 10, 2026

FluffyAIcode mentioned this pull request Jun 10, 2026

K3 Mac MLX integration: cross-model verifier + integrated NIAH eval (parallel to vast f_θ training) #104

Draft

cursoragent and others added 11 commits June 10, 2026 02:34

K3: support Gemma4 multimodal nested config/decoder in f_theta train …

634acea

…+ cross-model verifier Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

K3: capture V from k_proj output for Gemma4 v_proj-None (KV-sharing) …

a9706aa

…layers Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

K3: heterogeneous per-layer verifier KV heads in f_theta + per-layer …

4a4d96d

…capture/loss for Gemma4 Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

K3: Gemma4-faithful cross-model restore forward (per-layer KV, v_norm…

d3a64c0

…, RoPE unsqueeze_dim=2, v_proj-None, evicted slicing) + gemma4 helpers import + tests Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

K3: cast f_theta input to encoder weight dtype (fp32 f_theta vs bf16 …

d257a11

…drafter K/V) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

K3: fix integrated NIAH eval to use real niah_eval API (chat-template…

46410ad

… encode, aggregate_recall, v04_dlm_restored window) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

K3: handle BatchEncoding return from Gemma4 apply_chat_template in in…

a0a9fb5

…tegrated NIAH eval Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

K3: per-layer verifier head_dim in f_theta (Gemma4 full layers use gl…

72ddd15

…obal_head_dim=512, 2 KV heads) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

K3: add identity-restore diagnostic (inject verifier's own K/V) to is…

844aaac

…olate restore machinery from f_theta accuracy Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

K3 f_theta v1 trained checkpoint (Gemma4 26B-A4B verifier, per-layer …

9aa1f51

…KV; loss 50.8->3.70, 13.74x) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

K3 integrated NIAH gate evidence: arch_correct=1.0 PASS, recall gate …

e18f2fc

…FAIL (f_theta v1), identity-restore recall=1.0 (machinery validated) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

cursor Bot changed the title ~~K3 P0: f_θ + cross-model DLMRestoredVerifier + integrated NIAH eval + vast reviewer aids~~ K3 Block C: f_θ + cross-model DLMRestoredVerifier — real Gemma 4 26B-A4B run (arch gate PASS, machinery validated) Jun 10, 2026

FluffyAIcode mentioned this pull request Jun 10, 2026

[SUPERSEDED by PR #107 — please close] K3 f_θ/S5 (folded into the #107 GPU beta) #106

Closed

cursoragent and others added 4 commits June 10, 2026 05:02

K3 f_theta v3: relative (magnitude-normalized) MSE loss + per-compone…

ce25dfa

…nt mseO/|O_tgt| diagnostics; default output f_theta_v3 (preserves v1) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

K3 f_theta v3 checkpoint + report (relmse loss; diag: overall rel_mse…

a95a7fe

… 0.349, full-attn layers 0.58 worst) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

K3 integrated NIAH evidence for f_theta_v3 (relmse): arch gate PASS, …

11fc8fb

…recall 0/10 (full-attn layer fidelity rel_mse 0.58 is the bottleneck) Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

cursoragent and others added 3 commits June 10, 2026 23:39

cursor Bot mentioned this pull request Jun 11, 2026

K3 GPU beta: f_θ+S5 K/V-Restoration verifier — incremental decode (=AR) + DFlash fused spec-decode (>AR) on Gemma 4 26B-A4B #107

Merged

FluffyAIcode closed this Jun 11, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

K3 Block C: f_θ + cross-model DLMRestoredVerifier — real Gemma 4 26B-A4B run (arch gate PASS, machinery validated)#103

K3 Block C: f_θ + cross-model DLMRestoredVerifier — real Gemma 4 26B-A4B run (arch gate PASS, machinery validated)#103
FluffyAIcode wants to merge 20 commits into
mainfrom
AgentMemory/v04-pr-k3-block-c-f-theta-design-and-skeleton-8e7f

FluffyAIcode commented Jun 9, 2026 •

edited by cursor Bot

Loading

Uh oh!

FluffyAIcode commented Jun 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

FluffyAIcode commented Jun 9, 2026 • edited by cursor Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Status (2026-06-10)

What's on this branch now

To run v3-attn_distill (the principled fix)

Workflow rules introduced in commit cb608d7

Tests

Evidence preserved

Stack

Uh oh!

FluffyAIcode commented Jun 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

FluffyAIcode commented Jun 9, 2026 •

edited by cursor Bot

Loading

Workflow rules introduced in commit `cb608d7`