feat(mtp): aligned head cache - acceptance 74.8 -> 92.0 percent, depth 2 sticks
Implements upstream PR #990's "batched MTP cache commits ... maintaining cache
alignment between backbone and MTP head", which vMLX never had. The MTP head
was drafting from either a single fused pair (fresh-per-cycle) or a history
with a hole at every bonus token (retained) - both starve it of the context it
was trained with, and that, not head precision or trunk quantization or CRACK,
was the acceptance gap to the reference implementations.
Mechanics: every confirmed token's (backbone_hidden_i, token_{i+1}) pair is
committed through the head in the SAME forward that drafts the next token -
the level-0 draft samples from the last position, so a multi-token commit is
free. Chain pairs from deeper draft levels are trimmed after every verify, so
a rejected draft can never persist in the head's context (the poisoning that
ledger 343 measured). On reject, the confirmed prefix plus the correction is
committed with proper backbone hiddens; the head cache is never recreated.
Measured live in the app, Qwen3.8-27B-JANG_4D-CRACK, code workload,
replay-skip on:
acceptance 74.8 -> 92.0 percent (1627/1768)
conditional d2 0.86-0.92 - depth 2 HOLDS for the first time; the
controller climbed D1->D2->D3, kept D2, demoted D3 (0.646)
tokens/cycle 1.75 -> 2.4
speed 30.0 / 32.3 / 33.9 / 36.1 t/s vs 22.1 baseline (+45-63%)
Output verified correct (coherent Python, proper RLock usage, docstrings,
finish=stop, no fallback).
Default ON via VMLX_MTP_ALIGNED_HEAD_CACHE (0 reverts to fresh-per-cycle).
_step_native_mtp_head now reshapes tokens (1, -1) so the level-0 call accepts
the multi-token commit; mocks updated to record the drafting pair (last
position) and the depth-3 partial-reject test now pins retain+trim instead of
recreate. MTP suite 336 green. J
Jinho Jang committed
21b420beeaade01a2fa08a480bd22cbe61e94ae9
Parent: 2b4d453