Add files using upload-large-folder tool

README.md

@@ -0,0 +1,103 @@
+---
+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-mixed-4_8bit
+
+MLX (Apple Silicon) conversion of [deepseek-ai/DeepSeek-V4-Flash-0731](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-0731),
+quantized with **4-bit routed experts and 8-bit everything else** (affine group quant, group size 64).
+
+**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 | smallest; +5.2% vs 8-bit |
+
+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.
+
+## ⚠️ 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,116 @@
+{
+  "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
+  }
+}

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]

encoding/README.md

@@ -0,0 +1,174 @@
+# DeepSeek-V4 Encoding
+
+This document describes the prompt encoding format used by DeepSeek-V4 series models. The encoding handles multi-turn conversations, tool calling, extended thinking (reasoning), and quick instruction tasks.
+
+A self-contained reference implementation is provided in `encoding_dsv4.py`.
+
+## Quick Start
+
+```python
+from encoding_dsv4 import encode_messages, parse_message_from_completion_text
+
+# Encode a conversation
+messages = [
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": "What is 2+2?"},
+]
+prompt = encode_messages(messages, thinking_mode="thinking")
+# => "<｜begin▁of▁sentence｜>You are a helpful assistant.<｜User｜>What is 2+2?<｜Assistant｜><think>"
+
+# Parse model output back to structured message
+completion = "Simple arithmetic.</think>2 + 2 = 4.<｜end▁of▁sentence｜>"
+parsed = parse_message_from_completion_text(completion, thinking_mode="thinking")
+# => {"role": "assistant", "reasoning_content": "Simple arithmetic.", "content": "2 + 2 = 4.", "tool_calls": []}
+```
+
+> **Note:** The `parse_message_from_completion_text` function is designed to handle well-formatted model output only. It does not attempt to correct or recover from malformed output that the model might occasionally generate. For production use, additional error handling is recommended.
+
+## Message Format
+
+### Special Tokens
+
+| Token | Purpose |
+|-------|---------|
+| `<｜begin▁of▁sentence｜>` | Beginning of sequence (BOS) |
+| `<｜end▁of▁sentence｜>` | End of assistant turn (EOS) |
+| `<｜User｜>` | User turn prefix |
+| `<｜Assistant｜>` | Assistant turn prefix |
+| `<｜latest_reminder｜>` | Latest reminder (date, locale, etc.) |
+| `<think>` / `</think>` | Reasoning block delimiters |
+| `｜DSML｜` | DSML markup token |
+
+### Roles
+
+The encoding supports the following message roles: `system`, `user`, `assistant`, `tool`, `latest_reminder`, and `developer`.
+
+> **Note on the `developer` role:** The `developer` role is used exclusively in the internal search agent pipeline. It is not needed for general-purpose chat or tool-calling tasks, and the official API does not accept messages with this role.
+
+### Basic Chat
+
+A simple multi-turn conversation is encoded as:
+
+```
+<｜begin▁of▁sentence｜>{system_prompt}
+<｜User｜>{user_message}<｜Assistant｜></think>{response}<｜end▁of▁sentence｜>
+<｜User｜>{user_message_2}<｜Assistant｜></think>{response_2}<｜end▁of▁sentence｜>
+```
+
+- The BOS token is prepended at the very beginning of the conversation.
+- In **chat mode** (`thinking_mode="chat"`), `</think>` is placed right after `<｜Assistant｜>` to immediately close the thinking block, so the model generates content directly.
+
+### Interleaved Thinking Mode
+
+In **thinking mode** (`thinking_mode="thinking"`), the model produces explicit reasoning inside `<think>...</think>` blocks before responding.
+
+```
+<｜begin▁of▁sentence｜>{system_prompt}
+<｜User｜>{message}<｜Assistant｜><think>{reasoning}</think>{response}<｜end▁of▁sentence｜>
+```
+
+The `drop_thinking` parameter (default `True`) controls whether reasoning from earlier turns is preserved:
+
+- **Without tools**: `drop_thinking` takes effect. Reasoning content from assistant turns **before** the last user message is stripped. Only the final assistant turn retains its `<think>...</think>` block.
+- **With tools** (on system or developer message): `drop_thinking` is automatically disabled. All turns retain their reasoning, because tool-calling conversations require full context for the model to track multi-step reasoning across tool calls.
+
+### Tool Calling (DSML Format)
+
+Tools are defined on the `system` or `developer` message via the `tools` field (OpenAI-compatible format). When tools are present, the following schema block is injected into the system/user prompt:
+
+```
+## Tools
+
+You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following:
+
+<｜DSML｜tool_calls>
+<｜DSML｜invoke name="$TOOL_NAME">
+<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</｜DSML｜parameter>
+...
+</｜DSML｜invoke>
+<｜DSML｜invoke name="$TOOL_NAME2">
+...
+</｜DSML｜invoke>
+</｜DSML｜tool_calls>
+
+String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
+
+If thinking_mode is enabled (triggered by <think>), you MUST output your complete reasoning inside <think>...</think> BEFORE any tool calls or final response.
+
+Otherwise, output directly after </think> with tool calls or final response.
+
+### Available Tool Schemas
+
+{tool_definitions_json}
+
+You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
+```
+
+An actual tool call in the assistant turn looks like:
+
+```xml
+<｜DSML｜tool_calls>
+<｜DSML｜invoke name="function_name">
+<｜DSML｜parameter name="param" string="true">string_value</｜DSML｜parameter>
+<｜DSML｜parameter name="count" string="false">5</｜DSML｜parameter>
+</｜DSML｜invoke>
+</｜DSML｜tool_calls><｜end▁of▁sentence｜>
+```
+
+- `string="true"`: the parameter value is a raw string.
+- `string="false"`: the parameter value is JSON (number, boolean, array, object).
+
+Tool execution results are wrapped in `<tool_result>` tags within user messages:
+
+```
+<｜User｜><tool_result>{result_json}</tool_result><｜Assistant｜><think>...
+```
+
+When multiple tool results are present, they are sorted by the order of the corresponding `tool_calls` in the preceding assistant message.
+
+### Reasoning Effort
+
+In thinking mode, the `reasoning_effort` parameter selects one of three levels, which control how much deliberation the model spends before answering. The level is realized purely as a text prefix prepended at the very beginning of the prompt (before the system message); the rest of the encoding is identical across levels.
+
+| `reasoning_effort` | Prompt prefix |
+|:---|:---|
+| `"low"` (default) | none |
+| `"high"` | `Reasoning Effort: Absolute maximum ...` |
+| `"max"` | `Reasoning Effort: Beyond maximum ...` |
+
+`reasoning_effort` has no effect in chat mode (`thinking_mode="chat"`), where the model does not produce a reasoning block at all.
+
+The full prefix text for `"high"`:
+
+```
+Reasoning Effort: Absolute maximum with no shortcuts permitted.
+You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.
+Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.
+```
+
+And for `"max"`:
+
+```
+Reasoning Effort: Beyond maximum — exhaustive, relentless, and uncompromising.
+You MUST reason with the utmost depth and rigor, leaving absolutely nothing to chance: exhaustively decompose the problem into its most fundamental components, trace every causal chain to its root, and resolve the underlying cause rather than any surface symptom.
+Do not stop reasoning until you have independently verified the solution from multiple angles and are certain that no assumption remains unchecked and no error remains undiscovered.
+```
+
+### Quick Instruction Special Tokens
+
+Quick instruction tokens are used for auxiliary classification and generation tasks. They are appended to messages via the `"task"` field to trigger specialized model behavior for a single-token or short-form output.
+
+| Special Token | Description | Format |
+|:---|:---|:---|
+| `<｜action｜>` | Determines whether the user prompt requires a web search or can be answered directly. | `...<｜User｜>{prompt}<｜Assistant｜><think><｜action｜>` |
+| `<｜title｜>` | Generates a concise conversation title after the first assistant response. | `...<｜Assistant｜>{response}<｜end▁of▁sentence｜><｜title｜>` |
+| `<｜query｜>` | Generates search queries for the user prompt. | `...<｜User｜>{prompt}<｜query｜>` |
+| `<｜authority｜>` | Classifies the user prompt's demand for source authoritativeness. | `...<｜User｜>{prompt}<｜authority｜>` |
+| `<｜domain｜>` | Identifies the domain of the user prompt. | `...<｜User｜>{prompt}<｜domain｜>` |
+| `<｜extracted_url｜>` `<｜read_url｜>` | Determines whether each URL in the user prompt should be fetched and read. | `...<｜User｜>{prompt}<｜extracted_url｜>{url}<｜read_url｜>` |
+
+Usage in message format:
+
+- **`action`** on a user message: the `<｜action｜>` token is placed after the assistant prefix and thinking token, triggering a routing decision (e.g., "Search" or "Answer").
+- **Other tasks** (`query`, `authority`, `domain`, `read_url`) on a user message: the task token is appended directly after the user content.
+- **`title`** on an assistant message: the `<｜title｜>` token is appended after the assistant's EOS. The next assistant message provides the generated title.

