[jianc99]

Supersedes #526 (thanks @linson007); fixes #437 (thanks @ezoushen).

This keeps #526's same-forward check, with the penalty kernels' host
reference in place of its fp64 model, without the cold-score probe it
kept as a diagnostic or its run logs, and with the logits captured as
dispatch profiles are rather than through an observer callback.

The runtime oracle checked a penalized greedy transcript by scoring it
again: for each emitted token it cold-prefilled the context before it,
applied the penalties on the host to that prefill's logits, and required
the emitted token to trail the best option by no more than graphDrift,
3% of the values compared, beyond a tolerance of twice the unpenalized
transcript's drift and at least 0.1. Most decisions come from the verify
graph, and the prefill and verify graphs can rank differently the tokens
the penalties leave near a tie; how the kernels compile moves which
tokens those are. So the check, added to catch a count the runtime got
wrong, refused correct decisions.

linson007 reproduced such a refusal on an M5 Pro (#526): 1.4.0's oracle
on the Qwen3.6-35B-A3B UD-IQ3_XXS GGUF with bf16 KV fails under shader
validation at repetition 1.3 and presence 1.5 (reference_loss 0.205502),
and on the 64 GB M5 Pro the same oracle passes twice without validation
and fails under it with that loss. In #437, ezoushen's run of a 4-bit
MLX fine-tune of the 35B on an M5 Max running 1.3.2 failed without
validation and passed with it (0.636419); at that decision ezoushen
found the comma's raw logit 1.32 from verify and 2.00 from a cold
prefill of the same context, while the GPU's penalty arithmetic matched
a host reference over the whole vocabulary.

The oracle now checks each penalized greedy decision against the logits
of the forward that made it (PenaltyProbe): the first token each
penalized greedy run selects and every verify row its decode retains, in
each preemption mode. Every logit of such a row must be the one the
penalty kernel writes from the raw logit and the counts of the tokens
selected before it, which the probe rebuilds from the fixture rather
than reading the runtime's tables: the prompt, which only repetition
reads, the transcript, the pending anchor and the row's accepted draft
prefix. The selected token must be the best selectable one. The check is
the penalty kernels' host reference, referencePenalty and
penaltyMatches, which moves from the Metal test of those kernels to
dev/tests/engine/PenaltyReference.hpp for both tests: an untouched,
prompt-only, saturated or unselectable logit matches bitwise, and a
counted one within 2 ulps of the larger operand of its subtraction. The
reference now leaves an unselectable logit (NaN or -inf) as it is, as
penalize_logit does, where it turned -inf into -FLT_MAX, and
target-sampling checks that against both kernels. Each run the probe
covers prints penalty_same_forward=PASS and its rows.

The cold-score comparison goes, with graphDrift and the drift, tolerance
and losses it printed. Kept as a diagnostic, as #526 kept it, it would
assert nothing and still take up to 216 cold prefills for its nine
transcripts: 61 s of the 194 s that test-real's first invocation took on
the Qwen3.8-27B UD-Q4_K_M GGUF on the M5 Max (36 GB). Each penalty_probe
line keeps its setting and whether it changed the control's transcript,
as one setting must. The runs and their logs stay in #526, not in the
tree.

The runtime gives the oracle the logits as it gives dispatch profiles
(BackendInstrumentation::setDispatchProfiling and takeDispatchProfile):
while RuntimeInstrumentation::captureTargetLogits(runtime, true) holds,
each command that selects first tokens or verifies copies each batch
lane's rows in a host step before the policy rewrites them in place, row
0 of a selection or a verify's verified rows, and
takeTargetLogits(runtime) reads and clears the copies of the last such
command; pendingToken(runtime, request) names the anchor the request's
next decode emits first. Without capture a command adds no host step,
and all of it is behind SPLASH_RUNTIME_INSTRUMENTATION: the production
build has none of it, and no shader or default changes.

On the Ornith-1.5-122B, the check told a penalty mismatch from drift:
the cold-prefill comparison, as 1.4.0's oracle makes it, fails there
with int8 KV (loss 1.33) and passes with --kv-format bf16, while every
decision matches the logits of its own forward with either, so the
mismatch is int8 KV's drift between the graphs, not a fault of the
penalty kernels.

Tests: at the top of this stack, on the 64 GB M5 Pro, test-engine-cpu
and, under shader validation, check-native-metal (target-sampling among
it) pass, and test-real passes on the Qwen3.6-35B-A3B UD-Q4_K_M GGUF and
4-bit MLX, the Qwen3.8-27B 4-bit MLX, and the Ornith-1.5-122B UD-Q2_K_XL
GGUF with int8 KV and with --kv-format bf16, penalty_same_forward
passing with both. On the M5 Max (36 GB), test-real passes on the
Qwen3.8-27B UD-Q4_K_M GGUF, and the runtime oracle on the
Ornith-1.5-122B 4-bit MLX (language only, --expert-memory 16G); on a Mac
mini M4, check-native-metal and test-engine-cpu pass.

Earlier, with the change on 1.4.0 (1482f6b), on the M5 Max (36 GB) with
the 27B UD-Q4_K_M GGUF, test-real's first invocation (--ane-ffn-share 0)
passed in 128 s, where the build with the cold-score probe took 192 and
194 s and 1.4.0's oracle 190 s, with 1.4.0's TOKENS,
preemption_decisions and penalty_probe lines but for the losses, and 12
penalty_same_forward lines of 25 rows; its three invocations passed, and
so did the first under MTL_SHADER_VALIDATION=1, every logit of the
checked rows matching bitwise, or within 2 ulps where counted. Before
that, linson007's runs of #526's final code, with its fp64 reference,
passed on a 48 GB M5 Pro over four representative model and KV/state
combinations (1,137 rows).

Stacked on #561; review the commit after it.

🤖 Generated with Claude Code

---

^{Need help on this PR? Tag @codesmith-bot with what you need.}