omlx 0.5.7 as a library: Pro-0813 DSpark under mlx.launch TP4 (32-head ring GEMM)

[guruswami-ai]

Using omlx 0.5.7 as a library for DeepSeek-V4-Pro DSpark under mlx.launch (TP4)

Report from a four-node M3 Ultra 512 GB Thunderbolt 5 mesh. We are not running the omlx menu-bar / single-node HTTP server. We import the 0.5.7 overlay (deepseek_v4 + mlx_lm_mtp) under mlx.launch so Pro-0813's own DSpark drafter exists in distributed TP4.

It works. Interactive decode 15–23 tok/s, DSpark armed, out to 32 k prompt tokens, behind a small OpenAI server that DeepSeek Harness can drive. A failing unit-test task went through bash + edit and passed. This is the first trillion-class open-weight we have run on this mesh that is useful as an agent, not just a load test.

What omlx provided that stock mlx-lm 0.31.3 does not

• deepseek_v4 model class (mlx-lm PRs still unreleased).
• MTP overlay that keeps mtp.* (base sanitize drops them).
• BatchGenerator / GenerationBatch path where DSpark actually arms. stream_generate never does.
• Gather fallback in SparseCompressedAttention once the 64-head ring kernel is hidden.

Model.shard() still does not shard self.mtp. We added a matching shard_mtp() for the three stages (targets 58/59/60, depth 5, 32 heads/rank).

Ring GEMM: 64-head kernel vs 32-head TP4

Stock dspark_ring_gemm is compiled for M=64 (Flash single-node). Pro-0813 has 128 heads; TP4 leaves 32 per rank. Python turns the kernel on once pooled KV exceeds index_topk 1024 (~4 k prompt at compress_ratio 4). Then:

ValueError: Unsupported DSpark physical-ring GEMM layout

Hiding the symbol forces gather. After that:

prompt	decode tok/s	accept	TTFT
4 096	17.4	82%	14 s
16 384	14.7	76%	59 s
32 768	12.8	79%	122 s

We did compile a 32-head kernel (accept M in {32,64}, values BM 64→32). At 4 k: native 11.4 tok/s / 56% accept vs gather 10.9 / 70%. Then 16 k after a native switch died in JACCL Recv -12. Serving stays on gather. A correct official 32-head kernel would be the highest-leverage omlx change for Pro TP4.

DSpark is singleton-only

_omlx_mtp_rowwise_unsupported = True on DSpark (caches are singleton / committed-only). Measured: B=2 and B=4 have mtp_activated=0. B=4 aggregate ~29 tok/s is ~7 tok/s per stream, worse per user than one DSpark stream at 15.

If that flag is staying, it should be documented as a contract. If row-wise DSpark is planned, Pro TP4 is a good test case.

Other notes for this overlay

• set_mtp_active(True) must happen before mlx_lm.load.
• Shared mx.random.seed on every rank or sampled decode desynchronises and JACCL dies.
• We run MLX_METAL_FAST_SYNCH=0.
• JACCL Recv/Send -12 on the second generate in one mlx.launch process is still common. That is likely mlx/JACCL, not the Python overlay, but it is the main reason this is not a 24/7 server.

Ask

1. Ship a 32-head (or runtime-M) ring GEMM that matches gather numerically.
2. Keep the overlay usable as a library under mlx.launch. That is how this cluster run exists.
3. Document DSpark as singleton-only, or lift _omlx_mtp_rowwise_unsupported.

Happy to share the shard_mtp helper, the ring-gate, and the 32-head patch. Not asking you to adopt our HTTP server.

[guruswami-ai]

Recipe and scripts (convert / shard_mtp / BatchGenerator smoke / OpenAI serve). No converted weights — rebuild from deepseek-ai/DeepSeek-V4-Pro-0813.

https://huggingface.co/guruswami-ai/deepseek-v4-pro-0813-mlx-tp4-recipe

[guruswami-ai]

Interim: still gather; serve survived three HTTP turns

Serving stays on the gather fallback. The native 32-head dspark_ring_gemm is still unused: at 4 k it was faster but accept fell 70% → 56%, then a 16 k switch tore JACCL.

The singleton OpenAI path (omlx 0.5.7 as a library under mlx.launch) just took three sequential HTTP turns after a ~2.5 h idle. DSpark stayed on. Prompt-cache trim fired on turns 2–3.

The second-generate JACCL -12 was an undrained close(); drain is mx.synchronize() first. See mlx #4247. A correct 32-head (or M-from-tensor) kernel is still the decode leftover that is actually yours.

Recipe: https://huggingface.co/guruswami-ai/deepseek-v4-pro-0813-mlx-tp4-recipe

[guruswami-ai]

Follow-up, 15 August 2026

Same setup: omlx 0.5.7 as a library under mlx.launch, not the menu-bar server.

Pad-to-64 is not a substitute for a 32-head kernel

The original post asked for a 32-head dspark_ring_gemm (stock is M=64; Pro TP4 is 32 heads/rank). Before compiling another native, we wrapped the stock 64-head symbol: pad the head axis 32→64 with zeros, call the kernel, slice back.

Two fresh processes, same 4 096-token prompt, shared seed 20260814, 64 new greedy tokens, DSpark depth 5:

arm	accept	decode tok/s	greedy ids
gather (off)	70.8%	14.1	digest 1969909217
pad-to-64	71.2%	15.0	digest 902458483

Acceptance is the same. IDs diverge at generated token 10 (values vs symbols in the decoded prefix). Pad is slightly faster and wrong. Default stays gather. Native M=32 stays off until greedy ids match gather on scores, values, and accept.

A wrapper that only pads axis 1 still changes the numerics. Please do not treat pad-to-64 as proven lossless.

Context on gather

Hiding the 64-head symbol still unlocks long context. One-process climb, 32 new tokens:

prompt	decode tok/s	accept	TTFT
32 360	12.3	70%	112 s
64 711	14.8	83%	299 s
97 062	12.0	75%	507 s
129 339	12.2	76%	736 s

Model.shard() still skips self.mtp. Row-wise DSpark for B>1 is still unsupported. Those asks stand.

[guruswami-ai]

Parking Pro-0813 TP4 as a daily driver (17 August 2026). omlx 0.5.7-as-library still loads and DSpark still decodes on a fresh process. Long-lived JACCL is the blocker: ml-explore/mlx#4319 / FB24371487. Full note on mlx discussion 4247.