encoding/encoding_dsv4.py

@@ -0,0 +1,760 @@
+"""
+DeepSeek-V4 Encoding
+
+A self-contained implementation for encoding/decoding DeepSeek-V4 chat messages
+with tool calling, thinking mode, and quick instruction task support.
+"""
+
+from typing import Any, Dict, List, Union, Optional, Tuple
+import copy
+import json
+import re
+
+# ============================================================
+# Special Tokens
+# ============================================================
+
+bos_token: str = "<｜begin▁of▁sentence｜>"
+eos_token: str = "<｜end▁of▁sentence｜>"
+thinking_start_token: str = "<think>"
+thinking_end_token: str = "</think>"
+dsml_token: str = "｜DSML｜"
+
+USER_SP_TOKEN = "<｜User｜>"
+ASSISTANT_SP_TOKEN = "<｜Assistant｜>"
+LATEST_REMINDER_SP_TOKEN = "<｜latest_reminder｜>"
+
+# Task special tokens for internal classification tasks
+DS_TASK_SP_TOKENS = {
+    "action": "<｜action｜>",
+    "query": "<｜query｜>",
+    "authority": "<｜authority｜>",
+    "domain": "<｜domain｜>",
+    "title": "<｜title｜>",
+    "read_url": "<｜read_url｜>",
+}
+VALID_TASKS = set(DS_TASK_SP_TOKENS.keys())
+
+# ============================================================
+# Templates
+# ============================================================
+
+system_msg_template: str = "{content}"
+user_msg_template: str = "{content}"
+latest_reminder_msg_template: str = "{content}"
+assistant_msg_template: str = "{reasoning}{content}{tool_calls}" + eos_token
+assistant_msg_wo_eos_template: str = "{reasoning}{content}{tool_calls}"
+thinking_template: str = "{reasoning_content}"
+
+response_format_template: str = (
+    "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}"
+)
+tool_call_template: str = (
+    "<{dsml_token}invoke name=\"{name}\">\n{arguments}\n</{dsml_token}invoke>"
+)
+tool_calls_template = (
+    "<{dsml_token}{tc_block_name}>\n{tool_calls}\n</{dsml_token}{tc_block_name}>"
+)
+tool_calls_block_name: str = "tool_calls"
+
+tool_output_template: str = (
+    "<tool_result>{content}</tool_result>"
+)
+
+# Reasoning effort levels. In thinking mode, the prompt for the selected level is
+# prepended at the very beginning of the conversation. `low` is the default and
+# adds nothing.
+REASONING_EFFORT_PROMPTS: Dict[str, str] = {
+    "low": "",
+    "high": (
+        "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n"
+        "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n"
+        "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n"
+    ),
+    "max": (
+        "Reasoning Effort: Beyond maximum — exhaustive, relentless, and uncompromising.\n"
+        "You MUST reason with the utmost depth and rigor, leaving absolutely nothing to chance: exhaustively decompose the problem into its most fundamental components, trace every causal chain to its root, and resolve the underlying cause rather than any surface symptom.\n"
+        "Do not stop reasoning until you have independently verified the solution from multiple angles and are certain that no assumption remains unchecked and no error remains undiscovered.\n\n"
+    ),
+}
+DEFAULT_REASONING_EFFORT = "low"
+
+TOOLS_TEMPLATE = """## Tools
+
+You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following:
+
+<{dsml_token}tool_calls>
+<{dsml_token}invoke name="$TOOL_NAME">
+<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}parameter>
+...
+</{dsml_token}invoke>
+<{dsml_token}invoke name="$TOOL_NAME2">
+...
+</{dsml_token}invoke>
+</{dsml_token}tool_calls>
+
+String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
+
+If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response.
+
+Otherwise, output directly after {thinking_end_token} with tool calls or final response.
+
+### Available Tool Schemas
+
+{tool_schemas}
+
+You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
+"""
+
+# ============================================================
+# Utility Functions
+# ============================================================
+
+def to_json(value: Any) -> str:
+    """Serialize a value to JSON string."""
+    try:
+        return json.dumps(value, ensure_ascii=False)
+    except:
+        return json.dumps(value, ensure_ascii=True)
+
+
+def tools_from_openai_format(tools):
+    """Extract function definitions from OpenAI-format tool list."""
+    return [tool["function"] for tool in tools]
+
+
+def tool_calls_from_openai_format(tool_calls):
+    """Convert OpenAI-format tool calls to internal format."""
+    return [
+        {
+            "name": tool_call["function"]["name"],
+            "arguments": tool_call["function"]["arguments"],
+        }
+        for tool_call in tool_calls
+    ]
+
+
+def tool_calls_to_openai_format(tool_calls):
+    """Convert internal tool calls to OpenAI format."""
+    return [
+        {
+            "type": "function",
+            "function": {
+                "name": tool_call["name"],
+                "arguments": tool_call["arguments"],
+            }
+        }
+        for tool_call in tool_calls
+    ]
+
+
+def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str:
+    """
+    Encode tool call arguments into DSML parameter format.
+
+    Args:
+        tool_call: Dict with "name" and "arguments" (JSON string) keys.
+
+    Returns:
+        DSML-formatted parameter string.
+    """
+    p_dsml_template = '<{dsml_token}parameter name="{key}" string="{is_str}">{value}</{dsml_token}parameter>'
+    P_dsml_strs = []
+
+    try:
+        arguments = json.loads(tool_call["arguments"])
+    except Exception as err:
+        arguments = {"arguments": tool_call["arguments"]}
+
+    for k, v in arguments.items():
+        p_dsml_str = p_dsml_template.format(
+            dsml_token=dsml_token,
+            key=k,
+            is_str="true" if isinstance(v, str) else "false",
+            value=v if isinstance(v, str) else to_json(v),
+        )
+        P_dsml_strs.append(p_dsml_str)
+
+    return "\n".join(P_dsml_strs)
+
+
+def decode_dsml_to_arguments(tool_name: str, tool_args: Dict[str, Tuple[str, str]]) -> Dict[str, str]:
+    """
+    Decode DSML parameters back to a tool call dict.
+
+    Args:
+        tool_name: Name of the tool.
+        tool_args: Dict mapping param_name -> (value, is_string_flag).
+
+    Returns:
+        Dict with "name" and "arguments" (JSON string) keys.
+    """
+    def _decode_value(key: str, value: str, string: str):
+        if string == "true":
+            value = to_json(value)
+        return f"{to_json(key)}: {value}"
+
+    tool_args_json = "{" + ", ".join([_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]) + "}"
+    return dict(name=tool_name, arguments=tool_args_json)
+
+
+def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str:
+    """
+    Render tool schemas into the system prompt format.
+
+    Args:
+        tools: List of tool schema dicts (each with name, description, parameters).
+
+    Returns:
+        Formatted tools section string.
+    """
+    tools_json = [to_json(t) for t in tools]
+
+    return TOOLS_TEMPLATE.format(
+        tool_schemas="\n".join(tools_json),
+        dsml_token=dsml_token,
+        thinking_start_token=thinking_start_token,
+        thinking_end_token=thinking_end_token,
+    )
+
+
+def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
+    """Find the index of the last user/developer message."""
+    last_user_index = -1
+    for idx in range(len(messages) - 1, -1, -1):
+        if messages[idx].get("role") in ["user", "developer"]:
+            last_user_index = idx
+            break
+    return last_user_index
+
+
+# ============================================================
+# Message Rendering
+# ============================================================
+
+def render_message(index: int, messages: List[Dict[str, Any]], thinking_mode: str, drop_thinking: bool = True, reasoning_effort: Optional[str] = None) -> str:
+    """
+    Render a single message at the given index into its encoded string form.
+
+    This is the core function that converts each message in the conversation
+    into the DeepSeek-V4 format.
+
+    Args:
+        index: Index of the message to render.
+        messages: Full list of messages in the conversation.
+        thinking_mode: Either "chat" or "thinking".
+        drop_thinking: Whether to drop reasoning content from earlier turns.
+        reasoning_effort: Reasoning effort level, one of "low", "high", "max".
+            None is treated as "low".
+
+    Returns:
+        Encoded string for this message.
+    """
+    assert 0 <= index < len(messages)
+    assert thinking_mode in ["chat", "thinking"], f"Invalid thinking_mode `{thinking_mode}`"
+
+    prompt = ""
+    msg = messages[index]
+    last_user_idx = find_last_user_index(messages)
+
+    role = msg.get("role")
+    content = msg.get("content")
+    tools = msg.get("tools")
+    response_format = msg.get("response_format")
+    tool_calls = msg.get("tool_calls")
+    reasoning_content = msg.get("reasoning_content")
+    wo_eos = msg.get("wo_eos", False)
+
+    if tools:
+        tools = tools_from_openai_format(tools)
+    if tool_calls:
+        tool_calls = tool_calls_from_openai_format(tool_calls)
+
+    # Reasoning effort prefix (only at index 0 in thinking mode; "low" adds nothing)
+    reasoning_effort = reasoning_effort or DEFAULT_REASONING_EFFORT
+    assert reasoning_effort in REASONING_EFFORT_PROMPTS, \
+        f"Invalid reasoning effort: {reasoning_effort}, expected one of {list(REASONING_EFFORT_PROMPTS)}"
+    if index == 0 and thinking_mode == "thinking":
+        prompt += REASONING_EFFORT_PROMPTS[reasoning_effort]
+
+    if role == "system":
+        prompt += system_msg_template.format(content=content or "")
+        if tools:
+            prompt += "\n\n" + render_tools(tools)
+        if response_format:
+            prompt += "\n\n" + response_format_template.format(schema=to_json(response_format))
+
+    elif role == "developer":
+        assert content, f"Invalid message for role `{role}`: {msg}"
+
+        content_developer = USER_SP_TOKEN
+        content_developer += content
+
+        if tools:
+            content_developer += "\n\n" + render_tools(tools)
+        if response_format:
+            content_developer += "\n\n" + response_format_template.format(schema=to_json(response_format))
+
+        prompt += user_msg_template.format(content=content_developer)
+
+    elif role == "user":
+        prompt += USER_SP_TOKEN
+
+        # Handle content blocks (tool results mixed with text)
+        content_blocks = msg.get("content_blocks")
+        if content_blocks:
+            parts = []
+            for block in content_blocks:
+                block_type = block.get("type")
+                if block_type == "text":
+                    parts.append(block.get("text", ""))
+                elif block_type == "tool_result":
+                    tool_content = block.get("content", "")
+                    if isinstance(tool_content, list):
+                        text_parts = []
+                        for b in tool_content:
+                            if b.get("type") == "text":
+                                text_parts.append(b.get("text", ""))
+                            else:
+                                text_parts.append(f"[Unsupported {b.get('type')}]")
+                        tool_content = "\n\n".join(text_parts)
+                    parts.append(tool_output_template.format(content=tool_content))
+                else:
+                    parts.append(f"[Unsupported {block_type}]")
+            prompt += "\n\n".join(parts)
+        else:
+            prompt += content or ""
+
+    elif role == "latest_reminder":
+        prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format(content=content)
+
+    elif role == "tool":
+        raise NotImplementedError("deepseek_v4 merges tool messages into user; please preprocess with merge_tool_messages()")
+
+    elif role == "assistant":
+        thinking_part = ""
+        tc_content = ""
+
+        if tool_calls:
+            tc_list = [
+                tool_call_template.format(
+                    dsml_token=dsml_token,
+                    name=tc.get("name"),
+                    arguments=encode_arguments_to_dsml(tc)
+                )
+                for tc in tool_calls
+            ]
+            tc_content += '\n\n' + tool_calls_template.format(
+                dsml_token=dsml_token,
+                tool_calls="\n".join(tc_list),
+                tc_block_name=tool_calls_block_name,
+            )
+
+        summary_content = content or ""
+        rc = reasoning_content or ""
+
+        # Check if previous message has a task - if so, this is a task output (no thinking)
+        prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None
+
+        if thinking_mode == "thinking" and not prev_has_task:
+            if not drop_thinking or index > last_user_idx:
+                thinking_part = thinking_template.format(reasoning_content=rc) + thinking_end_token
+            else:
+                thinking_part = ""
+
+        if wo_eos:
+            prompt += assistant_msg_wo_eos_template.format(
+                reasoning=thinking_part,
+                content=summary_content,
+                tool_calls=tc_content,
+            )
+        else:
+            prompt += assistant_msg_template.format(
+                reasoning=thinking_part,
+                content=summary_content,
+                tool_calls=tc_content,
+            )
+    else:
+        raise NotImplementedError(f"Unknown role: {role}")
+
+    # Append transition tokens based on what follows
+    if index + 1 < len(messages) and messages[index + 1].get("role") not in ["assistant", "latest_reminder"]:
+        return prompt
+
+    task = messages[index].get("task")
+    if task is not None:
+        # Task special token for internal classification tasks
+        assert task in VALID_TASKS, f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}"
+        task_sp_token = DS_TASK_SP_TOKENS[task]
+
+        if task != "action":
+            # Non-action tasks: append task sp token directly after the message
+            prompt += task_sp_token
+        else:
+            # Action task: append Assistant + thinking token + action sp token
+            prompt += ASSISTANT_SP_TOKEN
+            prompt += thinking_end_token if thinking_mode != "thinking" else thinking_start_token
+            prompt += task_sp_token
+
+    elif messages[index].get("role") in ["user", "developer"]:
+        # Normal generation: append Assistant + thinking token
+        prompt += ASSISTANT_SP_TOKEN
+        if not drop_thinking and thinking_mode == "thinking":
+            prompt += thinking_start_token
+        elif drop_thinking and thinking_mode == "thinking" and index >= last_user_idx:
+            prompt += thinking_start_token
+        else:
+            prompt += thinking_end_token
+
+    return prompt
+
+
+# ============================================================
+# Preprocessing
+# ============================================================
+
+def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """
+    Merge tool messages into the preceding user message using content_blocks format.
+
+    DeepSeek-V4 does not have a standalone "tool" role; instead, tool results
+    are encoded as <tool_result> blocks within user messages.
+
+    This function converts a standard OpenAI-format conversation (with separate
+    "tool" role messages) into V4 format where tool results are merged into
+    user messages.
+
+    Args:
+        messages: List of message dicts in OpenAI format.
+
+    Returns:
+        Processed message list with tool messages merged into user messages.
+    """
+    merged: List[Dict[str, Any]] = []
+
+    for msg in messages:
+        msg = copy.deepcopy(msg)
+        role = msg.get("role")
+
+        if role == "tool":
+            # Convert tool message to a user message with tool_result block
+            tool_block = {
+                "type": "tool_result",
+                "tool_use_id": msg.get("tool_call_id", ""),
+                "content": msg.get("content", ""),
+            }
+            # Merge into previous message if it's already a user (merged tool)
+            if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1]:
+                merged[-1]["content_blocks"].append(tool_block)
+            else:
+                merged.append({
+                    "role": "user",
+                    "content_blocks": [tool_block],
+                })
+        elif role == "user":
+            text_block = {"type": "text", "text": msg.get("content", "")}
+            if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1] and merged[-1].get("task") is None:
+                merged[-1]["content_blocks"].append(text_block)
+            else:
+                new_msg = {
+                    "role": "user",
+                    "content": msg.get("content", ""),
+                    "content_blocks": [text_block],
+                }
+                # Preserve extra fields (task, wo_eos, mask, etc.)
+                for key in ("task", "wo_eos", "mask"):
+                    if key in msg:
+                        new_msg[key] = msg[key]
+                merged.append(new_msg)
+        else:
+            merged.append(msg)
+
+    return merged
+
+
+def sort_tool_results_by_call_order(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """
+    Sort tool_result blocks within user messages by the order of tool_calls
+    in the preceding assistant message.
+
+    Args:
+        messages: Preprocessed message list (after merge_tool_messages).
+
+    Returns:
+        Message list with sorted tool result blocks.
+    """
+    last_tool_call_order: Dict[str, int] = {}
+
+    for msg in messages:
+        role = msg.get("role")
+        if role == "assistant" and msg.get("tool_calls"):
+            last_tool_call_order = {}
+            for idx, tc in enumerate(msg["tool_calls"]):
+                tc_id = tc.get("id") or tc.get("function", {}).get("id", "")
+                if tc_id:
+                    last_tool_call_order[tc_id] = idx
+
+        elif role == "user" and msg.get("content_blocks"):
+            tool_blocks = [b for b in msg["content_blocks"] if b.get("type") == "tool_result"]
+            if len(tool_blocks) > 1 and last_tool_call_order:
+                sorted_blocks = sorted(
+                    tool_blocks,
+                    key=lambda b: last_tool_call_order.get(b.get("tool_use_id", ""), 0)
+                )
+                sorted_idx = 0
+                new_blocks = []
+                for block in msg["content_blocks"]:
+                    if block.get("type") == "tool_result":
+                        new_blocks.append(sorted_blocks[sorted_idx])
+                        sorted_idx += 1
+                    else:
+                        new_blocks.append(block)
+                msg["content_blocks"] = new_blocks
+
+    return messages
+
+
+# ============================================================
+# Main Encoding Function
+# ============================================================
+
+def encode_messages(
+    messages: List[Dict[str, Any]],
+    thinking_mode: str,
+    context: Optional[List[Dict[str, Any]]] = None,
+    drop_thinking: bool = True,
+    add_default_bos_token: bool = True,
+    reasoning_effort: Optional[str] = None,
+) -> str:
+    """
+    Encode a list of messages into the DeepSeek-V4 prompt format.
+
+    This is the main entry point for encoding conversations. It handles:
+    - BOS token insertion
+    - Thinking mode with optional reasoning content dropping
+    - Tool message merging into user messages
+    - Multi-turn conversation context
+
+    Args:
+        messages: List of message dicts to encode.
+        thinking_mode: Either "chat" or "thinking".
+        context: Optional preceding context messages (already encoded prefix).
+        drop_thinking: If True, drop reasoning_content from earlier assistant turns
+                      (only keep reasoning for messages after the last user message).
+        add_default_bos_token: Whether to prepend BOS token at conversation start.
+        reasoning_effort: Reasoning effort level, one of "low", "high", "max".
+            Only takes effect in thinking mode. None is treated as "low".
+
+    Returns:
+        The encoded prompt string.
+    """
+    context = context if context else []
+
+    # Preprocess: merge tool messages and sort tool results
+    messages = merge_tool_messages(messages)
+    messages = sort_tool_results_by_call_order(context + messages)[len(context):]
+    if context:
+        context = merge_tool_messages(context)
+        context = sort_tool_results_by_call_order(context)
+
+    full_messages = context + messages
+
+    prompt = bos_token if add_default_bos_token and len(context) == 0 else ""
+
+    # Resolve drop_thinking: if any message has tools defined, don't drop thinking
+    effective_drop_thinking = drop_thinking
+    if any(m.get("tools") for m in full_messages):
+        effective_drop_thinking = False
+
+    if thinking_mode == "thinking" and effective_drop_thinking:
+        full_messages = _drop_thinking_messages(full_messages)
+        # After dropping, recalculate how many messages to render
+        # (context may have shrunk too)
+        num_to_render = len(full_messages) - len(_drop_thinking_messages(context))
+        context_len = len(full_messages) - num_to_render
+    else:
+        num_to_render = len(messages)
+        context_len = len(context)
+
+    for idx in range(num_to_render):
+        prompt += render_message(
+            idx + context_len,
+            full_messages,
+            thinking_mode=thinking_mode,
+            drop_thinking=effective_drop_thinking,
+            reasoning_effort=reasoning_effort,
+        )
+
+    return prompt
+
+
+def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """
+    Drop reasoning_content and non-essential messages before the last user message.
+
+    Behavior:
+    - Messages with role in ["user", "system", "tool", "latest_reminder"] are always kept.
+    - Messages at or after the last user index are always kept.
+    - Assistant messages before the last user get reasoning_content removed.
+    - Developer messages before the last user are dropped entirely.
+    """
+    last_user_idx = find_last_user_index(messages)
+    result = []
+    keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"}
+
+    for idx, msg in enumerate(messages):
+        role = msg.get("role")
+        if role in keep_roles or idx >= last_user_idx:
+            result.append(msg)
+        elif role == "assistant":
+            msg = copy.copy(msg)
+            msg.pop("reasoning_content", None)
+            result.append(msg)
+        # developer and other roles before last_user_idx are dropped
+
+    return result
+
+
+# ============================================================
+# Parsing (Decoding model output)
+# ============================================================
+
+def _read_until_stop(index: int, text: str, stop: List[str]) -> Tuple[int, str, Optional[str]]:
+    """
+    Read text from index until one of the stop strings is found.
+
+    Returns:
+        Tuple of (new_index, content_before_stop, matched_stop_string_or_None).
+    """
+    min_pos = len(text)
+    matched_stop = None
+
+    for s in stop:
+        pos = text.find(s, index)
+        if pos != -1 and pos < min_pos:
+            min_pos = pos
+            matched_stop = s
+
+    if matched_stop:
+        content = text[index:min_pos]
+        return min_pos + len(matched_stop), content, matched_stop
+    else:
+        content = text[index:]
+        return len(text), content, None
+
+
+def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Dict[str, str]]]:
+    """
+    Parse DSML tool calls from text starting at the given index.
+
+    Args:
+        index: Starting position in text.
+        text: The full text to parse.
+
+    Returns:
+        Tuple of (new_index, last_stop_token, list_of_tool_call_dicts).
+        Each tool call dict has "name" and "arguments" keys.
+    """
+    tool_calls: List[Dict[str, Any]] = []
+    stop_token = None
+    tool_calls_end_token = f"</{dsml_token}{tool_calls_block_name}>"
+
+    while index < len(text):
+        index, _, stop_token = _read_until_stop(index, text, [f"<{dsml_token}invoke", tool_calls_end_token])
+        if _ != ">\n":
+            raise ValueError(f"Tool call format error: expected '>\\n' but got '{_}'")
+
+        if stop_token == tool_calls_end_token:
+            break
+
+        if stop_token is None:
+            raise ValueError("Missing special token in tool calls")
+
+        index, tool_name_content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"</{dsml_token}invoke"])
+
+        p_tool_name = re.findall(r'^\s*name="(.*?)">\n$', tool_name_content, flags=re.DOTALL)
+        if len(p_tool_name) != 1:
+            raise ValueError(f"Tool name format error: '{tool_name_content}'")
+        tool_name = p_tool_name[0]
+
+        tool_args: Dict[str, Tuple[str, str]] = {}
+        while stop_token == f"<{dsml_token}parameter":
+            index, param_content, stop_token = _read_until_stop(index, text, [f"/{dsml_token}parameter"])
+
+            param_kv = re.findall(r'^ name="(.*?)" string="(true|false)">(.*?)<$', param_content, flags=re.DOTALL)
+            if len(param_kv) != 1:
+                raise ValueError(f"Parameter format error: '{param_content}'")
+            param_name, string, param_value = param_kv[0]
+
+            if param_name in tool_args:
+                raise ValueError(f"Duplicate parameter name: '{param_name}'")
+            tool_args[param_name] = (param_value, string)
+
+            index, content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"</{dsml_token}invoke"])
+            if content != ">\n":
+                raise ValueError(f"Parameter format error: expected '>\\n' but got '{content}'")
+
+        tool_call = decode_dsml_to_arguments(tool_name=tool_name, tool_args=tool_args)
+        tool_calls.append(tool_call)
+
+    return index, stop_token, tool_calls
+
+
+def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[str, Any]:
+    """
+    Parse a model completion text into a structured assistant message.
+
+    This function takes the raw text output from the model (a single assistant turn)
+    and extracts:
+    - reasoning_content (thinking block)
+    - content (summary/response)
+    - tool_calls (if any)
+
+    NOTE: This function is designed to parse only correctly formatted strings and
+    will raise ValueError for malformed output.
+
+    Args:
+        text: The raw completion text (including EOS token).
+        thinking_mode: Either "chat" or "thinking".
+
+    Returns:
+        Dict with keys: "role", "content", "reasoning_content", "tool_calls".
+        tool_calls are in OpenAI format.
+    """
+    summary_content, reasoning_content, tool_calls = "", "", []
+    index, stop_token = 0, None
+    tool_calls_start_token = f"\n\n<{dsml_token}{tool_calls_block_name}"
+
+    is_thinking = thinking_mode == "thinking"
+    is_tool_calling = False
+
+    if is_thinking:
+        index, content_delta, stop_token = _read_until_stop(index, text, [thinking_end_token, tool_calls_start_token])
+        reasoning_content = content_delta
+        assert stop_token == thinking_end_token, "Invalid thinking format: missing </think>"
+
+    index, content_delta, stop_token = _read_until_stop(index, text, [eos_token, tool_calls_start_token])
+    summary_content = content_delta
+    if stop_token == tool_calls_start_token:
+        is_tool_calling = True
+    else:
+        assert stop_token == eos_token, "Invalid format: missing EOS token"
+
+    if is_tool_calling:
+        index, stop_token, tool_calls = parse_tool_calls(index, text)
+
+        index, tool_ends_text, stop_token = _read_until_stop(index, text, [eos_token])
+        assert not tool_ends_text, "Unexpected content after tool calls"
+
+    assert len(text) == index and stop_token in [eos_token, None], "Unexpected content at end"
+
+    for sp_token in [bos_token, eos_token, thinking_start_token, thinking_end_token, dsml_token]:
+        assert sp_token not in summary_content and sp_token not in reasoning_content, \
+            f"Unexpected special token '{sp_token}' in content"
+
+    return {
+        "role": "assistant",
+        "content": summary_content,
+        "reasoning_content": reasoning_content,
+        "tool_calls": tool_calls_to_openai_format(tool_calls)
+    }

