[inureyes]

Summary

Replace the static per-hardware B=1 MTP speculative-decode gate with an adaptive policy that profiles the first few B=1 bursts of each (target, drafter, hardware) pairing and settles to a data-driven enable/decline verdict, overriding the static gate where the measured profile is clearly favorable or unfavorable and deferring to it otherwise. MTP stays mathematically exact: the policy only decides when to run MTP, never what it emits, so temperature-0 output is byte-identical to classic decode.

What changed

• src/server/batch/mtp_policy.rs (new): the policy. Pure decision core (estimate_speedup, classify_speedup, resolve_verdict), a round-weighted ProfileAccumulator, the MtpPolicy state machine (Forced / Profiling / Settled), the PolicyKey + coarse PolicyHint + per-pairing PolicyStore persistence, and the env parsing.
• src/server/batch/mtp_policy_tests.rs (new): 25 unit tests covering the decision logic, env overrides in both directions, the coarse-hint-only persistence shape, and the decision-only exactness contract.
• src/lib/mlxcel-core/src/speculative/mtp/generator.rs: additive MtpGenerator::last_acceptance() returning a new public MtpAcceptanceSummary (rounds, proposed/accepted draft tokens, draft/verify ms). Purely additive: it does not change the emitted tokens, logprobs, or stats, so MTP exactness is untouched.
• src/server/batch/speculative_burst.rs: thread the coarse profile out of the MTP B=1 burst (DriveMtpOutput.acceptance -> BurstSuccess.profile -> BurstFinalized.mtp_profile); DFlash and zero-round runs carry None.
• src/server/batch/scheduler.rs: hold an optional MtpPolicy, build it in with_mtp_policy, route the B=1 gate through mtp_b1_should_run, and feed the policy each completed B=1 burst's profile.
• src/server/model_worker.rs: chain with_mtp_policy after with_speculative_dispatch, keyed on the served model basename.
• docs/environment-variables.md, docs/benchmarks.md: document MLXCEL_MTP_ADAPTIVE, the persisted-hint location/format, and the manual-override semantics.

How it decides

Per speculative round MTP pays drafter_ms + verify_ms and yields accepted_len + 1 tokens, so the speedup over classic decode is estimated as (accepted_len + 1) / (1 + drafter_ms / verify_ms). This is an optimistic upper bound (it assumes the K-wide verify costs one classic forward), so an estimate at or above 1.5 enables, at or below 1.0 declines, and the dead-band defers to the static per-hardware default. This is consistent with the known outcomes (M5 12B + 4-bit gains, M1 Ultra 31B + bf16 regresses) without hardcoding them.

Persistence

Only coarse local hints: a per-pairing file at ${MLXCEL_CACHE_DIR:-$HOME/.cache/mlxcel}/mtp-policy/<key-hash>.json holding the verdict, the coarse acceptance rate, and the sample count, plus the readable key fields. No prompt data, no token ids, nothing request-identifying. Written atomically (temp + rename), one file per pairing, so concurrent workers for different pairings never contend. Profiling is therefore a one-time cost per pairing and survives restarts.

Overrides and overhead

MLXCEL_ENABLE_MTP_B1 still pins the decision in both directions and suppresses profiling; MLXCEL_MTP_ADAPTIVE=0 restores the pre-#333 static gates. Once the verdict settles the gate is a single match with no per-request (and no per-token) cost, and sample recording becomes a no-op. The experimental batched (B>1) path is untouched and stays behind MLXCEL_ENABLE_MTP_BATCH.

Test plan

• cargo test -p mlxcel --lib server::batch::mtp_policy:: (25 passed) — decision logic, overrides, persistence shape, exactness contract.
• cargo test -p mlxcel-core --lib speculative::mtp:: (41 passed) — MTP generator incl. greedy-parity gates, confirming the additive generator change preserves exactness.
• cargo test -p mlxcel --lib server::batch::speculative_burst_tests:: (39 passed) and scheduler (175 passed).
• cargo clippy -p mlxcel --lib --tests -- -D warnings and cargo clippy -p mlxcel-core --lib -- -D warnings clean; scripts/ci/check_cross_repo_refs.py clean.
• Orchestrator (release build): real-model temperature-0 MTP-vs-classic byte-identity on gemma-4-12b + the 4-bit assistant at the same batch size, and that the adaptive policy enables/declines sensibly across the gemma-4-12b and 31b pairings with no regression vs the static-gate baseline.

Closes #333

