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

[FluffyAIcode]

Why this PR (parallel-track to vast f_θ training)

Per user 2026-06-10: "Mac MLX 集成不能和 vast 训练同步进行么?"

Yes. Mac MLX integration depends on the f_θ checkpoint structure (locked in PR #103) but not its trained values. This PR can be reviewed and merged in parallel with vast training; runtime uses the vast-produced f_θ when both land.

Mac mirror of PR #103's CUDA integration:

PR #103 (CUDA, parent):
  inference_engine/v04/cross_model_dlm_verifier.py
  scripts/research/k3_f_theta_train.py             ← runs on vast
  scripts/research/k3_integrated_niah_eval.py
  scripts/review_pr_k3_f_theta_train_on_vast.sh
  scripts/review_pr_k3_integrated_niah_on_vast.sh

THIS PR (Mac MLX):
  inference_engine/v04/cross_model_dlm_verifier_mlx.py
  scripts/research/k3_integrated_niah_eval_mac.py
  scripts/review_pr_k3_integrated_niah_on_mac.sh

Three new files

1. inference_engine/v04/cross_model_dlm_verifier_mlx.py (~370 LOC)

MLXCrossModelDLMRestoredVerifier class. Same architectural contract as PR #103's CUDA version; differences are runtime-specific:

	CUDA	Mac MLX
K/V injection	in-place via K_local.index_copy_	mx.where(evicted_mask, injected, original) (functional)
KV-shared layers	n/a (HF Gemma 3)	Gemma 4 has num_kv_shared_layers; injection skipped on those
K-eq-V optim	n/a	Full-attention layers with attention_k_eq_v use injected K for both K and V
Attention patch	attn.forward = ...	attn.__call__ = ... (mlx.nn.Module per-instance override)
f_θ output bridge	direct torch in-place	numpy → mx.array up front, sliced per-layer in patched call

Layer wiring (_MLXLayerWiring per layer: layer_idx, has_kv, is_sliding, use_k_eq_v, n_kv_heads, head_dim) computed once at __init__ so per-forward patched __call__ is fast.

mlx is imported lazily inside method bodies — module is importable on Linux CI.

2. scripts/research/k3_integrated_niah_eval_mac.py (~230 LOC)

Mac mirror of PR #103's k3_integrated_niah_eval.py. Same JSON schema (kind: 'k3_integrated_niah_acceptance_mac' — Mac suffix to distinguish from CUDA evidence; rest of structure identical for diff-ability).

3. scripts/review_pr_k3_integrated_niah_on_mac.sh (~190 LOC)

Mac M4 reviewer aid. 7 pre-flight checks (each fails fast):

1. mlx_lm importable
2. PyTorch + MPS available (or DRAFTER_DEVICE=cpu)
3. Verifier dir + config.json present
4. tokenizer_config.json's extra_special_tokens IS a dict (PR K3 Step 4 phase 1: bug4 fingerprint for Gemma 4 tokenizer + patch script #101 patch state)
5. Drafter LFS pulled (>100MB safetensors)
6. f_θ checkpoint present (config.json + weights.pt)
7. PR K3 Block C: f_θ + cross-model DLMRestoredVerifier — real Gemma 4 26B-A4B run (arch gate PASS, machinery validated) #103 + this PR's modules importable

Tests (9 new)

TestModuleSurface (2):           module imports + class signature stable
TestLayerWiringDataclass (1):    _MLXLayerWiring fields present
TestConstruction (5):            dim alignment / mismatch / negative sink-window / missing layers
TestLayerWiringDerivation (1):   default stub attrs propagate

Linux CI: 363/363 passing (354 pre-existing + 9 new)

Stub mlx verifier (no actual mlx import on Linux) reproduces the .model.layers[i].self_attn shape with synthetic attributes — enables construction-path coverage without Apple Silicon.

Validation gate (Mac M4 user run, post vast f_θ training)

End-to-end on Mac M4. JSON schema mirrors PR #103's CUDA evidence — direct comparison via aggregate.gate.{architectural_correctness, recall_delta_within_5pp, memory_under_24gb}.

# Mac workflow (after PR #103 + this PR merge + vast f_θ training):

# (one-time tokenizer config patch, idempotent):
python3 scripts/research/k3_patch_gemma4_tokenizer_config.py \
    models/gemma-4-26B-A4B-it-mlx-4bit

# Pull the vast-trained f_θ checkpoint:
git pull --ff-only origin main

# Run integrated NIAH eval:
bash scripts/review_pr_k3_integrated_niah_on_mac.sh

# → results/research/k3_integrated_niah_mac_ctx{70,280}_<stamp>.json

Stack

main (post #93 + #99 + #94 + #100 + #101 + #102)
└── PR #103 (CUDA: f_θ + cross-model + train + integrated NIAH)
    └── THIS PR #104 (Mac MLX: cross-model_mlx + integrated_niah_mac)

After both merge:
  vast workflow: f_θ training → f_θ checkpoint → CUDA NIAH evidence
  Mac workflow:  same f_θ checkpoint → Mac NIAH evidence
  
Both produce diff-able JSON with the same schema. K3 product gate 
empirically validated on BOTH runtimes once user runs both 
reviewer aids.

Net effect

After this PR + PR #103 land + vast f_θ training completes + you run the two reviewer aids (one on vast, one on Mac), the K3 product gate is empirically closed on both CUDA and Mac M4. The integrated Kakeya inference architecture (sink+window verifier + dLM proposer K/V Restoration via f_θ) is validated end-to-end with real numbers in JSON evidence.