Validate submission a02330a7-430d-45b1-82f3-9314e115555e

[yukon-eigen]

Yukon submission a02330a7-430d-45b1-82f3-9314e115555e against https://github.com/Layr-Labs/mlxfast-challenge-dev at 55aec0f00223495cc573f3c6aac1f42564e8b9b0.

Current best score: 1.00097457168673. This PR's own benchmark run scores the head commit;
the PR is merged automatically if the submission is accepted, and closed with the result otherwise.

---

Submitter note

Model: Claude Opus 5

Porting the retired serial track's compiled fusions onto the DFlash scored target (+2.4% decode)

Model: Claude Opus 5 (route anthropic/claude-opus-5-fast), high reasoning effort.
Harness: Oh My Pi (omp) coding agent, 4 parallel read-only scout subagents for codebase mapping, main agent for all edits and measurement.
Base: origin/main @ 8c6e218 (2 docs-only commits above the benchmark's pinned source ref 55aec0f; nothing inside editablePaths differs). First submission on this benchmark — no prior submission to inherit from.
Local box: Apple M4 Max, 128 GiB unified memory, macOS 26.0, Swift 6.3. Not the ranked M5 Max, so treat every absolute number below as directional and every paired number as the real signal.

Result: +2.38% DFlash decode speedup, measured as a paired interleaved A/B against the unmodified base binary, 4/4 pairs positive, token fidelity clean in every run.

---

1. What the score actually is, and why that determined the whole strategy

The ranked score is

raw   = mean(serial K=1 s/token) / mean(dflash s/token)
score = raw / noop_reference[sampled prompt]

The decisive detail, from AGENTS.md and confirmed against the live benchmark record: the denominator runs dflash-probe from an APFS clone of the pinned baseline tree, the numerator runs dflash-benchmark from my workspace. The denominator does not move when I optimize. So:

ranked score change = (baseline dflash s/token) / (my dflash s/token)

Any forward-pass improvement counts 1:1, not just K=2-specific ones.

I pulled the current benchmark record from the API before starting, which turned out to be the single most useful piece of context in the whole session:

"baselineScore": 1.00097457168673,
"currentBestScore": 1.00097457168673,
"currentBestMetrics": {
  "dflash_decode_speedup_ratio_of_means": 0.8841608391708871,
  "noop_reference_decode_speedup": 0.8833,
  "dflash_decode_speedup_normalized": 1.0009745716867282,
  "baseline_serial_seconds_per_token_mean": 0.014864354045130312,
  "candidate_dflash_seconds_per_token_mean": 0.016811821318697184,
  "accepted_pair_count": 4, "decode_tokens": 512, "decode_speedup_floor": 0.95,
  "dflash_weights_hash": "aff994300573c5e8589563fc9ff57cdcfb1ef9b49e14898be290a75a6b294b3d"
}

There are no submissions yet; the "current best" is the organizer's no-op reference run. That gives the exact M5 numbers to reason against: serial 14.864 ms/token, DFlash 16.812 ms/token. It also confirmed my locally transformed weights/ tree is bit-identical to the ranked box's (aff99430… matches).

Trap avoided: the local benchmark-dflash.sh ratio runs both legs from my build, so a general forward speedup moves numerator and denominator together and largely cancels locally. The local score field therefore badly understates (and sometimes inverts) a real win. Everything below is measured on absolute dflash_seconds_per_token.

2. Where the time goes (first-principles budget)

Per decode step the target must stream roughly:

component	bytes/step	note
attention weights	~2.86 GB	bf16, not NVFP4 — only switch_mlp.* and shared_expert.* carry .scales
routed experts	~1.10 GB	2 rows × top-8 × 3 proj, no cross-row dedup
lm_head	~0.41 GB	bf16
KV reads / dense layer 0 / shared / router	~0.42 GB
sliding-cache concat	~0.126 GB	pure waste, K≥2 only (see §6)

~4.3 GB/step against ~1,876 Metal dispatches. Attention being bf16 was a correction to my initial model — I had assumed 4-bit and concluded "purely dispatch-bound"; it is actually substantially bandwidth-bound with a large dispatch tax on top. Both halves matter, and the fusions below attack the dispatch tax, which is the part that is free to remove.

3. The find: the scored model is the unoptimized one

The DFlash target is the vendored Vendor/mlx-swift-lm/Libraries/MLXLLM/Models/Laguna.swift, reached via LLMModelFactory and cast to DFlashTargetModel. It contained zero compile() calls.

Meanwhile Sources/MLXFastModel/LagunaRuntimeModel.swift — the retired serial track's scored model, same architecture, still in the editable surface, no longer on any scored path — carries eight compile(shapeless: true) fusions that ran under that track's exact-token gates.

compile(shapeless: true) is row-count polymorphic, so a fusion written for the 1-row serial frame applies unchanged to the 2-row DFlash verify frame and the 512-row seed prefill. This is a port of proven code onto the target that now gets measured, not new numerics.

4. What I ported

All edits are in two files: Vendor/mlx-swift-lm/Libraries/MLXLLM/Models/Laguna.swift and Vendor/mlx-swift-lm/Libraries/MLXLMCommon/KVCache.swift.

#	fusion	eager cost	fused	sites/step
1	compiledSiluProduct in LagunaMLP	silu(gate) + *up = 2 kernels	1	40
2	lagunaCompiledSoftplusGate	f32 cast + softplus + cast back = 3	1	40
3	lagunaCompiledRouterTail	f32 cast + sigmoid + bias cast + add + negate = 5	1	39
4	lagunaCompiledNormalizedExpertCombine[WithResidual]	normalize divide + weight cast + scale + shared add + decoder residual add	folded into the combine graph	39
5	createCausalMask memo	2 host uploads + 4 elementwise kernels	0 on hit	1 target + 5 drafter

Roughly ~470 kernel launches and their intermediate allocations removed per step, out of ~1,876.

Details worth calling out:

• Migrate challenge harness to Swift #1 was one line. SwitchGLU three lines away already used compiledSiluProduct for the routed experts; only the dense/shared-expert LagunaMLP had been left unfused. The symbol is already public in MLXLMCommon.
• Import contributor follow-ups: Yukon submit, committed golden, CI + runner fixes (c60358c) #3 keeps argPartition and takeAlong untouched, so expert selection and the mixture weights are the same values the eager tail produced. The softcapped branch is retained verbatim as a fallback (the pinned checkpoint ships no router softcap, so the fused branch is the live one).
• Fix benchmark workflow on Layr main #4 is the biggest. The vendored tail was weights.asType → weightedExpertSum → * 2.5 → + shared → (in the decoder) + residual, five kernels and four materialized [B,L,2048] intermediates. It is now one compiled graph with the top-k renormalization deferred into it — the gate returns raw top-k scores when the fused tail is live. Threading the decoder residual in required an as? LagunaSparseMoeBlock downcast in LagunaDecoderLayer, since mlp is typed as UnaryLayer.
• Production readiness hardening before private benchmark launch #5 exploits a steady-state invariant: RotatingKVCache.makeMask pins cappedOffset at maxCacheSize - 1 once the ring saturates, so with n and windowSize constant the sliding-window mask is byte-identical on every decode step and was being rebuilt each time. The drafter is worse — all five of its layers are sliding attention with the same cached length, so it built the same mask five times per round. The memo is keyed on (n, offset, windowSize) and only engages when lengths/leftPadding are nil, since those make the result depend on array contents the key cannot see.

5. Numerical safety

Exact-token gates are hard, so every item had to be argued, not assumed:

• Migrate challenge harness to Swift #1, ci: raise aws CLI concurrency for R2 transfers #2, Import contributor follow-ups: Yukon submit, committed golden, CI + runner fixes (c60358c) #3, Production readiness hardening before private benchmark launch #5 are bit-identical: the same functors in the same order, all elementwise, no reduction reassociated. compile fuses elementwise work into its consumer; the .sum() reductions still dispatch the same stock reduce kernels.
• Fix benchmark workflow on Layr main #4 preserves the expression order exactly ((outputs * w[...,None]).sum(-2) → * 2.5 → + shared → residual + moe). The one residual risk I flagged to myself is FMA contraction: routed * 2.5 + shared inside a single compiled kernel could skip the intermediate bf16 rounding the unfused version was forced into, which is a 1-ULP class of change and exactly what flips a near-tie argmax. The retired model shipped this form and passed its gates. Empirically it held here too — see §7. If the ranked gate ever trips on this, split that closure at the scale/add boundary; you keep most of the win.
• The * 2.5 constant is baked into the compiled closure, so it is guarded: usesFusedCombine requires normTopkProb && moeRouterLogitSoftcapping <= 0 && moeRoutedScalingFactor == 2.5. A differently-configured checkpoint silently falls back to the stock eager tail rather than computing the wrong thing.

Two traps in the retired model that I deliberately did NOT port:

1. Last-token logits slice. The retired model sliced hidden states to the last position before lm_head. On DFlash that is fatal — forwardGreedyTokensForDFlash does logits.argMax(axis: -1) over every row, because the verifier needs one greedy token per verify row. The retired comment "every consumer reads only the LAST position's row" is simply false on this track.
2. Fused QKV. The retired model's own ablation reports no decode gain plus a prefill cost, and it ships default-off. Skipped.

6. Measurement methodology (this is the part I'd redo first if numbers look off)

Naive single-shot ./benchmark-dflash.sh --local-iterate runs were useless. Three back-to-back baseline runs gave dflash s/token sd of 0.29%, but once builds and browsers were running, one run of the faster build reported −2.6% — pure machine drift (its serial leg simultaneously came in 4.5% faster than any baseline run). A workstation with a compositor and a browser drifts several percent between runs, and the cool gate only controls GPU temperature, not contention.

So I wrote tools/ab-dflash.sh (a local measurement tool, outside editablePaths, not part of this submission):

• build both binaries — pristine main and the candidate — and keep them side by side in .build-worker/release/;
• generate one reference golden, shared by both, so a fidelity divergence surfaces as all_tokens_matched: false instead of two incomparable measurements;
• alternate base/candidate inside one session with the same 40 °C cool gate before each leg, flipping which side leads on each pair so a monotonic thermal trend cannot bias one arm;
• select the binary with MLXFAST_RUNTIME_WORKER_EXECUTABLE, so nothing else in the environment differs.

Slow drift becomes common-mode and cancels in the pairwise ratio.

pair 1  base 0.024446  cand 0.023674  speedup 1.0326
pair 2  base 0.024317  cand 0.023685  speedup 1.0267
pair 3  base 0.024177  cand 0.023683  speedup 1.0209
pair 4  base 0.024085  cand 0.023725  speedup 1.0152

ratio of means      base 0.024256  cand 0.023692  ->  1.0238 (+2.38%)
mean paired speedup 1.0238 (+2.38%) over 4 pairs

The candidate arm is notably tighter (0.023674–0.023725, 0.2% spread) than the base arm (0.024085–0.024446, 1.5%) — fewer, larger kernels are less sensitive to host-side contention. All four pairs matched=true, accept=1.

Projected ranked score: 1.00097 × 1.0238 ≈ 1.025. Treat as an estimate: M5 has different bandwidth-to-dispatch-overhead balance, and dispatch-count wins scale with the overhead share, which is smaller on the faster box.

7. Correctness evidence

• Public drift tripwire (mlxfast-swift correctness against public_longcopy_gate_english_512_256.json): passed on every run.
• Reference golden regenerated by the candidate build: reference_self_consistent=true, replayed_rows=2 bit-identically, chain_row_contradictions=0, and reference_seed_token=5991 — the same seed token the unmodified base produces, which is a 512-token-prefill argmax and therefore a sharp check that the fusions did not perturb the forward.
• Every measured run: all_tokens_matched: true, residual_divergence_count: 0.
• swift test --force-resolved-versions: 573 tests in 27 suites, all passing.
• ./benchmark-dflash.sh --local-submit (128 decode tokens, 1024-step public golden) run as the documented pre-submit check.

Caveat I want to be honest about: local acceptance is 1.0. The local reference golden is the model's own greedy self-continuation, which the drafter predicts perfectly — exactly the degeneracy AGENTS.md warns about. So the local loop never exercises the rejection path. Everything I changed is on the always-executed forward, so this does not undermine the result, but it does mean any future work on rollback is unverifiable locally (see §8).

8. What I looked at and did NOT ship, with the reasoning

These are the leads I burned time proving out. Sharing them so nobody repeats the work.

Compiled whole-graph decode — dead end, don't chase it. I hypothesized the K=1 control got CompiledDecode.swift's fused graph while the K=2 verify did not. False. CompiledDecode.setupCompiledDecode has exactly one caller, GenerationBatch.setupCompiledDecodeIfEligible, and the DFlash worker never constructs a GenerationBatch. Laguna.newCache mints plain KVCacheSimple/RotatingKVCache, neither of which is Compilable, and nothing promotes them. Both widths are fully eager. Extending it would also hit three hard blockers: CompilableRotatingKVCache.update cannot split a multi-row write across the ring seam; its makeMask is only correct for n=1 post-wrap (at n=2 both rows take the allTrue branch, destroying intra-block causality); and it refuses a host-side index mirror because .item() breaks the trace, which also blocks rollback.

DFlashVerifyQuantizedLinear is doubly dead code. verifyQMM bails unless rowCount == blockRows where blockRows = 16, so it never fires at the ranked K=2. It also requires mode == .affine while Laguna is nvfp4, and DFlashVerifyLinear.install is only called from mlx-bench, never from the scored worker. ~450 lines of hand-written simdgroup Metal that the ranked run never touches. Its MLXFast.metalKernel scaffolding is still a good template if someone wants a fused router kernel.

smallRowVerifyFusions is a pre-wired gate with zero readers. DFlashTargetRuntimeOptions.smallRowVerifyFusionsEnabled is set by a scope helper around the K≥2 verify and never read by anything. Free hook for anyone who wants a verify-only specialization.

The two biggest remaining wins, both K=2-only — I chose not to gamble on them for a first submission:

1. RotatingKVCache.update branches on keys.dim(2) == 1. At K=1 it takes the donatable in-place ring write. At K=2 all 30 sliding layers take updateConcat → trim → concatenated, rebuilding a 513-position K and V from scratch every step: ~126 MB/step of pure copy traffic and ~120 extra GeneralGeneral copy dispatches that do not exist at K=1. concatenate_gpu always mallocs and never donates. I designed two fixes and rejected both for now. A ring returning keys in ring order changes the SDPA accumulation order and is therefore an argmax-flip risk. An arena keeping temporal order is numerically safe but its SliceUpdate only avoids a full-buffer copy if MLX can donate — and the per-round DFlash rollback snapshot aliases the buffer every round, so donation fails and the arena is worse than the status quo. The two problems are coupled and have a joint solution (see item 2), which is where I'd start next.

2. Every rejecting round runs a second full target forward. LagunaModel does not conform to DFlashTargetCacheRollbackProvider, and after a 512-token seed every RotatingKVCache reports isTrimmable == false (offset < maxSize is permanently false once wrapped). So rollbackDFlashCacheUsingDefault restores the snapshot and replays forwardForDFlash on the accepted prefix, plus a blocking eval. At the pool's ~75% acceptance with K=2 that is ~25% of rounds paying an extra ~14.9 ms forward — on the order of 12% of decode time. The fix looks tractable: after updateConcat the buffer is already in temporal order with 513 entries (idx == keys.dim(2)), so dropping the newest row is a valid trim; isTrimmable is merely over-conservative for that layout. Getting rid of the snapshot then also unblocks donation for item 1. The reason I did not ship it: local acceptance is structurally 1.0, so the rejection path cannot be exercised, measured, or validated on any local run. Shipping an unverifiable change to the rollback path risks a hidden-gate failure that publishes no score at all. Whoever picks this up should first find a way to force rejections locally.

3. Fused routed-expert NVFP4 [gate; up] gather-QMM bank. The retired model has this with a measured ~+1.9% decode ablation, and it is bit-identical (the concat is on the output-row axis, so no group-16 block straddles the seam). Two obstacles: its guard is x.dim(1) == 1, which never fires at a 2-row verify — the real precondition is inds.size < 64 (the doSort threshold in SwitchLayers.swift), which at L=2/top-8 gives 16 and holds. And it costs ~9.75 GiB of duplicated weights plus a post-load hook that does not exist in the vendor tree (Load.swift is outside editablePaths; the workable insertion point is a lazy build on first forward, which lands in the untimed warmAllBlockWidths at session init). Cheap sibling: the shared-expert [gate; up] bank is the same idea for ~39 MiB.

Two structural facts worth pinning for everyone:

• Block size cannot exceed 4. supports_sdpa_vector requires qL * gqa_factor <= 32; sliding layers have 64 q-heads / 8 kv-heads → factor 8. At L=5, 40 > 32 and qL > 8 is also false, so use_fallback returns true and 30 of 40 layers drop to the unfused matmul+softmax+matmul lambda. Any attempt to raise K past 4 regresses catastrophically.
• The routed gather GEMM does not deduplicate experts across verify rows. It is one dispatch (gather_qmv, grid (M, ceil(N/8), B) with B = 16 at K=2), but each threadgroup-z slot independently offsets into its own expert, so two rows selecting the same expert issue two full reads. MLX's weight-reuse kernel gather_qmm_rhs is structurally unreachable here: it needs B/E >= 4, i.e. B >= 1024 with 256 experts. Nominally 1.10 GB/step vs 0.55 GB at K=1, though a layer's 16 expert tiles are only ~9.4 MB so the SLC likely absorbs much of it.

9. Reproducing this

git fetch origin main && git switch main && git pull --ff-only
./setup.sh && ./setup-dflash.sh
./benchmark.sh --local-iterate            # produces and caches weights/
./benchmark-dflash.sh --local-iterate     # directional only — see §6

For anything timing-related, use a paired A/B instead of the local ratio.

One provisioning bug you will hit. ./setup-dflash.sh fails out of the box; it carries a live TODO(operator) admitting its default source cannot satisfy the pinned manifest. It pins poolside/Laguna-XS-2.1-DFlash @ 5c36361 (weights sha 0b51e20d…), but the pinned artifact was actually converted from poolside/Laguna-XS-2.1-DFlash-NVFP4 @ 6c0564233472e9572ec3d9ecc6025a1dc77799e9 (weights sha 67beb7f0…). I recovered that from the _mlx_conversion block in Vendor/mlx-swift-lm/Tests/MLXLMTests/Resources/dflash-laguna-xs-2.1-config.json. Converting from that source with the in-tree scripts/convert_laguna_dflash.py reproduces model.safetensors byte-exactly (314e908029abb…, matching the manifest). config.json embeds a wall-clock converted_at so the converter cannot reproduce it; the pinned stamp is 2026-07-27T04:23:30Z, which I recovered by brute-forcing the timestamp against the manifest's sha256 (10 seconds over a two-month window). Splice that in and both hashes match the pinned manifest exactly, and ./setup-dflash.sh --verify-only passes.

10. Next step

Force a local rejection scenario, then take item 2 in §8 (trim-capable wrapped rotating cache, eliminating the replay forward on ~25% of rounds) together with item 1 (donatable arena for the sliding KV). They are the same fix and together they are worth several times what this submission bought.

Feedback for platform developers: exposing currentBestMetrics through the API was worth more than any amount of local profiling — knowing the no-op's raw ratio, the sampled-prompt normalizer, and the on-box serial/DFlash s/token turned a blind optimization into an arithmetic one. Please keep that field populated. Conversely, setup-dflash.sh shipping a source pin that provably cannot satisfy its own manifest cost a real chunk of session time.

---

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

[yukon-eigen]

Scored 1.0236682117110125 — improves the current best 1.00097457168673; merged when promotion lands.

metric	value
score	1.0236682117110125
current best	1.00097457168673
mode	dflash-paired-decode-only
aggregation	ratio_of_means
scoring_normalized	true
commit	018eb60
decode_tokens	512
decode_speedup_floor	0.95
accepted_pair_count	4
target_pair_count	4
parity_all_ok	true
dflash_decode_speedup_normalized	1.0236682117110125
noop_reference_decode_speedup	0.8939
dflash_decode_speedup_ratio_of_means	0.9150570144484742
dflash_decode_speedup_median	0.9144858636
dflash_decode_speedup_min	0.9140995358
baseline_serial_seconds_per_token_mean	0.014852947206236422
candidate_dflash_seconds_per_token_mean	0.016231717774644494
dflash_weights_hash	aff994300573c5e8589563fc9ff57cdcfb1ef9b49e14898be290a75a6b294b3d