How I turned one oversized DeepSeek optimization into three merged oMLX PRs

Ablation means removing one optimization at a time, measuring what changes, and keeping only the code whose benefit justifies its complexity.

TL;DR

I started with a combined DeepSeek V4 prefill optimization that measured 16.53% higher prompt throughput on my M3 Ultra. It also mixed six changes into a 3,111-addition pull request.

Instead of arguing that the benchmark justified the diff, I rebuilt the work against current oMLX, measured each optimization independently, and gave the maintainer enough evidence to decide which changes were worth carrying.

Three small optimizations survived that process and were merged as PR #2559, PR #2562, and PR #2563. The maintainer independently tested the exact three-PR set and measured 10.7% to 17.7% higher prompt throughput from 4K through 131K context, with decode speed unchanged.

The important result was not only faster prefill. It was turning a promising prototype into three changes that an upstream maintainer could review, benchmark, merge, or revert independently.

The fast patch that was not ready to merge

DeepSeek V4 Flash long-prompt prefill was spending most of its time in routed mixture-of-experts work and compressed attention. I profiled those paths and assembled six optimizations into draft PR #2553.

On a Mac Studio with an M3 Ultra and 512 GB of unified memory, the combined candidate improved a controlled 17,219-token cold-prefill workload:

Metric	Upstream	Combined candidate	Change
Median cold prefill, five runs	35.97 s	30.87 s	-14.18%
Median prompt throughput	478.67 tok/s	557.81 tok/s	+16.53%

The broader campaign also included an 864-item paired quality comparison, cache checks, real-model tests, and a complete repository test run. But a strong benchmark and a large evidence package did not make the implementation easy to review.

The PR had 3,111 additions across 21 files. Only 347 additions were production code. The rest included 621 lines of tests, a reproduction runner, a benchmark report, and 13 raw evidence files. More importantly, six interacting optimizations made it difficult to identify each change’s value, risk, and fallback boundary.

I posted my own scope correction, identified problems in the cross-chip assumptions, low-bit guard, mask handling, and evidence consistency, and stopped treating the combined patch as merge-ready. The maintainer agreed that the work was worth pursuing but asked for a smaller review surface.

Turn review feedback into an experiment

Splitting the files was not enough. The optimizations interacted, so each one needed to be measured rather than assigned a share of the combined result.

That is the purpose of a leave-one-out ablation: remove one change from the complete candidate, hold everything else fixed, and measure how much of the improvement disappears. It turns “the bundle is faster” into evidence about what each part actually contributes.

I rebuilt the campaign against current upstream and published the methodology in issue #2558. The controlled ablation held the model, prompt, cache policy, prefill step, MTP configuration, expected output, and request concurrency fixed. It recorded 95 full-model requests and eight warm-ups.

The longer experiment also exposed machine-speed drift. Instead of comparing measurements from different phases, I bracketed candidate runs with upstream controls and used same-phase or bracketing comparisons. Every timed request had to report zero cached tokens and return exactly READY.

The leave-one-out results changed which code was worth proposing:

Optimization	Marginal prompt-throughput contribution	Decision
Native ratio-128 compressed attention	10.18%	Merge
Fused hyperconnection expansion	About 2.45%	Hold: new kernel and low-bit evidence cost
MXFP4 large-block threshold	About 1.9%	Merge
Native ratio-4 all-pooled attention	0.5%–1.2%	Hold: mask and dispatch complexity
Skip redundant indexer scoring	About 0.4%	Merge
Fused route counting and block planning	0.22% in the current screen	Hold: custom-kernel maintenance cost

These percentages are not additive. They are leave-one-out measurements with feature interactions. Their purpose was to rank engineering value, not manufacture a combined headline number.

The maintainer explicitly selected three changes and asked that the other three stay out for now. That was a useful outcome: the investigation did not merely find fast code; it produced enough evidence to reject complexity that was not earning its maintenance cost.

The three changes that merged

1. Route ratio-128 prefill through native attention

PR #2559 addressed the largest measured bottleneck. Ratio-128 compressed attention was concatenating local and pooled key/value state before using generic dense attention. The change reused oMLX’s existing native local-plus-pooled kernel with sequential pooled indices.

The fast path was intentionally narrow. Custom masks, sub-4-bit checkpoints, unsupported shapes or dtypes, missing native symbols, and immediate kernel failures retain the reference path. DSpark single-token behavior and empty-pool attention remain unchanged.