encoding/test_encoding_dsv4.py

@@ -0,0 +1,89 @@
+"""
+Test suite for DeepSeek-V4 Encoding.
+
+Run: python test_encoding_dsv4.py
+"""
+
+import json
+import os
+
+from encoding_dsv4 import encode_messages, parse_message_from_completion_text
+
+TESTS_DIR = os.path.join(os.path.dirname(__file__), "tests")
+
+
+def test_case_1():
+    """Thinking mode with tool calls (multi-turn, tool results merged into user)."""
+    with open(os.path.join(TESTS_DIR, "test_input_1.json")) as f:
+        td = json.load(f)
+        messages = td["messages"]
+        messages[0]["tools"] = td["tools"]
+    gold = open(os.path.join(TESTS_DIR, "test_output_1.txt")).read()
+    prompt = encode_messages(messages, thinking_mode="thinking")
+    assert prompt == gold
+
+    # Parse: assistant turn with tool call
+    marker = "<｜Assistant｜><think>"
+    first_start = prompt.find(marker) + len(marker)
+    first_end = prompt.find("<｜User｜>", first_start)
+    parsed_tc = parse_message_from_completion_text(prompt[first_start:first_end], thinking_mode="thinking")
+    assert parsed_tc["reasoning_content"] == "The user wants to know the weather in Beijing. I should use the get_weather tool."
+    assert parsed_tc["content"] == ""
+    assert len(parsed_tc["tool_calls"]) == 1
+    assert parsed_tc["tool_calls"][0]["function"]["name"] == "get_weather"
+    assert json.loads(parsed_tc["tool_calls"][0]["function"]["arguments"]) == {"location": "Beijing", "unit": "celsius"}
+
+    # Parse: final assistant turn with content
+    last_start = prompt.rfind(marker) + len(marker)
+    parsed_final = parse_message_from_completion_text(prompt[last_start:], thinking_mode="thinking")
+    assert parsed_final["reasoning_content"] == "Got the weather data. Let me format a nice response."
+    assert "22°C" in parsed_final["content"]
+    assert parsed_final["tool_calls"] == []
+
+    print("  [PASS] case 1: thinking with tools (encode + parse)")
+
+
+def test_case_2():
+    """Thinking mode without tools (drop_thinking removes earlier reasoning)."""
+    messages = json.load(open(os.path.join(TESTS_DIR, "test_input_2.json")))
+    gold = open(os.path.join(TESTS_DIR, "test_output_2.txt")).read()
+    prompt = encode_messages(messages, thinking_mode="thinking")
+    assert prompt == gold
+
+    # Parse: last assistant turn
+    marker = "<｜Assistant｜><think>"
+    last_start = prompt.rfind(marker) + len(marker)
+    parsed = parse_message_from_completion_text(prompt[last_start:], thinking_mode="thinking")
+    assert parsed["reasoning_content"] == "The user asks about the capital of France. It is Paris."
+    assert parsed["content"] == "The capital of France is Paris."
+    assert parsed["tool_calls"] == []
+
+    # Verify drop_thinking: first assistant's reasoning should be absent
+    assert "The user said hello" not in prompt
+
+    print("  [PASS] case 2: thinking without tools (encode + parse)")
+
+
+def test_case_3():
+    """Interleaved thinking + search (developer with tools, latest_reminder)."""
+    messages = json.load(open(os.path.join(TESTS_DIR, "test_input_3.json")))
+    gold = open(os.path.join(TESTS_DIR, "test_output_3.txt")).read()
+    assert encode_messages(messages, thinking_mode="thinking") == gold
+    print("  [PASS] case 3: interleaved thinking + search")
+
+
+def test_case_4():
+    """Quick instruction task with latest_reminder (chat mode, action task)."""
+    messages = json.load(open(os.path.join(TESTS_DIR, "test_input_4.json")))
+    gold = open(os.path.join(TESTS_DIR, "test_output_4.txt")).read()
+    assert encode_messages(messages, thinking_mode="chat") == gold
+    print("  [PASS] case 4: quick instruction task")
+
+
+if __name__ == "__main__":
+    print("Running DeepSeek-V4 Encoding Tests...\n")
+    test_case_1()
+    test_case_2()
+    test_case_3()
+    test_case_4()
+    print("\nAll 4 tests passed!")

