jundot/omlx v0.5.4rc2

0.5.4rc2

on GitHub

This second release candidate adds support for two newly released models, DeepSeek V4 Flash 0731 with embedded DSpark Lightning MTP and Thinking Machines Inkling Small, plus Lightning MTP for Step-3.7-Flash, kernel-level Inkling decode work, and benchmark uploads for accelerated runs. Everything from 0.5.4rc1 is included, and the notes below only cover what changed since then.

Highlights since rc1

• DeepSeek V4 Flash 0731 with DSpark Lightning MTP. The 0731 checkpoints load natively with their embedded DSpark weights, and the existing Lightning MTP toggle picks the DSpark backend automatically when the checkpoint carries dspark_* configuration. oQ/oQe quantize the embedded mtp.0 through mtp.2 weights with expert importance collection, and Metal kernels cover the DSpark projections, routed experts, DSA indexer scoring, and block verification. (#2460)

  M3 Ultra 512 GB, DeepSeek-V4-Flash-0731-oQ4e-mtp at 4.36 bpw, single stream, TG128, MTP off, prefix and SSD cache disabled.

  Context	1K	4K	8K	16K	32K	64K	128K
  Prefill tok/s	531.0	486.5	493.2	492.8	485.5	468.5	438.6
  Decode tok/s	27.6	25.8	25.1	24.3	23.5	22.8	21.3

  Peak memory over the same sweep goes from 145.7 GiB at 1K to 164.0 GiB at 128K.

  Decode with Lightning MTP off and on, greedy:

  Workload	Off	On	Change	Acceptance
  Code 4K	25.5	47.4	+85.6%	92.7%
  Code 16K	24.3	43.7	+79.5%	88.9%
  Story 4K	25.8	32.9	+27.8%	63.7%
  Story 16K	24.3	24.6	+1.3%	44.0%

  Greedy token hashes were identical off and on for every code, English, Korean, and Japanese workload. Prose lands at 32 to 57% acceptance and does not consistently gain, so the adaptive controller parks speculation there and returns the request to standard decoding. The DSpark model and its runtime state add about 10.2 GiB of peak memory.

• Inkling Small. Thinking Machines Inkling Small (266B fine-grained MoE, text and vision) serves with the full cache subsystem, oQ/oQe, a parser for the model's channel format, and Lightning MTP on the checkpoint's built-in 8-depth head. Streaming oQ/oQe turns the 495 GB bf16 original into 141 GB at 4.59 bpw effective. Audio is not wired yet. Decode was then accelerated with a fused short convolution, a single simdgroup MoE routing kernel, shared experts folded into dense projections at load time, an affine Q4/G64 routed-down kernel, and stacked QKVR projections. (#2438, #2463)

  M3 Ultra 512 GB, Inkling-Small-oQ4e-mtp, exact code prompts, greedy, 512 generated tokens.

  Context	Prefill tok/s	Decode, MTP off	Decode, MTP on	Speedup
  4K	671	44.5	52.5	1.18x
  16K	641	41.7	51.3	1.23x

  Prefill was measured at initial support and decode after the fast paths, which raised standard decode by 16 to 20% on the same prompts. Story prompts ran 1.26x at 4K and 1.17x at 16K on that initial runtime. mtp_enabled stays opt-in per model. Inkling is also the first model mixing a sliceable KVCache with a non-sliceable ArraysCache in one layer, which exposed a store and restore mismatch in the prefix cache plus smaller gaps in the SSD cache and boundary snapshots, now fixed for every model with composite cache layers.

• Lightning MTP for Step-3.7-Flash. Step-3.7 uses the same DeepSeek-V3-style nextn design as the existing DeepSeek V4 and GLM-5.2 patches, but neither StepFun's own HF reference nor upstream mlx-lm implements MTP inference for it, so the forward pass is a direct port of StepFun's reference C++ implementation. This also fixed two silent bugs affecting any VLM-shaped nextn checkpoint: weight detection only emitted bare model.layers.{n}. prefixes, so the settings UI refused to let mtp_enabled be turned on at all, and step3p7 was missing from the MTP architecture allowlist. By @True2456 in #2450

  REAP-pruned, mixed-precision-quantized Step-3.7-Flash, same checkpoint and prompts, TG128.

  Context	Off	On	Speedup
  4K	52.7	64.4	1.22x
  8K	51.2	63.3	1.24x
  16K	48.8	61.1	1.25x
  32K	42.6	55.2	1.30x

• Accelerated runs now reach the benchmark leaderboard. Runs with MTP, TurboQuant, DFlash, SpecPrefill, or VLM MTP enabled were skipped on upload, so the leaderboard had no data for the features that matter most. They now upload with a normalized feature-flag list and a performance-only settings snapshot, model names go up exactly as oMLX shows them instead of having their quantization and MLX suffixes stripped, and each row carries CPU utilization split by performance and efficiency cluster, GPU utilization, thermal pressure, and memory sampled at 4 Hz over its own window, so a throttled run is identifiable from the result itself.

Other changes since rc1

• Homebrew installs with the grammar extra no longer segfault, the extra allowed an open-ended xgrammar that could resolve to a native pair crashing on import on macOS arm64. Reported by @zviratko in #2428
• Speed-priority prefill no longer rejects requests well below capacity, a short tail chunk was scaled up to a full step and charged as a multi-gigabyte admission cost. Reported by @funnygeeker in #2434
• One noisy prefill sample can no longer poison the transient memory estimate, samples more than 8x above the running estimate are now excluded from the blend. By @fparrav in #2414
• partial: true works on VLM-served models, the flag was validated and then discarded at render time, which also hit text-only conversations on any checkpoint carrying a vision_config. By @JimStenstrom in #2435, reported by @dougy86 in #2357
• Embedding compile state is decided honestly at load, the smoke test probed without an attention_mask, so mask-branching models logged compiled=True and then served eager. By @JimStenstrom in #2461, reported by @alantmiller in #2447
• Streamed /v1/completions responses keep one completion ID, per-frame IDs and the cmpl-keepalive sentinel made strict OpenAI accumulators discard later chunks. By @ShiroKSH in #2462
• Community Inkling MLX layouts load, pipenetwork and mlx-community checkpoints previously failed with unmapped parameters or an unsupported model type. Reported by @studioburnside in #2451
• DFlash warns on precision-mismatched draft pairs and reports speculation metrics, the Active Models card shows draft share, accepted tokens per cycle, and fallback reasons in all nine locales. By @FabioMalpezzi in #2445
• Deep Inkling MTP blocks stay warm, a lag-triggered keepalive refold keeps a later depth probe from paying the whole gap refold and drafting from stale rows.
• Step-3.7 checkpoint handling is hardened, oQ and model loading now share one nextn prefix resolver.
• A skipped benchmark upload no longer hangs its subscribers, the skip event was missing from the terminal SSE event set.

New Contributors

Thank you to @layer4down, @cebertowicz, @tobomobo, @gilbert-barajas, @1qh, @PowerSpy, @FabioMalpezzi, @ShiroKSH, @NirvanaCh7, and @True2456 for their first contributions to oMLX in the 0.5.4 release line.

Full Changelog: v0.5.4rc1...v0.5.4rc2