Validation included FP16 and BF16 comparisons against the causal reference, real high-bit and sub-4-bit checkpoint paths, cache-backed continuation, focused regression tests, and custom-kernel ABI probes. The controlled ablation attributed a 10.18% marginal prompt-throughput contribution to this slice.

2. Use the right MXFP4 block size for the production shape

PR #2562 was much smaller. A 2,048-token DeepSeek prefill chunk produces 12,288 expert routes, but the existing policy switched from the faster MXFP4 BM16 path to BM32 at 8,192 routes.

The change moved only the MXFP4 crossover to 16,384 routes. It did not change kernel math, and it left affine 2/3-bit and mixed block plans at their existing threshold. M5/NAX retained its stock path, while an environment override let other pre-NAX systems restore the previous crossover.

In the PR’s five-run M3 Ultra comparison, median prompt throughput increased from 479.08 to 487.91 tokens per second, a 1.84% improvement. Median prefill time decreased from 35.94 to 35.29 seconds.

3. Do not rank candidates when every candidate wins

PR #2563 removed redundant work. When every pooled row already fits inside index_topk, query projection, scoring, and top-k selection cannot eliminate anything.

The production change skips that scoring path, advances the indexer compressor and cache exactly as before, and constructs the same sequential indices directly. It also checks that the attention and indexer pooled-row counts have not diverged. Empty and larger pools keep the full indexer path.

The optimization contributed only about 0.4% in the ablation, but it removed demonstrably unnecessary work with a small production diff and an exact cache-state equivalence test. That made it a better maintenance trade than several faster but broader alternatives.

Independent upstream reproduction

Before merging, oMLX maintainer jundot benchmarked the three PRs together against oMLX main on another DeepSeek V4 Flash oQ4e configuration. The test used one cold request at a time, a 2,048-token prefill step, no MTP, and zero cached tokens.

Context	oMLX main	Main plus three PRs	Change
4,096	464.1 tok/s	546.3 tok/s	+17.7%
16,384	487.8 tok/s	559.6 tok/s	+14.7%
65,536	463.4 tok/s	525.1 tok/s	+13.3%
131,072	416.1 tok/s	460.8 tok/s	+10.7%

Decode throughput was unchanged, as expected for prefill-only changes. These were one-run-per-context measurements on an M3 Ultra, not a universal DeepSeek or Apple Silicon claim.

The maintainer also found a production-shaped constraint: the default paged SSD cache used 512-token chunks, which largely hid the improvement. Aligning it to 2,048 tokens recovered the gains with caching enabled. Minutes after merging the three PRs, he added a scheduler adjustment so PoolingCache models could reach the optimized shape in normal operation.

That follow-up mattered. A faster kernel path is not a product improvement if the scheduler never feeds it the shape it was designed to accelerate.

Evidence without polluting the review

The raw campaign was valuable, but it did not belong inside each production PR. I moved the normalized results, raw investigation records, correctness summaries, provenance, and checksums into a separate public evidence package.

That kept each PR focused on one implementation and its direct regression coverage while preserving an auditable trail behind the performance claims. It also let the broader campaign-level accuracy evidence remain available without incorrectly attributing it to only the final three-PR set.

AI-assisted, evidence-directed

I disclosed publicly that Sol xhigh generated the implementation code. That is important context for describing the work accurately.

My contribution was directing the engineering loop: choosing the workload, profiling the system, designing controlled comparisons, detecting bad experimental comparisons, defining fallback and correctness boundaries, curating reproducible evidence, deciding which changes not to pursue, and collaborating with the maintainer until the useful pieces could merge.

That is the AI-engineering skill I want this work to demonstrate. Generating code is only one part of changing a performance-critical system. The harder responsibility is producing trustworthy evidence, constraining the fast path, and knowing when an optimization does not justify its complexity.

What I learned

The original 16.53% result was real, but it was not yet a mergeable contribution. The better result came from making the work smaller:

• measure interacting optimizations independently;
• use bracketing controls when the machine changes speed during a long experiment;
• preserve reference paths for unsupported masks, quantizations, and hardware;
• keep raw evidence available without forcing reviewers to carry it in the production diff;
• let maintenance cost influence which benchmark wins become code; and
• treat upstream reproduction as stronger evidence than your own best run.

All three selected changes were merged into oMLX on August 10, 2026: b6811ed6, c48e1a82, and 4c6b5931.