Add files using upload-large-folder tool

README.md

@@ -0,0 +1,136 @@
+---
+license: mit
+base_model: deepseek-ai/DeepSeek-V4-Flash-0731
+base_model_relation: quantized
+pipeline_tag: text-generation
+library_name: mlx
+tags:
+- mlx
+- moe
+- deepseek
+- deepseek-v4
+---
+
+# DeepSeek-V4-Flash-MLX-REAP25
+
+MLX (Apple Silicon) conversion of [deepseek-ai/DeepSeek-V4-Flash-0731](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-0731),
+built from `mixed-4_8bit` (4-bit routed experts, 8-bit elsewhere) with **25% of routed experts pruned** by REAP saliency.
+
+**Code / loader:** [github.com/PipeNetwork/deepseek-v4-mlx](https://github.com/PipeNetwork/deepseek-v4-mlx)
+
+DeepSeek-V4-Flash is a **304B-parameter** sparse-MoE model. Its `deepseek_v4`
+architecture is **not** implemented in any released runtime — not `transformers`,
+and not `mlx-lm` (which carries `deepseek`, `_v2`, `_v3`, `_v32`). This repo bundles
+a from-scratch MLX port under `deepseek_v4_mlx/`, validated against DeepSeek's own
+reference implementation.
+
+## Builds
+
+| Variant | Size | wikitext-2 ppl | Notes |
+|---|---:|---:|---|
+| [8bit](https://huggingface.co/pipenetwork/DeepSeek-V4-Flash-MLX-8bit) | ~304 GB | 5.9878 | highest fidelity |
+| [6bit](https://huggingface.co/pipenetwork/DeepSeek-V4-Flash-MLX-6bit) | ~233 GB | 6.0260 | +0.6% vs 8-bit for 70 GB less |
+| [mixed-4_8bit](https://huggingface.co/pipenetwork/DeepSeek-V4-Flash-MLX-mixed-4_8bit) | ~165 GB | 6.1262 | **recommended** — 4-bit experts, 8-bit everything else |
+| [4bit](https://huggingface.co/pipenetwork/DeepSeek-V4-Flash-MLX-4bit) | ~162 GB | 6.3005 | superseded by `mixed-4_8bit` (3 GB more, 2.8% better) |
+| [REAP25](https://huggingface.co/pipenetwork/DeepSeek-V4-Flash-MLX-REAP25) | ~129 GB | 6.4025 | 25% of experts pruned — **matches 4-bit quality 33 GB smaller** |
+| [REAP37](https://huggingface.co/pipenetwork/DeepSeek-V4-Flash-MLX-REAP37) | ~111 GB | 6.9634 | 37% pruned; fits a 128 GB Mac |
+| [REAP50](https://huggingface.co/pipenetwork/DeepSeek-V4-Flash-MLX-REAP50) | ~93 GB | 7.7520 | 50% pruned; smallest build |
+
+Perplexity is teacher-forced over **wikitext-2-raw-v1 test**, 199,485 tokens in 195
+windows of 1024 — identical windows for every build. Differences are established by
+a **paired** bootstrap over shared windows, which removes window-difficulty variance:
+independent intervals overlap almost entirely and settle nothing, while the paired
+interval is ~10x tighter. Every gap quoted here is significant at 95%.
+
+### Why `mixed-4_8bit` is the recommended build
+
+The routed experts are **98.1%** of all quantized parameters, and they ship natively
+as FP4. Keeping them at 4 bits while raising everything else — attention,
+embeddings, shared experts, about 5B parameters — to 8 bits costs **3 GB** and
+recovers **56%** of the gap between 4-bit and 8-bit. Per parameter, the non-expert
+weights are roughly 50x more quantization-sensitive.
+
+Two variants were built, measured, and **not** published:
+
+* **group-size 32** (matching FP4's native 32-value scale blocks) is 0.5% *worse*
+  than group 64 and 18 GB larger. Finer scaling improves raw tensor reconstruction
+  but not prediction — the mismatch is that FP4's levels are non-uniform while
+  affine's are evenly spaced, which no group size fixes.
+* **3-bit experts** collapse the model entirely (perplexity 1.4e7, degenerate token
+  loops). 4-bit sits just above a cliff for this checkpoint, not in the middle of a
+  gentle curve.
+
+### REAP-pruned builds
+
+[REAP (Cerebras, arXiv:2510.13999)](https://arxiv.org/abs/2510.13999) drops the
+lowest-saliency routed experts, where saliency is the mean of
+`router_weight x ||expert_output||` over the tokens routed to that expert — its
+actual contribution, not how often it fires. The router renormalizes over survivors.
+
+**Only 40 of the 43 layers are pruned.** The first three route through a fixed
+`tid2eid` token-id table with no router to renormalize, so dropping an expert there
+would leave table entries pointing at nothing. They keep all 256, and the config
+records two expert counts which the loader resolves per layer.
+
+Saliency was profiled over **150k tokens** of English prose, ten other languages and
+real code — deliberately disjoint from the evaluation set, verified by 32-gram
+overlap (0 of 4687), so pruning decisions are not fit to the text they are judged on.
+
+| Build | Experts kept | Saliency retained | ppl | vs unpruned |
+|---|---:|---:|---:|---:|
+| `mixed-4_8bit` | 256 | 100% | 6.1262 | — |
+| `REAP25` | 192 | 96.2% | 6.4025 | +4.5% |
+| `REAP37` | 161 | 91.5% | 6.9634 | +13.7% |
+| `REAP50` | 128 | 83.9% | 7.7520 | +26.5% |
+
+**`REAP25` is the one to take** — its perplexity is statistically indistinguishable
+from plain 4-bit (paired ratio 1.016, CI [0.998, 1.038], crossing 1.0) at 33 GB less.
+
+Note how poorly saliency retention predicts damage: 96.2% retained still costs 4.5%,
+and the last 12 points of retention cost 22 points of perplexity. Treat it as a
+screen for whether pruning is viable, not as an estimate of what it will cost.
+
+## ⚠️ Loading requires the bundled loader
+
+```bash
+pip install mlx mlx-lm transformers
+```
+```python
+from deepseek_v4_mlx.load import load
+from deepseek_v4_mlx.generate import greedy_generate, load_tokenizer
+
+model, args = load("/path/to/this/repo")
+tok = load_tokenizer()          # DeepSeek-V3's tokenizer — see below
+ids = tok("The capital of France is")["input_ids"]
+print(tok.decode(greedy_generate(model, args, ids, max_new_tokens=32)))
+```
+
+**The base repo ships no tokenizer** — only `encoding/encoding_dsv4.py`, which renders
+prompt *strings*. DeepSeek-V3's tokenizer is the correct one: same padded `vocab_size`
+129280 over 128815 real tokens, `bos=0`/`eos=1` matching V4's `generation_config.json`,
+and every V4 special token present.
+
+## What this port implements
+
+| Mechanism | Notes |
+|---|---|
+| **Hyper-Connections** | the residual stream carries **4 parallel copies**, mixed by a Sinkhorn-normalized combination matrix |
+| **Hash-routed experts** | the first 3 MoE layers pick experts from a `tid2eid[vocab, 6]` table keyed on token id — routing independent of context |
+| **Learned KV compression** | gated pooling over spans, alternating ratio 4/128, overlapping windows at ratio 4 |
+| **Sparse-attention indexer** | scores compressed spans, keeps the top 512 for attention to read |
+| **MLA + grouped output LoRA** | single shared KV vector; `o_groups: 8` block-diagonal output projection |
+| **sqrt-softplus routing**, clamped SwiGLU, attention sinks | |
+
+The source checkpoint is **doubly pre-quantized** — FP8 `e4m3` with `e8m0` 128x128
+block scales for most matmuls, FP4 `e2m1` packed two-per-byte for routed experts — so
+conversion dequantizes both formats before requantizing. Both decode paths are
+bit-exact against an independent implementation.
+
+Dropped: the multi-token-prediction head and **DSpark** (speculative decode over
+5-token blocks in the last three layers). Both hang off the reference's `forward_spec`
+path and are unused for ordinary generation.
+
+`wo_a` is kept unquantized: it is a block-diagonal projection the model reshapes by
+group, and a packed quantized weight cannot be reshaped that way. Costs ~2% of build size.
+
+License: MIT (inherits the base model).

config.json

@@ -0,0 +1,122 @@
+{
+  "architectures": [
+    "DeepseekV4ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "eos_token_id": 1,
+  "hc_eps": 1e-06,
+  "hc_mult": 4,
+  "hc_sinkhorn_iters": 20,
+  "head_dim": 512,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "index_head_dim": 128,
+  "index_n_heads": 64,
+  "index_topk": 512,
+  "initializer_range": 0.02,
+  "max_position_embeddings": 1048576,
+  "model_type": "deepseek_v4",
+  "moe_intermediate_size": 2048,
+  "n_routed_experts": 256,
+  "n_shared_experts": 1,
+  "norm_topk_prob": true,
+  "num_attention_heads": 64,
+  "num_experts_per_tok": 6,
+  "num_hidden_layers": 43,
+  "num_hash_layers": 3,
+  "num_key_value_heads": 1,
+  "num_nextn_predict_layers": 1,
+  "o_groups": 8,
+  "o_lora_rank": 1024,
+  "q_lora_rank": 1024,
+  "qk_rope_head_dim": 64,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": {
+    "beta_fast": 32,
+    "beta_slow": 1,
+    "factor": 16,
+    "original_max_position_embeddings": 65536,
+    "type": "yarn"
+  },
+  "rope_theta": 10000,
+  "routed_scaling_factor": 1.5,
+  "scoring_func": "sqrtsoftplus",
+  "sliding_window": 128,
+  "swiglu_limit": 10.0,
+  "tie_word_embeddings": false,
+  "topk_method": "noaux_tc",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.57.1",
+  "use_cache": true,
+  "vocab_size": 129280,
+  "compress_rope_theta": 160000,
+  "compress_ratios": [
+    0,
+    0,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    0,
+    0,
+    0
+  ],
+  "dspark_block_size": 5,
+  "dspark_noise_token_id": 128799,
+  "dspark_target_layer_ids": [
+    40,
+    41,
+    42
+  ],
+  "dspark_markov_rank": 256,
+  "quantization": {
+    "group_size": 64,
+    "bits": 8,
+    "expert_bits": 4
+  },
+  "reap": {
+    "kept_experts": 192,
+    "ratio_pct": 25,
+    "hash_layers_keep": 256,
+    "scored_layers": 40
+  }
+}

deepseek_v4_mlx/attention.py

@@ -0,0 +1,239 @@
+"""DeepSeek-V4 attention — MLA over a sliding window plus selected compressed spans.
+
+Shape of one layer:
+
+* queries go through a low-rank bottleneck (``wq_a`` -> ``q_norm`` -> ``wq_b``),
+  then get a **weightless** per-head RMS normalization;
+* a single shared KV vector per position (MLA) — key and value are the same tensor;
+* every layer attends to the last ``window_size`` tokens; compression layers also
+  attend to a set of pooled spans, chosen by :class:`Indexer` at ratio 4 or taken
+  wholesale at ratio 128;
+* the output has the rotary embedding applied **in reverse** before a grouped
+  low-rank projection (``wo_a`` per group, then a shared ``wo_b``).
+
+Both paths live here. ``__call__`` prefills a whole sequence and, given a cache,
+leaves it in the state decode would have produced; ``step`` decodes one token
+against that cache. They share the q/kv/output projections so the two cannot drift
+apart.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .compressor import Compressor
+from .config import ModelArgs
+from .fakequant import fake_quant_kv
+from .indexer import Indexer
+from .layers import RMSNorm, apply_rotary_emb, precompute_freqs_cis
+from .sparse_attention import sparse_attn
+
+
+def window_topk_idxs(window_size: int, seqlen: int) -> mx.array:
+    """[s, w] indices of the local window for each query, -1 where out of range."""
+    base = mx.arange(seqlen).reshape(-1, 1)
+    w = min(seqlen, window_size)
+    idx = mx.maximum(base - window_size + 1, 0) + mx.arange(w).reshape(1, -1)
+    return mx.where(idx > base, mx.array(-1, mx.int32), idx.astype(mx.int32))
+
+
+def compress_topk_idxs(ratio: int, seqlen: int, offset: int) -> mx.array:
+    """[s, s//ratio] indices of every *complete* span, -1 where not yet visible."""
+    nb = seqlen // ratio
+    m = mx.broadcast_to(mx.arange(nb).reshape(1, -1), (seqlen, nb))
+    vis = m < (mx.arange(1, seqlen + 1).reshape(-1, 1) // ratio)
+    return mx.where(vis, (m + offset).astype(mx.int32), mx.array(-1, mx.int32))
+
+
+class Attention(nn.Module):
+    def __init__(self, layer_id: int, args: ModelArgs):
+        super().__init__()
+        self.layer_id = layer_id
+        self.dim = args.dim
+        self.n_heads = args.n_heads
+        self.head_dim = args.head_dim
+        self.rope_head_dim = args.rope_head_dim
+        self.q_lora_rank = args.q_lora_rank
+        self.o_lora_rank = args.o_lora_rank
+        self.n_groups = args.o_groups
+        self.window_size = args.window_size
+        self.compress_ratio = args.compress_ratio(layer_id)
+        self.eps = args.norm_eps
+        self.softmax_scale = args.head_dim ** -0.5
+
+        self.attn_sink = mx.zeros((self.n_heads,), dtype=mx.float32)
+        self.wq_a = nn.Linear(self.dim, self.q_lora_rank, bias=False)
+        self.q_norm = RMSNorm(self.q_lora_rank, self.eps)
+        self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False)
+        self.wkv = nn.Linear(self.dim, self.head_dim, bias=False)
+        self.kv_norm = RMSNorm(self.head_dim, self.eps)
+        self.wo_a = nn.Linear(self.n_heads * self.head_dim // self.n_groups,
+                              self.n_groups * self.o_lora_rank, bias=False)
+        self.wo_b = nn.Linear(self.n_groups * self.o_lora_rank, self.dim, bias=False)
+
+        if self.compress_ratio:
+            self.compressor = Compressor(args, self.compress_ratio, self.head_dim)
+            self.indexer = Indexer(args, self.compress_ratio) if self.compress_ratio == 4 else None
+
+        # compression layers use YaRN on a longer rope base; pure sliding-window
+        # layers disable YaRN entirely and use the standard base
+        if self.compress_ratio:
+            orig_len, theta = args.original_seq_len, args.compress_rope_theta
+        else:
+            orig_len, theta = 0, args.rope_theta
+        self._rope = (args.rope_head_dim, orig_len, theta, args.rope_factor,
+                      args.beta_fast, args.beta_slow)
+        self._cos = None
+        self._sin = None
+
+    def _freqs(self, seqlen: int):
+        if self._cos is None or self._cos.shape[0] < seqlen:
+            rd, orig_len, theta, factor, bf, bs = self._rope
+            self._cos, self._sin = precompute_freqs_cis(rd, seqlen, orig_len, theta,
+                                                        factor, bf, bs)
+        return self._cos[:seqlen], self._sin[:seqlen]
+
+    def _project_q(self, x, cos, sin, pos0: int):
+        """Shared query path: low-rank, weightless per-head RMS, rope."""
+        bsz, seqlen, _ = x.shape
+        rd = self.rope_head_dim
+        qr = self.q_norm(self.wq_a(x))
+        q = self.wq_b(qr).reshape(bsz, seqlen, self.n_heads, self.head_dim)
+        qf = q.astype(mx.float32)
+        q = (qf * mx.rsqrt(mx.mean(mx.square(qf), axis=-1, keepdims=True) + self.eps)).astype(q.dtype)
+        c, s = cos[pos0:pos0 + seqlen], sin[pos0:pos0 + seqlen]
+        q = mx.concatenate([q[..., :-rd], apply_rotary_emb(q[..., -rd:], c, s)], axis=-1)
+        return q, qr
+
+    def _project_kv(self, x, cos, sin, pos0: int):
+        """Shared KV path: single MLA vector, rope on the tail, FP8 on the rest."""
+        rd = self.rope_head_dim
+        seqlen = x.shape[1]
+        kv = self.kv_norm(self.wkv(x))
+        c, s = cos[pos0:pos0 + seqlen], sin[pos0:pos0 + seqlen]
+        kv = mx.concatenate([kv[..., :-rd], apply_rotary_emb(kv[..., -rd:], c, s)], axis=-1)
+        return fake_quant_kv(kv, rd, 64)
+
+    def _out(self, o, cos, sin, pos0: int, bsz: int, seqlen: int, dtype):
+        """Inverse rope, then the grouped low-rank projection.
+
+        ``dtype`` is the caller's working precision — hardcoding bf16 here silently
+        downcasts an fp32 forward.
+        """
+        rd = self.rope_head_dim
+        c, s = cos[pos0:pos0 + seqlen], sin[pos0:pos0 + seqlen]
+        o = mx.concatenate([o[..., :-rd], apply_rotary_emb(o[..., -rd:], c, s, inverse=True)],
+                           axis=-1)
+        o = o.reshape(bsz, seqlen, self.n_groups, -1)
+        wo_a = self.wo_a.weight.reshape(self.n_groups, self.o_lora_rank, -1)
+        o = mx.einsum("bsgd,grd->bsgr", o.astype(mx.float32), wo_a.astype(mx.float32))
+        return self.wo_b(o.reshape(bsz, seqlen, -1).astype(dtype))
+
+    def step(self, x: mx.array, cache) -> mx.array:
+        """Decode one token against ``cache``. ``x`` is [b, 1, dim]."""
+        bsz = x.shape[0]
+        pos = cache.offset
+        cos, sin = self._freqs(pos + 1)
+
+        q, qr = self._project_q(x, cos, sin, pos)
+        kv = self._project_kv(x, cos, sin, pos)
+        cache.put_window(pos, kv[:, 0])
+
+        idxs = cache.window_idxs(pos)[None, None, :]
+        if self.compress_ratio:
+            if self.indexer is not None:
+                cidx = self.indexer.step(x, qr, cos, sin, pos, cache, self.window_size)
+            else:
+                ci = cache.compressed_idxs(pos)
+                cidx = ci[None, None, :] if ci.size else None
+            pooled = self.compressor.step(x, pos, cache.comp, cos, sin)
+            if pooled is not None:
+                cache.put_compressed(pos // self.compress_ratio, pooled[:, 0])
+            if cidx is not None and cidx.shape[-1] > 0:
+                idxs = mx.concatenate([mx.broadcast_to(idxs, (bsz, 1, idxs.shape[-1])),
+                                       mx.broadcast_to(cidx, (bsz, 1, cidx.shape[-1]))], axis=-1)
+
+        idxs = mx.broadcast_to(idxs, (bsz, 1, idxs.shape[-1]))
+        o = sparse_attn(q, cache.kv, self.attn_sink, idxs, self.softmax_scale)
+        out = self._out(o, cos, sin, pos, bsz, 1, x.dtype)
+        cache.offset = pos + 1
+        return out
+
+    def __call__(self, x: mx.array, cache=None) -> mx.array:
+        if cache is not None and x.shape[1] == 1 and cache.offset > 0:
+            return self.step(x, cache)
+        bsz, seqlen, _ = x.shape
+        rd = self.rope_head_dim
+        cos, sin = self._freqs(seqlen)
+
+        # --- queries: low-rank, then a weightless per-head RMS ---
+        qr = self.q_norm(self.wq_a(x))
+        q = self.wq_b(qr).reshape(bsz, seqlen, self.n_heads, self.head_dim)
+        qf = q.astype(mx.float32)
+        q = (qf * mx.rsqrt(mx.mean(mx.square(qf), axis=-1, keepdims=True) + self.eps)).astype(q.dtype)
+        q = mx.concatenate([q[..., :-rd], apply_rotary_emb(q[..., -rd:], cos, sin)], axis=-1)
+
+        # --- shared KV (MLA) ---
+        kv = self.kv_norm(self.wkv(x))
+        kv = mx.concatenate([kv[..., :-rd], apply_rotary_emb(kv[..., -rd:], cos, sin)], axis=-1)
+        kv = fake_quant_kv(kv, rd, 64)
+
+        idxs = mx.broadcast_to(window_topk_idxs(self.window_size, seqlen)[None],
+                               (bsz, seqlen, min(seqlen, self.window_size)))
+
+        if self.compress_ratio:
+            offset = seqlen                       # compressed entries follow the window KV
+            if self.indexer is not None:
+                cidx = self.indexer(x, qr, cos, sin, offset)
+            else:
+                cidx = mx.broadcast_to(
+                    compress_topk_idxs(self.compress_ratio, seqlen, offset)[None],
+                    (bsz, seqlen, seqlen // self.compress_ratio))
+            if cidx is not None:
+                idxs = mx.concatenate([idxs, cidx], axis=-1)
+            kv_c = self.compressor(x, cos, sin)
+            if kv_c is not None:
+                kv = mx.concatenate([kv, kv_c], axis=1)
+
+        o = sparse_attn(q, kv, self.attn_sink, idxs, self.softmax_scale)
+
+        if cache is not None:
+            self._seed_cache(cache, x, kv, seqlen, cos, sin)
+
+        return self._out(o, cos, sin, 0, bsz, seqlen, x.dtype)
+
+    def _seed_cache(self, cache, x, kv, seqlen, cos, sin):
+        """Leave the cache in exactly the state decode would have produced.
+
+        Prefill pools spans in one shot, so the compressors' rolling state has to be
+        replayed. Replay starts at the **last complete span**, not merely at the
+        remainder: for overlapping (ratio-4) compressors, closing that span is what
+        moves it into the "previous span" half of the state, which the next pooling
+        reads. Starting at the remainder would leave that half empty and quietly
+        change the first few decoded tokens.
+
+        Slots beyond the replayed remainder keep stale values, which is harmless —
+        a span only closes once every slot in the current half has been rewritten.
+        """
+        cache.put_window_prefill(kv[:, :seqlen])
+        if not self.compress_ratio:
+            cache.offset = seqlen
+            return
+
+        ratio = self.compress_ratio
+        n_span = seqlen // ratio
+        if n_span:
+            cache.put_compressed_prefill(kv[:, seqlen:seqlen + n_span])
+            if self.indexer is not None:
+                idx_kv = self.indexer.compressor(x, cos, sin)
+                if idx_kv is not None:
+                    cache._adopt(idx_kv.dtype)
+                    cache.idx_kv[:, :idx_kv.shape[1]] = idx_kv
+
+        start = max(0, (n_span - 1) * ratio)
+        for p in range(start, seqlen):
+            self.compressor.step(x[:, p:p + 1], p, cache.comp, cos, sin)
+            if self.indexer is not None:
+                self.indexer.compressor.step(x[:, p:p + 1], p, cache.idx_comp, cos, sin)
+        cache.offset = seqlen

deepseek_v4_mlx/cache.py

@@ -0,0 +1,142 @@
+"""Incremental decode state for DeepSeek-V4.
+
+Three things persist between tokens, and they have different shapes:
+
+* **the sliding window** — a *circular* buffer of ``window_size`` KV entries.
+  Position ``p`` lives at slot ``p % window_size``, so reading it back in
+  chronological order means rotating the slot list, not slicing it.
+* **compressed spans** — appended after the window in the same buffer, one entry
+  per completed span at slot ``window_size + p // ratio``.
+* **the compressor's partial span** — the tokens seen since the last span closed,
+  held as running ``kv``/``score`` state. Ratio-4 layers keep two spans' worth,
+  because their pooling overlaps the previous span.
+
+The indexer keeps its own compressor state and its own compressed-KV buffer,
+independent of attention's.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+NEG_INF = -1e30
+
+
+class CompressorCache:
+    """Rolling state for one compressor's in-progress span."""
+
+    def __init__(self, bsz: int, ratio: int, head_dim: int, overlap: bool):
+        coff = 1 + overlap
+        self.ratio = ratio
+        self.overlap = overlap
+        self.head_dim = head_dim
+        self.kv_state = mx.zeros((bsz, coff * ratio, coff * head_dim), dtype=mx.float32)
+        self.score_state = mx.full((bsz, coff * ratio, coff * head_dim), NEG_INF,
+                                   dtype=mx.float32)
+
+    def write(self, slot: int, kv: mx.array, score: mx.array):
+        """Record one token's contribution at ``slot``."""
+        self.kv_state[:, slot] = kv
+        self.score_state[:, slot] = score
+
+    def shift(self):
+        """After closing an overlapping span, the current half becomes the previous."""
+        r = self.ratio
+        self.kv_state[:, :r] = self.kv_state[:, r:]
+        self.score_state[:, :r] = self.score_state[:, r:]
+
+
+class LayerCache:
+    """Per-layer decode state: window + compressed KV, plus compressor state."""
+
+    def __init__(self, bsz: int, args, layer_id: int, max_seq_len: int,
+                 dtype=mx.bfloat16):
+        self.window = args.window_size
+        self.ratio = args.compress_ratio(layer_id)
+        self.offset = 0
+        self.dtype = dtype
+
+        n_comp = (max_seq_len // self.ratio) if self.ratio else 0
+        self.kv = mx.zeros((bsz, self.window + n_comp, args.head_dim), dtype=dtype)
+
+        if self.ratio:
+            self.comp = CompressorCache(bsz, self.ratio, args.head_dim, self.ratio == 4)
+        else:
+            self.comp = None
+
+        if self.ratio == 4:
+            self.idx_comp = CompressorCache(bsz, self.ratio, args.index_head_dim, True)
+            self.idx_kv = mx.zeros((bsz, n_comp, args.index_head_dim), dtype=dtype)
+        else:
+            self.idx_comp = None
+            self.idx_kv = None
+
+    def _adopt(self, dtype):
+        """Match the precision actually flowing through the model.
+
+        Allocating bf16 while the forward runs in fp32 truncates every cached KV on
+        write. The absolute error is small, but it flips discrete decisions
+        downstream — expert top-k and indexer span selection — so decode drifts away
+        from full recomputation far faster than the rounding alone suggests.
+        """
+        if dtype == self.dtype:
+            return
+        self.dtype = dtype
+        self.kv = self.kv.astype(dtype)
+        if self.idx_kv is not None:
+            self.idx_kv = self.idx_kv.astype(dtype)
+
+    # ---- window ----
+
+    def put_window(self, pos: int, kv: mx.array):
+        self._adopt(kv.dtype)
+        self.kv[:, pos % self.window] = kv
+
+    def put_window_prefill(self, kv: mx.array):
+        """Seed the circular buffer from a prefill of ``s`` tokens.
+
+        Keeps the last ``window`` entries, placed so that position ``p`` sits at
+        slot ``p % window`` — the same invariant decode maintains.
+        """
+        self._adopt(kv.dtype)
+        s = kv.shape[1]
+        if s <= self.window:
+            self.kv[:, :s] = kv
+            return
+        tail = kv[:, -self.window:]                 # positions s-window .. s-1
+        cut = s % self.window                       # where position s-1 lands + 1
+        self.kv[:, cut:self.window] = tail[:, :self.window - cut]
+        self.kv[:, :cut] = tail[:, self.window - cut:]
+
+    def put_compressed(self, span: int, kv: mx.array):
+        self._adopt(kv.dtype)
+        self.kv[:, self.window + span] = kv
+
+    def put_compressed_prefill(self, kv: mx.array):
+        self._adopt(kv.dtype)
+        n = kv.shape[1]
+        self.kv[:, self.window:self.window + n] = kv
+
+    # ---- index lists for one decoding step ----
+
+    def window_idxs(self, pos: int) -> mx.array:
+        """Window slots in chronological order for a query at ``pos``."""
+        w = self.window
+        if pos >= w - 1:
+            p = pos % w
+            return mx.concatenate([mx.arange(p + 1, w), mx.arange(0, p + 1)]).astype(mx.int32)
+        return mx.concatenate([mx.arange(pos + 1),
+                               mx.full((w - pos - 1,), -1)]).astype(mx.int32)
+
+    def compressed_idxs(self, pos: int) -> mx.array:
+        """Every span completed by ``pos``, addressed in the shared buffer."""
+        n = (pos + 1) // self.ratio
+        if n <= 0:
+            return mx.zeros((0,), dtype=mx.int32)
+        return (mx.arange(n) + self.window).astype(mx.int32)
+
+
+def make_cache(args, bsz: int = 1, max_seq_len: int | None = None):
+    """One :class:`LayerCache` per layer."""
+    max_seq_len = max_seq_len or 4096
+    return [LayerCache(bsz, args, i, max_seq_len) for i in range(args.n_layers)]

deepseek_v4_mlx/compressor.py

@@ -0,0 +1,154 @@
+"""Learned KV compression — V4 pools spans of tokens into single cache entries.
+
+Instead of caching every position, compression layers pool ``compress_ratio``
+consecutive tokens into one entry using a *learned gate*: a projection scores each
+position, the scores are softmaxed across the span, and the KV vectors are summed
+under those weights. A per-slot positional embedding (``ape``) lets the gate know
+where inside the span a token sat.
+
+Ratio-4 layers additionally **overlap**: each pooled entry sees 2x ratio positions
+— the current span plus the preceding one — so information is not lost at span
+boundaries. That is what the doubled ``wkv`` / ``wgate`` output width is for: the
+first ``head_dim`` channels carry the overlapping (previous-span) contribution and
+the second ``head_dim`` the current one.
+
+Only the prefill path (``start_pos == 0``) lives here. Incremental decode
+maintains rolling state buffers and is added alongside the attention cache.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .config import ModelArgs
+from .fakequant import fake_quant_fp4, fake_quant_kv
+from .layers import RMSNorm, apply_rotary_emb, hadamard_transform
+
+NEG_INF = -1e30
+
+
+class Compressor(nn.Module):
+    """Pools ``compress_ratio`` tokens into one KV entry via a learned softmax gate."""
+
+    def __init__(self, args: ModelArgs, compress_ratio: int = 4,
+                 head_dim: int = 512, rotate: bool = False):
+        super().__init__()
+        self.dim = args.dim
+        self.head_dim = head_dim
+        self.rope_head_dim = args.rope_head_dim
+        self.compress_ratio = compress_ratio
+        self.overlap = compress_ratio == 4
+        self.rotate = rotate
+        coff = 1 + self.overlap
+
+        self.ape = mx.zeros((compress_ratio, coff * head_dim), dtype=mx.float32)
+        self.wkv = nn.Linear(self.dim, coff * head_dim, bias=False)
+        self.wgate = nn.Linear(self.dim, coff * head_dim, bias=False)
+        self.norm = RMSNorm(head_dim, args.norm_eps)
+
+    def _overlap_transform(self, t: mx.array, value: float) -> mx.array:
+        """[b, n, ratio, 2d] -> [b, n, 2*ratio, d].
+
+        Slots ``[ratio:]`` take the *second* half of the channels from this span;
+        slots ``[:ratio]`` take the *first* half from the **previous** span. The
+        first span has no predecessor, so its leading slots stay at ``value``
+        (0 for KV, -inf for scores, so the softmax ignores them).
+        """
+        b, n, r, _ = t.shape
+        d = self.head_dim
+        cur = t[..., d:]                       # [b, n, ratio, d]
+        prev = t[..., :d]                      # [b, n, ratio, d]
+        pad = mx.full((b, 1, r, d), value, dtype=t.dtype)
+        prev_shifted = mx.concatenate([pad, prev[:, :-1]], axis=1)
+        return mx.concatenate([prev_shifted, cur], axis=2)   # [b, n, 2*ratio, d]
+
+    def __call__(self, x: mx.array, cos: mx.array, sin: mx.array):
+        """Prefill: compress ``x`` [b, s, dim] into [b, s//ratio, head_dim].
+
+        Returns ``None`` when the sequence is shorter than one full span, matching
+        the reference's ``should_compress`` short-circuit.
+        """
+        bsz, seqlen, _ = x.shape
+        ratio, d, rd = self.compress_ratio, self.head_dim, self.rope_head_dim
+        dtype = x.dtype
+
+        if seqlen < ratio:
+            return None
+
+        xf = x.astype(mx.float32)
+        kv = self.wkv(xf)
+        score = self.wgate(xf)
+
+        remainder = seqlen % ratio
+        cutoff = seqlen - remainder
+        if remainder:
+            kv = kv[:, :cutoff]
+            score = score[:, :cutoff]
+
+        nb = cutoff // ratio
+        kv = kv.reshape(bsz, nb, ratio, -1)
+        score = score.reshape(bsz, nb, ratio, -1) + self.ape
+
+        if self.overlap:
+            kv = self._overlap_transform(kv, 0.0)
+            score = self._overlap_transform(score, NEG_INF)
+
+        kv = mx.sum(kv * mx.softmax(score, axis=2), axis=2)   # [b, nb, d]
+        kv = self.norm(kv.astype(dtype))
+
+        # one rope position per pooled span, taken at the span's first token
+        c = cos[:cutoff:ratio]
+        s = sin[:cutoff:ratio]
+        rot = apply_rotary_emb(kv[..., -rd:], c, s)
+        kv = mx.concatenate([kv[..., :-rd], rot], axis=-1)
+
+        return self._finish(kv, rd)
+
+    def _finish(self, kv: mx.array, rd: int) -> mx.array:
+        # QAT simulation: the indexer's compressor rotates then fakes FP4 over the
+        # whole vector; the attention compressor fakes FP8 on the non-rope half only.
+        if self.rotate:
+            kv = hadamard_transform(kv)
+            return fake_quant_fp4(kv, 32)
+        return fake_quant_kv(kv, rd, 64)
+
+    def step(self, x: mx.array, pos: int, cache, cos: mx.array, sin: mx.array):
+        """Decode one token. Returns the pooled entry when ``pos`` closes a span.
+
+        ``x`` is [b, 1, dim]. Between span boundaries this only accumulates state
+        and returns ``None`` — the caller writes nothing to the KV buffer.
+        """
+        ratio, d, rd = self.compress_ratio, self.head_dim, self.rope_head_dim
+        dtype = x.dtype
+        xf = x.astype(mx.float32)
+        kv = self.wkv(xf)[:, 0]                       # [b, coff*d]
+        score = self.wgate(xf)[:, 0] + self.ape[pos % ratio]
+
+        if self.overlap:
+            cache.write(ratio + pos % ratio, kv, score)
+        else:
+            cache.write(pos % ratio, kv, score)
+
+        if (pos + 1) % ratio != 0:
+            return None
+
+        if self.overlap:
+            # first half of the channels from the previous span, second half from this
+            kv_s = mx.concatenate([cache.kv_state[:, :ratio, :d],
+                                   cache.kv_state[:, ratio:, d:]], axis=1)
+            sc_s = mx.concatenate([cache.score_state[:, :ratio, :d],
+                                   cache.score_state[:, ratio:, d:]], axis=1)
+        else:
+            kv_s, sc_s = cache.kv_state, cache.score_state
+
+        pooled = mx.sum(kv_s * mx.softmax(sc_s, axis=1), axis=1, keepdims=True)
+        if self.overlap:
+            cache.shift()
+
+        pooled = self.norm(pooled.astype(dtype))
+        # the span's rope position is its first token
+        p = pos + 1 - ratio
+        rot = apply_rotary_emb(pooled[..., -rd:], cos[p:p + 1], sin[p:p + 1])
+        pooled = mx.concatenate([pooled[..., :-rd], rot], axis=-1)
+        return self._finish(pooled, rd)

deepseek_v4_mlx/config.py

@@ -0,0 +1,152 @@
+"""Configuration for DeepSeek-V4 (MLX port).
+
+Parses the HuggingFace ``config.json`` layout published with
+``deepseek-ai/DeepSeek-V4-Flash-0731``. The reference ``inference/config.json``
+uses different key names for the same values (``dim`` vs ``hidden_size``,
+``n_layers`` vs ``num_hidden_layers``, ...), so both spellings are accepted —
+handy for running the reference and this port off one file during parity work.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+def _get(d: dict, *names, default=None):
+    for n in names:
+        if n in d and d[n] is not None:
+            return d[n]
+    return default
+
+
+@dataclass
+class ModelArgs:
+    vocab_size: int = 129280
+    dim: int = 4096
+    n_layers: int = 43
+    n_hash_layers: int = 3
+    moe_inter_dim: int = 2048
+
+    # attention (MLA)
+    n_heads: int = 64
+    head_dim: int = 512
+    rope_head_dim: int = 64
+    q_lora_rank: int = 1024
+    o_lora_rank: int = 1024
+    o_groups: int = 8
+    window_size: int = 128
+    norm_eps: float = 1e-6
+
+    # MoE
+    n_routed_experts: int = 256
+    n_shared_experts: int = 1
+    n_activated_experts: int = 6
+    score_func: str = "sqrtsoftplus"
+    route_scale: float = 1.5
+    swiglu_limit: float = 10.0
+
+    # KV compression: one ratio per layer (0 = none, pure sliding window)
+    compress_ratios: tuple = ()
+    compress_rope_theta: float = 160000.0
+
+    # rope / YaRN
+    original_seq_len: int = 65536
+    rope_theta: float = 10000.0
+    rope_factor: float = 16.0
+    beta_fast: int = 32
+    beta_slow: int = 1
+    max_seq_len: int = 1048576
+
+    # sparse-attention indexer
+    index_n_heads: int = 64
+    index_head_dim: int = 128
+    index_topk: int = 512
+
+    # hyper-connections
+    hc_mult: int = 4
+    hc_sinkhorn_iters: int = 20
+    hc_eps: float = 1e-6
+
+    # REAP: experts kept in the *scored* layers (None = unpruned)
+    reap_kept_experts: int | None = None
+
+    # side-paths dropped for inference (kept for provenance)
+    n_mtp_layers: int = 0
+    dspark_block_size: int = 0
+    dspark_target_layer_ids: tuple = ()
+
+    @property
+    def nope_head_dim(self) -> int:
+        return self.head_dim - self.rope_head_dim
+
+    def compress_ratio(self, layer_id: int) -> int:
+        """Compression ratio for a layer; 0 means pure sliding-window attention."""
+        if layer_id < len(self.compress_ratios):
+            return int(self.compress_ratios[layer_id])
+        return 0
+
+    def is_hash_layer(self, layer_id: int) -> bool:
+        """First ``n_hash_layers`` route experts by token-id lookup, not by score."""
+        return layer_id < self.n_hash_layers
+
+    def n_experts_for(self, layer_id: int) -> int:
+        """Routed-expert count for a layer — REAP prunes only the scored layers.
+
+        Hash layers index a fixed ``tid2eid`` table, so dropping an expert there
+        would leave table entries pointing at nothing. They keep the full set, and
+        a pruned checkpoint therefore has *two* expert counts. Storing one global
+        ``n_routed_experts`` would silently build the hash layers too small and
+        load garbage into them.
+        """
+        if self.reap_kept_experts and not self.is_hash_layer(layer_id):
+            return self.reap_kept_experts
+        return self.n_routed_experts
+
+    @classmethod
+    def from_dict(cls, c: dict) -> "ModelArgs":
+        rope = _get(c, "rope_scaling", default={}) or {}
+        return cls(
+            vocab_size=_get(c, "vocab_size", default=129280),
+            dim=_get(c, "hidden_size", "dim", default=4096),
+            n_layers=_get(c, "num_hidden_layers", "n_layers", default=43),
+            n_hash_layers=_get(c, "num_hash_layers", "n_hash_layers", default=0),
+            moe_inter_dim=_get(c, "moe_intermediate_size", "moe_inter_dim", default=2048),
+            n_heads=_get(c, "num_attention_heads", "n_heads", default=64),
+            head_dim=_get(c, "head_dim", default=512),
+            rope_head_dim=_get(c, "qk_rope_head_dim", "rope_head_dim", default=64),
+            q_lora_rank=_get(c, "q_lora_rank", default=1024),
+            o_lora_rank=_get(c, "o_lora_rank", default=1024),
+            o_groups=_get(c, "o_groups", default=8),
+            window_size=_get(c, "sliding_window", "window_size", default=128),
+            norm_eps=_get(c, "rms_norm_eps", "norm_eps", default=1e-6),
+            n_routed_experts=_get(c, "n_routed_experts", default=256),
+            n_shared_experts=_get(c, "n_shared_experts", default=1),
+            n_activated_experts=_get(c, "num_experts_per_tok", "n_activated_experts", default=6),
+            score_func=_get(c, "scoring_func", "score_func", default="sqrtsoftplus"),
+            route_scale=_get(c, "routed_scaling_factor", "route_scale", default=1.5),
+            swiglu_limit=_get(c, "swiglu_limit", default=0.0),
+            compress_ratios=tuple(_get(c, "compress_ratios", default=()) or ()),
+            compress_rope_theta=_get(c, "compress_rope_theta", default=160000.0),
+            original_seq_len=_get(rope, "original_max_position_embeddings",
+                                  default=_get(c, "original_seq_len", default=65536)),
+            rope_theta=_get(c, "rope_theta", default=10000.0),
+            rope_factor=_get(rope, "factor", default=_get(c, "rope_factor", default=16.0)),
+            beta_fast=_get(rope, "beta_fast", default=_get(c, "beta_fast", default=32)),
+            beta_slow=_get(rope, "beta_slow", default=_get(c, "beta_slow", default=1)),
+            max_seq_len=_get(c, "max_position_embeddings", "max_seq_len", default=1048576),
+            index_n_heads=_get(c, "index_n_heads", default=64),
+            index_head_dim=_get(c, "index_head_dim", default=128),
+            index_topk=_get(c, "index_topk", default=512),
+            hc_mult=_get(c, "hc_mult", default=4),
+            hc_sinkhorn_iters=_get(c, "hc_sinkhorn_iters", default=20),
+            hc_eps=_get(c, "hc_eps", default=1e-6),
+            reap_kept_experts=(c.get("reap") or {}).get("kept_experts"),
+            n_mtp_layers=_get(c, "num_nextn_predict_layers", "n_mtp_layers", default=0),
+            dspark_block_size=_get(c, "dspark_block_size", default=0),
+            dspark_target_layer_ids=tuple(_get(c, "dspark_target_layer_ids", default=()) or ()),
+        )
+
+    @property
+    def raw(self) -> dict[str, Any]:
+        return {"model_type": "deepseek_v4"}

deepseek_v4_mlx/convert.py

@@ -0,0 +1,218 @@
+"""Streaming DeepSeek-V4 checkpoint -> MLX, with optional affine quantization.
+
+The published weights are FP8/FP4 (see :mod:`.dequant`), so conversion is
+dequantize-then-requantize rather than a straight cast. It runs one source shard
+at a time and the layout cooperates: no shard spans more than one layer, so a
+layer's 256 experts are always resident together and can be stacked without
+buffering the model.
+
+Three transforms are applied:
+
+* **dequantize** — every ``(weight, scale)`` pair to bf16; FP4 for routed experts,
+  FP8 for everything else;
+* **stack experts** — the checkpoint stores ``ffn.experts.{0..255}.w{1,2,3}``
+  separately; MLX's ``SwitchGLU`` wants one ``[n_experts, out, in]`` tensor per
+  projection (``w1 -> gate_proj``, ``w3 -> up_proj``, ``w2 -> down_proj``);
+* **drop the side-paths** — ``mtp.*`` covers both the multi-token-prediction head
+  and DSpark, neither of which is used for ordinary generation.
+
+``bits=None`` writes bf16; ``bits in {2,3,4,5,6,8}`` affine-quantizes the matmul
+weights, leaving norms, the router, ``tid2eid``, the compressor projections and the
+Hyper-Connection parameters in their original precision.
+"""
+
+from __future__ import annotations
+
+import glob
+import json
+import os
+import re
+import shutil
+
+import mlx.core as mx
+
+from .dequant import dequant_pair, is_fp4_expert
+
+_EXPERT_RE = re.compile(r"^(layers\.\d+\.ffn)\.experts\.(\d+)\.(w[123])$")
+_PROJ = {"w1": "gate_proj", "w3": "up_proj", "w2": "down_proj"}
+
+_SHARD_CAP_BYTES = 5_000_000_000
+
+# Leaves that get affine-quantized; everything else keeps its source precision.
+#
+# `wo_a` is deliberately excluded. It is a *block-diagonal* projection — input
+# group g multiplies only block g of the weight — which attention implements by
+# reshaping the weight to [n_groups, o_lora_rank, -1]. After nn.quantize a
+# Linear's `.weight` is the packed uint32 tensor, so that reshape would carve up
+# packed bytes rather than logical weights. Keeping it unquantized costs ~1.4B
+# parameters (about 2% of a 4-bit build) and keeps the reshape meaningful. Moving
+# it to a SwitchLinear would recover that, at the cost of a layout change.
+_QUANT_LEAVES = {"wq_a", "wq_b", "wkv", "wo_b", "w1", "w2", "w3",
+                 "gate_proj", "up_proj", "down_proj", "embed", "head"}
+
+
+def should_drop(name: str) -> bool:
+    """MTP and DSpark live under `mtp.` and are inference-irrelevant here."""
+    return name.startswith("mtp.")
+
+
+def is_routed_expert(out_name: str) -> bool:
+    """The stacked routed-expert matmuls — 98% of all quantized parameters.
+
+    They ship natively as FP4 (16 levels), which 4-bit affine reproduces almost
+    exactly, so precision spent above 4 bits here buys little while costing nearly
+    the whole build size. Everything else is ~5B parameters and effectively free to
+    keep at high precision.
+    """
+    return ".ffn.experts." in out_name and ".shared_experts." not in out_name
+
+
+def is_quant_target(out_name: str, in_dim: int, group_size: int) -> bool:
+    if not out_name.endswith(".weight"):
+        return False
+    leaf = out_name[: -len(".weight")].rsplit(".", 1)[-1]
+    if leaf not in _QUANT_LEAVES:
+        return False
+    # the router gate and tid2eid share the `gate` prefix but are not in _QUANT_LEAVES
+    return in_dim % group_size == 0
+
+
+def bits_for(out_name: str, bits: int, expert_bits: int | None) -> int:
+    """Per-tensor bit width. ``expert_bits`` splits routed experts from the rest."""
+    if expert_bits is not None and is_routed_expert(out_name):
+        return expert_bits
+    return bits
+
+
+def _collect(tensors: dict) -> dict:
+    """Dequantize pairs, stack experts, drop side-paths. Returns MLX-named bf16."""
+    out: dict[str, mx.array] = {}
+    experts: dict[str, dict[int, mx.array]] = {}
+
+    names = [n for n in tensors if not should_drop(n)]
+    for name in names:
+        if name.endswith(".scale"):
+            continue                                   # consumed with its weight
+        arr = tensors[name]
+        base = name[: -len(".weight")] if name.endswith(".weight") else name
+        scale_name = base + ".scale"
+        if scale_name in tensors:
+            arr = dequant_pair(base, arr, tensors[scale_name], mx.bfloat16)
+
+        m = _EXPERT_RE.match(base)
+        if m:
+            prefix, idx, w = m.group(1), int(m.group(2)), m.group(3)
+            key = f"{prefix}.experts.{_PROJ[w]}.weight"
+            experts.setdefault(key, {})[idx] = arr
+            continue
+
+        if base.endswith("tid2eid"):
+            out[base] = arr.astype(mx.int32)           # stored int64, only needs int32
+            continue
+
+        out[name if name.endswith(".weight") else base] = arr
+
+    for key, byidx in experts.items():
+        stacked = mx.stack([byidx[i] for i in sorted(byidx)], axis=0)
+        out[key] = stacked
+    return out
+
+
+def convert(src: str, dst: str, bits: int | None = None, group_size: int = 64,
+            out_dtype=mx.bfloat16, verbose: bool = True,
+            expert_bits: int | None = None) -> str:
+    os.makedirs(dst, exist_ok=True)
+    index = json.load(open(os.path.join(src, "model.safetensors.index.json")))
+    weight_map = index["weight_map"]
+
+    shard_to_names: dict[str, list[str]] = {}
+    for n, s in weight_map.items():
+        shard_to_names.setdefault(s, []).append(n)
+
+    out_index: dict[str, str] = {}
+    buf: dict[str, mx.array] = {}
+    nbytes = 0
+    shard_id = 0
+    dropped = 0
+
+    def flush():
+        nonlocal buf, nbytes, shard_id
+        if not buf:
+            return
+        shard_id += 1
+        fname = f"model-{shard_id:05d}.safetensors"
+        mx.save_safetensors(os.path.join(dst, fname), buf, metadata={"format": "mlx"})
+        for k in buf:
+            out_index[k] = fname
+        buf = {}
+        nbytes = 0
+
+    for shard in sorted(shard_to_names):
+        tensors = mx.load(os.path.join(src, shard))
+        dropped += sum(1 for n in tensors if should_drop(n))
+        converted = _collect(tensors)
+        for out_name, arr in converted.items():
+            if bits is not None and is_quant_target(out_name, arr.shape[-1], group_size):
+                b = bits_for(out_name, bits, expert_bits)
+                qw, sc, bi = mx.quantize(arr, group_size=group_size, bits=b)
+                base = out_name[: -len(".weight")]
+                items = ((out_name, qw), (base + ".scales", sc), (base + ".biases", bi))
+            else:
+                keep_fp32 = arr.dtype == mx.float32 or arr.dtype == mx.int32
+                items = ((out_name, arr if keep_fp32 else arr.astype(out_dtype)),)
+            for k, v in items:
+                mx.eval(v)
+                buf[k] = v
+                nbytes += v.nbytes
+            if nbytes >= _SHARD_CAP_BYTES:
+                flush()
+        del tensors, converted
+        if verbose:
+            print(f"[convert] {shard} -> {shard_id} shards written", flush=True)
+    flush()
+
+    _finalize(dst, out_index, shard_id)
+    _write_config(src, dst, bits, group_size, expert_bits)
+    _copy_aux(src, dst)
+    if verbose:
+        print(f"[convert] dropped {dropped} mtp/dspark tensors")
+    return dst
+
+
+def _finalize(dst, out_index, n):
+    remap = {}
+    for i in range(1, n + 1):
+        old, new = f"model-{i:05d}.safetensors", f"model-{i:05d}-of-{n:05d}.safetensors"
+        if os.path.exists(os.path.join(dst, old)):
+            os.rename(os.path.join(dst, old), os.path.join(dst, new))
+        remap[old] = new
+    wm = {k: remap[v] for k, v in out_index.items()}
+    total = sum(os.path.getsize(os.path.join(dst, f)) for f in set(wm.values()))
+    json.dump({"metadata": {"total_size": total}, "weight_map": wm},
+              open(os.path.join(dst, "model.safetensors.index.json"), "w"), indent=2)
+
+
+def _write_config(src, dst, bits, group_size, expert_bits=None):
+    cfg = json.load(open(os.path.join(src, "config.json")))
+    # the source quantization_config describes FP8/FP4 weights that no longer exist
+    cfg.pop("quantization_config", None)
+    cfg.pop("expert_dtype", None)
+    if bits is not None:
+        q = {"group_size": group_size, "bits": bits}
+        if expert_bits is not None:
+            q["expert_bits"] = expert_bits
+        cfg["quantization"] = q
+    json.dump(cfg, open(os.path.join(dst, "config.json"), "w"), indent=2)
+
+
+def _copy_aux(src, dst):
+    for pat in ("tokenizer*", "*.json", "*.jinja", "*.txt", "encoding"):
+        for p in glob.glob(os.path.join(src, pat)):
+            base = os.path.basename(p)
+            if base in ("config.json", "model.safetensors.index.json"):
+                continue
+            target = os.path.join(dst, base)
+            if os.path.isdir(p):
+                shutil.copytree(p, target, dirs_exist_ok=True)
+            else:
+                shutil.copy2(p, target)

deepseek_v4_mlx/dequant.py

@@ -0,0 +1,92 @@
+"""Dequantize DeepSeek-V4's natively-quantized checkpoint to bf16.
+
+The published checkpoint is quantized in *two* different formats, and neither is
+a format MLX can consume directly:
+
+  * **Most matmul weights** — FP8 ``e4m3`` stored as ``uint8`` ``[out, in]``, with
+    one ``e8m0`` (exponent-only) scale per ``128x128`` block:
+    ``scale[ceil(out/128), ceil(in/128)]``.
+  * **Routed experts** (``expert_dtype: "fp4"``) — FP4 ``e2m1`` packed **two values
+    per byte** into ``int8 [out, in//2]``, with one ``e8m0`` scale per 32 values
+    along the input dim: ``scale[out, in//32]``.
+
+Everything else (norms, gate weights, compressor projections, embeddings, the LM
+head, the HC parameters) is already bf16 or fp32 and passes through untouched.
+
+The `e8m0` format is exponent-only: byte ``b`` denotes ``2**(b - 127)``, with
+``b == 0`` reserved for zero. That makes every scale an exact power of two, so
+applying it is lossless in fp32.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+FP8_BLOCK = 128   # e4m3 weights: 128x128 scale blocks
+FP4_BLOCK = 32    # e2m1 experts: 32 values per scale, along the input dim
+
+# FP4 e2m1 code -> value. 1 sign bit, 2 exponent bits, 1 mantissa bit; the 8
+# representable magnitudes are {0, .5, 1, 1.5, 2, 3, 4, 6} with a sign.
+# Transcribed from FP4_TABLE in the reference inference/convert.py — note code 8
+# ("negative zero") maps to +0.0 there, so it does the same here.
+_E2M1 = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
+         0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]
+
+
+def e8m0_to_float(scale_u8: mx.array) -> mx.array:
+    """Decode exponent-only fp8 scales: byte b -> 2**(b-127), with 0 -> 0."""
+    e = scale_u8.astype(mx.int32)
+    val = mx.power(mx.array(2.0, mx.float32), (e - 127).astype(mx.float32))
+    return mx.where(e == 0, mx.array(0.0, mx.float32), val)
+
+
+def _expand_blocks(scale: mx.array, rows: int, cols: int, br: int, bc: int) -> mx.array:
+    """Broadcast a per-block scale grid up to full [rows, cols], then crop.
+
+    Crop matters: block counts use ceil, so the expanded grid can overhang the
+    weight when a dimension is not a multiple of the block size.
+    """
+    s = mx.repeat(mx.repeat(scale, br, axis=0), bc, axis=1)
+    return s[:rows, :cols]
+
+
+def dequant_fp8(weight_u8: mx.array, scale_u8: mx.array, dtype=mx.bfloat16) -> mx.array:
+    """FP8 e4m3 ``[out, in]`` + e8m0 ``[out/128, in/128]`` block scales -> ``dtype``."""
+    w = mx.from_fp8(weight_u8, mx.float32)
+    rows, cols = w.shape
+    s = _expand_blocks(e8m0_to_float(scale_u8), rows, cols, FP8_BLOCK, FP8_BLOCK)
+    return (w * s).astype(dtype)
+
+
+def unpack_fp4(packed_i8: mx.array) -> mx.array:
+    """Unpack ``int8 [out, in//2]`` of two e2m1 nibbles into float32 ``[out, in]``.
+
+    Low nibble holds the even (lower) input index, high nibble the odd one — the
+    same ordering torch uses for ``float4_e2m1fn_x2``.
+    """
+    b = packed_i8.view(mx.uint8).astype(mx.int32)
+    lo = b & 0x0F
+    hi = (b >> 4) & 0x0F
+    lut = mx.array(_E2M1, mx.float32)
+    vals = mx.stack([lut[lo], lut[hi]], axis=-1)      # [out, in//2, 2]
+    return vals.reshape(vals.shape[0], -1)            # [out, in]
+
+
+def dequant_fp4(weight_i8: mx.array, scale_u8: mx.array, dtype=mx.bfloat16) -> mx.array:
+    """Packed FP4 e2m1 ``[out, in//2]`` + e8m0 ``[out, in//32]`` scales -> ``dtype``."""
+    w = unpack_fp4(weight_i8)                          # [out, in]
+    rows, cols = w.shape
+    s = _expand_blocks(e8m0_to_float(scale_u8), rows, cols, 1, FP4_BLOCK)
+    return (w * s).astype(dtype)
+
+
+def is_fp4_expert(name: str) -> bool:
+    """Routed-expert matmuls are the FP4 ones; shared experts are FP8."""
+    return ".experts." in name and ".shared_experts." not in name
+
+
+def dequant_pair(name: str, weight: mx.array, scale: mx.array, dtype=mx.bfloat16) -> mx.array:
+    """Dequantize one (weight, scale) pair, choosing the format from the name."""
+    if is_fp4_expert(name):
+        return dequant_fp4(weight, scale, dtype)
+    return dequant_fp8(weight, scale, dtype)

deepseek_v4_mlx/fakequant.py

@@ -0,0 +1,95 @@
+"""Activation fake-quantization — the inference-time half of DeepSeek-V4's QAT.
+
+V4 was quantization-aware trained, and the reference *simulates* that quantization
+during inference: several activations are rounded through FP8 (or FP4) and back to
+bf16 before use. These are the ``inplace=True`` calls to ``act_quant`` /
+``fp4_act_quant`` in ``inference/model.py``, and they are not optional — skipping
+them changes the numbers the model was trained to see.
+
+Both follow the same shape:
+
+  1. take the absolute max over each block along the last axis,
+  2. round the scale **up to a power of two**: ``2**ceil(log2(amax / dtype_max))``
+     (this is what ``scale_fmt="ue8m0"`` means),
+  3. divide, round to the target format, multiply back.
+
+Applied at:
+
+  * ``kv[..., :-rope_dim]`` in attention and the compressor — FP8, 64-wide blocks;
+    the rope channels stay bf16 so positional precision is not degraded,
+  * indexer queries and its compressed KV — FP4, 32-wide blocks, after a Hadamard
+    rotation.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+FP8_MAX = 448.0    # e4m3fn
+FP4_MAX = 6.0      # e2m1
+
+# e2m1 magnitudes; used to round a scaled value to the nearest representable FP4.
+_E2M1_POS = mx.array([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=mx.float32)
+
+
+def _pow2_scale(amax: mx.array, dtype_max: float) -> mx.array:
+    """``2**ceil(log2(amax / dtype_max))``, matching ``fast_round_scale``.
+
+    Guarded so an all-zero block yields scale 1 rather than 0 or a NaN.
+    """
+    r = amax / dtype_max
+    r = mx.maximum(r, 1e-38)
+    return mx.power(mx.array(2.0, mx.float32), mx.ceil(mx.log2(r)))
+
+
+def _blockify(x: mx.array, block: int):
+    n = x.shape[-1]
+    if n % block:
+        raise ValueError(f"last dim {n} not divisible by block {block}")
+    return x.reshape(*x.shape[:-1], n // block, block)
+
+
+def fake_quant_fp8(x: mx.array, block: int = 64) -> mx.array:
+    """Round ``x`` through e4m3 with per-block power-of-two scales, back to input dtype."""
+    dtype = x.dtype
+    shape = x.shape
+    xb = _blockify(x.astype(mx.float32), block)
+    amax = mx.max(mx.abs(xb), axis=-1, keepdims=True)
+    s = _pow2_scale(amax, FP8_MAX)
+    q = mx.from_fp8(mx.to_fp8(xb / s), mx.float32) * s
+    return q.reshape(shape).astype(dtype)
+
+
+def fake_quant_fp4(x: mx.array, block: int = 32) -> mx.array:
+    """Round ``x`` through e2m1 with per-block power-of-two scales, back to input dtype.
+
+    MLX has no fp4 cast, so rounding is done by snapping |v| to the nearest of the
+    eight e2m1 magnitudes (ties-to-even is not modelled; the reference kernel casts
+    via hardware, which rounds to nearest-even, so extreme midpoints can differ by
+    one code — well below the noise this simulation represents).
+    """
+    dtype = x.dtype
+    shape = x.shape
+    xb = _blockify(x.astype(mx.float32), block)
+    amax = mx.max(mx.abs(xb), axis=-1, keepdims=True)
+    s = _pow2_scale(amax, FP4_MAX)
+    v = mx.clip(xb / s, -FP4_MAX, FP4_MAX)
+
+    sign = mx.sign(v)
+    mag = mx.abs(v)
+    # nearest of the 8 magnitudes: compare against midpoints
+    lut = _E2M1_POS
+    diffs = mx.abs(mag[..., None] - lut.reshape(*([1] * mag.ndim), -1))
+    nearest = mx.argmin(diffs, axis=-1)
+    q = sign * lut[nearest]
+    return (q * s).reshape(shape).astype(dtype)
+
+
+def fake_quant_kv(kv: mx.array, rope_dim: int, block: int = 64) -> mx.array:
+    """FP8-simulate the non-rope channels of a KV vector, leaving rope channels alone.
+
+    Mirrors ``act_quant(kv[..., :-rd], 64, ..., inplace=True)`` in the reference,
+    which quantizes a *slice* in place.
+    """
+    nope = fake_quant_fp8(kv[..., :-rope_dim], block)
+    return mx.concatenate([nope, kv[..., -rope_dim:]], axis=-1)

deepseek_v4_mlx/generate.py

@@ -0,0 +1,37 @@
+"""Greedy generation with an incremental cache."""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+from .cache import make_cache
+
+
+def greedy_generate(model, args, prompt_ids, max_new_tokens: int = 32,
+                    eos_id: int | None = 1, max_seq_len: int | None = None,
+                    verbose: bool = False):
+    """Prefill the prompt, then decode one token at a time against the cache."""
+    ids = list(prompt_ids)
+    cache = make_cache(args, bsz=1,
+                       max_seq_len=max_seq_len or (len(ids) + max_new_tokens + 8))
+
+    logits = model(mx.array([ids]), last_logit_only=True, cache=cache)
+    nxt = int(mx.argmax(logits[0, -1]).item())
+    ids.append(nxt)
+    if verbose:
+        print(f"[gen] prefilled {len(prompt_ids)} tokens", flush=True)
+
+    for _ in range(max_new_tokens - 1):
+        if eos_id is not None and nxt == eos_id:
+            break
+        logits = model(mx.array([[nxt]]), last_logit_only=True, cache=cache)
+        nxt = int(mx.argmax(logits[0, -1]).item())
+        ids.append(nxt)
+    return ids
+
+
+def load_tokenizer(path_or_repo: str = "deepseek-ai/DeepSeek-V3"):
+    """DeepSeek-V4 ships no tokenizer; V3's matches (same padded 129280 vocab,
+    same bos/eos, all V4 special tokens present)."""
+    from transformers import AutoTokenizer
+    return AutoTokenizer.from_pretrained(path_or_repo, trust_remote_code=True)

deepseek_v4_mlx/hyper_connections.py

@@ -0,0 +1,118 @@
+"""Hyper-Connections (HC) — DeepSeek-V4's replacement for the plain residual stream.
+
+Instead of carrying one hidden state, a V4 block carries ``hc_mult`` (=4) parallel
+copies. Each sub-layer:
+
+  * ``hc_pre``  reduces the 4 copies to 1 via learned per-copy weights, and also
+    emits the ``post`` (expand) weights and a ``comb`` (combination) matrix,
+  * the sub-layer (attention or MoE) runs on that single reduced state,
+  * ``hc_post`` writes the result back out to 4 copies, mixing in the previous 4
+    via ``comb``.
+
+All three weight sets come from one projection of the flattened state, split and
+normalized by :func:`split_sinkhorn`. ``comb`` is made doubly-stochastic-ish by
+Sinkhorn iteration (alternating row/column normalization), which is what keeps the
+4 streams from collapsing into each other.
+
+Ported from the reference ``hc_split_sinkhorn_kernel`` in ``inference/kernel.py``
+and ``Block.hc_pre`` / ``Block.hc_post`` in ``inference/model.py``. The kernel is
+pure elementwise + reductions, so this is a faithful transcription rather than an
+approximation.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+
+def split_sinkhorn(mixes: mx.array, hc_scale: mx.array, hc_base: mx.array,
+                   hc_mult: int = 4, sinkhorn_iters: int = 20, eps: float = 1e-6):
+    """Split one projection into (pre, post, comb).
+
+    ``mixes``: [..., (2 + hc) * hc]  -- the projected+normalized state
+    ``hc_scale``: [3]                -- one scale per output group
+    ``hc_base``:  [(2 + hc) * hc]    -- one bias per output element
+
+    Returns ``pre`` [..., hc], ``post`` [..., hc], ``comb`` [..., hc, hc].
+
+    Layout matches the reference kernel exactly: the first ``hc`` entries feed
+    ``pre``, the next ``hc`` feed ``post``, and the remaining ``hc*hc`` form
+    ``comb`` in row-major order.
+    """
+    hc = hc_mult
+    mixes = mixes.astype(mx.float32)
+    scale = hc_scale.astype(mx.float32)
+    base = hc_base.astype(mx.float32)
+
+    # pre: sigmoid(x*s0 + b) + eps          -- strictly positive reduce weights
+    pre = mx.sigmoid(mixes[..., :hc] * scale[0] + base[:hc]) + eps
+    # post: 2*sigmoid(x*s1 + b)             -- expand weights, mean 1 at init
+    post = 2.0 * mx.sigmoid(mixes[..., hc:2 * hc] * scale[1] + base[hc:2 * hc])
+
+    comb = mixes[..., 2 * hc:] * scale[2] + base[2 * hc:]
+    comb = comb.reshape(*comb.shape[:-1], hc, hc)
+
+    # Row softmax, then one column normalization, then (iters-1) full Sinkhorn
+    # sweeps. The eps terms are inside the divisions in the reference kernel;
+    # keeping them there matters because comb rows can get very small.
+    comb = mx.softmax(comb, axis=-1) + eps
+    comb = comb / (comb.sum(axis=-2, keepdims=True) + eps)
+    for _ in range(sinkhorn_iters - 1):
+        comb = comb / (comb.sum(axis=-1, keepdims=True) + eps)
+        comb = comb / (comb.sum(axis=-2, keepdims=True) + eps)
+
+    return pre, post, comb
+
+
+def hc_pre(x: mx.array, hc_fn: mx.array, hc_scale: mx.array, hc_base: mx.array,
+           hc_mult: int, sinkhorn_iters: int, norm_eps: float, hc_eps: float):
+    """Reduce the ``hc`` residual copies to one state.
+
+    ``x``: [b, s, hc, d] -> returns ([b, s, d], post [b,s,hc], comb [b,s,hc,hc]).
+
+    The projection input is the *flattened* [b, s, hc*d] state scaled by its RMS
+    reciprocal — note the reference multiplies by rsqrt AFTER the linear, which is
+    equivalent here only because rsqrt is per-token; it is transcribed in that
+    order to keep the arithmetic bit-for-bit comparable.
+    """
+    shape = x.shape
+    dtype = x.dtype
+    xf = x.reshape(*shape[:2], -1).astype(mx.float32)
+    rsqrt = mx.rsqrt(mx.mean(mx.square(xf), axis=-1, keepdims=True) + norm_eps)
+    mixes = (xf @ hc_fn.astype(mx.float32).T) * rsqrt
+    pre, post, comb = split_sinkhorn(mixes, hc_scale, hc_base, hc_mult, sinkhorn_iters, hc_eps)
+    y = mx.sum(pre[..., None] * xf.reshape(shape), axis=2)
+    return y.astype(dtype), post, comb
+
+
+def hc_post(x: mx.array, residual: mx.array, post: mx.array, comb: mx.array) -> mx.array:
+    """Expand the sub-layer output back to ``hc`` copies and mix in the residual.
+
+    ``x``: [b, s, d]; ``residual``: [b, s, hc, d]; returns [b, s, hc, d].
+
+    Output copy ``k`` is ``post[k]*x + sum_j comb[j, k] * residual[j]`` — residual is
+    indexed by ``j``, the axis being summed, so it must sit on axis 2 of the product
+    (the reference's ``residual.unsqueeze(-2)``). Broadcasting it onto axis 3 instead
+    silently computes ``(sum_j comb[j, k]) * residual[k]``, which has the right shape
+    and stays finite, so only a value comparison catches it.
+    """
+    prod = comb[..., None] * residual[..., :, None, :]     # [b, s, j, k, d]
+    out = post[..., None] * x[..., None, :] + mx.sum(prod, axis=2)
+    return out.astype(x.dtype)
+
+
+def hc_head(x: mx.array, hc_fn: mx.array, hc_scale: mx.array, hc_base: mx.array,
+            norm_eps: float, hc_eps: float) -> mx.array:
+    """Final reduction before the LM head: like :func:`hc_pre` but weights only.
+
+    Uses the raw sigmoid gate without the Sinkhorn split, matching
+    ``Block.hc_head`` in the reference.
+    """
+    shape = x.shape
+    dtype = x.dtype
+    xf = x.reshape(*shape[:2], -1).astype(mx.float32)
+    rsqrt = mx.rsqrt(mx.mean(mx.square(xf), axis=-1, keepdims=True) + norm_eps)
+    mixes = (xf @ hc_fn.astype(mx.float32).T) * rsqrt
+    pre = mx.sigmoid(mixes * hc_scale + hc_base) + hc_eps
+    y = mx.sum(pre[..., None] * xf.reshape(shape), axis=2)
+    return y.astype(dtype)

deepseek_v4_mlx/indexer.py

@@ -0,0 +1,113 @@
+"""The sparse-attention indexer — a second attention that decides what the first reads.
+
+Each ratio-4 layer runs this alongside its sliding window. It keeps its own
+:class:`Compressor` (the Hadamard-rotated, FP4 variant), scores every compressed
+span against the current queries, and returns the indices of the top
+``index_topk`` spans. Attention then reads its local window **plus exactly those**.
+
+Scoring is deliberately cheap and low-precision: queries are Hadamard-rotated and
+FP4-simulated, scores are ReLU'd before a learned per-head weighting collapses the
+head axis. Positions the query may not see are marked ``-1`` rather than dropped,
+so the returned index block stays rectangular for the gather downstream.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .compressor import Compressor
+from .config import ModelArgs
+from .fakequant import fake_quant_fp4
+from .layers import apply_rotary_emb, hadamard_transform
+
+NEG_INF = -1e30
+
+
+class Indexer(nn.Module):
+    def __init__(self, args: ModelArgs, compress_ratio: int = 4):
+        super().__init__()
+        self.n_heads = args.index_n_heads
+        self.head_dim = args.index_head_dim
+        self.rope_head_dim = args.rope_head_dim
+        self.index_topk = args.index_topk
+        self.compress_ratio = compress_ratio
+        self.softmax_scale = self.head_dim ** -0.5
+
+        self.wq_b = nn.Linear(args.q_lora_rank, self.n_heads * self.head_dim, bias=False)
+        self.weights_proj = nn.Linear(args.dim, self.n_heads, bias=False)
+        self.compressor = Compressor(args, compress_ratio, self.head_dim, rotate=True)
+
+    def __call__(self, x: mx.array, qr: mx.array, cos: mx.array, sin: mx.array,
+                 offset: int):
+        """Prefill: choose compressed spans for every query position.
+
+        ``x``  [b, s, dim]        — hidden state, feeds the compressor and weighting
+        ``qr`` [b, s, q_lora_rank] — the shared low-rank query, before wq_b
+        ``offset`` — where compressed entries begin in attention's KV cache, so the
+        returned indices address that cache directly.
+
+        Returns ``[b, s, k]`` int32 indices, with ``-1`` marking "not visible".
+        """
+        bsz, seqlen, _ = x.shape
+        ratio, rd = self.compress_ratio, self.rope_head_dim
+
+        q = self.wq_b(qr).reshape(bsz, seqlen, self.n_heads, self.head_dim)
+        rot = apply_rotary_emb(q[..., -rd:], cos[:seqlen], sin[:seqlen])
+        q = mx.concatenate([q[..., :-rd], rot], axis=-1)
+        q = hadamard_transform(q)
+        q = fake_quant_fp4(q, 32)
+
+        kv = self.compressor(x, cos, sin)          # [b, nb, head_dim] or None
+        if kv is None:
+            return None
+        nb = kv.shape[1]
+
+        # per-head score, ReLU'd, then collapsed by a learned per-head weight
+        w = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5)
+        scores = mx.einsum("bshd,btd->bsht", q.astype(mx.float32), kv.astype(mx.float32))
+        scores = mx.maximum(scores, 0.0) * w[..., None].astype(mx.float32)
+        scores = mx.sum(scores, axis=2)            # [b, s, nb]
+
+        # span t is visible to query s only once the span is complete: t < (s+1)//ratio
+        vis = mx.arange(nb)[None, :] < (mx.arange(1, seqlen + 1)[:, None] // ratio)
+        scores = scores + mx.where(vis, 0.0, NEG_INF)[None]
+
+        k = min(self.index_topk, nb)
+        # argpartition yields unsigned indices; go to int32 before the -1 sentinel
+        idx = mx.argpartition(-scores, k - 1, axis=-1)[..., :k].astype(mx.int32)
+
+        # invisible spans may still be picked when fewer than k are visible; mark them
+        picked_vis = mx.take_along_axis(
+            mx.broadcast_to(vis[None], (bsz, seqlen, nb)), idx, axis=-1)
+        return mx.where(picked_vis, idx + offset, mx.array(-1, mx.int32))
+
+    def step(self, x: mx.array, qr: mx.array, cos: mx.array, sin: mx.array,
+             pos: int, cache, offset: int):
+        """Decode one token: extend the compressed index, then pick top-k spans."""
+        rd = self.rope_head_dim
+        bsz = x.shape[0]
+
+        pooled = self.compressor.step(x, pos, cache.idx_comp, cos, sin)
+        if pooled is not None:
+            cache._adopt(pooled.dtype)
+            cache.idx_kv[:, pos // self.compress_ratio] = pooled[:, 0]
+
+        nb = (pos + 1) // self.compress_ratio
+        if nb <= 0:
+            return None
+
+        q = self.wq_b(qr).reshape(bsz, 1, self.n_heads, self.head_dim)
+        rot = apply_rotary_emb(q[..., -rd:], cos[pos:pos + 1], sin[pos:pos + 1])
+        q = mx.concatenate([q[..., :-rd], rot], axis=-1)
+        q = fake_quant_fp4(hadamard_transform(q), 32)
+
+        kv = cache.idx_kv[:, :nb]
+        w = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5)
+        scores = mx.einsum("bshd,btd->bsht", q.astype(mx.float32), kv.astype(mx.float32))
+        scores = mx.maximum(scores, 0.0) * w[..., None].astype(mx.float32)
+        scores = mx.sum(scores, axis=2)                        # [b, 1, nb]
+
+        k = min(self.index_topk, nb)
+        idx = mx.argpartition(-scores, k - 1, axis=-1)[..., :k].astype(mx.int32)
+        return idx + offset

deepseek_v4_mlx/layers.py

@@ -0,0 +1,132 @@
+"""Core DeepSeek-V4 layers: RMSNorm, YaRN rope, and the Hadamard rotation.
+
+Transcribed from ``inference/model.py``. Two details are easy to get wrong and are
+called out where they occur:
+
+* rope is applied to the **last** ``rope_head_dim`` channels only, pairing
+  *adjacent* elements as (real, imag) — not the split-halves convention used by
+  Llama-style models;
+* attention applies the rope **inverse** to its output (``apply_rotary_emb(...,
+  inverse=True)``), which is why the conjugate path exists here at all.
+"""
+
+from __future__ import annotations
+
+import math
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class RMSNorm(nn.Module):
+    """RMSNorm computed in fp32, returning the input dtype (matches reference)."""
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = mx.ones((dim,), dtype=mx.float32)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        dtype = x.dtype
+        xf = x.astype(mx.float32)
+        var = mx.mean(mx.square(xf), axis=-1, keepdims=True)
+        xf = xf * mx.rsqrt(var + self.eps)
+        return (self.weight * xf).astype(dtype)
+
+
+def precompute_freqs_cis(dim: int, seqlen: int, original_seq_len: int, base: float,
+                         factor: float, beta_fast: int, beta_slow: int):
+    """YaRN-scaled rotary frequencies, returned as (cos, sin) of shape [seqlen, dim//2].
+
+    MLX has no complex dtype in the ops we need, so the reference's ``polar``
+    representation is carried as an explicit cos/sin pair.
+
+    ``original_seq_len == 0`` disables YaRN entirely, which the reference uses for
+    the pure sliding-window layers.
+    """
+    def find_correction_dim(num_rotations, dim, base, max_seq_len):
+        return dim * math.log(max_seq_len / (num_rotations * 2 * math.pi)) / (2 * math.log(base))
+
+    def find_correction_range(low_rot, high_rot, dim, base, max_seq_len):
+        low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len))
+        high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len))
+        return max(low, 0), min(high, dim - 1)
+
+    def linear_ramp_factor(mn, mx_, d):
+        if mn == mx_:
+            mx_ += 0.001
+        lin = (mx.arange(d, dtype=mx.float32) - mn) / (mx_ - mn)
+        return mx.clip(lin, 0, 1)
+
+    freqs = 1.0 / (base ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim))
+    if original_seq_len > 0:
+        low, high = find_correction_range(beta_fast, beta_slow, dim, base, original_seq_len)
+        smooth = 1 - linear_ramp_factor(low, high, dim // 2)
+        freqs = freqs / factor * (1 - smooth) + freqs * smooth
+
+    t = mx.arange(seqlen, dtype=mx.float32)
+    ang = t[:, None] * freqs[None, :]          # [seqlen, dim//2]
+    return mx.cos(ang), mx.sin(ang)
+
+
+def apply_rotary_emb(x: mx.array, cos: mx.array, sin: mx.array, inverse: bool = False) -> mx.array:
+    """Rotate ``x`` [..., s, (h,) d] using adjacent-pair complex multiplication.
+
+    ``cos``/``sin`` are [s, d//2]. ``inverse=True`` conjugates (de-rotates), which
+    attention uses on its output.
+    """
+    dtype = x.dtype
+    xf = x.astype(mx.float32)
+    shape = xf.shape
+    xf = xf.reshape(*shape[:-1], shape[-1] // 2, 2)
+    xr, xi = xf[..., 0], xf[..., 1]
+
+    # broadcast [s, d//2] against [b, s, d//2] or [b, s, h, d//2]
+    if xr.ndim == 4:
+        c = cos[None, :, None, :]
+        s = sin[None, :, None, :]
+    else:
+        c = cos[None, :, :]
+        s = sin[None, :, :]
+    if inverse:
+        s = -s
+
+    out = mx.stack([xr * c - xi * s, xr * s + xi * c], axis=-1)
+    return out.reshape(shape).astype(dtype)
+
+
+def hadamard_transform(x: mx.array, scale: float | None = None) -> mx.array:
+    """Fast Walsh-Hadamard transform along the last axis (power-of-two length).
+
+    Stands in for ``fast_hadamard_transform.hadamard_transform``, which the indexer
+    uses to spread information across channels before FP4 quantization. The
+    reference always calls it with ``scale = d ** -0.5``; that is the default here.
+    """
+    d = x.shape[-1]
+    if d & (d - 1):
+        raise ValueError(f"hadamard_transform needs a power-of-two last dim, got {d}")
+    if scale is None:
+        scale = d ** -0.5
+
+    dtype = x.dtype
+    y = x.astype(mx.float32)
+    lead = y.shape[:-1]
+    h = 1
+    while h < d:
+        # pair up elements h apart and apply the 2x2 butterfly
+        y = y.reshape(*lead, d // (2 * h), 2, h)
+        a, b = y[..., 0, :], y[..., 1, :]
+        y = mx.stack([a + b, a - b], axis=-2)
+        y = y.reshape(*lead, d)
+        h *= 2
+    return (y * scale).astype(dtype)
+
+
+def swiglu(gate: mx.array, up: mx.array, limit: float = 0.0) -> mx.array:
+    """SwiGLU with V4's asymmetric clamp: ``up`` two-sided, ``gate`` upper-only."""
+    g = gate.astype(mx.float32)
+    u = up.astype(mx.float32)
+    if limit > 0:
+        u = mx.clip(u, -limit, limit)
+        g = mx.minimum(g, limit)
+    return (g * mx.sigmoid(g)) * u

deepseek_v4_mlx/load.py

@@ -0,0 +1,68 @@
+"""Load a converted DeepSeek-V4 MLX build."""
+
+from __future__ import annotations
+
+import glob
+import json
+import os
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten
+
+from .config import ModelArgs
+from .convert import bits_for, is_quant_target
+from .model import Model
+
+
+def quant_predicate(group_size: int, bits: int, expert_bits: int | None):
+    """Quantize exactly what the converter did, at the same per-tensor bit width.
+
+    Returning a dict rather than True lets a mixed build load: MLX applies the
+    per-module ``{"group_size", "bits"}`` it returns. Answering plain ``True`` would
+    quantize every module at the top-level ``bits`` and silently mis-read the
+    experts, which are stored at ``expert_bits``.
+    """
+    def pred(path, module):
+        if not hasattr(module, "to_quantized"):
+            return False
+        w = getattr(module, "weight", None)
+        if w is None:
+            return False
+        name = path + ".weight"
+        if not is_quant_target(name, w.shape[-1], group_size):
+            return False
+        return {"group_size": group_size, "bits": bits_for(name, bits, expert_bits)}
+    return pred
+
+
+def load(path: str, lazy: bool = False):
+    cfg = json.load(open(os.path.join(path, "config.json")))
+    args = ModelArgs.from_dict(cfg)
+    model = Model(args)
+
+    q = cfg.get("quantization")
+    if q:
+        nn.quantize(model, group_size=q["group_size"], bits=q["bits"],
+                    class_predicate=quant_predicate(q["group_size"], q["bits"],
+                                                    q.get("expert_bits")))
+
+    loaded = set()
+    for shard in sorted(glob.glob(os.path.join(path, "*.safetensors"))):
+        w = mx.load(shard)
+        model.load_weights(list(w.items()), strict=False)
+        if not lazy:
+            # materialize this shard now (bounded graph) rather than building one
+            # enormous eval over every parameter at the end
+            mx.eval(list(w.values()))
+        loaded.update(w.keys())
+        del w
+
+    expected = {k for k, _ in tree_flatten(model.parameters())}
+    missing = expected - loaded
+    if missing:
+        raise ValueError(f"{len(missing)} params missing from checkpoint, "
+                         f"e.g. {sorted(missing)[:3]}")
+
+    model.eval()
+    return model, args

deepseek_v4_mlx/model.py

@@ -0,0 +1,102 @@
+"""DeepSeek-V4 block and full model.
+
+The block is where Hyper-Connections meet the sub-layers: instead of
+``x = x + attn(norm(x))`` it is
+
+    residual = x                      # [b, s, hc, d]
+    x, post, comb = hc_pre(x, ...)    # 4 copies -> 1
+    x = attn(attn_norm(x))
+    x = hc_post(x, residual, post, comb)   # 1 -> 4, mixing the previous 4
+
+and the same again around the MoE. The MoE additionally takes ``input_ids``,
+because hash-routed layers look their experts up by token id.
+
+The model expands the embedding to ``hc_mult`` identical copies on the way in and
+collapses them with a single gate (``hc_head``) on the way out.
+
+DSpark and the MTP head are deliberately absent: they hang off the reference's
+``forward_spec`` path and are not used for ordinary generation.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .attention import Attention
+from .config import ModelArgs
+from .hyper_connections import hc_head, hc_post, hc_pre
+from .layers import RMSNorm
+from .moe import MoE
+
+
+class Block(nn.Module):
+    def __init__(self, layer_id: int, args: ModelArgs):
+        super().__init__()
+        self.layer_id = layer_id
+        self.norm_eps = args.norm_eps
+        self.hc_mult = args.hc_mult
+        self.hc_iters = args.hc_sinkhorn_iters
+        self.hc_eps = args.hc_eps
+
+        self.attn = Attention(layer_id, args)
+        self.ffn = MoE(layer_id, args)
+        self.attn_norm = RMSNorm(args.dim, args.norm_eps)
+        self.ffn_norm = RMSNorm(args.dim, args.norm_eps)
+
+        mix_hc = (2 + args.hc_mult) * args.hc_mult
+        hc_dim = args.hc_mult * args.dim
+        self.hc_attn_fn = mx.zeros((mix_hc, hc_dim), dtype=mx.float32)
+        self.hc_ffn_fn = mx.zeros((mix_hc, hc_dim), dtype=mx.float32)
+        self.hc_attn_base = mx.zeros((mix_hc,), dtype=mx.float32)
+        self.hc_ffn_base = mx.zeros((mix_hc,), dtype=mx.float32)
+        self.hc_attn_scale = mx.zeros((3,), dtype=mx.float32)
+        self.hc_ffn_scale = mx.zeros((3,), dtype=mx.float32)
+
+    def __call__(self, x: mx.array, input_ids: mx.array | None = None,
+                 cache=None) -> mx.array:
+        residual = x
+        h, post, comb = hc_pre(x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base,
+                               self.hc_mult, self.hc_iters, self.norm_eps, self.hc_eps)
+        h = self.attn(self.attn_norm(h), cache)
+        x = hc_post(h, residual, post, comb)
+
+        residual = x
+        h, post, comb = hc_pre(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base,
+                               self.hc_mult, self.hc_iters, self.norm_eps, self.hc_eps)
+        h = self.ffn(self.ffn_norm(h), input_ids)
+        return hc_post(h, residual, post, comb)
+
+
+class Model(nn.Module):
+    """embed -> expand to hc copies -> blocks -> collapse -> norm -> logits."""
+
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.hc_mult = args.hc_mult
+        self.embed = nn.Embedding(args.vocab_size, args.dim)
+        self.layers = [Block(i, args) for i in range(args.n_layers)]
+        self.norm = RMSNorm(args.dim, args.norm_eps)
+        self.head = nn.Linear(args.dim, args.vocab_size, bias=False)
+
+        self.hc_head_fn = mx.zeros((args.hc_mult, args.hc_mult * args.dim), dtype=mx.float32)
+        self.hc_head_base = mx.zeros((args.hc_mult,), dtype=mx.float32)
+        self.hc_head_scale = mx.zeros((1,), dtype=mx.float32)
+
+    def __call__(self, input_ids: mx.array, last_logit_only: bool = False,
+                 cache=None) -> mx.array:
+        h = self.embed(input_ids)
+        # the four Hyper-Connection streams start as identical copies
+        h = mx.broadcast_to(h[:, :, None, :], (*h.shape[:2], self.hc_mult, h.shape[-1]))
+
+        for i, layer in enumerate(self.layers):
+            h = layer(h, input_ids, cache[i] if cache is not None else None)
+
+        h = hc_head(h, self.hc_head_fn, self.hc_head_scale, self.hc_head_base,
+                    self.args.norm_eps, self.args.hc_eps)
+        h = self.norm(h)
+        if last_logit_only:
+            h = h[:, -1:]
+        # the reference keeps the LM head in fp32
+        return self.head(h.astype(mx.float32))

deepseek_v4_mlx/moe.py

@@ -0,0 +1,140 @@
+"""DeepSeek-V4 MoE: hash-or-score routing, clamped-SwiGLU experts, one shared expert.
+
+Two things differ from every other MoE in ``mlx-lm``:
+
+* **The first ``n_hash_layers`` layers do not score experts at all.** They read the
+  six expert indices straight out of a ``tid2eid[vocab_size, n_activated]`` table
+  indexed by *token id*, so routing there is independent of context.
+* **The selection bias does not reach the routing weights.** Scores are captured
+  *before* the bias is added; the bias only reorders the top-k. Applying it to the
+  weights too is the obvious mistake and would be invisible on a smoke test.
+
+Routing weights are applied after ``down_proj`` rather than before it as in the
+reference. ``down_proj`` is linear and the weight is one scalar per token, so
+``w2(w * h) == w * w2(h)``; doing it afterwards is what lets the batched
+:class:`SwitchGLU` gather-matmul be used instead of a 256-iteration Python loop.
+
+That rewrite sums each token's expert outputs, whereas the reference scatters them.
+The two agree as long as a token's expert indices are distinct — which holds for
+this checkpoint: every sampled ``tid2eid`` row contains 6 distinct experts, and
+scored routing takes a top-k, which cannot repeat.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx_lm.models.switch_layers import SwitchGLU
+
+from .config import ModelArgs
+
+
+class ClampedSwiGLU(nn.Module):
+    """SwiGLU with V4's asymmetric clamp: ``up`` two-sided, ``gate`` upper-only.
+
+    Called as ``activation(x_up, x_gate)`` to match ``SwitchGLU``.
+    """
+
+    def __init__(self, limit: float = 0.0):
+        super().__init__()
+        self.limit = limit
+
+    def __call__(self, x, gate):
+        if self.limit > 0:
+            x = mx.clip(x, -self.limit, self.limit)
+            gate = mx.minimum(gate, self.limit)
+        return nn.silu(gate) * x
+
+
+class Gate(nn.Module):
+    """Routes each token to ``n_activated_experts`` experts.
+
+    ``hash=True`` layers use the ``tid2eid`` lookup and carry no bias; scored
+    layers carry a per-expert selection bias.
+    """
+
+    def __init__(self, layer_id: int, args: ModelArgs):
+        super().__init__()
+        self.topk = args.n_activated_experts
+        self.score_func = args.score_func
+        self.route_scale = args.route_scale
+        self.hash = args.is_hash_layer(layer_id)
+        n_exp = args.n_experts_for(layer_id)
+        self.weight = mx.zeros((n_exp, args.dim))
+        if self.hash:
+            self.tid2eid = mx.zeros((args.vocab_size, args.n_activated_experts), dtype=mx.int32)
+        else:
+            self.bias = mx.zeros((n_exp,), dtype=mx.float32)
+
+    def __call__(self, x: mx.array, input_ids: mx.array | None = None):
+        scores = x.astype(mx.float32) @ self.weight.astype(mx.float32).T
+        if self.score_func == "softmax":
+            scores = mx.softmax(scores, axis=-1)
+        elif self.score_func == "sigmoid":
+            scores = mx.sigmoid(scores)
+        else:  # sqrtsoftplus
+            scores = mx.sqrt(nn.softplus(scores))
+
+        original_scores = scores
+        if not self.hash:
+            # bias shifts *selection* only — weights are read from original_scores
+            scores = scores + self.bias
+            indices = mx.argpartition(-scores, self.topk - 1, axis=-1)[..., :self.topk]
+            # argpartition does not order within the partition; sort so the expert
+            # order matches topk() and the gathered weights line up.
+            part = mx.take_along_axis(scores, indices, axis=-1)
+            order = mx.argsort(-part, axis=-1)
+            indices = mx.take_along_axis(indices, order, axis=-1)
+        else:
+            indices = self.tid2eid[input_ids]
+
+        weights = mx.take_along_axis(original_scores, indices, axis=-1)
+        if self.score_func != "softmax":
+            weights = weights / mx.sum(weights, axis=-1, keepdims=True)
+        weights = weights * self.route_scale
+        return weights, indices
+
+
+class MoE(nn.Module):
+    """256 routed experts (top-6) plus 1 always-on shared expert."""
+
+    def __init__(self, layer_id: int, args: ModelArgs):
+        super().__init__()
+        self.dim = args.dim
+        self.gate = Gate(layer_id, args)
+        act = ClampedSwiGLU(args.swiglu_limit)
+        self.experts = SwitchGLU(args.dim, args.moe_inter_dim,
+                                 args.n_experts_for(layer_id), activation=act)
+        self.shared_experts = SharedExpert(args.dim, args.moe_inter_dim, args.swiglu_limit)
+
+    def __call__(self, x: mx.array, input_ids: mx.array | None = None) -> mx.array:
+        shape = x.shape
+        xf = x.reshape(-1, self.dim)
+        ids = input_ids.reshape(-1) if input_ids is not None else None
+
+        weights, indices = self.gate(xf, ids)
+        y = self.experts(xf, indices)                       # [tokens, topk, dim]
+        y = (y * weights[..., None].astype(y.dtype)).sum(axis=-2)
+        y = y + self.shared_experts(xf)
+        return y.reshape(shape).astype(x.dtype)
+
+
+class SharedExpert(nn.Module):
+    """The always-on expert. Same clamped SwiGLU, no routing weight."""
+
+    def __init__(self, dim: int, inter_dim: int, limit: float = 0.0):
+        super().__init__()
+        self.w1 = nn.Linear(dim, inter_dim, bias=False)
+        self.w2 = nn.Linear(inter_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, inter_dim, bias=False)
+        self.limit = limit
+
+    def __call__(self, x: mx.array) -> mx.array:
+        dtype = x.dtype
+        gate = self.w1(x).astype(mx.float32)
+        up = self.w3(x).astype(mx.float32)
+        if self.limit > 0:
+            up = mx.clip(up, -self.limit, self.limit)
+            gate = mx.minimum(gate, self.limit)
+        h = nn.silu(gate) * up
+        return self.w2(h.astype(dtype))

deepseek_v4_mlx/sparse_attention.py

@@ -0,0 +1,71 @@
+"""Gather-based sparse attention — the MLX stand-in for the TileLang kernel.
+
+Each query position attends to an explicit list of KV indices rather than a
+contiguous range. Semantics transcribed from ``sparse_attn_kernel`` in
+``inference/kernel.py``:
+
+* ``kv`` is a **single** shared vector per position — MLA, so key and value are
+  the same tensor, not two projections;
+* an index of ``-1`` means "not visible": its logit is ``-inf`` and it contributes
+  nothing;
+* the learned ``attn_sink`` enters the softmax **denominator only**. It carries no
+  value vector, so a head whose real logits are all small attends to essentially
+  nothing rather than being forced to distribute weight over whatever is present.
+
+The reference computes this with an online (FlashAttention-style) softmax for
+memory reasons. A plain max-subtracted softmax is numerically equivalent; queries
+are chunked here so the gathered ``[b, chunk, k, d]`` block stays bounded.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+NEG_INF = -1e30
+
+
+def _gather_kv(kv: mx.array, idx: mx.array) -> mx.array:
+    """``kv`` [b, n, d], ``idx`` [b, m, k] -> [b, m, k, d], row 0 for negative idx."""
+    b, n, d = kv.shape
+    flat = kv.reshape(b * n, d)
+    safe = mx.maximum(idx, 0).astype(mx.int32)
+    base = (mx.arange(b, dtype=mx.int32) * n).reshape(b, 1, 1)
+    return flat[(safe + base).reshape(-1)].reshape(*idx.shape, d)
+
+
+def sparse_attn(q: mx.array, kv: mx.array, attn_sink: mx.array, topk_idxs: mx.array,
+                softmax_scale: float, chunk: int = 256) -> mx.array:
+    """Attend each query to its own list of KV positions.
+
+    ``q``          [b, m, h, d]
+    ``kv``         [b, n, d]      shared key/value (MLA)
+    ``attn_sink``  [h]            per-head null-attention logit
+    ``topk_idxs``  [b, m, k]      int32, -1 = masked
+    returns        [b, m, h, d]
+    """
+    b, m, h, d = q.shape
+    sink = attn_sink.astype(mx.float32).reshape(1, 1, h, 1)
+
+    outs = []
+    for start in range(0, m, chunk):
+        stop = min(start + chunk, m)
+        qc = q[:, start:stop].astype(mx.float32)          # [b, c, h, d]
+        ic = topk_idxs[:, start:stop]                     # [b, c, k]
+        kvc = _gather_kv(kv, ic).astype(mx.float32)       # [b, c, k, d]
+
+        # [b, c, h, k]
+        logits = mx.einsum("bchd,bckd->bchk", qc, kvc) * softmax_scale
+        valid = (ic >= 0)[:, :, None, :]
+        logits = mx.where(valid, logits, NEG_INF)
+
+        mmax = mx.max(logits, axis=-1, keepdims=True)             # [b, c, h, 1]
+        # the sink competes with the real logits for the running max
+        mmax = mx.maximum(mmax, sink)
+        w = mx.exp(logits - mmax)
+        w = mx.where(valid, w, 0.0)
+        denom = mx.sum(w, axis=-1, keepdims=True) + mx.exp(sink - mmax)
+
+        o = mx.einsum("bchk,bckd->bchd", w, kvc) / denom
+        outs.append(o.astype(q.dtype))
+
+    return mx.concatenate(outs, axis=1) if len(outs) > 1 else outs[0]

generation_config.json

@@ -0,0 +1,9 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 1,
+  "do_sample": true,
+  "temperature": 1.0,
+  "top_p": 1.0,
+  "transformers_version": "4.46.3"
+}

model-00001-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:69c564c2ec6e4a688d7460cdc87bf5e9268b8a1c8a95807c70b9d5d049669a2c
+size 5436806040

model-00002-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d56f47f0d43b87b265347a2bc6224c1f8ee42ea26f07d168819f09e7a985c84
+size 5355630769

model-00003-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8ce3abddaee4ff1e739746b4fb9f07e8e2830072d979c5d444b5fc063a30cf84
+size 5032677715

model-00004-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e17646e2ebb2f7d3076b926c68b4d826241b0232780f160e1307a574ef46efbd
+size 5020874379

model-00005-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80cee940cc1b8275796f1b7e5b21b1e297c22d75e54c007fa67b8fdcbef8bd7d
+size 5643415454

model-00006-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c7047e1a3b400b438181ec89cf58a34817c02aaa33ce7732b6db00454b612324
+size 5023215191

model-00007-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b6189b80d7742493395872cf2c895222849eff90950a3b2631b5222103fe83f0
+size 5699008910

model-00008-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1026a8ecd7dcf739660bf91a8db9c4461ba4aa61e6de4360367d27630e00e5f4
+size 5004876411

model-00009-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:81b8e752cd30250358cbbe42268c8eff90ff5aabb57a532a5e008155dc056834
+size 5417984314

model-00010-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4e56fd8c4fe18036336a8a57252d134e3da67b94680b55f06def4e249f84823c
+size 5000256653

model-00011-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:983caf5380dc6f369338f7cb47166cf91c021e1009e397ad5f85ca2c32458e01
+size 5029122176

model-00012-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d406af482af4b3a2d06065fabb6cba2dd0282c4220679bff04488c16ccfc0b9c
+size 5603569790

model-00013-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:abb96ea31791fc15440c101a5c8f2cf36ce1a188060ffabb047fa024426f4988
+size 5053432424

model-00014-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a6f8679d4085160fb8bba4b7b812988ad7542fd97682c264ea05940f6c088bb
+size 5063649550

model-00015-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:02e03eba5e371337adf5f6bc011feea5f7629a86ae7b4b1cad8ee523cfe37467
+size 5578617059

model-00016-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b79881e04d45a988e7f9c0ae091552f473bc35441ea8c58dfcac491742aba2e
+size 5742264789

model-00017-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca3b4e2a6334c10c695425c820caac128cdb96e1cca917412d9943158ab6f7fb
+size 5246779122

model-00018-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2299145c9c09e8199c3aea5a7e4ad3487f8097b727c7224c7df887daee23bc5f
+size 5226205499

model-00019-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3e37047aa189522115706bb05a838a11c687dd99e36c3ea0d2ca18abf366a7ff
+size 5104788615

model-00020-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e95a73c4de94ec15729b64e120df35ec6cada6407cb5342524f6a128905313c7
+size 5686325586

model-00021-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c9178c606f1fe02b564f959659e73a33da75f85c2881b8f5c9c7ecf5424ed3b
+size 5732777176

model-00022-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:33490aa69fcd26c407c5dc438d1f46809f46a1a155ee99013e7e14dbe866b398
+size 5003800685

model-00023-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:199606bb4212972c64a25fe10f5ad1c3e25f74a5061b34470003fa39aa101995
+size 5008867619

model-00024-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd36364122fe7ee8e945225f9af580caaa94ef11c55ddf0ef2eea114b3000cf6
+size 5692721355

model-00025-of-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3fb1cc7fe173e967e1d4bc191b82fb32da1c9574f98b3a5179e26803184e8fca
+size 1579230717

tokenizer_config.json

@@ -0,0 +1,34 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<｜begin▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "clean_up_tokenization_spaces": false,
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "<｜end▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "legacy": true,
+  "model_max_length": 1048576,
+  "pad_token": {
+    "__type": "AddedToken",
+    "content": "<｜end▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sp_model_kwargs": {},
+  "unk_token": null,
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}