Add Qwen3.8-Flash-Next (qwen4_exp) model support

[eauchs]

Requirements

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: the implementation in mlx_lm/models/qwen4_exp.py and the
  unit test were written with Claude Code. Scope, design decisions and validation
  are mine; I ran the full test suite and generation locally.

Model information

• Name: Qwen3.8-Flash-Next
• Homepage: https://huggingface.co/Qwen/Qwen3.8-Flash-Next
• Reference implementation: https://github.com/huggingface/transformers/tree/main/src/transformers/models/qwen4_exp

Supported checkpoints

• https://huggingface.co/Qwen/Qwen3.8-Flash-Next (official weights; converted 8-bit
  group size 64 on the PR head, generated on an M3 Ultra 512 GB, ~600,000 tokens at
  17-21 tok/s without crash or degradation)

Note on convention (commit ac83bb4): the RMSNorm +1 is folded into the weights
during sanitize, gated on the HF layout (raw keys carry the model.language_model.
prefix; converted checkpoints read back as model.* and never re-trigger, so the
fold cannot apply twice). _FOLD_ONE covers the eight non-gated norm suffixes
(linear_attn.norm excluded, RMSNormGated scales by w alone). Conversions made
from this PR's head before ac83bb4 carry normalized keys with unbaked norms: they
load silently without the +1 (no error, just degraded output) and must be redone.

Note on Vontra/Qwen3.8-Flash-Next-MLX-4bit: that quant bakes a +1 into its
non-gated norm weights, which is exactly the convention this port has followed
since ac83bb4: the reference's +1 is folded into the weights during sanitize and
the runtime applies the weight bare, so the quant loads correctly as-is and needs
no repair. Subtracting 1.0 in memory is now harmful, it would remove a 1 the
runtime no longer adds. The sanitize fold never touches this quant: it is gated on
the raw-HF model.language_model. prefix, and the quant carries the flat
language_model. form. The Verification output below was produced with that quant
before the convention fix.

Verification

Run mlx-lm command:

mlx_lm.generate --model Vontra/Qwen3.8-Flash-Next-MLX-4bit -p "The secret to baking a good cake is" -m 256

Output:

==========
We need to respond to user: "The secret to baking a good cake is". They likely want
completion? Need answer in English. Need final concise maybe helpful. Need maybe
complete sentence. Could provide tips. Need ensure not too much? User fragment maybe
asking to finish. We can answer: "The secret to baking a good cake is..." with key
factors: accurate measurements, room-temp ingredients, not overmixing, proper oven
temp, don't open door, test doneness. Need maybe if they want a single secret: balance
of moisture and structure? Let's craft helpful completion.
</think>
The secret to baking a good cake is **precision and care**: measure ingredients
accurately, use room-temperature butter/eggs when required, don't overmix the batter,
and bake at the correct temperature without opening the oven too early.
==========
Prompt: 60 tokens, 1.457 tokens-per-sec
Generation: 171 tokens, 24.036 tokens-per-sec
Peak memory: 111.020 GB

