[pmichalas]

Requirements

• I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: The model code and the test were written with AI (Claude). I reviewed them, and checked the port against an independent fp32 PyTorch implementation (details below).

Model information

• Name: NVIDIA Nemotron-Labs Dory 1.0 4B (dory), a looped nGPT-style reasoning model (16 input layers, 40 shared layers run n_recurrent_loops times, 16 output layers; 72 layers alternating attention and MLP, sliding window 1024 with a full-attention layer every 8)
• Homepage: https://huggingface.co/nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16
• Reference implementation: not in Transformers; written from the vLLM implementation ([Model] Add Dory support vllm-project/vllm#60149) and cross-checked against llama.cpp

Supported checkpoints

• NVIDIA's original checkpoint: https://huggingface.co/nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16 (sanitize fuses the q/k/v/gate and MLP gate/up weights when loading)
• https://huggingface.co/netprojects/nemotron-dory1.0-4b-mlx-8bit
• https://huggingface.co/netprojects/nemotron-dory1.0-4b-mlx-bf16

Verification

Run mlx-lm command:

mlx_lm.generate --model netprojects/nemotron-dory1.0-4b-mlx-8bit -p "The secret to baking a good cake is" -m 1024

Output (shortened; the model reasons first, and the 1024-token limit stopped it before the final answer):

==========
The user is asking "The secret to baking a good cake is". This is an open-ended question. I need to provide a helpful, accurate, and comprehensive answer about cake baking.

The user didn't specify a particular type of cake (chocolate, vanilla, sponges, etc.), so I should give a general answer that covers the fundamental principles of good cake baking, while perhaps mentioning that different cakes have different needs.

Key elements of good cake baking generally include:
1. Quality ingredients (fresh eggs, good flour, etc.)
2. Proper measurements (especially leavening agents like baking powder/soda)
3. Creaming method (for butter cakes) - beating butter and sugar together to incorporate air
4. Not overmixing (to avoid too much gluten)
...

Let me draft a response:

"The secret to baking a good cake isn't a single trick—it's a combination of technique, ingredients, and patience. Here are the key principles that usually make the difference:

1. **Room temperature ingredients**: Let your butter, eggs, and milk sit out for 30–60 minutes. ...
2. **Creaming is crucial**: For butter cakes, beating softened butter and sugar together until light and fluffy incorporates air, which helps the cake rise and stay tender.
3. **Measure accurately**: Spoon dry ingredients into your measuring cup and level them off with a knife. ...
...
==========
Prompt: 24 tokens, 16.343 tokens-per-sec
Generation: 1024 tokens, 18.231 tokens-per-sec
Peak memory: 4.551 GB

Other checks:

• test_dory added to tests/test_models.py (passes).
• Last-token logits on the real weights agree with an independent fp32 PyTorch implementation (top-1 agreement 100%, KL 0.001 to 0.003, 2 loops, 55 tokens). 8-bit vs bf16: top-1 agreement 98%, KL 0.003.
• Sliding-window caches checked with the window shrunk to 64 (chunked prefill and token-by-token decode vs one pass): top-1 agreement at least 98.7%.

Note on the tokenizer: the released Transformers (5.20.1) does not know dory and rejects the mlp entries in layer_types, so AutoTokenizer fails on this checkpoint. This PR therefore adds a fallback in tokenizer_utils.load that loads the fast tokenizer directly from the checkpoint's tokenizer files when AutoTokenizer fails. Token ids match the raw tokenizer.json on the texts I tried.

Reference output (if you have access to CUDA hardware)

No CUDA hardware, not run.

[pmichalas]

Independent M1 Pro numbers from nace-ai/kapitsa#296: with this PR, Dory 4-bit decode holds ~28 tok/s at 1K context and ~19 tok/s at 32K (−31%).

On an M5 Max (128 GB), plain mlx-lm at this PR (208065f), 8-bit, 2 loops, per-loop KV, temp 0, no speculation, medians of 3 runs (1 GiB MLX cache limit, which costs ~5–7% of prefill speed; decode unaffected):

Prompt → reply	Prompt tok/s	Gen tok/s	Peak
25 → 505	402	68.0	4.39 GB
4,096 → 192	2198	64.4	8.97 GB
16,384 → 192	2101	59.6	8.97 GB
32,768 → 192	1874	53.0	8.97 GB

So 8-bit decode drops ~22% from a short prompt to 32K. The PR runner and the dory-mlx package agree within ~3%. Earlier single run with the snsnc 4-bit weights, short prompt: 103 tok/s decode, 2.5 GB peak.

Edit: updated with a rerun using exact 4K/16K/32K prompts and medians of 3.

Happy to add anything else you need for review.

[xquantize]

Tested this PR on an M3 Pro (18 GB) with bf16, 8-bit, 6-bit and 4-bit builds at k = 1, 2, 4. Works well; full results and charts are on the model card here: https://huggingface.co/xquantize/nemotron-labs-dory1.0-4b-mlx-6bit

One thing I noticed in last_loop mode: make_cache only uses RotatingKVCache when the mode is per_loop, so in last_loop every attention layer gets a full KVCache, including the 16 input/output layers that never loop. With an 8,000-token prompt (6-bit), last_loop used 1.21 GB of cache at every k, more than per_loop even at k=4 (0.42 / 0.65 / 1.11 GB for k = 1 / 2 / 4).

If only the recurrent layers need a full cache for the slot-overwriting, something like this might work (untested):

recurrent = a.n_input_layers <= i < a.n_input_layers + a.n_recurrent_layers
if a.layer_types[i] == "sliding_attention" and (use_sliding or not recurrent):
    caches[idx] = RotatingKVCache(max_size=a.sliding_window)

Related question: in last_loop mode, is the 1,024-token window still enforced by the attention mask for the recurrent sliding layers once the context exceeds it? My quality checks (99.8% top-1 vs per_loop) only used passages shorter than the window, so I couldn't test that.

Other small notes: the transformers tokenizer warning about fix_mistral_regex appears to be a false alarm (identical token IDs with and without the fix, and against the raw tokenizer.json, across 8 test strings). Also, the quantization error does compound slightly with loops at 4-bit (KL 0.113 → 0.128 from k=1 to k=4), but not at 6- or 8-bit.

[pmichalas]

Thanks @xquantize, good catch, and thanks for the multi-quant testing.

Pushed 986ab63: in last_loop mode the non-recurrent sliding layers (input/output) now get RotatingKVCache, as you suggested. The recurrent layers keep a full KVCache, since later loops rewind cache.offset and overwrite the same slots.

Is the window still enforced in last_loop? Yes. The sliding layers call create_attention_mask(x, cache, window_size=sliding_window), and KVCache.make_mask builds a windowed causal mask whenever window_size is set, including for a single decode token (create_causal_mask(N, offset, window_size=...) adds linds < rinds + window_size). So the change saves memory without changing outputs: with the 8-bit build, a ~3K-token prompt gives the same 64 greedy tokens before and after.

KV cache after prefill (sum of keys + values), 8-bit, M5 Max:

Mode	Context	Before	After
last_loop, k = 1 / 2 / 4	8K	1.21 GB	0.96 GB
last_loop, k = 1 / 2 / 4	32K	4.83 GB	3.37 GB
per_loop, k = 1 / 2 / 4 (unchanged)	8K	0.64 / 1.00 / 1.71 GB
per_loop, k = 1 / 2 / 4 (unchanged)	32K	1.55 / 2.41 / 4.13 GB

The recurrent layers still hold most of the last_loop cache (24 of its 36 caches), so the saving is about 21% at 8K and 30% at 32K. Decode at 8K, last_loop, k = 2: 58.9 → 60.2 tok/s (median of 3).

Tests added in tests/test_models.py:

• test_dory_last_loop_cache_types: checks the cache type per layer in both modes.
• test_dory_last_loop_past_window: chunked prefill plus token-by-token decoding well past an 8-token window. per_loop matches a full uncached pass. last_loop matches the same run with all-KVCache caches.
  • last_loop is not compared with a full pass, because it keeps only the last loop's K/V for the recurrent layers, so it differs from a full pass by design (~1e-3 on the tiny model).
  • The tiny model sets its nGPT scale parameters to initializer_range, as they are stored in a checkpoint. With all-ones the attention logits are scaled by roughly 1e5 and nothing matches, even in per_loop.

Thanks also for the tokenizer and 4-bit KL notes.

[pmichalas]

Pushed a small fix for a flaky test (test_dory_sliding_window_cache now seeds before building the layer; 50/50 passes on mlx 0.32.4 and main). Also checked against current main: rebases cleanly, and Dory tests and real-weight parity pass on mlx main (which includes #4734); the fused sliding-window path also matches Dory's window mask exactly, so it can be adopted once it ships. Could a maintainer approve the CI workflow run? Thanks!

[pmichalas]

Quick status update:

• Merged current main into the branch (no conflicts). The Dory tests pass on MLX 0.32.4; the rest of tests/test_models.py gives the same results as on main. Parity against the independent fp32 PyTorch reference is unchanged (real weights, 2 loops: top-1 100%, KL 0.0017 bf16 / 0.0035 8-bit).
• Converted builds that load with this model file:
  • 8-bit and bf16
  • New: a 4-bit build made from NVIDIA's own NVFP4 checkpoint, in MLX's nvfp4 format (2.6 GB; WikiText-2 perplexity at 2 loops 11.59 vs 11.62 for bf16).
• An independent measurement on a 16 GB M1 Pro with this PR is in nace-ai/kapitsa#296.

The CI workflow is still waiting for approval. Could a maintainer approve the run when convenient? Happy to make any changes.