encoding/tests/test_input_1.json

@@ -0,0 +1,81 @@
+{
+    "tools": [
+        {
+            "type": "function",
+            "function": {
+                "name": "get_weather",
+                "description": "Get the weather for a specific location",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "location": {
+                            "type": "string",
+                            "description": "The city name"
+                        },
+                        "unit": {
+                            "type": "string",
+                            "enum": ["celsius", "fahrenheit"],
+                            "description": "Temperature unit"
+                        }
+                    },
+                    "required": ["location"]
+                }
+            }
+        },
+        {
+            "type": "function",
+            "function": {
+                "name": "search",
+                "description": "Search the web for information",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "query": {
+                            "type": "string",
+                            "description": "Search query"
+                        },
+                        "num_results": {
+                            "type": "integer",
+                            "description": "Number of results to return"
+                        }
+                    },
+                    "required": ["query"]
+                }
+            }
+        }
+    ],
+    "messages": [
+        {
+            "role": "system",
+            "content": "You are a helpful assistant."
+        },
+        {
+            "role": "user",
+            "content": "What's the weather in Beijing?"
+        },
+        {
+            "role": "assistant",
+            "reasoning_content": "The user wants to know the weather in Beijing. I should use the get_weather tool.",
+            "tool_calls": [
+                {
+                    "id": "call_001",
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "arguments": "{\"location\": \"Beijing\", \"unit\": \"celsius\"}"
+                    }
+                }
+            ]
+        },
+        {
+            "role": "tool",
+            "tool_call_id": "call_001",
+            "content": "{\"temperature\": 22, \"condition\": \"sunny\", \"humidity\": 45}"
+        },
+        {
+            "role": "assistant",
+            "reasoning_content": "Got the weather data. Let me format a nice response.",
+            "content": "The weather in Beijing is currently sunny with a temperature of 22°C and 45% humidity."
+        }
+    ]
+}

