K3 Step 3b: Mac M4 cross-runtime DFlash speculative decoding eval (MLX verifier + PyTorch drafter)

[FluffyAIcode]

Why this PR (Step 3b)

Mac M4 mirror of PR #93's CUDA spec decode eval (k3_dflash_specdecode_eval.py), crossing the MLX verifier ↔ PyTorch drafter runtime boundary so the acceptance metric is directly comparable to PR #93's CUDA evidence.

User's verifier load + forward verified end-to-end 2026-06-09 (commit 6f7f516):

load:    6.27s, 15.26 GB driver memory
forward: 8 tokens / 10.33s = 0.77 tok/s
peak:    15.49 GB  ← 8.5 GB headroom for the drafter on Mac M4 24 GB

Step 3b builds on this verified baseline by adding the spec decode loop.

Architecture

MLX 4-bit verifier (mlx_lm) on Apple Silicon
          ↕  scripts/research/k3_dflash_mlx_bridge.py
PyTorch DFlash drafter on MPS / CPU (PR #93 native impl)

The verifier MUST be 4-bit MLX on Mac M4 24 GB — bf16 won't fit the 26B model. The drafter (0.43B) stays PyTorch (PR #93's DFlashDrafter is pure torch and runs on MPS comfortably).

Three new files

1. scripts/research/k3_dflash_mlx_bridge.py (~290 LOC)

Cross-runtime bridge primitives:

mx_to_torch(x, dtype, device)              # mlx.array → torch.Tensor (numpy intermediate)
torch_to_mx(t)                              # torch.Tensor → mlx.array
MLXVerifierAuxProvider(model, aux_layers)   # PR #93's AuxHiddenProvider contract
build_mlx_verifier_callbacks(model, ...)    # (embed_fn, lm_head_fn) for drafter
mlx_verify_block(model, committed, draft)   # Mac equivalent of CUDA verify_block

MLXVerifierAuxProvider replays Gemma4TextModel.__call__ layer-loop with intermediate capture at aux_layer_ids (PR #93's +1 shift convention). Verified compatible with mlx_lm 0.31.3's layer-loop structure.

2. scripts/research/k3_dflash_specdecode_eval_mac.py (~280 LOC)

Mac variant of PR #93's eval. Same prompt corpora (PROMPTS + HELD_OUT_PROMPTS). Same output JSON schema (schema_version: 1, kind: 'k3_dflash_specdecode_acceptance_mac').

Spec decode loop:

for prompt:
    while len(generated) < max_new_tokens:
        aux_ctx, bonus = MLXVerifierAuxProvider(committed)
        proposal = DFlashProposer.propose_block(committed, L, ...)
        accepted, _ = mlx_verify_block(committed, [bonus] + proposal.tokens)
        commit accepted; bonus is always-correct
    ar_out = mlx_greedy_ar(committed)  # for lossless check

Aggregate report includes both reference_humaneval (vLLM PR #41703: 0.447 / 7.7) and reference_cuda_held_out (PR #93's k3_dflash_specdecode_corpus_heldout.json: 0.107 / 2.45) for direct comparison in the same JSON.

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

Mac M4 reviewer aid. Six pre-flight checks each fail fast with actionable error:

#	Check
1	mlx_lm importable
2	PyTorch + MPS available (or DRAFTER_DEVICE=cpu fallback)
3	Verifier dir present + has config.json
4	Drafter source local-path validation (mirror PR #99's fail-fast)
5	tokenizer_config.json's extra_special_tokens IS a dict (PR #101 patch state)
6	PR #93's DFlashDrafter importable (main is current)

Env knobs: VERIFIER_PATH, DRAFTER_ID, DRAFTER_DEVICE, N_PROMPTS, BLOCK_SIZE, NUM_STEPS, MAX_NEW_TOKENS, HELD_OUT.

Tests (Linux CI: 12 new tests)

tests/research/test_k3_dflash_mlx_bridge.py

Covers the testable surface — mx_to_torch bridge utility using numpy stand-ins for mx.array. The MLX-touching paths require Mac M4 to validate end-to-end.

TestMxToTorch (7):
  ✓ basic shape preserved (rank-3+)
  ✓ dtype override (fp32 → fp64)
  ✓ device default 'cpu' / explicit 'cpu'
  ✓ value preservation (allclose)
  ✓ COPY INDEPENDENCE — mutating numpy source must NOT mutate torch
    result. Critical for spec decode correctness; otherwise verifier
    hidden state could leak through the bridge.
  ✓ int dtype passthrough (logits/argmax)
  ✓ high-rank tensor handling [B, H, T, D, ...]

TestModuleInterface (5):
  ✓ public name surface stable
  ✓ MLXVerifierAuxProvider.__init__ signature stable
  ✓ build_mlx_verifier_callbacks signature stable
  ✓ mlx_verify_block signature stable
  ✓ __all__ contains public names

tests/inference_engine/v04/ + tests/research/: 327 passed

Validation gate (Mac M4 user run)

End-to-end on Mac M4. Produces JSON with same schema as PR #93's CUDA evidence. Acceptance rate is expected to be in the same neighbourhood as the CUDA held-out (~10.7% / 2.45 length, with the same training-data-corpus-size limitation), since the verifier and drafter weights are identical and only the runtime differs.

# After this PR + PR #100 + PR #101 merge:

# (one-time, idempotent — if already run, skipped):
python3 scripts/research/k3_patch_gemma4_tokenizer_config.py \
    models/gemma-4-26B-A4B-it-mlx-4bit

# held-out eval (directly comparable to PR #93 CUDA held-out):
HELD_OUT=1 bash scripts/review_pr_k3_dflash_specdecode_on_mac.sh

# in-domain eval:
bash scripts/review_pr_k3_dflash_specdecode_on_mac.sh

Stack

main (post #93 + #99 + #94 merged)
├── PR #100 (Step 3a — DFlashProposer platform-aware peak memory) [draft]
│   └── THIS PR #102 (Step 3b — Mac MLX spec decode eval) [draft]
└── PR #101 (Step 4 phase 1 — Gemma 4 tokenizer patch) [draft]

PR #100 is required for Step 3b (DFlashProposer.propose_block needs platform-aware peak memory on MPS). PR #101 is required at runtime on Mac (extra_special_tokens patch). Both unblock Mac MLX critical path together.

Outstanding work

• Step 5 — K2.A backport PR (dtype hotfix / transpose hotfix / HF Cache contract from PR [SUPERSEDED BY #93] K3 Mac smoke: alignment baseline default + proposer K/V capture (K2.A fixes need backport) #98 superseded into K2.A.1 / K2.A.2 lineage)
• Step 6 — alignment training corpus expansion (vast time; user's --n-prompts 64 → 256 → 1024 path to push held-out acceptance toward reference 44.7%)

Net effect

After this PR + PR #100 + PR #101 land, the user can run bash scripts/review_pr_k3_dflash_specdecode_on_mac.sh and get production-grade cross-model speculative decoding evidence on Mac M4 with the same metric structure as PR #93's CUDA evidence. No fallback, no mock, no fake — real verifier weights + real drafter weights + real spec decode loop + real lossless_vs_AR check.