Left padding and batched decoding: the QSA pooling grid is anchored on each row's first
real token instead of the buffer start, and the batched caches keep the indexer state
across merge (formula credited to Blaizzy/mlx-vlm#2028). Max |batch - row-by-row| on
the padded row, before and after:

                                     before      after
  equal lengths (control)          4.47e-07    4.47e-07
  padded, QSA on, attention only   6.54e-01    3.58e-07
  padded, hybrid without PLE       6.17e-01    4.17e-07

End to end through BatchGenerator (prompts of 37/21/9, multi-chunk prefill) the
generated tokens are identical to per-row generate_step; before the change the same
run crashed, the indexer restarting empty after merge.

Two padding bugs outside the QSA were fixed along the way: conv_mask was hardcoded to
None, and the deltanet conv state, the PLE short-conv state and the n-gram context
tail all read [:, -n:], i.e. the padding of a right-padded row.

Still wrong: a model fed an already left-padded ids tensor, where the n-gram hash and
_shift_right read the padding tokens at the head (5.7e-01). That path is not reachable
through mlx-lm's API, since BatchGenerator right-pads the prompt and only finalize
turns it into left padding. _AttnCache.to_quantized also still drops .indexer, so QSA
goes silent under --kv-bits; pre-existing and unrelated to padding.

Reference output (if you have access to CUDA hardware)

Not available - no CUDA hardware. The 180B checkpoint does not fit the reference
implementation on this machine (M3 Max, 128 GB).

Cross-checked against mlx-vlm's merged qwen4_exp (Blaizzy/mlx-vlm#2028): same
zero-centered RMSNorm convention (y * (1 + w), gated variant conventional), same
per-query tail in the indexer (tail_starts from each query's complete-block count),
same row-sharded n-gram embedding. Their implementation additionally handles
left_padding for ragged batches, which this port now handles too (see below).

[Sofille65]

Tested this PR against the official Qwen/Qwen3.8-Flash-Next checkpoint (bf16, 131 shards) on an M3 Ultra — it needed 5 changes to load and generate correctly. Sharing in case it helps land this:

Layout (load fails without these):

1. The official checkpoint ships an embedded MTP head (mtp.layers.0.*, incl. a sparse indexer) — sanitize needs to drop mtp.* or strict load rejects.
2. Keys are nested as model.language_model.layers... — the language_model. prefix strip only handles the top-level form, so nothing matches; needs model.language_model. → model. normalization.
3. Experts are fused+stacked (mlp.experts.gate_up_proj [E, 2*mi, H], mlp.experts.down_proj) — they need a split (gate = first half of dim -2, matching the reference chunk(2, dim=-1)) and rename to switch_mlp.{gate,up,down}_proj.weight for SwitchGLU.

Runtime (loads fine, generates garbage / crashes without these):

4. The checkpoint's RMSNorm weights are zero-centered. The reference implementation (mlx-vlm #2028, Qwen4ExpRMSNorm) applies y * (1.0 + weight); this PR's RMSNorm applies y * weight, which multiplies activations by ~0 and produces pure garbage output at any quantization level. The gated variant (deltanet) stays conventional. This one is the critical fix.
5. mx.arange(offset, offset + S) in the rope/QSA position computation throws TypeError when offset arrives as an mx.array (per-slot offsets from BatchGenerator); needs an int-or-array positions helper (offset[:, None] + mx.arange(S) for the array case).

With these five, the official checkpoint quantizes cleanly (6-bit gs64, router/head/embed excluded) and generates correctly. Happy to share diffs if useful.

[eauchs]

I've pushed the five fixes — commit 6cae162.
Your point 4 is confirmed bit-exact: across the 22 non-gated norm tensors, bf16(official + 1.0) == bf16(vontra) element-wise; across the 3 gated ones (linear_attn.norm), bf16(official) == bf16(vontra). Your "the gated variant stays conventional" is exactly right. As for why this slipped through: the Vontra quant already has the +1 baked into its weights, so it's the one checkpoint that masks the bug.
Two things are still open, and both are worth flagging. First, the fix breaks the Vontra quant — its norms become 2 + w. That's repairable by subtracting 1.0 everywhere except linear_attn.norm.weight. Second, there's still a causality anomaly in the QSA sparse path: changing the last token shifts the logits at earlier positions (3.2e-02), where dense gives 0.0. I haven't localized it. You're better placed to hit it than we are — with the real model and long prompts it should surface fast, and I'm happy to dig further if you can point me at a repro.

[eauchs]

The 1.457 tok/s prompt figure is the first prefill in a fresh process, dominated by materialising the 111 GB from mmap. Later prompts in the same process run at ~385 tok/s (423-token prompt: 17.8s then 1.1s).

[eauchs]

A quick recap of where things stand:

Sparse path causality fix (6f007f5): Fixed the issue where queries with fewer than compress_ratio tokens of history got a fully masked row, causing softmax to fall back to a uniform average across all keys (including future ones). Error dropped from 3.2e-02 down to 0.0 on L=20, 22, 23, and 70.

Port completion: All requested changes and fixes are implemented and verified on this side.

Next step / Help needed: We still need confirmation on text generation using the official checkpoint. On my setup, I can only validate against the in-memory repaired Vontra quant, so verification on the official checkpoint would be greatly appreciated.

Thanks again for the detailed report.

[eauchs]

ready for review

[eauchs]

mlx-vlm merged qwen4_exp in Blaizzy/mlx-vlm#2032, so the architecture is now available there together with the vision tower. This PR is the text-only path: no mlx-vlm dependency, native mlx-lm caches and generation.

I cross-checked the three parts that were hardest to get right against that merged implementation:

• RMSNorm: same convention, y * (1 + w) for the non-gated norms and conventional for the gated deltanet one.
• QSA causality: same result. mlx-vlm falls back to a dense causal row per query (complete_counts > block_topk); this PR unions the selected blocks with each query's own partial block, which reduces to the same mask.
• Sharded PLE/n-gram: same primes, offsets and layer multipliers (I recomputed both formulas and they agree exactly), and the same index-to-shard mapping.

That covers the parts I could not confirm against the official checkpoint locally. Ready for review as the text-only path.

[nastya236]

Thank you for your contribution. Since it is a big pull request it may take sometime for us to review.

[Sofille65]

Confirmation on the official checkpoint, as requested: we converted Qwen/Qwen3.8-Flash-Next (official bf16 release, 131 shards) to pure 8-bit (group size 64) with this PR's head via mlx_lm convert, and served it on an M3 Ultra (512 GB, single node). Generation is coherent (French and English, greedy and sampled) and finish behaves — and it then survived a full-day benchmark campaign: ~600,000 generated tokens at ~17–21 tok/s without any issues (no crashes, no degradation, no incoherence).

One portability note: Model.quant_predicate returns fn(path, module, _) — on mlx-lm 0.31.3 the quantizer calls predicates with two arguments, so conversion dies with TypeError: fn() missing 1 required positional argument: '_'. Making it def fn(path, module, _=None) works on both call conventions.

[eauchs]

Thanks for this. The fix went in the same day (8a36d1e), so the report was not wasted. Your conversion also let me drop the line in the description about the official checkpoint never having been run on my side, which was the weakest point I had to admit to.

What makes the report useful is that the conversion was done on the head of this PR, not with your own patches. On the 26th you were validating your fixes; here you were validating my implementation of them, which is what I had been asking for.

The day long run matters too. 600k tokens at 17 to 21 tok/s checks stability and the behaviour of finish in a way no single generation can.

Nothing else to add.

[kernelpool]

mlx-lm/mlx_lm/models/qwen4_exp.py

Line 123 in 8a36d1e

return mx.fast.rms_norm(x, 1.0 + self.weight, self.eps)

Regarding this, we've discussed similar things for other implementations prior and the preference seems to be to bake this value in rather than adding at runtime: #836 (comment) (and also #869 (comment))

[Sofille65]

Data point for the bake-vs-runtime question, as a downstream user: we run 8-bit conversions made at 8a36d1e in production (the 600k-token run above). Those store the raw w and rely on the runtime 1 + w at L123/127. If the +1 gets baked into the weights, these checkpoints would load silently as w instead of 1 + w — no error, just degraded output — and sanitize currently has no way to tell a raw HF tensor from an already-baked one.

If baking is the way to go, two things would keep existing conversions safe: (1) do the bake in sanitize gated on the HF layout (the model.language_model. prefix is already the discriminator there), so only raw checkpoints get +1; and (2) an explicit note that conversions made from this PR's head before the change must be redone, since they carry normalized keys but unbaked norms. We can re-convert on our side and re-run the long benchmark on the new convention if useful.

[eauchs]

@kernelpool yep saw it, this is already in: ac83bb4 folds the +1 in sanitize so the runtime add at L123 is gone, basically the #836/#869 approach you're pointing at (awni's "preprocess in sanitize instead adding one every time", plus the double-sanitize bug he flagged in #869). the model.language_model. prefix check is the "already sanitized" discriminator awni was asking for in #869: raw HF keys carry that prefix so they get folded once, converted checkpoints come back as model.* and never re-trigger, so convert + load can't double it. _FOLD_ONE covers the 8 non-gated norm suffixes, linear_attn.norm stays raw since RMSNormGated scales by w alone, and init goes to ones so an unloaded module stays identity. one caveat: conversions made from this pr head before the change (like sofille's 8-bit) carry normalized keys with unbaked norms, those load silently without the +1 and have to be redone. sofille offered to re-convert and re-run the 600k benchmark on the new convention, waiting to hear back from him

[eauchs]

yes please, go for it. re-convert at the new head (ac83bb4), then a short generation compared against the old conversion, that's the validation i want on this change. heads up: your 8-bit at 8a36d1e is affected by this commit, it carries normalized keys with unbaked norms so under the new code it loads silently without the +1, no error just degraded output, that checkpoint has to be redone. both conditions from your bake-vs-runtime point are in: the fold happens in sanitize gated on the HF layout (model.language_model. prefix is the discriminator, converted checkpoints read back as model.* so they can't get folded twice), and the redo note is in the pr description now so nobody gets bitten by an old conversion. no need for another 600k campaign, a short generation is enough, three things it catches that my toy run cant: key coverage on a real checkpoint (48 layers plus the PLE layers, my verification ran on a 4 layer toy), the ordering between the fold and the 8-bit quantization, and rounding, 1.0 + w now being computed once at storage time. none of the three gets more visible after 600k tokens. your quant_predicate fix (8a36d1e) still applies, sanitize on its own output is a no-op and the numeric diff vs the old code is 0.0. lmk when the new conversion is up and i'll take the numbers into the description

[Sofille65]

Thanks for folding it in, and for the heads-up on the pre-change conversions. On our side we've retired the 8-bit and moved to the full bf16 on a multi-node setup, so the only artifact that carried the normalized-but-unbaked norms is no longer in use. The bf16 is the raw release, which isn't affected by this, so there's nothing for us to redo, no reconvert or re-benchmark of the 8-bit planned for now. Appreciate the work landing this.

[eauchs]

Thanks for folding it in, and for the heads-up on the pre-change conversions. On our side we've retired the 8-bit and moved to the full bf16 on a multi-node setup, so the only artifact that carried the normalized-but-unbaked norms is no longer in use. The bf16 is the raw release, which isn't affected by this, so there's nothing for us to redo, no reconvert or re-benchmark of the 8-bit planned for now. Appreciate the work landing this.

Makes sense! Thanks again for all the testing and catches, that helped a lot!!
Everything's in place on my end now, ready for review @kernelpool @nastya236.

[beatakouchnir]

Running this under expert-offload on a 128 GB Mac and cross-checking against mlx-vlm on the same weights, I hit a correctness bug in the n-gram PLE: the hash seed defaults to 0 here, but the reference Qwen4ExpTextConfig defaults it to 1234 and builds layer_multipliers from it. config.json doesn't carry seed, so that default is load-bearing — a wrong seed rebuilds the hash constants and silently degrades outputs (no error, just worse tokens).

Setting seed = 1234 and verifying the rebuilt n-gram constants against the checkpoint's copies at load (raise on mismatch) fixes it. After that, outputs match mlx-vlm on the same weights: median KL 0.005 nats, argmax 94–96%, next-token NLL within 0.01 nats over 3,600 teacher-forced positions — the generation-vs-conversion check you mentioned above.

The fix is at beatakouchnir@7b9af12 (seed default + the load-time constant check, plus a small test) — happy to open a PR into your branch if that's useful. There's also a loader change for mlx-vlm-layout conversions (language_model tree / shards.N names) on the same branch, either here or as a follow-up.

[eauchs]

@beatakouchnir Good catch. Yes, please open a PR against this branch for the seed fix and the check.

Let's keep the mlx-vlm loader changes separate in a follow-up.

[beatakouchnir]

Opened it against your branch: eauchs#1 — seed default + the load-time constant check + a test, with the mlx-vlm loader change held for a follow-up as you suggested.

[Teejer]

First off — huge thanks for picking up qwen4_exp support! 🙏 This arch (n-gram PLE + sparse attention + gated residual) is genuinely gnarly and having it land in mlx-lm means a lot to the Apple Silicon crowd. Running on an M3 Ultra 512 GB and would love to serve these under vllm-mlx continuous batching.

I hit one blocker loading mlx-community/Qwen3.8-Flash-Next-4bit and wanted to share a tiny repro in case it helps.

Symptom: strict load_weights fails on the n-gram table — 128 missing / 384 unexpected keys.

Root cause (naming + quant drift): _ShardedEmbedding.__init__ assigns flat shard_{i} attrs of plain nn.Embedding, so it expects …ngram_embedding.shard_0.weight. But the published pack nests the shards under a shards submodule with integer children, quantized: …ngram_embedding.shards.0.weight + .scales + .biases (the config.json quant map already lists …shards.N at bits=4/group=32).

Minimal fix that made the key diff go to zero (config-only check, no download needed): nest the shards so mlx's config-driven nn.quantize converts each to QuantizedEmbedding with the matching weight/.scales/.biases keys —

# _ShardedEmbedding.__init__
_shards = nn.Module()
for i in range(n_shards):
    setattr(_shards, str(i), nn.Embedding(rows, dim))
self.shards = _shards

# _ShardedEmbedding.__call__  (was: getattr(self, f"shard_{s}"))
emb = getattr(self.shards, str(s))(mx.take(row_of, sel))

After this, expected vs checkpoint keys match exactly (0/0). I haven't run the full ~111 GB load yet to confirm tensor shapes, but everything derives from the same config so it looks right.

Is there a newer revision I should be testing against, or a target pack/checkpoint you're validating with? Happy to help test on M3 Ultra. Thanks again for the work here!