encoding/tests/test_input_2.json

@@ -0,0 +1,24 @@
+[
+  {
+    "role": "system",
+    "content": "You are a helpful assistant."
+  },
+  {
+    "role": "user",
+    "content": "Hello"
+  },
+  {
+    "role": "assistant",
+    "reasoning_content": "The user said hello, I should greet back.",
+    "content": "Hi there! How can I help you?"
+  },
+  {
+    "role": "user",
+    "content": "What is the capital of France?"
+  },
+  {
+    "role": "assistant",
+    "reasoning_content": "The user asks about the capital of France. It is Paris.",
+    "content": "The capital of France is Paris."
+  }
+]

encoding/tests/test_input_3.json

@@ -0,0 +1,159 @@
+[
+  {
+    "role": "system",
+    "content": "该助手为DeepSeek，由深度求索公司创造。"
+  },
+  {
+    "role": "latest_reminder",
+    "content": "2026-02-21,星期六,广州,App,中文"
+  },
+  {
+    "role": "developer",
+    "content": "小柴胡冲剂和布洛芬能一起吃吗？\n\nCITATION FORMAT: 【{cursor_id}†L{start_line_id}(-L{end_line_id})?】",
+    "tools": [
+      {
+        "type": "function",
+        "function": {
+          "name": "search",
+          "description": "Web search. Split multiple queries with '||'.",
+          "parameters": {
+            "type": "object",
+            "properties": {
+              "queries": {
+                "type": "string",
+                "description": "query1||query2"
+              }
+            },
+            "required": [
+              "queries"
+            ],
+            "additionalProperties": false,
+            "$schema": "http://json-schema.org/draft-07/schema#"
+          }
+        }
+      },
+      {
+        "type": "function",
+        "function": {
+          "name": "open",
+          "description": "Batch open IDs (format 【{id}†...】) or URLs.",
+          "parameters": {
+            "type": "object",
+            "properties": {
+              "open_list": {
+                "type": "array",
+                "items": {
+                  "type": "object",
+                  "properties": {
+                    "id": {
+                      "description": "ID or URL",
+                      "anyOf": [
+                        {
+                          "type": "integer"
+                        },
+                        {
+                          "type": "string"
+                        }
+                      ],
+                      "default": -1
+                    },
+                    "cursor": {
+                      "type": "integer",
+                      "description": "",
+                      "default": -1
+                    },
+                    "loc": {
+                      "type": "integer",
+                      "description": "Start line",
+                      "default": -1
+                    },
+                    "num_lines": {
+                      "type": "integer",
+                      "description": "",
+                      "default": -1
+                    },
+                    "view_source": {
+                      "type": "boolean",
+                      "description": "",
+                      "default": false
+                    }
+                  },
+                  "additionalProperties": false
+                },
+                "description": ""
+              }
+            },
+            "required": [
+              "open_list"
+            ],
+            "additionalProperties": false,
+            "$schema": "http://json-schema.org/draft-07/schema#"
+          }
+        }
+      },
+      {
+        "type": "function",
+        "function": {
+          "name": "find",
+          "description": "Find exact text pattern in pages.",
+          "parameters": {
+            "type": "object",
+            "properties": {
+              "find_list": {
+                "type": "array",
+                "items": {
+                  "type": "object",
+                  "properties": {
+                    "pattern": {
+                      "type": "string",
+                      "description": ""
+                    },
+                    "cursor": {
+                      "type": "integer",
+                      "description": "",
+                      "default": -1
+                    }
+                  },
+                  "required": [
+                    "pattern"
+                  ],
+                  "additionalProperties": false
+                },
+                "description": ""
+              }
+            },
+            "required": [
+              "find_list"
+            ],
+            "additionalProperties": false,
+            "$schema": "http://json-schema.org/draft-07/schema#"
+          }
+        }
+      }
+    ]
+  },
+  {
+    "role": "assistant",
+    "content": "",
+    "reasoning_content": "用户想知道小柴胡冲剂和布洛芬能否一起服用。",
+    "tool_calls": [
+      {
+        "type": "function",
+        "function": {
+          "name": "search",
+          "arguments": "{\"queries\": \"小柴胡冲剂 布洛芬 相互作用 一起吃\"}"
+        }
+      }
+    ]
+  },
+  {
+    "role": "tool",
+    "content": "[0]"
+  },
+  {
+    "role": "assistant",
+    "content": "请及时就医。",
+    "reasoning_content": "现在开始组织回答。",
+    "tool_calls": []
+  }
+]