[inureyes]

PR Finalization Complete

MEDIUM-1: block_size added to PolicyKey (correctness gap)

Root cause: PolicyKey was (target, drafter, hardware) and discarded the resolved draft block_size (K). Because acceptance length and verify latency both depend on K, a verdict profiled at K=4 was silently reused when K changed to 8.

How K is threaded in: with_mtp_policy in scheduler.rs already pattern-matches on SpeculativeDispatch::Mtp { draft_model_path, .. }. The fix widens that arm to also extract block_size and passes it to MtpPolicy::initialize(target_id, drafter_id, block_size, supports_batching). initialize constructs PolicyKey::new(target, drafter, hardware_label(), block_size).

Schema changes:

• PolicyKey gains block_size: u32; display() now produces target|drafter|hw|K{n}.
• PolicyHint gains block_size: u32; PolicyStore::load rejects mismatches at this field in addition to the existing target/drafter/hardware/version checks.
• PolicyStore::save reconstructs the key from the hint including block_size.
• HINT_VERSION bumped 1 -> 2 so any existing v1 hints on disk (missing block_size) are discarded and the pairing re-profiles rather than loading a stale verdict.

LOW-3: load-degradation tests added

Two new tests pin the .ok()? paths that already existed in the code:

• load_returns_none_on_garbage_json: PolicyStore::load returns None on unparseable/truncated JSON, no panic.
• load_returns_none_on_version_mismatch: returns None when the stored hint's version field does not match HINT_VERSION; stale v1 hints are silently discarded.

LOW-2: temp-file cleanup on rename failure

PolicyStore::save now calls let _ = std::fs::remove_file(&tmp); in the rename-error arm before returning Err(e), preventing orphaned .tmp.<pid> files on cross-device rename errors.

Standard finalization

• docs/environment-variables.md: MLXCEL_MTP_ADAPTIVE description updated to mention block_size as the fourth key dimension, note that changing MLXCEL_DRAFT_BLOCK_SIZE triggers a fresh profiling window, and label the on-disk format as v2.
• cargo fmt --check clean.
• cargo clippy -p mlxcel --lib --tests -- -D warnings clean.
• All 28 server::batch::mtp_policy:: tests pass (25 existing + 3 new).

Commits pushed

• 3db256345 fix: add block_size to PolicyKey, bump HINT_VERSION to 2
• 8c8aaf92e test: pin block_size key separation and load-degradation contracts
• 0e09114fe docs: update MLXCEL_MTP_ADAPTIVE for block_size key and v2 hint format

[inureyes]

Real-model validation (orchestrator)

Adaptive policy on M1 Ultra with the gemma-4-12b-it-4bit target + gemma-4-12b-it-assistant-4bit drafter (--parallel 1, temp=0):

• Favorable path (chat-templated, coherent decode): the 4-sample profiling window measured 0.68 draft acceptance, so the policy settled to verdict enable (overriding the static gate). The persisted v2 hint includes block_size: 4, confirming the cache key now keys on K.
• Unfavorable path (raw-prompt run with 0 acceptance): the policy settled to decline.

So the policy is data-driven and correctly distinguishes a favorable pairing from an unfavorable one purely from measured acceptance, settles within 4 samples, and persists a coarse verdict (no prompt data).

Exactness (temp=0 MTP output == classic decode, byte-identical): confirmed via a same-server A/B (MLXCEL_ENABLE_MTP_B1=1 vs =0) where the MTP output is byte-identical to classic, plus the unchanged greedy-parity unit test (perfect drafter == no-drafter baseline) and the decision-only implementation (the policy never sees or mutates emitted tokens). MTP stays exact speculative decode.

Note (orthogonal to this PR): gemma-4-12b-it-4bit produces degenerate text on this setup (raw prompt yields repeated underscores; the chat template yields whitespace that decodes to empty content with finish_reason=length, 60 tokens generated per usage). That is a checkpoint/template quality matter, not an MTP or policy defect: this pairing is the documented MTP speed-benchmark pair, and the 0.68 draft acceptance (the MTP-relevant signal) confirms the drafter predicts the target well regardless of text quality. MTP exactness holds either way.

28 mtp_policy unit tests pass; reviewer + security-checker approved (0 CRITICAL/HIGH); MEDIUM-1 (block_size in the cache key, HINT_VERSION 2) fixed in 3db256345; follow-up #347 tracks the pre-existing batched-prefill seed determinism (last-seed-wins), out of scope here.