Negative result: dual-ANE shared-expert offload for DeepSeek V4 misses correctness and whole-prefill gates

[DiscoStew6082]

I tested a dual-ANE fused shared-expert path for deepseek-ai/DeepSeek-V4-Flash-0731 on a Mac Studio M3 Ultra. I am sharing the negative result because the design looked promising on paper, but a gated one-layer prototype showed that it was not worth compiling and integrating across all 43 layers.

This is not a claim that ANE acceleration cannot help DeepSeek V4. It is a scoped result for this checkpoint, this shared-expert seam, this INT8 ANE conversion path, and a fixed 2,048-token prefill shape.

Short version

• The dual-ANE shared branch was fast enough to fit comfortably inside the routed GPU branch: 6.60 ms versus 39.36 ms mean.
• The real ANE shared output was not accurate enough: cosine similarity 0.999119 and relative L2 error 4.091%.
• Concurrent routed-GPU plus shared-ANE execution reduced one-layer MoE wall time from 38.41 ms to 34.98 ms, an 8.9% MoE improvement, but missed my predeclared 10% gate.
• Projected over the complete 43-layer prefill, the measured improvement was only 4.11%.
• Even a generous zero-overhead upper bound that removes the complete GPU shared-expert time from every layer is only 6.46%, leaving too little margin for a repeatable gain above 5% after integration overhead.
• I therefore stopped before building the 43-layer procedure banks or running service-level benchmarks for this primitive.

Environment

• Host: Mac Studio, Apple M3 Ultra, 512 GB unified memory
• OS: macOS 26.6.1
• oMLX: 0.6.1, source commit 6307e396acfdfff7b75d3ceb9b4756272fb03b6b
• MLX: 0.32.0
• mlx-lm: 0.31.3
• mlx-vlm: 0.6.3
• Model: deepseek-ai/DeepSeek-V4-Flash-0731
• Representative main-model layer: 21
• Input: real normalized MoE activation, shape 1 x 2048 x 4096, BF16
• Repetitions: six warmed, alternating measurements per mode

The installed checkpoint is mixed quantization: the always-on shared experts use MXFP8, while the routed experts use packed MXFP4 and select six of 256 experts per token. This made it possible to test the real shared-expert weights without constructing an all-MXFP8 checkpoint.

Why this seam looked promising

The routed and shared branches are independent after receiving the normalized MoE input:

1. The GPU keeps the existing router, sorted route plan, six MXFP4 routed experts, score weighting, and scatter/unsort path.
2. The 2,048-channel shared expert is split into two 1,024-channel hidden shards.
3. Each ANE execution path computes a complete shard: gate projection, up projection, limited SwiGLU, and a partial down projection to the full 4,096 output width.
4. The two partial ANE outputs are summed and then added, unweighted, to the routed GPU result.

The fused graph used DeepSeek's exact limited SwiGLU:

gate = min(gate, 10)
up = clip(up, -10, 10)
hidden = silu(gate) * up

The Python/reference shard algebra, clamp behavior beyond both limits, output shape/dtype, unsupported-shape fallback, repeated use, procedure selection, and finite-output checks all passed. The native extension rebuilt successfully, and 220 relevant tests passed.

Predeclared one-layer gates

I set the continuation gates before compiling the full bank:

Gate	Requirement
Shared-output cosine similarity	>= 0.9999
Shared-output relative L2 error	<= 1%
ANE shared time / routed GPU time	<= 1.05
Concurrent MoE time / sequential MoE time	<= 0.90
Projected whole-prefill improvement	>= 7%
Projected additional resident memory for 43 layers	< about 2.5 GB

The 7% projection gate was intended to leave enough margin to retain a repeatable improvement above 5% after integration overhead.

Correctness result

Check	Result	Decision
Shared-output cosine similarity	0.99911898	Fail
Shared-output relative L2 error	4.091%	Fail
Full-MoE cosine similarity	0.99975210	Informational
Full-MoE relative L2 error	1.972%	Informational
Shape and BF16 dtype	Exact	Pass
NaN or infinity	None	Pass