encoding/tests/test_input_4.json

@@ -0,0 +1,28 @@
+[
+  {
+    "role": "system",
+    "content": "该助手为DeepSeek-V3，由深度求索公司创造。\n今天是2025年10月17日，星期五。"
+  },
+  {
+    "role": "latest_reminder",
+    "content": "2024-11-15,上海市,App,中文"
+  },
+  {
+    "role": "user",
+    "content": "热海大滚锅是世界著名温泉吗"
+  },
+  {
+    "role": "assistant",
+    "content": "关于热海大滚锅是否是世界著名温泉，可以这样总结：它在**中国乃至全球的地热奇观中占有重要地位**，但“世界著名”的称号更侧重于它作为**独特的地质现象和旅游景点**，而非普遍的温泉疗养体验。\n\n为了让你快速了解，我整理了一个简要的表格：\n\n| 维度 | 热海大滚锅的地位与特点 |\n| :--- | :--- |\n| **地质奇观** | **世界罕见**的高温地热系统，泉眼多、水温高、形态多样。 |\n| **国内声誉** | **中国三大地热区之一**，**国家5A级旅游景区**，自古闻名（徐霞客曾游历并记载）。 |\n| **国际知名度** | 在一些旅游平台被提及为“世界六大温泉”之一，但此说法流传不广，其国际声誉更多建立在地质独特性上。 |\n| **核心体验** | **观赏地热奇观**（如97℃沸腾的“大滚锅”）、**体验温泉煮鸡蛋**。 |\n\n### 💡 游玩攻略与温馨提示\n\n如果你计划前往热海大滚锅，这里有一些实用信息供你参考：\n\n- **门票与开放时间**：\n    - **门票**：景区门票约为**50元/人**。如果选择包含温泉沐浴的套餐，价格会更高，例如约**288元**。\n    - **开放时间**：景区一般**08:00-18:00**开放，但具体时间可能变动，建议提前核实。\n\n- **特色体验**：\n    - **温泉煮鸡蛋**：这几乎是必试项目。可以在景区门口购买用草绳串起的生鸡蛋（约5-8元/串），然后到“大滚锅”旁的指定区域蒸煮，几分钟便可熟食，趣味十足。\n    - **金汤足浴**：可以直接用从“大滚锅”流出的温泉水泡脚，缓解旅途疲劳。\n\n- **注意事项**：\n    - **安全第一**：“大滚锅”水温极高，务必遵守游览规则，在指定区域内观赏，切勿随意触碰泉水。\n    - **规划行程**：建议为热海景区预留**3-4小时**的游览时间。景区内步道不走回头路，出入口有观光车接送。\n\n希望这些信息能帮助你更好地了解热海大滚锅。如果你对腾冲的其他景点或者行程规划有更多疑问，我很乐意提供进一步的信息。",
+    "mask": 1
+  },
+  {
+    "role": "user",
+    "content": "世界著名温泉有哪些",
+    "task": "action"
+  },
+  {
+    "role": "assistant",
+    "content": "Search"
+  }
+]

