perf(speculative): single-slice the MTP accept hook and drop the drafter's materialized causal mask

[inureyes]

Summary

The two Phase 4 structural cleanups from #1185, plus the measurement that closes out the drafter side of that issue.

• accept_verified_tokens built its paired-hidden block with a chain of per-position concatenate calls. The positions are the contiguous range starting at keep (drafts keep..accepted, then position accepted for the bonus), so one slice of verify_hidden replaces n slices plus n-1 intermediate buffers.
• forward_hidden_stack materialized an [s, s+offset] additive causal mask on every multi-token call. Multi-token forwards now route through the shared causal_attention wrapper, which derives the same bottom-right alignment from the cache K length and selects MLX's native "causal" mask mode (and the M5 NAX causal kernel where present) without building a mask array. Single-token draft steps keep the maskless dispatch byte-for-byte unchanged.

The caching the issue suggested for the mask does not apply: the offset advances every round, so an (s, offset) key never repeats. Removing the materialization is the version of that cleanup that actually exists.

Related issues

Refs #1185 (Phase 4).

Type of change

• perf — structural cleanup on the drafter's per-round path

Why this cannot move the output

Drafter numerics only influence acceptance (layer.rs module docs): the target verifies every proposal, so emitted tokens are the target's own greedy choices regardless of what the drafter proposes. Verified directly below anyway, and on this run even the drafted blocks came out identical.

Measurements

Apple M5 Max 128 GB, macOS 26.6.1, qwen3.8-27b-4bit + qwen3.8-27b-mtp-4bit, block 3, temperature 0, offline CLI. Arms: main at 9e2c6675 (before) and this branch (after). ABBA blocks, two warm-ups discarded, 6 samples per arm, 8 s cooldowns.

Byte identity, 400 tokens: classic == MTP(before) == MTP(after). The two MTP arms also produced identical round statistics (113 rounds, acceptance 0.8274), so the causal-mode switch did not perturb the drafter's proposals on this run.

Throughput, 300 tokens (tok/s): before median 54.81 (52.93-55.90), after median 55.27 (53.52-56.29). +0.8%, inside the spread. This is a wash, as expected at a ~370-token context where the mask is small; the change is structural, not a headline number.

Round-loop diagnostics (one logged run per arm, same 113 rounds):

per round	before	after
draft_ms	2.740	2.727
accept_hook_ms	2.769	2.744
verify_forward_ms	41.85	42.74

The drafter-side terms move at the sub-1% level; the verify difference is run noise on an untouched path.

Gemma 4 regression check (gemma-4-12b-it-4bit + 4-bit assistant, block 5, 300 tokens, branch binary): 71.05 tok/s, acceptance 0.7841, emitted_per_verify 3.6463 — the same chain figure PR #1214 recorded for this pairing. The gemma4_assistant drafter and the shared round loop are untouched by this PR.

Phase 2a, considered and declined

Phase 2a (fused LM-head + argmax) targeted 21.4% of a drafter step that #1203 has since cut from 10.5 ms to 2.7 ms per round. The remaining prize is ~0.6 ms per forward, ~1.2 ms per round, under 2% of a 62 ms round, and it costs a custom quantized-matmul-with-argmax Metal kernel that has to survive every future MLX pin bump. Declined on that ratio; recorded here so the issue can close over it.

Test plan

• cargo fmt --check
• cargo clippy --workspace --all-targets --features metal,accelerate -- -D warnings
• cargo test --release -p mlxcel-core --features metal,accelerate qwen3_5_mtp (22 tests) and speculative (129 tests)
• Real checkpoint validation as above: byte identity, ABBA A/B, round diagnostics, Gemma 4 smoke