The ANE constants were produced by dequantizing the checkpoint's MXFP8 weights and converting them to per-row INT8. The exact source of the remaining error could be in that conversion or the compiled execution path; I did not continue isolating it after the independent projected-value gate also failed. FP16 constants might improve fidelity, but they would increase the 43-layer memory cost and cannot improve the perfect-hide upper bound.

Timing result

Complete wall time over six warmed alternating repetitions:

Mode	Mean	Median	Min-max
GPU shared expert	5.395 ms	5.545 ms	4.948-5.680 ms
Dual-ANE shared expert	6.599 ms	6.588 ms	5.716-7.120 ms
GPU routed branch	39.362 ms	33.737 ms	33.077-53.766 ms
Existing sequential MoE	38.411 ms	38.424 ms	38.301-38.485 ms
Concurrent routed GPU + dual-ANE shared	34.980 ms	34.990 ms	34.479-35.566 ms

The host intervals placed the ANE call inside the routed GPU interval on every repetition. That establishes concurrent host-side work, but it does not prove physical GPU/ANE execution overlap. I did not find a supported device-level timeline source that could provide independent GPU and ANE start/end timestamps.

The full warmed 2,048-token model forward averaged 3,593.38 ms.

Measured projection:

43 * (38.4108 - 34.9801) / 3593.38 = 4.11%

Generous perfect-hide upper bound:

43 * 5.3954 / 3593.38 = 6.46%

Because stream or event tuning cannot exceed the perfect-hide bound, I treated this as a decisive no-go for full integration.

Memory result

I changed the ANE procedure-bank design to use one reusable input/output workspace per ANE bank under a strict one-in-flight invariant, instead of allocating mutable IOSurfaces per procedure.

Procedures per ANE	Compile time	Measured physical-footprint delta
1	0.168 s	49,152 B
4	0.463 s	48,644,168 B
8	0.845 s	80,035,936 B

These measurements deliberately reused the same layer-21 weights. The private compiler may deduplicate identical constants, so this is only a lower bound and not a defensible projection for 43 unique layers. The 43 shared experts contain about 1.08 GB of INT8 weight data before compiler expansion. I did not compile unique 43-layer banks merely to refine a memory estimate after the correctness and value gates had already failed.

Related attention-projection result

This shared-expert experiment followed a reconstructed DeepSeek attention-projection ANE path. That path produced repeatable long-context prompt-throughput gains of 1.57% at 4K, 3.50% at 16K, and 3.06% at 32K, while adding about 8 GB of physical footprint. Its isolated hybrid projection was slower than GPU-only at every tested split. That result is why I moved the seam to the whole MoE shared branch rather than continuing to tune attention fractions.

What I think this establishes

• The mixed MXFP4/MXFP8 checkpoint is sufficient for developing and testing the shared-expert primitive; an all-MXFP8 model is not required.
• The dual-ANE branch has enough isolated throughput to be plausibly hideable behind routed GPU work.
• Exact hidden-shard algebra and exact limited-SwiGLU clamps do not guarantee acceptable end-to-end numerical fidelity after ANE weight conversion and compilation.
• For this captured layer, the shared expert's measured wall-time contribution is too small to provide a robust whole-prefill win, even under an unrealistically favorable overlap bound.
• A separate MLX stream plus asynchronous evaluation is not, by itself, proof of device-level GPU/ANE overlap.

It does not establish that every DeepSeek V4 ANE strategy is unhelpful, nor that a different ANE format or a larger primitive could not work.

Questions for maintainers and other experimenters

1. Is there a supported or known-safe way to obtain device-level GPU and ANE start/end timestamps for this kind of concurrent execution?
2. Is there a higher-fidelity ANE constant/weight path that avoids the observed INT8 error without an unreasonable compiled-memory increase?
3. Does the 6.46% perfect-hide bound match how you would evaluate this seam, or is there an important whole-model scheduling effect that this one-layer projection misses?
4. Would the local experimental patch and raw JSON measurements be useful to others? If so, I can clean them into a separate opt-in branch; I have not opened a PR because the continuation gates failed.

I am mainly posting this so others considering the same seemingly straightforward shared-expert offload can start from measured constraints rather than rebuilding the entire 43-layer bank first.