encoding/tests/test_output_1.txt

@@ -0,0 +1,36 @@
+<｜begin▁of▁sentence｜>You are a helpful assistant.
+
+## Tools
+
+You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following:
+
+<｜DSML｜tool_calls>
+<｜DSML｜invoke name="$TOOL_NAME">
+<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</｜DSML｜parameter>
+...
+</｜DSML｜invoke>
+<｜DSML｜invoke name="$TOOL_NAME2">
+...
+</｜DSML｜invoke>
+</｜DSML｜tool_calls>
+
+String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
+
+If thinking_mode is enabled (triggered by <think>), you MUST output your complete reasoning inside <think>...</think> BEFORE any tool calls or final response.
+
+Otherwise, output directly after </think> with tool calls or final response.
+
+### Available Tool Schemas
+
+{"name": "get_weather", "description": "Get the weather for a specific location", "parameters": {"type": "object", "properties": {"location": {"type": "string", "description": "The city name"}, "unit": {"type": "string", "enum": ["celsius", "fahrenheit"], "description": "Temperature unit"}}, "required": ["location"]}}
+{"name": "search", "description": "Search the web for information", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "Search query"}, "num_results": {"type": "integer", "description": "Number of results to return"}}, "required": ["query"]}}
+
+You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
+<｜User｜>What's the weather in Beijing?<｜Assistant｜><think>The user wants to know the weather in Beijing. I should use the get_weather tool.</think>
+
+<｜DSML｜tool_calls>
+<｜DSML｜invoke name="get_weather">
+<｜DSML｜parameter name="location" string="true">Beijing</｜DSML｜parameter>
+<｜DSML｜parameter name="unit" string="true">celsius</｜DSML｜parameter>
+</｜DSML｜invoke>
+</｜DSML｜tool_calls><｜end▁of▁sentence｜><｜User｜><tool_result>{"temperature": 22, "condition": "sunny", "humidity": 45}</tool_result><｜Assistant｜><think>Got the weather data. Let me format a nice response.</think>The weather in Beijing is currently sunny with a temperature of 22°C and 45% humidity.<｜end▁of▁sentence｜>

encoding/tests/test_output_2.txt

@@ -0,0 +1 @@
+<｜begin▁of▁sentence｜>You are a helpful assistant.<｜User｜>Hello<｜Assistant｜></think>Hi there! How can I help you?<｜end▁of▁sentence｜><｜User｜>What is the capital of France?<｜Assistant｜><think>The user asks about the capital of France. It is Paris.</think>The capital of France is Paris.<｜end▁of▁sentence｜>

encoding/tests/test_output_3.txt

@@ -0,0 +1,38 @@
+<｜begin▁of▁sentence｜>该助手为DeepSeek，由深度求索公司创造。<｜latest_reminder｜>2026-02-21,星期六,广州,App,中文<｜User｜>小柴胡冲剂和布洛芬能一起吃吗？
+
+CITATION FORMAT: 【{cursor_id}†L{start_line_id}(-L{end_line_id})?】
+
+## Tools
+
+You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following:
+
+<｜DSML｜tool_calls>
+<｜DSML｜invoke name="$TOOL_NAME">
+<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</｜DSML｜parameter>
+...
+</｜DSML｜invoke>
+<｜DSML｜invoke name="$TOOL_NAME2">
+...
+</｜DSML｜invoke>
+</｜DSML｜tool_calls>
+
+String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
+
+If thinking_mode is enabled (triggered by <think>), you MUST output your complete reasoning inside <think>...</think> BEFORE any tool calls or final response.
+
+Otherwise, output directly after </think> with tool calls or final response.
+
+### Available Tool Schemas
+
+{"name": "search", "description": "Web search. Split multiple queries with '||'.", "parameters": {"type": "object", "properties": {"queries": {"type": "string", "description": "query1||query2"}}, "required": ["queries"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
+{"name": "open", "description": "Batch open IDs (format 【{id}†...】) or URLs.", "parameters": {"type": "object", "properties": {"open_list": {"type": "array", "items": {"type": "object", "properties": {"id": {"description": "ID or URL", "anyOf": [{"type": "integer"}, {"type": "string"}], "default": -1}, "cursor": {"type": "integer", "description": "", "default": -1}, "loc": {"type": "integer", "description": "Start line", "default": -1}, "num_lines": {"type": "integer", "description": "", "default": -1}, "view_source": {"type": "boolean", "description": "", "default": false}}, "additionalProperties": false}, "description": ""}}, "required": ["open_list"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
+{"name": "find", "description": "Find exact text pattern in pages.", "parameters": {"type": "object", "properties": {"find_list": {"type": "array", "items": {"type": "object", "properties": {"pattern": {"type": "string", "description": ""}, "cursor": {"type": "integer", "description": "", "default": -1}}, "required": ["pattern"], "additionalProperties": false}, "description": ""}}, "required": ["find_list"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
+
+You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
+<｜Assistant｜><think>用户想知道小柴胡冲剂和布洛芬能否一起服用。</think>
+
+<｜DSML｜tool_calls>
+<｜DSML｜invoke name="search">
+<｜DSML｜parameter name="queries" string="true">小柴胡冲剂 布洛芬 相互作用 一起吃</｜DSML｜parameter>
+</｜DSML｜invoke>
+</｜DSML｜tool_calls><｜end▁of▁sentence｜><｜User｜><tool_result>[0]</tool_result><｜Assistant｜><think>现在开始组织回答。</think>请及时就医。<｜end▁of▁sentence｜>

encoding/tests/test_output_4.txt

@@ -0,0 +1,29 @@
+<｜begin▁of▁sentence｜>该助手为DeepSeek-V3，由深度求索公司创造。
+今天是2025年10月17日，星期五。<｜latest_reminder｜>2024-11-15,上海市,App,中文<｜User｜>热海大滚锅是世界著名温泉吗<｜Assistant｜></think>关于热海大滚锅是否是世界著名温泉，可以这样总结：它在**中国乃至全球的地热奇观中占有重要地位**，但“世界著名”的称号更侧重于它作为**独特的地质现象和旅游景点**，而非普遍的温泉疗养体验。
+
+为了让你快速了解，我整理了一个简要的表格：
+
+| 维度 | 热海大滚锅的地位与特点 |
+| :--- | :--- |
+| **地质奇观** | **世界罕见**的高温地热系统，泉眼多、水温高、形态多样。 |
+| **国内声誉** | **中国三大地热区之一**，**国家5A级旅游景区**，自古闻名（徐霞客曾游历并记载）。 |
+| **国际知名度** | 在一些旅游平台被提及为“世界六大温泉”之一，但此说法流传不广，其国际声誉更多建立在地质独特性上。 |
+| **核心体验** | **观赏地热奇观**（如97℃沸腾的“大滚锅”）、**体验温泉煮鸡蛋**。 |
+
+### 💡 游玩攻略与温馨提示
+
+如果你计划前往热海大滚锅，这里有一些实用信息供你参考：
+
+- **门票与开放时间**：
+    - **门票**：景区门票约为**50元/人**。如果选择包含温泉沐浴的套餐，价格会更高，例如约**288元**。
+    - **开放时间**：景区一般**08:00-18:00**开放，但具体时间可能变动，建议提前核实。
+
+- **特色体验**：
+    - **温泉煮鸡蛋**：这几乎是必试项目。可以在景区门口购买用草绳串起的生鸡蛋（约5-8元/串），然后到“大滚锅”旁的指定区域蒸煮，几分钟便可熟食，趣味十足。
+    - **金汤足浴**：可以直接用从“大滚锅”流出的温泉水泡脚，缓解旅途疲劳。
+
+- **注意事项**：
+    - **安全第一**：“大滚锅”水温极高，务必遵守游览规则，在指定区域内观赏，切勿随意触碰泉水。
+    - **规划行程**：建议为热海景区预留**3-4小时**的游览时间。景区内步道不走回头路，出入口有观光车接送。
+
+希望这些信息能帮助你更好地了解热海大滚锅。如果你对腾冲的其他景点或者行程规划有更多疑问，我很乐意提供进一步的信息。<｜end▁of▁sentence｜><｜User｜>世界著名温泉有哪些<｜Assistant｜></think><｜action｜>Search<｜end▁of▁sentence｜>

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-00033.safetensors

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

model-00004-of-00033.safetensors

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

model-00005-of-00033.safetensors

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

model-00011-of-00033.safetensors

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

model-00012-of-00033.safetensors

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

model-00013-of-00033.safetensors

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

model-00014-of-00033.safetensors

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

model-00016-of-00033.safetensors

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

model-00017-of-00033.safetensors

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

model-00018-of-00033.safetensors

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

model-00019-of-00033.safetensors

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

model-00020-of-00033.safetensors

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

model-00021-of-00033.safetensors

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

model-00024-of-00033.safetensors

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

model-00025-of-00033.safetensors

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

model-00026-of-00033.safetensors

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

model-00032-of-00033.safetensors

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

model-00033-of-00033.safetensors

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

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"
+}